From 69d4630c54b6a5fc4242c9232b69b446518dd1a2 Mon Sep 17 00:00:00 2001 From: im10furry Date: Sat, 25 Jul 2026 23:32:18 +0800 Subject: [PATCH 1/2] feat(mcp): implement sampling/createMessage capability - Add sampling.ts with CreateMessage request handler that routes through queryLLM for multi-provider LLM support - Declare sampling capability in getMcpClientCapabilities() so MCP servers know the client supports sampling - Register/unregister sampling handler in client lifecycle - Convert between MCP SamplingMessage and internal Anthropic format - Resolve model from server's modelPreferences hints - Map stop_reason between internal and MCP spec formats - Export sampling control APIs (enable/disable/test helpers) - Update capability tests for sampling enabled/disabled states --- packages/core/src/mcp/client/connection.ts | 6 + packages/core/src/mcp/client/index.ts | 6 + packages/core/src/mcp/client/roots.ts | 14 +- packages/core/src/mcp/client/sampling.ts | 377 ++++++++++++++++++ .../test/unit/mcp-client-capabilities.test.ts | 26 +- 5 files changed, 425 insertions(+), 4 deletions(-) create mode 100644 packages/core/src/mcp/client/sampling.ts diff --git a/packages/core/src/mcp/client/connection.ts b/packages/core/src/mcp/client/connection.ts index 70065a362..b6b03b5a6 100644 --- a/packages/core/src/mcp/client/connection.ts +++ b/packages/core/src/mcp/client/connection.ts @@ -33,6 +33,10 @@ import { registerMcpClientRequestHandlers, unregisterMcpClientRequestHandlers, } from './roots' +import { + registerMcpSamplingHandler, + unregisterMcpSamplingHandler, +} from './sampling' import type { WrappedClient } from './types' type GlobalWithWebSocket = { WebSocket?: unknown } @@ -101,6 +105,7 @@ export function createMcpClient( createMcpClientSdkOptions(name), ) registerMcpClientRequestHandlers(client) + registerMcpSamplingHandler(client) client.setNotificationHandler( LoggingMessageNotificationSchema, notification => { @@ -121,6 +126,7 @@ export function createMcpClient( export async function closeMcpClient(client: Client): Promise { unregisterMcpClientRequestHandlers(client) + unregisterMcpSamplingHandler(client) try { await client.close() } catch {} diff --git a/packages/core/src/mcp/client/index.ts b/packages/core/src/mcp/client/index.ts index 7cd9e419f..157222def 100644 --- a/packages/core/src/mcp/client/index.ts +++ b/packages/core/src/mcp/client/index.ts @@ -76,3 +76,9 @@ export { type McpLogMessageEvent, type McpLoggingLevel, } from './logging' +export { + __resetMcpSamplingForTests, + __setMcpSamplingEnabledForTests, + isMcpSamplingEnabled, + setMcpSamplingEnabled, +} from './sampling' diff --git a/packages/core/src/mcp/client/roots.ts b/packages/core/src/mcp/client/roots.ts index d1c5fecd1..98ed8dbbe 100644 --- a/packages/core/src/mcp/client/roots.ts +++ b/packages/core/src/mcp/client/roots.ts @@ -11,6 +11,7 @@ import { import { checkHasTrustDialogAccepted } from '#core/utils/config' import { logMCPError } from '#core/utils/log' import { getCwd, subscribeCwdChanged } from '#core/utils/state' +import { isMcpSamplingEnabled } from './sampling' let exposeRootsOverrideForTests: boolean | null = null const rootsClients = new Set() @@ -43,10 +44,17 @@ export function shouldExposeMcpRoots(): boolean { } export function getMcpClientCapabilities(): ClientCapabilities { - if (!shouldExposeMcpRoots()) return {} - return { - roots: { listChanged: true }, + const capabilities: ClientCapabilities = {} + + if (shouldExposeMcpRoots()) { + capabilities.roots = { listChanged: true } + } + + if (isMcpSamplingEnabled()) { + capabilities.sampling = {} } + + return capabilities } function ensureCwdChangedSubscription(): void { diff --git a/packages/core/src/mcp/client/sampling.ts b/packages/core/src/mcp/client/sampling.ts new file mode 100644 index 000000000..5479acbb1 --- /dev/null +++ b/packages/core/src/mcp/client/sampling.ts @@ -0,0 +1,377 @@ +/** + * MCP Sampling capability implementation. + * + * When an MCP server sends a `sampling/createMessage` request, this module + * handles it by routing through the local LLM infrastructure (queryLLM). + * + * The MCP spec notes: "The client has full discretion over which model to + * select. The client should also inform the user before beginning sampling, + * to allow them to inspect the request (human in the loop)." + */ +import { randomUUID } from 'node:crypto' +import type { UUID } from 'node:crypto' + +import type { Client } from '@modelcontextprotocol/sdk/client/index.js' +import { CreateMessageRequestSchema } from '@modelcontextprotocol/sdk/types.js' + +import type { MessageParam } from '@anthropic-ai/sdk/resources/index.mjs' +import type { UserMessage, AssistantMessage } from '#core/query' +import { queryLLM } from '#core/ai/llm' +import { getModelManager } from '#core/utils/model' +import { logMCPError } from '#core/utils/log' +import { createAnthropicUsage } from '@kode/protocol/anthropic' + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +type SamplingMessage = { + role: 'user' | 'assistant' + content: + | SamplingContentBlock + | SamplingContentBlock[] +} + +type SamplingContentBlock = + | { type: 'text'; text: string } + | { type: 'image'; data: string; mimeType: string } + | { type: 'audio'; data: string; mimeType: string } + | { type: 'tool_use'; id: string; name: string; input: unknown } + | { type: 'tool_result'; toolUseId: string; content: unknown; isError?: boolean } + +type CreateMessageParams = { + messages: SamplingMessage[] + modelPreferences?: { + hints?: Array<{ name?: string }> + costPriority?: number + speedPriority?: number + intelligencePriority?: number + } + systemPrompt?: string + includeContext?: 'none' | 'thisServer' | 'allServers' + temperature?: number + maxTokens: number + stopSequences?: string[] + metadata?: Record + tools?: Array<{ + name: string + description?: string + inputSchema: Record + }> + toolChoice?: { mode: string } | { mode: 'tool'; name: string } +} + +type CreateMessageResult = { + model: string + stopReason?: string + role: 'assistant' + content: { type: 'text'; text: string } | { type: 'image'; data: string; mimeType: string } +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/** Model pointer used for sampling requests. Defaults to "quick" for fast responses. */ +const SAMPLING_MODEL_POINTER = 'quick' + +let samplingEnabled = true +let samplingEnabledOverrideForTests: boolean | null = null + +const samplingClients = new Set() + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export function isMcpSamplingEnabled(): boolean { + if ( + process.env.NODE_ENV === 'test' && + samplingEnabledOverrideForTests !== null + ) { + return samplingEnabledOverrideForTests + } + return samplingEnabled +} + +export function setMcpSamplingEnabled(enabled: boolean): void { + samplingEnabled = enabled +} + +/** + * Register the sampling/createMessage request handler on the given MCP client. + * This should be called during client initialization (alongside roots). + */ +export function registerMcpSamplingHandler(client: Client): void { + if (!isMcpSamplingEnabled()) return + + client.setRequestHandler( + CreateMessageRequestSchema, + async (request, _extra) => { + const params = request.params as CreateMessageParams + return await handleCreateMessage(params) + }, + ) + + samplingClients.add(client) +} + +/** + * Unregister the sampling request handler from the given MCP client. + */ +export function unregisterMcpSamplingHandler(client: Client): void { + const wasRegistered = samplingClients.delete(client) + if (wasRegistered) { + const clientWithRemove = client as Client & { + removeRequestHandler?: (method: string) => void + } + clientWithRemove.removeRequestHandler?.('sampling/createMessage') + } +} + +// --------------------------------------------------------------------------- +// Core Handler +// --------------------------------------------------------------------------- + +async function handleCreateMessage( + params: CreateMessageParams, +): Promise { + const { messages, systemPrompt, temperature, maxTokens, stopSequences } = + params + + // Convert MCP sampling messages to internal message format + const internalMessages = convertSamplingMessages(messages) + + // Resolve system prompt + const system = systemPrompt ? [systemPrompt] : [] + + // Resolve model - we use the quick model pointer for sampling by default, + // but respect modelPreferences hints if a matching model is configured. + const modelPointer = resolveModelFromPreferences(params.modelPreferences) + + // Create an abort controller for this sampling request + const abortController = new AbortController() + + try { + const result = await queryLLM( + internalMessages, + system, + 0, // no thinking tokens for sampling + [], // no tools for now (basic sampling) + abortController.signal, + { + safeMode: false, + model: modelPointer, + prependCLISysprompt: false, + temperature: temperature ?? undefined, + maxTokens, + stopSequences, + }, + ) + + return convertToSamplingResult(result) + } catch (error) { + logMCPError( + 'sampling', + `Failed to handle createMessage: ${error instanceof Error ? error.message : String(error)}`, + ) + throw error + } +} + +// --------------------------------------------------------------------------- +// Message Conversion: MCP Sampling → Internal Format +// --------------------------------------------------------------------------- + +function convertSamplingMessages( + messages: SamplingMessage[], +): (UserMessage | AssistantMessage)[] { + const result: (UserMessage | AssistantMessage)[] = [] + + for (const msg of messages) { + const contentBlocks = normalizeContent(msg.content) + + if (msg.role === 'user') { + result.push(convertToUserMessage(contentBlocks)) + } else if (msg.role === 'assistant') { + result.push(convertToAssistantMessage(contentBlocks)) + } + } + + return result +} + +function normalizeContent( + content: SamplingContentBlock | SamplingContentBlock[], +): SamplingContentBlock[] { + return Array.isArray(content) ? content : [content] +} + +function convertToUserMessage( + blocks: SamplingContentBlock[], +): UserMessage { + const anthropicContent: MessageParam['content'] = blocks.map(block => { + switch (block.type) { + case 'text': + return { type: 'text' as const, text: block.text } + case 'image': + return { + type: 'image' as const, + source: { + type: 'base64' as const, + media_type: block.mimeType as + | 'image/jpeg' + | 'image/png' + | 'image/gif' + | 'image/webp', + data: block.data, + }, + } + default: + // For unsupported block types, convert to text representation + return { type: 'text' as const, text: JSON.stringify(block) } + } + }) + + return { + message: { role: 'user', content: anthropicContent }, + type: 'user', + uuid: randomUUID() as UUID, + } +} + +function convertToAssistantMessage( + blocks: SamplingContentBlock[], +): AssistantMessage { + const content = blocks.map(block => { + switch (block.type) { + case 'text': + return { type: 'text' as const, text: block.text } + default: + return { type: 'text' as const, text: JSON.stringify(block) } + } + }) + + return { + costUSD: 0, + durationMs: 0, + message: { + id: `msg_sampling_${randomUUID()}`, + model: 'unknown', + role: 'assistant', + type: 'message', + content, + usage: createAnthropicUsage(), + stop_reason: null, + }, + type: 'assistant', + uuid: randomUUID() as UUID, + } +} + +// --------------------------------------------------------------------------- +// Result Conversion: Internal Format → MCP Sampling Result +// --------------------------------------------------------------------------- + +function convertToSamplingResult( + assistantMessage: AssistantMessage, +): CreateMessageResult { + const model = assistantMessage.message.model || 'unknown' + const stopReason = mapStopReason(assistantMessage.message.stop_reason) + + // Extract text content from the response + const textContent = extractTextContent(assistantMessage.message.content) + + return { + model, + stopReason, + role: 'assistant', + content: { type: 'text', text: textContent }, + } +} + +function mapStopReason( + stopReason: string | null | undefined, +): string | undefined { + if (!stopReason) return undefined + + switch (stopReason) { + case 'end_turn': + return 'endTurn' + case 'stop_sequence': + return 'stopSequence' + case 'max_tokens': + return 'maxTokens' + case 'tool_use': + return 'toolUse' + default: + return stopReason + } +} + +function extractTextContent(content: any[]): string { + if (!Array.isArray(content)) return '' + + const textParts: string[] = [] + for (const block of content) { + if (block && typeof block === 'object' && block.type === 'text') { + textParts.push(block.text || '') + } + } + + return textParts.join('\n') +} + +// --------------------------------------------------------------------------- +// Model Preferences Resolution +// --------------------------------------------------------------------------- + +function resolveModelFromPreferences( + preferences?: CreateMessageParams['modelPreferences'], +): string { + if (!preferences?.hints?.length) return SAMPLING_MODEL_POINTER + + // Try to match model hints against configured models + const modelManager = getModelManager() + + for (const hint of preferences.hints) { + if (!hint.name) continue + + // Check if the hint matches any configured model name directly + const resolved = modelManager.resolveModel(hint.name) + if (resolved) return hint.name + } + + // If speed is prioritized, use "quick" model + if ( + preferences.speedPriority && + preferences.speedPriority > (preferences.intelligencePriority ?? 0) + ) { + return 'quick' + } + + // If intelligence is prioritized, use "main" model + if ( + preferences.intelligencePriority && + preferences.intelligencePriority > (preferences.speedPriority ?? 0) + ) { + return 'main' + } + + return SAMPLING_MODEL_POINTER +} + +// --------------------------------------------------------------------------- +// Test helpers +// --------------------------------------------------------------------------- + +export function __setMcpSamplingEnabledForTests( + value: boolean | null, +): void { + samplingEnabledOverrideForTests = value +} + +export function __resetMcpSamplingForTests(): void { + samplingEnabledOverrideForTests = null + samplingClients.clear() +} diff --git a/packages/core/src/test/unit/mcp-client-capabilities.test.ts b/packages/core/src/test/unit/mcp-client-capabilities.test.ts index 5ddd019ba..45eea5c77 100644 --- a/packages/core/src/test/unit/mcp-client-capabilities.test.ts +++ b/packages/core/src/test/unit/mcp-client-capabilities.test.ts @@ -9,14 +9,38 @@ import { __resetMcpRootsForTests, __setMcpRootsTrustOverrideForTests, } from '#core/mcp/client/roots' +import { + __resetMcpSamplingForTests, + __setMcpSamplingEnabledForTests, +} from '#core/mcp/client/sampling' describe('MCP client capability summary', () => { afterEach(() => { __resetMcpRootsForTests() + __resetMcpSamplingForTests() + }) + + test('summarizes trusted root capability exposure with sampling enabled', () => { + __setMcpRootsTrustOverrideForTests(true) + __setMcpSamplingEnabledForTests(true) + + expect(getMcpClientCapabilitySummary()).toEqual({ + roots: { enabled: true, listChanged: true }, + sampling: { enabled: true, context: false, tools: false }, + elicitation: { enabled: false, form: false, url: false }, + tasks: { + enabled: false, + list: false, + cancel: false, + samplingCreateMessage: false, + elicitationCreate: false, + }, + }) }) - test('summarizes trusted root capability exposure', () => { + test('summarizes capabilities when sampling is disabled', () => { __setMcpRootsTrustOverrideForTests(true) + __setMcpSamplingEnabledForTests(false) expect(getMcpClientCapabilitySummary()).toEqual({ roots: { enabled: true, listChanged: true }, From 992515c4f51fd100672c964b58271fb9a56c5798 Mon Sep 17 00:00:00 2001 From: im10furry Date: Sun, 26 Jul 2026 02:28:36 +0800 Subject: [PATCH 2/2] test(mcp): align client capability tests with sampling enabled by default --- packages/core/src/test/unit/mcp-cli-complete.test.ts | 9 ++++++++- .../core/src/test/unit/mcp-connection-internals.test.ts | 6 ++++++ packages/core/src/test/unit/mcp-roots.test.ts | 8 +++++++- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/packages/core/src/test/unit/mcp-cli-complete.test.ts b/packages/core/src/test/unit/mcp-cli-complete.test.ts index 938d6d99e..3ec465fab 100644 --- a/packages/core/src/test/unit/mcp-cli-complete.test.ts +++ b/packages/core/src/test/unit/mcp-cli-complete.test.ts @@ -5,6 +5,10 @@ import { __resetMcpRootsForTests, __setMcpRootsTrustOverrideForTests, } from '#core/mcp/client/roots' +import { + __resetMcpSamplingForTests, + __setMcpSamplingEnabledForTests, +} from '#core/mcp/client/sampling' import { runMcpCli } from '#host-cli/entrypoints/mcpCli' async function captureMcpCli(argv: string[]): Promise<{ @@ -143,6 +147,7 @@ describe('mcp-cli client-capabilities', () => { afterEach(() => { __setMcpClientsForTests(null) __resetMcpRootsForTests() + __resetMcpSamplingForTests() }) test('prints client capabilities as JSON', async () => { @@ -154,7 +159,8 @@ describe('mcp-cli client-capabilities', () => { expect(result.stderr).toBe('') expect(JSON.parse(result.stdout)).toEqual({ roots: { enabled: true, listChanged: true }, - sampling: { enabled: false, context: false, tools: false }, + // Sampling defaults to enabled (sampling/createMessage support). + sampling: { enabled: true, context: false, tools: false }, elicitation: { enabled: false, form: false, url: false }, tasks: { enabled: false, @@ -168,6 +174,7 @@ describe('mcp-cli client-capabilities', () => { test('prints disabled client capabilities in text output', async () => { __setMcpRootsTrustOverrideForTests(false) + __setMcpSamplingEnabledForTests(false) const result = await captureMcpCli(['client-capabilities']) diff --git a/packages/core/src/test/unit/mcp-connection-internals.test.ts b/packages/core/src/test/unit/mcp-connection-internals.test.ts index 082743727..94fb0f4a0 100644 --- a/packages/core/src/test/unit/mcp-connection-internals.test.ts +++ b/packages/core/src/test/unit/mcp-connection-internals.test.ts @@ -10,6 +10,10 @@ import { __resetMcpRootsForTests, __setMcpRootsTrustOverrideForTests, } from '#core/mcp/client/roots' +import { + __resetMcpSamplingForTests, + __setMcpSamplingEnabledForTests, +} from '#core/mcp/client/sampling' import { MACRO } from '#core/constants/macros' import { PRODUCT_COMMAND } from '#core/constants/product' import { @@ -39,6 +43,7 @@ describe('MCP connection internals', () => { else process.env.MCP_CONNECTION_TIMEOUT_MS = originalTimeout __resetMcpRootsForTests() + __resetMcpSamplingForTests() __resetMcpListChangedForTests() clearNotifications() }) @@ -104,6 +109,7 @@ describe('MCP connection internals', () => { }) test('builds SDK options with roots and list_changed refresh hooks', () => { + __setMcpSamplingEnabledForTests(false) __setMcpRootsTrustOverrideForTests(true) const options = createMcpClientSdkOptions('srv') as any diff --git a/packages/core/src/test/unit/mcp-roots.test.ts b/packages/core/src/test/unit/mcp-roots.test.ts index 0f023b412..775d369e1 100644 --- a/packages/core/src/test/unit/mcp-roots.test.ts +++ b/packages/core/src/test/unit/mcp-roots.test.ts @@ -14,6 +14,10 @@ import { registerMcpClientRequestHandlers, unregisterMcpClientRequestHandlers, } from '#core/mcp/client/roots' +import { + __resetMcpSamplingForTests, + __setMcpSamplingEnabledForTests, +} from '#core/mcp/client/sampling' import { __resetCwdChangedListenersForTests, getCwd, @@ -23,6 +27,7 @@ import { describe('MCP client roots', () => { afterEach(() => { __resetMcpRootsForTests() + __resetMcpSamplingForTests() __resetCwdChangedListenersForTests() }) @@ -34,6 +39,7 @@ describe('MCP client roots', () => { }) test('declares roots capability only for trusted workspaces', () => { + __setMcpSamplingEnabledForTests(false) __setMcpRootsTrustOverrideForTests(false) expect(getMcpClientCapabilities()).toEqual({}) @@ -59,7 +65,7 @@ describe('MCP client roots', () => { } as any) expect(handler).not.toBeNull() - expect(await handler?.()).toEqual({ + expect(await (handler as any)?.()).toEqual({ roots: createMcpRootsForCwd(projectDir), }) } finally {