ccs/tests/unit/proxy/messages-route.test.ts

import { afterEach, beforeEach, describe, expect, it } from 'bun:test';
import { EventEmitter } from 'events';
import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';
import { PassThrough } from 'stream';
import { Agent } from 'undici';
import { resolveOpenAICompatProfileConfig } from '../../../src/proxy/profile-router';
import {
  attachDisconnectAbortHandlers,
  buildUpstreamAgentTimeouts,
  handleProxyMessagesRequest as handleProxyMessagesRequestImpl,
} from '../../../src/proxy/server/messages-route';
import {
  closeUpstreamDispatcher,
  isBunRuntime,
  toLogErrorInfo,
} from '../../../src/proxy/server/upstream-transport';
import { loadSettings } from '../../../src/utils/config-manager';

class FakeSocket extends EventEmitter {
  destroyed = false;
}

class FakeRequest extends PassThrough {
  headers: Record<string, string>;
  method = 'POST';
  url = '/v1/messages';
  socket = new FakeSocket();

  constructor(headers: Record<string, string> = {}) {
    super();
    this.headers = headers;
  }
}

class FakeResponse extends PassThrough {
  headers = new Map<string, string>();
  statusCode = 200;
  socket = new FakeSocket();

  setHeader(name: string, value: string): this {
    this.headers.set(name.toLowerCase(), value);
    return this;
  }

  writeHead(statusCode: number, headers: Record<string, string>): this {
    this.statusCode = statusCode;
    for (const [name, value] of Object.entries(headers)) {
      this.setHeader(name, value);
    }
    return this;
  }
}

let originalCcsHome: string | undefined;
let originalFetch: typeof globalThis.fetch;
let originalAgentClose: typeof Agent.prototype.close;
let tempDir: string;

function writeSettings(profileName: string, env: Record<string, string>): string {
  const settingsPath = path.join(tempDir, '.ccs', `${profileName}.settings.json`);
  fs.writeFileSync(settingsPath, JSON.stringify({ env }, null, 2), 'utf8');
  return settingsPath;
}

function buildProfile(profileName: string) {
  const settingsPath = path.join(tempDir, '.ccs', `${profileName}.settings.json`);
  const profile = resolveOpenAICompatProfileConfig(
    profileName,
    settingsPath,
    loadSettings(settingsPath).env || {}
  );
  expect(profile).toBeTruthy();
  return profile!;
}

function handleProxyMessagesRequest(
  req: Parameters<typeof handleProxyMessagesRequestImpl>[0],
  res: Parameters<typeof handleProxyMessagesRequestImpl>[1],
  profile: Parameters<typeof handleProxyMessagesRequestImpl>[2],
  expectedAuthToken: string,
  insecureDispatcher?: Parameters<typeof handleProxyMessagesRequestImpl>[4]
) {
  return handleProxyMessagesRequestImpl(
    req,
    res,
    profile,
    expectedAuthToken,
    insecureDispatcher,
    undefined,
    globalThis.fetch
  );
}

beforeEach(() => {
  originalCcsHome = process.env.CCS_HOME;
  originalFetch = globalThis.fetch;
  originalAgentClose = Agent.prototype.close;
  tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-messages-route-'));
  fs.mkdirSync(path.join(tempDir, '.ccs'), { recursive: true });
  process.env.CCS_HOME = tempDir;

  const profiles = {
    hf: writeSettings('hf', {
      ANTHROPIC_BASE_URL: 'https://router.huggingface.co/v1',
      ANTHROPIC_AUTH_TOKEN: 'hf_token',
      ANTHROPIC_MODEL: 'hf-default',
      CCS_DROID_PROVIDER: 'generic-chat-completion-api',
    }),
    deepseek: writeSettings('deepseek', {
      ANTHROPIC_BASE_URL: 'https://api.deepseek.com/v1',
      ANTHROPIC_AUTH_TOKEN: 'deepseek_token',
      ANTHROPIC_MODEL: 'deepseek-chat',
      CCS_DROID_PROVIDER: 'generic-chat-completion-api',
      CCS_OPENAI_PROXY_INSECURE: 'true',
    }),
    search: writeSettings('search', {
      ANTHROPIC_BASE_URL: 'https://search.example.com/v1',
      ANTHROPIC_AUTH_TOKEN: 'search_token',
      ANTHROPIC_MODEL: 'sonar-pro',
      CCS_DROID_PROVIDER: 'generic-chat-completion-api',
    }),
    mm: writeSettings('mm', {
      ANTHROPIC_BASE_URL: 'https://api.minimax.io/v1',
      ANTHROPIC_AUTH_TOKEN: 'minimax_token',
      ANTHROPIC_MODEL: 'MiniMax-M2.7',
      CCS_DROID_PROVIDER: 'openai',
    }),
    kimic: writeSettings('kimic', {
      ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/',
      ANTHROPIC_AUTH_TOKEN: 'kimi_token',
      ANTHROPIC_MODEL: 'kimi-k2p7-coding',
      CCS_DROID_PROVIDER: 'generic-chat-completion-api',
    }),
  };

  fs.writeFileSync(
    path.join(tempDir, '.ccs', 'config.json'),
    JSON.stringify({ profiles, proxy: { routing: {} } }, null, 2),
    'utf8'
  );
});

