Zion Boggan
repos/TreeTrace/src/mcp.js
zionboggan.com ↗
198 lines · javascript
History for this file →
1
import { createInterface } from 'node:readline';
2
import { resolve } from 'node:path';
3
import { parseArgs, loadRedactedTree, detectProjectName, assertClean } from './cli.js';
4
import { renderHandoff } from './handoff.js';
5
import { renderLessonsMarkdown, analyzeTree, renderRejectionsJson } from './analyze.js';
6
import { renderSecurityReport } from './security-report.js';
7
import { renderHallucinationsJson } from './hallucinate.js';
8
import { renderJson } from './render-json.js';
9
import { SCHEMA_VERSION } from './config.js';
10
import { TreetraceError, ExitCode } from './util.js';
11
 
12
const PROTOCOL_VERSION = '2024-11-05';
13
const MAX_REQUEST_BYTES = 1048576;
14
 
15
const TOOL_DEFS = [
16
  {
17
    name: 'handoff',
18
    description: 'Continuation brief for the next agent: goal, accepted decisions, constraints, and dead ends. Read only.',
19
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
20
  },
21
  {
22
    name: 'lessons',
23
    description: 'Accepted constraints and repeated corrections distilled from the session lineage. Read only.',
24
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
25
  },
26
  {
27
    name: 'security_summary',
28
    description: 'Evidence-backed security-sensitive touches, test changes, risky commands, and hallucinated references. Read only.',
29
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
30
  },
31
  {
32
    name: 'eval_candidates',
33
    description: 'Compact regression cases derived from session corrections and hallucinated references. Read only.',
34
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
35
  },
36
  {
37
    name: 'tree',
38
    description: 'Full prompt-lineage tree as canonical JSON (nodes, stats, analysis). The structured counterpart to the Markdown reports. Read only.',
39
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
40
  },
41
  {
42
    name: 'rejections_summary',
43
    description: 'Typed rejection / refusal / decline events captured on the session (tool declines, interrupts, permission denials, tool errors, model refusals). Read only.',
44
    inputSchema: { type: 'object', properties: {}, additionalProperties: false },
45
  },
46
];
47
 
48
export async function startMcpServer({ argv, version }, io = {}) {
49
  const input = io.input || process.stdin;
50
  const output = io.output || process.stdout;
51
  const opts = parseArgs((argv || []).filter((a) => a !== 'mcp' && a !== '--mcp'));
52
  if (opts.stdin) {
53
    throw new TreetraceError(
54
      'treetrace mcp does not support --stdin: stdin is the JSON-RPC transport for the MCP server. ' +
55
        'Point the server at a project with --dir, or import a transcript with --file.',
56
      ExitCode.USAGE
57
    );
58
  }
59
  const projectDir = resolve(opts.dir || process.cwd());
60
  const projectName = detectProjectName(projectDir);
61
 
62
  let cache = null;
63
  let inFlight = null;
64
  const ensureTree = async () => {
65
    if (cache) return cache;
66
    if (!inFlight) {
67
      inFlight = (async () => {
68
        const { tree, decisions } = await loadRedactedTree(opts, projectDir, projectName, () => {}, { forceAuto: true });
69
        cache = { tree, decisions, renderOpts: { projectName, version, projectDir, generatedAt: new Date().toISOString() } };
70
        return cache;
71
      })().finally(() => {
72
        inFlight = null;
73
      });
74
    }
75
    return inFlight;
76
  };
77
 
78
  return new Promise((resolveServer) => {
79
    const rl = createInterface({ input, crlfDelay: Infinity });
80
    const send = (msg) => output.write(`${JSON.stringify(msg)}\n`);
81
 
82
    rl.on('line', async (line) => {
83
      const text = line.trim();
84
      if (!text) return;
85
      if (text.length > MAX_REQUEST_BYTES) {
86
        send({ jsonrpc: '2.0', id: null, error: { code: -32600, message: 'Invalid Request: request exceeds size limit' } });
87
        return;
88
      }
89
      let req;
90
      try {
91
        req = JSON.parse(text);
92
      } catch {
93
        send({ jsonrpc: '2.0', id: null, error: { code: -32700, message: 'Parse error' } });
94
        return;
95
      }
96
      if (Array.isArray(req)) {
97
        send({ jsonrpc: '2.0', id: null, error: { code: -32600, message: 'Invalid Request: JSON-RPC batch requests are not supported' } });
98
        return;
99
      }
100
      try {
101
        await handle(req, send, ensureTree, version);
102
      } catch (err) {
103
        if (isRequestWithId(req)) {
104
          send({
105
            jsonrpc: '2.0',
106
            id: req.id,
107
            error: { code: -32603, message: `Internal error: ${err && err.message ? err.message : 'unknown'}` },
108
          });
109
        }
110
      }
111
    });
112
    rl.on('close', () => resolveServer());
113
  });
114
}
115
 
