Fix Codex workspace identity edge cases

Monter · craft-agents-bot · Monter · commit 72c606114cbe · 2026-03-23T07:24:45.000+01:00
Co-Authored-By: Craft Agent &lt;agents-noreply@craft.do&gt;
diff --git a/Sources/CodexBar/Providers/Codex/CodexProviderRuntime.swift b/Sources/CodexBar/Providers/Codex/CodexProviderRuntime.swift
@@ -9,6 +9,7 @@ final class CodexProviderRuntime: ProviderRuntime {
         let cookieSource: ProviderCookieSource
         let hasManualCookieHeader: Bool
         let hasTokenAccounts: Bool
+        let selectedTokenAccountID: UUID?
 
         var hasManualCredentials: Bool {
             self.hasManualCookieHeader || self.hasTokenAccounts
@@ -20,7 +21,8 @@ final class CodexProviderRuntime: ProviderRuntime {
             return Self(
                 cookieSource: settings.codexCookieSource,
                 hasManualCookieHeader: !manualHeader.isEmpty,
-                hasTokenAccounts: !settings.tokenAccounts(for: .codex).isEmpty)
+                hasTokenAccounts: !settings.tokenAccounts(for: .codex).isEmpty,
+                selectedTokenAccountID: settings.selectedTokenAccount(for: .codex)?.id)
         }
     }
 
