feat(proxy): improve parity, rate limiting, and token scheduling

src/tests/unit/gemini-controller.integration.test.ts

@@ -31,7 +31,14 @@ describe('GeminiController Integration', () => {
       expect.objectContaining({
         models: expect.arrayContaining([
           expect.objectContaining({
-            name: 'models/gemini-2.5-flash',
+            name: 'models/gemini-3-flash',
+            description: '',
+            inputTokenLimit: 128000,
+            outputTokenLimit: 8192,
+            supportedGenerationMethods: ['generateContent', 'countTokens'],
+          }),
+          expect.objectContaining({
+            name: 'models/gemini-3.1-pro-high',
             description: '',
             inputTokenLimit: 128000,
             outputTokenLimit: 8192,
@@ -74,13 +81,13 @@ describe('GeminiController Integration', () => {
     const reply = createReplyMock();
 
     await controller.modelAction(
-      'models/gemini-2.5-flash:generateContent',
+      'models/gemini-3.1-pro-high:generateContent',
       { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] } as any,
       reply as any,
     );
 
     expect(proxyService.handleGeminiGenerateContent).toHaveBeenCalledWith(
-      'models/gemini-2.5-flash',
+      'models/gemini-3.1-pro-high',
       expect.any(Object),
     );
     expect(reply.status).toHaveBeenCalledWith(200);

src/tests/unit/parity-rollout-guard.test.ts

@@ -0,0 +1,106 @@
+import { beforeEach, describe, expect, it } 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';
+
+function createProxyConfig(overrides: Partial<ProxyConfig>): ProxyConfig {
+  return {
+    ...DEFAULT_APP_CONFIG.proxy,
+    ...overrides,
+    upstream_proxy: {
+      ...DEFAULT_APP_CONFIG.proxy.upstream_proxy,
+      ...(overrides.upstream_proxy ?? {}),
+    },
+  };
+}
+
+function seedTokens(service: TokenManagerService): void {
+  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',
+      },
+    ],
+    [
+      'acc-2',
+      {
+        account_id: 'acc-2',
+        email: '[email protected]',
+        access_token: 'token-2',
+        refresh_token: 'refresh-2',
+        token_type: 'Bearer',
+        expires_in: 3600,
+        expiry_timestamp: nowSec + 3600,
+        project_id: 'project-2',
+      },
+    ],
+  ]);
+}
+
+describe('Parity rollout guard replay', () => {
+  let service: TokenManagerService;
+
+  beforeEach(() => {
+    service = new TokenManagerService();
+    seedTokens(service);
+  });
+
+  it('supports kill-switch rollback from parity path to baseline path', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    const paritySelected = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(paritySelected?.id).toBe('acc-2');
+
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: true,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    (service as any).currentIndex = 0;
+    const rollbackSelected = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(rollbackSelected?.id).toBe('acc-1');
+  });
+
+  it('enforces no-go threshold after shadow mismatch and blocks parity enablement', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: false,
+        parity_shadow_enabled: true,
+        parity_no_go_mismatch_rate: 0,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+
+    const shadowSelection = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(shadowSelection?.id).toBe('acc-1');
+
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_shadow_enabled: false,
+        parity_kill_switch: false,
+        preferred_account_id: 'acc-2',
+      }),
+    );
+    (service as any).currentIndex = 0;
+    const blockedSelection = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(blockedSelection?.id).toBe('acc-1');
+  });
+});
+

src/tests/unit/proxy-controller.integration.test.ts

@@ -200,6 +200,73 @@ describe('ProxyController Integration', () => {
     );
   });
 
