feat: Add "state" command to pull and push workspace state (#1264)

It's possible for a workspace to become in an invalid state.
This is something we'll detect for jobs, and allow monitoring of.

These commands will allow admins to manually reconcile state.

cli/root.go

@@ -80,6 +80,7 @@ func Root() *cobra.Command {
 		server(),
 		show(),
 		start(),
+		state(),
 		stop(),
 		ssh(),
 		templates(),

cli/state.go

@@ -0,0 +1,121 @@
+package cli
+
+import (
+	"io"
+	"os"
+	"time"
+
+	"github.com/spf13/cobra"
+
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
+)
+
+func state() *cobra.Command {
+	cmd := &cobra.Command{
+		Use: "state",
+	}
+	cmd.AddCommand(statePull(), statePush())
+	return cmd
+}
+
+func statePull() *cobra.Command {
+	var buildName string
+	cmd := &cobra.Command{
+		Use:  "pull <workspace> [file]",
+		Args: cobra.MinimumNArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			client, err := createClient(cmd)
+			if err != nil {
+				return err
+			}
+			organization, err := currentOrganization(cmd, client)
+			if err != nil {
+				return err
+			}
+
+			workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, args[0])
+			if err != nil {
+				return err
+			}
+			var build codersdk.WorkspaceBuild
+			if buildName == "latest" {
+				build = workspace.LatestBuild
+			} else {
+				build, err = client.WorkspaceBuildByName(cmd.Context(), workspace.ID, buildName)
+				if err != nil {
+					return err
+				}
+			}
+
+			state, err := client.WorkspaceBuildState(cmd.Context(), build.ID)
+			if err != nil {
+				return err
+			}
+
+			if len(args) < 2 {
+				cmd.Println(string(state))
+				return nil
+			}
+
+			return os.WriteFile(args[1], state, 0600)
+		},
+	}
+	cmd.Flags().StringVarP(&buildName, "build", "b", "latest", "Specify a workspace build to target by name.")
+	return cmd
+}
+
+func statePush() *cobra.Command {
+	var buildName string
+	cmd := &cobra.Command{
+		Use:  "push <workspace> <file>",
+		Args: cobra.ExactArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			client, err := createClient(cmd)
+			if err != nil {
+				return err
+			}
+			organization, err := currentOrganization(cmd, client)
+			if err != nil {
+				return err
+			}
+
+			workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, args[0])
+			if err != nil {
+				return err
+			}
+			var build codersdk.WorkspaceBuild
+			if buildName == "latest" {
+				build = workspace.LatestBuild
+			} else {
+				build, err = client.WorkspaceBuildByName(cmd.Context(), workspace.ID, buildName)
+				if err != nil {
+					return err
+				}
+			}
+
+			var state []byte
+			if args[1] == "-" {
+				state, err = io.ReadAll(cmd.InOrStdin())
+			} else {
+				state, err = os.ReadFile(args[1])
+			}
+			if err != nil {
+				return err
+			}
+
+			before := time.Now()
+			build, err = client.CreateWorkspaceBuild(cmd.Context(), workspace.ID, codersdk.CreateWorkspaceBuildRequest{
+				TemplateVersionID: build.TemplateVersionID,
+				Transition:        build.Transition,
+				ProvisionerState:  state,
+			})
+			if err != nil {
+				return err
+			}
+			return cliui.WorkspaceBuild(cmd.Context(), cmd.OutOrStderr(), client, build.ID, before)
+		},
+	}
+	cmd.Flags().StringVarP(&buildName, "build", "b", "latest", "Specify a workspace build to target by name.")
+	return cmd
+}

cli/state_test.go

