Add OAuth account proxy controls (#307)

* Add oauth account proxy controls

* Refine oauth proxy review fixes

* Fix oauth clear-proxy exchange fallback

src/server/services/oauth/antigravityProvider.ts

@@ -237,7 +237,7 @@ export const antigravityOauthProvider: OAuthProviderDefinition = {
     requiresProjectId: false,
     supportsDirectAccountRouting: true,
     supportsCloudValidation: true,
-    supportsNativeProxy: false,
+    supportsNativeProxy: true,
   },
   site: {
     name: 'Google Antigravity OAuth',

src/server/services/oauth/requestProxy.ts

@@ -1,6 +1,7 @@
 import { eq } from 'drizzle-orm';
 import { db, schema } from '../../db/index.js';
 import { getOAuthProviderDefinition } from './providers.js';
+import { getProxyUrlFromExtraConfig } from '../accountExtraConfig.js';
 import { resolveSiteProxyUrlByRequestUrl } from '../siteProxy.js';
 
 export async function resolveOauthProviderProxyUrl(provider: string): Promise<string | null> {
@@ -9,11 +10,16 @@ export async function resolveOauthProviderProxyUrl(provider: string): Promise<st
   return resolveSiteProxyUrlByRequestUrl(definition.site.url);
 }
 
-export async function resolveOauthAccountProxyUrl(siteId: number | null | undefined): Promise<string | null> {
-  if (!siteId || siteId <= 0) return null;
+export async function resolveOauthAccountProxyUrl(input: {
+  siteId?: number | null;
+  extraConfig?: string | null;
+}): Promise<string | null> {
+  const accountProxyUrl = getProxyUrlFromExtraConfig(input.extraConfig);
+  if (accountProxyUrl) return accountProxyUrl;
+  if (!input.siteId || input.siteId <= 0) return null;
   const site = await db.select({
     url: schema.sites.url,
-  }).from(schema.sites).where(eq(schema.sites.id, siteId)).get();
+  }).from(schema.sites).where(eq(schema.sites.id, input.siteId)).get();
   if (!site?.url) return null;
   return resolveSiteProxyUrlByRequestUrl(site.url);
 }

src/server/services/oauth/service.ts

@@ -1,6 +1,6 @@
 import { and, desc, eq, inArray, isNull, or, sql } from 'drizzle-orm';
 import { db, schema } from '../../db/index.js';
-import { mergeAccountExtraConfig } from '../accountExtraConfig.js';
+import { getProxyUrlFromExtraConfig, mergeAccountExtraConfig } from '../accountExtraConfig.js';
 import { refreshModelsForAccount } from '../modelService.js';
 import * as routeRefreshWorkflow from '../routeRefreshWorkflow.js';
 import {
@@ -200,6 +200,7 @@ async function upsertOauthAccount(input: {
     providerData?: Record<string, unknown>;
   };
   rebindAccountId?: number;
+  proxyUrl?: string | null;
 }) {
   const site = await ensureOauthSite(input.definition);
   const existing = await findExistingOauthAccount({
@@ -228,6 +229,7 @@ async function upsertOauthAccount(input: {
   });
   const extraConfig = mergeAccountExtraConfig(existing?.extraConfig, {
     credentialMode: 'session',
+    ...(input.proxyUrl !== undefined ? { proxyUrl: input.proxyUrl } : {}),
     oauth: buildStoredOauthState(oauth),
   });
 
@@ -284,6 +286,7 @@ export async function startOauthProviderFlow(input: {
   provider: string;
   rebindAccountId?: number;
   projectId?: string;
+  proxyUrl?: string | null;
   requestOrigin?: string;
 }) {
   const definition = getOAuthProviderDefinition(input.provider);
@@ -300,6 +303,7 @@ export async function startOauthProviderFlow(input: {
     redirectUri,
     rebindAccountId: input.rebindAccountId,
     projectId: input.projectId,
+    proxyUrl: input.proxyUrl,
   });
   return {
     provider: input.provider,
@@ -355,19 +359,22 @@ export async function handleOauthCallback(input: {
   }
 
   try {
-    const proxyUrl = await resolveOauthProviderProxyUrl(input.provider);
+    const resolvedProxyUrl = session.proxyUrl == null
+      ? await resolveOauthProviderProxyUrl(input.provider)
+      : session.proxyUrl;
     const exchange = await definition.exchangeAuthorizationCode({
       code,
       state: input.state,
       redirectUri: session.redirectUri,
       codeVerifier: session.codeVerifier,
       projectId: session.projectId,
-      proxyUrl,
+      proxyUrl: resolvedProxyUrl,
     });
     const { account, site, created, previousAccount } = await upsertOauthAccount({
       definition,
       exchange,
       rebindAccountId: session.rebindAccountId,
+      proxyUrl: session.proxyUrl,
     });
     if (!account) {
       markOauthSessionError(input.state, 'failed to persist oauth account');
@@ -534,6 +541,7 @@ export async function listOauthConnections(options: {
       routeChannelCount: routeChannelCountByAccount.get(row.accounts.id) || 0,
       lastModelSyncAt: oauth.lastModelSyncAt,
       lastModelSyncError: oauth.lastModelSyncError,
+      proxyUrl: getProxyUrlFromExtraConfig(row.accounts.extraConfig),
       site: {
         id: row.sites.id,
         name: row.sites.name,
@@ -567,7 +575,11 @@ export async function refreshOauthConnectionQuota(accountId: number) {
   return { success: true, quota };
 }
 
-export async function startOauthRebindFlow(accountId: number, requestOrigin?: string) {
+export async function startOauthRebindFlow(
+  accountId: number,
+  options?: { requestOrigin?: string; proxyUrl?: string | null },
+) {
+  const { requestOrigin, proxyUrl } = options ?? {};
   const account = await db.select().from(schema.accounts)
     .where(eq(schema.accounts.id, accountId))
     .get();
@@ -582,6 +594,9 @@ export async function startOauthRebindFlow(accountId: number, requestOrigin?: st
     provider: oauth.provider,
     rebindAccountId: accountId,
     projectId: oauth.projectId,
+    proxyUrl: proxyUrl !== undefined
+      ? proxyUrl
+      : (getProxyUrlFromExtraConfig(account.extraConfig) ?? undefined),
     requestOrigin,
   });
 }
@@ -637,7 +652,10 @@ export async function refreshOauthAccessToken(accountId: number) {
       projectId: oauth.projectId,
       providerData: oauth.providerData,
     },
-    proxyUrl: await resolveOauthAccountProxyUrl(account.siteId),
+    proxyUrl: await resolveOauthAccountProxyUrl({
+      siteId: account.siteId,
+      extraConfig: account.extraConfig,
+    }),
   });
   const nextOauth = buildOauthInfoFromAccount(account, {
     provider: oauth.provider,

src/server/services/oauth/sessionStore.ts

@@ -16,6 +16,7 @@ export type OAuthSessionRecord = {
   error?: string;
   rebindAccountId?: number;
   projectId?: string;
+  proxyUrl?: string | null;
 };
 
 export interface OAuthSessionStore {
@@ -24,6 +25,7 @@ export interface OAuthSessionStore {
     redirectUri: string;
     rebindAccountId?: number;
     projectId?: string;
+    proxyUrl?: string | null;
   }): OAuthSessionRecord;
   get(state: string): OAuthSessionRecord | null;
   markSuccess(state: string, patch: { accountId: number; siteId: number }): OAuthSessionRecord | null;
@@ -65,6 +67,7 @@ class MemoryOAuthSessionStore implements OAuthSessionStore {
     redirectUri: string;
     rebindAccountId?: number;
     projectId?: string;
+    proxyUrl?: string | null;
   }): OAuthSessionRecord {
     this.pruneExpiredSessions();
     const state = toBase64Url(randomBytes(24));
@@ -82,6 +85,7 @@ class MemoryOAuthSessionStore implements OAuthSessionStore {
       expiresAt,
       rebindAccountId: input.rebindAccountId,
       projectId: input.projectId,
+      proxyUrl: input.proxyUrl,
     };
     this.sessions.set(state, record);
     return record;
@@ -132,6 +136,7 @@ export function createOauthSession(input: {
   redirectUri: string;
   rebindAccountId?: number;
   projectId?: string;
+  proxyUrl?: string | null;
 }): OAuthSessionRecord {
   return oauthSessionStore.create(input);
 }