Enhance Pulsar Function Management with Session ID Support

- Updated `PulsarFunctionManagedMcpTools` to accept a `sessionId` parameter, allowing for session-based tool management.
- Modified the `NewPulsarFunctionManager` function to include `sessionId`, improving the handling of session-specific tools.
- Enhanced the `updateFunctions` method to manage tools based on the presence of a session ID, ensuring proper addition and deletion of session tools.
- Added comments for better clarity on the changes and their implications for session management.

Author: freeznet

pkg/cmd/mcp/sse.go

@@ -76,7 +76,8 @@ func runSseServer(configOpts *ServerOptions) error {
 	}
 
 	// add Pulsar Functions as MCP tools
-	mcpServer.PulsarFunctionManagedMcpTools(configOpts.ReadOnly, configOpts.Features)
+	// SSE is not support session-based tools, so we pass an empty sessionId
+	mcpServer.PulsarFunctionManagedMcpTools(configOpts.ReadOnly, configOpts.Features, "")
 
 	sseServer := server.NewSSEServer(
 		mcpServer.MCPServer,

pkg/mcp/internal/pftools/manager.go

@@ -62,7 +62,7 @@ type Server struct {
 }
 
 // NewPulsarFunctionManager creates a new PulsarFunctionManager
-func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerOptions) (*PulsarFunctionManager, error) {
+func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerOptions, sessionId string) (*PulsarFunctionManager, error) {
 	// Get Pulsar client and admin client
 	if snServer.PulsarSession == nil {
 		return nil, fmt.Errorf("Pulsar session not found in context")
@@ -104,6 +104,7 @@ func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerO
 		circuitBreakers:   make(map[string]*CircuitBreaker),
 		tenantNamespaces:  options.TenantNamespaces,
 		strictExport:      options.StrictExport,
+		sessionId:         sessionId,
 	}
 
 	return manager, nil
@@ -197,9 +198,23 @@ func (m *PulsarFunctionManager) updateFunctions() {
 		}
 
 		if changed {
-			m.mcpServer.DeleteTools(fnTool.Tool.Name)
+			if m.sessionId != "" {
+				err := m.mcpServer.DeleteSessionTools(m.sessionId, fnTool.Tool.Name)
+				if err != nil {
+					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+				}
+			} else {
+				m.mcpServer.DeleteTools(fnTool.Tool.Name)
+			}
+		}
+		if m.sessionId != "" {
+			err := m.mcpServer.AddSessionTool(m.sessionId, fnTool.Tool, m.handleToolCall(fnTool))
+			if err != nil {
+				log.Printf("Failed to add tool %s to session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+			}
+		} else {
+			m.mcpServer.AddTool(fnTool.Tool, m.handleToolCall(fnTool))
 		}
-		m.mcpServer.AddTool(fnTool.Tool, m.handleToolCall(fnTool))
 
 		// Add function to map
 		m.mutex.Lock()
@@ -217,7 +232,14 @@ func (m *PulsarFunctionManager) updateFunctions() {
 	m.mutex.Lock()
 	for fullName, fnTool := range m.fnToToolMap {
 		if !seenFunctions[fullName] {
-			m.mcpServer.DeleteTools(fnTool.Tool.Name)
+			if m.sessionId != "" {
+				err := m.mcpServer.DeleteSessionTools(m.sessionId, fnTool.Tool.Name)
+				if err != nil {
+					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+				}
+			} else {
+				m.mcpServer.DeleteTools(fnTool.Tool.Name)
+			}
 			delete(m.fnToToolMap, fullName)
 			log.Printf("Removed function %s from MCP tools [%s]", fullName, fnTool.Tool.Name)
 		}

pkg/mcp/internal/pftools/types.go

@@ -47,6 +47,7 @@ type PulsarFunctionManager struct {
 	circuitBreakers   map[string]*CircuitBreaker
 	tenantNamespaces  []string
 	strictExport      bool
+	sessionId         string
 }
 
 type FunctionTool struct {

pkg/mcp/pulsar_functions_as_tools.go

@@ -51,7 +51,7 @@ func StopAllPulsarFunctionManagers() {
 	log.Println("All Pulsar Function managers stopped")
 }
 
-func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string) {
+func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string, sessionId string) {
 	if !slices.Contains(features, string(FeatureAll)) &&
 		!slices.Contains(features, string(FeatureFunctionsAsTools)) &&
 		!slices.Contains(features, string(FeatureStreamNativeCloud)) {
@@ -106,7 +106,7 @@ func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string)
 		Logger:        s.logger,
 	}
 
-	manager, err := pftools.NewPulsarFunctionManager(pftoolsServer, readOnly, options)
+	manager, err := pftools.NewPulsarFunctionManager(pftoolsServer, readOnly, options, sessionId)
 	if err != nil {
 		log.Printf("Failed to create Pulsar Function manager: %v", err)
 		return