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?

This is being worked on currently; disregard for the moment.

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

spikecurtis · 2025-05-02T10:43:47Z

cli/agent.go

+
+				select {
+				case <-ctx.Done():
+					logger.Warn(ctx, "agent shutting down", slog.Error(ctx.Err()), slog.Error(context.Cause(ctx)))


Warn seems alarmist; lots of legit reasons to shut down. Also, it's unexpected to include more than one call to slog.Error

spikecurtis · 2025-05-02T10:49:24Z

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


Why is a 24 hour timeout desirable? A prebuilt workspace could remain up without a reinit for much longer than this, no?

Also, this is a quiet side effect to the client: leaving it in a different state than it started in.

spikecurtis · 2025-05-02T11:01:45Z

codersdk/agentsdk/agentsdk.go

+		if err != nil {
+			return nil, xerrors.Errorf("unmarshal reinit response: %w", err)
+		}
+		select {


At this point you've got the event; no reason to check the context for cancelation again, just return it.

spikecurtis · 2025-05-02T11:05:24Z

codersdk/agentsdk/agentsdk.go

+		if err != nil {
+			return nil, xerrors.Errorf("failed to read server-sent event: %w", err)
+		}
+		if sse.Type != codersdk.ServerSentEventTypeData {


Ignoring Ping seems OK, but probably shouldn't ignore errors.

spikecurtis · 2025-05-02T11:07:55Z

codersdk/agentsdk/agentsdk.go

+	go func() {
+		for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
+			logger.Debug(ctx, "waiting for agent reinitialization instructions")
+			reinitEvent, err := client.WaitForReinit(ctx)


We have this SSE endpoint that can stream multiple events --- so why are we hanging up after the first event just to redial the endpoint?

spikecurtis · 2025-05-02T12:14:09Z

coderd/prebuilds/claim.go

+	var once sync.Once
+	cancel := func() {
+		once.Do(func() {
+			cancelSub()


I checked PGPubsub, and there isn't any guarantee that canceling the subscription will guarantee the callback is not in progress or about to be called once more. So, it's not safe to close the channel.

Fortunately for this particular application, the reader on the other side is not expecting the channel to be closed anyway.

spikecurtis · 2025-05-02T12:14:36Z

coderd/prebuilds/claim.go

+	default:
+	}
+
+	workspaceClaims := make(chan agentsdk.ReinitializationEvent, 1)


I can't see any reason this channel should be buffered.

spikecurtis · 2025-05-02T12:17:59Z

coderd/workspaces.go

@@ -718,7 +727,7 @@ func createWorkspace(
 			Initiator(initiatorID).
 			ActiveVersion().
 			RichParameterValues(req.RichParameterValues).
-			TemplateVersionPresetID(req.TemplateVersionPresetID)
+			RunningAgentAuthTokens(agentTokensByAgentID)


agentTokensByAgentID never seems to be populated.

spikecurtis · 2025-05-02T12:29:59Z

coderd/provisionerdserver/provisionerdserver.go

+	// reused. Context: the agent token is often used in immutable attributes of workspace resource (e.g. VM/container)
+	// to initialize the agent, so if that value changes it will necessitate a replacement of that resource, thus
+	// obviating the whole point of the prebuild.
+	RunningAgentAuthTokens map[uuid.UUID]string `json:"running_agent_auth_tokens"`


I'm very nervous about putting secret credentials into the job input, which is returned when you query the job, making it much more likely to leak. I'd feel much better if provisionerdserver queried them from the database based on the agent IDs, rather than duplicating the creds in the database.

spikecurtis · 2025-05-02T12:42:57Z

coderd/provisionerdserver/provisionerdserver.go

@@ -1733,6 +1744,21 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 		if err != nil {
 			return nil, xerrors.Errorf("update workspace: %w", err)
 		}
+
+		if input.PrebuildClaimedByUser != uuid.Nil {


This feels narrow and fragile that we need to pass this imperative direction to provisionerdserver to send a signal specifically for prebuilds. It seems like a generally useful signal to know that a build completed for a given workspace... why not just always send it?

Even if we wanted to send the new owner ID in the message, there is no need to get it from the job input, we've already queried the workspace.

Signed-off-by: Danny Kopping <dannykopping@gmail.com>

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

SasSwart added 2 commits May 2, 2025 12:41

dont require ids to InsertPresetParameters

61784c9

dont require ids to InsertPresetParameters

604eb27

spikecurtis reviewed May 2, 2025

View reviewed changes

fix: set the running agent token

bf4d2cf

Signed-off-by: Danny Kopping <dannykopping@gmail.com>

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

dannykopping May 2, 2025

johnstcn May 2, 2025

SasSwart May 2, 2025

johnstcn May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis May 2, 2025

spikecurtis 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

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