fix(cli): ssh: auto-update workspace (#11773)

mtojek · web-flow · commit 77a4792ecd33 · 2024-01-23T18:01:44.000+01:00
diff --git a/cli/restart.go b/cli/restart.go
@@ -64,7 +64,7 @@ func (r *RootCmd) restart() *clibase.Cmd {
 			// It's possible for a workspace build to fail due to the template requiring starting
 			// workspaces with the active version.
 			if cerr, ok := codersdk.AsError(err); ok && cerr.StatusCode() == http.StatusForbidden {
-				_, _ = fmt.Fprintln(inv.Stdout, "Failed to restart with the template version from your last build. Policy may require you to restart with the current active template version.")
+				_, _ = fmt.Fprintln(inv.Stdout, "Unable to restart the workspace with the template version from the last build. Policy may require you to restart with the current active template version.")
 				build, err = startWorkspace(inv, client, workspace, parameterFlags, WorkspaceUpdate)
 				if err != nil {
 					return xerrors.Errorf("start workspace with active template version: %w", err)
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
@@ -575,13 +576,27 @@ func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *
 					codersdk.WorkspaceStatusStopped,
 				)
 		}
-		// startWorkspace based on the last build parameters.
+
+		// Start workspace based on the last build parameters.
+		// It's possible for a workspace build to fail due to the template requiring starting
+		// workspaces with the active version.
 		_, _ = fmt.Fprintf(inv.Stderr, "Workspace was stopped, starting workspace to allow connecting to %q...\n", workspace.Name)
-		build, err := startWorkspace(inv, client, workspace, workspaceParameterFlags{}, WorkspaceStart)
+		_, err = startWorkspace(inv, client, workspace, workspaceParameterFlags{}, WorkspaceStart)
+		if cerr, ok := codersdk.AsError(err); ok && cerr.StatusCode() == http.StatusForbidden {
+			_, _ = fmt.Fprintln(inv.Stdout, "Unable to start the workspace with template version from last build. The auto-update policy may require you to restart with the current active template version.")
+			_, err = startWorkspace(inv, client, workspace, workspaceParameterFlags{}, WorkspaceUpdate)
+			if err != nil {
+				return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("start workspace with active template version: %w", err)
+			}
+		} else if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("start workspace with current template version: %w", err)
+		}
+
+		// Refresh workspace state so that `outdated`, `build`,`template_*` fields are up-to-date.
+		workspace, err = namedWorkspace(ctx, client, workspaceParts[0])
 		if err != nil {
-			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("unable to start workspace: %w", err)
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
 		}
-		workspace.LatestBuild = build
 	}
 	if workspace.LatestBuild.Job.CompletedAt == nil {
 		err := cliui.WorkspaceBuild(ctx, inv.Stderr, client, workspace.LatestBuild.ID)
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -133,6 +133,71 @@ func TestSSH(t *testing.T) {
 		pty.WriteLine("exit")
 		<-cmdDone
 	})
