fix(chatd): wait for startup scripts before returning from create_workspace (#22498)

The `create_workspace` tool waited for the workspace build to succeed and the agent to become connectable, but did not wait for the agent's startup scripts (e.g. git clone) to finish. This caused agents to attempt file operations on repositories that hadn't been cloned yet. Add a waitForStartupScripts step that polls the agent's lifecycle_state via GetWorkspaceAgentLifecycleStateByID until it transitions out of created/starting into a terminal state (ready, start_error, or start_timeout). The tool now only returns success once the workspace is fully initialized. If the scripts fail or time out, the tool still returns (non-fatal) with an appropriate agent_status so the model knows something went wrong. Created using thingies (Opus 4.6 Max)
2026-09-24 15:04:27 +08:00 · 2026-03-05 09:42:12 +00:00
parent 9f83eb1544
commit 4d0d187806
3 changed files with 255 additions and 60 deletions
@@ -30,6 +30,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisioner/echo"
+	proto "github.com/coder/coder/v2/provisionersdk/proto"
 	"github.com/coder/coder/v2/testutil"
 )

@@ -619,7 +620,7 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) {
 		Database:                   db,
 		ReplicaID:                  uuid.New(),
 		Pubsub:                     ps,
-		PendingChatAcquireInterval: testutil.WaitSuperLong,
+		PendingChatAcquireInterval: testutil.WaitLong,
 		InFlightChatStaleAfter:     staleAfter,
 	})
 	t.Cleanup(func() {
@@ -733,7 +734,7 @@ func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) {
 		Database:                   db,
 		ReplicaID:                  uuid.New(),
 		Pubsub:                     ps,
-		PendingChatAcquireInterval: testutil.WaitSuperLong,
+		PendingChatAcquireInterval: testutil.WaitLong,
 		InFlightChatStaleAfter:     500 * time.Millisecond,
 	})
 	t.Cleanup(func() {
@@ -969,11 +970,20 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) {
 	user := coderdtest.CreateFirstUser(t, client)

 	agentToken := uuid.NewString()
+	// Add a startup script so the agent spends time in the
+	// "starting" lifecycle state. This lets us verify that
+	// create_workspace waits for scripts to finish.
 	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
 		Parse:          echo.ParseComplete,
 		ProvisionPlan:  echo.PlanComplete,
 		ProvisionApply: echo.ApplyComplete,
-		ProvisionGraph: echo.ProvisionGraphWithAgent(agentToken),
+		ProvisionGraph: echo.ProvisionGraphWithAgent(agentToken, func(g *proto.GraphComplete) {
+			g.Resources[0].Agents[0].Scripts = []*proto.Script{{
+				DisplayName: "setup",
+				Script:      "sleep 5",
+				RunOnStart:  true,
+			}}
+		}),
 	})
 	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
 	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
@@ -1082,6 +1092,20 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) {
 	}
 	require.True(t, foundCreateWorkspaceResult, "expected create_workspace tool result message")

+	// Verify that the tool waited for startup scripts to
+	// complete. The agent should be in "ready" state by the
+	// time create_workspace returns its result.
+	workspace, err = client.Workspace(ctx, workspaceID)
+	require.NoError(t, err)
+	var agentLifecycle codersdk.WorkspaceAgentLifecycle
+	for _, res := range workspace.LatestBuild.Resources {
+		for _, agt := range res.Agents {
+			agentLifecycle = agt.LifecycleState
+		}
+	}
+	require.Equal(t, codersdk.WorkspaceAgentLifecycleReady, agentLifecycle,
+		"agent should be ready after create_workspace returns; startup scripts were not awaited")
+
 	require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2))
 	streamedCallsMu.Lock()
 	recordedStreamCalls := append([][]chattest.OpenAIMessage(nil), streamedCalls...)
@@ -1123,7 +1147,7 @@ func newTestServer(
 		Database:                   db,
 		ReplicaID:                  replicaID,
 		Pubsub:                     ps,
-		PendingChatAcquireInterval: testutil.WaitSuperLong,
+		PendingChatAcquireInterval: testutil.WaitLong,
 	})
 	t.Cleanup(func() {
 		require.NoError(t, server.Close())
@@ -1330,7 +1354,7 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T)
 		ReplicaID:                  uuid.New(),
 		Pubsub:                     ps,
 		PendingChatAcquireInterval: 10 * time.Millisecond,
-		InFlightChatStaleAfter:     testutil.WaitSuperLong,
+		InFlightChatStaleAfter:     testutil.WaitLong,
 	})
 	t.Cleanup(func() {
 		require.NoError(t, serverA.Close())
@@ -1383,7 +1407,7 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T)
 		ReplicaID:                  uuid.New(),
 		Pubsub:                     ps,
 		PendingChatAcquireInterval: 10 * time.Millisecond,
-		InFlightChatStaleAfter:     testutil.WaitSuperLong,
+		InFlightChatStaleAfter:     testutil.WaitLong,
 	})
 	t.Cleanup(func() {
 		require.NoError(t, serverB.Close())
@@ -3,6 +3,7 @@ package chattool
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"strings"
 	"sync"
@@ -37,6 +38,13 @@ const (
 	// agentPingTimeout is the timeout for a single agent ping
 	// when checking whether an existing workspace is alive.
 	agentPingTimeout = 5 * time.Second
+	// startupScriptTimeout is the maximum time to wait for the
+	// workspace agent's startup scripts to finish after the agent
+	// is reachable.
+	startupScriptTimeout = 10 * time.Minute
+	// startupScriptPollInterval is how often we check the agent's
+	// lifecycle state while waiting for startup scripts.
+	startupScriptPollInterval = 2 * time.Second
 )

 // CreateWorkspaceFn creates a workspace for the given owner.
@@ -194,27 +202,24 @@ func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool {
 				})
 			}

