package chattool_test import ( "context" "database/sql" "encoding/json" "sync" "sync/atomic" "testing" "time" "charm.land/fantasy" "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" "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/dbfake" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/httpapi/httperror" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" sdkproto "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" ) func TestStartWorkspace(t *testing.T) { t.Parallel() t.Run("NoWorkspace", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "test-no-workspace", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.Contains(t, resp.Content, "no workspace") }) t.Run("AlreadyRunning", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-already-running", }) agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil } tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) require.Nil(t, result["build_id"], "build_id should not be present when workspace was already running") require.Equal(t, true, result["no_build"], "no_build should be true when workspace was already running") }) t.Run("AlreadyRunningPrefersChatSuffixAgent", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).WithAgent(func(agents []*sdkproto.Agent) []*sdkproto.Agent { agents[0].Name = "dev" return append(agents, &sdkproto.Agent{ Id: uuid.NewString(), Name: "dev-coderd-chat", Auth: &sdkproto.Agent_Token{Token: uuid.NewString()}, Env: map[string]string{}, }) }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, }).Do() ws := wsResp.Workspace now := time.Now().UTC() preferredAgentID := uuid.Nil for _, agent := range wsResp.Agents { if agent.Name == "dev-coderd-chat" { preferredAgentID = agent.ID } err := db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ ID: agent.ID, LifecycleState: database.WorkspaceAgentLifecycleStateReady, StartedAt: sql.NullTime{Time: now, Valid: true}, ReadyAt: sql.NullTime{Time: now, Valid: true}, }) require.NoError(t, err) } require.NotEqual(t, uuid.Nil, preferredAgentID) chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-preferred-agent", }) var connectedAgentID uuid.UUID agentConnFn := func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { connectedAgentID = agentID return nil, func() {}, nil } tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.Equal(t, preferredAgentID, connectedAgentID) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) }) t.Run("AlreadyRunningWithoutAgentsReturnsNoAgent", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).WithAgent(func(_ []*sdkproto.Agent) []*sdkproto.Agent { return nil }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-no-agent", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatal("AgentConnFn should not be called when no agents exist") return nil, func() {}, nil }, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) require.Equal(t, "no_agent", result["agent_status"]) }) t.Run("AlreadyRunningPreservesAgentSelectionError", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).WithAgent(func(agents []*sdkproto.Agent) []*sdkproto.Agent { agents[0].Name = "alpha-coderd-chat" return append(agents, &sdkproto.Agent{ Id: uuid.NewString(), Name: "beta-coderd-chat", Auth: &sdkproto.Agent_Token{Token: uuid.NewString()}, Env: map[string]string{}, }) }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-selection-error", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatal("AgentConnFn should not be called when agent selection fails") return nil, func() {}, nil }, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) require.Equal(t, "selection_error", result["agent_status"]) require.Contains(t, result["agent_error"], "multiple agents match the chat suffix") }) t.Run("StoppedWorkspace", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) // Create a completed "stop" build so the workspace is stopped. wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-stopped-workspace", }) var startCalled bool var startBuildID uuid.UUID startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { startCalled = true require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition) require.Equal(t, ws.ID, wsID) require.Empty(t, req.RichParameterValues, "no parameters should be forwarded for bare start") // Simulate start by inserting a new completed "start" build. buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, BuildNumber: 2, }).Do() startBuildID = buildResp.Build.ID return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil } agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil } tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: startFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.True(t, startCalled, "expected StartFn to be called") var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) require.Equal(t, startBuildID.String(), result["build_id"]) require.Nil(t, result["no_build"], "no_build should not be set when a build was triggered") }) t.Run("StoppedWorkspaceReportsAutoUpdate", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-stopped-workspace-auto-update", }) startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition) require.Equal(t, ws.ID, wsID) buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, BuildNumber: 2, }).Do() return codersdk.WorkspaceBuild{ ID: buildResp.Build.ID, TemplateVersionID: uuid.New(), }, nil } tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, true, result["updated_to_active_version"]) require.Equal(t, "template requires active versions", result["update_reason"]) require.Contains(t, result["message"], "updated to the active template version") }) t.Run("PassesParameters", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-passes-parameters", }) expectedParams := []codersdk.WorkspaceBuildParameter{ {Name: "region", Value: "us-east-1"}, {Name: "size", Value: "large"}, } startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition) require.Equal(t, ws.ID, wsID) require.ElementsMatch(t, expectedParams, req.RichParameterValues) buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, BuildNumber: 2, }).Do() return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil } tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: `{"parameters":{"region":"us-east-1","size":"large"}}`}) require.NoError(t, err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, true, result["started"]) }) t.Run("ManualUpdateRequired", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-manual-update-required", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { return codersdk.WorkspaceBuild{}, httperror.NewResponseError(400, codersdk.Response{ Message: "The workspace needs the template's active version before it can start. Use read_template with this workspace's template_id to inspect the active version's required parameters, then retry start_workspace with a parameters object that supplies any missing or changed values.", 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{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.False(t, resp.IsError) require.NotContains(t, resp.Content, "start workspace:") var result struct { Error string `json:"error"` Detail string `json:"detail"` TemplateID string `json:"template_id"` Validations []codersdk.ValidationError `json:"validations"` } require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Contains(t, result.Error, "read_template") require.Contains(t, result.Error, "retry start_workspace") require.Equal(t, ws.TemplateID.String(), result.TemplateID) 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) }) t.Run("ResponderErrorWithoutValidationsOmitsTemplateID", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-responder-error-without-validations", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { return codersdk.WorkspaceBuild{}, httperror.NewResponseError(502, codersdk.Response{ Message: "workspace start failed", Detail: "temporary provisioner outage", }) }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.False(t, resp.IsError) var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) require.Equal(t, "workspace start failed", result["error"]) require.Equal(t, "temporary provisioner outage", result["detail"]) _, hasTemplateID := result["template_id"] require.False(t, hasTemplateID) }) t.Run("InProgressBuild", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) // Create a workspace with a build that is still running. wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, }).Starting().Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-in-progress-build", }) // Wrap the DB so we know exactly when the tool reads // the job status. The interceptor signals AFTER the // first GetProvisionerJobByID read completes, so the // main goroutine can safely complete the build knowing // the tool already observed Running. jobRead := make(chan struct{}, 1) wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead} agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil } var onChatUpdatedCalled atomic.Bool tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for an in-progress build") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), OnChatUpdated: func(_ database.Chat) { onChatUpdatedCalled.Store(true) }, }) // Run tool.Run in a goroutine. It will see the job as // Running and enter waitForBuild which polls every 2s. type toolResult struct { resp fantasy.ToolResponse err error } done := make(chan toolResult, 1) go func() { resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) done <- toolResult{resp, err} }() // Wait for the tool to read the job status (Running). testutil.TryReceive(ctx, t, jobRead) // Now complete the build. The next poll in waitForBuild // will see Succeeded and return the build ID. now := time.Now().UTC() require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{ ID: wsResp.Build.JobID, UpdatedAt: now, CompletedAt: sql.NullTime{Time: now, Valid: true}, })) res := testutil.TryReceive(ctx, t, done) require.NoError(t, res.err) resp := res.resp var result map[string]any require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) started, ok := result["started"].(bool) require.True(t, ok) require.True(t, started) require.Equal(t, wsResp.Build.ID.String(), result["build_id"]) require.True(t, onChatUpdatedCalled.Load(), "OnChatUpdated should be called to notify frontend of build ID") }) t.Run("FailedBuildQuota", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) orgResp := dbfake.Organization(t, db). EveryoneAllowance(40). Members(user). Do() org := orgResp.Org // Create a workspace with a build that is still running. wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, DailyCost: 40, }).Starting().Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-failed-build", }) authzDB := dbauthz.New( db, rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), testAccessControlStorePointer(), ) jobRead := make(chan struct{}, 1) wrappedDB := &jobInterceptStore{Store: authzDB, jobRead: jobRead} tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil }, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for an in-progress build") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) type toolResult struct { resp fantasy.ToolResponse err error } done := make(chan toolResult, 1) go func() { resp, err := tool.Run( dbauthz.AsChatd(ctx), fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}, ) done <- toolResult{resp, err} }() // Wait for the tool to observe the running job. testutil.TryReceive(ctx, t, jobRead) // Fail the build. now := time.Now().UTC() require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{ ID: wsResp.Build.JobID, UpdatedAt: now, CompletedAt: sql.NullTime{Time: now, Valid: true}, Error: sql.NullString{String: "insufficient quota", Valid: true}, ErrorCode: sql.NullString{ String: string(codersdk.InsufficientQuota), Valid: true, }, })) res := testutil.TryReceive(ctx, t, done) require.NoError(t, res.err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(res.resp.Content), &result)) require.Contains(t, result["error"], "waiting for in-progress build") require.Equal(t, string(codersdk.InsufficientQuota), result["error_code"]) require.Equal(t, "Workspace quota reached", result["title"]) require.Contains(t, result["message"], "workspace quota is full") require.Equal(t, wsResp.Build.ID.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, res.resp.IsError, "quota responses must not set IsError; chatprompt strips structured fields from error responses") }) t.Run("StartTriggeredBuildFailure", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) // Create a stopped workspace with a succeeded stop transition. wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-triggered-generic-build-failure", }) var startBuildJobID uuid.UUID var startBuildID uuid.UUID startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition) require.Equal(t, ws.ID, wsID) buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStart, BuildNumber: 2, }).Starting().Do() startBuildJobID = buildResp.Build.JobID startBuildID = buildResp.Build.ID return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil } jobRead := make(chan struct{}, 2) wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead} tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil }, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), }) type toolResult struct { resp fantasy.ToolResponse err error } done := make(chan toolResult, 1) go func() { resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) done <- toolResult{resp, err} }() testutil.TryReceive(ctx, t, jobRead) testutil.TryReceive(ctx, t, jobRead) now := time.Now().UTC() require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{ ID: startBuildJobID, UpdatedAt: now, CompletedAt: sql.NullTime{Time: now, Valid: true}, Error: sql.NullString{String: "terraform apply failed", Valid: true}, })) res := testutil.TryReceive(ctx, t, done) require.NoError(t, res.err) var result map[string]any require.NoError(t, json.Unmarshal([]byte(res.resp.Content), &result)) require.Contains(t, result["error"], "workspace start build failed") require.Equal(t, startBuildID.String(), result["build_id"]) require.NotContains(t, result, "error_code") require.NotContains(t, result, "quota") require.False(t, res.resp.IsError, "buildToolResponse must not set IsError; chatprompt strips structured fields from error responses") }) t.Run("DeletedWorkspace", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) // Create a workspace that has been soft-deleted. wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, Deleted: true, }).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionDelete, }).Do() ws := wsResp.Workspace chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-deleted-workspace", }) tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for deleted workspace") return codersdk.WorkspaceBuild{}, nil }, WorkspaceMu: &sync.Mutex{}, }) resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"}) require.NoError(t, err) require.Contains(t, resp.Content, "workspace was deleted") }) } // seedModelConfig inserts a provider and model config for testing. func seedModelConfig( t *testing.T, db database.Store, ) database.ChatModelConfig { t.Helper() dbgen.ChatProvider(t, db, database.ChatProvider{}) return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ IsDefault: true, }) } // jobInterceptStore wraps a database.Store and signals a // channel after the first GetProvisionerJobByID read completes. // This lets the test synchronize: the tool observes the Running // job status before the main goroutine completes the build. type jobInterceptStore struct { database.Store jobRead chan struct{} } func (s *jobInterceptStore) GetProvisionerJobByID(ctx context.Context, id uuid.UUID) (database.ProvisionerJob, error) { result, err := s.Store.GetProvisionerJobByID(ctx, id) select { case s.jobRead <- struct{}{}: default: } return result, err } func testAccessControlStorePointer() *atomic.Pointer[dbauthz.AccessControlStore] { acs := &atomic.Pointer[dbauthz.AccessControlStore]{} var store dbauthz.AccessControlStore = dbauthz.AGPLTemplateAccessControlStore{} acs.Store(&store) return acs }