Skip to content

Commit 652f19a

Browse files
daniel-lxsdaniel-lxs
committed
fix: properly finalize MCP tool calls in native protocol mode
- Simplify MCP tool schema to pass arguments directly (no toolInputProps wrapper) - Extract server_name/tool_name from function name (mcp_serverName_toolName) - Add finalizeRawChunks() call after stream ends to properly convert MCP tools - Add dynamic MCP tool validation against mcp group in mode permissions - Fix NativeToolCallParser to support string names for dynamic MCP tools
1 parent bfecce5 commit 652f19a

File tree

5 files changed

+121
-39
lines changed

5 files changed

+121
-39
lines changed

‎src/core/assistant-message/NativeToolCallParser.ts‎

Lines changed: 36 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,12 @@ export type ToolCallStreamEvent = ApiStreamToolCallStartChunk | ApiStreamToolCal
4747
*/
4848
export class NativeToolCallParser {
4949
// Streaming state management for argument accumulation (keyed by tool call id)
50+
// Note: name is string to accommodate dynamic MCP tools (mcp_serverName_toolName)
5051
private static streamingToolCalls = new Map<
5152
string,
5253
{
5354
id: string
54-
name: ToolName
55+
name: string
5556
argumentsAccumulator: string
5657
}
5758
>()
@@ -194,8 +195,9 @@ export class NativeToolCallParser {
194195
/**
195196
* Start streaming a new tool call.
196197
* Initializes tracking for incremental argument parsing.
198+
* Accepts string to support both ToolName and dynamic MCP tools (mcp_serverName_toolName).
197199
*/
198-
public static startStreamingToolCall(id: string, name: ToolName): void {
200+
public static startStreamingToolCall(id: string, name: string): void {
199201
this.streamingToolCalls.set(id, {
200202
id,
201203
name,
@@ -235,6 +237,11 @@ export class NativeToolCallParser {
235237
// Accumulate the JSON string
236238
toolCall.argumentsAccumulator += chunk
237239

240+
// For dynamic MCP tools, we don't return partial updates - wait for final
241+
if (toolCall.name.startsWith("mcp_")) {
242+
return null
243+
}
244+
238245
// Parse whatever we can from the incomplete JSON!
239246
// partial-json-parser extracts partial values (strings, arrays, objects) immediately
240247
try {
@@ -243,7 +250,7 @@ export class NativeToolCallParser {
243250
// Create partial ToolUse with extracted values
244251
return this.createPartialToolUse(
245252
toolCall.id,
246-
toolCall.name,
253+
toolCall.name as ToolName,
247254
partialArgs || {},
248255
true, // partial
249256
)
@@ -266,9 +273,10 @@ export class NativeToolCallParser {
266273
}
267274

268275
// Parse the complete accumulated JSON
276+
// Cast to any for the name since parseToolCall handles both ToolName and dynamic MCP tools
269277
const finalToolUse = this.parseToolCall({
270278
id: toolCall.id,
271-
name: toolCall.name,
279+
name: toolCall.name as ToolName,
272280
arguments: toolCall.argumentsAccumulator,
273281
})
274282

@@ -592,6 +600,15 @@ export class NativeToolCallParser {
592600
}
593601
break
594602

603+
case "access_mcp_resource":
604+
if (args.server_name !== undefined && args.uri !== undefined) {
605+
nativeArgs = {
606+
server_name: args.server_name,
607+
uri: args.uri,
608+
} as NativeArgsFor<TName>
609+
}
610+
break
611+
595612
default:
596613
break
597614
}
@@ -623,16 +640,23 @@ export class NativeToolCallParser {
623640
*/
624641
public static parseDynamicMcpTool(toolCall: { id: string; name: string; arguments: string }): McpToolUse | null {
625642
try {
626-
const args = JSON.parse(toolCall.arguments)
643+
// Parse the arguments - these are the actual tool arguments passed directly
644+
const args = JSON.parse(toolCall.arguments || "{}")
645+
646+
// Extract server_name and tool_name from the tool name itself
647+
// Format: mcp_serverName_toolName
648+
const nameParts = toolCall.name.split("_")
649+
if (nameParts.length < 3 || nameParts[0] !== "mcp") {
650+
console.error(`Invalid dynamic MCP tool name format: ${toolCall.name}`)
651+
return null
652+
}
627653

628-
// Extract server_name and tool_name from the arguments
629-
// The dynamic tool schema includes these as const properties
630-
const serverName = args.server_name
631-
const toolName = args.tool_name
632-
const toolInputProps = args.toolInputProps || {}
654+
// Server name is the second part, tool name is everything after
655+
const serverName = nameParts[1]
656+
const toolName = nameParts.slice(2).join("_")
633657

634658
if (!serverName || !toolName) {
635-
console.error(`Missing server_name or tool_name in dynamic MCP tool`)
659+
console.error(`Could not extract server_name or tool_name from: ${toolCall.name}`)
636660
return null
637661
}
638662

@@ -643,7 +667,7 @@ export class NativeToolCallParser {
643667
name: toolCall.name,
644668
serverName,
645669
toolName,
646-
arguments: toolInputProps,
670+
arguments: args,
647671
partial: false,
648672
}
649673

‎src/core/prompts/tools/native-tools/mcp_server.ts‎

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -29,37 +29,21 @@ export function getMcpServerTools(mcpHub?: McpHub): OpenAI.Chat.ChatCompletionTo
2929
const toolInputProps = originalSchema?.properties ?? {}
3030
const toolInputRequired = (originalSchema?.required ?? []) as string[]
3131

32-
// Create a proper JSON Schema object for toolInputProps
33-
const toolInputPropsSchema: Record<string, any> = {
32+
// Build parameters directly from the tool's input schema.
33+
// The server_name and tool_name are encoded in the function name itself
34+
// (e.g., mcp_serverName_toolName), so they don't need to be in the arguments.
35+
const parameters: OpenAI.FunctionParameters = {
3436
type: "object",
3537
properties: toolInputProps,
3638
additionalProperties: false,
3739
}
3840

3941
// Only add required if there are required fields
4042
if (toolInputRequired.length > 0) {
41-
toolInputPropsSchema.required = toolInputRequired
43+
parameters.required = toolInputRequired
4244
}
4345

44-
// Build parameters with all properties defined before adding required array
45-
const parameters = {
46-
type: "object",
47-
properties: {
48-
toolInputProps: toolInputPropsSchema,
49-
server_name: {
50-
type: "string",
51-
const: server.name,
52-
},
53-
tool_name: {
54-
type: "string",
55-
const: tool.name,
56-
},
57-
},
58-
required: ["server_name", "tool_name", "toolInputProps"],
59-
additionalProperties: false,
60-
} as OpenAI.FunctionParameters
61-
62-
// Use triple underscores as separator to allow underscores in tool and server names
46+
// Use mcp_ prefix to identify dynamic MCP tools
6347
const toolDefinition: OpenAI.Chat.ChatCompletionTool = {
6448
type: "function",
6549
function: {

‎src/core/task/Task.ts‎

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2538,6 +2538,37 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
25382538
}
25392539
}
25402540

2541+
// Finalize any remaining streaming tool calls that weren't explicitly ended
2542+
// This is critical for MCP tools which need tool_call_end events to be properly
2543+
// converted from ToolUse to McpToolUse via finalizeStreamingToolCall()
2544+
const finalizeEvents = NativeToolCallParser.finalizeRawChunks()
2545+
for (const event of finalizeEvents) {
2546+
if (event.type === "tool_call_end") {
2547+
// Finalize the streaming tool call
2548+
const finalToolUse = NativeToolCallParser.finalizeStreamingToolCall(event.id)
2549+
2550+
if (finalToolUse) {
2551+
// Store the tool call ID
2552+
;(finalToolUse as any).id = event.id
2553+
2554+
// Get the index and replace partial with final
2555+
const toolUseIndex = this.streamingToolCallIndices.get(event.id)
2556+
if (toolUseIndex !== undefined) {
2557+
this.assistantMessageContent[toolUseIndex] = finalToolUse
2558+
}
2559+
2560+
// Clean up tracking
2561+
this.streamingToolCallIndices.delete(event.id)
2562+
2563+
// Mark that we have new content to process
2564+
this.userMessageContentReady = false
2565+
2566+
// Present the finalized tool call
2567+
presentAssistantMessage(this)
2568+
}
2569+
}
2570+
}
2571+
25412572
// Create a copy of current token values to avoid race conditions
25422573
const currentTokens = {
25432574
input: inputTokens,
@@ -2879,17 +2910,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
28792910
for (const block of toolUseBlocks) {
28802911
if (block.type === "mcp_tool_use") {
28812912
// McpToolUse already has the original tool name (e.g., "mcp_serverName_toolName")
2913+
// The arguments are the raw tool arguments (matching the simplified schema)
28822914
const mcpBlock = block as import("../../shared/tools").McpToolUse
28832915
if (mcpBlock.id) {
28842916
assistantContent.push({
28852917
type: "tool_use" as const,
28862918
id: mcpBlock.id,
28872919
name: mcpBlock.name, // Original dynamic name
2888-
input: {
2889-
server_name: mcpBlock.serverName,
2890-
tool_name: mcpBlock.toolName,
2891-
toolInputProps: mcpBlock.arguments,
2892-
},
2920+
input: mcpBlock.arguments, // Direct tool arguments
28932921
})
28942922
}
28952923
} else {

‎src/core/tools/__tests__/validateToolUse.spec.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,42 @@ describe("mode-validator", () => {
103103
})
104104
})
105105

106+
describe("dynamic MCP tools", () => {
107+
it("allows dynamic MCP tools when mcp group is in mode groups", () => {
108+
// Code mode has mcp group, so dynamic MCP tools should be allowed
109+
expect(isToolAllowedForMode("mcp_context7_resolve-library-id", codeMode, [])).toBe(true)
110+
expect(isToolAllowedForMode("mcp_serverName_toolName", codeMode, [])).toBe(true)
111+
})
112+
113+
it("disallows dynamic MCP tools when mcp group is not in mode groups", () => {
114+
const customModes: ModeConfig[] = [
115+
{
116+
slug: "no-mcp-mode",
117+
name: "No MCP Mode",
118+
roleDefinition: "Custom role",
119+
groups: ["read", "edit"] as const,
120+
},
121+
]
122+
// Custom mode without mcp group should not allow dynamic MCP tools
123+
expect(isToolAllowedForMode("mcp_context7_resolve-library-id", "no-mcp-mode", customModes)).toBe(false)
124+
expect(isToolAllowedForMode("mcp_serverName_toolName", "no-mcp-mode", customModes)).toBe(false)
125+
})
126+
127+
it("allows dynamic MCP tools in custom mode with mcp group", () => {
128+
const customModes: ModeConfig[] = [
129+
{
130+
slug: "custom-mcp-mode",
131+
name: "Custom MCP Mode",
132+
roleDefinition: "Custom role",
133+
groups: ["read", "mcp"] as const,
134+
},
135+
]
136+
expect(isToolAllowedForMode("mcp_context7_resolve-library-id", "custom-mcp-mode", customModes)).toBe(
137+
true,
138+
)
139+
})
140+
})
141+
106142
describe("tool requirements", () => {
107143
it("respects tool requirements when provided", () => {
108144
const requirements = { apply_diff: false }

‎src/shared/modes.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,10 @@ export function isToolAllowedForMode(
176176
if (ALWAYS_AVAILABLE_TOOLS.includes(tool as any)) {
177177
return true
178178
}
179+
180+
// Check if this is a dynamic MCP tool (mcp_serverName_toolName)
181+
// These should be allowed if the mcp group is allowed for the mode
182+
const isDynamicMcpTool = tool.startsWith("mcp_")
179183
if (experiments && Object.values(EXPERIMENT_IDS).includes(tool as ExperimentId)) {
180184
if (!experiments[tool]) {
181185
return false
@@ -204,6 +208,12 @@ export function isToolAllowedForMode(
204208

205209
const groupConfig = TOOL_GROUPS[groupName]
206210

211+
// Check if this is a dynamic MCP tool and the mcp group is allowed
212+
if (isDynamicMcpTool && groupName === "mcp") {
213+
// Dynamic MCP tools are allowed if the mcp group is in the mode's groups
214+
return true
215+
}
216+
207217
// If the tool isn't in this group's tools, continue to next group
208218
if (!groupConfig.tools.includes(tool)) {
209219
continue

0 commit comments

Comments
 (0)