Skip to content

Add session recovery for stale backend sessions #56

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andrmnl wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add session recovery for stale backend sessions #56

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e8553c
Fix Mcp-Session-Id header dropped on SSE responses
andrmnl Apr 6, 2026
fb48ebc
Add session recovery for stale backend sessions
andrmnl Apr 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 100 additions & 28 deletions internal/server/streamable_proxy.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ package server
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"maps"
"net/http"
Expand All @@ -13,8 +15,10 @@ import (
"github.com/stainless-api/mcp-front/internal/log"
)

var errNoBackendSessionID = errors.New("backend did not return Mcp-Session-Id header")

// forwardStreamablePostToBackend handles POST requests for streamable-http transport
func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, config *config.MCPClientConfig) {
func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, cfg *config.MCPClientConfig) {
body, err := io.ReadAll(r.Body)
if err != nil {
log.LogErrorWithFields("streamable_proxy", "Failed to read request body", map[string]any{
Expand All @@ -24,42 +28,49 @@ func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter,
return
}

req, err := http.NewRequestWithContext(ctx, http.MethodPost, config.URL, bytes.NewReader(body))
resp, err := sendStreamablePost(ctx, body, r.Header, cfg)
if err != nil {
log.LogErrorWithFields("streamable_proxy", "Failed to create backend request", map[string]any{
log.LogErrorWithFields("streamable_proxy", "Backend request failed", map[string]any{
"error": err.Error(),
"url": cfg.URL,
})
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "Failed to create request")
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend request failed")
return
}

// Copy relevant headers from original request, excluding hop-by-hop and sensitive headers
copyRequestHeaders(req.Header, r.Header)
// If 404, the backend session is stale (e.g. backend restarted).
// Re-initialize a fresh session with the backend and retry.
if resp.StatusCode == http.StatusNotFound && r.Header.Get("Mcp-Session-Id") != "" {
_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()

// Add configured headers (e.g., auth headers)
for k, v := range config.Headers {
req.Header.Set(k, v)
}
log.LogInfoWithFields("streamable_proxy", "Backend session stale, re-initializing", map[string]any{
"backendURL": cfg.URL,
})

req.Header.Set("Accept", "application/json, text/event-stream")
newSessionID, err := initBackendSession(ctx, r.Header, cfg)
if err != nil {
log.LogErrorWithFields("streamable_proxy", "Failed to re-initialize backend session", map[string]any{
"error": err.Error(),
"backendURL": cfg.URL,
})
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend session recovery failed")
return
}

log.LogDebugWithFields("streamable_proxy", "Forwarding POST to backend", map[string]any{
"backendURL": config.URL,
"method": r.Method,
"headers": config.Headers,
})
// Retry the original request with the new backend session
retryHeaders := r.Header.Clone()
retryHeaders.Set("Mcp-Session-Id", newSessionID)

client := &http.Client{
Timeout: config.Timeout,
}
resp, err := client.Do(req)
if err != nil {
log.LogErrorWithFields("streamable_proxy", "Backend request failed", map[string]any{
"error": err.Error(),
"url": config.URL,
})
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend request failed")
return
resp, err = sendStreamablePost(ctx, body, retryHeaders, cfg)
if err != nil {
log.LogErrorWithFields("streamable_proxy", "Backend request failed after session recovery", map[string]any{
"error": err.Error(),
"url": cfg.URL,
})
jsonrpc.WriteError(w, nil, jsonrpc.InternalError, "backend request failed")
return
}
}
defer resp.Body.Close()

Expand All @@ -71,7 +82,7 @@ func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter,
})

for k, v := range resp.Header {
if k == "Content-Type" || k == "Cache-Control" || k == "Connection" {
if k == "Content-Type" || k == "Cache-Control" || k == "Connection" || k == "Mcp-Session-Id" {
w.Header()[k] = v
}
}
Expand All @@ -97,3 +108,64 @@ func forwardStreamablePostToBackend(ctx context.Context, w http.ResponseWriter,
}
}
}

// sendStreamablePost sends a single POST request to the backend.
func sendStreamablePost(ctx context.Context, body []byte, srcHeaders http.Header, cfg *config.MCPClientConfig) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodPost, cfg.URL, bytes.NewReader(body))
if err != nil {
return nil, err
}

copyRequestHeaders(req.Header, srcHeaders)
for k, v := range cfg.Headers {
req.Header.Set(k, v)
}
req.Header.Set("Accept", "application/json, text/event-stream")

return (&http.Client{Timeout: cfg.Timeout}).Do(req)
}

// initBackendSession creates a fresh MCP session with the backend by sending
// initialize + notifications/initialized, and returns the new Mcp-Session-Id.
func initBackendSession(ctx context.Context, srcHeaders http.Header, cfg *config.MCPClientConfig) (string, error) {
headers := srcHeaders.Clone()
headers.Del("Mcp-Session-Id")

// Step 1: send initialize
initBody := []byte(`{"jsonrpc":"2.0","id":"_mcp_front_reinit","method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"mcp-front","version":"1.0.0"}}}`)

resp, err := sendStreamablePost(ctx, initBody, headers, cfg)
if err != nil {
return "", err
}
_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()

if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("backend initialize returned status %d", resp.StatusCode)
}

newSessionID := resp.Header.Get("Mcp-Session-Id")
if newSessionID == "" {
return "", errNoBackendSessionID
}

log.LogInfoWithFields("streamable_proxy", "New backend session established", map[string]any{
"backendURL": cfg.URL,
"sessionID": newSessionID,
})

// Step 2: send initialized notification
notifyBody := []byte(`{"jsonrpc":"2.0","method":"notifications/initialized"}`)
notifyHeaders := headers.Clone()
notifyHeaders.Set("Mcp-Session-Id", newSessionID)

resp, err = sendStreamablePost(ctx, notifyBody, notifyHeaders, cfg)
if err != nil {
return newSessionID, nil // session exists, notification failure is non-fatal
}
Comment on lines +163 to +166

@gemini-code-assist gemini-code-assist Bot Apr 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If the notifications/initialized call fails, the function returns the newSessionID and a nil error. While the comment states this is non-fatal, it might lead to issues if the backend expects this notification to finalize the session state. Consider if this should at least be logged as a warning, or if the error should be propagated to ensure the session is fully ready before use.

_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()

return newSessionID, nil
}
Loading