Remove group label from workload (#1173)

Author: amirejaz

cmd/thv/app/run.go

@@ -80,10 +80,9 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	// Get debug mode flag
 	debugMode, _ := cmd.Flags().GetBool("debug")
 
-	// Build the run configuration
-	runnerConfig, err := BuildRunnerConfig(ctx, &runFlags, serverOrImage, cmdArgs, debugMode, cmd)
-	if err != nil {
-		return err
+	workloadName := runFlags.Name
+	if workloadName == "" {
+		workloadName = serverOrImage
 	}
 
 	if runFlags.Group != "" {
@@ -93,12 +92,12 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		}
 
 		// Check if the workload is already in a group
-		group, err := groupManager.GetWorkloadGroup(ctx, runFlags.Name)
+		group, err := groupManager.GetWorkloadGroup(ctx, workloadName)
 		if err != nil {
 			return fmt.Errorf("failed to get workload group: %v", err)
 		}
 		if group != nil && group.Name != runFlags.Group {
-			return fmt.Errorf("workload '%s' is already in group '%s'", runFlags.Name, group.Name)
+			return fmt.Errorf("workload '%s' is already in group '%s'", workloadName, group.Name)
 		}
 
 		// Validate that the group specified exists
@@ -111,6 +110,12 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	// Build the run configuration
+	runnerConfig, err := BuildRunnerConfig(ctx, &runFlags, serverOrImage, cmdArgs, debugMode, cmd)
+	if err != nil {
+		return err
+	}
+
 	// Create container runtime
 	rt, err := container.NewFactory().Create(ctx)
 	if err != nil {

pkg/errors/errors.go

@@ -25,6 +25,9 @@ const (
 	// ErrContainerAlreadyRunning is returned when a container is already running
 	ErrContainerAlreadyRunning = "container_already_running"
 
+	// ErrRunConfigNotFound is returned when a run configuration is not found
+	ErrRunConfigNotFound = "run_config_not_found"
+
 	// ErrTransport is returned when there is an error with the transport
 	ErrTransport = "transport"
 
@@ -99,6 +102,11 @@ func NewContainerAlreadyRunningError(message string, cause error) *Error {
 	return NewError(ErrContainerAlreadyRunning, message, cause)
 }
 
+// NewRunConfigNotFoundError creates a new run configuration not found error
+func NewRunConfigNotFoundError(message string, cause error) *Error {
+	return NewError(ErrRunConfigNotFound, message, cause)
+}
+
 // NewTransportError creates a new transport error
 func NewTransportError(message string, cause error) *Error {
 	return NewError(ErrTransport, message, cause)
@@ -150,6 +158,12 @@ func IsContainerAlreadyRunning(err error) bool {
 	return ok && e.Type == ErrContainerAlreadyRunning
 }
 
+// IsRunConfigNotFound checks if the error is a run configuration not found error
+func IsRunConfigNotFound(err error) bool {
+	e, ok := err.(*Error)
+	return ok && e.Type == ErrRunConfigNotFound
+}
+
 // IsTransport checks if the error is a transport error
 func IsTransport(err error) bool {
 	e, ok := err.(*Error)

pkg/groups/group_test.go

@@ -13,8 +13,6 @@ import (
 
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/state/mocks"
-	"github.com/stacklok/toolhive/pkg/workloads"
-	workloadmocks "github.com/stacklok/toolhive/pkg/workloads/mocks"
 )
 
 func init() {
@@ -103,7 +101,7 @@ func TestManager_Create(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore, workloadManager: nil}
+			manager := &manager{store: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -178,7 +176,7 @@ func TestManager_Get(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore, workloadManager: nil}
+			manager := &manager{store: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -284,7 +282,7 @@ func TestManager_List(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore, workloadManager: nil}
+			manager := &manager{store: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -370,7 +368,7 @@ func TestManager_Delete(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore, workloadManager: nil}
+			manager := &manager{store: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -444,7 +442,7 @@ func TestManager_Exists(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			manager := &manager{store: mockStore, workloadManager: nil}
+			manager := &manager{store: mockStore}
 
 			// Set up mock expectations
 			tt.setupMock(mockStore)
@@ -469,66 +467,18 @@ func TestManager_GetWorkloadGroup(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name                 string
-		workloadName         string
-		setupMock            func(*mocks.MockStore)
-		setupWorkloadManager func(*workloadmocks.MockManager)
-		workloadGroup        string // The group the workload belongs to (empty if none)
-		expectError          bool
-		expectedGroup        *Group
-		errorMsg             string
+		name          string
+		workloadName  string
+		expectError   bool
+		expectedGroup *Group
+		errorMsg      string
 	}{
 		{
-			name:          "workload found in group",
-			workloadName:  "test-workload",
-			workloadGroup: testGroupName,
-			setupMock: func(mock *mocks.MockStore) {
-				// Mock GetReader for the group
-				mock.EXPECT().
-					GetReader(gomock.Any(), testGroupName).
-					Return(&mockReadCloser{data: `{"name":"` + testGroupName + `"}`}, nil)
-			},
-			setupWorkloadManager: func(mock *workloadmocks.MockManager) {
-				// Mock GetWorkload to return a workload in a group
-				mock.EXPECT().
-					GetWorkload(gomock.Any(), "test-workload").
-					Return(workloads.Workload{Name: "test-workload", Group: testGroupName}, nil)
-			},
-			expectError:   false,
-			expectedGroup: &Group{Name: testGroupName},
-		},
-		{
-			name:          "workload not found in any group",
+			name:          "workload not found",
 			workloadName:  "nonexistent-workload",
-			workloadGroup: "",
-			setupMock: func(_ *mocks.MockStore) {
-				// No mock expectations needed since workload has no group
-			},
-			setupWorkloadManager: func(mock *workloadmocks.MockManager) {
-				// Mock GetWorkload to return a workload with no group
-				mock.EXPECT().
-					GetWorkload(gomock.Any(), "nonexistent-workload").
-					Return(workloads.Workload{Name: "nonexistent-workload", Group: ""}, nil)
-			},
 			expectError:   false,
 			expectedGroup: nil,
 		},
-		{
-			name:          "workload manager fails to get workload",
-			workloadName:  "test-workload",
-			workloadGroup: "",
-			setupMock: func(_ *mocks.MockStore) {
-				// No mock expectations needed since workload manager will fail
-			},
-			setupWorkloadManager: func(mock *workloadmocks.MockManager) {
-				// Mock GetWorkload to fail
-				mock.EXPECT().
-					GetWorkload(gomock.Any(), "test-workload").
-					Return(workloads.Workload{}, errors.New("workload not found"))
-			},
-			expectError:   false,
-			expectedGroup: nil, // When workload manager fails, we return nil (not in any group)
-		},
 	}
 
 	for _, tt := range tests {
@@ -539,12 +489,7 @@ func TestManager_GetWorkloadGroup(t *testing.T) {
 			defer ctrl.Finish()
 
 			mockStore := mocks.NewMockStore(ctrl)
-			mockWorkloadManager := workloadmocks.NewMockManager(ctrl)
-			manager := &manager{store: mockStore, workloadManager: mockWorkloadManager}
-
-			// Set up mock expectations
-			tt.setupMock(mockStore)
-			tt.setupWorkloadManager(mockWorkloadManager)
+			manager := &manager{store: mockStore}
 
 			// Call the method
 			group, err := manager.GetWorkloadGroup(context.Background(), tt.workloadName)
@@ -580,22 +525,3 @@ func (m *mockWriteCloser) Write(p []byte) (n int, err error) {
 func (*mockWriteCloser) Close() error {
 	return nil
 }
-
-// mockReadCloser implements io.ReadCloser for testing
-type mockReadCloser struct {
-	data string
-	pos  int
-}
-
-func (m *mockReadCloser) Read(p []byte) (n int, err error) {
-	if m.pos >= len(m.data) {
-		return 0, io.EOF
-	}
-	n = copy(p, m.data[m.pos:])
-	m.pos += n
-	return n, nil
-}
-
-func (*mockReadCloser) Close() error {
-	return nil
-}

pkg/groups/manager.go

@@ -5,14 +5,14 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/stacklok/toolhive/pkg/errors"
+	"github.com/stacklok/toolhive/pkg/runner"
 	"github.com/stacklok/toolhive/pkg/state"
-	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
 // manager implements the Manager interface
 type manager struct {
-	store           state.Store
-	workloadManager workloads.Manager
+	store state.Store
 }
 
 // NewManager creates a new group manager
@@ -22,23 +22,7 @@ func NewManager() (Manager, error) {
 		return nil, fmt.Errorf("failed to create group state store: %w", err)
 	}
 
-	workloadManager, err := workloads.NewManager(context.Background())
-	if err != nil {
-		return nil, fmt.Errorf("failed to create workload manager: %w", err)
-	}
-
-	return &manager{store: store, workloadManager: workloadManager}, nil
-}
-
-// NewManagerWithWorkloadManager creates a new group manager with a custom workload manager
-// This is useful for testing with mocks
-func NewManagerWithWorkloadManager(workloadManager workloads.Manager) (Manager, error) {
-	store, err := state.NewGroupConfigStore("toolhive")
-	if err != nil {
-		return nil, fmt.Errorf("failed to create group state store: %w", err)
-	}
-
-	return &manager{store: store, workloadManager: workloadManager}, nil
+	return &manager{store: store}, nil
 }
 
 // Create creates a new group with the given name
@@ -108,22 +92,21 @@ func (m *manager) Exists(ctx context.Context, name string) (bool, error) {
 
 // GetWorkloadGroup returns the group that a workload belongs to, if any
 func (m *manager) GetWorkloadGroup(ctx context.Context, workloadName string) (*Group, error) {
-	// Use the workload manager to get the workload and check its group label
-
-	// Get the workload details
-	workload, err := m.workloadManager.GetWorkload(ctx, workloadName)
+	runnerInstance, err := runner.LoadState(ctx, workloadName)
 	if err != nil {
-		// If workload not found, it's not in any group
-		return nil, nil
+		if errors.IsRunConfigNotFound(err) {
+			return nil, nil
+		}
+		return nil, err
 	}
 
-	// Check if the workload has a group
-	if workload.Group == "" {
-		return nil, nil // Workload is not in any group
+	// If the workload has no group, return nil
+	if runnerInstance.Config.Group == "" {
+		return nil, nil
 	}
 
 	// Get the group details
-	return m.Get(ctx, workload.Group)
+	return m.Get(ctx, runnerInstance.Config.Group)
 }
 
 // saveGroup saves the group to the group state store

pkg/labels/labels.go

@@ -40,9 +40,8 @@ const (
 	LabelToolHiveValue = "true"
 )
 
-// AddStandardLabelsWithGroup adds standard labels to a container, including group if provided
-func AddStandardLabelsWithGroup(labels map[string]string, containerName, containerBaseName, transportType string,
-	port int, groupName string) {
+// AddStandardLabels adds standard labels to a container
+func AddStandardLabels(labels map[string]string, containerName, containerBaseName, transportType string, port int) {
 	labels[LabelToolHive] = LabelToolHiveValue
 	labels[LabelName] = containerName
 	labels[LabelBaseName] = containerBaseName
@@ -51,14 +50,6 @@ func AddStandardLabelsWithGroup(labels map[string]string, containerName, contain
 
 	// TODO: In the future, we'll support different tool types beyond just "mcp"
 	labels[LabelToolType] = "mcp"
-	if groupName != "" {
-		labels[LabelGroup] = groupName
-	}
-}
-
-// AddStandardLabels adds standard labels to a container
-func AddStandardLabels(labels map[string]string, containerName, containerBaseName, transportType string, port int) {
-	AddStandardLabelsWithGroup(labels, containerName, containerBaseName, transportType, port, "")
 }
 
 // AddNetworkLabels adds network-related labels to a network
@@ -148,7 +139,6 @@ func IsStandardToolHiveLabel(key string) bool {
 		LabelPort,
 		LabelToolType,
 		LabelNetworkIsolation,
-		LabelGroup,
 	}
 
 	for _, standardLabel := range standardLabels {