diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 560653fe75..a8e944c016 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -4468,7 +4468,7 @@ func TestPatchChat(t *testing.T) { t.Run("WorkspaceBinding", func(t *testing.T) { t.Parallel() - t.Run("BindValidWorkspace", func(t *testing.T) { + t.Run("BindExistingExternalWorkspace", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) @@ -4482,6 +4482,8 @@ func TestPatchChat(t *testing.T) { workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OrganizationID: firstUser.OrganizationID, OwnerID: firstUser.UserID, + }).Seed(database.WorkspaceBuild{ + HasExternalAgent: sql.NullBool{Bool: true, Valid: true}, }).WithAgent().Do() chat := createStoredChat( ctx, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index e6bc005b8a..ff5ba1f3f4 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -143,8 +143,27 @@ var ( "connection to the workspace agent timed out. " + "The workspace may need to be restarted from the Coder dashboard", ) + errChatExternalAgentUnavailable = xerrors.New("external workspace agent unavailable") ) +type chatExternalAgentUnavailableError struct { + message string +} + +func (e chatExternalAgentUnavailableError) Error() string { + return e.message +} + +func (chatExternalAgentUnavailableError) Is(target error) bool { + return target == errChatExternalAgentUnavailable +} + +func newChatExternalAgentUnavailableError(agent database.WorkspaceAgent) error { + return chatExternalAgentUnavailableError{ + message: chattool.ExternalAgentUnavailableMessage(agent), + } +} + // Server handles background processing of pending chats. type Server struct { cancel context.CancelFunc @@ -764,6 +783,46 @@ func isAgentUnreachable(now time.Time, agent database.WorkspaceAgent, inactiveTi status.Status == database.WorkspaceAgentStatusTimeout } +func (c *turnWorkspaceContext) externalAgentError( + ctx context.Context, + agent database.WorkspaceAgent, + fallback error, +) error { + isExternal, err := chattool.IsExternalWorkspaceAgent(ctx, c.server.db, agent) + if err != nil || !isExternal { + return fallback + } + return newChatExternalAgentUnavailableError(agent) +} + +func (c *turnWorkspaceContext) externalAgentPreflightError( + ctx context.Context, + chatSnapshot database.Chat, + agent database.WorkspaceAgent, +) error { + // Mirror the cache-hit gate: only short-circuit on clearly offline + // states (Disconnected/Timeout). Connecting is allowed through so + // an external agent the user just started can still connect inside + // the normal dial window. + if !isAgentUnreachable(c.server.clock.Now(), agent, c.server.agentInactiveDisconnectTimeout) { + return nil + } + + isExternal, err := chattool.IsExternalWorkspaceAgent(ctx, c.server.db, agent) + if err != nil || !isExternal || !chatSnapshot.WorkspaceID.Valid { + return nil + } + + // Stale agent bindings rely on dialWithLazyValidation to discover + // replacement agents, so only skip the dial when this agent is still + // the latest selected chat agent for the workspace. + latestAgentID, err := c.latestWorkspaceAgentID(ctx, chatSnapshot.WorkspaceID.UUID) + if err != nil || latestAgentID != agent.ID { + return nil + } + return newChatExternalAgentUnavailableError(agent) +} + func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspacesdk.AgentConn, error) { if c.server.agentConnFn == nil { return nil, xerrors.New("workspace agent connector is not configured") @@ -793,7 +852,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces // next tool call. } else if isAgentUnreachable(c.server.clock.Now(), freshAgent, c.server.agentInactiveDisconnectTimeout) { c.clearCachedWorkspaceState() - return nil, errChatAgentDisconnected + return nil, c.externalAgentError(ctx, freshAgent, errChatAgentDisconnected) } } return currentConn, nil @@ -806,6 +865,9 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces if err != nil { return nil, err } + if err := c.externalAgentPreflightError(ctx, chatSnapshot, agent); err != nil { + return nil, err + } // Wrap the dial in a timeout to bound the time spent // waiting for an unreachable agent. The timeout scopes @@ -833,7 +895,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces // canceled (e.g. ErrInterrupted), its error must // propagate unchanged so the chatloop can detect it. if ctx.Err() == nil && errors.Is(context.Cause(dialCtx), errChatDialTimeout) { - return nil, errChatDialTimeout + return nil, c.externalAgentError(ctx, agent, errChatDialTimeout) } return nil, err } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 093d2da517..b23ec23d18 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -1808,6 +1808,7 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes workspaceID := uuid.New() staleAgentID := uuid.New() + resourceID := uuid.New() chat := database.Chat{ ID: uuid.New(), WorkspaceID: uuid.NullUUID{ @@ -1820,7 +1821,7 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes }, } - staleAgent := database.WorkspaceAgent{ID: staleAgentID} + staleAgent := database.WorkspaceAgent{ID: staleAgentID, ResourceID: resourceID} db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), staleAgentID). Return(staleAgent, nil). @@ -1828,6 +1829,12 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil). Times(1) + db.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID). + Return(database.WorkspaceResource{ + ID: resourceID, + Type: chattool.ExternalAgentResourceType, + }, nil). + AnyTimes() server := &Server{ db: db, @@ -1852,6 +1859,7 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes gotConn, err := workspaceCtx.getWorkspaceConn(ctx) require.Nil(t, gotConn) require.ErrorIs(t, err, errChatHasNoWorkspaceAgent) + require.NotErrorIs(t, err, errChatExternalAgentUnavailable) workspaceCtx.mu.Lock() defer workspaceCtx.mu.Unlock() @@ -4538,12 +4546,10 @@ func TestGetWorkspaceConn_DialTimeoutParentCanceled(t *testing.T) { require.ErrorIs(t, err, context.Canceled) } -func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { - // Regression test: a non-timeout dial error (e.g. auth - // failure) with the parent context still alive must NOT be - // converted to errChatDialTimeout. Before the fix, - // dialCancel() poisoned dialCtx.Err(), causing all errors - // to be misclassified. +func TestGetWorkspaceConn_PreflightExternalAgentTimedOut(t *testing.T) { + // External agent never connected and the connection window has + // elapsed (Timeout). Preflight must short-circuit before any + // dial attempt and return the external-agent error. t.Parallel() ctrl := gomock.NewController(t) @@ -4551,6 +4557,151 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { workspaceID := uuid.New() agentID := uuid.New() + resourceID := uuid.New() + agent := database.WorkspaceAgent{ + ID: agentID, + Name: "main", + ResourceID: resourceID, + CreatedAt: time.Now().Add(-10 * time.Minute), + ConnectionTimeoutSeconds: 60, + } + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(agent, nil). + Times(1) + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{agent}, nil). + Times(1) + db.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID). + Return(database.WorkspaceResource{ + ID: resourceID, + Type: chattool.ExternalAgentResourceType, + }, nil). + Times(1) + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: defaultDialTimeout, + } + server.agentConnFn = func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + t.Fatal("unexpected agent dial for external agent preflight") + return nil, nil, xerrors.New("unexpected agent dial") + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil }, + } + defer workspaceCtx.close() + + ctx := testutil.Context(t, testutil.WaitMedium) + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + require.Nil(t, gotConn) + require.ErrorIs(t, err, errChatExternalAgentUnavailable) + require.Equal(t, chattool.ExternalAgentUnavailableMessage(agent), err.Error()) +} + +func TestGetWorkspaceConn_PreflightExternalAgentConnectingDials(t *testing.T) { + // External agent in the Connecting state (never connected yet, + // still inside ConnectionTimeoutSeconds) must fall through to the + // dial so the user can succeed in the same turn if they just + // started the agent on their host. + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + resourceID := uuid.New() + agent := database.WorkspaceAgent{ + ID: agentID, + Name: "main", + ResourceID: resourceID, + CreatedAt: time.Now().Add(-1 * time.Second), + ConnectionTimeoutSeconds: 600, + } + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(agent, nil). + Times(1) + + conn := agentconnmock.NewMockAgentConn(ctrl) + conn.EXPECT().SetExtraHeaders(gomock.Any()).Times(1) + + dialed := false + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: defaultDialTimeout, + } + server.agentConnFn = func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { + dialed = true + require.Equal(t, agentID, id) + return conn, func() {}, nil + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil }, + } + defer workspaceCtx.close() + + ctx := testutil.Context(t, testutil.WaitMedium) + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + require.NoError(t, err) + require.Same(t, conn, gotConn) + require.True(t, dialed, "preflight must let Connecting external agents reach the dial") +} + +func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { + // Regression test: a non-timeout dial error (e.g. auth + // failure) with the parent context still alive must NOT be + // converted to errChatDialTimeout or masked as external-agent + // unavailability. + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + resourceID := uuid.New() chat := database.Chat{ ID: uuid.New(), WorkspaceID: uuid.NullUUID{ @@ -4564,7 +4715,8 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { } connectedAgent := database.WorkspaceAgent{ - ID: agentID, + ID: agentID, + ResourceID: resourceID, FirstConnectedAt: sql.NullTime{ Time: time.Now().Add(-1 * time.Minute), Valid: true, @@ -4585,6 +4737,12 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{connectedAgent}, nil). AnyTimes() + db.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID). + Return(database.WorkspaceResource{ + ID: resourceID, + Type: chattool.ExternalAgentResourceType, + }, nil). + AnyTimes() dialErr := xerrors.New("authentication failed") server := &Server{ @@ -4613,9 +4771,11 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) gotConn, err := workspaceCtx.getWorkspaceConn(ctx) require.Nil(t, gotConn) - // Must NOT be misclassified as a dial timeout. + // Must NOT be misclassified as a dial timeout or external-agent outage. require.NotErrorIs(t, err, errChatDialTimeout) + require.NotErrorIs(t, err, errChatExternalAgentUnavailable) // The original dial error should propagate. + require.ErrorIs(t, err, dialErr) require.ErrorContains(t, err, "authentication failed") } diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index 5d96731247..2766733569 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -2,6 +2,7 @@ package chattool import ( "context" + "database/sql" "errors" "fmt" "strings" @@ -171,6 +172,16 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option ), nil } + hasExternalAgent, externalAgentErr := templateHasExternalAgent(ctx, db, tmpl) + if externalAgentErr != nil { + return fantasy.NewTextErrorResponse( + xerrors.Errorf("look up template version: %w", externalAgentErr).Error(), + ), nil + } + if hasExternalAgent { + return fantasy.NewTextErrorResponse(createWorkspaceExternalAgentMessage), nil + } + var ttlMs *int64 raw, err := db.GetChatWorkspaceTTL(ctx) if err != nil { @@ -291,7 +302,7 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option // Select the chat agent so follow-up tools wait on the // intended workspace agent. - workspaceAgentID := uuid.Nil + selectedAgent := database.WorkspaceAgent{} agents, agentErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID) if agentErr == nil { if len(agents) == 0 { @@ -302,14 +313,14 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option result["agent_status"] = "selection_error" result["agent_error"] = selectErr.Error() } else { - workspaceAgentID = selected.ID + selectedAgent = selected } } } // Wait for the agent to come online and startup scripts to finish. - if workspaceAgentID != uuid.Nil { - agentStatus := waitForAgentReady(ctx, db, workspaceAgentID, options.AgentConnFn) + if selectedAgent.ID != uuid.Nil { + agentStatus := waitForAgentReady(ctx, db, selectedAgent, options.AgentConnFn) for k, v := range agentStatus { result[k] = v } @@ -443,7 +454,7 @@ func (o CreateWorkspaceOptions) checkExistingWorkspace( ) selected = agents[0] } - for k, v := range waitForAgentReady(ctx, db, selected.ID, agentConnFn) { + for k, v := range waitForAgentReady(ctx, db, selected, agentConnFn) { result[k] = v } } @@ -484,13 +495,13 @@ func (o CreateWorkspaceOptions) checkExistingWorkspace( switch status.Status { case database.WorkspaceAgentStatusConnected: result["message"] = "workspace is already running and recently connected" - for k, v := range waitForAgentReady(ctx, db, selected.ID, nil) { + for k, v := range waitForAgentReady(ctx, db, selected, nil) { result[k] = v } return existingWorkspaceResult{Result: result, Done: true} case database.WorkspaceAgentStatusConnecting: result["message"] = "workspace exists and the agent is still connecting" - for k, v := range waitForAgentReady(ctx, db, selected.ID, agentConnFn) { + for k, v := range waitForAgentReady(ctx, db, selected, agentConnFn) { result[k] = v } return existingWorkspaceResult{Result: result, Done: true} @@ -567,16 +578,48 @@ func waitForBuild( } } +func templateHasExternalAgent( + ctx context.Context, + db database.Store, + tmpl database.Template, +) (bool, error) { + version, err := db.GetTemplateVersionByID(ctx, tmpl.ActiveVersionID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return false, nil + } + return false, err + } + return version.HasExternalAgent.Valid && version.HasExternalAgent.Bool, nil +} + +// externalAgentReadyError returns the external-agent-specific error +// message when agent belongs to an external resource, or the empty +// string otherwise. Errors looking up the resource are treated as +// non-external so the caller falls back to the dial error. +func externalAgentReadyError( + ctx context.Context, + db database.Store, + agent database.WorkspaceAgent, +) string { + isExternal, err := IsExternalWorkspaceAgent(ctx, db, agent) + if err != nil || !isExternal { + return "" + } + return ExternalAgentUnavailableMessage(agent) +} + // 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, db database.Store, - agentID uuid.UUID, + agent database.WorkspaceAgent, agentConnFn AgentConnFunc, ) map[string]any { result := map[string]any{} + agentID := agent.ID // Phase 1: retry connecting to the agent. if agentConnFn != nil { @@ -601,7 +644,16 @@ func waitForAgentReady( select { case <-agentCtx.Done(): result["agent_status"] = "not_ready" - result["agent_error"] = lastErr.Error() + // External agents may need user action on a different + // host. Surface that guidance instead of the raw dial + // error after the retry window has elapsed. The retry + // loop itself is unchanged, so a Connecting external + // agent still gets the full window to come online. + if msg := externalAgentReadyError(ctx, db, agent); msg != "" { + result["agent_error"] = msg + } else { + result["agent_error"] = lastErr.Error() + } return result case <-ticker.C: } diff --git a/coderd/x/chatd/chattool/createworkspace_test.go b/coderd/x/chatd/chattool/createworkspace_test.go index 7557bf4c01..224ed3c753 100644 --- a/coderd/x/chatd/chattool/createworkspace_test.go +++ b/coderd/x/chatd/chattool/createworkspace_test.go @@ -25,13 +25,22 @@ import ( "github.com/coder/coder/v2/codersdk/workspacesdk" ) +func newCreateWorkspaceMockStore(ctrl *gomock.Controller) *dbmock.MockStore { + db := dbmock.NewMockStore(ctrl) + db.EXPECT(). + GetTemplateVersionByID(gomock.Any(), gomock.Any()). + Return(database.TemplateVersion{}, sql.ErrNoRows). + AnyTimes() + return db +} + func TestWaitForAgentReady(t *testing.T) { t.Parallel() t.Run("AgentConnectsAndLifecycleReady", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) agentID := uuid.New() // Mock returns Ready lifecycle state. @@ -46,14 +55,14 @@ func TestWaitForAgentReady(t *testing.T) { return nil, func() {}, nil } - result := waitForAgentReady(context.Background(), db, agentID, connFn) + result := waitForAgentReady(context.Background(), db, database.WorkspaceAgent{ID: agentID}, connFn) require.Empty(t, result) }) t.Run("AgentConnectTimeout", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) agentID := uuid.New() // AgentConnFn always fails - context will timeout. @@ -65,15 +74,90 @@ func TestWaitForAgentReady(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - result := waitForAgentReady(ctx, db, agentID, connFn) + result := waitForAgentReady(ctx, db, database.WorkspaceAgent{ID: agentID}, connFn) require.Equal(t, "not_ready", result["agent_status"]) require.NotEmpty(t, result["agent_error"]) }) + t.Run("ExternalAgentTimeoutMessage", func(t *testing.T) { + // External agent retry loop should still run for the full + // window. When it eventually times out, the error message + // should be the external-agent-specific guidance, not the + // raw dial error. + t.Parallel() + ctrl := gomock.NewController(t) + db := newCreateWorkspaceMockStore(ctrl) + agentID := uuid.New() + resourceID := uuid.New() + agent := database.WorkspaceAgent{ + ID: agentID, + ResourceID: resourceID, + } + + db.EXPECT(). + GetWorkspaceResourceByID(gomock.Any(), resourceID). + Return(database.WorkspaceResource{ + ID: resourceID, + Type: ExternalAgentResourceType, + }, nil) + + attempts := 0 + connFn := func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { + attempts++ + require.Equal(t, agentID, id) + return nil, nil, context.DeadlineExceeded + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result := waitForAgentReady(ctx, db, agent, connFn) + require.GreaterOrEqual(t, attempts, 1) + require.Equal(t, "not_ready", result["agent_status"]) + require.Equal(t, ExternalAgentUnavailableMessage(agent), result["agent_error"]) + }) + + t.Run("ExternalAgentEventuallyConnects", func(t *testing.T) { + // External agent that fails the first dial but succeeds on + // the second attempt must not be short-circuited; the user + // may have just started the agent on their host. + t.Parallel() + ctrl := gomock.NewController(t) + db := newCreateWorkspaceMockStore(ctrl) + agentID := uuid.New() + resourceID := uuid.New() + agent := database.WorkspaceAgent{ + ID: agentID, + ResourceID: resourceID, + } + + // Mock returns Ready lifecycle so phase 2 exits cleanly. + db.EXPECT(). + GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID). + Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + }, nil) + + attempts := 0 + connFn := func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { + attempts++ + require.Equal(t, agentID, id) + if attempts == 1 { + return nil, nil, context.DeadlineExceeded + } + return nil, func() {}, nil + } + + result := waitForAgentReady(context.Background(), db, agent, connFn) + require.Equal(t, 2, attempts, "second attempt must run for Connecting external agents") + require.NotContains(t, result, "agent_status", "successful late connect must not surface not_ready") + require.NotContains(t, result, "agent_error") + }) + t.Run("AgentConnectsButStartupFails", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) agentID := uuid.New() // Mock returns StartError lifecycle state. @@ -87,7 +171,7 @@ func TestWaitForAgentReady(t *testing.T) { return nil, func() {}, nil } - result := waitForAgentReady(context.Background(), db, agentID, connFn) + result := waitForAgentReady(context.Background(), db, database.WorkspaceAgent{ID: agentID}, connFn) require.Equal(t, "startup_scripts_failed", result["startup_scripts"]) require.Equal(t, "start_error", result["lifecycle_state"]) }) @@ -95,7 +179,7 @@ func TestWaitForAgentReady(t *testing.T) { t.Run("NilAgentConnFn", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) agentID := uuid.New() // Mock returns Ready lifecycle state. @@ -105,16 +189,31 @@ func TestWaitForAgentReady(t *testing.T) { LifecycleState: database.WorkspaceAgentLifecycleStateReady, }, nil) - result := waitForAgentReady(context.Background(), db, agentID, nil) + result := waitForAgentReady(context.Background(), db, database.WorkspaceAgent{ID: 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, nil, ctx.Err() + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result := waitForAgentReady(ctx, nil, database.WorkspaceAgent{ID: uuid.New()}, connFn) + require.Equal(t, "not_ready", result["agent_status"]) + require.NotEmpty(t, result["agent_error"]) + }) } func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -223,7 +322,7 @@ func TestCreateWorkspace_ReturnsSelectionErrorImmediately(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -327,7 +426,7 @@ func TestCreateWorkspace_PostCreationBuildFailure(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -426,7 +525,7 @@ func TestCreateWorkspace_PostCreationQuotaFailure(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -673,7 +772,7 @@ func TestCreateWorkspace_ResponderErrorPreservesStructuredFields(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -779,7 +878,7 @@ func TestCreateWorkspace_GlobalTTL(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -883,7 +982,7 @@ func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() chatOrgID := uuid.New() @@ -940,10 +1039,73 @@ func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) { require.Contains(t, resp.Content, "organization") } +func TestCreateWorkspace_BlocksExternalTemplate(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + ownerID := uuid.New() + orgID := uuid.New() + chatID := uuid.New() + templateID := uuid.New() + activeVersionID := uuid.New() + + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). + GetAuthorizationUserRoles(gomock.Any(), ownerID). + Return(database.GetAuthorizationUserRolesRow{ + ID: ownerID, + Roles: []string{}, + Groups: []string{}, + Status: database.UserStatusActive, + }, nil) + db.EXPECT(). + GetTemplateByID(gomock.Any(), templateID). + Return(database.Template{ + ID: templateID, + OrganizationID: orgID, + ActiveVersionID: activeVersionID, + }, nil) + db.EXPECT(). + GetTemplateVersionByID(gomock.Any(), activeVersionID). + Return(database.TemplateVersion{ + ID: activeVersionID, + HasExternalAgent: sql.NullBool{ + Bool: true, + Valid: true, + }, + }, nil) + + createCalled := false + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ + OwnerID: ownerID, + CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { + createCalled = true + return codersdk.Workspace{}, nil + }, + WorkspaceMu: &sync.Mutex{}, + Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + }) + + input := fmt.Sprintf(`{"template_id":%q}`, templateID.String()) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{ + ID: "call-1", + Name: "create_workspace", + Input: input, + }) + require.NoError(t, err) + require.True(t, resp.IsError) + require.False(t, createCalled, "CreateFn must not be called for external template") + require.Equal(t, createWorkspaceExternalAgentMessage, resp.Content) +} + func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -993,7 +1155,7 @@ func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) { func TestCheckExistingWorkspace_InProgressBuildReturnsBuildID(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1088,7 +1250,7 @@ func TestCheckExistingWorkspace_InProgressBuildReturnsBuildID(t *testing.T) { func TestCheckExistingWorkspace_InProgressBuildFailureReturnsBuildID(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1169,7 +1331,7 @@ func TestCheckExistingWorkspace_InProgressBuildFailureReturnsBuildID(t *testing. func TestCheckExistingWorkspace_ConnectingAgentWaits(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1248,7 +1410,7 @@ func TestCheckExistingWorkspace_DeadAgentAllowsCreation(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1281,7 +1443,7 @@ func TestWaitForBuild_CanceledJob(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() @@ -1374,7 +1536,7 @@ func TestWaitForBuild_CanceledJob(t *testing.T) { func TestCheckExistingWorkspace_StoppedWorkspace(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1402,7 +1564,7 @@ func TestCheckExistingWorkspace_StoppedWorkspace(t *testing.T) { func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() @@ -1497,7 +1659,7 @@ func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() templateID := uuid.New() @@ -1645,7 +1807,7 @@ func setupCreateWorkspacePresetTest(t *testing.T) createWorkspacePresetTestSetup t.Helper() ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) + db := newCreateWorkspaceMockStore(ctrl) s := createWorkspacePresetTestSetup{ DB: db, diff --git a/coderd/x/chatd/chattool/external_agents.go b/coderd/x/chatd/chattool/external_agents.go new file mode 100644 index 0000000000..20ed1a2d87 --- /dev/null +++ b/coderd/x/chatd/chattool/external_agents.go @@ -0,0 +1,47 @@ +package chattool + +import ( + "context" + + "github.com/google/uuid" + + "github.com/coder/coder/v2/coderd/database" +) + +// ExternalAgentResourceType is the Terraform resource type for externally +// managed agents. +const ExternalAgentResourceType = "coder_external_agent" + +const createWorkspaceExternalAgentMessage = "create_workspace cannot create workspaces from templates with externally managed agents. " + + "Use list_templates to choose a different template, or if the user wants " + + "to use an external workspace, they should create it and start it up fully " + + "themselves first, then attach it to this chat" + +const externalAgentNotConnectedMessage = "workspace uses an externally managed agent that has not connected yet. " + + "The user needs to start the workspace externally and make sure the " + + "external agent is connected, then try again" + +const externalAgentDisconnectedMessage = "workspace uses an externally managed agent that is currently offline. " + + "The user needs to reconnect the external agent on its host, then try again" + +// ExternalAgentUnavailableMessage explains how to make an externally managed +// agent usable based on its connection history. +func ExternalAgentUnavailableMessage(agent database.WorkspaceAgent) string { + if agent.FirstConnectedAt.Valid { + return externalAgentDisconnectedMessage + } + return externalAgentNotConnectedMessage +} + +// IsExternalWorkspaceAgent reports whether agent belongs to an external +// resource. +func IsExternalWorkspaceAgent(ctx context.Context, db database.Store, agent database.WorkspaceAgent) (bool, error) { + if db == nil || agent.ResourceID == uuid.Nil { + return false, nil + } + resource, err := db.GetWorkspaceResourceByID(ctx, agent.ResourceID) + if err != nil { + return false, err + } + return resource.Type == ExternalAgentResourceType, nil +} diff --git a/coderd/x/chatd/chattool/startworkspace.go b/coderd/x/chatd/chattool/startworkspace.go index af388c2f2b..16d1d1f9be 100644 --- a/coderd/x/chatd/chattool/startworkspace.go +++ b/coderd/x/chatd/chattool/startworkspace.go @@ -307,7 +307,7 @@ func waitForAgentAndRespond( } setBuildID(result, buildID) setNoBuild(result, buildID) - for k, v := range waitForAgentReady(ctx, db, selected.ID, agentConnFn) { + for k, v := range waitForAgentReady(ctx, db, selected, agentConnFn) { result[k] = v } return result