feat: removes mcp servers from client config if autodiscovery is enabled (#133)

Signed-off-by: ChrisJBurns <29541485+ChrisJBurns@users.noreply.github.com>

Author: ChrisJBurns

cmd/thv/rm.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"github.com/stacklok/toolhive/pkg/client"
+	"github.com/stacklok/toolhive/pkg/config"
 	"github.com/stacklok/toolhive/pkg/container"
 	"github.com/stacklok/toolhive/pkg/labels"
 	"github.com/stacklok/toolhive/pkg/logger"
@@ -29,6 +31,7 @@ func init() {
 	rmCmd.Flags().BoolVarP(&rmForce, "force", "f", false, "Force removal of a running container")
 }
 
+//nolint:gocyclo // This function is complex but manageable
 func rmCmdFunc(_ *cobra.Command, args []string) error {
 	// Get container name
 	containerName := args[0]
@@ -101,5 +104,47 @@ func rmCmdFunc(_ *cobra.Command, args []string) error {
 	}
 
 	logger.Log.Info(fmt.Sprintf("Container %s removed", containerName))
+
+	if shouldRemoveClientConfig() {
+		if err := removeClientConfigurations(containerName); err != nil {
+			logger.Log.Error(fmt.Sprintf("Warning: Failed to remove client configurations: %v", err))
+		} else {
+			logger.Log.Info(fmt.Sprintf("Client configurations for %s removed", containerName))
+		}
+	}
+
+	return nil
+}
+
+func shouldRemoveClientConfig() bool {
+	c := config.GetConfig()
+	return len(c.Clients.RegisteredClients) > 0 && c.Clients.AutoDiscovery
+}
+
+// updateClientConfigurations updates client configuration files with the MCP server URL
+func removeClientConfigurations(containerName string) error {
+	// Find client configuration files
+	configs, err := client.FindClientConfigs()
+	if err != nil {
+		return fmt.Errorf("failed to find client configurations: %w", err)
+	}
+
+	if len(configs) == 0 {
+		logger.Log.Info("No client configuration files found")
+		return nil
+	}
+
+	for _, c := range configs {
+		logger.Log.Info(fmt.Sprintf("Removing MCP server from client configuration: %s", c.Path))
+
+		// Remove the MCP server configuration with locking
+		if err := c.DeleteConfigWithLock(containerName, c.Editor); err != nil {
+			logger.Log.Warn(fmt.Sprintf("Warning: Failed to remove MCP server from client configurationn %s: %v", c.Path, err))
+			continue
+		}
+
+		logger.Log.Info(fmt.Sprintf("Successfully removed MCP server from client configuration: %s", c.Path))
+	}
+
 	return nil
 }

pkg/client/config.go

@@ -336,6 +336,67 @@ func (c *ConfigFile) SaveWithLock(serverName, url string, editor ConfigEditor) e
 	return nil
 }
 
