Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ enum CodexAuthenticatedHTTPTransport {
return response.data
case 401:
throw CodexOAuthFetchError.unauthorized
case 403:
throw CodexOAuthFetchError.forbidden
default:
let body = String(data: response.data, encoding: .utf8)
throw CodexOAuthFetchError.serverError(response.statusCode, body)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -367,6 +367,7 @@ public struct CodexUsageResponse: Decodable, Sendable {

public enum CodexOAuthFetchError: LocalizedError, Sendable {
case unauthorized
case forbidden
case invalidResponse
case serverError(Int, String?)
case networkError(Error)
Expand All @@ -375,6 +376,8 @@ public enum CodexOAuthFetchError: LocalizedError, Sendable {
switch self {
case .unauthorized:
return "Codex OAuth token expired or invalid. Run `codex login` to re-authenticate."
case .forbidden:
return "Codex account cannot access the selected workspace (HTTP 403)."
case .invalidResponse:
return "Invalid response from Codex usage API."
case let .serverError(code, message):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ struct CodexPATFetchStrategy: ProviderFetchStrategy {
switch fetchError {
case .unauthorized:
return true
case .invalidResponse, .serverError, .networkError:
case .forbidden, .invalidResponse, .serverError, .networkError:
return false
}
}
Expand Down
50 changes: 37 additions & 13 deletions Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ public enum CodexProviderDescriptor {
let oauthWithNativeRefresh: [any ProviderFetchStrategy] = [oauth, CodexOAuthNativeRefreshCLIStrategy()]
let autoStrategies: [any ProviderFetchStrategy] = context.codexWorkspaceID == nil
? [pat, oauth, cli]
: [pat, oauth]
: [pat, oauth, CodexOAuthNativeRefreshCLIStrategy()]

switch context.sourceMode {
case .oauth:
Expand Down Expand Up @@ -338,38 +338,53 @@ struct CodexCLIUsageStrategy: ProviderFetchStrategy {
}
}

/// Explicit OAuth may recover stale native credentials through the Codex CLI, without allowing
/// OAuth may recover stale native credentials through the Codex CLI, without allowing
/// missing or external credentials to silently switch sources.
struct CodexOAuthNativeRefreshCLIStrategy: ProviderFetchStrategy {
typealias CredentialRefresher = @Sendable (ProviderFetchContext) async throws -> Void

let id: String = "codex.oauth-native-refresh-cli"
let kind: ProviderFetchKind = .cli
private let binaryResolver: @Sendable (ProviderFetchContext) -> String?
private let credentialRefresher: CredentialRefresher

init(
binaryResolver: @escaping @Sendable (ProviderFetchContext) -> String? = {
CodexCLIUsageStrategy.resolvedBinary(env: $0.env)
},
credentialRefresher: @escaping CredentialRefresher = {
try await $0.fetcher.refreshNativeCodexCredentials()
})
{
self.binaryResolver = binaryResolver
self.credentialRefresher = credentialRefresher
}

func isAvailable(_ context: ProviderFetchContext) async -> Bool {
// The Codex CLI app-server has no supported way to receive CodexBar's selected managed
// workspace account header. Falling back to it would therefore report the auth.json
// workspace under a different selected workspace. Keep this path unavailable until the
// owner CLI can carry that scope explicitly.
guard context.codexWorkspaceID == nil,
context.sourceMode == .oauth,
guard context.sourceMode == .auto || context.sourceMode == .oauth,
self.binaryResolver(context) != nil,
let credentials = try? CodexOAuthCredentialsStore.loadForUsage(
env: context.env,
allowExternalSources: context.settings?.codex?.allowExternalOAuthSources == true)
else { return false }
return credentials.source == .codexHome && credentials.needsRefresh
return credentials.source == .codexHome && !credentials.isAPIKey
}

func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
try await CodexCLIUsageStrategy().fetch(context)
// Let the owner CLI rotate and persist its native tokens inside this exact CODEX_HOME.
// Then reload them and perform the normal OAuth request, which preserves CodexBar's
// selected managed-workspace header and account-ownership checks.
try await self.credentialRefresher(context)
let credentials = try CodexOAuthCredentialsStore.loadForUsage(
env: context.env,
allowExternalSources: context.settings?.codex?.allowExternalOAuthSources == true)
guard credentials.source == .codexHome else {
throw CodexOAuthCredentialsError.readOnlySource
}
guard !credentials.isAPIKey else {
throw CodexOAuthCredentialsError.missingTokens
}
return try await CodexOAuthFetchStrategy.fetch(context: context, credentials: credentials)
}

func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
Expand All @@ -394,7 +409,7 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy {
return try await Self.fetch(context: context, credentials: credentials)
}

private static func fetch(
fileprivate static func fetch(
context: ProviderFetchContext,
credentials initialCredentials: CodexOAuthCredentials) async throws -> ProviderFetchResult
{
Expand Down Expand Up @@ -477,7 +492,16 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy {
} else {
false
}
guard context.sourceMode == .auto || (context.sourceMode == .oauth && isExplicitNativeRefresh) else {
let isUnauthorized = if let fetchError = error as? CodexOAuthFetchError,
case .unauthorized = fetchError
{
true
} else {
false
}
guard context.sourceMode == .auto
|| (context.sourceMode == .oauth && (isExplicitNativeRefresh || isUnauthorized))
else {
return false
}
// Auto mode may launch the CLI as the next strategy. Keep that fallback
Expand All @@ -488,7 +512,7 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy {
switch fetchError {
case .unauthorized:
return true
case .invalidResponse, .serverError, .networkError:
case .forbidden, .invalidResponse, .serverError, .networkError:
return false
}
}
Expand Down
65 changes: 65 additions & 0 deletions Sources/CodexBarCore/UsageFetcher.swift
Original file line number Diff line number Diff line change
Expand Up @@ -867,6 +867,7 @@ private enum RPCRequestRaceResult<Value: Sendable>: Sendable {
private final class CodexRPCClient: @unchecked Sendable {
// Provider-specific by design: Codex RPC owns its dedicated subprocess log category.
private static let log = CodexBarLog.logger(LogCategories.provider(.codex, scope: "rpc"))
private static let accountRefreshTimeoutSeconds: TimeInterval = 30
private let process = Process()
private let stdin = RPCChildProcessInput()
private let stdoutPipe = Pipe()
Expand Down Expand Up @@ -987,6 +988,13 @@ private final class CodexRPCClient: @unchecked Sendable {
return try self.decodeResult(from: message)
}

func refreshAccount() async throws {
_ = try await self.request(
method: "account/read",
params: ["refreshToken": true],
timeout: Self.accountRefreshTimeoutSeconds)
}

func fetchRateLimits() async throws -> RPCRateLimitsResponse {
let message = try await self.request(method: "account/rateLimits/read")
return try self.decodeResult(from: message)
Expand Down Expand Up @@ -1135,6 +1143,45 @@ private final class CodexRPCClient: @unchecked Sendable {

// MARK: - Public fetcher used by the app

private actor CodexNativeCredentialRefreshCoordinator {
private struct Entry {
let id: UUID
let task: Task<Void, Error>
}

static let shared = CodexNativeCredentialRefreshCoordinator()

private var inFlightByHome: [String: Entry] = [:]

func refresh(
home: String,
operation: @escaping @Sendable () async throws -> Void) async throws
{
if let existing = self.inFlightByHome[home] {
try await existing.task.value
return
}

let id = UUID()
let task = Task {
try await operation()
}
Comment on lines +1166 to +1168

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 cancellation while coalescing credential refreshes

When an in-flight provider refresh is superseded—for example after switching accounts or manually refreshing again—ProviderRefreshCoordinator cancels the old fetch and waits for it before starting the replacement. This unstructured Task does not inherit that cancellation, and awaiting task.value is not cancellation-aware, so the old app-server request can remain alive for the full initialization plus 30-second renewal timeout and block the replacement refresh. Make the coordinator waiter-aware so cancellation releases the canceled caller and cancels the subprocess when no other waiter still needs it.

Useful? React with 👍 / 👎.

self.inFlightByHome[home] = Entry(id: id, task: task)
do {
try await task.value
self.clear(home: home, id: id)
} catch {
self.clear(home: home, id: id)
throw error
}
}

private func clear(home: String, id: UUID) {
guard self.inFlightByHome[home]?.id == id else { return }
self.inFlightByHome[home] = nil
}
}

public struct UsageFetcher: Sendable {
private let environment: [String: String]
private let initializeTimeoutSeconds: TimeInterval
Expand Down Expand Up @@ -1172,6 +1219,24 @@ public struct UsageFetcher: Sendable {
return usage
}

/// Ask the credential-owning Codex app-server to renew the scoped native auth file.
/// This intentionally does not consume app-server usage because it cannot carry CodexBar's
/// selected managed-workspace header.
func refreshNativeCodexCredentials() async throws {
let home = CodexHomeScope.ambientHomeURL(env: self.environment).standardizedFileURL.path
try await CodexNativeCredentialRefreshCoordinator.shared.refresh(home: home) {
let rpc = try CodexRPCClient(
arguments: self.codexArguments,
environment: self.environment,
initializeTimeoutSeconds: self.initializeTimeoutSeconds,
requestTimeoutSeconds: self.requestTimeoutSeconds,
resolveExecutable: self.codexExecutableResolver)
defer { rpc.shutdown() }
try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4")
try await rpc.refreshAccount()
}
}

public func loadLatestCLIAccountSnapshot() async throws -> CodexCLIAccountSnapshot {
let rpc = try CodexRPCClient(
arguments: self.codexArguments,
Expand Down
40 changes: 18 additions & 22 deletions Tests/CodexBarTests/CodexOAuthExpiryPipelineTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -41,12 +41,10 @@ struct CodexOAuthExpiryPipelineTests {
return
}
} else {
guard case let .serverError(statusCode, body) = error else {
Issue.record("A permission denial must retain its HTTP status")
guard case .forbidden = error else {
Issue.record("Expected a terminal permission denial")
return
}
#expect(statusCode == 403)
#expect(body == "fixture refusal")
}
#expect(CodexOAuthFetchStrategy().shouldFallback(on: error, context: context) == (code == 401))
#expect(CodexPATFetchStrategy().shouldFallback(on: error, context: context) == (code == 401))
Expand Down Expand Up @@ -98,25 +96,18 @@ struct CodexOAuthExpiryPipelineTests {
throw URLError(.cancelled)
}
let recovery = CodexOAuthNativeRefreshCLIStrategy(binaryResolver: { _ in "/fixture/codex" })
#expect(await recovery.isAvailable(context) == (mode == .oauth && !managed))
#expect(await recovery.isAvailable(context))
let outcome = await CodexAuthenticatedHTTPTransport.$overrideForTesting.withValue(transport) {
await Self.pipeline.fetch(context: context, provider: .codex)
}
if managed {
guard case let .failure(error) = outcome.result,
case .nativeRefreshRequired = error as? CodexOAuthCredentialsError
else {
Issue.record("Managed scope must retain nativeRefreshRequired without CLI recovery")
continue
}
} else {
guard case let .failure(error) = outcome.result, error is CLISelected else {
Issue.record("Native refresh must be handed to the CLI")
continue
}
#expect(outcome.attempts.last?.strategyID
== (mode == .auto ? "codex.cli" : "codex.oauth-native-refresh-cli"))
guard case let .failure(error) = outcome.result, error is CLISelected else {
Issue.record("Native refresh must be handed to the credential-owning CLI")
continue
}
let expectedRecoveryID = managed || mode == .oauth
? "codex.oauth-native-refresh-cli"
: "codex.cli"
#expect(outcome.attempts.last?.strategyID == expectedRecoveryID)
#expect(await transport.requests().isEmpty)
try fixture.expectUnchanged()
}
Expand All @@ -143,19 +134,24 @@ struct CodexOAuthExpiryPipelineTests {
await Self.pipeline.fetch(context: context, provider: .codex)
}
let unauthorized = failure == "401"
let expectsCLI = unauthorized && mode == .auto && !managed
let forbidden = failure == "403"
let expectsCLI = unauthorized
guard case let .failure(error) = outcome.result else {
Issue.record("An expiry hint cannot authenticate a rejected token")
return
}
if expectsCLI {
#expect(error is CLISelected)
#expect(outcome.attempts.last?.strategyID == "codex.cli")
let expectedRecoveryID = managed || mode == .oauth
? "codex.oauth-native-refresh-cli"
: "codex.cli"
#expect(outcome.attempts.last?.strategyID == expectedRecoveryID)
} else {
let oauthError = try #require(error as? CodexOAuthFetchError)
switch oauthError {
case .unauthorized: #expect(unauthorized)
case .serverError: #expect(failure == "403" || failure == "500")
case .forbidden: #expect(forbidden)
case .serverError: #expect(failure == "500")
case .invalidResponse: #expect(failure == "decode")
case .networkError: #expect(failure == "network")
}
Expand Down
Loading
Loading