From 63c1325ad5f52dbce8e9193563a91e1c3962049f Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Tue, 26 Aug 2025 15:24:42 +0100 Subject: [PATCH 1/2] feat(cli): add exp task create command (#19492) Partially implements https://github.com/coder/internal/issues/893 This isn't the full implementation of `coder exp tasks create` as defined in the issue, but it is the minimum required to create a task. --- cli/exp_task.go | 1 + cli/exp_taskcreate.go | 127 +++++++++++++++++++++ cli/exp_taskcreate_test.go | 227 +++++++++++++++++++++++++++++++++++++ 3 files changed, 355 insertions(+) create mode 100644 cli/exp_taskcreate.go create mode 100644 cli/exp_taskcreate_test.go diff --git a/cli/exp_task.go b/cli/exp_task.go index 81316d155000d..860f7b954f47f 100644 --- a/cli/exp_task.go +++ b/cli/exp_task.go @@ -14,6 +14,7 @@ func (r *RootCmd) tasksCommand() *serpent.Command { }, Children: []*serpent.Command{ r.taskList(), + r.taskCreate(), }, } return cmd diff --git a/cli/exp_taskcreate.go b/cli/exp_taskcreate.go new file mode 100644 index 0000000000000..b23da632a12c2 --- /dev/null +++ b/cli/exp_taskcreate.go @@ -0,0 +1,127 @@ +package cli + +import ( + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/serpent" +) + +func (r *RootCmd) taskCreate() *serpent.Command { + var ( + orgContext = NewOrganizationContext() + client = new(codersdk.Client) + + templateName string + templateVersionName string + presetName string + taskInput string + ) + + return &serpent.Command{ + Use: "create [template]", + Short: "Create an experimental task", + Middleware: serpent.Chain( + serpent.RequireRangeArgs(0, 1), + r.InitClient(client), + ), + Options: serpent.OptionSet{ + { + Flag: "input", + Env: "CODER_TASK_INPUT", + Value: serpent.StringOf(&taskInput), + Required: true, + }, + { + Env: "CODER_TASK_TEMPLATE_NAME", + Value: serpent.StringOf(&templateName), + }, + { + Env: "CODER_TASK_TEMPLATE_VERSION", + Value: serpent.StringOf(&templateVersionName), + }, + { + Flag: "preset", + Env: "CODER_TASK_PRESET_NAME", + Value: serpent.StringOf(&presetName), + Default: PresetNone, + }, + }, + Handler: func(inv *serpent.Invocation) error { + var ( + ctx = inv.Context() + expClient = codersdk.NewExperimentalClient(client) + + templateVersionID uuid.UUID + templateVersionPresetID uuid.UUID + ) + + organization, err := orgContext.Selected(inv, client) + if err != nil { + return xerrors.Errorf("get current organization: %w", err) + } + + if len(inv.Args) > 0 { + templateName, templateVersionName, _ = strings.Cut(inv.Args[0], "@") + } + + if templateName == "" { + return xerrors.Errorf("template name not provided") + } + + if templateVersionName != "" { + templateVersion, err := client.TemplateVersionByOrganizationAndName(ctx, organization.ID, templateName, templateVersionName) + if err != nil { + return xerrors.Errorf("get template version: %w", err) + } + + templateVersionID = templateVersion.ID + } else { + template, err := client.TemplateByName(ctx, organization.ID, templateName) + if err != nil { + return xerrors.Errorf("get template: %w", err) + } + + templateVersionID = template.ActiveVersionID + } + + if presetName != PresetNone { + templatePresets, err := client.TemplateVersionPresets(ctx, templateVersionID) + if err != nil { + return xerrors.Errorf("get template presets: %w", err) + } + + preset, err := resolvePreset(templatePresets, presetName) + if err != nil { + return xerrors.Errorf("resolve preset: %w", err) + } + + templateVersionPresetID = preset.ID + } + + workspace, err := expClient.CreateTask(ctx, codersdk.Me, codersdk.CreateTaskRequest{ + TemplateVersionID: templateVersionID, + TemplateVersionPresetID: templateVersionPresetID, + Prompt: taskInput, + }) + if err != nil { + return xerrors.Errorf("create task: %w", err) + } + + _, _ = fmt.Fprintf( + inv.Stdout, + "The task %s has been created at %s!\n", + cliui.Keyword(workspace.Name), + cliui.Timestamp(time.Now()), + ) + + return nil + }, + } +} diff --git a/cli/exp_taskcreate_test.go b/cli/exp_taskcreate_test.go new file mode 100644 index 0000000000000..7a4a4bfb5a43e --- /dev/null +++ b/cli/exp_taskcreate_test.go @@ -0,0 +1,227 @@ +package cli_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/serpent" +) + +func TestTaskCreate(t *testing.T) { + t.Parallel() + + var ( + organizationID = uuid.New() + templateID = uuid.New() + templateVersionID = uuid.New() + templateVersionPresetID = uuid.New() + ) + + templateAndVersionFoundHandler := func(t *testing.T, ctx context.Context, templateName, templateVersionName, presetName, prompt string) http.HandlerFunc { + t.Helper() + + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template/versions/my-template-version", organizationID): + httpapi.Write(ctx, w, http.StatusOK, codersdk.TemplateVersion{ + ID: templateVersionID, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template", organizationID): + httpapi.Write(ctx, w, http.StatusOK, codersdk.Template{ + ID: templateID, + ActiveVersionID: templateVersionID, + }) + case fmt.Sprintf("/api/v2/templateversions/%s/presets", templateVersionID): + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Preset{ + { + ID: templateVersionPresetID, + Name: presetName, + }, + }) + case "/api/experimental/tasks/me": + var req codersdk.CreateTaskRequest + if !httpapi.Read(ctx, w, r, &req) { + return + } + + assert.Equal(t, prompt, req.Prompt, "prompt mismatch") + assert.Equal(t, templateVersionID, req.TemplateVersionID, "template version mismatch") + + if presetName == "" { + assert.Equal(t, uuid.Nil, req.TemplateVersionPresetID, "expected no template preset id") + } else { + assert.Equal(t, templateVersionPresetID, req.TemplateVersionPresetID, "template version preset id mismatch") + } + + httpapi.Write(ctx, w, http.StatusCreated, codersdk.Workspace{ + Name: "task-wild-goldfish-27", + }) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + } + + tests := []struct { + args []string + env []string + expectError string + expectOutput string + handler func(t *testing.T, ctx context.Context) http.HandlerFunc + }{ + { + args: []string{"my-template@my-template-version", "--input", "my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + env: []string{"CODER_TASK_TEMPLATE_VERSION=my-template-version"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"--input", "my custom prompt"}, + env: []string{"CODER_TASK_TEMPLATE_NAME=my-template", "CODER_TASK_TEMPLATE_VERSION=my-template-version"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + env: []string{"CODER_TASK_TEMPLATE_NAME=my-template", "CODER_TASK_TEMPLATE_VERSION=my-template-version", "CODER_TASK_INPUT=my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt", "--preset", "my-preset"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + env: []string{"CODER_TASK_PRESET_NAME=my-preset"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt", "--preset", "not-real-preset"}, + expectError: `preset "not-real-preset" not found`, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template@not-real-template-version", "--input", "my custom prompt"}, + expectError: httpapi.ResourceNotFoundResponse.Message, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template/versions/not-real-template-version", organizationID): + httpapi.ResourceNotFound(w) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + }, + }, + { + args: []string{"not-real-template", "--input", "my custom prompt"}, + expectError: httpapi.ResourceNotFoundResponse.Message, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/not-real-template", organizationID): + httpapi.ResourceNotFound(w) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + }, + }, + } + + for _, tt := range tests { + t.Run(strings.Join(tt.args, ","), func(t *testing.T) { + t.Parallel() + + var ( + ctx = testutil.Context(t, testutil.WaitShort) + srv = httptest.NewServer(tt.handler(t, ctx)) + client = new(codersdk.Client) + args = []string{"exp", "task", "create"} + sb strings.Builder + err error + ) + + t.Cleanup(srv.Close) + + client.URL, err = url.Parse(srv.URL) + require.NoError(t, err) + + inv, root := clitest.New(t, append(args, tt.args...)...) + inv.Environ = serpent.ParseEnviron(tt.env, "") + inv.Stdout = &sb + inv.Stderr = &sb + clitest.SetupConfig(t, client, root) + + err = inv.WithContext(ctx).Run() + if tt.expectError == "" { + assert.NoError(t, err) + } else { + assert.ErrorContains(t, err, tt.expectError) + } + + assert.Contains(t, sb.String(), tt.expectOutput) + }) + } +} From ef0d74fb750f6e4c342c9ed12fc1ae630b4ea69b Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Tue, 26 Aug 2025 09:26:11 -0500 Subject: [PATCH 2/2] chore: improve performance of 'GetLatestWorkspaceBuildsByWorkspaceIDs' (#19452) Closes https://github.com/coder/internal/issues/716 This prevents a scan over the entire `workspace_build` table by removing a `join`. This is still imperfect as we are still scanning over the number of builds for the workspaces in the arguments. Ideally we would have some index or something precomputed. Then we could skip scanning over the builds for the correct workspaces that are not the latest. --- coderd/database/querier_test.go | 73 +++++++++++++++++++++ coderd/database/queries.sql.go | 23 +++---- coderd/database/queries/workspacebuilds.sql | 24 +++---- 3 files changed, 92 insertions(+), 28 deletions(-) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 18c10d6388f37..a8b3c186edd8b 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -32,6 +32,7 @@ import ( "github.com/coder/coder/v2/coderd/provisionerdserver" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" + "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisionersdk" "github.com/coder/coder/v2/testutil" @@ -6579,3 +6580,75 @@ func TestWorkspaceBuildDeadlineConstraint(t *testing.T) { } } } + +// TestGetLatestWorkspaceBuildsByWorkspaceIDs populates the database with +// workspaces and builds. It then tests that +// GetLatestWorkspaceBuildsByWorkspaceIDs returns the latest build for some +// subset of the workspaces. +func TestGetLatestWorkspaceBuildsByWorkspaceIDs(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + org := dbgen.Organization(t, db, database.Organization{}) + admin := dbgen.User(t, db, database.User{}) + + tv := dbfake.TemplateVersion(t, db). + Seed(database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: admin.ID, + }). + Do() + + users := make([]database.User, 5) + wrks := make([][]database.WorkspaceTable, len(users)) + exp := make(map[uuid.UUID]database.WorkspaceBuild) + for i := range users { + users[i] = dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: users[i].ID, + OrganizationID: org.ID, + }) + + // Each user gets 2 workspaces. + wrks[i] = make([]database.WorkspaceTable, 2) + for wi := range wrks[i] { + wrks[i][wi] = dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tv.Template.ID, + OwnerID: users[i].ID, + }) + + // Choose a deterministic number of builds per workspace + // No more than 5 builds though, that would be excessive. + for j := int32(1); int(j) <= (i+wi)%5; j++ { + wb := dbfake.WorkspaceBuild(t, db, wrks[i][wi]). + Seed(database.WorkspaceBuild{ + WorkspaceID: wrks[i][wi].ID, + BuildNumber: j + 1, + }). + Do() + + exp[wrks[i][wi].ID] = wb.Build // Save the final workspace build + } + } + } + + // Only take half the users. And only take 1 workspace per user for the test. + // The others are just noice. This just queries a subset of workspaces and builds + // to make sure the noise doesn't interfere with the results. + assertWrks := wrks[:len(users)/2] + ctx := testutil.Context(t, testutil.WaitLong) + ids := slice.Convert[[]database.WorkspaceTable, uuid.UUID](assertWrks, func(pair []database.WorkspaceTable) uuid.UUID { + return pair[0].ID + }) + + require.Greater(t, len(ids), 0, "expected some workspace ids for test") + builds, err := db.GetLatestWorkspaceBuildsByWorkspaceIDs(ctx, ids) + require.NoError(t, err) + for _, b := range builds { + expB, ok := exp[b.WorkspaceID] + require.Truef(t, ok, "unexpected workspace build for workspace id %s", b.WorkspaceID) + require.Equalf(t, expB.ID, b.ID, "unexpected workspace build id for workspace id %s", b.WorkspaceID) + require.Equal(t, expB.BuildNumber, b.BuildNumber, "unexpected build number") + } +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2f56b422f350b..014c433cab690 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -18983,20 +18983,15 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, w } const getLatestWorkspaceBuildsByWorkspaceIDs = `-- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many -SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost, wb.max_deadline, wb.template_version_preset_id, wb.has_ai_task, wb.ai_task_sidebar_app_id, wb.has_external_agent, wb.initiator_by_avatar_url, wb.initiator_by_username, wb.initiator_by_name -FROM ( - SELECT - workspace_id, MAX(build_number) as max_build_number - FROM - workspace_build_with_user AS workspace_builds - WHERE - workspace_id = ANY($1 :: uuid [ ]) - GROUP BY - workspace_id -) m -JOIN - workspace_build_with_user AS wb -ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number +SELECT + DISTINCT ON (workspace_id) + id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, template_version_preset_id, has_ai_task, ai_task_sidebar_app_id, has_external_agent, initiator_by_avatar_url, initiator_by_username, initiator_by_name +FROM + workspace_build_with_user AS workspace_builds +WHERE + workspace_id = ANY($1 :: uuid [ ]) +ORDER BY + workspace_id, build_number DESC -- latest first ` func (q *sqlQuerier) GetLatestWorkspaceBuildsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceBuild, error) { diff --git a/coderd/database/queries/workspacebuilds.sql b/coderd/database/queries/workspacebuilds.sql index 6c020f5a97f50..0736c5514b3f7 100644 --- a/coderd/database/queries/workspacebuilds.sql +++ b/coderd/database/queries/workspacebuilds.sql @@ -76,20 +76,16 @@ LIMIT 1; -- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many -SELECT wb.* -FROM ( - SELECT - workspace_id, MAX(build_number) as max_build_number - FROM - workspace_build_with_user AS workspace_builds - WHERE - workspace_id = ANY(@ids :: uuid [ ]) - GROUP BY - workspace_id -) m -JOIN - workspace_build_with_user AS wb -ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number; +SELECT + DISTINCT ON (workspace_id) + * +FROM + workspace_build_with_user AS workspace_builds +WHERE + workspace_id = ANY(@ids :: uuid [ ]) +ORDER BY + workspace_id, build_number DESC -- latest first +; -- name: InsertWorkspaceBuild :exec INSERT INTO