-			// Wait for the agent to come online.
-			if workspaceAgentID != uuid.Nil && options.AgentConnFn != nil {
-				if err := waitForAgent(ctx, options.AgentConnFn, workspaceAgentID); err != nil {
-					// Non-fatal: the workspace was created
-					// successfully, the agent just isn't ready
-					// yet. The model can retry.
-					return toolResponse(map[string]any{
-						"created":        true,
-						"workspace_name": workspace.FullName(),
-						"agent_status":   "not_ready",
-						"agent_error":    err.Error(),
-					}), nil
+			// Wait for the agent to come online and startup scripts to finish.
+			if workspaceAgentID != uuid.Nil {
+				agentStatus := waitForAgentReady(ctx, options.DB, workspaceAgentID, options.AgentConnFn)
+				result := map[string]any{
+					"created":        true,
+					"workspace_name": workspace.FullName(),
 				}
+				for k, v := range agentStatus {
+					result[k] = v
+				}
+				return toolResponse(result), nil
 			}

 			return toolResponse(map[string]any{
 				"created":        true,
 				"workspace_name": workspace.FullName(),
 			}), nil
-		},
-	)
+		})
 }

 // checkExistingWorkspace checks whether the chat already has a usable
@@ -268,34 +273,42 @@ func checkExistingWorkspace(
 				"existing workspace build failed: %w", err,
 			)
 		}
-		return map[string]any{
+		result := map[string]any{
 			"created":        false,
 			"workspace_name": ws.Name,
 			"status":         "already_exists",
-			"message":        "workspace was already being built and is now ready",
-		}, true, nil
+			"message":        "workspace build completed",
+		}
+		agents, agentsErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID)
+		if agentsErr == nil && len(agents) > 0 {
+			for k, v := range waitForAgentReady(ctx, db, agents[0].ID, agentConnFn) {
+				result[k] = v
+			}
+		}
+		return result, true, nil

 	case database.ProvisionerJobStatusSucceeded:
 		// Build succeeded — check if agent is reachable.
 		agents, agentsErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID)
 		if agentsErr == nil && len(agents) > 0 && agentConnFn != nil {
-			pingCtx, cancel := context.WithTimeout(
-				ctx, agentPingTimeout,
-			)
-			defer cancel()
-
-			conn, release, connErr := agentConnFn(
-				pingCtx, agents[0].ID,
-			)
+			pingCtx, cancel := context.WithTimeout(ctx, agentPingTimeout)
+			conn, release, connErr := agentConnFn(pingCtx, agents[0].ID)
+			cancel()
 			if connErr == nil {
 				release()
 				_ = conn
-				return map[string]any{
+				// Agent is reachable; wait for startup scripts.
+				result := map[string]any{
 					"created":        false,
 					"workspace_name": ws.Name,
 					"status":         "already_exists",
 					"message":        "workspace is already running and reachable",
-				}, true, nil
+				}
+				// Pass nil for agentConnFn since we already confirmed connectivity.
+				for k, v := range waitForAgentReady(ctx, db, agents[0].ID, nil) {
+					result[k] = v
+				}
+				return result, true, nil
 			}
 			// Agent unreachable — workspace is dead, allow
 			// creation.
@@ -365,40 +378,88 @@ func waitForBuild(
 	}
 }

