changes from review

Author: yrobla

rfcs/THV-0038-session-scoped-client-lifecycle.md

@@ -616,7 +616,7 @@ The default factory implementation follows this pattern:
    - **Client initialization includes MCP handshake**: Each client sends `InitializeRequest` to its backend, and the backend responds with capabilities and its own `Mcp-Session-Id`. The client stores the session ID for protocol compliance (includes it in subsequent request headers).
    - **Capture backend session IDs**: Factory also captures each backend's session ID (via `client.SessionID()`) for observability, storing them in a map to pass to the session
    - **Performance requirement**: Use parallel initialization (e.g., `errgroup` with bounded concurrency) to avoid sequential latency accumulation. Connection initialization (TCP handshake + TLS negotiation + MCP protocol handshake) can take tens to hundreds of milliseconds per backend depending on network latency and backend responsiveness. With 20 backends, sequential initialization could easily exceed acceptable session creation latency.
-   - **Bounded concurrency**: Limit parallel goroutines (e.g., 10 concurrent initializations) to avoid resource exhaustion. This limit is **per-session-creation** (not global), implemented as a semaphore inside the factory. It should be a configurable vMCP server-level parameter (e.g., `max_backend_init_concurrency`, default: 10). Operators with many backends on a fast private network can raise it; resource-constrained deployments or backends with expensive initialization should lower it. A global limit across concurrent session creations is not necessary — the per-session semaphore already bounds the worst case per event.
+   - **Bounded concurrency**: Limit parallel goroutines (e.g., 10 concurrent initializations) per session to avoid resource exhaustion. This limit is **per-session-creation**, implemented as a semaphore inside the factory via a configurable vMCP server-level parameter (`max_backend_init_concurrency`, default: 10). Operators should note that the aggregate system load is protected by the global `TOOLHIVE_MAX_SESSIONS` limit (see [Resource Exhaustion & DoS Protection](#concurrency--resource-safety)). For additional safety during traffic spikes, the factory may optionally implement a global initialization semaphore (e.g., `max_global_backend_init_concurrency`, default: 100) to cap the total number of simultaneous connection attempts across all active session creations (preventing, for example, 100 concurrent session requests from triggering 1,000 backend initializations).
    - **Per-backend timeout**: Apply context timeout (e.g., 5s per backend) so one slow backend doesn't block session creation
    - **Partial initialization**: If some backends fail, log warnings and continue with successfully initialized backends (failed backends not added to clients map)
    - Clients are connection-ready and stateful (each maintains its backend session for protocol use)
@@ -706,6 +706,7 @@ The default session implementation stores:
   - Used for: logging, metrics, health checks, debugging, and explicit session cleanup
   - Updated when clients are re-initialized (e.g., after backend session expiration)
 - RWMutex for thread-safe access (read lock for queries/calls, write lock for Close)
+- `singleflight.Group` (or per-backend locks) to coordinate concurrent re-initialization of backend sessions without stalling the whole session
 
 **Backend session ID lifecycle management:**
 
@@ -965,6 +966,9 @@ The session **continues operating** with remaining backends. If a backend client
 func (s *Session) CallTool(ctx context.Context, name string, arguments map[string]any) (*ToolResult, error) {
     backend := s.routingTable.Lookup(name)
     client := s.clients[backend.ID]
+    if client == nil {
+        return nil, fmt.Errorf("no client found for backend %s", backend.ID)
+    }
 
     // Call backend client - if it fails, return the error
     result, err := client.CallTool(ctx, name, arguments)
@@ -1055,9 +1059,19 @@ When vMCP detects backend session expiration (404 or "session expired" error), a
 func (s *Session) CallTool(ctx context.Context, name string, arguments map[string]any) (*ToolResult, error) {
     backend := s.routingTable.Lookup(name)
 
-    // Hold RLock only while reading from shared state; release before network I/O.
+    // Check if session is closed and increment in-flight counter
     s.mu.RLock()
+    if s.closed {
+        s.mu.RUnlock()
+        return nil, fmt.Errorf("session closed")
+    }
     client := s.clients[backend.ID]
+    if client == nil {
+        s.mu.RUnlock()
+        return nil, fmt.Errorf("no client found for backend %s", backend.ID)
+    }
+    s.inFlight.Add(1)
+    defer s.inFlight.Done()
     s.mu.RUnlock()
 
     result, err := client.CallTool(ctx, name, arguments)
@@ -1067,19 +1081,23 @@ func (s *Session) CallTool(ctx context.Context, name string, arguments map[strin
         log.Warn("Backend %s session expired, re-initializing", backend.ID)
 
         // Re-initialize the backend session (acquires write lock internally).
-        if reinitErr := s.reinitializeBackend(ctx, backend.ID); reinitErr != nil {
+        // Pass the failed client to avoid redundant re-initialization if multiple
+        // goroutines detect the failure simultaneously.
+        if reinitErr := s.reinitializeBackend(ctx, backend.ID, client); reinitErr != nil {
             return nil, fmt.Errorf("failed to reinitialize backend %s: %w",
                 backend.ID, reinitErr)
         }
 
         // Retry the operation once with the re-initialized session.
-        // No further retries: if the new session also fails, surface the error
-        // immediately to avoid infinite loops or resource exhaustion.
         // Re-read under RLock: reinitializeBackend replaced the client under write lock.
         s.mu.RLock()
         client = s.clients[backend.ID]
         s.mu.RUnlock()
 
+        if client == nil {
+            return nil, fmt.Errorf("no client found for backend %s after re-initialization", backend.ID)
+        }
+
         result, err = client.CallTool(ctx, name, arguments)
         if err != nil {
             return nil, fmt.Errorf("backend %s failed after re-initialization: %w",
@@ -1091,38 +1109,53 @@ func (s *Session) CallTool(ctx context.Context, name string, arguments map[strin
     return result, err
 }
 
-func (s *Session) reinitializeBackend(ctx context.Context, backendID string) error {
-    s.mu.Lock()
-    defer s.mu.Unlock()
-
-    backend := s.getBackendConfig(backendID)
-
-    // Close old client (with expired backend session)
-    if oldClient := s.clients[backendID]; oldClient != nil {
-        oldClient.Close()
-    }
-
-    // Create NEW client (triggers new Initialize handshake with backend)
-    // This sends InitializeRequest to backend MCP server
-    newClient, err := s.clientFactory.CreateClient(ctx, backend, s.identity)
-    if err != nil {
-        return fmt.Errorf("client creation failed: %w", err)
+func (s *Session) reinitializeBackend(ctx context.Context, backendID string, failedClient *Client) error {
+    // 1. Double-check if re-init is still needed under read lock
+    s.mu.RLock()
+    if s.clients[backendID] != failedClient {
+        s.mu.RUnlock()
+        return nil
     }
+    s.mu.RUnlock()
 
-    // Backend responds with new session ID (in Mcp-Session-Id header)
-    // The client stores it for subsequent requests (protocol requirement)
+    // 2. Coordinate concurrent re-initialization (e.g., using singleflight)
+    // to ensure only one network initialization proceeds for this backendID.
+    _, err, _ := s.reinitGroup.Do(backendID, func() (interface{}, error) {
+        // Re-check after acquiring re-initialization group ownership
+        s.mu.RLock()
+        if s.clients[backendID] != failedClient {
+            s.mu.RUnlock()
+            return nil, nil
+        }
+        s.mu.RUnlock()
 
-    // Capture the new backend session ID for observability
-    newBackendSessionID := newClient.SessionID()  // Assuming client exposes this method
+        backend := s.getBackendConfig(backendID)
 
-    // Update session state
-    oldSessionID := s.backendSessions[backendID]
-    s.backendSessions[backendID] = newBackendSessionID
-    s.clients[backendID] = newClient
+        // 3. Create NEW client (network I/O happens OUTSIDE the session-wide lock)
+        newClient, err := s.clientFactory.CreateClient(ctx, backend, s.identity)
+        if err != nil {
+            return nil, fmt.Errorf("client creation failed: %w", err)
+        }
 
-    log.Info("Backend %s session re-initialized: old=%s, new=%s",
-        backendID, oldSessionID, newBackendSessionID)
-    return nil
+        // 4. Update session state under write lock (short critical section)
+        s.mu.Lock()
+        oldSessionID := s.backendSessions[backendID]
+        newBackendSessionID := newClient.SessionID()
+        s.backendSessions[backendID] = newBackendSessionID
+        s.clients[backendID] = newClient
+        s.mu.Unlock()
+
+        // 5. Close the old client OUTSIDE any locks
+        if failedClient != nil {
+            failedClient.Close()
+        }
+        
+        log.Info("Backend %s session re-initialized: old=%s, new=%s",
+            backendID, oldSessionID, newBackendSessionID)
+        return nil, nil
+    })
+    
+    return err
 }
 ```
 
@@ -1148,7 +1181,7 @@ func (s *Session) reinitializeBackend(ctx context.Context, backendID string) err
    - **Approach**: Best effort - attempt keepalive, gracefully handle backends that don't support it
    - **Configuration**: Enable per backend, configurable interval (default: 5 min)
 
-   The preferred keepalive method is the MCP spec-defined `ping` protocol request, which is side-effect-free and supported by all compliant servers; explicit tool calls should only be used as a fallback. Keepalive failures must not affect healthy sessions — after N consecutive failures the feature should be disabled for that backend, with a periodic probe to re-enable on recovery. Keepalive should default to disabled for stateless backends or where TTL alignment already covers the session lifetime. The keepalive goroutine must hold the backend lock to avoid races with session re-initialization. Operators should be able to observe keepalive health via per-backend metrics covering attempt counts, failure reasons, and auto-disable events.
+   The preferred keepalive method is the MCP spec-defined `ping` protocol request, which is side-effect-free and supported by all compliant servers; explicit tool calls should only be used as a fallback. Keepalive failures must not affect healthy sessions — after N consecutive failures the feature should be disabled for that backend, with a periodic probe to re-enable on recovery. Keepalive should default to disabled for stateless backends or where TTL alignment already covers the session lifetime. The keepalive goroutine must use the same in-flight counter (`sync.WaitGroup`) approach as other operations to avoid races with session re-initialization while ensuring no locks are held during network I/O. Operators should be able to observe keepalive health via per-backend metrics covering attempt counts, failure reasons, and auto-disable events.
 
 2. **Session TTL alignment**:
    - Configure backend session TTLs longer than vMCP session TTL
@@ -1231,10 +1264,11 @@ For initial implementation, we assume most backends use long-lived credentials (
 
 **Required for production deployment**:
 1. **Session binding to authentication token**:
-   - Store a cryptographic hash of the original authentication token (e.g., `SHA256(bearerToken)`) in the session during creation
-   - On each request, validate that the current auth token hash matches the session's bound token hash
-   - If mismatch, reject with "session authentication mismatch" error and terminate session
-   - This prevents stolen session IDs from being used with different credentials
+   - Store a secure cryptographic hash of the original authentication token in the session during creation. To prevent offline attacks if session state is leaked (e.g., from Redis/Valkey), prefer a keyed hash (e.g., `HMAC-SHA256` with a server-managed secret and a per-session salt).
+   - On each request, validate the current auth token against the session's bound hash using a constant-time comparison to prevent timing attacks.
+   - If mismatch, reject with "session authentication mismatch" error and terminate session.
+   - Ensure the hash value is treated as sensitive and is never logged or exposed in traces.
+   - This prevents stolen session IDs from being used with different credentials.
 
 2. **TLS-only enforcement**:
    - Require TLS for all vMCP connections (prevent session ID interception)
@@ -1503,7 +1537,7 @@ if config.SessionManagementV2 {
 - Verify old code path still works when flag is disabled (no regressions)
 
 **Security (blocking for production rollout)**:
-- Implement token hash binding during `CreateSession()`: store `SHA256(bearerToken)` in the session, validate on each subsequent request, reject with "session authentication mismatch" and terminate on mismatch (see Security Considerations → Session Hijacking Prevention). This must be completed before the feature flag is enabled in any production environment.
+- Implement token hash binding during `CreateSession()`: store a secure keyed hash (e.g., `HMAC-SHA256` with a server-managed secret and a per-session salt) of the original authentication token in the session. Validate this hash on each subsequent request using constant-time comparison. Reject with "session authentication mismatch" and terminate the session on mismatch (see Security Considerations → Session Hijacking Prevention). This must be completed before the feature flag is enabled in any production environment.
 
 **Files Modified**:
 - `pkg/vmcp/server/server.go` - Add conditional logic based on `sessionManagementV2` flag
@@ -1537,7 +1571,7 @@ if config.SessionManagementV2 {
 - `pkg/vmcp/server/server.go` - Remove old code path and feature flag conditionals
 - `pkg/vmcp/discovery/middleware.go` - Delete (replaced by SessionFactory)
 - `pkg/vmcp/client/client.go` - Remove `httpBackendClient` (replaced by Session ownership)
-- `pkg/transport/session/manager.go` - Update `DeleteExpired()` to call `session.Close()` before removing from storage (fixes resource leak). Because the storage layer operates on the `Session` interface (which has no `Close()` method), use an optional interface check: `if closer, ok := sess.(io.Closer); ok { closer.Close() }`. This avoids adding `Close()` to the base interface (which would require all existing session types to implement it) while still dispatching cleanup to sessions that carry resources.
+- `pkg/transport/session/manager.go` - Update `DeleteExpired()` to call `session.Close()` before removing from storage (fixes resource leak). Because the storage layer operates on the base `transportsession.Session` interface (which lacks a `Close()` method, unlike the vMCP-specific `Session` interface defined in this RFC), use an optional interface check: `if closer, ok := sess.(io.Closer); ok { _ = closer.Close() }`. This avoids adding `Close()` to the base `transportsession.Session` interface (which would require all existing session types to implement it) while still dispatching cleanup to sessions that carry active resources.
 - Delete old `VMCPSession` implementation files
 
 **Rationale**: Once the new code path is validated, delete the old code entirely rather than maintaining both paths. This avoids technical debt and ongoing maintenance burden.