116
function isRequestWithId(req) {
117
  return Boolean(req) && typeof req === 'object' && !Array.isArray(req) && 'id' in req;
118
}
119
 
120
async function handle(req, send, ensureTree, version) {
121
  const hasId = isRequestWithId(req);
122
  if (!req || req.jsonrpc !== '2.0' || typeof req.method !== 'string') {
123
    if (hasId) send({ jsonrpc: '2.0', id: req.id, error: { code: -32600, message: 'Invalid Request' } });
124
    return;
125
  }
126
  const isNotification = !hasId;
127
  const reply = (result) => { if (!isNotification) send({ jsonrpc: '2.0', id: req.id, result }); };
128
  const fail = (code, message) => { if (!isNotification) send({ jsonrpc: '2.0', id: req.id, error: { code, message } }); };
129
 
130
  switch (req.method) {
131
    case 'initialize':
132
      reply({
133
        protocolVersion: PROTOCOL_VERSION,
134
        capabilities: { tools: {} },
135
        serverInfo: { name: 'treetrace', version: version || '0.0.0' },
136
      });
137
      return;
138
    case 'notifications/initialized':
139
    case 'initialized':
140
      return;
141
    case 'ping':
142
      reply({});
143
      return;
144
    case 'tools/list':
145
      reply({ tools: TOOL_DEFS });
146
      return;
147
    case 'tools/call': {
148
      const params = req.params || {};
149
      const name = params.name;
150
      const def = TOOL_DEFS.find((t) => t.name === name);
151
      if (!def) {
152
        fail(-32602, `Unknown tool: ${name}`);
153
        return;
154
      }
155
      const args = params.arguments;
156
      if (args !== undefined && args !== null) {
157
        if (typeof args !== 'object' || Array.isArray(args) || Object.keys(args).length > 0) {
158
          fail(-32602, `Tool ${name} accepts no arguments`);
159
          return;
160
        }
161
      }
162
      const { tree, decisions, renderOpts } = await ensureTree();
163
      const text = renderTool(name, tree, renderOpts);
164
      assertClean(text, decisions, `mcp tool ${name}`);
165
      reply({ content: [{ type: 'text', text }], isError: false });
166
      return;
167
    }
168
    default:
169
      fail(-32601, `Method not found: ${req.method}`);
170
  }
171
}
172
 
173
function renderTool(name, tree, renderOpts) {
174
  switch (name) {
175
    case 'handoff':
176
      return renderHandoff(tree, renderOpts);
177
    case 'lessons':
178
      return renderLessonsMarkdown(tree, renderOpts);
179
    case 'security_summary':
180
      return renderSecurityReport(tree, renderOpts.projectDir || null, renderOpts);
181
    case 'eval_candidates': {
182
      const analysis = analyzeTree(tree);
183
      const hall = renderHallucinationsJson(tree, renderOpts.projectDir || null, renderOpts);
184
      const payload = {
185
        schemaVersion: SCHEMA_VERSION,
186
        evalCandidates: analysis.evalCandidates,
187
        hallucinationEvalCandidates: hall.hallucinations.map((h) => h.evalCandidate),
188
      };
189
      return JSON.stringify(payload, null, 2);
190
    }
191
    case 'tree':
192
      return JSON.stringify(renderJson(tree, renderOpts), null, 2);
193
    case 'rejections_summary':
194
      return JSON.stringify(renderRejectionsJson(tree, renderOpts), null, 2);
195
    default:
196
      return '';
197
  }
198
}