fix: Add command to reconnecting PTY (#1860)

This fixes #1708 and opens the door for PTYs to execute
non-shell commands!

agent/agent.go

@@ -456,8 +456,8 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
 
 	// The ID format is referenced in conn.go.
 	// <uuid>:<height>:<width>
-	idParts := strings.Split(rawID, ":")
-	if len(idParts) != 3 {
+	idParts := strings.SplitN(rawID, ":", 4)
+	if len(idParts) != 4 {
 		a.logger.Warn(ctx, "client sent invalid id format", slog.F("raw-id", rawID))
 		return
 	}
@@ -489,7 +489,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
 		}
 	} else {
 		// Empty command will default to the users shell!
-		cmd, err := a.createCommand(ctx, "", nil)
+		cmd, err := a.createCommand(ctx, idParts[3], nil)
 		if err != nil {
 			a.logger.Warn(ctx, "create reconnecting pty command", slog.Error(err))
 			return

agent/agent_test.go

@@ -221,7 +221,7 @@ func TestAgent(t *testing.T) {
 
 		conn := setupAgent(t, agent.Metadata{}, 0)
 		id := uuid.NewString()
-		netConn, err := conn.ReconnectingPTY(id, 100, 100)
+		netConn, err := conn.ReconnectingPTY(id, 100, 100, "/bin/bash")
 		require.NoError(t, err)
 		bufRead := bufio.NewReader(netConn)
 
@@ -259,7 +259,7 @@ func TestAgent(t *testing.T) {
 		expectLine(matchEchoOutput)
 
 		_ = netConn.Close()
-		netConn, err = conn.ReconnectingPTY(id, 100, 100)
+		netConn, err = conn.ReconnectingPTY(id, 100, 100, "/bin/bash")
 		require.NoError(t, err)
 		bufRead = bufio.NewReader(netConn)
 

agent/conn.go

@@ -34,8 +34,10 @@ type Conn struct {
 
 // ReconnectingPTY returns a connection serving a TTY that can
 // be reconnected to via ID.
-func (c *Conn) ReconnectingPTY(id string, height, width uint16) (net.Conn, error) {
-	channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d", id, height, width), &peer.ChannelOptions{
+//
+// The command is optional and defaults to start a shell.
+func (c *Conn) ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error) {
+	channel, err := c.CreateChannel(context.Background(), fmt.Sprintf("%s:%d:%d:%s", id, height, width, command), &peer.ChannelOptions{
 		Protocol: ProtocolReconnectingPTY,
 	})
 	if err != nil {

coderd/workspaceagents.go

@@ -381,7 +381,7 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer agentConn.Close()
-	ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width))
+	ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width), "")
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
 		return