+// DeleteConfigWithLock safely removes the MCP server configuration in the file
+// It acquires a lock, reads the latest content, applies the change, and saves the file
+func (c *ConfigFile) DeleteConfigWithLock(serverName string, editor ConfigEditor) error {
+	// Create a lock file
+	fileLock := flock.New(c.Path + ".lock")
+
+	// Create a context with timeout
+	ctx, cancel := context.WithTimeout(context.Background(), lockTimeout)
+	defer cancel()
+
+	// Try to acquire the lock with a timeout
+	locked, err := fileLock.TryLockContext(ctx, 100*time.Millisecond)
+	if err != nil {
+		return fmt.Errorf("failed to acquire lock: %w", err)
+	}
+	if !locked {
+		return fmt.Errorf("failed to acquire lock: timeout after %v", lockTimeout)
+	}
+	defer fileLock.Unlock()
+
+	// Read the latest content from the file
+	latestConfig, err := readConfigFile(c.Path)
+	if err != nil {
+		return fmt.Errorf("failed to read latest config: %w", err)
+	}
+
+	// Apply our change to the latest content
+	if err := editor.RemoveServer(&latestConfig, serverName); err != nil {
+		return fmt.Errorf("failed to update latest config: %w", err)
+	}
+
+	// Determine format based on file extension
+	ext := strings.ToLower(filepath.Ext(c.Path))
+
+	var data []byte
+
+	if IsYAML(ext) {
+		// Marshal YAML
+		data, err = yaml.Marshal(latestConfig.Contents)
+		if err != nil {
+			return fmt.Errorf("failed to marshal YAML: %w", err)
+		}
+	} else {
+		// Default to JSON
+		data, err = json.MarshalIndent(latestConfig.Contents, "", "  ")
+		if err != nil {
+			return fmt.Errorf("failed to marshal JSON: %w", err)
+		}
+	}
+
+	// Write file
+	if err := os.WriteFile(c.Path, data, 0600); err != nil {
+		return fmt.Errorf("failed to write file: %w", err)
+	}
+
+	// Update our in-memory representation to match the file
+	c.Contents = latestConfig.Contents
+
+	return nil
+}
+
 // GenerateMCPServerURL generates the URL for an MCP server
 func GenerateMCPServerURL(host string, port int, containerName string) string {
 	// The URL format is: http://host:port/sse#container-name

pkg/client/config_test.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/transport/ssecommon"
 )
 
@@ -62,7 +63,7 @@ func setupTestConfig(t *testing.T, testName string) (string, string, ConfigFile)
 }
 
 // getMCPServers reads the config file and returns the mcpServers map
