co-browser
diff --git a/‎internal/app/app.go
Copy file name to clipboard
+45-9Lines changed: 45 additions & 9 deletions b/‎internal/app/app.go
Copy file name to clipboard
+45-9Lines changed: 45 additions & 9 deletions
diff --git a/‎internal/mcp/config.go
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎internal/mcp/config.go
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎internal/mcp/server.go
Copy file name to clipboard
+73-115Lines changed: 73 additions & 115 deletions b/‎internal/mcp/server.go
Copy file name to clipboard
+73-115Lines changed: 73 additions & 115 deletions
@@ -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
 
@@ -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
 
@@ -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()
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ type RemoteMCPServer struct {`
`13`	`13`	URL string `json:"url"`
`14`	`14`	Name string `json:"name"`
`15`	`15`	Description string `json:"description"`
	`16`	+ ID int64 `json:"-"` // Add ID field, exclude from JSON config if needed
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`// MCPServerConfig holds the configuration for the MCP server`