Skip to content

Commit dbc1751

Browse files
jsourcebotjsourcebot
committed
Refine MCP tool name sanitization
1 parent edd372a commit dbc1751

12 files changed

Lines changed: 277 additions & 116 deletions

File tree

‎packages/web/src/ee/features/chat/agent.test.ts‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,62 @@ beforeEach(() => {
242242
});
243243

244244
describe('createMessageStream approval continuation', () => {
245+
test('streams raw MCP tool names for client display', async () => {
246+
const { getConnectedMcpClients } = await import('@/ee/features/chat/mcp/mcpClientFactory');
247+
const { getMcpTools } = await import('@/ee/features/chat/mcp/mcpToolSets');
248+
vi.mocked(getConnectedMcpClients).mockResolvedValueOnce([
249+
{ serverId: 'server-backstage', serverName: 'Backstage' },
250+
] as never);
251+
vi.mocked(getMcpTools).mockResolvedValueOnce({
252+
tools: {},
253+
failedServers: [],
254+
serverFaviconUrls: {
255+
backstage: 'https://backstage.example.com/favicon.ico',
256+
},
257+
toolDisplayNames: {
258+
'mcp_backstage__catalog_query-catalog-entities': 'catalog.query-catalog-entities',
259+
},
260+
cleanup: vi.fn(),
261+
});
262+
mockAi.streamText.mockReturnValue(createFakeStreamResult());
263+
264+
await createMessageStream({
265+
chatId: 'chat-id',
266+
messages: [createUserMessage()],
267+
selectedRepos: [],
268+
disabledMcpServerIds: [],
269+
prisma: {},
270+
model: {},
271+
modelName: 'test-model',
272+
promptCacheStrategy: noopStrategy,
273+
onFinish: vi.fn(),
274+
onError: () => 'error',
275+
userId: 'user-id',
276+
orgId: 1,
277+
} as unknown as Parameters<typeof createMessageStream>[0]);
278+
279+
const execute = mockAi.latestCreateUIMessageStreamOptions?.execute;
280+
if (!execute) {
281+
throw new Error('Expected createUIMessageStream to capture execute callback.');
282+
}
283+
284+
const write = vi.fn();
285+
await execute({
286+
writer: {
287+
merge: vi.fn(),
288+
write,
289+
},
290+
});
291+
292+
expect(write).toHaveBeenCalledWith({
293+
type: 'data-mcp-tool',
294+
data: {
295+
modelToolName: 'mcp_backstage__catalog_query-catalog-entities',
296+
rawToolName: 'catalog.query-catalog-entities',
297+
},
298+
});
299+
});
300+
245301
test.each([
246302
['dynamic', dynamicApprovalRespondedPart],
247303
['static', staticApprovalRespondedPart],

‎packages/web/src/ee/features/chat/agent.ts‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,12 @@ export const createMessageStream = async ({
338338
data: { sanitizedName, faviconUrl },
339339
});
340340
},
341+
onMcpToolDiscovered: (modelToolName, rawToolName) => {
342+
writer.write({
343+
type: 'data-mcp-tool',
344+
data: { modelToolName, rawToolName },
345+
});
346+
},
341347
onMcpServerFailed: (serverName) => {
342348
writer.write({
343349
type: 'data-mcp-failed-server',
@@ -470,6 +476,7 @@ interface AgentOptions {
470476
inputSources: Source[];
471477
onWriteSource: (source: Source) => void;
472478
onMcpServerDiscovered: (sanitizedName: string, faviconUrl: string) => void;
479+
onMcpToolDiscovered: (modelToolName: string, rawToolName: string) => void;
473480
onMcpServerFailed: (serverName: string) => void;
474481
traceId: string;
475482
chatId: string;
@@ -489,6 +496,7 @@ const createAgentStream = async ({
489496
disabledMcpServerIds,
490497
onWriteSource,
491498
onMcpServerDiscovered,
499+
onMcpToolDiscovered,
492500
onMcpServerFailed,
493501
traceId,
494502
chatId,
@@ -525,7 +533,7 @@ const createAgentStream = async ({
525533
}))
526534
).filter((source) => source !== undefined);
527535

528-
let mcpToolSetsObj: McpToolsResult = { tools: {}, failedServers: [], serverFaviconUrls: {}, cleanup: async () => {} };
536+
let mcpToolSetsObj: McpToolsResult = { tools: {}, failedServers: [], serverFaviconUrls: {}, toolDisplayNames: {}, cleanup: async () => {} };
529537
if (userId && orgId && await hasEntitlement('ask') && disabledMcpServerIds !== undefined) {
530538
try {
531539
const allMcpClients = await getConnectedMcpClients(prisma, userId, orgId);
@@ -539,6 +547,9 @@ const createAgentStream = async ({
539547
for (const [sanitizedName, faviconUrl] of Object.entries(mcpToolSetsObj.serverFaviconUrls)) {
540548
onMcpServerDiscovered(sanitizedName, faviconUrl);
541549
}
550+
for (const [modelToolName, rawToolName] of Object.entries(mcpToolSetsObj.toolDisplayNames)) {
551+
onMcpToolDiscovered(modelToolName, rawToolName);
552+
}
542553

543554
if (mcpClients.length > 0) {
544555
logger.info(`Connected to ${mcpClients.length} external MCP server(s): ${mcpClients.map(c => c.serverName).join(', ')}`);

‎packages/web/src/ee/features/chat/components/chatThread/chatThread.tsx‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ import { duplicateChat } from '@/features/chat/actions';
2828
import { generateAndUpdateChatNameFromMessage } from '@/ee/features/chat/actions';
2929
import { isServiceError } from '@/lib/utils';
3030
import { NotConfiguredErrorBanner } from '@/features/chat/components/notConfiguredErrorBanner';
31-
import { McpServerIconContext, McpServerIconMap } from '../../mcpServerIconContext';
31+
import { McpServerIconContext, McpServerIconMap, McpToolNameContext, McpToolNameMap } from '../../mcpServerIconContext';
3232
import { ToolApprovalProvider } from '../../toolApprovalContext';
3333
import useCaptureEvent from '@/hooks/useCaptureEvent';
3434
import { SignInPromptBanner } from './signInPromptBanner';
@@ -104,6 +104,18 @@ export const ChatThread = ({
104104
return map;
105105
});
106106

107+
const [mcpToolNameMap, setMcpToolNameMap] = useState<McpToolNameMap>(() => {
108+
const map: McpToolNameMap = {};
109+
initialMessages?.forEach((message) => {
110+
message.parts
111+
.filter((part) => part.type === 'data-mcp-tool')
112+
.forEach((part) => {
113+
map[part.data.modelToolName] = part.data.rawToolName;
114+
});
115+
});
116+
return map;
117+
});
118+
107119
const [failedMcpServers, setFailedMcpServers] = useState<string[]>(() => {
108120
const names: string[] = [];
109121
initialMessages?.forEach((message) => {
@@ -173,6 +185,12 @@ export const ChatThread = ({
173185
[dataPart.data.sanitizedName]: dataPart.data.faviconUrl,
174186
}));
175187
}
188+
if (dataPart.type === 'data-mcp-tool') {
189+
setMcpToolNameMap((prev) => ({
190+
...prev,
191+
[dataPart.data.modelToolName]: dataPart.data.rawToolName,
192+
}));
193+
}
176194
if (dataPart.type === 'data-mcp-failed-server') {
177195
setFailedMcpServers((prev) => {
178196
if (prev.includes(dataPart.data.serverName)) {
@@ -385,6 +403,7 @@ export const ChatThread = ({
385403
return (
386404
<ToolApprovalProvider value={addToolApprovalResponse}>
387405
<McpServerIconContext.Provider value={mcpServerIconMap}>
406+
<McpToolNameContext.Provider value={mcpToolNameMap}>
388407
<ChatPaneDropzone
389408
className="flex flex-col flex-1 min-h-0 w-full"
390409
onFilesDropped={(files) => chatBoxRef.current?.addFiles(files)}
@@ -532,6 +551,7 @@ export const ChatThread = ({
532551
)}
533552
</div>
534553
</ChatPaneDropzone>
554+
</McpToolNameContext.Provider>
535555
</McpServerIconContext.Provider>
536556
</ToolApprovalProvider>
537557
);

‎packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,7 @@ export const StepPartRenderer = ({ part, toolTokenUsageMap }: { part: SBChatMess
536536
return null;
537537
case 'data-source':
538538
case 'data-mcp-server':
539+
case 'data-mcp-tool':
539540
case 'data-mcp-failed-server':
540541
case 'data-attachment':
541542
case 'file':

‎packages/web/src/ee/features/chat/components/chatThread/toolApprovalBanner.tsx‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@
22

33
import { Button } from "@/components/ui/button";
44
import { McpFavicon } from "@/ee/features/chat/mcp/components/mcpFavicon";
5-
import { useMcpServerIconMap } from "@/ee/features/chat/mcpServerIconContext";
5+
import { McpToolNameMap, useMcpServerIconMap, useMcpToolNameMap } from "@/ee/features/chat/mcpServerIconContext";
66
import { useToolApproval } from "@/ee/features/chat/toolApprovalContext";
77
import { SBChatToolPart } from "@/features/chat/utils";
88
import { cn } from "@/lib/utils";
99
import { getToolName } from "ai";
1010
import { ChevronRight } from "lucide-react";
1111
import { useCallback, useState } from "react";
12-
import { parseMcpToolName } from "./tools/mcpToolComponent";
12+
import { getMcpToolDisplayParts } from "./tools/mcpToolComponent";
1313
import { JsonHighlighter } from "./tools/jsonHighlighter";
1414

1515
export type ApprovalRequestedToolPart = SBChatToolPart & {
@@ -23,6 +23,7 @@ interface ToolApprovalBannerProps {
2323
export const ToolApprovalBanner = ({ parts }: ToolApprovalBannerProps) => {
2424
const addToolApprovalResponse = useToolApproval();
2525
const iconMap = useMcpServerIconMap();
26+
const rawToolNames = useMcpToolNameMap();
2627

2728
if (parts.length === 0) {
2829
return null;
@@ -36,6 +37,7 @@ export const ToolApprovalBanner = ({ parts }: ToolApprovalBannerProps) => {
3637
part={part}
3738
addToolApprovalResponse={addToolApprovalResponse}
3839
iconMap={iconMap}
40+
rawToolNames={rawToolNames}
3941
/>
4042
))}
4143
</div>
@@ -46,17 +48,17 @@ const ToolApprovalItem = ({
4648
part,
4749
addToolApprovalResponse,
4850
iconMap,
51+
rawToolNames,
4952
}: {
5053
part: ApprovalRequestedToolPart;
5154
addToolApprovalResponse: ReturnType<typeof useToolApproval>;
5255
iconMap: Record<string, string | undefined>;
56+
rawToolNames: McpToolNameMap;
5357
}) => {
5458
const [isExpanded, setIsExpanded] = useState(false);
5559
const partToolName = getToolName(part);
56-
const parsed = parseMcpToolName(partToolName);
57-
const serverName = parsed?.serverName ?? partToolName;
58-
const toolName = parsed?.toolName ?? partToolName;
59-
const faviconUrl = parsed ? iconMap[parsed.serverName] : undefined;
60+
const display = getMcpToolDisplayParts(partToolName, rawToolNames);
61+
const faviconUrl = display.serverName ? iconMap[display.serverName] : undefined;
6062

6163
const requestText = JSON.stringify(part.input, null, 2);
6264

@@ -83,13 +85,13 @@ const ToolApprovalItem = ({
8385
>
8486
<McpFavicon faviconUrl={faviconUrl} className="w-4 h-4" />
8587
<span className="text-sm text-foreground truncate">
86-
{parsed ? (
88+
{display.serverName ? (
8789
<>
88-
Agent wants to use <span className="font-medium">{toolName}</span> from <span className="font-medium">{serverName}</span>
90+
Agent wants to use <span className="font-medium">{display.toolName}</span> from <span className="font-medium">{display.serverName}</span>
8991
</>
9092
) : (
9193
<>
92-
Agent wants to use <span className="font-medium">{toolName}</span>
94+
Agent wants to use <span className="font-medium">{display.toolName}</span>
9395
</>
9496
)}
9597
</span>

‎packages/web/src/ee/features/chat/components/chatThread/tools/mcpToolComponent.test.tsx‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
import { render, screen } from '@testing-library/react';
2+
import type { DynamicToolUIPart } from 'ai';
3+
import { describe, expect, test } from 'vitest';
4+
import { McpToolNameContext } from '@/ee/features/chat/mcpServerIconContext';
5+
import { getMcpToolDisplayParts, McpToolComponent } from './mcpToolComponent';
6+
7+
describe('getMcpToolDisplayParts', () => {
8+
test('maps provider-safe MCP tool names back to raw tool names for display', () => {
9+
expect(getMcpToolDisplayParts(
10+
'mcp_backstage__catalog_query-catalog-entities',
11+
{
12+
'mcp_backstage__catalog_query-catalog-entities': 'catalog.query-catalog-entities',
13+
},
14+
)).toEqual({
15+
serverName: 'backstage',
16+
toolName: 'catalog.query-catalog-entities',
17+
displayName: 'backstage: catalog.query-catalog-entities',
18+
});
19+
});
20+
21+
test('falls back to the provider-safe name for older messages without metadata', () => {
22+
expect(getMcpToolDisplayParts('mcp_backstage__catalog_query-catalog-entities')).toEqual({
23+
serverName: 'backstage',
24+
toolName: 'catalog_query-catalog-entities',
25+
displayName: 'backstage: catalog_query-catalog-entities',
26+
});
27+
});
28+
});
29+
30+
describe('McpToolComponent', () => {
31+
test('renders the raw MCP tool name when display metadata is available', () => {
32+
const part = {
33+
type: 'dynamic-tool',
34+
toolName: 'mcp_backstage__catalog_query-catalog-entities',
35+
toolCallId: 'tool-call-1',
36+
state: 'approval-requested',
37+
input: { filter: 'kind=component' },
38+
} as DynamicToolUIPart;
39+
40+
render(
41+
<McpToolNameContext.Provider value={{
42+
'mcp_backstage__catalog_query-catalog-entities': 'catalog.query-catalog-entities',
43+
}}>
44+
<McpToolComponent part={part} />
45+
</McpToolNameContext.Provider>
46+
);
47+
48+
expect(screen.getByText('backstage: catalog.query-catalog-entities')).toBeTruthy();
49+
expect(screen.getByText('Request (backstage: catalog.query-catalog-entities)')).toBeTruthy();
50+
expect(screen.queryByText('backstage: catalog_query-catalog-entities')).toBeNull();
51+
});
52+
});

0 commit comments

Comments
 (0)