afterEach(() => {
  globalThis.fetch = originalFetch;
  Agent.prototype.close = originalAgentClose;
  if (originalCcsHome !== undefined) {
    process.env.CCS_HOME = originalCcsHome;
  } else {
    delete process.env.CCS_HOME;
  }
  fs.rmSync(tempDir, { recursive: true, force: true });
});

describe('attachDisconnectAbortHandlers', () => {
  it('cleans up registered listeners after the request completes', () => {
    const req = new FakeRequest();
    const res = new FakeResponse();
    const controller = new AbortController();

    const cleanup = attachDisconnectAbortHandlers(req as never, res as never, controller, () => {});

    expect(req.listenerCount('aborted')).toBe(1);
    expect(req.listenerCount('close')).toBe(0);
    expect(req.socket.listenerCount('close')).toBe(1);
    expect(res.listenerCount('close')).toBe(0);
    expect(res.socket.listenerCount('close')).toBe(1);

    cleanup();

    expect(req.listenerCount('aborted')).toBe(0);
    expect(req.socket.listenerCount('close')).toBe(0);
    expect(res.socket.listenerCount('close')).toBe(0);
  });

  it('aborts at most once when disconnect signals race each other', () => {
    const req = new FakeRequest();
    const res = new FakeResponse();
    const controller = new AbortController();
    let disconnectCount = 0;

    const cleanup = attachDisconnectAbortHandlers(req as never, res as never, controller, () => {
      disconnectCount += 1;
    });

    req.emit('aborted');
    req.socket.emit('close');
    res.socket.emit('close');

    expect(controller.signal.aborted).toBe(true);
    expect(disconnectCount).toBe(1);

    cleanup();
  });
});

describe('buildUpstreamAgentTimeouts', () => {
  afterEach(() => {
    delete process.env.CCS_OPENAI_PROXY_REQUEST_TIMEOUT_MS;
  });

  it('keeps undici per-phase timeouts above the default request timeout', () => {
    const timeouts = buildUpstreamAgentTimeouts();
    expect(timeouts.headersTimeout).toBeGreaterThan(600_000);
    expect(timeouts.bodyTimeout).toBeGreaterThan(600_000);
  });

  it('tracks CCS_OPENAI_PROXY_REQUEST_TIMEOUT_MS overrides', () => {
    process.env.CCS_OPENAI_PROXY_REQUEST_TIMEOUT_MS = '120000';
    const timeouts = buildUpstreamAgentTimeouts();
    expect(timeouts.headersTimeout).toBeGreaterThan(120_000);
    expect(timeouts.bodyTimeout).toBeGreaterThan(120_000);
  });
});

describe('closeUpstreamDispatcher', () => {
  it('closes an owned dispatcher once across repeated cleanup', async () => {
    let closeCalls = 0;
    const dispatcher = {
      close: async () => {
        closeCalls += 1;
      },
    } as never;

    await Promise.all([
      closeUpstreamDispatcher(dispatcher),
      closeUpstreamDispatcher(dispatcher),
      closeUpstreamDispatcher(dispatcher),
    ]);

    expect(closeCalls).toBe(1);
  });
});

describe('toLogErrorInfo', () => {
  it('preserves nested transport cause metadata for diagnostics', () => {
    const cause = Object.assign(new Error('invalid onError method'), {
      name: 'InvalidArgumentError',
      code: 'UND_ERR_INVALID_ARG',
    });
    const error = new TypeError('fetch failed', { cause });

    expect(toLogErrorInfo(error)).toEqual({
      name: 'TypeError',
      message: 'fetch failed',
      cause: {
        name: 'InvalidArgumentError',
        message: 'invalid onError method',
        code: 'UND_ERR_INVALID_ARG',
      },
    });
  });
});

