coder · SasSwart · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025 · Apr 24, 2025
diff --git a/agent/agent.go b/agent/agent.go
@@ -36,6 +36,9 @@ import (
 	"tailscale.com/util/clientmetric"
 
 	"cdr.dev/slog"
+
+	"github.com/coder/retry"
+
 	"github.com/coder/clistat"
 	"github.com/coder/coder/v2/agent/agentcontainers"
 	"github.com/coder/coder/v2/agent/agentexec"
@@ -53,7 +56,6 @@ import (
 	"github.com/coder/coder/v2/tailnet"
 	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/quartz"
-	"github.com/coder/retry"
 )
 
 const (
@@ -363,9 +365,11 @@ func (a *agent) runLoop() {
 		if ctx.Err() != nil {
 			// Context canceled errors may come from websocket pings, so we
 			// don't want to use `errors.Is(err, context.Canceled)` here.
+			a.logger.Warn(ctx, "runLoop exited with error", slog.Error(ctx.Err()))
 			return
 		}
 		if a.isClosed() {
+			a.logger.Warn(ctx, "runLoop exited because agent is closed")
 			return
 		}
 		if errors.Is(err, io.EOF) {
@@ -1046,7 +1050,11 @@ func (a *agent) run() (retErr error) {
 		return a.statsReporter.reportLoop(ctx, aAPI)
 	})
 
-	return connMan.wait()
+	err = connMan.wait()
+	if err != nil {
+		a.logger.Warn(context.Background(), "connection manager errored", slog.Error(err))
+	}
+	return err
 }
 
 // handleManifest returns a function that fetches and processes the manifest

diff --git a/cli/agent.go b/cli/agent.go
@@ -25,6 +25,8 @@ import (
 	"cdr.dev/slog/sloggers/sloghuman"
 	"cdr.dev/slog/sloggers/slogjson"
 	"cdr.dev/slog/sloggers/slogstackdriver"
+	"github.com/coder/serpent"
+
 	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agentexec"
 	"github.com/coder/coder/v2/agent/agentssh"
@@ -33,7 +35,6 @@ import (
 	"github.com/coder/coder/v2/cli/clilog"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
-	"github.com/coder/serpent"
 )
 
 func (r *RootCmd) workspaceAgent() *serpent.Command {
@@ -62,8 +63,10 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 		// This command isn't useful to manually execute.
 		Hidden: true,
 		Handler: func(inv *serpent.Invocation) error {
-			ctx, cancel := context.WithCancel(inv.Context())
-			defer cancel()
+			ctx, cancel := context.WithCancelCause(inv.Context())
+			defer func() {
+				cancel(xerrors.New("agent exited"))
+			}()
 
 			var (
 				ignorePorts = map[int]string{}
@@ -280,7 +283,6 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 				return xerrors.Errorf("add executable to $PATH: %w", err)
 			}
 
-			prometheusRegistry := prometheus.NewRegistry()
 			subsystemsRaw := inv.Environ.Get(agent.EnvAgentSubsystem)
 			subsystems := []codersdk.AgentSubsystem{}
 			for _, s := range strings.Split(subsystemsRaw, ",") {
@@ -324,45 +326,69 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 				logger.Info(ctx, "agent devcontainer detection not enabled")
 			}
 
-			agnt := agent.New(agent.Options{
-				Client:        client,
-				Logger:        logger,
-				LogDir:        logDir,
-				ScriptDataDir: scriptDataDir,
-				// #nosec G115 - Safe conversion as tailnet listen port is within uint16 range (0-65535)
-				TailnetListenPort: uint16(tailnetListenPort),
-				ExchangeToken: func(ctx context.Context) (string, error) {
-					if exchangeToken == nil {
-						return client.SDK.SessionToken(), nil
-					}
-					resp, err := exchangeToken(ctx)
-					if err != nil {
-						return "", err
-					}
-					client.SetSessionToken(resp.SessionToken)
-					return resp.SessionToken, nil
-				},
-				EnvironmentVariables: environmentVariables,
-				IgnorePorts:          ignorePorts,
-				SSHMaxTimeout:        sshMaxTimeout,
-				Subsystems:           subsystems,
-
-				PrometheusRegistry: prometheusRegistry,
-				BlockFileTransfer:  blockFileTransfer,
-				Execer:             execer,
-
-				ExperimentalDevcontainersEnabled: experimentalDevcontainersEnabled,
-			})
-
-			promHandler := agent.PrometheusMetricsHandler(prometheusRegistry, logger)
-			prometheusSrvClose := ServeHandler(ctx, logger, promHandler, prometheusAddress, "prometheus")
-			defer prometheusSrvClose()
-
-			debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")
-			defer debugSrvClose()
-
-			<-ctx.Done()
-			return agnt.Close()
+			reinitEvents := agentsdk.WaitForReinitLoop(ctx, logger, client)
+
+			var (
+				lastErr  error
+				mustExit bool
+			)
+			for {
+				prometheusRegistry := prometheus.NewRegistry()
+
+				agnt := agent.New(agent.Options{
+					Client:        client,
+					Logger:        logger,
+					LogDir:        logDir,
+					ScriptDataDir: scriptDataDir,
+					// #nosec G115 - Safe conversion as tailnet listen port is within uint16 range (0-65535)
+					TailnetListenPort: uint16(tailnetListenPort),
+					ExchangeToken: func(ctx context.Context) (string, error) {
+						if exchangeToken == nil {
+							return client.SDK.SessionToken(), nil
+						}
+						resp, err := exchangeToken(ctx)
+						if err != nil {
+							return "", err
+						}
+						client.SetSessionToken(resp.SessionToken)
+						return resp.SessionToken, nil
+					},
+					EnvironmentVariables: environmentVariables,
+					IgnorePorts:          ignorePorts,
+					SSHMaxTimeout:        sshMaxTimeout,
+					Subsystems:           subsystems,
+
+					PrometheusRegistry:               prometheusRegistry,
+					BlockFileTransfer:                blockFileTransfer,
+					Execer:                           execer,
+					ExperimentalDevcontainersEnabled: experimentalDevcontainersEnabled,
+				})
+
+				promHandler := agent.PrometheusMetricsHandler(prometheusRegistry, logger)
+				prometheusSrvClose := ServeHandler(ctx, logger, promHandler, prometheusAddress, "prometheus")
+
+				debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")
+
+				select {
+				case <-ctx.Done():
+					logger.Warn(ctx, "agent shutting down", slog.Error(ctx.Err()), slog.Error(context.Cause(ctx)))
+					mustExit = true
+				case event := <-reinitEvents:
+					logger.Warn(ctx, "agent received instruction to reinitialize",
+						slog.F("user_id", event.UserID), slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
+				}
+
+				lastErr = agnt.Close()
+				debugSrvClose()
+				prometheusSrvClose()
+
+				if mustExit {
+					break
+				}
+
+				logger.Info(ctx, "agent reinitializing")
+			}
+			return lastErr
 		},
 	}
 

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -21,7 +21,10 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/prebuilds"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 	"github.com/coder/coder/v2/testutil"
@@ -321,6 +324,63 @@ func TestWorkspaceAgent(t *testing.T) {
 	})
 }
 
+func TestAgent_Prebuild(t *testing.T) {
+	t.Parallel()
+
+	db, pubsub := dbtestutil.NewDB(t)
+	client := coderdtest.New(t, &coderdtest.Options{
+		Database: db,
+		Pubsub:   pubsub,
+	})
+	user := coderdtest.CreateFirstUser(t, client)
+	presetID := uuid.New()
+	tv := dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{
+		OrganizationID: user.OrganizationID,
+		CreatedBy:      user.UserID,
+	}).Preset(database.TemplateVersionPreset{
+		ID: presetID,
+	}).Do()
+	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OwnerID:    prebuilds.SystemUserID,
+		TemplateID: tv.Template.ID,
+	}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
+		a[0].Scripts = []*proto.Script{
+			{
+				DisplayName: "Prebuild Test Script",
+				Script:      "sleep 5", // Make reinitialization take long enough to assert that it happened
+				RunOnStart:  true,
+			},
+		}
+		return a
+	}).Do()
+
+	// Spin up an agent
+	logDir := t.TempDir()
+	inv, _ := clitest.New(t,
+		"agent",
+		"--auth", "token",
+		"--agent-token", r.AgentToken,
+		"--agent-url", client.URL.String(),
+		"--log-dir", logDir,
+	)
+	clitest.Start(t, inv)
+
+	// Check that the agent is in a happy steady state
+	waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID)
+	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.AgentsNotReady)
+
+	// Check that reinitialization completed
+	waiter.WaitFor(coderdtest.AgentsReady)
+}
+
 func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
 	if len(rs) < 1 {
 		return false

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json