refactor(mcp): shared trim util + global response size-budget net (open-metadata#28764)

* refactor(mcp): shared response-trim + params utils, global size-budget net, error-message null-guards

* refactor(mcp): expose serializeWithinBudget so Collate dispatcher shares the size-budget net

* refactor(mcp): move McpParams/McpResponseTrim to util package, guard RCA error messages

Author: Vishnuujain

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

@@ -11,6 +11,7 @@
 import java.util.Set;
 import java.util.function.Predicate;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.mcp.util.McpResponseTrim;
 import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage;
 import org.openmetadata.schema.utils.JsonUtils;
 import org.openmetadata.service.limits.Limits;
@@ -26,6 +27,9 @@ public class DefaultToolContext {
   private static final int STATUS_TOO_MANY_REQUESTS = 429;
   private static final int STATUS_INTERNAL_ERROR = 500;
   private static final int STATUS_GATEWAY_TIMEOUT = 504;
+  private static final String OVERSIZED_ADVICE =
+      "Response exceeded the size limit and was withheld. Narrow your request — use a more specific "
+          + "query, request fewer results, or fetch a single entity by its fullyQualifiedName.";
 
   public DefaultToolContext() {}
 
@@ -145,7 +149,7 @@ public CallToolOutcome callToolWithMetadata(
 
       return new CallToolOutcome(
           McpSchema.CallToolResult.builder()
-              .content(List.of(new McpSchema.TextContent(JsonUtils.pojoToJson(result))))
+              .content(List.of(new McpSchema.TextContent(serializeWithinBudget(result, toolName))))
               .isError(false)
               .build(),
           elapsedMs(startNanos),
@@ -160,7 +164,8 @@ public CallToolOutcome callToolWithMetadata(
                           JsonUtils.pojoToJson(
                               Map.of(
                                   "error",
-                                  String.format("Authorization error: %s", ex.getMessage()),
+                                  String.format(
+                                      "Authorization error: %s", McpResponseTrim.safeMessage(ex)),
                                   "statusCode",
                                   STATUS_FORBIDDEN)))))
               .isError(true)
@@ -177,7 +182,8 @@ public CallToolOutcome callToolWithMetadata(
                           JsonUtils.pojoToJson(
                               Map.of(
                                   "error",
-                                  String.format("Error executing tool: %s", ex.getMessage()),
+                                  String.format(
+                                      "Error executing tool: %s", McpResponseTrim.safeMessage(ex)),
                                   "statusCode",
                                   resolveStatusCode(ex))))))
               .isError(true)
@@ -300,6 +306,32 @@ private static long elapsedMs(long startNanos) {
     return (System.nanoTime() - startNanos) / 1_000_000L;
   }
 
