Fix session handling

Author: fjakobs

experimental/apps-mcp/lib/mcp/middleware_test.go

@@ -224,7 +224,7 @@ func TestServerMiddleware(t *testing.T) {
 		Name:    "test-server",
 		Version: "1.0.0",
 	}
-	server := mcp.NewServer(impl, nil)
+	server := mcp.NewServer(impl, nil, nil)
 
 	var executionOrder []string
 
@@ -264,7 +264,7 @@ func TestServerSessionPersistence(t *testing.T) {
 		Name:    "test-server",
 		Version: "1.0.0",
 	}
-	server := mcp.NewServer(impl, nil)
+	server := mcp.NewServer(impl, nil, nil)
 
 	// Add middleware that increments a counter
 	server.AddMiddlewareFunc(func(ctx *mcp.MiddlewareContext, next mcp.NextFunc) (*mcp.CallToolResult, error) {

experimental/apps-mcp/lib/mcp/server.go

@@ -30,11 +30,15 @@ type serverTool struct {
 }
 
 // NewServer creates a new MCP server.
-func NewServer(impl *Implementation, options any) *Server {
+// If sess is nil, a new session will be created.
+func NewServer(impl *Implementation, options any, sess *session.Session) *Server {
+	if sess == nil {
+		sess = session.NewSession()
+	}
 	return &Server{
 		impl:    impl,
 		tools:   make(map[string]*serverTool),
-		session: session.NewSession(),
+		session: sess,
 	}
 }
 

experimental/apps-mcp/lib/server/server.go

@@ -31,11 +31,8 @@ func NewServer(ctx context.Context, cfg *mcp.Config) *Server {
 		Version: build.GetInfo().Version,
 	}
 
-	server := mcpsdk.NewServer(impl, nil)
 	sess := session.NewSession()
-
-	// Set enabled capabilities for this MCP server
-	sess.Set(session.CapabilitiesDataKey, []string{"apps"})
+	server := mcpsdk.NewServer(impl, nil, sess)
 
 	tracker, err := trajectory.NewTracker(ctx, sess, cfg)
 	if err != nil {
@@ -50,6 +47,9 @@ func NewServer(ctx context.Context, cfg *mcp.Config) *Server {
 
 	sess.SetTracker(tracker)
 
+	// Set enabled capabilities for this MCP server
+	sess.Set(session.CapabilitiesDataKey, []string{"apps"})
+
 	return &Server{
 		server:  server,
 		config:  cfg,
@@ -85,9 +85,6 @@ func (s *Server) RegisterTools(ctx context.Context) error {
 func (s *Server) registerCLIToolsProvider(ctx context.Context) error {
 	log.Info(ctx, "Registering CLI tools provider")
 
-	// Add session to context
-	ctx = session.WithSession(ctx, s.session)
-
 	provider, err := clitools.NewProvider(ctx, s.config, s.session)
 	if err != nil {
 		return err