@@ -0,0 +1,129 @@
+package cli_test
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/provisioner/echo"
+	"github.com/coder/coder/provisionersdk/proto"
+)
+
+func TestStatePull(t *testing.T) {
+	t.Parallel()
+	t.Run("File", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		wantState := []byte("some state")
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse: echo.ParseComplete,
+			Provision: []*proto.Provision_Response{{
+				Type: &proto.Provision_Response_Complete{
+					Complete: &proto.Provision_Complete{
+						State: wantState,
+					},
+				},
+			}},
+		})
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		statefilePath := filepath.Join(t.TempDir(), "state")
+		cmd, root := clitest.New(t, "state", "pull", workspace.Name, statefilePath)
+		clitest.SetupConfig(t, client, root)
+		err := cmd.Execute()
+		require.NoError(t, err)
+		gotState, err := os.ReadFile(statefilePath)
+		require.NoError(t, err)
+		require.Equal(t, wantState, gotState)
+	})
+	t.Run("Stdout", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		wantState := []byte("some state")
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse: echo.ParseComplete,
+			Provision: []*proto.Provision_Response{{
+				Type: &proto.Provision_Response_Complete{
+					Complete: &proto.Provision_Complete{
+						State: wantState,
+					},
+				},
+			}},
+		})
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		cmd, root := clitest.New(t, "state", "pull", workspace.Name)
+		var gotState bytes.Buffer
+		cmd.SetOut(&gotState)
+		clitest.SetupConfig(t, client, root)
+		err := cmd.Execute()
+		require.NoError(t, err)
+		require.Equal(t, wantState, bytes.TrimSpace(gotState.Bytes()))
+	})
+}
+
+func TestStatePush(t *testing.T) {
+	t.Parallel()
+	t.Run("File", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:     echo.ParseComplete,
+			Provision: echo.ProvisionComplete,
+		})
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		stateFile, err := os.CreateTemp(t.TempDir(), "")
+		require.NoError(t, err)
+		wantState := []byte("some magic state")
+		_, err = stateFile.Write(wantState)
+		require.NoError(t, err)
+		err = stateFile.Close()
+		require.NoError(t, err)
+		cmd, root := clitest.New(t, "state", "push", workspace.Name, stateFile.Name())
+		cmd.SetErr(io.Discard)
+		cmd.SetOut(io.Discard)
+		clitest.SetupConfig(t, client, root)
+		err = cmd.Execute()
+		require.NoError(t, err)
+	})
+
+	t.Run("Stdin", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:     echo.ParseComplete,
+			Provision: echo.ProvisionComplete,
+		})
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		cmd, root := clitest.New(t, "state", "push", "--build", workspace.LatestBuild.Name, workspace.Name, "-")
+		clitest.SetupConfig(t, client, root)
+		cmd.SetIn(strings.NewReader("some magic state"))
+		err := cmd.Execute()
+		require.NoError(t, err)
+	})
+}

coderd/coderd.go

@@ -268,6 +268,7 @@ func New(options *Options) (http.Handler, func()) {
 			r.Patch("/cancel", api.patchCancelWorkspaceBuild)
 			r.Get("/logs", api.workspaceBuildLogs)
 			r.Get("/resources", api.workspaceBuildResources)
+			r.Get("/state", api.workspaceBuildState)
 		})
 	})
 	r.NotFound(site.DefaultHandler().ServeHTTP)

coderd/provisionerdaemons.go

@@ -95,7 +95,7 @@ func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request
 		},
 	})
 	err = server.Serve(r.Context(), session)
-	if err != nil {
+	if err != nil && !xerrors.Is(err, io.EOF) {
 		api.Logger.Debug(r.Context(), "provisioner daemon disconnected", slog.Error(err))
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("serve: %s", err))
 		return

coderd/workspacebuilds.go

@@ -87,6 +87,14 @@ func (api *api) workspaceBuildLogs(rw http.ResponseWriter, r *http.Request) {
 	api.provisionerJobLogs(rw, r, job)
 }
 
+func (*api) workspaceBuildState(rw http.ResponseWriter, r *http.Request) {
+	workspaceBuild := httpmw.WorkspaceBuildParam(r)
+
+	rw.Header().Set("Content-Type", "application/json")
+	rw.WriteHeader(http.StatusOK)
+	_, _ = rw.Write(workspaceBuild.ProvisionerState)
+}
+
 func convertWorkspaceBuild(workspaceBuild database.WorkspaceBuild, job codersdk.ProvisionerJob) codersdk.WorkspaceBuild {
 	//nolint:unconvert
 	return codersdk.WorkspaceBuild{

coderd/workspacebuilds_test.go

@@ -168,3 +168,29 @@ func TestWorkspaceBuildLogs(t *testing.T) {
 	}
 	require.Fail(t, "example message never happened")
 }
+
+func TestWorkspaceBuildState(t *testing.T) {
+	t.Parallel()
+	client := coderdtest.New(t, nil)
+	user := coderdtest.CreateFirstUser(t, client)
+	coderdtest.NewProvisionerDaemon(t, client)
+	wantState := []byte("some kinda state")
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:           echo.ParseComplete,
+		ProvisionDryRun: echo.ProvisionComplete,
+		Provision: []*proto.Provision_Response{{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					State: wantState,
+				},
+			},
+		}},
+	})
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+	gotState, err := client.WorkspaceBuildState(context.Background(), workspace.LatestBuild.ID)
+	require.NoError(t, err)
+	require.Equal(t, wantState, gotState)
+}

