Skip to content

Commit 0cdde8a

Browse files
Copilotlpcox
Copilot
and
lpcox
committed
Add command-line flag and environment variable for payload size threshold
Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
1 parent 6eeedf6 commit 0cdde8a

5 files changed

Lines changed: 152 additions & 14 deletions

File tree

‎README.md‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -224,15 +224,16 @@ Available Commands:
224224
help Help about any command
225225
226226
Flags:
227-
-c, --config string Path to config file
228-
--config-stdin Read MCP server configuration from stdin (JSON format). When enabled, overrides --config
229-
--enable-difc Enable DIFC enforcement and session requirement (requires sys___init call before tool access)
230-
--env string Path to .env file to load environment variables
231-
-h, --help help for awmg
232-
-l, --listen string HTTP server listen address (default "127.0.0.1:3000")
233-
--log-dir string Directory for log files (falls back to stdout if directory cannot be created) (default "/tmp/gh-aw/mcp-logs")
234-
--payload-dir string Directory for storing large payload files (segmented by session ID) (default "/tmp/jq-payloads")
235-
--routed Run in routed mode (each backend at /mcp/<server>)
227+
-c, --config string Path to config file
228+
--config-stdin Read MCP server configuration from stdin (JSON format). When enabled, overrides --config
229+
--enable-difc Enable DIFC enforcement and session requirement (requires sys___init call before tool access)
230+
--env string Path to .env file to load environment variables
231+
-h, --help help for awmg
232+
-l, --listen string HTTP server listen address (default "127.0.0.1:3000")
233+
--log-dir string Directory for log files (falls back to stdout if directory cannot be created) (default "/tmp/gh-aw/mcp-logs")
234+
--payload-dir string Directory for storing large payload files (segmented by session ID) (default "/tmp/jq-payloads")
235+
--payload-size-threshold int Size threshold (in bytes) for storing payloads to disk. Payloads larger than this are stored, smaller ones returned inline (default 1024)
236+
--routed Run in routed mode (each backend at /mcp/<server>)
236237
--sequential-launch Launch MCP servers sequentially during startup (parallel launch is default)
237238
--unified Run in unified mode (all backends at /mcp)
238239
--validate-env Validate execution environment (Docker, env vars) before starting
@@ -269,6 +270,7 @@ When running locally (`run.sh`), these variables are optional (warnings shown if
269270
| `MODE` | Gateway mode flag | `--routed` |
270271
| `MCP_GATEWAY_LOG_DIR` | Log file directory (sets default for `--log-dir` flag) | `/tmp/gh-aw/mcp-logs` |
271272
| `MCP_GATEWAY_PAYLOAD_DIR` | Large payload storage directory (sets default for `--payload-dir` flag) | `/tmp/jq-payloads` |
273+
| `MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD` | Size threshold in bytes for payload storage (sets default for `--payload-size-threshold` flag) | `1024` |
272274

273275
### Docker Configuration
274276

‎config.example-payload-threshold.toml‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,35 @@
11
# Example MCP Gateway Configuration with Payload Size Threshold
22
# This demonstrates the configurable payload size threshold feature
3+
#
4+
# Configuration Priority (highest to lowest):
5+
# 1. Command-line flag: --payload-size-threshold 2048
6+
# 2. Environment variable: MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD=2048
7+
# 3. Config file: payload_size_threshold = 2048
8+
# 4. Default: 1024 bytes
39

410
[gateway]
511
port = 3000
612
api_key = "your-api-key-here"
713

814
# Payload directory for storing large tool responses
15+
# Can also be set via:
16+
# - Flag: --payload-dir /custom/path
17+
# - Env: MCP_GATEWAY_PAYLOAD_DIR=/custom/path
918
# Default: /tmp/jq-payloads
1019
payload_dir = "/tmp/jq-payloads"
1120

1221
# Payload size threshold (in bytes) for storing responses to disk
1322
# Payloads LARGER than this threshold are stored to disk and return metadata
1423
# Payloads SMALLER than or equal to this threshold are returned inline
24+
#
25+
# Can also be set via:
26+
# - Flag: --payload-size-threshold 2048
27+
# - Env: MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD=2048
1528
# Default: 1024 bytes (1KB)
1629
#
1730
# Examples:
18-
# payload_size_threshold = 1024 # 1KB - default, good for most use cases
1931
# payload_size_threshold = 512 # 512 bytes - more aggressive file storage
32+
# payload_size_threshold = 1024 # 1KB - default, good for most use cases
2033
# payload_size_threshold = 2048 # 2KB - fewer files, more inline responses
2134
# payload_size_threshold = 10240 # 10KB - minimal file storage
2235
payload_size_threshold = 1024

‎internal/cmd/flags_logging.go‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,27 +3,31 @@ package cmd
33
// Logging-related flags
44

55
import (
6+
"fmt"
67
"os"
78

89
"github.com/spf13/cobra"
910
)
1011

1112
// Logging flag defaults
1213
const (
13-
defaultLogDir = "/tmp/gh-aw/mcp-logs"
14-
defaultPayloadDir = "/tmp/jq-payloads"
14+
defaultLogDir = "/tmp/gh-aw/mcp-logs"
15+
defaultPayloadDir = "/tmp/jq-payloads"
16+
defaultPayloadSizeThreshold = 1024 // 1KB default threshold
1517
)
1618

1719
// Logging flag variables
1820
var (
19-
logDir string
20-
payloadDir string
21+
logDir string
22+
payloadDir string
23+
payloadSizeThreshold int
2124
)
2225

2326
func init() {
2427
RegisterFlag(func(cmd *cobra.Command) {
2528
cmd.Flags().StringVar(&logDir, "log-dir", getDefaultLogDir(), "Directory for log files (falls back to stdout if directory cannot be created)")
2629
cmd.Flags().StringVar(&payloadDir, "payload-dir", getDefaultPayloadDir(), "Directory for storing large payload files (segmented by session ID)")
30+
cmd.Flags().IntVar(&payloadSizeThreshold, "payload-size-threshold", getDefaultPayloadSizeThreshold(), "Size threshold (in bytes) for storing payloads to disk. Payloads larger than this are stored, smaller ones returned inline")
2731
})
2832
}
2933

@@ -44,3 +48,17 @@ func getDefaultPayloadDir() string {
4448
}
4549
return defaultPayloadDir
4650
}
51+
52+
// getDefaultPayloadSizeThreshold returns the default payload size threshold, checking
53+
// MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD environment variable first, then falling back to the hardcoded default
54+
func getDefaultPayloadSizeThreshold() int {
55+
if envThreshold := os.Getenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD"); envThreshold != "" {
56+
// Try to parse as integer
57+
var threshold int
58+
if _, err := fmt.Sscanf(envThreshold, "%d", &threshold); err == nil && threshold > 0 {
59+
return threshold
60+
}
61+
// Invalid value, use default
62+
}
63+
return defaultPayloadSizeThreshold
64+
}

