From dea368b5e7e2d813f9b5325b41bf0d82a7ad4608 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 18 Mar 2025 20:59:48 +0000 Subject: [PATCH 1/2] feat(agent/reconnectingpty): allow selecting backend type --- agent/reconnectingpty/reconnectingpty.go | 13 ++++++-- agent/reconnectingpty/server.go | 5 +-- cli/exp_rpty.go | 39 ++++++++++++++++-------- coderd/workspaceapps/proxy.go | 2 ++ codersdk/workspacesdk/agentconn.go | 2 ++ codersdk/workspacesdk/workspacesdk.go | 8 +++++ 6 files changed, 53 insertions(+), 16 deletions(-) diff --git a/agent/reconnectingpty/reconnectingpty.go b/agent/reconnectingpty/reconnectingpty.go index b5c4e0aaa0b39..4b5251ef31472 100644 --- 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 index 33ed76a73c60e..04bbdc7efb7b2 100644 --- 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 index ddfdc15ece58d..48074c7ef5fb9 100644 --- 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/coderd/workspaceapps/proxy.go b/coderd/workspaceapps/proxy.go index ab67e6c260349..836279b76191b 100644 --- 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 index ef0c292e010e9..8c4a3c169b564 100644 --- 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 index 08aabe9d5f699..e28579216d526 100644 --- 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) From f3595c3d6a09a1cc688a3407eb5e3f424f85550a Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 19 Mar 2025 21:23:06 +0000 Subject: [PATCH 2/2] fixup! feat(agent/reconnectingpty): allow selecting backend type --- cli/exp_rpty_test.go | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go index bfede8213d4c9..5089796f5ac3a 100644 --- 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)