mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: rank chat workspace templates (#25037)
closes CODAGT-203
## Summary
`list_templates` now returns a ranked shortlist with a recommendation,
so the chat agent can pick the right template the way a colleague would:
prefer what matches the request, what the user already uses, and what
the rest of the organization uses. Instead of teaching the model an enum
protocol in prompts, every result carries a fixed `next_step`
instruction telling the agent what to do.
## How list_templates works
1. **Fetch**: active, non-deprecated templates in the chat's
organization, filtered by the admin template allowlist, authorized as
the chat owner (no system escalation).
2. **Query relevance** (optional `query` argument): each template
receives the highest tier any of its fields matches, and a higher tier
always outranks a lower one regardless of usage:
| Tier | Match |
|------|-------|
| 4 | name or display name equals the query |
| 3 | name or display name starts with the query |
| 2 | name or display name contains the query |
| 1 | description contains the query (checked only when no name field
matched) |
| 0 | no match; the template is excluded |
Matching is case-insensitive and ignores spaces/hyphens/underscores
(`python gpu` matches `python-gpu`).
3. **Usage signals**: a new `GetTemplateRankingSignalsByOwnerID` query
returns, per template, the owner's active and recently-deleted workspace
counts within a 60-day window, the last in-window usage, and the count
of distinct developers with an active workspace (unclaimed prebuilds
excluded).
4. **Affinity score** (computed in Go, per template, from that
template's signals only):
```text
affinity = 10 x (active + 0.5 x deleted) x 0.5^(days_since_last_use /
14)
+ ln(1 + active_developers)
```
`active`/`deleted` are the owner's in-window workspace counts,
`days_since_last_use` is measured from the most recent in-window usage
(the personal term is zero without in-window usage), and
`active_developers` is the org-wide count. Personal usage carries 10x
the weight of org popularity; the confidence floor is the score of two
active developers (`ln 3`) and the required lead over the runner-up is
`ln 3 - ln 2`.
5. **Rank**: query tier first (when a query is present), then affinity
score, then name/ID for determinism. Results paginate 10 per page with
`next_page` present only when more exist.
## Recommendation contract
The result tells the agent what to do next instead of describing
confidence levels:
- `recommended_template_id` is present only when the top template is a
clear winner: the only available template, a decisive query match, or an
affinity score that clears a floor and leads the runner-up by a derived
margin.
- `next_step` is always present and is one of four fixed sentences: use
the recommendation, ask the user to choose, retry a query that matched
nothing, or report that no templates are available.
Per-template items carry raw evidence (`active_developers`,
`your_workspace_count`, `last_used_by_you`) rather than derived labels.
When signals fail to load, the tool logs and degrades to asking the user
unless the query alone is decisive.
Prompts and the `create_workspace`/`read_template` descriptions
reference the field through the `chattool.NextStepField` constant, so
the instruction lives in one place and cannot drift. `create_workspace`
remains idempotent and allowlist-enforced.
## Authorization
The signals query runs with the chat owner's permissions: reading the
owner's own workspaces plus a template-metadata read for the cross-user
popularity count. dbauthz rejects the call if any requested template is
not readable by the owner (covered by allow and deny method tests).
## Docs
Adds `docs/ai-coder/agents/tools/` explaining how agent tool calls work,
with `list_templates` ranking and the `next_step` contract as the first
documented tools.
This commit is contained in:
@@ -4592,6 +4592,45 @@ func (q *querier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateI
|
||||
return q.db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
|
||||
}
|
||||
|
||||
func (q *querier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) {
|
||||
// The personal signal reads only the owner's own workspaces.
|
||||
workspaceObj := rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String())
|
||||
if arg.OrganizationID != uuid.Nil {
|
||||
workspaceObj = workspaceObj.InOrg(arg.OrganizationID)
|
||||
} else {
|
||||
workspaceObj = workspaceObj.AnyOrganization()
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, workspaceObj); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// The cross-user popularity count is template metadata, not workspace
|
||||
// reads, so it only requires read access to every requested template.
|
||||
if len(arg.TemplateIDs) > 0 {
|
||||
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceTemplate.Type)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
|
||||
}
|
||||
authorizedTemplates, err := q.db.GetAuthorizedTemplates(ctx, database.GetTemplatesWithFilterParams{
|
||||
Deleted: false,
|
||||
OrganizationID: arg.OrganizationID,
|
||||
IDs: arg.TemplateIDs,
|
||||
}, prep)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
authorizedIDs := make(map[uuid.UUID]struct{}, len(authorizedTemplates))
|
||||
for _, template := range authorizedTemplates {
|
||||
authorizedIDs[template.ID] = struct{}{}
|
||||
}
|
||||
for _, templateID := range arg.TemplateIDs {
|
||||
if _, ok := authorizedIDs[templateID]; !ok {
|
||||
return nil, NotAuthorizedError{Err: xerrors.Errorf("not authorized to read template %s", templateID)}
|
||||
}
|
||||
}
|
||||
}
|
||||
return q.db.GetTemplateRankingSignalsByOwnerID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {
|
||||
if err := q.authorizeTemplateInsights(ctx, arg.TemplateIDs); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -3586,6 +3586,49 @@ func (s *MethodTestSuite) TestWorkspace() {
|
||||
// No asserts here because SQLFilter.
|
||||
check.Args(ws.OwnerID, emptyPreparedAuthorized{}).Asserts()
|
||||
}))
|
||||
s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetTemplateRankingSignalsByOwnerIDParams{
|
||||
OwnerID: uuid.New(),
|
||||
OrganizationID: uuid.New(),
|
||||
TemplateIDs: []uuid.UUID{uuid.New()},
|
||||
}
|
||||
dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{
|
||||
Deleted: false,
|
||||
OrganizationID: arg.OrganizationID,
|
||||
IDs: arg.TemplateIDs,
|
||||
}, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes()
|
||||
dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), arg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead)
|
||||
}))
|
||||
s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetTemplateRankingSignalsByOwnerIDParams{
|
||||
OwnerID: uuid.New(),
|
||||
TemplateIDs: []uuid.UUID{uuid.New()},
|
||||
}
|
||||
dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{
|
||||
Deleted: false,
|
||||
IDs: arg.TemplateIDs,
|
||||
}, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes()
|
||||
dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), arg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).AnyOrganization(), policy.ActionRead)
|
||||
}))
|
||||
s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
// Deny path: an unauthorized template ID rejects the call before the
|
||||
// query runs (no query expectation is registered).
|
||||
arg := database.GetTemplateRankingSignalsByOwnerIDParams{
|
||||
OwnerID: uuid.New(),
|
||||
OrganizationID: uuid.New(),
|
||||
TemplateIDs: []uuid.UUID{uuid.New(), uuid.New()},
|
||||
}
|
||||
dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{
|
||||
Deleted: false,
|
||||
OrganizationID: arg.OrganizationID,
|
||||
IDs: arg.TemplateIDs,
|
||||
}, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes()
|
||||
check.Args(arg).
|
||||
Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead).
|
||||
Errors(dbauthz.NotAuthorizedError{Err: xerrors.Errorf("not authorized to read template %s", arg.TemplateIDs[1])})
|
||||
}))
|
||||
s.Run("GetWorkspaceACLByID", s.Mocked(func(dbM *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
ws := testutil.Fake(s.T(), faker, database.Workspace{})
|
||||
dbM.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes()
|
||||
|
||||
@@ -192,6 +192,10 @@ func (s *MethodTestSuite) SubtestWithDB(db database.Store, testCaseF func(db dat
|
||||
testName := s.T().Name()
|
||||
names := strings.Split(testName, "/")
|
||||
methodName := names[len(names)-1]
|
||||
// Repeated subtests get "#NN" suffixes; count them under the base method.
|
||||
if baseMethodName, _, ok := strings.Cut(methodName, "#"); ok {
|
||||
methodName = baseMethodName
|
||||
}
|
||||
s.methodAccounting[methodName]++
|
||||
|
||||
fakeAuthorizer := &coderdtest.FakeAuthorizer{}
|
||||
|
||||
+8
@@ -2978,6 +2978,14 @@ func (m queryMetricsStore) GetTemplatePresetsWithPrebuilds(ctx context.Context,
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetTemplateRankingSignalsByOwnerID(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetTemplateRankingSignalsByOwnerID").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTemplateRankingSignalsByOwnerID").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetTemplateUsageStats(ctx, arg)
|
||||
|
||||
Generated
+15
@@ -5547,6 +5547,21 @@ func (mr *MockStoreMockRecorder) GetTemplatePresetsWithPrebuilds(ctx, templateID
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatePresetsWithPrebuilds", reflect.TypeOf((*MockStore)(nil).GetTemplatePresetsWithPrebuilds), ctx, templateID)
|
||||
}
|
||||
|
||||
// GetTemplateRankingSignalsByOwnerID mocks base method.
|
||||
func (m *MockStore) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetTemplateRankingSignalsByOwnerID", ctx, arg)
|
||||
ret0, _ := ret[0].([]database.GetTemplateRankingSignalsByOwnerIDRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetTemplateRankingSignalsByOwnerID indicates an expected call of GetTemplateRankingSignalsByOwnerID.
|
||||
func (mr *MockStoreMockRecorder) GetTemplateRankingSignalsByOwnerID(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplateRankingSignalsByOwnerID", reflect.TypeOf((*MockStore)(nil).GetTemplateRankingSignalsByOwnerID), ctx, arg)
|
||||
}
|
||||
|
||||
// GetTemplateUsageStats mocks base method.
|
||||
func (m *MockStore) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+6
@@ -776,6 +776,12 @@ type sqlcQuerier interface {
|
||||
// It also returns the number of desired instances for each preset.
|
||||
// If template_id is specified, only template versions associated with that template will be returned.
|
||||
GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]GetTemplatePresetsWithPrebuildsRow, error)
|
||||
// GetTemplateRankingSignalsByOwnerID returns raw template-ranking signals for
|
||||
// one owner: in-window active and recently-deleted workspace counts, the last
|
||||
// in-window usage, and distinct active developers per template. The affinity
|
||||
// score is computed in Go (see listtemplates.go) so the ranking policy and
|
||||
// its confidence thresholds live in one place.
|
||||
GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error)
|
||||
GetTemplateUsageStats(ctx context.Context, arg GetTemplateUsageStatsParams) ([]TemplateUsageStat, error)
|
||||
GetTemplateVersionByID(ctx context.Context, id uuid.UUID) (TemplateVersion, error)
|
||||
GetTemplateVersionByJobID(ctx context.Context, jobID uuid.UUID) (TemplateVersion, error)
|
||||
|
||||
Generated
+111
@@ -36996,6 +36996,117 @@ func (q *sqlQuerier) GetRegularWorkspaceCreateMetrics(ctx context.Context) ([]Ge
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getTemplateRankingSignalsByOwnerID = `-- name: GetTemplateRankingSignalsByOwnerID :many
|
||||
WITH org_usage AS (
|
||||
-- Distinct developers with a non-deleted workspace; the prebuilds system
|
||||
-- user is excluded so unclaimed prebuilds do not inflate popularity.
|
||||
SELECT
|
||||
w.template_id,
|
||||
COUNT(DISTINCT w.owner_id) AS org_devs
|
||||
FROM
|
||||
workspaces w
|
||||
WHERE
|
||||
w.template_id = ANY($1 :: uuid[])
|
||||
AND NOT w.deleted
|
||||
AND w.owner_id != $2 :: uuid
|
||||
AND CASE
|
||||
WHEN $3 :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN
|
||||
w.organization_id = $3
|
||||
ELSE true
|
||||
END
|
||||
GROUP BY
|
||||
w.template_id
|
||||
),
|
||||
user_usage AS (
|
||||
-- The owner's workspaces used within the lookback window, split into
|
||||
-- active and recently-deleted counts.
|
||||
SELECT
|
||||
w.template_id,
|
||||
COUNT(*) FILTER (WHERE NOT w.deleted) AS active_count,
|
||||
COUNT(*) FILTER (WHERE w.deleted) AS deleted_recent_count,
|
||||
MAX(w.last_used_at) :: timestamptz AS last_used_at
|
||||
FROM
|
||||
workspaces w
|
||||
WHERE
|
||||
w.owner_id = $4
|
||||
AND w.template_id = ANY($1 :: uuid[])
|
||||
AND w.last_used_at > $5 :: timestamptz
|
||||
AND CASE
|
||||
WHEN $3 :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN
|
||||
w.organization_id = $3
|
||||
ELSE true
|
||||
END
|
||||
GROUP BY
|
||||
w.template_id
|
||||
)
|
||||
SELECT
|
||||
t.template_id :: uuid AS template_id,
|
||||
COALESCE(u.active_count, 0) :: bigint AS active_count,
|
||||
COALESCE(u.deleted_recent_count, 0) :: bigint AS deleted_recent_count,
|
||||
u.last_used_at,
|
||||
COALESCE(o.org_devs, 0) :: bigint AS org_devs
|
||||
FROM
|
||||
unnest($1 :: uuid[]) AS t(template_id)
|
||||
LEFT JOIN user_usage u ON u.template_id = t.template_id
|
||||
LEFT JOIN org_usage o ON o.template_id = t.template_id
|
||||
`
|
||||
|
||||
type GetTemplateRankingSignalsByOwnerIDParams struct {
|
||||
TemplateIDs []uuid.UUID `db:"template_ids" json:"template_ids"`
|
||||
PrebuildsUserID uuid.UUID `db:"prebuilds_user_id" json:"prebuilds_user_id"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
|
||||
LookbackCutoff time.Time `db:"lookback_cutoff" json:"lookback_cutoff"`
|
||||
}
|
||||
|
||||
type GetTemplateRankingSignalsByOwnerIDRow struct {
|
||||
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
|
||||
ActiveCount int64 `db:"active_count" json:"active_count"`
|
||||
DeletedRecentCount int64 `db:"deleted_recent_count" json:"deleted_recent_count"`
|
||||
LastUsedAt sql.NullTime `db:"last_used_at" json:"last_used_at"`
|
||||
OrgDevs int64 `db:"org_devs" json:"org_devs"`
|
||||
}
|
||||
|
||||
// GetTemplateRankingSignalsByOwnerID returns raw template-ranking signals for
|
||||
// one owner: in-window active and recently-deleted workspace counts, the last
|
||||
// in-window usage, and distinct active developers per template. The affinity
|
||||
// score is computed in Go (see listtemplates.go) so the ranking policy and
|
||||
// its confidence thresholds live in one place.
|
||||
func (q *sqlQuerier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getTemplateRankingSignalsByOwnerID,
|
||||
pq.Array(arg.TemplateIDs),
|
||||
arg.PrebuildsUserID,
|
||||
arg.OrganizationID,
|
||||
arg.OwnerID,
|
||||
arg.LookbackCutoff,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []GetTemplateRankingSignalsByOwnerIDRow
|
||||
for rows.Next() {
|
||||
var i GetTemplateRankingSignalsByOwnerIDRow
|
||||
if err := rows.Scan(
|
||||
&i.TemplateID,
|
||||
&i.ActiveCount,
|
||||
&i.DeletedRecentCount,
|
||||
&i.LastUsedAt,
|
||||
&i.OrgDevs,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, i)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getWorkspaceACLByID = `-- name: GetWorkspaceACLByID :one
|
||||
SELECT
|
||||
group_acl as groups,
|
||||
|
||||
@@ -497,6 +497,65 @@ LEFT JOIN workspaces ON workspaces.template_id = templates.id AND workspaces.del
|
||||
WHERE templates.id = ANY(@template_ids :: uuid[])
|
||||
GROUP BY templates.id;
|
||||
|
||||
-- name: GetTemplateRankingSignalsByOwnerID :many
|
||||
-- GetTemplateRankingSignalsByOwnerID returns raw template-ranking signals for
|
||||
-- one owner: in-window active and recently-deleted workspace counts, the last
|
||||
-- in-window usage, and distinct active developers per template. The affinity
|
||||
-- score is computed in Go (see listtemplates.go) so the ranking policy and
|
||||
-- its confidence thresholds live in one place.
|
||||
WITH org_usage AS (
|
||||
-- Distinct developers with a non-deleted workspace; the prebuilds system
|
||||
-- user is excluded so unclaimed prebuilds do not inflate popularity.
|
||||
SELECT
|
||||
w.template_id,
|
||||
COUNT(DISTINCT w.owner_id) AS org_devs
|
||||
FROM
|
||||
workspaces w
|
||||
WHERE
|
||||
w.template_id = ANY(@template_ids :: uuid[])
|
||||
AND NOT w.deleted
|
||||
AND w.owner_id != @prebuilds_user_id :: uuid
|
||||
AND CASE
|
||||
WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN
|
||||
w.organization_id = @organization_id
|
||||
ELSE true
|
||||
END
|
||||
GROUP BY
|
||||
w.template_id
|
||||
),
|
||||
user_usage AS (
|
||||
-- The owner's workspaces used within the lookback window, split into
|
||||
-- active and recently-deleted counts.
|
||||
SELECT
|
||||
w.template_id,
|
||||
COUNT(*) FILTER (WHERE NOT w.deleted) AS active_count,
|
||||
COUNT(*) FILTER (WHERE w.deleted) AS deleted_recent_count,
|
||||
MAX(w.last_used_at) :: timestamptz AS last_used_at
|
||||
FROM
|
||||
workspaces w
|
||||
WHERE
|
||||
w.owner_id = @owner_id
|
||||
AND w.template_id = ANY(@template_ids :: uuid[])
|
||||
AND w.last_used_at > @lookback_cutoff :: timestamptz
|
||||
AND CASE
|
||||
WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN
|
||||
w.organization_id = @organization_id
|
||||
ELSE true
|
||||
END
|
||||
GROUP BY
|
||||
w.template_id
|
||||
)
|
||||
SELECT
|
||||
t.template_id :: uuid AS template_id,
|
||||
COALESCE(u.active_count, 0) :: bigint AS active_count,
|
||||
COALESCE(u.deleted_recent_count, 0) :: bigint AS deleted_recent_count,
|
||||
u.last_used_at,
|
||||
COALESCE(o.org_devs, 0) :: bigint AS org_devs
|
||||
FROM
|
||||
unnest(@template_ids :: uuid[]) AS t(template_id)
|
||||
LEFT JOIN user_usage u ON u.template_id = t.template_id
|
||||
LEFT JOIN org_usage o ON o.template_id = t.template_id;
|
||||
|
||||
-- name: InsertWorkspace :one
|
||||
INSERT INTO
|
||||
workspaces (
|
||||
|
||||
@@ -10844,7 +10844,7 @@ func TestChatSystemPrompt(t *testing.T) {
|
||||
const workspaceAwareness = `No workspace is attached to this chat yet.
|
||||
Do not create or start a workspace by default. Many requests can be completed using the conversation, provider tools such as web_search when available, or configured external MCP tools.
|
||||
Workspace tools such as execute, read_file, write_file, and edit_files require an attached workspace. Only call create_workspace or start_workspace when the user explicitly asks for a workspace-backed task, or when the task cannot be completed without inspecting, editing, or running files in a workspace.
|
||||
If a workspace is needed, use list_templates and read_template as needed before create_workspace.`
|
||||
If a workspace is needed, use list_templates before create_workspace and follow its next_step. Call read_template only when you need template parameter or preset details.`
|
||||
|
||||
updateChatSystemPrompt := func(t *testing.T, ctx context.Context, req codersdk.UpdateChatSystemPromptRequest) {
|
||||
t.Helper()
|
||||
|
||||
@@ -4587,7 +4587,15 @@ func TestWorkspaceDormant(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should be able to stop a workspace while it is dormant.
|
||||
coderdtest.MustTransitionWorkspace(t, client, workspace.ID, codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop)
|
||||
workspace = coderdtest.MustTransitionWorkspace(t, client, workspace.ID, codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop)
|
||||
testutil.Eventually(ctx, t, func(context.Context) bool {
|
||||
return auditor.Contains(t, database.AuditLog{
|
||||
ResourceID: workspace.LatestBuild.ID,
|
||||
ResourceType: database.ResourceTypeWorkspaceBuild,
|
||||
Action: database.AuditActionStop,
|
||||
StatusCode: http.StatusOK,
|
||||
})
|
||||
}, testutil.IntervalFast)
|
||||
|
||||
// Reset the auditor
|
||||
auditor.ResetLogs()
|
||||
|
||||
@@ -4199,6 +4199,8 @@ func (p *Server) appendRootChatTools(
|
||||
tools = append(tools,
|
||||
chattool.ListTemplates(p.db, opts.chat.OrganizationID, chattool.ListTemplatesOptions{
|
||||
OwnerID: opts.chat.OwnerID,
|
||||
Logger: p.logger,
|
||||
Clock: p.clock,
|
||||
AllowedTemplateIDs: p.chatTemplateAllowlist,
|
||||
}),
|
||||
chattool.ReadTemplate(p.db, opts.chat.OrganizationID, chattool.ReadTemplateOptions{
|
||||
|
||||
@@ -11053,6 +11053,171 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) {
|
||||
"create_workspace for blocked template should be rejected")
|
||||
}
|
||||
|
||||
func TestChatAsksUserWhenListTemplatesRequiresSelection(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
|
||||
var tplCode, tplDocker database.Template
|
||||
var callCount atomic.Int32
|
||||
var sawSelectionRule atomic.Bool
|
||||
var sawSelectionRequiredResult atomic.Bool
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
switch callCount.Add(1) {
|
||||
case 1:
|
||||
promptAndTools := string(req.RawBody)
|
||||
for _, message := range req.Messages {
|
||||
promptAndTools += "\n" + message.Content
|
||||
}
|
||||
if strings.Contains(promptAndTools, "follow its next_step") {
|
||||
sawSelectionRule.Store(true)
|
||||
}
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk("list_templates", `{}`),
|
||||
)
|
||||
case 2:
|
||||
if listTemplatesResultRequiresUserSelection(req.Messages) {
|
||||
sawSelectionRequiredResult.Store(true)
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks(
|
||||
"I found two templates, typescript-alpha and Docker Containers. Which template should I use?",
|
||||
)...,
|
||||
)
|
||||
}
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk("create_workspace",
|
||||
fmt.Sprintf(`{"template_id":%q}`, tplCode.ID.String())),
|
||||
)
|
||||
default:
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("Done.")...,
|
||||
)
|
||||
}
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL)
|
||||
tplCode = dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "code-2",
|
||||
DisplayName: "typescript-alpha",
|
||||
Description: "this is a long description",
|
||||
})
|
||||
tplDocker = dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "docker",
|
||||
DisplayName: "Docker Containers",
|
||||
Description: "Provision Docker containers as Coder workspaces",
|
||||
})
|
||||
|
||||
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
|
||||
cfg.CreateWorkspace = func(
|
||||
context.Context,
|
||||
uuid.UUID,
|
||||
codersdk.CreateWorkspaceRequest,
|
||||
) (codersdk.Workspace, error) {
|
||||
t.Error("create_workspace should not be called when list_templates requires user selection")
|
||||
return codersdk.Workspace{}, xerrors.New("unexpected create_workspace call")
|
||||
}
|
||||
})
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
APIKeyID: testAPIKeyID(t, db, user.ID),
|
||||
Title: "ask-template-selection-test",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("Create a workspace."),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var chatResult database.Chat
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
chatResult = got
|
||||
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
if chatResult.Status == database.ChatStatusError {
|
||||
require.FailNowf(t, "chat run failed", "last_error=%q", chatLastErrorMessage(chatResult.LastError))
|
||||
}
|
||||
|
||||
require.True(t, sawSelectionRule.Load(), "model request should include the next_step selection rule")
|
||||
require.True(t, sawSelectionRequiredResult.Load(), "model should receive a list_templates result requiring user selection")
|
||||
|
||||
messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
|
||||
ChatID: chat.ID,
|
||||
AfterID: 0,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var listTemplatesResult map[string]any
|
||||
var assistantText string
|
||||
var sawCreateWorkspaceResult bool
|
||||
for _, message := range messages {
|
||||
parts, parseErr := chatprompt.ParseContent(message)
|
||||
require.NoError(t, parseErr)
|
||||
for _, part := range parts {
|
||||
switch {
|
||||
case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == "list_templates":
|
||||
require.NoError(t, json.Unmarshal(part.Result, &listTemplatesResult))
|
||||
case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == "create_workspace":
|
||||
sawCreateWorkspaceResult = true
|
||||
case message.Role == database.ChatMessageRoleAssistant && part.Type == codersdk.ChatMessagePartTypeText:
|
||||
assistantText += part.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
require.NotNil(t, listTemplatesResult, "expected list_templates tool result")
|
||||
require.Equal(t, chattool.NextStepAskUser, listTemplatesResult["next_step"])
|
||||
require.NotContains(t, listTemplatesResult, "recommended_template_id")
|
||||
require.Contains(t, listTemplatesResult["templates"], any(map[string]any{
|
||||
"id": tplCode.ID.String(),
|
||||
"name": "code-2",
|
||||
"display_name": "typescript-alpha",
|
||||
"description": "this is a long description",
|
||||
}))
|
||||
require.Contains(t, listTemplatesResult["templates"], any(map[string]any{
|
||||
"id": tplDocker.ID.String(),
|
||||
"name": "docker",
|
||||
"display_name": "Docker Containers",
|
||||
"description": "Provision Docker containers as Coder workspaces",
|
||||
}))
|
||||
require.False(t, sawCreateWorkspaceResult, "agent should ask instead of calling create_workspace")
|
||||
require.Contains(t, assistantText, "Which template should I use?")
|
||||
}
|
||||
|
||||
func listTemplatesResultRequiresUserSelection(messages []chattest.OpenAIMessage) bool {
|
||||
for _, message := range messages {
|
||||
if message.Role != "tool" || !json.Valid([]byte(message.Content)) {
|
||||
continue
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal([]byte(message.Content), &result); err != nil {
|
||||
continue
|
||||
}
|
||||
if result["next_step"] == chattool.NextStepAskUser {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestCreateChatImmediatelyProcessesNewChat verifies that CreateChat
|
||||
// starts processing a new chat without waiting for the acquire ticker
|
||||
// to fire. The ticker interval is set to an hour so it never fires
|
||||
|
||||
@@ -77,8 +77,8 @@ type CreateWorkspaceOptions struct {
|
||||
type createWorkspaceArgs struct {
|
||||
TemplateID string `json:"template_id" description:"The UUIDv4 of the template to create the workspace from. Obtain this from list_templates."`
|
||||
Name string `json:"name,omitempty" description:"The name of the workspace to create. If not provided, a random name will be generated."`
|
||||
Parameters map[string]string `json:"parameters,omitempty" description:"Key-value pairs of template parameters to use when creating the workspace. Obtain available parameters from read_template."`
|
||||
PresetID string `json:"preset_id,omitempty" description:"The UUIDv4 of a template version preset to use. Obtain available presets from read_template. When provided, the preset's parameters are applied automatically and the workspace may claim a prebuilt instance for faster startup."`
|
||||
Parameters map[string]string `json:"parameters,omitempty" description:"Key-value pairs of template parameters to use when creating the workspace. Obtain available parameters from read_template when needed."`
|
||||
PresetID string `json:"preset_id,omitempty" description:"The UUIDv4 of a template version preset to use. Obtain available presets from read_template when needed. When provided, the preset's parameters are applied automatically and the workspace may claim a prebuilt instance for faster startup."`
|
||||
}
|
||||
|
||||
// CreateWorkspace returns a tool that creates a new workspace from a
|
||||
@@ -95,15 +95,12 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option
|
||||
"or when the user explicitly asks for one. Do not use this as a "+
|
||||
"default first step for requests answerable from conversation "+
|
||||
"context, provider tools, or external MCP tools. Requires a "+
|
||||
"template_id (from list_templates). Optionally provide "+
|
||||
"a name and parameter values (from read_template). "+
|
||||
"If no name is given, one will be generated. "+
|
||||
"Provide a preset_id (from read_template) to apply "+
|
||||
"preset parameters and potentially claim a prebuilt "+
|
||||
"workspace for faster startup. "+
|
||||
"This tool is idempotent. If the chat already has a "+
|
||||
"workspace that is building or running, the existing "+
|
||||
"workspace is returned.",
|
||||
"template_id from list_templates; follow its "+NextStepField+" "+
|
||||
"before calling. Optionally provide a name (one is generated if "+
|
||||
"omitted), parameter values, and a preset_id from read_template "+
|
||||
"to apply preset parameters and potentially claim a prebuilt "+
|
||||
"workspace for faster startup. Idempotent: if the chat already "+
|
||||
"has a workspace building or running, it is returned.",
|
||||
func(ctx context.Context, args createWorkspaceArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
if options.CreateFn == nil {
|
||||
return fantasy.NewTextErrorResponse("workspace creator is not configured"), nil
|
||||
|
||||
@@ -5,49 +5,127 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"maps"
|
||||
"math"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/httpmw"
|
||||
"github.com/coder/coder/v2/coderd/rbac"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const listTemplatesPageSize = 10
|
||||
const (
|
||||
listTemplatesPageSize = 10
|
||||
|
||||
// ListTemplatesOptions configures the list_templates tool.
|
||||
// Minimum active developers before organization popularity alone is a
|
||||
// confident recommendation.
|
||||
listTemplatesMinActiveDevelopersForRecommendation = 2
|
||||
|
||||
// Affinity ("frecency") parameters: recency-decayed personal usage plus
|
||||
// log-scaled organization popularity. The score is computed in Go so the
|
||||
// ranking policy and its confidence thresholds live in one place.
|
||||
listTemplatesLookbackDays = 60
|
||||
listTemplatesHalfLife = 14 * 24 * time.Hour
|
||||
listTemplatesPersonalWeight = 10.0
|
||||
listTemplatesOrgWeight = 1.0
|
||||
listTemplatesDeletedWeight = 0.5
|
||||
)
|
||||
|
||||
var (
|
||||
// Confidence floor: organization popularity alone is confident at the
|
||||
// active-developer minimum.
|
||||
minConfidentAffinityScore = listTemplatesOrgWeight * math.Log1p(listTemplatesMinActiveDevelopersForRecommendation)
|
||||
|
||||
// Required rank-1 lead over rank 2, derived so "2 developers versus 1"
|
||||
// recommends while "16 versus 15" does not.
|
||||
minConfidentGap = listTemplatesOrgWeight * (math.Log1p(listTemplatesMinActiveDevelopersForRecommendation) - math.Log1p(listTemplatesMinActiveDevelopersForRecommendation-1))
|
||||
)
|
||||
|
||||
// affinityScoreEpsilon absorbs float rounding at threshold boundaries.
|
||||
const affinityScoreEpsilon = 1e-9
|
||||
|
||||
func affinityScoreAtLeast(score, threshold float64) bool {
|
||||
return score >= threshold-affinityScoreEpsilon
|
||||
}
|
||||
|
||||
// NextStepField is the list_templates result field carrying the instruction
|
||||
// the model should follow next. Tool descriptions and prompts reference it
|
||||
// by name.
|
||||
const NextStepField = "next_step"
|
||||
|
||||
// Next-step instructions returned with every list_templates result.
|
||||
const (
|
||||
NextStepUseRecommended = "Use recommended_template_id with create_workspace. Call read_template first only if you need parameter or preset details."
|
||||
NextStepAskUser = "Do not call create_workspace yet. Ask the user to choose a template, unless they already named one."
|
||||
NextStepNoMatches = "No templates matched the query. Retry without a query or ask the user."
|
||||
NextStepNoTemplates = "No templates are available to this chat. Inform the user."
|
||||
)
|
||||
|
||||
const (
|
||||
queryScoreExactName = 4
|
||||
queryScoreNamePrefix = 3
|
||||
queryScoreNameContains = 2
|
||||
queryScoreDescriptionMatch = 1
|
||||
)
|
||||
|
||||
// ListTemplatesOptions configures the list_templates tool. OwnerID is
|
||||
// required; Clock defaults to a real clock when nil. AllowedTemplateIDs
|
||||
// optionally restricts which templates can be returned.
|
||||
type ListTemplatesOptions struct {
|
||||
OwnerID uuid.UUID
|
||||
Logger slog.Logger
|
||||
Clock quartz.Clock
|
||||
AllowedTemplateIDs func() map[uuid.UUID]bool
|
||||
}
|
||||
|
||||
type listTemplatesArgs struct {
|
||||
Query string `json:"query,omitempty" description:"Optional text to filter templates by name or description."`
|
||||
Page int `json:"page,omitempty" description:"Page number for pagination (starts at 1). Each page returns up to 10 templates."`
|
||||
Query string `json:"query,omitempty" description:"Optional text to filter templates by name, display name, or description."`
|
||||
Page int `json:"page,omitempty" description:"Page number (starts at 1)."`
|
||||
}
|
||||
|
||||
// ListTemplates returns a tool that lists available workspace templates.
|
||||
// The agent uses this to discover templates before creating a workspace.
|
||||
// Results are ordered by number of active developers (most popular first)
|
||||
// and paginated at 10 per page.
|
||||
// db must not be nil.
|
||||
type rankedTemplate struct {
|
||||
Template database.Template
|
||||
QueryScore int
|
||||
Signals templateRankingSignals
|
||||
AffinityScore float64
|
||||
}
|
||||
|
||||
// templateRankingSignals holds the per-template ranking inputs returned by
|
||||
// GetTemplateRankingSignalsByOwnerID.
|
||||
type templateRankingSignals struct {
|
||||
ActiveCount int64
|
||||
DeletedRecentCount int64
|
||||
LastUsedAt time.Time
|
||||
OrgDevs int64
|
||||
}
|
||||
|
||||
// ListTemplates returns a tool that lists workspace templates as a ranked
|
||||
// shortlist, ordered by query relevance, the user's recent usage, and
|
||||
// organization popularity. db must not be nil.
|
||||
func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemplatesOptions) fantasy.AgentTool {
|
||||
clock := options.Clock
|
||||
if clock == nil {
|
||||
clock = quartz.NewReal()
|
||||
}
|
||||
|
||||
return fantasy.NewAgentTool(
|
||||
"list_templates",
|
||||
"List available workspace templates. Optionally filter by a "+
|
||||
"search query matching template name or description. "+
|
||||
"Use this to find a template before creating a workspace. "+
|
||||
"Results are ordered by number of active developers (most popular first). "+
|
||||
"Returns 10 per page. Use the page parameter to paginate through results.",
|
||||
"List workspace templates as a ranked shortlist, optionally filtered "+
|
||||
"by a query matching template name, display name, or description. "+
|
||||
"Follow the "+NextStepField+" field in the result. Returns 10 per "+
|
||||
"page; fetch next_page only when no listed template fits the request.",
|
||||
func(ctx context.Context, args listTemplatesArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
ctx, err := asOwner(ctx, db, options.OwnerID)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
return fantasy.NewTextErrorResponse(xerrors.Errorf("authorize list_templates owner: %w", err).Error()), nil
|
||||
}
|
||||
|
||||
filterParams := database.GetTemplatesWithFilterParams{
|
||||
@@ -58,10 +136,6 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp
|
||||
Valid: true,
|
||||
},
|
||||
}
|
||||
query := strings.TrimSpace(args.Query)
|
||||
if query != "" {
|
||||
filterParams.FuzzyName = query
|
||||
}
|
||||
|
||||
var allowlist map[uuid.UUID]bool
|
||||
if options.AllowedTemplateIDs != nil {
|
||||
@@ -75,78 +149,284 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
|
||||
// Look up active developer counts so we can sort by popularity.
|
||||
templateIDs := make([]uuid.UUID, len(templates))
|
||||
for i, t := range templates {
|
||||
templateIDs[i] = t.ID
|
||||
}
|
||||
ownerCounts := make(map[uuid.UUID]int64)
|
||||
if len(templateIDs) > 0 {
|
||||
rows, countErr := db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs)
|
||||
query := strings.TrimSpace(args.Query)
|
||||
visibleTemplateCount := len(templates)
|
||||
ranked := scoreTemplateCandidates(templates, query)
|
||||
|
||||
if countErr == nil {
|
||||
for _, row := range rows {
|
||||
ownerCounts[row.TemplateID] = row.UniqueOwnersSum
|
||||
}
|
||||
}
|
||||
templateIDs := make([]uuid.UUID, len(ranked))
|
||||
for i, t := range ranked {
|
||||
templateIDs[i] = t.Template.ID
|
||||
}
|
||||
now := clock.Now()
|
||||
signalsByTemplate, signalsErr := loadTemplateRankingSignals(
|
||||
ctx, db, options.OwnerID, organizationID, templateIDs, now,
|
||||
)
|
||||
if signalsErr != nil {
|
||||
options.Logger.Warn(ctx, "failed to load template ranking signals",
|
||||
slog.F("owner_id", options.OwnerID),
|
||||
slog.F("organization_id", organizationID),
|
||||
slog.F("template_count", len(templateIDs)),
|
||||
slog.Error(signalsErr),
|
||||
)
|
||||
}
|
||||
|
||||
// Sort by active developer count descending.
|
||||
slices.SortStableFunc(templates, func(a, b database.Template) int {
|
||||
return cmp.Compare(ownerCounts[b.ID], ownerCounts[a.ID])
|
||||
})
|
||||
// Paginate.
|
||||
for i := range ranked {
|
||||
ranked[i].Signals = signalsByTemplate[ranked[i].Template.ID]
|
||||
ranked[i].AffinityScore = computeAffinityScore(ranked[i].Signals, now)
|
||||
}
|
||||
|
||||
rankTemplates(ranked, query)
|
||||
recommendedID, nextStep := selectTemplateRecommendation(
|
||||
ranked,
|
||||
visibleTemplateCount,
|
||||
signalsErr,
|
||||
)
|
||||
|
||||
page := args.Page
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
totalCount := len(templates)
|
||||
totalPages := (totalCount + listTemplatesPageSize - 1) / listTemplatesPageSize
|
||||
if totalPages == 0 {
|
||||
totalPages = 1
|
||||
}
|
||||
start := (page - 1) * listTemplatesPageSize
|
||||
end := start + listTemplatesPageSize
|
||||
if start > totalCount {
|
||||
start = totalCount
|
||||
}
|
||||
if end > totalCount {
|
||||
end = totalCount
|
||||
}
|
||||
pageTemplates := templates[start:end]
|
||||
totalCount := len(ranked)
|
||||
start := min((page-1)*listTemplatesPageSize, totalCount)
|
||||
end := min(start+listTemplatesPageSize, totalCount)
|
||||
|
||||
items := make([]map[string]any, 0, len(pageTemplates))
|
||||
for _, t := range pageTemplates {
|
||||
item := map[string]any{
|
||||
"id": t.ID.String(),
|
||||
"name": t.Name,
|
||||
"organization_id": t.OrganizationID.String(),
|
||||
}
|
||||
if display := strings.TrimSpace(t.DisplayName); display != "" {
|
||||
item["display_name"] = display
|
||||
}
|
||||
if desc := strings.TrimSpace(t.Description); desc != "" {
|
||||
item["description"] = truncateRunes(desc, 200)
|
||||
}
|
||||
if count, ok := ownerCounts[t.ID]; ok && count > 0 {
|
||||
item["active_developers"] = count
|
||||
}
|
||||
items = append(items, item)
|
||||
items := make([]map[string]any, 0, end-start)
|
||||
for _, t := range ranked[start:end] {
|
||||
items = append(items, templateItem(t))
|
||||
}
|
||||
|
||||
return toolResponse(map[string]any{
|
||||
result := map[string]any{
|
||||
"templates": items,
|
||||
"count": len(items),
|
||||
"page": page,
|
||||
"total_pages": totalPages,
|
||||
"total_count": totalCount,
|
||||
}), nil
|
||||
NextStepField: nextStep,
|
||||
}
|
||||
if end < totalCount {
|
||||
result["next_page"] = page + 1
|
||||
}
|
||||
if recommendedID != uuid.Nil {
|
||||
result["recommended_template_id"] = recommendedID.String()
|
||||
}
|
||||
return toolResponse(result), nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// asOwner sets up a dbauthz context for the given owner so that
|
||||
// subsequent database calls are scoped to what that user can access.
|
||||
func scoreTemplateCandidates(templates []database.Template, query string) []rankedTemplate {
|
||||
candidates := make([]rankedTemplate, 0, len(templates))
|
||||
for _, t := range templates {
|
||||
queryScore := templateQueryScore(t, query)
|
||||
if query != "" && queryScore == 0 {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, rankedTemplate{
|
||||
Template: t,
|
||||
QueryScore: queryScore,
|
||||
})
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
func loadTemplateRankingSignals(
|
||||
ctx context.Context,
|
||||
db database.Store,
|
||||
ownerID uuid.UUID,
|
||||
organizationID uuid.UUID,
|
||||
templateIDs []uuid.UUID,
|
||||
now time.Time,
|
||||
) (map[uuid.UUID]templateRankingSignals, error) {
|
||||
signals := make(map[uuid.UUID]templateRankingSignals)
|
||||
if len(templateIDs) == 0 {
|
||||
return signals, nil
|
||||
}
|
||||
|
||||
// Runs with the owner's permissions; no system escalation. See the
|
||||
// dbauthz GetTemplateRankingSignalsByOwnerID authorization notes.
|
||||
rows, err := db.GetTemplateRankingSignalsByOwnerID(ctx, database.GetTemplateRankingSignalsByOwnerIDParams{
|
||||
TemplateIDs: templateIDs,
|
||||
OwnerID: ownerID,
|
||||
OrganizationID: organizationID,
|
||||
PrebuildsUserID: database.PrebuildsSystemUserID,
|
||||
LookbackCutoff: now.Add(-listTemplatesLookbackDays * 24 * time.Hour),
|
||||
})
|
||||
if err != nil {
|
||||
return signals, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
s := templateRankingSignals{
|
||||
ActiveCount: row.ActiveCount,
|
||||
DeletedRecentCount: row.DeletedRecentCount,
|
||||
OrgDevs: row.OrgDevs,
|
||||
}
|
||||
if row.LastUsedAt.Valid {
|
||||
s.LastUsedAt = row.LastUsedAt.Time
|
||||
}
|
||||
signals[row.TemplateID] = s
|
||||
}
|
||||
return signals, nil
|
||||
}
|
||||
|
||||
// computeAffinityScore folds the raw signals into a single "frecency" score:
|
||||
// recency-decayed personal usage plus log-scaled organization popularity.
|
||||
func computeAffinityScore(s templateRankingSignals, now time.Time) float64 {
|
||||
personal := 0.0
|
||||
if !s.LastUsedAt.IsZero() {
|
||||
count := float64(s.ActiveCount) + listTemplatesDeletedWeight*float64(s.DeletedRecentCount)
|
||||
age := now.Sub(s.LastUsedAt)
|
||||
if age < 0 {
|
||||
age = 0
|
||||
}
|
||||
decay := math.Pow(0.5, float64(age)/float64(listTemplatesHalfLife))
|
||||
personal = listTemplatesPersonalWeight * count * decay
|
||||
}
|
||||
org := listTemplatesOrgWeight * math.Log1p(float64(s.OrgDevs))
|
||||
return personal + org
|
||||
}
|
||||
|
||||
// rankTemplates orders by query relevance (when a query is present), then
|
||||
// affinity score, then name and ID for determinism.
|
||||
func rankTemplates(ranked []rankedTemplate, query string) {
|
||||
slices.SortStableFunc(ranked, func(a, b rankedTemplate) int {
|
||||
if query != "" {
|
||||
if c := cmp.Compare(b.QueryScore, a.QueryScore); c != 0 {
|
||||
return c
|
||||
}
|
||||
}
|
||||
if c := cmp.Compare(b.AffinityScore, a.AffinityScore); c != 0 {
|
||||
return c
|
||||
}
|
||||
if c := cmp.Compare(a.Template.Name, b.Template.Name); c != 0 {
|
||||
return c
|
||||
}
|
||||
return cmp.Compare(a.Template.ID.String(), b.Template.ID.String())
|
||||
})
|
||||
}
|
||||
|
||||
// selectTemplateRecommendation returns the recommended template (uuid.Nil for
|
||||
// none) and the next-step instruction. A decisive query match recommends on
|
||||
// its own; otherwise the affinity score must clear a floor and lead the
|
||||
// runner-up by a margin.
|
||||
func selectTemplateRecommendation(
|
||||
ranked []rankedTemplate,
|
||||
visibleTemplateCount int,
|
||||
rankingSignalsErr error,
|
||||
) (uuid.UUID, string) {
|
||||
if len(ranked) == 0 {
|
||||
if visibleTemplateCount == 0 {
|
||||
return uuid.Nil, NextStepNoTemplates
|
||||
}
|
||||
return uuid.Nil, NextStepNoMatches
|
||||
}
|
||||
|
||||
top := ranked[0]
|
||||
if visibleTemplateCount == 1 && len(ranked) == 1 {
|
||||
return top.Template.ID, NextStepUseRecommended
|
||||
}
|
||||
|
||||
// A decisive query match recommends even when signals failed to load.
|
||||
if top.QueryScore > 0 && (len(ranked) == 1 || top.QueryScore > ranked[1].QueryScore) {
|
||||
return top.Template.ID, NextStepUseRecommended
|
||||
}
|
||||
|
||||
// Beyond a decisive query match, confidence comes from the affinity
|
||||
// score, so a failed signal load means asking the user.
|
||||
if rankingSignalsErr != nil {
|
||||
return uuid.Nil, NextStepAskUser
|
||||
}
|
||||
|
||||
// Query tie: break it with a clear affinity gap.
|
||||
if top.QueryScore > 0 {
|
||||
if len(ranked) > 1 && affinityScoreAtLeast(top.AffinityScore-ranked[1].AffinityScore, minConfidentGap) {
|
||||
return top.Template.ID, NextStepUseRecommended
|
||||
}
|
||||
return uuid.Nil, NextStepAskUser
|
||||
}
|
||||
|
||||
// No query: the affinity score alone decides.
|
||||
if !affinityScoreAtLeast(top.AffinityScore, minConfidentAffinityScore) {
|
||||
return uuid.Nil, NextStepAskUser
|
||||
}
|
||||
if len(ranked) > 1 &&
|
||||
affinityScoreAtLeast(ranked[1].AffinityScore, minConfidentAffinityScore) &&
|
||||
!affinityScoreAtLeast(top.AffinityScore-ranked[1].AffinityScore, minConfidentGap) {
|
||||
return uuid.Nil, NextStepAskUser
|
||||
}
|
||||
return top.Template.ID, NextStepUseRecommended
|
||||
}
|
||||
|
||||
func templateItem(t rankedTemplate) map[string]any {
|
||||
item := map[string]any{
|
||||
"id": t.Template.ID.String(),
|
||||
"name": t.Template.Name,
|
||||
}
|
||||
if display := strings.TrimSpace(t.Template.DisplayName); display != "" {
|
||||
item["display_name"] = display
|
||||
}
|
||||
if desc := strings.TrimSpace(t.Template.Description); desc != "" {
|
||||
item["description"] = truncateRunes(desc, 200)
|
||||
}
|
||||
if t.Signals.OrgDevs > 0 {
|
||||
item["active_developers"] = t.Signals.OrgDevs
|
||||
}
|
||||
if t.Signals.ActiveCount > 0 {
|
||||
item["your_workspace_count"] = t.Signals.ActiveCount
|
||||
}
|
||||
if !t.Signals.LastUsedAt.IsZero() {
|
||||
item["last_used_by_you"] = t.Signals.LastUsedAt.Format(time.RFC3339Nano)
|
||||
}
|
||||
return item
|
||||
}
|
||||
|
||||
func templateQueryScore(t database.Template, query string) int {
|
||||
query = normalizeTemplateSearch(query)
|
||||
queryCompact := compactTemplateSearch(query)
|
||||
if query == "" || queryCompact == "" {
|
||||
return 0
|
||||
}
|
||||
|
||||
best := 0
|
||||
for _, field := range []string{t.Name, t.DisplayName} {
|
||||
best = max(best, nameQueryScore(field, query, queryCompact))
|
||||
}
|
||||
if best > 0 {
|
||||
return best
|
||||
}
|
||||
desc := normalizeTemplateSearch(t.Description)
|
||||
if strings.Contains(desc, query) || strings.Contains(compactTemplateSearch(desc), queryCompact) {
|
||||
return queryScoreDescriptionMatch
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// nameQueryScore returns the relevance tier of a single name-like field:
|
||||
// exact match outranks prefix match, which outranks substring match.
|
||||
func nameQueryScore(field, query, queryCompact string) int {
|
||||
field = normalizeTemplateSearch(field)
|
||||
if field == "" {
|
||||
return 0
|
||||
}
|
||||
fieldCompact := compactTemplateSearch(field)
|
||||
switch {
|
||||
case field == query || fieldCompact == queryCompact:
|
||||
return queryScoreExactName
|
||||
case strings.HasPrefix(field, query) || strings.HasPrefix(fieldCompact, queryCompact):
|
||||
return queryScoreNamePrefix
|
||||
case strings.Contains(field, query) || strings.Contains(fieldCompact, queryCompact):
|
||||
return queryScoreNameContains
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func normalizeTemplateSearch(value string) string {
|
||||
return strings.ToLower(strings.TrimSpace(value))
|
||||
}
|
||||
|
||||
var templateSearchCompactReplacer = strings.NewReplacer(" ", "", "-", "", "_", "")
|
||||
|
||||
func compactTemplateSearch(value string) string {
|
||||
return templateSearchCompactReplacer.Replace(value)
|
||||
}
|
||||
|
||||
// asOwner sets up a dbauthz context scoped to what the owner can access.
|
||||
func asOwner(ctx context.Context, db database.Store, ownerID uuid.UUID) (context.Context, error) {
|
||||
actor, _, err := httpmw.UserRBACSubject(ctx, db, ownerID, rbac.ScopeAll)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
package chattool
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
)
|
||||
|
||||
func TestComputeAffinityScore(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
// No signals at all scores zero.
|
||||
require.Zero(t, computeAffinityScore(templateRankingSignals{}, now))
|
||||
|
||||
// With no personal usage the score collapses to the log-scaled org term.
|
||||
orgOnly := computeAffinityScore(templateRankingSignals{OrgDevs: 3}, now)
|
||||
require.InDelta(t, listTemplatesOrgWeight*math.Log1p(3), orgOnly, 1e-9)
|
||||
|
||||
// Org popularity is monotonic in the developer count.
|
||||
require.Greater(t,
|
||||
computeAffinityScore(templateRankingSignals{OrgDevs: 3}, now),
|
||||
computeAffinityScore(templateRankingSignals{OrgDevs: 1}, now),
|
||||
)
|
||||
|
||||
// Recency decay: the same usage counts more when it is more recent.
|
||||
recent := computeAffinityScore(templateRankingSignals{ActiveCount: 2, LastUsedAt: now.Add(-1 * 24 * time.Hour)}, now)
|
||||
stale := computeAffinityScore(templateRankingSignals{ActiveCount: 2, LastUsedAt: now.Add(-30 * 24 * time.Hour)}, now)
|
||||
require.Greater(t, recent, stale)
|
||||
|
||||
// Deleted workspaces contribute at reduced weight, so the same number of
|
||||
// active workspaces outscores deleted ones.
|
||||
last := now.Add(-1 * time.Hour)
|
||||
activeOnly := computeAffinityScore(templateRankingSignals{ActiveCount: 2, LastUsedAt: last}, now)
|
||||
deletedOnly := computeAffinityScore(templateRankingSignals{DeletedRecentCount: 2, LastUsedAt: last}, now)
|
||||
require.Greater(t, activeOnly, deletedOnly)
|
||||
require.Greater(t, deletedOnly, 0.0)
|
||||
|
||||
// A future last_used_at clamps the age to zero rather than amplifying.
|
||||
future := computeAffinityScore(templateRankingSignals{ActiveCount: 1, LastUsedAt: now.Add(time.Hour)}, now)
|
||||
atNow := computeAffinityScore(templateRankingSignals{ActiveCount: 1, LastUsedAt: now}, now)
|
||||
require.InDelta(t, atNow, future, 1e-9)
|
||||
}
|
||||
|
||||
func TestSelectTemplateRecommendation(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
loadErr := xerrors.New("signals failed to load")
|
||||
|
||||
t.Run("NoTemplatesAvailable", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(nil, 0, nil)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepNoTemplates, next)
|
||||
})
|
||||
|
||||
t.Run("QueryFiltersEverything", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(nil, 2, nil)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepNoMatches, next)
|
||||
})
|
||||
|
||||
t.Run("OnlyAvailable", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
only := uuid.New()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{{Template: database.Template{ID: only}}}, 1, loadErr,
|
||||
)
|
||||
require.Equal(t, only, id)
|
||||
require.Equal(t, NextStepUseRecommended, next)
|
||||
})
|
||||
|
||||
t.Run("DecisiveQueryRecommendsEvenWithLoadError", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
top := uuid.New()
|
||||
for _, err := range []error{nil, loadErr} {
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: top}, QueryScore: queryScoreExactName},
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreDescriptionMatch},
|
||||
}, 2, err,
|
||||
)
|
||||
require.Equal(t, top, id)
|
||||
require.Equal(t, NextStepUseRecommended, next)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("QueryTieBrokenByAffinityGap", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
top := uuid.New()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: top}, QueryScore: queryScoreNamePrefix, AffinityScore: 10, Signals: templateRankingSignals{ActiveCount: 1}},
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix, AffinityScore: 0},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, top, id)
|
||||
require.Equal(t, NextStepUseRecommended, next)
|
||||
})
|
||||
|
||||
t.Run("QueryTieWithSmallGapIsAmbiguous", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix, AffinityScore: 0.1},
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix, AffinityScore: 0},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
|
||||
t.Run("QueryTieWithLoadErrorAsksUser", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix},
|
||||
{Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix},
|
||||
}, 2, loadErr,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryNoSignal", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}},
|
||||
{Template: database.Template{ID: uuid.New()}},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryWeakSignalBelowFloor", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// One active developer scores ln(2), below the ln(3) floor.
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: math.Log1p(1), Signals: templateRankingSignals{OrgDevs: 1}},
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: 0},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryConfidentWhenLeadsRunnerUp", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
top := uuid.New()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: top}, AffinityScore: math.Log1p(3), Signals: templateRankingSignals{OrgDevs: 3}},
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: math.Log1p(1), Signals: templateRankingSignals{OrgDevs: 1}},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, top, id)
|
||||
require.Equal(t, NextStepUseRecommended, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryAmbiguousWhenBothClearFloorAndClose", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: 1.20, Signals: templateRankingSignals{OrgDevs: 2}},
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: 1.15, Signals: templateRankingSignals{OrgDevs: 2}},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryConfidentWhenBothClearFloorWithLargeGap", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
top := uuid.New()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: top}, AffinityScore: 2.0, Signals: templateRankingSignals{OrgDevs: 6}},
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: 1.2, Signals: templateRankingSignals{OrgDevs: 2}},
|
||||
}, 2, nil,
|
||||
)
|
||||
require.Equal(t, top, id)
|
||||
require.Equal(t, NextStepUseRecommended, next)
|
||||
})
|
||||
|
||||
t.Run("NoQueryLoadErrorAsksUser", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
id, next := selectTemplateRecommendation(
|
||||
[]rankedTemplate{
|
||||
{Template: database.Template{ID: uuid.New()}, AffinityScore: math.Log1p(3), Signals: templateRankingSignals{OrgDevs: 3}},
|
||||
{Template: database.Template{ID: uuid.New()}},
|
||||
}, 2, loadErr,
|
||||
)
|
||||
require.Equal(t, uuid.Nil, id)
|
||||
require.Equal(t, NextStepAskUser, next)
|
||||
})
|
||||
}
|
||||
@@ -3,7 +3,9 @@ package chattool_test
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
@@ -15,6 +17,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestListTemplates_OrganizationFilter(t *testing.T) {
|
||||
@@ -82,6 +85,9 @@ func TestListTemplates_OrganizationFilter(t *testing.T) {
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
templates := result["templates"].([]any)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, chattool.NextStepAskUser, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
})
|
||||
|
||||
t.Run("ReadTemplate_CrossOrgRejected", func(t *testing.T) {
|
||||
@@ -121,6 +127,586 @@ func TestListTemplates_OrganizationFilter(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
displayTemplate := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "tpl-42",
|
||||
DisplayName: "Data Science Lab",
|
||||
})
|
||||
descriptionTemplate := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "node-general",
|
||||
Description: "A JavaScript and TypeScript workspace.",
|
||||
})
|
||||
_ = dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unrelated",
|
||||
Description: "A plain Linux workspace.",
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
|
||||
result := runListTemplates(ctx, t, tool, `{"query":"Data Science"}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 1)
|
||||
require.Equal(t, displayTemplate.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, displayTemplate.ID.String(), result["recommended_template_id"])
|
||||
|
||||
result = runListTemplates(ctx, t, tool, `{"query":"TypeScript"}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Len(t, templates, 1)
|
||||
require.Equal(t, descriptionTemplate.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, descriptionTemplate.ID.String(), result["recommended_template_id"])
|
||||
|
||||
result = runListTemplates(ctx, t, tool, `{"query":"-"}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Empty(t, templates)
|
||||
require.Equal(t, chattool.NextStepNoMatches, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
|
||||
result = runListTemplates(ctx, t, tool, `{"query":"does-not-exist"}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Empty(t, templates)
|
||||
require.Equal(t, chattool.NextStepNoMatches, result["next_step"])
|
||||
_, ok = result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestListTemplates_QueryScoreTiers(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
exact := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python",
|
||||
})
|
||||
prefix := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python-alpha",
|
||||
})
|
||||
contains := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "go-python",
|
||||
})
|
||||
description := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "generic-dev",
|
||||
Description: "Python-capable general environment.",
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{"query":"python"}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 4)
|
||||
require.Equal(t, exact.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, prefix.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, contains.ID.String(), templates[2]["id"])
|
||||
require.Equal(t, description.ID.String(), templates[3]["id"])
|
||||
|
||||
hyphenated := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python-gpu",
|
||||
})
|
||||
result = runListTemplates(ctx, t, tool, `{"query":"python gpu"}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Len(t, templates, 1)
|
||||
require.Equal(t, hyphenated.ID.String(), templates[0]["id"])
|
||||
|
||||
descriptionHyphenated := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "ml-tools",
|
||||
Description: "Includes machine-learning libraries.",
|
||||
})
|
||||
result = runListTemplates(ctx, t, tool, `{"query":"machine learning"}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Len(t, templates, 1)
|
||||
require.Equal(t, descriptionHyphenated.ID.String(), templates[0]["id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_RanksAllCandidatesBeforePagination(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
var target database.Template
|
||||
for i := range 11 {
|
||||
tpl := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: fmt.Sprintf("template-%02d", i),
|
||||
})
|
||||
if i == 10 {
|
||||
target = tpl
|
||||
}
|
||||
}
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: target.ID,
|
||||
LastUsedAt: time.Now().Add(-time.Hour),
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 10)
|
||||
require.Equal(t, float64(1), result["page"])
|
||||
require.Equal(t, float64(2), result["next_page"])
|
||||
require.Equal(t, target.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, float64(1), templates[0]["your_workspace_count"])
|
||||
require.NotEmpty(t, templates[0]["last_used_by_you"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, target.ID.String(), result["recommended_template_id"])
|
||||
|
||||
result = runListTemplates(ctx, t, tool, `{"page":2}`)
|
||||
templates = listTemplateItems(t, result)
|
||||
require.Len(t, templates, 1)
|
||||
require.Equal(t, float64(2), result["page"])
|
||||
_, hasNextPage := result["next_page"]
|
||||
require.False(t, hasNextPage)
|
||||
}
|
||||
|
||||
func TestListTemplates_QueryRelevanceOutranksPersonalUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
target := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python-gpu",
|
||||
Description: "GPU workspace.",
|
||||
})
|
||||
used := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "generic-dev",
|
||||
Description: "Python-capable general environment.",
|
||||
})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: used.ID,
|
||||
LastUsedAt: time.Now().Add(-14 * 24 * time.Hour),
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{"query":"python"}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, target.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, used.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, target.ID.String(), result["recommended_template_id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_PersonalUsageBreaksEqualQueryScoreTie(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python-alpha",
|
||||
})
|
||||
used := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "python-beta",
|
||||
})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: used.ID,
|
||||
LastUsedAt: time.Now().Add(-time.Hour),
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{"query":"python"}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, used.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, used.ID.String(), result["recommended_template_id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_OrgPopularityFallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
popular := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "popular-template",
|
||||
})
|
||||
lessPopular := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "less-popular-template",
|
||||
})
|
||||
for range 2 {
|
||||
otherUser := dbgen.User(t, db, database.User{})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: otherUser.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: popular.ID,
|
||||
})
|
||||
}
|
||||
otherUser := dbgen.User(t, db, database.User{})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: otherUser.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: lessPopular.ID,
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, popular.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, float64(2), templates[0]["active_developers"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, popular.ID.String(), result["recommended_template_id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_WeakOrgPopularityDoesNotRecommend(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
usedByOne := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "used-by-one",
|
||||
})
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unused",
|
||||
})
|
||||
otherUser := dbgen.User(t, db, database.User{})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: otherUser.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: usedByOne.ID,
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, usedByOne.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, float64(1), templates[0]["active_developers"])
|
||||
require.Equal(t, chattool.NextStepAskUser, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestListTemplates_StalePersonalUsageDoesNotRecommend(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC)
|
||||
clock.Set(now).MustWait(ctx)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
oldUsage := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "old-usage",
|
||||
})
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unused",
|
||||
})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: oldUsage.ID,
|
||||
LastUsedAt: now.Add(-180 * 24 * time.Hour),
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
Clock: clock,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, oldUsage.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
// 180 days old is outside the 60-day lookback window.
|
||||
_, hasCount := templates[0]["your_workspace_count"]
|
||||
require.False(t, hasCount)
|
||||
require.Equal(t, chattool.NextStepAskUser, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestListTemplates_StaleFrequentPersonalUsageDoesNotRecommend(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC)
|
||||
clock.Set(now).MustWait(ctx)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
staleUsage := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "stale-usage",
|
||||
})
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unused",
|
||||
})
|
||||
// Stale usage decays out of the personal signal despite its frequency.
|
||||
for range 2 {
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: staleUsage.ID,
|
||||
LastUsedAt: now.Add(-180 * 24 * time.Hour),
|
||||
})
|
||||
}
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
Clock: clock,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, staleUsage.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, chattool.NextStepAskUser, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
_, hasCount := templates[0]["your_workspace_count"]
|
||||
require.False(t, hasCount)
|
||||
}
|
||||
|
||||
func TestListTemplates_RecentPersonalUsageRecommends(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC)
|
||||
clock.Set(now).MustWait(ctx)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
recentUsage := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "recent-usage",
|
||||
})
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unused",
|
||||
})
|
||||
// Recent in-window usage is a confident signal.
|
||||
for range 2 {
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: recentUsage.ID,
|
||||
LastUsedAt: now.Add(-2 * 24 * time.Hour),
|
||||
})
|
||||
}
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
Clock: clock,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, recentUsage.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
require.Equal(t, float64(2), templates[0]["your_workspace_count"])
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, recentUsage.ID.String(), result["recommended_template_id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_DeletedRecentPersonalUsageShowsEvidence(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC)
|
||||
clock.Set(now).MustWait(ctx)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
deletedUsage := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "deleted-usage",
|
||||
})
|
||||
unused := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "unused",
|
||||
})
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: deletedUsage.ID,
|
||||
LastUsedAt: now.Add(-2 * 24 * time.Hour),
|
||||
Deleted: true,
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
Clock: clock,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, deletedUsage.ID.String(), templates[0]["id"])
|
||||
require.Equal(t, unused.ID.String(), templates[1]["id"])
|
||||
require.NotEmpty(t, templates[0]["last_used_by_you"])
|
||||
_, hasActiveCount := templates[0]["your_workspace_count"]
|
||||
require.False(t, hasActiveCount)
|
||||
// Recent deleted usage alone clears the confidence floor.
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, deletedUsage.ID.String(), result["recommended_template_id"])
|
||||
}
|
||||
|
||||
func TestListTemplates_AmbiguousTopMatches(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
|
||||
_ = dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "go-alpha",
|
||||
})
|
||||
_ = dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
Name: "go-beta",
|
||||
})
|
||||
|
||||
tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
result := runListTemplates(ctx, t, tool, `{"query":"go"}`)
|
||||
templates := listTemplateItems(t, result)
|
||||
require.Len(t, templates, 2)
|
||||
require.Equal(t, chattool.NextStepAskUser, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
//nolint:tparallel,paralleltest // Subtests share a single DB and run sequentially.
|
||||
func TestTemplateAllowlistEnforcement(t *testing.T) {
|
||||
t.Parallel()
|
||||
@@ -187,6 +773,8 @@ func TestTemplateAllowlistEnforcement(t *testing.T) {
|
||||
require.Len(t, templates, 1)
|
||||
m := templates[0].(map[string]any)
|
||||
require.Equal(t, t1.ID.String(), m["id"].(string))
|
||||
require.Equal(t, chattool.NextStepUseRecommended, result["next_step"])
|
||||
require.Equal(t, t1.ID.String(), result["recommended_template_id"])
|
||||
})
|
||||
|
||||
t.Run("NoMatches", func(t *testing.T) {
|
||||
@@ -201,6 +789,9 @@ func TestTemplateAllowlistEnforcement(t *testing.T) {
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
templates := result["templates"].([]any)
|
||||
require.Empty(t, templates)
|
||||
require.Equal(t, chattool.NextStepNoTemplates, result["next_step"])
|
||||
_, ok := result["recommended_template_id"]
|
||||
require.False(t, ok)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -301,3 +892,117 @@ func TestTemplateAllowlistEnforcement(t *testing.T) {
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetTemplateRankingSignalsByOwnerID exercises the raw SQL signals query:
|
||||
// the lookback window, the active/deleted split, and excluding the prebuilds
|
||||
// system user from the organization developer count.
|
||||
func TestGetTemplateRankingSignalsByOwnerID(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
now := time.Now()
|
||||
lookbackCutoff := now.Add(-60 * 24 * time.Hour)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
otherUser := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
for _, u := range []uuid.UUID{user.ID, otherUser.ID} {
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: u, OrganizationID: org.ID})
|
||||
}
|
||||
|
||||
used := dbgen.Template(t, db, database.Template{OrganizationID: org.ID, CreatedBy: user.ID, Name: "used"})
|
||||
unused := dbgen.Template(t, db, database.Template{OrganizationID: org.ID, CreatedBy: user.ID, Name: "unused"})
|
||||
|
||||
activeLastUsedAt := now.Add(-2 * 24 * time.Hour)
|
||||
deletedLastUsedAt := now.Add(-3 * 24 * time.Hour)
|
||||
// Active, in-window workspace for the requesting user.
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID, OrganizationID: org.ID, TemplateID: used.ID,
|
||||
LastUsedAt: activeLastUsedAt,
|
||||
})
|
||||
// Recently-deleted, in-window workspace for the requesting user.
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID, OrganizationID: org.ID, TemplateID: used.ID,
|
||||
LastUsedAt: deletedLastUsedAt, Deleted: true,
|
||||
})
|
||||
// Outside the lookback window: excluded from in-window counts, still an org dev.
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID, OrganizationID: org.ID, TemplateID: used.ID,
|
||||
LastUsedAt: now.Add(-90 * 24 * time.Hour),
|
||||
})
|
||||
// Another developer's active workspace contributes to org popularity.
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: otherUser.ID, OrganizationID: org.ID, TemplateID: used.ID,
|
||||
LastUsedAt: now.Add(-1 * 24 * time.Hour),
|
||||
})
|
||||
// The prebuilds system user must be excluded from the org developer count.
|
||||
dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: database.PrebuildsSystemUserID, OrganizationID: org.ID, TemplateID: used.ID,
|
||||
LastUsedAt: now.Add(-1 * 24 * time.Hour),
|
||||
})
|
||||
|
||||
rows, err := db.GetTemplateRankingSignalsByOwnerID(ctx, database.GetTemplateRankingSignalsByOwnerIDParams{
|
||||
TemplateIDs: []uuid.UUID{used.ID, unused.ID},
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
PrebuildsUserID: database.PrebuildsSystemUserID,
|
||||
LookbackCutoff: lookbackCutoff,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
byTemplate := make(map[uuid.UUID]database.GetTemplateRankingSignalsByOwnerIDRow, len(rows))
|
||||
for _, row := range rows {
|
||||
byTemplate[row.TemplateID] = row
|
||||
}
|
||||
// The unnest LEFT JOIN returns a row for every requested template.
|
||||
require.Len(t, byTemplate, 2)
|
||||
|
||||
usedRow := byTemplate[used.ID]
|
||||
require.Equal(t, int64(1), usedRow.ActiveCount, "only the in-window active workspace counts")
|
||||
require.Equal(t, int64(1), usedRow.DeletedRecentCount, "the in-window deleted workspace counts")
|
||||
require.Equal(t, int64(2), usedRow.OrgDevs, "user and otherUser count; prebuilds user is excluded")
|
||||
require.True(t, usedRow.LastUsedAt.Valid)
|
||||
require.WithinDuration(t, activeLastUsedAt, usedRow.LastUsedAt.Time, time.Microsecond)
|
||||
|
||||
unusedRow := byTemplate[unused.ID]
|
||||
require.Equal(t, int64(0), unusedRow.ActiveCount)
|
||||
require.Equal(t, int64(0), unusedRow.DeletedRecentCount)
|
||||
require.Equal(t, int64(0), unusedRow.OrgDevs)
|
||||
require.False(t, unusedRow.LastUsedAt.Valid)
|
||||
}
|
||||
|
||||
func runListTemplates(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
tool fantasy.AgentTool,
|
||||
input string,
|
||||
) map[string]any {
|
||||
t.Helper()
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
||||
ID: uuid.NewString(),
|
||||
Name: "list_templates",
|
||||
Input: input,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError)
|
||||
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
return result
|
||||
}
|
||||
|
||||
func listTemplateItems(t *testing.T, result map[string]any) []map[string]any {
|
||||
t.Helper()
|
||||
|
||||
rawTemplates, ok := result["templates"].([]any)
|
||||
require.True(t, ok)
|
||||
templates := make([]map[string]any, 0, len(rawTemplates))
|
||||
for _, raw := range rawTemplates {
|
||||
template, ok := raw.(map[string]any)
|
||||
require.True(t, ok)
|
||||
templates = append(templates, template)
|
||||
}
|
||||
return templates
|
||||
}
|
||||
|
||||
@@ -23,16 +23,17 @@ type readTemplateArgs struct {
|
||||
}
|
||||
|
||||
// ReadTemplate returns a tool that retrieves details about a specific
|
||||
// template, including its configurable rich parameters. The agent
|
||||
// uses this after list_templates and before create_workspace.
|
||||
// template, including its configurable rich parameters. The agent uses
|
||||
// this after list_templates when it needs parameters or presets before
|
||||
// create_workspace.
|
||||
// db must not be nil.
|
||||
func ReadTemplate(db database.Store, organizationID uuid.UUID, options ReadTemplateOptions) fantasy.AgentTool {
|
||||
return fantasy.NewAgentTool(
|
||||
"read_template",
|
||||
"Get details about a workspace template, including its "+
|
||||
"configurable parameters and available presets. Use this "+
|
||||
"after finding a template with list_templates and before "+
|
||||
"creating a workspace with create_workspace.",
|
||||
"after list_templates when you need required parameter "+
|
||||
"details or preset IDs before create_workspace.",
|
||||
func(ctx context.Context, args readTemplateArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
templateIDStr := strings.TrimSpace(args.TemplateID)
|
||||
if templateIDStr == "" {
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
package chatd
|
||||
|
||||
import "github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
||||
|
||||
const defaultSystemPromptPlanPathBlockPlaceholder = "{{CODER_CHAT_PLAN_FILE_PATH_BLOCK}}"
|
||||
|
||||
const workspaceAttachedAwareness = "This chat is attached to a workspace. You can use workspace tools like execute, read_file, write_file, etc."
|
||||
@@ -9,7 +11,7 @@ Do not create or start a workspace by default. Many requests can be completed us
|
||||
Workspace tools such as execute, read_file, write_file, and edit_files require an attached workspace.`
|
||||
|
||||
const workspaceDetachedAwareness = workspaceDetachedAwarenessBase + ` Only call create_workspace or start_workspace when the user explicitly asks for a workspace-backed task, or when the task cannot be completed without inspecting, editing, or running files in a workspace.
|
||||
If a workspace is needed, use list_templates and read_template as needed before create_workspace.`
|
||||
If a workspace is needed, use list_templates before create_workspace and follow its ` + chattool.NextStepField + `. Call read_template only when you need template parameter or preset details.`
|
||||
|
||||
const workspaceDetachedNoCreateAwareness = workspaceDetachedAwarenessBase + ` This delegated chat cannot create or start a workspace. If workspace-backed work is required, report that need to the parent agent instead of trying workspace tools.`
|
||||
|
||||
@@ -104,6 +106,12 @@ Do not start with clarifying questions if the codebase or tools can answer them.
|
||||
Ask the minimum number of questions needed to define the scope together.
|
||||
</collaboration>
|
||||
|
||||
<workspace-template-selection>
|
||||
When no workspace is attached and you need to create one:
|
||||
- Call list_templates with concise search terms from the user's task, then follow its ` + chattool.NextStepField + `: use the recommended template, or ask the user to choose when none is recommended.
|
||||
- Call read_template only when you need parameter or preset details before create_workspace.
|
||||
</workspace-template-selection>
|
||||
|
||||
<planning>
|
||||
Propose a plan when:
|
||||
- The task is too ambiguous to implement with confidence.
|
||||
|
||||
Reference in New Issue
Block a user