coderd/workspaces.go

@@ -212,6 +212,11 @@ func (api *api) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
 		if err != nil {
 			return xerrors.Errorf("insert provisioner job: %w", err)
 		}
+		state := createBuild.ProvisionerState
+		if state == nil || len(state) == 0 {
+			state = priorHistory.ProvisionerState
+		}
+
 		workspaceBuild, err = db.InsertWorkspaceBuild(r.Context(), database.InsertWorkspaceBuildParams{
 			ID:                workspaceBuildID,
 			CreatedAt:         database.Now(),
@@ -220,7 +225,7 @@ func (api *api) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
 			TemplateVersionID: templateVersion.ID,
 			BeforeID:          priorHistoryID,
 			Name:              namesgenerator.GetRandomName(1),
-			ProvisionerState:  priorHistory.ProvisionerState,
+			ProvisionerState:  state,
 			InitiatorID:       apiKey.UserID,
 			Transition:        createBuild.Transition,
 			JobID:             provisionerJob.ID,

coderd/workspaces_test.go

@@ -130,6 +130,29 @@ func TestPostWorkspaceBuild(t *testing.T) {
 		require.Equal(t, build.ID.String(), firstBuild.AfterID.String())
 	})
 
+	t.Run("WithState", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		closeDaemon := coderdtest.NewProvisionerDaemon(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		_ = closeDaemon.Close()
+		wantState := []byte("something")
+		build, err := client.CreateWorkspaceBuild(context.Background(), workspace.ID, codersdk.CreateWorkspaceBuildRequest{
+			TemplateVersionID: template.ActiveVersionID,
+			Transition:        database.WorkspaceTransitionStart,
+			ProvisionerState:  wantState,
+		})
+		require.NoError(t, err)
+		gotState, err := client.WorkspaceBuildState(context.Background(), build.ID)
+		require.NoError(t, err)
+		require.Equal(t, wantState, gotState)
+	})
+
 	t.Run("Delete", func(t *testing.T) {
 		t.Parallel()
 		client := coderdtest.New(t, nil)

codersdk/workspacebuilds.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"time"
 
@@ -79,3 +80,16 @@ func (c *Client) WorkspaceBuildLogsBefore(ctx context.Context, build uuid.UUID,
 func (c *Client) WorkspaceBuildLogsAfter(ctx context.Context, build uuid.UUID, after time.Time) (<-chan ProvisionerJobLog, error) {
 	return c.provisionerJobLogsAfter(ctx, fmt.Sprintf("/api/v2/workspacebuilds/%s/logs", build), after)
 }
+
+// WorkspaceBuildState returns the provisioner state of the build.
+func (c *Client) WorkspaceBuildState(ctx context.Context, build uuid.UUID) ([]byte, error) {
+	res, err := c.request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspacebuilds/%s/state", build), nil)
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		return nil, readBodyAsError(res)
+	}
+	return io.ReadAll(res.Body)
+}

codersdk/workspaces.go

@@ -34,6 +34,7 @@ type CreateWorkspaceBuildRequest struct {
 	TemplateVersionID uuid.UUID                    `json:"template_version_id"`
 	Transition        database.WorkspaceTransition `json:"transition" validate:"oneof=create start stop delete,required"`
 	DryRun            bool                         `json:"dry_run"`
+	ProvisionerState  []byte                       `json:"state,omitempty"`
 }
 
 // Workspace returns a single workspace.