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

agent/reconnectingpty: allow specifying backend type
cli: exp rpty: automatically select backend based on command

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
+}

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)