feat(agent/reconnectingpty): allow selecting backend type (#17011)

johnstcn · web-flow · commit 68624092a499 · 2025-03-20T13:45:31.000Z
agent/reconnectingpty: allow specifying backend type
cli: exp rpty: automatically select backend based on command
diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go
@@ -32,6 +32,8 @@ type Options struct {
 	Timeout time.Duration
 	// Metrics tracks various error counters.
 	Metrics *prometheus.CounterVec
+	// BackendType specifies the ReconnectingPTY backend to use.
+	BackendType string
 }
 
 // ReconnectingPTY is a pty that can be reconnected within a timeout and to
@@ -64,13 +66,20 @@ func New(ctx context.Context, logger slog.Logger, execer agentexec.Execer, cmd *
 	// runs) but in CI screen often incorrectly claims the session name does not
 	// exist even though screen -list shows it.  For now, restrict screen to
 	// Linux.
-	backendType := "buffered"
+	autoBackendType := "buffered"
 	if runtime.GOOS == "linux" {
 		_, err := exec.LookPath("screen")
 		if err == nil {
-			backendType = "screen"
+			autoBackendType = "screen"
 		}
 	}
+	var backendType string
+	switch options.BackendType {
+	case "":
+		backendType = autoBackendType
+	default:
+		backendType = options.BackendType
+	}
 
 	logger.Info(ctx, "start reconnecting pty", slog.F("backend_type", backendType))
 
diff --git a/agent/reconnectingpty/server.go b/agent/reconnectingpty/server.go
@@ -207,8 +207,9 @@ func (s *Server) handleConn(ctx context.Context, logger slog.Logger, conn net.Co
 			s.commandCreator.Execer,
 			cmd,
 			&Options{
-				Timeout: s.timeout,
-				Metrics: s.errorsTotal,
+				Timeout:     s.timeout,
+				Metrics:     s.errorsTotal,
+				BackendType: msg.BackendType,
 			},
 		)
 
diff --git a/cli/exp_rpty.go b/cli/exp_rpty.go
@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
-	"fmt"
 	"io"
 	"os"
 	"strings"
@@ -15,6 +14,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
 	"github.com/coder/coder/v2/pty"
@@ -96,6 +96,7 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
 	} else {
 		reconnectID = uuid.New()
 	}
+
 	ws, agt, err := getWorkspaceAndAgent(ctx, inv, client, true, args.NamedWorkspace)
 	if err != nil {
 		return err
@@ -118,14 +119,6 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
 		}
 	}
 
-	if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{
-		FetchInterval: 0,
-		Fetch:         client.WorkspaceAgent,
-		Wait:          false,
-	}); err != nil {
-		return err
-	}
-
 	// Get the width and height of the terminal.
 	var termWidth, termHeight uint16
 	stdoutFile, validOut := inv.Stdout.(*os.File)
@@ -149,6 +142,15 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
 		}()
 	}
 
+	// If a user does not specify a command, we'll assume they intend to open an
+	// interactive shell.
+	var backend string
+	if isOneShotCommand(args.Command) {
+		// If the user specified a command, we'll prefer to use the buffered method.
+		// The screen backend is not well suited for one-shot commands.
+		backend = "buffered"
+	}
+
 	conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{
 		AgentID:       agt.ID,
 		Reconnect:     reconnectID,
@@ -157,14 +159,13 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
 		ContainerUser: args.ContainerUser,
 		Width:         termWidth,
 		Height:        termHeight,
+		BackendType:   backend,
 	})
 	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)
-	cliui.Infof(inv.Stderr, "Reconnect ID: %s", reconnectID)
 	closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
 		AgentID: agt.ID,
 		AppName: codersdk.UsageAppNameReconnectingPty,
@@ -210,7 +211,21 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
 	_, _ = io.Copy(inv.Stdout, conn)
 	cancel()
 	_ = conn.Close()
-	_, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n")
 
 	return nil
 }
