coder · SasSwart · Jun 23, 2026 · Jun 19, 2026 · Jun 22, 2026 · Jun 23, 2026
diff --git a/coderd/database/dbfake/dbfake.go b/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)
 

diff --git a/coderd/workspaceagents.go b/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
 		}
+		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
+		}
 
-		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,
+		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
 			}
-			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",
+			// 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("error", job.Error.String))
+				slog.F("latest_build_id", latestBuild.ID),
+				slog.F("latest_build_number", latestBuild.BuildNumber))
+			cancelSub()
 			httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
-				Message: "Claim build failed permanently.",
-				Detail:  job.Error.String,
+				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
 		}
-
-		// 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).
 	}
 
 	transmitter := agentsdk.NewSSEAgentReinitTransmitter(log, rw, r)

diff --git a/coderd/workspaceagents_test.go b/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).