-func getMCPServers(t *testing.T, configPath string) map[string]interface{} {
+func getMCPServers(t *testing.T, configPath string, editor ConfigEditor) map[string]interface{} {
 	t.Helper()
 
 	// Read the config file
@@ -71,27 +72,48 @@ func getMCPServers(t *testing.T, configPath string) map[string]interface{} {
 		t.Fatalf("Failed to read updated config file: %v", err)
 	}
 
-	// Check if the servers were updated correctly
-	mcpServers, ok := updatedConfig.Contents["mcpServers"].(map[string]interface{})
-	if !ok {
-		t.Fatalf("mcpServers is not a map")
+	_, vsOk := editor.(*VSCodeConfigEditor)
+	if vsOk {
+		updatedConfig.Editor = editor
+		mcpMap, ok := updatedConfig.Contents["mcp"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("mcp is not a map")
+		}
+
+		// Get servers child object
+		mcpServers, ok := mcpMap["servers"]
+		if !ok {
+			t.Fatalf("mcpServers is not a map")
+		}
+		return mcpServers.(map[string]interface{})
+	}
+
+	_, standardOk := editor.(*StandardConfigEditor)
+	if standardOk {
+		mcpServers, ok := updatedConfig.Contents["mcpServers"].(map[string]interface{})
+		if !ok {
+			t.Fatalf("mcpServers is not a map")
+		}
+		return mcpServers
 	}
 
-	return mcpServers
+	return nil
 }
 
 // testUpdateExistingServer tests updating an existing server
 func testUpdateExistingServer(t *testing.T, config ConfigFile, configPath string) {
 	t.Helper()
 	// Test updating an existing server with lock
 	expectedURL := "http://localhost:54321" + ssecommon.HTTPSSEEndpoint + "#test-container"
-	err := config.SaveWithLock("existing-server", expectedURL, &StandardConfigEditor{})
+	editor := &StandardConfigEditor{}
+
+	err := config.SaveWithLock("existing-server", expectedURL, editor)
 	if err != nil {
 		t.Fatalf("Failed to update MCP server config: %v", err)
 	}
 
 	// Get the updated servers
-	mcpServers := getMCPServers(t, configPath)
+	mcpServers := getMCPServers(t, configPath, editor)
 
 	// Check existing server
 	existingServer, ok := mcpServers["existing-server"].(map[string]interface{})
@@ -112,13 +134,14 @@ func testAddNewServer(t *testing.T, config ConfigFile, configPath string) {
 	t.Helper()
 	// Test adding a new server with lock
 	expectedURL := "http://localhost:9876" + ssecommon.HTTPSSEEndpoint + "#new-container"
-	err := config.SaveWithLock("new-server", expectedURL, &StandardConfigEditor{})
+	editor := &StandardConfigEditor{}
+	err := config.SaveWithLock("new-server", expectedURL, editor)
 	if err != nil {
 		t.Fatalf("Failed to add new MCP server config: %v", err)
 	}
 
 	// Get the updated servers
-	mcpServers := getMCPServers(t, configPath)
+	mcpServers := getMCPServers(t, configPath, editor)
 
 	// Check new server
 	newServer, ok := mcpServers["new-server"].(map[string]interface{})
@@ -134,11 +157,53 @@ func testAddNewServer(t *testing.T, config ConfigFile, configPath string) {
 	}
 }
 
+// testAddNewServer tests adding a new server
+func testRemovingServer(t *testing.T, config ConfigFile, configPath string, editor ConfigEditor) {
+	t.Helper()
+	// Test adding a new server with lock
+	expectedURL := "http://localhost:9876" + ssecommon.HTTPSSEEndpoint + "#new-container"
+	err := config.SaveWithLock("new-server", expectedURL, editor)
+	if err != nil {
+		t.Fatalf("Failed to add new MCP server config: %v", err)
+	}
+
+	// Get the updated servers
+	mcpServers := getMCPServers(t, configPath, editor)
+
+	// Check new server
+	newServer, ok := mcpServers["new-server"].(map[string]interface{})
+	if !ok {
+		t.Fatalf("new-server is not a map")
+	}
+	newURL, ok := newServer["url"].(string)
+	if !ok {
+		t.Fatalf("url is not a string")
+	}
+	if newURL != expectedURL {
+		t.Fatalf("Unexpected URL for new-server: %s, expected: %s", newURL, expectedURL)
+	}
+
+	// Remove the server
+	err = config.DeleteConfigWithLock("new-server", editor)
+	if err != nil {
+		t.Fatalf("Failed to remove MCP server config: %v", err)
+	}
+
+	mcpServersNew := getMCPServers(t, configPath, editor)
+
+	// Check that the server was removed
+	_, ok = mcpServersNew["new-server"]
+	if ok {
+		t.Fatalf("new-server is still in the config")
+	}
+}
+
 // testPreserveExistingConfig tests that existing configurations are preserved
 func testPreserveExistingConfig(t *testing.T, configPath string) {
 	t.Helper()
 	// Get the updated servers
-	mcpServers := getMCPServers(t, configPath)
+	editor := &StandardConfigEditor{}
+	mcpServers := getMCPServers(t, configPath, editor)
 
 	// Check postgres server (should be unchanged)
 	postgresServer, ok := mcpServers["postgres"].(map[string]interface{})
@@ -213,6 +278,41 @@ func TestUpdateMCPServerConfig(t *testing.T) {
 
 		testPreserveExistingConfig(t, configPath)
 	})
+
+}
+
+func TestRemoveMCPServerConfig(t *testing.T) {
+	t.Parallel()
+
+	logger.Initialize()
+
+	t.Run("RemoveExistingServerStandardEditor", func(t *testing.T) {
+		t.Parallel()
+
+		// Setup test environment for this subtest
+		tempDir, configPath, config := setupTestConfig(t, "remove")
+		t.Cleanup(func() {
+			if err := os.RemoveAll(tempDir); err != nil {
+				t.Logf("Failed to remove temp dir: %v", err)
+			}
+		})
+
+		testRemovingServer(t, config, configPath, &StandardConfigEditor{})
+	})
+
+	t.Run("RemoveExistingServerVSCodeEditor", func(t *testing.T) {
+		t.Parallel()
+
+		// Setup test environment for this subtest
+		tempDir, configPath, config := setupTestConfig(t, "remove")
+		t.Cleanup(func() {
+			if err := os.RemoveAll(tempDir); err != nil {
+				t.Logf("Failed to remove temp dir: %v", err)
+			}
+		})
+
+		testRemovingServer(t, config, configPath, &VSCodeConfigEditor{})
+	})
 }
 
 func TestGenerateMCPServerURL(t *testing.T) {