mcp-control

1import type { Register } from 'claude-code'
2
3import { enableText, learn, listText, serversFor } from './servers.js'
4import type { Catalog, Server } from './servers.js'
5
6/**
7 * One mod for MCP servers: Claude can list every configured server (on or off in this project),
8 * ask you with one click to enable one, and reconnect a server after editing it. Enable and
9 * reconnect run the built-in `/mcp enable|reconnect <server>` once the turn ends (a command can't
10 * run mid-turn), then a follow-up prompt lets Claude carry on with the new tools.
11 *
12 * `/mcp enable` edits this project's `disabledMcpServers` in ~/.claude.json, so whatever was on
13 * or off in the last conversation in a folder is how the next one starts.
14 */
15const P = 'mcp__mcp-control__'
16type Job = { kind: 'enable' | 'reconnect'; server: string }
17const S: { jobs: Job[]; then: string; catalog: Catalog; off: Server[]; timer?: { cancel: () => void }; epoch: number } = { jobs: [], then: '', catalog: {}, off: [], epoch: 0 }
18
19export const register: Register = on => {
20  on('session.start', async ($, e, next) => {
21    const r = await next(e)
22    S.timer?.cancel(); S.timer = undefined; S.epoch++
23    S.jobs = []; S.then = ''
24    S.catalog = ((await $.store.get('catalog').catch(() => undefined)) as Catalog | undefined) ?? {}
25    const tools = [
26      { name: 'list_mcp_servers', description: 'List every MCP server configured for this project, whether it is ON or OFF here, and what each does (its tools). Use it when a task might need a service you have no tools for (Jira, Confluence, a browser, logs, GitHub…).', inputSchema: { type: 'object', properties: {} } },
27      { name: 'enable_mcp', description: 'Turn ON an MCP server that is OFF in this project. The user approves with one click; it is then enabled right after this turn ends and you get a follow-up message, after which its tools work. It stays on for this project. Pass server (the name list_mcp_servers shows) and then (what to do with it).', inputSchema: { type: 'object', properties: { server: { type: 'string' }, then: { type: 'string', description: 'What to do once it is on, e.g. "search Jira for open PIQ bugs".' } }, required: ['server'] } },
28      { name: 'reconnect_mcp', description: 'Restart and reconnect an MCP server so its latest code and tools load (use after editing an MCP server you are testing). Happens right after this turn ends; you then get a follow-up message. Pass server (as /mcp lists it, or "all").', inputSchema: { type: 'object', properties: { server: { type: 'string' }, then: { type: 'string' } }, required: ['server'] } },
29    ]
30    for (const t of tools) await $.tool.register(t as never)
31    await remember($)
32    // What is OFF is fixed for the session except for what enable_mcp turns on, so read it once here.
33    try { S.off = (await servers($)).filter(s => s.off) } catch { S.off = [] }
34    $.ui.invalidate('tool.describe')
35    return r
36  })
37
38  on('classic.SessionStart', async ($, e, next) => {
39    const r = await next(e)
40    if (['clear', 'resume', 'fork'].includes(e.source)) {
41      S.timer?.cancel(); S.timer = undefined; S.epoch++; S.jobs = []; S.then = ''
42      S.off = (await servers($)).filter(s => s.off)
43      $.ui.invalidate('tool.describe')
44    }
45    return r
46  })
47
48  on('tool.call', { tool: `${P}list_mcp_servers` }, async $ => ({ result: listText(await servers($), S.catalog) }))
49
50  on('tool.call', { tool: `${P}enable_mcp` }, async ($, e) => {
51    const input = e as unknown as { server?: string; then?: string }
52    const name = String(input.server ?? '').trim()
53    const all = await servers($)
54    const s = all.find(x => x.name === name)
55    if (!s) return { result: `No MCP server named "${name}". Known: ${all.map(x => x.name).join(', ')}` }
56    if (!s.off) return { result: `"${name}" is already on in this project. If its tools are missing, call reconnect_mcp.` }
57    let answer = ''
58    try {
59      answer = await $.ui.ask(`Enable the "${name}" MCP server for this project?${input.then ? ` Claude wants to: ${String(input.then).slice(0, 120)}` : ''}`,
60        { header: 'MCP', options: ['Enable', 'Not now'] })
61    } catch { answer = '' }
62    if (answer !== 'Enable') return { result: `The user did not enable "${name}". Carry on without it, or say what it would be needed for.` }
63    queue({ kind: 'enable', server: name }, String(input.then ?? ''))
64    return { result: `Approved. "${name}" will be enabled as soon as this turn ends. End your turn now; a follow-up message will tell you when its tools are ready.` }
65  })
66
67  on('tool.call', { tool: `${P}reconnect_mcp` }, async ($, e) => {
68    const input = e as unknown as { server?: string; then?: string }
69    const server = String(input.server ?? '').trim()
70    if (!server) return { result: 'Give the server name as /mcp lists it, or "all".' }
71    queue({ kind: 'reconnect', server }, String(input.then ?? ''))
72    return { result: `Reconnect of "${server}" is scheduled for when this turn ends. End your turn now; you will get a follow-up message once it has reconnected, and its new tools will be available then.` }
73  })
74
75  // The OFF servers ride in enable_mcp's own description, and the tool is pinned out of tool
76  // search so Claude always sees it. No per-prompt call: Claude judges what a task needs. The text
77  // is cached with the prompt and changes only when a server is enabled (ui.invalidate below).
78  on('tool.describe', { tool: /^mcp__mcp-control__enable_mcp$/ }, async ($, e, next) => {
79    const r = await next(e)
80    return { ...r, description: enableText(S.off, S.catalog), isDeferred: false }
81  })
82
83  on('turn.complete', async ($, e, next) => {
84    const r = await next(e)
85    if (e.agentId || S.jobs.length === 0) return r
86    const jobs = S.jobs; const then = S.then; const epoch = S.epoch
87    S.jobs = []; S.then = ''
88    S.timer = $.clock.after(50, () => {
89      S.timer = undefined
90      void (async () => {
91        const lines: string[] = []
92        for (const j of jobs) {
93          if (epoch !== S.epoch) return
94          try {
95            const out = (await $.command.run({ command: 'mcp', args: `${j.kind} ${j.server}` })).text ?? ''
96            const state = await servers($)
97            S.off = state.filter(s => s.off)
98            $.ui.invalidate('tool.describe')
99            const connected = (await $.tool.list()).some(t => t.name.startsWith(`mcp__${j.server}__`))
100            const enabled = state.some(s => s.name === j.server && !s.off)
101            const failed = /\b(error|failed|unknown|not found|cannot|unable)\b/i.test(out)
102            lines.push(out.trim() || `Command completed for "${j.server}".`)
103            if (failed || (j.kind === 'enable' && !enabled)) lines.push(`"${j.server}" was not verified enabled. Do not use its tools yet.`)
104            else if (!connected) lines.push(`Tools for "${j.server}" are not verified available yet; check /mcp before continuing with them.`)
105          } catch (err) { lines.push(`${j.kind} ${j.server} failed: ${String(err).slice(0, 200)}`) }
106        }
107        if (epoch !== S.epoch) return
108        await remember($)
109        try {
110          await $.prompt.submit({ text: `MCP ${jobs.map(j => j.kind).join('/')} result: ${lines.join(' ')}${then ? `\nContinue: ${then}` : '\nContinue with the task.'}` })
111        } catch { /* headless run already ended: nothing to continue */ }
112      })()
113    })
114    return r
115  })
116}
117
118function queue(job: Job, then: string): void {
119  if (!S.jobs.some(j => j.kind === job.kind && j.server === job.server)) S.jobs.push(job)
120  if (then) S.then = S.then ? `${S.then}; ${then}` : then
121}
122
123/** The servers Claude Code knows for this project, read fresh (they change as you toggle them). */
124async function servers($: import('claude-code').EngineInterface): Promise<Server[]> {
125  const home = (await $.env.get('HOME')) ?? ''
126  const root = await $.session.root()
127  const json = async (path: string) => { try { return JSON.parse(await $.fs.read(path)) } catch { return {} } }
128  return serversFor(await json(`${home}/.claude.json`), await json(`${root}/.mcp.json`), root)
129}
130
131/** Record connected servers' tool names, so OFF servers can still be described later. */
132async function remember($: import('claude-code').EngineInterface): Promise<void> {
133  try {
134    const next = learn(S.catalog, await $.tool.list())
135    if (JSON.stringify(next) !== JSON.stringify(S.catalog)) { S.catalog = next; await $.store.set('catalog', next) }
136  } catch { /* best effort */ }
137}
138

