Skip to content

Commit 7af5213

Browse files
dgellowdgellow
committed
Add session configuration for timeouts and per-user limits
Allows configuring session timeout, cleanup interval, and max sessions per user through the config file instead of hardcoded values.
1 parent 8aaa7f7 commit 7af5213

File tree

10 files changed

+37
-23
lines changed

10 files changed

+37
-23
lines changed

‎integration/oauth_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -897,7 +897,7 @@ func TestToolAdvertisementWithUserTokens(t *testing.T) {
897897
assert.Equal(t, "token_required", errorInfo["code"], "Error code should be token_required")
898898

899899
errorMessage := errorInfo["message"].(string)
900-
assert.Contains(t, errorMessage, "Token Required", "Error should mention token required")
900+
assert.Contains(t, errorMessage, "token required", "Error should mention token required")
901901
assert.Contains(t, errorMessage, "/my/tokens", "Error should mention token setup URL")
902902
assert.Contains(t, errorMessage, "Test Service", "Error should mention service name")
903903

‎internal/client/client.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,7 @@ func (c *Client) wrapToolHandler(
436436
serverName,
437437
setupBaseURL,
438438
tokenSetup,
439-
"Configuration error: This service requires user tokens but OAuth is not properly configured.",
439+
"configuration error: this service requires user tokens but OAuth is not properly configured.",
440440
)
441441

442442
errorJSON, _ := json.Marshal(errorData)
@@ -451,17 +451,17 @@ func (c *Client) wrapToolHandler(
451451
var errorMessage string
452452
if tokenSetup != nil {
453453
errorMessage = fmt.Sprintf(
454-
"Token Required: %s requires a user token to access the API. "+
455-
"Please visit %s to set up your %s token. %s",
454+
"token required: %s requires a user token to access the API. "+
455+
"please visit %s to set up your %s token. %s",
456456
tokenSetup.DisplayName,
457457
tokenSetupURL,
458458
tokenSetup.DisplayName,
459459
tokenSetup.Instructions,
460460
)
461461
} else {
462462
errorMessage = fmt.Sprintf(
463-
"Token Required: This service requires a user token. "+
464-
"Please visit %s to configure your token.",
463+
"token required: this service requires a user token. "+
464+
"please visit %s to configure your token.",
465465
tokenSetupURL,
466466
)
467467
}

‎internal/config/load.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,12 @@ func ValidateConfig(config *Config) error {
142142
if config.Proxy.Sessions.Timeout > 0 && config.Proxy.Sessions.CleanupInterval > config.Proxy.Sessions.Timeout {
143143
internal.LogWarn("Session cleanup interval is greater than session timeout")
144144
}
145+
if config.Proxy.Sessions.MaxPerUser < 0 {
146+
return fmt.Errorf("proxy.sessions.maxPerUser cannot be negative")
147+
}
148+
if config.Proxy.Sessions.MaxPerUser == 0 {
149+
internal.LogWarn("Session maxPerUser is 0 (unlimited) - this may allow resource exhaustion")
150+
}
145151
}
146152

147153
return nil

‎internal/config/types.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ type MCPClientConfig struct {
116116
type SessionConfig struct {
117117
Timeout time.Duration
118118
CleanupInterval time.Duration
119+
MaxPerUser int
119120
}
120121

121122
// AdminConfig represents admin UI configuration

‎internal/config/unmarshal.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ func (o *OAuthAuthConfig) UnmarshalJSON(data []byte) error {
181181

182182
// Validate JWT secret length
183183
if len(o.JWTSecret) < 32 {
184-
return fmt.Errorf("JWT secret must be at least 32 bytes, got %d", len(o.JWTSecret))
184+
return fmt.Errorf("jwt secret must be at least 32 bytes, got %d", len(o.JWTSecret))
185185
}
186186

187187
// Validate encryption key if storage requires it
@@ -351,6 +351,7 @@ func (s *SessionConfig) UnmarshalJSON(data []byte) error {
351351
var raw struct {
352352
Timeout string `json:"timeout"`
353353
CleanupInterval string `json:"cleanupInterval"`
354+
MaxPerUser *int `json:"maxPerUser"` // Pointer to detect explicit 0
354355
}
355356

356357
if err := json.Unmarshal(data, &raw); err != nil {
@@ -375,5 +376,10 @@ func (s *SessionConfig) UnmarshalJSON(data []byte) error {
375376
s.CleanupInterval = interval
376377
}
377378

379+
// Set MaxPerUser if present (0 is a valid value, means no upper bound)
380+
if raw.MaxPerUser != nil {
381+
s.MaxPerUser = *raw.MaxPerUser
382+
}
383+
378384
return nil
379385
}

‎internal/config/unmarshal_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ func TestOAuthAuthConfig_ValidationErrors(t *testing.T) {
289289
"jwtSecret": {"$env": "SHORT_SECRET"}
290290
}`,
291291
envVars: map[string]string{"SHORT_SECRET": "too-short"},
292-
expectedError: "JWT secret must be at least 32 bytes",
292+
expectedError: "jwt secret must be at least 32 bytes",
293293
},
294294
{
295295
name: "user token in oauth field",

‎internal/inline/handler.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,12 @@ func NewHandler(name string, server MCPServer) *Handler {
3737

3838
// ServeHTTP handles HTTP requests for the inline MCP server
3939
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
40-
if r.URL.Path == "/sse" || r.URL.Path == "/"+h.name+"/sse" {
40+
switch r.URL.Path {
41+
case "/sse", "/" + h.name + "/sse":
4142
h.handleSSE(w, r)
42-
} else if r.URL.Path == "/message" || r.URL.Path == "/"+h.name+"/message" {
43+
case "/message", "/" + h.name + "/message":
4344
h.handleMessage(w, r)
44-
} else {
45+
default:
4546
jsonwriter.WriteNotFound(w, "Endpoint not found")
4647
}
4748
}

‎internal/server/mcp_handler.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ func (h *MCPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
9898
})
9999
h.handleStreamableGet(ctx, w, r, userEmail, serverConfig)
100100
} else {
101-
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
101+
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
102102
}
103103
} else {
104104
if h.isMessageRequest(r) {
@@ -162,7 +162,7 @@ func (h *MCPHandler) handleSSERequest(ctx context.Context, w http.ResponseWriter
162162
internal.LogErrorWithFields("mcp", "No shared SSE server configured for stdio server", map[string]any{
163163
"server": h.serverName,
164164
})
165-
jsonwriter.WriteInternalServerError(w, "Server misconfiguration")
165+
jsonwriter.WriteInternalServerError(w, "server misconfiguration")
166166
return
167167
}
168168

@@ -196,7 +196,7 @@ func (h *MCPHandler) handleMessageRequest(ctx context.Context, w http.ResponseWr
196196
if isStdioServer(config) {
197197
sessionID := r.URL.Query().Get("sessionId")
198198
if sessionID == "" {
199-
jsonrpc.WriteError(w, nil, jsonrpc.InvalidParams, "Missing sessionId")
199+
jsonrpc.WriteError(w, nil, jsonrpc.InvalidParams, "missing sessionId")
200200
return
201201
}
202202

@@ -221,14 +221,14 @@ func (h *MCPHandler) handleMessageRequest(ctx context.Context, w http.ResponseWr
221221
})
222222
// Per MCP spec: return HTTP 404 Not Found when session is terminated or not found
223223
// The response body MAY comprise a JSON-RPC error response
224-
jsonrpc.WriteErrorWithStatus(w, nil, jsonrpc.InvalidParams, "Session not found", http.StatusNotFound)
224+
jsonrpc.WriteErrorWithStatus(w, nil, jsonrpc.InvalidParams, "session not found", http.StatusNotFound)
225225
return
226226
}
227227
if h.sharedSSEServer == nil {
228228
internal.LogErrorWithFields("mcp", "No shared SSE server configured", map[string]any{
229229
"sessionID": sessionID,
230230
})
231-
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Server misconfiguration")
231+
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "server misconfiguration")
232232
return
233233
}
234234

@@ -293,7 +293,7 @@ func (h *MCPHandler) forwardMessageToBackend(ctx context.Context, w http.Respons
293293
"error": err.Error(),
294294
"server": h.serverName,
295295
})
296-
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Failed to read request")
296+
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "failed to read request")
297297
return
298298
}
299299

@@ -304,7 +304,7 @@ func (h *MCPHandler) forwardMessageToBackend(ctx context.Context, w http.Respons
304304
"server": h.serverName,
305305
"url": backendURL,
306306
})
307-
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Failed to create request")
307+
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "failed to create request")
308308
return
309309
}
310310

@@ -336,7 +336,7 @@ func (h *MCPHandler) forwardMessageToBackend(ctx context.Context, w http.Respons
336336
"server": h.serverName,
337337
"url": backendURL,
338338
})
339-
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Backend request failed")
339+
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend request failed")
340340
return
341341
}
342342
defer resp.Body.Close()

‎internal/server/mcp_handler_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ func TestForwardMessageToBackend_ErrorCases(t *testing.T) {
289289

290290
assert.NotNil(t, response.Error)
291291
assert.Equal(t, jsonrpc.InternalError, response.Error.Code)
292-
assert.Equal(t, "Backend request failed", response.Error.Message)
292+
assert.Equal(t, "backend request failed", response.Error.Message)
293293
})
294294

295295
t.Run("backend returns error status", func(t *testing.T) {
@@ -333,7 +333,7 @@ func TestForwardMessageToBackend_ErrorCases(t *testing.T) {
333333

334334
assert.NotNil(t, response.Error)
335335
assert.Equal(t, jsonrpc.InternalError, response.Error.Code)
336-
assert.Equal(t, "Failed to read request", response.Error.Message)
336+
assert.Equal(t, "failed to read request", response.Error.Message)
337337
})
338338
}
339339

@@ -529,7 +529,7 @@ func TestHandleStreamablePost(t *testing.T) {
529529

530530
assert.NotNil(t, response.Error)
531531
assert.Equal(t, jsonrpc.InternalError, response.Error.Code)
532-
assert.Equal(t, "Backend request failed", response.Error.Message)
532+
assert.Equal(t, "backend request failed", response.Error.Message)
533533
})
534534
}
535535

‎internal/server/streamable_proxy.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter,
5555
"error": err.Error(),
5656
"url": config.URL,
5757
})
58-
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Backend request failed")
58+
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend request failed")
5959
return
6060
}
6161
defer resp.Body.Close()

0 commit comments

Comments
 (0)