Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
254 changes: 254 additions & 0 deletions coderd/aitasks.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,17 @@
package coderd

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

"github.com/go-chi/chi/v5"
"github.com/google/uuid"
"golang.org/x/xerrors"

"cdr.dev/slog"

Expand All @@ -17,6 +20,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 +191,252 @@ 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. This method enforces that only AI task
// workspaces are given.
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.

// Enforce that only AI task workspaces are given.
for _, ws := range apiWorkspaces {
if ws.LatestBuild.HasAITask == nil || !*ws.LatestBuild.HasAITask {
return nil, xerrors.Errorf("workspace %s is not an AI task workspace", ws.ID)
}
}

// 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
}

// For now, taskID = workspaceID, once we have a task data model in
// the DB, we can change this lookup.
workspaceID := taskID
workspace, err := api.Database.GetWorkspaceByID(ctx, workspaceID)
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
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy