Merge branch 'main' into issue1063

Author: ChrisJBurns

.github/workflows/verify-swagger.yml renamed to .github/workflows/verify-gen.yml

@@ -1,4 +1,4 @@
-name: Swagger
+name: Codegen
 
 on:
   workflow_call:
@@ -8,7 +8,7 @@ permissions:
 
 jobs:
   swagger:
-    name: Verify Swagger
+    name: Verify code generation
     runs-on: ubuntu-latest
 
     steps:
@@ -21,16 +21,16 @@ jobs:
         with:
           version: 3.x
           repo-token: ${{ secrets.GITHUB_TOKEN }}
-      - name: Install swagger
-        run: task swagger-install
-      - name: Generate swagger files
-        run: task swagger-gen
+      - name: Install mockgen
+        run: task mock-install
+      - name: Generate code files
+        run: task gen
       - name: Check for changes
         run: |
           if ! git diff --exit-code docs/server/; then
-            echo "❌ Swagger files are not up to date!"
-            echo "Please run 'task swagger-gen' or 'swag init -g pkg/api/server.go --v3.1 -o docs/server' and commit the changes."
+            echo "❌ Generated code files are not up to date!"
+            echo "Please run 'task gen' and commit the changes."
             exit 1
           else
-            echo "✅ Swagger files are up to date!"
+            echo "✅ Generated code files are up to date!"
           fi

CLAUDE.md

@@ -41,8 +41,8 @@ task test-e2e
 # Run all tests (unit and e2e)
 task test-all
 
-# Generate OpenAPI/Swagger documentation
-task swagger-gen
+# Generate mocks
+task gen
 
 # Generate CLI documentation
 task docs

Taskfile.yml

@@ -8,20 +8,27 @@ includes:
 tasks:
   docs:
     desc: Regenerate the docs
