feat(fxmcpserver): Added MCP SSE server context hooks (ankorstore#347)

fxmcpserver/README.md

@@ -20,6 +20,7 @@
     * [Resource templates](#resource-templates)
     * [Prompts](#prompts)
     * [Tools](#tools)
+  * [Hooks](#hooks)
   * [Testing](#testing)
 <!-- TOC -->
 
@@ -37,6 +38,7 @@ This module provides an [MCP server](https://modelcontextprotocol.io/introductio
 - automatic requests logging and tracing (method, target, duration, ...)
 - automatic requests metrics (count and duration)
 - possibility to register MCP resources, resource templates, prompts and tools
+- possibility to register MCP SSE server context hooks
 - possibility to expose the MCP server via Stdio (local) and/or HTTP SSE (remote)
 
 ## Documentation
@@ -549,6 +551,63 @@ modules:
       capabilities:
         tools: true # to expose MCP tools (disabled by default)
 ```
+### Hooks
+
+This module offers the possibility to provide context hooks with [MCPSSEServerContextHook](server/sse/context.go) implementations, applied on each MCP SSE request.
+
+You can use the `AsMCPSSEServerMiddleware()` function to register an MCP SSE server middleware, or `AsMCPSSEServerMiddlewares()` to register several MCP SSE server middlewares at once.
+
+The dependencies of your MCP SSE server middlewares will be autowired.
+
+```go
+package main
+
+import (
+  "context"
+  "net/http"
+
+  "github.com/ankorstore/yokai/config"
+  "github.com/ankorstore/yokai/fxconfig"
+  "github.com/ankorstore/yokai/fxgenerate"
+  "github.com/ankorstore/yokai/fxlog"
+  "github.com/ankorstore/yokai/fxmcpserver"
+  "github.com/ankorstore/yokai/fxmetrics"
+  "github.com/ankorstore/yokai/fxtrace"
+  "github.com/mark3labs/mcp-go/mcp"
+  "github.com/mark3labs/mcp-go/server"
+  "go.uber.org/fx"
+)
+
+type ExampleHook struct {
+  config *config.Config
+}
+
+func NewExampleHook(config *config.Config) *ExampleHook {
+  return &ExampleHook{
+    config: config,
+  }
+}
+
+func (r *ExampleHook) Handle() server.SSEContextFunc {
+  return func(ctx context.Context, r *http.Request) context.Context {
+    return context.WithValue(ctx, "foo", "bar")
+  }
+}
+
+func main() {
+  fx.New(
+    fxconfig.FxConfigModule,
+    fxlog.FxLogModule,
+    fxtrace.FxTraceModule,
+    fxmetrics.FxMetricsModule,
+    fxgenerate.FxGenerateModule,
+    fxmcpserver.FxMCPServerModule,
+    fx.Options(
+      fxmcpserver.AsMCPSSEServerContextHook(NewExampleHook), // registers the NewExampleHook as MCP SSE server context hook
+    ),
+  ).Run()
+}
+```
 
 ### Testing
 

fxmcpserver/module.go

@@ -127,14 +127,20 @@ func ProvideMCPServer(p ProvideMCPServerParam) *server.MCPServer {
 // ProvideDefaultMCPSSEContextHandlerParam allows injection of the required dependencies in ProvideDefaultMCPSSEServerContextHandler.
 type ProvideDefaultMCPSSEContextHandlerParam struct {
 	fx.In
-	Generator      uuid.UuidGenerator
-	TracerProvider trace.TracerProvider
-	Logger         *log.Logger
+	Generator                uuid.UuidGenerator
+	TracerProvider           trace.TracerProvider
+	Logger                   *log.Logger
+	MCPSSEServerContextHooks []sse.MCPSSEServerContextHook `group:"mcp-sse-server-context-hooks"`
 }
 
 // ProvideDefaultMCPSSEServerContextHandler provides the default sse.MCPSSEServerContextHandler instance.
 func ProvideDefaultMCPSSEServerContextHandler(p ProvideDefaultMCPSSEContextHandlerParam) *sse.DefaultMCPSSEServerContextHandler {
-	return sse.NewDefaultMCPSSEServerContextHandler(p.Generator, p.TracerProvider, p.Logger)
+	return sse.NewDefaultMCPSSEServerContextHandler(
+		p.Generator,
+		p.TracerProvider,
+		p.Logger,
+		p.MCPSSEServerContextHooks...,
+	)
 }
 
 // ProvideDefaultMCPSSEServerFactoryParams allows injection of the required dependencies in ProvideDefaultMCPSSEServerFactory.

fxmcpserver/module_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/ankorstore/yokai/fxmcpserver"
 	"github.com/ankorstore/yokai/fxmcpserver/fxmcpservertest"
 	fs "github.com/ankorstore/yokai/fxmcpserver/server"
+	"github.com/ankorstore/yokai/fxmcpserver/testdata/hook"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/prompt"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/resource"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/resourcetemplate"
@@ -56,6 +57,7 @@ func TestMCPServerModule(t *testing.T) {
 			fxmcpserver.AsMCPServerPrompts(prompt.NewSimpleTestPrompt),
 			fxmcpserver.AsMCPServerResources(resource.NewSimpleTestResource),
 			fxmcpserver.AsMCPServerResourceTemplates(resourcetemplate.NewSimpleTestResourceTemplate),
+			fxmcpserver.AsMCPSSEServerContextHooks(hook.NewSimpleMCPSSEServerContextHook),
 			fxhealthcheck.AsCheckerProbe(fs.NewMCPServerProbe),
 		),
 		fx.Supply(fx.Annotate(context.Background(), fx.As(new(context.Context)))),
@@ -174,15 +176,15 @@ func TestMCPServerModule(t *testing.T) {
 
 	// send success prompts/get request
 	expectedRequest = `{"method":"prompts/get","params":{"name":"simple-test-prompt"}}`
-	expectedResponse = `{"description":"ok","messages":[{"role":"assistant","content":{"type":"text","text":"simple test prompt"}}]}`
+	expectedResponse = `{"description":"ok","messages":[{"role":"assistant","content":{"type":"text","text":"context hook value: bar"}}]}`
 
 	getPromptRequest := mcp.GetPromptRequest{}
 	getPromptRequest.Params.Name = "simple-test-prompt"
 
 	getPromptResult, err := testClient.GetPrompt(ctx, getPromptRequest)
 	assert.NoError(t, err)
 	assert.Equal(t, mcp.RoleAssistant, getPromptResult.Messages[0].Role)
-	assert.Equal(t, "simple test prompt", getPromptResult.Messages[0].Content.(mcp.TextContent).Text)
+	assert.Equal(t, "context hook value: bar", getPromptResult.Messages[0].Content.(mcp.TextContent).Text)
 
 	logtest.AssertHasLogRecord(t, logBuffer, map[string]any{
 		"level":        "info",

fxmcpserver/register.go

@@ -2,6 +2,7 @@ package fxmcpserver
 
 import (
 	"github.com/ankorstore/yokai/fxmcpserver/server"
+	"github.com/ankorstore/yokai/fxmcpserver/server/sse"
 	"go.uber.org/fx"
 )
 
@@ -92,3 +93,25 @@ func AsMCPServerResourceTemplates(constructors ...any) fx.Option {
 
 	return fx.Options(options...)
 }
+
+// AsMCPSSEServerContextHook registers an MCP SSE server context hook.
+func AsMCPSSEServerContextHook(constructor any) fx.Option {
+	return fx.Provide(
+		fx.Annotate(
+			constructor,
+			fx.As(new(sse.MCPSSEServerContextHook)),
+			fx.ResultTags(`group:"mcp-sse-server-context-hooks"`),
+		),
+	)
+}
+
+// AsMCPSSEServerContextHooks registers several MCP SSE server context hook.
+func AsMCPSSEServerContextHooks(constructors ...any) fx.Option {
+	options := []fx.Option{}
+
+	for _, constructor := range constructors {
+		options = append(options, AsMCPSSEServerContextHook(constructor))
+	}
+
+	return fx.Options(options...)
+}

fxmcpserver/register_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/ankorstore/yokai/fxmcpserver"
+	"github.com/ankorstore/yokai/fxmcpserver/testdata/hook"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/prompt"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/resource"
 	"github.com/ankorstore/yokai/fxmcpserver/testdata/resourcetemplate"
@@ -84,3 +85,21 @@ func TestAsMCPServerResourceTemplates(t *testing.T) {
 	assert.Equal(t, "fx.optionGroup", fmt.Sprintf("%T", reg))
 	assert.Implements(t, (*fx.Option)(nil), reg)
 }
+
+func TestAsMCPSSEServerContextHook(t *testing.T) {
+	t.Parallel()
+
+	reg := fxmcpserver.AsMCPSSEServerContextHook(hook.NewSimpleMCPSSEServerContextHook)
+
+	assert.Equal(t, "fx.provideOption", fmt.Sprintf("%T", reg))
+	assert.Implements(t, (*fx.Option)(nil), reg)
+}
+
+func TestAsMCPSSEServerContextHooks(t *testing.T) {
+	t.Parallel()
+
+	reg := fxmcpserver.AsMCPSSEServerContextHooks(hook.NewSimpleMCPSSEServerContextHook)
+
+	assert.Equal(t, "fx.optionGroup", fmt.Sprintf("%T", reg))
+	assert.Implements(t, (*fx.Option)(nil), reg)
+}

fxmcpserver/server/sse/context.go

@@ -16,6 +16,11 @@ import (
 
 var _ MCPSSEServerContextHandler = (*DefaultMCPSSEServerContextHandler)(nil)
 
+// MCPSSEServerContextHook is the interface for MCP SSE server context hooks.
+type MCPSSEServerContextHook interface {
+	Handle() server.SSEContextFunc
+}
+
 // MCPSSEServerContextHandler is the interface for MCP SSE server context handlers.
 type MCPSSEServerContextHandler interface {
 	Handle() server.SSEContextFunc
@@ -26,18 +31,21 @@ type DefaultMCPSSEServerContextHandler struct {
 	generator      uuid.UuidGenerator
 	tracerProvider ot.TracerProvider
 	logger         *log.Logger
+	contextHooks   []MCPSSEServerContextHook
 }
 
 // NewDefaultMCPSSEServerContextHandler returns a new DefaultMCPSSEServerContextHandler instance.
 func NewDefaultMCPSSEServerContextHandler(
 	generator uuid.UuidGenerator,
 	tracerProvider ot.TracerProvider,
 	logger *log.Logger,
+	contextHooks ...MCPSSEServerContextHook,
 ) *DefaultMCPSSEServerContextHandler {
 	return &DefaultMCPSSEServerContextHandler{
 		generator:      generator,
 		tracerProvider: tracerProvider,
 		logger:         logger,
+		contextHooks:   contextHooks,
 	}
 }
 
@@ -91,7 +99,14 @@ func (h *DefaultMCPSSEServerContextHandler) Handle() server.SSEContextFunc {
 
 		ctx = logger.WithContext(ctx)
 
-		// cancellation removal
-		return context.WithoutCancel(ctx)
+		// cancellation removal propagation
+		ctx = context.WithoutCancel(ctx)
+
+		// hooks propagation
+		for _, hook := range h.contextHooks {
+			ctx = hook.Handle()(ctx, req)
+		}
+
+		return ctx
 	}
 }

fxmcpserver/server/sse/context_test.go

@@ -8,6 +8,7 @@ import (
 
 	servercontext "github.com/ankorstore/yokai/fxmcpserver/server/context"
 	"github.com/ankorstore/yokai/fxmcpserver/server/sse"
+	"github.com/ankorstore/yokai/fxmcpserver/testdata/hook"
 	"github.com/ankorstore/yokai/log"
 	"github.com/ankorstore/yokai/log/logtest"
 	"github.com/stretchr/testify/assert"
@@ -27,11 +28,11 @@ func (m *generatorMock) Generate() string {
 func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 	t.Parallel()
 
-	t.Run("with provided session id and request id", func(t *testing.T) {
+	t.Run("with defaults", func(t *testing.T) {
 		t.Parallel()
 
 		gm := new(generatorMock)
-		gm.AssertNotCalled(t, "Generate")
+		gm.On("Generate").Return("test-request-id")
 
 		tp := trace.NewTracerProvider()
 
@@ -41,12 +42,11 @@ func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 
 		handler := sse.NewDefaultMCPSSEServerContextHandler(gm, tp, lg)
 
-		req := httptest.NewRequest(http.MethodGet, "/sse?sessionId=test-session-id", nil)
-		req.Header.Set("X-Request-Id", "test-request-id")
+		req := httptest.NewRequest(http.MethodGet, "/sse", nil)
 
 		ctx := handler.Handle()(context.Background(), req)
 
-		assert.Equal(t, "test-session-id", servercontext.CtxSessionID(ctx))
+		assert.Equal(t, "", servercontext.CtxSessionID(ctx))
 		assert.Equal(t, "test-request-id", servercontext.CtxRequestId(ctx))
 
 		span, ok := servercontext.CtxRootSpan(ctx).(trace.ReadWriteSpan)
@@ -62,7 +62,7 @@ func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 				assert.Equal(t, "sse", attr.Value.AsString())
 			}
 			if attr.Key == "mcp.sessionID" {
-				assert.Equal(t, "test-session-id", attr.Value.AsString())
+				assert.Equal(t, "", attr.Value.AsString())
 			}
 			if attr.Key == "mcp.requestID" {
 				assert.Equal(t, "test-request-id", attr.Value.AsString())
@@ -75,33 +75,36 @@ func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 			"level":        "info",
 			"system":       "mcpserver",
 			"mcpTransport": "sse",
-			"mcpSessionID": "test-session-id",
+			"mcpSessionID": "",
 			"mcpRequestID": "test-request-id",
 			"message":      "test log",
 		})
 
 		gm.AssertExpectations(t)
 	})
 
-	t.Run("without provided session id and request id", func(t *testing.T) {
+	t.Run("with provided session id and request id and middleware", func(t *testing.T) {
 		t.Parallel()
 
 		gm := new(generatorMock)
-		gm.On("Generate").Return("test-request-id")
+		gm.AssertNotCalled(t, "Generate")
 
 		tp := trace.NewTracerProvider()
 
 		lb := logtest.NewDefaultTestLogBuffer()
 		lg, err := log.NewDefaultLoggerFactory().Create(log.WithOutputWriter(lb))
 		assert.NoError(t, err)
 
-		handler := sse.NewDefaultMCPSSEServerContextHandler(gm, tp, lg)
+		hk := hook.NewSimpleMCPSSEServerContextHook()
 
-		req := httptest.NewRequest(http.MethodGet, "/sse", nil)
+		handler := sse.NewDefaultMCPSSEServerContextHandler(gm, tp, lg, hk)
+
+		req := httptest.NewRequest(http.MethodGet, "/sse?sessionId=test-session-id", nil)
+		req.Header.Set("X-Request-Id", "test-request-id")
 
 		ctx := handler.Handle()(context.Background(), req)
 
-		assert.Equal(t, "", servercontext.CtxSessionID(ctx))
+		assert.Equal(t, "test-session-id", servercontext.CtxSessionID(ctx))
 		assert.Equal(t, "test-request-id", servercontext.CtxRequestId(ctx))
 
 		span, ok := servercontext.CtxRootSpan(ctx).(trace.ReadWriteSpan)
@@ -117,7 +120,7 @@ func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 				assert.Equal(t, "sse", attr.Value.AsString())
 			}
 			if attr.Key == "mcp.sessionID" {
-				assert.Equal(t, "", attr.Value.AsString())
+				assert.Equal(t, "test-session-id", attr.Value.AsString())
 			}
 			if attr.Key == "mcp.requestID" {
 				assert.Equal(t, "test-request-id", attr.Value.AsString())
@@ -130,11 +133,14 @@ func TestDefaultMCPSSEServerContextHandler_Handle(t *testing.T) {
 			"level":        "info",
 			"system":       "mcpserver",
 			"mcpTransport": "sse",
-			"mcpSessionID": "",
+			"mcpSessionID": "test-session-id",
 			"mcpRequestID": "test-request-id",
 			"message":      "test log",
 		})
 
+		//nolint:forcetypeassert
+		assert.Equal(t, "bar", ctx.Value("foo").(string))
+
 		gm.AssertExpectations(t)
 	})
 }

fxmcpserver/server/stdio/context.go

@@ -78,7 +78,7 @@ func (h *DefaultMCPStdioServerContextHandler) Handle() server.StdioContextFunc {
 
 		ctx = logger.WithContext(ctx)
 
-		// cancellation removal
+		// cancellation removal propagation
 		return context.WithoutCancel(ctx)
 	}
 }

fxmcpserver/testdata/hook/simple.go

@@ -0,0 +1,20 @@
+package hook
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/mark3labs/mcp-go/server"
+)
+
+type SimpleMCPSSEServerContextHook struct{}
+
+func NewSimpleMCPSSEServerContextHook() *SimpleMCPSSEServerContextHook {
+	return &SimpleMCPSSEServerContextHook{}
+}
+
+func (p *SimpleMCPSSEServerContextHook) Handle() server.SSEContextFunc {
+	return func(ctx context.Context, r *http.Request) context.Context {
+		return context.WithValue(ctx, "foo", "bar")
+	}
+}