Add per-session prompt support (SessionWithPrompts)

mcpcompat/server/server.go

@@ -16,12 +16,13 @@
 //
 // The global registration path (AddTool/AddResource/AddPrompt served over the
 // stdio and HTTP transports) is fully functional and tested. The per-session
-// interfaces (SessionWithTools, SessionWithResources, SessionWithElicitation,
-// SessionIdManager) and the Hooks type are implemented and wired: per-session
-// tool/resource overlays set via SetSessionTools/SetSessionResources are
-// reconciled onto the session's live go-sdk server (syncSessionTools/
-// syncSessionResources), and the before-list/before-call hooks fire ahead of
-// SDK dispatch so ToolHive's lazy per-session tool injection runs first.
+// interfaces (SessionWithTools, SessionWithResources, SessionWithPrompts,
+// SessionWithElicitation, SessionIdManager) and the Hooks type are implemented
+// and wired: per-session tool/resource/prompt overlays set via SetSessionTools/
+// SetSessionResources/SetSessionPrompts are reconciled onto the session's live
+// go-sdk server (syncSessionTools/syncSessionResources/syncSessionPrompts), and
+// the before-list/before-call hooks fire ahead of SDK dispatch so ToolHive's
+// lazy per-session tool injection runs first.
 // Cross-replica session rehydration (Validate-driven lazy eviction) and the
 // Streamable HTTP transports are functional. See the notes on SessionWithTools
 // for the live-overlay reconciliation details.
@@ -403,12 +404,13 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
 // buildServer constructs a go-sdk Server from the globally-registered features
 // (AddTool/AddResource/AddPrompt).
 //
-// Per-session overlays (SessionWithTools/SessionWithResources) are NOT baked in
-// here: the streamable/SSE transports call this once per new client session (via
-// getServer) so each session gets its own go-sdk Server, and the registration
-// middleware installed by this function syncs that session's overlay tools and
-// resources onto its own server once the OnRegisterSession hooks have run. This
-// mirrors mcp-go, whose per-session tools were dispatched per connection.
+// Per-session overlays (SessionWithTools/SessionWithResources/
+// SessionWithPrompts) are NOT baked in here: the streamable/SSE transports call
+// this once per new client session (via getServer) so each session gets its own
+// go-sdk Server, and the registration middleware installed by this function
+// syncs that session's overlay tools, resources and prompts onto its own server
+// once the OnRegisterSession hooks have run. This mirrors mcp-go, whose
+// per-session tools were dispatched per connection.
 func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, error) {
 	s.mu.RLock()
 	tools := make(map[string]ServerTool, len(s.tools))

mcpcompat/server/server_internal_test.go

@@ -20,6 +20,7 @@ var (
 	_ ClientSession          = (*clientSession)(nil)
 	_ SessionWithTools       = (*clientSession)(nil)
 	_ SessionWithResources   = (*clientSession)(nil)
+	_ SessionWithPrompts     = (*clientSession)(nil)
 	_ SessionWithElicitation = (*clientSession)(nil)
 )
 
@@ -47,6 +48,16 @@ func TestClientSession_Store(t *testing.T) {
 		"file:///r": {Resource: mcp.Resource{URI: "file:///r"}},
 	})
 	assert.Contains(t, cs.GetSessionResources(), "file:///r")
+
+	cs.SetSessionPrompts(map[string]ServerPrompt{
+		"p": {Prompt: mcp.Prompt{Name: "p"}},
+	})
+	gotPrompts := cs.GetSessionPrompts()
+	require.Contains(t, gotPrompts, "p")
+
+	// GetSessionPrompts must return a copy (mutating it must not affect the store).
+	gotPrompts["p2"] = ServerPrompt{}
+	assert.NotContains(t, cs.GetSessionPrompts(), "p2")
 }
 
 // TestForgetSession_ClosesNotifChannel verifies that forgetSession closes the

mcpcompat/server/server_test.go

@@ -12,6 +12,7 @@ import (
 	"net/http/httptest"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -916,6 +917,140 @@ func TestCallUnknownTool_ErrorContainsToolName(t *testing.T) {
 		"unknown-tool error must not carry the empty tool name")
 }
 
