feat: reinitialize agents when a prebuilt workspace is claimed #17475

SasSwart · 2025-04-21T10:57:53Z

This pull request allows coder workspace agents to be reinitialized when a prebuilt workspace is claimed by a user. This facilitates the transfer of ownership between the anonymous prebuilds system user and the new owner of the workspace.

Only a single agent per prebuilt workspace is supported for now, but plumbing has already been done to facilitate the seamless transition to multi-agent support.

agent/metrics.go

…d agent reinitializes

cli/agent.go

coderd/agentapi/api.go

coderd/agentapi/manifest.go

coderd/coderdtest/coderdtest.go

coderd/workspaceagents.go

provisioner/terraform/executor.go

provisioner/terraform/provision.go

… conditions

coderd/provisionerdserver/provisionerdserver.go

coderd/provisionerdserver/provisionerdserver_test.go

dannykopping · 2025-04-28T09:19:06Z

coderd/provisionerdserver/provisionerdserver_test.go

+				// Complete the job, optionally triggering workspace agent reinitialization:
+
+				completedJob := proto.CompletedJob{
+					JobId: job.ID.String(),
+					Type: &proto.CompletedJob_WorkspaceBuild_{
+						WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{},
+					},
+				}
+				_, err = srv.CompleteJob(ctx, &completedJob)
+				require.NoError(t, err)


Why is this necessary?

can you elaborate by defining 'this'?

Does CompleteJob need to occur for the test to work?

ah. CompleteJob is the thing that publishes to the channel. Its the entire point of this test, yes. If CompleteJob ever stops publishing, this should fail.

coderd/provisionerdserver/provisionerdserver_test.go

…einit

coderd/prebuilds/claim.go

dannykopping · 2025-05-01T08:46:06Z

coderd/prebuilds/claim.go

+	}
+
+	go func() {
+		<-ctx.Done()


Howcome we wait for this to complete before cancel-ing? To protect messages from being passed to workspaceClaims after we close the channel.

we expect the user to defer cancel() where they call this. But we can't enforce that. As a defensive measure to protect against leaking the goroutine and pubsub connection, as well as the channel, we call cancel() ourselves when the context expires.

coderd/prebuilds/claim.go

dannykopping · 2025-05-01T08:50:16Z

coderd/prebuilds/claim.go

+}
+
+func StreamAgentReinitEvents(ctx context.Context, logger slog.Logger, rw http.ResponseWriter, r *http.Request, reinitEvents <-chan agentsdk.ReinitializationEvent) {
+	sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(rw, r)


We can fix this later, but I think this might be inappropriate at this layer.
This seems like a detail of the HTTP API, and we aren't testing this directly so I'd argue this code could go back to coderd/workspaceagents.go.

coderd/prebuilds/claim_test.go

coderd/provisionerdserver/provisionerdserver.go

coderd/workspaceagents.go

dannykopping · 2025-05-01T09:00:33Z

codersdk/agentsdk/agentsdk.go

-	Message string                 `json:"message"`
-	Reason  ReinitializationReason `json:"reason"`
+	WorkspaceID uuid.UUID
+	UserID      uuid.UUID


This bleeds the abstraction a bit because there's nothing prebuilds-specific in this func. The workspace's agent is what reinitializes, and the user should be irrelevant here. Do we need this?

…einit

SasSwart · 2025-05-02T07:33:54Z

provisionersdk/proto/provisioner.proto

@@ -294,7 +299,7 @@ message Metadata {
    string workspace_owner_login_type =  18;
    repeated Role workspace_owner_rbac_roles = 19;
    bool is_prebuild = 20;
-    string running_workspace_agent_token = 21;


self-review: it concerns me that running_workspace_agent_token made it into main. I'd love to undo that before this is released but it might already be too late.

running_workspace_agent_token is defunct and meant to be replaced by running_agent_auth_tokens.
The former only supports a single agent. The latter supports multiple agents.
Nothing will break if the former stays, but it pollutes the API

johnstcn · 2025-05-02T08:32:22Z

provisionersdk/proto/provisioner.proto

+message RunningAgentAuthToken {
+    string agent_id = 1;
+    string token = 2;
+}
+


When adding a new field, you need to increment proto.CurrentVersion. We tend to add a comment as well explaning the versioning history. See provisionerd/proto/version.go.

johnstcn · 2025-05-02T08:39:48Z

provisioner/terraform/provision.go

+	tokens := metadata.GetRunningAgentAuthTokens()
+	if len(tokens) == 1 {
+		env = append(env, provider.RunningAgentTokenEnvironmentVariable("")+"="+tokens[0].Token)
+	} else {
+		// Not currently supported, but added for forward-compatibility
+		for _, t := range tokens {
+			// 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)
+		}
+	}


Shouldn't changes here also introduce tests in provisioner_test.go?

johnstcn · 2025-05-02T08:42:41Z

codersdk/agentsdk/agentsdk_test.go

+		defer srv.Close()
+
+		requestCtx := testutil.Context(t, testutil.WaitShort)
+		req, err := http.NewRequestWithContext(requestCtx, "GET", srv.URL, nil)


Why not use srv.Client() here?

johnstcn · 2025-05-02T08:49:43Z

codersdk/agentsdk/agentsdk.go

+// - prebuild claimed: a prebuilt workspace is claimed, so the agent must reinitialize.
+func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, error) {
+	// TODO: allow configuring httpclient
+	c.SDK.HTTPClient.Timeout = time.Hour * 24


Does this essentially place an upper limit of 24 hours on the maximum lifetime of a prebuild?

johnstcn · 2025-05-02T08:51:19Z

coderd/database/queries/presets.sql

@@ -16,8 +18,9 @@ VALUES (

 -- name: InsertPresetParameters :many
 INSERT INTO
-	template_version_preset_parameters (template_version_preset_id, name, value)
+	template_version_preset_parameters (id, template_version_preset_id, name, value)


Why is it important to be able to control the ID of a template version preset?

when setting up test environments where one needs to assert based on the relationships between presets, their parameters and other parts of the DB, we need to be able to express that relationship in the test by setting the ids explicitly.

johnstcn · 2025-05-02T08:52:26Z

coderd/coderdtest/coderdtest.go

+//
+// waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
+// waiter.WaitFor(coderdtest.AgentsReady)
+type WaitForCriterium func(agent codersdk.WorkspaceAgent) bool


suggestion, nonblocking: rename to WaitForAgentFn for clarity

github-actions bot assigned SasSwart Apr 21, 2025

SasSwart changed the title ~~WIP: agent reinitialization~~ feat: reinitialize agents when a prebuilt workspace is claimed Apr 21, 2025

dannykopping mentioned this pull request Apr 21, 2025

feat: implement claiming of prebuilt workspaces #17458

Merged

dannykopping reviewed Apr 21, 2025

View reviewed changes

agent/metrics.go Outdated Show resolved Hide resolved

SasSwart force-pushed the jjs/prebuilds-agent-reinit branch from 35e4bf8 to 18da76e Compare April 23, 2025 13:49

evgeniy-scherbina force-pushed the yevhenii/512-claim-prebuild branch from fe569d4 to fcdbba8 Compare April 23, 2025 15:23

SasSwart force-pushed the jjs/prebuilds-agent-reinit branch from cc25406 to 26dbc3a Compare April 24, 2025 12:34

Base automatically changed from yevhenii/512-claim-prebuild to main April 24, 2025 13:39

SasSwart added 6 commits April 25, 2025 08:31

WIP: agent reinitialization

c09c9b9

fix assignment to nil map

476fe71

fix: ensure prebuilt workspace agent tokens are reused when a prebuil…

8c8bca6

…d agent reinitializes

test agent reinitialization

7ce4eea

remove defunct metric

52ac64e

Remove todo

362db7c

SasSwart force-pushed the jjs/prebuilds-agent-reinit branch from ec9ed29 to 362db7c Compare April 25, 2025 08:32

dannykopping reviewed Apr 25, 2025

View reviewed changes

SasSwart added 2 commits April 28, 2025 09:03

test that we trigger workspace agent reinitialization under the right…

dcc7379

… conditions

slight improvements to a test

ff66b3f

dannykopping reviewed Apr 28, 2025

View reviewed changes

review notes to improve legibility

efff5d9

dannykopping reviewed Apr 28, 2025

View reviewed changes

coderd/provisionerdserver/provisionerdserver_test.go Outdated Show resolved Hide resolved

SasSwart added 6 commits April 29, 2025 11:27

add an integration test for prebuilt workspace agent reinitialization

cebd5db

Merge remote-tracking branch 'origin/main' into jjs/prebuilds-agent-r…

2679138

…einit

enable the premium license in a prebuilds integration test

9feebef

encapsulate WaitForReinitLoop for easier testing

b117b5c

introduce unit testable abstraction layers

a22b414

test workspace claim pubsub

9bbd2c7

dannykopping reviewed May 1, 2025

View reviewed changes

SasSwart added 2 commits May 1, 2025 18:49

add tests for agent reinitialization

5804201

review notes

7e8dcee

SasSwart added 3 commits May 1, 2025 19:08

Merge remote-tracking branch 'origin/main' into jjs/prebuilds-agent-r…

725f97b

…einit

make fmt lint

a9b1567

remove go mod replace

21ee970

SasSwart marked this pull request as ready for review May 1, 2025 19:34

SasSwart requested review from spikecurtis and johnstcn as code owners May 1, 2025 19:34

SasSwart added 5 commits May 1, 2025 19:35

remove defunct logging

e54d7e7

update dependency on terraform-provider-coder

2799858

update dependency on terraform-provider-coder

1d93003

go mod tidy

763fc12

make -B gen

0f879c7

SasSwart commented May 2, 2025

View reviewed changes

johnstcn reviewed May 2, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: reinitialize agents when a prebuilt workspace is claimed #17475

feat: reinitialize agents when a prebuilt workspace is claimed #17475

SasSwart commented Apr 21, 2025

dannykopping Apr 28, 2025

SasSwart Apr 28, 2025

dannykopping Apr 28, 2025

SasSwart May 1, 2025

dannykopping May 1, 2025

SasSwart May 1, 2025

dannykopping May 1, 2025

dannykopping May 1, 2025

SasSwart May 2, 2025

SasSwart May 2, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

SasSwart May 2, 2025

johnstcn May 2, 2025

feat: reinitialize agents when a prebuilt workspace is claimed #17475

Are you sure you want to change the base?

feat: reinitialize agents when a prebuilt workspace is claimed #17475

Conversation

SasSwart commented Apr 21, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment