Skip to content

Commit a8e54eb

Browse files
committed
refactor(middleware): define PayloadMetadata as struct with constant for instructions
- Add PayloadMetadata struct for type-safe payload metadata responses - Extract PayloadTruncatedInstructions as a constant - Update applyJqSchema to return interface{} instead of JSON string - Update tests with helper functions for struct-to-map conversion
1 parent 6b2b196 commit a8e54eb

5 files changed

Lines changed: 92 additions & 232 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ You are an AI agent trying to verify a secret in a file hosted by the filesystem
99

1010
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.
1111

12+
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.
13+
1214
## Important Notes
1315

1416
- **Keep all outputs concise** - Use brief, factual statements

.github/workflows/large-payload-tester-README.md

Lines changed: 0 additions & 181 deletions
This file was deleted.

internal/middleware/jqschema.go

Lines changed: 39 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,22 @@ import (
1717

1818
var logMiddleware = logger.New("middleware:jqschema")
1919

20+
// PayloadTruncatedInstructions is the message returned to clients when a payload
21+
// has been truncated and saved to the filesystem
22+
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."
23+
24+
// PayloadMetadata represents the metadata response returned when a payload is too large
25+
// and has been saved to the filesystem
26+
type PayloadMetadata struct {
27+
QueryID string `json:"queryID"`
28+
PayloadPath string `json:"payloadPath"`
29+
Preview string `json:"preview"`
30+
Schema interface{} `json:"schema"`
31+
OriginalSize int `json:"originalSize"`
32+
Truncated bool `json:"truncated"`
33+
Instructions string `json:"instructions"`
34+
}
35+
2036
// jqSchemaFilter is the jq filter that transforms JSON to schema
2137
// This is the same logic as in gh-aw shared/jqschema.md
2238
const jqSchemaFilter = `
@@ -71,17 +87,18 @@ func generateRandomID() string {
7187
// applyJqSchema applies the jq schema transformation to JSON data
7288
// Uses pre-compiled query code for better performance (3-10x faster than parsing on each request)
7389
// Accepts a context for timeout and cancellation support
74-
func applyJqSchema(ctx context.Context, jsonData interface{}) (string, error) {
90+
// Returns the schema as an interface{} object (not a JSON string)
91+
func applyJqSchema(ctx context.Context, jsonData interface{}) (interface{}, error) {
7592
// Check if compilation succeeded at init time
7693
if jqSchemaCompileErr != nil {
77-
return "", jqSchemaCompileErr
94+
return nil, jqSchemaCompileErr
7895
}
7996

8097
// Run the pre-compiled query with context support (much faster than Parse+Run)
8198
iter := jqSchemaCode.RunWithContext(ctx, jsonData)
8299
v, ok := iter.Next()
83100
if !ok {
84-
return "", fmt.Errorf("jq schema filter returned no results")
101+
return nil, fmt.Errorf("jq schema filter returned no results")
85102
}
86103

87104
// Check for errors with type-specific handling
@@ -90,22 +107,17 @@ func applyJqSchema(ctx context.Context, jsonData interface{}) (string, error) {
90107
if haltErr, ok := err.(*gojq.HaltError); ok {
91108
// HaltError with nil value means clean halt (not an error)
92109
if haltErr.Value() == nil {
93-
return "", fmt.Errorf("jq schema filter halted cleanly with no output")
110+
return nil, fmt.Errorf("jq schema filter halted cleanly with no output")
94111
}
95112
// HaltError with non-nil value is an actual error
96-
return "", fmt.Errorf("jq schema filter halted with error (exit code %d): %w", haltErr.ExitCode(), err)
113+
return nil, fmt.Errorf("jq schema filter halted with error (exit code %d): %w", haltErr.ExitCode(), err)
97114
}
98115
// Generic error case
99-
return "", fmt.Errorf("jq schema filter error: %w", err)
100-
}
101-
102-
// Convert result to JSON
103-
schemaJSON, err := json.Marshal(v)
104-
if err != nil {
105-
return "", fmt.Errorf("failed to marshal schema result: %w", err)
116+
return nil, fmt.Errorf("jq schema filter error: %w", err)
106117
}
107118

108-
return string(schemaJSON), nil
119+
// Return the schema object directly (no JSON marshaling needed here)
120+
return v, nil
109121
}
110122

111123
// savePayload saves the payload to disk and returns the file path
@@ -234,7 +246,7 @@ func WrapToolHandler(
234246

235247
// Apply jq schema transformation
236248
logger.LogDebug("payload", "Applying jq schema transformation: tool=%s, queryID=%s", toolName, queryID)
237-
var schemaJSON string
249+
var schemaObj interface{}
238250
if schemaErr := func() error {
239251
// Unmarshal to interface{} for jq processing
240252
var jsonData interface{}
@@ -246,7 +258,7 @@ func WrapToolHandler(
246258
if err != nil {
247259
return err
248260
}
249-
schemaJSON = schema
261+
schemaObj = schema
250262
return nil
251263
}(); schemaErr != nil {
252264
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(
256268
return result, data, err
257269
}
258270

271+
// Calculate schema size for logging (marshal temporarily)
272+
schemaBytes, _ := json.Marshal(schemaObj)
259273
logger.LogDebug("payload", "Schema transformation completed: tool=%s, queryID=%s, schemaSize=%d bytes",
260-
toolName, queryID, len(schemaJSON))
274+
toolName, queryID, len(schemaBytes))
261275

262276
// Build the transformed response: first 500 chars + schema
263277
payloadStr := string(payloadJSON)
@@ -273,27 +287,22 @@ func WrapToolHandler(
273287
toolName, queryID, len(payloadStr))
274288
}
275289

276-
// Create rewritten response
277-
rewrittenResponse := map[string]interface{}{
278-
"queryID": queryID,
279-
"payloadPath": filePath,
280-
"preview": preview,
281-
"schema": schemaJSON,
282-
"originalSize": len(payloadJSON),
283-
"truncated": truncated,
290+
// Create rewritten response using the PayloadMetadata struct
291+
rewrittenResponse := PayloadMetadata{
292+
QueryID: queryID,
293+
PayloadPath: filePath,
294+
Preview: preview,
295+
Schema: schemaObj,
296+
OriginalSize: len(payloadJSON),
297+
Truncated: truncated,
298+
Instructions: PayloadTruncatedInstructions,
284299
}
285300

286301
logMiddleware.Printf("Rewritten response: tool=%s, queryID=%s, sessionID=%s, originalSize=%d, truncated=%v",
287302
toolName, queryID, sessionID, len(payloadJSON), truncated)
288303
logger.LogInfo("payload", "Created metadata response for client: tool=%s, queryID=%s, session=%s, payloadPath=%s, originalSize=%d bytes, truncated=%v",
289304
toolName, queryID, sessionID, filePath, len(payloadJSON), truncated)
290305

291-
// Parse the schema JSON string back to an object for cleaner display
292-
var schemaObj interface{}
293-
if err := json.Unmarshal([]byte(schemaJSON), &schemaObj); err == nil {
294-
rewrittenResponse["schema"] = schemaObj
295-
}
296-
297306
// Marshal the rewritten response to JSON for the Content field
298307
rewrittenJSON, marshalErr := json.Marshal(rewrittenResponse)
299308
if marshalErr != nil {

internal/middleware/jqschema_integration_test.go

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,21 @@ import (
1212
"github.com/stretchr/testify/require"
1313
)
1414

15+
// integrationPayloadMetadataToMap converts PayloadMetadata to map[string]interface{} for test assertions
16+
func integrationPayloadMetadataToMap(t *testing.T, data interface{}) map[string]interface{} {
17+
t.Helper()
18+
pm, ok := data.(PayloadMetadata)
19+
if !ok {
20+
t.Fatalf("expected PayloadMetadata, got %T", data)
21+
}
22+
jsonBytes, err := json.Marshal(pm)
23+
require.NoError(t, err)
24+
var result map[string]interface{}
25+
err = json.Unmarshal(jsonBytes, &result)
26+
require.NoError(t, err)
27+
return result
28+
}
29+
1530
// TestMiddlewareIntegration tests the complete middleware flow
1631
func TestMiddlewareIntegration(t *testing.T) {
1732
// Create temporary directory for test
@@ -88,8 +103,7 @@ func TestMiddlewareIntegration(t *testing.T) {
88103
assert.Len(t, queryIDFromContent, 32, "QueryID should be 32 hex characters")
89104

90105
// Verify response structure in data return value (for internal use)
91-
dataMap, ok := data.(map[string]interface{})
92-
require.True(t, ok, "Response should be a map")
106+
dataMap := integrationPayloadMetadataToMap(t, data)
93107

94108
// Check all required fields exist
95109
assert.Contains(t, dataMap, "queryID")
@@ -170,7 +184,7 @@ func TestMiddlewareIntegration(t *testing.T) {
170184
assert.False(t, dataMap["truncated"].(bool), "Should not be truncated for small payloads")
171185

172186
// Verify originalSize
173-
originalSize := dataMap["originalSize"].(int)
187+
originalSize := int(dataMap["originalSize"].(float64))
174188
assert.Greater(t, originalSize, 0, "Original size should be positive")
175189
}
176190

@@ -222,7 +236,7 @@ func TestMiddlewareWithLargePayload(t *testing.T) {
222236
}
223237

224238
// Also check data return value
225-
dataMap := data.(map[string]interface{})
239+
dataMap := integrationPayloadMetadataToMap(t, data)
226240

227241
// Verify truncation occurred
228242
truncated := dataMap["truncated"].(bool)
@@ -277,7 +291,7 @@ func TestMiddlewareDirectoryCreation(t *testing.T) {
277291
queryIDFromContent := contentMap["queryID"].(string)
278292

279293
// Also check data return value
280-
dataMap := data.(map[string]interface{})
294+
dataMap := integrationPayloadMetadataToMap(t, data)
281295
queryID := dataMap["queryID"].(string)
282296

283297
// Both should match

0 commit comments

Comments
 (0)