Skip to content

Commit 0553363

Browse files
harshachVishnuujain
harshach
authored and
Vishnuujain
committed
MCP Tool Usage (#28352)
* MCP Tool Usage * Update generated TypeScript types * Address PR review feedback on MCP usage tracking Reorder UA heuristic so VS Code wins over Claude CLI for composite User-Agents, refactor to a predicate list, and sanitise the resolved client name (trim, strip control chars, cap at 64 chars). Bound the schema field to match. Bound the latency aggregation lists in McpUsageResource with reservoir sampling so summary/per-tool percentile estimates stay valid without unbounded heap growth. Skip null-timestamp rows in the history loop and update the stale /history Swagger description to reflect the ok/fail shape. Convert CallToolOutcome to a Java record and update the recorder flow to use accessor methods. Fix the pre-existing regression in McpImpersonationTest where the mock still wired the legacy callTool path. Add DefaultToolContextTest with direct coverage for classifyException (all four ErrorCategory buckets, cause-chain walk, null message in chain) and the unknown-tool outcome. (cherry picked from commit 1dcf8dd) (cherry picked from commit 21aee7f)
1 parent e44c5b7 commit 0553363

12 files changed

Lines changed: 1069 additions & 107 deletions

File tree

‎openmetadata-mcp/src/main/java/org/openmetadata/mcp/AuthEnrichedMcpContextExtractor.java‎

Lines changed: 114 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,129 @@
33
import io.modelcontextprotocol.common.McpTransportContext;
44
import io.modelcontextprotocol.server.McpTransportContextExtractor;
55
import jakarta.servlet.http.HttpServletRequest;
6+
import java.util.HashMap;
7+
import java.util.List;
8+
import java.util.Locale;
69
import java.util.Map;
10+
import java.util.function.Predicate;
711
import org.openmetadata.service.security.JwtFilter;
812

913
public class AuthEnrichedMcpContextExtractor
1014
implements McpTransportContextExtractor<HttpServletRequest> {
1115
public static final String AUTHORIZATION_HEADER = "Authorization";
1216

17+
/**
18+
* Context key carrying the resolved client name (Claude Desktop / Cursor / VS Code / etc.)
19+
* derived from the User-Agent header. Stamped on every {@link
20+
* org.openmetadata.schema.entity.app.mcp.McpToolCallUsage} row so the Billing > MCP page can
21+
* group by client. Kept as a constant on this class so producer + consumer share the spelling.
22+
*/
23+
public static final String CLIENT_NAME = "Mcp-Client-Name";
24+
25+
private static final String USER_AGENT_HEADER = "User-Agent";
26+
27+
/**
28+
* Upper bound on the persisted client name. The value comes from a user-controlled
29+
* {@code User-Agent} header, so we cap it to prevent oversized strings or junk characters from
30+
* leaking into the {@code McpToolCallUsage} rows. Mirrors {@code maxLength} on the
31+
* {@code clientName} property in {@code mcpToolCallUsage.json}.
32+
*/
33+
static final int MAX_CLIENT_NAME_LENGTH = 64;
34+
35+
/**
36+
* Ordered list of (predicate, label) pairs used to classify a lower-cased User-Agent into a
37+
* human-readable client name. Order matters — VS Code is checked before Claude CLI so a UA
38+
* containing both {@code claude} and {@code code} substrings (e.g. a Claude extension hosted in
39+
* VS Code) is attributed to the host process rather than the standalone CLI.
40+
*/
41+
private static final List<UaMatcher> UA_MATCHERS =
42+
List.of(
43+
new UaMatcher(ua -> ua.contains("claude") && ua.contains("desktop"), "Claude Desktop"),
44+
new UaMatcher(
45+
ua ->
46+
ua.contains("vscode")
47+
|| ua.contains("vs code")
48+
|| ua.contains("visual studio code"),
49+
"VS Code"),
50+
new UaMatcher(
51+
ua ->
52+
ua.contains("claude-cli")
53+
|| ua.contains("claude-code")
54+
|| ua.contains("claude cli")
55+
|| ua.contains("claude code"),
56+
"Claude CLI"),
57+
new UaMatcher(ua -> ua.contains("cursor"), "Cursor"),
58+
new UaMatcher(ua -> ua.contains("zed"), "Zed"),
59+
new UaMatcher(ua -> ua.contains("windsurf"), "Windsurf"));
60+
61+
/** Pairing of a User-Agent substring predicate with the label shown in the dashboard. */
62+
private record UaMatcher(Predicate<String> matches, String label) {}
63+
1364
@Override
1465
public McpTransportContext extract(HttpServletRequest request) {
1566
String token = JwtFilter.extractToken(request.getHeader(AUTHORIZATION_HEADER));
16-
return McpTransportContext.create(Map.of(AUTHORIZATION_HEADER, token != null ? token : ""));
67+
String clientName = resolveClientName(request.getHeader(USER_AGENT_HEADER));
68+
Map<String, Object> values = new HashMap<>();
69+
values.put(AUTHORIZATION_HEADER, token != null ? token : "");
70+
if (clientName != null) {
71+
values.put(CLIENT_NAME, clientName);
72+
}
73+
return McpTransportContext.create(values);
74+
}
75+
76+
/**
77+
* Heuristic: classify the {@code User-Agent} into the labels the dashboard uses (Claude Desktop,
78+
* Cursor, VS Code, Claude CLI, etc.). MCP clients all set distinctive UAs — the table at
79+
* {@link #UA_MATCHERS} covers the common ones explicitly and we fall back to the raw product
80+
* token so a new client still surfaces in the breakdown without a code change. The fallback
81+
* value (and every label) is sanitised before return so the persisted value is bounded.
82+
*/
83+
static String resolveClientName(String userAgent) {
84+
String resolved = null;
85+
if (userAgent != null && !userAgent.isBlank()) {
86+
String ua = userAgent.toLowerCase(Locale.ROOT);
87+
resolved =
88+
UA_MATCHERS.stream()
89+
.filter(matcher -> matcher.matches().test(ua))
90+
.map(UaMatcher::label)
91+
.findFirst()
92+
.orElseGet(() -> fallbackProductToken(userAgent));
93+
}
94+
return sanitize(resolved);
95+
}
96+
97+
/**
98+
* First product token of the User-Agent (the bit before the first slash or space), capitalised.
99+
* Avoids leaking version strings into the dashboard while still surfacing unknown clients with a
100+
* human-readable label. Returns {@code null} when no usable token is present.
101+
*/
102+
private static String fallbackProductToken(String userAgent) {
103+
String head = userAgent.split("[\\s/]", 2)[0];
104+
String result = null;
105+
if (!head.isBlank()) {
106+
result = head.substring(0, 1).toUpperCase(Locale.ROOT) + head.substring(1);
107+
}
108+
return result;
109+
}
110+
111+
/**
112+
* Trims whitespace, strips ISO control characters, and caps the value to
113+
* {@link #MAX_CLIENT_NAME_LENGTH} characters. The User-Agent header is attacker-controlled so
114+
* the persisted value must be sanitised before it reaches the database or the dashboard.
115+
*/
116+
private static String sanitize(String value) {
117+
String result = null;
118+
if (value != null) {
119+
StringBuilder sb = new StringBuilder(value.length());
120+
value.codePoints().filter(cp -> !Character.isISOControl(cp)).forEach(sb::appendCodePoint);
121+
String stripped = sb.toString().trim();
122+
if (!stripped.isEmpty()) {
123+
result =
124+
stripped.length() > MAX_CLIENT_NAME_LENGTH
125+
? stripped.substring(0, MAX_CLIENT_NAME_LENGTH)
126+
: stripped;
127+
}
128+
}
129+
return result;
17130
}
18131
}

‎openmetadata-mcp/src/main/java/org/openmetadata/mcp/McpServer.java‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -250,14 +250,23 @@ protected McpStatelessServerFeatures.SyncToolSpecification getTool(McpSchema.Too
250250
CatalogSecurityContext securityContext =
251251
jwtFilter.getCatalogSecurityContext((String) context.get("Authorization"));
252252
String userName = securityContext.getUserPrincipal().getName();
253-
McpSchema.CallToolResult result = null;
253+
String clientName =
254+
(String)
255+
context.get(org.openmetadata.mcp.AuthEnrichedMcpContextExtractor.CLIENT_NAME);
256+
org.openmetadata.mcp.tools.DefaultToolContext.CallToolOutcome outcome = null;
254257
try {
255258
ImpersonationContext.setImpersonatedBy(getMcpBotName());
256-
result = toolContext.callTool(authorizer, limits, tool.name(), securityContext, req);
257-
return result;
259+
outcome =
260+
toolContext.callToolWithMetadata(
261+
authorizer, limits, tool.name(), securityContext, req);
262+
return outcome.result();
258263
} finally {
259-
boolean success = result != null && !Boolean.TRUE.equals(result.isError());
260-
McpUsageRecorder.record(tool.name(), userName, success);
264+
boolean success = outcome != null && !Boolean.TRUE.equals(outcome.result().isError());
265+
Long latencyMs = outcome != null ? outcome.latencyMs() : null;
266+
org.openmetadata.schema.entity.app.mcp.McpToolCallUsage.ErrorCategory category =
267+
outcome != null ? outcome.errorCategory() : null;
268+
McpUsageRecorder.record(
269+
tool.name(), userName, success, latencyMs, category, clientName);
261270
ImpersonationContext.clear();
262271
}
263272
});

