fix: portal JSX auto-mount and provenance session ID propagation (#193)

* fix: portal JSX auto-mount and provenance session ID propagation

JSX Renderer: agent-generated components that lack createRoot/ReactDOM.render
now auto-mount via blob module import with resolved bare specifiers. Self-mounting
content uses the existing import-map path. Error handling added for both paths.

Provenance: AwareHandler now propagates session ID through the Go request context
(WithAwareSessionID/AwareSessionID). MCPToolCallMiddleware falls back to this
context value when the MCP SDK returns the default "stdio" session ID, ensuring
provenance is correctly keyed per client session in Streamable HTTP mode.

* fix: harden JSX renderer, provenance logging, and session ID test coverage

JsxRenderer: fix module blob URL memory leak via useRef lifecycle tracking,
deduplicate import maps (single source of truth from BARE_IMPORT_MAP),
use textContent instead of innerHTML for error display, improve hasMountCode
with regex to reduce false positives, add class declarations to ensureExport.

Provenance: lower Harvest log from Info to Debug to reduce production noise.

Tests: remove misleading "keeps SDK session ID" subtest (was duplicate of
fallback test), add integration test wiring AwareHandler → Stateless
StreamableHTTP → MCPToolCallMiddleware → MCPProvenanceMiddleware to verify
end-to-end session ID propagation.

Author: cjimti

pkg/middleware/mcp.go

@@ -13,6 +13,8 @@ import (
 
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
+
+	pkgsession "github.com/txn2/mcp-data-platform/pkg/session"
 )
 
 const (
@@ -78,6 +80,13 @@ func MCPToolCallMiddleware(authenticator Authenticator, authorizer Authorizer, t
 			pc := NewPlatformContext(generateRequestID())
 			pc.ToolName = toolName
 			pc.SessionID = extractSessionID(req)
+			// Fall back to AwareHandler-managed session ID when the MCP SDK
+			// doesn't propagate one (SSE returns "", stateless mode may vary).
+			if pc.SessionID == defaultSessionID {
+				if awareID := pkgsession.AwareSessionID(ctx); awareID != "" {
+					pc.SessionID = awareID
+				}
+			}
 			pc.Transport = transport
 			pc.Source = "mcp"
 			ctx = buildToolCallContext(ctx, req, pc, toolkitLookup, toolName)

pkg/middleware/mcp_provenance.go

@@ -3,6 +3,7 @@ package middleware
 import (
 	"context"
 	"encoding/json"
+	"log/slog"
 	"sync"
 	"time"
 
@@ -53,8 +54,10 @@ func NewProvenanceTracker() *ProvenanceTracker {
 // Each session is capped at maxCallsPerSession entries (oldest are evicted).
 func (pt *ProvenanceTracker) Record(sessionID, toolName string, params map[string]any) {
 	if sessionID == "" {
+		slog.Debug("provenance: skipping record for empty session ID", "tool", toolName)
 		return
 	}
+	slog.Debug("provenance.record", "session_id", sessionID, "tool", toolName)
 
 	pt.mu.Lock()
 	defer pt.mu.Unlock()
@@ -80,6 +83,7 @@ func (pt *ProvenanceTracker) Harvest(sessionID string) []ProvenanceToolCall {
 
 	calls := pt.sessions[sessionID]
 	delete(pt.sessions, sessionID)
+	slog.Debug("provenance.harvest", "session_id", sessionID, "count", len(calls))
 	return calls
 }
 
@@ -136,6 +140,8 @@ func MCPProvenanceMiddleware(tracker *ProvenanceTracker, saveToolName string) mc
 			sessionID := ""
 			if pc != nil {
 				sessionID = pc.SessionID
+			} else {
+				slog.Warn("provenance: PlatformContext missing, cannot track tool call", "tool", toolName)
 			}
 
 			if toolName == saveToolName {

pkg/middleware/mcp_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 
 	"github.com/txn2/mcp-data-platform/pkg/registry"
+	pkgsession "github.com/txn2/mcp-data-platform/pkg/session"
 )
 
 // Test constants to avoid repeated string literals.
@@ -969,3 +970,73 @@ func TestMCPToolCallMiddleware_ConnectionOverride(t *testing.T) {
 		}
 	})
 }
+
+func TestMCPToolCallMiddleware_AwareSessionIDFallback(t *testing.T) {
+	authenticator := &mcpTestAuthenticator{
+		userInfo: &UserInfo{
+			UserID: mcpTestUserID,
+			Roles:  []string{mcpTestPersona},
+		},
+	}
+	authorizer := &mcpTestAuthorizer{authorized: true, personaName: mcpTestPersona}
+
+	mw := MCPToolCallMiddleware(authenticator, authorizer, nil, "http")
+
+	t.Run("uses AwareHandler session ID when SDK returns default", func(t *testing.T) {
+		next := func(ctx context.Context, _ string, _ mcp.Request) (mcp.Result, error) {
+			pc := GetPlatformContext(ctx)
+			if pc == nil {
+				t.Fatal(mcpTestPCExpected)
+			}
+			if pc.SessionID != "aware-session-abc" {
+				t.Errorf("expected SessionID 'aware-session-abc', got %q", pc.SessionID)
+			}
+			return &mcp.CallToolResult{
+				Content: []mcp.Content{&mcp.TextContent{Text: "ok"}},
+			}, nil
+		}
+
+		handler := mw(next)
+		req := newMCPTestRequest(mcpTestToolName)
+
+		// Set AwareHandler session ID in context (simulates what AwareHandler does).
+		ctx := pkgsession.WithAwareSessionID(context.Background(), "aware-session-abc")
+
+		_, err := handler(ctx, mcpTestMethod, req)
+		if err != nil {
+			t.Fatalf(mcpTestErrFmt, err)
+		}
+	})
+
+	// NOTE: A test for "SDK session ID takes priority over AwareHandler session
+	// ID" is not possible in a unit test because mcp.Session has unexported
+	// methods (sendingMethodInfos, receivingMethodInfos, etc.) that prevent
+	// external mocking. Constructing a *mcp.ServerSession with a real session
+	// ID requires internal SDK types (mcpConn implementing hasSessionID).
+	// The AwareHandler → middleware integration test in middleware_chain_test.go
+	// covers this path through a real Streamable HTTP transport.
+
+	t.Run("falls back to default when no AwareHandler session", func(t *testing.T) {
+		next := func(ctx context.Context, _ string, _ mcp.Request) (mcp.Result, error) {
+			pc := GetPlatformContext(ctx)
+			if pc == nil {
+				t.Fatal(mcpTestPCExpected)
+			}
+			if pc.SessionID != defaultSessionID {
+				t.Errorf("expected SessionID %q, got %q", defaultSessionID, pc.SessionID)
+			}
+			return &mcp.CallToolResult{
+				Content: []mcp.Content{&mcp.TextContent{Text: "ok"}},
+			}, nil
+		}
+
+		handler := mw(next)
+		req := newMCPTestRequest(mcpTestToolName)
+
+		// No AwareHandler session in context
+		_, err := handler(context.Background(), mcpTestMethod, req)
+		if err != nil {
+			t.Fatalf(mcpTestErrFmt, err)
+		}
+	})
+}

pkg/middleware/middleware_chain_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/txn2/mcp-data-platform/pkg/query"
 	"github.com/txn2/mcp-data-platform/pkg/registry"
 	"github.com/txn2/mcp-data-platform/pkg/semantic"
+	pkgsession "github.com/txn2/mcp-data-platform/pkg/session"
 	"github.com/txn2/mcp-data-platform/pkg/storage"
 	"github.com/txn2/mcp-data-platform/pkg/tuning"
 )
@@ -2188,5 +2189,134 @@ func TestWorkflowGating_BackwardCompat(t *testing.T) {
 	assertContentContainsText(t, result, "datahub_search")
 }
 
+// TestMiddlewareChain_AwareHandler_ProvenanceSessionID verifies that session
+// IDs propagated by AwareHandler through the Go request context reach
+// MCPToolCallMiddleware and MCPProvenanceMiddleware, so that provenance is
+// correctly keyed and harvested per-client in stateless Streamable HTTP mode.
+//
+// This is the integration test for the context propagation chain:
+//
+//	AwareHandler.handleInitialize → WithAwareSessionID(ctx)
+//	  → MCPToolCallMiddleware reads AwareSessionID(ctx) as fallback
+//	    → sets PlatformContext.SessionID
+//	      → MCPProvenanceMiddleware uses PlatformContext.SessionID for Record/Harvest
+func TestMiddlewareChain_AwareHandler_ProvenanceSessionID(t *testing.T) {
+	const (
+		queryToolName = "data_query"
+		saveToolName  = "save_artifact"
+	)
+
+	tracker := middleware.NewProvenanceTracker()
+
+	authenticator := &testAuthenticator{
+		userInfo: &middleware.UserInfo{
+			UserID: "prov-test-user",
+			Roles:  []string{chainTestAnalyst},
+		},
+	}
+	authorizer := &testAuthorizer{persona: chainTestAnalyst}
+
+	server := mcp.NewServer(&mcp.Implementation{
+		Name:    "test-aware-provenance",
+		Version: "v0.0.1",
+	}, nil)
+
+	// data_query: a normal tool whose calls should be recorded in provenance.
+	server.AddTool(&mcp.Tool{
+		Name:        queryToolName,
+		Description: "Run a query",
+		InputSchema: json.RawMessage(`{"type":"object","properties":{"sql":{"type":"string"}}}`),
+	}, func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{&mcp.TextContent{Text: "query result: 42"}},
+		}, nil
+	})
+
+	// save_artifact: reads provenance from context (injected by MCPProvenanceMiddleware)
+	// and returns the count so the test can assert it.
+	server.AddTool(&mcp.Tool{
+		Name:        saveToolName,
+		Description: "Save an artifact",
+		InputSchema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"}}}`),
+	}, func(ctx context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		calls := middleware.GetProvenanceToolCalls(ctx)
+		resp := fmt.Sprintf(`{"provenance_count":%d}`, len(calls))
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{&mcp.TextContent{Text: resp}},
+		}, nil
+	})
+
+	// Middleware order (innermost first, outermost last):
+	// 1. Provenance (innermost) — records tool calls, harvests on save_artifact
+	// 2. Auth (outermost) — creates PlatformContext with session ID
+	server.AddReceivingMiddleware(middleware.MCPProvenanceMiddleware(tracker, saveToolName))
+	server.AddReceivingMiddleware(middleware.MCPToolCallMiddleware(authenticator, authorizer, nil, "http"))
+
+	// Stateless Streamable HTTP handler — no SDK-managed sessions.
+	streamableHandler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
+		return server
+	}, &mcp.StreamableHTTPOptions{
+		Stateless: true,
+	})
+
+	// Wrap with AwareHandler to provide session ID propagation via context.
+	awareHandler := pkgsession.NewAwareHandler(streamableHandler, pkgsession.HandlerConfig{
+		Store: pkgsession.NewMemoryStore(10 * time.Minute),
+		TTL:   10 * time.Minute,
+	})
+
+	ts := httptest.NewServer(awareHandler)
+	defer ts.Close()
+
+	client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, nil)
+	ctx := context.Background()
+	clientSession, err := client.Connect(ctx, &mcp.StreamableClientTransport{Endpoint: ts.URL}, nil)
+	if err != nil {
+		t.Fatalf("connecting client: %v", err)
+	}
+	defer func() { _ = clientSession.Close() }()
+
+	// Call data_query — this should be recorded in provenance under the
+	// AwareHandler session ID (not the default "stdio").
+	_, err = clientSession.CallTool(ctx, &mcp.CallToolParams{
+		Name:      queryToolName,
+		Arguments: map[string]any{"sql": "SELECT 1"},
+	})
+	if err != nil {
+		t.Fatalf("calling data_query: %v", err)
+	}
+
+	// Call save_artifact — MCPProvenanceMiddleware harvests the provenance
+	// for this session and injects it into the context. The tool handler
+	// returns the count so we can verify the chain worked end-to-end.
+	result, err := clientSession.CallTool(ctx, &mcp.CallToolParams{
+		Name:      saveToolName,
+		Arguments: map[string]any{"name": "test-artifact"},
+	})
+	if err != nil {
+		t.Fatalf("calling save_artifact: %v", err)
+	}
+	if result.IsError {
+		t.Fatalf("save_artifact returned error: %v", result.Content)
+	}
+
+	// Verify provenance was captured (count > 0 means session IDs matched
+	// between Record and Harvest through the full AwareHandler → middleware chain).
+	found := false
+	for _, c := range result.Content {
+		if tc, ok := c.(*mcp.TextContent); ok {
+			if strings.Contains(tc.Text, `"provenance_count":1`) {
+				found = true
+				break
+			}
+			// Dump content for debugging if assertion fails
+			t.Logf("save_artifact content: %s", tc.Text)
+		}
+	}
+	if !found {
+		t.Fatal("expected provenance_count:1 — AwareHandler session ID did not propagate through middleware chain")
+	}
+}
+
 // Suppress unused import warnings for storage (used in EnrichmentConfig).
 var _ storage.Provider = nil

