fix: preserve ImageContent and AudioContent blocks through proxy (#368) (#369)

The proxy was JSON-serializing upstream CallToolResult and wrapping the
entire payload in a single TextContent block. This destroyed the native
content type of any ImageContent, AudioContent, or other non-text block,
forcing vision-capable LLMs to receive base64 as text — 50-250x more
token-consumptive than the native vision path.

Fix: introduce forwardContentResult() which walks CallToolResult.Content
and preserves non-text blocks unchanged while applying truncation only to
TextContent.Text. Applied in three sites:
  - handleCallToolVariant (call_tool_read/write/destructive)
  - handleCallTool (legacy handler)
  - makeDirectModeHandler (direct routing mode)

Tests:
  - internal/server/content_forward_test.go — unit tests covering
    image preservation, selective truncation, text-only case, and the
    non-CallToolResult fallback path
  - internal/server/e2e_content_forward_test.go — end-to-end test that
    spins up a mock upstream returning text+image+audio and verifies
    the downstream client receives all three native content types
    (fails before this fix with "expected 3 content blocks, got 1")
  - updated TestE2E_ToolCalling: it was inadvertently asserting the
    buggy double-wrapping behavior; updated to reflect native forwarding

Co-authored-by: Claude Code <noreply@anthropic.com>

Author: Dumbris

internal/server/content_forward.go

@@ -0,0 +1,114 @@
+package server
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/mark3labs/mcp-go/mcp"
+
+	"github.com/smart-mcp-proxy/mcpproxy-go/internal/truncate"
+)
+
+// forwardContentResult preserves non-text content blocks (ImageContent, AudioContent,
+// EmbeddedResource) from an upstream CallToolResult while applying truncation only to
+// TextContent blocks. This fixes issue #368 where all content types were being
+// serialized to JSON and wrapped in a single TextContent, destroying the ability of
+// vision-capable LLMs to process images efficiently.
+//
+// Parameters:
+//   - result: the upstream CallToolResult (passed as interface{} from Manager.CallTool)
+//   - truncator: applies size limits to text blocks only
+//   - toolName, args: passed to truncator for caching
+//
+// Returns:
+//   - forwarded: the CallToolResult to return downstream (with original content types)
+//   - textRepresentation: a text rendering of the full result for logging/metrics
+//   - wasTruncated: whether any TextContent block was truncated
+//   - cacheStored: whether the full response was stored in the cache (for audit)
+//
+// If result is not a *mcp.CallToolResult, it falls back to JSON-serializing the whole
+// thing into a TextContent block (legacy behavior).
+func forwardContentResult(result interface{}, truncator *truncate.Truncator, toolName string, args map[string]interface{}) (forwarded *mcp.CallToolResult, textRepresentation string, wasTruncated bool) {
+	ctr, ok := result.(*mcp.CallToolResult)
+	if !ok || ctr == nil {
+		// Fallback: not a CallToolResult (should not happen with current upstream chain,
+		// but guard against future changes). Use legacy JSON-wrap behavior.
+		jsonBytes, err := json.Marshal(result)
+		if err != nil {
+			return mcp.NewToolResultError(fmt.Sprintf("Failed to serialize result: %v", err)), "", false
+		}
+		text := string(jsonBytes)
+		if truncator != nil && truncator.ShouldTruncate(text) {
+			tr := truncator.Truncate(text, toolName, args)
+			text = tr.TruncatedContent
+			wasTruncated = true
+		}
+		return mcp.NewToolResultText(text), text, wasTruncated
+	}
+
+	// Walk the Content slice, truncating only TextContent blocks.
+	// Build a parallel text representation for logging/metrics.
+	newContent := make([]mcp.Content, 0, len(ctr.Content))
+	var textBuilder []string
+	for _, c := range ctr.Content {
+		switch tc := c.(type) {
+		case mcp.TextContent:
+			txt := tc.Text
+			if truncator != nil && truncator.ShouldTruncate(txt) {
+				tr := truncator.Truncate(txt, toolName, args)
+				txt = tr.TruncatedContent
+				wasTruncated = true
+			}
+			tc.Text = txt
+			newContent = append(newContent, tc)
+			textBuilder = append(textBuilder, txt)
+		case mcp.ImageContent:
+			// Preserve image blocks unchanged. For logging, emit a placeholder.
+			newContent = append(newContent, tc)
+			textBuilder = append(textBuilder, fmt.Sprintf("[image:%s len=%d]", tc.MIMEType, len(tc.Data)))
+		case mcp.AudioContent:
+			// Preserve audio blocks unchanged. For logging, emit a placeholder.
+			newContent = append(newContent, tc)
+			textBuilder = append(textBuilder, fmt.Sprintf("[audio:%s len=%d]", tc.MIMEType, len(tc.Data)))
+		default:
+			// Unknown content type (e.g., EmbeddedResource). Forward as-is and
+			// best-effort JSON encode for logging.
+			newContent = append(newContent, c)
+			if b, err := json.Marshal(c); err == nil {
+				textBuilder = append(textBuilder, string(b))
+			}
+		}
+	}
+
+	forwarded = &mcp.CallToolResult{
+		Result:            ctr.Result,
+		Content:           newContent,
+		StructuredContent: ctr.StructuredContent,
+		IsError:           ctr.IsError,
+	}
+	textRepresentation = joinTextParts(textBuilder)
+	return forwarded, textRepresentation, wasTruncated
+}
+
+// joinTextParts concatenates text parts with a newline separator.
+// Equivalent to strings.Join but avoids importing strings just for this.
+func joinTextParts(parts []string) string {
+	switch len(parts) {
+	case 0:
+		return ""
+	case 1:
+		return parts[0]
+	}
+	total := len(parts) - 1 // separators
+	for _, p := range parts {
+		total += len(p)
+	}
+	out := make([]byte, 0, total)
+	for i, p := range parts {
+		if i > 0 {
+			out = append(out, '\n')
+		}
+		out = append(out, p...)
+	}
+	return string(out)
+}