‎openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java‎

Lines changed: 166 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,11 @@
44

55
import io.modelcontextprotocol.spec.McpSchema;
66
import java.util.List;
7+
import java.util.Locale;
78
import java.util.Map;
9+
import java.util.function.Predicate;
810
import lombok.extern.slf4j.Slf4j;
11+
import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage;
912
import org.openmetadata.schema.utils.JsonUtils;
1013
import org.openmetadata.service.limits.Limits;
1114
import org.openmetadata.service.security.AuthorizationException;
@@ -32,6 +35,22 @@ public McpSchema.CallToolResult callTool(
3235
String toolName,
3336
CatalogSecurityContext securityContext,
3437
McpSchema.CallToolRequest request) {
38+
return callToolWithMetadata(authorizer, limits, toolName, securityContext, request).result();
39+
}
40+
41+
/**
42+
* Phase 3 entry point. Returns the tool result alongside the metadata the {@link
43+
* org.openmetadata.mcp.usage.McpUsageRecorder} needs (latency + error category). Kept as a
44+
* separate method so the legacy single-result signature stays available for external callers
45+
* that haven't migrated yet.
46+
*/
47+
public CallToolOutcome callToolWithMetadata(
48+
Authorizer authorizer,
49+
Limits limits,
50+
String toolName,
51+
CatalogSecurityContext securityContext,
52+
McpSchema.CallToolRequest request) {
53+
long startNanos = System.nanoTime();
3554
LOG.info(
3655
"Catalog Principal: {} is trying to call the tool: {}",
3756
securityContext.getUserPrincipal().getName(),
@@ -97,47 +116,159 @@ public McpSchema.CallToolResult callTool(
97116
result = new CreateDataProductTool().execute(authorizer, limits, securityContext, params);
98117
break;
99118
default:
100-
return McpSchema.CallToolResult.builder()
101-
.content(
102-
List.of(
103-
new McpSchema.TextContent(
104-
JsonUtils.pojoToJson(Map.of("error", "Unknown function: " + toolName)))))
105-
.isError(true)
106-
.build();
119+
return new CallToolOutcome(
120+
McpSchema.CallToolResult.builder()
121+
.content(
122+
List.of(
123+
new McpSchema.TextContent(
124+
JsonUtils.pojoToJson(
125+
Map.of("error", "Unknown function: " + toolName)))))
126+
.isError(true)
127+
.build(),
128+
elapsedMs(startNanos),
129+
McpToolCallUsage.ErrorCategory.VALIDATION);
107130
}
108131

109-
return McpSchema.CallToolResult.builder()
110-
.content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
111-
.isError(false)
112-
.build();
132+
return new CallToolOutcome(
133+
McpSchema.CallToolResult.builder()
134+
.content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
135+
.isError(false)
136+
.build(),
137+
elapsedMs(startNanos),
138+
null);
113139
} catch (AuthorizationException ex) {
114140
LOG.warn("Authorization error: {}", ex.getMessage());
115-
return McpSchema.CallToolResult.builder()
116-
.content(
117-
List.of(
118-
new McpSchema.TextContent(
119-
JsonUtils.pojoToJson(
120-
Map.of(
121-
"error",
122-
String.format("Authorization error: %s", ex.getMessage()),
123-
"statusCode",
124-
403)))))
125-
.isError(true)
126-
.build();
141+
return new CallToolOutcome(
142+
McpSchema.CallToolResult.builder()
143+
.content(
144+
List.of(
145+
new McpSchema.TextContent(
146+
JsonUtils.pojoToJson(
147+
Map.of(
148+
"error",
149+
String.format("Authorization error: %s", ex.getMessage()),
150+
"statusCode",
151+
403)))))
152+
.isError(true)
153+
.build(),
154+
elapsedMs(startNanos),
155+
McpToolCallUsage.ErrorCategory.AUTH);
127156
} catch (Exception ex) {
128157
LOG.error("Error executing tool '{}': {}", toolName, ex.getMessage(), ex);
129-
return McpSchema.CallToolResult.builder()
130-
.content(
131-
List.of(
132-
new McpSchema.TextContent(
133-
JsonUtils.pojoToJson(
134-
Map.of(
135-
"error",
136-
String.format("Error executing tool: %s", ex.getMessage()),
137-
"statusCode",
138-
500)))))
139-
.isError(true)
140-
.build();
158+
return new CallToolOutcome(
159+
McpSchema.CallToolResult.builder()
160+
.content(
161+
List.of(
162+
new McpSchema.TextContent(
163+
JsonUtils.pojoToJson(
164+
Map.of(
165+
"error",
166+
String.format("Error executing tool: %s", ex.getMessage()),
167+
"statusCode",
168+
500)))))
169+
.isError(true)
170+
.build(),
171+
elapsedMs(startNanos),
172+
classifyException(ex));
141173
}
142174
}
175+
176+
/**
177+
* Maps an arbitrary exception type to one of the {@link McpToolCallUsage.ErrorCategory} values.
178+
* Walks the cause chain because the tool wrappers usually rethrow framework errors wrapped in
179+
* a {@link RuntimeException}. Defaults to {@link McpToolCallUsage.ErrorCategory#INTERNAL} when
180+
* no specific bucket matches.
181+
*/
182+
static McpToolCallUsage.ErrorCategory classifyException(Throwable t) {
183+
McpToolCallUsage.ErrorCategory result = McpToolCallUsage.ErrorCategory.INTERNAL;
184+
Throwable cursor = t;
185+
while (cursor != null && result == McpToolCallUsage.ErrorCategory.INTERNAL) {
186+
McpToolCallUsage.ErrorCategory match = matchCategory(cursor);
187+
if (match != null) {
188+
result = match;
189+
} else {
190+
Throwable next = cursor.getCause();
191+
cursor = (next == null || next == cursor) ? null : next;
192+
}
193+
}
194+
return result;
195+
}
196+
197+
/**
198+
* Pairing of an exception (name, message) predicate with the bucket it should produce. Kept
199+
* as a static table so adding a new category (or extending an existing one with a new keyword)
200+
* is a one-line change rather than another {@code else if} branch.
201+
*/
202+
private record CategoryMatcher(
203+
Predicate<ExceptionMeta> matches, McpToolCallUsage.ErrorCategory category) {}
204+
205+
/** Lower-cased name + message pair so each matcher inspects both without re-parsing. */
206+
private record ExceptionMeta(String name, String message) {}
207+
208+
/**
209+
* Ordered category table. Check order matters: more specific patterns sit before broader ones so
210+
* a {@code RateLimitException} doesn't get caught by the generic message-substring rules below
211+
* it. {@code AUTH} sits above {@code VALIDATION} because some auth exceptions ({@code
212+
* AuthorizationException}) extend {@code IllegalArgumentException}-style hierarchies and would
213+
* otherwise be mis-bucketed.
214+
*/
215+
private static final List<CategoryMatcher> CATEGORY_MATCHERS =
216+
List.of(
217+
new CategoryMatcher(
218+
meta -> meta.name().contains("RateLimit") || meta.message().contains("rate limit"),
219+
McpToolCallUsage.ErrorCategory.RATE_LIMIT),
220+
new CategoryMatcher(
221+
meta ->
222+
meta.name().contains("Authorization")
223+
|| meta.name().contains("Forbidden")
224+
|| meta.name().contains("Unauthorized")
225+
|| meta.message().contains("forbidden")
226+
|| meta.message().contains("unauthorized")
227+
|| meta.message().contains("access denied")
228+
|| meta.message().contains("permission denied"),
229+
McpToolCallUsage.ErrorCategory.AUTH),
230+
new CategoryMatcher(
231+
meta ->
232+
meta.name().contains("Validation")
233+
|| meta.name().contains("IllegalArgument")
234+
|| meta.name().contains("BadRequest")
235+
|| meta.message().contains("invalid argument"),
236+
McpToolCallUsage.ErrorCategory.VALIDATION),
237+
new CategoryMatcher(
238+
meta ->
239+
meta.name().contains("Timeout")
240+
|| meta.message().contains("timeout")
241+
|| meta.message().contains("timed out"),
242+
McpToolCallUsage.ErrorCategory.TIMEOUT));
243+
244+
/**
245+
* Returns the category that matches the supplied throwable's name or message, or {@code null}
246+
* when no specific bucket applies. Kept separate from {@link #classifyException} so the
247+
* cause-chain walk reads as a single linear loop.
248+
*/
249+
private static McpToolCallUsage.ErrorCategory matchCategory(Throwable cursor) {
250+
ExceptionMeta meta =
251+
new ExceptionMeta(
252+
cursor.getClass().getSimpleName(),
253+
cursor.getMessage() == null ? "" : cursor.getMessage().toLowerCase(Locale.ROOT));
254+
return CATEGORY_MATCHERS.stream()
255+
.filter(matcher -> matcher.matches().test(meta))
256+
.map(CategoryMatcher::category)
257+
.findFirst()
258+
.orElse(null);
259+
}
260+
261+
private static long elapsedMs(long startNanos) {
262+
return (System.nanoTime() - startNanos) / 1_000_000L;
263+
}
264+
265+
/**
266+
* Phase 3 — tuple returned by {@link #callToolWithMetadata} so the MCP server can record the
267+
* call with full diagnostic detail without re-classifying the exception or re-measuring the
268+
* latency at its level.
269+
*/
270+
public record CallToolOutcome(
271+
McpSchema.CallToolResult result,
272+
long latencyMs,
273+
McpToolCallUsage.ErrorCategory errorCategory) {}
143274
}

0 commit comments

Comments
 (0)