Skip to content

Commit cb89bc1

Browse files
EmyrkEmyrk
authored
feat: restart stopped workspaces on ssh command (#11050)
* feat: autostart workspaces on ssh & port forward This is opt out by default. VScode ssh does not have this behavior
1 parent 1f7c63c commit cb89bc1

12 files changed

+170
-23
lines changed

‎cli/configssh.go

+24-4
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"path/filepath"
1414
"runtime"
1515
"sort"
16+
"strconv"
1617
"strings"
1718

1819
"github.com/cli/safeexec"
@@ -46,9 +47,10 @@ const (
4647
// sshConfigOptions represents options that can be stored and read
4748
// from the coder config in ~/.ssh/coder.
4849
type sshConfigOptions struct {
49-
waitEnum string
50-
userHostPrefix string
51-
sshOptions []string
50+
waitEnum string
51+
userHostPrefix string
52+
sshOptions []string
53+
disableAutostart bool
5254
}
5355

5456
// addOptions expects options in the form of "option=value" or "option value".
@@ -106,7 +108,7 @@ func (o sshConfigOptions) equal(other sshConfigOptions) bool {
106108
if !slices.Equal(opt1, opt2) {
107109
return false
108110
}
109-
return o.waitEnum == other.waitEnum && o.userHostPrefix == other.userHostPrefix
111+
return o.waitEnum == other.waitEnum && o.userHostPrefix == other.userHostPrefix && o.disableAutostart == other.disableAutostart
110112
}
111113

112114
func (o sshConfigOptions) asList() (list []string) {
@@ -116,6 +118,9 @@ func (o sshConfigOptions) asList() (list []string) {
116118
if o.userHostPrefix != "" {
117119
list = append(list, fmt.Sprintf("ssh-host-prefix: %s", o.userHostPrefix))
118120
}
121+
if o.disableAutostart {
122+
list = append(list, fmt.Sprintf("disable-autostart: %v", o.disableAutostart))
123+
}
119124
for _, opt := range o.sshOptions {
120125
list = append(list, fmt.Sprintf("ssh-option: %s", opt))
121126
}
@@ -392,6 +397,9 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
392397
if sshConfigOpts.waitEnum != "auto" {
393398
flags += " --wait=" + sshConfigOpts.waitEnum
394399
}
400+
if sshConfigOpts.disableAutostart {
401+
flags += " --disable-autostart=true"
402+
}
395403
defaultOptions = append(defaultOptions, fmt.Sprintf(
396404
"ProxyCommand %s --global-config %s ssh --stdio%s %s",
397405
escapedCoderBinary, escapedGlobalConfig, flags, workspaceHostname,
@@ -566,6 +574,13 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
566574
Default: "auto",
567575
Value: clibase.EnumOf(&sshConfigOpts.waitEnum, "yes", "no", "auto"),
568576
},
577+
{
578+
Flag: "disable-autostart",
579+
Description: "Disable starting the workspace automatically when connecting via SSH.",
580+
Env: "CODER_CONFIGSSH_DISABLE_AUTOSTART",
581+
Value: clibase.BoolOf(&sshConfigOpts.disableAutostart),
582+
Default: "false",
583+
},
569584
{
570585
Flag: "force-unix-filepaths",
571586
Env: "CODER_CONFIGSSH_UNIX_FILEPATHS",
@@ -602,6 +617,9 @@ func sshConfigWriteSectionHeader(w io.Writer, addNewline bool, o sshConfigOption
602617
if o.userHostPrefix != "" {
603618
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-host-prefix", o.userHostPrefix)
604619
}
620+
if o.disableAutostart {
621+
_, _ = fmt.Fprintf(&ow, "# :%s=%v\n", "disable-autostart", o.disableAutostart)
622+
}
605623
for _, opt := range o.sshOptions {
606624
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-option", opt)
607625
}
@@ -634,6 +652,8 @@ func sshConfigParseLastOptions(r io.Reader) (o sshConfigOptions) {
634652
o.userHostPrefix = parts[1]
635653
case "ssh-option":
636654
o.sshOptions = append(o.sshOptions, parts[1])
655+
case "disable-autostart":
656+
o.disableAutostart, _ = strconv.ParseBool(parts[1])
637657
default:
638658
// Unknown option, ignore.
639659
}

‎cli/ping.go

+1
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ func (r *RootCmd) ping() *clibase.Cmd {
4040
workspaceName := inv.Args[0]
4141
_, workspaceAgent, err := getWorkspaceAndAgent(
4242
ctx, inv, client,
43+
false, // Do not autostart for a ping.
4344
codersdk.Me, workspaceName,
4445
)
4546
if err != nil {

‎cli/portforward.go

+5-3
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ import (
2626

2727
func (r *RootCmd) portForward() *clibase.Cmd {
2828
var (
29-
tcpForwards []string // <port>:<port>
30-
udpForwards []string // <port>:<port>
29+
tcpForwards []string // <port>:<port>
30+
udpForwards []string // <port>:<port>
31+
disableAutostart bool
3132
)
3233
client := new(codersdk.Client)
3334
cmd := &clibase.Cmd{
@@ -76,7 +77,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
7677
return xerrors.New("no port-forwards requested")
7778
}
7879

79-
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
80+
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, codersdk.Me, inv.Args[0])
8081
if err != nil {
8182
return err
8283
}
@@ -180,6 +181,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
180181
Description: "Forward UDP port(s) from the workspace to the local machine. The UDP connection has TCP-like semantics to support stateful UDP protocols.",
181182
Value: clibase.StringArrayOf(&udpForwards),
182183
},
184+
sshDisableAutostartOption(clibase.BoolOf(&disableAutostart)),
183185
}
184186

185187
return cmd

‎cli/speedtest.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
3535
ctx, cancel := context.WithCancel(inv.Context())
3636
defer cancel()
3737

38-
_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
38+
_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, false, codersdk.Me, inv.Args[0])
3939
if err != nil {
4040
return err
4141
}

‎cli/ssh.go

+55-15
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"sync"
1515
"time"
1616

17+
"github.com/coder/retry"
1718
"github.com/gen2brain/beeep"
1819
"github.com/gofrs/flock"
1920
"github.com/google/uuid"
@@ -34,7 +35,6 @@ import (
3435
"github.com/coder/coder/v2/coderd/util/ptr"
3536
"github.com/coder/coder/v2/codersdk"
3637
"github.com/coder/coder/v2/cryptorand"
37-
"github.com/coder/retry"
3838
)
3939

4040
var (
@@ -44,15 +44,16 @@ var (
4444

4545
func (r *RootCmd) ssh() *clibase.Cmd {
4646
var (
47-
stdio bool
48-
forwardAgent bool
49-
forwardGPG bool
50-
identityAgent string
51-
wsPollInterval time.Duration
52-
waitEnum string
53-
noWait bool
54-
logDirPath string
55-
remoteForward string
47+
stdio bool
48+
forwardAgent bool
49+
forwardGPG bool
50+
identityAgent string
51+
wsPollInterval time.Duration
52+
waitEnum string
53+
noWait bool
54+
logDirPath string
55+
remoteForward string
56+
disableAutostart bool
5657
)
5758
client := new(codersdk.Client)
5859
cmd := &clibase.Cmd{
@@ -143,7 +144,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
143144
}
144145
}
145146

146-
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
147+
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, codersdk.Me, inv.Args[0])
147148
if err != nil {
148149
return err
149150
}
@@ -459,6 +460,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
459460
FlagShorthand: "R",
460461
Value: clibase.StringOf(&remoteForward),
461462
},
463+
sshDisableAutostartOption(clibase.BoolOf(&disableAutostart)),
462464
}
463465
return cmd
464466
}
@@ -530,9 +532,9 @@ startWatchLoop:
530532
}
531533

532534
// getWorkspaceAgent returns the workspace and agent selected using either the
533-
// `<workspace>[.<agent>]` syntax via `in` or picks a random workspace and agent
534-
// if `shuffle` is true.
535-
func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *codersdk.Client, userID string, in string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) { //nolint:revive
535+
// `<workspace>[.<agent>]` syntax via `in`.
536+
// If autoStart is true, the workspace will be started if it is not already running.
537+
func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *codersdk.Client, autostart bool, userID string, in string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) { //nolint:revive
536538
var (
537539
workspace codersdk.Workspace
538540
workspaceParts = strings.Split(in, ".")
@@ -545,7 +547,35 @@ func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *
545547
}
546548

547549
if workspace.LatestBuild.Transition != codersdk.WorkspaceTransitionStart {
548-
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("workspace must be in start transition to ssh")
550+
if !autostart {
551+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("workspace must be in start transition to ssh")
552+
}
553+
// Autostart the workspace for the user.
554+
// For some failure modes, return a better message.
555+
if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionDelete {
556+
// Any sort of deleting status, we should reject with a nicer error.
557+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is deleted", workspace.Name)
558+
}
559+
if workspace.LatestBuild.Job.Status == codersdk.ProvisionerJobFailed {
560+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{},
561+
xerrors.Errorf("workspace %q is in failed state, unable to autostart the workspace", workspace.Name)
562+
}
563+
// The workspace needs to be stopped before we can start it.
564+
// It cannot be in any pending or failed state.
565+
if workspace.LatestBuild.Status != codersdk.WorkspaceStatusStopped {
566+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{},
567+
xerrors.Errorf("workspace must be in start transition to ssh, was unable to autostart as the last build job is %q, expected %q",
568+
workspace.LatestBuild.Status,
569+
codersdk.WorkspaceStatusStopped,
570+
)
571+
}
572+
// startWorkspace based on the last build parameters.
573+
_, _ = fmt.Fprintf(inv.Stderr, "Workspace was stopped, starting workspace to allow connecting to %q...\n", workspace.Name)
574+
build, err := startWorkspace(inv, client, workspace, workspaceParameterFlags{}, WorkspaceStart)
575+
if err != nil {
576+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("unable to start workspace: %w", err)
577+
}
578+
workspace.LatestBuild = build
549579
}
550580
if workspace.LatestBuild.Job.CompletedAt == nil {
551581
err := cliui.WorkspaceBuild(ctx, inv.Stderr, client, workspace.LatestBuild.ID)
@@ -915,3 +945,13 @@ func (c *rawSSHCopier) Close() error {
915945
}
916946
return err
917947
}
948+
949+
func sshDisableAutostartOption(src *clibase.Bool) clibase.Option {
950+
return clibase.Option{
951+
Flag: "disable-autostart",
952+
Description: "Disable starting the workspace automatically when connecting via SSH.",
953+
Env: "CODER_SSH_DISABLE_AUTOSTART",
954+
Value: src,
955+
Default: "false",
956+
}
957+
}

‎cli/ssh_test.go

+45
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"testing"
2222
"time"
2323

24+
"github.com/google/uuid"
2425
"github.com/stretchr/testify/assert"
2526
"github.com/stretchr/testify/require"
2627
"golang.org/x/crypto/ssh"
@@ -38,7 +39,9 @@ import (
3839
"github.com/coder/coder/v2/coderd/database"
3940
"github.com/coder/coder/v2/coderd/database/dbfake"
4041
"github.com/coder/coder/v2/coderd/database/dbtestutil"
42+
"github.com/coder/coder/v2/coderd/rbac"
4143
"github.com/coder/coder/v2/codersdk"
44+
"github.com/coder/coder/v2/provisioner/echo"
4245
"github.com/coder/coder/v2/provisionersdk/proto"
4346
"github.com/coder/coder/v2/pty"
4447
"github.com/coder/coder/v2/pty/ptytest"
@@ -86,6 +89,48 @@ func TestSSH(t *testing.T) {
8689
pty.WriteLine("exit")
8790
<-cmdDone
8891
})
92+
t.Run("StartStoppedWorkspace", func(t *testing.T) {
93+
t.Parallel()
94+
95+
authToken := uuid.NewString()
96+
ownerClient := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
97+
owner := coderdtest.CreateFirstUser(t, ownerClient)
98+
client, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleTemplateAdmin())
99+
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, &echo.Responses{
100+
Parse: echo.ParseComplete,
101+
ProvisionPlan: echo.PlanComplete,
102+
ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
103+
})
104+
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
105+
template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
106+
workspace := coderdtest.CreateWorkspace(t, client, owner.OrganizationID, template.ID)
107+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
108+
// Stop the workspace
109+
workspaceBuild := coderdtest.CreateWorkspaceBuild(t, client, workspace, database.WorkspaceTransitionStop)
110+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceBuild.ID)
111+
112+
// SSH to the workspace which should autostart it
113+
inv, root := clitest.New(t, "ssh", workspace.Name)
114+
clitest.SetupConfig(t, client, root)
115+
pty := ptytest.New(t).Attach(inv)
116+
117+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
118+
defer cancel()
119+
120+
cmdDone := tGo(t, func() {
121+
err := inv.WithContext(ctx).Run()
122+
assert.NoError(t, err)
123+
})
124+
125+
// When the agent connects, the workspace was started, and we should
126+
// have access to the shell.
127+
_ = agenttest.New(t, client.URL, authToken)
128+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
129+
130+
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
131+
pty.WriteLine("exit")
132+
<-cmdDone
133+
})
89134
t.Run("ShowTroubleshootingURLAfterTimeout", func(t *testing.T) {
90135
t.Parallel()
91136

‎cli/testdata/coder_config-ssh_--help.golden

+3
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ OPTIONS:
2121
ProxyCommand. By default, the binary invoking this command ('config
2222
ssh') is used.
2323

24+
--disable-autostart bool, $CODER_CONFIGSSH_DISABLE_AUTOSTART (default: false)
25+
Disable starting the workspace automatically when connecting via SSH.
26+
2427
-n, --dry-run bool, $CODER_SSH_DRY_RUN
2528
Perform a trial run with no changes made, showing a diff at the end.
2629

‎cli/testdata/coder_port-forward_--help.golden

+3
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ USAGE:
3434
$ coder port-forward <workspace> --tcp 1.2.3.4:8080:8080
3535

3636
OPTIONS:
37+
--disable-autostart bool, $CODER_SSH_DISABLE_AUTOSTART (default: false)
38+
Disable starting the workspace automatically when connecting via SSH.
39+
3740
-p, --tcp string-array, $CODER_PORT_FORWARD_TCP
3841
Forward TCP port(s) from the workspace to the local machine.
3942

‎cli/testdata/coder_ssh_--help.golden

+3
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ USAGE:
66
Start a shell into a workspace
77

88
OPTIONS:
9+
--disable-autostart bool, $CODER_SSH_DISABLE_AUTOSTART (default: false)
10+
Disable starting the workspace automatically when connecting via SSH.
11+
912
-A, --forward-agent bool, $CODER_SSH_FORWARD_AGENT
1013
Specifies whether to forward the SSH agent specified in
1114
$SSH_AUTH_SOCK.

‎docs/cli/config-ssh.md

+10
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/cli/port-forward.md

+10
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)