feat: Use event bus for MCP status/tool updates

michiosw · michiosw · commit dd1e1f7c9fcb · 2025-04-14T17:14:03.000-07:00
Removes the API client dependency from ConnectionManager.

ConnectionManager now publishes ServerStatusChanged and ToolsUpdated
events directly to the event bus instead of calling back into the Web API.

BackendService subscribes to these events to handle DB updates.

InitModule now publishes ServerAddedEvent for existing servers at startup
to ensure ConnectionManager connects to them.
diff --git a/internal/app/app.go b/internal/app/app.go
@@ -68,9 +68,23 @@ var BackendModule = fx.Module("backend",
 	fx.Provide(
 		// Provide the backend service implementation, satisfying the backend.Service interface.
 		func(db database.DBInterface, bus events.Bus, logger log.Logger) backend.Service {
+			// Just return the service instance created by NewService.
 			return backend.NewService(db, bus, logger)
 		},
+		/* // Remove previous attempts
+		func(lc fx.Lifecycle, db database.DBInterface, bus events.Bus, logger log.Logger) backend.Service {
+			// ... creation, assertion, hook registration ...
+			return serviceInterface
+		},
+		*/
 	),
+	// Invoke the dedicated registration function from the backend package.
+	fx.Invoke(backend.RegisterEventHandlers),
+	/* // Remove previous attempts
+	fx.Invoke(func(bus events.Bus, serviceImpl *backend.service, logger log.Logger) { // Request concrete type
+		// ... subscription ...
+	}),
+	*/
 )
 
 // MCPClientModule provides the components responsible for connecting *to* remote MCP servers.
@@ -91,6 +105,8 @@ var MCPClientModule = fx.Module("mcp_client",
 	),
 	// Register lifecycle hooks for starting/stopping the MCP ConnectionManager
 	fx.Invoke(mcp.RegisterMCPServerHooks),
+	// Register event subscribers for the ConnectionManager
+	fx.Invoke(mcp.RegisterEventSubscribers),
 )
 
 // WebModule provides the HTTP server, API handlers, and UI handlers (if any).
@@ -116,57 +132,61 @@ var WebModule = fx.Module("web",
 	fx.Invoke(web.RegisterWebServerHooks),
 )
 
-// InitModule performs initial setup tasks, like seeding the database with default servers.
-// It runs once during application startup if needed.
+// InitModule performs initial setup tasks, like seeding the database with default servers
+// and triggering connections for existing servers.
 var InitModule = fx.Module("init",
-	fx.Invoke(func(bs backend.Service, logger log.Logger) {
+	// Use fx.Invoke to run initialization logic after dependencies are ready.
+	fx.Invoke(func(bs backend.Service, bus events.Bus, logger log.Logger) {
+		logger.Info().Msg("Running initialization logic...")
 		// Check if any servers already exist in the database.
 		servers, err := bs.ListMCPServers()
 		if err != nil {
-			logger.Error().Err(err).Msg("Failed to check existing servers during initialization")
-			// Decide if this should be fatal. For now, we continue, but defaults won't be added.
+			logger.Error().Err(err).Msg("Failed to list existing servers during initialization")
+			// Decide if this should be fatal. For now, we continue.
 			return
 		}
 
-		// Only add default servers if the database is empty.
+		// Add default servers ONLY if the database is empty.
 		if len(servers) == 0 {
 			logger.Info().Msg("No MCP servers found in the database. Adding default servers...")
-
 			// Define the default servers to add.
-			// TODO: Consider moving these defaults to configuration.
 			defaultServers := []struct {
 				name        string
 				url         string
 				description string
 			}{
 				{
 					name:        "Local Test Server",
-					url:         "https://0.0.0.0:8001/sse", // Assumes local test server runs on 8001
+					url:         "https://0.0.0.0:8001/sse",
 					description: "Local MCP test server",
 				},
 			}
 
 			// Add each default server via the backend service.
 			for _, server := range defaultServers {
-				_, err := bs.AddMCPServer(server.name, server.url) // Description is not stored via AddMCPServer currently
+				// AddMCPServer already publishes ServerAddedEvent, so we don't need to do it manually here.
+				_, err := bs.AddMCPServer(server.name, server.url)
 				if err != nil {
-					// Log error but continue trying to add others.
 					logger.Error().
 						Err(err).
 						Str("name", server.name).
 						Str("url", server.url).
-						// Str("description", server.description). // Not used by AddMCPServer
 						Msg("Failed to add default server")
-					continue // Skip logging success for this server
+					continue
 				}
 				logger.Info().
 					Str("name", server.name).
 					Str("url", server.url).
-					// Str("description", server.description). // Not stored
-					Msg("Added default server")
+					Msg("Added default server (will trigger ServerAddedEvent)")
 			}
 		} else {
-			logger.Info().Int("count", len(servers)).Msg("Servers already exist, skipping default server initialization")
+			// If servers already exist, publish ServerAddedEvent for each one
+			// to trigger the ConnectionManager.
+			logger.Info().Int("count", len(servers)).Msg("Existing servers found. Publishing ServerAddedEvent for each...")
+			for _, server := range servers {
+				logger.Debug().Int64("id", server.ID).Str("url", server.URL).Msg("Publishing ServerAddedEvent for existing server")
+				bus.Publish(events.NewServerAddedEvent(server))
+			}
 		}
 	}),
 )
diff --git a/internal/backend/service.go b/internal/backend/service.go
@@ -167,7 +167,6 @@ func (s *service) ProcessFetchedTools(serverID int64, fetchedTools []models.Fetc
 	}
 
 	s.logger.Info().Int("fetched", len(fetchedTools)).Int("added", addedCount).Int("updated", updatedCount).Int64("serverID", serverID).Msg("Processed fetched tools")
-	s.bus.Publish(events.NewToolsUpdatedEvent(serverID, len(fetchedTools))) // Existing event
 	return addedCount, updatedCount, nil
 }
 
@@ -180,9 +179,9 @@ func (s *service) UpdateMCPServerStatus(id int64, state models.ConnectionState,
 		return fmt.Errorf("failed to update server status: %w", err)
 	}
 
-	// Publish the status change event
-	s.bus.Publish(events.NewServerStatusChangedEvent(id, state, errStr))
-	s.logger.Debug().Int64("id", id).Str("state", string(state)).Msg("Updated server status and published ServerStatusChangedEvent")
+	// Publish the status change event - REMOVED - ConnectionManager publishes this now.
+	// s.bus.Publish(events.NewServerStatusChangedEvent(id, state, errStr))
+	s.logger.Debug().Int64("id", id).Str("state", string(state)).Msg("Updated server status in DB (event published by ConnectionManager)")
 
 	// Fetching the updated server is no longer strictly necessary just for publishing the event
 	// updatedServer, getErr := s.db.GetServerByID(id)
@@ -229,3 +228,122 @@ func (s *service) UpdateMCPServer(id int64, name, url string) (*models.MCPServer
 	// s.bus.Publish(events.NewServerUpdatedEvent(*updatedServer)) // TODO: Define and use a ServerUpdatedEvent
 	return updatedServer, nil
 }
+
+// --- Event Handlers ---
+
+// HandleServerStatusChanged processes ServerStatusChangedEvent received from the event bus.
+func (s *service) HandleServerStatusChanged(event events.Event) {
+	statusEvent, ok := event.(*events.ServerStatusChangedEvent)
+	if !ok {
+		s.logger.Error().Str("eventType", string(event.Type())).Msg("Received event of unexpected type in HandleServerStatusChanged")
+		return
+	}
+
+	// Need to find the server ID from the URL if ID is 0 (publisher might not know it)
+	serverID := statusEvent.ServerID
+	if serverID == 0 {
+		server, err := s.db.GetServerByURL(statusEvent.ServerURL)
+		if err != nil {
+			s.logger.Error().Err(err).Str("url", statusEvent.ServerURL).Msg("Error finding server by URL for status update event")
+			return
+		}
+		if server == nil {
+			s.logger.Warn().Str("url", statusEvent.ServerURL).Msg("Received status update event for unknown server URL")
+			return
+		}
+		serverID = server.ID
+		s.logger.Debug().Int64("serverID", serverID).Str("url", statusEvent.ServerURL).Msg("Mapped server URL to ID for status update event")
+	}
+
+	s.logger.Info().
+		Int64("serverID", serverID).
+		Str("url", statusEvent.ServerURL).
+		Str("newState", string(statusEvent.NewState)).
+		Msg("Handling ServerStatusChanged event")
+
+	// Call the existing DB update method
+	err := s.db.UpdateServerStatus(serverID, statusEvent.NewState, statusEvent.LastError, time.Now())
+	if err != nil {
+		s.logger.Error().Err(err).Int64("serverID", serverID).Msg("Failed to update DB from ServerStatusChangedEvent")
+		// Note: We don't re-publish the event here to avoid loops.
+	}
+}
+
+// HandleToolsUpdated processes ToolsUpdatedEvent received from the event bus.
+func (s *service) HandleToolsUpdated(event events.Event) {
+	toolsEvent, ok := event.(*events.ToolsUpdatedEvent)
+	if !ok {
+		s.logger.Error().Str("eventType", string(event.Type())).Msg("Received event of unexpected type in HandleToolsUpdated")
+		return
+	}
+
+	// Need to find the server ID from the URL if ID is 0
+	serverID := toolsEvent.ServerID
+	if serverID == 0 {
+		server, err := s.db.GetServerByURL(toolsEvent.ServerURL)
+		if err != nil {
+			s.logger.Error().Err(err).Str("url", toolsEvent.ServerURL).Msg("Error finding server by URL for tools update event")
+			return
+		}
+		if server == nil {
+			s.logger.Warn().Str("url", toolsEvent.ServerURL).Msg("Received tools update event for unknown server URL")
+			return
+		}
+		serverID = server.ID
+		s.logger.Debug().Int64("serverID", serverID).Str("url", toolsEvent.ServerURL).Msg("Mapped server URL to ID for tools update event")
+	}
+
+	s.logger.Info().
+		Int64("serverID", serverID).
+		Str("url", toolsEvent.ServerURL).
+		Int("toolCount", len(toolsEvent.Tools)).
+		Msg("Handling ToolsUpdated event")
+
+	// Process each tool using the DB upsert method
+	addedCount := 0
+	updatedCount := 0
+	hadError := false
+	for _, tool := range toolsEvent.Tools {
+		// Ensure the SourceServerID is set correctly, as the event might have had 0
+		tool.SourceServerID = serverID
+
+		wasAdded, upsertErr := s.db.UpsertTool(tool)
+		if upsertErr != nil {
+			s.logger.Error().Err(upsertErr).Str("externalID", tool.ExternalID).Int64("serverID", serverID).Msg("Error upserting tool from ToolsUpdatedEvent")
+			hadError = true
+			continue // Continue processing other tools
+		}
+		if wasAdded {
+			addedCount++
+		} else {
+			updatedCount++
+		}
+	}
+
+	if hadError {
+		s.logger.Warn().Int64("serverID", serverID).Msg("Encountered errors while processing tools from ToolsUpdatedEvent")
+	}
+
+	s.logger.Info().
+		Int64("serverID", serverID).
+		Int("added", addedCount).
+		Int("updated", updatedCount).
+		Msg("Finished processing ToolsUpdatedEvent")
+
+	// TODO: Potentially publish another event like 'ToolsProcessed' if other components
+	// need to know the DB update is complete (e.g., to trigger MCP server tool refresh).
+}
+
+// RegisterEventHandlers registers the necessary event handlers for the backend service.
+// It performs a type assertion to access the unexported handler methods.
+func RegisterEventHandlers(bus events.Bus, serviceInterface Service, logger log.Logger) {
+	serviceImpl, ok := serviceInterface.(*service) // Type assertion to concrete *service
+	if !ok {
+		logger.Fatal().Msg("Backend service provided to RegisterEventHandlers is not of expected type *service")
+		return // Or panic
+	}
+
+	logger.Info().Msg("Registering backend service event handlers...")
+	bus.Subscribe(events.ServerStatusChanged, serviceImpl.HandleServerStatusChanged)
+	bus.Subscribe(events.ToolsUpdated, serviceImpl.HandleToolsUpdated)
+}
diff --git a/internal/events/events.go b/internal/events/events.go
@@ -81,17 +81,19 @@ func NewServerRemovedEvent(serverID int64, serverURL string) *ServerRemovedEvent
 type ToolsUpdatedEvent struct {
 	baseEvent
 	ServerID     int64
-	FetchedCount int // How many tools were in the fetched list (can be 0)
-	// TODO: Optionally include added/updated counts if ProcessFetchedTools can return them
-	// TODO: Optionally include list of changed tool IDs/names?
+	ServerURL    string
+	FetchedCount int
+	Tools        []models.Tool
 }
 
 // NewToolsUpdatedEvent creates a new event for when tools are updated for a server
-func NewToolsUpdatedEvent(serverID int64, fetchedCount int) *ToolsUpdatedEvent {
+func NewToolsUpdatedEvent(serverID int64, serverURL string, tools []models.Tool) *ToolsUpdatedEvent {
 	return &ToolsUpdatedEvent{
 		baseEvent:    newBaseEvent(ToolsUpdated),
 		ServerID:     serverID,
-		FetchedCount: fetchedCount,
+		ServerURL:    serverURL,
+		FetchedCount: len(tools),
+		Tools:        tools,
 	}
 }
 
@@ -100,15 +102,17 @@ func NewToolsUpdatedEvent(serverID int64, fetchedCount int) *ToolsUpdatedEvent {
 type ServerStatusChangedEvent struct {
 	baseEvent
 	ServerID  int64
-	NewState  models.ConnectionState // The new state (connected, disconnected, failed)
-	LastError *string                // Optional error message if state is 'failed'
+	ServerURL string
+	NewState  models.ConnectionState
+	LastError *string
 }
 
 // NewServerStatusChangedEvent creates a new event for server status changes.
-func NewServerStatusChangedEvent(serverID int64, newState models.ConnectionState, lastError *string) *ServerStatusChangedEvent {
+func NewServerStatusChangedEvent(serverID int64, serverURL string, newState models.ConnectionState, lastError *string) *ServerStatusChangedEvent {
 	return &ServerStatusChangedEvent{
 		baseEvent: newBaseEvent(ServerStatusChanged),
 		ServerID:  serverID,
+		ServerURL: serverURL,
 		NewState:  newState,
 		LastError: lastError,
 	}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
