feat(proxy): improve quota fallback and proxy parity handling

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

@@ -97,4 +97,5 @@ describe('Proxy Parity Fixtures', () => {
       expect(output).toContain(token);
     }
   });
+
 });

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

@@ -12,6 +12,7 @@ const mockTokenManager = {
   recordParityError: vi.fn(),
   getModelOutputLimitForAccount: vi.fn(),
   getModelThinkingBudgetForAccount: vi.fn(),
+  resolveDynamicModelForAccount: vi.fn((_token: unknown, model: string) => model),
 };
 const mockGeminiClient = { streamGenerateInternal: vi.fn(), generateInternal: vi.fn() };
 
@@ -651,4 +652,43 @@ describe('ProxyService Protocol Parity Fixtures', () => {
     expect(output).toContain('"content":"final answer"');
     expect(output).toContain('data: [DONE]');
   });
+
+  it('propagates OpenAI-compatible upstream stream errors instead of completing with [DONE]', async () => {
+    const service = new TestableProxyService();
+    const stream = new EventEmitter();
+    const observable = (service as any).processStreamResponse(stream, 'gpt-4o-mini');
+    const chunks: string[] = [];
+
+    const streamResult = await new Promise<{ error?: Error; completed: boolean }>((resolve) => {
+      observable.subscribe({
+        next: (chunk: string) => {
+          chunks.push(chunk);
+        },
+        error: (error: unknown) =>
+          resolve({
+            completed: false,
+            error: error instanceof Error ? error : new Error(String(error)),
+          }),
+        complete: () => resolve({ completed: true }),
+      });
+
+      const payload = JSON.stringify({
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'partial output' }],
+            },
+          },
+        ],
+      });
+
+      stream.emit('data', Buffer.from(`data: ${payload}\n`));
+      stream.emit('error', new Error('socket hang up'));
+    });
+
+    expect(streamResult.completed).not.toBe(true);
+    expect(streamResult.error?.message).toContain('socket hang up');
+    expect(chunks.join('')).toContain('"content":"partial output"');
+    expect(chunks.join('')).not.toContain('data: [DONE]');
+  });
 });

src/tests/unit/token-manager-parity-scheduling.test.ts

@@ -1,7 +1,8 @@
-import { beforeEach, describe, expect, it } from 'vitest';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { DEFAULT_APP_CONFIG, ProxyConfig } from '../../types/config';
 import { setServerConfig } from '../../server/server-config';
 import { TokenManagerService } from '../../server/modules/proxy/token-manager.service';
+import { GoogleAPIService } from '../../services/GoogleAPIService';
 
 function createProxyConfig(overrides: Partial<ProxyConfig>): ProxyConfig {
   return {
@@ -56,6 +57,96 @@ describe('TokenManagerService parity scheduling replay', () => {
     seedTokens(service);
   });
 
+  it('rewrites gemini pro model to first available account candidate', () => {
+    const nowSec = Math.floor(Date.now() / 1000);
+    (service as any).tokens = new Map([
+      [
+        'acc-1',
+        {
+          account_id: 'acc-1',
+          email: '[email protected]',
+          access_token: 'token-1',
+          refresh_token: 'refresh-1',
+          token_type: 'Bearer',
+          expires_in: 3600,
+          expiry_timestamp: nowSec + 3600,
+          project_id: 'project-1',
+          session_id: 'session-1',
+          model_quotas: {
+            'gemini-3.1-pro-low': 80,
+          },
+          model_limits: {},
+          model_reset_times: {},
+          model_forwarding_rules: {},
+        },
+      ],
+    ]);
+
+    const resolved = service.resolveDynamicModelForAccount('acc-1', 'gemini-3-pro');
+    expect(resolved).toBe('gemini-3.1-pro-low');
+  });
+
+  it('keeps original model when dynamic rewrite is not applicable', () => {
+    const resolved = service.resolveDynamicModelForAccount('acc-1', 'gemini-3-flash');
+    expect(resolved).toBe('gemini-3-flash');
+  });
+
+  it('passes oauth_client_key when refreshing token and persists refreshed key', async () => {
+    const nowSec = Math.floor(Date.now() / 1000);
+    const tokenData = {
+      account_id: 'acc-1',
+      email: '[email protected]',
+      access_token: 'token-1',
+      refresh_token: 'refresh-1',
+      oauth_client_key: 'custom-client',
+      upstream_proxy_url: 'http://127.0.0.1:8080',
+      token_type: 'Bearer',
+      expires_in: 3600,
+      expiry_timestamp: nowSec - 1,
+      project_id: 'project-1',
+      session_id: 'session-1',
+      model_quotas: {},
+      model_limits: {},
+      model_reset_times: {},
+      model_forwarding_rules: {},
+    };
+
+    (service as any).tokens = new Map([['acc-1', tokenData]]);
+
+    const refreshSpy = vi.spyOn(GoogleAPIService, 'refreshAccessToken').mockResolvedValue({
+      access_token: 'token-new',
+      expires_in: 7200,
+      token_type: 'Bearer',
+      oauth_client_key: 'custom-fallback',
+    });
+    const persistSpy = vi.spyOn(service as any, 'persistTokenState').mockResolvedValue(undefined);
+
+    const selected = await (service as any).finalizeSelectedToken('acc-1', tokenData, nowSec);
+
+    expect(refreshSpy).toHaveBeenCalledWith(
+      'refresh-1',
+      'http://127.0.0.1:8080',
+      'custom-client',
+    );
+    expect(selected?.token.oauth_client_key).toBe('custom-fallback');
+    expect((service as any).tokens.get('acc-1')?.oauth_client_key).toBe('custom-fallback');
+
+    refreshSpy.mockRestore();
+    persistSpy.mockRestore();
+  });
+
+  it('keeps oauth_client_key unset for legacy account refreshed by enterprise client', () => {
+    const normalized = (service as any).normalizeRefreshedOauthClientKey(
+      {
+        oauth_client_key: undefined,
+        project_id: undefined,
+      },
+      'antigravity_enterprise',
+    );
+
+    expect(normalized).toBeUndefined();
+  });
+
   it('prioritizes preferred account in parity mode', async () => {
     setServerConfig(
       createProxyConfig({