fix: ensure prebuilt workspace agent tokens are reused when a prebuild agent reinitializes

SasSwart · SasSwart · commit 35e4bf87ffb1 · 2025-04-23T12:15:39.000Z
diff --git a/provisioner/terraform/executor.go b/provisioner/terraform/executor.go
@@ -261,17 +261,6 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
 	e.mut.Lock()
 	defer e.mut.Unlock()
 
-	// TODO: defunct?
-	// var isPrebuild bool
-	// for _, v := range env {
-	// 	if envName(v) == provider.IsPrebuildEnvironmentVariable() && envVar(v) == "true" {
-	// 		isPrebuild = true
-	// 		break
-	// 	}
-	// }
-
-	// _ = isPrebuild
-
 	planfilePath := getPlanFilePath(e.workdir)
 	args := []string{
 		"plan",
@@ -341,6 +330,68 @@ func onlyDataResources(sm tfjson.StateModule) tfjson.StateModule {
 	return filtered
 }
 
+func (e *executor) logResourceReplacements(ctx context.Context, plan *tfjson.Plan) {
+	if plan == nil {
+		return
+	}
+
+	if len(plan.ResourceChanges) == 0 {
+		return
+	}
+	var (
+		count        int
+		replacements = make(map[string][]string, len(plan.ResourceChanges))
+	)
+
+	for _, ch := range plan.ResourceChanges {
+		// No change, no problem!
+		if ch.Change == nil {
+			continue
+		}
+
+		// No-op change, no problem!
+		if ch.Change.Actions.NoOp() {
+			continue
+		}
+
+		// No replacements, no problem!
+		if len(ch.Change.ReplacePaths) == 0 {
+			continue
+		}
+
+		// Replacing our resources, no problem!
+		if strings.Index(ch.Type, "coder_") == 0 {
+			continue
+		}
+
+		for _, p := range ch.Change.ReplacePaths {
+			var path string
+			switch p := p.(type) {
+			case []interface{}:
+				segs := p
+				list := make([]string, 0, len(segs))
+				for _, s := range segs {
+					list = append(list, fmt.Sprintf("%v", s))
+				}
+				path = strings.Join(list, ".")
+			default:
+				path = fmt.Sprintf("%v", p)
+			}
+
+			replacements[ch.Address] = append(replacements[ch.Address], path)
+		}
+
+		count++
+	}
+
+	if count > 0 {
+		e.server.logger.Warn(ctx, "plan introduces resource changes", slog.F("count", count))
+		for n, p := range replacements {
+			e.server.logger.Warn(ctx, "resource will be replaced!", slog.F("name", n), slog.F("replacement_paths", strings.Join(p, ",")))
+		}
+	}
+}
+
 // planResources must only be called while the lock is held.
 func (e *executor) planResources(ctx, killCtx context.Context, planfilePath string) (*State, json.RawMessage, error) {
 	ctx, span := e.server.startTrace(ctx, tracing.FuncName())
@@ -351,6 +402,8 @@ func (e *executor) planResources(ctx, killCtx context.Context, planfilePath stri
 		return nil, nil, xerrors.Errorf("show terraform plan file: %w", err)
 	}
 
+	e.logResourceReplacements(ctx, plan)
+
 	rawGraph, err := e.graph(ctx, killCtx)
 	if err != nil {
 		return nil, nil, xerrors.Errorf("graph: %w", err)
diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
@@ -272,21 +272,16 @@ func provisionEnv(
 	)
 	if metadata.GetIsPrebuild() {
 		env = append(env, provider.IsPrebuildEnvironmentVariable()+"=true")
+	}
+	tokens := metadata.GetRunningAgentAuthTokens()
+	if len(tokens) == 1 {
+		env = append(env, provider.RunningAgentTokenEnvironmentVariable("")+"="+tokens[0].Token)
 	} else {
-		// TODO: provide an agentID to these functions so that we provide the right token
-		// for single agent support, we use the zero value "" as the agentID
-		// TODO: looks like we only provide agent tokens for reinit if metadata.GetIsPrebuild() is false
-		// check this for consistency wherever else we use the isPrebuild attribute from the Proto and where we use the env derived from it.
-		const singleAgentID = ""
-		tokens := metadata.GetRunningAgentAuthTokens()
-		var token string
 		for _, t := range tokens {
-			if t.AgentId == singleAgentID {
-				token = t.Token
-				break
-			}
+			// If there are multiple agents, provide all the tokens to terraform so that it can
+			// choose the correct one for each agent ID.
+			env = append(env, provider.RunningAgentTokenEnvironmentVariable(t.AgentId)+"="+t.Token)
 		}
-		env = append(env, provider.RunningAgentTokenEnvironmentVariable(singleAgentID)+"="+token)
 	}
 	for key, value := range provisionersdk.AgentScriptEnv() {
 		env = append(env, key+"="+value)
