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

Looks like it was added to main in #16951 but hasn't been released, so I guess it's fine to change the type of this index. Still, as Cian mentioned, you need to increment the protocol version in provisionerd/proto/version.go.

Is this change actually back-compatible though? Like, what will an external provisionerd that's an old version that doesn't understand this new field end up doing? I think it will fail to set the appropriate environment variable, and then the terraform provider will generate a new token for the Agent, which might invalidate the prebuild or cause the existing agent to get disconnected. Not good.

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?

codersdk/agentsdk/agentsdk.go

coderd/database/queries/presets.sql

coderd/coderdtest/coderdtest.go

cli/agent.go

codersdk/agentsdk/agentsdk.go

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?

Good point, but would you mind disregarding this given that we're going to replace this with the manifest stream by the next release?

fair enough

coderd/prebuilds/claim.go

coderd/workspaces.go

coderd/provisionerdserver/provisionerdserver.go

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 <[email protected]>

…einit

spikecurtis

The biggest remaining issue is around back compatibility for old versions of the provisionerd server.

coderd/prebuilds/claim.go

codersdk/agentsdk/agentsdk.go

coderd/workspaceagents.go

coderd/prebuilds/claim_test.go

spikecurtis · 2025-05-07T12:37:34Z

enterprise/coderd/workspaceagents_test.go

+		a[0].Scripts = []*proto.Script{
+			{
+				DisplayName: "Prebuild Test Script",
+				Script:      fmt.Sprintf("sleep 5; printenv | grep 'CODER_AGENT_TOKEN' >> %s; echo '---\n' >> %s", tempAgentLog.Name(), tempAgentLog.Name()), // Make reinitialization take long enough to assert that it happened


Writing the token doesn't seem like a good test since the token is not supposed to change between reinits. Can't we write the owner name or ID?

spikecurtis · 2025-05-07T12:39:37Z

enterprise/coderd/workspaceagents_test.go

+			Valid: true,
+			Time:  time.Now(),
+		},
+	})


Either you are running a provisioner daemon, in which case it should be responsible for completing the build job, or you are not, in which case, you need something else to send the pubsub kick that the build is done.

coderd/provisionerdserver/provisionerdserver_test.go

spikecurtis · 2025-05-07T12:49:51Z

coderd/provisionerdserver/provisionerdserver_test.go

+				_, err = srv.CompleteJob(ctx, &completedJob)
+				require.NoError(t, err)
+
+				select {


c.f. testutil.RequireReceive

spikecurtis · 2025-05-07T12:50:57Z

coderd/provisionerdserver/provisionerdserver.go

+					slog.F("workspace_id", workspace.ID), slog.Error(err))
+			}
+			for _, agent := range agents {
+				runningAgentAuthTokens = append(runningAgentAuthTokens, &sdkproto.RunningAgentAuthToken{


missing a test for this functionality.

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

dannykopping and others added 4 commits May 2, 2025 15:31

fix: set the running agent token

bf4d2cf

Signed-off-by: Danny Kopping <[email protected]>

fix: use http client without timeout like we do in connectRPCVersion

38b4f0d

Signed-off-by: Danny Kopping <[email protected]>

review notes

20df538

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

4bb3b68

…einit

spikecurtis reviewed May 7, 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

spikecurtis May 7, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

johnstcn May 2, 2025

spikecurtis May 2, 2025

SasSwart May 7, 2025

spikecurtis May 7, 2025

spikecurtis May 2, 2025

spikecurtis left a comment

spikecurtis May 7, 2025

spikecurtis May 7, 2025

spikecurtis May 7, 2025

spikecurtis May 7, 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

spikecurtis left a 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