‎internal/cmd/flags_logging_test.go‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package cmd
2+
3+
import (
4+
"os"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
)
9+
10+
func TestGetDefaultPayloadSizeThreshold(t *testing.T) {
11+
tests := []struct {
12+
name string
13+
envValue string
14+
expected int
15+
}{
16+
{
17+
name: "no env var - returns default",
18+
envValue: "",
19+
expected: defaultPayloadSizeThreshold,
20+
},
21+
{
22+
name: "valid env var",
23+
envValue: "2048",
24+
expected: 2048,
25+
},
26+
{
27+
name: "very large threshold",
28+
envValue: "10240",
29+
expected: 10240,
30+
},
31+
{
32+
name: "small threshold",
33+
envValue: "512",
34+
expected: 512,
35+
},
36+
{
37+
name: "invalid value - non-numeric",
38+
envValue: "invalid",
39+
expected: defaultPayloadSizeThreshold,
40+
},
41+
{
42+
name: "invalid value - negative",
43+
envValue: "-100",
44+
expected: defaultPayloadSizeThreshold,
45+
},
46+
{
47+
name: "invalid value - zero",
48+
envValue: "0",
49+
expected: defaultPayloadSizeThreshold,
50+
},
51+
}
52+
53+
for _, tt := range tests {
54+
t.Run(tt.name, func(t *testing.T) {
55+
// Set or unset environment variable
56+
if tt.envValue != "" {
57+
os.Setenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD", tt.envValue)
58+
defer os.Unsetenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD")
59+
} else {
60+
os.Unsetenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD")
61+
}
62+
63+
result := getDefaultPayloadSizeThreshold()
64+
assert.Equal(t, tt.expected, result, "Threshold should match expected value")
65+
})
66+
}
67+
}
68+
69+
func TestPayloadSizeThresholdFlagDefault(t *testing.T) {
70+
// Ensure environment is clean
71+
os.Unsetenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD")
72+
73+
result := getDefaultPayloadSizeThreshold()
74+
assert.Equal(t, 1024, result, "Default should be 1024 bytes")
75+
}
76+
77+
func TestPayloadSizeThresholdEnvVar(t *testing.T) {
78+
// Test that environment variable overrides default
79+
os.Setenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD", "4096")
80+
defer os.Unsetenv("MCP_GATEWAY_PAYLOAD_SIZE_THRESHOLD")
81+
82+
result := getDefaultPayloadSizeThreshold()
83+
assert.Equal(t, 4096, result, "Environment variable should override default")
84+
}

‎internal/cmd/root.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,27 @@ func run(cmd *cobra.Command, args []string) error {
222222
cfg.EnableDIFC = enableDIFC
223223
cfg.SequentialLaunch = sequentialLaunch
224224

225+
// Override gateway config with command-line flags
226+
if cfg.Gateway == nil {
227+
cfg.Gateway = &config.GatewayConfig{}
228+
}
229+
230+
// Apply payload directory flag (if different from default, it was explicitly set)
231+
if cmd.Flags().Changed("payload-dir") {
232+
cfg.Gateway.PayloadDir = payloadDir
233+
} else if payloadDir != "" && payloadDir != defaultPayloadDir {
234+
// Environment variable was set
235+
cfg.Gateway.PayloadDir = payloadDir
236+
}
237+
238+
// Apply payload size threshold flag (if different from default, it was explicitly set)
239+
if cmd.Flags().Changed("payload-size-threshold") {
240+
cfg.Gateway.PayloadSizeThreshold = payloadSizeThreshold
241+
} else if payloadSizeThreshold != defaultPayloadSizeThreshold {
242+
// Environment variable was set
243+
cfg.Gateway.PayloadSizeThreshold = payloadSizeThreshold
244+
}
245+
225246
if enableDIFC {
226247
log.Println("DIFC enforcement and session requirement enabled")
227248
} else {

0 commit comments

Comments
 (0)