Skip to content

Commit a5b52de

Browse files
abhipatel12abhipatel12
committed
addressing christian's feedback
1 parent e43928f commit a5b52de

5 files changed

Lines changed: 136 additions & 41 deletions

File tree

‎packages/core/src/policy/policy-engine.ts‎

Lines changed: 21 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,19 @@ import {
2525
hasRedirection,
2626
} from '../utils/shell-utils.js';
2727
import { getToolAliases } from '../tools/tool-names.js';
28-
import { MCP_TOOL_PREFIX } from '../tools/mcp-tool.js';
28+
import {
29+
MCP_TOOL_PREFIX,
30+
isMcpToolAnnotation,
31+
parseMcpToolName,
32+
} from '../tools/mcp-tool.js';
2933

3034
function isWildcardPattern(name: string): boolean {
3135
return name === '*' || name.includes('*');
3236
}
3337

3438
/**
3539
* Checks if a tool call matches a wildcard pattern.
36-
* Supports global (*) and the new explicit MCP (*mcp_serverName_**) format.
40+
* Supports global (*) and the explicit MCP (*mcp_serverName_**) format.
3741
*/
3842
function matchesWildcard(
3943
pattern: string,
@@ -91,7 +95,7 @@ function ruleMatches(
9195

9296
// Check tool name if specified
9397
if (rule.toolName) {
94-
// Support wildcard patterns: "serverName__*" matches "serverName__anyTool"
98+
// Support wildcard patterns: "mcp_serverName_*" matches "mcp_serverName_anyTool"
9599
if (rule.toolName === '*') {
96100
// Match all tools
97101
} else if (isWildcardPattern(rule.toolName)) {
@@ -349,18 +353,19 @@ export class PolicyEngine {
349353
serverName: string | undefined,
350354
toolAnnotations?: Record<string, unknown>,
351355
): Promise<CheckResult> {
352-
if (
353-
!serverName &&
354-
toolAnnotations &&
355-
typeof toolAnnotations['_serverName'] === 'string'
356-
) {
357-
serverName = toolAnnotations['_serverName'];
356+
// Case 1: Metadata injection is the primary and safest way to identify an MCP server.
357+
// If we have explicit `_serverName` metadata (usually injected by tool-registry for active tools), use it.
358+
if (!serverName && isMcpToolAnnotation(toolAnnotations)) {
359+
serverName = toolAnnotations._serverName;
358360
}
359361

360-
if (!serverName && toolCall.name?.startsWith(MCP_TOOL_PREFIX)) {
361-
const parts = toolCall.name.split('_');
362-
if (parts.length >= 3) {
363-
serverName = parts[1];
362+
// Case 2: Fallback for static FQN strings (e.g. from TOML policies or allowed/excluded settings strings).
363+
// These strings don't have active metadata objects associated with them during policy generation,
364+
// so we must extract the server name from the qualified `mcp_{server}_{tool}` format.
365+
if (!serverName && toolCall.name) {
366+
const parsed = parseMcpToolName(toolCall.name);
367+
if (parsed.serverName) {
368+
serverName = parsed.serverName;
364369
}
365370
}
366371

@@ -397,20 +402,12 @@ export class PolicyEngine {
397402
let matchedRule: PolicyRule | undefined;
398403
let decision: PolicyDecision | undefined;
399404

400-
// For tools with a server name, we want to try matching both the
401-
// original name and the fully qualified name (server__tool).
402405
// We also want to check legacy aliases for the tool name.
403406
const toolNamesToTry = toolCall.name ? getToolAliases(toolCall.name) : [];
404407

405408
const toolCallsToTry: FunctionCall[] = [];
406409
for (const name of toolNamesToTry) {
407410
toolCallsToTry.push({ ...toolCall, name });
408-
if (serverName && !name.includes('__')) {
409-
toolCallsToTry.push({
410-
...toolCall,
411-
name: `${serverName}__${name}`,
412-
});
413-
}
414411
}
415412

416413
for (const rule of this.rules) {
@@ -654,9 +651,9 @@ export class PolicyEngine {
654651

655652
for (const toolName of allToolNames) {
656653
const annotations = toolMetadata?.get(toolName);
657-
const rawServerName = annotations?.['_serverName'];
658-
const serverName =
659-
typeof rawServerName === 'string' ? rawServerName : undefined;
654+
const serverName = isMcpToolAnnotation(annotations)
655+
? annotations._serverName
656+
: undefined;
660657

661658
let staticallyExcluded = false;
662659
let matchFound = false;

‎packages/core/src/policy/toml-loader.ts‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import path from 'node:path';
2424
import toml from '@iarna/toml';
2525
import { z, type ZodError } from 'zod';
2626
import { isNodeError } from '../utils/errors.js';
27-
import { MCP_TOOL_PREFIX } from '../tools/mcp-tool.js';
27+
import { MCP_TOOL_PREFIX, formatMcpToolName } from '../tools/mcp-tool.js';
2828

2929
/**
3030
* Maximum Levenshtein distance to consider a name a likely typo of a built-in tool.
@@ -455,14 +455,11 @@ export async function loadPoliciesFromToml(
455455
let effectiveToolName: string | undefined = toolName;
456456
const mcpName = rule.mcpName;
457457

458-
if (mcpName === '*' && !effectiveToolName) {
459-
effectiveToolName = `${MCP_TOOL_PREFIX}*`;
460-
} else if (mcpName === '*') {
461-
effectiveToolName = `${MCP_TOOL_PREFIX}*_${effectiveToolName}`;
462-
} else if (mcpName && !effectiveToolName) {
463-
effectiveToolName = `${MCP_TOOL_PREFIX}${mcpName}_*`;
464-
} else if (mcpName) {
465-
effectiveToolName = `${MCP_TOOL_PREFIX}${mcpName}_${effectiveToolName}`;
458+
if (mcpName) {
459+
effectiveToolName = formatMcpToolName(
460+
mcpName,
461+
effectiveToolName,
462+
);
466463
}
467464

468465
const policyRule: PolicyRule = {

‎packages/core/src/tools/mcp-tool.test.ts‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,11 @@ import {
1515
type Mocked,
1616
} from 'vitest';
1717
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
18-
import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
18+
import {
19+
DiscoveredMCPTool,
20+
generateValidName,
21+
formatMcpToolName,
22+
} from './mcp-tool.js'; // Added getStringifiedResultForDisplay
1923
import { ToolConfirmationOutcome, type ToolResult } from './tools.js';
2024
import type { CallableTool, Part } from '@google/genai';
2125
import { ToolErrorType } from './tool-error.js';
@@ -80,6 +84,30 @@ describe('generateValidName', () => {
8084
);
8185
});
8286

87+
describe('formatMcpToolName', () => {
88+
it('should format a fully qualified name', () => {
89+
expect(formatMcpToolName('github', 'list_repos')).toBe(
90+
'mcp_github_list_repos',
91+
);
92+
});
93+
94+
it('should handle global wildcards', () => {
95+
expect(formatMcpToolName('*')).toBe('mcp_*');
96+
});
97+
98+
it('should handle tool-level wildcards', () => {
99+
expect(formatMcpToolName('github', '*')).toBe('mcp_github_*');
100+
});
101+
102+
it('should handle undefined toolName as a tool-level wildcard', () => {
103+
expect(formatMcpToolName('github')).toBe('mcp_github_*');
104+
});
105+
106+
it('should format explicitly global wildcard with specific tool', () => {
107+
expect(formatMcpToolName('*', 'list_repos')).toBe('mcp_*_list_repos');
108+
});
109+
});
110+
83111
describe('DiscoveredMCPTool', () => {
84112
const serverName = 'mock-mcp-server';
85113
const serverToolName = 'actual-server-tool-name';

‎packages/core/src/tools/mcp-tool.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,77 @@ export function isMcpToolName(name: string): boolean {
4141
return name.startsWith(MCP_TOOL_PREFIX);
4242
}
4343

44+
/**
45+
* Extracts the server name and tool name from a fully qualified MCP tool name.
46+
* Expected format: `mcp_{server_name}_{tool_name}`
47+
* @param name The fully qualified tool name.
48+
* @returns An object containing the extracted `serverName` and `toolName`, or
49+
* `undefined` properties if the name doesn't match the expected format.
50+
*/
51+
export function parseMcpToolName(name: string): {
52+
serverName?: string;
53+
toolName?: string;
54+
} {
55+
if (!isMcpToolName(name)) {
56+
return {};
57+
}
58+
// Remove the prefix
59+
const withoutPrefix = name.slice(MCP_TOOL_PREFIX.length);
60+
// The first segment is the server name, the rest is the tool name
61+
const match = withoutPrefix.match(/^([^_]+)_(.+)$/);
62+
if (match) {
63+
return {
64+
serverName: match[1],
65+
toolName: match[2],
66+
};
67+
}
68+
return {};
69+
}
70+
71+
/**
72+
* Assembles a fully qualified MCP tool name (or wildcard pattern) from its server and tool components.
73+
*
74+
* @param serverName The backend MCP server name (can be '*' for global wildcards).
75+
* @param toolName The name of the tool (can be undefined or '*' for tool-level wildcards).
76+
* @returns The fully qualified name (e.g., `mcp_server_tool`, `mcp_*`, `mcp_server_*`).
77+
*/
78+
export function formatMcpToolName(
79+
serverName: string,
80+
toolName?: string,
81+
): string {
82+
if (serverName === '*' && !toolName) {
83+
return `${MCP_TOOL_PREFIX}*`;
84+
} else if (serverName === '*') {
85+
return `${MCP_TOOL_PREFIX}*_${toolName}`;
86+
} else if (!toolName) {
87+
return `${MCP_TOOL_PREFIX}${serverName}_*`;
88+
} else {
89+
return `${MCP_TOOL_PREFIX}${serverName}_${toolName}`;
90+
}
91+
}
92+
93+
/**
94+
* Interface representing metadata annotations specific to an MCP tool.
95+
* Ensures strongly-typed access to server-level properties.
96+
*/
97+
export interface McpToolAnnotation extends Record<string, unknown> {
98+
_serverName: string;
99+
}
100+
101+
/**
102+
* Type guard to check if tool annotations implement McpToolAnnotation.
103+
*/
104+
export function isMcpToolAnnotation(
105+
annotation: unknown,
106+
): annotation is McpToolAnnotation {
107+
return (
108+
typeof annotation === 'object' &&
109+
annotation !== null &&
110+
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
111+
typeof (annotation as Record<string, unknown>)['_serverName'] === 'string'
112+
);
113+
}
114+
44115
type ToolParams = Record<string, unknown>;
45116

46117
// Discriminated union for MCP Content Blocks to ensure type safety.

‎packages/core/src/tools/tool-registry.ts‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -445,13 +445,15 @@ export class ToolRegistry {
445445
private buildToolMetadata(): Map<string, Record<string, unknown>> {
446446
const toolMetadata = new Map<string, Record<string, unknown>>();
447447
for (const [name, tool] of this.allKnownTools) {
448-
if (tool.toolAnnotations) {
449-
const metadata: Record<string, unknown> = { ...tool.toolAnnotations };
450-
// Include server name so the policy engine can resolve composite
451-
// wildcard patterns (e.g. "*__*") against unqualified tool names.
452-
if (tool instanceof DiscoveredMCPTool) {
453-
metadata['_serverName'] = tool.serverName;
454-
}
448+
const metadata: Record<string, unknown> = tool.toolAnnotations
449+
? { ...tool.toolAnnotations }
450+
: {};
451+
// Include server name so the policy engine can resolve composite
452+
// wildcard patterns (e.g. "*__*") against unqualified tool names.
453+
if (tool instanceof DiscoveredMCPTool) {
454+
metadata['_serverName'] = tool.serverName;
455+
}
456+
if (Object.keys(metadata).length > 0) {
455457
toolMetadata.set(name, metadata);
456458
}
457459
}

0 commit comments

Comments
 (0)