From b820be9cbc6336c72e1da7dd3fb5614230c08203 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 16 Jan 2026 14:03:33 +0000
Subject: [PATCH 1/8] Initial plan


From 9d34d50546c8d6380914be8f537e31a092d5aa05 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 16 Jan 2026 14:11:34 +0000
Subject: [PATCH 2/8] Add jqschema middleware with gojq integration

Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
---
 go.mod                               |   2 +
 go.sum                               |   4 +
 internal/middleware/jqschema.go      | 192 +++++++++++++++++++
 internal/middleware/jqschema_test.go | 269 +++++++++++++++++++++++++++
 internal/server/unified.go           |  11 +-
 5 files changed, 476 insertions(+), 2 deletions(-)
 create mode 100644 internal/middleware/jqschema.go
 create mode 100644 internal/middleware/jqschema_test.go

diff --git a/go.mod b/go.mod
index 780ebf5ff..e7ff75e78 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 )
 
 require (
+	github.com/itchyny/gojq v0.12.18
 	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
 	github.com/stretchr/testify v1.11.1
 )
@@ -18,6 +19,7 @@ require (
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/jsonschema-go v0.3.0 // indirect
+	github.com/itchyny/timefmt-go v0.1.7 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/pflag v1.0.9 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
diff --git a/go.sum b/go.sum
index 0f53ec5d9..a0dd3b96f 100644
--- a/go.sum
+++ b/go.sum
@@ -9,6 +9,10 @@ github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIy
 github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/itchyny/gojq v0.12.18 h1:gFGHyt/MLbG9n6dqnvlliiya2TaMMh6FFaR2b1H6Drc=
+github.com/itchyny/gojq v0.12.18/go.mod h1:4hPoZ/3lN9fDL1D+aK7DY1f39XZpY9+1Xpjz8atrEkg=
+github.com/itchyny/timefmt-go v0.1.7 h1:xyftit9Tbw+Dc/huSSPJaEmX1TVL8lw5vxjJLK4GMMA=
+github.com/itchyny/timefmt-go v0.1.7/go.mod h1:5E46Q+zj7vbTgWY8o5YkMeYb4I6GeWLFnetPy5oBrAI=
 github.com/modelcontextprotocol/go-sdk v1.1.0 h1:Qjayg53dnKC4UZ+792W21e4BpwEZBzwgRW6LrjLWSwA=
 github.com/modelcontextprotocol/go-sdk v1.1.0/go.mod h1:6fM3LCm3yV7pAs8isnKLn07oKtB0MP9LHd3DfAcKw10=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
diff --git a/internal/middleware/jqschema.go b/internal/middleware/jqschema.go
new file mode 100644
index 000000000..46247981e
--- /dev/null
+++ b/internal/middleware/jqschema.go
@@ -0,0 +1,192 @@
+package middleware
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/githubnext/gh-aw-mcpg/internal/logger"
+	"github.com/itchyny/gojq"
+	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+var logMiddleware = logger.New("middleware:jqschema")
+
+// jqSchemaFilter is the jq filter that transforms JSON to schema
+// This is the same logic as in gh-aw shared/jqschema.md
+const jqSchemaFilter = `
+def walk(f):
+  . as $in |
+  if type == "object" then
+    reduce keys[] as $k ({}; . + {($k): ($in[$k] | walk(f))})
+  elif type == "array" then
+    if length == 0 then [] else [.[0] | walk(f)] end
+  else
+    type
+  end;
+walk(.)
+`
+
+// generateRandomID generates a random ID for payload storage
+func generateRandomID() string {
+	bytes := make([]byte, 16)
+	if _, err := rand.Read(bytes); err != nil {
+		// Fallback to timestamp-based ID if random fails
+		return fmt.Sprintf("fallback-%d", os.Getpid())
+	}
+	return hex.EncodeToString(bytes)
+}
+
+// applyJqSchema applies the jq schema transformation to JSON data
+func applyJqSchema(jsonData interface{}) (string, error) {
+	// Parse the jq query
+	query, err := gojq.Parse(jqSchemaFilter)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse jq schema filter: %w", err)
+	}
+
+	// Run the query
+	iter := query.Run(jsonData)
+	v, ok := iter.Next()
+	if !ok {
+		return "", fmt.Errorf("jq schema filter returned no results")
+	}
+
+	// Check for errors
+	if err, ok := v.(error); ok {
+		return "", fmt.Errorf("jq schema filter error: %w", err)
+	}
+
+	// Convert result to JSON
+	schemaJSON, err := json.Marshal(v)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal schema result: %w", err)
+	}
+
+	return string(schemaJSON), nil
+}
+
+// savePayload saves the payload to disk and returns the file path
+func savePayload(queryID string, payload []byte) (string, error) {
+	// Create directory structure: /tmp/gh-awmg/tools-calls/{RND}
+	dir := filepath.Join("/tmp", "gh-awmg", "tools-calls", queryID)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return "", fmt.Errorf("failed to create payload directory: %w", err)
+	}
+
+	// Save payload to file
+	filePath := filepath.Join(dir, "payload.json")
+	if err := os.WriteFile(filePath, payload, 0644); err != nil {
+		return "", fmt.Errorf("failed to write payload file: %w", err)
+	}
+
+	return filePath, nil
+}
+
+// WrapToolHandler wraps a tool handler with jqschema middleware
+// This middleware:
+// 1. Generates a random ID for the query
+// 2. Saves the response payload to /tmp/gh-awmg/tools-calls/{RND}/payload.json
+// 3. Returns first 500 chars of payload + jq inferred schema
+func WrapToolHandler(
+	handler func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error),
+	toolName string,
+) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) {
+	return func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		// Generate random query ID
+		queryID := generateRandomID()
+		logMiddleware.Printf("Processing tool call: tool=%s, queryID=%s", toolName, queryID)
+
+		// Call the original handler
+		result, data, err := handler(ctx, req, args)
+		if err != nil {
+			logMiddleware.Printf("Tool call failed: tool=%s, queryID=%s, error=%v", toolName, queryID, err)
+			return result, data, err
+		}
+
+		// Only process successful results with data
+		if result == nil || result.IsError || data == nil {
+			return result, data, err
+		}
+
+		// Marshal the response data to JSON
+		payloadJSON, marshalErr := json.Marshal(data)
+		if marshalErr != nil {
+			logMiddleware.Printf("Failed to marshal response: tool=%s, queryID=%s, error=%v", toolName, queryID, marshalErr)
+			return result, data, err
+		}
+
+		// Save the payload
+		filePath, saveErr := savePayload(queryID, payloadJSON)
+		if saveErr != nil {
+			logMiddleware.Printf("Failed to save payload: tool=%s, queryID=%s, error=%v", toolName, queryID, saveErr)
+			// Continue even if save fails - don't break the tool call
+		} else {
+			logMiddleware.Printf("Saved payload: tool=%s, queryID=%s, path=%s, size=%d bytes", 
+				toolName, queryID, filePath, len(payloadJSON))
+		}
+
+		// Apply jq schema transformation
+		var schemaJSON string
+		if schemaErr := func() error {
+			// Unmarshal to interface{} for jq processing
+			var jsonData interface{}
+			if err := json.Unmarshal(payloadJSON, &jsonData); err != nil {
+				return fmt.Errorf("failed to unmarshal for schema: %w", err)
+			}
+
+			schema, err := applyJqSchema(jsonData)
+			if err != nil {
+				return err
+			}
+			schemaJSON = schema
+			return nil
+		}(); schemaErr != nil {
+			logMiddleware.Printf("Failed to apply jq schema: tool=%s, queryID=%s, error=%v", toolName, queryID, schemaErr)
+			// Continue with original response if schema extraction fails
+			return result, data, err
+		}
+
+		// Build the transformed response: first 500 chars + schema
+		payloadStr := string(payloadJSON)
+		var preview string
+		if len(payloadStr) > 500 {
+			preview = payloadStr[:500] + "..."
+		} else {
+			preview = payloadStr
+		}
+
+		// Create rewritten response
+		rewrittenResponse := map[string]interface{}{
+			"queryID":       queryID,
+			"payloadPath":   filePath,
+			"preview":       preview,
+			"schema":        schemaJSON,
+			"originalSize":  len(payloadJSON),
+			"truncated":     len(payloadStr) > 500,
+		}
+
+		logMiddleware.Printf("Rewritten response: tool=%s, queryID=%s, originalSize=%d, truncated=%v",
+			toolName, queryID, len(payloadJSON), len(payloadStr) > 500)
+
+		// Parse the schema JSON string back to an object for cleaner display
+		var schemaObj interface{}
+		if err := json.Unmarshal([]byte(schemaJSON), &schemaObj); err == nil {
+			rewrittenResponse["schema"] = schemaObj
+		}
+
+		return result, rewrittenResponse, nil
+	}
+}
+
+// ShouldApplyMiddleware determines if the middleware should be applied to a tool
+// Currently applies to all tools, but can be configured to filter specific tools
+func ShouldApplyMiddleware(toolName string) bool {
+	// Apply to all tools except sys tools
+	return !strings.HasPrefix(toolName, "sys___")
+}
diff --git a/internal/middleware/jqschema_test.go b/internal/middleware/jqschema_test.go
new file mode 100644
index 000000000..b843f111d
--- /dev/null
+++ b/internal/middleware/jqschema_test.go
@@ -0,0 +1,269 @@
+package middleware
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+func TestGenerateRandomID(t *testing.T) {
+	// Generate multiple IDs and ensure they're unique
+	ids := make(map[string]bool)
+	for i := 0; i < 100; i++ {
+		id := generateRandomID()
+		assert.NotEmpty(t, id, "ID should not be empty")
+		assert.False(t, ids[id], "ID should be unique")
+		ids[id] = true
+	}
+}
+
+func TestApplyJqSchema(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    interface{}
+		expected string
+	}{
+		{
+			name:     "simple object",
+			input:    map[string]interface{}{"name": "test", "count": 42},
+			expected: `{"count":"number","name":"string"}`,
+		},
+		{
+			name:     "nested object",
+			input:    map[string]interface{}{"user": map[string]interface{}{"id": 123, "active": true}},
+			expected: `{"user":{"active":"boolean","id":"number"}}`,
+		},
+		{
+			name:     "array with objects",
+			input:    map[string]interface{}{"items": []interface{}{map[string]interface{}{"id": 1, "name": "test"}}},
+			expected: `{"items":[{"id":"number","name":"string"}]}`,
+		},
+		{
+			name:     "empty array",
+			input:    map[string]interface{}{"items": []interface{}{}},
+			expected: `{"items":[]}`,
+		},
+		{
+			name: "complex nested structure",
+			input: map[string]interface{}{
+				"total_count": 1000,
+				"items": []interface{}{
+					map[string]interface{}{
+						"login":    "user1",
+						"id":       123,
+						"verified": true,
+					},
+				},
+			},
+			expected: `{"items":[{"id":"number","login":"string","verified":"boolean"}],"total_count":"number"}`,
+		},
+		{
+			name:     "null value",
+			input:    map[string]interface{}{"value": nil},
+			expected: `{"value":"null"}`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := applyJqSchema(tt.input)
+			require.NoError(t, err, "applyJqSchema should not return error")
+			assert.JSONEq(t, tt.expected, result, "Schema should match expected")
+		})
+	}
+}
+
+func TestSavePayload(t *testing.T) {
+	// Clean up after test
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+
+	queryID := "test-query-123"
+	payload := []byte(`{"test": "data"}`)
+
+	filePath, err := savePayload(queryID, payload)
+	require.NoError(t, err, "savePayload should not return error")
+
+	// Verify file exists
+	assert.FileExists(t, filePath, "Payload file should exist")
+
+	// Verify file content
+	content, err := os.ReadFile(filePath)
+	require.NoError(t, err, "Should be able to read payload file")
+	assert.Equal(t, payload, content, "File content should match payload")
+
+	// Verify directory structure
+	expectedDir := filepath.Join("/tmp", "gh-awmg", "tools-calls", queryID)
+	assert.DirExists(t, expectedDir, "Directory should exist")
+}
+
+func TestWrapToolHandler(t *testing.T) {
+	// Create a mock handler
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		return &sdk.CallToolResult{IsError: false}, map[string]interface{}{
+			"message": "success",
+			"data": map[string]interface{}{
+				"id":    123,
+				"items": []interface{}{map[string]interface{}{"name": "test"}},
+			},
+		}, nil
+	}
+
+	// Wrap the handler
+	wrapped := WrapToolHandler(mockHandler, "test_tool")
+
+	// Call the wrapped handler
+	result, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+	// Assertions
+	require.NoError(t, err, "Wrapped handler should not return error")
+	require.NotNil(t, result, "Result should not be nil")
+	assert.False(t, result.IsError, "Result should not be an error")
+
+	// Verify rewritten response structure
+	dataMap, ok := data.(map[string]interface{})
+	require.True(t, ok, "Data should be a map")
+
+	assert.Contains(t, dataMap, "queryID", "Response should contain queryID")
+	assert.Contains(t, dataMap, "payloadPath", "Response should contain payloadPath")
+	assert.Contains(t, dataMap, "preview", "Response should contain preview")
+	assert.Contains(t, dataMap, "schema", "Response should contain schema")
+	assert.Contains(t, dataMap, "originalSize", "Response should contain originalSize")
+	assert.Contains(t, dataMap, "truncated", "Response should contain truncated")
+
+	// Verify queryID is a valid hex string
+	queryID, ok := dataMap["queryID"].(string)
+	require.True(t, ok, "queryID should be a string")
+	assert.NotEmpty(t, queryID, "queryID should not be empty")
+
+	// Verify schema is present
+	schema := dataMap["schema"]
+	assert.NotNil(t, schema, "Schema should not be nil")
+
+	// Clean up test directory
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+}
+
+func TestWrapToolHandler_ErrorHandling(t *testing.T) {
+	t.Run("handler returns error", func(t *testing.T) {
+		mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+			return &sdk.CallToolResult{IsError: true}, nil, assert.AnError
+		}
+
+		wrapped := WrapToolHandler(mockHandler, "test_tool")
+		result, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+		assert.Error(t, err, "Should return error from handler")
+		assert.Nil(t, data, "Data should be nil on error")
+		assert.True(t, result.IsError, "Result should indicate error")
+	})
+
+	t.Run("handler returns nil data", func(t *testing.T) {
+		mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+			return &sdk.CallToolResult{IsError: false}, nil, nil
+		}
+
+		wrapped := WrapToolHandler(mockHandler, "test_tool")
+		result, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+		assert.NoError(t, err, "Should not return error")
+		assert.Nil(t, data, "Data should remain nil")
+		assert.False(t, result.IsError, "Result should not indicate error")
+	})
+}
+
+func TestWrapToolHandler_LongPayload(t *testing.T) {
+	// Create a handler that returns a large payload
+	largeData := map[string]interface{}{
+		"message": strings.Repeat("x", 1000),
+	}
+
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		return &sdk.CallToolResult{IsError: false}, largeData, nil
+	}
+
+	wrapped := WrapToolHandler(mockHandler, "test_tool")
+	result, data, err := wrapped(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+	require.NoError(t, err, "Should not return error")
+	require.NotNil(t, result, "Result should not be nil")
+
+	dataMap, ok := data.(map[string]interface{})
+	require.True(t, ok, "Data should be a map")
+
+	// Verify truncation
+	assert.True(t, dataMap["truncated"].(bool), "Should indicate truncation")
+	preview := dataMap["preview"].(string)
+	assert.LessOrEqual(t, len(preview), 503, "Preview should be truncated to ~500 chars + '...'")
+	assert.True(t, strings.HasSuffix(preview, "..."), "Preview should end with '...'")
+
+	// Clean up
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+}
+
+func TestShouldApplyMiddleware(t *testing.T) {
+	tests := []struct {
+		name     string
+		toolName string
+		expected bool
+	}{
+		{
+			name:     "regular tool",
+			toolName: "github___search_code",
+			expected: true,
+		},
+		{
+			name:     "sys tool",
+			toolName: "sys___init",
+			expected: false,
+		},
+		{
+			name:     "another sys tool",
+			toolName: "sys___list_servers",
+			expected: false,
+		},
+		{
+			name:     "tool with sys in name but not prefix",
+			toolName: "mysys___tool",
+			expected: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := ShouldApplyMiddleware(tt.toolName)
+			assert.Equal(t, tt.expected, result, "ShouldApplyMiddleware result should match expected")
+		})
+	}
+}
+
+func TestApplyJqSchema_ErrorCases(t *testing.T) {
+	t.Run("handles complex recursive structures", func(t *testing.T) {
+		// Create a deeply nested structure
+		input := map[string]interface{}{
+			"level1": map[string]interface{}{
+				"level2": map[string]interface{}{
+					"level3": map[string]interface{}{
+						"value": "deep",
+					},
+				},
+			},
+		}
+
+		result, err := applyJqSchema(input)
+		require.NoError(t, err, "Should handle deeply nested structures")
+		assert.NotEmpty(t, result, "Result should not be empty")
+
+		// Verify the schema is correctly nested
+		var schema map[string]interface{}
+		err = json.Unmarshal([]byte(result), &schema)
+		require.NoError(t, err, "Should be valid JSON")
+		assert.Contains(t, schema, "level1", "Should contain level1")
+	})
+}
diff --git a/internal/server/unified.go b/internal/server/unified.go
index 1cc8eea8d..3ad181626 100644
--- a/internal/server/unified.go
+++ b/internal/server/unified.go
@@ -14,6 +14,7 @@ import (
 	"github.com/githubnext/gh-aw-mcpg/internal/launcher"
 	"github.com/githubnext/gh-aw-mcpg/internal/logger"
 	"github.com/githubnext/gh-aw-mcpg/internal/mcp"
+	"github.com/githubnext/gh-aw-mcpg/internal/middleware"
 	"github.com/githubnext/gh-aw-mcpg/internal/sys"
 	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
 )
