Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 20, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dea368b
feat(agent/reconnectingpty): allow selecting backend type
johnstcn Mar 18, 2025
f3595c3
fixup! feat(agent/reconnectingpty): allow selecting backend type
johnstcn Mar 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions agent/reconnectingpty/reconnectingpty.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: can we call this plain? Otherwise it feels like we're leaking implementation details.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll do this as a follow-up refactor 👍

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))

Expand Down
5 changes: 3 additions & 2 deletions agent/reconnectingpty/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
)

Expand Down
39 changes: 27 additions & 12 deletions cli/exp_rpty.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"bufio"
"context"
"encoding/json"
"fmt"
"io"
"os"
"strings"
Expand All @@ -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"
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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"
}
Copy link
Member

@mafredri mafredri Mar 20, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: It'd be nice to see this as a CLI flag too, to (auto, plain, screen).

Side-note: I guess screen could even be called multiplexer to allow to swap out for e.g. tmux in the future.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I considered having it a CLI flag, but hesitated exposing this internal implementation detail. It does make sense, especially with your vim comment below.


conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{
AgentID: agt.ID,
Reconnect: reconnectID,
Expand All @@ -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,
Expand Down Expand Up @@ -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
}
Copy link
Member

@mafredri mafredri Mar 20, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has some interesting repercussions when it's neither a shell and not a one-shot command. Imagine doing vim /my/file.txt. Here we would return true but vim is one of the reasons we have the screen backend.

This is probably fine though, but I'm flagging this behavior as it's not ideal. At the very least worth a comment.

(The flag would alleviate this as you can enforce the behavior then.)

30 changes: 25 additions & 5 deletions cli/exp_rpty_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"

Expand All @@ -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)
Expand All @@ -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()

Expand Down Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions coderd/workspaceapps/proxy.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down Expand Up @@ -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))
Expand Down
2 changes: 2 additions & 0 deletions codersdk/workspacesdk/agentconn.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 8 additions & 0 deletions codersdk/workspacesdk/workspacesdk.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand Down
Loading