refactor(client): make ProxySingleMachineContext and ProxyMachinesContext package functions as well

cmd/uc/ps.go

@@ -187,7 +187,7 @@ func printContainers(containers []containerInfo) error {
 }
 
 func collectContainers(ctx context.Context, cli *client.Client) ([]containerInfo, error) {
-	listCtx := cli.ProxyMachinesContext(ctx, nil)
+	listCtx := client.ProxyMachinesContext(ctx, nil)
 
 	// List all service containers across all machines in the cluster.
 	machineContainers, err := cli.Docker.ListServiceContainers(

pkg/client/client.go

@@ -78,8 +78,8 @@ func (cli *Client) progressOut() *streams.Out {
 // ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines.
 // If namesOrIDs is nil or empty, all machines are included.
 // This triggers One2Many proxying, which always injects metadata into the response.
-func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
+func ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
-	md := metadata.New(nil)
+	md := outgoingMetadataWithoutProxyTargets(ctx)
 	if len(namesOrIDs) == 0 {
 		md.Append("machines", "*")
 	} else {
@@ -92,7 +92,26 @@ func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string
 // ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine.
 // This triggers One2One proxying, which does NOT inject metadata into the response.
 // Use this for requests that expect a single response message without metadata wrapper.
-func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
+func ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
-	md := metadata.Pairs("machine", nameOrID)
+	md := outgoingMetadataWithoutProxyTargets(ctx)
+	md.Set("machine", nameOrID)
 	return metadata.NewOutgoingContext(ctx, md)
 }
+
+func outgoingMetadataWithoutProxyTargets(ctx context.Context) metadata.MD {
+	md, _ := metadata.FromOutgoingContext(ctx)
+	md = md.Copy()
+	md.Delete("machine")
+	md.Delete("machines")
+	return md
+}
+
+// ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines.
+func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context {
+	return ProxyMachinesContext(ctx, namesOrIDs)
+}
+
+// ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine.
+func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context {
+	return ProxySingleMachineContext(ctx, nameOrID)
+}

pkg/client/client_test.go

@@ -0,0 +1,79 @@
+package client
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestProxySingleMachineContext(t *testing.T) {
+	original := metadata.Pairs(
+		"authorization", "token",
+		"machine", "old-machine",
+		"machines", "old-machine-a",
+		"machines", "old-machine-b",
+	)
+	ctx := metadata.NewOutgoingContext(context.Background(), original)
+
+	proxyCtx := ProxySingleMachineContext(ctx, "new-machine")
+
+	md, ok := metadata.FromOutgoingContext(proxyCtx)
+	require.True(t, ok)
+	require.Equal(t, metadata.Pairs(
+		"authorization", "token",
+		"machine", "new-machine",
+	), md)
+	require.Equal(t, metadata.Pairs(
+		"authorization", "token",
+		"machine", "old-machine",
+		"machines", "old-machine-a",
+		"machines", "old-machine-b",
+	), original)
+}
+
+func TestProxyMachinesContext(t *testing.T) {
+	tests := []struct {
+		name     string
+		machines []string
+		want     []string
+	}{
+		{
+			name:     "specified machines",
+			machines: []string{"machine-a", "machine-b"},
+			want:     []string{"machine-a", "machine-b"},
+		},
+		{
+			name: "all machines",
+			want: []string{"*"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			original := metadata.Pairs(
+				"authorization", "token",
+				"machine", "old-machine",
+				"machines", "old-machine-a",
+				"machines", "old-machine-b",
+			)
+			ctx := metadata.NewOutgoingContext(context.Background(), original)
+
+			proxyCtx := ProxyMachinesContext(ctx, tt.machines)
+
+			md, ok := metadata.FromOutgoingContext(proxyCtx)
+			require.True(t, ok)
+			require.Equal(t, metadata.MD{
+				"authorization": {"token"},
+				"machines":      tt.want,
+			}, md)
+			require.Equal(t, metadata.Pairs(
+				"authorization", "token",
+				"machine", "old-machine",
+				"machines", "old-machine-a",
+				"machines", "old-machine-b",
+			), original)
+		})
+	}
+}

pkg/client/container.go

@@ -74,7 +74,7 @@ func (cli *Client) createServiceContainerWithPull(
 	resp.Name = containerName
 
 	// Proxy Docker gRPC requests to the selected machine.
-	ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
+	ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
 
 	pw := progress.ContextWriter(ctx)
 	eventID := cliprogress.NewContainerEventID(ctx, containerName, machine.Machine.Name)
@@ -277,7 +277,7 @@ func (cli *Client) resolveContainerOperation(
 
 	eventID := cliprogress.ContainerEventID(ctx, ctr.Container.ServiceSpec.Name, ctr.Container.ID, ctr.MachineName)
 	return containerOperationContext{
-		ctx:         cli.ProxySingleMachineContext(ctx, ctr.MachineID),
+		ctx:         ProxySingleMachineContext(ctx, ctr.MachineID),
 		containerID: ctr.Container.ID,
 		eventID:     eventID,
 	}, nil
@@ -375,7 +375,7 @@ func (cli *Client) ExecContainer(
 	}
 
 	// Proxy Docker gRPC requests to the machine hosting the container
-	ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
+	ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
 
 	// Execute the command in the container
 	exitCode, err := cli.Docker.ExecContainer(ctx, machinedocker.ExecConfig{
@@ -452,7 +452,7 @@ func (cli *Client) WaitContainerHealthy(
 	}
 
 	// For containers with a health check, wait until Docker reports healthy or unhealthy.
-	mctx := cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
+	mctx := ProxySingleMachineContext(ctx, machine.Machine.Id)
 	mctx, cancel := context.WithTimeout(mctx, healthcheckTimeout(mc.Container.Config.Healthcheck))
 	defer cancel()
 	ticker := time.NewTicker(1 * time.Second)

pkg/client/image.go

@@ -59,7 +59,7 @@ func (cli *Client) InspectRemoteImage(ctx context.Context, id string) ([]api.Mac
 // it lists images on all machines.
 func (cli *Client) ListImages(ctx context.Context, filter api.ImageFilter) ([]api.MachineImages, error) {
 	// Broadcast the image list request to the specified machines or all machines if none specified.
-	listCtx := cli.ProxyMachinesContext(ctx, filter.Machines)
+	listCtx := ProxyMachinesContext(ctx, filter.Machines)
 
 	opts := image.ListOptions{Manifests: true}
 	if filter.Name != "" {

pkg/client/logs.go

@@ -102,7 +102,7 @@ func (cli *Client) ServiceLogs(
 func (cli *Client) ContainerLogs(
 	ctx context.Context, machineNameOrID string, containerID string, opts api.ServiceLogsOptions,
 ) (<-chan api.LogEntry, error) {
-	proxyCtx := cli.ProxySingleMachineContext(ctx, machineNameOrID)
+	proxyCtx := ProxySingleMachineContext(ctx, machineNameOrID)
 
 	req := &pb.LogsRequest{
 		Id:     containerID,
@@ -198,7 +198,7 @@ func (cli *Client) MachineLogs(
 func (cli *Client) systemServiceLogs(
 	ctx context.Context, machineID, service string, opts api.ServiceLogsOptions,
 ) (<-chan api.LogEntry, error) {
-	proxyCtx := cli.ProxySingleMachineContext(ctx, machineID)
+	proxyCtx := ProxySingleMachineContext(ctx, machineID)
 
 	req := &pb.LogsRequest{
 		Id:     service,

pkg/client/machine.go

@@ -78,7 +78,7 @@ func (cli *Client) ListMachines(ctx context.Context, filter *api.MachineFilter)
 func (cli *Client) UpdateMachine(
 	ctx context.Context, nameOrID string, req *pb.UpdateMachineRequest,
 ) (*pb.MachineInfo, error) {
-	ctx = cli.ProxySingleMachineContext(ctx, nameOrID)
+	ctx = ProxySingleMachineContext(ctx, nameOrID)
 	resp, err := cli.MachineClient.UpdateMachine(ctx, req)
 	if err != nil {
 		if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {

pkg/client/volume.go

@@ -27,7 +27,7 @@ func (cli *Client) CreateVolume(
 		return resp, fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err)
 	}
 	// Proxy Docker gRPC requests to the selected machine.
-	ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
+	ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
 
 	pw := progress.ContextWriter(ctx)
 	eventID := cliprogress.VolumeEventID(opts.Name, machine.Machine.Name)
@@ -56,7 +56,7 @@ func (cli *Client) ListVolumes(ctx context.Context, filter *api.VolumeFilter) ([
 		proxyMachines = filter.Machines
 	}
 
-	listCtx := cli.ProxyMachinesContext(ctx, proxyMachines)
+	listCtx := ProxyMachinesContext(ctx, proxyMachines)
 	machineVolumes, err := cli.Docker.ListVolumes(listCtx, volume.ListOptions{})
 	if err != nil {
 		return nil, err
@@ -107,7 +107,7 @@ func (cli *Client) RemoveVolume(ctx context.Context, machineNameOrID, volumeName
 		return fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err)
 	}
 	// Proxy Docker gRPC requests to the selected machine.
-	ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id)
+	ctx = ProxySingleMachineContext(ctx, machine.Machine.Id)
 
 	pw := progress.ContextWriter(ctx)
 	eventID := cliprogress.VolumeEventID(volumeName, machine.Machine.Name)