Add configurable payload size threshold with comprehensive tests

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>

Author: Copilot

internal/config/config_core.go

@@ -54,6 +54,11 @@ type GatewayConfig struct {
 
 	// PayloadDir is the directory for storing large payloads
 	PayloadDir string `toml:"payload_dir" json:"payload_dir,omitempty"`
+
+	// PayloadSizeThreshold is the size threshold (in bytes) for storing payloads to disk.
+	// Payloads larger than this threshold are stored to disk, smaller ones are returned inline.
+	// Default: 1024 bytes (1KB)
+	PayloadSizeThreshold int `toml:"payload_size_threshold" json:"payload_size_threshold,omitempty"`
 }
 
 // ServerConfig represents an individual MCP server configuration.

internal/config/config_payload.go

@@ -5,11 +5,21 @@ package config
 // DefaultPayloadDir is the default directory for storing large payloads.
 const DefaultPayloadDir = "/tmp/jq-payloads"
 
+// DefaultPayloadSizeThreshold is the default size threshold (in bytes) for storing payloads to disk.
+// Payloads larger than this threshold are stored to disk, smaller ones are returned inline.
+// Default: 1024 bytes (1KB)
+const DefaultPayloadSizeThreshold = 1024
+
 func init() {
-	// Register default setter for PayloadDir
+	// Register default setter for PayloadDir and PayloadSizeThreshold
 	RegisterDefaults(func(cfg *Config) {
-		if cfg.Gateway != nil && cfg.Gateway.PayloadDir == "" {
-			cfg.Gateway.PayloadDir = DefaultPayloadDir
+		if cfg.Gateway != nil {
+			if cfg.Gateway.PayloadDir == "" {
+				cfg.Gateway.PayloadDir = DefaultPayloadDir
+			}
+			if cfg.Gateway.PayloadSizeThreshold == 0 {
+				cfg.Gateway.PayloadSizeThreshold = DefaultPayloadSizeThreshold
+			}
 		}
 	})
 }

internal/middleware/jqschema.go

@@ -167,12 +167,14 @@ func savePayload(baseDir, sessionID, queryID string, payload []byte) (string, er
 // This middleware:
 // 1. Generates a random ID for the query
 // 2. Extracts session ID from context (or uses "default")
-// 3. Saves the response payload to {baseDir}/{sessionID}/{queryID}/payload.json
-// 4. Returns first 500 chars of payload + jq inferred schema
+// 3. If payload size > sizeThreshold: saves to {baseDir}/{sessionID}/{queryID}/payload.json and returns metadata
+// 4. If payload size <= sizeThreshold: returns original response directly (no file storage)
+// 5. For large payloads: returns first 500 chars of payload + jq inferred schema
 func WrapToolHandler(
 	handler func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error),
 	toolName string,
 	baseDir string,
+	sizeThreshold int,
 	getSessionID func(context.Context) 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) {
@@ -224,12 +226,24 @@ func WrapToolHandler(
 		}
 
 		payloadSize := len(payloadJSON)
-		logger.LogInfo("payload", "Response data marshaled to JSON: tool=%s, queryID=%s, size=%d bytes (%.2f KB, %.2f MB)",
-			toolName, queryID, payloadSize, float64(payloadSize)/1024, float64(payloadSize)/(1024*1024))
+		logger.LogInfo("payload", "Response data marshaled to JSON: tool=%s, queryID=%s, size=%d bytes (%.2f KB, %.2f MB), threshold=%d bytes",
+			toolName, queryID, payloadSize, float64(payloadSize)/1024, float64(payloadSize)/(1024*1024), sizeThreshold)
+
+		// Check if payload size is within threshold - if so, return original response directly
+		if payloadSize <= sizeThreshold {
+			logger.LogInfo("payload", "Payload size (%d bytes) is within threshold (%d bytes), returning inline without file storage: tool=%s, queryID=%s",
+				payloadSize, sizeThreshold, toolName, queryID)
+			logMiddleware.Printf("Payload within threshold: tool=%s, queryID=%s, size=%d bytes, threshold=%d bytes, returning inline",
+				toolName, queryID, payloadSize, sizeThreshold)
+			// Return the original result without modification
+			return result, data, err
+		}
 
-		// Save the payload
-		logger.LogInfo("payload", "Starting payload storage to filesystem: tool=%s, queryID=%s, session=%s, baseDir=%s",
-			toolName, queryID, sessionID, baseDir)
+		// Payload is larger than threshold - save to filesystem
+		logger.LogInfo("payload", "Payload size (%d bytes) exceeds threshold (%d bytes), saving to filesystem: tool=%s, queryID=%s, session=%s, baseDir=%s",
+			payloadSize, sizeThreshold, toolName, queryID, sessionID, baseDir)
+		logMiddleware.Printf("Payload exceeds threshold: tool=%s, queryID=%s, size=%d bytes, threshold=%d bytes, saving to disk",
+			toolName, queryID, payloadSize, sizeThreshold)
 
 		filePath, saveErr := savePayload(baseDir, sessionID, queryID, payloadJSON)
 		if saveErr != nil {

internal/middleware/jqschema_integration_test.go

@@ -67,7 +67,7 @@ func TestMiddlewareIntegration(t *testing.T) {
 	}
 
 	// Wrap with middleware
-	wrappedHandler := WrapToolHandler(mockHandler, "github___search_repositories", baseDir, testGetSessionID)
+	wrappedHandler := WrapToolHandler(mockHandler, "github___search_repositories", baseDir, 5, testGetSessionID)
 
 	// Call the wrapped handler
 	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{
@@ -211,7 +211,7 @@ func TestMiddlewareWithLargePayload(t *testing.T) {
 		}, nil
 	}
 
-	wrappedHandler := WrapToolHandler(mockHandler, "test_tool", baseDir, testGetSessionID)
+	wrappedHandler := WrapToolHandler(mockHandler, "test_tool", baseDir, 5, testGetSessionID)
 	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
 
 	require.NoError(t, err)
@@ -273,7 +273,7 @@ func TestMiddlewareDirectoryCreation(t *testing.T) {
 		return &sdk.CallToolResult{IsError: false}, map[string]interface{}{"test": "data"}, nil
 	}
 
-	wrappedHandler := WrapToolHandler(mockHandler, "test_tool", baseDir, testGetSessionID)
+	wrappedHandler := WrapToolHandler(mockHandler, "test_tool", baseDir, 5, testGetSessionID)
 	result, data, err := wrappedHandler(context.Background(), &sdk.CallToolRequest{}, map[string]interface{}{})
 
 	require.NoError(t, err)