draftgo-cli 1.0.4

package/src/mcp/client.js

@@ -0,0 +1,558 @@
+'use strict';
+
+const pkg = require('../../package.json');
+const {
+  DEFAULT_PROTOCOL_VERSION,
+  McpHttpError,
+  postJsonRpc,
+  redactText,
+} = require('./protocol');
+const { allWithAbort } = require('./parallel');
+
+const SAFE_TEST_TOOLS = [
+  'draftgo_project_overview',
+  'draftgo_resource_list',
+  'draftgo_api_search',
+  'draftgo_api_describe',
+  'draftgo_api_call',
+];
+
+const REQUIRED_DRAFTGO_TOOLS = [
+  'draftgo_api_search',
+  'draftgo_api_describe',
+  'draftgo_api_call',
+];
+
+const TOOL_ALIASES = Object.freeze({
+  draftgo_api_search: Object.freeze([
+    'draftgo_api_search', 'draftgo_operation_search', 'search_operations', 'search',
+  ]),
+  draftgo_api_describe: Object.freeze([
+    'draftgo_api_describe', 'draftgo_operation_describe', 'describe_operation', 'describe',
+  ]),
+  draftgo_api_call: Object.freeze([
+    'draftgo_api_invoke', 'draftgo_api_call', 'draftgo_operation_invoke', 'invoke_operation', 'invoke', 'call',
+  ]),
+});
+
+function diagnosticArguments(canonicalName) {
+  if (canonicalName === 'draftgo_resource_list') {
+    return { resource_type: 'pages', limit: 1 };
+  }
+  if (canonicalName === 'draftgo_api_search') return { query: 'listDbMeta', limit: 1 };
+  return {};
+}
+
+function diagnosticData(result) {
+  let value = result;
+  if (value && value.structuredContent) value = value.structuredContent;
+  else if (value && value.structured_content) value = value.structured_content;
+  if (value && value.data) value = value.data;
+  if (value && Object.prototype.hasOwnProperty.call(value, 'value')) value = value.value;
+  return value;
+}
+
+const CONNECTION_PROBE_OPERATION = 'getCurrentUserAPIKey';
+
+function findConnectionProbeOperation(result) {
+  const value = diagnosticData(result);
+  const items = value && Array.isArray(value.items) ? value.items : [];
+  return items.find((operation) => operation
+    && operation.operation_id === CONNECTION_PROBE_OPERATION
+    && String(operation.method || '').toUpperCase() === 'GET'
+    && String(operation.risk || '').toLowerCase() === 'low'
+    && operation.destructive !== true) || null;
+}
+
+function assertSafeConnectionProbe(description) {
+  const value = diagnosticData(description);
+  const operation = value && value.operation && typeof value.operation === 'object'
+    ? value.operation : value;
+  const properties = operation && operation.input_schema && operation.input_schema.properties || {};
+  const required = operation && operation.input_schema && operation.input_schema.required || [];
+  const inputKeys = Object.keys(properties).filter((key) => key !== 'confirm');
+  if (!operation || operation.operation_id !== CONNECTION_PROBE_OPERATION
+      || String(operation.method || '').toUpperCase() !== 'GET'
+      || String(operation.risk || '').toLowerCase() !== 'low'
+      || operation.destructive === true || required.length > 0 || inputKeys.length > 0) {
+    throw new McpRpcError(`DraftGo Registry does not expose a safe ${CONNECTION_PROBE_OPERATION} probe.`, {
+      code: -32601,
+    });
+  }
+  return operation;
+}
+
+class McpRpcError extends Error {
+  constructor(message, { code = -32000, data, id = null } = {}) {
+    super(message);
+    this.name = 'McpRpcError';
+    this.code = code;
+    this.data = data;
+    this.id = id;
+  }
+}
+
+const SESSION_INVALID_CODES = new Set([
+  -32002,
+  'SESSION_EXPIRED',
+  'SESSION_NOT_FOUND',
+  'MCP_SESSION_EXPIRED',
+  'MCP_SESSION_NOT_FOUND',
+]);
+
+function isSessionInvalidError(error) {
+  if (!error) return false;
+  if (Number(error.status) === 404) return true;
+  const code = error.code != null
+    ? error.code
+    : error.rpc && error.rpc.error && error.rpc.error.code;
+  if (SESSION_INVALID_CODES.has(code)) return true;
+  const message = String(error.message
+    || error.rpc && error.rpc.error && error.rpc.error.message
+    || '');
+  return /\b(?:mcp\s+)?session(?:\s+id)?\s+(?:was\s+)?(?:not\s+found|expired|invalid|unknown|lost)\b/i.test(message)
+    || /\b(?:server|mcp)\s+(?:is\s+)?(?:not initialized|uninitialized)\b/i.test(message)
+    || /\binitialize first\b/i.test(message);
+}
+
+function requestIdKeys(message) {
+  const messages = Array.isArray(message) ? message : [message];
+  return new Set(messages
+    .filter((item) => isObject(item) && typeof item.method === 'string' && item.id != null)
+    .map((item) => `${typeof item.id}:${JSON.stringify(item.id)}`));
+}
+
+function sessionInvalidResponse(message, expectedIds) {
+  const messages = Array.isArray(message) ? message : [message];
+  return messages.find((item) => isObject(item)
+    && item.error
+    && expectedIds.has(`${typeof item.id}:${JSON.stringify(item.id)}`)
+    && isSessionInvalidError(item.error)) || null;
+}
+
+function redactValue(value, secrets = [], seen = new WeakMap()) {
+  if (typeof value === 'string') return redactText(value, secrets);
+  if (!value || typeof value !== 'object') return value;
+  if (seen.has(value)) return seen.get(value);
+  if (Array.isArray(value)) {
+    const result = [];
+    seen.set(value, result);
+    for (const item of value) result.push(redactValue(item, secrets, seen));
+    return result;
+  }
+  const result = {};
+  seen.set(value, result);
+  for (const [key, item] of Object.entries(value)) {
+    result[key] = redactValue(item, secrets, seen);
+  }
+  return result;
+}
+
+function isObject(value) {
+  return !!value && typeof value === 'object' && !Array.isArray(value);
+}
+
+function sameId(left, right) {
+  return left === right;
+}
+
+function findInitializeRequest(message) {
+  const messages = Array.isArray(message) ? message : [message];
+  return messages.find((item) => isObject(item) && item.method === 'initialize' && item.id != null);
+}
+
+function findResponse(messages, id) {
+  return messages.find((message) => isObject(message)
+    && Object.prototype.hasOwnProperty.call(message, 'id')
+    && sameId(message.id, id)
+    && !message.method);
+}
+
+function toolMatches(name, candidate) {
+  const aliases = TOOL_ALIASES[candidate] || [candidate];
+  return aliases.some((alias) => name === alias || name.endsWith(`.${alias}`)
+    || name.endsWith(`/${alias}`) || name.endsWith(`:${alias}`));
+}
+
+class DraftGoMcpClient {
+  constructor(config, options = {}) {
+    if (!config || typeof config !== 'object') throw new TypeError('DraftGo MCP config is required.');
+    const token = String(config.token || config.sat || '').trim();
+    if (!token) throw new Error('DraftGo MCP config is missing API Key.');
+    if (!config.server && !config.mcp_url) throw new Error('DraftGo MCP config is missing server.');
+
+    this.config = { ...config, token };
+    this.secrets = [token];
+    this.sessionId = null;
+    this.protocolVersion = options.protocolVersion || DEFAULT_PROTOCOL_VERSION;
+    this.nextId = 1;
+    this.sessionGeneration = 0;
+    this.sessionRecovery = null;
+    this.sessionRecoveryCount = 0;
+    this.onMessage = typeof options.onMessage === 'function' ? options.onMessage : null;
+  }
+
+  async forward(message, options = {}) {
+    if (!isObject(message) && !Array.isArray(message)) {
+      throw new TypeError('MCP message must be a JSON object or batch.');
+    }
+
+    const initialize = findInitializeRequest(message);
+    const internalRecovery = options._sessionRecoveryInternal === true;
+    if (!initialize && !internalRecovery && this.sessionRecovery) {
+      await this.sessionRecovery;
+    }
+    if (initialize) this.sessionId = null;
+    const requestedVersion = initialize
+      && initialize.params
+      && initialize.params.protocolVersion;
+    const delivered = [];
+    const expectedIds = requestIdKeys(message);
+    const attemptedSessionId = initialize ? null : this.sessionId;
+    const attemptedGeneration = this.sessionGeneration;
+    const canRecover = !initialize
+      && !internalRecovery
+      && options._sessionRecoveryAttempted !== true
+      && attemptedSessionId != null;
+    let suppressedSessionError = null;
+
+    try {
+      await postJsonRpc(this.config, message, {
+        signal: options.signal,
+        timeoutMs: options.timeoutMs,
+        sessionId: attemptedSessionId,
+        protocolVersion: requestedVersion || this.protocolVersion,
+        onSession: (sessionId) => {
+          if (initialize || (this.sessionGeneration === attemptedGeneration
+              && this.sessionId === attemptedSessionId)) {
+            this.sessionId = sessionId;
+          }
+        },
+        onMessage: (remoteMessage) => {
+          const safe = redactValue(remoteMessage, this.secrets);
+          delivered.push(safe);
+          const invalid = canRecover && sessionInvalidResponse(safe, expectedIds);
+          if (invalid) {
+            suppressedSessionError = invalid;
+            return;
+          }
+          if (options._suppressCallbacks === true) return;
+          if (this.onMessage) this.onMessage(safe);
+          if (typeof options.onMessage === 'function') options.onMessage(safe);
+        },
+      });
+      if (suppressedSessionError) {
+        const remoteError = suppressedSessionError.error;
+        const error = new McpRpcError(
+          remoteError.message || `MCP error ${remoteError.code}`,
+          { code: remoteError.code, data: remoteError.data, id: suppressedSessionError.id },
+        );
+        error.rpc = suppressedSessionError;
+        throw error;
+      }
+    } catch (error) {
+      if (error && error.rpc) error.rpc = redactValue(error.rpc, this.secrets);
+      if (error && error.message) error.message = redactText(error.message, this.secrets);
+      if (canRecover && isSessionInvalidError(error)) {
+        await this.recoverSession(attemptedSessionId, attemptedGeneration, options);
+        return this.forward(message, { ...options, _sessionRecoveryAttempted: true });
+      }
+      throw error;
+    }
+
+    if (initialize) {
+      const response = findResponse(delivered, initialize.id);
+      const negotiated = response && response.result && response.result.protocolVersion;
+      if (negotiated) this.protocolVersion = negotiated;
+    }
+    return delivered;
+  }
+
+  async recoverSession(failedSessionId, failedGeneration, options = {}) {
+    if (this.sessionGeneration !== failedGeneration) return;
+    if (this.sessionId != null && this.sessionId !== failedSessionId) return;
+    if (this.sessionRecovery) return this.sessionRecovery;
+
+    const recoveryNumber = this.sessionRecoveryCount + 1;
+    this.sessionId = null;
+    const recovery = (async () => {
+      if (typeof options.onSessionRecovery === 'function') {
+        options.onSessionRecovery('started', { attempt: recoveryNumber });
+      }
+      try {
+        await this.initialize({
+          id: `draftgo-session-recovery-${recoveryNumber}`,
+          protocolVersion: this.protocolVersion,
+          timeoutMs: options.timeoutMs,
+          signal: options.signal,
+          _sessionRecoveryInternal: true,
+          _suppressCallbacks: true,
+        });
+        this.sessionRecoveryCount = recoveryNumber;
+        if (typeof options.onSessionRecovery === 'function') {
+          options.onSessionRecovery('succeeded', { attempt: recoveryNumber });
+        }
+      } catch (error) {
+        this.sessionId = null;
+        if (typeof options.onSessionRecovery === 'function') {
+          options.onSessionRecovery('failed', { attempt: recoveryNumber });
+        }
+        const message = redactText(error && error.message ? error.message : error, this.secrets);
+        const wrapped = new McpRpcError(`DraftGo MCP session recovery failed: ${message}`, {
+          code: error && error.code != null ? error.code : -32000,
+          data: error && error.data,
+        });
+        wrapped.status = Number(error && error.status) || 0;
+        wrapped.sessionRecovery = true;
+        wrapped.cause = error;
+        throw wrapped;
+      }
+    })();
+    this.sessionRecovery = recovery;
+    try {
+      await recovery;
+    } finally {
+      if (this.sessionRecovery === recovery) this.sessionRecovery = null;
+    }
+  }
+
+  async request(method, params, options = {}) {
+    const id = options.id == null ? this.nextId++ : options.id;
+    const message = { jsonrpc: '2.0', id, method };
+    if (params !== undefined) message.params = params;
+    const messages = await this.forward(message, options);
+    const response = findResponse(messages, id);
+    if (!response) {
+      throw new McpRpcError(`DraftGo MCP returned no response for ${method}.`, { id });
+    }
+    if (response.error) {
+      throw new McpRpcError(
+        redactText(response.error.message || `MCP error ${response.error.code}`, this.secrets),
+        {
+          code: response.error.code,
+          data: redactValue(response.error.data, this.secrets),
+          id,
+        },
+      );
+    }
+    return response.result;
+  }
+
+  async notify(method, params, options = {}) {
+    const message = { jsonrpc: '2.0', method };
+    if (params !== undefined) message.params = params;
+    await this.forward(message, options);
+  }
+
+  async initialize(options = {}) {
+    if (typeof options.onStage === 'function') options.onStage('initialize', 'started');
+    let result;
+    try {
+      result = await this.request('initialize', {
+        protocolVersion: options.protocolVersion || this.protocolVersion,
+        capabilities: options.capabilities || {},
+        clientInfo: options.clientInfo || {
+          name: 'draftgo-cli',
+          version: pkg.version,
+        },
+      }, options);
+    } catch (error) {
+      if (typeof options.onStage === 'function') options.onStage('initialize', 'failed');
+      throw error;
+    }
+    if (typeof options.onStage === 'function') options.onStage('initialize', 'succeeded', result);
+    if (result && result.protocolVersion) this.protocolVersion = result.protocolVersion;
+    try {
+      await this.notify('notifications/initialized', undefined, options);
+    } catch (error) {
+      if (typeof options.onStage === 'function') options.onStage('initialized', 'failed');
+      throw error;
+    }
+    if (typeof options.onStage === 'function') options.onStage('initialized', 'succeeded');
+    this.sessionGeneration += 1;
+    return result;
+  }
+
+  async toolsList(cursor, options = {}) {
+    const params = cursor ? { cursor } : {};
+    return this.request('tools/list', params, options);
+  }
+
+  async listAllTools(options = {}) {
+    const tools = [];
+    const seenCursors = new Set();
+    let cursor;
+    const maxPages = Number(options.maxPages || 100);
+    for (let page = 0; page < maxPages; page += 1) {
+      const result = await this.toolsList(cursor, options);
+      if (result && Array.isArray(result.tools)) tools.push(...result.tools);
+      const next = result && result.nextCursor;
+      if (!next) return tools;
+      if (seenCursors.has(next)) throw new McpRpcError('DraftGo MCP repeated a tools/list cursor.');
+      seenCursors.add(next);
+      cursor = next;
+    }
+    throw new McpRpcError(`DraftGo MCP tools/list exceeded ${maxPages} pages.`);
+  }
+
+  async toolsCall(name, args = {}, options = {}) {
+    const result = await this.request('tools/call', { name, arguments: args }, options);
+    if (result && result.isError === true) {
+      const structured = result.structuredContent || result.structured_content || {};
+      const text = Array.isArray(result.content)
+        ? result.content.find((part) => part && part.type === 'text' && typeof part.text === 'string')
+        : null;
+      const message = structured.message || structured.error && structured.error.message
+        || text && text.text || `DraftGo tool failed: ${name}`;
+      throw new McpRpcError(redactText(String(message).slice(0, 4096), this.secrets), {
+        code: structured.code || structured.error && structured.error.code || -32000,
+        data: redactValue(structured, this.secrets),
+      });
+    }
+    return result;
+  }
+
+  async testConnection(options = {}) {
+    const initialized = await this.initialize(options);
+    if (typeof options.onStage === 'function') options.onStage('tools/list', 'started');
+    let tools;
+    try {
+      tools = await this.listAllTools(options);
+    } catch (error) {
+      if (typeof options.onStage === 'function') options.onStage('tools/list', 'failed');
+      throw error;
+    }
+    if (typeof options.onStage === 'function') options.onStage('tools/list', 'succeeded', { count: tools.length });
+    const requiredTools = options.requiredTools || REQUIRED_DRAFTGO_TOOLS;
+    const missing = requiredTools.filter((expected) => !tools.some((tool) =>
+      tool && typeof tool.name === 'string' && toolMatches(tool.name, expected)));
+    if (missing.length) {
+      throw new McpRpcError(`DraftGo MCP is missing required tools: ${missing.join(', ')}.`, { code: -32601 });
+    }
+    const safeTools = options.safeTools || SAFE_TEST_TOOLS;
+    const hasDiagnosticOverrides = Boolean(options.requiredTools || options.safeTools);
+    const firstSafeTool = tools.find((tool) => tool && typeof tool.name === 'string'
+      && safeTools.some((canonical) => toolMatches(tool.name, canonical)));
+    if (!firstSafeTool) {
+      throw new McpRpcError(
+        `DraftGo MCP exposes no safe diagnostic tool (${safeTools.join(', ')}).`,
+        { code: -32601 },
+      );
+    }
+
+    const selected = safeTools.map((canonical) => ({
+      canonical,
+      tool: tools.find((candidate) => candidate && typeof candidate.name === 'string'
+        && toolMatches(candidate.name, canonical)),
+    })).filter((entry) => entry.tool);
+    const byCanonical = Object.fromEntries(selected.map((entry) => [entry.canonical, entry.tool]));
+    const overrideCanonical = safeTools.find((canonical) => toolMatches(firstSafeTool.name, canonical));
+    const defaultPlan = hasDiagnosticOverrides ? [{
+      label: overrideCanonical,
+      canonical: overrideCanonical,
+      tool: firstSafeTool,
+      args: diagnosticArguments(overrideCanonical),
+    }] : [
+      { label: 'api:catalog', canonical: 'draftgo_api_search', tool: byCanonical.draftgo_api_search, args: { query: CONNECTION_PROBE_OPERATION, limit: 100 } },
+    ];
+    const runPlan = (plan, startIndex = 0) => allWithAbort(plan.map((entry, offset) =>
+      async (queryOptions) => {
+        if (typeof options.onStage === 'function') options.onStage(`tools/call:${entry.label}`, 'started');
+        let result;
+        try {
+          result = await this.toolsCall(
+            entry.tool.name,
+            options.toolArguments && startIndex + offset === 0
+              ? options.toolArguments
+              : entry.args,
+            queryOptions,
+          );
+        } catch (error) {
+          if (typeof options.onStage === 'function') options.onStage(`tools/call:${entry.label}`, 'failed');
+          throw error;
+        }
+        if (typeof options.onStage === 'function') options.onStage(`tools/call:${entry.label}`, 'succeeded');
+        return ({
+        label: entry.label,
+        canonical: entry.canonical,
+        name: entry.tool.name,
+        arguments: options.toolArguments && startIndex + offset === 0
+          ? options.toolArguments
+          : entry.args,
+        result,
+      });
+      }), options);
+    const tested = await runPlan(defaultPlan);
+    if (!hasDiagnosticOverrides) {
+      const discovery = tested.find((entry) => entry.label === 'api:catalog');
+      const catalog = diagnosticData(discovery && discovery.result);
+      const operation = findConnectionProbeOperation(discovery && discovery.result);
+      if (!operation || !operation.operation_id) {
+        throw new McpRpcError(`DraftGo Registry returned no safe ${CONNECTION_PROBE_OPERATION} operation.`, {
+          code: -32601,
+        });
+      }
+      const operationID = String(operation.operation_id);
+      const revision = String(catalog.registry_revision || '').trim();
+      if (!revision) {
+        throw new McpRpcError('DraftGo api_search did not return registry_revision.');
+      }
+      const [described] = await runPlan([{
+        label: 'api:describe',
+        canonical: 'draftgo_api_describe',
+        tool: byCanonical.draftgo_api_describe,
+        args: { operation_id: operationID, registry_revision: revision },
+      }], tested.length);
+      tested.push(described);
+      assertSafeConnectionProbe(described.result);
+      tested.push(...await runPlan([{
+        label: 'api:invoke',
+        canonical: 'draftgo_api_call',
+        tool: byCanonical.draftgo_api_call,
+        args: { operation_id: operationID, registry_revision: revision },
+      }], tested.length));
+    }
+    return {
+      initialized,
+      protocolVersion: this.protocolVersion,
+      sessionId: this.sessionId,
+      sessionRecoveries: this.sessionRecoveryCount,
+      tools,
+      testedTool: tested[0].name,
+      toolResult: tested[0].result,
+      testedTools: [...new Set(tested.map((entry) => entry.name))],
+      testedCalls: tested.map((entry) => entry.label),
+      diagnosticResults: tested,
+    };
+  }
+}
+
+async function testConnection(config, options = {}) {
+  const client = new DraftGoMcpClient(config, options);
+  return client.testConnection(options);
+}
+
+async function callTool(config, name, args = {}, options = {}) {
+  const client = options.client || new DraftGoMcpClient(config, options);
+  if (!options.initialized) await client.initialize(options);
+  return client.toolsCall(name, args, options);
+}
+
+module.exports = {
+  DraftGoMcpClient,
+  McpClient: DraftGoMcpClient,
+  McpRpcError,
+  McpHttpError,
+  SAFE_TEST_TOOLS,
+  REQUIRED_DRAFTGO_TOOLS,
+  TOOL_ALIASES,
+  diagnosticArguments,
+  diagnosticData,
+  findConnectionProbeOperation,
+  assertSafeConnectionProbe,
+  isSessionInvalidError,
+  redactValue,
+  testConnection,
+  callTool,
+  toolMatches,
+};