diff --git a/packages/cli/src/__tests__/runtime-host-capability-provider-command.test.ts b/packages/cli/src/__tests__/runtime-host-capability-provider-command.test.ts index 68c2d10e5b..c059040227 100644 --- a/packages/cli/src/__tests__/runtime-host-capability-provider-command.test.ts +++ b/packages/cli/src/__tests__/runtime-host-capability-provider-command.test.ts @@ -18,11 +18,25 @@ */ import assert from 'node:assert/strict'; +import { readFile } from 'node:fs/promises'; import { test } from 'node:test'; import type { McpToolBinding } from '@maka/core/mcp'; import type { McpClientManager } from '@maka/mcp'; +import { createMcpCapabilityProvider as createPureMcpCapabilityProvider } from '../mcp-capability-provider.js'; import { createMcpCapabilityProvider } from '../runtime-host-capability-provider-command.js'; +test('TUI MCP control keeps its capability provider on the pure import boundary', async () => { + const [tuiSource, providerSource] = await Promise.all([ + readFile(new URL('../tui-mcp-control.js', import.meta.url), 'utf8'), + readFile(new URL('../mcp-capability-provider.js', import.meta.url), 'utf8'), + ]); + + assert.equal(createMcpCapabilityProvider, createPureMcpCapabilityProvider); + assert.match(tuiSource, /from ['"]\.\/mcp-capability-provider\.js['"]/u); + assert.doesNotMatch(tuiSource, /@maka\/runtime-host\/server/u); + assert.doesNotMatch(providerSource, /@maka\/runtime-host\/server/u); +}); + test('MCP capability publication freezes an accepted callable tool snapshot', async () => { let accepted = false; const binding = 'binding-1' as McpToolBinding; diff --git a/packages/cli/src/mcp-capability-provider.ts b/packages/cli/src/mcp-capability-provider.ts new file mode 100644 index 0000000000..d1b8d34608 --- /dev/null +++ b/packages/cli/src/mcp-capability-provider.ts @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { createHash } from 'node:crypto'; +import type { McpBoundTool, McpToolBinding } from '@maka/core/mcp'; +import type { McpClientManager } from '@maka/mcp'; +import type { ClientCapabilityProvider } from '@maka/runtime-host/client'; +import { + CLIENT_CAPABILITY_MAX_TOOLS, + CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER, + decodeClientCapabilityReplaceInput, + type ClientCapabilityCallResult, + type ClientCapabilityOffer, +} from '@maka/runtime-host/protocol'; +import type { McpCallResult, McpToolDescriptor } from '@maka/core/mcp'; + +const CAPABILITY_VERSION = '0'; + +export function createMcpCapabilityProvider( + manager: Pick, +): ClientCapabilityProvider | undefined { + const toolSnapshot = manager.toolSnapshot(); + const tools = [...toolSnapshot.tools].sort( + (left, right) => + left.descriptor.serverId.localeCompare(right.descriptor.serverId) || + left.descriptor.name.localeCompare(right.descriptor.name), + ); + const toolCount = tools.length; + if (toolCount === 0) return undefined; + if (toolCount > CLIENT_CAPABILITY_MAX_TOOLS) { + throw new Error( + `MCP capability provider exposes ${toolCount} tools; the limit is ${CLIENT_CAPABILITY_MAX_TOOLS}`, + ); + } + + const projectedIdentities = new Set(); + const projected: Array<{ + readonly source: McpBoundTool; + readonly descriptor: ReturnType; + }> = []; + for (const source of tools) { + const descriptor = projectMcpTool( + source.descriptor, + capabilityEntityId(source.descriptor.serverId), + ); + const identity = `${descriptor.serverId}\0${descriptor.name}`; + if (projectedIdentities.has(identity)) { + throw new Error('MCP tools collide after Client Capability identity normalization'); + } + projectedIdentities.add(identity); + projected.push({ source, descriptor }); + } + + const bindings = new Map(); + const offers: ClientCapabilityOffer[] = []; + for (let offset = 0; offset < projected.length; offset += CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER) { + const chunk = projected.slice(offset, offset + CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER); + const offerId = mcpOfferId(chunk, offset / CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER); + const servers = new Set(chunk.map(({ source }) => source.descriptor.serverId)); + offers.push({ + offerId, + version: CAPABILITY_VERSION, + affinity: 'session', + hostPathAccess: 'none', + label: + servers.size === 1 + ? `MCP: ${chunk[0]?.source.descriptor.serverId ?? 'tools'}`.slice(0, 128) + : `MCP tools (${servers.size} servers)`, + description: 'Use tools provided by connected MCP servers.', + tools: chunk.map(({ descriptor }) => descriptor), + }); + for (const { source, descriptor } of chunk) { + bindings.set( + capabilityBindingKey(offerId, descriptor.serverId, descriptor.name), + source.binding, + ); + } + } + const canonical = decodeClientCapabilityReplaceInput({ + registrationId: '00000000-0000-4000-8000-000000000000', + offers, + }); + + return { + offers: () => canonical.offers, + call: async (frame, options) => { + const binding = bindings.get( + capabilityBindingKey(frame.offerId, frame.serverId, frame.toolName), + ); + if (!binding) throw new Error('MCP capability is not part of the published snapshot'); + await options.accept(); + return projectMcpResult( + await manager.callTool(binding, frame.arguments, { signal: options.signal }), + ); + }, + }; +} + +function projectMcpTool(tool: McpToolDescriptor, wireServerId: string) { + return { + serverId: wireServerId, + name: capabilityEntityId(tool.name), + ...(tool.description ? { description: tool.description } : {}), + inputSchema: structuredClone(tool.inputSchema), + ...(tool.annotations ? { annotations: { ...tool.annotations } } : {}), + }; +} + +function capabilityEntityId(value: string): string { + if (/^[A-Za-z0-9_-]{1,128}$/u.test(value)) return value; + const label = value.replace(/[^A-Za-z0-9_-]+/gu, '_').slice(0, 103) || 'mcp'; + const digest = createHash('sha256').update(value).digest('hex').slice(0, 24); + return `${label}_${digest}`; +} + +function projectMcpResult(result: McpCallResult): ClientCapabilityCallResult { + return { + content: result.content.map((block) => structuredClone(block)), + ...(result.structuredContent === undefined + ? {} + : { structuredContent: structuredClone(result.structuredContent) }), + }; +} + +function mcpOfferId(tools: readonly { readonly source: McpBoundTool }[], chunk: number): string { + const hash = createHash('sha256').update(`mcp-capability-offer-v0\0${chunk}`); + for (const { source } of tools) + hash + .update('\0') + .update(source.descriptor.serverId) + .update('\0') + .update(source.descriptor.name); + return `mcp_${hash.digest('hex').slice(0, 24)}_${chunk}`; +} + +function capabilityBindingKey(offerId: string, serverId: string, toolName: string): string { + return `${offerId}\0${serverId}\0${toolName}`; +} diff --git a/packages/cli/src/runtime-host-capability-provider-command.ts b/packages/cli/src/runtime-host-capability-provider-command.ts index ce68ed2d0c..3b9a320469 100644 --- a/packages/cli/src/runtime-host-capability-provider-command.ts +++ b/packages/cli/src/runtime-host-capability-provider-command.ts @@ -21,12 +21,6 @@ import { createHash } from 'node:crypto'; import { readFile } from 'node:fs/promises'; import { homedir } from 'node:os'; import { dirname, join, resolve } from 'node:path'; -import type { - McpBoundTool, - McpCallResult, - McpToolBinding, - McpToolDescriptor, -} from '@maka/core/mcp'; import { createCredentialMcpOAuthStorage, McpClientManager } from '@maka/mcp'; import { createFileCredentialStore } from '@maka/storage/credential-store'; import { normalizeMcpConfig } from '@maka/storage/mcp-config-store'; @@ -36,20 +30,15 @@ import { remoteRuntimeHostUnavailableError, RuntimeHostPermanentReconnectError, startRuntimeHostCapabilityProviderService, - type ClientCapabilityProvider, type RuntimeHostConnection, } from '@maka/runtime-host/client'; import { - CLIENT_CAPABILITY_MAX_TOOLS, - CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER, - decodeClientCapabilityReplaceInput, INTERACTIVE_RUNTIME_HOST_COMPOSITION_ID, RUNTIME_HOST_COMPATIBILITY_EPOCH, RUNTIME_HOST_PROTOCOL_VERSION, - type ClientCapabilityCallResult, - type ClientCapabilityOffer, } from '@maka/runtime-host/protocol'; import { runRuntimeHostProcessLifecycle } from '@maka/runtime-host/server'; +import { createMcpCapabilityProvider } from './mcp-capability-provider.js'; const DEFAULT_CREDENTIAL_ENV = 'MAKA_RUNTIME_HOST_ACCESS_CREDENTIAL'; const CAPABILITY_VERSION = '0'; @@ -199,126 +188,7 @@ async function connectRemoteCapabilityProvider(input: { throw new Error('Runtime Host is draining'); } -export function createMcpCapabilityProvider( - manager: Pick, -): ClientCapabilityProvider | undefined { - const toolSnapshot = manager.toolSnapshot(); - const tools = [...toolSnapshot.tools].sort( - (left, right) => - left.descriptor.serverId.localeCompare(right.descriptor.serverId) || - left.descriptor.name.localeCompare(right.descriptor.name), - ); - const toolCount = tools.length; - if (toolCount === 0) return undefined; - if (toolCount > CLIENT_CAPABILITY_MAX_TOOLS) { - throw new Error( - `MCP capability provider exposes ${toolCount} tools; the limit is ${CLIENT_CAPABILITY_MAX_TOOLS}`, - ); - } - - const projectedIdentities = new Set(); - const projected: Array<{ - readonly source: McpBoundTool; - readonly descriptor: ReturnType; - }> = []; - for (const source of tools) { - const descriptor = projectMcpTool( - source.descriptor, - capabilityEntityId(source.descriptor.serverId), - ); - const identity = `${descriptor.serverId}\0${descriptor.name}`; - if (projectedIdentities.has(identity)) { - throw new Error('MCP tools collide after Client Capability identity normalization'); - } - projectedIdentities.add(identity); - projected.push({ source, descriptor }); - } - - const bindings = new Map(); - const offers: ClientCapabilityOffer[] = []; - for (let offset = 0; offset < projected.length; offset += CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER) { - const chunk = projected.slice(offset, offset + CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER); - const offerId = mcpOfferId(chunk, offset / CLIENT_CAPABILITY_MAX_TOOLS_PER_OFFER); - const servers = new Set(chunk.map(({ source }) => source.descriptor.serverId)); - offers.push({ - offerId, - version: CAPABILITY_VERSION, - affinity: 'session', - hostPathAccess: 'none', - label: - servers.size === 1 - ? `MCP: ${chunk[0]?.source.descriptor.serverId ?? 'tools'}`.slice(0, 128) - : `MCP tools (${servers.size} servers)`, - description: 'Use tools provided by connected MCP servers.', - tools: chunk.map(({ descriptor }) => descriptor), - }); - for (const { source, descriptor } of chunk) { - bindings.set( - capabilityBindingKey(offerId, descriptor.serverId, descriptor.name), - source.binding, - ); - } - } - const canonical = decodeClientCapabilityReplaceInput({ - registrationId: '00000000-0000-4000-8000-000000000000', - offers, - }); - - return { - offers: () => canonical.offers, - call: async (frame, options) => { - const binding = bindings.get( - capabilityBindingKey(frame.offerId, frame.serverId, frame.toolName), - ); - if (!binding) throw new Error('MCP capability is not part of the published snapshot'); - await options.accept(); - return projectMcpResult( - await manager.callTool(binding, frame.arguments, { signal: options.signal }), - ); - }, - }; -} - -function projectMcpTool(tool: McpToolDescriptor, wireServerId: string) { - return { - serverId: wireServerId, - name: capabilityEntityId(tool.name), - ...(tool.description ? { description: tool.description } : {}), - inputSchema: structuredClone(tool.inputSchema), - ...(tool.annotations ? { annotations: { ...tool.annotations } } : {}), - }; -} - -function capabilityEntityId(value: string): string { - if (/^[A-Za-z0-9_-]{1,128}$/u.test(value)) return value; - const label = value.replace(/[^A-Za-z0-9_-]+/gu, '_').slice(0, 103) || 'mcp'; - const digest = createHash('sha256').update(value).digest('hex').slice(0, 24); - return `${label}_${digest}`; -} - -function projectMcpResult(result: McpCallResult): ClientCapabilityCallResult { - return { - content: result.content.map((block) => structuredClone(block)), - ...(result.structuredContent === undefined - ? {} - : { structuredContent: structuredClone(result.structuredContent) }), - }; -} - -function mcpOfferId(tools: readonly { readonly source: McpBoundTool }[], chunk: number): string { - const hash = createHash('sha256').update(`mcp-capability-offer-v0\0${chunk}`); - for (const { source } of tools) - hash - .update('\0') - .update(source.descriptor.serverId) - .update('\0') - .update(source.descriptor.name); - return `mcp_${hash.digest('hex').slice(0, 24)}_${chunk}`; -} - -function capabilityBindingKey(offerId: string, serverId: string, toolName: string): string { - return `${offerId}\0${serverId}\0${toolName}`; -} +export { createMcpCapabilityProvider } from './mcp-capability-provider.js'; function reportRefreshFailure(error: unknown): void { const message = error instanceof Error ? error.message : String(error); diff --git a/packages/cli/src/tui-mcp-control.ts b/packages/cli/src/tui-mcp-control.ts index 2fb3dcd49d..ff91bd6c6c 100644 --- a/packages/cli/src/tui-mcp-control.ts +++ b/packages/cli/src/tui-mcp-control.ts @@ -26,7 +26,7 @@ import type { RuntimeHostConnectionAvailability, RuntimeHostReconnectingConnection, } from '@maka/runtime-host/client'; -import { createMcpCapabilityProvider } from './runtime-host-capability-provider-command.js'; +import { createMcpCapabilityProvider } from './mcp-capability-provider.js'; const RUNTIME_HOST_CREDENTIAL_ENV = 'MAKA_RUNTIME_HOST_ACCESS_CREDENTIAL';