+// TestSessionPrompts_ListAndGet_EndToEnd is the regression anchor for
+// per-session prompts (SessionWithPrompts): a prompt injected onto a session via
+// the OnRegisterSession hook's SetSessionPrompts must be served by BOTH
+// prompts/list AND prompts/get. Before SessionWithPrompts existed, prompts/get
+// returned -32602 (InvalidParams) for a session-injected prompt because nothing
+// registered it onto the session's go-sdk server. This drives the whole
+// server->go-sdk->client path over Streamable HTTP and asserts the prompt is
+// listed, gettable (returning the handler's distinctive result), that a global
+// prompt is merged alongside it, and that a SECOND session whose hook injects
+// nothing does NOT see the per-session prompt (session isolation).
+func TestSessionPrompts_ListAndGet_EndToEnd(t *testing.T) {
+	t.Parallel()
+	ctx := t.Context()
+
+	const (
+		sessionPromptName = "session-prompt"
+		globalPromptName  = "global-prompt"
+		promptReply       = "session prompt reply"
+	)
+
+	// injectPrompts controls whether the register hook installs the per-session
+	// prompt; the second client below flips it off to assert isolation.
+	var injectPrompts atomic.Bool
+	injectPrompts.Store(true)
+
+	hooks := &server.Hooks{}
+	hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) {
+		if !injectPrompts.Load() {
+			return
+		}
+		swp, ok := s.(server.SessionWithPrompts)
+		require.True(t, ok, "session must implement SessionWithPrompts")
+		swp.SetSessionPrompts(map[string]server.ServerPrompt{
+			sessionPromptName: {
+				Prompt: mcp.Prompt{Name: sessionPromptName, Description: "per-session prompt"},
+				Handler: func(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
+					return &mcp.GetPromptResult{
+						Description: "per-session prompt",
+						Messages: []mcp.PromptMessage{
+							mcp.NewPromptMessage(mcp.RoleUser, mcp.NewTextContent(promptReply)),
+						},
+					}, nil
+				},
+			},
+		})
+	})
+
+	srv := server.NewMCPServer("prompt-server", testClientVersion,
+		server.WithPromptCapabilities(true),
+		server.WithHooks(hooks),
+	)
+	// A cheap global prompt to assert list returns global+session merged.
+	srv.AddPrompt(mcp.Prompt{Name: globalPromptName, Description: "global prompt"},
+		func(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
+			return &mcp.GetPromptResult{
+				Messages: []mcp.PromptMessage{
+					mcp.NewPromptMessage(mcp.RoleUser, mcp.NewTextContent("global reply")),
+				},
+			}, nil
+		})
+
+	httpSrv := server.NewStreamableHTTPServer(srv)
+	ts := httptest.NewServer(httpSrv)
+	t.Cleanup(ts.Close)
+
+	c, err := client.NewStreamableHttpClient(ts.URL)
+	require.NoError(t, err)
+	require.NoError(t, c.Start(ctx))
+	t.Cleanup(func() { _ = c.Close() })
+
+	initRes, err := c.Initialize(ctx, mcp.InitializeRequest{
+		Params: mcp.InitializeParams{
+			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
+			ClientInfo:      mcp.Implementation{Name: testClientName, Version: testClientVersion},
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, initRes.Capabilities.Prompts, "prompts capability must be advertised")
+
+	// prompts/list must return both the global and the session-injected prompt.
+	list, err := c.ListPrompts(ctx, mcp.ListPromptsRequest{})
+	require.NoError(t, err)
+	names := make([]string, 0, len(list.Prompts))
+	for _, p := range list.Prompts {
+		names = append(names, p.Name)
+	}
+	assert.Contains(t, names, sessionPromptName, "prompts/list must include the per-session prompt")
+	assert.Contains(t, names, globalPromptName, "prompts/list must include the global prompt")
+
+	// prompts/get on the session-injected prompt must succeed (NOT -32602) and
+	// return the handler's distinctive result.
+	got, err := c.GetPrompt(ctx, mcp.GetPromptRequest{
+		Params: mcp.GetPromptParams{Name: sessionPromptName},
+	})
+	require.NoError(t, err, "prompts/get on a session-injected prompt must succeed (regression: was -32602)")
+	require.Len(t, got.Messages, 1)
+	txt, ok := mcp.AsTextContent(got.Messages[0].Content)
+	require.True(t, ok)
+	assert.Equal(t, promptReply, txt.Text, "prompts/get must return the session prompt handler's result")
+
+	// Session isolation: a SECOND client whose hook injects nothing must not see
+	// the per-session prompt.
+	injectPrompts.Store(false)
+	c2, err := client.NewStreamableHttpClient(ts.URL)
+	require.NoError(t, err)
+	require.NoError(t, c2.Start(ctx))
+	t.Cleanup(func() { _ = c2.Close() })
+
+	_, err = c2.Initialize(ctx, mcp.InitializeRequest{
+		Params: mcp.InitializeParams{
+			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
+			ClientInfo:      mcp.Implementation{Name: testClientName, Version: testClientVersion},
+		},
+	})
+	require.NoError(t, err)
+
+	list2, err := c2.ListPrompts(ctx, mcp.ListPromptsRequest{})
+	require.NoError(t, err)
+	names2 := make([]string, 0, len(list2.Prompts))
+	for _, p := range list2.Prompts {
+		names2 = append(names2, p.Name)
+	}
+	assert.NotContains(t, names2, sessionPromptName,
+		"a session without an injected prompt must not see another session's prompt")
+	assert.Contains(t, names2, globalPromptName, "the global prompt must still be visible to the second session")
+
+	// prompts/get for the un-injected per-session prompt must fail on the second
+	// session (the regression's original -32602 behavior is correct HERE).
+	_, err = c2.GetPrompt(ctx, mcp.GetPromptRequest{
+		Params: mcp.GetPromptParams{Name: sessionPromptName},
+	})
+	require.Error(t, err, "prompts/get for a prompt not injected on this session must fail")
+}
+
 // TestRequestElicitation_NoActiveSession is a fast unit-level test (no HTTP
 // server) exercising the ErrNoActiveSession guard: calling RequestElicitation
 // outside any session's request context returns ErrNoActiveSession.

mcpcompat/server/session.go

@@ -53,6 +53,15 @@ type SessionWithResources interface {
 	SetSessionResources(resources map[string]ServerResource)
 }
 
+// SessionWithPrompts is a ClientSession that carries per-session prompts.
+type SessionWithPrompts interface {
+	ClientSession
+	// GetSessionPrompts returns the session's prompts. Thread-safe.
+	GetSessionPrompts() map[string]ServerPrompt
+	// SetSessionPrompts sets the session's prompts. Thread-safe.
+	SetSessionPrompts(prompts map[string]ServerPrompt)
+}
+
 // SessionIdManager governs MCP session ID lifecycle. It mirrors mcp-go's
 // server.SessionIdManager so ToolHive's implementation can be supplied via
 // WithSessionIdManager.
@@ -81,20 +90,22 @@ type clientSession struct {
 	goSession  atomic.Pointer[gosdk.ServerSession]
 
 	// owner and boundServer are set when the session's go-sdk server is bound
-	// (at registration). They let SetSessionTools/SetSessionResources reconcile
-	// the per-session overlay onto the live go-sdk server at runtime.
+	// (at registration). They let SetSessionTools/SetSessionResources/
+	// SetSessionPrompts reconcile the per-session overlay onto the live go-sdk
+	// server at runtime.
 	//
 	// owner is an atomic.Pointer (not a plain field) because SetSessionTools/
-	// SetSessionResources may run from any goroutine (test code, before-hooks on
-	// a live request) concurrently with registerAndSync/bindRehydratedSession,
-	// which set owner on the session's connection goroutine. boundServer uses the
-	// same atomic pattern for the same reason.
+	// SetSessionResources/SetSessionPrompts may run from any goroutine (test
+	// code, before-hooks on a live request) concurrently with registerAndSync/
+	// bindRehydratedSession, which set owner on the session's connection
+	// goroutine. boundServer uses the same atomic pattern for the same reason.
 	owner       atomic.Pointer[MCPServer]
 	boundServer atomic.Pointer[gosdk.Server]
 
 	mu        sync.RWMutex
 	tools     map[string]ServerTool
 	resources map[string]ServerResource
+	prompts   map[string]ServerPrompt
 	// sdkToolNames tracks the tool names this session has added to its go-sdk
 	// server, so a later SetSessionTools can remove the ones that went away.
 	sdkToolNames map[string]struct{}
@@ -171,6 +182,30 @@ func (c *clientSession) SetSessionResources(resources map[string]ServerResource)
 	}
 }
 