diff --git a/Sources/CodexBarCore/OpenAIDashboardModels.swift b/Sources/CodexBarCore/OpenAIDashboardModels.swift
@@ -118,6 +118,7 @@ extension OpenAIDashboardSnapshot {
     public func toUsageSnapshot(
         provider: UsageProvider = .codex,
         accountEmail: String? = nil,
+        accountOrganization: String? = nil,
         accountPlan: String? = nil) -> UsageSnapshot?
     {
         guard let primaryLimit else { return nil }
@@ -126,7 +127,7 @@ extension OpenAIDashboardSnapshot {
         let identity = ProviderIdentitySnapshot(
             providerID: provider,
             accountEmail: resolvedEmail,
-            accountOrganization: nil,
+            accountOrganization: accountOrganization,
             loginMethod: resolvedPlan)
         return UsageSnapshot(
             primary: primaryLimit,
diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift
@@ -92,6 +92,13 @@ public struct OpenAIDashboardBrowserCookieImporter {
         let log: (String) -> Void
     }
 
+    private struct TargetMatchContext {
+        let targetEmail: String
+        let targetWorkspaceLabel: String?
+        let candidateLabel: String
+        let log: (String) -> Void
+    }
+
     private static let cookieDomains = ["chatgpt.com", "openai.com"]
     private static let cookieClient = BrowserCookieClient()
     private static let cookieImportOrder: BrowserCookieImportOrder =
@@ -213,6 +220,7 @@ public struct OpenAIDashboardBrowserCookieImporter {
         switch await self.evaluateCandidate(
             candidate,
             targetEmail: normalizedTarget,
+            targetWorkspaceLabel: targetWorkspaceLabel,
             allowAnyAccount: allowAnyAccount,
             log: log)
         {
@@ -421,6 +429,7 @@ public struct OpenAIDashboardBrowserCookieImporter {
         switch await self.evaluateCandidate(
             candidate,
             targetEmail: context.targetEmail,
+            targetWorkspaceLabel: context.targetWorkspaceLabel,
             allowAnyAccount: context.allowAnyAccount,
             log: context.log)
         {
@@ -475,11 +484,17 @@ public struct OpenAIDashboardBrowserCookieImporter {
     private func evaluateCandidate(
         _ candidate: Candidate,
         targetEmail: String?,
+        targetWorkspaceLabel: String?,
         allowAnyAccount: Bool,
         log: @escaping (String) -> Void) async -> CandidateEvaluation
     {
         log("Trying candidate \(candidate.label) (\(candidate.cookies.count) cookies)")
 
+        let resolvedWorkspaceLabel = self.resolveWorkspaceLabel(from: candidate.cookies)
+        if let resolvedWorkspaceLabel, !resolvedWorkspaceLabel.isEmpty {
+            log("Candidate \(candidate.label) workspace: \(resolvedWorkspaceLabel)")
+        }
+
         let apiEmail = await self.fetchSignedInEmailFromAPI(cookies: candidate.cookies, logger: log)
         if let apiEmail {
             log("Candidate \(candidate.label) API email: \(apiEmail)")
@@ -488,7 +503,15 @@ public struct OpenAIDashboardBrowserCookieImporter {
         // Prefer the API email when available (fast; avoids WebKit hydration/timeout risks).
         if let apiEmail, !apiEmail.isEmpty {
             if let targetEmail {
-                if apiEmail.lowercased() == targetEmail.lowercased() {
+                if self.matchesTarget(
+                    signedInEmail: apiEmail,
+                    candidateWorkspaceLabel: resolvedWorkspaceLabel,
+                    context: TargetMatchContext(
+                        targetEmail: targetEmail,
+                        targetWorkspaceLabel: targetWorkspaceLabel,
+                        candidateLabel: candidate.label,
+                        log: log))
+                {
                     return .match(candidate: candidate, signedInEmail: apiEmail)
                 }
                 return .mismatch(candidate: candidate, signedInEmail: apiEmail)
@@ -515,7 +538,15 @@ public struct OpenAIDashboardBrowserCookieImporter {
             let resolvedEmail = signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines)
             if let resolvedEmail, !resolvedEmail.isEmpty {
                 if let targetEmail {
-                    if resolvedEmail.lowercased() == targetEmail.lowercased() {
+                    if self.matchesTarget(
+                        signedInEmail: resolvedEmail,
+                        candidateWorkspaceLabel: resolvedWorkspaceLabel,
+                        context: TargetMatchContext(
+                            targetEmail: targetEmail,
+                            targetWorkspaceLabel: targetWorkspaceLabel,
+                            candidateLabel: candidate.label,
+                            log: log))
+                    {
                         return .match(candidate: candidate, signedInEmail: resolvedEmail)
                     }
                     return .mismatch(candidate: candidate, signedInEmail: resolvedEmail)
@@ -544,6 +575,81 @@ public struct OpenAIDashboardBrowserCookieImporter {
         return false
     }
 
+    private func matchesTarget(
+        signedInEmail: String,
+        candidateWorkspaceLabel: String?,
+        context: TargetMatchContext) -> Bool
+    {
+        guard signedInEmail.lowercased() == context.targetEmail.lowercased() else { return false }
+
+        let normalizedTargetWorkspace = self.normalizeWorkspaceLabel(context.targetWorkspaceLabel)
+        guard let normalizedTargetWorkspace else { return true }
+
+        let normalizedCandidateWorkspace = self.normalizeWorkspaceLabel(candidateWorkspaceLabel)
+        guard let normalizedCandidateWorkspace else {
+            context.log(
+                "Candidate \(context.candidateLabel) matched email but workspace is unknown; " +
+                    "expected \(normalizedTargetWorkspace)")
+            return false
+        }
+
+        if normalizedCandidateWorkspace == normalizedTargetWorkspace {
+            return true
+        }
+
+        context.log(
+            "Candidate \(context.candidateLabel) matched email but workspace mismatched " +
+                "(candidate=\(normalizedCandidateWorkspace), target=\(normalizedTargetWorkspace))")
+        return false
+    }
+
+    private func normalizeWorkspaceLabel(_ label: String?) -> String? {
+        let trimmed = label?.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard let trimmed, !trimmed.isEmpty else { return nil }
+        return trimmed.lowercased()
+    }
+
+    public func _normalizeWorkspaceLabelForTesting(_ label: String?) -> String? {
+        self.normalizeWorkspaceLabel(label)
+    }
+
+    private func resolveWorkspaceLabel(from cookies: [HTTPCookie]) -> String? {
+        guard let accountID = cookies.first(where: { $0.name == "_account" })?.value,
+              !accountID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
+              let sessionCookie = cookies.first(where: { $0.name == "oai-client-auth-session" })?.value,
+              let payload = self.decodeBase64URLJSONPayload(fromCookieValue: sessionCookie),
+              let json = try? JSONSerialization.jsonObject(with: payload) as? [String: Any],
+              let workspaces = json["workspaces"] as? [[String: Any]]
+        else {
+            return nil
+        }
+
+        guard let workspace = workspaces.first(where: {
+            ($0["id"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) == accountID
+        }) else {
+            return nil
+        }
+
+        let name = (workspace["name"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines)
+        if let name, !name.isEmpty { return name }
+
+        let kind = (workspace["kind"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
+        if kind == "personal" { return "Personal" }
+        return nil
+    }
+
+    private func decodeBase64URLJSONPayload(fromCookieValue value: String) -> Data? {
+        let prefix = value.split(separator: ".", maxSplits: 1, omittingEmptySubsequences: false).first
+            .map(String.init) ?? value
+        guard !prefix.isEmpty else { return nil }
+        var base64 = prefix.replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
+        let remainder = base64.count % 4
+        if remainder != 0 {
+            base64 += String(repeating: "=", count: 4 - remainder)
+        }
+        return Data(base64Encoded: base64)
+    }
+
     private func handleMismatch(
         candidate: Candidate,
         signedInEmail: String,
@@ -552,13 +658,13 @@ public struct OpenAIDashboardBrowserCookieImporter {
     {
         log("Candidate \(candidate.label) mismatch (\(signedInEmail)); continuing browser search")
         diagnostics.mismatches.append(FoundAccount(sourceLabel: candidate.label, email: signedInEmail))
-        // Mismatch still means we found a valid signed-in session. Persist it keyed by its email so if
-        // the user switches Codex accounts later, we can reuse this session immediately without another
-        // Keychain prompt.
+        // Mismatch still means we found a valid signed-in session. Persist it keyed by the
+        // candidate's resolved email/workspace so later account switches can reuse the right
+        // browser state without collapsing same-email workspaces together.
         await self.persistCookies(
             candidate: candidate,
             accountEmail: signedInEmail,
-            workspaceLabel: nil,
+            workspaceLabel: self.resolveWorkspaceLabel(from: candidate.cookies),
             logger: log)
     }
 
diff --git a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift
@@ -128,7 +128,11 @@ extension CodexWebDashboardStrategy {
             logger.append(line)
         }
         let dashboard = try await Self.fetchOpenAIWebDashboard(request, logger: log)
-        guard let usage = dashboard.toUsageSnapshot(provider: .codex, accountEmail: request.accountEmail) else {
+        guard let usage = dashboard.toUsageSnapshot(
+            provider: .codex,
+            accountEmail: request.accountEmail,
+            accountOrganization: request.workspaceLabel)
+        else {
             throw OpenAIWebCodexError.missingUsage
         }
         let credits = dashboard.toCreditsSnapshot()
diff --git a/Tests/CodexBarTests/OpenAIDashboardBrowserCookieImporterTests.swift b/Tests/CodexBarTests/OpenAIDashboardBrowserCookieImporterTests.swift
@@ -13,4 +13,11 @@ struct OpenAIDashboardBrowserCookieImporterTests {
         #expect(msg.contains("Safari=a@example.com"))
         #expect(msg.contains("Chrome=b@example.com"))
     }
+
+    @Test @MainActor
+    func `normalize workspace label trims and lowercases`() {
+        let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: BrowserDetection(cacheTTL: 0))
+        #expect(importer._normalizeWorkspaceLabelForTesting("  Team Workspace  ") == "team workspace")
+        #expect(importer._normalizeWorkspaceLabelForTesting(nil) == nil)
+    }
 }
diff --git a/Tests/CodexBarTests/OpenAIDashboardParserTests.swift b/Tests/CodexBarTests/OpenAIDashboardParserTests.swift
@@ -156,4 +156,29 @@ struct OpenAIDashboardParserTests {
         let snapshot = try decoder.decode(OpenAIDashboardSnapshot.self, from: Data(json.utf8))
         #expect(snapshot.usageBreakdown.isEmpty)
     }
+
+    @Test
+    func `to usage snapshot preserves account organization override`() {
+        let snapshot = OpenAIDashboardSnapshot(
+            signedInEmail: "user@example.com",
+            codeReviewRemainingPercent: nil,
+            creditEvents: [],
+            dailyBreakdown: [],
+            usageBreakdown: [],
+            creditsPurchaseURL: nil,
+            primaryLimit: RateWindow(usedPercent: 10, windowMinutes: 300, resetsAt: nil, resetDescription: nil),
+            secondaryLimit: nil,
+            creditsRemaining: nil,
+            accountPlan: "Plus",
+            updatedAt: Date(timeIntervalSince1970: 1_700_000_000))
+
+        let usage = snapshot.toUsageSnapshot(
+            provider: .codex,
+            accountEmail: "user@example.com",
+            accountOrganization: "Team Workspace")
+
+        #expect(usage?.accountEmail(for: .codex) == "user@example.com")
+        #expect(usage?.accountOrganization(for: .codex) == "Team Workspace")
+        #expect(usage?.loginMethod(for: .codex) == "Plus")
+    }
 }
diff --git a/Tests/CodexBarTests/TokenAccountEnvironmentPrecedenceTests.swift b/Tests/CodexBarTests/TokenAccountEnvironmentPrecedenceTests.swift
@@ -97,6 +97,23 @@ struct TokenAccountEnvironmentPrecedenceTests {
         #expect(snapshot.manualCookieHeader == "Cookie: a=b")
     }
 
+    @Test
+    func `codex token account selection preserves workspace label in app settings snapshot`() throws {
+        let settings = Self.makeSettingsStore(suite: "TokenAccountEnvironmentPrecedenceTests-codex-workspace")
+        settings.addTokenAccount(
+            provider: .codex,
+            label: "user@example.com — Team Workspace",
+            token: "Cookie: a=b")
+
+        let account = try #require(settings.selectedTokenAccount(for: .codex))
+        let snapshot = settings.codexSettingsSnapshot(tokenOverride: TokenAccountOverride(
+            provider: .codex,
+            account: account))
+
+        #expect(snapshot.accountEmail == "user@example.com")
+        #expect(snapshot.workspaceLabel == "Team Workspace")
+    }
+
     @Test
     func `claude OAuth token account overrides environment in app environment builder`() {
         let settings = Self.makeSettingsStore(suite: "TokenAccountEnvironmentPrecedenceTests-claude-app")

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,11 @@ extension CodexWebDashboardStrategy {`
`128`	`128`	`logger.append(line)`
`129`	`129`	`}`
`130`	`130`	`let dashboard = try await Self.fetchOpenAIWebDashboard(request, logger: log)`
`131`		`- guard let usage = dashboard.toUsageSnapshot(provider: .codex, accountEmail: request.accountEmail) else {`
	`131`	`+ guard let usage = dashboard.toUsageSnapshot(`
	`132`	`+ provider: .codex,`
	`133`	`+ accountEmail: request.accountEmail,`
	`134`	`+ accountOrganization: request.workspaceLabel)`
	`135`	`+ else {`
`132`	`136`	`throw OpenAIWebCodexError.missingUsage`
`133`	`137`	`}`
`134`	`138`	`let credits = dashboard.toCreditsSnapshot()`
Original file line number	Diff line number	Diff line change
`@@ -13,4 +13,11 @@ struct OpenAIDashboardBrowserCookieImporterTests {`
`13`	`13`	`#expect(msg.contains("Safari=a@example.com"))`
`14`	`14`	`#expect(msg.contains("Chrome=b@example.com"))`
`15`	`15`	`}`
	`16`	`+`
	`17`	`+ @Test @MainActor`
	`18`	+ func `normalize workspace label trims and lowercases`() {
	`19`	`+ let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: BrowserDetection(cacheTTL: 0))`
	`20`	`+ #expect(importer._normalizeWorkspaceLabelForTesting(" Team Workspace ") == "team workspace")`
	`21`	`+ #expect(importer._normalizeWorkspaceLabelForTesting(nil) == nil)`
	`22`	`+ }`
`16`	`23`	`}`