describe('handleProxyMessagesRequest', () => {
  it('dispatches secure upstream fetches with an explicit dispatcher (not undici defaults)', async () => {
    const activeProfile = buildProfile('hf');
    let capturedDispatcher: unknown;

    globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
      capturedDispatcher = (init as RequestInit & { dispatcher?: unknown })?.dispatcher;
      throw new Error('upstream exploded');
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(req as never, res as never, activeProfile, 'local-token');
    req.end(
      JSON.stringify({
        model: 'hf-default',
        stream: true,
        messages: [{ role: 'user', content: 'stay secure' }],
      })
    );
    await pending;

    // A bare global-dispatcher fallback would reapply undici's 300s
    // headersTimeout/bodyTimeout and undercut the proxy's request timeout.
    if (isBunRuntime()) {
      expect(capturedDispatcher).toBeUndefined();
    } else {
      expect(capturedDispatcher).toBeDefined();
    }
    expect(res.statusCode).toBe(502);
  });

  it('auto-passes through Kimi requests and preserves the coding-agent User-Agent', async () => {
    const activeProfile = buildProfile('kimic');
    let capturedInput: RequestInfo | URL | undefined;
    let capturedInit: RequestInit | undefined;
    const rawBody = JSON.stringify({
      model: 'kimi-k2p7-coding',
      max_tokens: 50,
      messages: [{ role: 'user', content: 'Say hi in 2 words' }],
    });

    globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
      capturedInput = input;
      capturedInit = init;
      return new Response(
        JSON.stringify({
          id: 'msg_1',
          type: 'message',
          role: 'assistant',
          content: [{ type: 'text', text: 'Hello there!' }],
          model: 'kimi-k2p7-coding',
          stop_reason: 'end_turn',
        }),
        { status: 200, headers: { 'content-type': 'application/json' } }
      );
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
      'user-agent': 'claude-cli/2.1.170',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token'
    );
    req.end(rawBody);
    await pending;

    expect(String(capturedInput)).toBe('https://api.kimi.com/coding/v1/messages');
    expect(capturedInit?.body).toBe(rawBody);
    expect((capturedInit?.headers as Record<string, string>)['User-Agent']).toBe(
      'claude-cli/2.1.170'
    );
    expect(res.statusCode).toBe(200);
  });

  it('uses routed-profile passthrough for explicit profile:model selectors', async () => {
    const activeProfile = buildProfile('hf');
    let capturedInput: RequestInfo | URL | undefined;
    let capturedInit: RequestInit | undefined;
    const rawBody = JSON.stringify({
      model: 'kimic:kimi-k2p7-coding',
      max_tokens: 50,
      messages: [{ role: 'user', content: 'Say hi' }],
    });

    globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
      capturedInput = input;
      capturedInit = init;
      return new Response(
        JSON.stringify({
          id: 'msg_2',
          type: 'message',
          role: 'assistant',
          content: [{ type: 'text', text: 'Hi' }],
          model: 'kimi-k2p7-coding',
          stop_reason: 'end_turn',
        }),
        { status: 200, headers: { 'content-type': 'application/json' } }
      );
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
      'x-stainless-user-agent': 'claude-cli/2.1.170',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token'
    );
    req.end(rawBody);
    await pending;

    expect(String(capturedInput)).toBe('https://api.kimi.com/coding/v1/messages');
    expect(capturedInit?.body).toBe(rawBody);
    expect((capturedInit?.headers as Record<string, string>)['User-Agent']).toBe(
      'claude-cli/2.1.170'
    );
    expect(res.statusCode).toBe(200);
  });

  it('keeps default OpenAI-compatible requests on chat completions with CCS User-Agent', async () => {
    const activeProfile = buildProfile('hf');
    let capturedInput: RequestInfo | URL | undefined;
    let capturedInit: RequestInit | undefined;

    globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
      capturedInput = input;
      capturedInit = init;
      return new Response(
        JSON.stringify({
          id: 'chatcmpl_1',
          object: 'chat.completion',
          created: 1,
          model: 'hf-default',
          choices: [{ index: 0, message: { role: 'assistant', content: 'ok' } }],
        }),
        { status: 200, headers: { 'content-type': 'application/json' } }
      );
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
      'user-agent': 'claude-cli/2.1.170',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token'
    );
    req.end(
      JSON.stringify({
        model: 'hf-default',
        messages: [{ role: 'user', content: 'stay translated' }],
      })
    );
    await pending;

    expect(String(capturedInput)).toBe('https://router.huggingface.co/v1/chat/completions');
    expect(JSON.parse(String(capturedInit?.body))).toMatchObject({
      model: 'hf-default',
      messages: [{ role: 'user', content: 'stay translated' }],
    });
    expect((capturedInit?.headers as Record<string, string>)['User-Agent']).toBe(
      'CCS-OpenAI-Compat-Proxy/1.0'
    );
    expect(res.statusCode).toBe(200);
  });

  it('uses a per-request insecure dispatcher for routed profiles and closes it on failure', async () => {
    const activeProfile = buildProfile('hf');
    const sharedDispatcher = { name: 'shared-insecure-dispatcher' } as never;
    let closeCalls = 0;
    let capturedDispatcher: unknown;

    Agent.prototype.close = function close(): Promise<void> {
      closeCalls += 1;
      return Promise.resolve();
    };

    globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
      capturedDispatcher = (init as RequestInit & { dispatcher?: unknown })?.dispatcher;
      throw new Error('upstream exploded');
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token',
      sharedDispatcher
    );
    req.end(
      JSON.stringify({
        model: 'deepseek:deepseek-chat',
        stream: true,
        messages: [{ role: 'user', content: 'route this' }],
      })
    );
    await pending;

    if (isBunRuntime()) {
      expect(capturedDispatcher).toBeUndefined();
    } else {
      expect(capturedDispatcher).toBeInstanceOf(Agent);
    }
    expect(capturedDispatcher).not.toBe(sharedDispatcher);
    expect(closeCalls).toBe(isBunRuntime() ? 0 : 1);
    expect(res.statusCode).toBe(502);
  });

  it('reuses the shared insecure dispatcher when the request stays on the active profile', async () => {
    const activeProfile = buildProfile('hf');
    const sharedDispatcher = { name: 'shared-insecure-dispatcher' } as never;
    let closeCalls = 0;
    let capturedDispatcher: unknown;

    Agent.prototype.close = function close(): Promise<void> {
      closeCalls += 1;
      return Promise.resolve();
    };

    globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
      capturedDispatcher = (init as RequestInit & { dispatcher?: unknown })?.dispatcher;
      throw new Error('upstream exploded');
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token',
      sharedDispatcher
    );
    req.end(
      JSON.stringify({
        model: 'hf-default',
        stream: true,
        messages: [{ role: 'user', content: 'stay local' }],
      })
    );
    await pending;

    expect(capturedDispatcher).toBe(sharedDispatcher);
    expect(closeCalls).toBe(0);
    expect(res.statusCode).toBe(502);
  });

  it('moves system messages into the first user message for MiniMax OpenAI-compatible upstreams', async () => {
    const activeProfile = buildProfile('mm');
    let capturedBody: unknown;

    globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
      capturedBody = JSON.parse(String(init?.body));
      return new Response(
        JSON.stringify({
          id: 'chatcmpl_1',
          object: 'chat.completion',
          created: 1,
          model: 'MiniMax-M2.7',
          choices: [{ index: 0, message: { role: 'assistant', content: 'ok' } }],
        }),
        { status: 200, headers: { 'content-type': 'application/json' } }
      );
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token'
    );
    req.end(
      JSON.stringify({
        model: 'MiniMax-M2.7',
        stream: false,
        messages: [
          { role: 'system', content: 'Use Turkish.' },
          { role: 'user', content: 'bu hangi model' },
        ],
      })
    );
    await pending;

    expect(capturedBody).toMatchObject({
      model: 'MiniMax-M2.7',
      messages: [{ role: 'user', content: 'Use Turkish.\n\nbu hangi model' }],
    });
    expect((capturedBody as { messages: Array<{ role: string }> }).messages).not.toContainEqual(
      expect.objectContaining({ role: 'system' })
    );
  });

  it('strips blank system messages for MiniMax OpenAI-compatible upstreams', async () => {
    const activeProfile = buildProfile('mm');
    let capturedBody: unknown;

    globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => {
      capturedBody = JSON.parse(String(init?.body));
      return new Response(
        JSON.stringify({
          id: 'chatcmpl_1',
          object: 'chat.completion',
          created: 1,
          model: 'MiniMax-M2.7',
          choices: [{ index: 0, message: { role: 'assistant', content: 'ok' } }],
        }),
        { status: 200, headers: { 'content-type': 'application/json' } }
      );
    }) as typeof globalThis.fetch;

    const req = new FakeRequest({
      'x-api-key': 'local-token',
    });
    const res = new FakeResponse();
    const pending = handleProxyMessagesRequest(
      req as never,
      res as never,
      activeProfile,
      'local-token'
    );
    req.end(
      JSON.stringify({
        model: 'MiniMax-M2.7',
        stream: false,
        messages: [
          { role: 'system', content: [{ type: 'text', text: '' }] },
          { role: 'user', content: 'bu hangi model' },
        ],
      })
    );
    await pending;

    expect(capturedBody).toMatchObject({
      model: 'MiniMax-M2.7',
      messages: [{ role: 'user', content: 'bu hangi model' }],
    });
    expect((capturedBody as { messages: Array<{ role: string }> }).messages).not.toContainEqual(
      expect.objectContaining({ role: 'system' })
    );
  });
});