fix(upstream): re-initialize a terminated Streamable HTTP session and retry once (Spec 113-e)

docs/setup.md

@@ -757,6 +757,10 @@ Get-Content -Path "$env:LOCALAPPDATA\mcpproxy\logs\main.log" -Wait
 tail -f ~/Library/Logs/mcpproxy/main.log | grep -E "(github-server|oauth|error)"
 ```
 
+**Upstream session lost (HTTP 404):**
+
+A remote Streamable HTTP server that restarts or expires its session answers HTTP 404 to the session id MCPProxy holds. MCPProxy re-initializes the MCP session on the same connection instead of marking the server as errored. List, prompt and health-ping requests are retried once. A `tools/call` is retried once only for read-only tools (`readOnlyHint: true`) whose description and schema did not change; any other tool returns an error saying the session was re-established and the call was not repeated, and the next call uses the new session. Each re-initialization is logged at info level (`Upstream session terminated (HTTP 404); re-initialized in place`) and counted as `session_reinit_count` in the server's connection status. If the re-initialization or the retry also fails, the normal reconnect path applies.
+
 ## Advanced Configuration
 
 **📚 For complete configuration reference:** See [Configuration Documentation](configuration.md) for all available options.

internal/upstream/core/session_reinit.go

@@ -0,0 +1,70 @@
+package core
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/mark3labs/mcp-go/mcp"
+	"go.uber.org/zap"
+)
+
+// Spec 113-e (G8): helpers that let the managed layer re-initialize a
+// Streamable HTTP MCP session in place after the upstream answered HTTP 404
+// (mcp-go: transport.ErrSessionTerminated) for a session id it no longer
+// knows, without tearing down the connection or touching connectionEpoch.
+
+// SessionSnapshot reports the Streamable HTTP transport's current session id
+// (empty when the transport has none, e.g. mcp-go cleared it after a 404, or
+// when the transport is stdio/SSE) and whether the negotiated protocol is the
+// stateless 2026-07-28 era, where sessions do not exist and no re-init applies.
+func (c *Client) SessionSnapshot() (id string, modern bool) {
+	c.mu.RLock()
+	cl := c.client
+	info := c.serverInfo
+	c.mu.RUnlock()
+
+	if info != nil {
+		modern = mcp.IsModernProtocol(info.ProtocolVersion)
+	}
+	if cl == nil {
+		return "", modern
+	}
+	if sh, ok := cl.GetTransport().(*transport.StreamableHTTP); ok {
+		return sh.GetSessionId(), modern
+	}
+	return "", modern
+}
+
+// ReinitializeSession performs a fresh initialize + notifications/initialized
+// handshake on the existing mcp-go client and transport, so the upstream
+// issues a new session id. It does not Close/Start the transport and does not
+// change the connection generation. The legacy-era pin (Spec 058 FR-027)
+// applies exactly as on the first handshake.
+func (c *Client) ReinitializeSession(ctx context.Context) error {
+	c.mu.RLock()
+	cl := c.client
+	c.mu.RUnlock()
+	if cl == nil || !c.IsConnected() {
+		return fmt.Errorf("client not connected")
+	}
+
+	req := mcp.InitializeRequest{}
+	req.Params.ProtocolVersion = mcp.LATEST_LEGACY_PROTOCOL_VERSION
+	req.Params.ClientInfo = mcp.Implementation{Name: "mcpproxy-go", Version: "1.0.0"}
+	req.Params.Capabilities = mcp.ClientCapabilities{}
+
+	res, err := cl.Initialize(ctx, req)
+	if err != nil {
+		return fmt.Errorf("session re-initialize failed: %w", err)
+	}
+	if mcp.IsModernProtocol(res.ProtocolVersion) {
+		return fmt.Errorf("upstream answered MCP protocol version %s on session re-initialize (spec 058 FR-027)", res.ProtocolVersion)
+	}
+
+	c.mu.Lock()
+	c.serverInfo = res
+	c.mu.Unlock()
+	c.logger.Debug("MCP session re-initialized", zap.String("server", c.config.Name))
+	return nil
+}

internal/upstream/core/session_reinit_mcpgo_test.go

@@ -0,0 +1,171 @@
+package core
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+
+	"github.com/smart-mcp-proxy/mcpproxy-go/internal/config"
+	"github.com/smart-mcp-proxy/mcpproxy-go/internal/secret"
+)
+
+// sessionUpstream is a minimal legacy-era Streamable HTTP upstream with real
+// session semantics: initialize issues a new session id; every other request
+// must carry a live one or it is answered 404 before anything executes.
+type sessionUpstream struct {
+	mu    sync.Mutex
+	seq   int
+	live  map[string]bool
+	inits int
+}
+
+func newSessionUpstream() *sessionUpstream { return &sessionUpstream{live: map[string]bool{}} }
+
+// forget makes the upstream drop every session it has issued (server restart).
+func (u *sessionUpstream) forget() {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	u.live = map[string]bool{}
+}
+
+func (u *sessionUpstream) handler() http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			w.WriteHeader(http.StatusMethodNotAllowed)
+			return
+		}
+		body, _ := io.ReadAll(r.Body)
+		var req struct {
+			ID     json.RawMessage `json:"id"`
+			Method string          `json:"method"`
+			Params struct {
+				ProtocolVersion string `json:"protocolVersion"`
+			} `json:"params"`
+		}
+		_ = json.Unmarshal(body, &req)
+		write := func(res map[string]any) {
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": json.RawMessage(req.ID), "result": res})
+		}
+		if req.Method == "initialize" {
+			u.mu.Lock()
+			u.inits++
+			u.seq++
+			sid := fmt.Sprintf("sess-%08d", u.seq)
+			u.live[sid] = true
+			u.mu.Unlock()
+			w.Header().Set("Mcp-Session-Id", sid)
+			write(map[string]any{"protocolVersion": req.Params.ProtocolVersion, "capabilities": map[string]any{"tools": map[string]any{}}, "serverInfo": map[string]any{"name": "sess", "version": "1"}})
+			return
+		}
+		u.mu.Lock()
+		ok := u.live[r.Header.Get("Mcp-Session-Id")]
+		u.mu.Unlock()
+		if !ok {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+		switch req.Method {
+		case "tools/list":
+			write(map[string]any{"tools": []any{}})
+		default:
+			if len(req.ID) == 0 {
+				w.WriteHeader(http.StatusAccepted)
+				return
+			}
+			write(map[string]any{})
+		}
+	})
+}
+
+// T220 / FR-088: pins the mcp-go v1.0.0 behaviour the re-init design relies on.
+func TestMCPGo_SessionTerminated_ReinitializeOnSameTransport(t *testing.T) {
+	up := newSessionUpstream()
+	srv := httptest.NewServer(up.handler())
+	defer srv.Close()
+
+	tr, err := transport.NewStreamableHTTP(srv.URL)
+	require.NoError(t, err)
+	cl := client.NewClient(tr)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	require.NoError(t, cl.Start(ctx))
+	t.Cleanup(func() { _ = cl.Close() })
+
+	initReq := mcp.InitializeRequest{}
+	initReq.Params.ProtocolVersion = mcp.LATEST_LEGACY_PROTOCOL_VERSION
+	initReq.Params.ClientInfo = mcp.Implementation{Name: "t", Version: "1"}
+	_, err = cl.Initialize(ctx, initReq)
+	require.NoError(t, err)
+	first := tr.GetSessionId()
+	require.NotEmpty(t, first)
+
+	// 404 on a non-initialize POST -> ErrSessionTerminated, session id cleared.
+	up.forget()
+	_, err = cl.ListTools(ctx, mcp.ListToolsRequest{})
+	require.Error(t, err)
+	assert.True(t, errors.Is(err, transport.ErrSessionTerminated), "got %v", err)
+	assert.Empty(t, tr.GetSessionId(), "mcp-go clears the session id on 404")
+
+	// 404 with no session id at all is also ErrSessionTerminated (the next
+	// request goes out without a header, which this upstream rejects).
+	_, err = cl.ListTools(ctx, mcp.ListToolsRequest{})
+	require.Error(t, err)
+	assert.True(t, errors.Is(err, transport.ErrSessionTerminated), "got %v", err)
+
+	// A second Initialize on the same client/transport stores the new id.
+	_, err = cl.Initialize(ctx, initReq)
+	require.NoError(t, err)
+	second := tr.GetSessionId()
+	require.NotEmpty(t, second)
+	assert.NotEqual(t, first, second)
+	_, err = cl.ListTools(ctx, mcp.ListToolsRequest{})
+	require.NoError(t, err)
+}
+
+func TestCoreSessionSnapshotAndReinitialize(t *testing.T) {
+	disableOAuthForTest(t)
+	up := newSessionUpstream()
+	srv := httptest.NewServer(up.handler())
+	defer srv.Close()
+
+	cfg := &config.ServerConfig{Name: "sess", Protocol: "streamable-http", URL: srv.URL, Enabled: true}
+	c, err := NewClient("sess", cfg, zap.NewNop(), nil, nil, nil, secret.NewResolver())
+	require.NoError(t, err)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	require.NoError(t, c.Connect(ctx))
+	t.Cleanup(func() { _ = c.Disconnect() })
+
+	id, modern := c.SessionSnapshot()
+	assert.NotEmpty(t, id)
+	assert.False(t, modern)
+
+	up.forget()
+	_, err = c.ListTools(ctx)
+	require.Error(t, err)
+	assert.True(t, errors.Is(err, transport.ErrSessionTerminated), "got %v", err)
+	id2, _ := c.SessionSnapshot()
+	assert.Empty(t, id2)
+
+	require.NoError(t, c.ReinitializeSession(ctx))
+	id3, _ := c.SessionSnapshot()
+	assert.NotEmpty(t, id3)
+	assert.NotEqual(t, id, id3)
+	_, err = c.ListTools(ctx)
+	require.NoError(t, err)
+}

internal/upstream/managed/client.go

@@ -105,6 +105,9 @@ type Client struct {
 	// (hand-constructed clients in tests).
 	toolInvoker toolCaller
 
+	// sess is the Spec 113-e session re-init state (single-flight, baseline, counter).
+	sess sessionReinit
+
 	// ambiguousProbeInFlight gates the async liveness probe fired after an
 	// ambiguous tools/call cancellation (GH #965) so a burst of canceled calls
 	// results in at most one probe against the upstream.
@@ -579,6 +582,10 @@ func (mc *Client) Connect(ctx context.Context) error {
 	mc.connectionEpoch.Store(nextConnectionEpoch())
 	mc.epochMu.Unlock()
 
+	// Spec 113-e FR-080: remember the session id this connect obtained, so a
+	// transport that loses it before the first managed request can recover.
+	_, _, _ = mc.sessionState()
+
 	// Transition to ready state only if not already ready
 	if mc.StateManager.GetState() != types.StateReady {
 		mc.StateManager.TransitionTo(types.StateReady)
@@ -838,6 +845,8 @@ func (mc *Client) GetConnectionStatus() map[string]interface{} {
 		"should_retry": mc.ShouldRetry(),
 		"retry_count":  info.RetryCount,
 		"server_name":  info.ServerName,
+		// Spec 113-e FR-085: in-place Streamable HTTP session re-initializations.
+		"session_reinit_count": mc.SessionReinitCount(),
 	}
 
 	if info.LastError != nil {
@@ -1062,7 +1071,7 @@ func (mc *Client) runListToolsAsLeader(listCtx context.Context, release func() b
 	}()
 
 	listEpoch := mc.connectionEpoch.Load()
-	tools, err := mc.coreClient.ListTools(listCtx)
+	tools, err := mc.listToolsUpstream(listCtx)
 	mc.publishListToolsResult(tools, err)
 
 	if err != nil {
@@ -1201,8 +1210,17 @@ func (mc *Client) callTool(ctx context.Context, toolName string, args map[string
 	// wait (a queued call may outlive a reconnect), so a transport failure is
 	// only ever charged to the session that produced it (RC4-STDIO-001).
 	callEpoch := mc.connectionEpoch.Load()
-	result, err := invoker.CallTool(core.WithConnectionGeneration(ctx, callEpoch), toolName, args)
+	result, err := mc.callToolWithSession(core.WithConnectionGeneration(ctx, callEpoch), invoker, toolName, args, expectedEpoch)
 	if err != nil {
+		if errors.Is(err, ErrSessionReestablished) || errors.Is(err, ErrConnectionGenerationChanged) {
+			// Spec 113-e: the session was re-established and the call was
+			// deliberately not repeated. Not a connection failure.
+			mc.logger.Warn("Tool call not repeated after upstream session re-initialize",
+				zap.String("server", mc.GetConfig().Name),
+				zap.String("tool", toolName),
+				zap.Error(err))
+			return nil, err
+		}
 		mc.recordCallToolOAuthSignal(toolName, err)
 		// A 429 answered to a tools/call is the same instruction as one answered
 		// to connect (#1040). Recording it here does NOT mark the server
@@ -1752,7 +1770,7 @@ func (mc *Client) performHealthCheck() {
 	// the transport (the ping fails with "transport closed") and only then
 	// bumps the epoch, so the verdict below must land on this generation only.
 	pingEpoch := mc.connectionEpoch.Load()
-	err := prober.Ping(ctx)
+	err := mc.probeLiveness(ctx, prober)
 
 	if err != nil {
 		// Pick up any rate-limit hint this ping's response carried, BEFORE the
@@ -2480,7 +2498,7 @@ func (mc *Client) GetCachedToolCount(ctx context.Context) (int, error) {
 	// Fetch fresh tool count with timeout. Publish the result so any concurrent
 	// ListTools waiter coalesced behind us receives the real tools list.
 	countEpoch := mc.connectionEpoch.Load()
-	tools, err := mc.coreClient.ListTools(listCtx)
+	tools, err := mc.listToolsUpstream(listCtx)
 	mc.publishListToolsResult(tools, err)
 	if err != nil {
 		mc.logger.Debug("Tool count fetch failed, returning cached value",

internal/upstream/managed/prompts.go

@@ -14,7 +14,11 @@ func (mc *Client) ListPrompts(ctx context.Context) ([]mcp.Prompt, error) {
 		return nil, fmt.Errorf("client not connected (state: %s)", mc.StateManager.GetState().String())
 	}
 
-	prompts, err := mc.coreClient.ListPrompts(ctx)
+	var prompts []mcp.Prompt
+	err := mc.withSession(ctx, string(mcp.MethodPromptsList), func() (e error) {
+		prompts, e = mc.coreClient.ListPrompts(ctx)
+		return e
+	})
 	if err != nil {
 		mc.logger.Error("Failed to list prompts",
 			zap.String("server", mc.GetConfig().Name),
@@ -33,7 +37,11 @@ func (mc *Client) GetPrompt(ctx context.Context, name string, args map[string]st
 	}
 
 	promptEpoch := mc.connectionEpoch.Load()
-	result, err := mc.coreClient.GetPrompt(ctx, name, args)
+	var result *mcp.GetPromptResult
+	err := mc.withSession(ctx, string(mcp.MethodPromptsGet), func() (e error) {
+		result, e = mc.coreClient.GetPrompt(ctx, name, args)
+		return e
+	})
 	if err != nil {
 		if mc.isConnectionError(err) {
 			// Guarded: concurrent failures on one dead transport mark it