feat(cli): add experimental rpty command

coder · johnstcn · Feb 26, 2025 · Feb 25, 2025 · Feb 25, 2025 · Feb 25, 2025
commit a9ff0456991b57e93e0b4e5a7e9b3bc984bc09b6
diff --git a/cli/exp.go b/cli/exp.go
@@ -14,6 +14,7 @@ func (r *RootCmd) expCmd() *serpent.Command {
 			r.scaletestCmd(),
 			r.errorExample(),
 			r.promptExample(),
+			r.rptyCommand(),
 		},
 	}
 	return cmd

diff --git a/cli/exp_rpty.go b/cli/exp_rpty.go
@@ -0,0 +1,217 @@
+package cli
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/google/uuid"
+	"github.com/mattn/go-isatty"
+	"golang.org/x/term"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+	"github.com/coder/coder/v2/pty"
+	"github.com/coder/serpent"
+)
+
+func (r *RootCmd) rptyCommand() *serpent.Command {
+	var (
+		client = new(codersdk.Client)
+		args   handleRPTYArgs
+	)
+
+	cmd := &serpent.Command{
+		Handler: func(inv *serpent.Invocation) error {
+			if r.disableDirect {
+				return xerrors.New("direct connections are disabled, but you can try websocat ;-)")
+			}
+			args.NamedWorkspace = inv.Args[0]
+			args.Command = inv.Args[1:]
+			return handleRPTY(inv, client, args)
+		},
+		Long: "Establish an RPTY session with a workspace/agent. This uses the same mechanism as the Web Terminal.",
+		Middleware: serpent.Chain(
+			serpent.RequireRangeArgs(1, -1),
+			r.InitClient(client),
+		),
+		Options: []serpent.Option{
+			{
+				Name:          "container",
+				Description:   "The container name or ID to connect to.",
+				Flag:          "container",
+				FlagShorthand: "c",
+				Default:       "",
+				Value:         serpent.StringOf(&args.Container),
+			},
+			{
+				Name:          "container-user",
+				Description:   "The user to connect as.",
+				Flag:          "container-user",
+				FlagShorthand: "u",
+				Default:       "",
+				Value:         serpent.StringOf(&args.ContainerUser),
+			},
+			{
+				Name:          "reconnect",
+				Description:   "The reconnect ID to use.",
+				Flag:          "reconnect",
+				FlagShorthand: "r",
+				Default:       "",
+				Value:         serpent.StringOf(&args.ReconnectID),
+			},
+		},
+		Short: "Establish an RPTY session with a workspace/agent.",
+		Use:   "rpty",
+	}
+
+	return cmd
+}
+
+type handleRPTYArgs struct {
+	Command        []string
+	Container      string
+	ContainerUser  string
+	NamedWorkspace string
+	ReconnectID    string
+}
+
+func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPTYArgs) error {
+	ctx, cancel := context.WithCancel(inv.Context())
+	defer cancel()
+
+	var reconnectID uuid.UUID
+	if args.ReconnectID != "" {
+		rid, err := uuid.Parse(args.ReconnectID)
+		if err != nil {
+			return xerrors.Errorf("invalid reconnect ID: %w", err)
+		}
+		reconnectID = rid
+	} else {
+		reconnectID = uuid.New()
+	}
+	ws, agt, err := getWorkspaceAndAgent(ctx, inv, client, true, args.NamedWorkspace)
+	if err != nil {
+		return err
+	}
+
+	var ctID string
+	if args.Container != "" {
+		cts, err := client.WorkspaceAgentListContainers(ctx, agt.ID, nil)
+		if err != nil {
+			return err
+		}
+		for _, ct := range cts.Containers {
+			if ct.FriendlyName == args.Container || ct.ID == args.Container {
+				ctID = ct.ID
+				break
+			}
+		}
+		if ctID == "" {
+			return xerrors.Errorf("container %q not found", args.Container)
+		}
+	}
+
+	if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{
+		FetchInterval: 0,
+		Fetch:         client.WorkspaceAgent,
+		FetchLogs:     client.WorkspaceAgentLogsAfter,
+		Wait:          false,
+	}); err != nil {
+		return err
+	}
+
+	// Get the width and height of the terminal.
+	var termWidth, termHeight uint16
+	stdoutFile, validOut := inv.Stdout.(*os.File)
+	if validOut && isatty.IsTerminal(stdoutFile.Fd()) {
+		w, h, err := term.GetSize(int(stdoutFile.Fd()))
+		if err == nil {
+			//nolint: gosec
+			termWidth, termHeight = uint16(w), uint16(h)
+		}
+	}
+
+	// Set stdin to raw mode so that control characters work.
+	stdinFile, validIn := inv.Stdin.(*os.File)
+	if validIn && isatty.IsTerminal(stdinFile.Fd()) {
+		inState, err := pty.MakeInputRaw(stdinFile.Fd())
+		if err != nil {
+			return xerrors.Errorf("failed to set input terminal to raw mode: %w", err)
+		}
+		defer func() {
+			_ = pty.RestoreTerminal(stdinFile.Fd(), inState)
+		}()
+	}
+
+	conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{
+		AgentID:       agt.ID,
+		Reconnect:     reconnectID,
+		Command:       strings.Join(args.Command, " "),
+		Container:     ctID,
+		ContainerUser: args.ContainerUser,
+		Width:         termWidth,
+		Height:        termHeight,
+	})
+	if err != nil {
+		return xerrors.Errorf("open reconnecting PTY: %w", err)
+	}
+	defer conn.Close()
+
+	cliui.Infof(inv.Stderr, "Connected to %s (agent id: %s)", args.NamedWorkspace, agt.ID)
+	closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
+		AgentID: agt.ID,
+		AppName: codersdk.UsageAppNameReconnectingPty,
+	})
+	defer closeUsage()
+
+	stdinDone := make(chan struct{})
+	stdoutDone := make(chan struct{})
+	stderrDone := make(chan struct{})
+	done := make(chan struct{})
+
+	go func() {
+		defer close(stdinDone)
+		// This is how we send commands to the agent.
+		br := bufio.NewScanner(inv.Stdin)
+		// Split on bytes, otherwise you have to send a newline to flush the buffer.
+		br.Split(bufio.ScanBytes)
+		je := json.NewEncoder(conn)
+		for br.Scan() {
+			if err := je.Encode(map[string]string{
+				"data": br.Text(),
+			}); err != nil {
+				return
+			}
+		}
+	}()
+	go func() {
+		defer func() {
+			close(stdoutDone)
+		}()
+		_, _ = io.Copy(inv.Stdout, conn)
+	}()
+	go func() {
+		defer func() {
+			close(stderrDone)
+		}()
+		_, _ = io.Copy(inv.Stderr, conn)
+	}()
+	go func() {
+		defer close(done)
+		<-stdoutDone
+		<-stderrDone
+		_ = conn.Close()
+		_, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n")
+	}()
+
+	<-done
+
+	return nil
+}
diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go
@@ -0,0 +1,111 @@
+package cli_test
+
+import (
+	"fmt"
+	"runtime"
+	"testing"
+
+	"github.com/ory/dockertest/v3"
+	"github.com/ory/dockertest/v3/docker"
+
+	"github.com/coder/coder/v2/agent"
+	"github.com/coder/coder/v2/agent/agenttest"
+	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/pty/ptytest"
+	"github.com/coder/coder/v2/testutil"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExpRpty(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		inv, root := clitest.New(t, "exp", "rpty", workspace.Name)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		_ = agenttest.New(t, client.URL, agentToken)
+		_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
+
+		pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name))
+		pty.WriteLine("exit")
+		<-cmdDone
+	})
+
+	t.Run("NotFound", func(t *testing.T) {
+		t.Parallel()
+
+		client, _, _ := setupWorkspaceForAgent(t)
+		inv, root := clitest.New(t, "exp", "rpty", "not-found")
+		clitest.SetupConfig(t, client, root)
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		err := inv.WithContext(ctx).Run()
+		require.ErrorContains(t, err, "not found")
+	})
+
+	t.Run("Container", func(t *testing.T) {
+		t.Parallel()
+		// Skip this test on non-Linux platforms since it requires Docker
+		if runtime.GOOS != "linux" {
+			t.Skip("Skipping test on non-Linux platform")
+		}
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+		pool, err := dockertest.NewPool("")
+		require.NoError(t, err, "Could not connect to docker")
+		ct, err := pool.RunWithOptions(&dockertest.RunOptions{
+			Repository: "busybox",
+			Tag:        "latest",
+			Cmd:        []string{"sleep", "infnity"},
+		}, func(config *docker.HostConfig) {
+			config.AutoRemove = true
+			config.RestartPolicy = docker.RestartPolicy{Name: "no"}
+		})
+		require.NoError(t, err, "Could not start container")
+		// Wait for container to start
+		require.Eventually(t, func() bool {
+			ct, ok := pool.ContainerByName(ct.Container.Name)
+			return ok && ct.Container.State.Running
+		}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
+		t.Cleanup(func() {
+			err := pool.Purge(ct)
+			require.NoError(t, err, "Could not stop container")
+		})
+
+		inv, root := clitest.New(t, "exp", "rpty", workspace.Name, "-c", ct.Container.ID)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
+			o.ExperimentalContainersEnabled = true
+		})
+		_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
+
+		pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name))
+		pty.ExpectMatch(" #")
+		pty.WriteLine("hostname")
+		pty.ExpectMatch(ct.Container.Config.Hostname)
+		pty.WriteLine("exit")
+		<-cmdDone
+	})
+}