Skip to content

feat: persist AI task state in template imports & workspace builds #18449

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 22 commits into from
Jun 24, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
9fa7a93
chore: go.mod
dannykopping Jun 18, 2025
d01a050
feat: proto & db changes to persist coder_ai_task state
dannykopping Jun 18, 2025
4383350
feat: store coder_ai_task resource state
dannykopping Jun 18, 2025
957fb08
chore: rename ai_tasks_sidebar_app_id to ai_task_sidebar_app_id
dannykopping Jun 19, 2025
91c8f0c
chore: update to d9b0f892f9c2b3768b9647b089900b0628cb7828, prompt nam…
dannykopping Jun 19, 2025
efe5da9
chore: update provisioner API version comments
dannykopping Jun 19, 2025
5e0d28a
chore: move prompt param validation to ConvertState with other valida…
dannykopping Jun 19, 2025
57da39f
chore: unify tf testdata versions & fix tests
dannykopping Jun 20, 2025
c6c7b5c
chore: dbauthz
dannykopping Jun 20, 2025
4816ed5
chore: fix TestWorkspaceAgentAppHealth test after agent ID field now …
dannykopping Jun 20, 2025
e48df2e
chore: add tests
dannykopping Jun 20, 2025
0b2b732
chore: correct bad rebase
dannykopping Jun 23, 2025
5edee05
chore: undo changes to terraform testdata, so as to not complicate th…
dannykopping Jun 23, 2025
3614a5b
chore: fixing tests
dannykopping Jun 23, 2025
d7e47f9
chore: ensure sidebar app ID must be set when has_ai_task is set
dannykopping Jun 23, 2025
f4111df
chore: test multiple coder_ai_tasks cannot be configured
dannykopping Jun 23, 2025
cc0e290
chore: fix CI
dannykopping Jun 23, 2025
138d865
chore: minor refactor
dannykopping Jun 23, 2025
1aba76b
chore: add clarifying/defensive parentheses
dannykopping Jun 23, 2025
6454ffe
chore: fix conflict
dannykopping Jun 23, 2025
bce4a88
chore: review comments
dannykopping Jun 24, 2025
9a453e8
Merge branch 'main' of github.com:/coder/coder into dk/coder-ai-task-res
dannykopping Jun 24, 2025
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
3 changes: 2 additions & 1 deletion cli/testdata/coder_list_--output_json.golden
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,8 @@
"available": 0,
"most_recently_seen": null
},
"template_version_preset_id": null
"template_version_preset_id": null,
"has_ai_task": false
},
"latest_app_status": null,
"outdated": false,
Expand Down
40 changes: 40 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -4518,6 +4518,28 @@ func (q *querier) UpdateTemplateScheduleByID(ctx context.Context, arg database.U
return update(q.log, q.auth, fetch, q.db.UpdateTemplateScheduleByID)(ctx, arg)
}

func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
// An actor is allowed to update the template version AI task flag if they are authorized to update the template.
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
if err != nil {
return err
}
var obj rbac.Objecter
if !tv.TemplateID.Valid {
obj = rbac.ResourceTemplate.InOrg(tv.OrganizationID)
} else {
tpl, err := q.db.GetTemplateByID(ctx, tv.TemplateID.UUID)
if err != nil {
return err
}
obj = tpl
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
return err
}
return q.db.UpdateTemplateVersionAITaskByJobID(ctx, arg)
}

func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
// An actor is allowed to update the template version if they are authorized to update the template.
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
Expand Down Expand Up @@ -4874,6 +4896,24 @@ func (q *querier) UpdateWorkspaceAutostart(ctx context.Context, arg database.Upd
return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceAutostart)(ctx, arg)
}

func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
if err != nil {
return err
}

workspace, err := q.db.GetWorkspaceByID(ctx, build.WorkspaceID)
if err != nil {
return err
}

err = q.authorizeContext(ctx, policy.ActionUpdate, workspace.RBACObject())
if err != nil {
return err
}
return q.db.UpdateWorkspaceBuildAITaskByID(ctx, arg)
}

// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
Expand Down
52 changes: 52 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1391,6 +1391,24 @@ func (s *MethodTestSuite) TestTemplate() {
ID: t1.ID,
}).Asserts(t1, policy.ActionUpdate)
}))
s.Run("UpdateTemplateVersionAITaskByJobID", s.Subtest(func(db database.Store, check *expects) {
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
o := dbgen.Organization(s.T(), db, database.Organization{})
u := dbgen.User(s.T(), db, database.User{})
_ = dbgen.OrganizationMember(s.T(), db, database.OrganizationMember{OrganizationID: o.ID, UserID: u.ID})
t := dbgen.Template(s.T(), db, database.Template{OrganizationID: o.ID, CreatedBy: u.ID})
job := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{OrganizationID: o.ID})
_ = dbgen.TemplateVersion(s.T(), db, database.TemplateVersion{
OrganizationID: o.ID,
CreatedBy: u.ID,
JobID: job.ID,
TemplateID: uuid.NullUUID{UUID: t.ID, Valid: true},
})
check.Args(database.UpdateTemplateVersionAITaskByJobIDParams{
JobID: job.ID,
HasAITask: sql.NullBool{Bool: true, Valid: true},
}).Asserts(t, policy.ActionUpdate)
}))
s.Run("UpdateTemplateWorkspacesLastUsedAt", s.Subtest(func(db database.Store, check *expects) {
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
t1 := dbgen.Template(s.T(), db, database.Template{})
Expand Down Expand Up @@ -3050,6 +3068,40 @@ func (s *MethodTestSuite) TestWorkspace() {
Deadline: b.Deadline,
}).Asserts(w, policy.ActionUpdate)
}))
s.Run("UpdateWorkspaceBuildAITaskByID", s.Subtest(func(db database.Store, check *expects) {
u := dbgen.User(s.T(), db, database.User{})
o := dbgen.Organization(s.T(), db, database.Organization{})
tpl := dbgen.Template(s.T(), db, database.Template{
OrganizationID: o.ID,
CreatedBy: u.ID,
})
tv := dbgen.TemplateVersion(s.T(), db, database.TemplateVersion{
TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true},
OrganizationID: o.ID,
CreatedBy: u.ID,
})
w := dbgen.Workspace(s.T(), db, database.WorkspaceTable{
TemplateID: tpl.ID,
OrganizationID: o.ID,
OwnerID: u.ID,
})
j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{
Type: database.ProvisionerJobTypeWorkspaceBuild,
})
b := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{
JobID: j.ID,
WorkspaceID: w.ID,
TemplateVersionID: tv.ID,
})
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: b.JobID})
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
app := dbgen.WorkspaceApp(s.T(), db, database.WorkspaceApp{AgentID: agt.ID})
check.Args(database.UpdateWorkspaceBuildAITaskByIDParams{
HasAITask: sql.NullBool{Bool: true, Valid: true},
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
ID: b.ID,
}).Asserts(w, policy.ActionUpdate)
}))
s.Run("SoftDeleteWorkspaceByID", s.Subtest(func(db database.Store, check *expects) {
u := dbgen.User(s.T(), db, database.User{})
o := dbgen.Organization(s.T(), db, database.Organization{})
Expand Down
28 changes: 24 additions & 4 deletions coderd/database/dbgen/dbgen.go
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,9 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
t.Helper()

buildID := takeFirst(orig.ID, uuid.New())
jobID := takeFirst(orig.JobID, uuid.New())
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
sidebarAppID := takeFirst(orig.AITaskSidebarAppID, uuid.NullUUID{})
var build database.WorkspaceBuild
err := db.InTx(func(db database.Store) error {
err := db.InsertWorkspaceBuild(genCtx, database.InsertWorkspaceBuildParams{
Expand All @@ -401,7 +404,7 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
BuildNumber: takeFirst(orig.BuildNumber, 1),
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
JobID: takeFirst(orig.JobID, uuid.New()),
JobID: jobID,
ProvisionerState: takeFirstSlice(orig.ProvisionerState, []byte{}),
Deadline: takeFirst(orig.Deadline, dbtime.Now().Add(time.Hour)),
MaxDeadline: takeFirst(orig.MaxDeadline, time.Time{}),
Expand All @@ -410,7 +413,6 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
UUID: uuid.UUID{},
Valid: false,
}),
HasAITask: orig.HasAITask,
})
if err != nil {
return err
Expand All @@ -424,6 +426,15 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
require.NoError(t, err)
}

if hasAITask.Valid {
require.NoError(t, db.UpdateWorkspaceBuildAITaskByID(genCtx, database.UpdateWorkspaceBuildAITaskByIDParams{
HasAITask: hasAITask,
SidebarAppID: sidebarAppID,
UpdatedAt: dbtime.Now(),
ID: buildID,
}))
}

build, err = db.GetWorkspaceBuildByID(genCtx, buildID)
if err != nil {
return err
Expand Down Expand Up @@ -971,6 +982,8 @@ func ExternalAuthLink(t testing.TB, db database.Store, orig database.ExternalAut

func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVersion) database.TemplateVersion {
var version database.TemplateVersion
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
jobID := takeFirst(orig.JobID, uuid.New())
err := db.InTx(func(db database.Store) error {
versionID := takeFirst(orig.ID, uuid.New())
err := db.InsertTemplateVersion(genCtx, database.InsertTemplateVersionParams{
Expand All @@ -982,15 +995,22 @@ func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVers
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
Message: orig.Message,
Readme: takeFirst(orig.Readme, testutil.GetRandomName(t)),
JobID: takeFirst(orig.JobID, uuid.New()),
JobID: jobID,
CreatedBy: takeFirst(orig.CreatedBy, uuid.New()),
SourceExampleID: takeFirst(orig.SourceExampleID, sql.NullString{}),
HasAITask: orig.HasAITask,
})
if err != nil {
return err
}

if hasAITask.Valid {
require.NoError(t, db.UpdateTemplateVersionAITaskByJobID(genCtx, database.UpdateTemplateVersionAITaskByJobIDParams{
JobID: jobID,
HasAITask: hasAITask,
UpdatedAt: dbtime.Now(),
}))
}

version, err = db.GetTemplateVersionByID(genCtx, versionID)
if err != nil {
return err
Expand Down
51 changes: 49 additions & 2 deletions coderd/database/dbmem/dbmem.go
Original file line number Diff line number Diff line change
Expand Up @@ -9465,7 +9465,6 @@ func (q *FakeQuerier) InsertTemplateVersion(_ context.Context, arg database.Inse
JobID: arg.JobID,
CreatedBy: arg.CreatedBy,
SourceExampleID: arg.SourceExampleID,
HasAITask: arg.HasAITask,
}
q.templateVersions = append(q.templateVersions, version)
return nil
Expand Down Expand Up @@ -10103,7 +10102,6 @@ func (q *FakeQuerier) InsertWorkspaceBuild(_ context.Context, arg database.Inser
MaxDeadline: arg.MaxDeadline,
Reason: arg.Reason,
TemplateVersionPresetID: arg.TemplateVersionPresetID,
HasAITask: arg.HasAITask,
}
q.workspaceBuilds = append(q.workspaceBuilds, workspaceBuild)
return nil
Expand Down Expand Up @@ -11308,6 +11306,26 @@ func (q *FakeQuerier) UpdateTemplateScheduleByID(_ context.Context, arg database
return sql.ErrNoRows
}

func (q *FakeQuerier) UpdateTemplateVersionAITaskByJobID(_ context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
if err := validateDatabaseType(arg); err != nil {
return err
}

q.mutex.Lock()
defer q.mutex.Unlock()

for index, templateVersion := range q.templateVersions {
if templateVersion.JobID != arg.JobID {
continue
}
templateVersion.HasAITask = arg.HasAITask
templateVersion.UpdatedAt = arg.UpdatedAt
q.templateVersions[index] = templateVersion
return nil
}
return sql.ErrNoRows
}

func (q *FakeQuerier) UpdateTemplateVersionByID(_ context.Context, arg database.UpdateTemplateVersionByIDParams) error {
if err := validateDatabaseType(arg); err != nil {
return err
Expand Down Expand Up @@ -12003,6 +12021,35 @@ func (q *FakeQuerier) UpdateWorkspaceAutostart(_ context.Context, arg database.U
return sql.ErrNoRows
}

func (q *FakeQuerier) UpdateWorkspaceBuildAITaskByID(_ context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
if arg.HasAITask.Bool && !arg.SidebarAppID.Valid {
return xerrors.Errorf("ai_task_sidebar_app_id is required when has_ai_task is true")
}
if !arg.HasAITask.Valid && arg.SidebarAppID.Valid {
return xerrors.Errorf("ai_task_sidebar_app_id is can only be set when has_ai_task is true")
}

err := validateDatabaseType(arg)
if err != nil {
return err
}

q.mutex.Lock()
defer q.mutex.Unlock()

for index, workspaceBuild := range q.workspaceBuilds {
if workspaceBuild.ID != arg.ID {
continue
}
workspaceBuild.HasAITask = arg.HasAITask
workspaceBuild.AITaskSidebarAppID = arg.SidebarAppID
workspaceBuild.UpdatedAt = dbtime.Now()
q.workspaceBuilds[index] = workspaceBuild
return nil
}
return sql.ErrNoRows
}

func (q *FakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
if err := validateDatabaseType(arg); err != nil {
return err
Expand Down
14 changes: 14 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

28 changes: 28 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 4 additions & 3 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion coderd/database/foreign_key_constraint.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading