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)

Author: harshach

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

@@ -3,16 +3,129 @@
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.server.McpTransportContextExtractor;
 import jakarta.servlet.http.HttpServletRequest;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
+import java.util.function.Predicate;
 import org.openmetadata.service.security.JwtFilter;
 
 public class AuthEnrichedMcpContextExtractor
     implements McpTransportContextExtractor<HttpServletRequest> {
   public static final String AUTHORIZATION_HEADER = "Authorization";
 
+  /**
+   * Context key carrying the resolved client name (Claude Desktop / Cursor / VS Code / etc.)
+   * derived from the User-Agent header. Stamped on every {@link
+   * org.openmetadata.schema.entity.app.mcp.McpToolCallUsage} row so the Billing > MCP page can
+   * group by client. Kept as a constant on this class so producer + consumer share the spelling.
+   */
+  public static final String CLIENT_NAME = "Mcp-Client-Name";
+
+  private static final String USER_AGENT_HEADER = "User-Agent";
+
+  /**
+   * Upper bound on the persisted client name. The value comes from a user-controlled
+   * {@code User-Agent} header, so we cap it to prevent oversized strings or junk characters from
+   * leaking into the {@code McpToolCallUsage} rows. Mirrors {@code maxLength} on the
+   * {@code clientName} property in {@code mcpToolCallUsage.json}.
+   */
+  static final int MAX_CLIENT_NAME_LENGTH = 64;
+
+  /**
+   * Ordered list of (predicate, label) pairs used to classify a lower-cased User-Agent into a
+   * human-readable client name. Order matters — VS Code is checked before Claude CLI so a UA
+   * containing both {@code claude} and {@code code} substrings (e.g. a Claude extension hosted in
+   * VS Code) is attributed to the host process rather than the standalone CLI.
+   */
+  private static final List<UaMatcher> UA_MATCHERS =
+      List.of(
+          new UaMatcher(ua -> ua.contains("claude") && ua.contains("desktop"), "Claude Desktop"),
+          new UaMatcher(
+              ua ->
+                  ua.contains("vscode")
+                      || ua.contains("vs code")
+                      || ua.contains("visual studio code"),
+              "VS Code"),
+          new UaMatcher(
+              ua ->
+                  ua.contains("claude-cli")
+                      || ua.contains("claude-code")
+                      || ua.contains("claude cli")
+                      || ua.contains("claude code"),
+              "Claude CLI"),
+          new UaMatcher(ua -> ua.contains("cursor"), "Cursor"),
+          new UaMatcher(ua -> ua.contains("zed"), "Zed"),
+          new UaMatcher(ua -> ua.contains("windsurf"), "Windsurf"));
+
+  /** Pairing of a User-Agent substring predicate with the label shown in the dashboard. */
+  private record UaMatcher(Predicate<String> matches, String label) {}
+
   @Override
   public McpTransportContext extract(HttpServletRequest request) {
     String token = JwtFilter.extractToken(request.getHeader(AUTHORIZATION_HEADER));
-    return McpTransportContext.create(Map.of(AUTHORIZATION_HEADER, token != null ? token : ""));
+    String clientName = resolveClientName(request.getHeader(USER_AGENT_HEADER));
+    Map<String, Object> values = new HashMap<>();
+    values.put(AUTHORIZATION_HEADER, token != null ? token : "");
+    if (clientName != null) {
+      values.put(CLIENT_NAME, clientName);
+    }
+    return McpTransportContext.create(values);
+  }
+
+  /**
+   * Heuristic: classify the {@code User-Agent} into the labels the dashboard uses (Claude Desktop,
+   * Cursor, VS Code, Claude CLI, etc.). MCP clients all set distinctive UAs — the table at
+   * {@link #UA_MATCHERS} covers the common ones explicitly and we fall back to the raw product
+   * token so a new client still surfaces in the breakdown without a code change. The fallback
+   * value (and every label) is sanitised before return so the persisted value is bounded.
+   */
+  static String resolveClientName(String userAgent) {
+    String resolved = null;
+    if (userAgent != null && !userAgent.isBlank()) {
+      String ua = userAgent.toLowerCase(Locale.ROOT);
+      resolved =
+          UA_MATCHERS.stream()
+              .filter(matcher -> matcher.matches().test(ua))
+              .map(UaMatcher::label)
+              .findFirst()
+              .orElseGet(() -> fallbackProductToken(userAgent));
+    }
+    return sanitize(resolved);
+  }
+
+  /**
+   * First product token of the User-Agent (the bit before the first slash or space), capitalised.
+   * Avoids leaking version strings into the dashboard while still surfacing unknown clients with a
+   * human-readable label. Returns {@code null} when no usable token is present.
+   */
+  private static String fallbackProductToken(String userAgent) {
+    String head = userAgent.split("[\\s/]", 2)[0];
+    String result = null;
+    if (!head.isBlank()) {
+      result = head.substring(0, 1).toUpperCase(Locale.ROOT) + head.substring(1);
+    }
+    return result;
+  }
+
+  /**
+   * Trims whitespace, strips ISO control characters, and caps the value to
+   * {@link #MAX_CLIENT_NAME_LENGTH} characters. The User-Agent header is attacker-controlled so
+   * the persisted value must be sanitised before it reaches the database or the dashboard.
+   */
+  private static String sanitize(String value) {
+    String result = null;
+    if (value != null) {
+      StringBuilder sb = new StringBuilder(value.length());
+      value.codePoints().filter(cp -> !Character.isISOControl(cp)).forEach(sb::appendCodePoint);
+      String stripped = sb.toString().trim();
+      if (!stripped.isEmpty()) {
+        result =
+            stripped.length() > MAX_CLIENT_NAME_LENGTH
+                ? stripped.substring(0, MAX_CLIENT_NAME_LENGTH)
+                : stripped;
+      }
+    }
+    return result;
   }
 }

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

