Skip to content

Commit ebcfae2

Browse files
code-ashercode-asher
authored
feat: add task create, list, status, and delete MCP tools (coder#19901)
1 parent 0993dcf commit ebcfae2

File tree

3 files changed

+633
-3
lines changed

3 files changed

+633
-3
lines changed

‎coderd/database/dbfake/dbfake.go‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/coder/coder/v2/coderd/rbac"
2525
"github.com/coder/coder/v2/coderd/telemetry"
2626
"github.com/coder/coder/v2/coderd/wspubsub"
27+
"github.com/coder/coder/v2/codersdk"
2728
"github.com/coder/coder/v2/provisionersdk"
2829
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
2930
)
@@ -55,6 +56,7 @@ type WorkspaceBuildBuilder struct {
5556
params []database.WorkspaceBuildParameter
5657
agentToken string
5758
dispo workspaceBuildDisposition
59+
taskAppID uuid.UUID
5860
}
5961

6062
type workspaceBuildDisposition struct {
@@ -117,6 +119,23 @@ func (b WorkspaceBuildBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []
117119
return b
118120
}
119121

122+
func (b WorkspaceBuildBuilder) WithTask() WorkspaceBuildBuilder {
123+
//nolint: revive // returns modified struct
124+
b.taskAppID = uuid.New()
125+
return b.Params(database.WorkspaceBuildParameter{
126+
Name: codersdk.AITaskPromptParameterName,
127+
Value: "list me",
128+
}).WithAgent(func(a []*sdkproto.Agent) []*sdkproto.Agent {
129+
a[0].Apps = []*sdkproto.App{
130+
{
131+
Id: b.taskAppID.String(),
132+
Slug: "vcode",
133+
},
134+
}
135+
return a
136+
})
137+
}
138+
120139
func (b WorkspaceBuildBuilder) Starting() WorkspaceBuildBuilder {
121140
//nolint: revive // returns modified struct
122141
b.dispo.starting = true
@@ -134,6 +153,14 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
134153
b.seed.ID = uuid.New()
135154
b.seed.JobID = jobID
136155

156+
if b.taskAppID != uuid.Nil {
157+
b.seed.HasAITask = sql.NullBool{
158+
Bool: true,
159+
Valid: true,
160+
}
161+
b.seed.AITaskSidebarAppID = uuid.NullUUID{UUID: b.taskAppID, Valid: true}
162+
}
163+
137164
resp := WorkspaceResponse{
138165
AgentToken: b.agentToken,
139166
}

‎codersdk/toolsdk/toolsdk.go‎

Lines changed: 250 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ const (
5050
ToolNameWorkspaceEditFile = "coder_workspace_edit_file"
5151
ToolNameWorkspaceEditFiles = "coder_workspace_edit_files"
5252
ToolNameWorkspacePortForward = "coder_workspace_port_forward"
53+
ToolNameCreateTask = "coder_create_task"
54+
ToolNameDeleteTask = "coder_delete_task"
55+
ToolNameListTasks = "coder_list_tasks"
56+
ToolNameGetTaskStatus = "coder_get_task_status"
5357
)
5458

5559
func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) {
@@ -223,6 +227,10 @@ var All = []GenericTool{
223227
WorkspaceEditFile.Generic(),
224228
WorkspaceEditFiles.Generic(),
225229
WorkspacePortForward.Generic(),
230+
CreateTask.Generic(),
231+
DeleteTask.Generic(),
232+
ListTasks.Generic(),
233+
GetTaskStatus.Generic(),
226234
}
227235

228236
type ReportTaskArgs struct {
@@ -344,7 +352,7 @@ is provisioned correctly and the agent can connect to the control plane.
344352
Properties: map[string]any{
345353
"user": map[string]any{
346354
"type": "string",
347-
"description": "Username or ID of the user to create the workspace for. Use the `me` keyword to create a workspace for the authenticated user.",
355+
"description": userDescription("create a workspace"),
348356
},
349357
"template_version_id": map[string]any{
350358
"type": "string",
@@ -1393,8 +1401,6 @@ type WorkspaceLSResponse struct {
13931401
Contents []WorkspaceLSFile `json:"contents"`
13941402
}
13951403

1396-
const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used."
1397-
13981404
var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{
13991405
Tool: aisdk.Tool{
14001406
Name: ToolNameWorkspaceLS,
@@ -1750,6 +1756,237 @@ var WorkspacePortForward = Tool[WorkspacePortForwardArgs, WorkspacePortForwardRe
17501756
},
17511757
}
17521758

1759+
type CreateTaskArgs struct {
1760+
Input string `json:"input"`
1761+
TemplateVersionID string `json:"template_version_id"`
1762+
TemplateVersionPresetID string `json:"template_version_preset_id"`
1763+
User string `json:"user"`
1764+
}
1765+
1766+
var CreateTask = Tool[CreateTaskArgs, codersdk.Task]{
1767+
Tool: aisdk.Tool{
1768+
Name: ToolNameCreateTask,
1769+
Description: `Create a task.`,
1770+
Schema: aisdk.Schema{
1771+
Properties: map[string]any{
1772+
"input": map[string]any{
1773+
"type": "string",
1774+
"description": "Input/prompt for the task.",
1775+
},
1776+
"template_version_id": map[string]any{
1777+
"type": "string",
1778+
"description": "ID of the template version to create the task from.",
1779+
},
1780+
"template_version_preset_id": map[string]any{
1781+
"type": "string",
1782+
"description": "Optional ID of the template version preset to create the task from.",
1783+
},
1784+
"user": map[string]any{
1785+
"type": "string",
1786+
"description": userDescription("create a task"),
1787+
},
1788+
},
1789+
Required: []string{"input", "template_version_id"},
1790+
},
1791+
},
1792+
UserClientOptional: true,
1793+
Handler: func(ctx context.Context, deps Deps, args CreateTaskArgs) (codersdk.Task, error) {
1794+
if args.Input == "" {
1795+
return codersdk.Task{}, xerrors.New("input is required")
1796+
}
1797+
1798+
tvID, err := uuid.Parse(args.TemplateVersionID)
1799+
if err != nil {
1800+
return codersdk.Task{}, xerrors.New("template_version_id must be a valid UUID")
1801+
}
1802+
1803+
var tvPresetID uuid.UUID
1804+
if args.TemplateVersionPresetID != "" {
1805+
tvPresetID, err = uuid.Parse(args.TemplateVersionPresetID)
1806+
if err != nil {
1807+
return codersdk.Task{}, xerrors.New("template_version_preset_id must be a valid UUID")
1808+
}
1809+
}
1810+
1811+
if args.User == "" {
1812+
args.User = codersdk.Me
1813+
}
1814+
1815+
expClient := codersdk.NewExperimentalClient(deps.coderClient)
1816+
task, err := expClient.CreateTask(ctx, args.User, codersdk.CreateTaskRequest{
1817+
Input: args.Input,
1818+
TemplateVersionID: tvID,
1819+
TemplateVersionPresetID: tvPresetID,
1820+
})
1821+
if err != nil {
1822+
return codersdk.Task{}, xerrors.Errorf("create task: %w", err)
1823+
}
1824+
1825+
return task, nil
1826+
},
1827+
}
1828+
1829+
type DeleteTaskArgs struct {
1830+
TaskID string `json:"task_id"`
1831+
}
1832+
1833+
var DeleteTask = Tool[DeleteTaskArgs, codersdk.Response]{
1834+
Tool: aisdk.Tool{
1835+
Name: ToolNameDeleteTask,
1836+
Description: `Delete a task.`,
1837+
Schema: aisdk.Schema{
1838+
Properties: map[string]any{
1839+
"task_id": map[string]any{
1840+
"type": "string",
1841+
"description": taskIDDescription("delete"),
1842+
},
1843+
},
1844+
Required: []string{"task_id"},
1845+
},
1846+
},
1847+
UserClientOptional: true,
1848+
Handler: func(ctx context.Context, deps Deps, args DeleteTaskArgs) (codersdk.Response, error) {
1849+
if args.TaskID == "" {
1850+
return codersdk.Response{}, xerrors.New("task_id is required")
1851+
}
1852+
1853+
expClient := codersdk.NewExperimentalClient(deps.coderClient)
1854+
1855+
var owner string
1856+
id, err := uuid.Parse(args.TaskID)
1857+
if err == nil {
1858+
task, err := expClient.TaskByID(ctx, id)
1859+
if err != nil {
1860+
return codersdk.Response{}, xerrors.Errorf("get task %q: %w", args.TaskID, err)
1861+
}
1862+
owner = task.OwnerName
1863+
} else {
1864+
ws, err := normalizedNamedWorkspace(ctx, deps.coderClient, args.TaskID)
1865+
if err != nil {
1866+
return codersdk.Response{}, xerrors.Errorf("get task workspace %q: %w", args.TaskID, err)
1867+
}
1868+
owner = ws.OwnerName
1869+
id = ws.ID
1870+
}
1871+
1872+
err = expClient.DeleteTask(ctx, owner, id)
1873+
if err != nil {
1874+
return codersdk.Response{}, xerrors.Errorf("delete task: %w", err)
1875+
}
1876+
1877+
return codersdk.Response{
1878+
Message: "Task deleted successfully",
1879+
}, nil
1880+
},
1881+
}
1882+
1883+
type ListTasksArgs struct {
1884+
Status string `json:"status"`
1885+
User string `json:"user"`
1886+
}
1887+
1888+
type ListTasksResponse struct {
1889+
Tasks []codersdk.Task `json:"tasks"`
1890+
}
1891+
1892+
var ListTasks = Tool[ListTasksArgs, ListTasksResponse]{
1893+
Tool: aisdk.Tool{
1894+
Name: ToolNameListTasks,
1895+
Description: `List tasks.`,
1896+
Schema: aisdk.Schema{
1897+
Properties: map[string]any{
1898+
"status": map[string]any{
1899+
"type": "string",
1900+
"description": "Optional filter by task status.",
1901+
},
1902+
"user": map[string]any{
1903+
"type": "string",
1904+
"description": userDescription("list tasks"),
1905+
},
1906+
},
1907+
Required: []string{},
1908+
},
1909+
},
1910+
UserClientOptional: true,
1911+
Handler: func(ctx context.Context, deps Deps, args ListTasksArgs) (ListTasksResponse, error) {
1912+
if args.User == "" {
1913+
args.User = codersdk.Me
1914+
}
1915+
1916+
expClient := codersdk.NewExperimentalClient(deps.coderClient)
1917+
tasks, err := expClient.Tasks(ctx, &codersdk.TasksFilter{
1918+
Owner: args.User,
1919+
Status: args.Status,
1920+
})
1921+
if err != nil {
1922+
return ListTasksResponse{}, xerrors.Errorf("list tasks: %w", err)
1923+
}
1924+
1925+
return ListTasksResponse{
1926+
Tasks: tasks,
1927+
}, nil
1928+
},
1929+
}
1930+
1931+
type GetTaskStatusArgs struct {
1932+
TaskID string `json:"task_id"`
1933+
}
1934+
1935+
type GetTaskStatusResponse struct {
1936+
Status codersdk.WorkspaceStatus `json:"status"`
1937+
State *codersdk.TaskStateEntry `json:"state"`
1938+
}
1939+
1940+
var GetTaskStatus = Tool[GetTaskStatusArgs, GetTaskStatusResponse]{
1941+
Tool: aisdk.Tool{
1942+
Name: ToolNameGetTaskStatus,
1943+
Description: `Get the status of a task.`,
1944+
Schema: aisdk.Schema{
1945+
Properties: map[string]any{
1946+
"task_id": map[string]any{
1947+
"type": "string",
1948+
"description": taskIDDescription("get"),
1949+
},
1950+
},
1951+
Required: []string{"task_id"},
1952+
},
1953+
},
1954+
UserClientOptional: true,
1955+
Handler: func(ctx context.Context, deps Deps, args GetTaskStatusArgs) (GetTaskStatusResponse, error) {
1956+
if args.TaskID == "" {
1957+
return GetTaskStatusResponse{}, xerrors.New("task_id is required")
1958+
}
1959+
1960+
expClient := codersdk.NewExperimentalClient(deps.coderClient)
1961+
1962+
id, err := uuid.Parse(args.TaskID)
1963+
if err != nil {
1964+
ws, err := normalizedNamedWorkspace(ctx, deps.coderClient, args.TaskID)
1965+
if err != nil {
1966+
return GetTaskStatusResponse{}, xerrors.Errorf("get task workspace %q: %w", args.TaskID, err)
1967+
}
1968+
id = ws.ID
1969+
}
1970+
1971+
task, err := expClient.TaskByID(ctx, id)
1972+
if err != nil {
1973+
return GetTaskStatusResponse{}, xerrors.Errorf("get task %q: %w", args.TaskID, err)
1974+
}
1975+
1976+
return GetTaskStatusResponse{
1977+
Status: task.Status,
1978+
State: task.CurrentState,
1979+
}, nil
1980+
},
1981+
}
1982+
1983+
// normalizedNamedWorkspace normalizes the workspace name before getting the
1984+
// workspace by name.
1985+
func normalizedNamedWorkspace(ctx context.Context, client *codersdk.Client, name string) (codersdk.Workspace, error) {
1986+
// Maybe namedWorkspace should itself call NormalizeWorkspaceInput?
1987+
return namedWorkspace(ctx, client, NormalizeWorkspaceInput(name))
1988+
}
1989+
17531990
// NormalizeWorkspaceInput converts workspace name input to standard format.
17541991
// Handles the following input formats:
17551992
// - workspace → workspace
@@ -1810,3 +2047,13 @@ func newAgentConn(ctx context.Context, client *codersdk.Client, workspace string
18102047
}
18112048
return conn, nil
18122049
}
2050+
2051+
const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used."
2052+
2053+
func taskIDDescription(action string) string {
2054+
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)
2055+
}
2056+
2057+
func userDescription(action string) string {
2058+
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)
2059+
}

0 commit comments

Comments
 (0)