-// waitForAgent retries connecting to the workspace agent until it
-// succeeds or the timeout expires.
-func waitForAgent(
+// waitForAgentReady waits for the workspace agent to become
+// reachable and for its startup scripts to finish. It returns
+// status fields suitable for merging into a tool response.
+func waitForAgentReady(
 	ctx context.Context,
-	agentConnFn AgentConnFunc,
+	db database.Store,
 	agentID uuid.UUID,
-) error {
-	agentCtx, cancel := context.WithTimeout(ctx, agentConnectTimeout)
-	defer cancel()
+	agentConnFn AgentConnFunc,
+) map[string]any {
+	result := map[string]any{}

-	ticker := time.NewTicker(agentRetryInterval)
-	defer ticker.Stop()
+	// Phase 1: retry connecting to the agent.
+	if agentConnFn != nil {
+		agentCtx, agentCancel := context.WithTimeout(ctx, agentConnectTimeout)
+		defer agentCancel()

-	var lastErr error
-	for {
-		attemptCtx, attemptCancel := context.WithTimeout(agentCtx, agentAttemptTimeout)
-		conn, release, err := agentConnFn(attemptCtx, agentID)
-		attemptCancel()
-		if err == nil {
-			release()
-			_ = conn
-			return nil
-		}
-		lastErr = err
+		ticker := time.NewTicker(agentRetryInterval)
+		defer ticker.Stop()

-		select {
-		case <-agentCtx.Done():
-			return xerrors.Errorf(
-				"timed out waiting for workspace agent: %w",
-				lastErr,
-			)
-		case <-ticker.C:
+		var lastErr error
+		for {
+			attemptCtx, attemptCancel := context.WithTimeout(agentCtx, agentAttemptTimeout)
+			conn, release, err := agentConnFn(attemptCtx, agentID)
+			attemptCancel()
+			if err == nil {
+				release()
+				_ = conn
+				break
+			}
+			lastErr = err
+
+			select {
+			case <-agentCtx.Done():
+				result["agent_status"] = "not_ready"
+				result["agent_error"] = lastErr.Error()
+				return result
+			case <-ticker.C:
+			}
 		}
 	}
+
+	// Phase 2: poll lifecycle until startup scripts finish.
+	if db != nil {
+		scriptCtx, scriptCancel := context.WithTimeout(ctx, startupScriptTimeout)
+		defer scriptCancel()
+
+		ticker := time.NewTicker(startupScriptPollInterval)
+		defer ticker.Stop()
+
+		var lastState database.WorkspaceAgentLifecycleState
+		for {
+			row, err := db.GetWorkspaceAgentLifecycleStateByID(scriptCtx, agentID)
+			if err == nil {
+				lastState = row.LifecycleState
+				switch lastState {
+				case database.WorkspaceAgentLifecycleStateCreated,
+					database.WorkspaceAgentLifecycleStateStarting:
+					// Still in progress, keep polling.
+				case database.WorkspaceAgentLifecycleStateReady:
+					return result
+				default:
+					// Terminal non-ready state.
+					result["startup_scripts"] = "startup_scripts_failed"
+					result["lifecycle_state"] = string(lastState)
+					return result
+				}
+			}
+
+			select {
+			case <-scriptCtx.Done():
+				if errors.Is(scriptCtx.Err(), context.DeadlineExceeded) {
+					result["startup_scripts"] = "startup_scripts_timeout"
+				} else {
+					result["startup_scripts"] = "startup_scripts_unknown"
+				}
+				return result
+			case <-ticker.C:
+			}
+		}
+	}
+
+	return result
 }

 func generatedWorkspaceName(seed string) string {
@@ -0,0 +1,110 @@
+package chattool //nolint:testpackage // Uses internal symbols.
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbmock"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+)
+
+func TestWaitForAgentReady(t *testing.T) {
+	t.Parallel()
+
+	t.Run("AgentConnectsAndLifecycleReady", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		agentID := uuid.New()
+
+		// Mock returns Ready lifecycle state.
+		db.EXPECT().
+			GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
+			Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
+				LifecycleState: database.WorkspaceAgentLifecycleStateReady,
+			}, nil)
+
+		// AgentConnFn succeeds immediately.
+		connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
+			return nil, func() {}, nil
+		}
+
+		result := waitForAgentReady(context.Background(), db, agentID, connFn)
+		require.Empty(t, result)
+	})
+
+	t.Run("AgentConnectTimeout", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		agentID := uuid.New()
+
+		// AgentConnFn always fails - context will timeout.
+		connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
+			return nil, nil, context.DeadlineExceeded
+		}
+
+		// Use a context that's already canceled to avoid waiting.
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		result := waitForAgentReady(ctx, db, agentID, connFn)
+		require.Equal(t, "not_ready", result["agent_status"])
+		require.NotEmpty(t, result["agent_error"])
+	})
+
+	t.Run("AgentConnectsButStartupFails", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		agentID := uuid.New()
+
+		// Mock returns StartError lifecycle state.
+		db.EXPECT().
+			GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
+			Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
+				LifecycleState: database.WorkspaceAgentLifecycleStateStartError,
+			}, nil)
+
+		connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
+			return nil, func() {}, nil
+		}
+
+		result := waitForAgentReady(context.Background(), db, agentID, connFn)
+		require.Equal(t, "startup_scripts_failed", result["startup_scripts"])
+		require.Equal(t, "start_error", result["lifecycle_state"])
+	})
+
+	t.Run("NilAgentConnFn", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		agentID := uuid.New()
+
+		// Mock returns Ready lifecycle state.
+		db.EXPECT().
+			GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
+			Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
+				LifecycleState: database.WorkspaceAgentLifecycleStateReady,
+			}, nil)
+
+		result := waitForAgentReady(context.Background(), db, agentID, nil)
+		require.Empty(t, result)
+	})
+
+	t.Run("NilDB", func(t *testing.T) {
+		t.Parallel()
+
+		connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
+			return nil, func() {}, nil
+		}
+
+		result := waitForAgentReady(context.Background(), nil, uuid.New(), connFn)
+		require.Empty(t, result)
+	})
+}