refactor(middleware): define PayloadMetadata as struct with constant …

.github/agentics/large-payload-tester.md

@@ -9,6 +9,8 @@ You are an AI agent trying to verify a secret in a file hosted by the filesystem
 
 Use the filesystem MCP server to access a file called `large-test-file.json`, which contains a field `secret_reference`. Verify that the secret value matches the content of file `secret.txt` which should also be accessed from the filesystem MCP server.
 
+If a payload is too large to return over MCP, the server will return a path in the local filesystem to the payload file instead. Use the path in the local filesystem to access the full payload and extract the secret.
+
 ## Important Notes
 
 - **Keep all outputs concise** - Use brief, factual statements

internal/middleware/jqschema.go

@@ -17,6 +17,22 @@ import (
 
 var logMiddleware = logger.New("middleware:jqschema")
 
+// PayloadTruncatedInstructions is the message returned to clients when a payload
+// has been truncated and saved to the filesystem
+const PayloadTruncatedInstructions = "The payload was too large for an MCP response. The full response can be accessed through the local file system at the payloadPath."
+
+// PayloadMetadata represents the metadata response returned when a payload is too large
+// and has been saved to the filesystem
+type PayloadMetadata struct {
+	QueryID      string      `json:"queryID"`
+	PayloadPath  string      `json:"payloadPath"`
+	Preview      string      `json:"preview"`
+	Schema       interface{} `json:"schema"`
+	OriginalSize int         `json:"originalSize"`
+	Truncated    bool        `json:"truncated"`
+	Instructions string      `json:"instructions"`
+}
+
 // jqSchemaFilter is the jq filter that transforms JSON to schema
 // This is the same logic as in gh-aw shared/jqschema.md
 const jqSchemaFilter = `
@@ -71,17 +87,18 @@ func generateRandomID() string {
 // applyJqSchema applies the jq schema transformation to JSON data
 // Uses pre-compiled query code for better performance (3-10x faster than parsing on each request)
 // Accepts a context for timeout and cancellation support
-func applyJqSchema(ctx context.Context, jsonData interface{}) (string, error) {
+// Returns the schema as an interface{} object (not a JSON string)
+func applyJqSchema(ctx context.Context, jsonData interface{}) (interface{}, error) {
 	// Check if compilation succeeded at init time
 	if jqSchemaCompileErr != nil {
-		return "", jqSchemaCompileErr
+		return nil, jqSchemaCompileErr
 	}
 
 	// Run the pre-compiled query with context support (much faster than Parse+Run)
 	iter := jqSchemaCode.RunWithContext(ctx, jsonData)
 	v, ok := iter.Next()
 	if !ok {
-		return "", fmt.Errorf("jq schema filter returned no results")
+		return nil, fmt.Errorf("jq schema filter returned no results")
 	}
 
 	// Check for errors with type-specific handling
@@ -90,22 +107,17 @@ func applyJqSchema(ctx context.Context, jsonData interface{}) (string, error) {
 		if haltErr, ok := err.(*gojq.HaltError); ok {
 			// HaltError with nil value means clean halt (not an error)
 			if haltErr.Value() == nil {
-				return "", fmt.Errorf("jq schema filter halted cleanly with no output")
+				return nil, fmt.Errorf("jq schema filter halted cleanly with no output")
 			}
 			// HaltError with non-nil value is an actual error
-			return "", fmt.Errorf("jq schema filter halted with error (exit code %d): %w", haltErr.ExitCode(), err)
+			return nil, fmt.Errorf("jq schema filter halted with error (exit code %d): %w", haltErr.ExitCode(), err)
 		}
 		// Generic error case
-		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 nil, fmt.Errorf("jq schema filter error: %w", err)
 	}
 
-	return string(schemaJSON), nil
+	// Return the schema object directly (no JSON marshaling needed here)
+	return v, nil
 }
 
 // savePayload saves the payload to disk and returns the file path
@@ -234,7 +246,7 @@ func WrapToolHandler(
 
 		// Apply jq schema transformation
 		logger.LogDebug("payload", "Applying jq schema transformation: tool=%s, queryID=%s", toolName, queryID)
-		var schemaJSON string
+		var schemaObj interface{}
 		if schemaErr := func() error {
 			// Unmarshal to interface{} for jq processing
 			var jsonData interface{}
@@ -246,7 +258,7 @@ func WrapToolHandler(
 			if err != nil {
 				return err
 			}
-			schemaJSON = schema
+			schemaObj = schema
 			return nil
 		}(); schemaErr != nil {
 			logMiddleware.Printf("Failed to apply jq schema: tool=%s, queryID=%s, sessionID=%s, error=%v", toolName, queryID, sessionID, schemaErr)
@@ -256,8 +268,10 @@ func WrapToolHandler(
 			return result, data, err
 		}
 
+		// Calculate schema size for logging (marshal temporarily)
+		schemaBytes, _ := json.Marshal(schemaObj)
 		logger.LogDebug("payload", "Schema transformation completed: tool=%s, queryID=%s, schemaSize=%d bytes",
-			toolName, queryID, len(schemaJSON))
+			toolName, queryID, len(schemaBytes))
 
 		// Build the transformed response: first 500 chars + schema
 		payloadStr := string(payloadJSON)
@@ -273,27 +287,22 @@ func WrapToolHandler(
 				toolName, queryID, len(payloadStr))
 		}
 
-		// Create rewritten response
-		rewrittenResponse := map[string]interface{}{
-			"queryID":      queryID,
-			"payloadPath":  filePath,
-			"preview":      preview,
-			"schema":       schemaJSON,
-			"originalSize": len(payloadJSON),
-			"truncated":    truncated,
+		// Create rewritten response using the PayloadMetadata struct
+		rewrittenResponse := PayloadMetadata{
+			QueryID:      queryID,
+			PayloadPath:  filePath,
+			Preview:      preview,
+			Schema:       schemaObj,
+			OriginalSize: len(payloadJSON),
+			Truncated:    truncated,
+			Instructions: PayloadTruncatedInstructions,
 		}
 
 		logMiddleware.Printf("Rewritten response: tool=%s, queryID=%s, sessionID=%s, originalSize=%d, truncated=%v",
 			toolName, queryID, sessionID, len(payloadJSON), truncated)
 		logger.LogInfo("payload", "Created metadata response for client: tool=%s, queryID=%s, session=%s, payloadPath=%s, originalSize=%d bytes, truncated=%v",
 			toolName, queryID, sessionID, filePath, len(payloadJSON), truncated)
 
-		// 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
-		}
-
 		// Marshal the rewritten response to JSON for the Content field
 		rewrittenJSON, marshalErr := json.Marshal(rewrittenResponse)
 		if marshalErr != nil {

internal/middleware/jqschema_integration_test.go

@@ -12,6 +12,21 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// integrationPayloadMetadataToMap converts PayloadMetadata to map[string]interface{} for test assertions
+func integrationPayloadMetadataToMap(t *testing.T, data interface{}) map[string]interface{} {
+	t.Helper()
+	pm, ok := data.(PayloadMetadata)
+	if !ok {
+		t.Fatalf("expected PayloadMetadata, got %T", data)
+	}
+	jsonBytes, err := json.Marshal(pm)
+	require.NoError(t, err)
+	var result map[string]interface{}
+	err = json.Unmarshal(jsonBytes, &result)
+	require.NoError(t, err)
+	return result
+}
+
 // TestMiddlewareIntegration tests the complete middleware flow
 func TestMiddlewareIntegration(t *testing.T) {
 	// Create temporary directory for test
@@ -88,8 +103,7 @@ func TestMiddlewareIntegration(t *testing.T) {
 	assert.Len(t, queryIDFromContent, 32, "QueryID should be 32 hex characters")
 
 	// Verify response structure in data return value (for internal use)
-	dataMap, ok := data.(map[string]interface{})
-	require.True(t, ok, "Response should be a map")
+	dataMap := integrationPayloadMetadataToMap(t, data)
 
 	// Check all required fields exist
 	assert.Contains(t, dataMap, "queryID")
@@ -170,7 +184,7 @@ func TestMiddlewareIntegration(t *testing.T) {
 	assert.False(t, dataMap["truncated"].(bool), "Should not be truncated for small payloads")
 
 	// Verify originalSize
-	originalSize := dataMap["originalSize"].(int)
+	originalSize := int(dataMap["originalSize"].(float64))
 	assert.Greater(t, originalSize, 0, "Original size should be positive")
 }
 
@@ -222,7 +236,7 @@ func TestMiddlewareWithLargePayload(t *testing.T) {
 	}
 
 	// Also check data return value
-	dataMap := data.(map[string]interface{})
+	dataMap := integrationPayloadMetadataToMap(t, data)
 
 	// Verify truncation occurred
 	truncated := dataMap["truncated"].(bool)
@@ -277,7 +291,7 @@ func TestMiddlewareDirectoryCreation(t *testing.T) {
 	queryIDFromContent := contentMap["queryID"].(string)
 
 	// Also check data return value
-	dataMap := data.(map[string]interface{})
+	dataMap := integrationPayloadMetadataToMap(t, data)
 	queryID := dataMap["queryID"].(string)
 
 	// Both should match