+func (c *clientSession) GetSessionPrompts() map[string]ServerPrompt {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	out := make(map[string]ServerPrompt, len(c.prompts))
+	for k, v := range c.prompts {
+		out[k] = v
+	}
+	return out
+}
+
+func (c *clientSession) SetSessionPrompts(prompts map[string]ServerPrompt) {
+	c.mu.Lock()
+	c.prompts = make(map[string]ServerPrompt, len(prompts))
+	for k, v := range prompts {
+		c.prompts[k] = v
+	}
+	c.mu.Unlock()
+	if srv := c.boundServer.Load(); srv != nil {
+		if owner := c.owner.Load(); owner != nil {
+			owner.syncSessionPrompts(srv, c)
+		}
+	}
+}
+
 // sessionContextKey is the context key under which the ClientSession is stored.
 type sessionContextKey struct{}
 
@@ -197,7 +232,7 @@ func (s *MCPServer) sessionFor(id string) *clientSession {
 
 // registerAndSync registers the session for the given go-sdk ServerSession,
 // firing the OnRegisterSession hooks exactly once and then reconciling any
-// per-session tool/resource overlay the hooks installed onto srv (the go-sdk
+// per-session tool/resource/prompt overlay the hooks installed onto srv (the go-sdk
 // server bound to this session). It is invoked from the initialize dispatch
 // middleware (matching mcp-go's on-initialize timing) and, defensively, from the
 // InitializedHandler; the once-guard makes the second call a cheap no-op.
@@ -220,12 +255,13 @@ func (s *MCPServer) registerAndSync(ctx context.Context, ss *gosdk.ServerSession
 	if s.hooks != nil {
 		s.hooks.registerSession(ctx, cs)
 	}
-	// The hooks may have installed per-session tools/resources via
-	// SetSessionTools/SetSessionResources; those calls reconcile onto srv
-	// themselves now that boundServer is set. Sync once more here to cover any
-	// overlay set before the server was bound.
+	// The hooks may have installed per-session tools/resources/prompts via
+	// SetSessionTools/SetSessionResources/SetSessionPrompts; those calls
+	// reconcile onto srv themselves now that boundServer is set. Sync once more
+	// here to cover any overlay set before the server was bound.
 	s.syncSessionTools(srv, cs)
 	s.syncSessionResources(srv, cs)
+	s.syncSessionPrompts(srv, cs)
 }
 
 // syncSessionTools reconciles the session's tool overlay onto its go-sdk server:
@@ -308,6 +344,28 @@ func (s *MCPServer) syncSessionResources(srv *gosdk.Server, cs *clientSession) {
 	}
 }
 
