1 files changed, 158 insertions, 49 deletions
diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts
index 9e814bba..3dd62e2b 100644
--- a/packages/core/src/tools/mcp-tool.ts
+++ b/packages/core/src/tools/mcp-tool.ts
@@ -22,6 +22,40 @@ import {
 
 type ToolParams = Record<string, unknown>;
 
+// Discriminated union for MCP Content Blocks to ensure type safety.
+type McpTextBlock = {
+  type: 'text';
+  text: string;
+};
+
+type McpMediaBlock = {
+  type: 'image' | 'audio';
+  mimeType: string;
+  data: string;
+};
+
+type McpResourceBlock = {
+  type: 'resource';
+  resource: {
+    text?: string;
+    blob?: string;
+    mimeType?: string;
+  };
+};
+
+type McpResourceLinkBlock = {
+  type: 'resource_link';
+  uri: string;
+  title?: string;
+  name?: string;
+};
+
+type McpContentBlock =
+  | McpTextBlock
+  | McpMediaBlock
+  | McpResourceBlock
+  | McpResourceLinkBlock;
+
 export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
   private static readonly allowlist: Set<string> = new Set();
 
@@ -114,70 +148,145 @@ export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
       },
     ];
 
-    const responseParts: Part[] = await this.mcpTool.callTool(functionCalls);
+    const rawResponseParts = await this.mcpTool.callTool(functionCalls);
+    const transformedParts = transformMcpContentToParts(rawResponseParts);
 
     return {
-      llmContent: responseParts,
-      returnDisplay: getStringifiedResultForDisplay(responseParts),
+      llmContent: transformedParts,
+      returnDisplay: getStringifiedResultForDisplay(rawResponseParts),
     };
   }
 }
 
+function transformTextBlock(block: McpTextBlock): Part {
+  return { text: block.text };
+}
+
+function transformImageAudioBlock(
+  block: McpMediaBlock,
+  toolName: string,
+): Part[] {
+  return [
+    {
+      text: `[Tool '${toolName}' provided the following ${
+        block.type
+      } data with mime-type: ${block.mimeType}]`,
+    },
+    {
+      inlineData: {
+        mimeType: block.mimeType,
+        data: block.data,
+      },
+    },
+  ];
+}
+
+function transformResourceBlock(
+  block: McpResourceBlock,
+  toolName: string,
+): Part | Part[] | null {
+  const resource = block.resource;
+  if (resource?.text) {
+    return { text: resource.text };
+  }
+  if (resource?.blob) {
+    const mimeType = resource.mimeType || 'application/octet-stream';
+    return [
+      {
+        text: `[Tool '${toolName}' provided the following embedded resource with mime-type: ${mimeType}]`,
+      },
+      {
+        inlineData: {
+          mimeType,
+          data: resource.blob,
+        },
+      },
+    ];
+  }
+  return null;
+}
+
+function transformResourceLinkBlock(block: McpResourceLinkBlock): Part {
+  return {
+    text: `Resource Link: ${block.title || block.name} at ${block.uri}`,
+  };
+}
+
 /**
- * Processes an array of `Part` objects, primarily from a tool's execution result,
- * to generate a user-friendly string representation, typically for display in a CLI.
- *
- * The `result` array can contain various types of `Part` objects:
- * 1. `FunctionResponse` parts:
- *    - If the `response.content` of a `FunctionResponse` is an array consisting solely
- *      of `TextPart` objects, their text content is concatenated into a single string.
- *      This is to present simple textual outputs directly.
- *    - If `response.content` is an array but contains other types of `Part` objects (or a mix),
- *      the `content` array itself is preserved. This handles structured data like JSON objects or arrays
- *      returned by a tool.
- *    - If `response.content` is not an array or is missing, the entire `functionResponse`
- *      object is preserved.
- * 2. Other `Part` types (e.g., `TextPart` directly in the `result` array):
- *    - These are preserved as is.
- *
- * All processed parts are then collected into an array, which is JSON.stringify-ed
- * with indentation and wrapped in a markdown JSON code block.
+ * Transforms the raw MCP content blocks from the SDK response into a
+ * standard GenAI Part array.
+ * @param sdkResponse The raw Part[] array from `mcpTool.callTool()`.
+ * @returns A clean Part[] array ready for the scheduler.
  */