+    deps: [swagger-install]
     cmds:
       - rm -rf docs/cli/*
       - go run cmd/help/main.go --dir docs/cli
+      - swag init -g pkg/api/server.go --v3.1 -o docs/server
 
   swagger-install:
     desc: Install the swag tool for OpenAPI/Swagger generation
     cmds:
       - go install github.com/swaggo/swag/v2/cmd/swag@latest
 
-  swagger-gen:
-    desc: Generate OpenAPI/Swagger documentation
-    deps: [swagger-install]
+  mock-install:
+    desc: Install the mockgen tool for mock generation
     cmds:
-      - swag init -g pkg/api/server.go --v3.1 -o docs/server
+      - go install go.uber.org/mock/mockgen@latest
+
+  gen:
+    desc: Generate mock files using go generate
+    deps: [mock-install]
+    cmds:
+      - go generate ./...
 
   lint:
     desc: Run linting tools
@@ -53,6 +60,7 @@ tasks:
 
   test:
     desc: Run unit tests (excluding e2e tests)
+    deps: [gen]
     cmds:
       - task: test-unixlike
         platforms: [linux, darwin]
@@ -128,6 +136,7 @@ tasks:
 
   build:
     desc: Build the binary
+    deps: [gen]
     vars:
       VERSION:
         sh: git describe --tags --always --dirty --match "v*" || echo "dev"

cmd/thv-proxyrunner/app/run.go

@@ -266,7 +266,10 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to create RunConfig: %v", err)
 	}
 
-	workloadManager := workloads.NewManagerFromRuntime(rt)
+	workloadManager, err := workloads.NewManagerFromRuntime(rt)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %v", err)
+	}
 	return workloadManager.RunWorkload(ctx, runConfig)
 }
 

cmd/thv/app/client.go

@@ -12,6 +12,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/config"
 	"github.com/stacklok/toolhive/pkg/core"
 	"github.com/stacklok/toolhive/pkg/groups"
+	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
@@ -104,6 +105,8 @@ func init() {
 	// TODO: Re-enable when group functionality is complete
 	//clientRegisterCmd.Flags().StringSliceVar(
 	//	&groupNames, "group", []string{"default"}, "Only register workloads from specified groups")
+	//clientRemoveCmd.Flags().StringSliceVar(
+	//	&groupNames, "group", []string{}, "Remove client from specified groups (if not set, removes all workloads from the client)")
 }
 
 func clientStatusCmdFunc(_ *cobra.Command, _ []string) error {
@@ -208,21 +211,7 @@ func clientRemoveCmdFunc(cmd *cobra.Command, args []string) error {
 			clientType)
 	}
 
-	ctx := cmd.Context()
-
-	manager, err := client.NewManager(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to create client manager: %w", err)
-	}
-
-	err = manager.UnregisterClients(ctx, []client.Client{
-		{Name: client.MCPClient(clientType)},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to remove client %s: %w", clientType, err)
-	}
-
-	return nil
+	return performClientRemoval(cmd.Context(), client.Client{Name: client.MCPClient(clientType)}, groupNames)
 }
 
 func listRegisteredClientsCmdFunc(_ *cobra.Command, _ []string) error {
@@ -339,3 +328,114 @@ func registerClientsGlobally(
 
 	return nil
 }
+
+func performClientRemoval(ctx context.Context, clientToRemove client.Client, groupNames []string) error {
+	clientManager, err := client.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create client manager: %w", err)
+	}
+
+	workloadManager, err := workloads.NewManager(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create workload manager: %w", err)
+	}
+
+	runningWorkloads, err := workloadManager.ListWorkloads(ctx, false)
+	if err != nil {
+		return fmt.Errorf("failed to list running workloads: %w", err)
+	}
+
+	groupManager, err := groups.NewManager()
+	if err != nil {
+		return fmt.Errorf("failed to create group manager: %w", err)
+	}
+
+	if len(groupNames) > 0 {
+		return removeClientFromGroups(ctx, clientToRemove, groupNames, runningWorkloads, groupManager, clientManager)
+	}
+
+	return removeClientGlobally(ctx, clientToRemove, runningWorkloads, groupManager, clientManager)
+}
+
+func removeClientFromGroups(
+	ctx context.Context,
+	clientToRemove client.Client,
+	groupNames []string,
+	runningWorkloads []core.Workload,
+	groupManager groups.Manager,
+	clientManager client.Manager,
+) error {
+	fmt.Printf("Filtering workloads to groups: %v\n", groupNames)
+
+	// Remove client from specific groups only
+	filteredWorkloads, err := workloads.FilterByGroups(runningWorkloads, groupNames)
+	if err != nil {
+		return fmt.Errorf("failed to filter workloads by groups: %w", err)
+	}
+
+	// Remove the workloads from the client's configuration file
+	err = clientManager.UnregisterClients(ctx, []client.Client{clientToRemove}, filteredWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister client: %w", err)
+	}
+
+	// Remove the client from the groups
+	err = groupManager.UnregisterClients(ctx, groupNames, []string{string(clientToRemove.Name)})
+	if err != nil {
+		return fmt.Errorf("failed to unregister client from groups: %w", err)
+	}
+
+	fmt.Printf("Successfully removed client %s from groups: %v\n", clientToRemove.Name, groupNames)
+
+	return nil
+}
+
+func removeClientGlobally(
+	ctx context.Context,
+	clientToRemove client.Client,
+	runningWorkloads []core.Workload,
+	groupManager groups.Manager,
+	clientManager client.Manager,
+) error {
+	// Remove the workloads from the client's configuration file
+	err := clientManager.UnregisterClients(ctx, []client.Client{clientToRemove}, runningWorkloads)
+	if err != nil {
+		return fmt.Errorf("failed to unregister client: %w", err)
+	}
+
+	allGroups, err := groupManager.List(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to list groups: %w", err)
+	}
+
+	if len(allGroups) > 0 {
+		// Remove client from all groups first
+		allGroupNames := make([]string, len(allGroups))
+		for i, group := range allGroups {
+			allGroupNames[i] = group.Name
+		}
+
+		err = groupManager.UnregisterClients(ctx, allGroupNames, []string{string(clientToRemove.Name)})
+		if err != nil {
+			return fmt.Errorf("failed to unregister client from groups: %w", err)
+		}
+	}
+
+	// Remove client from global registered clients list
+	err = config.UpdateConfig(func(c *config.Config) {
+		for i, registeredClient := range c.Clients.RegisteredClients {
+			if registeredClient == string(clientToRemove.Name) {
+				// Remove client from slice
+				c.Clients.RegisteredClients = append(c.Clients.RegisteredClients[:i], c.Clients.RegisteredClients[i+1:]...)
+				logger.Infof("Successfully unregistered client: %s\n", clientToRemove.Name)
+				return
+			}
+		}
+		logger.Warnf("Client %s was not found in registered clients list", clientToRemove.Name)
+	})
+	if err != nil {
+		return fmt.Errorf("failed to update configuration for client %s: %w", clientToRemove.Name, err)
+	}
+
+	return nil
+}

cmd/thv/app/mcp.go

@@ -35,6 +35,9 @@ func newMCPCommand() *cobra.Command {
 		Long:  `The mcp command provides subcommands to interact with MCP (Model Context Protocol) servers for debugging purposes.`,
 	}
 
+	// Add serve subcommand
+	cmd.AddCommand(newMCPServeCommand())
+
 	// Create list command
 	listCmd := &cobra.Command{
 		Use:   "list [tools|resources|prompts]",