fix: resolve dual React instance in JSX renderer and propagate provenance session IDs (#194)

* fix: resolve dual React instance in JSX renderer and propagate provenance session IDs

The JSX sandbox iframe used ?bundle on all esm.sh imports, causing each
package (react, react-dom, react/jsx-runtime) to bundle its own copy of
React internals. When createRoot mounted on one instance but components
used hooks from another, useState saw a null dispatcher — blank page.

Fix: remove ?bundle from react/react-dom (let esm.sh deduplicate via
shared module URLs), use ?bundle&external=react,react-dom on leaf
packages (recharts, lucide-react). Add unhandledrejection handler so
module load errors surface instead of silent blank pages.

Also adds WithReplacedSessionID/ReplacedSessionID to the session handler
for provenance tracking, and harvestProvenance helper to the provenance
middleware for merging both original and replacement session IDs.

* fix: escape </script> in JSX renderer to prevent HTML breakout

Transformed JSX injected into the iframe's <script> block could contain
literal </script> strings, breaking the HTML structure. Add
escapeScriptClose() to replace </script with <\/script (valid JS, does
not terminate the HTML tag). Also switch the error path to use
textContent via JSON.stringify instead of raw interpolation.

Author: cjimti

pkg/middleware/mcp_provenance.go

@@ -8,6 +8,8 @@ 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.
@@ -122,6 +124,17 @@ 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
+}
+
 // MCPProvenanceMiddleware tracks tool calls per session and injects
 // accumulated provenance into the context when save_artifact is called.
 func MCPProvenanceMiddleware(tracker *ProvenanceTracker, saveToolName string) mcp.Middleware {
@@ -145,7 +158,7 @@ func MCPProvenanceMiddleware(tracker *ProvenanceTracker, saveToolName string) mc
 			}
 
 			if toolName == saveToolName {
-				calls := tracker.Harvest(sessionID)
+				calls := harvestProvenance(ctx, tracker, sessionID)
 				ctx = WithProvenanceToolCalls(ctx, calls)
 				return next(ctx, method, req)
 			}

pkg/middleware/mcp_provenance_test.go

@@ -9,6 +9,8 @@ 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.
@@ -238,6 +240,73 @@ 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,6 +15,10 @@ 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 {
@@ -30,6 +34,22 @@ 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"
@@ -142,6 +162,7 @@ func (h *AwareHandler) handleExisting(w http.ResponseWriter, r *http.Request, se
 		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)
 			return
 		}

pkg/session/handler_test.go

@@ -442,17 +442,19 @@ func TestAwareSessionID_Roundtrip(t *testing.T) {
 	assert.Equal(t, "test-session-123", got)
 }
 
-// contextCapturingHandler captures the AwareSessionID from the request context.
+// contextCapturingHandler captures the AwareSessionID and ReplacedSessionID from the request context.
 type contextCapturingHandler struct {
-	mu             sync.Mutex
-	awareSessionID string
-	capturedCalled bool
+	mu                sync.Mutex
+	awareSessionID    string
+	replacedSessionID string
+	capturedCalled    bool
 }
 
 func (h *contextCapturingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	h.mu.Lock()
 	h.capturedCalled = true
 	h.awareSessionID = AwareSessionID(r.Context())
+	h.replacedSessionID = ReplacedSessionID(r.Context())
 	h.mu.Unlock()
 	w.WriteHeader(http.StatusOK)
 }
@@ -531,3 +533,69 @@ func TestHandler_ConcurrentAccess(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestReplacedSessionID_EmptyContext(t *testing.T) {
+	got := ReplacedSessionID(context.Background())
+	assert.Empty(t, got, "plain context should return empty string")
+}
+
+func TestReplacedSessionID_Roundtrip(t *testing.T) {
+	ctx := WithReplacedSessionID(context.Background(), "old-session-abc")
+	got := ReplacedSessionID(ctx)
+	assert.Equal(t, "old-session-abc", got)
+}
+
+func TestHandler_SessionReplacement_SetsReplacedSessionID(t *testing.T) {
+	store := NewMemoryStore(handlerTestTTL)
+	capture := &contextCapturingHandler{}
+	handler := NewAwareHandler(capture, HandlerConfig{
+		Store: store,
+		TTL:   handlerTestTTL,
+	})
+
+	// Create an expired session
+	sess := newTestSession("old-session-for-replace", -time.Second)
+	sess.UserID = hashToken("replace-token")
+	require.NoError(t, store.Create(context.Background(), sess))
+
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodPost, handlerTestPath, http.NoBody)
+	req.Header.Set(sessionIDHeader, "old-session-for-replace")
+	req.Header.Set(handlerTestAuthHeader, "Bearer replace-token")
+	w := httptest.NewRecorder()
+
+	handler.ServeHTTP(w, req)
+
+	capture.mu.Lock()
+	defer capture.mu.Unlock()
+	assert.True(t, capture.capturedCalled, "inner handler should be called")
+	assert.NotEmpty(t, capture.awareSessionID, "new session ID should be set")
+	assert.NotEqual(t, "old-session-for-replace", capture.awareSessionID)
+	assert.Equal(t, "old-session-for-replace", capture.replacedSessionID,
+		"replaced session ID should carry the old session ID")
+}
+
+func TestHandler_NormalSession_NoReplacedSessionID(t *testing.T) {
+	store := NewMemoryStore(handlerTestTTL)
+	capture := &contextCapturingHandler{}
+	handler := NewAwareHandler(capture, HandlerConfig{
+		Store: store,
+		TTL:   handlerTestTTL,
+	})
+
+	// Create a valid (non-expired) session
+	sess := newTestSession("normal-session", handlerTestTTL)
+	sess.UserID = ""
+	require.NoError(t, store.Create(context.Background(), sess))
+
+	req := httptest.NewRequestWithContext(context.Background(), http.MethodPost, handlerTestPath, http.NoBody)
+	req.Header.Set(sessionIDHeader, "normal-session")
+	w := httptest.NewRecorder()
+
+	handler.ServeHTTP(w, req)
+
+	capture.mu.Lock()
+	defer capture.mu.Unlock()
+	assert.True(t, capture.capturedCalled)
+	assert.Empty(t, capture.replacedSessionID,
+		"normal sessions should not have a replaced session ID")
+}