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
 mustExit = true
 case event := <-reinitEvents:
 logger.Info(ctx, "agent received instruction to reinitialize",
 slog.F("user_id", event.UserID), slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
 slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
 }

 lastErr = agnt.Close()
diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go

 func (p PubsubWorkspaceClaimPublisher) PublishWorkspaceClaim(claim agentsdk.ReinitializationEvent) error {
 channel := agentsdk.PrebuildClaimedChannel(claim.WorkspaceID)
 if err := p.ps.Publish(channel, []byte(claim.UserID.String())); err != nil {
 if err := p.ps.Publish(channel, []byte(claim.Reason)); err != nil {
 return xerrors.Errorf("failed to trigger prebuilt workspace agent reinitialization: %w", err)
 }
 return nil
 default:
 }

 cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
 claimantID, err := uuid.ParseBytes(id)
 if err != nil {
 p.logger.Error(ctx, "invalid prebuild claimed channel payload", slog.F("input", string(id)))
 return
 }
 cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, reason []byte) {
 claim := agentsdk.ReinitializationEvent{
 UserID:      claimantID,
 WorkspaceID: workspaceID,
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 Reason:      agentsdk.ReinitializationReason(reason),
 }

 select {
diff --git a/coderd/prebuilds/claim_test.go b/coderd/prebuilds/claim_test.go
 logger := testutil.Logger(t)
 ps := pubsub.NewInMemory()
 workspaceID := uuid.New()
 userID := uuid.New()
 reinitEvents := make(chan agentsdk.ReinitializationEvent, 1)
 publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
 listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, logger)
 defer cancel()

 claim := agentsdk.ReinitializationEvent{
 UserID:      userID,
 WorkspaceID: workspaceID,
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 }
 err = publisher.PublishWorkspaceClaim(claim)
 require.NoError(t, err)

 gotUserID := testutil.RequireReceive(testutil.Context(t, testutil.WaitShort), t, reinitEvents)
 require.Equal(t, userID, gotUserID.UserID)
 gotEvent := testutil.RequireReceive(ctx, t, reinitEvents)
 require.Equal(t, workspaceID, gotEvent.WorkspaceID)
 require.Equal(t, claim.Reason, gotEvent.Reason)
 })

 t.Run("fail to publish claim", func(t *testing.T) {

 publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
 claim := agentsdk.ReinitializationEvent{
 UserID:      uuid.New(),
 WorkspaceID: uuid.New(),
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 }
 claims := make(chan agentsdk.ReinitializationEvent, 1) // Buffer to avoid messing with goroutines in the rest of the test

 workspaceID := uuid.New()
 userID := uuid.New()
 cancelFunc, err := listener.ListenForWorkspaceClaims(context.Background(), workspaceID, claims)
 require.NoError(t, err)
 defer cancelFunc()

 // Publish a claim
 channel := agentsdk.PrebuildClaimedChannel(workspaceID)
 err = ps.Publish(channel, []byte(userID.String()))
 reason := agentsdk.ReinitializeReasonPrebuildClaimed
 err = ps.Publish(channel, []byte(reason))
 require.NoError(t, err)

 // Verify we receive the claim
 select {
 case claim, ok := <-claims:
 require.True(t, ok, "received on a closed channel")
 require.Equal(t, userID, claim.UserID)
 require.Equal(t, workspaceID, claim.WorkspaceID)
 require.Equal(t, agentsdk.ReinitializeReasonPrebuildClaimed, claim.Reason)
 case <-time.After(testutil.WaitSuperLong): // TODO: revert to waitshort
 t.Fatal("timeout waiting for claim")
 }
 ctx := testutil.Context(t, testutil.WaitShort)
 claim := testutil.RequireReceive(ctx, t, claims)
 require.Equal(t, workspaceID, claim.WorkspaceID)
 require.Equal(t, reason, claim.Reason)
 })

 t.Run("ignores claim events for other workspaces", func(t *testing.T) {

 // Publish a claim for a different workspace
 channel := agentsdk.PrebuildClaimedChannel(otherWorkspaceID)
 err = ps.Publish(channel, []byte(uuid.New().String()))
 err = ps.Publish(channel, []byte(agentsdk.ReinitializeReasonPrebuildClaimed))
 require.NoError(t, err)

 // Verify we don't receive the claim
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
 _, err = tc.acquire(ctx, srv)
 require.ErrorContains(t, err, "sql: no rows in result set")
 })
 for _, prebuiltWorkspace := range []bool{false, true} {
 prebuiltWorkspace := prebuiltWorkspace
 for _, prebuiltWorkspaceBuildStage := range []sdkproto.PrebuiltWorkspaceBuildStage{
 sdkproto.PrebuiltWorkspaceBuildStage_NONE,
 sdkproto.PrebuiltWorkspaceBuildStage_CREATE,
 sdkproto.PrebuiltWorkspaceBuildStage_CLAIM,
 } {
 prebuiltWorkspaceBuildStage := prebuiltWorkspaceBuildStage
 t.Run(tc.name+"_WorkspaceBuildJob", func(t *testing.T) {
 t.Parallel()
 // Set the max session token lifetime so we can assert we
 OwnerID:        user.ID,
 OrganizationID: pd.OrganizationID,
 })
 build :=dbgen.WorkspaceBuild(t, db,database.WorkspaceBuild{
 build := database.WorkspaceBuild{
 WorkspaceID:       workspace.ID,
 BuildNumber:       1,
 JobID:             uuid.New(),
 TemplateVersionID: version.ID,
 Transition:        database.WorkspaceTransitionStart,
 Reason:            database.BuildReasonInitiator,
 })
 var buildState sdkproto.PrebuiltWorkspaceBuildStage
 if prebuiltWorkspace {
 buildState = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
 }
 _ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
 ID:             build.ID,
 input := provisionerdserver.WorkspaceProvisionJob{
 WorkspaceBuildID: build.ID,
 }
 dbJob := database.ProvisionerJob{
 ID:             build.JobID,
 OrganizationID: pd.OrganizationID,
 InitiatorID:    user.ID,
 Provisioner:    database.ProvisionerTypeEcho,
 StorageMethod:  database.ProvisionerStorageMethodFile,
 FileID:         file.ID,
 Type:           database.ProvisionerJobTypeWorkspaceBuild,
 Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
 WorkspaceBuildID:            build.ID,
 PrebuiltWorkspaceBuildStage: buildState,
 })),
 })
 Input:          must(json.Marshal(input)),
 }
 if prebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
 dbgen.WorkspaceBuild(t, db, build)
 prebuildJob := dbgen.ProvisionerJob(t, db, ps, dbJob)
 resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 JobID: prebuildJob.ID,
 })
 dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
 ResourceID: resource.ID,
 AuthToken:  uuid.New(),
 })
 build.BuildNumber = 2
 input.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CLAIM
 dbJob.Input = must(json.Marshal(input))
 }
 build = dbgen.WorkspaceBuild(t, db, build)
 dbJob = dbgen.ProvisionerJob(t, db, ps, dbJob)

 startPublished := make(chan struct{})
 var closed bool

 <-startPublished

 got, err := json.Marshal(job.Type)
 got, err := json.Marshal(dbJob.Type)
 require.NoError(t, err)

 // Validate that a session token is generated during the job.
 WorkspaceBuildId:              build.ID.String(),
 WorkspaceOwnerLoginType:       string(user.LoginType),
 WorkspaceOwnerRbacRoles:       []*sdkproto.Role{{Name: rbac.RoleOrgMember(), OrgId: pd.OrganizationID.String()}, {Name: "member", OrgId: ""}, {Name: rbac.RoleOrgAuditor(), OrgId: pd.OrganizationID.String()}},
 }
 if prebuiltWorkspace {
 wantedMetadata.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
 PrebuiltWorkspaceBuildStage:   prebuiltWorkspaceBuildStage,
 }

 slices.SortFunc(wantedMetadata.WorkspaceOwnerRbacRoles, func(a, b *sdkproto.Role) int {
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
 func TestReinit(t *testing.T) {
 t.Parallel()

 t.Run("claimed workspaces are reinitialized", func(t *testing.T) {
 t.Parallel()

 db, ps := dbtestutil.NewDB(t)
 pubsubSpy := pubsubReinitSpy{
 Pubsub:        ps,
 subscriptions: make(chan string),
 }
 client := coderdtest.New(t, &coderdtest.Options{
 Database: db,
 Pubsub:   pubsubSpy,
 })
 user := coderdtest.CreateFirstUser(t, client)

 r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
 OrganizationID: user.OrganizationID,
 OwnerID:        user.UserID,
 }).WithAgent().Do()
 db, ps := dbtestutil.NewDB(t)
 pubsubSpy := pubsubReinitSpy{
 Pubsub:     ps,
 subscribed: make(chan string),
 }
 client := coderdtest.New(t, &coderdtest.Options{
 Database: db,
 Pubsub:   &pubsubSpy,
 })
 user := coderdtest.CreateFirstUser(t, client)

 agentCtx := testutil.Context(t, testutil.WaitShort)
 agentClient := agentsdk.New(client.URL)
 agentClient.SetSessionToken(r.AgentToken)
 r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
 OrganizationID: user.OrganizationID,
 OwnerID:        user.UserID,
 }).WithAgent().Do()
 pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)

 agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
 go func() {
 reinitEvent, err := agentClient.WaitForReinit(agentCtx)
 assert.NoError(t, err)
 agentReinitializedCh <- reinitEvent
 }()
 agentCtx := testutil.Context(t, testutil.WaitShort)
 agentClient := agentsdk.New(client.URL)
 agentClient.SetSessionToken(r.AgentToken)

 // We need to subscribe before we publish, lest we miss the event
 for subscription := range pubsubSpy.subscriptions {
 if subscription == agentsdk.PrebuildClaimedChannel(r.Workspace.ID) {
 break
 }
 }
 agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
 go func() {
 reinitEvent, err := agentClient.WaitForReinit(agentCtx)
 assert.NoError(t, err)
 agentReinitializedCh <- reinitEvent
 }()

 err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
 WorkspaceID: r.Workspace.ID,
 UserID:      user.UserID,
 })
 require.NoError(t, err)
 // We need to subscribe before we publish, lest we miss the event
 ctx := testutil.Context(t, testutil.WaitShort)
 testutil.TryReceive(ctx, t, pubsubSpy.subscribed) // Wait for the appropriate subscription

 ctx := testutil.Context(t, testutil.WaitShort)
 reinitEvent :=testutil.TryReceive(ctx, t, agentReinitializedCh)
 require.NotNil(t, reinitEvent)
 require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
 // Now that we're subscribed, publish the event
 err :=prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
 WorkspaceID: r.Workspace.ID,
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 })
 require.NoError(t, err)

 ctx = testutil.Context(t, testutil.WaitShort)
 reinitEvent := testutil.TryReceive(ctx, t, agentReinitializedCh)
 require.NotNil(t, reinitEvent)
 require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
 }

 type pubsubReinitSpy struct {
 pubsub.Pubsub
 subscriptions chan string
 subscribed    chan string
 expectedEvent string
 }

 func (p pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
 p.subscriptions <- event
 func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
 if p.expectedEvent != "" && event == p.expectedEvent {
 close(p.subscribed)
 }
 return p.Pubsub.Subscribe(event, listener)
 }
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go

 type ReinitializationEvent struct {
 WorkspaceID uuid.UUID
 UserID      uuid.UUID
 Reason      ReinitializationReason `json:"reason"`
 }

 return xerrors.Errorf("failed to create sse transmitter: %w", err)
 }

 // Prevent handler from returning until the sender is closed.
 defer func() {
 // Block returning until the ServerSentEventSender is closed
 // to avoid a race condition where we might write or flush to rw after the handler returns.
 <-sseSenderClosed
 }()