+	t.Run("RequireActiveVersion", func(t *testing.T) {
+		t.Parallel()
+
+		authToken := uuid.NewString()
+		ownerClient := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+		owner := coderdtest.CreateFirstUser(t, ownerClient)
+		client, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleMember())
+
+		echoResponses := &echo.Responses{
+			Parse:          echo.ParseComplete,
+			ProvisionPlan:  echo.PlanComplete,
+			ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+		}
+
+		version := coderdtest.CreateTemplateVersion(t, ownerClient, owner.OrganizationID, echoResponses)
+		coderdtest.AwaitTemplateVersionJobCompleted(t, ownerClient, version.ID)
+		template := coderdtest.CreateTemplate(t, ownerClient, owner.OrganizationID, version.ID)
+
+		workspace := coderdtest.CreateWorkspace(t, client, owner.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+			cwr.AutomaticUpdates = codersdk.AutomaticUpdatesAlways
+		})
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+
+		// Stop the workspace
+		workspaceBuild := coderdtest.CreateWorkspaceBuild(t, client, workspace, database.WorkspaceTransitionStop)
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceBuild.ID)
+
+		// Update template version
+		version = coderdtest.UpdateTemplateVersion(t, ownerClient, owner.OrganizationID, echoResponses, template.ID)
+		coderdtest.AwaitTemplateVersionJobCompleted(t, ownerClient, version.ID)
+		err := ownerClient.UpdateActiveTemplateVersion(context.Background(), template.ID, codersdk.UpdateActiveTemplateVersion{
+			ID: version.ID,
+		})
+		require.NoError(t, err)
+
+		// SSH to the workspace which should auto-update and autostart it
+		inv, root := clitest.New(t, "ssh", workspace.Name)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		// When the agent connects, the workspace was started, and we should
+		// have access to the shell.
+		_ = agenttest.New(t, client.URL, authToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
+		pty.WriteLine("exit")
+		<-cmdDone
+
+		// Double-check if workspace's template version is up-to-date
+		workspace, err = client.Workspace(context.Background(), workspace.ID)
+		require.NoError(t, err)
+		assert.Equal(t, version.ID, workspace.TemplateActiveVersionID)
+		assert.Equal(t, workspace.TemplateActiveVersionID, workspace.LatestBuild.TemplateVersionID)
+		assert.False(t, workspace.Outdated)
+	})
+
 	t.Run("ShowTroubleshootingURLAfterTimeout", func(t *testing.T) {
 		t.Parallel()
 
diff --git a/cli/start.go b/cli/start.go
@@ -49,7 +49,7 @@ func (r *RootCmd) start() *clibase.Cmd {
 				// It's possible for a workspace build to fail due to the template requiring starting
 				// workspaces with the active version.
 				if cerr, ok := codersdk.AsError(err); ok && cerr.StatusCode() == http.StatusForbidden {
-					_, _ = fmt.Fprintln(inv.Stdout, "Failed to restart with the template version from your last build. Policy may require you to restart with the current active template version.")
+					_, _ = fmt.Fprintln(inv.Stdout, "Unable to start the workspace with the template version from the last build. Policy may require you to restart with the current active template version.")
 					build, err = startWorkspace(inv, client, workspace, parameterFlags, WorkspaceUpdate)
 					if err != nil {
 						return xerrors.Errorf("start workspace with active template version: %w", err)
@@ -79,6 +79,7 @@ func (r *RootCmd) start() *clibase.Cmd {
 
 func buildWorkspaceStartRequest(inv *clibase.Invocation, client *codersdk.Client, workspace codersdk.Workspace, parameterFlags workspaceParameterFlags, action WorkspaceCLIAction) (codersdk.CreateWorkspaceBuildRequest, error) {
 	version := workspace.LatestBuild.TemplateVersionID
+
 	if workspace.AutomaticUpdates == codersdk.AutomaticUpdatesAlways || action == WorkspaceUpdate {
 		version = workspace.TemplateActiveVersionID
 		if version != workspace.LatestBuild.TemplateVersionID {
diff --git a/enterprise/cli/start_test.go b/enterprise/cli/start_test.go
@@ -159,8 +159,16 @@ func TestStart(t *testing.T) {
 
 						ws = coderdtest.MustWorkspace(t, c.Client, ws.ID)
 						require.Equal(t, c.ExpectedVersion, ws.LatestBuild.TemplateVersionID)
-						if initialTemplateVersion != ws.LatestBuild.TemplateVersionID {
-							require.Contains(t, buf.String(), "Failed to restart with the template version from your last build. Policy may require you to restart with the current active template version.")
+						if initialTemplateVersion == ws.LatestBuild.TemplateVersionID {
+							return
+						}
+
+						if cmd == "start" {
+							require.Contains(t, buf.String(), "Unable to start the workspace with the template version from the last build")
+						}
+
+						if cmd == "restart" {
+							require.Contains(t, buf.String(), "Unable to restart the workspace with the template version from the last build")
 						}
 					})
 				}
