package chattool import ( "context" "database/sql" "encoding/json" "fmt" "sync" "testing" "time" "charm.land/fantasy" "github.com/google/uuid" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/xerrors" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/httpapi/httperror" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" "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 TestCreateWorkspaceDescriptionDelaysWorkspaceCreation(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) tool := CreateWorkspace(db, uuid.New(), uuid.New(), CreateWorkspaceOptions{}) info := tool.Info() require.Contains(t, info.Description, "Create a new workspace from a template only when workspace-backed") require.Contains(t, info.Description, "user explicitly asks") require.Contains(t, info.Description, "Do not use this as a default first step") } func TestWaitForAgentReady(t *testing.T) { t.Parallel() t.Run("AgentConnectsAndLifecycleReady", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(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, database.WorkspaceAgent{ID: agentID}, connFn) require.Empty(t, result) }) t.Run("AgentConnectTimeout", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(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, 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 := newCreateWorkspaceMockStore(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, database.WorkspaceAgent{ID: 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 := newCreateWorkspaceMockStore(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, 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 := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() fallbackAgentID := uuid.New() chatAgentID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ID: chatID}, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusSucceeded, }, nil) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{ {ID: fallbackAgentID, Name: "dev", DisplayOrder: 0}, {ID: chatAgentID, Name: "dev-coderd-chat", DisplayOrder: 1}, }, nil) db.EXPECT(). GetWorkspaceAgentLifecycleStateByID(gomock.Any(), chatAgentID). Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ LifecycleState: database.WorkspaceAgentLifecycleStateReady, }, nil) var connectedAgentID uuid.UUID createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } agentConnFn := func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { connectedAgentID = agentID return nil, func() {}, nil } tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-chat-agent"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.NotEmpty(t, resp.Content) require.Equal(t, chatAgentID, connectedAgentID) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, buildID.String(), result["build_id"]) } func TestCreateWorkspace_ReturnsSelectionErrorImmediately(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusSucceeded, }, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, BuildID: uuid.NullUUID{UUID: buildID, Valid: true}, AgentID: uuid.NullUUID{}, }). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{ {ID: uuid.New(), Name: "alpha-coderd-chat", DisplayOrder: 0}, {ID: uuid.New(), Name: "beta-coderd-chat", DisplayOrder: 1}, }, nil) tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil }, AgentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatal("AgentConnFn should not be called when agent selection fails") return nil, nil, xerrors.New("unexpected agent dial") }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-selection-error"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, true, result["created"]) require.Equal(t, "testuser/test-selection-error", result["workspace_name"]) require.Equal(t, "selection_error", result["agent_status"]) require.Contains(t, result["agent_error"], "multiple agents match the chat suffix") require.Equal(t, buildID.String(), result["build_id"]) } func TestCreateWorkspace_PostCreationBuildFailure(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ID: chatID}, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) // waitForBuild fetches the build by ID. db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) // waitForBuild polls the provisioner job. Return Failed. db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusFailed, Error: sql.NullString{String: "terraform apply failed", Valid: true}, }, nil) createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-build-fail"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Contains(t, result["error"], "workspace build failed") require.Equal(t, buildID.String(), result["build_id"]) require.NotContains(t, result, "error_code", "generic build failures must not surface a quota error_code") require.NotContains(t, result, "quota", "generic build failures must not surface quota details") require.False(t, resp.IsError, "buildToolResponse must not set IsError; chatprompt strips structured fields from error responses") } func TestCreateWorkspace_PostCreationQuotaFailure(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ID: chatID}, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusFailed, Error: sql.NullString{String: "insufficient quota", Valid: true}, ErrorCode: sql.NullString{ String: string(codersdk.InsufficientQuota), Valid: true, }, }, nil) db.EXPECT(). GetQuotaConsumedForUser(gomock.Any(), database.GetQuotaConsumedForUserParams{ OwnerID: ownerID, OrganizationID: orgID, }). Return(int64(40), nil) db.EXPECT(). GetQuotaAllowanceForUser(gomock.Any(), database.GetQuotaAllowanceForUserParams{ UserID: ownerID, OrganizationID: orgID, }). Return(int64(40), nil) createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-quota-fail"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, string(codersdk.InsufficientQuota), result["error_code"]) require.Equal(t, "Workspace quota reached", result["title"]) require.Contains(t, result["error"], "workspace build failed") require.Contains(t, result["message"], "workspace quota is full") require.Contains(t, result["message"], "Delete a workspace") require.Contains(t, result["message"], "raise your group quota allowance") require.NotContains(t, result, "next_steps") require.Equal(t, buildID.String(), result["build_id"]) quota, ok := result["quota"].(map[string]any) require.True(t, ok) require.Equal(t, float64(40), quota["credits_consumed"]) require.Equal(t, float64(40), quota["budget"]) require.False(t, resp.IsError, "quota responses must not set IsError; chatprompt strips structured fields from error responses") } func TestCreateWorkspace_ExistingBuildQuotaFailure(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() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ ID: ownerID, Roles: []string{}, Groups: []string{}, Status: database.UserStatusActive, }, nil) db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) db.EXPECT(). GetWorkspaceByID(gomock.Any(), workspaceID). Return(database.Workspace{ ID: workspaceID, Name: "existing-quota-workspace", OrganizationID: orgID, }, nil) db.EXPECT(). GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) firstJob := db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusRunning, }, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, BuildID: uuid.NullUUID{UUID: buildID, Valid: true}, AgentID: uuid.NullUUID{}, }). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusFailed, Error: sql.NullString{String: "insufficient quota", Valid: true}, ErrorCode: sql.NullString{ String: string(codersdk.InsufficientQuota), Valid: true, }, }, nil). After(firstJob) ownerCtx := ownerContextMatcher{ownerID: ownerID} db.EXPECT(). GetQuotaConsumedForUser(ownerCtx, database.GetQuotaConsumedForUserParams{ OwnerID: ownerID, OrganizationID: orgID, }). Return(int64(40), nil) db.EXPECT(). GetQuotaAllowanceForUser(ownerCtx, database.GetQuotaAllowanceForUserParams{ UserID: ownerID, OrganizationID: orgID, }). Return(int64(40), nil) tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { t.Fatal("CreateFn should not be called when an existing build is in progress") return codersdk.Workspace{}, nil }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-existing-quota-fail"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, string(codersdk.InsufficientQuota), result["error_code"]) require.Equal(t, "Workspace quota reached", result["title"]) require.Contains(t, result["error"], "existing workspace build failed") require.Contains(t, result["message"], "could not start this workspace") require.Contains(t, result["message"], "workspace quota is full") require.Equal(t, buildID.String(), result["build_id"]) quota, ok := result["quota"].(map[string]any) require.True(t, ok) require.Equal(t, float64(40), quota["credits_consumed"]) require.Equal(t, float64(40), quota["budget"]) require.False(t, resp.IsError) } func TestCreateWorkspace_ResponderErrorPreservesStructuredFields(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{}, httperror.NewResponseError(400, codersdk.Response{ Message: "missing required parameter", Detail: "region must be set before the workspace can start", Validations: []codersdk.ValidationError{{ Field: "region", Detail: "region must be set before the workspace can start", }}, }) }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-structured-error"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.False(t, resp.IsError) var result struct { Error string `json:"error"` Detail string `json:"detail"` Validations []codersdk.ValidationError `json:"validations"` } require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, "missing required parameter", result.Error) require.Equal(t, "region must be set before the workspace can start", result.Detail) require.Equal(t, []codersdk.ValidationError{{ Field: "region", Detail: "region must be set before the workspace can start", }}, result.Validations) } func TestCreateWorkspace_GlobalTTL(t *testing.T) { t.Parallel() tests := []struct { name string ttlReturn string ttlErr error wantTTLMs *int64 }{ { name: "PositiveTTL", ttlReturn: "2h", wantTTLMs: ptr.Ref(int64(2 * time.Hour / time.Millisecond)), }, { name: "ZeroTTLUsesTemplateDefault", ttlReturn: "0s", wantTTLMs: nil, }, { name: "DBError_FallsBackToNil", ttlReturn: "", ttlErr: xerrors.New("db error"), wantTTLMs: nil, }, { name: "InvalidStoredValue_FallsBackToNil", ttlReturn: "not-a-duration", wantTTLMs: nil, }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ID: chatID}, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return(tc.ttlReturn, tc.ttlErr) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusSucceeded, }, nil) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil) var capturedReq codersdk.CreateWorkspaceRequest createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { capturedReq = req return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-ws-%s"}`, templateID.String(), tc.name) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.NotEmpty(t, resp.Content) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, buildID.String(), result["build_id"]) if tc.wantTTLMs != nil { require.NotNil(t, capturedReq.TTLMillis) require.Equal(t, *tc.wantTTLMs, *capturedReq.TTLMillis) } else { require.Nil(t, capturedReq.TTLMillis) } }) } } func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() chatOrgID := uuid.New() templateOrgID := uuid.New() // Different org. templateID := uuid.New() chatID := uuid.New() // Chat exists but has no workspace binding. db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{}, }, 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: templateOrgID, Name: "wrong-org-template", }, nil) createCalled := false tool := CreateWorkspace(db, chatOrgID, 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.False(t, createCalled, "CreateFn must not be called for cross-org template") 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 := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() agentID := uuid.New() now := time.Now().UTC() expectExistingWorkspaceLookup( db, chatID, workspaceID, jobID, "existing-workspace", database.ProvisionerJobStatusSucceeded, database.WorkspaceTransitionStart, ) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{{ ID: agentID, Name: "dev", CreatedAt: now.Add(-time.Minute), FirstConnectedAt: validNullTime(now.Add(-45 * time.Second)), LastConnectedAt: validNullTime(now.Add(-5 * time.Second)), }}, nil) db.EXPECT(). GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID). Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ LifecycleState: database.WorkspaceAgentLifecycleStateReady, }, nil) connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatalf("unexpected agent dial for connected workspace") return nil, nil, xerrors.New("unexpected agent dial") } options := testCheckExistingWorkspaceOptions(connFn) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) require.Equal(t, "already_exists", check.Result["status"]) require.Equal(t, "existing-workspace", check.Result["workspace_name"]) require.Equal(t, "workspace is already running and recently connected", check.Result["message"]) } func TestCheckExistingWorkspace_InProgressBuildReturnsBuildID(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() // GetChatByID returns a chat linked to a workspace. db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) // GetWorkspaceByID returns a non-deleted workspace. db.EXPECT(). GetWorkspaceByID(gomock.Any(), workspaceID). Return(database.Workspace{ ID: workspaceID, Name: "building-workspace", }, nil) // GetLatestWorkspaceBuildByWorkspaceID is called once in // checkExistingWorkspace. waitForBuild now uses // GetWorkspaceBuildByID to track the specific build. db.EXPECT(). GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) // First GetProvisionerJobByID (in checkExistingWorkspace) returns // Running, triggering waitForBuild. The second call (waitForBuild's // first poll) returns Succeeded so the loop exits immediately. firstJob := db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusRunning, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusSucceeded, }, nil). After(firstJob) // The in-progress path now publishes the build ID before // waitForBuild. db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, BuildID: uuid.NullUUID{UUID: buildID, Valid: true}, AgentID: uuid.NullUUID{}, }). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) // After waitForBuild completes, checkExistingWorkspace fetches // agents. Return empty to keep the test focused on build_id. db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil) options := testCheckExistingWorkspaceOptions(nil) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) require.Equal(t, false, check.Result["created"]) require.Equal(t, "already_exists", check.Result["status"]) require.Equal(t, buildID.String(), check.Result["build_id"]) require.Equal(t, "building-workspace", check.Result["workspace_name"]) require.Equal(t, "workspace build completed", check.Result["message"]) } func TestCheckExistingWorkspace_InProgressBuildFailureReturnsBuildID(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) db.EXPECT(). GetWorkspaceByID(gomock.Any(), workspaceID). Return(database.Workspace{ ID: workspaceID, Name: "failing-workspace", }, nil) db.EXPECT(). GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, Transition: database.WorkspaceTransitionStart, }, nil) // First call returns Running (triggers waitForBuild), second // returns Failed so waitForBuild returns an error. firstJob := db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusRunning, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusFailed, }, nil). After(firstJob) // The in-progress path publishes the build ID before // waitForBuild. db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, BuildID: uuid.NullUUID{UUID: buildID, Valid: true}, AgentID: uuid.NullUUID{}, }). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) options := testCheckExistingWorkspaceOptions(nil) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.Error(t, check.BuildErr) require.Contains(t, check.BuildErr.Error(), "existing workspace build failed") require.Equal(t, buildID, check.BuildID) require.Equal(t, buildFailureActionStart, check.BuildAction) require.NoError(t, check.Err) } func TestCheckExistingWorkspace_ConnectingAgentWaits(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() agentID := uuid.New() now := time.Now().UTC() connectCalls := 0 expectExistingWorkspaceLookup( db, chatID, workspaceID, jobID, "existing-workspace", database.ProvisionerJobStatusSucceeded, database.WorkspaceTransitionStart, ) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{{ ID: agentID, Name: "dev", CreatedAt: now, ConnectionTimeoutSeconds: 60, }}, nil) db.EXPECT(). GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID). Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ LifecycleState: database.WorkspaceAgentLifecycleStateReady, }, nil) connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { connectCalls++ return nil, func() {}, nil } options := testCheckExistingWorkspaceOptions(connFn) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) require.Equal(t, 1, connectCalls) require.Equal(t, "already_exists", check.Result["status"]) require.Equal(t, "workspace exists and the agent is still connecting", check.Result["message"]) } func TestCheckExistingWorkspace_DeadAgentAllowsCreation(t *testing.T) { t.Parallel() tests := []struct { name string agent database.WorkspaceAgent }{ { name: "Disconnected", agent: database.WorkspaceAgent{ ID: uuid.New(), Name: "disconnected", CreatedAt: time.Now().UTC().Add(-2 * time.Minute), FirstConnectedAt: validNullTime(time.Now().UTC().Add(-2 * time.Minute)), LastConnectedAt: validNullTime(time.Now().UTC().Add(-time.Minute)), }, }, { name: "TimedOut", agent: database.WorkspaceAgent{ ID: uuid.New(), Name: "timed-out", CreatedAt: time.Now().UTC().Add(-2 * time.Second), ConnectionTimeoutSeconds: 1, }, }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() expectExistingWorkspaceLookup( db, chatID, workspaceID, jobID, "existing-workspace", database.ProvisionerJobStatusSucceeded, database.WorkspaceTransitionStart, ) db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{tc.agent}, nil) options := testCheckExistingWorkspaceOptions(nil) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.False(t, check.Done) require.Nil(t, check.Result) }) } } func TestWaitForBuild_CanceledJob(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() orgID := uuid.New() chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ID: chatID}, nil) db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). 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, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) // waitForBuild fetches the build by ID. db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) // waitForBuild polls the provisioner job. Return Canceled. db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusCanceled, }, nil) createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-build-cancel"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Contains(t, result["error"], "build was canceled") require.Equal(t, buildID.String(), result["build_id"]) require.False(t, resp.IsError, "buildToolResponse must not set IsError; chatprompt strips structured fields from error responses") } func TestCheckExistingWorkspace_StoppedWorkspace(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() expectExistingWorkspaceLookup( db, chatID, workspaceID, jobID, "stopped-workspace", database.ProvisionerJobStatusSucceeded, database.WorkspaceTransitionStop, ) options := testCheckExistingWorkspaceOptions(nil) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.True(t, check.Done) require.NoError(t, check.Err) require.Equal(t, "stopped", check.Result["status"]) require.Contains(t, check.Result["message"], "start_workspace") } func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) chatID := uuid.New() workspaceID := uuid.New() // Mock GetChatByID returns a chat linked to a workspace. db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) // Mock GetWorkspaceByID returns a soft-deleted workspace. db.EXPECT(). GetWorkspaceByID(gomock.Any(), workspaceID). Return(database.Workspace{ ID: workspaceID, Deleted: true, }, nil) options := testCheckExistingWorkspaceOptions(nil) check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.False(t, check.Done, "should allow creation for deleted workspace") require.Nil(t, check.Result) } func testCheckExistingWorkspaceOptions( agentConnFn AgentConnFunc, ) CreateWorkspaceOptions { return CreateWorkspaceOptions{ AgentConnFn: agentConnFn, AgentInactiveDisconnectTimeout: 30 * time.Second, } } type ownerContextMatcher struct { ownerID uuid.UUID } func (m ownerContextMatcher) Matches(v any) bool { ctx, ok := v.(context.Context) if !ok { return false } actor, ok := dbauthz.ActorFromContext(ctx) return ok && actor.ID == m.ownerID.String() } func (ownerContextMatcher) String() string { return "context with owner actor" } func expectExistingWorkspaceLookup( db *dbmock.MockStore, chatID uuid.UUID, workspaceID uuid.UUID, jobID uuid.UUID, workspaceName string, jobStatus database.ProvisionerJobStatus, transition database.WorkspaceTransition, ) { db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) db.EXPECT(). GetWorkspaceByID(gomock.Any(), workspaceID). Return(database.Workspace{ ID: workspaceID, Name: workspaceName, }, nil) db.EXPECT(). GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). Return(database.WorkspaceBuild{ WorkspaceID: workspaceID, JobID: jobID, Transition: transition, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: jobStatus, }, nil) } func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) ownerID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() chatID := uuid.New() jobID := uuid.New() buildID := uuid.New() // checkExistingWorkspace calls GetChatByID first. Return a chat // with no workspace so the tool proceeds to creation. 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) // Org check: GetTemplateByID returns a template in the // same org (uuid.Nil matches our organizationID param). db.EXPECT(). GetTemplateByID(gomock.Any(), templateID). Return(database.Template{ ID: templateID, OrganizationID: uuid.Nil, }, nil) db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) // UpdateChatWorkspaceBinding — triggers first OnChatUpdated. db.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) // waitForBuild: fetch build, then poll job as completed. db.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), buildID). Return(database.WorkspaceBuild{ ID: buildID, WorkspaceID: workspaceID, JobID: jobID, }, nil) db.EXPECT(). GetProvisionerJobByID(gomock.Any(), jobID). Return(database.ProvisionerJob{ ID: jobID, JobStatus: database.ProvisionerJobStatusSucceeded, CompletedAt: validNullTime(time.Now()), }, nil) // GetChatByID — called after waitForBuild for second OnChatUpdated. db.EXPECT(). GetChatByID(gomock.Any(), chatID). Return(database.Chat{ ID: chatID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) // Agent lookup after build completes — return empty so we skip // agent selection and waitForAgentReady. db.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil) var mu sync.Mutex var callbackChats []database.Chat createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, Name: req.Name, OwnerName: "testuser", LatestBuild: codersdk.WorkspaceBuild{ ID: buildID, }, }, nil } tool := CreateWorkspace(db, uuid.Nil, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), OnChatUpdated: func(chat database.Chat) { mu.Lock() callbackChats = append(callbackChats, chat) mu.Unlock() }, }) input := fmt.Sprintf(`{"template_id":%q,"name":"test-callback"}`, templateID.String()) resp, err := tool.Run(context.Background(), fantasy.ToolCall{ ID: "call-1", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.False(t, resp.IsError) mu.Lock() defer mu.Unlock() require.Len(t, callbackChats, 2, "OnChatUpdated should fire twice: once on binding, once after build completes") // Both callbacks should carry the workspace ID. for i, chat := range callbackChats { require.True(t, chat.WorkspaceID.Valid, "callback %d should have workspace ID", i) require.Equal(t, workspaceID, chat.WorkspaceID.UUID) } } func validNullTime(t time.Time) sql.NullTime { return sql.NullTime{Time: t, Valid: true} } // createWorkspacePresetTestSetup holds common test dependencies // for create_workspace preset tests. type createWorkspacePresetTestSetup struct { DB *dbmock.MockStore OwnerID uuid.UUID OrgID uuid.UUID TemplateID uuid.UUID ChatID uuid.UUID WorkspaceID uuid.UUID BuildID uuid.UUID AgentID uuid.UUID } // setupCreateWorkspacePresetTest creates common mock expectations // for preset-related create_workspace tests. It sets up RBAC, // template lookup, TTL, and chat lookup. func setupCreateWorkspacePresetTest(t *testing.T) createWorkspacePresetTestSetup { t.Helper() ctrl := gomock.NewController(t) db := newCreateWorkspaceMockStore(ctrl) s := createWorkspacePresetTestSetup{ DB: db, OwnerID: uuid.New(), OrgID: uuid.New(), TemplateID: uuid.New(), ChatID: uuid.New(), WorkspaceID: uuid.New(), BuildID: uuid.New(), AgentID: uuid.New(), } // RBAC. db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), s.OwnerID). Return(database.GetAuthorizationUserRolesRow{ ID: s.OwnerID, Username: "testuser", Status: "active", }, nil) // Template lookup. db.EXPECT(). GetTemplateByID(gomock.Any(), s.TemplateID). Return(database.Template{ ID: s.TemplateID, OrganizationID: s.OrgID, Name: "test-template", ActiveVersionID: uuid.New(), }, nil) // Chat workspace TTL. db.EXPECT(). GetChatWorkspaceTTL(gomock.Any()). Return("", sql.ErrNoRows) // Check for existing workspace (no existing). db.EXPECT(). GetChatByID(gomock.Any(), s.ChatID). Return(database.Chat{ID: s.ChatID}, nil) return s } // expectSuccessfulBuild adds mock expectations for a successful // build, agent lookup, and agent lifecycle check. func (s createWorkspacePresetTestSetup) expectSuccessfulBuild() { s.DB.EXPECT(). UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). Return(database.Chat{ID: s.ChatID}, nil) s.DB.EXPECT(). GetWorkspaceBuildByID(gomock.Any(), s.BuildID). Return(database.WorkspaceBuild{ ID: s.BuildID, JobID: uuid.New(), }, nil) s.DB.EXPECT(). GetProvisionerJobByID(gomock.Any(), gomock.Any()). Return(database.ProvisionerJob{ JobStatus: database.ProvisionerJobStatusSucceeded, }, nil) s.DB.EXPECT(). GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), s.WorkspaceID). Return([]database.WorkspaceAgent{{ ID: s.AgentID, Name: "main", }}, nil) s.DB.EXPECT(). GetWorkspaceAgentLifecycleStateByID(gomock.Any(), s.AgentID). Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ LifecycleState: database.WorkspaceAgentLifecycleStateReady, }, nil) } func TestCreateWorkspace_WithPresetID(t *testing.T) { t.Parallel() s := setupCreateWorkspacePresetTest(t) s.expectSuccessfulBuild() presetID := uuid.New() var capturedReq codersdk.CreateWorkspaceRequest createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { capturedReq = req return codersdk.Workspace{ ID: s.WorkspaceID, Name: req.Name, LatestBuild: codersdk.WorkspaceBuild{ ID: s.BuildID, }, }, nil } agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil } tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, CreateFn: createFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf( `{"template_id":%q,"preset_id":%q,"name":"test-ws"}`, s.TemplateID.String(), presetID.String(), ) ctx := context.Background() resp, err := tool.Run(ctx, fantasy.ToolCall{ ID: "call-preset", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.False(t, resp.IsError, "unexpected error: %s", resp.Content) require.Equal(t, presetID, capturedReq.TemplateVersionPresetID, "expected preset ID to be set on CreateWorkspaceRequest") } func TestCreateWorkspace_InvalidPresetID(t *testing.T) { t.Parallel() s := setupCreateWorkspacePresetTest(t) tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { t.Fatal("CreateFn should not be called with invalid preset_id") return codersdk.Workspace{}, nil }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf( `{"template_id":%q,"preset_id":"not-a-uuid","name":"test-ws"}`, s.TemplateID.String(), ) ctx := context.Background() resp, err := tool.Run(ctx, fantasy.ToolCall{ ID: "call-bad-preset", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.True(t, resp.IsError) require.Contains(t, resp.Content, "invalid preset_id") } func TestCreateWorkspace_WithPresetAndParams(t *testing.T) { t.Parallel() s := setupCreateWorkspacePresetTest(t) s.expectSuccessfulBuild() presetID := uuid.New() var capturedReq codersdk.CreateWorkspaceRequest createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { capturedReq = req return codersdk.Workspace{ ID: s.WorkspaceID, Name: req.Name, LatestBuild: codersdk.WorkspaceBuild{ ID: s.BuildID, }, }, nil } agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil } tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, CreateFn: createFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) input := fmt.Sprintf( `{"template_id":%q,"preset_id":%q,"name":"test-ws","parameters":{"region":"us-east"}}`, s.TemplateID.String(), presetID.String(), ) ctx := context.Background() resp, err := tool.Run(ctx, fantasy.ToolCall{ ID: "call-preset-params", Name: "create_workspace", Input: input, }) require.NoError(t, err) require.False(t, resp.IsError, "unexpected error: %s", resp.Content) // Verify preset ID is set. require.Equal(t, presetID, capturedReq.TemplateVersionPresetID, "expected preset ID to be set") // Verify parameters are also populated. require.Len(t, capturedReq.RichParameterValues, 1, "expected rich parameter values to be set") require.Equal(t, "region", capturedReq.RichParameterValues[0].Name) require.Equal(t, "us-east", capturedReq.RichParameterValues[0].Value) }