proxy.js

const express = require('express');
const { Readable } = require('node:stream');
const { pipeline } = require('node:stream/promises');
const { fetch: defaultFetch } = require('undici');

const DEFAULT_TARGET_URL = 'https://cn2.router.team/v1/chat/completions';
const DEFAULT_TIMEOUT_MS = 90_000;
const DEFAULT_BODY_LIMIT = '10mb';

const REQUEST_HEADERS = new Set([
  'accept',
  'authorization',
  'content-type',
  'openai-organization',
  'openai-project',
]);

// Undici transparently decodes compressed bodies, and transformed SSE bodies have
// a different length. Forwarding either original header corrupts the response.
const OMITTED_RESPONSE_HEADERS = new Set([
  'connection',
  'content-encoding',
  'content-length',
  'keep-alive',
  'proxy-authenticate',
  'proxy-authorization',
  'te',
  'trailer',
  'transfer-encoding',
  'upgrade',
]);

function getPositiveNumber(value, fallback) {
  const parsed = Number(value);
  return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
}

function getRequestHeaders(headers) {
  const result = {};

  for (const [name, value] of Object.entries(headers)) {
    if (REQUEST_HEADERS.has(name.toLowerCase()) && value !== undefined) {
      result[name] = value;
    }
  }

  return result;
}

function copyResponseHeaders(headers, res) {
  for (const [name, value] of headers) {
    if (!OMITTED_RESPONSE_HEADERS.has(name.toLowerCase())) {
      res.setHeader(name, value);
    }
  }
}

function patchCompletionPayload(data) {
  if (!data || typeof data !== 'object' || Array.isArray(data)) return data;

  // MiMo treats the presence of `error` as a failure and reads `error.message`.
  // Remove empty sentinels from successful chunks, but preserve real errors.
  if (data.error == null) delete data.error;

  if (!Array.isArray(data.choices)) return data;

  for (const choice of data.choices) {
    const toolCalls = choice?.delta?.tool_calls;
    if (!Array.isArray(toolCalls)) continue;

    for (const toolCall of toolCalls) {
      if (
        toolCall
        && typeof toolCall === 'object'
        && toolCall.function === undefined
      ) {
        toolCall.function = {};
      }
    }
  }

  return data;
}

function transformSseEvent(rawEvent) {
  const lines = rawEvent.split(/\r\n|\r|\n/);
  const dataLineIndexes = [];
  const dataParts = [];

  for (let index = 0; index < lines.length; index += 1) {
    const line = lines[index];
    const colonIndex = line.indexOf(':');
    const field = colonIndex === -1 ? line : line.slice(0, colonIndex);

    if (field !== 'data') continue;

    let value = colonIndex === -1 ? '' : line.slice(colonIndex + 1);
    if (value.startsWith(' ')) value = value.slice(1);
    dataLineIndexes.push(index);
    dataParts.push(value);
  }

  if (dataLineIndexes.length === 0) return `${lines.join('\n')}\n\n`;

  const payload = dataParts.join('\n').trim();
  if (payload === '[DONE]') return `${lines.join('\n')}\n\n`;

  let replacement;
  try {
    replacement = `data: ${JSON.stringify(patchCompletionPayload(JSON.parse(payload)))}`;
  } catch {
    return `${lines.join('\n')}\n\n`;
  }

  const dataIndexes = new Set(dataLineIndexes);
  const transformedLines = [];
  let dataWritten = false;

  for (let index = 0; index < lines.length; index += 1) {
    if (!dataIndexes.has(index)) {
      transformedLines.push(lines[index]);
    } else if (!dataWritten) {
      transformedLines.push(replacement);
      dataWritten = true;
    }
  }

  return `${transformedLines.join('\n')}\n\n`;
}

function takeNextSseEvent(buffer) {
  function lineEndingLengthAt(index) {
    if (buffer[index] === '\n') return 1;
    if (buffer[index] !== '\r') return 0;
    return buffer[index + 1] === '\n' ? 2 : 1;
  }

  for (let index = 0; index < buffer.length; index += 1) {
    const firstLength = lineEndingLengthAt(index);
    if (firstLength === 0) continue;

    const secondLength = lineEndingLengthAt(index + firstLength);
    if (secondLength === 0) {
      index += firstLength - 1;
      continue;
    }

    return {
      event: buffer.slice(0, index),
      rest: buffer.slice(index + firstLength + secondLength),
    };
  }

  return null;
}

async function writeWithBackpressure(res, value) {
  if (res.destroyed || res.writableEnded) throw new Error('Client disconnected');
  if (res.write(value)) return;

  await new Promise((resolve, reject) => {
    const cleanup = () => {
      res.off('close', onClose);
      res.off('drain', onDrain);
      res.off('error', onError);
    };
    const onClose = () => {
      cleanup();
      reject(new Error('Client disconnected'));
    };
    const onDrain = () => {
      cleanup();
      resolve();
    };
    const onError = error => {
      cleanup();
      reject(error);
    };

    res.once('close', onClose);
    res.once('drain', onDrain);
    res.once('error', onError);
  });
}

