fixup! feat(cli): add capability for SSH command to connect to a running container

johnstcn · johnstcn · commit 56538a975f8b · 2025-02-27T15:21:48.000Z
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -61,6 +61,14 @@ const (
 	// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
 	// This is stripped from any commands being executed, and is counted towards connection stats.
 	MagicSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
+	// ContainerEnvironmentVariable is used to specify the target container for an SSH connection.
+	// This is stripped from any commands being executed.
+	// Only available if CODER_AGENT_DEVCONTAINERS_ENABLE=true.
+	ContainerEnvironmentVariable = "CODER_CONTAINER"
+	// ContainerUserEnvironmentVariable is used to specify the container user for
+	// an SSH connection.
+	// Only available if CODER_AGENT_DEVCONTAINERS_ENABLE=true.
+	ContainerUserEnvironmentVariable = "CODER_CONTAINER_USER"
 )
 
 // MagicSessionType enums.
@@ -330,17 +338,17 @@ func (s *sessionCloseTracker) Close() error {
 
 func extractContainerInfo(env []string) (container, containerUser string, filteredEnv []string) {
 	for _, kv := range env {
-		if strings.HasPrefix(kv, "CODER_CONTAINER=") {
-			container = strings.TrimPrefix(kv, "CODER_CONTAINER=")
+		if strings.HasPrefix(kv, ContainerEnvironmentVariable+"=") {
+			container = strings.TrimPrefix(kv, ContainerEnvironmentVariable+"=")
 		}
 
-		if strings.HasPrefix(kv, "CODER_CONTAINER_USER=") {
-			containerUser = strings.TrimPrefix(kv, "CODER_CONTAINER_USER=")
+		if strings.HasPrefix(kv, ContainerUserEnvironmentVariable+"=") {
+			containerUser = strings.TrimPrefix(kv, ContainerUserEnvironmentVariable+"=")
 		}
 	}
 
 	return container, containerUser, slices.DeleteFunc(env, func(kv string) bool {
-		return strings.HasPrefix(kv, "CODER_CONTAINER=") || strings.HasPrefix(kv, "CODER_CONTAINER_USER=")
+		return strings.HasPrefix(kv, ContainerEnvironmentVariable+"=") || strings.HasPrefix(kv, ContainerUserEnvironmentVariable+"=")
 	})
 }
 
@@ -536,7 +544,6 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, env []str
 		ptyLabel = "yes"
 	}
 
-	// plumb in envinfoer here to modify command for container exec?
 	var ei usershell.EnvInfoer
 	var err error
 	if s.config.ExperimentalContainersEnabled && container != "" {
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -34,6 +34,7 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
+	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/cli/cliutil"
 	"github.com/coder/coder/v2/coderd/autobuild/notify"
@@ -476,8 +477,8 @@ func (r *RootCmd) ssh() *serpent.Command {
 
 			if container != "" {
 				for k, v := range map[string]string{
-					"CODER_CONTAINER":      container,
-					"CODER_CONTAINER_USER": containerUser,
+					agentssh.ContainerEnvironmentVariable:     container,
+					agentssh.ContainerUserEnvironmentVariable: containerUser,
 				} {
 					if err := sshSession.Setenv(k, v); err != nil {
 						return xerrors.Errorf("setenv: %w", err)
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -1994,6 +1994,20 @@ func TestSSH_Container(t *testing.T) {
 		err := inv.WithContext(ctx).Run()
 		require.ErrorContains(t, err, "container not found:")
 	})
+
+	t.Run("NotEnabled", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		_ = agenttest.New(t, client.URL, agentToken)
+		_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
+
+		inv, root := clitest.New(t, "ssh", workspace.Name, "-c", uuid.NewString())
+		clitest.SetupConfig(t, client, root)
+		err := inv.WithContext(ctx).Run()
+		require.ErrorContains(t, err, "container not found:")
+	})
 }
 
 // tGoContext runs fn in a goroutine passing a context that will be
