opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
51591e3d59232f0fb0ec07ea8a6fad808a63ae10
coder/coderd/x/chatd/chattool/listtemplates_test.go
T
Jaayden Halko bc44cdda75 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.
2026-06-18 06:41:47 +01:00

1009 lines
35 KiB
Go
Raw Blame History

package chattool_test
import (
"context"
"encoding/json"
"fmt"
"testing"
"time"
"charm.land/fantasy"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"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) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
orgA := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: orgA.ID,
})
orgB := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: orgB.ID,
})
tAlpha := dbgen.Template(t, db, database.Template{
OrganizationID: orgA.ID,
CreatedBy: user.ID,
Name: "alpha",
})
tBeta := dbgen.Template(t, db, database.Template{
OrganizationID: orgB.ID,
CreatedBy: user.ID,
Name: "beta",
})
t.Run("ScopedToOrgA", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
tool := chattool.ListTemplates(db, orgA.ID, chattool.ListTemplatesOptions{
OwnerID: user.ID,
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "org-a", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
require.False(t, resp.IsError)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
templates := result["templates"].([]any)
require.Len(t, templates, 1)
m := templates[0].(map[string]any)
require.Equal(t, tAlpha.ID.String(), m["id"].(string))
})
t.Run("NilOrgReturnsBoth", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{
OwnerID: user.ID,
// Pass uuid.Nil to skip org filtering.
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "nil-org", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
require.False(t, resp.IsError)
var result map[string]any
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) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
// Tool scoped to orgA, but requesting a template in orgB.
tool := chattool.ReadTemplate(db, orgA.ID, chattool.ReadTemplateOptions{
OwnerID: user.ID,
})
input := `{"template_id":"` + tBeta.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "cross-org", Name: "read_template", Input: input})
require.NoError(t, err)
require.True(t, resp.IsError)
require.Contains(t, resp.Content, "not found")
})
t.Run("ReadTemplate_SameOrgAllowed", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
// Tool scoped to orgA, requesting a template in orgA.
tool := chattool.ReadTemplate(db, orgA.ID, chattool.ReadTemplateOptions{
OwnerID: user.ID,
})
input := `{"template_id":"` + tAlpha.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "same-org", Name: "read_template", 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))
tmplInfo := result["template"].(map[string]any)
require.Equal(t, tAlpha.ID.String(), tmplInfo["id"].(string))
})
}
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()
ctx := testutil.Context(t, testutil.WaitLong)
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,
})
t1 := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
Name: "template-alpha",
})
t2 := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
Name: "template-beta",
})
t.Run("ListTemplates", func(t *testing.T) {
t.Run("NoAllowlist", func(t *testing.T) {
tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{
OwnerID: user.ID,
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c1", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
templates := result["templates"].([]any)
require.Len(t, templates, 2)
})
t.Run("EmptyAllowlist", func(t *testing.T) {
tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{} },
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c2", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
templates := result["templates"].([]any)
require.Len(t, templates, 2)
})
t.Run("OneMatch", func(t *testing.T) {
tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c3", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
templates := result["templates"].([]any)
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) {
tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} },
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c4", Name: "list_templates", Input: "{}"})
require.NoError(t, err)
var result map[string]any
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)
})
})
t.Run("ReadTemplate", func(t *testing.T) {
t.Run("Allowed", func(t *testing.T) {
tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
})
input := `{"template_id":"` + t1.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c5", Name: "read_template", 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))
tmplInfo := result["template"].(map[string]any)
require.Equal(t, t1.ID.String(), tmplInfo["id"].(string))
})
t.Run("Disallowed", func(t *testing.T) {
tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} },
})
input := `{"template_id":"` + t2.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c6", Name: "read_template", Input: input})
require.NoError(t, err)
require.True(t, resp.IsError)
require.Contains(t, resp.Content, "not found")
})
t.Run("NoAllowlist", func(t *testing.T) {
tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{
OwnerID: user.ID,
})
input := `{"template_id":"` + t2.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c7", Name: "read_template", Input: input})
require.NoError(t, err)
require.False(t, resp.IsError)
})
})
t.Run("CreateWorkspace", func(t *testing.T) {
t.Run("Allowed", func(t *testing.T) {
// CreateWorkspace requires a real chat row so the existing
// workspace lookup can fall through to creation.
model := seedModelConfig(t, db)
chat, err := db.InsertChat(ctx, database.InsertChatParams{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: model.ID,
Title: "allowed-create",
Status: database.ChatStatusWaiting,
ClientType: database.ChatClientTypeApi,
})
require.NoError(t, err)
createCalled := false
tool := chattool.CreateWorkspace(db, org.ID, chat.ID, chattool.CreateWorkspaceOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} },
CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
createCalled = true
return codersdk.Workspace{}, nil
},
})
input := `{"template_id":"` + t1.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c8a", Name: "create_workspace", Input: input})
require.NoError(t, err)
require.True(t, createCalled, "CreateFn should be called for allowed template")
// We don't assert resp.IsError here because CreateWorkspace
// does additional work (asOwner, workspace lookup) that
// depends on full RBAC setup. The key assertion is that
// the allowlist gate passed and CreateFn was invoked.
_ = resp
})
t.Run("Disallowed", func(t *testing.T) {
var createCalled bool
tool := chattool.CreateWorkspace(db, org.ID, uuid.New(), chattool.CreateWorkspaceOptions{
OwnerID: user.ID,
AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t2.ID: true} },
CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
createCalled = true
t.Fatal("CreateFn should not be called for blocked template")
return codersdk.Workspace{}, nil
},
})
input := `{"template_id":"` + t1.ID.String() + `"}`
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "c8", Name: "create_workspace", Input: input})
require.NoError(t, err)
require.True(t, resp.IsError)
require.Contains(t, resp.Content, "template not available for chat workspaces")
require.False(t, createCalled, "CreateFn should not be called for blocked template")
})
})
}
// 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
}
Reference in New Issue View Git Blame Copy Permalink