async function pipeSseBody(body, res) {
  const reader = body.getReader();
  const decoder = new TextDecoder('utf-8');
  let buffer = '';

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    buffer += decoder.decode(value, { stream: true });

    while (true) {
      const next = takeNextSseEvent(buffer);
      if (!next) break;
      buffer = next.rest;
      await writeWithBackpressure(res, transformSseEvent(next.event));
    }
  }

  buffer += decoder.decode();

  while (true) {
    const next = takeNextSseEvent(buffer);
    if (!next) break;
    buffer = next.rest;
    await writeWithBackpressure(res, transformSseEvent(next.event));
  }

  // Preserve a final upstream event even if the connection closes without a blank line.
  if (buffer.length > 0) {
    await writeWithBackpressure(res, transformSseEvent(buffer));
  }

  res.end();
}

function createApp(options = {}) {
  const app = express();
  const fetchImpl = options.fetch ?? defaultFetch;
  const logger = options.logger ?? console;
  const chatTargetUrl = options.targetUrl ?? process.env.UPSTREAM_URL ?? DEFAULT_TARGET_URL;
  const responsesTargetUrl = options.responsesTargetUrl
    ?? process.env.RESPONSES_UPSTREAM_URL
    ?? new URL('/v1/responses', chatTargetUrl).toString();
  const timeoutMs = getPositiveNumber(
    options.timeoutMs ?? process.env.UPSTREAM_TIMEOUT_MS,
    DEFAULT_TIMEOUT_MS,
  );
  const bodyLimit = options.bodyLimit ?? process.env.REQUEST_BODY_LIMIT ?? DEFAULT_BODY_LIMIT;

  // Validate configuration during startup instead of failing every request later.
  new URL(chatTargetUrl);
  new URL(responsesTargetUrl);

  const parseBody = express.raw({ type: () => true, limit: bodyLimit });

  const createProxyHandler = (upstreamUrl, { transformCompletionSse = false } = {}) => async (req, res) => {
    const controller = new AbortController();
    let timedOut = false;

    const timeout = setTimeout(() => {
      timedOut = true;
      controller.abort(new Error(`Upstream timed out after ${timeoutMs}ms`));
    }, timeoutMs);
    timeout.unref?.();

    const clearUpstreamTimeout = () => clearTimeout(timeout);

    const abortUpstream = () => controller.abort(new Error('Client disconnected'));
    const abortIfIncomplete = () => {
      if (!res.writableEnded) abortUpstream();
    };

    req.once('aborted', abortUpstream);
    res.once('close', abortIfIncomplete);

    try {
      const upstreamRes = await fetchImpl(upstreamUrl, {
        method: 'POST',
        headers: getRequestHeaders(req.headers),
        body: req.body?.length ? req.body : undefined,
        signal: controller.signal,
      });

      // The timeout limits how long the upstream may take to start responding.
      // Streaming responses can legitimately remain open much longer.
      clearUpstreamTimeout();

      res.status(upstreamRes.status);
      copyResponseHeaders(upstreamRes.headers, res);

      if (!upstreamRes.body) {
        res.end();
        return;
      }

      const contentType = upstreamRes.headers.get('content-type')?.toLowerCase() ?? '';
      const isEventStream = contentType.includes('text/event-stream');
      if (isEventStream) res.flushHeaders();

      if (isEventStream && transformCompletionSse) {
        await pipeSseBody(upstreamRes.body, res);
      } else {
        await pipeline(Readable.fromWeb(upstreamRes.body), res);
      }
    } catch (error) {
      const clientDisconnected = controller.signal.aborted && !timedOut;
      if (!clientDisconnected) {
        logger.error('Proxy request failed:', error.message);
      }

      if (!res.headersSent) {
        res.status(timedOut ? 504 : 502).json({
          error: {
            message: timedOut ? 'Upstream request timed out' : 'Upstream request failed',
            type: 'proxy_error',
          },
        });
      } else if (!res.writableEnded) {
        res.destroy();
      }
    } finally {
      clearUpstreamTimeout();
      req.off('aborted', abortUpstream);
      res.off('close', abortIfIncomplete);
    }
  };

  app.post(
    ['/v1/chat/completions', '/chat/completions'],
    parseBody,
    createProxyHandler(chatTargetUrl, { transformCompletionSse: true }),
  );
  app.post(
    ['/v1/responses', '/responses'],
    parseBody,
    createProxyHandler(responsesTargetUrl),
  );

  app.all(
    ['/v1/chat/completions', '/chat/completions', '/v1/responses', '/responses'],
    (_req, res) => {
      res.setHeader('allow', 'POST');
      res.status(405).send('Method Not Allowed: Only POST is supported');
    },
  );

  return app;
}

function startServer(options = {}) {
  const app = createApp(options);
  const port = getPositiveNumber(options.port ?? process.env.PORT, 3000);
  const host = options.host ?? process.env.HOST ?? '0.0.0.0';

  return app.listen(port, host, () => {
    console.log(`Proxy listening on ${host}:${port}`);
  });
}

if (require.main === module) startServer();

module.exports = {
  createApp,
  patchCompletionPayload,
  startServer,
  transformSseEvent,
};