From ebcfae27a2b6cca88a10d37a21f0e1b5b8f90b78 Mon Sep 17 00:00:00 2001 From: Asher Date: Wed, 1 Oct 2025 13:39:45 -0800 Subject: [PATCH] feat: add task create, list, status, and delete MCP tools (#19901) --- coderd/database/dbfake/dbfake.go | 27 +++ codersdk/toolsdk/toolsdk.go | 253 +++++++++++++++++++++- codersdk/toolsdk/toolsdk_test.go | 356 +++++++++++++++++++++++++++++++ 3 files changed, 633 insertions(+), 3 deletions(-) diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go index 6d99005fb3..a4742a09ff 100644 --- a/coderd/database/dbfake/dbfake.go +++ b/coderd/database/dbfake/dbfake.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/telemetry" "github.com/coder/coder/v2/coderd/wspubsub" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisionersdk" sdkproto "github.com/coder/coder/v2/provisionersdk/proto" ) @@ -55,6 +56,7 @@ type WorkspaceBuildBuilder struct { params []database.WorkspaceBuildParameter agentToken string dispo workspaceBuildDisposition + taskAppID uuid.UUID } type workspaceBuildDisposition struct { @@ -117,6 +119,23 @@ func (b WorkspaceBuildBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) [] return b } +func (b WorkspaceBuildBuilder) WithTask() WorkspaceBuildBuilder { + //nolint: revive // returns modified struct + b.taskAppID = uuid.New() + return b.Params(database.WorkspaceBuildParameter{ + Name: codersdk.AITaskPromptParameterName, + Value: "list me", + }).WithAgent(func(a []*sdkproto.Agent) []*sdkproto.Agent { + a[0].Apps = []*sdkproto.App{ + { + Id: b.taskAppID.String(), + Slug: "vcode", + }, + } + return a + }) +} + func (b WorkspaceBuildBuilder) Starting() WorkspaceBuildBuilder { //nolint: revive // returns modified struct b.dispo.starting = true @@ -134,6 +153,14 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse { b.seed.ID = uuid.New() b.seed.JobID = jobID + if b.taskAppID != uuid.Nil { + b.seed.HasAITask = sql.NullBool{ + Bool: true, + Valid: true, + } + b.seed.AITaskSidebarAppID = uuid.NullUUID{UUID: b.taskAppID, Valid: true} + } + resp := WorkspaceResponse{ AgentToken: b.agentToken, } diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go index acc7a9b102..63e17c2f03 100644 --- a/codersdk/toolsdk/toolsdk.go +++ b/codersdk/toolsdk/toolsdk.go @@ -50,6 +50,10 @@ const ( ToolNameWorkspaceEditFile = "coder_workspace_edit_file" ToolNameWorkspaceEditFiles = "coder_workspace_edit_files" ToolNameWorkspacePortForward = "coder_workspace_port_forward" + ToolNameCreateTask = "coder_create_task" + ToolNameDeleteTask = "coder_delete_task" + ToolNameListTasks = "coder_list_tasks" + ToolNameGetTaskStatus = "coder_get_task_status" ) func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) { @@ -223,6 +227,10 @@ var All = []GenericTool{ WorkspaceEditFile.Generic(), WorkspaceEditFiles.Generic(), WorkspacePortForward.Generic(), + CreateTask.Generic(), + DeleteTask.Generic(), + ListTasks.Generic(), + GetTaskStatus.Generic(), } type ReportTaskArgs struct { @@ -344,7 +352,7 @@ is provisioned correctly and the agent can connect to the control plane. Properties: map[string]any{ "user": map[string]any{ "type": "string", - "description": "Username or ID of the user to create the workspace for. Use the `me` keyword to create a workspace for the authenticated user.", + "description": userDescription("create a workspace"), }, "template_version_id": map[string]any{ "type": "string", @@ -1393,8 +1401,6 @@ type WorkspaceLSResponse struct { Contents []WorkspaceLSFile `json:"contents"` } -const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used." - var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{ Tool: aisdk.Tool{ Name: ToolNameWorkspaceLS, @@ -1750,6 +1756,237 @@ var WorkspacePortForward = Tool[WorkspacePortForwardArgs, WorkspacePortForwardRe }, } +type CreateTaskArgs struct { + Input string `json:"input"` + TemplateVersionID string `json:"template_version_id"` + TemplateVersionPresetID string `json:"template_version_preset_id"` + User string `json:"user"` +} + +var CreateTask = Tool[CreateTaskArgs, codersdk.Task]{ + Tool: aisdk.Tool{ + Name: ToolNameCreateTask, + Description: `Create a task.`, + Schema: aisdk.Schema{ + Properties: map[string]any{ + "input": map[string]any{ + "type": "string", + "description": "Input/prompt for the task.", + }, + "template_version_id": map[string]any{ + "type": "string", + "description": "ID of the template version to create the task from.", + }, + "template_version_preset_id": map[string]any{ + "type": "string", + "description": "Optional ID of the template version preset to create the task from.", + }, + "user": map[string]any{ + "type": "string", + "description": userDescription("create a task"), + }, + }, + Required: []string{"input", "template_version_id"}, + }, + }, + UserClientOptional: true, + Handler: func(ctx context.Context, deps Deps, args CreateTaskArgs) (codersdk.Task, error) { + if args.Input == "" { + return codersdk.Task{}, xerrors.New("input is required") + } + + tvID, err := uuid.Parse(args.TemplateVersionID) + if err != nil { + return codersdk.Task{}, xerrors.New("template_version_id must be a valid UUID") + } + + var tvPresetID uuid.UUID + if args.TemplateVersionPresetID != "" { + tvPresetID, err = uuid.Parse(args.TemplateVersionPresetID) + if err != nil { + return codersdk.Task{}, xerrors.New("template_version_preset_id must be a valid UUID") + } + } + + if args.User == "" { + args.User = codersdk.Me + } + + expClient := codersdk.NewExperimentalClient(deps.coderClient) + task, err := expClient.CreateTask(ctx, args.User, codersdk.CreateTaskRequest{ + Input: args.Input, + TemplateVersionID: tvID, + TemplateVersionPresetID: tvPresetID, + }) + if err != nil { + return codersdk.Task{}, xerrors.Errorf("create task: %w", err) + } + + return task, nil + }, +} + +type DeleteTaskArgs struct { + TaskID string `json:"task_id"` +} + +var DeleteTask = Tool[DeleteTaskArgs, codersdk.Response]{ + Tool: aisdk.Tool{ + Name: ToolNameDeleteTask, + Description: `Delete a task.`, + Schema: aisdk.Schema{ + Properties: map[string]any{ + "task_id": map[string]any{ + "type": "string", + "description": taskIDDescription("delete"), + }, + }, + Required: []string{"task_id"}, + }, + }, + UserClientOptional: true, + Handler: func(ctx context.Context, deps Deps, args DeleteTaskArgs) (codersdk.Response, error) { + if args.TaskID == "" { + return codersdk.Response{}, xerrors.New("task_id is required") + } + + expClient := codersdk.NewExperimentalClient(deps.coderClient) + + var owner string + id, err := uuid.Parse(args.TaskID) + if err == nil { + task, err := expClient.TaskByID(ctx, id) + if err != nil { + return codersdk.Response{}, xerrors.Errorf("get task %q: %w", args.TaskID, err) + } + owner = task.OwnerName + } else { + ws, err := normalizedNamedWorkspace(ctx, deps.coderClient, args.TaskID) + if err != nil { + return codersdk.Response{}, xerrors.Errorf("get task workspace %q: %w", args.TaskID, err) + } + owner = ws.OwnerName + id = ws.ID + } + + err = expClient.DeleteTask(ctx, owner, id) + if err != nil { + return codersdk.Response{}, xerrors.Errorf("delete task: %w", err) + } + + return codersdk.Response{ + Message: "Task deleted successfully", + }, nil + }, +} + +type ListTasksArgs struct { + Status string `json:"status"` + User string `json:"user"` +} + +type ListTasksResponse struct { + Tasks []codersdk.Task `json:"tasks"` +} + +var ListTasks = Tool[ListTasksArgs, ListTasksResponse]{ + Tool: aisdk.Tool{ + Name: ToolNameListTasks, + Description: `List tasks.`, + Schema: aisdk.Schema{ + Properties: map[string]any{ + "status": map[string]any{ + "type": "string", + "description": "Optional filter by task status.", + }, + "user": map[string]any{ + "type": "string", + "description": userDescription("list tasks"), + }, + }, + Required: []string{}, + }, + }, + UserClientOptional: true, + Handler: func(ctx context.Context, deps Deps, args ListTasksArgs) (ListTasksResponse, error) { + if args.User == "" { + args.User = codersdk.Me + } + + expClient := codersdk.NewExperimentalClient(deps.coderClient) + tasks, err := expClient.Tasks(ctx, &codersdk.TasksFilter{ + Owner: args.User, + Status: args.Status, + }) + if err != nil { + return ListTasksResponse{}, xerrors.Errorf("list tasks: %w", err) + } + + return ListTasksResponse{ + Tasks: tasks, + }, nil + }, +} + +type GetTaskStatusArgs struct { + TaskID string `json:"task_id"` +} + +type GetTaskStatusResponse struct { + Status codersdk.WorkspaceStatus `json:"status"` + State *codersdk.TaskStateEntry `json:"state"` +} + +var GetTaskStatus = Tool[GetTaskStatusArgs, GetTaskStatusResponse]{ + Tool: aisdk.Tool{ + Name: ToolNameGetTaskStatus, + Description: `Get the status of a task.`, + Schema: aisdk.Schema{ + Properties: map[string]any{ + "task_id": map[string]any{ + "type": "string", + "description": taskIDDescription("get"), + }, + }, + Required: []string{"task_id"}, + }, + }, + UserClientOptional: true, + Handler: func(ctx context.Context, deps Deps, args GetTaskStatusArgs) (GetTaskStatusResponse, error) { + if args.TaskID == "" { + return GetTaskStatusResponse{}, xerrors.New("task_id is required") + } + + expClient := codersdk.NewExperimentalClient(deps.coderClient) + + id, err := uuid.Parse(args.TaskID) + if err != nil { + ws, err := normalizedNamedWorkspace(ctx, deps.coderClient, args.TaskID) + if err != nil { + return GetTaskStatusResponse{}, xerrors.Errorf("get task workspace %q: %w", args.TaskID, err) + } + id = ws.ID + } + + task, err := expClient.TaskByID(ctx, id) + if err != nil { + return GetTaskStatusResponse{}, xerrors.Errorf("get task %q: %w", args.TaskID, err) + } + + return GetTaskStatusResponse{ + Status: task.Status, + State: task.CurrentState, + }, nil + }, +} + +// normalizedNamedWorkspace normalizes the workspace name before getting the +// workspace by name. +func normalizedNamedWorkspace(ctx context.Context, client *codersdk.Client, name string) (codersdk.Workspace, error) { + // Maybe namedWorkspace should itself call NormalizeWorkspaceInput? + return namedWorkspace(ctx, client, NormalizeWorkspaceInput(name)) +} + // NormalizeWorkspaceInput converts workspace name input to standard format. // Handles the following input formats: // - workspace → workspace @@ -1810,3 +2047,13 @@ func newAgentConn(ctx context.Context, client *codersdk.Client, workspace string } return conn, nil } + +const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used." + +func taskIDDescription(action string) string { + return fmt.Sprintf("ID or workspace identifier in the format [owner/]workspace[.agent] for the task to %s. If an owner is not specified, the authenticated user is used.", action) +} + +func userDescription(action string) string { + return fmt.Sprintf("Username or ID of the user for which to %s. Omit or use the `me` keyword to %s for the authenticated user.", action, action) +} diff --git a/codersdk/toolsdk/toolsdk_test.go b/codersdk/toolsdk/toolsdk_test.go index f89f22e008..8e0cc09fff 100644 --- a/codersdk/toolsdk/toolsdk_test.go +++ b/codersdk/toolsdk/toolsdk_test.go @@ -2,6 +2,7 @@ package toolsdk_test import ( "context" + "database/sql" "encoding/json" "fmt" "os" @@ -791,6 +792,361 @@ func TestTools(t *testing.T) { }) } }) + + t.Run("WorkspaceCreateTask", func(t *testing.T) { + t.Parallel() + + presetID := uuid.New() + // nolint:gocritic // This is in a test package and does not end up in the build + aiTV := dbfake.TemplateVersion(t, store).Seed(database.TemplateVersion{ + OrganizationID: owner.OrganizationID, + CreatedBy: member.ID, + HasAITask: sql.NullBool{ + Bool: true, + Valid: true, + }, + }).Preset(database.TemplateVersionPreset{ + ID: presetID, + DesiredInstances: sql.NullInt32{ + Int32: 1, + Valid: true, + }, + }).Do() + + tests := []struct { + name string + args toolsdk.CreateTaskArgs + error string + }{ + { + name: "OK", + args: toolsdk.CreateTaskArgs{ + TemplateVersionID: aiTV.TemplateVersion.ID.String(), + Input: "do a barrel roll", + User: "me", + }, + }, + { + name: "NoUser", + args: toolsdk.CreateTaskArgs{ + TemplateVersionID: aiTV.TemplateVersion.ID.String(), + Input: "do another barrel roll", + }, + }, + { + name: "NoInput", + args: toolsdk.CreateTaskArgs{ + TemplateVersionID: aiTV.TemplateVersion.ID.String(), + }, + error: "input is required", + }, + { + name: "NotTaskTemplate", + args: toolsdk.CreateTaskArgs{ + TemplateVersionID: r.TemplateVersion.ID.String(), + Input: "do yet another barrel roll", + }, + error: "Template does not have required parameter \"AI Prompt\"", + }, + { + name: "WithPreset", + args: toolsdk.CreateTaskArgs{ + TemplateVersionID: r.TemplateVersion.ID.String(), + TemplateVersionPresetID: presetID.String(), + Input: "not enough barrel rolls", + }, + error: "Template does not have required parameter \"AI Prompt\"", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tb, err := toolsdk.NewDeps(memberClient) + require.NoError(t, err) + + _, err = testTool(t, toolsdk.CreateTask, tb, tt.args) + if tt.error != "" { + require.Error(t, err) + require.ErrorContains(t, err, tt.error) + } else { + require.NoError(t, err) + } + }) + } + }) + + t.Run("WorkspaceDeleteTask", func(t *testing.T) { + t.Parallel() + + // nolint:gocritic // This is in a test package and does not end up in the build + aiTV := dbfake.TemplateVersion(t, store).Seed(database.TemplateVersion{ + OrganizationID: owner.OrganizationID, + CreatedBy: member.ID, + HasAITask: sql.NullBool{ + Bool: true, + Valid: true, + }, + }).Do() + + // nolint:gocritic // This is in a test package and does not end up in the build + ws1 := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + Name: "delete-task-workspace-1", + OrganizationID: owner.OrganizationID, + OwnerID: member.ID, + TemplateID: aiTV.Template.ID, + }).WithTask().Do() + + // nolint:gocritic // This is in a test package and does not end up in the build + _ = dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + Name: "delete-task-workspace-2", + OrganizationID: owner.OrganizationID, + OwnerID: member.ID, + TemplateID: aiTV.Template.ID, + }).WithTask().Do() + + tests := []struct { + name string + args toolsdk.DeleteTaskArgs + error string + }{ + { + name: "ByUUID", + args: toolsdk.DeleteTaskArgs{ + TaskID: ws1.Workspace.ID.String(), + }, + }, + { + name: "ByWorkspaceIdentifier", + args: toolsdk.DeleteTaskArgs{ + TaskID: "delete-task-workspace-2", + }, + }, + { + name: "NoID", + args: toolsdk.DeleteTaskArgs{}, + error: "task_id is required", + }, + { + name: "NoTaskByID", + args: toolsdk.DeleteTaskArgs{ + TaskID: uuid.New().String(), + }, + error: "Resource not found", + }, + { + name: "NoTaskByWorkspaceIdentifier", + args: toolsdk.DeleteTaskArgs{ + TaskID: "non-existent", + }, + error: "Resource not found", + }, + { + name: "ExistsButNotATask", + args: toolsdk.DeleteTaskArgs{ + TaskID: r.Workspace.ID.String(), + }, + error: "Resource not found", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tb, err := toolsdk.NewDeps(memberClient) + require.NoError(t, err) + + _, err = testTool(t, toolsdk.DeleteTask, tb, tt.args) + if tt.error != "" { + require.Error(t, err) + require.ErrorContains(t, err, tt.error) + } else { + require.NoError(t, err) + } + }) + } + }) + + t.Run("WorkspaceListTasks", func(t *testing.T) { + t.Parallel() + + taskClient, taskUser := coderdtest.CreateAnotherUserMutators(t, client, owner.OrganizationID, nil) + + // nolint:gocritic // This is in a test package and does not end up in the build + aiTV := dbfake.TemplateVersion(t, store).Seed(database.TemplateVersion{ + OrganizationID: owner.OrganizationID, + CreatedBy: owner.UserID, + HasAITask: sql.NullBool{ + Bool: true, + Valid: true, + }, + }).Do() + + // This task should not show up since listing is user-scoped. + // nolint:gocritic // This is in a test package and does not end up in the build + _ = dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + Name: "list-task-workspace-member", + OrganizationID: owner.OrganizationID, + OwnerID: member.ID, + TemplateID: aiTV.Template.ID, + }).WithTask().Do() + + // These tasks should show up. + for i := range 5 { + // nolint:gocritic // This is in a test package and does not end up in the build + var transition database.WorkspaceTransition + if i == 0 { + // nolint:gocritic // This is in a test package and does not end up in the build + transition = database.WorkspaceTransitionStop + } + // nolint:gocritic // This is in a test package and does not end up in the build + _ = dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + Name: fmt.Sprintf("list-task-workspace-%d", i), + OrganizationID: owner.OrganizationID, + OwnerID: taskUser.ID, + TemplateID: aiTV.Template.ID, + }).Seed(database.WorkspaceBuild{Transition: transition}).WithTask().Do() + } + + tests := []struct { + name string + args toolsdk.ListTasksArgs + expected []string + error string + }{ + { + name: "ListAllOwned", + args: toolsdk.ListTasksArgs{}, + expected: []string{ + "list-task-workspace-0", + "list-task-workspace-1", + "list-task-workspace-2", + "list-task-workspace-3", + "list-task-workspace-4", + }, + }, + { + name: "ListFiltered", + args: toolsdk.ListTasksArgs{ + Status: "stopped", + }, + expected: []string{ + "list-task-workspace-0", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tb, err := toolsdk.NewDeps(taskClient) + require.NoError(t, err) + + res, err := testTool(t, toolsdk.ListTasks, tb, tt.args) + if tt.error != "" { + require.Error(t, err) + require.ErrorContains(t, err, tt.error) + } else { + require.NoError(t, err) + require.Len(t, res.Tasks, len(tt.expected)) + for _, task := range res.Tasks { + require.Contains(t, tt.expected, task.Name) + } + } + }) + } + }) + + t.Run("WorkspaceGetTask", func(t *testing.T) { + t.Parallel() + + // nolint:gocritic // This is in a test package and does not end up in the build + aiTV := dbfake.TemplateVersion(t, store).Seed(database.TemplateVersion{ + OrganizationID: owner.OrganizationID, + CreatedBy: member.ID, + HasAITask: sql.NullBool{ + Bool: true, + Valid: true, + }, + }).Do() + + // nolint:gocritic // This is in a test package and does not end up in the build + ws1 := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + Name: "get-task-workspace-1", + OrganizationID: owner.OrganizationID, + OwnerID: member.ID, + TemplateID: aiTV.Template.ID, + }).WithTask().Do() + + tests := []struct { + name string + args toolsdk.GetTaskStatusArgs + expected codersdk.WorkspaceStatus + error string + }{ + { + name: "ByUUID", + args: toolsdk.GetTaskStatusArgs{ + TaskID: ws1.Workspace.ID.String(), + }, + expected: codersdk.WorkspaceStatusRunning, + }, + { + name: "ByWorkspaceIdentifier", + args: toolsdk.GetTaskStatusArgs{ + TaskID: "get-task-workspace-1", + }, + expected: codersdk.WorkspaceStatusRunning, + }, + { + name: "NoID", + args: toolsdk.GetTaskStatusArgs{}, + error: "task_id is required", + }, + { + name: "NoTaskByID", + args: toolsdk.GetTaskStatusArgs{ + TaskID: uuid.New().String(), + }, + error: "Resource not found", + }, + { + name: "NoTaskByWorkspaceIdentifier", + args: toolsdk.GetTaskStatusArgs{ + TaskID: "non-existent", + }, + error: "Resource not found", + }, + { + name: "ExistsButNotATask", + args: toolsdk.GetTaskStatusArgs{ + TaskID: r.Workspace.ID.String(), + }, + error: "Resource not found", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tb, err := toolsdk.NewDeps(memberClient) + require.NoError(t, err) + + res, err := testTool(t, toolsdk.GetTaskStatus, tb, tt.args) + if tt.error != "" { + require.Error(t, err) + require.ErrorContains(t, err, tt.error) + } else { + require.NoError(t, err) + require.Equal(t, tt.expected, res.Status) + } + }) + } + }) } // TestedTools keeps track of which tools have been tested.