feat: OAuth authentication for AI providers

package.json

@@ -66,7 +66,8 @@
     "appId": "com.dortort.keystone",
     "productName": "Keystone",
     "directories": {
-      "output": "release"
+      "output": "release",
+      "buildResources": "resources"
     },
     "files": [
       "dist/**/*"
@@ -90,4 +91,4 @@
       ]
     }
   }
-}
+}

src/agents/providers/GoogleAdapter.ts

@@ -1,13 +1,41 @@
 import { BaseLLMClient } from './BaseLLMClient'
 import type { ChatMessage, ChatOptions } from '@shared/types/provider'
 
+interface GoogleAuthApiKey {
+  apiKey: string
+}
+
+interface GoogleAuthOAuth {
+  oauthToken: string
+}
+
+type GoogleAuth = GoogleAuthApiKey | GoogleAuthOAuth
+
 export class GoogleAdapter extends BaseLLMClient {
-  private apiKey: string
+  private auth: GoogleAuth
   private baseUrl = 'https://generativelanguage.googleapis.com/v1beta'
 
-  constructor(apiKey: string) {
+  constructor(auth: string | GoogleAuth) {
     super()
-    this.apiKey = apiKey
+    this.auth = typeof auth === 'string' ? { apiKey: auth } : auth
+  }
+
+  updateOAuthToken(token: string): void {
+    this.auth = { oauthToken: token }
+  }
+
+  private getAuthHeaders(): Record<string, string> {
+    if ('oauthToken' in this.auth) {
+      return { Authorization: `Bearer ${this.auth.oauthToken}` }
+    }
+    return { 'x-goog-api-key': this.auth.apiKey }
+  }
+
+  private getUrl(model: string): string {
+    const base = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`
+    // When using API key (not OAuth), append key as query param is not needed
+    // since we send it via header. Just return base URL.
+    return base
   }
 
   async *chat(messages: ChatMessage[], options?: ChatOptions): AsyncIterable<string> {
@@ -35,13 +63,13 @@ export class GoogleAdapter extends BaseLLMClient {
       }
     }
 
-    const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`
+    const url = this.getUrl(model)
 
     const response = await fetch(url, {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
-        'x-goog-api-key': this.apiKey,
+        ...this.getAuthHeaders(),
       },
       body: JSON.stringify(body),
     })

src/agents/providers/OpenAIAdapter.ts

@@ -1,13 +1,45 @@
 import { BaseLLMClient } from './BaseLLMClient'
 import type { ChatMessage, ChatOptions } from '@shared/types/provider'
 
+interface OpenAIAuthApiKey {
+  apiKey: string
+}
+
+interface OpenAIAuthOAuth {
+  oauthToken: string
+  accountId?: string
+}
+
+type OpenAIAuth = OpenAIAuthApiKey | OpenAIAuthOAuth
+
 export class OpenAIAdapter extends BaseLLMClient {
-  private apiKey: string
+  private auth: OpenAIAuth
   private baseUrl = 'https://api.openai.com/v1'
 
-  constructor(apiKey: string) {
+  constructor(auth: string | OpenAIAuth) {
     super()
-    this.apiKey = apiKey
+    this.auth = typeof auth === 'string' ? { apiKey: auth } : auth
+  }
+
+  updateOAuthToken(token: string, accountId?: string): void {
+    this.auth = { oauthToken: token, accountId }
+  }
+
+  private getHeaders(): Record<string, string> {
+    const headers: Record<string, string> = {
+      'Content-Type': 'application/json',
+    }
+
+    if ('oauthToken' in this.auth) {
+      headers['Authorization'] = `Bearer ${this.auth.oauthToken}`
+      if (this.auth.accountId) {
+        headers['chatgpt-account-id'] = this.auth.accountId
+      }
+    } else {
+      headers['Authorization'] = `Bearer ${this.auth.apiKey}`
+    }
+
+    return headers
   }
 
   async *chat(messages: ChatMessage[], options?: ChatOptions): AsyncIterable<string> {
@@ -31,10 +63,7 @@ export class OpenAIAdapter extends BaseLLMClient {
 
     const response = await fetch(`${this.baseUrl}/chat/completions`, {
       method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: `Bearer ${this.apiKey}`,
-      },
+      headers: this.getHeaders(),
       body: JSON.stringify(body),
     })
 

src/agents/providers/ProviderManager.ts

@@ -1,4 +1,4 @@
-import type { ProviderType, ProviderConfig } from '@shared/types/provider'
+import type { ProviderType, ProviderConfig, LegacyProviderConfig } from '@shared/types/provider'
 import { BaseLLMClient } from './BaseLLMClient'
 import { AnthropicAdapter } from './AnthropicAdapter'
 import { OpenAIAdapter } from './OpenAIAdapter'
@@ -8,21 +8,36 @@
   private providers: Map<ProviderType, BaseLLMClient> = new Map()
   private activeProvider: ProviderType | null = null
 
-  configure(config: ProviderConfig): void {
+  configure(config: ProviderConfig | LegacyProviderConfig): void {
     let client: BaseLLMClient
 
+    // Normalize: legacy configs lack authMethod field
+    const authMethod = 'authMethod' in config ? config.authMethod : 'apiKey'
+
     switch (config.type) {
       case 'anthropic':
-        client = new AnthropicAdapter(config.apiKey)
+        // Anthropic only supports API keys
+        if (authMethod === 'oauth') {
+          throw new Error('Anthropic does not support OAuth authentication')
+        }
+        client = new AnthropicAdapter('apiKey' in config ? config.apiKey : '')
         break
       case 'openai':
-        client = new OpenAIAdapter(config.apiKey)
+        if (authMethod === 'oauth' && 'oauthToken' in config) {
+          client = new OpenAIAdapter({ oauthToken: config.oauthToken, accountId: config.accountId })
+        } else {
+          client = new OpenAIAdapter('apiKey' in config ? config.apiKey : '')
+        }
         break
       case 'google':
-        client = new GoogleAdapter(config.apiKey)
+        if (authMethod === 'oauth' && 'oauthToken' in config) {
+          client = new GoogleAdapter({ oauthToken: config.oauthToken })
+        } else {
+          client = new GoogleAdapter('apiKey' in config ? config.apiKey : '')
+        }
         break
       default:
-        throw new Error(`Unknown provider: ${config.type}`)
+        throw new Error(`Unknown provider: ${(config as any).type}`)
     }
 
     this.providers.set(config.type, client)
@@ -31,6 +46,17 @@
     }
   }
 
+  updateOAuthToken(provider: ProviderType, token: string, accountId?: string): void {
+    const client = this.providers.get(provider)
+    if (!client) return
+
+    if (provider === 'openai' && client instanceof OpenAIAdapter) {
+      client.updateOAuthToken(token, accountId)
+    } else if (provider === 'google' && client instanceof GoogleAdapter) {
+      client.updateOAuthToken(token)
+    }
+  }
+
   setActive(type: ProviderType): void {
     if (!this.providers.has(type)) {
       throw new Error(`Provider ${type} not configured`)
@@ -60,4 +86,12 @@
   getConfiguredProviders(): ProviderType[] {
     return Array.from(this.providers.keys())
   }
+
+  removeProvider(type: ProviderType): void {
+    this.providers.delete(type)
+    if (this.activeProvider === type) {
+      const remaining = Array.from(this.providers.keys())
+      this.activeProvider = remaining.length > 0 ? remaining[0] : null
+    }
+  }
 }

src/main/ipc/ai.router.ts

@@ -144,16 +144,29 @@ export const aiRouter = router({
     .input(
       z.object({
         type: z.enum(['openai', 'anthropic', 'google']),
-        apiKey: z.string().min(1),
+        apiKey: z.string().min(1).optional(),
+        authMethod: z.enum(['apiKey', 'oauth']).optional(),
+        oauthToken: z.string().optional(),
+        accountId: z.string().optional(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      // Configure the in-memory provider
-      ctx.providerManager.configure({ type: input.type, apiKey: input.apiKey })
-      ctx.providerManager.setActive(input.type)
+      const authMethod = input.authMethod || 'apiKey'
+
+      if (authMethod === 'oauth' && input.oauthToken) {
+        ctx.providerManager.configure({
+          type: input.type,
+          authMethod: 'oauth',
+          oauthToken: input.oauthToken,
+          accountId: input.accountId,
+        })
+      } else if (input.apiKey) {
+        ctx.providerManager.configure({ type: input.type, authMethod: 'apiKey', apiKey: input.apiKey })
+        ctx.settingsService.setApiKey(input.type, input.apiKey)
+        ctx.settingsService.setAuthMethod(input.type, 'apiKey')
+      }
 
-      // Persist the API key and active provider selection
-      ctx.settingsService.setApiKey(input.type, input.apiKey)
+      ctx.providerManager.setActive(input.type)
       ctx.settingsService.setActiveProvider(input.type)
 
       return {

src/main/ipc/context.ts

@@ -1,11 +1,13 @@
 import type { Context } from './trpc'
+import type { ProviderType } from '@shared/types/provider'
 import { DatabaseService } from '../services/DatabaseService'
 import { ProjectService } from '../services/ProjectService'
 import { DocumentService } from '../services/DocumentService'
 import { ThreadService } from '../services/ThreadService'
 import { ProviderManager } from '../../agents/providers/ProviderManager'
 import { Orchestrator } from '../../agents/orchestrator/Orchestrator'
 import { SettingsService } from '../services/SettingsService'
+import { OAuthService } from '../services/OAuthService'
 
 let _context: Context | null = null
 
@@ -20,13 +22,58 @@ export async function createContext(): Promise<Context> {
   const settingsService = new SettingsService()
   const providerManager = new ProviderManager()
   const orchestrator = new Orchestrator()
+  const oauthService = new OAuthService()
 
-  // Restore any previously configured providers from persisted settings
+  // Wire OAuth token refresh callback
+  oauthService.onTokenRefresh((provider, tokens) => {
+    settingsService.setOAuthTokens(provider, tokens)
+    providerManager.updateOAuthToken(provider, tokens.accessToken, tokens.accountId)
+  })
+
+  // Restore API-key-based providers from persisted settings
   for (const provider of settingsService.getConfiguredProviders()) {
-    const apiKey = settingsService.getApiKey(provider)
-    if (apiKey) {
+    const authMethod = settingsService.getAuthMethod(provider)
+    if (authMethod === 'apiKey') {
+      const apiKey = settingsService.getApiKey(provider)
+      if (apiKey) {
+        try {
+          providerManager.configure({ type: provider as ProviderType, authMethod: 'apiKey', apiKey })
+        } catch {
+          // Skip invalid providers silently
+        }
+      }
+    }
+  }
+
+  // Restore OAuth-based providers from persisted settings
+  for (const provider of settingsService.getOAuthConfiguredProviders()) {
+    const tokens = settingsService.getOAuthTokens(provider)
+    const authMethod = settingsService.getAuthMethod(provider)
+    if (tokens && authMethod === 'oauth') {
       try {
-        providerManager.configure({ type: provider as 'openai' | 'anthropic' | 'google', apiKey })
+        // Check if token is expired — attempt immediate refresh if so
+        if (tokens.expiresAt < Date.now() && tokens.refreshToken) {
+          const refreshed = await oauthService.refreshToken(provider as ProviderType, tokens)
+          if (refreshed) {
+            settingsService.setOAuthTokens(provider, refreshed)
+            providerManager.configure({
+              type: provider as ProviderType,
+              authMethod: 'oauth',
+              oauthToken: refreshed.accessToken,
+              accountId: refreshed.accountId,
+            })
+            oauthService.scheduleRefresh(provider as ProviderType, refreshed)
+          }
+          // If refresh fails, skip — user will need to re-authenticate
+        } else {
+          providerManager.configure({
+            type: provider as ProviderType,
+            authMethod: 'oauth',
+            oauthToken: tokens.accessToken,
+            accountId: tokens.accountId,
+          })
+          oauthService.scheduleRefresh(provider as ProviderType, tokens)
+        }
       } catch {
         // Skip invalid providers silently
       }
@@ -37,12 +84,12 @@ export async function createContext(): Promise<Context> {
   const activeProvider = settingsService.getActiveProvider()
   if (activeProvider && providerManager.isConfigured()) {
     try {
-      providerManager.setActive(activeProvider as 'openai' | 'anthropic' | 'google')
+      providerManager.setActive(activeProvider as ProviderType)
     } catch {
       // Active provider may no longer be configured
     }
   }
 
-  _context = { db, projectService, documentService, threadService, settingsService, providerManager, orchestrator }
+  _context = { db, projectService, documentService, threadService, settingsService, providerManager, orchestrator, oauthService }
   return _context
 }