opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(coderd): use task data model when creating a new task (#20275)

Browse Source 
Updates coder/internal#976
This commit is contained in:
Mathias Fredriksson
2025-10-23 19:12:09 +03:00
committed by GitHub
parent 0f342ecc04
commit 5c802c2627
27 changed files with 1006 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/aitasks.go
+174 -39
View File
@@ -3,7 +3,6 @@ package coderd
import (
"context"
"database/sql"
"errors"
"fmt"
"net"
"net/http"
@@ -18,6 +17,7 @@ import (
"cdr.dev/slog"
"github.com/coder/coder/v2/coderd/audit"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/httpapi/httperror"
"github.com/coder/coder/v2/coderd/httpmw"
@@ -96,31 +96,54 @@ func (api *API) aiTasksPrompts(rw http.ResponseWriter, r *http.Request) {
// This endpoint creates a new task for the given user.
func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
var (
ctx = r.Context()
apiKey = httpmw.APIKey(r)
auditor = api.Auditor.Load()
mems = httpmw.OrganizationMembersParam(r)
ctx = r.Context()
apiKey = httpmw.APIKey(r)
auditor = api.Auditor.Load()
mems = httpmw.OrganizationMembersParam(r)
taskResourceInfo = audit.AdditionalFields{}
)
if mems.User != nil {
taskResourceInfo.WorkspaceOwner = mems.User.Username
}
aReq, commitAudit := audit.InitRequest[database.TaskTable](rw, &audit.RequestParams{
Audit: *auditor,
Log: api.Logger,
Request: r,
Action: database.AuditActionCreate,
AdditionalFields: taskResourceInfo,
})
defer commitAudit()
var req codersdk.CreateTaskRequest
if !httpapi.Read(ctx, rw, r, &req) {
return
}
hasAITask, err := api.Database.GetTemplateVersionHasAITask(ctx, req.TemplateVersionID)
// Fetch the template version to verify access and whether or not it has an
// AI task.
templateVersion, err := api.Database.GetTemplateVersionByID(ctx, req.TemplateVersionID)
if err != nil {
if errors.Is(err, sql.ErrNoRows) || rbac.IsUnauthorizedError(err) {
httpapi.ResourceNotFound(rw)
if httpapi.Is404Error(err) {
// Avoid using httpapi.ResourceNotFound() here because this is an
// input error and 404 would be confusing.
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Template version not found or you do not have access to this resource",
})
return
}
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching whether the template version has an AI task.",
Message: "Internal error fetching template version.",
Detail: err.Error(),
})
return
}
if !hasAITask {
aReq.UpdateOrganizationID(templateVersion.OrganizationID)
if !templateVersion.HasAITask.Valid || !templateVersion.HasAITask.Bool {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: fmt.Sprintf(`Template does not have required parameter %q`, codersdk.AITaskPromptParameterName),
})
@@ -177,23 +200,12 @@ func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
} else {
// A task can still be created if the caller can read the organization
// member. The organization is required, which can be sourced from the
// template.
// templateVersion.
//
// TODO: This code gets called twice for each workspace build request.
// This is inefficient and costs at most 2 extra RTTs to the DB.
// This can be optimized. It exists as it is now for code simplicity.
// The most common case is to create a workspace for 'Me'. Which does
// not enter this code branch.
template, err := requestTemplate(ctx, createReq, api.Database)
if err != nil {
httperror.WriteResponseError(ctx, rw, err)
return
}
// If the caller can find the organization membership in the same org
// as the template, then they can continue.
orgIndex := slices.IndexFunc(mems.Memberships, func(mem httpmw.OrganizationMember) bool {
return mem.OrganizationID == template.OrganizationID
return mem.OrganizationID == templateVersion.OrganizationID
})
if orgIndex == -1 {
httpapi.ResourceNotFound(rw)
@@ -206,26 +218,148 @@ func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
Username: member.Username,
AvatarURL: member.AvatarURL,
}
// Update workspace owner information for audit in case it changed.
taskResourceInfo.WorkspaceOwner = owner.Username
}
aReq, commitAudit := audit.InitRequest[database.WorkspaceTable](rw, &audit.RequestParams{
Audit: *auditor,
Log: api.Logger,
Request: r,
Action: database.AuditActionCreate,
AdditionalFields: audit.AdditionalFields{
WorkspaceOwner: owner.Username,
// Track insert from preCreateInTX.
var dbTaskTable database.TaskTable
// Ensure an audit log is created for the workspace creation event.
aReqWS, commitAuditWS := audit.InitRequest[database.WorkspaceTable](rw, &audit.RequestParams{
Audit: *auditor,
Log: api.Logger,
Request: r,
Action: database.AuditActionCreate,
AdditionalFields: taskResourceInfo,
OrganizationID: templateVersion.OrganizationID,
})
defer commitAuditWS()
workspace, err := createWorkspace(ctx, aReqWS, apiKey.UserID, api, owner, createReq, r, &createWorkspaceOptions{
// Before creating the workspace, ensure that this task can be created.
preCreateInTX: func(ctx context.Context, tx database.Store) error {
// Create task record in the database before creating the workspace so that
// we can request that the workspace be linked to it after creation.
dbTaskTable, err = tx.InsertTask(ctx, database.InsertTaskParams{
OrganizationID: templateVersion.OrganizationID,
OwnerID: owner.ID,
Name: taskName,
WorkspaceID: uuid.NullUUID{}, // Will be set after workspace creation.
TemplateVersionID: templateVersion.ID,
TemplateParameters: []byte("{}"),
Prompt: req.Input,
CreatedAt: dbtime.Time(api.Clock.Now()),
})
if err != nil {
return httperror.NewResponseError(http.StatusInternalServerError, codersdk.Response{
Message: "Internal error creating task.",
Detail: err.Error(),
})
}
return nil
},
// After the workspace is created, ensure that the task is linked to it.
postCreateInTX: func(ctx context.Context, tx database.Store, workspace database.Workspace) error {
// Update the task record with the workspace ID after creation.
dbTaskTable, err = tx.UpdateTaskWorkspaceID(ctx, database.UpdateTaskWorkspaceIDParams{
ID: dbTaskTable.ID,
WorkspaceID: uuid.NullUUID{
UUID: workspace.ID,
Valid: true,
},
})
if err != nil {
return httperror.NewResponseError(http.StatusInternalServerError, codersdk.Response{
Message: "Internal error updating task.",
Detail: err.Error(),
})
}
return nil
},
})
defer commitAudit()
w, err := createWorkspace(ctx, aReq, apiKey.UserID, api, owner, createReq, r)
if err != nil {
httperror.WriteResponseError(ctx, rw, err)
return
}
task := taskFromWorkspace(w, req.Input)
httpapi.Write(ctx, rw, http.StatusCreated, task)
aReq.New = dbTaskTable
// Fetch the task to get the additional columns from the view.
dbTask, err := api.Database.GetTaskByID(ctx, dbTaskTable.ID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching task.",
Detail: err.Error(),
})
return
}
httpapi.Write(ctx, rw, http.StatusCreated, taskFromDBTaskAndWorkspace(dbTask, workspace))
}
// taskFromDBTaskAndWorkspace creates a codersdk.Task response from the task
// database record and workspace.
func taskFromDBTaskAndWorkspace(dbTask database.Task, ws codersdk.Workspace) codersdk.Task {
var taskAgentLifecycle *codersdk.WorkspaceAgentLifecycle
var taskAgentHealth *codersdk.WorkspaceAgentHealth
// If we have an agent ID from the task, find the agent details in the
// workspace.
if dbTask.WorkspaceAgentID.Valid {
findTaskAgentLoop:
for _, resource := range ws.LatestBuild.Resources {
for _, agent := range resource.Agents {
if agent.ID == dbTask.WorkspaceAgentID.UUID {
taskAgentLifecycle = &agent.LifecycleState
taskAgentHealth = &agent.Health
break findTaskAgentLoop
}
}
}
}
// Ignore 'latest app status' if it is older than the latest build and the
// latest build is a 'start' transition. This ensures that you don't show a
// stale app status from a previous build. For stop transitions, there is
// still value in showing the latest app status.
var currentState *codersdk.TaskStateEntry
if ws.LatestAppStatus != nil {
if ws.LatestBuild.Transition != codersdk.WorkspaceTransitionStart || ws.LatestAppStatus.CreatedAt.After(ws.LatestBuild.CreatedAt) {
currentState = &codersdk.TaskStateEntry{
Timestamp: ws.LatestAppStatus.CreatedAt,
State: codersdk.TaskState(ws.LatestAppStatus.State),
Message: ws.LatestAppStatus.Message,
URI: ws.LatestAppStatus.URI,
}
}
}
return codersdk.Task{
ID: dbTask.ID,
OrganizationID: dbTask.OrganizationID,
OwnerID: dbTask.OwnerID,
OwnerName: ws.OwnerName,
Name: dbTask.Name,
TemplateID: ws.TemplateID,
TemplateVersionID: dbTask.TemplateVersionID,
TemplateName: ws.TemplateName,
TemplateDisplayName: ws.TemplateDisplayName,
TemplateIcon: ws.TemplateIcon,
WorkspaceID: dbTask.WorkspaceID,
WorkspaceBuildNumber: dbTask.WorkspaceBuildNumber.Int32,
WorkspaceStatus: ws.LatestBuild.Status,
WorkspaceAgentID: dbTask.WorkspaceAgentID,
WorkspaceAgentLifecycle: taskAgentLifecycle,
WorkspaceAgentHealth: taskAgentHealth,
WorkspaceAppID: dbTask.WorkspaceAppID,
InitialPrompt: dbTask.Prompt,
Status: codersdk.TaskStatus(dbTask.Status),
CurrentState: currentState,
CreatedAt: dbTask.CreatedAt,
UpdatedAt: ws.UpdatedAt,
}
}
func taskFromWorkspace(ws codersdk.Workspace, initialPrompt string) codersdk.Task {
@@ -253,9 +387,10 @@ func taskFromWorkspace(ws codersdk.Workspace, initialPrompt string) codersdk.Tas
}
}
// Ignore 'latest app status' if it is older than the latest build and the latest build is a 'start' transition.
// This ensures that you don't show a stale app status from a previous build.
// For stop transitions, there is still value in showing the latest app status.
// Ignore 'latest app status' if it is older than the latest build and the
// latest build is a 'start' transition. This ensures that you don't show a
// stale app status from a previous build. For stop transitions, there is
// still value in showing the latest app status.
var currentState *codersdk.TaskStateEntry
if ws.LatestAppStatus != nil {
if ws.LatestBuild.Transition != codersdk.WorkspaceTransitionStart || ws.LatestAppStatus.CreatedAt.After(ws.LatestBuild.CreatedAt) {
@@ -295,7 +430,7 @@ func taskFromWorkspace(ws codersdk.Workspace, initialPrompt string) codersdk.Tas
CreatedAt: ws.CreatedAt,
UpdatedAt: ws.UpdatedAt,
InitialPrompt: initialPrompt,
Status: ws.LatestBuild.Status,
WorkspaceStatus: ws.LatestBuild.Status,
CurrentState: currentState,
}
}
coderd/aitasks_test.go
+225 -8
View File
@@ -240,6 +240,8 @@ func TestTasks(t *testing.T) {
t.Run("List", func(t *testing.T) {
t.Parallel()
t.Skip("TODO(mafredri): Remove, fixed down-stack!")
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
user := coderdtest.CreateFirstUser(t, client)
ctx := testutil.Context(t, testutil.WaitLong)
@@ -266,12 +268,14 @@ func TestTasks(t *testing.T) {
assert.Equal(t, workspace.Name, got.Name, "task name should map from workspace name")
assert.Equal(t, workspace.ID, got.WorkspaceID.UUID, "workspace id should match")
// Status should be populated via app status or workspace status mapping.
assert.NotEmpty(t, got.Status, "task status should not be empty")
assert.NotEmpty(t, got.WorkspaceStatus, "task status should not be empty")
})
t.Run("Get", func(t *testing.T) {
t.Parallel()
t.Skip("TODO(mafredri): Remove, fixed down-stack!")
var (
client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
ctx = testutil.Context(t, testutil.WaitLong)
@@ -315,7 +319,7 @@ func TestTasks(t *testing.T) {
assert.Equal(t, workspace.Name, task.Name, "task name should map from workspace name")
assert.Equal(t, wantPrompt, task.InitialPrompt, "task prompt should match the AI Prompt parameter")
assert.Equal(t, workspace.ID, task.WorkspaceID.UUID, "workspace id should match")
assert.NotEmpty(t, task.Status, "task status should not be empty")
assert.NotEmpty(t, task.WorkspaceStatus, "task status should not be empty")
// Stop the workspace
coderdtest.MustTransitionWorkspace(t, client, workspace.ID, codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop)
@@ -339,6 +343,8 @@ func TestTasks(t *testing.T) {
t.Run("Delete", func(t *testing.T) {
t.Parallel()
t.Skip("TODO(mafredri): Remove, fixed down-stack!")
t.Run("OK", func(t *testing.T) {
t.Parallel()
@@ -354,7 +360,8 @@ func TestTasks(t *testing.T) {
Input: "delete me",
})
require.NoError(t, err)
ws, err := client.Workspace(ctx, task.ID)
require.True(t, task.WorkspaceID.Valid, "task should have a workspace ID")
ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
@@ -363,7 +370,7 @@ func TestTasks(t *testing.T) {
// Poll until the workspace is deleted.
for {
dws, derr := client.DeletedWorkspace(ctx, task.ID)
dws, derr := client.DeletedWorkspace(ctx, task.WorkspaceID.UUID)
if derr == nil && dws.LatestBuild.Status == codersdk.WorkspaceStatusDeleted {
break
}
@@ -434,7 +441,8 @@ func TestTasks(t *testing.T) {
Input: "delete me not",
})
require.NoError(t, err)
ws, err := client.Workspace(ctx, task.ID)
require.True(t, task.WorkspaceID.Valid, "task should have a workspace ID")
ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
@@ -458,6 +466,8 @@ func TestTasks(t *testing.T) {
t.Run("Send", func(t *testing.T) {
t.Parallel()
t.Skip("TODO(mafredri): Remove, fixed down-stack!")
t.Run("IntegrationOK", func(t *testing.T) {
t.Parallel()
@@ -645,6 +655,8 @@ func TestTasks(t *testing.T) {
t.Run("Logs", func(t *testing.T) {
t.Parallel()
t.Skip("TODO(mafredri): Remove, fixed down-stack!")
t.Run("OK", func(t *testing.T) {
t.Parallel()
@@ -791,7 +803,7 @@ func TestTasksCreate(t *testing.T) {
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: "AI Prompt", Type: "string"}},
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
@@ -864,7 +876,7 @@ func TestTasksCreate(t *testing.T) {
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: "AI Prompt", Type: "string"}},
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
@@ -960,7 +972,212 @@ func TestTasksCreate(t *testing.T) {
var sdkErr *codersdk.Error
require.Error(t, err)
require.ErrorAsf(t, err, &sdkErr, "error should be of type *codersdk.Error")
assert.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
assert.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
})
t.Run("TaskTableCreatedAndLinked", func(t *testing.T) {
t.Parallel()
var (
ctx = testutil.Context(t, testutil.WaitShort)
taskPrompt = "Create a REST API"
)
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
user := coderdtest.CreateFirstUser(t, client)
// Create a template with AI task support to test the new task data model.
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
})
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
expClient := codersdk.NewExperimentalClient(client)
task, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
TemplateVersionID: template.ActiveVersionID,
Input: taskPrompt,
})
require.NoError(t, err)
require.True(t, task.WorkspaceID.Valid)
ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
// Verify that the task was created in the tasks table with the correct
// fields. This ensures the data model properly separates task records
// from workspace records.
dbCtx := dbauthz.AsSystemRestricted(ctx)
dbTask, err := db.GetTaskByID(dbCtx, task.ID)
require.NoError(t, err)
assert.Equal(t, user.OrganizationID, dbTask.OrganizationID)
assert.Equal(t, user.UserID, dbTask.OwnerID)
assert.Equal(t, task.Name, dbTask.Name)
assert.True(t, dbTask.WorkspaceID.Valid)
assert.Equal(t, ws.ID, dbTask.WorkspaceID.UUID)
assert.Equal(t, version.ID, dbTask.TemplateVersionID)
assert.Equal(t, taskPrompt, dbTask.Prompt)
assert.False(t, dbTask.DeletedAt.Valid)
// Verify the bidirectional relationship works by looking up the task
// via workspace ID.
dbTaskByWs, err := db.GetTaskByWorkspaceID(dbCtx, ws.ID)
require.NoError(t, err)
assert.Equal(t, dbTask.ID, dbTaskByWs.ID)
})
t.Run("TaskWithCustomName", func(t *testing.T) {
t.Parallel()
var (
ctx = testutil.Context(t, testutil.WaitShort)
taskPrompt = "Build a dashboard"
taskName = "my-custom-task"
)
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
user := coderdtest.CreateFirstUser(t, client)
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
})
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
expClient := codersdk.NewExperimentalClient(client)
task, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
TemplateVersionID: template.ActiveVersionID,
Input: taskPrompt,
Name: taskName,
})
require.NoError(t, err)
require.Equal(t, taskName, task.Name)
// Verify the custom name is preserved in the database record.
dbCtx := dbauthz.AsSystemRestricted(ctx)
dbTask, err := db.GetTaskByID(dbCtx, task.ID)
require.NoError(t, err)
assert.Equal(t, taskName, dbTask.Name)
})
t.Run("MultipleTasksForSameUser", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
user := coderdtest.CreateFirstUser(t, client)
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
})
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
expClient := codersdk.NewExperimentalClient(client)
task1, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
TemplateVersionID: template.ActiveVersionID,
Input: "First task",
Name: "task-1",
})
require.NoError(t, err)
task2, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
TemplateVersionID: template.ActiveVersionID,
Input: "Second task",
Name: "task-2",
})
require.NoError(t, err)
// Verify both tasks are stored independently and can be listed together.
dbCtx := dbauthz.AsSystemRestricted(ctx)
tasks, err := db.ListTasks(dbCtx, database.ListTasksParams{
OwnerID: user.UserID,
OrganizationID: uuid.Nil,
})
require.NoError(t, err)
require.GreaterOrEqual(t, len(tasks), 2)
taskIDs := make(map[uuid.UUID]bool)
for _, task := range tasks {
taskIDs[task.ID] = true
}
assert.True(t, taskIDs[task1.ID], "task1 should be in the list")
assert.True(t, taskIDs[task2.ID], "task2 should be in the list")
})
t.Run("TaskLinkedToCorrectTemplateVersion", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
user := coderdtest.CreateFirstUser(t, client)
version1 := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
})
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version1.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version1.ID)
version2 := coderdtest.UpdateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
ProvisionApply: echo.ApplyComplete,
ProvisionPlan: []*proto.Response{
{Type: &proto.Response_Plan{Plan: &proto.PlanComplete{
Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}},
HasAiTasks: true,
}}},
},
}, template.ID)
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version2.ID)
expClient := codersdk.NewExperimentalClient(client)
// Create a task using version 2 to verify the template_version_id is
// stored correctly.
task, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
TemplateVersionID: version2.ID,
Input: "Use version 2",
})
require.NoError(t, err)
// Verify the task references the correct template version, not just the
// active one.
dbCtx := dbauthz.AsSystemRestricted(ctx)
dbTask, err := db.GetTaskByID(dbCtx, task.ID)
require.NoError(t, err)
assert.Equal(t, version2.ID, dbTask.TemplateVersionID, "task should be linked to version 2")
})
}
coderd/apidoc/docs.go
Generated
+48 -10
View File
@@ -17768,19 +17768,15 @@ const docTemplate = `{
"status": {
"enum": [
"pending",
"starting",
"running",
"stopping",
"stopped",
"failed",
"canceling",
"canceled",
"deleting",
"deleted"
"initializing",
"active",
"paused",
"unknown",
"error"
],
"allOf": [
{
"$ref": "#/definitions/codersdk.WorkspaceStatus"
"$ref": "#/definitions/codersdk.TaskStatus"
}
]
},
@@ -17797,6 +17793,10 @@ const docTemplate = `{
"template_name": {
"type": "string"
},
"template_version_id": {
"type": "string",
"format": "uuid"
},
"updated_at": {
"type": "string",
"format": "date-time"
@@ -17833,6 +17833,25 @@ const docTemplate = `{
"$ref": "#/definitions/uuid.NullUUID"
}
]
},
"workspace_status": {
"enum": [
"pending",
"starting",
"running",
"stopping",
"stopped",
"failed",
"canceling",
"canceled",
"deleting",
"deleted"
],
"allOf": [
{
"$ref": "#/definitions/codersdk.WorkspaceStatus"
}
]
}
}
},
@@ -17917,6 +17936,25 @@ const docTemplate = `{
}
}
},
"codersdk.TaskStatus": {
"type": "string",
"enum": [
"pending",
"initializing",
"active",
"paused",
"unknown",
"error"
],
"x-enum-varnames": [
"TaskStatusPending",
"TaskStatusInitializing",
"TaskStatusActive",
"TaskStatusPaused",
"TaskStatusUnknown",
"TaskStatusError"
]
},
"codersdk.TelemetryConfig": {
"type": "object",
"properties": {
coderd/apidoc/swagger.json
Generated
+48 -10
View File
@@ -16256,19 +16256,15 @@
"status": {
"enum": [
"pending",
"starting",
"running",
"stopping",
"stopped",
"failed",
"canceling",
"canceled",
"deleting",
"deleted"
"initializing",
"active",
"paused",
"unknown",
"error"
],
"allOf": [
{
"$ref": "#/definitions/codersdk.WorkspaceStatus"
"$ref": "#/definitions/codersdk.TaskStatus"
}
]
},
@@ -16285,6 +16281,10 @@
"template_name": {
"type": "string"
},
"template_version_id": {
"type": "string",
"format": "uuid"
},
"updated_at": {
"type": "string",
"format": "date-time"
@@ -16321,6 +16321,25 @@
"$ref": "#/definitions/uuid.NullUUID"
}
]
},
"workspace_status": {
"enum": [
"pending",
"starting",
"running",
"stopping",
"stopped",
"failed",
"canceling",
"canceled",
"deleting",
"deleted"
],
"allOf": [
{
"$ref": "#/definitions/codersdk.WorkspaceStatus"
}
]
}
}
},
@@ -16394,6 +16413,25 @@
}
}
},
"codersdk.TaskStatus": {
"type": "string",
"enum": [
"pending",
"initializing",
"active",
"paused",
"unknown",
"error"
],
"x-enum-varnames": [
"TaskStatusPending",
"TaskStatusInitializing",
"TaskStatusActive",
"TaskStatusPaused",
"TaskStatusUnknown",
"TaskStatusError"
]
},
"codersdk.TelemetryConfig": {
"type": "object",
"properties": {
coderd/database/dbauthz/dbauthz.go
+24
View File
@@ -5009,6 +5009,30 @@ func (q *querier) UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg
return q.db.UpdateTailnetPeerStatusByCoordinator(ctx, arg)
}
func (q *querier) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
// An actor is allowed to update the workspace ID of a task if they are the
// owner of the task and workspace or have the appropriate permissions.
task, err := q.db.GetTaskByID(ctx, arg.ID)
if err != nil {
return database.TaskTable{}, err
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, task.RBACObject()); err != nil {
return database.TaskTable{}, err
}
ws, err := q.db.GetWorkspaceByID(ctx, arg.WorkspaceID.UUID)
if err != nil {
return database.TaskTable{}, err
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, ws.RBACObject()); err != nil {
return database.TaskTable{}, err
}
return q.db.UpdateTaskWorkspaceID(ctx, arg)
}
func (q *querier) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
fetch := func(ctx context.Context, arg database.UpdateTemplateACLByIDParams) (database.Template, error) {
return q.db.GetTemplateByID(ctx, arg.ID)
coderd/database/dbauthz/dbauthz_test.go
+14
View File
@@ -2395,6 +2395,20 @@ func (s *MethodTestSuite) TestTasks() {
check.Args(arg).Asserts(task, policy.ActionUpdate).Returns(database.TaskWorkspaceApp{})
}))
s.Run("UpdateTaskWorkspaceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
task := testutil.Fake(s.T(), faker, database.Task{})
ws := testutil.Fake(s.T(), faker, database.Workspace{})
arg := database.UpdateTaskWorkspaceIDParams{
ID: task.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
}
dbm.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil).AnyTimes()
dbm.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes()
dbm.EXPECT().UpdateTaskWorkspaceID(gomock.Any(), arg).Return(database.TaskTable{}, nil).AnyTimes()
check.Args(arg).Asserts(task, policy.ActionUpdate, ws, policy.ActionUpdate).Returns(database.TaskTable{})
}))
s.Run("GetTaskByWorkspaceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
task := testutil.Fake(s.T(), faker, database.Task{})
task.WorkspaceID = uuid.NullUUID{UUID: uuid.New(), Valid: true}
coderd/database/dbmetrics/querymetrics.go
+7
View File
@@ -3064,6 +3064,13 @@ func (m queryMetricsStore) UpdateTailnetPeerStatusByCoordinator(ctx context.Cont
return r0
}
func (m queryMetricsStore) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
start := time.Now()
r0, r1 := m.s.UpdateTaskWorkspaceID(ctx, arg)
m.queryLatencies.WithLabelValues("UpdateTaskWorkspaceID").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
start := time.Now()
err := m.s.UpdateTemplateACLByID(ctx, arg)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -6578,6 +6578,21 @@ func (mr *MockStoreMockRecorder) UpdateTailnetPeerStatusByCoordinator(ctx, arg a
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTailnetPeerStatusByCoordinator", reflect.TypeOf((*MockStore)(nil).UpdateTailnetPeerStatusByCoordinator), ctx, arg)
}
// UpdateTaskWorkspaceID mocks base method.
func (m *MockStore) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpdateTaskWorkspaceID", ctx, arg)
ret0, _ := ret[0].(database.TaskTable)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpdateTaskWorkspaceID indicates an expected call of UpdateTaskWorkspaceID.
func (mr *MockStoreMockRecorder) UpdateTaskWorkspaceID(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskWorkspaceID", reflect.TypeOf((*MockStore)(nil).UpdateTaskWorkspaceID), ctx, arg)
}
// UpdateTemplateACLByID mocks base method.
func (m *MockStore) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
m.ctrl.T.Helper()
coderd/database/querier.go
+1
View File
@@ -657,6 +657,7 @@ type sqlcQuerier interface {
UpdateProvisionerJobWithCompleteWithStartedAtByID(ctx context.Context, arg UpdateProvisionerJobWithCompleteWithStartedAtByIDParams) error
UpdateReplica(ctx context.Context, arg UpdateReplicaParams) (Replica, error)
UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg UpdateTailnetPeerStatusByCoordinatorParams) error
UpdateTaskWorkspaceID(ctx context.Context, arg UpdateTaskWorkspaceIDParams) (TaskTable, error)
UpdateTemplateACLByID(ctx context.Context, arg UpdateTemplateACLByIDParams) error
UpdateTemplateAccessControlByID(ctx context.Context, arg UpdateTemplateAccessControlByIDParams) error
UpdateTemplateActiveVersionByID(ctx context.Context, arg UpdateTemplateActiveVersionByIDParams) error
coderd/database/querier_test.go
+184
View File
@@ -7540,3 +7540,187 @@ func TestListTasks(t *testing.T) {
})
}
}
func TestUpdateTaskWorkspaceID(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
// Create organization, users, template, and template version.
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
template := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
OrganizationID: org.ID,
TemplateID: uuid.NullUUID{Valid: true, UUID: template.ID},
CreatedBy: user.ID,
})
// Create another template for mismatch test.
template2 := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
tests := []struct {
name string
setupTask func(t *testing.T) database.Task
setupWS func(t *testing.T) database.WorkspaceTable
wantErr bool
wantNoRow bool
}{
{
name: "successful update with matching template",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: false,
},
{
name: "task already has workspace_id",
setupTask: func(t *testing.T) database.Task {
existingWS := dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{Valid: true, UUID: existingWS.ID},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: true, // No row should be returned because WHERE condition fails.
},
{
name: "template mismatch between task and workspace",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{}, // NULL workspace_id
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template2.ID, // Different template, JOIN will fail.
})
},
wantErr: false,
wantNoRow: true, // No row should be returned because JOIN condition fails.
},
{
name: "task does not exist",
setupTask: func(t *testing.T) database.Task {
return database.Task{
ID: uuid.New(), // Non-existent task ID.
}
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: true,
},
{
name: "workspace does not exist",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return database.WorkspaceTable{
ID: uuid.New(), // Non-existent workspace ID.
}
},
wantErr: false,
wantNoRow: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
task := tt.setupTask(t)
workspace := tt.setupWS(t)
updatedTask, err := db.UpdateTaskWorkspaceID(ctx, database.UpdateTaskWorkspaceIDParams{
ID: task.ID,
WorkspaceID: uuid.NullUUID{Valid: true, UUID: workspace.ID},
})
if tt.wantErr {
require.Error(t, err)
return
}
if tt.wantNoRow {
require.ErrorIs(t, err, sql.ErrNoRows)
return
}
require.NoError(t, err)
require.Equal(t, task.ID, updatedTask.ID)
require.True(t, updatedTask.WorkspaceID.Valid)
require.Equal(t, workspace.ID, updatedTask.WorkspaceID.UUID)
require.Equal(t, task.OrganizationID, updatedTask.OrganizationID)
require.Equal(t, task.OwnerID, updatedTask.OwnerID)
require.Equal(t, task.Name, updatedTask.Name)
require.Equal(t, task.TemplateVersionID, updatedTask.TemplateVersionID)
// Verify the update persisted by fetching the task again.
fetchedTask, err := db.GetTaskByID(ctx, task.ID)
require.NoError(t, err)
require.True(t, fetchedTask.WorkspaceID.Valid)
require.Equal(t, workspace.ID, fetchedTask.WorkspaceID.UUID)
})
}
}
coderd/database/queries.sql.go
+43
View File
@@ -12730,6 +12730,49 @@ func (q *sqlQuerier) ListTasks(ctx context.Context, arg ListTasksParams) ([]Task
return items, nil
}
const updateTaskWorkspaceID = `-- name: UpdateTaskWorkspaceID :one
UPDATE
tasks
SET
workspace_id = $2
FROM
workspaces w
JOIN
template_versions tv
ON
tv.template_id = w.template_id
WHERE
tasks.id = $1
AND tasks.workspace_id IS NULL
AND w.id = $2
AND tv.id = tasks.template_version_id
RETURNING
tasks.id, tasks.organization_id, tasks.owner_id, tasks.name, tasks.workspace_id, tasks.template_version_id, tasks.template_parameters, tasks.prompt, tasks.created_at, tasks.deleted_at
`
type UpdateTaskWorkspaceIDParams struct {
ID uuid.UUID `db:"id" json:"id"`
WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"`
}
func (q *sqlQuerier) UpdateTaskWorkspaceID(ctx context.Context, arg UpdateTaskWorkspaceIDParams) (TaskTable, error) {
row := q.db.QueryRowContext(ctx, updateTaskWorkspaceID, arg.ID, arg.WorkspaceID)
var i TaskTable
err := row.Scan(
&i.ID,
&i.OrganizationID,
&i.OwnerID,
&i.Name,
&i.WorkspaceID,
&i.TemplateVersionID,
&i.TemplateParameters,
&i.Prompt,
&i.CreatedAt,
&i.DeletedAt,
)
return i, err
}
const upsertTaskWorkspaceApp = `-- name: UpsertTaskWorkspaceApp :one
INSERT INTO task_workspace_apps
(task_id, workspace_build_number, workspace_agent_id, workspace_app_id)
coderd/database/queries/tasks.sql
+19
View File
@@ -5,6 +5,25 @@ VALUES
(gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7, $8)
RETURNING *;
-- name: UpdateTaskWorkspaceID :one
UPDATE
tasks
SET
workspace_id = $2
FROM
workspaces w
JOIN
template_versions tv
ON
tv.template_id = w.template_id
WHERE
tasks.id = $1
AND tasks.workspace_id IS NULL
AND w.id = $2
AND tv.id = tasks.template_version_id
RETURNING
tasks.*;
-- name: UpsertTaskWorkspaceApp :one
INSERT INTO task_workspace_apps
(task_id, workspace_build_number, workspace_agent_id, workspace_app_id)
coderd/workspaces.go
+35 -2
View File
@@ -388,7 +388,7 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
AvatarURL: member.AvatarURL,
}
w, err := createWorkspace(ctx, aReq, apiKey.UserID, api, owner, req, r)
w, err := createWorkspace(ctx, aReq, apiKey.UserID, api, owner, req, r, nil)
if err != nil {
httperror.WriteResponseError(ctx, rw, err)
return
@@ -484,7 +484,7 @@ func (api *API) postUserWorkspaces(rw http.ResponseWriter, r *http.Request) {
defer commitAudit()
w, err := createWorkspace(ctx, aReq, apiKey.UserID, api, owner, req, r)
w, err := createWorkspace(ctx, aReq, apiKey.UserID, api, owner, req, r, nil)
if err != nil {
httperror.WriteResponseError(ctx, rw, err)
return
@@ -499,6 +499,15 @@ type workspaceOwner struct {
AvatarURL string
}
type createWorkspaceOptions struct {
// preCreateInTX is a function that is called within the transaction, before
// the workspace is created.
preCreateInTX func(ctx context.Context, tx database.Store) error
// postCreateInTX is a function that is called within the transaction, after
// the workspace is created but before the workspace build is created.
postCreateInTX func(ctx context.Context, tx database.Store, workspace database.Workspace) error
}
func createWorkspace(
ctx context.Context,
auditReq *audit.Request[database.WorkspaceTable],
@@ -507,7 +516,12 @@ func createWorkspace(
owner workspaceOwner,
req codersdk.CreateWorkspaceRequest,
r *http.Request,
opts *createWorkspaceOptions,
) (codersdk.Workspace, error) {
if opts == nil {
opts = &createWorkspaceOptions{}
}
template, err := requestTemplate(ctx, req, api.Database)
if err != nil {
return codersdk.Workspace{}, err
@@ -636,6 +650,16 @@ func createWorkspace(
claimedWorkspace *database.Workspace
)
// If a preCreate hook is provided, execute it before creating or
// claiming the workspace. This can be used to perform additional
// setup or validation before the workspace is created (e.g. task
// creation).
if opts.preCreateInTX != nil {
if err := opts.preCreateInTX(ctx, db); err != nil {
return xerrors.Errorf("workspace preCreate failed: %w", err)
}
}
// Use injected Clock to allow time mocking in tests
now := dbtime.Time(api.Clock.Now())
@@ -729,6 +753,15 @@ func createWorkspace(
return xerrors.Errorf("get workspace by ID: %w", err)
}
// If the postCreate hook is provided, execute it. This can be used to
// perform additional actions after the workspace has been created, like
// linking the workspace to a task.
if opts.postCreateInTX != nil {
if err := opts.postCreateInTX(ctx, db, workspace); err != nil {
return xerrors.Errorf("workspace postCreate failed: %w", err)
}
}
builder := wsbuilder.New(workspace, database.WorkspaceTransitionStart, *api.BuildUsageChecker.Load()).
Reason(database.BuildReasonInitiator).
Initiator(initiatorID).