memind/deepseek-no-think-proxy.mjs

import http from 'node:http';
import path from 'node:path';
import { fileURLToPath } from 'node:url';
import { Agent, fetch as undiciFetch } from 'undici';

export const DEFAULT_DEEPSEEK_NO_THINK_PORT = 18036;
export const DEFAULT_DEEPSEEK_UPSTREAM = 'https://api.deepseek.com';
export const DEFAULT_MOONSHOT_UPSTREAM = 'https://api.moonshot.cn';

const insecureDispatcher = new Agent({
  connect: { rejectUnauthorized: false },
});

function envFlag(value, defaultValue = false) {
  if (value == null || String(value).trim() === '') return defaultValue;
  return ['1', 'true', 'yes', 'on'].includes(String(value).trim().toLowerCase());
}

export function deepseekDisableThinkingEnabled(env = process.env) {
  const raw = env.MEMIND_DEEPSEEK_DISABLE_THINKING;
  if (raw != null && String(raw).trim() !== '') {
    return envFlag(raw, false);
  }
  const runtimeProfile = String(env.MEMIND_RUNTIME_PROFILE ?? '').trim().toLowerCase();
  return runtimeProfile === 'local' || runtimeProfile === 'split-service';
}

export function moonshotToolSchemaCompatEnabled(env = process.env) {
  const raw = env.MEMIND_MOONSHOT_TOOL_SCHEMA_COMPAT;
  if (raw != null && String(raw).trim() !== '') {
    return envFlag(raw, false);
  }
  const runtimeProfile = String(env.MEMIND_RUNTIME_PROFILE ?? '').trim().toLowerCase();
  return runtimeProfile === 'local' || runtimeProfile === 'split-service';
}

export function resolveDeepseekNoThinkListenPort(env = process.env) {
  const port = Number(env.MEMIND_DEEPSEEK_NO_THINK_PORT ?? DEFAULT_DEEPSEEK_NO_THINK_PORT);
  return Number.isFinite(port) && port > 0 ? port : DEFAULT_DEEPSEEK_NO_THINK_PORT;
}

export function resolveDeepseekNoThinkProxyBaseUrl(env = process.env) {
  const explicit = String(env.MEMIND_DEEPSEEK_NO_THINK_BASE_URL ?? '').trim();
  if (explicit) return explicit.replace(/\/$/, '');
  const host = String(
    env.MEMIND_DEEPSEEK_NO_THINK_HOST
      ?? env.MEMIND_GOOSED_HOST_GATEWAY
      ?? 'host.docker.internal',
  ).trim() || 'host.docker.internal';
  const port = resolveDeepseekNoThinkListenPort(env);
  return `http://${host}:${port}/v1`;
}

export function resolveMoonshotCompatProxyBaseUrl(env = process.env) {
  const explicit = String(env.MEMIND_MOONSHOT_COMPAT_BASE_URL ?? '').trim();
  if (explicit) return explicit.replace(/\/$/, '');
  const host = String(
    env.MEMIND_DEEPSEEK_NO_THINK_HOST
      ?? env.MEMIND_GOOSED_HOST_GATEWAY
      ?? 'host.docker.internal',
  ).trim() || 'host.docker.internal';
  const port = resolveDeepseekNoThinkListenPort(env);
  return `http://${host}:${port}/moonshot/v1`;
}

export function resolveDeepseekUpstreamBase(env = process.env) {
  const raw = String(
    env.MEMIND_DEEPSEEK_UPSTREAM
      ?? env.DEEPSEEK_API_BASE_URL
      ?? DEFAULT_DEEPSEEK_UPSTREAM,
  ).trim().replace(/\/$/, '');
  return raw.replace(/\/v1$/i, '') || DEFAULT_DEEPSEEK_UPSTREAM;
}

export function resolveMoonshotUpstreamBase(env = process.env) {
  const raw = String(
    env.MEMIND_MOONSHOT_UPSTREAM
      ?? DEFAULT_MOONSHOT_UPSTREAM,
  ).trim().replace(/\/$/, '');
  return raw.replace(/\/v1$/i, '') || DEFAULT_MOONSHOT_UPSTREAM;
}

export function isMoonshotApiUrl(value) {
  try {
    const hostname = new URL(String(value ?? '')).hostname.toLowerCase();
    return hostname === 'api.moonshot.cn' || hostname === 'api.moonshot.ai';
  } catch {
    return false;
  }
}

/**
 * DeepSeek V4 defaults to thinking mode. Tool rounds then require reasoning_content
 * to be replayed; goosed can drop it and fail with HTTP 400. Force-disable thinking
 * unless the caller already set an explicit thinking object.
 */
export function injectDeepseekThinkingDisabled(body) {
  if (!body || typeof body !== 'object' || Array.isArray(body)) {
    return { body, injected: false };
  }
  if (body.thinking?.type === 'disabled') {
    return { body, injected: false };
  }
  return {
    body: {
      ...body,
      thinking: { type: 'disabled' },
    },
    injected: true,
  };
}

