Don't attempt to remove proxy containers when network isolation is off (#994)

This extends the work started in #967 by explicitly setting the network isolation label to false when a new workload is created with network isolation disabled.

Workloads created before this change will not have the label, and in order to ensure that it is cleaned up properly - we treat all of those legacy workloads as having network isolation enabled. To minimize the impact of this, we log the proxy container not found errors at debug level.

Author: dmjb

pkg/container/docker/client.go

@@ -466,11 +466,9 @@ func (c *Client) DeployWorkload(
 		return "", 0, fmt.Errorf("failed to create external networks: %v", err)
 	}
 
+	networkIsolation := false
 	if isolateNetwork {
-		// Add a label to the MCP server indicating network isolation.
-		// This allows other methods to determine whether it needs to care
-		// about ingress/egress/dns containers.
-		lb.AddNetworkIsolationLabel(labels)
+		networkIsolation = true
 
 		internalNetworkLabels := map[string]string{}
 		lb.AddNetworkLabels(internalNetworkLabels, networkName)
@@ -516,6 +514,11 @@ func (c *Client) DeployWorkload(
 		return "", 0, fmt.Errorf("failed to generate port bindings: %v", err)
 	}
 
+	// Add a label to the MCP server indicating network isolation.
+	// This allows other methods to determine whether it needs to care
+	// about ingress/egress/dns containers.
+	lb.AddNetworkIsolationLabel(labels, networkIsolation)
+
 	containerId, err := c.createMcpContainer(
 		ctx,
 		name,
@@ -639,15 +642,12 @@ func (c *Client) StopWorkload(ctx context.Context, workloadID string) error {
 	}
 
 	// If network isolation is not enabled, then there is nothing else to do.
-	// TODO: This check is currently commented out because we need to ensure
-	// that workloads created by older versions of ToolHive get cleaned up
-	// properly. Once we are confident that all workloads with network isolation
-	// have the label set, we can uncomment this check.
-	/*
-		if !lb.HasNetworkIsolation(info.Labels) {
-			return nil
-		}
-	*/
+	// NOTE: This check treats all workloads created before the introduction of
+	// this label as having network isolation enabled. This is to ensure that they
+	// get cleaned up properly during stop/rm.
+	if !lb.HasNetworkIsolation(info.Labels) {
+		return nil
+	}
 
 	// remove / from container name
 	containerName := strings.TrimPrefix(info.Name, "/")
@@ -668,11 +668,11 @@ func (c *Client) StopWorkload(ctx context.Context, workloadID string) error {
 func (c *Client) stopProxyContainer(ctx context.Context, containerName string, timeoutSeconds int) {
 	containerId, err := c.findExistingContainer(ctx, containerName)
 	if err != nil {
-		logger.Warnf("Failed to find internal container %s: %v", containerName, err)
+		logger.Debugf("Failed to find internal container %s: %v", containerName, err)
 	} else {
 		err = c.client.ContainerStop(ctx, containerId, container.StopOptions{Timeout: &timeoutSeconds})
 		if err != nil {
-			logger.Warnf("Failed to stop internal container %s: %v", containerName, err)
+			logger.Debugf("Failed to stop internal container %s: %v", containerName, err)
 		}
 	}
 }
@@ -729,16 +729,13 @@ func (c *Client) RemoveWorkload(ctx context.Context, workloadID string) error {
 	}
 
 	// If network isolation is not enabled, then there is nothing else to do.
-	// TODO: This check is currently commented out because we need to ensure
-	// that workloads created by older versions of ToolHive get cleaned up
-	// properly. Once we are confident that all workloads with network isolation
-	// have the label set, we can uncomment this check.
-	// If network isolation is not enabled, then there is nothing else to do.
-	/*
-		if containerResponse.Config != nil && !lb.HasNetworkIsolation(containerResponse.Config.Labels) {
-			return nil
-		}
-	*/
+	// NOTE: This check treats all workloads created before the introduction of
+	// this label as having network isolation enabled. This is to ensure that they
+	// get cleaned up properly during stop/rm. There may be some spurious warnings
+	// from the following code, but they can be ignored.
+	if containerResponse.Config != nil && !lb.HasNetworkIsolation(containerResponse.Config.Labels) {
+		return nil
+	}
 
 	// remove egress, ingress, and dns containers
 	suffixes := []string{"egress", "ingress", "dns"}
@@ -747,7 +744,7 @@ func (c *Client) RemoveWorkload(ctx context.Context, workloadID string) error {
 		containerName := fmt.Sprintf("%s-%s", containerName, suffix)
 		containerId, err := c.findExistingContainer(ctx, containerName)
 		if err != nil {
-			logger.Warnf("Failed to find %s container %s: %v", suffix, containerName, err)
+			logger.Debugf("Failed to find %s container %s: %v", suffix, containerName, err)
 			continue
 		}
 		if containerId == "" {

pkg/labels/labels.go

@@ -4,6 +4,7 @@ package labels
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
 )
 
@@ -34,9 +35,6 @@ const (
 
 	// LabelToolHiveValue is the value for the LabelToolHive label
 	LabelToolHiveValue = "true"
-
-	// LabelNetworkIsolationValue is the value for the LabelNetworkIsolation label
-	LabelNetworkIsolationValue = "true"
 )
 
 // AddStandardLabels adds standard labels to a container
@@ -59,8 +57,8 @@ func AddNetworkLabels(labels map[string]string, networkName string) {
 }
 
 // AddNetworkIsolationLabel adds the network isolation label to a container
-func AddNetworkIsolationLabel(labels map[string]string) {
-	labels[LabelNetworkIsolation] = LabelNetworkIsolationValue
+func AddNetworkIsolationLabel(labels map[string]string, networkIsolation bool) {
+	labels[LabelNetworkIsolation] = strconv.FormatBool(networkIsolation)
 }
 
 // FormatToolHiveFilter formats a filter for ToolHive containers
@@ -77,7 +75,10 @@ func IsToolHiveContainer(labels map[string]string) bool {
 // HasNetworkIsolation checks if a container has network isolation enabled.
 func HasNetworkIsolation(labels map[string]string) bool {
 	value, ok := labels[LabelNetworkIsolation]
-	return ok && strings.ToLower(value) == LabelNetworkIsolationValue
+	// If the label is not present, assume that network isolation is enabled.
+	// This is to ensure that workloads created before this label was introduced
+	// will be properly cleaned up during stop/rm.
+	return !ok || strings.ToLower(value) == "true"
 }
 
 // GetContainerName gets the container name from labels

pkg/labels/labels_test.go

@@ -314,15 +314,22 @@ func TestHasNetworkIsolation(t *testing.T) {
 		{
 			name: "Network isolation enabled",
 			labels: map[string]string{
-				LabelNetworkIsolation: LabelNetworkIsolationValue,
+				LabelNetworkIsolation: "true",
 			},
 			expected: true,
 		},
 		{
-			name:     "Network isolation not enabled",
-			labels:   map[string]string{},
+			name: "Network isolation disabled",
+			labels: map[string]string{
+				LabelNetworkIsolation: "false",
+			},
 			expected: false,
 		},
+		{
+			name:     "Legacy workload without label",
+			labels:   map[string]string{},
+			expected: true,
+		},
 	}
 
 	for _, tc := range tests {