feat: refine OAuth proxy controls and quota UI (#433)

* feat: refine oauth management proxy controls

* fix: preserve oauth refresh state and tighten oauth rate limiting

* fix: satisfy codeql oauth rate limiting

* fix: address oauth review comments

* fix: move oauth sensitive limits into prehandlers

* fix: cap oauth quota refresh batch size

src/server/proxy-core/surfaces/chatSurface.ts

@@ -559,6 +559,7 @@ export async function handleChatSurfaceRequest(
             modelName,
             parsedUsage,
             upstreamUsagePresent,
+            upstreamHeaders: upstream.headers,
             requestStartedAtMs: startTime,
             isStream: true,
             firstByteLatencyMs,
@@ -922,6 +923,7 @@ export async function handleChatSurfaceRequest(
         modelName,
         parsedUsage,
         upstreamUsagePresent,
+        upstreamHeaders: upstream.headers,
         requestStartedAtMs: startTime,
         isStream: false,
         firstByteLatencyMs,

src/server/proxy-core/surfaces/openAiResponsesSurface.ts

@@ -771,6 +771,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
               modelName,
               parsedUsage,
               upstreamUsagePresent,
+              upstreamHeaders: upstream.headers,
               requestStartedAtMs: startTime,
               isStream: true,
               firstByteLatencyMs,
@@ -1221,6 +1222,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
             modelName,
             parsedUsage,
             upstreamUsagePresent,
+            upstreamHeaders: upstream.headers,
             requestStartedAtMs: startTime,
             isStream: false,
             firstByteLatencyMs,

src/server/proxy-core/surfaces/sharedSurface.test.ts

@@ -16,6 +16,7 @@ const reportProxyAllFailedMock = vi.fn();
 const reportTokenExpiredMock = vi.fn();
 const isTokenExpiredErrorMock = vi.fn();
 const shouldRetryProxyRequestMock = vi.fn();
+const recordOauthQuotaHeadersSnapshotMock = vi.fn();
 const recordOauthQuotaResetHintMock = vi.fn();
 const recordSuccessMock = vi.fn();
 const resolveProxyUsageWithSelfLogFallbackMock = vi.fn();
@@ -89,6 +90,7 @@ vi.mock('../../services/proxyRetryPolicy.js', () => ({
 }));
 
 vi.mock('../../services/oauth/quota.js', () => ({
+  recordOauthQuotaHeadersSnapshot: (...args: unknown[]) => recordOauthQuotaHeadersSnapshotMock(...args),
   recordOauthQuotaResetHint: (...args: unknown[]) => recordOauthQuotaResetHintMock(...args),
 }));
 
@@ -126,6 +128,7 @@ describe('selectSurfaceChannelForAttempt', () => {
     reportTokenExpiredMock.mockReset();
     isTokenExpiredErrorMock.mockReset();
     shouldRetryProxyRequestMock.mockReset();
+    recordOauthQuotaHeadersSnapshotMock.mockReset();
     recordOauthQuotaResetHintMock.mockReset();
     recordSuccessMock.mockReset();
     resolveProxyUsageWithSelfLogFallbackMock.mockReset();
@@ -1076,6 +1079,59 @@ describe('selectSurfaceChannelForAttempt', () => {
     }));
   });
 
+  it('captures codex quota headers from successful upstream responses as best-effort bookkeeping', async () => {
+    resolveProxyUsageWithSelfLogFallbackMock.mockResolvedValue({
+      promptTokens: 10,
+      completionTokens: 5,
+      totalTokens: 15,
+      recoveredFromSelfLog: false,
+      estimatedCostFromQuota: 0,
+      selfLogBillingMeta: null,
+      usageSource: 'upstream',
+    });
+    resolveProxyLogBillingMock.mockResolvedValue({
+      estimatedCost: 0.12,
+      billingDetails: null,
+    });
+    recordOauthQuotaHeadersSnapshotMock.mockResolvedValue(null);
+    const logSuccess = vi.fn().mockResolvedValue(undefined);
+
+    const { recordSurfaceSuccess } = await import('./sharedSurface.js');
+    await recordSurfaceSuccess({
+      selected: {
+        channel: { id: 11, routeId: 22 },
+        account: { id: 33, username: 'oauth-user' },
+        site: { id: 44, url: 'https://upstream.example.com', name: 'Codex OAuth' },
+        tokenValue: 'live-token',
+        tokenName: 'default',
+        actualModel: 'upstream-model',
+      },
+      requestedModel: 'gpt-5.2',
+      modelName: 'upstream-model',
+      parsedUsage: {
+        promptTokens: 10,
+        completionTokens: 5,
+        totalTokens: 15,
+      },
+      requestStartedAtMs: 1000,
+      latencyMs: 250,
+      retryCount: 0,
+      upstreamPath: '/v1/responses',
+      upstreamHeaders: new Headers({
+        'x-codex-primary-used-percent': '61',
+        'x-codex-secondary-used-percent': '13',
+      }),
+      logSuccess,
+    });
+
+    await vi.waitFor(() => {
+      expect(recordOauthQuotaHeadersSnapshotMock).toHaveBeenCalledWith({
+        accountId: 33,
+        headers: expect.any(Headers),
+      });
+    });
+  });
+
   it('treats success metrics as best-effort when requested', async () => {
     resolveProxyUsageWithSelfLogFallbackMock.mockRejectedValueOnce(new Error('billing failed'));
     const logSuccess = vi.fn().mockResolvedValue(undefined);

src/server/proxy-core/surfaces/sharedSurface.ts

@@ -14,7 +14,7 @@ import { insertProxyLog } from '../../services/proxyLogStore.js';
 import { dispatchRuntimeRequest } from '../../services/runtimeDispatch.js';
 import type { BuiltEndpointRequest } from '../orchestration/endpointFlow.js';
 import { buildUpstreamUrl } from '../orchestration/upstreamRequest.js';
-import { recordOauthQuotaResetHint } from '../../services/oauth/quota.js';
+import { recordOauthQuotaHeadersSnapshot, recordOauthQuotaResetHint } from '../../services/oauth/quota.js';
 import { refreshOauthAccessTokenSingleflight } from '../../services/oauth/refreshSingleflight.js';
 import { proxyChannelCoordinator } from '../../services/proxyChannelCoordinator.js';
 import { readRuntimeResponseText } from '../executors/types.js';
@@ -331,6 +331,7 @@ export async function recordSurfaceSuccess(input: {
   modelName: string;
   parsedUsage: SurfaceUsageSummary;
   upstreamUsagePresent?: boolean;
+  upstreamHeaders?: { get(name: string): string | null } | null;
   requestStartedAtMs: number;
   isStream?: boolean | null;
   firstByteLatencyMs?: number | null;
@@ -451,6 +452,15 @@ export async function recordSurfaceSuccess(input: {
     upstreamPath: input.upstreamPath,
   });
 
+  if (input.upstreamHeaders) {
+    void recordOauthQuotaHeadersSnapshot({
+      accountId: input.selected.account.id,
+      headers: input.upstreamHeaders,
+    }).catch((error) => {
+      console.warn('[proxy/shared] failed to record oauth quota headers', error);
+    });
+  }
+
   return {
     resolvedUsage,
     estimatedCost,