diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 075c83ecc0..de603b1a3e 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -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()) diff --git a/coderd/chatd/chattool/createworkspace.go b/coderd/chatd/chattool/createworkspace.go index 752aa22c19..644b1d51b3 100644 --- a/coderd/chatd/chattool/createworkspace.go +++ b/coderd/chatd/chattool/createworkspace.go @@ -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 { diff --git a/coderd/chatd/chattool/createworkspace_test.go b/coderd/chatd/chattool/createworkspace_test.go new file mode 100644 index 0000000000..3b8c914eb4 --- /dev/null +++ b/coderd/chatd/chattool/createworkspace_test.go @@ -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) + }) +}