function decodeJsonPointerToken(value) {
  return String(value ?? '').replace(/~1/g, '/').replace(/~0/g, '~');
}

function resolveLocalSchemaRef(root, ref) {
  if (!String(ref ?? '').startsWith('#/')) return null;
  const tokens = String(ref)
    .slice(2)
    .split('/')
    .map(decodeJsonPointerToken);
  let current = root;
  for (const token of tokens) {
    if (!current || typeof current !== 'object' || !(token in current)) return null;
    current = current[token];
  }
  return current;
}

/**
 * Moonshot rejects otherwise valid local JSON Schema references with
 * "detected infinite recursion". Inline local references before forwarding
 * tools so Kimi receives an equivalent, self-contained parameter schema.
 */
export function flattenLocalJsonSchemaRefs(schema) {
  if (!schema || typeof schema !== 'object' || Array.isArray(schema)) return schema;
  const root = schema;

  function visit(value, activeRefs = new Set()) {
    if (Array.isArray(value)) return value.map((item) => visit(item, activeRefs));
    if (!value || typeof value !== 'object') return value;

    const ref = typeof value.$ref === 'string' ? value.$ref : '';
    if (ref.startsWith('#/')) {
      const target = resolveLocalSchemaRef(root, ref);
      const siblings = Object.fromEntries(
        Object.entries(value).filter(([key]) => key !== '$ref'),
      );
      if (!target || activeRefs.has(ref)) {
        return visit({
          type: 'object',
          description: String(
            siblings.description ?? 'Recursive schema value',
          ),
        }, activeRefs);
      }
      const nextRefs = new Set(activeRefs);
      nextRefs.add(ref);
      const resolved = visit(target, nextRefs);
      return {
        ...(resolved && typeof resolved === 'object' && !Array.isArray(resolved)
          ? resolved
          : {}),
        ...visit(siblings, activeRefs),
      };
    }

    const output = {};
    for (const [key, child] of Object.entries(value)) {
      if (key === '$schema' || key === '$defs' || key === 'definitions') continue;
      output[key] = visit(child, activeRefs);
    }
    return output;
  }

  return visit(schema);
}

export function sanitizeMoonshotToolSchemas(body) {
  if (!body || typeof body !== 'object' || Array.isArray(body) || !Array.isArray(body.tools)) {
    return { body, sanitized: 0 };
  }
  let sanitized = 0;
  const tools = body.tools.map((tool) => {
    const parameters = tool?.function?.parameters;
    if (!parameters || typeof parameters !== 'object') return tool;
    const serialized = JSON.stringify(parameters);
    if (!serialized.includes('"$ref"') && !serialized.includes('"$defs"')) return tool;
    sanitized += 1;
    return {
      ...tool,
      function: {
        ...tool.function,
        parameters: flattenLocalJsonSchemaRefs(parameters),
      },
    };
  });
  return sanitized > 0
    ? { body: { ...body, tools }, sanitized }
    : { body, sanitized: 0 };
}

function readRequestBody(req) {
  return new Promise((resolve, reject) => {
    const chunks = [];
    req.on('data', (chunk) => chunks.push(chunk));
    req.on('end', () => resolve(Buffer.concat(chunks)));
    req.on('error', reject);
  });
}

function buildUpstreamUrl(upstreamBase, reqUrl) {
  const incoming = new URL(reqUrl || '/', 'http://127.0.0.1');
  let pathname = incoming.pathname || '/';
  // Accept both /v1/chat/completions and /chat/completions from OpenAI-compatible clients.
  if (!pathname.startsWith('/v1/') && pathname !== '/v1') {
    pathname = pathname === '/' ? '/v1' : `/v1${pathname}`;
  }
  return `${upstreamBase}${pathname}${incoming.search}`;
}

export function decodedUpstreamResponseHeaders(headers) {
  const responseHeaders = {};
  headers?.forEach?.((value, key) => {
    const normalized = String(key).toLowerCase();
    // undici transparently decodes gzip/br but keeps the upstream encoding
    // headers. Forwarding them makes the downstream client decode plaintext
    // again and silently lose the SSE stream.
    if (
      normalized === 'transfer-encoding'
      || normalized === 'content-encoding'
      || normalized === 'content-length'
      || normalized === 'connection'
    ) {
      return;
    }
    responseHeaders[key] = value;
  });
  return responseHeaders;
}