+  /**
+   * Serializes a tool result once and, only when it exceeds {@link
+   * McpResponseTrim#MAX_RESPONSE_CHARS}, replaces it with a generic {@code truncated:true} envelope.
+   * This is the dispatch-level floor that bounds tools without their own per-tool trim ({@code
+   * get_entity_details}, {@code get_test_definitions}) and backstops the rest. The happy path
+   * serializes exactly once; the re-serialization runs only on the rare oversized path.
+   *
+   * <p>Public so the Collate dispatcher ({@code CollateToolContext}), which builds its own success
+   * result for Collate-only tools, applies the same floor instead of re-implementing it.
+   */
+  public static String serializeWithinBudget(Object result, String toolName) {
+    String serialized = JsonUtils.pojoToJson(result);
+    if (serialized.length() > McpResponseTrim.MAX_RESPONSE_CHARS) {
+      LOG.warn(
+          "[MCP] tool '{}' response {} chars exceeds {} budget; returning truncation envelope",
+          toolName,
+          serialized.length(),
+          McpResponseTrim.MAX_RESPONSE_CHARS);
+      Map<String, Object> capped =
+          McpResponseTrim.oversizedEnvelope(
+              serialized.length(), Map.of("tool", toolName), OVERSIZED_ADVICE);
+      serialized = JsonUtils.pojoToJson(capped);
+    }
+    return serialized;
+  }
+
   /**
    * 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

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

@@ -7,6 +7,7 @@
 import java.util.List;
 import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.mcp.util.McpResponseTrim;
 import org.openmetadata.schema.utils.JsonUtils;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.limits.Limits;
@@ -42,8 +43,7 @@ public class GetEntityTool implements McpTool {
           "tagSources",
           "descriptionSources",
           "columnDescriptionStatus",
-          "descriptionStatus",
-          "embeddings");
+          "descriptionStatus");
 
   @Override
   public Map<String, Object> execute(
@@ -74,6 +74,7 @@ private static Map<String, Object> cleanEntityResponse(Map<String, Object> entit
     }
     Map<String, Object> cleaned = new HashMap<>(entityData);
     EXCLUDE_FIELDS.forEach(cleaned::remove);
+    McpResponseTrim.VECTOR_NOISE_FIELDS.forEach(cleaned::remove);
     return cleaned;
   }
 

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

@@ -13,6 +13,8 @@
 import java.util.Set;
 import java.util.UUID;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.mcp.util.McpParams;
+import org.openmetadata.mcp.util.McpResponseTrim;
 import org.openmetadata.schema.type.ColumnLineage;
 import org.openmetadata.schema.type.Edge;
 import org.openmetadata.schema.type.EntityLineage;
@@ -43,13 +45,6 @@ public class GetLineageTool implements McpTool {
   private static final int DEFAULT_DEPTH = 3;
   // Maximum depth to prevent exponential response growth (lineage graphs can explode)
   private static final int MAX_DEPTH = 10;
-  // SQL is the single heaviest field; keep the gist of the transform, cap the size
-  private static final int SQL_MAX_LENGTH = 500;
-  // Free-text markdown (pipeline / edge descriptions) is capped for the same reason as SQL:
-  // a single long pipeline doc string can be shared across many edges and reintroduce bloat
-  private static final int TEXT_MAX_LENGTH = 500;
-  // Final safety net mirroring SearchMetadataTool: even slimmed, a wide graph can blow the limit
-  private static final int MAX_RESPONSE_CHARS = 100_000;
   private static final String RELATIONSHIP_SQL = "sql";
 
   @JsonInclude(JsonInclude.Include.NON_NULL)
@@ -94,9 +89,9 @@ public Map<String, Object> execute(
         securityContext,
         new OperationContext(entityType, MetadataOperation.VIEW_BASIC),
         new ResourceContext<>(entityType));
-    int upstreamDepth = parseDepthParameter(params.get("upstreamDepth"), DEFAULT_DEPTH);
-    int downstreamDepth = parseDepthParameter(params.get("downstreamDepth"), DEFAULT_DEPTH);
-    boolean includeColumnLineage = parseBooleanParameter(params.get("includeColumnLineage"));
+    int upstreamDepth = clampDepth(McpParams.getInt(params, "upstreamDepth", DEFAULT_DEPTH));
+    int downstreamDepth = clampDepth(McpParams.getInt(params, "downstreamDepth", DEFAULT_DEPTH));
+    boolean includeColumnLineage = McpParams.getBoolean(params, "includeColumnLineage", false);
     LOG.info(
         "Getting lineage for entity type: {}, FQN: {}, upstreamDepth: {}, downstreamDepth: {}, "
             + "includeColumnLineage: {}",
@@ -212,11 +207,7 @@ private static String relationshipType(EntityReference pipeline) {
   }
 
   private static String truncateText(String text) {
-    String result = text;
-    if (text != null && text.length() > TEXT_MAX_LENGTH) {
-      result = text.substring(0, TEXT_MAX_LENGTH) + "...";
-    }
-    return result;
+    return McpResponseTrim.truncate(text, McpResponseTrim.TEXT_MAX_LENGTH);
   }
 
   private static String sourceValue(LineageDetails details) {
@@ -227,8 +218,8 @@ private static SqlText truncateSqlQuery(LineageDetails details) {
     String sql = details != null ? details.getSqlQuery() : null;
     SqlText result = new SqlText(null, null);
     if (sql != null) {
-      boolean tooLong = sql.length() > SQL_MAX_LENGTH;
-      String value = tooLong ? sql.substring(0, SQL_MAX_LENGTH) + "..." : sql;
+      boolean tooLong = sql.length() > McpResponseTrim.SQL_MAX_LENGTH;
+      String value = McpResponseTrim.truncate(sql, McpResponseTrim.SQL_MAX_LENGTH);
       result = new SqlText(value, tooLong ? Boolean.TRUE : null);
     }
     return result;
@@ -253,9 +244,9 @@ private static String refName(EntityReference ref) {
   @VisibleForTesting
   static Map<String, Object> enforceSizeBudget(SlimLineage slim) {
     Map<String, Object> response = JsonUtils.getMap(slim);
-    int responseSize = JsonUtils.pojoToJson(response).length();
+    int responseSize = McpResponseTrim.serializedLength(response);
     Map<String, Object> result = response;
-    if (responseSize > MAX_RESPONSE_CHARS) {
+    if (responseSize > McpResponseTrim.MAX_RESPONSE_CHARS) {
       result = oversizedHint(slim, responseSize);
     }
     return result;
@@ -276,44 +267,19 @@ private static Map<String, Object> oversizedHint(SlimLineage slim, int size) {
         String.format(
             "Lineage response exceeded %d characters (was %d). Reduce upstreamDepth/downstreamDepth,"
                 + " or keep includeColumnLineage disabled, to get a smaller graph.",
-            MAX_RESPONSE_CHARS, size));
+            McpResponseTrim.MAX_RESPONSE_CHARS, size));
     return hint;
   }
 
-  private static boolean parseBooleanParameter(Object value) {
-    boolean result = false;
-    if (value instanceof Boolean bool) {
-      result = bool;
-    } else if (value instanceof String string) {
-      result = Boolean.parseBoolean(string);
-    }
-    return result;
-  }
-
   /**
-   * Parses depth parameter with default value and enforces maximum limit to prevent excessive
-   * response sizes that could overwhelm LLM context.
+   * Clamps a requested depth into {@code [1, MAX_DEPTH]} to prevent excessive response sizes that
+   * could overwhelm LLM context. Parsing is delegated to {@link McpParams}; the valid range is
+   * specific to this tool, so the clamp stays here.
    */
-  private static int parseDepthParameter(Object depthObj, int defaultValue) {
-    int depth = defaultValue;
-    if (depthObj instanceof Number number) {
-      depth = number.intValue();
-    } else if (depthObj instanceof String string) {
-      depth = parseDepthString(string, defaultValue);
-    }
+  private static int clampDepth(int depth) {
     return Math.min(Math.max(depth, 1), MAX_DEPTH);
   }
 
-  private static int parseDepthString(String value, int defaultValue) {
-    int depth = defaultValue;
-    try {
-      depth = Integer.parseInt(value);
-    } catch (NumberFormatException e) {
-      depth = defaultValue;
-    }
-    return depth;
-  }
-
   @Override
   public Map<String, Object> execute(
       Authorizer authorizer,

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

@@ -15,6 +15,8 @@
 import java.util.Map;
 import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.mcp.util.McpParams;
+import org.openmetadata.mcp.util.McpResponseTrim;
 import org.openmetadata.schema.api.lineage.LineageDirection;
 import org.openmetadata.schema.api.lineage.SearchLineageRequest;
 import org.openmetadata.schema.api.lineage.SearchLineageResult;
@@ -36,13 +38,10 @@ public class RootCauseAnalysisTool implements McpTool {
 
   private static final int DEFAULT_DEPTH = 3;
   private static final int MAX_DEPTH = 10;
-  // Slimming budgets mirror GetLineageTool so RCA's lineage-derived payload stays within
+  // Slimming budgets come from McpResponseTrim so RCA's lineage-derived payload stays within
   // LLM/MCP context limits. The backend (searchDataQualityLineage / searchLineageWithDirection)
   // is shared with the UI LineageResource and is never touched — we only transform the
   // in-memory result before returning it to the MCP client.
-  private static final int SQL_MAX_LENGTH = 500;
-  private static final int TEXT_MAX_LENGTH = 500;
-  private static final int MAX_RESPONSE_CHARS = 100_000;
   private static final String RELATIONSHIP_SQL = "sql";
 
   @Override
@@ -52,12 +51,12 @@ public Map<String, Object> execute(
       Map<String, Object> parameters) {
     String fqn = (String) parameters.get("fqn");
     String entityType = (String) parameters.getOrDefault("entityType", "table");
-    int upstreamDepth = clampDepth(parseIntParam(parameters.get("upstreamDepth"), DEFAULT_DEPTH));
+    int upstreamDepth = clampDepth(McpParams.getInt(parameters, "upstreamDepth", DEFAULT_DEPTH));
     int downstreamDepth =
-        clampDepth(parseIntParam(parameters.get("downstreamDepth"), DEFAULT_DEPTH));
+        clampDepth(McpParams.getInt(parameters, "downstreamDepth", DEFAULT_DEPTH));
     String queryFilter = (String) parameters.get("queryFilter");
-    boolean includeDeleted = parseBooleanParam(parameters.get("includeDeleted"), false);
-    boolean includeColumns = parseBooleanParam(parameters.get("includeColumnLineage"), false);
+    boolean includeDeleted = McpParams.getBoolean(parameters, "includeDeleted", false);
+    boolean includeColumns = McpParams.getBoolean(parameters, "includeColumnLineage", false);
 
     if (fqn == null || fqn.trim().isEmpty()) {
       throw new IllegalArgumentException("Parameter 'fqn' is required and cannot be empty");
@@ -81,11 +80,12 @@ public Map<String, Object> execute(
       return analyze(request);
     } catch (IOException e) {
       LOG.error("IOException during root cause analysis for entity: {}", fqn, e);
-      throw new RuntimeException("Failed to perform root cause analysis: " + e.getMessage(), e);
+      throw new RuntimeException(
+          "Failed to perform root cause analysis: " + McpResponseTrim.safeMessage(e), e);
     } catch (Exception e) {
       LOG.error("Unexpected error during root cause analysis for entity: {}", fqn, e);
       throw new RuntimeException(
-          "Unexpected error during root cause analysis: " + e.getMessage(), e);
+          "Unexpected error during root cause analysis: " + McpResponseTrim.safeMessage(e), e);
     }
   }
 
@@ -172,7 +172,8 @@ private Map<String, Object> buildDownstreamAnalysis(RcaRequest request) {
       addDownstreamEdges(downstreamAnalysis, downstreamResult, request.includeColumns());
     } catch (Exception e) {
       LOG.warn("Failed to perform downstream impact analysis for entity: {}", request.fqn(), e);
-      downstreamAnalysis.put("error", "Failed to analyze downstream impact: " + e.getMessage());
+      downstreamAnalysis.put(
+          "error", "Failed to analyze downstream impact: " + McpResponseTrim.safeMessage(e));
     }
     return downstreamAnalysis;
   }
@@ -315,34 +316,30 @@ private static String relationshipType(Object pipeline) {
 
   private static void applyDescription(Map<String, Object> slim, Object description) {
     if (description instanceof String text && !text.isEmpty()) {
-      slim.put("description", truncate(text, TEXT_MAX_LENGTH));
+      slim.put("description", McpResponseTrim.truncate(text, McpResponseTrim.TEXT_MAX_LENGTH));
     }
   }
 
   private static void applySqlQuery(Map<String, Object> slim, Object sqlQuery) {
     if (sqlQuery instanceof String sql && !sql.isEmpty()) {
-      slim.put("sqlQuery", truncate(sql, SQL_MAX_LENGTH));
-      if (sql.length() > SQL_MAX_LENGTH) {
+      slim.put("sqlQuery", McpResponseTrim.truncate(sql, McpResponseTrim.SQL_MAX_LENGTH));
+      if (sql.length() > McpResponseTrim.SQL_MAX_LENGTH) {
         slim.put("sqlTruncated", Boolean.TRUE);
       }
     }
   }
 
   private static void truncateDescriptionInPlace(Map<String, Object> map) {
     Object description = map.get("description");
-    if (description instanceof String text && text.length() > TEXT_MAX_LENGTH) {
-      map.put("description", truncate(text, TEXT_MAX_LENGTH));
+    if (description instanceof String text && text.length() > McpResponseTrim.TEXT_MAX_LENGTH) {
+      map.put("description", McpResponseTrim.truncate(text, McpResponseTrim.TEXT_MAX_LENGTH));
     }
   }
 
-  private static String truncate(String value, int maxLength) {
-    return value.length() > maxLength ? value.substring(0, maxLength) + "..." : value;
-  }
-
   @VisibleForTesting
   static Map<String, Object> enforceSizeBudget(Map<String, Object> result) {
     Map<String, Object> output = result;
-    if (JsonUtils.pojoToJson(result).length() > MAX_RESPONSE_CHARS) {
+    if (McpResponseTrim.serializedLength(result) > McpResponseTrim.MAX_RESPONSE_CHARS) {
       output = oversizedHint(result);
     }
     return output;
@@ -417,36 +414,6 @@ private static void putIfPresent(Map<String, Object> map, String key, Object val
     }
   }
 
-  private static int parseIntParam(Object value, int defaultValue) {
-    if (value == null) {
-      return defaultValue;
-    }
-    if (value instanceof Number number) {
-      return number.intValue();
-    }
-    if (value instanceof String string) {
-      try {
-        return Integer.parseInt(string);
-      } catch (NumberFormatException e) {
-        return defaultValue;
-      }
-    }
-    return defaultValue;
-  }
-
-  private static boolean parseBooleanParam(Object value, boolean defaultValue) {
-    if (value == null) {
-      return defaultValue;
-    }
-    if (value instanceof Boolean bool) {
-      return bool;
-    }
-    if (value instanceof String string) {
-      return "true".equalsIgnoreCase(string);
-    }
-    return defaultValue;
-  }
-
   @Override
   public Map<String, Object> execute(
       Authorizer authorizer,