feat(codersdk/toolsdk): add template_version_id to coder_create_workspace_build (#17364)

johnstcn · web-flow · commit 272edba1d873 · 2025-04-14T17:29:43.000+01:00
The `coder_create_workspace_build` tool was missing the ability to
change the template version.
diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go
@@ -348,6 +348,11 @@ is provisioned correctly and the agent can connect to the control plane.
 					"transition": map[string]any{
 						"type":        "string",
 						"description": "The transition to perform. Must be one of: start, stop, delete",
+						"enum":        []string{"start", "stop", "delete"},
+					},
+					"template_version_id": map[string]any{
+						"type":        "string",
+						"description": "(Optional) The template version ID to use for the workspace build. If not provided, the previously built version will be used.",
 					},
 				},
 				Required: []string{"workspace_id", "transition"},
@@ -366,9 +371,17 @@ is provisioned correctly and the agent can connect to the control plane.
 			if !ok {
 				return codersdk.WorkspaceBuild{}, xerrors.New("transition must be a string")
 			}
-			return client.CreateWorkspaceBuild(ctx, workspaceID, codersdk.CreateWorkspaceBuildRequest{
+			templateVersionID, err := uuidFromArgs(args, "template_version_id")
+			if err != nil {
+				return codersdk.WorkspaceBuild{}, err
+			}
+			cbr := codersdk.CreateWorkspaceBuildRequest{
 				Transition: codersdk.WorkspaceTransition(rawTransition),
-			})
+			}
+			if templateVersionID != uuid.Nil {
+				cbr.TemplateVersionID = templateVersionID
+			}
+			return client.CreateWorkspaceBuild(ctx, workspaceID, cbr)
 		},
 	}
 
@@ -1240,7 +1253,11 @@ func workspaceAppStatusSlugFromContext(ctx context.Context) (string, bool) {
 }
 
 func uuidFromArgs(args map[string]any, key string) (uuid.UUID, error) {
-	raw, ok := args[key].(string)
+	argKey, ok := args[key]
+	if !ok {
+		return uuid.Nil, nil // No error if key is not present
+	}
+	raw, ok := argKey.(string)
 	if !ok {
 		return uuid.Nil, xerrors.Errorf("%s must be a string", key)
 	}
diff --git a/codersdk/toolsdk/toolsdk_test.go b/codersdk/toolsdk/toolsdk_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -154,6 +155,8 @@ func TestTools(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, codersdk.WorkspaceTransitionStop, result.Transition)
 			require.Equal(t, r.Workspace.ID, result.WorkspaceID)
+			require.Equal(t, r.TemplateVersion.ID, result.TemplateVersionID)
+			require.Equal(t, codersdk.WorkspaceTransitionStop, result.Transition)
 
 			// Important: cancel the build. We don't run any provisioners, so this
 			// will remain in the 'pending' state indefinitely.
@@ -172,11 +175,58 @@ func TestTools(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, codersdk.WorkspaceTransitionStart, result.Transition)
 			require.Equal(t, r.Workspace.ID, result.WorkspaceID)
+			require.Equal(t, r.TemplateVersion.ID, result.TemplateVersionID)
+			require.Equal(t, codersdk.WorkspaceTransitionStart, result.Transition)
 
 			// Important: cancel the build. We don't run any provisioners, so this
 			// will remain in the 'pending' state indefinitely.
 			require.NoError(t, client.CancelWorkspaceBuild(ctx, result.ID))
 		})
+
+		t.Run("TemplateVersionChange", func(t *testing.T) {
+			ctx := testutil.Context(t, testutil.WaitShort)
+			ctx = toolsdk.WithClient(ctx, memberClient)
+
+			// Get the current template version ID before updating
+			workspace, err := memberClient.Workspace(ctx, r.Workspace.ID)
+			require.NoError(t, err)
+			originalVersionID := workspace.LatestBuild.TemplateVersionID
+
+			// Create a new template version to update to
+			newVersion := dbfake.TemplateVersion(t, store).
+				// nolint:gocritic // This is in a test package and does not end up in the build
+				Seed(database.TemplateVersion{
+					OrganizationID: owner.OrganizationID,
+					CreatedBy:      owner.UserID,
+					TemplateID:     uuid.NullUUID{UUID: r.Template.ID, Valid: true},
+				}).Do()
+
+			// Update to new version
+			updateBuild, err := testTool(ctx, t, toolsdk.CreateWorkspaceBuild, map[string]any{
+				"workspace_id":        r.Workspace.ID.String(),
+				"transition":          "start",
+				"template_version_id": newVersion.TemplateVersion.ID.String(),
+			})
+			require.NoError(t, err)
+			require.Equal(t, codersdk.WorkspaceTransitionStart, updateBuild.Transition)
+			require.Equal(t, r.Workspace.ID.String(), updateBuild.WorkspaceID.String())
+			require.Equal(t, newVersion.TemplateVersion.ID.String(), updateBuild.TemplateVersionID.String())
+			// Cancel the build so it doesn't remain in the 'pending' state indefinitely.
+			require.NoError(t, client.CancelWorkspaceBuild(ctx, updateBuild.ID))
+
+			// Roll back to the original version
+			rollbackBuild, err := testTool(ctx, t, toolsdk.CreateWorkspaceBuild, map[string]any{
+				"workspace_id":        r.Workspace.ID.String(),
+				"transition":          "start",
+				"template_version_id": originalVersionID.String(),
+			})
+			require.NoError(t, err)
+			require.Equal(t, codersdk.WorkspaceTransitionStart, rollbackBuild.Transition)
+			require.Equal(t, r.Workspace.ID.String(), rollbackBuild.WorkspaceID.String())
+			require.Equal(t, originalVersionID.String(), rollbackBuild.TemplateVersionID.String())
+			// Cancel the build so it doesn't remain in the 'pending' state indefinitely.
+			require.NoError(t, client.CancelWorkspaceBuild(ctx, rollbackBuild.ID))
+		})
 	})
 
 	t.Run("ListTemplateVersionParameters", func(t *testing.T) {
