[Repo Assist] perf(middleware): skip json.Marshal for small text responses (fast path)

internal/middleware/jqschema.go

@@ -34,6 +34,14 @@ const PayloadPreviewSize = 500
 // has been truncated and saved to the filesystem
 const PayloadTruncatedInstructions = "The payload was too large for an MCP response. The complete original response data is saved as a JSON file at payloadPath. The file contains valid JSON that can be parsed directly. The payloadSchema shows the structure and types of fields in the full response, but not the actual values. To access the full data with all values, read and parse the JSON file at payloadPath."
 
+// fastPathOverheadBound is the conservative upper bound (in bytes) for the JSON
+// envelope that wraps the first text content item when data is serialised.
+// The outer structure is roughly: {"content":[{"type":"text","text":"..."}],"isError":false}
+// which adds ~52 bytes, plus up to ~1 KiB for any appended DIFC notice items.
+// Using 4 KiB as the margin ensures we never skip disk-storage for a payload that
+// actually exceeds the threshold due to envelope overhead.
+const fastPathOverheadBound = 4 * 1024
+
 // PayloadMetadata represents the metadata response returned when a payload is too large
 // and has been saved to the filesystem
 type PayloadMetadata struct {
@@ -589,6 +597,55 @@ func wrapToolHandler(
 
 		result, data = tryApplyToolResponseFilter(ctx, filterCode, result, data, toolName, queryID)
 
+		// Fast path: when the handler (or filter) sets a text content item, its byte
+		// length is a reliable lower bound for json.Marshal(data). The outer JSON
+		// envelope (content array wrapper, isError flag, any DIFC notice items) adds
+		// at most fastPathOverheadBound bytes of overhead. If the text is clearly
+		// within the threshold, skip the expensive json.Marshal allocation and return
+		// immediately.
+		//
+		// This optimisation is only applied when the threshold is large enough to
+		// absorb the overhead margin (sizeThreshold > fastPathOverheadBound), so it
+		// does not affect exact-byte-count test scenarios with small thresholds.
+		if sizeThreshold > fastPathOverheadBound && len(result.Content) > 0 {
+			if tc, ok := result.Content[0].(*sdk.TextContent); ok {
+				if env, ok := data.(map[string]interface{}); ok {
+					// Only apply fast path when the backing payload is the simple one-text-item MCP envelope.
+					if len(env) <= 2 {
+						onlyKnownKeys := true
+						for k := range env {
+							if k != "content" && k != "isError" {
+								onlyKnownKeys = false
+								break
+							}
+						}
+						if onlyKnownKeys {
+							var first map[string]interface{}
+							switch c := env["content"].(type) {
+							case []interface{}:
+								if len(c) == 1 {
+									first, _ = c[0].(map[string]interface{})
+								}
+							case []map[string]interface{}:
+								if len(c) == 1 {
+									first = c[0]
+								}
+							}
+							if first != nil {
+								if itemType, _ := first["type"].(string); itemType == "text" {
+									if text, _ := first["text"].(string); text == tc.Text && len(text) <= sizeThreshold-fastPathOverheadBound {
+										logMiddleware.Printf("fast-path: content_text_len=%d within threshold-%d=%d, skipping json.Marshal: tool=%s, queryID=%s",
+											len(text), fastPathOverheadBound, sizeThreshold-fastPathOverheadBound, toolName, queryID)
+										return result, data, nil
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+
 		// Marshal the response data to JSON
 		payloadJSON, marshalErr := json.Marshal(data)
 		if marshalErr != nil {

internal/middleware/jqschema_bench_test.go

@@ -2,9 +2,11 @@ package middleware
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/itchyny/gojq"
+	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
 // BenchmarkApplyJqSchema_CompiledCode benchmarks the current implementation
@@ -120,7 +122,69 @@ func BenchmarkCompileVsParse(b *testing.B) {
 	})
 }
 
-// BenchmarkPreviewCreation benchmarks the preview string creation for large payloads.
+// BenchmarkWrapToolHandler_FastPath compares the fast-path (text content small enough
+// to skip json.Marshal) against the normal marshal-and-size-check path.
+//
+// Run with: go test -bench=BenchmarkWrapToolHandler_FastPath -benchmem ./internal/middleware/
+func BenchmarkWrapToolHandler_FastPath(b *testing.B) {
+	baseDir := b.TempDir()
+
+	sizes := []struct {
+		name    string
+		textLen int
+	}{
+		{"1KB", 1 * 1024},
+		{"10KB", 10 * 1024},
+		{"100KB", 100 * 1024},
+	}
+
+	for _, tt := range sizes {
+		innerText := strings.Repeat("x", tt.textLen)
+		innerData := map[string]interface{}{
+			"content": []interface{}{
+				map[string]interface{}{"type": "text", "text": innerText},
+			},
+			"isError": false,
+		}
+
+		makeHandler := func() func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+			return func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+				return &sdk.CallToolResult{
+					Content: []sdk.Content{&sdk.TextContent{Text: innerText}},
+				}, innerData, nil
+			}
+		}
+
+		// threshold >> text size: fast path always fires
+		fastThreshold := tt.textLen + fastPathOverheadBound + 512*1024
+		fastWrapped := WrapToolHandler(makeHandler(), "bench_tool", baseDir, "", fastThreshold, func(ctx context.Context) string { return "bench-session" })
+
+		// Threshold tuned so fast path does not fire (textLen > threshold-fastPathOverheadBound),
+		// but the payload remains under the threshold so we benchmark marshal overhead without disk I/O.
+		slowThreshold := tt.textLen + fastPathOverheadBound - 1
+		slowWrapped := WrapToolHandler(makeHandler(), "bench_tool", baseDir, "", slowThreshold, func(ctx context.Context) string { return "bench-session" })
+		b.Run("fast-path/"+tt.name, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+			ctx := context.Background()
+			req := &sdk.CallToolRequest{}
+			for i := 0; i < b.N; i++ {
+				_, _, _ = fastWrapped(ctx, req, nil)
+			}
+		})
+
+		b.Run("marshal-path/"+tt.name, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+			ctx := context.Background()
+			req := &sdk.CallToolRequest{}
+			for i := 0; i < b.N; i++ {
+				_, _, _ = slowWrapped(ctx, req, nil)
+			}
+		})
+	}
+}
+
 // The optimized version slices the byte slice before converting to string, avoiding
 // a full allocation of the entire (potentially multi-MB) payload as a string.
 func BenchmarkPreviewCreation(b *testing.B) {

internal/middleware/jqschema_test.go

@@ -1746,3 +1746,96 @@ func TestWrapToolHandler_JsonNumberData(t *testing.T) {
 	require.NoError(t, json.Unmarshal([]byte(textContent.Text), &meta))
 	assert.Equal(t, "number", meta["payloadSchema"], "schema for json.Number data should be 'number'")
 }
+
+// panicOnMarshalBool is a test helper used to ensure json.Marshal(data) is not called
+// when the WrapToolHandler fast path fires.
+type panicOnMarshalBool bool
+
+func (panicOnMarshalBool) MarshalJSON() ([]byte, error) {
+	panic("unexpected json.Marshal(data) on fast path")
+}
+
+// TestWrapToolHandler_FastPath_SkipsMarshal verifies the fast-path optimisation:
+// when the handler returns a result whose first content item is a TextContent whose
+// byte length is clearly within the threshold (i.e. threshold - fastPathOverheadBound),
+// the middleware returns the original result without performing a json.Marshal.
+func TestWrapToolHandler_FastPath_SkipsMarshal(t *testing.T) {
+	t.Parallel()
+	baseDir := t.TempDir()
+
+	const innerText = `{"repo":"owner/test","stars":42}`
+	// Construct the MCP envelope that a real backend would return.
+	innerData := map[string]interface{}{
+		"content": []interface{}{
+			map[string]interface{}{"type": "text", "text": innerText},
+		},
+		"isError": panicOnMarshalBool(false),
+	}
+
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		result := &sdk.CallToolResult{
+			Content: []sdk.Content{&sdk.TextContent{Text: innerText}},
+			IsError: false,
+		}
+		return result, innerData, nil
+	}
+
+	// threshold is well above len(innerText) + fastPathOverheadBound, so fast path fires.
+	threshold := len(innerText) + fastPathOverheadBound + 1024
+	wrapped := WrapToolHandler(mockHandler, "test_tool", baseDir, "", threshold, testGetSessionID)
+	result, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, nil)
+
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	// Fast path must return the original result unchanged (no metadata wrapping).
+	require.Len(t, result.Content, 1)
+	tc, ok := result.Content[0].(*sdk.TextContent)
+	require.True(t, ok, "fast path must preserve TextContent")
+	assert.Equal(t, innerText, tc.Text, "fast path must not rewrite text content")
+
+	// data must be the original envelope map, not a PayloadMetadata.
+	_, isMetadata := data.(PayloadMetadata)
+	assert.False(t, isMetadata, "fast path must not produce PayloadMetadata")
+
+	// No payload file should have been written.
+	entries, readErr := os.ReadDir(baseDir)
+	require.NoError(t, readErr)
+	assert.Empty(t, entries, "fast path must not create payload files")
+}
+
+// TestWrapToolHandler_FastPath_FallsThroughAtBoundary verifies that the fast-path
+// does NOT fire when the threshold is smaller than the text length plus the overhead
+// margin, ensuring the full json.Marshal path still runs in that case.
+func TestWrapToolHandler_FastPath_FallsThroughAtBoundary(t *testing.T) {
+	t.Parallel()
+	baseDir := t.TempDir()
+
+	// Use a threshold <= fastPathOverheadBound so the fast-path guard is skipped.
+	// The code path falls through to json.Marshal, measures the real size, and
+	// triggers payload storage (threshold 0 forces storage for any non-empty payload).
+	threshold := 0
+	const innerText = `{"repo":"owner/test"}`
+	innerData := map[string]interface{}{
+		"content": []interface{}{
+			map[string]interface{}{"type": "text", "text": innerText},
+		},
+		"isError": false,
+	}
+
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		result := &sdk.CallToolResult{
+			Content: []sdk.Content{&sdk.TextContent{Text: innerText}},
+			IsError: false,
+		}
+		return result, innerData, nil
+	}
+
+	wrapped := WrapToolHandler(mockHandler, "test_tool", baseDir, "", threshold, testGetSessionID)
+	_, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, nil)
+
+	require.NoError(t, err)
+	// With threshold 0, every payload exceeds the threshold; metadata must be returned.
+	_, isMetadata := data.(PayloadMetadata)
+	assert.True(t, isMetadata, "threshold 0 must always trigger storage, not fast-path")
+}