+  it('supports OpenAI responses compatibility endpoint in stream mode with SSE headers', async () => {
+    const stream = of('data: {"id":"chatcmpl_resp_stream"}\n\n');
+    const proxyService = {
+      handleChatCompletions: vi.fn().mockResolvedValue(stream),
+    };
+
+    const controller = new ProxyController(proxyService as any);
+    const reply = createReplyMock();
+
+    await controller.responses(
+      {
+        model: 'gpt-4o',
+        instructions: 'stream output',
+        input: 'hello',
+        stream: true,
+      },
+      reply as any,
+    );
+
+    expect(reply.header).toHaveBeenCalledWith('Content-Type', 'text/event-stream');
+    expect(reply.header).toHaveBeenCalledWith('Cache-Control', 'no-cache');
+    expect(reply.header).toHaveBeenCalledWith('Connection', 'keep-alive');
+    expect(reply.send).toHaveBeenCalledWith(stream);
+  });
+
+  it('normalizes web_search_call in /v1/responses into builtin_web_search tool messages', async () => {
+    const proxyService = {
+      handleChatCompletions: vi.fn().mockResolvedValue({
+        id: 'chatcmpl_resp_search',
+        object: 'chat.completion',
+        created: 1700000002,
+        model: 'gpt-4o',
+        choices: [{ index: 0, finish_reason: 'stop', message: { content: 'done' } }],
+        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
+      }),
+    };
+
+    const controller = new ProxyController(proxyService as any);
+    const reply = createReplyMock();
+
+    await controller.responses(
+      {
+        model: 'gpt-4o',
+        input: [
+          {
+            type: 'web_search_call',
+            call_id: 'call_search_1',
+            action: { query: 'gemini api' },
+          },
+          {
+            type: 'function_call_output',
+            call_id: 'call_search_1',
+            output: { content: 'search result' },
+          },
+        ],
+      },
+      reply as any,
+    );
+
+    const callArg = proxyService.handleChatCompletions.mock.calls[0][0];
+    const assistantMessage = callArg.messages.find((message: { role: string }) => message.role === 'assistant');
+    const toolMessage = callArg.messages.find((message: { role: string }) => message.role === 'tool');
+
+    expect(assistantMessage?.tool_calls?.[0]?.function?.name).toBe('builtin_web_search');
+    expect(toolMessage?.name).toBe('builtin_web_search');
+  });
+
   it('supports image generations endpoint', async () => {
     const proxyService = {
       handleChatCompletions: vi.fn().mockResolvedValue({

src/tests/unit/streaming-state.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, beforeEach } from 'vitest';
-import { StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper';
+import { PartProcessor, StreamingState } from '../../lib/antigravity/ClaudeStreamingMapper';
 
 describe('StreamingState', () => {
   let state: StreamingState;
@@ -58,4 +58,48 @@ describe('StreamingState', () => {
       expect(state.getErrorCount()).toBe(2);
     });
   });
+
+  describe('stream aggregation compatibility', () => {
+    it('emits tool_use stop reason when functionCall appears in stream', () => {
+      const processor = new PartProcessor(state);
+      const functionChunks = processor.process({
+        functionCall: {
+          name: 'builtin_web_search',
+          args: { query: 'gemini docs' },
+          id: 'call_stream_1',
+        },
+      });
+      const finishChunks = state.emitFinish('STOP', {
+        promptTokenCount: 2,
+        candidatesTokenCount: 3,
+      } as any);
+
+      const output = [...functionChunks, ...finishChunks].join('');
+      expect(output).toContain('"type":"tool_use"');
+      expect(output).toContain('"stop_reason":"tool_use"');
+      expect(output).toContain('"message_stop"');
+    });
+
+    it('aggregates grounding metadata into final text block', () => {
+      state.webSearchQuery = 'gemini api';
+      state.groundingChunks = [
+        {
+          web: {
+            title: 'Gemini API Docs',
+            uri: 'https://example.com/gemini',
+          },
+        },
+      ];
+
+      const chunks = state.emitFinish('STOP', {
+        promptTokenCount: 1,
+        candidatesTokenCount: 1,
+      } as any);
+      const output = chunks.join('');
+
+      expect(output).toContain('Searched for you');
+      expect(output).toContain('Citations');
+      expect(output).toContain('https://example.com/gemini');
+    });
+  });
 });

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

@@ -0,0 +1,166 @@
+import { beforeEach, describe, expect, it } 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';
+
+function createProxyConfig(overrides: Partial<ProxyConfig>): ProxyConfig {
+  return {
+    ...DEFAULT_APP_CONFIG.proxy,
+    ...overrides,
+    upstream_proxy: {
+      ...DEFAULT_APP_CONFIG.proxy.upstream_proxy,
+      ...(overrides.upstream_proxy ?? {}),
+    },
+  };
+}
+
+function seedTokens(service: TokenManagerService): void {
+  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',
+      },
+    ],
+    [
+      'acc-2',
+      {
+        account_id: 'acc-2',
+        email: '[email protected]',
+        access_token: 'token-2',
+        refresh_token: 'refresh-2',
+        token_type: 'Bearer',
+        expires_in: 3600,
+        expiry_timestamp: nowSec + 3600,
+        project_id: 'project-2',
+        session_id: 'session-2',
+      },
+    ],
+  ]);
+}
+
+describe('TokenManagerService parity scheduling replay', () => {
+  let service: TokenManagerService;
+
+  beforeEach(() => {
+    service = new TokenManagerService();
+    seedTokens(service);
+  });
+
+  it('prioritizes preferred account in parity mode', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'balance',
+        preferred_account_id: 'acc-2',
+      }),
+    );
+
+    const token = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(token?.id).toBe('acc-2');
+  });
+
+  it('rotates sticky account when limited in balance mode', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'balance',
+        preferred_account_id: '',
+      }),
+    );
+
+    const first = await service.getNextToken({
+      sessionKey: 'openai:user-1',
+      model: 'gemini-2.5-flash',
+    });
+    expect(first?.id).toBe('acc-1');
+
+    await service.markFromUpstreamError({
+      accountIdOrEmail: 'acc-1',
+      status: 429,
+      model: 'gemini-2.5-flash',
+      body: JSON.stringify({
+        error: {
+          details: [{ reason: 'RATE_LIMIT_EXCEEDED' }],
+        },
+      }),
+    });
+
+    const second = await service.getNextToken({
+      sessionKey: 'openai:user-1',
+      model: 'gemini-2.5-flash',
+    });
+    expect(second?.id).toBe('acc-2');
+  });
+
+  it('applies model-level lock for quota exhausted only on the same model', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: true,
+        parity_kill_switch: false,
+        scheduling_mode: 'performance-first',
+      }),
+    );
+
+    await service.markFromUpstreamError({
+      accountIdOrEmail: 'acc-1',
+      status: 429,
+      model: 'gemini-2.5-flash',
+      body: JSON.stringify({
+        error: {
+          details: [{ reason: 'QUOTA_EXHAUSTED', metadata: { quotaResetDelay: '30s' } }],
+        },
+      }),
+    });
+
+    const sameModel = await service.getNextToken({ model: 'gemini-2.5-flash' });
+    expect(sameModel?.id).toBe('acc-2');
+
+    const otherModel = await service.getNextToken({
+      model: 'gemini-2.5-pro',
+      excludeAccountIds: ['acc-2'],
+    });
+    expect(otherModel?.id).toBe('acc-1');
+  });
+
+  it('falls back to excluded pool when retry exclusions would empty all candidates', async () => {
+    setServerConfig(
+      createProxyConfig({
+        parity_enabled: false,
+        parity_kill_switch: false,
+      }),
+    );
+
+    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',
+        },
+      ],
+    ]);
+
+    const token = await service.getNextToken({ excludeAccountIds: ['acc-1'] });
+    expect(token?.id).toBe('acc-1');
+  });
+});