From 5bd3f26f92677ce719d89cb7533f7b45bc0dccd2 Mon Sep 17 00:00:00 2001 From: leonzh Date: Fri, 19 Jun 2026 13:47:41 +0000 Subject: [PATCH 1/2] fix(coderd): only send prebuild claim reinit for the claim build The durable claim check in workspaceAgentReinit pre-seeded a prebuild_claimed reinitialization event whenever a workspace's first build was created by the prebuilds system user and its latest build succeeded. For workspaces claimed long ago this re-fired on every /reinit reconnection, restarting the agent (dropping SSH and IDE sessions) each time, and could loop indefinitely when the restart took down the workspace container. Only seed the event while the latest build is the claim build itself, as recorded in the build job's input, mirroring the check the provisioner server uses when publishing the claim event. Workspaces with any later build now receive a 409 so agents stop polling, the same as regular workspaces. --- coderd/database/dbfake/dbfake.go | 13 +++- coderd/workspaceagents.go | 105 ++++++++++++++++++++----------- coderd/workspaceagents_test.go | 47 ++++++++++++++ 3 files changed, 126 insertions(+), 39 deletions(-) diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go index 0b859a4fb1c66..82b66f504aaf7 100644 --- 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 index 0dc91010ccfab..f430c753c8743 100644 --- 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" @@ -1540,14 +1541,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)) @@ -1560,43 +1560,72 @@ 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). + // 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). + case latestBuild.InitiatorID == database.PrebuildsSystemUserID: + // The workspace owner has changed but the claim build has + // not been created yet. Fall through to the transmitter; + // 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("latest_build_id", latestBuild.ID), + slog.F("latest_build_number", latestBuild.BuildNumber)) 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, + 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 index 583332ebbaa80..1e52d1e35cff5 100644 --- 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). From ad010e0c479d96dc8510eac256175d8b8d37e806 Mon Sep 17 00:00:00 2001 From: leonzh Date: Mon, 22 Jun 2026 15:16:45 +0000 Subject: [PATCH 2/2] Include job_id in reinit log and clarify switch case comments --- coderd/workspaceagents.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index f430c753c8743..3b24fdeca17d6 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -1596,9 +1596,9 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) { }) 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 + // 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 @@ -1606,7 +1606,7 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) { // handles it). case latestBuild.InitiatorID == database.PrebuildsSystemUserID: // The workspace owner has changed but the claim build has - // not been created yet. Fall through to the transmitter; + // 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: @@ -1617,6 +1617,7 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) { // 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()