Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
242 changes: 242 additions & 0 deletions coderd/aitasks.go
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
package coderd

import (
"context"
"database/sql"
"errors"
"fmt"
"net/http"
"slices"
"strings"

"github.com/go-chi/chi/v5"
"github.com/google/uuid"

"cdr.dev/slog"
Expand All @@ -17,6 +19,8 @@ import (
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/coderd/searchquery"
"github.com/coder/coder/v2/coderd/taskname"
"github.com/coder/coder/v2/codersdk"
)
Expand Down Expand Up @@ -186,3 +190,241 @@ func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
defer commitAudit()
createWorkspace(ctx, aReq, apiKey.UserID, api, owner, createReq, rw, r)
}

// tasksFromWorkspaces converts a slice of API workspaces into tasks, fetching
// prompts and mapping status/state.
func (api *API) tasksFromWorkspaces(ctx context.Context, apiWorkspaces []codersdk.Workspace) ([]codersdk.Task, error) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we fine with this not validating that the passed in workspaces are tasks? I understand that where it is currently called from we only pass tasks in, I'm just curious.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably no harm to filter out workspaces that are not "tasks", wdyt?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be an error if passed to this function. Handling exists in both list and get endpoints currently, before calling this, but I can add a guard here as well or try to consolidate some more of the code into this method.

// Fetch prompts for each workspace build and map by build ID.
buildIDs := make([]uuid.UUID, 0, len(apiWorkspaces))
for _, ws := range apiWorkspaces {
buildIDs = append(buildIDs, ws.LatestBuild.ID)
}
parameters, err := api.Database.GetWorkspaceBuildParametersByBuildIDs(ctx, buildIDs)
if err != nil {
return nil, err
}
promptsByBuildID := make(map[uuid.UUID]string, len(parameters))
for _, p := range parameters {
if p.Name == codersdk.AITaskPromptParameterName {
promptsByBuildID[p.WorkspaceBuildID] = p.Value
}
}

tasks := make([]codersdk.Task, 0, len(apiWorkspaces))
for _, ws := range apiWorkspaces {
var currentState *codersdk.TaskStateEntry
if ws.LatestAppStatus != nil {
currentState = &codersdk.TaskStateEntry{
Timestamp: ws.LatestAppStatus.CreatedAt,
State: codersdk.TaskState(ws.LatestAppStatus.State),
Message: ws.LatestAppStatus.Message,
URI: ws.LatestAppStatus.URI,
}
}
tasks = append(tasks, codersdk.Task{
ID: ws.ID,
OrganizationID: ws.OrganizationID,
OwnerID: ws.OwnerID,
Name: ws.Name,
TemplateID: ws.TemplateID,
WorkspaceID: uuid.NullUUID{Valid: true, UUID: ws.ID},
CreatedAt: ws.CreatedAt,
UpdatedAt: ws.UpdatedAt,
InitialPrompt: promptsByBuildID[ws.LatestBuild.ID],
Status: ws.LatestBuild.Status,
CurrentState: currentState,
})
}

return tasks, nil
}

// tasksListResponse wraps a list of experimental tasks.
//
// Experimental: Response shape is experimental and may change.
type tasksListResponse struct {
Tasks []codersdk.Task `json:"tasks"`
Count int `json:"count"`
}

// tasksList is an experimental endpoint to list AI tasks by mapping
// workspaces to a task-shaped response.
func (api *API) tasksList(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
apiKey := httpmw.APIKey(r)

// Support standard pagination/filters for workspaces.
page, ok := ParsePagination(rw, r)
if !ok {
return
}
queryStr := r.URL.Query().Get("q")
filter, errs := searchquery.Workspaces(ctx, api.Database, queryStr, page, api.AgentInactiveDisconnectTimeout)
if len(errs) > 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid workspace search query.",
Validations: errs,
})
return
}

// Ensure that we only include AI task workspaces in the results.
filter.HasAITask = sql.NullBool{Valid: true, Bool: true}

if filter.OwnerUsername == "me" || filter.OwnerUsername == "" {
filter.OwnerID = apiKey.UserID
filter.OwnerUsername = ""
}

prepared, err := api.HTTPAuth.AuthorizeSQLFilter(r, policy.ActionRead, rbac.ResourceWorkspace.Type)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error preparing sql filter.",
Detail: err.Error(),
})
return
}

// Order with requester's favorites first, include summary row.
filter.RequesterID = apiKey.UserID
filter.WithSummary = true

workspaceRows, err := api.Database.GetAuthorizedWorkspaces(ctx, filter, prepared)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspaces.",
Detail: err.Error(),
})
return
}
if len(workspaceRows) == 0 {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspaces.",
Detail: "Workspace summary row is missing.",
})
return
}
if len(workspaceRows) == 1 {
httpapi.Write(ctx, rw, http.StatusOK, tasksListResponse{
Tasks: []codersdk.Task{},
Count: 0,
})
return
}

// Skip summary row.
workspaceRows = workspaceRows[:len(workspaceRows)-1]

workspaces := database.ConvertWorkspaceRows(workspaceRows)

// Gather associated data and convert to API workspaces.
data, err := api.workspaceData(ctx, workspaces)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace resources.",
Detail: err.Error(),
})
return
}
apiWorkspaces, err := convertWorkspaces(apiKey.UserID, workspaces, data)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error converting workspaces.",
Detail: err.Error(),
})
return
}

tasks, err := api.tasksFromWorkspaces(ctx, apiWorkspaces)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching task prompts and states.",
Detail: err.Error(),
})
return
}

httpapi.Write(ctx, rw, http.StatusOK, tasksListResponse{
Tasks: tasks,
Count: len(tasks),
})
}

// taskGet is an experimental endpoint to fetch a single AI task by ID
// (workspace ID). It returns a synthesized task response including
// prompt and status.
func (api *API) taskGet(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
apiKey := httpmw.APIKey(r)

idStr := chi.URLParam(r, "id")
taskID, err := uuid.Parse(idStr)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: fmt.Sprintf("Invalid UUID %q for task ID.", idStr),
})
return
}

workspace, err := api.Database.GetWorkspaceByID(ctx, taskID)
if httpapi.Is404Error(err) {
httpapi.ResourceNotFound(rw)
return
}
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace.",
Detail: err.Error(),
})
return
}

data, err := api.workspaceData(ctx, []database.Workspace{workspace})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace resources.",
Detail: err.Error(),
})
return
}
if len(data.builds) == 0 || len(data.templates) == 0 {
httpapi.ResourceNotFound(rw)
return
}
if data.builds[0].HasAITask == nil || !*data.builds[0].HasAITask {
httpapi.ResourceNotFound(rw)
return
}

appStatus := codersdk.WorkspaceAppStatus{}
if len(data.appStatuses) > 0 {
appStatus = data.appStatuses[0]
}

ws, err := convertWorkspace(
apiKey.UserID,
workspace,
data.builds[0],
data.templates[0],
api.Options.AllowWorkspaceRenames,
appStatus,
)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error converting workspace.",
Detail: err.Error(),
})
return
}

tasks, err := api.tasksFromWorkspaces(ctx, []codersdk.Workspace{ws})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching task prompt and state.",
Detail: err.Error(),
})
return
}

httpapi.Write(ctx, rw, http.StatusOK, tasks[0])
}
Loading
Loading