feat(fxmcpserver): Updated context handling

fxmcpserver/module.go

@@ -149,12 +149,9 @@
 }
 
 // ProvideMCPSSEServerParam allows injection of the required dependencies in ProvideMCPSSEServer.
-//
-//nolint:containedctx
 type ProvideMCPSSEServerParam struct {
 	fx.In
 	LifeCycle                  fx.Lifecycle
-	Context                    context.Context
 	Logger                     *log.Logger
 	Config                     *config.Config
 	MCPServer                  *server.MCPServer
@@ -169,12 +166,14 @@
 		server.WithSSEContextFunc(p.MCPSSEServerContextHandler.Handle()),
 	)
 
+	sseServerCtx := p.Logger.WithContext(context.Background())
+
 	if p.Config.GetBool("modules.mcp.server.transport.sse.expose") {
 		p.LifeCycle.Append(fx.Hook{
 			OnStart: func(context.Context) error {
 				if !p.Config.IsTestEnv() {
-					//nolint:contextcheck,errcheck
-					go sseServer.Start(p.Context)
+					//nolint:errcheck
+					go sseServer.Start(sseServerCtx)
 				}
 
 				return nil
@@ -224,12 +223,9 @@
 }
 
 // ProvideMCPStdioServerParam allows injection of the required dependencies in ProvideMCPStdioServer.
-//
-//nolint:containedctx
 type ProvideMCPStdioServerParam struct {
 	fx.In
 	LifeCycle                    fx.Lifecycle
-	Context                      context.Context
 	Logger                       *log.Logger
 	Config                       *config.Config
 	MCPServer                    *server.MCPServer
@@ -244,12 +240,14 @@
 		server.WithStdioContextFunc(p.MCPStdioServerContextHandler.Handle()),
 	)
 
+	stdioServerCtx := p.Logger.WithContext(context.Background())
+
 	if p.Config.GetBool("modules.mcp.server.transport.stdio.expose") {
 		p.LifeCycle.Append(fx.Hook{
 			OnStart: func(context.Context) error {
 				if !p.Config.IsTestEnv() {
-					//nolint:contextcheck,errcheck
-					go stdioServer.Start(p.Context)
+					//nolint:errcheck
+					go stdioServer.Start(stdioServerCtx)
 				}
 
 				return nil

fxmcpserver/server/sse/context.go

@@ -43,21 +43,21 @@ func NewDefaultMCPSSEServerContextHandler(
 
 // Handle returns the handler func.
 func (h *DefaultMCPSSEServerContextHandler) Handle() server.SSEContextFunc {
-	return func(ctx context.Context, r *http.Request) context.Context {
+	return func(ctx context.Context, req *http.Request) context.Context {
 		// start time propagation
 		ctx = fsc.WithStartTime(ctx, time.Now())
 
 		// sessionId propagation
-		sID := r.URL.Query().Get("sessionId")
+		sID := req.URL.Query().Get("sessionId")
 
 		ctx = fsc.WithSessionID(ctx, sID)
 
 		// requestId propagation
-		rID := r.Header.Get("X-Request-Id")
+		rID := req.Header.Get("X-Request-Id")
 
 		if rID == "" {
 			rID = h.generator.Generate()
-			r.Header.Set("X-Request-Id", rID)
+			req.Header.Set("X-Request-Id", rID)
 		}
 
 		ctx = fsc.WithRequestID(ctx, rID)
@@ -89,6 +89,9 @@ func (h *DefaultMCPSSEServerContextHandler) Handle() server.SSEContextFunc {
 			Str("mcpRequestID", rID).
 			Logger()
 
-		return logger.WithContext(ctx)
+		ctx = logger.WithContext(ctx)
+
+		// cancellation removal
+		return context.WithoutCancel(ctx)
 	}
 }

fxmcpserver/server/stdio/context.go

@@ -76,6 +76,9 @@ func (h *DefaultMCPStdioServerContextHandler) Handle() server.StdioContextFunc {
 			Str("mcpRequestID", rID).
 			Logger()
 
-		return logger.WithContext(ctx)
+		ctx = logger.WithContext(ctx)
+
+		// cancellation removal
+		return context.WithoutCancel(ctx)
 	}
 }