fix(coderd): only send prebuild claim reinit for the claim build (#26548)

## Problem

#23108 made prebuild claim delivery durable: when an agent connects to
`/api/v2/workspaceagents/me/reinit?wait=true`, the handler checks
whether the workspace's first build was created by the prebuilds system
user and whether its latest build succeeded, and if so pre-seeds a
`prebuild_claimed` reinitialization event in case the original pubsub
event was missed.

The check does not verify that the latest build is the claim build, so
it keeps firing for the rest of the workspace's life. Any workspace that
was claimed from a prebuild receives a spurious "prebuild claimed"
reinit every time its agent (re)opens the `/reinit` connection: after
every agent restart, every coderd deploy or replica restart, and every
dropped SSE connection. Each one shuts the agent down and reinitializes
it, killing SSH/IDE sessions and re-running startup scripts. In our
deployment, where most workspaces are claimed from prebuilds, this
caused fleet-wide "agent disconnected" blips whenever a coderd replica
restarted, and a few workspaces whose container exits when the agent
restarts went into a restart loop every 15-60 minutes. The agent-side
dedup (`lastOwnerID` in `cli/agent.go`) only suppresses the second event
within one agent process, so every new agent process takes at least one
spurious restart.

## Fix

Only seed the reinitialization event while the latest build is the claim
build itself, determined from the build job's input
(`prebuilt_workspace_stage`), the same signal `provisionerdserver` uses
when publishing the claim event:

- Latest build is the claim build: behavior unchanged (seed when the job
succeeded, 409 when it failed permanently, wait on pubsub while it is in
progress).
- Latest build is still a prebuilds-initiated build (claim build not
created yet): fall through to the pubsub subscription, which delivers
the claim event when the claim build completes.
- Latest build is any later user-initiated build: the claim was already
handled, so return 409 and the agent stops polling, the same as a
regular workspace.

`dbfake` gains a `MarkPrebuiltWorkspaceClaim()` builder option so tests
can model claim builds' job input, and the existing `TestReinit` claim
subtests now use it. A new subtest covers the long-claimed workspace
case.

One deliberate behavior change worth calling out: if a claim build fails
and the owner retries with another start build, the handler now returns
409 for that retry build rather than seeding a reinit. This matches the
existing treatment of failed claim builds as terminal for the reinit
poller.

## Verification

- `go test ./coderd/ -run TestReinit` against Postgres 17: all subtests
pass, including the new `workspace claimed in the past gets 409` case.
- `gofmt`, `go vet`, and `golangci-lint` (v1.64.8) are clean on the
touched packages.
- The fix mirrors behavior validated by hand against an affected
deployment: for a long-claimed workspace, `/reinit?wait=true` returned
the seeded `prebuild_claimed` event on every connection before the
change and a 409 afterwards.

Note: this branch was prepared in an environment without the full local
toolchain, so the repo's pre-commit hook (`make pre-commit`) was not run
locally; relying on CI for the full gen/fmt/lint suite. Opening as a
draft mainly to report the issue and propose a fix; happy to rework it
to the maintainers' preferred approach.

---------

Co-authored-by: Sas Swart <sas.swart.cdk@gmail.com>

coderd/database/dbfake/dbfake.go

@@ -69,6 +69,8 @@ type WorkspaceBuildBuilder struct {
 	jobErrorCode string // Error code for failed jobs
 
 	provisionerState []byte
+
+	prebuiltWorkspaceBuildStage sdkproto.PrebuiltWorkspaceBuildStage
 }
 
 // BuilderOption is a functional option for customizing job timestamps
@@ -149,6 +151,14 @@ func (b WorkspaceBuildBuilder) ProvisionerState(state []byte) WorkspaceBuildBuil
 	return b
 }
 
+// MarkPrebuiltWorkspaceClaim marks the build's provisioner job as the claim
+// of a prebuilt workspace, mirroring wsbuilder.MarkPrebuiltWorkspaceClaim.
+func (b WorkspaceBuildBuilder) MarkPrebuiltWorkspaceClaim() WorkspaceBuildBuilder {
+	//nolint: revive // returns modified struct
+	b.prebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CLAIM
+	return b
+}
+
 func (b WorkspaceBuildBuilder) Resource(resource ...*sdkproto.Resource) WorkspaceBuildBuilder {
 	//nolint: revive // returns modified struct
 	b.resources = append(b.resources, resource...)
@@ -368,7 +378,8 @@ func (b WorkspaceBuildBuilder) doInTX() WorkspaceResponse {
 
 	// Create a provisioner job for the build!
 	payload, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
-		WorkspaceBuildID: b.seed.ID,
+		WorkspaceBuildID:            b.seed.ID,
+		PrebuiltWorkspaceBuildStage: b.prebuiltWorkspaceBuildStage,
 	})
 	require.NoError(b.t, err)
 