diff --git a/codersdk/agentsdk/agentsdk_test.go b/codersdk/agentsdk/agentsdk_test.go
 t.Parallel()

 eventToSend := agentsdk.ReinitializationEvent{
 UserID:      uuid.New(),
 WorkspaceID: uuid.New(),
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 }
 t.Parallel()

 eventToSend := agentsdk.ReinitializationEvent{
 UserID:      uuid.New(),
 WorkspaceID: uuid.New(),
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 }
 t.Parallel()

 eventToSend := agentsdk.ReinitializationEvent{
 UserID:      uuid.New(),
 WorkspaceID: uuid.New(),
 Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -377,7 +377,7 @@ func (r RootCmd) workspaceAgent() serpent.Command {
		mustExit = true
		case event := <-reinitEvents:
		logger.Info(ctx, "agent received instruction to reinitialize",
		slog.F("user_id", event.UserID), slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
		slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
		}

		lastErr = agnt.Close()
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -22,7 +22,7 @@ type PubsubWorkspaceClaimPublisher struct {

		func (p PubsubWorkspaceClaimPublisher) PublishWorkspaceClaim(claim agentsdk.ReinitializationEvent) error {
		channel := agentsdk.PrebuildClaimedChannel(claim.WorkspaceID)
		if err := p.ps.Publish(channel, []byte(claim.UserID.String())); err != nil {
		if err := p.ps.Publish(channel, []byte(claim.Reason)); err != nil {
		return xerrors.Errorf("failed to trigger prebuilt workspace agent reinitialization: %w", err)
		}
		return nil
Expand All		@@ -48,16 +48,10 @@ func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Conte
		default:
		}

		cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
		claimantID, err := uuid.ParseBytes(id)
		if err != nil {
		p.logger.Error(ctx, "invalid prebuild claimed channel payload", slog.F("input", string(id)))
		return
		}
		cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, reason []byte) {
		claim := agentsdk.ReinitializationEvent{
		UserID: claimantID,
		WorkspaceID: workspaceID,
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		Reason: agentsdk.ReinitializationReason(reason),
		}

		select {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -25,7 +25,6 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
		logger := testutil.Logger(t)
		ps := pubsub.NewInMemory()
		workspaceID := uuid.New()
		userID := uuid.New()
		reinitEvents := make(chan agentsdk.ReinitializationEvent, 1)
		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, logger)
Expand All		@@ -35,15 +34,15 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
		defer cancel()

		claim := agentsdk.ReinitializationEvent{
		UserID: userID,
		WorkspaceID: workspaceID,
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		}
		err = publisher.PublishWorkspaceClaim(claim)
		require.NoError(t, err)

		gotUserID := testutil.RequireReceive(testutil.Context(t, testutil.WaitShort), t, reinitEvents)
		require.Equal(t, userID, gotUserID.UserID)
		gotEvent := testutil.RequireReceive(ctx, t, reinitEvents)
		require.Equal(t, workspaceID, gotEvent.WorkspaceID)
		require.Equal(t, claim.Reason, gotEvent.Reason)
		})

		t.Run("fail to publish claim", func(t *testing.T) {
Expand All		@@ -53,7 +52,6 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {

		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
		claim := agentsdk.ReinitializationEvent{
		UserID: uuid.New(),
		WorkspaceID: uuid.New(),
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		}
Expand All		@@ -74,26 +72,21 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
		claims := make(chan agentsdk.ReinitializationEvent, 1) // Buffer to avoid messing with goroutines in the rest of the test

		workspaceID := uuid.New()
		userID := uuid.New()
		cancelFunc, err := listener.ListenForWorkspaceClaims(context.Background(), workspaceID, claims)
		require.NoError(t, err)
		defer cancelFunc()

		// Publish a claim
		channel := agentsdk.PrebuildClaimedChannel(workspaceID)
		err = ps.Publish(channel, []byte(userID.String()))
		reason := agentsdk.ReinitializeReasonPrebuildClaimed
		err = ps.Publish(channel, []byte(reason))
		require.NoError(t, err)

		// Verify we receive the claim
		select {
		case claim, ok := <-claims:
		require.True(t, ok, "received on a closed channel")
		require.Equal(t, userID, claim.UserID)
		require.Equal(t, workspaceID, claim.WorkspaceID)
		require.Equal(t, agentsdk.ReinitializeReasonPrebuildClaimed, claim.Reason)
		case <-time.After(testutil.WaitSuperLong): // TODO: revert to waitshort
		t.Fatal("timeout waiting for claim")
		}
		ctx := testutil.Context(t, testutil.WaitShort)
		claim := testutil.RequireReceive(ctx, t, claims)
		require.Equal(t, workspaceID, claim.WorkspaceID)
		require.Equal(t, reason, claim.Reason)
		})

		t.Run("ignores claim events for other workspaces", func(t *testing.T) {
SasSwart marked this conversation as resolved. Show resolvedHide resolved
Expand All		@@ -111,7 +104,7 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {

		// Publish a claim for a different workspace
		channel := agentsdk.PrebuildClaimedChannel(otherWorkspaceID)
		err = ps.Publish(channel, []byte(uuid.New().String()))
		err = ps.Publish(channel, []byte(agentsdk.ReinitializeReasonPrebuildClaimed))
		require.NoError(t, err)

		// Verify we don't receive the claim
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -169,8 +169,12 @@ func TestAcquireJob(t *testing.T) {
		_, err = tc.acquire(ctx, srv)
		require.ErrorContains(t, err, "sql: no rows in result set")
		})
		for _, prebuiltWorkspace := range []bool{false, true} {
		prebuiltWorkspace := prebuiltWorkspace
		for _, prebuiltWorkspaceBuildStage := range []sdkproto.PrebuiltWorkspaceBuildStage{
		sdkproto.PrebuiltWorkspaceBuildStage_NONE,
		sdkproto.PrebuiltWorkspaceBuildStage_CREATE,
		sdkproto.PrebuiltWorkspaceBuildStage_CLAIM,
		} {
		prebuiltWorkspaceBuildStage := prebuiltWorkspaceBuildStage
		t.Run(tc.name+"_WorkspaceBuildJob", func(t *testing.T) {
		t.Parallel()
		// Set the max session token lifetime so we can assert we
Expand DownExpand Up		@@ -294,31 +298,43 @@ func TestAcquireJob(t *testing.T) {
		OwnerID: user.ID,
		OrganizationID: pd.OrganizationID,
		})
		build :=dbgen.WorkspaceBuild(t, db,database.WorkspaceBuild{
		build := database.WorkspaceBuild{
		WorkspaceID: workspace.ID,
		BuildNumber: 1,
		JobID: uuid.New(),
		TemplateVersionID: version.ID,
		Transition: database.WorkspaceTransitionStart,
		Reason: database.BuildReasonInitiator,
		})
		var buildState sdkproto.PrebuiltWorkspaceBuildStage
		if prebuiltWorkspace {
		buildState = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
		}
		_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
		ID: build.ID,
		input := provisionerdserver.WorkspaceProvisionJob{
		WorkspaceBuildID: build.ID,
		}
		dbJob := database.ProvisionerJob{
		ID: build.JobID,
		OrganizationID: pd.OrganizationID,
		InitiatorID: user.ID,
		Provisioner: database.ProvisionerTypeEcho,
		StorageMethod: database.ProvisionerStorageMethodFile,
		FileID: file.ID,
		Type: database.ProvisionerJobTypeWorkspaceBuild,
		Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
		WorkspaceBuildID: build.ID,
		PrebuiltWorkspaceBuildStage: buildState,
		})),
		})
		Input: must(json.Marshal(input)),
		}
		if prebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
		dbgen.WorkspaceBuild(t, db, build)
		prebuildJob := dbgen.ProvisionerJob(t, db, ps, dbJob)
		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
		JobID: prebuildJob.ID,
		})
		dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
		ResourceID: resource.ID,
		AuthToken: uuid.New(),
		})
		build.BuildNumber = 2
		input.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CLAIM
		dbJob.Input = must(json.Marshal(input))
		}
		build = dbgen.WorkspaceBuild(t, db, build)
		dbJob = dbgen.ProvisionerJob(t, db, ps, dbJob)

		startPublished := make(chan struct{})
		var closed bool