1/** Pure helpers: which MCP servers exist, which are off in this project, what each one does. */
2
3export type Server = { name: string; scope: 'user' | 'project' | 'mcp.json'; off: boolean; hint: string }
4export type Catalog = Record<string, string[]>
5
6type Json = Record<string, unknown>
7const obj = (v: unknown): Json => (v && typeof v === 'object' && !Array.isArray(v) ? (v as Json) : {})
8const list = (v: unknown): string[] => (Array.isArray(v) ? v.map(String) : [])
9
10/** A short "what is it" from the server's own config: the command or URL, minus noise. */
11function configHint(cfg: unknown): string {
12  const c = obj(cfg)
13  const raw = [c.url, c.command, ...list(c.args)].filter(Boolean).join(' ')
14  return raw.replace(/\/Users\/[^/\s]+/g, '~').replace(/\s+/g, ' ').slice(0, 90)
15}
16
17/** Every server Claude Code would offer in `root`, and whether it is switched off there. */
18export function serversFor(claudeJson: unknown, mcpJson: unknown, root: string): Server[] {
19  const c = obj(claudeJson)
20  const p = obj(obj(c.projects)[root])
21  const off = new Set([...list(p.disabledMcpServers), ...list(p.disabledMcpjsonServers)])
22  const out: Server[] = []
23  const add = (servers: unknown, scope: Server['scope']) => {
24    for (const [name, cfg] of Object.entries(obj(servers))) {
25      if (out.some(s => s.name === name)) continue
26      out.push({ name, scope, off: off.has(name), hint: configHint(cfg) })
27    }
28  }
29  add(p.mcpServers, 'project')
30  add(obj(mcpJson).mcpServers, 'mcp.json')
31  add(c.mcpServers, 'user')
32  return out.sort((a, b) => Number(b.off) - Number(a.off) || a.name.localeCompare(b.name))
33}
34
35/** The tool-name spelling of a server name: `mcp__<this>__tool`. */
36export const toolKey = (name: string) => name.replace(/[^A-Za-z0-9_-]/g, '_')
37
38/** Connected servers' tool names, from `$.tool.list()`, merged into what earlier sessions saw. */
39export function learn(prev: Catalog, tools: readonly { name: string; mcp: boolean }[]): Catalog {
40  const seen: Catalog = {}
41  for (const t of tools) {
42    const m = /^mcp__(.+?)__(.+)$/.exec(t.name)
43    if (!t.mcp || !m) continue
44    ;(seen[m[1]!] ??= []).push(m[2]!)
45  }
46  return { ...prev, ...seen }
47}
48
49/** One line per server: what it does, from the tools seen before, else its config. */
50export function describe(s: Server, catalog: Catalog): string {
51  const tools = catalog[toolKey(s.name)]
52  const what = tools?.length
53    ? `${tools.length} tools: ${tools.slice(0, 8).join(', ')}${tools.length > 8 ? ', …' : ''}`
54    : s.hint || 'no description yet'
55  return `${s.name} [${s.off ? 'OFF' : 'on'}] ${what}`
56}
57
58export function listText(servers: readonly Server[], catalog: Catalog): string {
59  if (servers.length === 0) return 'No MCP servers are configured.'
60  const off = servers.filter(s => s.off)
61  return [
62    `${servers.length} MCP servers configured, ${off.length} switched off in this project.`,
63    ...servers.map(s => '- ' + describe(s, catalog)),
64    off.length ? 'To use an OFF one, call enable_mcp with its name (the user approves with one click).' : '',
65  ].filter(Boolean).join('\n')
66}
67
68/** enable_mcp's description: what it does, then every server that is OFF here and what it offers. */
69export function enableText(off: readonly Server[], catalog: Catalog): string {
70  const head = 'Turn ON an MCP server that is switched OFF in this project. The user approves with one click; it is enabled right after this turn ends and a follow-up message tells you when its tools work. It stays on for this project. Pass server and then (what to do with it).'
71  if (off.length === 0) return `${head}\nEvery configured MCP server is already ON here.`
72  return [head,
73    'When a task needs a service you have no tools for (issue tracker, wiki, browser, logs, CI, code host…), check this list before saying you cannot: if one fits, call this tool.',
74    'OFF in this project:',
75    ...off.map(s => '- ' + describe(s, catalog).replace(' [OFF]', '')),
76  ].join('\n')
77}
78