test agent reinitialization

coder · SasSwart · May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025
commit 7ce4eea2af98c805bd61ade3be14086c2e684f4a
diff --git a/cli/agent.go b/cli/agent.go
@@ -413,7 +413,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					break
 				}
 
-				logger.Info(ctx, "reinitializing...")
+				logger.Info(ctx, "agent reinitializing")
 			}
 			return lastErr
 		},

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -21,7 +21,9 @@ 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/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 +323,56 @@ 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)
+	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OrganizationID: user.OrganizationID,
+		OwnerID:        user.UserID,
+	}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
+		a[0].Scripts = []*proto.Script{
+			{
+				DisplayName: "Prebuild Test Script",
+				Script:      "sleep 5", // Make reinitiazation 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.AgentReady)
+
+	// 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.AgentNotReady)
+
+	// Check that reinitialization completed
+	waiter.WaitFor(coderdtest.AgentReady)
+}
+
 func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
 	if len(rs) < 1 {
 		return false

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -1105,6 +1105,59 @@ func (w WorkspaceAgentWaiter) MatchResources(m func([]codersdk.WorkspaceResource
 	return w
 }
 
+type WaitForCriterium func(agent codersdk.WorkspaceAgent) bool
+
+func AgentReady(agent codersdk.WorkspaceAgent) bool {
+	return agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady
+}
+
+func AgentNotReady(agent codersdk.WorkspaceAgent) bool {
+	return !AgentReady(agent)
+}
+
+func (w WorkspaceAgentWaiter) WaitFor(criteria ...WaitForCriterium) {
+	w.t.Helper()
+
+	agentNamesMap := make(map[string]struct{}, len(w.agentNames))
+	for _, name := range w.agentNames {
+		agentNamesMap[name] = struct{}{}
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+
+	w.t.Logf("waiting for workspace agents (workspace %s)", w.workspaceID)
+	require.Eventually(w.t, func() bool {
+		var err error
+		workspace, err := w.client.Workspace(ctx, w.workspaceID)
+		if err != nil {
+			return false
+		}
+		if workspace.LatestBuild.Job.CompletedAt == nil {
+			return false
+		}
+		if workspace.LatestBuild.Job.CompletedAt.IsZero() {
+			return false
+		}
+
+		for _, resource := range workspace.LatestBuild.Resources {
+			for _, agent := range resource.Agents {
+				if len(w.agentNames) > 0 {
+					if _, ok := agentNamesMap[agent.Name]; !ok {
+						continue
+					}
+				}
+				for _, criterium := range criteria {
+					if !criterium(agent) {
+						return false
+					}
+				}
+			}
+		}
+		return true
+	}, testutil.WaitLong, testutil.IntervalMedium)
+}
+
 // Wait waits for the agent(s) to connect and fails the test if they do not within testutil.WaitLong
 func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
 	w.t.Helper()

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -689,10 +689,13 @@ func createWorkspace(
 				api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
 					slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
 			}
-			agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
-			for _, agent := range agents {
-				agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
+			if len(agents) > 1 {
+				return xerrors.Errorf("multiple agents are not yet supported in prebuilt workspaces")
 			}
+			// agentTokensByAgentID = make(map[uuid.UUID]string, len(agents))
+			// for _, agent := range agents {
+			// 	agentTokensByAgentID[agent.ID] = agent.AuthToken.String()
+			// }
 		}
 
 		// We have to refetch the workspace for the joined in fields.

diff --git a/go.sum b/go.sum
@@ -921,8 +921,6 @@ github.com/coder/tailscale v1.1.1-0.20250422090654-5090e715905e h1:nope/SZfoLB9M
 github.com/coder/tailscale v1.1.1-0.20250422090654-5090e715905e/go.mod h1:1ggFFdHTRjPRu9Yc1yA7nVHBYB50w9Ce7VIXNqcW6Ko=
 github.com/coder/terraform-config-inspect v0.0.0-20250107175719-6d06d90c630e h1:JNLPDi2P73laR1oAclY6jWzAbucf70ASAvf5mh2cME0=
 github.com/coder/terraform-config-inspect v0.0.0-20250107175719-6d06d90c630e/go.mod h1:Gz/z9Hbn+4KSp8A2FBtNszfLSdT2Tn/uAKGuVqqWmDI=
-github.com/coder/terraform-provider-coder/v2 v2.4.0-pre1.0.20250417100258-c86bb5c3ddcd h1:FsIG6Fd0YOEK7D0Hl/CJywRA+Y6Gd5RQbSIa2L+/BmE=
-github.com/coder/terraform-provider-coder/v2 v2.4.0-pre1.0.20250417100258-c86bb5c3ddcd/go.mod h1:56/KdGYaA+VbwXJbTI8CA57XPfnuTxN8rjxbR34PbZw=
 github.com/coder/trivy v0.0.0-20250409153844-e6b004bc465a h1:yryP7e+IQUAArlycH4hQrjXQ64eRNbxsV5/wuVXHgME=
 github.com/coder/trivy v0.0.0-20250409153844-e6b004bc465a/go.mod h1:dDvq9axp3kZsT63gY2Znd1iwzfqDq3kXbQnccIrjRYY=
 github.com/coder/websocket v1.8.13 h1:f3QZdXy7uGVz+4uCJy2nTZyM0yTBj8yANEHhqlXZ9FE=
-Original file line number
+Diff line change
@@ Expand Up / @@ -413,7 +413,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command { @@
     					break
     				}
-    				logger.Info(ctx, "reinitializing...")
+    				logger.Info(ctx, "agent reinitializing")
     			}
     			return lastErr
     		},
@@ Expand Down @@