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:
Jaayden Halko
2026-06-18 06:41:47 +01:00
committed by GitHub
parent ea1379d42c
commit bc44cdda75
21 changed files with 1975 additions and 94 deletions
+8 -11
View File
@@ -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
+355 -75
View File
@@ -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
}
+5 -4
View File
@@ -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 == "" {