export function createDeepseekNoThinkProxy({
  upstreamBase = resolveDeepseekUpstreamBase(),
  moonshotUpstreamBase = resolveMoonshotUpstreamBase(),
  fetchImpl = undiciFetch,
  logger = console,
} = {}) {
  const normalizedUpstream = String(upstreamBase || DEFAULT_DEEPSEEK_UPSTREAM).replace(/\/$/, '');
  const normalizedMoonshotUpstream = String(
    moonshotUpstreamBase || DEFAULT_MOONSHOT_UPSTREAM,
  ).replace(/\/$/, '');

  async function handle(req, res) {
    if (req.method === 'GET' && (req.url === '/health' || req.url === '/healthz')) {
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        ok: true,
        upstream: normalizedUpstream,
        moonshotUpstream: normalizedMoonshotUpstream,
      }));
      return;
    }

    if (req.method === 'OPTIONS') {
      res.writeHead(204, {
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Headers': '*',
        'Access-Control-Allow-Methods': 'GET,POST,OPTIONS',
      });
      res.end();
      return;
    }

    const incoming = new URL(req.url || '/', 'http://127.0.0.1');
    const isMoonshot = incoming.pathname === '/moonshot'
      || incoming.pathname.startsWith('/moonshot/');
    const routedRequestUrl = isMoonshot
      ? `${incoming.pathname.slice('/moonshot'.length) || '/'}${incoming.search}`
      : req.url;
    const upstreamUrl = buildUpstreamUrl(
      isMoonshot ? normalizedMoonshotUpstream : normalizedUpstream,
      routedRequestUrl,
    );
    const headers = { ...req.headers };
    delete headers.host;
    delete headers['content-length'];

    let body;
    if (req.method !== 'GET' && req.method !== 'HEAD') {
      const raw = await readRequestBody(req);
      if (raw.length > 0 && /chat\/completions/i.test(upstreamUrl)) {
        try {
          const parsed = JSON.parse(raw.toString('utf8'));
          if (isMoonshot) {
            const schemaResult = sanitizeMoonshotToolSchemas(parsed);
            const thinkingResult = injectDeepseekThinkingDisabled(schemaResult.body);
            body = Buffer.from(JSON.stringify(thinkingResult.body), 'utf8');
            if (schemaResult.sanitized > 0) {
              logger.info?.(
                `[moonshot-compat] inlined local refs in ${schemaResult.sanitized} tool schema(s) model=${String(parsed?.model ?? '')}`,
              );
            }
            if (thinkingResult.injected) {
              logger.info?.(
                `[moonshot-compat] injected thinking.disabled model=${String(parsed?.model ?? '')}`,
              );
            }
          } else {
            const { body: next, injected } = injectDeepseekThinkingDisabled(parsed);
            body = Buffer.from(JSON.stringify(next), 'utf8');
            if (injected) {
              logger.info?.(
                `[deepseek-no-think] injected thinking.disabled model=${String(parsed?.model ?? '')}`,
              );
            }
          }
        } catch {
          body = raw;
        }
      } else {
        body = raw;
      }
    }

    const upstream = await fetchImpl(upstreamUrl, {
      method: req.method,
      headers,
      body,
      dispatcher: upstreamUrl.startsWith('https://') ? insecureDispatcher : undefined,
    });

    const responseHeaders = decodedUpstreamResponseHeaders(upstream.headers);
    res.writeHead(upstream.status, responseHeaders);
    if (!upstream.body) {
      res.end();
      return;
    }
    const reader = upstream.body.getReader();
    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        res.write(Buffer.from(value));
      }
      res.end();
    } catch (err) {
      logger.warn?.('[deepseek-no-think] stream error', err instanceof Error ? err.message : err);
      res.end();
    }
  }

  return { handle, upstreamBase: normalizedUpstream };
}

export function startDeepseekNoThinkProxy({
  port = resolveDeepseekNoThinkListenPort(),
  host = '0.0.0.0',
  upstreamBase = resolveDeepseekUpstreamBase(),
  moonshotUpstreamBase = resolveMoonshotUpstreamBase(),
  fetchImpl = undiciFetch,
  logger = console,
} = {}) {
  const proxy = createDeepseekNoThinkProxy({
    upstreamBase,
    moonshotUpstreamBase,
    fetchImpl,
    logger,
  });
  const server = http.createServer((req, res) => {
    proxy.handle(req, res).catch((err) => {
      logger.error?.('[deepseek-no-think] request failed', err instanceof Error ? err.message : err);
      if (!res.headersSent) {
        res.writeHead(502, { 'Content-Type': 'application/json' });
      }
      res.end(JSON.stringify({
        error: {
          message: err instanceof Error ? err.message : 'deepseek no-think proxy failed',
        },
      }));
    });
  });
  return new Promise((resolve, reject) => {
    server.once('error', reject);
    server.listen(port, host, () => {
      logger.info?.(
        `[llm-compat-proxy] listening on http://${host}:${port} -> ${proxy.upstreamBase}`,
      );
      resolve(server);
    });
  });
}

const isMain = Boolean(
  process.env.MEMIND_DEEPSEEK_PROXY_ENTRYPOINT === '1'
  &&
  process.argv[1]
  && fileURLToPath(import.meta.url) === path.resolve(process.argv[1]),
);
if (isMain) {
  startDeepseekNoThinkProxy().catch((err) => {
    console.error(err);
    process.exit(1);
  });
}