dsh-search-mcp/lib/client.js

/**
 * MCP transport layer for one search.
 *
 * Desktop 0.2 ships the split MCP client package (v2). Import that package so
 * link:/Desktop profiles can resolve it from the host graph without a local
 * `node_modules` copy of the legacy monolith SDK.
 *
 * HTTP egress:
 * 1. DSH policy (`proxyRouteFor`) when `@deepseek-ai/dsh-http-proxy` resolves.
 * 2. Else env + undici global dispatcher (DSH boot installs proxy from HTTP_PROXY).
 * 3. Else DNS-pinned direct dial (SSRF-safe).
 */
import { Client, StreamableHTTPClientTransport } from '@modelcontextprotocol/client';
import { StdioClientTransport } from '@modelcontextprotocol/client/stdio';
import { WebError } from '@deepseek-ai/dsh-web';
import { Agent, fetch as undiciFetch } from 'undici';
import { clampSearchResults } from './catalog.js';
import { resolveEgressRoute } from './proxy-env.js';
import {
  isNonPublicIpLiteral,
  parseHttpEndpoint,
  validateHttpEndpoint,
} from './url-policy.js';

/** Run one search through a resolved server entry. */
export async function callMcpSearch(server, key, args, signal) {
  if (!server.toolName) {
    throw new WebError(
      `search-mcp server "${server.id}": no MCP tool name (set "toolName" or pick a known kind)`,
      'WEB_PROVIDER_ERROR',
    );
  }

  let runtime;
  // Do not opt into MCP 2026 era negotiation. SDK default is legacy (plain
  // 2025 initialize, no server/discover probe). Bailian/DashScope answers
  // that probe with HTTP 500 and connect fails before any search call.
  const client = new Client(
    { name: 'dsh-search-mcp', version: '0.2.42' },
    { capabilities: {} },
  );
  try {
    runtime = server.transport === 'stdio'
      ? { transport: stdioTransport(server, key), close: async () => {} }
      : await httpRuntime(server, key, signal);
    await race(client.connect(runtime.transport), signal, `connect to "${server.id}"`);
    const callArgs = { query: args.query };
    if (server.countArg.length > 0 && args.maxResults !== undefined) {
      callArgs[server.countArg] = clampSearchResults(server, args.maxResults);
    }
    const result = await race(
      client.callTool({ name: server.toolName, arguments: callArgs }),
      signal,
      `call "${server.id}" tool "${server.toolName}"`,
    );
    if (result.isError) {
      throw new WebError(
        `search-mcp: MCP server "${server.id}" tool "${server.toolName}" reported an error`,
        'WEB_PROVIDER_ERROR',
      );
    }
    return result;
  } catch (error) {
    if (error instanceof WebError) throw error;
    if (signal?.aborted) throw aborted(`complete request for "${server.id}"`);
    const detail = error?.name === 'SearchMcpUrlPolicyError' ? `: ${error.message}` : '';
    throw new WebError(
      `search-mcp server "${server.id}" request failed${detail}`,
      'WEB_PROVIDER_ERROR',
    );
  } finally {
    try {
      await client.close();
    } catch {
      // The connection is already gone.
    }
    try {
      await runtime?.close();
    } catch {
      // The dedicated dispatcher has no shared state to recover.
    }
  }
}

/** Build streamable-http transport: proxy when policy/env says so, else DNS-pin. */
async function httpRuntime(server, key, signal) {
  const parsed = parseHttpEndpoint(server.url);
  const url = new URL(parsed.url);
  const headers = {};
  if (key !== undefined && key.length > 0 && server.authParam.length > 0) {
    const value = `${server.authPrefix ?? ''}${key}`;
    if (server.authStyle === 'query') url.searchParams.set(server.authParam, value);
    else if (server.authStyle === 'header') headers[server.authParam] = value;
  }

  const route = await resolveEgressRoute(url, isNonPublicIpLiteral);
  if (route.proxied) {
    return buildTransport(url, headers, signal, route.dispatcher, route.close ?? (async () => {}));
  }

  const validated = await validateHttpEndpoint(server.url, { signal });
  const agent = new Agent({
    connect: { lookup: validated.lookup },
    connections: validated.addresses.length,
    maxRedirections: 0,
  });
  return buildTransport(url, headers, signal, agent, () => agent.close());
}

function buildTransport(url, headers, signal, dispatcher, close) {
  const expectedOrigin = url.origin;
  const secureFetch = async (input, init = {}) => {
    const requestUrl = new URL(typeof input === 'string' || input instanceof URL ? input : input.url);
    if (requestUrl.origin !== expectedOrigin) {
      throw new Error('search-mcp URL policy: request origin changed after validation');
    }
    const fetchOptions = {
      ...init,
      redirect: 'error',
      ...(signal !== undefined ? { signal: combineSignals(signal, init.signal) } : {}),
    };
    if (dispatcher !== undefined) fetchOptions.dispatcher = dispatcher;
    return undiciFetch(input, fetchOptions);
  };

  return {
    transport: new StreamableHTTPClientTransport(url, {
      fetch: secureFetch,
      requestInit: {
        headers,
        redirect: 'error',
        ...(signal !== undefined ? { signal } : {}),
      },
    }),
    close,
  };
}

/** Build a stdio transport; the authParam name doubles as the env var name. */
function stdioTransport(server, key) {
  const env = { ...process.env };
  if (key !== undefined && key.length > 0 && server.authParam.length > 0) {
    env[server.authParam] = `${server.authPrefix ?? ''}${key}`;
  }
  return new StdioClientTransport({
    command: server.command,
    args: server.args ?? [],
    env,
  });
}

/** Race a protocol operation against the caller/timeout abort signal. */
function race(promise, signal, stage) {
  if (signal === undefined) return promise;
  if (signal.aborted) throw aborted(stage);
  return new Promise((resolve, reject) => {
    const onAbort = () => {
      signal.removeEventListener('abort', onAbort);
      reject(aborted(stage));
    };
    signal.addEventListener('abort', onAbort, { once: true });
    promise.then(
      (value) => {
        signal.removeEventListener('abort', onAbort);
        resolve(value);
      },
      (error) => {
        signal.removeEventListener('abort', onAbort);
        reject(error);
      },
    );
  });
}

function combineSignals(base, request) {
  if (request === undefined || request === null || request === base) return base;
  return AbortSignal.any([base, request]);
}

function aborted(stage) {
  return new WebError(`search-mcp: aborted while trying to ${stage}`, 'WEB_ABORTED');
}