Expand DownExpand Up		@@ -352,7 +368,7 @@ func TestAcquireJob(t *testing.T) {

		<-startPublished

		got, err := json.Marshal(job.Type)
		got, err := json.Marshal(dbJob.Type)
		require.NoError(t, err)

		// Validate that a session token is generated during the job.
Expand DownExpand Up		@@ -385,9 +401,7 @@ func TestAcquireJob(t *testing.T) {
		WorkspaceBuildId: build.ID.String(),
		WorkspaceOwnerLoginType: string(user.LoginType),
		WorkspaceOwnerRbacRoles: []*sdkproto.Role{{Name: rbac.RoleOrgMember(), OrgId: pd.OrganizationID.String()}, {Name: "member", OrgId: ""}, {Name: rbac.RoleOrgAuditor(), OrgId: pd.OrganizationID.String()}},
		}
		if prebuiltWorkspace {
		wantedMetadata.PrebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CREATE
		PrebuiltWorkspaceBuildStage: prebuiltWorkspaceBuildStage,
		}

		slices.SortFunc(wantedMetadata.WorkspaceOwnerRbacRoles, func(a, b *sdkproto.Role) int {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2647,62 +2647,60 @@ func TestAgentConnectionInfo(t *testing.T) {
		func TestReinit(t *testing.T) {
		t.Parallel()

		t.Run("claimed workspaces are reinitialized", func(t *testing.T) {
		t.Parallel()

		db, ps := dbtestutil.NewDB(t)
		pubsubSpy := pubsubReinitSpy{
		Pubsub: ps,
		subscriptions: make(chan string),
		}
		client := coderdtest.New(t, &coderdtest.Options{
		Database: db,
		Pubsub: pubsubSpy,
		})
		user := coderdtest.CreateFirstUser(t, client)

		r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
		OrganizationID: user.OrganizationID,
		OwnerID: user.UserID,
		}).WithAgent().Do()
		db, ps := dbtestutil.NewDB(t)
		pubsubSpy := pubsubReinitSpy{
		Pubsub: ps,
		subscribed: make(chan string),
		}
		client := coderdtest.New(t, &coderdtest.Options{
		Database: db,
		Pubsub: &pubsubSpy,
		})
		user := coderdtest.CreateFirstUser(t, client)

		agentCtx := testutil.Context(t, testutil.WaitShort)
		agentClient := agentsdk.New(client.URL)
		agentClient.SetSessionToken(r.AgentToken)
		r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
		OrganizationID: user.OrganizationID,
		OwnerID: user.UserID,
		}).WithAgent().Do()
		pubsubSpy.expectedEvent = agentsdk.PrebuildClaimedChannel(r.Workspace.ID)
SasSwart marked this conversation as resolved. Show resolvedHide resolved

		agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
		go func() {
		reinitEvent, err := agentClient.WaitForReinit(agentCtx)
		assert.NoError(t, err)
		agentReinitializedCh <- reinitEvent
		}()
		agentCtx := testutil.Context(t, testutil.WaitShort)
		agentClient := agentsdk.New(client.URL)
		agentClient.SetSessionToken(r.AgentToken)

		// We need to subscribe before we publish, lest we miss the event
		for subscription := range pubsubSpy.subscriptions {
		if subscription == agentsdk.PrebuildClaimedChannel(r.Workspace.ID) {
		break
		}
		}
		agentReinitializedCh := make(chan *agentsdk.ReinitializationEvent)
		go func() {
		reinitEvent, err := agentClient.WaitForReinit(agentCtx)
		assert.NoError(t, err)
		agentReinitializedCh <- reinitEvent
		}()

		err := prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
		WorkspaceID: r.Workspace.ID,
		UserID: user.UserID,
		})
		require.NoError(t, err)
		// We need to subscribe before we publish, lest we miss the event
		ctx := testutil.Context(t, testutil.WaitShort)
		testutil.TryReceive(ctx, t, pubsubSpy.subscribed) // Wait for the appropriate subscription

		ctx := testutil.Context(t, testutil.WaitShort)
		reinitEvent :=testutil.TryReceive(ctx, t, agentReinitializedCh)
		require.NotNil(t, reinitEvent)
		require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
		// Now that we're subscribed, publish the event
		err :=prebuilds.NewPubsubWorkspaceClaimPublisher(ps).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
		WorkspaceID: r.Workspace.ID,
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		})
		require.NoError(t, err)

		ctx = testutil.Context(t, testutil.WaitShort)
		reinitEvent := testutil.TryReceive(ctx, t, agentReinitializedCh)
		require.NotNil(t, reinitEvent)
		require.Equal(t, r.Workspace.ID, reinitEvent.WorkspaceID)
		}

		type pubsubReinitSpy struct {
		pubsub.Pubsub
		subscriptions chan string
		subscribed chan string
		expectedEvent string
		}

		func (p pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
		p.subscriptions <- event
		func (p *pubsubReinitSpy) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
		if p.expectedEvent != "" && event == p.expectedEvent {
		close(p.subscribed)
		}
		return p.Pubsub.Subscribe(event, listener)
		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -698,7 +698,6 @@ const (

		type ReinitializationEvent struct {
		WorkspaceID uuid.UUID
		UserID uuid.UUID
		Reason ReinitializationReason `json:"reason"`
		}

Expand DownExpand Up		@@ -806,8 +805,9 @@ func (s *SSEAgentReinitTransmitter) Transmit(ctx context.Context, reinitEvents <
		return xerrors.Errorf("failed to create sse transmitter: %w", err)
		}

		// Prevent handler from returning until the sender is closed.
		defer func() {
		// Block returning until the ServerSentEventSender is closed
		// to avoid a race condition where we might write or flush to rw after the handler returns.
		<-sseSenderClosed
		}()
SasSwart marked this conversation as resolved. Show resolvedHide resolved

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -22,7 +22,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
		t.Parallel()

		eventToSend := agentsdk.ReinitializationEvent{
		UserID: uuid.New(),
		WorkspaceID: uuid.New(),
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		}
Expand DownExpand Up		@@ -56,7 +55,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
		t.Parallel()

		eventToSend := agentsdk.ReinitializationEvent{
		UserID: uuid.New(),
		WorkspaceID: uuid.New(),
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		}
Expand DownExpand Up		@@ -92,7 +90,6 @@ func TestStreamAgentReinitEvents(t *testing.T) {
		t.Parallel()

		eventToSend := agentsdk.ReinitializationEvent{
		UserID: uuid.New(),
		WorkspaceID: uuid.New(),
		Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
		}
Expand Down