-function getStringifiedResultForDisplay(result: Part[]) {
-  if (!result || result.length === 0) {
-    return '```json\n[]\n```';
+function transformMcpContentToParts(sdkResponse: Part[]): Part[] {
+  const funcResponse = sdkResponse?.[0]?.functionResponse;
+  const mcpContent = funcResponse?.response?.content as McpContentBlock[];
+  const toolName = funcResponse?.name || 'unknown tool';
+
+  if (!Array.isArray(mcpContent)) {
+    return [{ text: '[Error: Could not parse tool response]' }];
   }
 
-  const processFunctionResponse = (part: Part) => {
-    if (part.functionResponse) {
-      const responseContent = part.functionResponse.response?.content;
-      if (responseContent && Array.isArray(responseContent)) {
-        // Check if all parts in responseContent are simple TextParts
-        const allTextParts = responseContent.every(
-          (p: Part) => p.text !== undefined,
-        );
-        if (allTextParts) {
-          return responseContent.map((p: Part) => p.text).join('');
-        }
-        // If not all simple text parts, return the array of these content parts for JSON stringification
-        return responseContent;
+  const transformed = mcpContent.flatMap(
+    (block: McpContentBlock): Part | Part[] | null => {
+      switch (block.type) {
+        case 'text':
+          return transformTextBlock(block);
+        case 'image':
+        case 'audio':
+          return transformImageAudioBlock(block, toolName);
+        case 'resource':
+          return transformResourceBlock(block, toolName);
+        case 'resource_link':
+          return transformResourceLinkBlock(block);
+        default:
+          return null;
       }
+    },
+  );
 
-      // If no content, or not an array, or not a functionResponse, stringify the whole functionResponse part for inspection
-      return part.functionResponse;
-    }
-    return part; // Fallback for unexpected structure or non-FunctionResponsePart
-  };
+  return transformed.filter((part): part is Part => part !== null);
+}
 
-  const processedResults =
-    result.length === 1
-      ? processFunctionResponse(result[0])
-      : result.map(processFunctionResponse);
-  if (typeof processedResults === 'string') {
-    return processedResults;
+/**
+ * Processes the raw response from the MCP tool to generate a clean,
+ * human-readable string for display in the CLI. It summarizes non-text
+ * content and presents text directly.
+ *
+ * @param rawResponse The raw Part[] array from the GenAI SDK.
+ * @returns A formatted string representing the tool's output.
+ */
+function getStringifiedResultForDisplay(rawResponse: Part[]): string {
+  const mcpContent = rawResponse?.[0]?.functionResponse?.response
+    ?.content as McpContentBlock[];
+
+  if (!Array.isArray(mcpContent)) {
+    return '```json\n' + JSON.stringify(rawResponse, null, 2) + '\n```';
   }
 
-  return '```json\n' + JSON.stringify(processedResults, null, 2) + '\n```';
+  const displayParts = mcpContent.map((block: McpContentBlock): string => {
+    switch (block.type) {
+      case 'text':
+        return block.text;
+      case 'image':
+        return `[Image: ${block.mimeType}]`;
+      case 'audio':
+        return `[Audio: ${block.mimeType}]`;
+      case 'resource_link':
+        return `[Link to ${block.title || block.name}: ${block.uri}]`;
+      case 'resource':
+        if (block.resource?.text) {
+          return block.resource.text;
+        }
+        return `[Embedded Resource: ${
+          block.resource?.mimeType || 'unknown type'
+        }]`;
+      default:
+        return `[Unknown content type: ${(block as { type: string }).type}]`;
+    }
+  });
+
+  return displayParts.join('\n');
 }
 
 /** Visible for testing */