From e08ee25471e9313e62685a9934baeff176184671 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Fri, 17 Jul 2026 16:52:20 +0300 Subject: [PATCH] Add per-session prompt support (SessionWithPrompts) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shim exposed SessionWithTools and SessionWithResources but no per-session prompt equivalent, so a consumer that injects prompts per session (e.g. ToolHive vMCP, which registers aggregated backend prompts on each client session) could list them only globally — prompts/get returned -32602 "unknown prompt" for every session-injected prompt. Add SessionWithPrompts, mirroring SessionWithResources exactly: the interface, the clientSession prompt map guarded by the shared mutex, GetSessionPrompts/SetSessionPrompts (copy-on-read/write), and syncSessionPrompts, which registers each prompt onto the session's own go-sdk Server via srv.AddPrompt. Because each client session has its own go-sdk Server, go-sdk then serves the prompt from both prompts/list and prompts/get. Reconcile is wired at both timing points resources use: the registerAndSync sweep and the live SetSessionPrompts path. Additive only: a new exported interface plus new methods and one field on the unexported clientSession. Servers that never call SetSessionPrompts are byte-identical. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01UthZzKery5m7GBNcXYKALe --- mcpcompat/server/server.go | 26 +++-- mcpcompat/server/server_internal_test.go | 11 ++ mcpcompat/server/server_test.go | 135 +++++++++++++++++++++++ mcpcompat/server/session.go | 83 ++++++++++++-- 4 files changed, 231 insertions(+), 24 deletions(-) diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index 111833a..e398467 100644 --- a/mcpcompat/server/server.go +++ b/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)) diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 390325f..596bde7 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/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 diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go index 97adf87..616b3fc 100644 --- a/mcpcompat/server/server_test.go +++ b/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. diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go index 493554c..61b94f1 100644 --- a/mcpcompat/server/session.go +++ b/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)