@@ -250,14 +250,23 @@ protected McpStatelessServerFeatures.SyncToolSpecification getTool(McpSchema.Too
           CatalogSecurityContext securityContext =
               jwtFilter.getCatalogSecurityContext((String) context.get("Authorization"));
           String userName = securityContext.getUserPrincipal().getName();
-          McpSchema.CallToolResult result = null;
+          String clientName =
+              (String)
+                  context.get(org.openmetadata.mcp.AuthEnrichedMcpContextExtractor.CLIENT_NAME);
+          org.openmetadata.mcp.tools.DefaultToolContext.CallToolOutcome outcome = null;
           try {
             ImpersonationContext.setImpersonatedBy(getMcpBotName());
-            result = toolContext.callTool(authorizer, limits, tool.name(), securityContext, req);
-            return result;
+            outcome =
+                toolContext.callToolWithMetadata(
+                    authorizer, limits, tool.name(), securityContext, req);
+            return outcome.result();
           } finally {
-            boolean success = result != null && !Boolean.TRUE.equals(result.isError());
-            McpUsageRecorder.record(tool.name(), userName, success);
+            boolean success = outcome != null && !Boolean.TRUE.equals(outcome.result().isError());
+            Long latencyMs = outcome != null ? outcome.latencyMs() : null;
+            org.openmetadata.schema.entity.app.mcp.McpToolCallUsage.ErrorCategory category =
+                outcome != null ? outcome.errorCategory() : null;
+            McpUsageRecorder.record(
+                tool.name(), userName, success, latencyMs, category, clientName);
             ImpersonationContext.clear();
           }
         });

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

@@ -4,8 +4,11 @@
 
 import io.modelcontextprotocol.spec.McpSchema;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
+import java.util.function.Predicate;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage;
 import org.openmetadata.schema.utils.JsonUtils;
 import org.openmetadata.service.limits.Limits;
 import org.openmetadata.service.security.AuthorizationException;