+// syncSessionPrompts adds the session's prompt overlay onto its go-sdk server.
+// go-sdk serves both prompts/list and prompts/get from a prompt registered via
+// AddPrompt, so a per-session prompt injected here is enumerable AND gettable
+// without any shim-side dispatch. Prompts are add-only here (ToolHive sets them
+// once at registration). Unlike AddTool, go-sdk's AddPrompt performs no
+// schema/URI validation and cannot panic, so no recover wrapper is needed
+// (mirroring syncSessionResources rather than the tool path).
+func (s *MCPServer) syncSessionPrompts(srv *gosdk.Server, cs *clientSession) {
+	cs.mu.RLock()
+	defer cs.mu.RUnlock()
+	for name, sp := range cs.prompts {
+		gp := &gosdk.Prompt{}
+		if err := jsonConvert(sp.Prompt, gp); err != nil {
+			if s.logger != nil {
+				s.logger.Warn("skipping per-session prompt: conversion failed", "prompt", name, "error", err)
+			}
+			continue
+		}
+		srv.AddPrompt(gp, s.wrapPromptHandler(sp.Handler))
+	}
+}
+
 // isLocalSession reports whether the session ID was initialized on this server
 // instance (see MCPServer.localSessions).
 func (s *MCPServer) isLocalSession(id string) bool {
@@ -337,7 +395,8 @@ func (s *MCPServer) forgetSession(id string) {
 // two-phase creation), and cross-replica capability projection is driven by the
 // before-list/before-call hooks (ToolHive's lazy per-session tool injection),
 // not by OnRegisterSession. Binding owner+boundServer here lets those hooks'
-// SetSessionTools/SetSessionResources reconcile the overlay onto srv.
+// SetSessionTools/SetSessionResources/SetSessionPrompts reconcile the overlay
+// onto srv.
 func (s *MCPServer) bindRehydratedSession(id string, ss *gosdk.ServerSession, srv *gosdk.Server) {
 	cs := s.sessionFor(id)
 	cs.goSession.Store(ss)