fix: revive expired sessions with same ID to preserve provenance and dedup (#195)

* fix: revive expired sessions with same ID to preserve provenance and dedup

When a client reconnects after session expiry, handleExisting was calling
handleInitialize which generated a new random session ID. But MCP clients
(e.g. Claude Desktop) don't update their stored session ID from response
headers, so every subsequent request triggered another replacement —
breaking provenance tracking (0 tool calls) and enrichment dedup (never
fires).

Replace handleInitialize call with reviveSession that recreates the
session using the client's original ID. Delete-then-create avoids unique
constraint violations from expired-but-not-yet-cleaned rows.

* fix: remove dead ReplacedSessionID code and add upsert for revive race safety

Remove ReplacedSessionID/WithReplacedSessionID infrastructure that became
dead code after switching to same-ID session revival. The replaced-session
context value was never populated in the revive path, making the
harvestProvenance fallback and all related tests unreachable.

Add ON CONFLICT upsert to the Postgres session store's Create method so
concurrent revive attempts on the same expired session ID don't fail with
unique constraint violations.

Author: cjimti

pkg/middleware/mcp_provenance.go

@@ -8,8 +8,6 @@ import (
 	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
-
-	pkgsession "github.com/txn2/mcp-data-platform/pkg/session"
 )
 
 // provenanceContextKey is the context key for provenance tool calls.
@@ -124,15 +122,9 @@ func summarizeParams(params map[string]any) string {
 	return s
 }
 
-// harvestProvenance collects tool calls from the current session and, if a
-// session replacement occurred, also from the old (replaced) session.
-func harvestProvenance(ctx context.Context, tracker *ProvenanceTracker, sessionID string) []ProvenanceToolCall {
-	calls := tracker.Harvest(sessionID)
-	if replacedID := pkgsession.ReplacedSessionID(ctx); replacedID != "" {
-		oldCalls := tracker.Harvest(replacedID)
-		calls = append(oldCalls, calls...)
-	}
-	return calls
+// harvestProvenance collects tool calls from the current session.
+func harvestProvenance(_ context.Context, tracker *ProvenanceTracker, sessionID string) []ProvenanceToolCall {
+	return tracker.Harvest(sessionID)
 }
 
 // MCPProvenanceMiddleware tracks tool calls per session and injects

pkg/middleware/mcp_provenance_test.go

@@ -9,8 +9,6 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-
-	pkgsession "github.com/txn2/mcp-data-platform/pkg/session"
 )
 
 // newTestServerRequest creates a test server request for tools/call.
@@ -240,73 +238,6 @@ func TestProvenanceTracker_CleanupBefore(t *testing.T) {
 	assert.Len(t, calls, 1)
 }
 
-func TestMCPProvenanceMiddleware_SessionReplacement(t *testing.T) {
-	tracker := NewProvenanceTracker()
-
-	// Record tool calls under the "old" session (the one that will expire).
-	tracker.Record("old-session", "trino_query", map[string]any{"sql": "SELECT 1"})
-	tracker.Record("old-session", "datahub_search", map[string]any{"query": "sales"})
-
-	var capturedCalls []ProvenanceToolCall
-	base := func(ctx context.Context, _ string, _ mcp.Request) (mcp.Result, error) {
-		capturedCalls = GetProvenanceToolCalls(ctx)
-		return &mcp.CallToolResult{}, nil
-	}
-
-	handler := MCPProvenanceMiddleware(tracker, "save_artifact")(base)
-
-	req := newTestServerRequest(&mcp.CallToolParamsRaw{
-		Name: "save_artifact",
-	})
-
-	// The new session has no tool calls, but the replaced session ID carries
-	// the old session's provenance.
-	ctx := WithPlatformContext(context.Background(), &PlatformContext{SessionID: "new-session"})
-	ctx = pkgsession.WithReplacedSessionID(ctx, "old-session")
-
-	_, err := handler(ctx, methodToolsCall, req)
-	require.NoError(t, err)
-
-	// Both old-session tool calls should be recovered.
-	require.Len(t, capturedCalls, 2, "should recover provenance from replaced session")
-	assert.Equal(t, "trino_query", capturedCalls[0].ToolName)
-	assert.Equal(t, "datahub_search", capturedCalls[1].ToolName)
-
-	// Both sessions should be cleared after harvest.
-	assert.Nil(t, tracker.Harvest("old-session"))
-	assert.Nil(t, tracker.Harvest("new-session"))
-}
-
-func TestMCPProvenanceMiddleware_SessionReplacementMergesBoth(t *testing.T) {
-	tracker := NewProvenanceTracker()
-
-	// Old session has 2 calls, new session has 1 call.
-	tracker.Record("old-sess", "tool_a", nil)
-	tracker.Record("old-sess", "tool_b", nil)
-	tracker.Record("new-sess", "tool_c", nil)
-
-	var capturedCalls []ProvenanceToolCall
-	base := func(ctx context.Context, _ string, _ mcp.Request) (mcp.Result, error) {
-		capturedCalls = GetProvenanceToolCalls(ctx)
-		return &mcp.CallToolResult{}, nil
-	}
-
-	handler := MCPProvenanceMiddleware(tracker, "save_artifact")(base)
-
-	req := newTestServerRequest(&mcp.CallToolParamsRaw{Name: "save_artifact"})
-	ctx := WithPlatformContext(context.Background(), &PlatformContext{SessionID: "new-sess"})
-	ctx = pkgsession.WithReplacedSessionID(ctx, "old-sess")
-
-	_, err := handler(ctx, methodToolsCall, req)
-	require.NoError(t, err)
-
-	// Old calls come first, then new session calls.
-	require.Len(t, capturedCalls, 3)
-	assert.Equal(t, "tool_a", capturedCalls[0].ToolName)
-	assert.Equal(t, "tool_b", capturedCalls[1].ToolName)
-	assert.Equal(t, "tool_c", capturedCalls[2].ToolName)
-}
-
 func TestExtractToolParams_NilCases(t *testing.T) {
 	// Request with nil arguments.
 	req := newTestServerRequest(&mcp.CallToolParamsRaw{Name: "test"})

pkg/session/handler.go

@@ -15,10 +15,6 @@ import (
 // awareSessionKey is the context key for the AwareHandler session ID.
 type awareSessionKey struct{}
 
-// replacedSessionKey is the context key for the old session ID that was replaced
-// during session recovery (expired session → new session with auth credentials).
-type replacedSessionKey struct{}
-
 // AwareSessionID returns the session ID set by AwareHandler, or "".
 func AwareSessionID(ctx context.Context) string {
 	if id, ok := ctx.Value(awareSessionKey{}).(string); ok {
@@ -34,22 +30,6 @@ func WithAwareSessionID(ctx context.Context, sessionID string) context.Context {
 	return context.WithValue(ctx, awareSessionKey{}, sessionID)
 }
 
-// ReplacedSessionID returns the old session ID that was replaced during session
-// recovery, or "" if no replacement occurred.
-func ReplacedSessionID(ctx context.Context) string {
-	if id, ok := ctx.Value(replacedSessionKey{}).(string); ok {
-		return id
-	}
-	return ""
-}
-
-// WithReplacedSessionID returns a context carrying the old session ID that was
-// replaced. This allows downstream middleware (e.g. provenance) to recover data
-// recorded under the old session.
-func WithReplacedSessionID(ctx context.Context, oldSessionID string) context.Context {
-	return context.WithValue(ctx, replacedSessionKey{}, oldSessionID)
-}
-
 const (
 	// sessionIDHeader is the MCP session header name.
 	sessionIDHeader = "Mcp-Session-Id"
@@ -63,6 +43,9 @@ const (
 	// slogKeyError is the slog attribute key for error values.
 	slogKeyError = "error"
 
+	// httpErrInternal is the response body for HTTP 500 errors.
+	httpErrInternal = "internal server error"
+
 	// touchTimeout is the maximum time for async session touch operations.
 	touchTimeout = 5 * time.Second
 )
@@ -112,7 +95,7 @@ func (h *AwareHandler) handleInitialize(w http.ResponseWriter, r *http.Request)
 	sessionID, err := generateSessionID()
 	if err != nil {
 		slog.Error("session: failed to generate ID", slogKeyError, err)
-		http.Error(w, "internal server error", http.StatusInternalServerError)
+		http.Error(w, httpErrInternal, http.StatusInternalServerError)
 		return
 	}
 
@@ -128,7 +111,7 @@ func (h *AwareHandler) handleInitialize(w http.ResponseWriter, r *http.Request)
 
 	if err := h.store.Create(r.Context(), sess); err != nil {
 		slog.Error("session: failed to create", slogKeyError, err)
-		http.Error(w, "internal server error", http.StatusInternalServerError)
+		http.Error(w, httpErrInternal, http.StatusInternalServerError)
 		return
 	}
 
@@ -151,19 +134,26 @@ func (h *AwareHandler) handleExisting(w http.ResponseWriter, r *http.Request, se
 	sess, err := h.store.Get(r.Context(), sessionID)
 	if err != nil {
 		slog.Error("session: store error", slogKeyError, err)
-		http.Error(w, "internal server error", http.StatusInternalServerError)
+		http.Error(w, httpErrInternal, http.StatusInternalServerError)
 		return
 	}
 	if sess == nil {
 		// Session expired or was cleaned up. If the request carries auth
-		// credentials, create a replacement session transparently instead
-		// of forcing the client to re-initialize (which the Go SDK client
-		// does not do automatically).
+		// credentials, revive the session using the SAME ID. Clients like
+		// Claude Desktop do not update their stored session ID from
+		// response headers, so generating a new ID would cause every
+		// subsequent request to trigger another revive — breaking
+		// provenance tracking and enrichment dedup.
 		if extractToken(r) != "" {
-			slog.Info("session: expired, creating replacement",
-				"old_session_id", sanitizeLogValue(sessionID)) // #nosec G706 -- sessionID sanitized via sanitizeLogValue
-			r = r.WithContext(WithReplacedSessionID(r.Context(), sessionID))
-			h.handleInitialize(w, r)
+			slog.Info("session: reviving expired session",
+				"session_id", sanitizeLogValue(sessionID)) // #nosec G706 -- sessionID sanitized via sanitizeLogValue
+			if err := h.reviveSession(r.Context(), sessionID, r); err != nil {
+				slog.Error("session: failed to revive", slogKeyError, err)
+				http.Error(w, httpErrInternal, http.StatusInternalServerError)
+				return
+			}
+			r = r.WithContext(WithAwareSessionID(r.Context(), sessionID))
+			h.inner.ServeHTTP(w, r)
 			return
 		}
 		http.Error(w, "session not found or expired", http.StatusNotFound)
@@ -200,6 +190,28 @@ func (h *AwareHandler) handleDelete(w http.ResponseWriter, r *http.Request) {
 	h.inner.ServeHTTP(w, r)
 }
 
+// reviveSession recreates an expired or missing session using the same ID.
+// This ensures session stability when clients don't update their Mcp-Session-Id
+// from response headers (e.g. Claude Desktop). The expired row (if any) is
+// deleted first to avoid unique constraint violations in the database store.
+func (h *AwareHandler) reviveSession(ctx context.Context, sessionID string, r *http.Request) error {
+	// Remove expired-but-not-yet-cleaned row (if any) to avoid INSERT conflict.
+	_ = h.store.Delete(ctx, sessionID)
+
+	now := time.Now()
+	if err := h.store.Create(ctx, &Session{
+		ID:           sessionID,
+		UserID:       hashToken(extractToken(r)),
+		CreatedAt:    now,
+		LastActiveAt: now,
+		ExpiresAt:    now.Add(h.ttl),
+		State:        make(map[string]any),
+	}); err != nil {
+		return fmt.Errorf("reviving session: %w", err)
+	}
+	return nil
+}
+
 // validateOwnership checks that the request token matches the session owner.
 // Anonymous sessions (empty UserID) skip this check.
 func validateOwnership(sess *Session, r *http.Request) bool {