@@ -32,6 +35,22 @@ public McpSchema.CallToolResult callTool(
       String toolName,
       CatalogSecurityContext securityContext,
       McpSchema.CallToolRequest request) {
+    return callToolWithMetadata(authorizer, limits, toolName, securityContext, request).result();
+  }
+
+  /**
+   * Phase 3 entry point. Returns the tool result alongside the metadata the {@link
+   * org.openmetadata.mcp.usage.McpUsageRecorder} needs (latency + error category). Kept as a
+   * separate method so the legacy single-result signature stays available for external callers
+   * that haven't migrated yet.
+   */
+  public CallToolOutcome callToolWithMetadata(
+      Authorizer authorizer,
+      Limits limits,
+      String toolName,
+      CatalogSecurityContext securityContext,
+      McpSchema.CallToolRequest request) {
+    long startNanos = System.nanoTime();
     LOG.info(
         "Catalog Principal: {} is trying to call the tool: {}",
         securityContext.getUserPrincipal().getName(),
@@ -97,47 +116,159 @@ public McpSchema.CallToolResult callTool(
           result = new CreateDataProductTool().execute(authorizer, limits, securityContext, params);
           break;
         default:
-          return McpSchema.CallToolResult.builder()
-              .content(
-                  List.of(
-                      new McpSchema.TextContent(
-                          JsonUtils.pojoToJson(Map.of("error", "Unknown function: " + toolName)))))
-              .isError(true)
-              .build();
+          return new CallToolOutcome(
+              McpSchema.CallToolResult.builder()
+                  .content(
+                      List.of(
+                          new McpSchema.TextContent(
+                              JsonUtils.pojoToJson(
+                                  Map.of("error", "Unknown function: " + toolName)))))
+                  .isError(true)
+                  .build(),
+              elapsedMs(startNanos),
+              McpToolCallUsage.ErrorCategory.VALIDATION);
       }
 
-      return McpSchema.CallToolResult.builder()
-          .content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
-          .isError(false)
-          .build();
+      return new CallToolOutcome(
+          McpSchema.CallToolResult.builder()
+              .content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
+              .isError(false)
+              .build(),
+          elapsedMs(startNanos),
+          null);
     } catch (AuthorizationException ex) {
       LOG.warn("Authorization error: {}", ex.getMessage());
-      return McpSchema.CallToolResult.builder()
-          .content(
-              List.of(
-                  new McpSchema.TextContent(
-                      JsonUtils.pojoToJson(
-                          Map.of(
-                              "error",
-                              String.format("Authorization error: %s", ex.getMessage()),
-                              "statusCode",
-                              403)))))
-          .isError(true)
-          .build();
+      return new CallToolOutcome(
+          McpSchema.CallToolResult.builder()
+              .content(
+                  List.of(
+                      new McpSchema.TextContent(
+                          JsonUtils.pojoToJson(
+                              Map.of(
+                                  "error",
+                                  String.format("Authorization error: %s", ex.getMessage()),
+                                  "statusCode",
+                                  403)))))
+              .isError(true)
+              .build(),
+          elapsedMs(startNanos),
+          McpToolCallUsage.ErrorCategory.AUTH);
     } catch (Exception ex) {
       LOG.error("Error executing tool '{}': {}", toolName, ex.getMessage(), ex);
-      return McpSchema.CallToolResult.builder()
-          .content(
-              List.of(
-                  new McpSchema.TextContent(
-                      JsonUtils.pojoToJson(
-                          Map.of(
-                              "error",
-                              String.format("Error executing tool: %s", ex.getMessage()),
-                              "statusCode",
-                              500)))))
-          .isError(true)
-          .build();
+      return new CallToolOutcome(
+          McpSchema.CallToolResult.builder()
+              .content(
+                  List.of(
+                      new McpSchema.TextContent(
+                          JsonUtils.pojoToJson(
+                              Map.of(
+                                  "error",
+                                  String.format("Error executing tool: %s", ex.getMessage()),
+                                  "statusCode",
+                                  500)))))
+              .isError(true)
+              .build(),
+          elapsedMs(startNanos),
+          classifyException(ex));
     }
   }
