feat(proxy): add upstream error and rate-limit tracking support

src/tests/unit/proxy-parity-fixtures.test.ts

@@ -10,6 +10,8 @@ const mockTokenManager = {
   getNextToken: vi.fn(),
   markAsRateLimited: vi.fn(),
   markAsForbidden: vi.fn(),
+  markFromUpstreamError: vi.fn(),
+  recordParityError: vi.fn(),
 };
 const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
 

src/tests/unit/proxy-retry-mock.test.ts

@@ -8,6 +8,10 @@ const mockTokenManager = {
   getNextToken: vi.fn(),
   markAsRateLimited: vi.fn(),
   markAsForbidden: vi.fn(),
+  markFromUpstreamError: vi.fn(),
+  recordParityError: vi.fn(),
+  getModelOutputLimitForAccount: vi.fn(),
+  getModelThinkingBudgetForAccount: vi.fn(),
 };
 const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
 
@@ -27,7 +31,7 @@ class TestableProxyService extends ProxyService {
   }
 
   public testModelHeaders(model: string): Record<string, string> {
-    return (this as any).buildModelSpecificHeaders(model);
+    return (this as any).createModelSpecificHeaders(model);
   }
 }
 
@@ -54,7 +58,7 @@ describe('ProxyService Empty Stream Retry Logic', () => {
 
   it('classifies retry matrix consistently', () => {
     const service = new TestableProxyService();
-    const classify = (message: string) => (service as any).classifyUpstreamError(message);
+    const classify = (message: string) => (service as any).classifyUpstreamFailure(message);
 
     expect(classify('401 unauthorized token')).toEqual({
       retry: true,
@@ -401,6 +405,22 @@ describe('ProxyService Empty Stream Retry Logic', () => {
     expect(internalPayload).not.toHaveProperty('sessionId');
   });
 
+  it('normalizes Gemini 3.1 preview alias to Gemini 3.1 Pro High for upstream', async () => {
+    const service = new TestableProxyService();
+    mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1'));
+    mockGeminiClient.generateInternal.mockResolvedValue({
+      candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }],
+      usageMetadata: { totalTokenCount: 5 },
+    });
+
+    await service.handleGeminiGenerateContent('models/gemini-3.1-pro-preview', {
+      contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
+    } as any);
+
+    const internalPayload = mockGeminiClient.generateInternal.mock.calls[0][0];
+    expect(internalPayload.model).toBe('gemini-3.1-pro-high');
+  });
+
   it('strips non-parity Gemini usage metadata fields', async () => {
     const service = new TestableProxyService();
     mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1'));

src/tests/unit/rate-limit-tracker.test.ts

@@ -0,0 +1,108 @@
+import { describe, expect, it } from 'vitest';
+import { RateLimitReason, RateLimitTracker } from '../../server/modules/proxy/rate-limit-tracker';
+
+describe('RateLimitTracker parity replay', () => {
+  it('uses Retry-After header before body/default', () => {
+    const tracker = new RateLimitTracker();
+    const info = tracker.parseAndMarkFromError({
+      accountId: 'acc-1',
+      status: 429,
+      retryAfter: '30',
+      body: JSON.stringify({
+        error: {
+          details: [{ reason: 'RATE_LIMIT_EXCEEDED' }],
+          retry_after: 1,
+        },
+      }),
+      model: 'gemini-2.5-pro',
+      backoffSteps: [60, 300, 1800, 7200],
+    });
+
+    expect(info).not.toBeNull();
+    expect(info?.reason).toBe(RateLimitReason.RateLimitExceeded);
+    expect(info?.retryAfterSec).toBe(30);
+    expect(tracker.isRateLimited('acc-1')).toBe(true);
+  });
+
+  it('uses model-level key for quota exhausted', () => {
+    const tracker = new RateLimitTracker();
+    tracker.parseAndMarkFromError({
+      accountId: 'acc-2',
+      status: 429,
+      body: JSON.stringify({
+        error: {
+          details: [
+            {
+              reason: 'QUOTA_EXHAUSTED',
+              metadata: { quotaResetDelay: '42s' },
+            },
+          ],
+        },
+      }),
+      model: 'gemini-2.5-flash',
+      backoffSteps: [60, 300, 1800, 7200],
+    });
+
+    expect(tracker.isRateLimited('acc-2', 'gemini-2.5-flash')).toBe(true);
+    expect(tracker.isRateLimited('acc-2', 'gemini-2.5-pro')).toBe(false);
+  });
+
+  it('uses backoff steps when no header/body retry hint', () => {
+    const tracker = new RateLimitTracker();
+    const steps = [60, 300, 1800, 7200];
+
+    const first = tracker.parseAndMarkFromError({
+      accountId: 'acc-3',
+      status: 429,
+      body: JSON.stringify({
+        error: { details: [{ reason: 'QUOTA_EXHAUSTED' }] },
+      }),
+      backoffSteps: steps,
+    });
+    const second = tracker.parseAndMarkFromError({
+      accountId: 'acc-3',
+      status: 429,
+      body: JSON.stringify({
+        error: { details: [{ reason: 'QUOTA_EXHAUSTED' }] },
+      }),
+      backoffSteps: steps,
+    });
+
+    expect(first?.retryAfterSec).toBe(60);
+    expect(second?.retryAfterSec).toBe(300);
+  });
+
+  it('parses MODEL_CAPACITY_EXHAUSTED and RetryInfo.retryDelay from 503 payload', () => {
+    const tracker = new RateLimitTracker();
+    const info = tracker.parseAndMarkFromError({
+      accountId: 'acc-4',
+      status: 503,
+      model: 'gemini-3.1-pro-high',
+      body: JSON.stringify({
+        error: {
+          code: 503,
+          message: 'No capacity available for model gemini-3.1-pro-high on the server',
+          status: 'UNAVAILABLE',
+          details: [
+            {
+              '@type': 'type.googleapis.com/google.rpc.ErrorInfo',
+              reason: 'MODEL_CAPACITY_EXHAUSTED',
+              domain: 'cloudcode-pa.googleapis.com',
+              metadata: { model: 'gemini-3.1-pro-high' },
+            },
+            {
+              '@type': 'type.googleapis.com/google.rpc.RetryInfo',
+              retryDelay: '30s',
+            },
+          ],
+        },
+      }),
+      backoffSteps: [60, 300, 1800, 7200],
+    });
+
+    expect(info).not.toBeNull();
+    expect(info?.reason).toBe(RateLimitReason.ModelCapacityExhausted);
+    expect(info?.retryAfterSec).toBe(30);
+    expect(tracker.isRateLimited('acc-4')).toBe(true);
+  });
+});