May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025 · Apr 24, 2025 · Apr 24, 2025
diff --git a/cli/agent.go b/cli/agent.go
 Handler: func(inv *serpent.Invocation) error {
 ctx, cancel := context.WithCancelCause(inv.Context())
 defer func() {
 cancel(xerrors.New("defer"))
 cancel(xerrors.New("agent exited"))
 }()

 var (
 // TODO: timeout ok?
 reinitCtx, reinitCancel := context.WithTimeout(context.Background(), time.Hour*24)
 defer reinitCancel()
 reinitEvents := make(chan agentsdk.ReinitializationResponse)
 reinitEvents := make(chan agentsdk.ReinitializationEvent)

 go func() {
 // Retry to wait for reinit, main context cancels the retrier.
diff --git a/cli/agent_test.go b/cli/agent_test.go

 // Check that the agent is in a happy steady state
 waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
 waiter.WaitFor(coderdtest.AgentReady)
 waiter.WaitFor(coderdtest.AgentsReady)

 // Trigger reinitialization
 channel := agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
 err := pubsub.Publish(channel, []byte(user.UserID.String()))
 require.NoError(t, err)

 // Check that the agent reinitializes
 waiter.WaitFor(coderdtest.AgentNotReady)
 waiter.WaitFor(coderdtest.AgentsNotReady)

 // Check that reinitialization completed
 waiter.WaitFor(coderdtest.AgentReady)
 waiter.WaitFor(coderdtest.AgentsReady)
 }

 func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
 return w
 }

 // WaitForCriterium represents a boolean assertion to be made against each agent
 // that a given WorkspaceAgentWaited knows about. Each WaitForCriterium should apply
 // the check to a single agent, but it should be named for plural, because `func (w WorkspaceAgentWaiter) WaitFor`
 // applies the check to all agents that it is aware of. This ensures that the public API of the waiter
 // reads correctly. For example:
 //
 // waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
 // waiter.WaitFor(coderdtest.AgentsReady)
 type WaitForCriterium func(agent codersdk.WorkspaceAgent) bool

 func AgentReady(agent codersdk.WorkspaceAgent) bool {
 // AgentsReady checks that the latest lifecycle state of an agent is "Ready".
 func AgentsReady(agent codersdk.WorkspaceAgent) bool {
 return agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady
 }

 func AgentNotReady(agent codersdk.WorkspaceAgent) bool {
 return !AgentReady(agent)
 // AgentsReady checks that the latest lifecycle state of an agent is anything except "Ready".
 func AgentsNotReady(agent codersdk.WorkspaceAgent) bool {
 return !AgentsReady(agent)
 }

 func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForCriterium) {
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go

 channel := agentsdk.PrebuildClaimedChannel(workspace.ID)
 if err := s.Pubsub.Publish(channel, []byte(input.PrebuildClaimedByUser.String())); err != nil {
 s.Logger.Error(ctx, "failed topublish message to instruct prebuilt workspace agent reinitialization", slog.Error(err))
 s.Logger.Error(ctx, "failed totrigger prebuilt workspace agent reinitialization", slog.Error(err))
 }
 }
 case *proto.CompletedJob_TemplateDryRun_:
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
 t.Run(tc.name, func(t *testing.T) {
 t.Parallel()

 //Setup starting state:
 //GIVEN an enqueued provisioner job and its dependencies:

 userID := uuid.New()

 srv, db, ps, pd := setup(t, false, &overrides{})

 buildID := uuid.New()
 scheduledJobInput := provisionerdserver.WorkspaceProvisionJob{
 jobInput := provisionerdserver.WorkspaceProvisionJob{
 WorkspaceBuildID: buildID,
 }
 if tc.shouldReinitializeAgent { // This is the key lever in the test
 scheduledJobInput.PrebuildClaimedByUser = userID
 // GIVEN the enqueued provisioner job is for a workspace being claimed by a user:
 jobInput.PrebuildClaimedByUser = userID
 }
 input, err := json.Marshal(scheduledJobInput)
 input, err := json.Marshal(jobInput)
 require.NoError(t, err)

 ctx := testutil.Context(t, testutil.WaitShort)
 })
 require.NoError(t, err)

 //Subscribe to workspace reinitialization:
 //GIVEN something is listening to process workspace reinitialization:

 // If the job needs to reinitialize agents for the workspace,
 // check that the instruction to do so was enqueued
 eventName := agentsdk.PrebuildClaimedChannel(workspace.ID)
 gotChan := make(chan []byte, 1)
 reinitChan := make(chan []byte, 1)
 cancel, err := ps.Subscribe(eventName, func(inner context.Context, userIDMessage []byte) {
 gotChan <- userIDMessage
 reinitChan <- userIDMessage
 })
 require.NoError(t, err)
 defer cancel()

 //Complete thejob, optionally triggering workspace agent reinitialization:
 //WHEN thejop is completed

 completedJob := proto.CompletedJob{
 JobId: job.ID.String(),
 require.NoError(t, err)

 select {
 case userIDMessage := <-gotChan:
 case userIDMessage := <-reinitChan:
 // THEN workspace agent reinitialization instruction was received:
 gotUserID, err := uuid.ParseBytes(userIDMessage)
 require.NoError(t, err)
 require.True(t, tc.shouldReinitializeAgent)
 require.Equal(t, userID, gotUserID)
 case <-ctx.Done():
 // THEN workspace agent reinitialization instruction was not received.
 require.False(t, tc.shouldReinitializeAgent)
 }
 })
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
 Type: codersdk.ServerSentEventTypePing,
 })

 // Expand with future use-cases for agent reinitialization.
 for {
 select {
 case <-ctx.Done():
 return
 case user := <-prebuildClaims:
 err = sseSendEvent(codersdk.ServerSentEvent{
 Type: codersdk.ServerSentEventTypeData,
 Data: agentsdk.ReinitializationResponse{
 Data: agentsdk.ReinitializationEvent{
 Message: fmt.Sprintf("prebuild claimed by user: %s", user),
 Reason:  agentsdk.ReinitializeReasonPrebuildClaimed,
 },
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
 if len(agents) > 1 {
 return xerrors.Errorf("multiple agents are not yet supported in prebuilt workspaces")
 }
 // agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
 // for _, agent := range agents {
 // agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
 // }
 }

 // We have to refetch the workspace for the joined in fields.
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
 ReinitializeReasonPrebuildClaimed ReinitializationReason = "prebuild_claimed"
 )

 typeReinitializationResponse struct {
 typeReinitializationEvent struct {
 Message string                 `json:"message"`
 Reason  ReinitializationReason `json:"reason"`
 }
 // - ping: ignored, keepalive
 // - prebuild claimed: a prebuilt workspace is claimed, so the agent must reinitialize.
 // NOTE: the caller is responsible for closing the events chan.
 func (c *Client) WaitForReinit(ctx context.Context, events chan<-ReinitializationResponse) error {
 func (c *Client) WaitForReinit(ctx context.Context, events chan<-ReinitializationEvent) error {
 // TODO: allow configuring httpclient
 c.SDK.HTTPClient.Timeout = time.Hour * 24

 if sse.Type != codersdk.ServerSentEventTypeData {
 continue
 }
 varreinitResp ReinitializationResponse
 varreinitEvent ReinitializationEvent
 b, ok := sse.Data.([]byte)
 if !ok {
 return xerrors.Errorf("expected data as []byte, got %T", sse.Data)
 }
 err = json.Unmarshal(b, &reinitResp)
 err = json.Unmarshal(b, &reinitEvent)
 if err != nil {
 return xerrors.Errorf("unmarshal reinit response: %w", err)
 }
 select {
 case <-ctx.Done():
 return ctx.Err()
 case events <-reinitResp:
 case events <-reinitEvent:
 }
 }
 }
diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
 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.
Original file line number	Diff line number	Diff line change
Expand Up		@@ -68,7 +68,7 @@ func (r RootCmd) workspaceAgent() serpent.Command {
		Handler: func(inv *serpent.Invocation) error {
		ctx, cancel := context.WithCancelCause(inv.Context())
		defer func() {
		cancel(xerrors.New("defer"))
		cancel(xerrors.New("agent exited"))
		}()

		var (
Expand DownExpand Up		@@ -335,7 +335,7 @@ func (r RootCmd) workspaceAgent() serpent.Command {
		// TODO: timeout ok?
		reinitCtx, reinitCancel := context.WithTimeout(context.Background(), time.Hour*24)
		defer reinitCancel()
		reinitEvents := make(chan agentsdk.ReinitializationResponse)
		reinitEvents := make(chan agentsdk.ReinitializationEvent)

		go func() {
		// Retry to wait for reinit, main context cancels the retrier.
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -359,18 +359,18 @@ func TestAgent_Prebuild(t *testing.T) {

		// Check that the agent is in a happy steady state
		waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
		waiter.WaitFor(coderdtest.AgentReady)
		waiter.WaitFor(coderdtest.AgentsReady)

		// Trigger reinitialization
		channel := agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
		err := pubsub.Publish(channel, []byte(user.UserID.String()))
		require.NoError(t, err)

		// Check that the agent reinitializes
		waiter.WaitFor(coderdtest.AgentNotReady)
		waiter.WaitFor(coderdtest.AgentsNotReady)

		// Check that reinitialization completed
		waiter.WaitFor(coderdtest.AgentReady)
		waiter.WaitFor(coderdtest.AgentsReady)
		}

		func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1105,14 +1105,24 @@ func (w WorkspaceAgentWaiter) MatchResources(m func([]codersdk.WorkspaceResource
		return w
		}

		// WaitForCriterium represents a boolean assertion to be made against each agent
		// that a given WorkspaceAgentWaited knows about. Each WaitForCriterium should apply
		// the check to a single agent, but it should be named for plural, because `func (w WorkspaceAgentWaiter) WaitFor`
		// applies the check to all agents that it is aware of. This ensures that the public API of the waiter
		// reads correctly. For example:
		//
		// waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
		// waiter.WaitFor(coderdtest.AgentsReady)
		type WaitForCriterium func(agent codersdk.WorkspaceAgent) bool
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved

		func AgentReady(agent codersdk.WorkspaceAgent) bool {
		// AgentsReady checks that the latest lifecycle state of an agent is "Ready".
		func AgentsReady(agent codersdk.WorkspaceAgent) bool {
		return agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady
		}

		func AgentNotReady(agent codersdk.WorkspaceAgent) bool {
		return !AgentReady(agent)
		// AgentsReady checks that the latest lifecycle state of an agent is anything except "Ready".
		func AgentsNotReady(agent codersdk.WorkspaceAgent) bool {
		return !AgentsReady(agent)
		}

		func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForCriterium) {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1752,7 +1752,7 @@ func (s server) CompleteJob(ctx context.Context, completed proto.CompletedJob)

		channel := agentsdk.PrebuildClaimedChannel(workspace.ID)
		if err := s.Pubsub.Publish(channel, []byte(input.PrebuildClaimedByUser.String())); err != nil {
		s.Logger.Error(ctx, "failed topublish message to instruct prebuilt workspace agent reinitialization", slog.Error(err))
		s.Logger.Error(ctx, "failed totrigger prebuilt workspace agent reinitialization", slog.Error(err))
		}
		}
		case *proto.CompletedJob_TemplateDryRun_:
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1771,20 +1771,21 @@ func TestCompleteJob(t *testing.T) {
		t.Run(tc.name, func(t *testing.T) {
dannykopping marked this conversation as resolved. Show resolvedHide resolved
		t.Parallel()

		//Setup starting state:
		//GIVEN an enqueued provisioner job and its dependencies:

		userID := uuid.New()

		srv, db, ps, pd := setup(t, false, &overrides{})

		buildID := uuid.New()
		scheduledJobInput := provisionerdserver.WorkspaceProvisionJob{
		jobInput := provisionerdserver.WorkspaceProvisionJob{
		WorkspaceBuildID: buildID,
		}
		if tc.shouldReinitializeAgent { // This is the key lever in the test
		scheduledJobInput.PrebuildClaimedByUser = userID
		// GIVEN the enqueued provisioner job is for a workspace being claimed by a user:
		jobInput.PrebuildClaimedByUser = userID
		}
		input, err := json.Marshal(scheduledJobInput)
		input, err := json.Marshal(jobInput)
		require.NoError(t, err)

		ctx := testutil.Context(t, testutil.WaitShort)
Expand DownExpand Up		@@ -1821,19 +1822,17 @@ func TestCompleteJob(t *testing.T) {
		})
		require.NoError(t, err)

		//Subscribe to workspace reinitialization:
		//GIVEN something is listening to process workspace reinitialization:

		// If the job needs to reinitialize agents for the workspace,
		// check that the instruction to do so was enqueued
		eventName := agentsdk.PrebuildClaimedChannel(workspace.ID)
		gotChan := make(chan []byte, 1)
		reinitChan := make(chan []byte, 1)
		cancel, err := ps.Subscribe(eventName, func(inner context.Context, userIDMessage []byte) {
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved
		gotChan <- userIDMessage
		reinitChan <- userIDMessage
		})
		require.NoError(t, err)
		defer cancel()

		//Complete thejob, optionally triggering workspace agent reinitialization:
		//WHEN thejop is completed
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved

		completedJob := proto.CompletedJob{
		JobId: job.ID.String(),
Expand All		@@ -1845,12 +1844,14 @@ func TestCompleteJob(t *testing.T) {
		require.NoError(t, err)

		select {
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved
		case userIDMessage := <-gotChan:
		case userIDMessage := <-reinitChan:
		// THEN workspace agent reinitialization instruction was received:
		gotUserID, err := uuid.ParseBytes(userIDMessage)
		require.NoError(t, err)
		require.True(t, tc.shouldReinitializeAgent)
		require.Equal(t, userID, gotUserID)
		case <-ctx.Done():
		// THEN workspace agent reinitialization instruction was not received.
		require.False(t, tc.shouldReinitializeAgent)
		}
		})
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1233,15 +1233,14 @@ func (api API) workspaceAgentReinit(rw http.ResponseWriter, r http.Request) {
		Type: codersdk.ServerSentEventTypePing,
		})

		// Expand with future use-cases for agent reinitialization.
		for {
		select {
		case <-ctx.Done():
		return
		case user := <-prebuildClaims:
		err = sseSendEvent(codersdk.ServerSentEvent{
		Type: codersdk.ServerSentEventTypeData,
		Data: agentsdk.ReinitializationResponse{
		Data: agentsdk.ReinitializationEvent{
		Message: fmt.Sprintf("prebuild claimed by user: %s", user),
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		},
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -692,10 +692,6 @@ func createWorkspace(
		if len(agents) > 1 {
		return xerrors.Errorf("multiple agents are not yet supported in prebuilt workspaces")
Copy link Contributor dannykoppingApr 25, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. It feels wrong to error here. While we don't support multiple agents, we shouldn't fail the call to`createWorkspace`, but rather prevent the ability to claim prebuilds. Right now I think this would be unrecoverable until the template is changed. Can you please add a test for this scenario to validate this behaviour? We could hard-fail this, though, on the template import side when an admin tries to use prebuilds with multiple agents (in a follow-up PR).
		}
		// agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
		// for _, agent := range agents {
		// agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
		// }
		}

		// We have to refetch the workspace for the joined in fields.
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -694,7 +694,7 @@ const (
		ReinitializeReasonPrebuildClaimed ReinitializationReason = "prebuild_claimed"
		)

		typeReinitializationResponse struct {
		typeReinitializationEvent struct {
		Message string `json:"message"`
		Reason ReinitializationReason `json:"reason"`
		}
Expand All		@@ -707,7 +707,7 @@ func PrebuildClaimedChannel(id uuid.UUID) string {
		// - ping: ignored, keepalive
		// - prebuild claimed: a prebuilt workspace is claimed, so the agent must reinitialize.
		// NOTE: the caller is responsible for closing the events chan.
		func (c *Client) WaitForReinit(ctx context.Context, events chan<-ReinitializationResponse) error {
		func (c *Client) WaitForReinit(ctx context.Context, events chan<-ReinitializationEvent) error {
		// TODO: allow configuring httpclient
		c.SDK.HTTPClient.Timeout = time.Hour * 24
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved

Expand DownExpand Up		@@ -737,19 +737,19 @@ func (c *Client) WaitForReinit(ctx context.Context, events chan<- Reinitializati
		if sse.Type != codersdk.ServerSentEventTypeData {
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved
		continue
		}
		varreinitResp ReinitializationResponse
		varreinitEvent ReinitializationEvent
		b, ok := sse.Data.([]byte)
		if !ok {
		return xerrors.Errorf("expected data as []byte, got %T", sse.Data)
		}
		err = json.Unmarshal(b, &reinitResp)
		err = json.Unmarshal(b, &reinitEvent)
		if err != nil {
		return xerrors.Errorf("unmarshal reinit response: %w", err)
		}
		select {
SasSwart marked this conversation as resolved. OutdatedShow resolvedHide resolved
		case <-ctx.Done():
		return ctx.Err()
		case events <-reinitResp:
		case events <-reinitEvent:
		}
		}
		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -277,6 +277,7 @@ func provisionEnv(
		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 {
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		// If there are multiple agents, provide all the tokens to terraform so that it can
		// choose the correct one for each agent ID.
Expand Down