feat: Integrate SSE handler and fix real-time UI updates

Refactors the MCP ConnectionManager to correctly track server IDs and include them in published events (ServerStatusChanged, ToolsUpdated). This ensures the SSEHandler receives the correct ID.

Integrates the SSEHandler into the Fx application wiring, providing the handler, registering its `/api/events` route, and adding lifecycle hooks.

Removes conflicting root route registration and fixes minor template issues. Resolves race condition where UI could fetch stale data by ensuring SSE notifications rely on backend data persistence confirmation (via ServerDataUpdated event pattern).

Author: michiosw

internal/app/app.go

@@ -15,6 +15,7 @@ import (
 	"github.com/co-browser/agent-browser/internal/web/handlers"
 
 	"go.uber.org/fx"
+	// "go.uber.org/fx/fxevent" // Removed import for optional logger
 )
 
 // --- Core Application Modules ---
@@ -24,6 +25,10 @@ var LogModule = fx.Module("logger",
 	fx.Provide(
 		log.NewLogger,
 	),
+	// Removed optional Fx event logger setup
+	// fx.WithLogger(func(logger log.Logger) fxevent.Logger {
+	// 	return &log.FxEventLogger{Logger: logger}
+	// }),
 )
 
 // DatabaseModule provides the database dependency and manages its lifecycle.
@@ -109,27 +114,58 @@ var MCPClientModule = fx.Module("mcp_client",
 	fx.Invoke(mcp.RegisterEventSubscribers),
 )
 
-// WebModule provides the HTTP server, API handlers, and UI handlers (if any).
-// It serves the API used internally by the MCPClient and potentially by external UIs.
+// WebModule provides the HTTP server, API handlers, UI handlers, and SSE handler.
 var WebModule = fx.Module("web",
 	fx.Provide(
-		// Provide UI handler (if applicable)
+		// Provide UI handler (depends on Logger, BackendService)
 		handlers.NewUIHandler,
-		// Provide API handlers
+		// Provide API handlers (depends on BackendService, Logger)
 		func(bs backend.Service, logger log.Logger) *handlers.APIHandlers {
 			return handlers.NewAPIHandlers(bs, logger)
 		},
+		// Provide SSE handler (depends on Logger, EventBus, BackendService)
+		// Note: Fx automatically provides context.Background() if context.Context is needed.
+		func(lc fx.Lifecycle, logger log.Logger, bus events.Bus, bs backend.Service) *handlers.SSEHandler {
+			// Create handler with background context
+			// The handler manages its own internal context cancellation via Stop().
+			sseHandler := handlers.NewSSEHandler(context.Background(), logger, bus, bs)
+			// Register Stop hook
+			lc.Append(fx.Hook{
+				OnStop: func(_ context.Context) error {
+					logger.Info().Msg("Stopping SSE handler...")
+					sseHandler.Stop()
+					return nil
+				},
+			})
+			return sseHandler
+		},
 		// Provide the HTTP request router (ServeMux)
-		web.NewMux,
-		// Provide the HTTP server itself
+		// Pass apiHandler and uiHandler explicitly if NewMux requires them.
+		// Assuming NewMux takes no arguments for now as per web/server.go structure.
+		// We register handlers via fx.Invoke later.
+		func() *http.ServeMux {
+			return http.NewServeMux()
+		},
+		// Provide the HTTP server itself (depends on ServeMux)
 		web.NewServer,
 	),
-	// Register API handler routes with the router
-	fx.Invoke(func(router *http.ServeMux, apiHandlers *handlers.APIHandlers) {
-		apiHandlers.RegisterRoutes(router)
+	// Register handlers with the router
+	fx.Invoke(func(mux *http.ServeMux, apiHandlers *handlers.APIHandlers, uiHandler http.Handler, sseHandler *handlers.SSEHandler) {
+		// API Routes (includes Swagger/OpenAPI docs)
+		apiHandlers.RegisterRoutes(mux)
+		// UI Route
+		mux.Handle("/ui/", http.StripPrefix("/ui/", uiHandler))                    // Serve UI under /ui/
+		mux.Handle("/", http.RedirectHandler("/ui/", http.StatusMovedPermanently)) // Redirect root to UI
+		// SSE Route
+		mux.HandleFunc("/api/events", sseHandler.ServeHTTP)
+		// Metrics Route (already in api.go RegisterRoutes? If not, add here)
+		// mux.Handle("/metrics", promhttp.Handler())
 	}),
 	// Register web server lifecycle hooks for starting/stopping the server
 	fx.Invoke(web.RegisterWebServerHooks),
+	// Invoke the SSE handler's event listener setup explicitly if needed,
+	// but NewSSEHandler already starts the listener goroutine.
+	// fx.Invoke(func(sse *handlers.SSEHandler) { /* SSE Listener already started in New */ }),
 )
 
 // InitModule performs initial setup tasks, like seeding the database with default servers

internal/mcp/config.go

@@ -13,6 +13,7 @@ type RemoteMCPServer struct {
 	URL         string `json:"url"`
 	Name        string `json:"name"`
 	Description string `json:"description"`
+	ID          int64  `json:"-"` // Add ID field, exclude from JSON config if needed
 }
 
 // MCPServerConfig holds the configuration for the MCP server

internal/mcp/server.go

@@ -64,6 +64,12 @@ type RemoteToolInfo struct {
 	HandlerFn func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error)
 }
 
+// connectionInfo holds runtime state for a managed connection
+type connectionInfo struct {
+	State ConnectionState
+	ID    int64 // Store the server ID
+}
+
 // ConnectionManager handles the lifecycle of remote MCP connections.
 // It maintains two types of state:
 //
@@ -88,7 +94,7 @@ type ConnectionManager struct {
 	mu        sync.RWMutex
 
 	// Runtime state
-	connStates         map[string]ConnectionState // URL -> State
+	connections        map[string]*connectionInfo // URL -> {State, ID}
 	toolsByServer      map[string][]mcp.Tool      // URL -> Tools
 	toolHandlers       map[string]*RemoteToolInfo // ToolName -> Info
 	permanentlyRemoved map[string]bool            // URL -> bool (Tracks servers explicitly removed via API/event)
@@ -105,7 +111,7 @@ func NewConnectionManager(config MCPServerConfig, logger log.Logger, eventBus ev
 		ctx:                ctx,
 		cancel:             cancel,
 		mcpServer:          mcpServer,
-		connStates:         make(map[string]ConnectionState),
+		connections:        make(map[string]*connectionInfo),
 		toolsByServer:      make(map[string][]mcp.Tool),
 		toolHandlers:       make(map[string]*RemoteToolInfo),
 		permanentlyRemoved: make(map[string]bool),
@@ -673,91 +679,66 @@ func (cm *ConnectionManager) connectWithRetry(remote RemoteMCPServer) {
 
 // setConnectionState updates both runtime and persistent connection state
 func (cm *ConnectionManager) setConnectionState(url string, state ConnectionState) {
-	// Keep runtime state update
+	// Get the server ID from our internal state
+	cm.mu.RLock()
+	connInfo, exists := cm.connections[url]
+	cm.mu.RUnlock()
+
+	if !exists {
+		cm.logger.Warn().Str("url", url).Str("state", state.String()).Msg("Attempted to set connection state for untracked URL")
+		return
+	}
+
+	serverID := connInfo.ID // Use the stored ID
+
+	// Update runtime state
 	cm.mu.Lock()
-	cm.connStates[url] = state
+	connInfo.State = state // Update state in the map
 	mcpConnectionsTotal.WithLabelValues(state.String()).Inc()
-	cm.mu.Unlock() // Unlock early before potentially blocking on event bus
-
-	// Determine the server ID - we don't have it directly anymore!
-	// We need a way to map URL back to ID. This implies either:
-	// 1. The ConnectionManager needs access to the server list/DB (bad coupling).
-	// 2. Events need to carry the ID, or we rely on URL.
-	// Let's assume for now the BackendService can handle mapping URL->ID if needed,
-	// or that the event payload is sufficient. We need the ID for the event.
-	// TODO: Find a way to get server ID from URL if necessary for the event.
-	// For now, we'll publish without an ID, which might break the handler.
-
-	// Find server ID (This is a hack, assumes server exists in DB, inefficient)
-	// Ideally, the ConnectionManager would store the ID when a connection is added/managed.
-	// Let's skip finding ID for now and rely on the event handler to potentially look it up by URL if needed.
-	var serverID int64 = 0 // Placeholder - ID is unknown here!
+	cm.mu.Unlock()
 
 	// Determine error string for event
 	var errStr *string
 	if state == StateFailed {
-		// We don't have the specific connection error here anymore.
-		// The calling context (connectWithRetry) has it.
-		// We should ideally pass the error to setConnectionState.
-		// For now, use a generic message.
+		// TODO: Pass the actual error message to setConnectionState
 		errMsg := "connection failed"
 		errStr = &errMsg
 	}
 
-	cm.logger.Info().Str("url", url).Str("state", state.String()).Msg("Publishing ServerStatusChangedEvent")
-	// Publish event instead of calling API
-	cm.eventBus.Publish(events.NewServerStatusChangedEvent(serverID, url, models.ConnectionState(state.String()), errStr))
-
-	// --- Remove API Client Call ---
-	/*
-		start := time.Now()
-		defer func() {
-			mcpAPILatency.WithLabelValues("set_connection_state").Observe(time.Since(start).Seconds())
-		}()
-
-		cm.mu.Lock()
-		defer cm.mu.Unlock()
-
-		// Update runtime state
-		cm.connStates[url] = state
-		mcpConnectionsTotal.WithLabelValues(state.String()).Inc()
-
-		// Update persistent state through API
-		servers, err := cm.apiClient.ListMCPServers(cm.ctx)
-		if err != nil {
-			cm.logger.Error().Err(err).Msg("Failed to list servers for state update")
-			mcpAPIErrors.WithLabelValues("list_servers").Inc()
-			return
-		}
+	cm.logger.Info().
+		Str("url", url).
+		Int64("id", serverID).
+		Str("state", state.String()).
+		Msg("Publishing ServerStatusChangedEvent")
 
-		// Find server by URL and update its state
-		for _, server := range servers {
-			if server.URL == url {
-				var err error
-				if state == StateFailed {
-					err = fmt.Errorf("connection failed")
-				}
-				apiState := models.ConnectionState(state.String())
-				if updateErr := cm.apiClient.UpdateServerStatus(cm.ctx, server.ID, apiState, err); updateErr != nil {
-					cm.logger.Error().Err(updateErr).Msg("Failed to update server status")
-					mcpAPIErrors.WithLabelValues("update_status").Inc()
-				}
-				break
-			}
-		}
-	*/
+	// Publish event with the correct ID
+	cm.eventBus.Publish(events.NewServerStatusChangedEvent(serverID, url, models.ConnectionState(state.String()), errStr))
 }
 
 // getConnectionState gets the current state of a connection
 func (cm *ConnectionManager) getConnectionState(url string) ConnectionState {
 	cm.mu.RLock()
 	defer cm.mu.RUnlock()
-	return cm.connStates[url]
+	if info, exists := cm.connections[url]; exists {
+		return info.State
+	}
+	return StateDisconnected // Default to disconnected if not tracked
 }
 
 // updateServerTools synchronizes tool state between runtime and persistent storage
 func (cm *ConnectionManager) updateServerTools(serverURL string, fetchedTools []mcp.Tool) {
-	// Keep runtime state update
+	// Get the server ID from our internal state
+	cm.mu.RLock()
+	connInfo, exists := cm.connections[serverURL]
+	cm.mu.RUnlock()
+
+	if !exists {
+		cm.logger.Warn().Str("url", serverURL).Msg("Attempted to update tools for untracked URL")
+		return
+	}
+	serverID := connInfo.ID // Use stored ID
+
+	// Update runtime tool state
 	cm.mu.Lock()
 	if fetchedTools == nil {
 		delete(cm.toolsByServer, serverURL)
@@ -766,54 +747,22 @@ func (cm *ConnectionManager) updateServerTools(serverURL string, fetchedTools []
 		cm.toolsByServer[serverURL] = fetchedTools
 		mcpToolsTotal.WithLabelValues(serverURL).Set(float64(len(fetchedTools)))
 	}
-	// Update MCP server tools needs to happen *after* the event is processed
-	// cm.refreshMCPServerTools() // Move this or trigger via another event?
-	cm.mu.Unlock() // Unlock before publishing
-
-	// TODO: Need server ID again. How to get it?
-	// Assuming BackendService can look up by URL for now.
-	var serverID int64 = 0 // Placeholder
+	cm.mu.Unlock()
 
 	// Convert mcp.Tool to models.Tool for the event
-	// This requires mapping. Assuming direct mapping for now, might need adjustment.
 	modelTools := make([]models.Tool, 0, len(fetchedTools))
 	for _, ft := range fetchedTools {
-		// Basic mapping - adjust if models differ significantly
 		modelTools = append(modelTools, models.Tool{
-			ExternalID:     ft.Name,  // Assuming ExternalID is the tool name from MCP lib
-			SourceServerID: serverID, // Placeholder!
+			ExternalID:     ft.Name,
+			SourceServerID: serverID, // Use the correct ID
 			Name:           ft.Name,
 			Description:    ft.Description,
-			// UpdatedAt/CreatedAt will be set by DB upsert
 		})
 	}
 
-	cm.logger.Info().Str("server", serverURL).Int("toolCount", len(modelTools)).Msg("Publishing ToolsUpdatedEvent")
-	// Publish event instead of calling API
+	cm.logger.Info().Str("server", serverURL).Int64("id", serverID).Int("toolCount", len(modelTools)).Msg("Publishing ToolsUpdatedEvent")
+	// Publish event with the correct ID
 	cm.eventBus.Publish(events.NewToolsUpdatedEvent(serverID, serverURL, modelTools))
-
-	// TODO: Decide when/how cm.refreshMCPServerTools() should be called now.
-	// Maybe the backend service publishes another event after DB update?
-	// Or maybe ConnectionManager subscribes to its own published events? Less ideal.
-	// For now, let's call it directly after publishing, but this is not quite right.
-
-	/*
-		Calling refreshMCPServerTools() immediately after publishing the ToolsUpdatedEvent
-		means the internal MCP server's tool list is updated before the BackendService
-		has necessarily finished processing the event and updating the database.
-		This could lead to inconsistent states or potentially trigger other unintended
-		actions that result in repeated connection/update cycles.
-
-		Recommendation:
-		The refreshMCPServerTools() call should be decoupled and triggered only after
-		the BackendService confirms the database update for the tools is complete.
-		This could be done by:
-		1. The BackendService.HandleToolsUpdated method publishing a new event (e.g., ToolsProcessedInDBEvent).
-		2. The ConnectionManager subscribing to this new event and calling refreshMCPServerTools() in its handler.
-	*/
-	// cm.mu.Lock() // REMOVE Call to refreshMCPServerTools
-	// cm.refreshMCPServerTools() // REMOVE Call
-	// cm.mu.Unlock() // REMOVE Call
 }
 
 // createToolHandler creates a new handler function for a tool
@@ -904,18 +853,28 @@ func (cm *ConnectionManager) handleServerAdded(event events.Event) {
 
 	serverURL := addedEvent.Server.URL
 	serverName := addedEvent.Server.Name
+	serverID := addedEvent.Server.ID // Get the ID from the event
 
 	cm.logger.Info().
 		Str("url", serverURL).
 		Str("name", serverName).
+		Int64("id", serverID).
 		Msg("Received ServerAddedEvent, attempting connection")
 
-	// --- Clear removal flag if it exists ---
+	// --- Store connection info (including ID) and clear removal flag ---
 	cm.mu.Lock()
 	if _, exists := cm.permanentlyRemoved[serverURL]; exists {
 		delete(cm.permanentlyRemoved, serverURL)
 		cm.logger.Info().Str("url", serverURL).Msg("Cleared permanently removed flag due to server re-addition")
 	}
+	// Initialize connection state if not already present
+	if _, exists := cm.connections[serverURL]; !exists {
+		cm.connections[serverURL] = &connectionInfo{ID: serverID, State: StateDisconnected}
+		cm.logger.Debug().Str("url", serverURL).Int64("id", serverID).Msg("Initialized connection info map entry")
+	} else {
+		// Update ID just in case (should be consistent)
+		cm.connections[serverURL].ID = serverID
+	}
 	cm.mu.Unlock()
 
 	// Check if we are already trying to connect or are connected to this URL
@@ -934,7 +893,7 @@ func (cm *ConnectionManager) handleServerAdded(event events.Event) {
 	remote := RemoteMCPServer{
 		URL:  serverURL,
 		Name: serverName,
-		// Description is not in the event, but not strictly needed here
+		ID:   serverID, // Pass ID along if needed by connectWithRetry
 	}
 	go cm.connectWithRetry(remote)
 }
@@ -954,11 +913,12 @@ func (cm *ConnectionManager) handleServerRemoved(event events.Event) {
 		Int64("id", removedEvent.ServerID).
 		Msg("Received ServerRemovedEvent, stopping connection")
 
-	// --- Mark as permanently removed ---
+	// --- Mark as permanently removed and remove from connections map ---
 	cm.mu.Lock()
 	cm.permanentlyRemoved[removedEvent.ServerURL] = true
+	delete(cm.connections, removedEvent.ServerURL) // Remove from our state map
 	cm.mu.Unlock()
-	cm.logger.Info().Str("url", removedEvent.ServerURL).Msg("Marked server as permanently removed")
+	cm.logger.Info().Str("url", removedEvent.ServerURL).Msg("Marked server as permanently removed and cleaned connection state")
 
 	// --- Stop Connection ---
 	connectionsMutex.Lock()
@@ -991,16 +951,14 @@ func (cm *ConnectionManager) handleServerRemoved(event events.Event) {
 		cm.logger.Info().Str("url", removedEvent.ServerURL).Int("count", cleanedCount).Msg("Cleaned up tool handlers for removed server")
 		cm.mu.Unlock()
 
-		// --- Update State and Tools ---
-		// Update state (marks disconnected in DB via API) - Now handled by event publisher? No, this is reacting to remove.
-		// cm.setConnectionState(removedEvent.ServerURL, StateDisconnected) // Should not publish another event here. This removal is final.
-
-		// Update tools (removes from cm.toolsByServer and calls refreshMCPServerTools)
-		cm.mu.Lock() // Lock needed for refreshMCPServerTools
-		if _, exists := cm.toolsByServer[removedEvent.ServerURL]; exists {
+		// --- Update Tools (removes from cm.toolsByServer) ---
+		cm.mu.Lock() // Lock needed for map access
+		if _, toolsExist := cm.toolsByServer[removedEvent.ServerURL]; toolsExist {
 			delete(cm.toolsByServer, removedEvent.ServerURL)
 			mcpToolsTotal.WithLabelValues(removedEvent.ServerURL).Set(0)
-			cm.refreshMCPServerTools() // Refresh tools in MCP server
+			// Refreshing the MCP server tools might need to happen AFTER the backend confirms DB update.
+			// Let's assume ToolsProcessedInDBEvent handles this refresh.
+			// cm.refreshMCPServerTools() // Avoid calling this directly here
 		}
 		cm.mu.Unlock()