internal/server/content_forward_test.go

@@ -0,0 +1,123 @@
+package server
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/smart-mcp-proxy/mcpproxy-go/internal/truncate"
+)
+
+// TestForwardContentResult_PreservesImageContent verifies that an ImageContent
+// block from upstream is forwarded unchanged to the downstream client.
+// Regression test for issue #368.
+func TestForwardContentResult_PreservesImageContent(t *testing.T) {
+	upstream := &mcp.CallToolResult{
+		Content: []mcp.Content{
+			mcp.NewTextContent("Here is your image:"),
+			mcp.NewImageContent("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwC", "image/png"),
+		},
+	}
+	truncator := truncate.NewTruncator(0) // disabled
+
+	forwarded, text, truncated := forwardContentResult(upstream, truncator, "test:tool", nil)
+
+	require.NotNil(t, forwarded)
+	require.Equal(t, 2, len(forwarded.Content), "both content blocks must be forwarded")
+	assert.False(t, truncated)
+
+	// First block: text preserved
+	tc, ok := forwarded.Content[0].(mcp.TextContent)
+	require.True(t, ok, "block 0 should remain TextContent")
+	assert.Equal(t, "Here is your image:", tc.Text)
+
+	// Second block: image preserved as native type
+	ic, ok := forwarded.Content[1].(mcp.ImageContent)
+	require.True(t, ok, "block 1 should remain ImageContent (not serialized to text)")
+	assert.Equal(t, "image/png", ic.MIMEType)
+	assert.Equal(t, "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwC", ic.Data)
+
+	// Text representation used for logging should reference both blocks
+	assert.Contains(t, text, "Here is your image:")
+	assert.Contains(t, text, "[image:image/png")
+}
+
+// TestForwardContentResult_TruncatesOnlyText verifies that truncation applies
+// to TextContent but leaves ImageContent and AudioContent untouched regardless
+// of their size.
+func TestForwardContentResult_TruncatesOnlyText(t *testing.T) {
+	// Build a very large base64 payload to show it survives truncation
+	bigData := strings.Repeat("A", 10000)
+	bigText := strings.Repeat("x", 2000)
+
+	upstream := &mcp.CallToolResult{
+		Content: []mcp.Content{
+			mcp.NewTextContent(bigText),
+			mcp.NewImageContent(bigData, "image/png"),
+			mcp.NewAudioContent(bigData, "audio/wav"),
+		},
+	}
+	// Truncator with a 500-char limit
+	truncator := truncate.NewTruncator(500)
+
+	forwarded, _, truncated := forwardContentResult(upstream, truncator, "test:tool", nil)
+	require.NotNil(t, forwarded)
+	require.Equal(t, 3, len(forwarded.Content))
+	assert.True(t, truncated, "text block should be marked as truncated")
+
+	// Text was truncated
+	tc, ok := forwarded.Content[0].(mcp.TextContent)
+	require.True(t, ok)
+	assert.Less(t, len(tc.Text), len(bigText), "text block should be shorter after truncation")
+
+	// Image unchanged
+	ic, ok := forwarded.Content[1].(mcp.ImageContent)
+	require.True(t, ok)
+	assert.Equal(t, bigData, ic.Data, "image data must be forwarded byte-for-byte")
+
+	// Audio unchanged
+	ac, ok := forwarded.Content[2].(mcp.AudioContent)
+	require.True(t, ok)
+	assert.Equal(t, bigData, ac.Data, "audio data must be forwarded byte-for-byte")
+}
+
+// TestForwardContentResult_TextOnlyNoTruncation exercises the common case of a
+// small text-only response. Verifies the result is forwarded unchanged.
+func TestForwardContentResult_TextOnlyNoTruncation(t *testing.T) {
+	upstream := &mcp.CallToolResult{
+		Content: []mcp.Content{
+			mcp.NewTextContent("small result"),
+		},
+	}
+	truncator := truncate.NewTruncator(0)
+
+	forwarded, text, truncated := forwardContentResult(upstream, truncator, "test:tool", nil)
+	require.NotNil(t, forwarded)
+	require.Equal(t, 1, len(forwarded.Content))
+	assert.False(t, truncated)
+	assert.Equal(t, "small result", text)
+
+	tc, ok := forwarded.Content[0].(mcp.TextContent)
+	require.True(t, ok)
+	assert.Equal(t, "small result", tc.Text)
+}
+
+// TestForwardContentResult_Fallback verifies that if result is not a
+// *mcp.CallToolResult (e.g., nil or some other interface value), the function
+// falls back to legacy JSON-wrapping behavior without panicking.
+func TestForwardContentResult_Fallback(t *testing.T) {
+	// Case 1: nil — should not panic, returns a JSON "null" text wrapper
+	forwarded, _, _ := forwardContentResult(nil, truncate.NewTruncator(0), "t", nil)
+	require.NotNil(t, forwarded)
+	require.Equal(t, 1, len(forwarded.Content))
+
+	// Case 2: a plain map — legacy JSON marshal path
+	forwarded, text, _ := forwardContentResult(map[string]string{"key": "value"}, truncate.NewTruncator(0), "t", nil)
+	require.NotNil(t, forwarded)
+	require.Equal(t, 1, len(forwarded.Content))
+	assert.Contains(t, text, "key")
+	assert.Contains(t, text, "value")
+}