fix: fetch the previous agent when we need its token for prebuilt wor…

…kspaces
coder · SasSwart · May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025
commit 146b15857c5f5637168897dddb94940d94fa6a66
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -3012,6 +3012,15 @@ func (q *querier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, crea
 	return q.db.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt)
 }
 
+func (q *querier) GetWorkspaceAgentsByBuildID(ctx context.Context, arg database.GetWorkspaceAgentsByBuildIDParams) ([]database.WorkspaceAgent, error) {
+	_, err := q.GetWorkspaceByID(ctx, arg.WorkspaceID)
+	if err != nil {
+		return nil, err
+	}
+
+	return q.db.GetWorkspaceAgentsByBuildID(ctx, arg)
+}
+
 // GetWorkspaceAgentsByResourceIDs
 // The workspace/job is already fetched.
 func (q *querier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgent, error) {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -7641,6 +7641,30 @@ func (q *FakeQuerier) GetWorkspaceAgentUsageStatsAndLabels(_ context.Context, cr
 	return stats, nil
 }
 
+func (q *FakeQuerier) GetWorkspaceAgentsByBuildID(ctx context.Context, arg database.GetWorkspaceAgentsByBuildIDParams) ([]database.WorkspaceAgent, error) {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return nil, err
+	}
+
+	build, err := q.GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx, database.GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams(arg))
+	if err != nil {
+		return nil, err
+	}
+
+	resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, build.JobID)
+	if err != nil {
+		return nil, err
+	}
+
+	var resourceIDs []uuid.UUID
+	for _, resource := range resources {
+		resourceIDs = append(resourceIDs, resource.ID)
+	}
+
+	return q.GetWorkspaceAgentsByResourceIDs(ctx, resourceIDs)
+}
+
 func (q *FakeQuerier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, resourceIDs []uuid.UUID) ([]database.WorkspaceAgent, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql
@@ -252,6 +252,19 @@ WHERE
 			wb.workspace_id = @workspace_id :: uuid
 	);
 
+-- name: GetWorkspaceAgentsByBuildID :many
+SELECT
+	workspace_agents.*
+FROM
+	workspace_agents
+JOIN
+	workspace_resources ON workspace_agents.resource_id = workspace_resources.id
+JOIN
+	workspace_builds ON workspace_resources.job_id = workspace_builds.job_id
+WHERE
+	workspace_builds.workspace_id = @workspace_id :: uuid AND
+	workspace_builds.build_number = @build_number :: int;
+
 -- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one
 SELECT
 	sqlc.embed(workspaces),

diff --git a/coderd/prebuilds/claim_test.go b/coderd/prebuilds/claim_test.go
@@ -18,20 +18,19 @@ import (
 
 func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
 	t.Parallel()
-	t.Run("publish claim", func(t *testing.T) {
+	t.Run("published claim is received by a listener for the same workspace", func(t *testing.T) {
 		t.Parallel()
 
+		ctx := testutil.Context(t, testutil.WaitShort)
+		logger := testutil.Logger(t)
 		ps := pubsub.NewInMemory()
-		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
-
 		workspaceID := uuid.New()
 		userID := uuid.New()
+		reinitEvents := make(chan agentsdk.ReinitializationEvent, 1)
+		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, logger)
 
-		userIDCh := make(chan uuid.UUID, 1)
-		channel := agentsdk.PrebuildClaimedChannel(workspaceID)
-		cancel, err := ps.Subscribe(channel, func(ctx context.Context, message []byte) {
-			userIDCh <- uuid.MustParse(string(message))
-		})
+		cancel, err := listener.ListenForWorkspaceClaims(ctx, workspaceID, reinitEvents)
 		require.NoError(t, err)
 		defer cancel()
 
@@ -43,8 +42,8 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
 		err = publisher.PublishWorkspaceClaim(claim)
 		require.NoError(t, err)
 
-		gotUserID := testutil.TryReceive(testutil.Context(t, testutil.WaitShort), t, userIDCh)
-		require.Equal(t, userID, gotUserID)
+		gotUserID := testutil.RequireReceive(testutil.Context(t, testutil.WaitShort), t, reinitEvents)
+		require.Equal(t, userID, gotUserID.UserID)
 	})
 
 	t.Run("fail to publish claim", func(t *testing.T) {

diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -627,7 +627,11 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 			// 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.
-			agents, err := s.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID)
+			agents, err := s.Database.GetWorkspaceAgentsByBuildID(ctx, database.GetWorkspaceAgentsByBuildIDParams{
+				WorkspaceID: workspace.ID,
+				BuildNumber: 1,
+			})
+
 			if err != nil {
 				s.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
 					slog.F("workspace_id", workspace.ID), slog.Error(err))

diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -23,6 +23,7 @@ import (
 	"storj.io/drpc"
 
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/coderd/prebuilds"
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/quartz"
@@ -1823,17 +1824,12 @@ func TestCompleteJob(t *testing.T) {
 				require.NoError(t, err)
 
 				// GIVEN something is listening to process workspace reinitialization:
-
-				eventName := agentsdk.PrebuildClaimedChannel(workspace.ID)
-				reinitChan := make(chan []byte, 1)
-				cancel, err := ps.Subscribe(eventName, func(inner context.Context, userIDMessage []byte) {
-					reinitChan <- userIDMessage
-				})
+				reinitChan := make(chan agentsdk.ReinitializationEvent, 1) // Buffered to simplify test structure
+				cancel, err := prebuilds.NewPubsubWorkspaceClaimListener(ps, testutil.Logger(t)).ListenForWorkspaceClaims(ctx, workspace.ID, reinitChan)
 				require.NoError(t, err)
 				defer cancel()
 
 				// WHEN the job is completed
-
 				completedJob := proto.CompletedJob{
 					JobId: job.ID.String(),
 					Type: &proto.CompletedJob_WorkspaceBuild_{
@@ -1844,13 +1840,11 @@ func TestCompleteJob(t *testing.T) {
 				require.NoError(t, err)
 
 				select {
-				case userIDMessage := <-reinitChan:
+				case reinitEvent := <-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():
+					require.Equal(t, userID, reinitEvent.UserID)
+				default:
 					// THEN workspace agent reinitialization instruction was not received.
 					require.False(t, tc.shouldReinitializeAgent)
 				}
@@ -2953,3 +2947,13 @@ func (s *fakeStream) cancel() {
 	s.canceled = true
 	s.c.Broadcast()
 }
+
+type pubsubReinitSpy struct {
+	pubsub.Pubsub
+	subscriptions chan string
+}
+
+func (p pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
+	p.subscriptions <- event
+	return p.Pubsub.Subscribe(event, listener)
+}
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -1227,6 +1227,8 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
 		log.Info(ctx, "agent reinitialization subscription closed", slog.F("workspace_agent_id", workspaceAgent.ID))
 	case errors.Is(err, agentsdk.ErrTransmissionTargetClosed):
 		log.Info(ctx, "agent connection closed", slog.F("workspace_agent_id", workspaceAgent.ID))
+	case errors.Is(err, context.Canceled):
+		log.Info(ctx, "agent reinitialization", slog.Error(err))
 	case err != nil:
 		log.Error(ctx, "failed to stream agent reinit events", slog.Error(err))
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{