+
+  /**
+   * Maps an arbitrary exception type to one of the {@link McpToolCallUsage.ErrorCategory} values.
+   * Walks the cause chain because the tool wrappers usually rethrow framework errors wrapped in
+   * a {@link RuntimeException}. Defaults to {@link McpToolCallUsage.ErrorCategory#INTERNAL} when
+   * no specific bucket matches.
+   */
+  static McpToolCallUsage.ErrorCategory classifyException(Throwable t) {
+    McpToolCallUsage.ErrorCategory result = McpToolCallUsage.ErrorCategory.INTERNAL;
+    Throwable cursor = t;
+    while (cursor != null && result == McpToolCallUsage.ErrorCategory.INTERNAL) {
+      McpToolCallUsage.ErrorCategory match = matchCategory(cursor);
+      if (match != null) {
+        result = match;
+      } else {
+        Throwable next = cursor.getCause();
+        cursor = (next == null || next == cursor) ? null : next;
+      }
+    }
+    return result;
+  }
+
+  /**
+   * Pairing of an exception (name, message) predicate with the bucket it should produce. Kept
+   * as a static table so adding a new category (or extending an existing one with a new keyword)
+   * is a one-line change rather than another {@code else if} branch.
+   */
+  private record CategoryMatcher(
+      Predicate<ExceptionMeta> matches, McpToolCallUsage.ErrorCategory category) {}
+
+  /** Lower-cased name + message pair so each matcher inspects both without re-parsing. */
+  private record ExceptionMeta(String name, String message) {}
+
+  /**
+   * Ordered category table. Check order matters: more specific patterns sit before broader ones so
+   * a {@code RateLimitException} doesn't get caught by the generic message-substring rules below
+   * it. {@code AUTH} sits above {@code VALIDATION} because some auth exceptions ({@code
+   * AuthorizationException}) extend {@code IllegalArgumentException}-style hierarchies and would
+   * otherwise be mis-bucketed.
+   */
+  private static final List<CategoryMatcher> CATEGORY_MATCHERS =
+      List.of(
+          new CategoryMatcher(
+              meta -> meta.name().contains("RateLimit") || meta.message().contains("rate limit"),
+              McpToolCallUsage.ErrorCategory.RATE_LIMIT),
+          new CategoryMatcher(
+              meta ->
+                  meta.name().contains("Authorization")
+                      || meta.name().contains("Forbidden")
+                      || meta.name().contains("Unauthorized")
+                      || meta.message().contains("forbidden")
+                      || meta.message().contains("unauthorized")
+                      || meta.message().contains("access denied")
+                      || meta.message().contains("permission denied"),
+              McpToolCallUsage.ErrorCategory.AUTH),
+          new CategoryMatcher(
+              meta ->
+                  meta.name().contains("Validation")
+                      || meta.name().contains("IllegalArgument")
+                      || meta.name().contains("BadRequest")
+                      || meta.message().contains("invalid argument"),
+              McpToolCallUsage.ErrorCategory.VALIDATION),
+          new CategoryMatcher(
+              meta ->
+                  meta.name().contains("Timeout")
+                      || meta.message().contains("timeout")
+                      || meta.message().contains("timed out"),
+              McpToolCallUsage.ErrorCategory.TIMEOUT));
+
+  /**
+   * Returns the category that matches the supplied throwable's name or message, or {@code null}
+   * when no specific bucket applies. Kept separate from {@link #classifyException} so the
+   * cause-chain walk reads as a single linear loop.
+   */
+  private static McpToolCallUsage.ErrorCategory matchCategory(Throwable cursor) {
+    ExceptionMeta meta =
+        new ExceptionMeta(
+            cursor.getClass().getSimpleName(),
+            cursor.getMessage() == null ? "" : cursor.getMessage().toLowerCase(Locale.ROOT));
+    return CATEGORY_MATCHERS.stream()
+        .filter(matcher -> matcher.matches().test(meta))
+        .map(CategoryMatcher::category)
+        .findFirst()
+        .orElse(null);
+  }
+
+  private static long elapsedMs(long startNanos) {
+    return (System.nanoTime() - startNanos) / 1_000_000L;
+  }
+
+  /**
+   * Phase 3 — tuple returned by {@link #callToolWithMetadata} so the MCP server can record the
+   * call with full diagnostic detail without re-classifying the exception or re-measuring the
+   * latency at its level.
+   */
+  public record CallToolOutcome(
+      McpSchema.CallToolResult result,
+      long latencyMs,
+      McpToolCallUsage.ErrorCategory errorCategory) {}
 }