coderd/workspaceagents.go

@@ -37,6 +37,7 @@ import (
 	"github.com/coder/coder/v2/coderd/httpmw/loggermw"
 	"github.com/coder/coder/v2/coderd/jwtutils"
 	"github.com/coder/coder/v2/coderd/prebuilds"
+	"github.com/coder/coder/v2/coderd/provisionerdserver"
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/rbac/policy"
 	"github.com/coder/coder/v2/coderd/telemetry"
@@ -1537,14 +1538,13 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		// This workspace was a prebuild that got claimed. Check if
-		// the claim build completed successfully before sending
-		// reinit. We assume the latest build is the claim build
-		// (build 2). If a third build (e.g. a restart) starts
-		// between the claim and the agent's reconnection, this
-		// would check that build instead. The window is extremely
-		// small in practice, and a restart would trigger its own
-		// reinit path.
+		// This workspace was a prebuild that got claimed. The seeded
+		// reinit below recovers a claim event that was missed while
+		// the agent's /reinit connection was down. It only applies
+		// while the latest build is the claim build itself, which the
+		// build's provisioner job input records, mirroring the check
+		// the provisioner server uses when publishing the claim
+		// event.
 		latestBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
 		if err != nil {
 			log.Error(ctx, "failed to get latest workspace build", slog.Error(err))
@@ -1557,43 +1557,73 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
 			httpapi.InternalServerError(rw, xerrors.New("failed to get provisioner job"))
 			return
 		}
-
-		if job.CompletedAt.Valid && !job.Error.Valid {
-			// Claim build succeeded — cancel the pubsub
-			// subscription (no longer needed) and swap in a
-			// pre-seeded channel so the transmitter delivers
-			// exactly one reinit event.
-			cancelSub()
-			seeded := make(chan agentsdk.ReinitializationEvent, 1)
-			seeded <- agentsdk.ReinitializationEvent{
-				WorkspaceID: workspace.ID,
-				Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
-				OwnerID:     workspace.OwnerID,
-			}
-			reinitEvents = seeded
-		} else if job.CompletedAt.Valid && job.Error.Valid {
-			// Claim build failed permanently. Return 409 so the
-			// agent treats this as terminal and stops retrying
-			// (WaitForReinitLoop exits on any 409).
-			cancelSub()
-			log.Warn(ctx, "claim build failed",
-				slog.F("job_id", job.ID),
-				slog.F("error", job.Error.String))
-			httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
-				Message: "Claim build failed permanently.",
-				Detail:  job.Error.String,
-			})
+		var jobInput provisionerdserver.WorkspaceProvisionJob
+		if err := json.Unmarshal(job.Input, &jobInput); err != nil {
+			log.Error(ctx, "failed to unmarshal provisioner job input", slog.Error(err))
+			httpapi.InternalServerError(rw, xerrors.New("failed to unmarshal provisioner job input"))
 			return
 		}
 
