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,
-		LogLevel:         b.logLevel,
-		IsPrebuild:       b.prebuild,
+		WorkspaceBuildID:        workspaceBuildID,
+		LogLevel:                b.logLevel,
+		IsPrebuild:              b.prebuild,
+		RunningWorkspaceAgentID: b.runningWorkspaceAgentID,
 	})
 	if err != nil {
 		return nil, nil, nil, BuildError{