pkg/session/handler.go

@@ -12,6 +12,24 @@ import (
 	"time"
 )
 
+// awareSessionKey is the context key for the AwareHandler session ID.
+type awareSessionKey struct{}
+
+// AwareSessionID returns the session ID set by AwareHandler, or "".
+func AwareSessionID(ctx context.Context) string {
+	if id, ok := ctx.Value(awareSessionKey{}).(string); ok {
+		return id
+	}
+	return ""
+}
+
+// WithAwareSessionID returns a context carrying the given session ID.
+// This is used by AwareHandler internally and exposed for middleware that
+// needs to read the session ID via AwareSessionID.
+func WithAwareSessionID(ctx context.Context, sessionID string) context.Context {
+	return context.WithValue(ctx, awareSessionKey{}, sessionID)
+}
+
 const (
 	// sessionIDHeader is the MCP session header name.
 	sessionIDHeader = "Mcp-Session-Id"
@@ -101,6 +119,7 @@ func (h *AwareHandler) handleInitialize(w http.ResponseWriter, r *http.Request)
 		sessionID:      sessionID,
 	}
 	r.Header.Set(sessionIDHeader, sessionID)
+	r = r.WithContext(WithAwareSessionID(r.Context(), sessionID))
 	h.inner.ServeHTTP(sw, r)
 }
 
@@ -145,6 +164,7 @@ func (h *AwareHandler) handleExisting(w http.ResponseWriter, r *http.Request, se
 		}
 	}()
 
+	r = r.WithContext(WithAwareSessionID(r.Context(), sessionID))
 	h.inner.ServeHTTP(w, r)
 }