Skip to content
Merged
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
4 changes: 4 additions & 0 deletions Config/ContextPanel.entitlements
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,9 @@
<string>group.com.shinycomputers.contextpanel</string>
<string>MM5YXC7T6E.group.com.shinycomputers.contextpanel</string>
</array>
<key>keychain-access-groups</key>
<array>
<string>MM5YXC7T6E.com.shinycomputers.contextpanel.credentials</string>
</array>
</dict>
</plist>
4 changes: 4 additions & 0 deletions Config/ContextPanelAppStore.entitlements
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,9 @@
<string>group.com.shinycomputers.contextpanel</string>
<string>MM5YXC7T6E.group.com.shinycomputers.contextpanel</string>
</array>
<key>keychain-access-groups</key>
<array>
<string>MM5YXC7T6E.com.shinycomputers.contextpanel.credentials</string>
</array>
</dict>
</plist>
4 changes: 4 additions & 0 deletions Config/ContextPanelRefreshAgent.entitlements
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,9 @@
<string>group.com.shinycomputers.contextpanel</string>
<string>MM5YXC7T6E.group.com.shinycomputers.contextpanel</string>
</array>
<key>keychain-access-groups</key>
<array>
<string>MM5YXC7T6E.com.shinycomputers.contextpanel.credentials</string>
</array>
</dict>
</plist>
4 changes: 4 additions & 0 deletions Config/ContextPanelRefreshAgentAppStore.entitlements
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,9 @@
<string>group.com.shinycomputers.contextpanel</string>
<string>MM5YXC7T6E.group.com.shinycomputers.contextpanel</string>
</array>
<key>keychain-access-groups</key>
<array>
<string>MM5YXC7T6E.com.shinycomputers.contextpanel.credentials</string>
</array>
</dict>
</plist>
7 changes: 6 additions & 1 deletion Sources/ContextPanelCore/AccountConfigurationStore.swift
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ public struct AccountConfigurationStore: Sendable {
public enum AccountConnectorFactory {
public static func connectors(
from document: AccountConfigurationDocument,
credentialStore: (any ProviderCredentialStore)? = nil,
environment: [String: String] = ProcessInfo.processInfo.environment,
geminiMetadataFileLoader: @escaping @Sendable (String) throws -> String = { path in
try String(contentsOfFile: NSString(string: path).expandingTildeInPath, encoding: .utf8)
Expand Down Expand Up @@ -204,7 +205,7 @@ public enum AccountConnectorFactory {
accountName: account.displayName,
clientID: metadata.clientID,
clientSecret: metadata.clientSecret
)])
)], credentialStore: credentialStore, credentialKey: geminiCredentialKey(for: account))
case .claudeLocalStatus:
return ClaudeLocalStatusConnector(accounts: [ClaudeAccountConfiguration(
accountName: account.displayName,
Expand All @@ -230,6 +231,10 @@ public enum AccountConnectorFactory {
return GeminiOAuthClientMetadata(clientID: clientID, clientSecret: clientSecret)
}

public static func geminiCredentialKey(for account: LocalProviderAccountConfiguration) -> ProviderCredentialKey {
ProviderCredentialKey(provider: .google, accountID: account.id, kind: "gemini-oauth")
}

private static func codexAuthPath(for authPath: String) -> String {
let expanded = NSString(string: authPath).expandingTildeInPath
let url = URL(fileURLWithPath: expanded)
Expand Down
17 changes: 15 additions & 2 deletions Sources/ContextPanelCore/GeminiCodeAssistQuota.swift
Original file line number Diff line number Diff line change
Expand Up @@ -215,17 +215,23 @@ public struct GeminiCodeAssistConnector: ProviderConnector {
private let accounts: [GeminiAccountConfiguration]
private let httpClient: any ConnectorHTTPClient
private let fileLoader: @Sendable (String) throws -> Data
private let credentialStore: (any ProviderCredentialStore)?
private let credentialKey: ProviderCredentialKey?

public init(
accounts: [GeminiAccountConfiguration],
httpClient: any ConnectorHTTPClient = URLSessionConnectorHTTPClient(),
fileLoader: @escaping @Sendable (String) throws -> Data = { path in
try Data(contentsOf: URL(fileURLWithPath: NSString(string: path).expandingTildeInPath))
}
},
credentialStore: (any ProviderCredentialStore)? = nil,
credentialKey: ProviderCredentialKey? = nil
) {
self.accounts = accounts
self.httpClient = httpClient
self.fileLoader = fileLoader
self.credentialStore = credentialStore
self.credentialKey = credentialKey
}

public func refresh(now: Date) async -> ConnectorRefreshResult {
Expand All @@ -241,7 +247,7 @@ public struct GeminiCodeAssistConnector: ProviderConnector {
let localAccountID = ConnectorRedactor.localAccountID(provider: provider, path: account.authPath)

do {
let credentials = try JSONDecoder().decode(GeminiOAuthCredentials.self, from: try fileLoader(account.authPath))
let credentials = try JSONDecoder().decode(GeminiOAuthCredentials.self, from: try credentialData(for: account))
let accessToken = try await refreshedAccessToken(credentials: credentials, account: account)
let loadResponse = try await loadCodeAssist(accessToken: accessToken, endpoint: account.codeAssistEndpoint)
guard let project = loadResponse.cloudaicompanionProject, !project.isEmpty else {
Expand Down Expand Up @@ -276,6 +282,13 @@ public struct GeminiCodeAssistConnector: ProviderConnector {
}
}

private func credentialData(for account: GeminiAccountConfiguration) throws -> Data {
if let credentialStore, let credentialKey, let data = try credentialStore.data(for: credentialKey) {
return data
}
return try fileLoader(account.authPath)
}

private func refreshedAccessToken(credentials: GeminiOAuthCredentials, account: GeminiAccountConfiguration) async throws -> String {
guard let refreshToken = credentials.refreshToken, !refreshToken.isEmpty else {
throw ConnectorError.invalidAuth("Gemini OAuth file does not contain a refresh token")
Expand Down
123 changes: 123 additions & 0 deletions Sources/ContextPanelCore/ProviderCredentialStore.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
import Foundation
import Security

public protocol ProviderCredentialStore: Sendable {
func data(for key: ProviderCredentialKey) throws -> Data?
func store(_ data: Data, for key: ProviderCredentialKey) throws
func remove(_ key: ProviderCredentialKey) throws
}

public struct ProviderCredentialKey: Hashable, Sendable {
public let provider: Provider
public let accountID: String
public let kind: String

public init(provider: Provider, accountID: String, kind: String) {
self.provider = provider
self.accountID = accountID
self.kind = kind
}

var service: String {
"com.shinycomputers.contextpanel.\(provider.rawValue).\(kind)"
}

var account: String {
accountID
}
}

public enum ProviderCredentialStoreError: LocalizedError, Equatable, Sendable {
case keychainStatus(operation: String, status: OSStatus)

public var errorDescription: String? {
switch self {
case let .keychainStatus(operation, status):
"Keychain \(operation) failed with status \(status)"
}
}
}

public struct KeychainProviderCredentialStore: ProviderCredentialStore {
public static let defaultAccessGroup = "MM5YXC7T6E.com.shinycomputers.contextpanel.credentials"

private let accessGroup: String?

public init(accessGroup: String? = defaultAccessGroup) {
self.accessGroup = accessGroup
}

public func data(for key: ProviderCredentialKey) throws -> Data? {
var query = baseQuery(for: key)
query[kSecReturnData as String] = true
query[kSecMatchLimit as String] = kSecMatchLimitOne

var item: CFTypeRef?
let status = SecItemCopyMatching(query as CFDictionary, &item)
if status == errSecItemNotFound { return nil }
guard status == errSecSuccess else {
throw ProviderCredentialStoreError.keychainStatus(operation: "read", status: status)
}
return item as? Data
}

public func store(_ data: Data, for key: ProviderCredentialKey) throws {
var query = baseQuery(for: key)
let attributes: [String: Any] = [
kSecValueData as String: data,
kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly,
]

let updateStatus = SecItemUpdate(query as CFDictionary, attributes as CFDictionary)
if updateStatus == errSecSuccess { return }
guard updateStatus == errSecItemNotFound else {
throw ProviderCredentialStoreError.keychainStatus(operation: "update", status: updateStatus)
}

query.merge(attributes) { _, new in new }
let addStatus = SecItemAdd(query as CFDictionary, nil)
guard addStatus == errSecSuccess else {
throw ProviderCredentialStoreError.keychainStatus(operation: "store", status: addStatus)
}
}

public func remove(_ key: ProviderCredentialKey) throws {
let status = SecItemDelete(baseQuery(for: key) as CFDictionary)
guard status == errSecSuccess || status == errSecItemNotFound else {
throw ProviderCredentialStoreError.keychainStatus(operation: "delete", status: status)
}
}

private func baseQuery(for key: ProviderCredentialKey) -> [String: Any] {
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: key.service,
kSecAttrAccount as String: key.account,
]
if let accessGroup {
query[kSecAttrAccessGroup as String] = accessGroup
}
return query
}
}

public final class InMemoryProviderCredentialStore: ProviderCredentialStore, @unchecked Sendable {
private var values: [ProviderCredentialKey: Data]
private let lock = NSLock()

public init(values: [ProviderCredentialKey: Data] = [:]) {
self.values = values
}

public func data(for key: ProviderCredentialKey) throws -> Data? {
lock.withLock { values[key] }
}

public func store(_ data: Data, for key: ProviderCredentialKey) throws {
lock.withLock { values[key] = data }
}

public func remove(_ key: ProviderCredentialKey) throws {
lock.withLock { _ = values.removeValue(forKey: key) }
}
}
16 changes: 13 additions & 3 deletions Sources/ContextPanelCore/SnapshotRefreshService.swift
Original file line number Diff line number Diff line change
Expand Up @@ -158,10 +158,16 @@ public struct SnapshotRefreshRunner: Sendable {
public struct SnapshotRefreshService: Sendable {
private let accountStore: AccountConfigurationStore
private let stores: SnapshotRefreshStores
private let credentialStore: (any ProviderCredentialStore)?

public init(accountStore: AccountConfigurationStore, stores: SnapshotRefreshStores) {
public init(
accountStore: AccountConfigurationStore,
stores: SnapshotRefreshStores,
credentialStore: (any ProviderCredentialStore)? = nil
) {
self.accountStore = accountStore
self.stores = stores
self.credentialStore = credentialStore
}

public static func appDefault() -> SnapshotRefreshService {
Expand All @@ -170,7 +176,8 @@ public struct SnapshotRefreshService: Sendable {
configurationURL: ContextPanelLocations.accountConfigurationURL(),
fallbackConfigurationURL: ContextPanelLocations.legacyAccountConfigurationURL()
),
stores: .appDefault()
stores: .appDefault(),
credentialStore: KeychainProviderCredentialStore()
)
}

Expand All @@ -188,7 +195,10 @@ public struct SnapshotRefreshService: Sendable {

public func refresh(now: Date = Date()) async throws -> SnapshotRefreshOutcome {
let accountDocument = accountStore.load(now: now).document
let connectors = AccountConnectorFactory.connectors(from: accountDocument)
let connectors = AccountConnectorFactory.connectors(
from: accountDocument,
credentialStore: credentialStore
)
let refreshResult = await ProviderConnectorRuntime(connectors: connectors).refreshAll(now: now)
return try saveMerged(refreshResult: refreshResult, savedAt: now)
}
Expand Down
21 changes: 21 additions & 0 deletions Sources/ContextPanelPreview/ContextPanelPreviewApp.swift
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,10 @@ struct SettingsPane: View {
Text(account.isEnabled ? "Enabled" : "Disabled")
.font(.system(size: 11, weight: .semibold))
.foregroundStyle(account.isEnabled ? CPTheme.statusColor(.healthy) : CPTheme.tertiaryText)
if model.canImportCredentials(for: account) {
Button("Import") { model.importCredentials(for: account) }
.controlSize(.small)
}
}
Text(model.detailText(for: account))
.font(.system(size: 11))
Expand Down Expand Up @@ -212,6 +216,7 @@ final class SettingsPaneModel: ObservableObject {
private let resetPrimerSettingsStore = ResetPrimerSettingsStore(
settingsURL: ContextPanelLocations.resetPrimerSettingsURL(appGroupID: ContextPanelLocations.appGroupID)
)
private let credentialStore = KeychainProviderCredentialStore()

private var widgetPreferenceStores: WidgetDisplayPreferencesStoreSet {
WidgetDisplayPreferencesStoreSet(stores: [
Expand Down Expand Up @@ -290,6 +295,22 @@ final class SettingsPaneModel: ObservableObject {
saveResetPrimerSettings(updated)
}

func canImportCredentials(for account: LocalProviderAccountConfiguration) -> Bool {
account.connectorKind == .geminiCodeAssist && account.authPath != nil
}

func importCredentials(for account: LocalProviderAccountConfiguration) {
guard let authPath = account.authPath else { return }
do {
let expanded = NSString(string: authPath).expandingTildeInPath
let data = try Data(contentsOf: URL(fileURLWithPath: expanded))
try credentialStore.store(data, for: AccountConnectorFactory.geminiCredentialKey(for: account))
errorMessage = nil
} catch {
errorMessage = error.localizedDescription
}
}

private func saveResetPrimerSettings(_ updated: ResetPrimerSettings) {
do {
try resetPrimerSettingsStore.save(updated)
Expand Down
48 changes: 48 additions & 0 deletions Tests/ContextPanelCoreTests/ProviderConnectorTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,54 @@ import Testing
#expect(http.requests[2].body.flatMap { String(data: $0, encoding: .utf8) }?.contains("project-secret") == true)
}

@Test func geminiConnectorPrefersImportedCredentials() async throws {
let importedCredentials = #"{"refresh_token":"imported-refresh"}"#.data(using: .utf8)!
let fileCredentials = #"{"refresh_token":"file-refresh"}"#.data(using: .utf8)!
let key = ProviderCredentialKey(provider: .google, accountID: "gemini", kind: "gemini-oauth")
let store = InMemoryProviderCredentialStore(values: [key: importedCredentials])
let http = StubHTTPClient(responses: [
ConnectorHTTPResponse(statusCode: 200, data: #"{"access_token":"access"}"#.data(using: .utf8)!),
ConnectorHTTPResponse(statusCode: 200, data: #"{"cloudaicompanionProject":"project","currentTier":{"name":"Gemini Code Assist"}}"#.data(using: .utf8)!),
ConnectorHTTPResponse(statusCode: 200, data: #"{"buckets":[]}"#.data(using: .utf8)!),
])
let connector = GeminiCodeAssistConnector(
accounts: [GeminiAccountConfiguration(authPath: "/tmp/gemini.json", accountName: "Gemini", clientID: "client", clientSecret: "secret")],
httpClient: http,
fileLoader: { _ in fileCredentials },
credentialStore: store,
credentialKey: key
)

_ = await connector.refresh(now: Date(timeIntervalSince1970: 0))

let body = try #require(http.requests.first?.body.flatMap { String(data: $0, encoding: .utf8) })
#expect(body.contains("imported-refresh"))
#expect(body.contains("file-refresh") == false)
}

@Test func geminiConnectorFallsBackToAuthFileWhenImportedCredentialsAreMissing() async throws {
let fileCredentials = #"{"refresh_token":"file-refresh"}"#.data(using: .utf8)!
let key = ProviderCredentialKey(provider: .google, accountID: "gemini", kind: "gemini-oauth")
let store = InMemoryProviderCredentialStore()
let http = StubHTTPClient(responses: [
ConnectorHTTPResponse(statusCode: 200, data: #"{"access_token":"access"}"#.data(using: .utf8)!),
ConnectorHTTPResponse(statusCode: 200, data: #"{"cloudaicompanionProject":"project","currentTier":{"name":"Gemini Code Assist"}}"#.data(using: .utf8)!),
ConnectorHTTPResponse(statusCode: 200, data: #"{"buckets":[]}"#.data(using: .utf8)!),
])
let connector = GeminiCodeAssistConnector(
accounts: [GeminiAccountConfiguration(authPath: "/tmp/gemini.json", accountName: "Gemini", clientID: "client", clientSecret: "secret")],
httpClient: http,
fileLoader: { _ in fileCredentials },
credentialStore: store,
credentialKey: key
)

_ = await connector.refresh(now: Date(timeIntervalSince1970: 0))

let body = try #require(http.requests.first?.body.flatMap { String(data: $0, encoding: .utf8) })
#expect(body.contains("file-refresh"))
}

@Test func claudeConnectorReportsUnknownLiveAllowanceFromLocalStatus() async throws {
let auth = #"{"loggedIn":true,"authMethod":"claude.ai","apiProvider":"firstParty","subscriptionType":"pro"}"#.data(using: .utf8)!
let stats = #"{"version":3,"lastComputedDate":"2026-04-26","dailyActivity":[],"modelUsage":{},"totalSessions":2,"totalMessages":3}"#.data(using: .utf8)!
Expand Down