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

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
co-authored by Danny Kopping
parent fcbdd1a28e
commit 425ee6fa55
38 changed files with 2184 additions and 449 deletions
+55
View File
@@ -35,6 +35,7 @@ import (
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/coderd/httpmw/loggermw"
"github.com/coder/coder/v2/coderd/jwtutils"
"github.com/coder/coder/v2/coderd/prebuilds"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/coderd/telemetry"
@@ -1183,6 +1184,60 @@ func (api *API) workspaceAgentPostLogSource(rw http.ResponseWriter, r *http.Requ
httpapi.Write(ctx, rw, http.StatusCreated, apiSource)
}
// @Summary Get workspace agent reinitialization
// @ID get-workspace-agent-reinitialization
// @Security CoderSessionToken
// @Produce json
// @Tags Agents
// @Success 200 {object} agentsdk.ReinitializationEvent
// @Router /workspaceagents/me/reinit [get]
func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
// Allow us to interrupt watch via cancel.
ctx, cancel := context.WithCancel(r.Context())
defer cancel()
r = r.WithContext(ctx) // Rewire context for SSE cancellation.
workspaceAgent := httpmw.WorkspaceAgent(r)
log := api.Logger.Named("workspace_agent_reinit_watcher").With(
slog.F("workspace_agent_id", workspaceAgent.ID),
)
workspace, err := api.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
if err != nil {
log.Error(ctx, "failed to retrieve workspace from agent token", slog.Error(err))
httpapi.InternalServerError(rw, xerrors.New("failed to determine workspace from agent token"))
}
log.Info(ctx, "agent waiting for reinit instruction")
reinitEvents := make(chan agentsdk.ReinitializationEvent)
cancel, err = prebuilds.NewPubsubWorkspaceClaimListener(api.Pubsub, log).ListenForWorkspaceClaims(ctx, workspace.ID, reinitEvents)
if err != nil {
log.Error(ctx, "subscribe to prebuild claimed channel", slog.Error(err))
httpapi.InternalServerError(rw, xerrors.New("failed to subscribe to prebuild claimed channel"))
return
}
defer cancel()
transmitter := agentsdk.NewSSEAgentReinitTransmitter(log, rw, r)
err = transmitter.Transmit(ctx, reinitEvents)
switch {
case errors.Is(err, agentsdk.ErrTransmissionSourceClosed):
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{
Message: "Internal error streaming agent reinitialization events.",
Detail: err.Error(),
})
}
}
// convertProvisionedApps converts applications that are in the middle of provisioning process.
// It means that they may not have an agent or workspace assigned (dry-run job).
func convertProvisionedApps(dbApps []database.WorkspaceApp) []codersdk.WorkspaceApp {