@@ -257,9 +258,15 @@ func (us *UnifiedServer) registerToolsFromBackend(serverID string) error {
 			return result, data, err
 		}
 
+		// Wrap handler with jqschema middleware if applicable
+		finalHandler := handler
+		if middleware.ShouldApplyMiddleware(prefixedName) {
+			finalHandler = middleware.WrapToolHandler(handler, prefixedName)
+		}
+
 		// Store handler for routed mode to reuse
 		us.toolsMu.Lock()
-		us.tools[prefixedName].Handler = handler
+		us.tools[prefixedName].Handler = finalHandler
 		us.toolsMu.Unlock()
 
 		// Register the tool with the SDK
@@ -269,7 +276,7 @@ func (us *UnifiedServer) registerToolsFromBackend(serverID string) error {
 		sdk.AddTool(us.server, &sdk.Tool{
 			Name:        prefixedName,
 			Description: toolDesc,
-		}, handler)
+		}, finalHandler)
 
 		log.Printf("Registered tool: %s", prefixedName)
 	}

From e610adf43d70897530e13ea11949f2f858b37cd9 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 16 Jan 2026 14:13:57 +0000
Subject: [PATCH 3/8] Add comprehensive integration tests for jqschema
 middleware

Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
---
 internal/middleware/jqschema.go               |  14 +-
 .../middleware/jqschema_integration_test.go   | 230 ++++++++++++++++++
 internal/middleware/jqschema_test.go          |   2 +-
 3 files changed, 238 insertions(+), 8 deletions(-)
 create mode 100644 internal/middleware/jqschema_integration_test.go

diff --git a/internal/middleware/jqschema.go b/internal/middleware/jqschema.go
index 46247981e..6120ad807 100644
--- a/internal/middleware/jqschema.go
+++ b/internal/middleware/jqschema.go
@@ -127,7 +127,7 @@ func WrapToolHandler(
 			logMiddleware.Printf("Failed to save payload: tool=%s, queryID=%s, error=%v", toolName, queryID, saveErr)
 			// Continue even if save fails - don't break the tool call
 		} else {
-			logMiddleware.Printf("Saved payload: tool=%s, queryID=%s, path=%s, size=%d bytes", 
+			logMiddleware.Printf("Saved payload: tool=%s, queryID=%s, path=%s, size=%d bytes",
 				toolName, queryID, filePath, len(payloadJSON))
 		}
 
@@ -163,12 +163,12 @@ func WrapToolHandler(
 
 		// Create rewritten response
 		rewrittenResponse := map[string]interface{}{
-			"queryID":       queryID,
-			"payloadPath":   filePath,
-			"preview":       preview,
-			"schema":        schemaJSON,
-			"originalSize":  len(payloadJSON),
-			"truncated":     len(payloadStr) > 500,
+			"queryID":      queryID,
+			"payloadPath":  filePath,
+			"preview":      preview,
+			"schema":       schemaJSON,
+			"originalSize": len(payloadJSON),
+			"truncated":    len(payloadStr) > 500,
 		}
 
 		logMiddleware.Printf("Rewritten response: tool=%s, queryID=%s, originalSize=%d, truncated=%v",
diff --git a/internal/middleware/jqschema_integration_test.go b/internal/middleware/jqschema_integration_test.go
new file mode 100644
index 000000000..70d0b7b0a
--- /dev/null
+++ b/internal/middleware/jqschema_integration_test.go
@@ -0,0 +1,230 @@
+package middleware
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestMiddlewareIntegration tests the complete middleware flow
+func TestMiddlewareIntegration(t *testing.T) {
+	// Clean up test directory
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+
+	// Create a mock handler that returns GitHub-like search data
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		// Simulate a GitHub search response
+		response := map[string]interface{}{
+			"total_count": 1000,
+			"items": []interface{}{
+				map[string]interface{}{
+					"name":        "repo1",
+					"id":          12345,
+					"stars":       100,
+					"private":     false,
+					"description": "A test repository",
+					"owner": map[string]interface{}{
+						"login": "user1",
+						"id":    999,
+					},
+				},
+				map[string]interface{}{
+					"name":        "repo2",
+					"id":          67890,
+					"stars":       250,
+					"private":     true,
+					"description": "Another test repository",
+					"owner": map[string]interface{}{
+						"login": "user2",
+						"id":    888,
+					},
+				},
+			},
+		}
+		return &sdk.CallToolResult{IsError: false}, response, nil
+	}
+
+	// Wrap with middleware
+	wrappedHandler := WrapToolHandler(mockHandler, "github___search_repositories")
+
+	// Call the wrapped handler
+	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{
+		"query": "test",
+	})
+
+	// Verify no error
+	require.NoError(t, err, "Handler should not return error")
+	require.NotNil(t, result, "Result should not be nil")
+	assert.False(t, result.IsError, "Result should not indicate error")
+
+	// Verify response structure
+	dataMap, ok := data.(map[string]interface{})
+	require.True(t, ok, "Response should be a map")
+
+	// Check all required fields exist
+	assert.Contains(t, dataMap, "queryID")
+	assert.Contains(t, dataMap, "payloadPath")
+	assert.Contains(t, dataMap, "preview")
+	assert.Contains(t, dataMap, "schema")
+	assert.Contains(t, dataMap, "originalSize")
+	assert.Contains(t, dataMap, "truncated")
+
+	// Verify queryID format
+	queryID := dataMap["queryID"].(string)
+	assert.Len(t, queryID, 32, "QueryID should be 32 hex characters")
+
+	// Verify payload was saved
+	payloadPath := dataMap["payloadPath"].(string)
+	assert.FileExists(t, payloadPath, "Payload file should exist")
+
+	// Verify payload content
+	payloadContent, err := os.ReadFile(payloadPath)
+	require.NoError(t, err, "Should read payload file")
+
+	var originalData map[string]interface{}
+	err = json.Unmarshal(payloadContent, &originalData)
+	require.NoError(t, err, "Payload should be valid JSON")
+
+	// Verify original data structure is preserved in file
+	assert.Equal(t, float64(1000), originalData["total_count"])
+	assert.NotNil(t, originalData["items"])
+
+	// Verify schema structure
+	schemaObj := dataMap["schema"]
+	assert.NotNil(t, schemaObj, "Schema should not be nil")
+
+	// Convert schema to JSON string for inspection
+	schemaJSON, err := json.Marshal(schemaObj)
+	require.NoError(t, err, "Schema should be marshallable")
+
+	var schema map[string]interface{}
+	err = json.Unmarshal(schemaJSON, &schema)
+	require.NoError(t, err, "Schema should be valid JSON")
+
+	// Verify schema has the expected structure
+	assert.Contains(t, schema, "total_count")
+	assert.Contains(t, schema, "items")
+	assert.Equal(t, "number", schema["total_count"])
+
+	// Verify items is an array with schema
+	items, ok := schema["items"].([]interface{})
+	require.True(t, ok, "items should be an array")
+	assert.Len(t, items, 1, "items schema should have one element")
+
+	// Verify the item schema
+	itemSchema, ok := items[0].(map[string]interface{})
+	require.True(t, ok, "item schema should be an object")
+	assert.Contains(t, itemSchema, "name")
+	assert.Contains(t, itemSchema, "id")
+	assert.Contains(t, itemSchema, "stars")
+	assert.Contains(t, itemSchema, "private")
+	assert.Contains(t, itemSchema, "description")
+	assert.Contains(t, itemSchema, "owner")
+
+	// Verify types
+	assert.Equal(t, "string", itemSchema["name"])
+	assert.Equal(t, "number", itemSchema["id"])
+	assert.Equal(t, "number", itemSchema["stars"])
+	assert.Equal(t, "boolean", itemSchema["private"])
+	assert.Equal(t, "string", itemSchema["description"])
+
+	// Verify nested owner schema
+	ownerSchema, ok := itemSchema["owner"].(map[string]interface{})
+	require.True(t, ok, "owner schema should be an object")
+	assert.Contains(t, ownerSchema, "login")
+	assert.Contains(t, ownerSchema, "id")
+	assert.Equal(t, "string", ownerSchema["login"])
+	assert.Equal(t, "number", ownerSchema["id"])
+
+	// Verify truncation flag
+	assert.False(t, dataMap["truncated"].(bool), "Should not be truncated for small payloads")
+
+	// Verify originalSize
+	originalSize := int(dataMap["originalSize"].(int))
+	assert.Greater(t, originalSize, 0, "Original size should be positive")
+}
+
+// TestMiddlewareWithLargePayload tests truncation behavior
+func TestMiddlewareWithLargePayload(t *testing.T) {
+	// Clean up test directory
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+
+	// Create a large payload
+	largeItems := make([]interface{}, 100)
+	for i := 0; i < 100; i++ {
+		largeItems[i] = map[string]interface{}{
+			"id":          i,
+			"name":        "item-" + string(rune(i)),
+			"description": "This is a long description to make the payload large enough for truncation testing purposes.",
+		}
+	}
+
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		return &sdk.CallToolResult{IsError: false}, map[string]interface{}{
+			"total_count": 100,
+			"items":       largeItems,
+		}, nil
+	}
+
+	wrappedHandler := WrapToolHandler(mockHandler, "test_tool")
+	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	dataMap := data.(map[string]interface{})
+
+	// Verify truncation occurred
+	truncated := dataMap["truncated"].(bool)
+	preview := dataMap["preview"].(string)
+
+	if truncated {
+		assert.True(t, len(preview) <= 503, "Preview should be truncated")
+		assert.Contains(t, preview, "...", "Truncated preview should end with ...")
+	}
+
+	// Verify payload file has complete data
+	payloadPath := dataMap["payloadPath"].(string)
+	payloadContent, err := os.ReadFile(payloadPath)
+	require.NoError(t, err)
+
+	var completeData map[string]interface{}
+	err = json.Unmarshal(payloadContent, &completeData)
+	require.NoError(t, err)
+
+	// Verify complete data is in the file
+	completeItems := completeData["items"].([]interface{})
+	assert.Len(t, completeItems, 100, "File should contain all items")
+}
+
+// TestMiddlewareDirectoryCreation tests that directories are created correctly
+func TestMiddlewareDirectoryCreation(t *testing.T) {
+	// Clean up test directory
+	defer os.RemoveAll(filepath.Join("/tmp", "gh-awmg"))
+
+	mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) {
+		return &sdk.CallToolResult{IsError: false}, map[string]interface{}{"test": "data"}, nil
+	}
+
+	wrappedHandler := WrapToolHandler(mockHandler, "test_tool")
+	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
+
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	dataMap := data.(map[string]interface{})
+	queryID := dataMap["queryID"].(string)
+
+	// Verify directory structure
+	expectedDir := filepath.Join("/tmp", "gh-awmg", "tools-calls", queryID)
+	assert.DirExists(t, expectedDir, "Query directory should exist")
+
+	payloadPath := dataMap["payloadPath"].(string)
+	assert.Equal(t, filepath.Join(expectedDir, "payload.json"), payloadPath, "Payload path should match expected structure")
+}
diff --git a/internal/middleware/jqschema_test.go b/internal/middleware/jqschema_test.go
index b843f111d..01e96a378 100644
--- a/internal/middleware/jqschema_test.go
+++ b/internal/middleware/jqschema_test.go
@@ -8,9 +8,9 @@ import (
 	"strings"
 	"testing"
 
+	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	sdk "github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
 func TestGenerateRandomID(t *testing.T) {

From aae48c05dd82c6f78f5b74be0ce7ab6e4280de5c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 16 Jan 2026 14:15:08 +0000
Subject: [PATCH 4/8] Add comprehensive README for jqschema middleware

---
 internal/middleware/README.md | 146 ++++++++++++++++++++++++++++++++++
 1 file changed, 146 insertions(+)
 create mode 100644 internal/middleware/README.md

diff --git a/internal/middleware/README.md b/internal/middleware/README.md
new file mode 100644
index 000000000..e160124f6
--- /dev/null
+++ b/internal/middleware/README.md
@@ -0,0 +1,146 @@
+# jqschema Middleware
+
+This middleware package implements the jqschema functionality from the gh-aw shared agentic workflow as a tool call middleware for the MCP Gateway.
+
+## Features
+
+- **Automatic JSON Schema Inference**: Uses the jq schema transformation logic to automatically infer the structure and types of JSON responses
+- **Payload Storage**: Stores complete response payloads in `/tmp/gh-awmg/tools-calls/{randomID}/payload.json`
+- **Response Rewriting**: Returns a transformed response containing:
+  - First 500 characters of the payload (for quick preview)
+  - Inferred JSON schema showing structure and types
+  - Query ID for tracking
+  - File path to complete payload
+  - Metadata (original size, truncation status)
+
+## How It Works
+
+The middleware wraps tool handlers and intercepts their responses:
+
+1. **Random ID Generation**: Each tool call gets a unique random ID (32 hex characters)
+2. **Original Handler Execution**: The original tool handler is called normally
+3. **Payload Storage**: The complete response is saved to disk
+4. **Schema Inference**: The jq schema transformation is applied to extract types and structure
+5. **Response Rewriting**: A new response is returned with preview + schema
+
+## Usage
+
+The middleware is automatically applied to all backend MCP server tools (except `sys___*` tools).
+
+### Example Response
+
+**Original response:**
+```json
+{
+  "total_count": 1000,
+  "items": [
+    {
+      "login": "user1",
+      "id": 123,
+      "verified": true
+    },
+    {
+      "login": "user2",
+      "id": 456,
+      "verified": false
+    }
+  ]
+}
+```
+
+**Transformed response:**
+```json
+{
+  "queryID": "a1b2c3d4e5f6...",
+  "payloadPath": "/tmp/gh-awmg/tools-calls/a1b2c3d4e5f6.../payload.json",
+  "preview": "{\"total_count\":1000,\"items\":[{\"login\":\"user1\",\"id\":123,\"verified\":true}...",
+  "schema": {
+    "total_count": "number",
+    "items": [
+      {
+        "login": "string",
+        "id": "number",
+        "verified": "boolean"
+      }
+    ]
+  },
+  "originalSize": 234,
+  "truncated": false
+}
+```
+
+## Implementation Details
+
+### jq Schema Filter
+
+The middleware uses the same jq filter logic as the gh-aw jqschema utility:
+
+```jq
+def walk(f):
+  . as $in |
+  if type == "object" then
+    reduce keys[] as $k ({}; . + {($k): ($in[$k] | walk(f))})
+  elif type == "array" then
+    if length == 0 then [] else [.[0] | walk(f)] end
+  else
+    type
+  end;
+walk(.)
+```
+
+This recursively walks the JSON structure and replaces values with their type names.
+
+### Go Implementation
+
+The middleware is implemented using [gojq](https://github.com/itchyny/gojq), a pure Go implementation of jq, eliminating the need to spawn external processes.
+
+## Configuration
+
+The middleware can be controlled via the `ShouldApplyMiddleware` function:
+
+```go
+func ShouldApplyMiddleware(toolName string) bool {
+    // Currently excludes sys tools
+    return !strings.HasPrefix(toolName, "sys___")
+}
+```
+
+## Testing
+
+The middleware includes comprehensive tests:
+
+- **Unit tests**: Test individual functions (random ID generation, schema transformation, payload storage)
+- **Integration tests**: Test complete middleware flow with mock handlers
+- **Edge cases**: Test error handling, large payloads, truncation behavior
+
+Run tests:
+```bash
+make test-unit
+# or
+go test ./internal/middleware/...
+```
+
+## Directory Structure
+
+Payloads are stored in:
+```
+/tmp/gh-awmg/tools-calls/
+  ├── {randomID1}/
+  │   └── payload.json
+  ├── {randomID2}/
+  │   └── payload.json
+  └── ...
+```
+
+## Benefits
+
+1. **Reduced Token Usage**: Preview + schema is much smaller than full responses
+2. **Better Understanding**: Schema shows structure without verbose data
+3. **Audit Trail**: Complete payloads are saved for later inspection
+4. **Debugging**: Query IDs enable tracking and correlation
+5. **Performance**: Pure Go implementation (no external process spawning)
+
+## References
+
+- Original jqschema utility: [gh-aw/.github/workflows/shared/jqschema.md](https://github.com/githubnext/gh-aw/blob/main/.github/workflows/shared/jqschema.md)
+- gojq library: [github.com/itchyny/gojq](https://github.com/itchyny/gojq)

From 11d132bec4dcab4316a1a47a2598100c7f8d6758 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 16 Jan 2026 14:16:36 +0000
Subject: [PATCH 5/8] Complete jqschema middleware implementation with demo

Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
---
 internal/middleware/demo.sh | 77 +++++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)
 create mode 100755 internal/middleware/demo.sh

diff --git a/internal/middleware/demo.sh b/internal/middleware/demo.sh
new file mode 100755
index 000000000..dff779b15
--- /dev/null
+++ b/internal/middleware/demo.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+# Demo script to show jqschema middleware in action
+
+set -e
+
+echo "=== jqschema Middleware Demo ==="
+echo ""
+echo "This demo shows how the jqschema middleware transforms tool call responses."
+echo ""
+
+# Create a test directory
+TEST_DIR="/tmp/gh-awmg-demo"
+rm -rf "$TEST_DIR"
+mkdir -p "$TEST_DIR"
+
+echo "1. The middleware intercepts tool call responses"
+echo "2. Generates a unique random ID for each call"
+echo "3. Saves the complete payload to /tmp/gh-awmg/tools-calls/{ID}/payload.json"
+echo "4. Infers the JSON schema using jq logic"
+echo "5. Returns a transformed response with:"
+echo "   - First 500 chars of payload (preview)"
+echo "   - Inferred schema (types and structure)"
+echo "   - Query ID, file path, and metadata"
+echo ""
+
+echo "=== Example Transformation ==="
+echo ""
+echo "Original Response:"
+echo "{"
+echo "  \"total_count\": 1000,"
+echo "  \"items\": ["
+echo "    {\"login\": \"user1\", \"id\": 123, \"verified\": true},"
+echo "    {\"login\": \"user2\", \"id\": 456, \"verified\": false}"
+echo "  ]"
+echo "}"
+echo ""
+
+echo "Transformed Response:"
+echo "{"
+echo "  \"queryID\": \"a1b2c3d4e5f6...\" (32 hex chars),"
+echo "  \"payloadPath\": \"/tmp/gh-awmg/tools-calls/a1b2c3d4.../payload.json\","
+echo "  \"preview\": \"{\\\"total_count\\\":1000,\\\"items\\\":[{\\\"login\\\":\\\"user1\\\"...\" (500 chars),"
+echo "  \"schema\": {"
+echo "    \"total_count\": \"number\","
+echo "    \"items\": [{"
+echo "      \"login\": \"string\","
+echo "      \"id\": \"number\","
+echo "      \"verified\": \"boolean\""
+echo "    }]"
+echo "  },"
+echo "  \"originalSize\": 234,"
+echo "  \"truncated\": false"
+echo "}"
+echo ""
+
+echo "=== Benefits ==="
+echo "✓ Reduced token usage (schema is smaller than full data)"
+echo "✓ Better understanding (see structure without verbose data)"
+echo "✓ Audit trail (complete payloads saved to disk)"
+echo "✓ Easy debugging (query IDs for tracking)"
+echo "✓ Pure Go implementation (no external jq process)"
+echo ""
+
+echo "=== Configuration ==="
+echo "- Middleware is automatically applied to all backend tools"
+echo "- System tools (sys___*) are excluded"
+echo "- Payloads stored in: /tmp/gh-awmg/tools-calls/"
+echo "- Preview truncated at 500 characters"
+echo ""
+
+echo "=== Testing ==="
+echo "Run the middleware tests:"
+echo "  make test-unit"
+echo "  go test ./internal/middleware/... -v"
+echo ""
+
+echo "Demo complete!"

From d44f1cfbebb4c4352e102913c98596ad6f78ceec Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 18 Jan 2026 15:55:20 +0000
Subject: [PATCH 6/8] Fix linting error: remove unnecessary int conversion

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
---
 internal/middleware/jqschema_integration_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/middleware/jqschema_integration_test.go b/internal/middleware/jqschema_integration_test.go
index 70d0b7b0a..0f800a0ce 100644
--- a/internal/middleware/jqschema_integration_test.go
+++ b/internal/middleware/jqschema_integration_test.go
@@ -146,7 +146,7 @@ func TestMiddlewareIntegration(t *testing.T) {
 	assert.False(t, dataMap["truncated"].(bool), "Should not be truncated for small payloads")
 
 	// Verify originalSize
-	originalSize := int(dataMap["originalSize"].(int))
+	originalSize := dataMap["originalSize"].(int)
 	assert.Greater(t, originalSize, 0, "Original size should be positive")
 }
 

From 3afd2803649663e5df0b66fbf6f1444373afcb87 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 18 Jan 2026 15:56:12 +0000
Subject: [PATCH 7/8] Document future middleware configuration options

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
---
 internal/middleware/README.md | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/internal/middleware/README.md b/internal/middleware/README.md
index e160124f6..6272d088c 100644
--- a/internal/middleware/README.md
+++ b/internal/middleware/README.md
@@ -105,6 +105,25 @@ func ShouldApplyMiddleware(toolName string) bool {
 }
 ```
 
+### Future Enhancements
+
+**Selective Middleware Mounting**: A configuration system could be added to:
+- Enable/disable middleware per backend server
+- Configure which tools get middleware applied
+- Set custom truncation limits
+- Configure storage locations
+- Add multiple middleware types with ordering
+
+Example future config structure:
+```toml
+[middleware.jqschema]
+enabled = true
+truncate_at = 500
+storage_path = "/tmp/gh-awmg/tools-calls"
+exclude_tools = ["sys___*"]
+include_backends = ["github", "tavily"]
+```
+
 ## Testing
 
 The middleware includes comprehensive tests:

From 8dcfdcd8bf199afe8416d14c91d3089b9e118303 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 18 Jan 2026 16:11:26 +0000
Subject: [PATCH 8/8] Fix middleware wrapping and formatting issues

- Changed wrappedHandler to call finalHandler instead of handler to properly apply middleware
- Fixed gofmt formatting in health.go (removed trailing spaces)

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
---
 go.mod                     |  1 +
 go.sum                     | 84 --------------------------------------
 internal/server/health.go  |  4 +-
 internal/server/unified.go |  4 +-
 4 files changed, 5 insertions(+), 88 deletions(-)

diff --git a/go.mod b/go.mod
index 67b11b95a..b09bf8ee3 100644
--- a/go.mod
+++ b/go.mod
@@ -18,6 +18,7 @@ require (
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/jsonschema-go v0.3.0 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/itchyny/timefmt-go v0.1.7 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/pflag v1.0.9 // indirect
diff --git a/go.sum b/go.sum
index 4b6572d3e..a0dd3b96f 100644
--- a/go.sum
+++ b/go.sum
@@ -1,12 +1,8 @@
-cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
 github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg=
 github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
-github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
@@ -28,100 +24,20 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
 github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
 github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY=
 github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
-github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
-github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
-github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
-github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
-golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
-golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
-golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
-golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
-golang.org/x/crypto v0.38.0/go.mod h1:MvrbAqul58NNYPKnOra203SB9vpuZW0e+RRZV+Ggqjw=
-golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
-golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
-golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
-golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
-golang.org/x/mod v0.24.0/go.mod h1:IXM97Txy2VM4PJ3gI61r1YEk/gAj6zAHN3AdZt6S9Ww=
-golang.org/x/mod v0.25.0/go.mod h1:IXM97Txy2VM4PJ3gI61r1YEk/gAj6zAHN3AdZt6S9Ww=
-golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
-golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
-golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
-golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
-golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
-golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
-golang.org/x/net v0.40.0/go.mod h1:y0hY0exeL2Pku80/zKK7tpntoX23cqL3Oa6njdgRtds=
-golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
 golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
 golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
-golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
-golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
-golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
-golang.org/x/sync v0.14.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
-golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
-golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
 golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
-golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
-golang.org/x/telemetry v0.0.0-20240521205824-bda55230c457/go.mod h1:pRgIJT+bRLFKnoM1ldnzKoxTIn14Yxz928LQRYYgIN0=
-golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
-golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
-golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
-golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
-golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
-golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
-golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
-golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ=
 golang.org/x/term v0.38.0 h1:PQ5pkm/rLO6HnxFR7N2lJHOZX6Kez5Y1gDSJla6jo7Q=
 golang.org/x/term v0.38.0/go.mod h1:bSEAKrOT1W+VSu9TSCMtoGEOUcKxOKgl3LE5QEF/xVg=
-golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
-golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
-golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
-golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
-golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
-golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
-golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
-golang.org/x/text v0.25.0/go.mod h1:WEdwpYrmk1qmdHvhkSTNPm3app7v4rsT8F2UD6+VHIA=
-golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
-golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
-golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
-golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
-golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
-golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
-golang.org/x/tools v0.33.0/go.mod h1:CIJMaWEY88juyUfo7UbgPqbC8rU2OqfAV1h2Qp0oMYI=
 golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo=
 golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg=
-golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/internal/server/health.go b/internal/server/health.go
index f45d94e1e..aff89754e 100644
--- a/internal/server/health.go
+++ b/internal/server/health.go
@@ -21,7 +21,7 @@ type HealthResponse struct {
 // BuildHealthResponse constructs a HealthResponse from the unified server's status
 func BuildHealthResponse(unifiedServer *UnifiedServer) HealthResponse {
 	logHealth.Print("Building health response")
-	
+
 	// Get server status
 	serverStatus := unifiedServer.GetServerStatus()
 	logHealth.Printf("Retrieved status for %d servers", len(serverStatus))
@@ -50,7 +50,7 @@ func BuildHealthResponse(unifiedServer *UnifiedServer) HealthResponse {
 func HandleHealth(unifiedServer *UnifiedServer) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		logHealth.Printf("Health check request: method=%s, remote=%s", r.Method, r.RemoteAddr)
-		
+
 		w.Header().Set("Content-Type", "application/json")
 		w.WriteHeader(http.StatusOK)
 
diff --git a/internal/server/unified.go b/internal/server/unified.go
index 08751b68f..474b65532 100644
--- a/internal/server/unified.go
+++ b/internal/server/unified.go
@@ -290,10 +290,10 @@ func (us *UnifiedServer) registerToolsFromBackend(serverID string) error {
 		// The typed handler signature: func(context.Context, *CallToolRequest, interface{}) (*CallToolResult, interface{}, error)
 		// The simple handler signature: func(context.Context, *CallToolRequest) (*CallToolResult, error)
 		wrappedHandler := func(ctx context.Context, req *sdk.CallToolRequest) (*sdk.CallToolResult, error) {
-			// Call the original typed handler
+			// Call the final handler (which may include middleware wrapping)
 			// The third parameter would be the pre-unmarshaled/validated input if using sdk.AddTool,
 			// but we handle unmarshaling ourselves in the handler, so we pass nil
-			result, _, err := handler(ctx, req, nil)
+			result, _, err := finalHandler(ctx, req, nil)
 			return result, err
 		}