WIP CODER_AGENT_TOKEN reuse

Signed-off-by: Danny Kopping <danny@coder.com>

coderd/provisionerdserver/provisionerdserver.go

@@ -500,6 +500,17 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 		for _, group := range ownerGroups {
 			ownerGroupNames = append(ownerGroupNames, group.Group.Name)
 		}
+		var runningWorkspaceAgentToken string
+		if input.RunningWorkspaceAgentID != uuid.Nil {
+			agent, err := s.Database.GetWorkspaceAgentByID(ctx, input.RunningWorkspaceAgentID)
+			if err != nil {
+				s.Logger.Warn(ctx, "failed to retrieve running workspace agent by ID; this may affect prebuilds",
+					slog.F("workspace_agent_id", input.RunningWorkspaceAgentID),
+					slog.F("job_id", job.ID))
+			} else {
+				runningWorkspaceAgentToken = agent.AuthToken.String()
+			}
+		}
 
 		msg, err := json.Marshal(wspubsub.WorkspaceEvent{
 			Kind:        wspubsub.WorkspaceEventKindStateChange,
@@ -622,6 +633,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 					WorkspaceBuildId:              workspaceBuild.ID.String(),
 					WorkspaceOwnerLoginType:       string(owner.LoginType),
 					IsPrebuild:                    input.IsPrebuild,
+					RunningWorkspaceAgentToken:    runningWorkspaceAgentToken,
 				},
 				LogLevel: input.LogLevel,
 			},
@@ -2347,7 +2359,14 @@ type WorkspaceProvisionJob struct {
 	WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
 	DryRun           bool      `json:"dry_run"`
 	IsPrebuild       bool      `json:"is_prebuild,omitempty"`
-	LogLevel         string    `json:"log_level,omitempty"`
+	// RunningWorkspaceAgentID is *only* used for prebuilds. We pass it down when we want to rebuild a prebuilt workspace
+	// but not generate a new agent token. The provisionerdserver will retrieve this token and push it down to
+	// the provisioner (and ultimately to the `coder_agent` resource in the Terraform provider) where it will be
+	// reused. Context: the agent token is often used in immutable attributes of workspace resource (e.g. VM/container)
+	// to initialize the agent, so if that value changes it will necessitate a replacement of that resource, thus
+	// obviating the whole point of the prebuild.
+	RunningWorkspaceAgentID uuid.UUID `json:"running_workspace_agent_id"`
+	LogLevel                string    `json:"log_level,omitempty"`
 }
 
 // TemplateVersionDryRunJob is the payload for the "template_version_dry_run" job type.

coderd/workspaces.go

@@ -627,6 +627,8 @@ func createWorkspace(
 		provisionerJob     *database.ProvisionerJob
 		workspaceBuild     *database.WorkspaceBuild
 		provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
+
+		runningWorkspaceAgentID uuid.UUID
 	)
 	err = api.Database.InTx(func(db database.Store) error {
 		var claimedWorkspace *database.Workspace
@@ -663,8 +665,17 @@ func createWorkspace(
 				api.Logger.Warn(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
 				goto regularPath
 			}
-
 			claimedWorkspace = &lookup
+
+			agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ownerCtx, claimedWorkspace.ID)
+			if err != nil {
+				api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
+					slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
+			}
+			if len(agents) >= 1 {
+				// TODO: handle multiple agents
+				runningWorkspaceAgentID = agents[0].ID
+			}
 		}
 
 	regularPath:
@@ -710,7 +721,8 @@ func createWorkspace(
 			Reason(database.BuildReasonInitiator).
 			Initiator(initiatorID).
 			ActiveVersion().
-			RichParameterValues(req.RichParameterValues)
+			RichParameterValues(req.RichParameterValues).
+			RunningWorkspaceAgentID(runningWorkspaceAgentID)
 		if req.TemplateVersionID != uuid.Nil {
 			builder = builder.VersionID(req.TemplateVersionID)
 		}

coderd/wsbuilder/wsbuilder.go

@@ -73,6 +73,7 @@ type Builder struct {
 	parameterNames               *[]string
 	parameterValues              *[]string
 	prebuild                     bool
+	runningWorkspaceAgentID      uuid.UUID
 
 	verifyNoLegacyParametersOnce bool
 }
@@ -175,6 +176,13 @@ func (b Builder) MarkPrebuild() Builder {
 	return b
 }
 
+// RunningWorkspaceAgentID is only used for prebuilds; see the associated field in `provisionerdserver.WorkspaceProvisionJob`.
+func (b Builder) RunningWorkspaceAgentID(id uuid.UUID) Builder {
+	// nolint: revive
+	b.runningWorkspaceAgentID = id
+	return b
+}
+
 // SetLastWorkspaceBuildInTx prepopulates the Builder's cache with the last workspace build.  This allows us
 // to avoid a repeated database query when the Builder's caller also needs the workspace build, e.g. auto-start &
 // auto-stop.
@@ -300,9 +308,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
 
 	workspaceBuildID := uuid.New()
 	input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
-		WorkspaceBuildID: workspaceBuildID,
+		WorkspaceBuildID:        workspaceBuildID,
-		LogLevel:         b.logLevel,
+		LogLevel:                b.logLevel,
-		IsPrebuild:       b.prebuild,
+		IsPrebuild:              b.prebuild,
+		RunningWorkspaceAgentID: b.runningWorkspaceAgentID,
 	})
 	if err != nil {
 		return nil, nil, nil, BuildError{

provisioner/terraform/provision.go

@@ -263,7 +263,9 @@ func provisionEnv(
 		"CODER_WORKSPACE_BUILD_ID="+metadata.GetWorkspaceBuildId(),
 	)
 	if metadata.GetIsPrebuild() {
-		env = append(env, "CODER_WORKSPACE_IS_PREBUILD=true")
+		env = append(env, provider.IsPrebuildEnvironmentVariable()+"=true")
+	} else {
+		env = append(env, provider.RunningAgentTokenEnvironmentVariable()+"="+metadata.GetRunningWorkspaceAgentToken())
 	}
 	for key, value := range provisionersdk.AgentScriptEnv() {
 		env = append(env, key+"="+value)

provisioner/terraform/safeenv.go

@@ -30,6 +30,14 @@ func envName(env string) string {
 	return ""
 }
 
+func envVar(env string) string {
+	parts := strings.SplitN(env, "=", 1)
+	if len(parts) > 0 {
+		return parts[1]
+	}
+	return ""
+}
+
 func isCanarySet(env []string) bool {
 	for _, e := range env {
 		if envName(e) == unsafeEnvCanary {

provisionersdk/proto/provisioner.proto

@@ -276,6 +276,7 @@ message Metadata {
     string workspace_build_id = 17;
     string workspace_owner_login_type =  18;
     bool   is_prebuild = 19;
+    string running_workspace_agent_token = 20;
 }
 
 // Config represents execution configuration shared by all subsequent requests in the Session

site/e2e/provisionerGenerated.ts

@@ -290,6 +290,7 @@ export interface Metadata {
   workspaceBuildId: string;
   workspaceOwnerLoginType: string;
   isPrebuild: boolean;
+  runningWorkspaceAgentToken: string;
 }
 
 /** Config represents execution configuration shared by all subsequent requests in the Session */
@@ -965,6 +966,9 @@ export const Metadata = {
     if (message.isPrebuild === true) {
       writer.uint32(152).bool(message.isPrebuild);
     }
+    if (message.runningWorkspaceAgentToken !== "") {
+      writer.uint32(162).string(message.runningWorkspaceAgentToken);
+    }
     return writer;
   },
 };