synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-03 16:13:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: reinitialize agents when a prebuilt workspace is claimed (#17475)

Browse Source 
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.

---------

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: Danny Kopping <dannykopping@gmail.com>
This commit is contained in:
Sas Swart
2025-05-14 14:15:36 +02:00
committed by GitHub
parent fcbdd1a28e
commit 425ee6fa55
38 changed files with 2184 additions and 449 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
coderd/provisionerdserver/provisionerdserver.go
View File

@ -40,12 +40,14 @@ import (
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/externalauth"
"github.com/coder/coder/v2/coderd/notifications"
"github.com/coder/coder/v2/coderd/prebuilds"
"github.com/coder/coder/v2/coderd/promoauth"
"github.com/coder/coder/v2/coderd/schedule"
"github.com/coder/coder/v2/coderd/telemetry"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/coderd/wspubsub"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/provisioner"
"github.com/coder/coder/v2/provisionerd/proto"
"github.com/coder/coder/v2/provisionersdk"
@ -647,6 +649,30 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
}
}
runningAgentAuthTokens := []*sdkproto.RunningAgentAuthToken{}
if input.PrebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
// runningAgentAuthTokens are *only* used for prebuilds. We fetch them when we want to rebuild a prebuilt workspace
// but not generate new agent tokens. The provisionerdserver will push them down to
// the provisioner (and ultimately to the `coder_agent` resource in the Terraform provider) where they will be
// 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.GetWorkspaceAgentsByWorkspaceAndBuildNumber(ctx, database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
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))
}
for _, agent := range agents {
runningAgentAuthTokens = append(runningAgentAuthTokens, &sdkproto.RunningAgentAuthToken{
AgentId: agent.ID.String(),
Token: agent.AuthToken.String(),
})
}
}
protoJob.Type = &proto.AcquiredJob_WorkspaceBuild_{
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
WorkspaceBuildId: workspaceBuild.ID.String(),
@ -676,6 +702,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
WorkspaceBuildId: workspaceBuild.ID.String(),
WorkspaceOwnerLoginType: string(owner.LoginType),
WorkspaceOwnerRbacRoles: ownerRbacRoles,
RunningAgentAuthTokens: runningAgentAuthTokens,
PrebuiltWorkspaceBuildStage: input.PrebuiltWorkspaceBuildStage,
},
LogLevel: input.LogLevel,
@ -1812,6 +1839,19 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
if err != nil {
return nil, xerrors.Errorf("update workspace: %w", err)
}
if input.PrebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
s.Logger.Info(ctx, "workspace prebuild successfully claimed by user",
slog.F("workspace_id", workspace.ID))
err = prebuilds.NewPubsubWorkspaceClaimPublisher(s.Pubsub).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
WorkspaceID: workspace.ID,
Reason: agentsdk.ReinitializeReasonPrebuildClaimed,
})
if err != nil {
s.Logger.Error(ctx, "failed to publish workspace claim event", slog.Error(err))
}
}
case *proto.CompletedJob_TemplateDryRun_:
for _, resource := range jobType.TemplateDryRun.Resources {
s.Logger.Info(ctx, "inserting template dry-run job resource",
@ -1955,6 +1995,7 @@ func InsertWorkspacePresetAndParameters(ctx context.Context, db database.Store,
}
}
dbPreset, err := tx.InsertPreset(ctx, database.InsertPresetParams{
ID: uuid.New(),
TemplateVersionID: templateVersionID,
Name: protoPreset.Name,
CreatedAt: t,