feat: reuse agent tokens when a prebuilt agent reinitializes

docs/data-sources/workspace_preset.md

@@ -41,6 +41,10 @@ data "coder_workspace_preset" "example" {
 - `parameters` (Map of String) Workspace parameters that will be set by the workspace preset. For simple templates that only need prebuilds, you may define a preset with zero parameters. Because workspace parameters may change between Coder template versions, preset parameters are allowed to define values for parameters that do not exist in the current template version.
 - `prebuilds` (Block Set, Max: 1) Prebuilt workspace configuration related to this workspace preset. Coder will build and maintain workspaces in reserve based on this configuration. When a user creates a new workspace using a preset, they will be assigned a prebuilt workspace, instead of waiting for a new workspace to build. (see [below for nested schema](#nestedblock--prebuilds))
 
+### Optional
+
+- `prebuilds` (Block Set, Max: 1) Prebuilds of the workspace preset. (see [below for nested schema](#nestedblock--prebuilds))
+
 ### Read-Only
 
 - `id` (String) The preset ID is automatically generated and may change between runs. It is recommended to use the `name` attribute to identify the preset.

provider/agent.go

@@ -2,13 +2,16 @@ package provider
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"path/filepath"
 	"reflect"
 	"strings"
 
 	"github.com/google/uuid"
 	"github.com/hashicorp/go-cty/cty"
+	"github.com/hashicorp/terraform-plugin-log/tflog"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
@@ -22,10 +25,38 @@ func agentResource() *schema.Resource {
 		SchemaVersion: 1,
 
 		Description: "Use this resource to associate an agent.",
-		CreateContext: func(_ context.Context, resourceData *schema.ResourceData, i interface{}) diag.Diagnostics {
-			// This should be a real authentication token!
-			resourceData.SetId(uuid.NewString())
-			err := resourceData.Set("token", uuid.NewString())
+		CreateContext: func(ctx context.Context, resourceData *schema.ResourceData, i interface{}) diag.Diagnostics {
+			agentID := uuid.NewString()
+			resourceData.SetId(agentID)
+
+			// Most of the time, we will generate a new token for the agent.
+			// In the case of a prebuilt workspace being claimed, we will override with
+			// an existing token provided below.
+			token := uuid.NewString()
+
+			// If isPrebuild is true, then this workspace was built by the prebuilds system.
+			// This does not determine whether the workspace has been claimed by a user.
+			// At this point, it may or may not have been claimed.
+			isPrebuild := helpers.OptionalEnv(IsPrebuildEnvironmentVariable()) == "true"
+			// existingToken should only have been set if isPrebuild is true, because we only
+			// reuse the token when a prebuilt workspace is being claimed.
+			existingToken := helpers.OptionalEnv(RunningAgentTokenEnvironmentVariable(agentID))
+			logFields := map[string]interface{}{
+				"agent_id":       agentID,
+				"is_prebuild":    isPrebuild,
+				"token_provided": existingToken != "",
+			}
+			if isPrebuild && existingToken != "" {
+				// check if a token was already generated for this agent.
+				// If so, this workspace is in the process of being claimed
+				// and we should reuse the token. If not, we use a new token as usual.
+				tflog.Info(ctx, "using provided agent token for prebuild", logFields)
+				token = existingToken
+			} else {
+				tflog.Info(ctx, "using a new agent token", logFields)
+			}
+
+			err := resourceData.Set("token", token)
 			if err != nil {
 				return diag.FromErr(err)
 			}
@@ -469,3 +500,16 @@ func updateInitScript(resourceData *schema.ResourceData, i interface{}) diag.Dia
 	}
 	return nil
 }
+
+// RunningAgentTokenEnvironmentVariable returns the name of the environment variable
+// that contains the token for the running agent. This is used for prebuilds, where
+// we want to reuse the same token for the next iteration of a workspace agent before
+// and after the workspace was claimed by a user.
+//
+// agentID is unused for now, but will be used as soon as we support multiple agents.
+func RunningAgentTokenEnvironmentVariable(agentID string) string {
+	agentID = "" // remove this once we need to support multiple agents per prebuilt workspace.
+
+	sum := sha256.Sum256([]byte(agentID))
+	return "CODER_RUNNING_WORKSPACE_AGENT_TOKEN_" + hex.EncodeToString(sum[:])
+}

provider/workspace.go

@@ -140,6 +140,12 @@ func workspaceDataSource() *schema.Resource {
 	}
 }
 
+// IsPrebuildEnvironmentVariable returns the name of the environment
+// variable that indicates whether the workspace was prebuilt. The value of
+// this environment variable should be set to "true" if the workspace is prebuilt.
+// Any other values, including "false" and "" will be interpreted to mean that the
+// workspace is not prebuilt. If the workspace is prebuilt, it may or may not yet
+// have been claimed by a user.
 func IsPrebuildEnvironmentVariable() string {
 	return "CODER_WORKSPACE_IS_PREBUILD"
 }