+
+var knownShells = []string{"ash", "bash", "csh", "dash", "fish", "ksh", "powershell", "pwsh", "zsh"}
+
+func isOneShotCommand(cmd []string) bool {
+	// If the command is empty, we'll assume the user wants to open a shell.
+	if len(cmd) == 0 {
+		return false
+	}
+	// If the command is a single word, and that word is a known shell, we'll
+	// assume the user wants to open a shell.
+	if len(cmd) == 1 && slice.Contains(knownShells, cmd[0]) {
+		return false
+	}
+	return true
+}
diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go
@@ -1,10 +1,10 @@
 package cli_test
 
 import (
-	"fmt"
 	"runtime"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/ory/dockertest/v3"
 	"github.com/ory/dockertest/v3/docker"
 
@@ -23,7 +23,7 @@ import (
 func TestExpRpty(t *testing.T) {
 	t.Parallel()
 
-	t.Run("OK", func(t *testing.T) {
+	t.Run("DefaultCommand", func(t *testing.T) {
 		t.Parallel()
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
@@ -41,11 +41,33 @@ func TestExpRpty(t *testing.T) {
 		_ = 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("Command", func(t *testing.T) {
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		randStr := uuid.NewString()
+		inv, root := clitest.New(t, "exp", "rpty", workspace.Name, "echo", randStr)
+		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(randStr)
+		<-cmdDone
+	})
+
 	t.Run("NotFound", func(t *testing.T) {
 		t.Parallel()
 
@@ -103,8 +125,6 @@ func TestExpRpty(t *testing.T) {
 			assert.NoError(t, err)
 		})
 
-		pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name))
-		pty.ExpectMatch("Reconnect ID: ")
 		pty.ExpectMatch(" #")
 		pty.WriteLine("hostname")
 		pty.ExpectMatch(ct.Container.Config.Hostname)
diff --git a/coderd/workspaceapps/proxy.go b/coderd/workspaceapps/proxy.go
@@ -655,6 +655,7 @@ func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 	width := parser.UInt(values, 80, "width")
 	container := parser.String(values, "", "container")
 	containerUser := parser.String(values, "", "container_user")
+	backendType := parser.String(values, "", "backend_type")
 	if len(parser.Errors) > 0 {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message:     "Invalid query parameters.",
@@ -695,6 +696,7 @@ func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 	ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect, uint16(height), uint16(width), r.URL.Query().Get("command"), func(arp *workspacesdk.AgentReconnectingPTYInit) {
 		arp.Container = container
 		arp.ContainerUser = containerUser
+		arp.BackendType = backendType
 	})
 	if err != nil {
 		log.Debug(ctx, "dial reconnecting pty server in workspace agent", slog.Error(err))
diff --git a/codersdk/workspacesdk/agentconn.go b/codersdk/workspacesdk/agentconn.go
@@ -100,6 +100,8 @@ type AgentReconnectingPTYInit struct {
 	// This can be a username or UID, depending on the underlying implementation.
 	// This is ignored if Container is not set.
 	ContainerUser string
+
+	BackendType string
 }
 
 // AgentReconnectingPTYInitOption is a functional option for AgentReconnectingPTYInit.
diff --git a/codersdk/workspacesdk/workspacesdk.go b/codersdk/workspacesdk/workspacesdk.go
@@ -318,6 +318,11 @@ type WorkspaceAgentReconnectingPTYOpts struct {
 	// CODER_AGENT_DEVCONTAINERS_ENABLE set to "true".
 	Container     string
 	ContainerUser string
+
+	// BackendType is the type of backend to use for the PTY. If not set, the
+	// workspace agent will attempt to determine the preferred backend type.
+	// Supported values are "screen" and "buffered".
+	BackendType string
 }
 
 // AgentReconnectingPTY spawns a PTY that reconnects using the token provided.
@@ -339,6 +344,9 @@ func (c *Client) AgentReconnectingPTY(ctx context.Context, opts WorkspaceAgentRe
 	if opts.ContainerUser != "" {
 		q.Set("container_user", opts.ContainerUser)
 	}
+	if opts.BackendType != "" {
+		q.Set("backend_type", opts.BackendType)
+	}
 	// If we're using a signed token, set the query parameter.
 	if opts.SignedToken != "" {
 		q.Set(codersdk.SignedAppTokenQueryParameter, opts.SignedToken)

Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,8 @@ type AgentReconnectingPTYInit struct {`
`100`	`100`	`// This can be a username or UID, depending on the underlying implementation.`
`101`	`101`	`// This is ignored if Container is not set.`
`102`	`102`	`ContainerUser string`
	`103`	`+`
	`104`	`+ BackendType string`
`103`	`105`	`}`
`104`	`106`
`105`	`107`	`// AgentReconnectingPTYInitOption is a functional option for AgentReconnectingPTYInit.`