-		// Claim build still in progress — fall through to the
-		// transmitter. The pubsub subscription (set up above)
-		// will deliver the event when the build completes
-		// successfully. Note: FailJob does not publish a claim
-		// event, so a failed in-progress build will leave the
-		// agent blocking here until it disconnects and
-		// reconnects (at which point the durable check above
-		// handles it).
+		switch {
+		case jobInput.PrebuiltWorkspaceBuildStage.IsPrebuiltWorkspaceClaim():
+			if job.CompletedAt.Valid && !job.Error.Valid {
+				// Claim build succeeded: cancel the pubsub
+				// subscription (no longer needed) and swap in a
+				// pre-seeded channel so the transmitter delivers
+				// exactly one reinit event.
+				cancelSub()
+				seeded := make(chan agentsdk.ReinitializationEvent, 1)
+				seeded <- agentsdk.ReinitializationEvent{
+					WorkspaceID: workspace.ID,
+					Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+					OwnerID:     workspace.OwnerID,
+				}
+				reinitEvents = seeded
+			} else if job.CompletedAt.Valid && job.Error.Valid {
+				// Claim build failed permanently. Return 409 so the
+				// agent treats this as terminal and stops retrying
+				// (WaitForReinitLoop exits on any 409).
+				cancelSub()
+				log.Warn(ctx, "claim build failed",
+					slog.F("job_id", job.ID),
+					slog.F("error", job.Error.String))
+				httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
+					Message: "Claim build failed permanently.",
+					Detail:  job.Error.String,
+				})
+				return
+			}
+			// Claim build still in progress: proceed to the
+			// transmitter below. The pubsub subscription (set up
+			// above) will deliver the event when the build completes
+			// successfully. Note: FailJob does not publish a claim
+			// event, so a failed in-progress build will leave the
+			// agent blocking here until it disconnects and
+			// reconnects (at which point the durable check above
+			// handles it).
+		case latestBuild.InitiatorID == database.PrebuildsSystemUserID:
+			// The workspace owner has changed but the claim build has
+			// not been created yet. Proceed to the transmitter below;
+			// the pubsub subscription set up above delivers the claim
+			// event once the claim build completes.
+		default:
+			// The latest build is a user-initiated build other than
+			// the claim build, so the claim has already been handled.
+			// Re-sending the reinit event here would needlessly
+			// restart the agent of a long-claimed workspace on every
+			// reconnection. Return 409 so the agent stops polling,
+			// the same as a regular workspace.
+			log.Debug(ctx, "prebuild claim already handled, stopping reinit polling",
+				slog.F("job_id", job.ID),
+				slog.F("latest_build_id", latestBuild.ID),
+				slog.F("latest_build_number", latestBuild.BuildNumber))
+			cancelSub()
+			httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
+				Message: "Workspace is not a prebuilt workspace waiting to be claimed.",
+				Detail:  "The prebuild claim for this workspace has already been handled by an earlier build.",
+			})
+			return
+		}
 	}
 
 	transmitter := agentsdk.NewSSEAgentReinitTransmitter(log, rw, r)

coderd/workspaceagents_test.go

@@ -3464,6 +3464,7 @@ func TestReinit(t *testing.T) {
 				InitiatorID:       claimerID,
 				Transition:        database.WorkspaceTransitionStart,
 			}).
+			MarkPrebuiltWorkspaceClaim().
 			WithAgent()
 		if !complete {
 			builder = builder.Starting()
@@ -3562,6 +3563,52 @@ func TestReinit(t *testing.T) {
 		require.Equal(t, user.UserID, reinitEvent.OwnerID)
 	})
 
+	// Verifies that the durable claim check only applies while the
+	// latest build is the claim build. A workspace that was claimed
+	// in the past and has since had user-initiated builds must get a
+	// 409 instead of another reinit, otherwise its agent would be
+	// restarted on every /reinit reconnection for the rest of the
+	// workspace's life.
+	t.Run("workspace claimed in the past gets 409", func(t *testing.T) {
+		t.Parallel()
+
+		db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t)
+		client := coderdtest.New(t, &coderdtest.Options{
+			Database: db,
+			Pubsub:   ps,
+		})
+		user := coderdtest.CreateFirstUser(t, client)
+
+		// Create an unclaimed prebuild (build 1, completed) and claim
+		// it (build 2, completed).
+		r := setupPrebuildWorkspace(t, db, user.OrganizationID)
+		claimPrebuild(t, db, sqlDB, r.Workspace, user.UserID, r.TemplateVersion.ID, true)
+
+		// A later build initiated by the owner (e.g. a restart) means
+		// the claim has already been handled.
+		ws := r.Workspace
+		ws.OwnerID = user.UserID
+		laterR := dbfake.WorkspaceBuild(t, db, ws).
+			Seed(database.WorkspaceBuild{
+				TemplateVersionID: r.TemplateVersion.ID,
+				BuildNumber:       3,
+				InitiatorID:       user.UserID,
+				Transition:        database.WorkspaceTransitionStart,
+			}).
+			WithAgent().
+			Do()
+
+		agentCtx := testutil.Context(t, testutil.WaitShort)
+		agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(laterR.AgentToken))
+
+		// WaitForReinit should return an error wrapping a 409.
+		_, err := agentClient.WaitForReinit(agentCtx)
+		require.Error(t, err)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusConflict, sdkErr.StatusCode())
+	})
+
 	// Verifies that when the claim build completed with an error,
 	// the handler returns 409 so the agent treats it as terminal
 	// and stops retrying (WaitForReinitLoop exits on any 409).