From 35f980a47712cf5056b83c8a0df437e68850e899 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 7 May 2026 12:06:06 +0000 Subject: [PATCH 01/24] feat(coderd): rank chat workspace templates --- coderd/database/dbauthz/dbauthz.go | 13 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 + coderd/database/querier.go | 1 + coderd/database/queries.sql.go | 54 +++ coderd/database/queries/workspaces.sql | 18 + coderd/x/chatd/chattool/createworkspace.go | 15 +- coderd/x/chatd/chattool/listtemplates.go | 339 +++++++++++++++--- coderd/x/chatd/chattool/listtemplates_test.go | 266 ++++++++++++++ coderd/x/chatd/chattool/readtemplate.go | 11 +- coderd/x/chatd/prompt.go | 8 + 11 files changed, 682 insertions(+), 66 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 9badded7e03..ac795a8d273 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5077,6 +5077,19 @@ func (q *querier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, return q.db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs) } +func (q *querier) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { + obj := rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()) + if arg.OrganizationID != uuid.Nil { + obj = obj.InOrg(arg.OrganizationID) + } else { + obj = obj.AnyOrganization() + } + if err := q.authorizeContext(ctx, policy.ActionRead, obj); err != nil { + return nil, err + } + return q.db.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg) +} + func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceWorkspace.Type) if err != nil { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 125e86b2a4c..1d46b092fd1 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3440,6 +3440,14 @@ func (m queryMetricsStore) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context return r0, r1 } +func (m queryMetricsStore) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { + start := time.Now() + r0, r1 := m.s.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg) + m.queryLatencies.WithLabelValues("GetWorkspaceUsageGroupedByTemplateIDForOwner").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceUsageGroupedByTemplateIDForOwner").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { start := time.Now() r0, r1 := m.s.GetWorkspaces(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index bfb29d8559b..2e0fcd42d41 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -6437,6 +6437,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUniqueOwnerCountByTemplateIDs", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUniqueOwnerCountByTemplateIDs), ctx, templateIds) } +// GetWorkspaceUsageGroupedByTemplateIDForOwner mocks base method. +func (m *MockStore) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWorkspaceUsageGroupedByTemplateIDForOwner", ctx, arg) + ret0, _ := ret[0].([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWorkspaceUsageGroupedByTemplateIDForOwner indicates an expected call of GetWorkspaceUsageGroupedByTemplateIDForOwner. +func (mr *MockStoreMockRecorder) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUsageGroupedByTemplateIDForOwner", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUsageGroupedByTemplateIDForOwner), ctx, arg) +} + // GetWorkspaces mocks base method. func (m *MockStore) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 795c9a7af19..02b8212c70f 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -833,6 +833,7 @@ type sqlcQuerier interface { GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceResource, error) GetWorkspaceResourcesCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceResource, error) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, templateIds []uuid.UUID) ([]GetWorkspaceUniqueOwnerCountByTemplateIDsRow, error) + GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) // build_params is used to filter by build parameters if present. // It has to be a CTE because the set returning function 'unnest' cannot // be used in a WHERE clause. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index c18e34bb19c..cbefcfff43e 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -32389,6 +32389,60 @@ func (q *sqlQuerier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Conte return items, nil } +const getWorkspaceUsageGroupedByTemplateIDForOwner = `-- name: GetWorkspaceUsageGroupedByTemplateIDForOwner :many +SELECT + template_id, + COUNT(*) AS workspace_count, + MAX(last_used_at)::timestamptz AS last_used_at +FROM + workspaces +WHERE + owner_id = $1 + AND deleted = false + AND CASE + WHEN $2 :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN + organization_id = $2 + ELSE true + END + AND template_id = ANY($3 :: uuid[]) +GROUP BY template_id +` + +type GetWorkspaceUsageGroupedByTemplateIDForOwnerParams struct { + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + TemplateIDs []uuid.UUID `db:"template_ids" json:"template_ids"` +} + +type GetWorkspaceUsageGroupedByTemplateIDForOwnerRow struct { + TemplateID uuid.UUID `db:"template_id" json:"template_id"` + WorkspaceCount int64 `db:"workspace_count" json:"workspace_count"` + LastUsedAt time.Time `db:"last_used_at" json:"last_used_at"` +} + +func (q *sqlQuerier) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceUsageGroupedByTemplateIDForOwner, arg.OwnerID, arg.OrganizationID, pq.Array(arg.TemplateIDs)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetWorkspaceUsageGroupedByTemplateIDForOwnerRow + for rows.Next() { + var i GetWorkspaceUsageGroupedByTemplateIDForOwnerRow + if err := rows.Scan(&i.TemplateID, &i.WorkspaceCount, &i.LastUsedAt); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getWorkspaces = `-- name: GetWorkspaces :many WITH build_params AS ( diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql index 5269ea8fba5..972f4a54271 100644 --- a/coderd/database/queries/workspaces.sql +++ b/coderd/database/queries/workspaces.sql @@ -497,6 +497,24 @@ LEFT JOIN workspaces ON workspaces.template_id = templates.id AND workspaces.del WHERE templates.id = ANY(@template_ids :: uuid[]) GROUP BY templates.id; +-- name: GetWorkspaceUsageGroupedByTemplateIDForOwner :many +SELECT + template_id, + COUNT(*) AS workspace_count, + MAX(last_used_at)::timestamptz AS last_used_at +FROM + workspaces +WHERE + owner_id = @owner_id + AND deleted = false + AND CASE + WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN + organization_id = @organization_id + ELSE true + END + AND template_id = ANY(@template_ids :: uuid[]) +GROUP BY template_id; + -- name: InsertWorkspace :one INSERT INTO workspaces ( diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index 4a95a6e639a..eb58328c624 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -73,10 +73,10 @@ 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."` + TemplateID string `json:"template_id" description:"The UUIDv4 of the template to create the workspace from. Obtain this from list_templates recommended_template_id or a ranked template."` 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 @@ -89,10 +89,11 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option return fantasy.NewAgentTool( "create_workspace", "Create a new workspace from a template. 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 "+ + "template_id from list_templates. Use recommended_template_id "+ + "or rank 1 when list_templates reports a confident choice. "+ + "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 "+ diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 3c6d31c1b02..55abb3ee4b9 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -7,6 +7,7 @@ import ( "maps" "slices" "strings" + "time" "charm.land/fantasy" "github.com/google/uuid" @@ -20,6 +21,13 @@ import ( const listTemplatesPageSize = 10 +const ( + listTemplatesHintOnlyAvailable = "only_available_template" + listTemplatesHintHighConfidence = "high_confidence_recommendation" + listTemplatesHintAmbiguous = "ambiguous_top_matches" + listTemplatesHintNoConfidence = "no_confident_match" +) + // ListTemplatesOptions configures the list_templates tool. type ListTemplatesOptions struct { OwnerID uuid.UUID @@ -27,23 +35,39 @@ type ListTemplatesOptions struct { } 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 for pagination (starts at 1). Each page returns up to 10 ranked templates."` +} + +type rankedTemplate struct { + Template database.Template + QueryScore int + ActiveDevelopers int64 + Usage templateUsage + Rank int +} + +type templateUsage struct { + WorkspaceCount int64 + LastUsedAt time.Time } // 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. +// Results are ranked before pagination using query relevance, current-user +// usage, and organization-wide popularity. // db must not be nil. func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemplatesOptions) fantasy.AgentTool { 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 available workspace templates as a ranked shortlist. "+ + "Optionally provide a search query matching template name, "+ + "display name, or description. Use recommended_template_id "+ + "or rank 1 as the default choice when selection_hint is "+ + "only_available_template or high_confidence_recommendation. "+ + "Do not paginate unless the returned templates do not fit the "+ + "request, selection_hint reports ambiguity or no confident match, "+ + "or the user asked to browse templates. Returns 10 per page.", func(ctx context.Context, args listTemplatesArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { ctx, err := asOwner(ctx, db, options.OwnerID) if err != nil { @@ -58,10 +82,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,32 +95,35 @@ 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 + query := strings.TrimSpace(args.Query) + visibleTemplateCount := len(templates) + ranked := candidateRankedTemplates(templates, query) + + templateIDs := make([]uuid.UUID, len(ranked)) + for i, t := range ranked { + templateIDs[i] = t.Template.ID } - ownerCounts := make(map[uuid.UUID]int64) - if len(templateIDs) > 0 { - rows, countErr := db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs) - - if countErr == nil { - for _, row := range rows { - ownerCounts[row.TemplateID] = row.UniqueOwnersSum - } - } + ownerCounts := loadTemplateActiveDeveloperCounts(ctx, db, templateIDs) + usageByTemplate := loadTemplateUsage( + ctx, db, options.OwnerID, organizationID, templateIDs, + ) + + for i := range ranked { + ranked[i].ActiveDevelopers = ownerCounts[ranked[i].Template.ID] + ranked[i].Usage = usageByTemplate[ranked[i].Template.ID] } - // Sort by active developer count descending. - slices.SortStableFunc(templates, func(a, b database.Template) int { - return cmp.Compare(ownerCounts[b.ID], ownerCounts[a.ID]) - }) + rankTemplates(ranked, query) + selectionHint, recommendedID, recommendationReason := selectTemplateRecommendation( + ranked, visibleTemplateCount, + ) + // Paginate. page := args.Page if page < 1 { page = 1 } - totalCount := len(templates) + totalCount := len(ranked) totalPages := (totalCount + listTemplatesPageSize - 1) / listTemplatesPageSize if totalPages == 0 { totalPages = 1 @@ -113,38 +136,246 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp if end > totalCount { end = totalCount } - pageTemplates := templates[start:end] + pageTemplates := ranked[start:end] 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 = append(items, templateItem(t, recommendedID)) } - return toolResponse(map[string]any{ - "templates": items, - "count": len(items), - "page": page, - "total_pages": totalPages, - "total_count": totalCount, - }), nil + result := map[string]any{ + "templates": items, + "count": len(items), + "page": page, + "total_pages": totalPages, + "total_count": totalCount, + "selection_hint": selectionHint, + "recommendation_reason": recommendationReason, + } + if recommendedID != uuid.Nil { + result["recommended_template_id"] = recommendedID.String() + } + return toolResponse(result), nil }, ) } +func candidateRankedTemplates(templates []database.Template, query string) []rankedTemplate { + ranked := make([]rankedTemplate, 0, len(templates)) + for _, t := range templates { + queryScore := templateQueryScore(t, query) + if query != "" && queryScore == 0 { + continue + } + ranked = append(ranked, rankedTemplate{ + Template: t, + QueryScore: queryScore, + }) + } + return ranked +} + +func loadTemplateActiveDeveloperCounts( + ctx context.Context, + db database.Store, + templateIDs []uuid.UUID, +) map[uuid.UUID]int64 { + ownerCounts := make(map[uuid.UUID]int64) + if len(templateIDs) == 0 { + return ownerCounts + } + + // Templates are already filtered with the owner's permissions. The + // aggregate count query requires system read because it spans workspace + // owners, but it only receives IDs the owner can already see. + rows, err := db.GetWorkspaceUniqueOwnerCountByTemplateIDs(dbauthz.AsSystemRestricted(ctx), templateIDs) //nolint:gocritic // see above + if err != nil { + return ownerCounts + } + for _, row := range rows { + ownerCounts[row.TemplateID] = row.UniqueOwnersSum + } + return ownerCounts +} + +func loadTemplateUsage( + ctx context.Context, + db database.Store, + ownerID uuid.UUID, + organizationID uuid.UUID, + templateIDs []uuid.UUID, +) map[uuid.UUID]templateUsage { + usageByTemplate := make(map[uuid.UUID]templateUsage) + if ownerID == uuid.Nil || len(templateIDs) == 0 { + return usageByTemplate + } + + rows, err := db.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams{ + OwnerID: ownerID, + OrganizationID: organizationID, + TemplateIDs: templateIDs, + }) + if err != nil { + return usageByTemplate + } + for _, row := range rows { + usageByTemplate[row.TemplateID] = templateUsage{ + WorkspaceCount: row.WorkspaceCount, + LastUsedAt: row.LastUsedAt, + } + } + return usageByTemplate +} + +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.Usage.WorkspaceCount, a.Usage.WorkspaceCount); c != 0 { + return c + } + if c := b.Usage.LastUsedAt.Compare(a.Usage.LastUsedAt); c != 0 { + return c + } + if c := cmp.Compare(b.ActiveDevelopers, a.ActiveDevelopers); c != 0 { + return c + } + if c := strings.Compare(a.Template.Name, b.Template.Name); c != 0 { + return c + } + return strings.Compare(a.Template.ID.String(), b.Template.ID.String()) + }) + + for i := range ranked { + ranked[i].Rank = i + 1 + } +} + +func selectTemplateRecommendation( + ranked []rankedTemplate, + visibleTemplateCount int, +) (string, uuid.UUID, string) { + if len(ranked) == 0 { + return listTemplatesHintNoConfidence, uuid.Nil, "no_matching_templates" + } + + top := ranked[0] + if visibleTemplateCount == 1 && len(ranked) == 1 { + return listTemplatesHintOnlyAvailable, top.Template.ID, "only_available_template" + } + if !templateHasRankingSignal(top) { + return listTemplatesHintNoConfidence, uuid.Nil, "no_ranking_signal" + } + if len(ranked) > 1 && templatesAreAmbiguous(top, ranked[1]) { + return listTemplatesHintAmbiguous, uuid.Nil, "top_templates_are_ambiguous" + } + return listTemplatesHintHighConfidence, top.Template.ID, rankReason(top) +} + +func templatesAreAmbiguous(a, b rankedTemplate) bool { + return a.QueryScore == b.QueryScore && + a.Usage.WorkspaceCount == b.Usage.WorkspaceCount && + a.Usage.LastUsedAt.Equal(b.Usage.LastUsedAt) && + a.ActiveDevelopers == b.ActiveDevelopers +} + +func templateHasRankingSignal(t rankedTemplate) bool { + return t.QueryScore > 0 || t.Usage.WorkspaceCount > 0 || t.ActiveDevelopers > 0 +} + +func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { + item := map[string]any{ + "id": t.Template.ID.String(), + "name": t.Template.Name, + "organization_id": t.Template.OrganizationID.String(), + "rank": t.Rank, + "rank_reason": rankReason(t), + } + 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.ActiveDevelopers > 0 { + item["active_developers"] = t.ActiveDevelopers + } + if t.Usage.WorkspaceCount > 0 { + item["your_workspace_count"] = t.Usage.WorkspaceCount + item["last_used_by_you"] = t.Usage.LastUsedAt.Format(time.RFC3339Nano) + } + if t.Template.ID == recommendedID { + item["recommended"] = true + } + return item +} + +func rankReason(t rankedTemplate) string { + switch { + case t.QueryScore > 0 && t.Usage.WorkspaceCount > 0: + return "matches_query_and_used_by_you" + case t.QueryScore > 0: + return "matches_query" + case t.Usage.WorkspaceCount > 0: + return "used_by_you" + case t.ActiveDevelopers > 0: + return "popular_in_org" + default: + return "ordered_by_name" + } +} + +func templateQueryScore(t database.Template, query string) int { + query = normalizeTemplateSearch(query) + if query == "" { + return 0 + } + + queryCompact := compactTemplateSearch(query) + for _, field := range []string{t.Name, t.DisplayName} { + field = normalizeTemplateSearch(field) + if field == "" { + continue + } + if field == query || compactTemplateSearch(field) == queryCompact { + return 4 + } + } + for _, field := range []string{t.Name, t.DisplayName} { + field = normalizeTemplateSearch(field) + if field == "" { + continue + } + if strings.HasPrefix(field, query) || strings.HasPrefix(compactTemplateSearch(field), queryCompact) { + return 3 + } + } + for _, field := range []string{t.Name, t.DisplayName} { + field = normalizeTemplateSearch(field) + if field == "" { + continue + } + if strings.Contains(field, query) || strings.Contains(compactTemplateSearch(field), queryCompact) { + return 2 + } + } + if strings.Contains(normalizeTemplateSearch(t.Description), query) { + return 1 + } + return 0 +} + +func normalizeTemplateSearch(value string) string { + return strings.ToLower(strings.TrimSpace(value)) +} + +func compactTemplateSearch(value string) string { + return strings.ReplaceAll(value, " ", "") +} + // asOwner sets up a dbauthz context for the given owner so that // subsequent database calls are scoped to what that user can access. func asOwner(ctx context.Context, db database.Store, ownerID uuid.UUID) (context.Context, error) { diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index 0cf25d2c432..66c2cf6e8e4 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -3,7 +3,9 @@ package chattool_test import ( "context" "encoding/json" + "fmt" "testing" + "time" "charm.land/fantasy" "github.com/google/uuid" @@ -121,6 +123,231 @@ 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: "data-science", + 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, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, displayTemplate.ID.String(), result["recommended_template_id"]) + require.Equal(t, "matches_query", templates[0]["rank_reason"]) + + 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, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, descriptionTemplate.ID.String(), result["recommended_template_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.Date(2026, 5, 1, 12, 0, 0, 0, time.UTC), + }) + + 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(11), result["total_count"]) + require.Equal(t, float64(2), result["total_pages"]) + require.Equal(t, target.ID.String(), templates[0]["id"]) + require.Equal(t, float64(1), templates[0]["rank"]) + require.Equal(t, float64(1), templates[0]["your_workspace_count"]) + require.NotEmpty(t, templates[0]["last_used_by_you"]) + require.Equal(t, true, templates[0]["recommended"]) + require.Equal(t, "used_by_you", templates[0]["rank_reason"]) + require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, target.ID.String(), result["recommended_template_id"]) +} + +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.Date(2026, 5, 2, 12, 0, 0, 0, time.UTC), + }) + + 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, "matches_query", templates[0]["rank_reason"]) + require.Equal(t, "matches_query_and_used_by_you", templates[1]["rank_reason"]) + require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, target.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, "popular_in_org", templates[0]["rank_reason"]) + require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, popular.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, "ambiguous_top_matches", result["selection_hint"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) + _, ok = templates[0]["recommended"] + require.False(t, ok) +} + //nolint:tparallel,paralleltest // Subtests share a single DB and run sequentially. func TestTemplateAllowlistEnforcement(t *testing.T) { t.Parallel() @@ -187,6 +414,10 @@ 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, "only_available_template", result["selection_hint"]) + require.Equal(t, t1.ID.String(), result["recommended_template_id"]) + require.Equal(t, true, m["recommended"]) + require.Equal(t, float64(1), m["rank"]) }) t.Run("NoMatches", func(t *testing.T) { @@ -301,3 +532,38 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) }) } + +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 +} diff --git a/coderd/x/chatd/chattool/readtemplate.go b/coderd/x/chatd/chattool/readtemplate.go index 09179237bab..b790e2c4d1c 100644 --- a/coderd/x/chatd/chattool/readtemplate.go +++ b/coderd/x/chatd/chattool/readtemplate.go @@ -19,20 +19,21 @@ type ReadTemplateOptions struct { } type readTemplateArgs struct { - TemplateID string `json:"template_id" description:"The UUIDv4 of the template to read details for. Obtain this from list_templates."` + TemplateID string `json:"template_id" description:"The UUIDv4 of the template to read details for. Obtain this from list_templates recommended_template_id or a ranked template."` } // 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 == "" { diff --git a/coderd/x/chatd/prompt.go b/coderd/x/chatd/prompt.go index 23c42fcb9b0..c295ccfacec 100644 --- a/coderd/x/chatd/prompt.go +++ b/coderd/x/chatd/prompt.go @@ -82,6 +82,14 @@ Do not start with clarifying questions if the codebase or tools can answer them. Ask the minimum number of questions needed to define the scope together. + +When no workspace is attached and you need to create one: +- Call list_templates with concise search terms from the user's task when the task suggests a language, framework, image, or environment. +- Treat recommended_template_id, or rank 1 when selection_hint is only_available_template or high_confidence_recommendation, as the default template unless the user asked for a different template. +- Do not paginate unless selection_hint is ambiguous_top_matches or no_confident_match, no returned template fits the request, or the user asked to browse or compare templates. +- Call read_template before create_workspace when you need parameter names, required parameter values, or preset IDs. Otherwise use create_workspace with the selected template_id and defaults. + + Propose a plan when: - The task is too ambiguous to implement with confidence. From f8882fce4e6ceb4d6f0423f9d4c2f790ee288c0e Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 7 May 2026 12:21:39 +0000 Subject: [PATCH 02/24] test(coderd/database): cover workspace usage authz --- coderd/database/dbauthz/dbauthz_test.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 795a0e66416..fe688d28c57 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4764,6 +4764,15 @@ func (s *MethodTestSuite) TestSystemFunctions() { dbm.EXPECT().GetWorkspaceUniqueOwnerCountByTemplateIDs(gomock.Any(), ids).Return([]database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow{}, nil).AnyTimes() check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead) })) + s.Run("GetWorkspaceUsageGroupedByTemplateIDForOwner", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams{ + OwnerID: uuid.New(), + OrganizationID: uuid.New(), + TemplateIDs: []uuid.UUID{uuid.New()}, + } + dbm.EXPECT().GetWorkspaceUsageGroupedByTemplateIDForOwner(gomock.Any(), arg).Return([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead) + })) s.Run("GetWorkspaceAgentScriptsByAgentIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { ids := []uuid.UUID{uuid.New()} dbm.EXPECT().GetWorkspaceAgentScriptsByAgentIDs(gomock.Any(), ids).Return([]database.GetWorkspaceAgentScriptsByAgentIDsRow{}, nil).AnyTimes() From 34c7facc78a79e7c24e2892f1c42e45976e70943 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 15 May 2026 10:57:49 +0000 Subject: [PATCH 03/24] fix(coderd): address template ranking feedback --- coderd/database/dbauthz/dbauthz.go | 4 +- coderd/database/dbauthz/dbauthz_test.go | 18 +- coderd/database/dbmetrics/querymetrics.go | 8 +- coderd/database/dbmock/dbmock.go | 14 +- coderd/database/models.go | 2 +- coderd/database/querier.go | 4 +- coderd/database/queries.sql.go | 21 +- coderd/database/queries/workspaces.sql | 2 +- coderd/x/chatd/chatd.go | 1 + coderd/x/chatd/chattool/listtemplates.go | 210 +++++++++++++----- coderd/x/chatd/chattool/listtemplates_test.go | 210 +++++++++++++++++- 11 files changed, 392 insertions(+), 102 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index ac795a8d273..9a660961cc7 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5077,7 +5077,7 @@ func (q *querier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, return q.db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs) } -func (q *querier) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { +func (q *querier) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { obj := rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()) if arg.OrganizationID != uuid.Nil { obj = obj.InOrg(arg.OrganizationID) @@ -5087,7 +5087,7 @@ func (q *querier) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Conte if err := q.authorizeContext(ctx, policy.ActionRead, obj); err != nil { return nil, err } - return q.db.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg) + return q.db.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg) } func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index fe688d28c57..86cceca4704 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3110,6 +3110,15 @@ func (s *MethodTestSuite) TestWorkspace() { // No asserts here because SQLFilter. check.Args(ws.OwnerID, emptyPreparedAuthorized{}).Asserts() })) + s.Run("GetWorkspaceUsageGroupedByTemplateIDByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams{ + OwnerID: uuid.New(), + OrganizationID: uuid.New(), + TemplateIDs: []uuid.UUID{uuid.New()}, + } + dbm.EXPECT().GetWorkspaceUsageGroupedByTemplateIDByOwnerID(gomock.Any(), arg).Return([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead) + })) s.Run("GetWorkspaceACLByID", s.Mocked(func(dbM *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ws := testutil.Fake(s.T(), faker, database.Workspace{}) dbM.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes() @@ -4764,15 +4773,6 @@ func (s *MethodTestSuite) TestSystemFunctions() { dbm.EXPECT().GetWorkspaceUniqueOwnerCountByTemplateIDs(gomock.Any(), ids).Return([]database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow{}, nil).AnyTimes() check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead) })) - s.Run("GetWorkspaceUsageGroupedByTemplateIDForOwner", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - arg := database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams{ - OwnerID: uuid.New(), - OrganizationID: uuid.New(), - TemplateIDs: []uuid.UUID{uuid.New()}, - } - dbm.EXPECT().GetWorkspaceUsageGroupedByTemplateIDForOwner(gomock.Any(), arg).Return([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow{}, nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead) - })) s.Run("GetWorkspaceAgentScriptsByAgentIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { ids := []uuid.UUID{uuid.New()} dbm.EXPECT().GetWorkspaceAgentScriptsByAgentIDs(gomock.Any(), ids).Return([]database.GetWorkspaceAgentScriptsByAgentIDsRow{}, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1d46b092fd1..a444dc588c3 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3440,11 +3440,11 @@ func (m queryMetricsStore) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context return r0, r1 } -func (m queryMetricsStore) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { +func (m queryMetricsStore) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { start := time.Now() - r0, r1 := m.s.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg) - m.queryLatencies.WithLabelValues("GetWorkspaceUsageGroupedByTemplateIDForOwner").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceUsageGroupedByTemplateIDForOwner").Inc() + r0, r1 := m.s.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg) + m.queryLatencies.WithLabelValues("GetWorkspaceUsageGroupedByTemplateIDByOwnerID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceUsageGroupedByTemplateIDByOwnerID").Inc() return r0, r1 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 2e0fcd42d41..66b042c5f57 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -6437,19 +6437,19 @@ func (mr *MockStoreMockRecorder) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUniqueOwnerCountByTemplateIDs", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUniqueOwnerCountByTemplateIDs), ctx, templateIds) } -// GetWorkspaceUsageGroupedByTemplateIDForOwner mocks base method. -func (m *MockStore) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { +// GetWorkspaceUsageGroupedByTemplateIDByOwnerID mocks base method. +func (m *MockStore) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetWorkspaceUsageGroupedByTemplateIDForOwner", ctx, arg) - ret0, _ := ret[0].([]database.GetWorkspaceUsageGroupedByTemplateIDForOwnerRow) + ret := m.ctrl.Call(m, "GetWorkspaceUsageGroupedByTemplateIDByOwnerID", ctx, arg) + ret0, _ := ret[0].([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetWorkspaceUsageGroupedByTemplateIDForOwner indicates an expected call of GetWorkspaceUsageGroupedByTemplateIDForOwner. -func (mr *MockStoreMockRecorder) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, arg any) *gomock.Call { +// GetWorkspaceUsageGroupedByTemplateIDByOwnerID indicates an expected call of GetWorkspaceUsageGroupedByTemplateIDByOwnerID. +func (mr *MockStoreMockRecorder) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUsageGroupedByTemplateIDForOwner", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUsageGroupedByTemplateIDForOwner), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUsageGroupedByTemplateIDByOwnerID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUsageGroupedByTemplateIDByOwnerID), ctx, arg) } // GetWorkspaces mocks base method. diff --git a/coderd/database/models.go b/coderd/database/models.go index 65e6d5a1420..5dc17bd6905 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -1,6 +1,6 @@ // Code generated by sqlc. DO NOT EDIT. // versions: -// sqlc v1.30.0 +// sqlc v1.31.1 package database diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 02b8212c70f..8b98cc243f4 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1,6 +1,6 @@ // Code generated by sqlc. DO NOT EDIT. // versions: -// sqlc v1.30.0 +// sqlc v1.31.1 package database @@ -833,7 +833,7 @@ type sqlcQuerier interface { GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceResource, error) GetWorkspaceResourcesCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceResource, error) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, templateIds []uuid.UUID) ([]GetWorkspaceUniqueOwnerCountByTemplateIDsRow, error) - GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) + GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) // build_params is used to filter by build parameters if present. // It has to be a CTE because the set returning function 'unnest' cannot // be used in a WHERE clause. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cbefcfff43e..c19a11529cb 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1,6 +1,6 @@ // Code generated by sqlc. DO NOT EDIT. // versions: -// sqlc v1.30.0 +// sqlc v1.31.1 package database @@ -12850,8 +12850,9 @@ RETURNING id func (q *sqlQuerier) DeleteLicense(ctx context.Context, id int32) (int32, error) { row := q.db.QueryRowContext(ctx, deleteLicense, id) - err := row.Scan(&id) - return id, err + var id_2 int32 + err := row.Scan(&id_2) + return id_2, err } const getLicenseByID = `-- name: GetLicenseByID :one @@ -32389,7 +32390,7 @@ func (q *sqlQuerier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Conte return items, nil } -const getWorkspaceUsageGroupedByTemplateIDForOwner = `-- name: GetWorkspaceUsageGroupedByTemplateIDForOwner :many +const getWorkspaceUsageGroupedByTemplateIDByOwnerID = `-- name: GetWorkspaceUsageGroupedByTemplateIDByOwnerID :many SELECT template_id, COUNT(*) AS workspace_count, @@ -32408,27 +32409,27 @@ WHERE GROUP BY template_id ` -type GetWorkspaceUsageGroupedByTemplateIDForOwnerParams struct { +type GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams struct { OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` TemplateIDs []uuid.UUID `db:"template_ids" json:"template_ids"` } -type GetWorkspaceUsageGroupedByTemplateIDForOwnerRow struct { +type GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow struct { TemplateID uuid.UUID `db:"template_id" json:"template_id"` WorkspaceCount int64 `db:"workspace_count" json:"workspace_count"` LastUsedAt time.Time `db:"last_used_at" json:"last_used_at"` } -func (q *sqlQuerier) GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDForOwnerParams) ([]GetWorkspaceUsageGroupedByTemplateIDForOwnerRow, error) { - rows, err := q.db.QueryContext(ctx, getWorkspaceUsageGroupedByTemplateIDForOwner, arg.OwnerID, arg.OrganizationID, pq.Array(arg.TemplateIDs)) +func (q *sqlQuerier) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceUsageGroupedByTemplateIDByOwnerID, arg.OwnerID, arg.OrganizationID, pq.Array(arg.TemplateIDs)) if err != nil { return nil, err } defer rows.Close() - var items []GetWorkspaceUsageGroupedByTemplateIDForOwnerRow + var items []GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow for rows.Next() { - var i GetWorkspaceUsageGroupedByTemplateIDForOwnerRow + var i GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow if err := rows.Scan(&i.TemplateID, &i.WorkspaceCount, &i.LastUsedAt); err != nil { return nil, err } diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql index 972f4a54271..c860b7b0afe 100644 --- a/coderd/database/queries/workspaces.sql +++ b/coderd/database/queries/workspaces.sql @@ -497,7 +497,7 @@ LEFT JOIN workspaces ON workspaces.template_id = templates.id AND workspaces.del WHERE templates.id = ANY(@template_ids :: uuid[]) GROUP BY templates.id; --- name: GetWorkspaceUsageGroupedByTemplateIDForOwner :many +-- name: GetWorkspaceUsageGroupedByTemplateIDByOwnerID :many SELECT template_id, COUNT(*) AS workspace_count, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index ff551bb76f6..de5d59934ea 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -5952,6 +5952,7 @@ func (p *Server) appendRootChatTools( tools = append(tools, chattool.ListTemplates(p.db, opts.chat.OrganizationID, chattool.ListTemplatesOptions{ OwnerID: opts.chat.OwnerID, + Logger: p.logger, AllowedTemplateIDs: p.chatTemplateAllowlist, }), chattool.ReadTemplate(p.db, opts.chat.OrganizationID, chattool.ReadTemplateOptions{ diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 55abb3ee4b9..6d50e7b3a8e 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -4,6 +4,7 @@ import ( "cmp" "context" "database/sql" + "errors" "maps" "slices" "strings" @@ -13,13 +14,20 @@ import ( "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" ) -const listTemplatesPageSize = 10 +const ( + listTemplatesPageSize = 10 + + listTemplatesMinPersonalWorkspacesForRecommendation = 2 + listTemplatesMinActiveDevelopersForRecommendation = 2 + listTemplatesRecentUsageWindow = 90 * 24 * time.Hour +) const ( listTemplatesHintOnlyAvailable = "only_available_template" @@ -28,9 +36,17 @@ const ( listTemplatesHintNoConfidence = "no_confident_match" ) +const ( + queryScoreExactName = 4 + queryScoreNamePrefix = 3 + queryScoreNameContains = 2 + queryScoreDescriptionMatch = 1 +) + // ListTemplatesOptions configures the list_templates tool. type ListTemplatesOptions struct { OwnerID uuid.UUID + Logger slog.Logger AllowedTemplateIDs func() map[uuid.UUID]bool } @@ -52,6 +68,13 @@ type templateUsage struct { LastUsedAt time.Time } +type templateRankSignals struct { + QueryScore int + WorkspaceCount int64 + LastUsedAtUnixNano int64 + ActiveDevelopers int64 +} + // ListTemplates returns a tool that lists available workspace templates. // The agent uses this to discover templates before creating a workspace. // Results are ranked before pagination using query relevance, current-user @@ -97,16 +120,30 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp query := strings.TrimSpace(args.Query) visibleTemplateCount := len(templates) - ranked := candidateRankedTemplates(templates, query) + ranked := scoreTemplateCandidates(templates, query) templateIDs := make([]uuid.UUID, len(ranked)) for i, t := range ranked { templateIDs[i] = t.Template.ID } - ownerCounts := loadTemplateActiveDeveloperCounts(ctx, db, templateIDs) - usageByTemplate := loadTemplateUsage( + ownerCounts, ownerCountsErr := loadTemplateActiveDeveloperCounts(ctx, db, templateIDs) + if ownerCountsErr != nil { + options.Logger.Warn(ctx, "failed to load template active developer counts", + slog.F("template_count", len(templateIDs)), + slog.Error(ownerCountsErr), + ) + } + usageByTemplate, usageErr := loadTemplateUsage( ctx, db, options.OwnerID, organizationID, templateIDs, ) + if usageErr != nil { + options.Logger.Warn(ctx, "failed to load template usage", + slog.F("owner_id", options.OwnerID), + slog.F("organization_id", organizationID), + slog.F("template_count", len(templateIDs)), + slog.Error(usageErr), + ) + } for i := range ranked { ranked[i].ActiveDevelopers = ownerCounts[ranked[i].Template.ID] @@ -115,7 +152,10 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp rankTemplates(ranked, query) selectionHint, recommendedID, recommendationReason := selectTemplateRecommendation( - ranked, visibleTemplateCount, + ranked, + visibleTemplateCount, + errors.Join(ownerCountsErr, usageErr), + time.Now(), ) // Paginate. @@ -144,13 +184,14 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp } result := map[string]any{ - "templates": items, - "count": len(items), - "page": page, - "total_pages": totalPages, - "total_count": totalCount, - "selection_hint": selectionHint, - "recommendation_reason": recommendationReason, + "templates": items, + "count": len(items), + "page": page, + "total_pages": totalPages, + "total_count": totalCount, + "available_template_count": visibleTemplateCount, + "selection_hint": selectionHint, + "recommendation_reason": recommendationReason, } if recommendedID != uuid.Nil { result["recommended_template_id"] = recommendedID.String() @@ -160,29 +201,29 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp ) } -func candidateRankedTemplates(templates []database.Template, query string) []rankedTemplate { - ranked := make([]rankedTemplate, 0, len(templates)) +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 } - ranked = append(ranked, rankedTemplate{ + candidates = append(candidates, rankedTemplate{ Template: t, QueryScore: queryScore, }) } - return ranked + return candidates } func loadTemplateActiveDeveloperCounts( ctx context.Context, db database.Store, templateIDs []uuid.UUID, -) map[uuid.UUID]int64 { +) (map[uuid.UUID]int64, error) { ownerCounts := make(map[uuid.UUID]int64) if len(templateIDs) == 0 { - return ownerCounts + return ownerCounts, nil } // Templates are already filtered with the owner's permissions. The @@ -190,12 +231,12 @@ func loadTemplateActiveDeveloperCounts( // owners, but it only receives IDs the owner can already see. rows, err := db.GetWorkspaceUniqueOwnerCountByTemplateIDs(dbauthz.AsSystemRestricted(ctx), templateIDs) //nolint:gocritic // see above if err != nil { - return ownerCounts + return ownerCounts, err } for _, row := range rows { ownerCounts[row.TemplateID] = row.UniqueOwnersSum } - return ownerCounts + return ownerCounts, nil } func loadTemplateUsage( @@ -204,19 +245,19 @@ func loadTemplateUsage( ownerID uuid.UUID, organizationID uuid.UUID, templateIDs []uuid.UUID, -) map[uuid.UUID]templateUsage { +) (map[uuid.UUID]templateUsage, error) { usageByTemplate := make(map[uuid.UUID]templateUsage) if ownerID == uuid.Nil || len(templateIDs) == 0 { - return usageByTemplate + return usageByTemplate, nil } - rows, err := db.GetWorkspaceUsageGroupedByTemplateIDForOwner(ctx, database.GetWorkspaceUsageGroupedByTemplateIDForOwnerParams{ + rows, err := db.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams{ OwnerID: ownerID, OrganizationID: organizationID, TemplateIDs: templateIDs, }) if err != nil { - return usageByTemplate + return usageByTemplate, err } for _, row := range rows { usageByTemplate[row.TemplateID] = templateUsage{ @@ -224,29 +265,22 @@ func loadTemplateUsage( LastUsedAt: row.LastUsedAt, } } - return usageByTemplate + return usageByTemplate, nil } 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.Usage.WorkspaceCount, a.Usage.WorkspaceCount); c != 0 { + if c := compareTemplateRankSignals( + templateRankSignalsFor(a), + templateRankSignalsFor(b), + query, + ); c != 0 { return c } - if c := b.Usage.LastUsedAt.Compare(a.Usage.LastUsedAt); c != 0 { + if c := cmp.Compare(a.Template.Name, b.Template.Name); c != 0 { return c } - if c := cmp.Compare(b.ActiveDevelopers, a.ActiveDevelopers); c != 0 { - return c - } - if c := strings.Compare(a.Template.Name, b.Template.Name); c != 0 { - return c - } - return strings.Compare(a.Template.ID.String(), b.Template.ID.String()) + return cmp.Compare(a.Template.ID.String(), b.Template.ID.String()) }) for i := range ranked { @@ -254,15 +288,51 @@ func rankTemplates(ranked []rankedTemplate, query string) { } } +func templateRankSignalsFor(t rankedTemplate) templateRankSignals { + return templateRankSignals{ + QueryScore: t.QueryScore, + WorkspaceCount: t.Usage.WorkspaceCount, + LastUsedAtUnixNano: templateRankTime(t.Usage.LastUsedAt), + ActiveDevelopers: t.ActiveDevelopers, + } +} + +func templateRankTime(t time.Time) int64 { + if t.IsZero() { + return 0 + } + return t.UnixNano() +} + +func compareTemplateRankSignals(a, b templateRankSignals, query string) int { + if query != "" { + if c := cmp.Compare(b.QueryScore, a.QueryScore); c != 0 { + return c + } + } + if c := cmp.Compare(b.WorkspaceCount, a.WorkspaceCount); c != 0 { + return c + } + if c := cmp.Compare(b.LastUsedAtUnixNano, a.LastUsedAtUnixNano); c != 0 { + return c + } + return cmp.Compare(b.ActiveDevelopers, a.ActiveDevelopers) +} + func selectTemplateRecommendation( ranked []rankedTemplate, visibleTemplateCount int, + rankingSignalsErr error, + now time.Time, ) (string, uuid.UUID, string) { if len(ranked) == 0 { return listTemplatesHintNoConfidence, uuid.Nil, "no_matching_templates" } top := ranked[0] + if rankingSignalsErr != nil { + return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" + } if visibleTemplateCount == 1 && len(ranked) == 1 { return listTemplatesHintOnlyAvailable, top.Template.ID, "only_available_template" } @@ -272,27 +342,44 @@ func selectTemplateRecommendation( if len(ranked) > 1 && templatesAreAmbiguous(top, ranked[1]) { return listTemplatesHintAmbiguous, uuid.Nil, "top_templates_are_ambiguous" } - return listTemplatesHintHighConfidence, top.Template.ID, rankReason(top) + if !templateHasConfidentRankingSignal(top, now) { + return listTemplatesHintNoConfidence, uuid.Nil, "weak_ranking_signal" + } + return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) } func templatesAreAmbiguous(a, b rankedTemplate) bool { - return a.QueryScore == b.QueryScore && - a.Usage.WorkspaceCount == b.Usage.WorkspaceCount && - a.Usage.LastUsedAt.Equal(b.Usage.LastUsedAt) && - a.ActiveDevelopers == b.ActiveDevelopers + return templateRankSignalsFor(a) == templateRankSignalsFor(b) } func templateHasRankingSignal(t rankedTemplate) bool { - return t.QueryScore > 0 || t.Usage.WorkspaceCount > 0 || t.ActiveDevelopers > 0 + signals := templateRankSignalsFor(t) + return signals.QueryScore > 0 || signals.WorkspaceCount > 0 || signals.ActiveDevelopers > 0 +} + +func templateHasConfidentRankingSignal(t rankedTemplate, now time.Time) bool { + signals := templateRankSignalsFor(t) + if signals.QueryScore > 0 { + return true + } + if signals.WorkspaceCount >= listTemplatesMinPersonalWorkspacesForRecommendation { + return true + } + if signals.WorkspaceCount > 0 && + !t.Usage.LastUsedAt.IsZero() && + now.Sub(t.Usage.LastUsedAt) <= listTemplatesRecentUsageWindow { + return true + } + return signals.ActiveDevelopers >= listTemplatesMinActiveDevelopersForRecommendation } func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { item := map[string]any{ - "id": t.Template.ID.String(), - "name": t.Template.Name, - "organization_id": t.Template.OrganizationID.String(), - "rank": t.Rank, - "rank_reason": rankReason(t), + "id": t.Template.ID.String(), + "name": t.Template.Name, + "organization_id": t.Template.OrganizationID.String(), + "rank": t.Rank, + "relevance_signals": relevanceSignals(t), } if display := strings.TrimSpace(t.Template.DisplayName); display != "" { item["display_name"] = display @@ -313,15 +400,16 @@ func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { return item } -func rankReason(t rankedTemplate) string { +func relevanceSignals(t rankedTemplate) string { + signals := templateRankSignalsFor(t) switch { - case t.QueryScore > 0 && t.Usage.WorkspaceCount > 0: + case signals.QueryScore > 0 && signals.WorkspaceCount > 0: return "matches_query_and_used_by_you" - case t.QueryScore > 0: + case signals.QueryScore > 0: return "matches_query" - case t.Usage.WorkspaceCount > 0: + case signals.WorkspaceCount > 0: return "used_by_you" - case t.ActiveDevelopers > 0: + case signals.ActiveDevelopers > 0: return "popular_in_org" default: return "ordered_by_name" @@ -341,7 +429,7 @@ func templateQueryScore(t database.Template, query string) int { continue } if field == query || compactTemplateSearch(field) == queryCompact { - return 4 + return queryScoreExactName } } for _, field := range []string{t.Name, t.DisplayName} { @@ -350,7 +438,7 @@ func templateQueryScore(t database.Template, query string) int { continue } if strings.HasPrefix(field, query) || strings.HasPrefix(compactTemplateSearch(field), queryCompact) { - return 3 + return queryScoreNamePrefix } } for _, field := range []string{t.Name, t.DisplayName} { @@ -359,11 +447,11 @@ func templateQueryScore(t database.Template, query string) int { continue } if strings.Contains(field, query) || strings.Contains(compactTemplateSearch(field), queryCompact) { - return 2 + return queryScoreNameContains } } if strings.Contains(normalizeTemplateSearch(t.Description), query) { - return 1 + return queryScoreDescriptionMatch } return 0 } @@ -372,8 +460,10 @@ func normalizeTemplateSearch(value string) string { return strings.ToLower(strings.TrimSpace(value)) } +var templateSearchCompactReplacer = strings.NewReplacer(" ", "", "-", "", "_", "") + func compactTemplateSearch(value string) string { - return strings.ReplaceAll(value, " ", "") + return templateSearchCompactReplacer.Replace(value) } // asOwner sets up a dbauthz context for the given owner so that diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index 66c2cf6e8e4..5dd1417da9f 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -84,6 +84,10 @@ 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, "no_confident_match", result["selection_hint"]) + require.Equal(t, "no_ranking_signal", result["recommendation_reason"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) }) t.Run("ReadTemplate_CrossOrgRejected", func(t *testing.T) { @@ -163,7 +167,7 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Equal(t, displayTemplate.ID.String(), templates[0]["id"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) require.Equal(t, displayTemplate.ID.String(), result["recommended_template_id"]) - require.Equal(t, "matches_query", templates[0]["rank_reason"]) + require.Equal(t, "matches_query", templates[0]["relevance_signals"]) result = runListTemplates(ctx, t, tool, `{"query":"TypeScript"}`) templates = listTemplateItems(t, result) @@ -171,6 +175,69 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Equal(t, descriptionTemplate.ID.String(), templates[0]["id"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) require.Equal(t, descriptionTemplate.ID.String(), result["recommended_template_id"]) + + result = runListTemplates(ctx, t, tool, `{"query":"does-not-exist"}`) + templates = listTemplateItems(t, result) + require.Empty(t, templates) + require.Equal(t, float64(0), result["total_count"]) + require.Equal(t, float64(3), result["available_template_count"]) + require.Equal(t, "no_confident_match", result["selection_hint"]) + require.Equal(t, "no_matching_templates", result["recommendation_reason"]) +} + +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"]) } func TestListTemplates_RanksAllCandidatesBeforePagination(t *testing.T) { @@ -199,7 +266,7 @@ func TestListTemplates_RanksAllCandidatesBeforePagination(t *testing.T) { OwnerID: user.ID, OrganizationID: org.ID, TemplateID: target.ID, - LastUsedAt: time.Date(2026, 5, 1, 12, 0, 0, 0, time.UTC), + LastUsedAt: time.Now().Add(-time.Hour), }) tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{ @@ -215,7 +282,7 @@ func TestListTemplates_RanksAllCandidatesBeforePagination(t *testing.T) { require.Equal(t, float64(1), templates[0]["your_workspace_count"]) require.NotEmpty(t, templates[0]["last_used_by_you"]) require.Equal(t, true, templates[0]["recommended"]) - require.Equal(t, "used_by_you", templates[0]["rank_reason"]) + require.Equal(t, "used_by_you", templates[0]["relevance_signals"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) require.Equal(t, target.ID.String(), result["recommended_template_id"]) } @@ -258,12 +325,53 @@ func TestListTemplates_QueryRelevanceOutranksPersonalUsage(t *testing.T) { 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, "matches_query", templates[0]["rank_reason"]) - require.Equal(t, "matches_query_and_used_by_you", templates[1]["rank_reason"]) + require.Equal(t, "matches_query", templates[0]["relevance_signals"]) + require.Equal(t, "matches_query_and_used_by_you", templates[1]["relevance_signals"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) 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, "matches_query_and_used_by_you", templates[0]["relevance_signals"]) + require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, used.ID.String(), result["recommended_template_id"]) +} + func TestListTemplates_OrgPopularityFallback(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -308,11 +416,97 @@ func TestListTemplates_OrgPopularityFallback(t *testing.T) { 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, "popular_in_org", templates[0]["rank_reason"]) + require.Equal(t, "popular_in_org", templates[0]["relevance_signals"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) 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, "no_confident_match", result["selection_hint"]) + require.Equal(t, "weak_ranking_signal", result["recommendation_reason"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) +} + +func TestListTemplates_StalePersonalUsageDoesNotRecommend(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, + }) + + 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: time.Now().Add(-180 * 24 * 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, 2) + require.Equal(t, oldUsage.ID.String(), templates[0]["id"]) + require.Equal(t, unused.ID.String(), templates[1]["id"]) + require.Equal(t, float64(1), templates[0]["your_workspace_count"]) + require.Equal(t, "no_confident_match", result["selection_hint"]) + require.Equal(t, "weak_ranking_signal", result["recommendation_reason"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) +} + func TestListTemplates_AmbiguousTopMatches(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -432,6 +626,10 @@ 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, "no_confident_match", result["selection_hint"]) + require.Equal(t, "no_matching_templates", result["recommendation_reason"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) }) }) From b8ae6390e81071ff40d2b2f515b285f028a8b5c4 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 15 May 2026 15:14:49 +0000 Subject: [PATCH 04/24] fix(coderd): follow up on template ranking feedback --- coderd/x/chatd/chatd.go | 1 + coderd/x/chatd/chattool/listtemplates.go | 28 ++++++-- .../chattool/listtemplates_internal_test.go | 57 ++++++++++++++++ coderd/x/chatd/chattool/listtemplates_test.go | 65 ++++++++++++++++++- 4 files changed, 145 insertions(+), 6 deletions(-) create mode 100644 coderd/x/chatd/chattool/listtemplates_internal_test.go diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index de5d59934ea..b9d76cdfe5a 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -5953,6 +5953,7 @@ func (p *Server) appendRootChatTools( chattool.ListTemplates(p.db, opts.chat.OrganizationID, chattool.ListTemplatesOptions{ OwnerID: opts.chat.OwnerID, Logger: p.logger, + Clock: p.clock, AllowedTemplateIDs: p.chatTemplateAllowlist, }), chattool.ReadTemplate(p.db, opts.chat.OrganizationID, chattool.ReadTemplateOptions{ diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 6d50e7b3a8e..e909b41f964 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -19,6 +19,7 @@ import ( "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 ( @@ -47,6 +48,7 @@ const ( type ListTemplatesOptions struct { OwnerID uuid.UUID Logger slog.Logger + Clock quartz.Clock AllowedTemplateIDs func() map[uuid.UUID]bool } @@ -81,6 +83,11 @@ type templateRankSignals struct { // usage, and organization-wide 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 as a ranked shortlist. "+ @@ -155,7 +162,7 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp ranked, visibleTemplateCount, errors.Join(ownerCountsErr, usageErr), - time.Now(), + clock.Now(), ) // Paginate. @@ -330,12 +337,15 @@ func selectTemplateRecommendation( } top := ranked[0] - if rankingSignalsErr != nil { - return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" - } if visibleTemplateCount == 1 && len(ranked) == 1 { return listTemplatesHintOnlyAvailable, top.Template.ID, "only_available_template" } + if rankingSignalsErr != nil { + if templateHasDecisiveQuerySignal(ranked) { + return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) + } + return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" + } if !templateHasRankingSignal(top) { return listTemplatesHintNoConfidence, uuid.Nil, "no_ranking_signal" } @@ -357,6 +367,13 @@ func templateHasRankingSignal(t rankedTemplate) bool { return signals.QueryScore > 0 || signals.WorkspaceCount > 0 || signals.ActiveDevelopers > 0 } +func templateHasDecisiveQuerySignal(ranked []rankedTemplate) bool { + if len(ranked) == 0 || ranked[0].QueryScore == 0 { + return false + } + return len(ranked) == 1 || ranked[0].QueryScore > ranked[1].QueryScore +} + func templateHasConfidentRankingSignal(t rankedTemplate, now time.Time) bool { signals := templateRankSignalsFor(t) if signals.QueryScore > 0 { @@ -450,7 +467,8 @@ func templateQueryScore(t database.Template, query string) int { return queryScoreNameContains } } - if strings.Contains(normalizeTemplateSearch(t.Description), query) { + desc := normalizeTemplateSearch(t.Description) + if strings.Contains(desc, query) || strings.Contains(compactTemplateSearch(desc), queryCompact) { return queryScoreDescriptionMatch } return 0 diff --git a/coderd/x/chatd/chattool/listtemplates_internal_test.go b/coderd/x/chatd/chattool/listtemplates_internal_test.go new file mode 100644 index 00000000000..655c275b9ed --- /dev/null +++ b/coderd/x/chatd/chattool/listtemplates_internal_test.go @@ -0,0 +1,57 @@ +package chattool + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" +) + +func TestSelectTemplateRecommendationRankingSignalsUnavailable(t *testing.T) { + t.Parallel() + + enrichmentErr := xerrors.New("enrichment failed") + now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC) + + onlyTemplateID := uuid.New() + hint, recommendedID, reason := selectTemplateRecommendation( + []rankedTemplate{{Template: database.Template{ID: onlyTemplateID}}}, + 1, + enrichmentErr, + now, + ) + require.Equal(t, listTemplatesHintOnlyAvailable, hint) + require.Equal(t, onlyTemplateID, recommendedID) + require.Equal(t, "only_available_template", reason) + + topID := uuid.New() + hint, recommendedID, reason = selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: topID}, QueryScore: queryScoreExactName}, + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreDescriptionMatch}, + }, + 2, + enrichmentErr, + now, + ) + require.Equal(t, listTemplatesHintHighConfidence, hint) + require.Equal(t, topID, recommendedID) + require.Equal(t, "matches_query", reason) + + hint, recommendedID, reason = selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, + }, + 2, + enrichmentErr, + now, + ) + require.Equal(t, listTemplatesHintNoConfidence, hint) + require.Equal(t, uuid.Nil, recommendedID) + require.Equal(t, "ranking_signals_unavailable", reason) +} diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index 5dd1417da9f..04d066012b1 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -17,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) { @@ -238,6 +239,17 @@ func TestListTemplates_QueryScoreTiers(t *testing.T) { 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) { @@ -467,6 +479,9 @@ func TestListTemplates_WeakOrgPopularityDoesNotRecommend(t *testing.T) { 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{}) @@ -489,11 +504,12 @@ func TestListTemplates_StalePersonalUsageDoesNotRecommend(t *testing.T) { OwnerID: user.ID, OrganizationID: org.ID, TemplateID: oldUsage.ID, - LastUsedAt: time.Now().Add(-180 * 24 * time.Hour), + 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) @@ -507,6 +523,53 @@ func TestListTemplates_StalePersonalUsageDoesNotRecommend(t *testing.T) { require.False(t, ok) } +func TestListTemplates_PersonalUsageCountRecommendsStaleTemplate(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", + }) + 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, float64(2), templates[0]["your_workspace_count"]) + require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, staleUsage.ID.String(), result["recommended_template_id"]) +} + func TestListTemplates_AmbiguousTopMatches(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) From 4d357dea9f21cf6ee022c31b96f0f7c65810696b Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 1 Jun 2026 10:48:52 +0000 Subject: [PATCH 05/24] fix(coderd): address template ranking feedback --- coderd/workspaces_test.go | 10 ++++- coderd/x/chatd/chattool/listtemplates.go | 43 ++++++++++++++----- .../chattool/listtemplates_internal_test.go | 43 +++++++++++++++++-- coderd/x/chatd/chattool/listtemplates_test.go | 16 ++++++- 4 files changed, 97 insertions(+), 15 deletions(-) diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index b03253b76ba..557da144177 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -4587,7 +4587,15 @@ func TestWorkspaceDormant(t *testing.T) { require.NoError(t, err) // Should be able to stop a workspace while it is dormant. - coderdtest.MustTransitionWorkspace(t, client, workspace.ID, codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop) + workspace = coderdtest.MustTransitionWorkspace(t, client, workspace.ID, codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop) + testutil.Eventually(ctx, t, func(context.Context) bool { + return auditor.Contains(t, database.AuditLog{ + ResourceID: workspace.LatestBuild.ID, + ResourceType: database.ResourceTypeWorkspaceBuild, + Action: database.AuditActionStop, + StatusCode: http.StatusOK, + }) + }, testutil.IntervalFast) // Reset the auditor auditor.ResetLogs() diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index e909b41f964..26881fe5fd0 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -4,7 +4,6 @@ import ( "cmp" "context" "database/sql" - "errors" "maps" "slices" "strings" @@ -77,6 +76,15 @@ type templateRankSignals struct { ActiveDevelopers int64 } +type templateRankingSignalErrors struct { + ActiveDeveloperCounts error + Usage error +} + +func (e templateRankingSignalErrors) hasAny() bool { + return e.ActiveDeveloperCounts != nil || e.Usage != nil +} + // ListTemplates returns a tool that lists available workspace templates. // The agent uses this to discover templates before creating a workspace. // Results are ranked before pagination using query relevance, current-user @@ -161,7 +169,10 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp selectionHint, recommendedID, recommendationReason := selectTemplateRecommendation( ranked, visibleTemplateCount, - errors.Join(ownerCountsErr, usageErr), + templateRankingSignalErrors{ + ActiveDeveloperCounts: ownerCountsErr, + Usage: usageErr, + }, clock.Now(), ) @@ -329,7 +340,7 @@ func compareTemplateRankSignals(a, b templateRankSignals, query string) int { func selectTemplateRecommendation( ranked []rankedTemplate, visibleTemplateCount int, - rankingSignalsErr error, + rankingSignalErrors templateRankingSignalErrors, now time.Time, ) (string, uuid.UUID, string) { if len(ranked) == 0 { @@ -340,10 +351,15 @@ func selectTemplateRecommendation( if visibleTemplateCount == 1 && len(ranked) == 1 { return listTemplatesHintOnlyAvailable, top.Template.ID, "only_available_template" } - if rankingSignalsErr != nil { + if rankingSignalErrors.hasAny() { if templateHasDecisiveQuerySignal(ranked) { return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) } + if rankingSignalErrors.Usage == nil && + templateHasConfidentPersonalUsageSignal(top, now) && + (len(ranked) == 1 || !templatesAreAmbiguous(top, ranked[1])) { + return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) + } return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" } if !templateHasRankingSignal(top) { @@ -374,17 +390,21 @@ func templateHasDecisiveQuerySignal(ranked []rankedTemplate) bool { return len(ranked) == 1 || ranked[0].QueryScore > ranked[1].QueryScore } +func templateHasConfidentPersonalUsageSignal(t rankedTemplate, now time.Time) bool { + if t.Usage.WorkspaceCount >= listTemplatesMinPersonalWorkspacesForRecommendation { + return true + } + return t.Usage.WorkspaceCount > 0 && + !t.Usage.LastUsedAt.IsZero() && + now.Sub(t.Usage.LastUsedAt) <= listTemplatesRecentUsageWindow +} + func templateHasConfidentRankingSignal(t rankedTemplate, now time.Time) bool { signals := templateRankSignalsFor(t) if signals.QueryScore > 0 { return true } - if signals.WorkspaceCount >= listTemplatesMinPersonalWorkspacesForRecommendation { - return true - } - if signals.WorkspaceCount > 0 && - !t.Usage.LastUsedAt.IsZero() && - now.Sub(t.Usage.LastUsedAt) <= listTemplatesRecentUsageWindow { + if templateHasConfidentPersonalUsageSignal(t, now) { return true } return signals.ActiveDevelopers >= listTemplatesMinActiveDevelopersForRecommendation @@ -440,6 +460,9 @@ func templateQueryScore(t database.Template, query string) int { } queryCompact := compactTemplateSearch(query) + if queryCompact == "" { + return 0 + } for _, field := range []string{t.Name, t.DisplayName} { field = normalizeTemplateSearch(field) if field == "" { diff --git a/coderd/x/chatd/chattool/listtemplates_internal_test.go b/coderd/x/chatd/chattool/listtemplates_internal_test.go index 655c275b9ed..6ee20a28b9e 100644 --- a/coderd/x/chatd/chattool/listtemplates_internal_test.go +++ b/coderd/x/chatd/chattool/listtemplates_internal_test.go @@ -15,13 +15,17 @@ func TestSelectTemplateRecommendationRankingSignalsUnavailable(t *testing.T) { t.Parallel() enrichmentErr := xerrors.New("enrichment failed") + enrichmentErrors := templateRankingSignalErrors{ + ActiveDeveloperCounts: enrichmentErr, + Usage: enrichmentErr, + } now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC) onlyTemplateID := uuid.New() hint, recommendedID, reason := selectTemplateRecommendation( []rankedTemplate{{Template: database.Template{ID: onlyTemplateID}}}, 1, - enrichmentErr, + enrichmentErrors, now, ) require.Equal(t, listTemplatesHintOnlyAvailable, hint) @@ -35,20 +39,53 @@ func TestSelectTemplateRecommendationRankingSignalsUnavailable(t *testing.T) { {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreDescriptionMatch}, }, 2, - enrichmentErr, + enrichmentErrors, now, ) require.Equal(t, listTemplatesHintHighConfidence, hint) require.Equal(t, topID, recommendedID) require.Equal(t, "matches_query", reason) + personalUsageID := uuid.New() + hint, recommendedID, reason = selectTemplateRecommendation( + []rankedTemplate{ + { + Template: database.Template{ID: personalUsageID}, + Usage: templateUsage{ + WorkspaceCount: 3, + LastUsedAt: now.Add(-180 * 24 * time.Hour), + }, + }, + {Template: database.Template{ID: uuid.New()}}, + }, + 2, + templateRankingSignalErrors{ActiveDeveloperCounts: enrichmentErr}, + now, + ) + require.Equal(t, listTemplatesHintHighConfidence, hint) + require.Equal(t, personalUsageID, recommendedID) + require.Equal(t, "used_by_you", reason) + + hint, recommendedID, reason = selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: uuid.New()}, ActiveDevelopers: 2}, + {Template: database.Template{ID: uuid.New()}}, + }, + 2, + templateRankingSignalErrors{Usage: enrichmentErr}, + now, + ) + require.Equal(t, listTemplatesHintNoConfidence, hint) + require.Equal(t, uuid.Nil, recommendedID) + require.Equal(t, "ranking_signals_unavailable", reason) + hint, recommendedID, reason = selectTemplateRecommendation( []rankedTemplate{ {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, }, 2, - enrichmentErr, + enrichmentErrors, now, ) require.Equal(t, listTemplatesHintNoConfidence, hint) diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index 04d066012b1..f6568313263 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -142,7 +142,7 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { displayTemplate := dbgen.Template(t, db, database.Template{ OrganizationID: org.ID, CreatedBy: user.ID, - Name: "data-science", + Name: "tpl-42", DisplayName: "Data Science Lab", }) descriptionTemplate := dbgen.Template(t, db, database.Template{ @@ -177,6 +177,14 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) 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, float64(0), result["total_count"]) + require.Equal(t, float64(3), result["available_template_count"]) + require.Equal(t, "no_confident_match", result["selection_hint"]) + require.Equal(t, "no_matching_templates", result["recommendation_reason"]) + result = runListTemplates(ctx, t, tool, `{"query":"does-not-exist"}`) templates = listTemplateItems(t, result) require.Empty(t, templates) @@ -297,6 +305,12 @@ func TestListTemplates_RanksAllCandidatesBeforePagination(t *testing.T) { require.Equal(t, "used_by_you", templates[0]["relevance_signals"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) 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"]) + require.Equal(t, float64(11), templates[0]["rank"]) } func TestListTemplates_QueryRelevanceOutranksPersonalUsage(t *testing.T) { From 9d1ee5096db2d1fd7c15aabeb3a390173d703dc0 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 1 Jun 2026 13:47:25 +0000 Subject: [PATCH 06/24] fix(coderd/x/chatd): require template selection on weak matches --- coderd/x/chatd/chatd_test.go | 174 ++++++++++++++++++ coderd/x/chatd/chattool/createworkspace.go | 4 + coderd/x/chatd/chattool/listtemplates.go | 9 + coderd/x/chatd/chattool/listtemplates_test.go | 3 + coderd/x/chatd/prompt.go | 1 + 5 files changed, 191 insertions(+) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 9e67c102301..bbfc2dc23e4 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -8577,6 +8577,180 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) { "create_workspace for blocked template should be rejected") } +func TestChatAsksUserWhenListTemplatesRequiresSelection(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + + var tplCode, tplDocker database.Template + var callCount atomic.Int32 + var sawHardRule atomic.Bool + var sawSelectionRequiredResult atomic.Bool + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + switch callCount.Add(1) { + case 1: + promptAndTools := string(req.RawBody) + for _, message := range req.Messages { + promptAndTools += "\n" + message.Content + } + if strings.Contains(promptAndTools, "If user_selection_required is true") && + strings.Contains(promptAndTools, "do not call create_workspace") { + sawHardRule.Store(true) + } + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("list_templates", `{}`), + ) + case 2: + if listTemplatesResultRequiresUserSelection(req.Messages) { + sawSelectionRequiredResult.Store(true) + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks( + "I found two templates, typescript-alpha and Docker Containers. Which template should I use?", + )..., + ) + } + + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("create_workspace", + fmt.Sprintf(`{"template_id":%q}`, tplCode.ID.String())), + ) + default: + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("Done.")..., + ) + } + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + tplCode = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user.ID, + Name: "code-2", + DisplayName: "typescript-alpha", + Description: "this is a long description", + }) + tplDocker = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user.ID, + Name: "docker", + DisplayName: "Docker Containers", + Description: "Provision Docker containers as Coder workspaces", + }) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.CreateWorkspace = func( + context.Context, + uuid.UUID, + codersdk.CreateWorkspaceRequest, + ) (codersdk.Workspace, error) { + t.Error("create_workspace should not be called when list_templates requires user selection") + return codersdk.Workspace{}, xerrors.New("unexpected create_workspace call") + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "ask-template-selection-test", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Create a workspace."), + }, + }) + require.NoError(t, err) + + var chatResult database.Chat + require.Eventually(t, func() bool { + got, getErr := db.GetChatByID(ctx, chat.ID) + if getErr != nil { + return false + } + chatResult = got + return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError + }, testutil.WaitLong, testutil.IntervalFast) + + if chatResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat run failed", "last_error=%q", chatLastErrorMessage(chatResult.LastError)) + } + + require.True(t, sawHardRule.Load(), "model request should include the user-selection hard rule") + require.True(t, sawSelectionRequiredResult.Load(), "model should receive a list_templates result requiring user selection") + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + + var listTemplatesResult map[string]any + var assistantText string + var sawCreateWorkspaceResult bool + for _, message := range messages { + parts, parseErr := chatprompt.ParseContent(message) + require.NoError(t, parseErr) + for _, part := range parts { + switch { + case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == "list_templates": + require.NoError(t, json.Unmarshal(part.Result, &listTemplatesResult)) + case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == "create_workspace": + sawCreateWorkspaceResult = true + case message.Role == database.ChatMessageRoleAssistant && part.Type == codersdk.ChatMessagePartTypeText: + assistantText += part.Text + } + } + } + + require.NotNil(t, listTemplatesResult, "expected list_templates tool result") + require.Equal(t, "no_confident_match", listTemplatesResult["selection_hint"]) + require.Equal(t, "no_ranking_signal", listTemplatesResult["recommendation_reason"]) + require.Equal(t, true, listTemplatesResult["user_selection_required"]) + require.NotContains(t, listTemplatesResult, "recommended_template_id") + require.Contains(t, listTemplatesResult["templates"], any(map[string]any{ + "id": tplCode.ID.String(), + "name": "code-2", + "organization_id": org.ID.String(), + "display_name": "typescript-alpha", + "description": "this is a long description", + "rank": float64(1), + "relevance_signals": "ordered_by_name", + })) + require.Contains(t, listTemplatesResult["templates"], any(map[string]any{ + "id": tplDocker.ID.String(), + "name": "docker", + "organization_id": org.ID.String(), + "display_name": "Docker Containers", + "description": "Provision Docker containers as Coder workspaces", + "rank": float64(2), + "relevance_signals": "ordered_by_name", + })) + require.False(t, sawCreateWorkspaceResult, "agent should ask instead of calling create_workspace") + require.Contains(t, assistantText, "Which template should I use?") +} + +func listTemplatesResultRequiresUserSelection(messages []chattest.OpenAIMessage) bool { + for _, message := range messages { + if message.Role != "tool" || !json.Valid([]byte(message.Content)) { + continue + } + + var result map[string]any + if err := json.Unmarshal([]byte(message.Content), &result); err != nil { + continue + } + required, _ := result["user_selection_required"].(bool) + if result["selection_hint"] == "no_confident_match" && required { + return true + } + } + return false +} + // TestSignalWakeImmediateAcquisition verifies that CreateChat triggers // immediate processing via signalWake without waiting for the polling // ticker to fire. The ticker interval is set to an hour so it never diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index 78460b53452..010ce9a19de 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -96,6 +96,10 @@ func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, option "context, provider tools, or external MCP tools. Requires a "+ "template_id from list_templates. Use recommended_template_id "+ "or rank 1 when list_templates reports a confident choice. "+ + "If list_templates returned user_selection_required true, "+ + "or selection_hint no_confident_match or ambiguous_top_matches, "+ + "do not call create_workspace. Ask the user to choose a "+ + "template unless the user already explicitly selected one. "+ "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 "+ diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 26881fe5fd0..0e7c127b6eb 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -103,6 +103,10 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp "display name, or description. Use recommended_template_id "+ "or rank 1 as the default choice when selection_hint is "+ "only_available_template or high_confidence_recommendation. "+ + "If user_selection_required is true, or selection_hint is "+ + "no_confident_match or ambiguous_top_matches, do not call "+ + "create_workspace. Ask the user to choose a template unless "+ + "the user already explicitly selected one. "+ "Do not paginate unless the returned templates do not fit the "+ "request, selection_hint reports ambiguity or no confident match, "+ "or the user asked to browse templates. Returns 10 per page.", @@ -210,6 +214,7 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp "available_template_count": visibleTemplateCount, "selection_hint": selectionHint, "recommendation_reason": recommendationReason, + "user_selection_required": userSelectionRequired(selectionHint), } if recommendedID != uuid.Nil { result["recommended_template_id"] = recommendedID.String() @@ -374,6 +379,10 @@ func selectTemplateRecommendation( return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) } +func userSelectionRequired(selectionHint string) bool { + return selectionHint == listTemplatesHintAmbiguous || selectionHint == listTemplatesHintNoConfidence +} + func templatesAreAmbiguous(a, b rankedTemplate) bool { return templateRankSignalsFor(a) == templateRankSignalsFor(b) } diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index f6568313263..cfae13aa214 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -167,6 +167,7 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Len(t, templates, 1) require.Equal(t, displayTemplate.ID.String(), templates[0]["id"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) + require.Equal(t, false, result["user_selection_required"]) require.Equal(t, displayTemplate.ID.String(), result["recommended_template_id"]) require.Equal(t, "matches_query", templates[0]["relevance_signals"]) @@ -183,6 +184,7 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Equal(t, float64(0), result["total_count"]) require.Equal(t, float64(3), result["available_template_count"]) require.Equal(t, "no_confident_match", result["selection_hint"]) + require.Equal(t, true, result["user_selection_required"]) require.Equal(t, "no_matching_templates", result["recommendation_reason"]) result = runListTemplates(ctx, t, tool, `{"query":"does-not-exist"}`) @@ -191,6 +193,7 @@ func TestListTemplates_QueryMatchesDisplayNameAndDescription(t *testing.T) { require.Equal(t, float64(0), result["total_count"]) require.Equal(t, float64(3), result["available_template_count"]) require.Equal(t, "no_confident_match", result["selection_hint"]) + require.Equal(t, true, result["user_selection_required"]) require.Equal(t, "no_matching_templates", result["recommendation_reason"]) } diff --git a/coderd/x/chatd/prompt.go b/coderd/x/chatd/prompt.go index f02006a9aa1..1254ca6e8d3 100644 --- a/coderd/x/chatd/prompt.go +++ b/coderd/x/chatd/prompt.go @@ -108,6 +108,7 @@ Ask the minimum number of questions needed to define the scope together. When no workspace is attached and you need to create one: - Call list_templates with concise search terms from the user's task when the task suggests a language, framework, image, or environment. - Treat recommended_template_id, or rank 1 when selection_hint is only_available_template or high_confidence_recommendation, as the default template unless the user asked for a different template. +- If user_selection_required is true, or selection_hint is no_confident_match or ambiguous_top_matches, do not call create_workspace. Ask the user to choose a template unless the user already explicitly selected one. - Do not paginate unless selection_hint is ambiguous_top_matches or no_confident_match, no returned template fits the request, or the user asked to browse or compare templates. - Call read_template before create_workspace when you need parameter names, required parameter values, or preset IDs. Otherwise use create_workspace with the selected template_id and defaults. From 4c20d3f9d0cc08750c4cdc78b8c59b320efb532b Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 4 Jun 2026 12:43:22 +0000 Subject: [PATCH 07/24] fix(coderd): rank chat templates by frecency score Replace the lexicographic template-ranking comparator in list_templates with a frecency score (frequency discounted by recency), per reviewer feedback. - Add GetTemplateRankingSignalsByOwnerID, returning the user's recent active and recently-deleted workspace counts, last usage, and the count of distinct active developers in the org. Recently-deleted workspaces now contribute (recovering history the deleted=false filter discarded), scoped to a lookback window, and the prebuilds system user is excluded from the org popularity count. Replaces GetWorkspaceUsageGroupedByTemplateIDByOwnerID. - Compute the affinity score in Go (Wp*(active + Wd*deleted)*0.5^(age/half_life) + Wo*ln(1+org_devs)) because sqlc cannot reliably compile the parameterized decay expression; the query returns the raw signals. Weights, half-life, and lookback are explicit constants. - Recommendation confidence is now a single score comparison: a decisive query match recommends on its own, otherwise the top score must clear a floor derived from the active-developer minimum and lead the runner-up by a derived margin. Stale-but-frequent usage no longer recommends. - Replace the AsSystemRestricted call for the cross-user org count with a narrow dbauthz wrapper checking workspace-owner read plus a template-metadata read. - Clarify list_templates/read_template guidance in the detached prompt. --- coderd/database/dbauthz/dbauthz.go | 40 +- coderd/database/dbauthz/dbauthz_test.go | 21 +- coderd/database/dbmetrics/querymetrics.go | 16 +- coderd/database/dbmock/dbmock.go | 30 +- coderd/database/querier.go | 11 +- coderd/database/queries.sql.go | 172 ++++++--- coderd/database/queries/workspaces.sql | 78 +++- coderd/x/chatd/chattool/listtemplates.go | 363 +++++++++--------- .../chattool/listtemplates_internal_test.go | 250 ++++++++---- coderd/x/chatd/chattool/listtemplates_test.go | 146 ++++++- coderd/x/chatd/prompt.go | 2 +- 11 files changed, 759 insertions(+), 370 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index ca6e0203c6e..11b60746967 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4362,6 +4362,33 @@ func (q *querier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateI return q.db.GetTemplatePresetsWithPrebuilds(ctx, templateID) } +func (q *querier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) { + // The personal frecency signal reads the owner's own workspaces, which a + // user can always read. + workspaceObj := rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()) + if arg.OrganizationID != uuid.Nil { + workspaceObj = workspaceObj.InOrg(arg.OrganizationID) + } else { + workspaceObj = workspaceObj.AnyOrganization() + } + if err := q.authorizeContext(ctx, policy.ActionRead, workspaceObj); err != nil { + return nil, err + } + // The org-popularity signal is a cross-user COUNT(DISTINCT owner_id) that we + // treat as template popularity metadata, not as permission to read other + // users' workspaces. Callers only ever pass template IDs already authorized + // via GetTemplatesWithFilter, so a template read check is the minimal, + // intentional authorization here instead of a system escalation. + templateObj := rbac.ResourceTemplate.AnyOrganization() + if arg.OrganizationID != uuid.Nil { + templateObj = rbac.ResourceTemplate.InOrg(arg.OrganizationID) + } + if err := q.authorizeContext(ctx, policy.ActionRead, templateObj); err != nil { + return nil, err + } + return q.db.GetTemplateRankingSignalsByOwnerID(ctx, arg) +} + func (q *querier) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) { if err := q.authorizeTemplateInsights(ctx, arg.TemplateIDs); err != nil { return nil, err @@ -5394,19 +5421,6 @@ func (q *querier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, return q.db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs) } -func (q *querier) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { - obj := rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()) - if arg.OrganizationID != uuid.Nil { - obj = obj.InOrg(arg.OrganizationID) - } else { - obj = obj.AnyOrganization() - } - if err := q.authorizeContext(ctx, policy.ActionRead, obj); err != nil { - return nil, err - } - return q.db.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg) -} - func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceWorkspace.Type) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 63b650adda7..d9a61495f0d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3373,14 +3373,27 @@ func (s *MethodTestSuite) TestWorkspace() { // No asserts here because SQLFilter. check.Args(ws.OwnerID, emptyPreparedAuthorized{}).Asserts() })) - s.Run("GetWorkspaceUsageGroupedByTemplateIDByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - arg := database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams{ + s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + argOrg := database.GetTemplateRankingSignalsByOwnerIDParams{ OwnerID: uuid.New(), OrganizationID: uuid.New(), TemplateIDs: []uuid.UUID{uuid.New()}, } - dbm.EXPECT().GetWorkspaceUsageGroupedByTemplateIDByOwnerID(gomock.Any(), arg).Return([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow{}, nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead) + dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), argOrg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() + check.Args(argOrg).Asserts( + rbac.ResourceWorkspace.WithOwner(argOrg.OwnerID.String()).InOrg(argOrg.OrganizationID), policy.ActionRead, + rbac.ResourceTemplate.InOrg(argOrg.OrganizationID), policy.ActionRead, + ) + + argNoOrg := database.GetTemplateRankingSignalsByOwnerIDParams{ + OwnerID: uuid.New(), + TemplateIDs: []uuid.UUID{uuid.New()}, + } + dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), argNoOrg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() + check.Args(argNoOrg).Asserts( + rbac.ResourceWorkspace.WithOwner(argNoOrg.OwnerID.String()).AnyOrganization(), policy.ActionRead, + rbac.ResourceTemplate.AnyOrganization(), policy.ActionRead, + ) })) s.Run("GetWorkspaceACLByID", s.Mocked(func(dbM *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ws := testutil.Fake(s.T(), faker, database.Workspace{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 765bf0abc9c..7e91bf68ead 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2793,6 +2793,14 @@ func (m queryMetricsStore) GetTemplatePresetsWithPrebuilds(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) { + start := time.Now() + r0, r1 := m.s.GetTemplateRankingSignalsByOwnerID(ctx, arg) + m.queryLatencies.WithLabelValues("GetTemplateRankingSignalsByOwnerID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTemplateRankingSignalsByOwnerID").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) { start := time.Now() r0, r1 := m.s.GetTemplateUsageStats(ctx, arg) @@ -3641,14 +3649,6 @@ func (m queryMetricsStore) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context return r0, r1 } -func (m queryMetricsStore) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { - start := time.Now() - r0, r1 := m.s.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg) - m.queryLatencies.WithLabelValues("GetWorkspaceUsageGroupedByTemplateIDByOwnerID").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceUsageGroupedByTemplateIDByOwnerID").Inc() - return r0, r1 -} - func (m queryMetricsStore) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { start := time.Now() r0, r1 := m.s.GetWorkspaces(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 081ca4462d1..72af46e7634 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5220,6 +5220,21 @@ func (mr *MockStoreMockRecorder) GetTemplatePresetsWithPrebuilds(ctx, templateID return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatePresetsWithPrebuilds", reflect.TypeOf((*MockStore)(nil).GetTemplatePresetsWithPrebuilds), ctx, templateID) } +// GetTemplateRankingSignalsByOwnerID mocks base method. +func (m *MockStore) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg database.GetTemplateRankingSignalsByOwnerIDParams) ([]database.GetTemplateRankingSignalsByOwnerIDRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTemplateRankingSignalsByOwnerID", ctx, arg) + ret0, _ := ret[0].([]database.GetTemplateRankingSignalsByOwnerIDRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetTemplateRankingSignalsByOwnerID indicates an expected call of GetTemplateRankingSignalsByOwnerID. +func (mr *MockStoreMockRecorder) GetTemplateRankingSignalsByOwnerID(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplateRankingSignalsByOwnerID", reflect.TypeOf((*MockStore)(nil).GetTemplateRankingSignalsByOwnerID), ctx, arg) +} + // GetTemplateUsageStats mocks base method. func (m *MockStore) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) { m.ctrl.T.Helper() @@ -6825,21 +6840,6 @@ func (mr *MockStoreMockRecorder) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUniqueOwnerCountByTemplateIDs", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUniqueOwnerCountByTemplateIDs), ctx, templateIds) } -// GetWorkspaceUsageGroupedByTemplateIDByOwnerID mocks base method. -func (m *MockStore) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetWorkspaceUsageGroupedByTemplateIDByOwnerID", ctx, arg) - ret0, _ := ret[0].([]database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetWorkspaceUsageGroupedByTemplateIDByOwnerID indicates an expected call of GetWorkspaceUsageGroupedByTemplateIDByOwnerID. -func (mr *MockStoreMockRecorder) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceUsageGroupedByTemplateIDByOwnerID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceUsageGroupedByTemplateIDByOwnerID), ctx, arg) -} - // GetWorkspaces mocks base method. func (m *MockStore) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 12592fe61e5..20a555526f8 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -716,6 +716,16 @@ type sqlcQuerier interface { // It also returns the number of desired instances for each preset. // If template_id is specified, only template versions associated with that template will be returned. GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]GetTemplatePresetsWithPrebuildsRow, error) + // GetTemplateRankingSignalsByOwnerID returns the raw ranking signals for the + // given templates relative to a single owner: how many active and recently + // deleted workspaces the owner used within the lookback window, when the + // template was last used, and how many distinct developers in the organization + // currently have a non-deleted workspace on it. The affinity score itself is + // computed in Go (see listtemplates.go); the parameterized recency-decay math + // cannot be expressed through sqlc reliably, so this query returns the exact + // raw signals the score is built from. The lookback window is applied with a + // caller-computed cutoff timestamp. + GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error) GetTemplateUsageStats(ctx context.Context, arg GetTemplateUsageStatsParams) ([]TemplateUsageStat, error) GetTemplateVersionByID(ctx context.Context, id uuid.UUID) (TemplateVersion, error) GetTemplateVersionByJobID(ctx context.Context, jobID uuid.UUID) (TemplateVersion, error) @@ -901,7 +911,6 @@ type sqlcQuerier interface { GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceResource, error) GetWorkspaceResourcesCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceResource, error) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, templateIds []uuid.UUID) ([]GetWorkspaceUniqueOwnerCountByTemplateIDsRow, error) - GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) // build_params is used to filter by build parameters if present. // It has to be a CTE because the set returning function 'unnest' cannot // be used in a WHERE clause. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 65e066f25b5..b0acd5db7b4 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -34767,6 +34767,124 @@ func (q *sqlQuerier) GetRegularWorkspaceCreateMetrics(ctx context.Context) ([]Ge return items, nil } +const getTemplateRankingSignalsByOwnerID = `-- name: GetTemplateRankingSignalsByOwnerID :many +WITH org_usage AS ( + -- org_usage measures how many distinct developers currently have a + -- non-deleted workspace on each template. The prebuilds system user is + -- excluded so unclaimed prebuilds do not inflate popularity. + SELECT + w.template_id, + COUNT(DISTINCT w.owner_id) AS org_devs + FROM + workspaces w + WHERE + w.template_id = ANY($1 :: uuid[]) + AND NOT w.deleted + AND w.owner_id != $2 :: uuid + AND CASE + WHEN $3 :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN + w.organization_id = $3 + ELSE true + END + GROUP BY + w.template_id +), +user_usage AS ( + -- user_usage counts workspaces owned by the requesting user within the + -- lookback window, splitting active from recently deleted so deleted + -- history can be counted at reduced weight. The window is keyed on + -- last_used_at. + SELECT + w.template_id, + COUNT(*) FILTER (WHERE NOT w.deleted) AS active_count, + COUNT(*) FILTER (WHERE w.deleted) AS deleted_recent_count, + MAX(w.last_used_at) :: timestamptz AS last_used_at + FROM + workspaces w + WHERE + w.owner_id = $4 + AND w.template_id = ANY($1 :: uuid[]) + AND w.last_used_at > $5 :: timestamptz + AND CASE + WHEN $3 :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN + w.organization_id = $3 + ELSE true + END + GROUP BY + w.template_id +) +SELECT + t.template_id :: uuid AS template_id, + COALESCE(u.active_count, 0) :: bigint AS active_count, + COALESCE(u.deleted_recent_count, 0) :: bigint AS deleted_recent_count, + u.last_used_at, + COALESCE(o.org_devs, 0) :: bigint AS org_devs +FROM + unnest($1 :: uuid[]) AS t(template_id) +LEFT JOIN user_usage u ON u.template_id = t.template_id +LEFT JOIN org_usage o ON o.template_id = t.template_id +` + +type GetTemplateRankingSignalsByOwnerIDParams struct { + TemplateIDs []uuid.UUID `db:"template_ids" json:"template_ids"` + PrebuildsUserID uuid.UUID `db:"prebuilds_user_id" json:"prebuilds_user_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + LookbackCutoff time.Time `db:"lookback_cutoff" json:"lookback_cutoff"` +} + +type GetTemplateRankingSignalsByOwnerIDRow struct { + TemplateID uuid.UUID `db:"template_id" json:"template_id"` + ActiveCount int64 `db:"active_count" json:"active_count"` + DeletedRecentCount int64 `db:"deleted_recent_count" json:"deleted_recent_count"` + LastUsedAt sql.NullTime `db:"last_used_at" json:"last_used_at"` + OrgDevs int64 `db:"org_devs" json:"org_devs"` +} + +// GetTemplateRankingSignalsByOwnerID returns the raw ranking signals for the +// given templates relative to a single owner: how many active and recently +// deleted workspaces the owner used within the lookback window, when the +// template was last used, and how many distinct developers in the organization +// currently have a non-deleted workspace on it. The affinity score itself is +// computed in Go (see listtemplates.go); the parameterized recency-decay math +// cannot be expressed through sqlc reliably, so this query returns the exact +// raw signals the score is built from. The lookback window is applied with a +// caller-computed cutoff timestamp. +func (q *sqlQuerier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error) { + rows, err := q.db.QueryContext(ctx, getTemplateRankingSignalsByOwnerID, + pq.Array(arg.TemplateIDs), + arg.PrebuildsUserID, + arg.OrganizationID, + arg.OwnerID, + arg.LookbackCutoff, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetTemplateRankingSignalsByOwnerIDRow + for rows.Next() { + var i GetTemplateRankingSignalsByOwnerIDRow + if err := rows.Scan( + &i.TemplateID, + &i.ActiveCount, + &i.DeletedRecentCount, + &i.LastUsedAt, + &i.OrgDevs, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getWorkspaceACLByID = `-- name: GetWorkspaceACLByID :one SELECT group_acl as groups, @@ -35146,60 +35264,6 @@ func (q *sqlQuerier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Conte return items, nil } -const getWorkspaceUsageGroupedByTemplateIDByOwnerID = `-- name: GetWorkspaceUsageGroupedByTemplateIDByOwnerID :many -SELECT - template_id, - COUNT(*) AS workspace_count, - MAX(last_used_at)::timestamptz AS last_used_at -FROM - workspaces -WHERE - owner_id = $1 - AND deleted = false - AND CASE - WHEN $2 :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN - organization_id = $2 - ELSE true - END - AND template_id = ANY($3 :: uuid[]) -GROUP BY template_id -` - -type GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams struct { - OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` - OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` - TemplateIDs []uuid.UUID `db:"template_ids" json:"template_ids"` -} - -type GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow struct { - TemplateID uuid.UUID `db:"template_id" json:"template_id"` - WorkspaceCount int64 `db:"workspace_count" json:"workspace_count"` - LastUsedAt time.Time `db:"last_used_at" json:"last_used_at"` -} - -func (q *sqlQuerier) GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx context.Context, arg GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams) ([]GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow, error) { - rows, err := q.db.QueryContext(ctx, getWorkspaceUsageGroupedByTemplateIDByOwnerID, arg.OwnerID, arg.OrganizationID, pq.Array(arg.TemplateIDs)) - if err != nil { - return nil, err - } - defer rows.Close() - var items []GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow - for rows.Next() { - var i GetWorkspaceUsageGroupedByTemplateIDByOwnerIDRow - if err := rows.Scan(&i.TemplateID, &i.WorkspaceCount, &i.LastUsedAt); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const getWorkspaces = `-- name: GetWorkspaces :many WITH build_params AS ( diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql index c860b7b0afe..225e6cacdca 100644 --- a/coderd/database/queries/workspaces.sql +++ b/coderd/database/queries/workspaces.sql @@ -497,23 +497,71 @@ LEFT JOIN workspaces ON workspaces.template_id = templates.id AND workspaces.del WHERE templates.id = ANY(@template_ids :: uuid[]) GROUP BY templates.id; --- name: GetWorkspaceUsageGroupedByTemplateIDByOwnerID :many +-- name: GetTemplateRankingSignalsByOwnerID :many +-- GetTemplateRankingSignalsByOwnerID returns the raw ranking signals for the +-- given templates relative to a single owner: how many active and recently +-- deleted workspaces the owner used within the lookback window, when the +-- template was last used, and how many distinct developers in the organization +-- currently have a non-deleted workspace on it. The affinity score itself is +-- computed in Go (see listtemplates.go); the parameterized recency-decay math +-- cannot be expressed through sqlc reliably, so this query returns the exact +-- raw signals the score is built from. The lookback window is applied with a +-- caller-computed cutoff timestamp. +WITH org_usage AS ( + -- org_usage measures how many distinct developers currently have a + -- non-deleted workspace on each template. The prebuilds system user is + -- excluded so unclaimed prebuilds do not inflate popularity. + SELECT + w.template_id, + COUNT(DISTINCT w.owner_id) AS org_devs + FROM + workspaces w + WHERE + w.template_id = ANY(@template_ids :: uuid[]) + AND NOT w.deleted + AND w.owner_id != @prebuilds_user_id :: uuid + AND CASE + WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN + w.organization_id = @organization_id + ELSE true + END + GROUP BY + w.template_id +), +user_usage AS ( + -- user_usage counts workspaces owned by the requesting user within the + -- lookback window, splitting active from recently deleted so deleted + -- history can be counted at reduced weight. The window is keyed on + -- last_used_at. + SELECT + w.template_id, + COUNT(*) FILTER (WHERE NOT w.deleted) AS active_count, + COUNT(*) FILTER (WHERE w.deleted) AS deleted_recent_count, + MAX(w.last_used_at) :: timestamptz AS last_used_at + FROM + workspaces w + WHERE + w.owner_id = @owner_id + AND w.template_id = ANY(@template_ids :: uuid[]) + AND w.last_used_at > @lookback_cutoff :: timestamptz + AND CASE + WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000' :: uuid THEN + w.organization_id = @organization_id + ELSE true + END + GROUP BY + w.template_id +) SELECT - template_id, - COUNT(*) AS workspace_count, - MAX(last_used_at)::timestamptz AS last_used_at + t.template_id :: uuid AS template_id, + COALESCE(u.active_count, 0) :: bigint AS active_count, + COALESCE(u.deleted_recent_count, 0) :: bigint AS deleted_recent_count, + u.last_used_at, + COALESCE(o.org_devs, 0) :: bigint AS org_devs FROM - workspaces -WHERE - owner_id = @owner_id - AND deleted = false - AND CASE - WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN - organization_id = @organization_id - ELSE true - END - AND template_id = ANY(@template_ids :: uuid[]) -GROUP BY template_id; + unnest(@template_ids :: uuid[]) AS t(template_id) +LEFT JOIN user_usage u ON u.template_id = t.template_id +LEFT JOIN org_usage o ON o.template_id = t.template_id; -- name: InsertWorkspace :one INSERT INTO diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 0e7c127b6eb..715d2261608 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -5,6 +5,7 @@ import ( "context" "database/sql" "maps" + "math" "slices" "strings" "time" @@ -24,11 +25,55 @@ import ( const ( listTemplatesPageSize = 10 - listTemplatesMinPersonalWorkspacesForRecommendation = 2 - listTemplatesMinActiveDevelopersForRecommendation = 2 - listTemplatesRecentUsageWindow = 90 * 24 * time.Hour + // listTemplatesMinActiveDevelopersForRecommendation is the organization + // popularity floor: a template needs at least this many active developers + // before organization popularity on its own is a confident recommendation. + listTemplatesMinActiveDevelopersForRecommendation = 2 + + // The following constants parameterize the affinity score, a "frecency" + // signal (frequency discounted by recency). The personal term is the count + // of the user's recent workspaces (active plus a fraction of + // recently-deleted) multiplied by a recency decay; the organization term is + // a log-scaled active-developer count. Only the ratio of the personal to + // organization weight matters. They are deliberately explicit so the + // ranking can be calibrated as ranking-quality signal accrues. + // + // The score is computed in Go (computeAffinityScore) rather than SQL + // because sqlc cannot reliably compile the parameterized decay expression; + // see GetTemplateRankingSignalsByOwnerID. Keeping the score and the + // confidence thresholds in the same place also avoids Postgres-versus-Go + // floating-point differences at confidence boundaries. + listTemplatesLookbackDays = 60 + listTemplatesHalfLife = 14 * 24 * time.Hour + listTemplatesPersonalWeight = 10.0 + listTemplatesOrgWeight = 1.0 + listTemplatesDeletedWeight = 0.5 ) +var ( + // minConfidentAffinityScore preserves today's floor: organization + // popularity alone is confident once a template reaches the active-developer + // minimum. math.Log1p(n) == ln(1+n) is exactly the organization term of the + // affinity score, so the threshold and the score stay float-consistent. + minConfidentAffinityScore = listTemplatesOrgWeight * math.Log1p(listTemplatesMinActiveDevelopersForRecommendation) + + // minConfidentGap requires rank 1 to lead rank 2 by at least the score + // difference between "min" and "min-1" active developers before + // recommending when both clear the floor. It is derived, not tuned, so + // "2 developers versus 1" still recommends while "16 versus 15" does not. + minConfidentGap = listTemplatesOrgWeight * (math.Log1p(listTemplatesMinActiveDevelopersForRecommendation) - math.Log1p(listTemplatesMinActiveDevelopersForRecommendation-1)) +) + +// affinityScoreEpsilon absorbs floating-point rounding so a score sitting +// exactly on a threshold boundary counts as meeting it. +const affinityScoreEpsilon = 1e-9 + +// affinityScoreAtLeast reports whether score meets threshold within the +// comparison epsilon. +func affinityScoreAtLeast(score, threshold float64) bool { + return score >= threshold-affinityScoreEpsilon +} + const ( listTemplatesHintOnlyAvailable = "only_available_template" listTemplatesHintHighConfidence = "high_confidence_recommendation" @@ -57,32 +102,30 @@ type listTemplatesArgs struct { } type rankedTemplate struct { - Template database.Template - QueryScore int - ActiveDevelopers int64 - Usage templateUsage - Rank int + Template database.Template + QueryScore int + Signals templateRankingSignals + AffinityScore float64 + Rank int } -type templateUsage struct { - WorkspaceCount int64 - LastUsedAt time.Time +// templateRankingSignals holds the raw, per-template ranking inputs returned by +// GetTemplateRankingSignalsByOwnerID. ActiveCount and DeletedRecentCount are the +// user's in-window workspace counts; LastUsedAt is the most recent usage within +// the window (zero when there is none); OrgDevs is the count of distinct active +// developers in the organization. +type templateRankingSignals struct { + ActiveCount int64 + DeletedRecentCount int64 + LastUsedAt time.Time + OrgDevs int64 } -type templateRankSignals struct { - QueryScore int - WorkspaceCount int64 - LastUsedAtUnixNano int64 - ActiveDevelopers int64 -} - -type templateRankingSignalErrors struct { - ActiveDeveloperCounts error - Usage error -} - -func (e templateRankingSignalErrors) hasAny() bool { - return e.ActiveDeveloperCounts != nil || e.Usage != nil +// hasPersonalUsage reports whether the user used the template within the +// lookback window, counting recently-deleted workspaces so deleted history is +// still treated as personal usage. +func (s templateRankingSignals) hasPersonalUsage() bool { + return s.ActiveCount+s.DeletedRecentCount > 0 } // ListTemplates returns a tool that lists available workspace templates. @@ -145,39 +188,29 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp for i, t := range ranked { templateIDs[i] = t.Template.ID } - ownerCounts, ownerCountsErr := loadTemplateActiveDeveloperCounts(ctx, db, templateIDs) - if ownerCountsErr != nil { - options.Logger.Warn(ctx, "failed to load template active developer counts", - slog.F("template_count", len(templateIDs)), - slog.Error(ownerCountsErr), - ) - } - usageByTemplate, usageErr := loadTemplateUsage( - ctx, db, options.OwnerID, organizationID, templateIDs, + now := clock.Now() + signalsByTemplate, signalsErr := loadTemplateRankingSignals( + ctx, db, options.OwnerID, organizationID, templateIDs, now, ) - if usageErr != nil { - options.Logger.Warn(ctx, "failed to load template usage", + 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(usageErr), + slog.Error(signalsErr), ) } for i := range ranked { - ranked[i].ActiveDevelopers = ownerCounts[ranked[i].Template.ID] - ranked[i].Usage = usageByTemplate[ranked[i].Template.ID] + ranked[i].Signals = signalsByTemplate[ranked[i].Template.ID] + ranked[i].AffinityScore = computeAffinityScore(ranked[i].Signals, now) } rankTemplates(ranked, query) selectionHint, recommendedID, recommendationReason := selectTemplateRecommendation( ranked, visibleTemplateCount, - templateRankingSignalErrors{ - ActiveDeveloperCounts: ownerCountsErr, - Usage: usageErr, - }, - clock.Now(), + signalsErr, ) // Paginate. @@ -239,65 +272,78 @@ func scoreTemplateCandidates(templates []database.Template, query string) []rank return candidates } -func loadTemplateActiveDeveloperCounts( +func loadTemplateRankingSignals( ctx context.Context, db database.Store, + ownerID uuid.UUID, + organizationID uuid.UUID, templateIDs []uuid.UUID, -) (map[uuid.UUID]int64, error) { - ownerCounts := make(map[uuid.UUID]int64) + now time.Time, +) (map[uuid.UUID]templateRankingSignals, error) { + signals := make(map[uuid.UUID]templateRankingSignals) if len(templateIDs) == 0 { - return ownerCounts, nil - } - - // Templates are already filtered with the owner's permissions. The - // aggregate count query requires system read because it spans workspace - // owners, but it only receives IDs the owner can already see. - rows, err := db.GetWorkspaceUniqueOwnerCountByTemplateIDs(dbauthz.AsSystemRestricted(ctx), templateIDs) //nolint:gocritic // see above + return signals, nil + } + + // The templates were already authorized with the owner's permissions by + // GetTemplatesWithFilter. GetTemplateRankingSignalsByOwnerID authorizes the + // owner reading their own workspaces plus a template-metadata read for the + // cross-user popularity count, so no system escalation is needed here. + 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 ownerCounts, err + return signals, err } for _, row := range rows { - ownerCounts[row.TemplateID] = row.UniqueOwnersSum + 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 ownerCounts, nil + return signals, nil } -func loadTemplateUsage( - ctx context.Context, - db database.Store, - ownerID uuid.UUID, - organizationID uuid.UUID, - templateIDs []uuid.UUID, -) (map[uuid.UUID]templateUsage, error) { - usageByTemplate := make(map[uuid.UUID]templateUsage) - if ownerID == uuid.Nil || len(templateIDs) == 0 { - return usageByTemplate, nil - } - - rows, err := db.GetWorkspaceUsageGroupedByTemplateIDByOwnerID(ctx, database.GetWorkspaceUsageGroupedByTemplateIDByOwnerIDParams{ - OwnerID: ownerID, - OrganizationID: organizationID, - TemplateIDs: templateIDs, - }) - if err != nil { - return usageByTemplate, err - } - for _, row := range rows { - usageByTemplate[row.TemplateID] = templateUsage{ - WorkspaceCount: row.WorkspaceCount, - LastUsedAt: row.LastUsedAt, +// computeAffinityScore folds the raw signals into a single "frecency" score: +// the personal workspace count (active plus a fraction of recently-deleted) +// multiplied by a recency decay, plus a log-scaled organization-popularity +// term. When the user has no in-window usage the personal term is zero and the +// score collapses to 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 } - return usageByTemplate, nil + org := listTemplatesOrgWeight * math.Log1p(float64(s.OrgDevs)) + return personal + org } +// rankTemplates orders templates by query relevance first (only when a query is +// present), then by affinity score, with template name and ID as deterministic +// tiebreakers. func rankTemplates(ranked []rankedTemplate, query string) { slices.SortStableFunc(ranked, func(a, b rankedTemplate) int { - if c := compareTemplateRankSignals( - templateRankSignalsFor(a), - templateRankSignalsFor(b), - query, - ); c != 0 { + 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 { @@ -311,42 +357,15 @@ func rankTemplates(ranked []rankedTemplate, query string) { } } -func templateRankSignalsFor(t rankedTemplate) templateRankSignals { - return templateRankSignals{ - QueryScore: t.QueryScore, - WorkspaceCount: t.Usage.WorkspaceCount, - LastUsedAtUnixNano: templateRankTime(t.Usage.LastUsedAt), - ActiveDevelopers: t.ActiveDevelopers, - } -} - -func templateRankTime(t time.Time) int64 { - if t.IsZero() { - return 0 - } - return t.UnixNano() -} - -func compareTemplateRankSignals(a, b templateRankSignals, query string) int { - if query != "" { - if c := cmp.Compare(b.QueryScore, a.QueryScore); c != 0 { - return c - } - } - if c := cmp.Compare(b.WorkspaceCount, a.WorkspaceCount); c != 0 { - return c - } - if c := cmp.Compare(b.LastUsedAtUnixNano, a.LastUsedAtUnixNano); c != 0 { - return c - } - return cmp.Compare(b.ActiveDevelopers, a.ActiveDevelopers) -} - +// selectTemplateRecommendation decides whether to recommend the top-ranked +// template or ask the user to choose. Query relevance is the primary signal: a +// decisive query match recommends on its own. Otherwise confidence comes from +// the affinity score, which must clear a floor and lead the runner-up by a +// margin before recommending. func selectTemplateRecommendation( ranked []rankedTemplate, visibleTemplateCount int, - rankingSignalErrors templateRankingSignalErrors, - now time.Time, + rankingSignalsErr error, ) (string, uuid.UUID, string) { if len(ranked) == 0 { return listTemplatesHintNoConfidence, uuid.Nil, "no_matching_templates" @@ -356,26 +375,41 @@ func selectTemplateRecommendation( if visibleTemplateCount == 1 && len(ranked) == 1 { return listTemplatesHintOnlyAvailable, top.Template.ID, "only_available_template" } - if rankingSignalErrors.hasAny() { - if templateHasDecisiveQuerySignal(ranked) { - return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) - } - if rankingSignalErrors.Usage == nil && - templateHasConfidentPersonalUsageSignal(top, now) && - (len(ranked) == 1 || !templatesAreAmbiguous(top, ranked[1])) { - return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) - } - return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" + + // A decisive query match (strictly outscoring the runner-up, or the only + // match) is a confident recommendation on its own, even when the affinity + // signals failed to load. + if top.QueryScore > 0 && (len(ranked) == 1 || top.QueryScore > ranked[1].QueryScore) { + return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) } - if !templateHasRankingSignal(top) { - return listTemplatesHintNoConfidence, uuid.Nil, "no_ranking_signal" + + // Without a decisive query tier the affinity score decides confidence, so an + // unreliable (failed) signal load means we must ask the user. + if rankingSignalsErr != nil { + return listTemplatesHintNoConfidence, uuid.Nil, "ranking_signals_unavailable" } - if len(ranked) > 1 && templatesAreAmbiguous(top, ranked[1]) { + + // Query present but the top two tie on relevance: break the tie with the + // affinity score when the gap is clear, otherwise ask the user. + if top.QueryScore > 0 { + if len(ranked) > 1 && affinityScoreAtLeast(top.AffinityScore-ranked[1].AffinityScore, minConfidentGap) { + return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) + } return listTemplatesHintAmbiguous, uuid.Nil, "top_templates_are_ambiguous" } - if !templateHasConfidentRankingSignal(top, now) { + + // No query: recommend purely on the affinity score. + if !affinityScoreAtLeast(top.AffinityScore, minConfidentAffinityScore) { + if top.AffinityScore <= 0 { + return listTemplatesHintNoConfidence, uuid.Nil, "no_ranking_signal" + } return listTemplatesHintNoConfidence, uuid.Nil, "weak_ranking_signal" } + if len(ranked) > 1 && + affinityScoreAtLeast(ranked[1].AffinityScore, minConfidentAffinityScore) && + !affinityScoreAtLeast(top.AffinityScore-ranked[1].AffinityScore, minConfidentGap) { + return listTemplatesHintAmbiguous, uuid.Nil, "top_templates_are_ambiguous" + } return listTemplatesHintHighConfidence, top.Template.ID, relevanceSignals(top) } @@ -383,42 +417,6 @@ func userSelectionRequired(selectionHint string) bool { return selectionHint == listTemplatesHintAmbiguous || selectionHint == listTemplatesHintNoConfidence } -func templatesAreAmbiguous(a, b rankedTemplate) bool { - return templateRankSignalsFor(a) == templateRankSignalsFor(b) -} - -func templateHasRankingSignal(t rankedTemplate) bool { - signals := templateRankSignalsFor(t) - return signals.QueryScore > 0 || signals.WorkspaceCount > 0 || signals.ActiveDevelopers > 0 -} - -func templateHasDecisiveQuerySignal(ranked []rankedTemplate) bool { - if len(ranked) == 0 || ranked[0].QueryScore == 0 { - return false - } - return len(ranked) == 1 || ranked[0].QueryScore > ranked[1].QueryScore -} - -func templateHasConfidentPersonalUsageSignal(t rankedTemplate, now time.Time) bool { - if t.Usage.WorkspaceCount >= listTemplatesMinPersonalWorkspacesForRecommendation { - return true - } - return t.Usage.WorkspaceCount > 0 && - !t.Usage.LastUsedAt.IsZero() && - now.Sub(t.Usage.LastUsedAt) <= listTemplatesRecentUsageWindow -} - -func templateHasConfidentRankingSignal(t rankedTemplate, now time.Time) bool { - signals := templateRankSignalsFor(t) - if signals.QueryScore > 0 { - return true - } - if templateHasConfidentPersonalUsageSignal(t, now) { - return true - } - return signals.ActiveDevelopers >= listTemplatesMinActiveDevelopersForRecommendation -} - func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { item := map[string]any{ "id": t.Template.ID.String(), @@ -433,12 +431,16 @@ func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { if desc := strings.TrimSpace(t.Template.Description); desc != "" { item["description"] = truncateRunes(desc, 200) } - if t.ActiveDevelopers > 0 { - item["active_developers"] = t.ActiveDevelopers + if t.Signals.OrgDevs > 0 { + item["active_developers"] = t.Signals.OrgDevs } - if t.Usage.WorkspaceCount > 0 { - item["your_workspace_count"] = t.Usage.WorkspaceCount - item["last_used_by_you"] = t.Usage.LastUsedAt.Format(time.RFC3339Nano) + // your_workspace_count exposes only active workspaces so deleted history is + // not surfaced to the model, though it still contributes to the score. + 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) + } } if t.Template.ID == recommendedID { item["recommended"] = true @@ -447,15 +449,16 @@ func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { } func relevanceSignals(t rankedTemplate) string { - signals := templateRankSignalsFor(t) + hasQuery := t.QueryScore > 0 + hasPersonal := t.Signals.hasPersonalUsage() switch { - case signals.QueryScore > 0 && signals.WorkspaceCount > 0: + case hasQuery && hasPersonal: return "matches_query_and_used_by_you" - case signals.QueryScore > 0: + case hasQuery: return "matches_query" - case signals.WorkspaceCount > 0: + case hasPersonal: return "used_by_you" - case signals.ActiveDevelopers > 0: + case t.Signals.OrgDevs > 0: return "popular_in_org" default: return "ordered_by_name" diff --git a/coderd/x/chatd/chattool/listtemplates_internal_test.go b/coderd/x/chatd/chattool/listtemplates_internal_test.go index 6ee20a28b9e..4dc9479608f 100644 --- a/coderd/x/chatd/chattool/listtemplates_internal_test.go +++ b/coderd/x/chatd/chattool/listtemplates_internal_test.go @@ -1,6 +1,7 @@ package chattool import ( + "math" "testing" "time" @@ -11,84 +12,183 @@ import ( "github.com/coder/coder/v2/coderd/database" ) -func TestSelectTemplateRecommendationRankingSignalsUnavailable(t *testing.T) { +func TestComputeAffinityScore(t *testing.T) { t.Parallel() - enrichmentErr := xerrors.New("enrichment failed") - enrichmentErrors := templateRankingSignalErrors{ - ActiveDeveloperCounts: enrichmentErr, - Usage: enrichmentErr, - } now := time.Date(2026, 5, 15, 12, 0, 0, 0, time.UTC) - onlyTemplateID := uuid.New() - hint, recommendedID, reason := selectTemplateRecommendation( - []rankedTemplate{{Template: database.Template{ID: onlyTemplateID}}}, - 1, - enrichmentErrors, - now, - ) - require.Equal(t, listTemplatesHintOnlyAvailable, hint) - require.Equal(t, onlyTemplateID, recommendedID) - require.Equal(t, "only_available_template", reason) - - topID := uuid.New() - hint, recommendedID, reason = selectTemplateRecommendation( - []rankedTemplate{ - {Template: database.Template{ID: topID}, QueryScore: queryScoreExactName}, - {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreDescriptionMatch}, - }, - 2, - enrichmentErrors, - now, - ) - require.Equal(t, listTemplatesHintHighConfidence, hint) - require.Equal(t, topID, recommendedID) - require.Equal(t, "matches_query", reason) - - personalUsageID := uuid.New() - hint, recommendedID, reason = selectTemplateRecommendation( - []rankedTemplate{ - { - Template: database.Template{ID: personalUsageID}, - Usage: templateUsage{ - WorkspaceCount: 3, - LastUsedAt: now.Add(-180 * 24 * time.Hour), - }, - }, - {Template: database.Template{ID: uuid.New()}}, - }, - 2, - templateRankingSignalErrors{ActiveDeveloperCounts: enrichmentErr}, - now, - ) - require.Equal(t, listTemplatesHintHighConfidence, hint) - require.Equal(t, personalUsageID, recommendedID) - require.Equal(t, "used_by_you", reason) - - hint, recommendedID, reason = selectTemplateRecommendation( - []rankedTemplate{ - {Template: database.Template{ID: uuid.New()}, ActiveDevelopers: 2}, - {Template: database.Template{ID: uuid.New()}}, - }, - 2, - templateRankingSignalErrors{Usage: enrichmentErr}, - now, - ) - require.Equal(t, listTemplatesHintNoConfidence, hint) - require.Equal(t, uuid.Nil, recommendedID) - require.Equal(t, "ranking_signals_unavailable", reason) - - hint, recommendedID, reason = selectTemplateRecommendation( - []rankedTemplate{ - {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, - {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, - }, - 2, - enrichmentErrors, - now, + // 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), ) - require.Equal(t, listTemplatesHintNoConfidence, hint) - require.Equal(t, uuid.Nil, recommendedID) - require.Equal(t, "ranking_signals_unavailable", reason) + + // 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("NoMatches", func(t *testing.T) { + t.Parallel() + hint, id, reason := selectTemplateRecommendation(nil, 0, nil) + require.Equal(t, listTemplatesHintNoConfidence, hint) + require.Equal(t, uuid.Nil, id) + require.Equal(t, "no_matching_templates", reason) + }) + + t.Run("OnlyAvailable", func(t *testing.T) { + t.Parallel() + only := uuid.New() + hint, id, reason := selectTemplateRecommendation( + []rankedTemplate{{Template: database.Template{ID: only}}}, 1, loadErr, + ) + require.Equal(t, listTemplatesHintOnlyAvailable, hint) + require.Equal(t, only, id) + require.Equal(t, "only_available_template", reason) + }) + + t.Run("DecisiveQueryRecommendsEvenWithLoadError", func(t *testing.T) { + t.Parallel() + top := uuid.New() + for _, err := range []error{nil, loadErr} { + hint, id, reason := selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: top}, QueryScore: queryScoreExactName}, + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreDescriptionMatch}, + }, 2, err, + ) + require.Equal(t, listTemplatesHintHighConfidence, hint) + require.Equal(t, top, id) + require.Equal(t, "matches_query", reason) + } + }) + + t.Run("QueryTieBrokenByAffinityGap", func(t *testing.T) { + t.Parallel() + top := uuid.New() + hint, id, reason := 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, listTemplatesHintHighConfidence, hint) + require.Equal(t, top, id) + require.Equal(t, "matches_query_and_used_by_you", reason) + }) + + t.Run("QueryTieWithSmallGapIsAmbiguous", func(t *testing.T) { + t.Parallel() + hint, id, _ := 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, listTemplatesHintAmbiguous, hint) + require.Equal(t, uuid.Nil, id) + }) + + t.Run("QueryTieWithLoadErrorIsUnavailable", func(t *testing.T) { + t.Parallel() + hint, id, reason := selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, + {Template: database.Template{ID: uuid.New()}, QueryScore: queryScoreNamePrefix}, + }, 2, loadErr, + ) + require.Equal(t, listTemplatesHintNoConfidence, hint) + require.Equal(t, uuid.Nil, id) + require.Equal(t, "ranking_signals_unavailable", reason) + }) + + t.Run("NoQueryNoSignal", func(t *testing.T) { + t.Parallel() + hint, _, reason := selectTemplateRecommendation( + []rankedTemplate{ + {Template: database.Template{ID: uuid.New()}}, + {Template: database.Template{ID: uuid.New()}}, + }, 2, nil, + ) + require.Equal(t, listTemplatesHintNoConfidence, hint) + require.Equal(t, "no_ranking_signal", reason) + }) + + t.Run("NoQueryWeakSignalBelowFloor", func(t *testing.T) { + t.Parallel() + // One active developer scores ln(2), below the ln(3) floor. + hint, _, reason := 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, listTemplatesHintNoConfidence, hint) + require.Equal(t, "weak_ranking_signal", reason) + }) + + t.Run("NoQueryConfidentWhenLeadsRunnerUp", func(t *testing.T) { + t.Parallel() + top := uuid.New() + hint, id, reason := 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, listTemplatesHintHighConfidence, hint) + require.Equal(t, top, id) + require.Equal(t, "popular_in_org", reason) + }) + + t.Run("NoQueryAmbiguousWhenBothClearFloorAndClose", func(t *testing.T) { + t.Parallel() + hint, id, _ := 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, listTemplatesHintAmbiguous, hint) + require.Equal(t, uuid.Nil, id) + }) + + t.Run("NoQueryLoadErrorIsUnavailable", func(t *testing.T) { + t.Parallel() + hint, id, reason := 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, listTemplatesHintNoConfidence, hint) + require.Equal(t, uuid.Nil, id) + require.Equal(t, "ranking_signals_unavailable", reason) + }) } diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index cfae13aa214..cd2c290af11 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -343,7 +343,7 @@ func TestListTemplates_QueryRelevanceOutranksPersonalUsage(t *testing.T) { OwnerID: user.ID, OrganizationID: org.ID, TemplateID: used.ID, - LastUsedAt: time.Date(2026, 5, 2, 12, 0, 0, 0, time.UTC), + LastUsedAt: time.Now().Add(-14 * 24 * time.Hour), }) tool := chattool.ListTemplates(db, org.ID, chattool.ListTemplatesOptions{ @@ -533,14 +533,17 @@ func TestListTemplates_StalePersonalUsageDoesNotRecommend(t *testing.T) { require.Len(t, templates, 2) require.Equal(t, oldUsage.ID.String(), templates[0]["id"]) require.Equal(t, unused.ID.String(), templates[1]["id"]) - require.Equal(t, float64(1), templates[0]["your_workspace_count"]) + // The 180-day-old workspace is outside the 60-day lookback window, so it no + // longer counts as in-window personal usage. + _, hasCount := templates[0]["your_workspace_count"] + require.False(t, hasCount) require.Equal(t, "no_confident_match", result["selection_hint"]) require.Equal(t, "weak_ranking_signal", result["recommendation_reason"]) _, ok := result["recommended_template_id"] require.False(t, ok) } -func TestListTemplates_PersonalUsageCountRecommendsStaleTemplate(t *testing.T) { +func TestListTemplates_StaleFrequentPersonalUsageDoesNotRecommend(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) clock := quartz.NewMock(t) @@ -564,6 +567,9 @@ func TestListTemplates_PersonalUsageCountRecommendsStaleTemplate(t *testing.T) { CreatedBy: user.ID, Name: "unused", }) + // Two workspaces used 180 days ago. Frequency no longer dominates recency: + // usage outside the lookback window decays out of the personal signal, so a + // frequently-but-stalely-used template is no longer a confident match. for range 2 { dbgen.Workspace(t, db, database.WorkspaceTable{ OwnerID: user.ID, @@ -582,9 +588,64 @@ func TestListTemplates_PersonalUsageCountRecommendsStaleTemplate(t *testing.T) { 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, "no_confident_match", result["selection_hint"]) + require.Equal(t, "weak_ranking_signal", result["recommendation_reason"]) + _, ok := result["recommended_template_id"] + require.False(t, ok) + // The stale workspaces fall outside the lookback window, so no in-window + // personal count is surfaced. + _, 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", + }) + // Two workspaces used two days ago. Recent, in-window usage is a confident + // signal: this is the frecency improvement over the old count-only ranking. + 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, "used_by_you", templates[0]["relevance_signals"]) require.Equal(t, "high_confidence_recommendation", result["selection_hint"]) - require.Equal(t, staleUsage.ID.String(), result["recommended_template_id"]) + require.Equal(t, recentUsage.ID.String(), result["recommended_template_id"]) } func TestListTemplates_AmbiguousTopMatches(t *testing.T) { @@ -811,6 +872,83 @@ 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"}) + + // Active, in-window workspace for the requesting user. + dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, OrganizationID: org.ID, TemplateID: used.ID, + LastUsedAt: now.Add(-2 * 24 * time.Hour), + }) + // 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: now.Add(-3 * 24 * time.Hour), Deleted: true, + }) + // Non-deleted but outside the lookback window: it must not count toward the + // in-window active count, though it still keeps the user in the org count. + 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) + + 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, diff --git a/coderd/x/chatd/prompt.go b/coderd/x/chatd/prompt.go index 1254ca6e8d3..0fb503032b4 100644 --- a/coderd/x/chatd/prompt.go +++ b/coderd/x/chatd/prompt.go @@ -9,7 +9,7 @@ Do not create or start a workspace by default. Many requests can be completed us Workspace tools such as execute, read_file, write_file, and edit_files require an attached workspace.` const workspaceDetachedAwareness = workspaceDetachedAwarenessBase + ` Only call create_workspace or start_workspace when the user explicitly asks for a workspace-backed task, or when the task cannot be completed without inspecting, editing, or running files in a workspace. -If a workspace is needed, use list_templates and read_template as needed before create_workspace.` +If a workspace is needed, use list_templates before create_workspace. Call read_template only when you need template parameter or preset details.` const workspaceDetachedNoCreateAwareness = workspaceDetachedAwarenessBase + ` This delegated chat cannot create or start a workspace. If workspace-backed work is required, report that need to the parent agent instead of trying workspace tools.` From dfc293d07c54b75a598393731b1a5f95995dd145 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 5 Jun 2026 11:26:14 +0000 Subject: [PATCH 08/24] fix(coderd): address template ranking review feedback Address the latest PR review feedback for frecency-based template ranking: - Authorize template ranking signals against the exact candidate template IDs using the same prepared-filter semantics as GetTemplatesWithFilter, so users with template ACL access keep ranking signals without broad org-wide template read. - Split the dbauthz mock coverage for org-scoped and any-organization calls; normalize duplicate subtest suffixes in the dbauthz method-test harness. - Surface deleted-only personal usage evidence in list_templates output with a recently-deleted count and last-used timestamp. - Assert the raw SQL query returns the maximum last_used_at value. - Clarify detached prompt guidance for user_selection_required and conditional read_template usage. - Document ListTemplatesOptions requirements and defaults. --- coderd/database/dbauthz/dbauthz.go | 33 ++++++++--- coderd/database/dbauthz/dbauthz_test.go | 30 +++++----- coderd/database/dbauthz/setup_test.go | 3 + coderd/x/chatd/chattool/listtemplates.go | 15 +++-- coderd/x/chatd/chattool/listtemplates_test.go | 55 ++++++++++++++++++- coderd/x/chatd/prompt.go | 2 +- 6 files changed, 108 insertions(+), 30 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 11b60746967..1af2a21395e 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4377,14 +4377,31 @@ func (q *querier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg da // The org-popularity signal is a cross-user COUNT(DISTINCT owner_id) that we // treat as template popularity metadata, not as permission to read other // users' workspaces. Callers only ever pass template IDs already authorized - // via GetTemplatesWithFilter, so a template read check is the minimal, - // intentional authorization here instead of a system escalation. - templateObj := rbac.ResourceTemplate.AnyOrganization() - if arg.OrganizationID != uuid.Nil { - templateObj = rbac.ResourceTemplate.InOrg(arg.OrganizationID) - } - if err := q.authorizeContext(ctx, policy.ActionRead, templateObj); err != nil { - return nil, err + // via GetTemplatesWithFilter, and we verify those exact IDs here with the + // same prepared-filter semantics so ACL-only template readers keep their + // ranking signals without requiring broad org-wide template read. + if len(arg.TemplateIDs) > 0 { + prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceTemplate.Type) + if err != nil { + return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err) + } + authorizedTemplates, err := q.db.GetAuthorizedTemplates(ctx, database.GetTemplatesWithFilterParams{ + Deleted: false, + OrganizationID: arg.OrganizationID, + IDs: arg.TemplateIDs, + }, prep) + if err != nil { + return nil, err + } + authorizedIDs := make(map[uuid.UUID]struct{}, len(authorizedTemplates)) + for _, template := range authorizedTemplates { + authorizedIDs[template.ID] = struct{}{} + } + for _, templateID := range arg.TemplateIDs { + if _, ok := authorizedIDs[templateID]; !ok { + return nil, NotAuthorizedError{Err: xerrors.Errorf("not authorized to read template %s", templateID)} + } + } } return q.db.GetTemplateRankingSignalsByOwnerID(ctx, arg) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index d9a61495f0d..cd57e5c0d55 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3374,26 +3374,30 @@ func (s *MethodTestSuite) TestWorkspace() { check.Args(ws.OwnerID, emptyPreparedAuthorized{}).Asserts() })) s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - argOrg := database.GetTemplateRankingSignalsByOwnerIDParams{ + arg := database.GetTemplateRankingSignalsByOwnerIDParams{ OwnerID: uuid.New(), OrganizationID: uuid.New(), TemplateIDs: []uuid.UUID{uuid.New()}, } - dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), argOrg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() - check.Args(argOrg).Asserts( - rbac.ResourceWorkspace.WithOwner(argOrg.OwnerID.String()).InOrg(argOrg.OrganizationID), policy.ActionRead, - rbac.ResourceTemplate.InOrg(argOrg.OrganizationID), policy.ActionRead, - ) - - argNoOrg := database.GetTemplateRankingSignalsByOwnerIDParams{ + dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{ + Deleted: false, + OrganizationID: arg.OrganizationID, + IDs: arg.TemplateIDs, + }, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes() + dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), arg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead) + })) + s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.GetTemplateRankingSignalsByOwnerIDParams{ OwnerID: uuid.New(), TemplateIDs: []uuid.UUID{uuid.New()}, } - dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), argNoOrg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() - check.Args(argNoOrg).Asserts( - rbac.ResourceWorkspace.WithOwner(argNoOrg.OwnerID.String()).AnyOrganization(), policy.ActionRead, - rbac.ResourceTemplate.AnyOrganization(), policy.ActionRead, - ) + dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{ + Deleted: false, + IDs: arg.TemplateIDs, + }, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes() + dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), arg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).AnyOrganization(), policy.ActionRead) })) s.Run("GetWorkspaceACLByID", s.Mocked(func(dbM *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ws := testutil.Fake(s.T(), faker, database.Workspace{}) diff --git a/coderd/database/dbauthz/setup_test.go b/coderd/database/dbauthz/setup_test.go index bab2cac91cf..5bf5163b91e 100644 --- a/coderd/database/dbauthz/setup_test.go +++ b/coderd/database/dbauthz/setup_test.go @@ -192,6 +192,9 @@ func (s *MethodTestSuite) SubtestWithDB(db database.Store, testCaseF func(db dat testName := s.T().Name() names := strings.Split(testName, "/") methodName := names[len(names)-1] + if baseMethodName, _, ok := strings.Cut(methodName, "#"); ok { + methodName = baseMethodName + } s.methodAccounting[methodName]++ fakeAuthorizer := &coderdtest.FakeAuthorizer{} diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 715d2261608..1f6b85fe2bc 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -88,7 +88,9 @@ const ( queryScoreDescriptionMatch = 1 ) -// ListTemplatesOptions configures the list_templates tool. +// ListTemplatesOptions configures the list_templates tool. OwnerID is required. +// Logger may be zero-valued; 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 @@ -434,13 +436,14 @@ func templateItem(t rankedTemplate, recommendedID uuid.UUID) map[string]any { if t.Signals.OrgDevs > 0 { item["active_developers"] = t.Signals.OrgDevs } - // your_workspace_count exposes only active workspaces so deleted history is - // not surfaced to the model, though it still contributes to the score. 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) - } + } + if t.Signals.DeletedRecentCount > 0 { + item["your_recently_deleted_workspace_count"] = t.Signals.DeletedRecentCount + } + if t.Signals.hasPersonalUsage() && !t.Signals.LastUsedAt.IsZero() { + item["last_used_by_you"] = t.Signals.LastUsedAt.Format(time.RFC3339Nano) } if t.Template.ID == recommendedID { item["recommended"] = true diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index cd2c290af11..f957cb5e9e7 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -648,6 +648,54 @@ func TestListTemplates_RecentPersonalUsageRecommends(t *testing.T) { 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.Equal(t, "used_by_you", templates[0]["relevance_signals"]) + require.Equal(t, float64(1), templates[0]["your_recently_deleted_workspace_count"]) + require.NotEmpty(t, templates[0]["last_used_by_you"]) + _, hasActiveCount := templates[0]["your_workspace_count"] + require.False(t, hasActiveCount) +} + func TestListTemplates_AmbiguousTopMatches(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -893,15 +941,17 @@ func TestGetTemplateRankingSignalsByOwnerID(t *testing.T) { 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: now.Add(-2 * 24 * time.Hour), + 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: now.Add(-3 * 24 * time.Hour), Deleted: true, + LastUsedAt: deletedLastUsedAt, Deleted: true, }) // Non-deleted but outside the lookback window: it must not count toward the // in-window active count, though it still keeps the user in the org count. @@ -941,6 +991,7 @@ func TestGetTemplateRankingSignalsByOwnerID(t *testing.T) { 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) diff --git a/coderd/x/chatd/prompt.go b/coderd/x/chatd/prompt.go index 0fb503032b4..d5ccedd7c31 100644 --- a/coderd/x/chatd/prompt.go +++ b/coderd/x/chatd/prompt.go @@ -9,7 +9,7 @@ Do not create or start a workspace by default. Many requests can be completed us Workspace tools such as execute, read_file, write_file, and edit_files require an attached workspace.` const workspaceDetachedAwareness = workspaceDetachedAwarenessBase + ` Only call create_workspace or start_workspace when the user explicitly asks for a workspace-backed task, or when the task cannot be completed without inspecting, editing, or running files in a workspace. -If a workspace is needed, use list_templates before create_workspace. Call read_template only when you need template parameter or preset details.` +If a workspace is needed, use list_templates before create_workspace. If list_templates returns user_selection_required or a no_confident_match or ambiguous_top_matches selection_hint, ask the user to choose before create_workspace. Call read_template only when you need template parameter or preset details.` const workspaceDetachedNoCreateAwareness = workspaceDetachedAwarenessBase + ` This delegated chat cannot create or start a workspace. If workspace-backed work is required, report that need to the parent agent instead of trying workspace tools.` From 5a7641ff662e0683346ee2800a47d360c00a509a Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 5 Jun 2026 11:52:00 +0000 Subject: [PATCH 09/24] fix(coderd/x/chatd/chattool): add list_templates auth error context Address CRF-26 by adding list_templates operation context to the user-facing asOwner authorization error response. --- coderd/x/chatd/chattool/listtemplates.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 1f6b85fe2bc..480bd237c72 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -158,7 +158,7 @@ func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemp 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{ From 42fed81df5a3e98082ff714a417bb3d08e50b603 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 5 Jun 2026 12:18:48 +0000 Subject: [PATCH 10/24] test(coderd): update chat system prompt expectation Keep TestChatSystemPrompt's detached workspace awareness expectation in sync with the updated list_templates guidance. --- coderd/exp_chats_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index fdbc1160d8a..0d93978a1e8 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -10631,7 +10631,7 @@ func TestChatSystemPrompt(t *testing.T) { const workspaceAwareness = `No workspace is attached to this chat yet. Do not create or start a workspace by default. Many requests can be completed using the conversation, provider tools such as web_search when available, or configured external MCP tools. Workspace tools such as execute, read_file, write_file, and edit_files require an attached workspace. Only call create_workspace or start_workspace when the user explicitly asks for a workspace-backed task, or when the task cannot be completed without inspecting, editing, or running files in a workspace. -If a workspace is needed, use list_templates and read_template as needed before create_workspace.` +If a workspace is needed, use list_templates before create_workspace. If list_templates returns user_selection_required or a no_confident_match or ambiguous_top_matches selection_hint, ask the user to choose before create_workspace. Call read_template only when you need template parameter or preset details.` updateChatSystemPrompt := func(t *testing.T, ctx context.Context, req codersdk.UpdateChatSystemPromptRequest) { t.Helper() From d7fe770829dad0fb2598a05b91d3222380c3f2c8 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 9 Jun 2026 13:53:25 +0000 Subject: [PATCH 11/24] chore: update comments --- coderd/database/querier.go | 7 ++++--- coderd/database/queries.sql.go | 7 ++++--- coderd/database/queries/workspaces.sql | 7 ++++--- coderd/x/chatd/chattool/listtemplates.go | 11 ++++++----- 4 files changed, 18 insertions(+), 14 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 284ae0a7482..f958e008931 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -731,9 +731,10 @@ type sqlcQuerier interface { // deleted workspaces the owner used within the lookback window, when the // template was last used, and how many distinct developers in the organization // currently have a non-deleted workspace on it. The affinity score itself is - // computed in Go (see listtemplates.go); the parameterized recency-decay math - // cannot be expressed through sqlc reliably, so this query returns the exact - // raw signals the score is built from. The lookback window is applied with a + // computed in Go (see listtemplates.go) because sqlc type inference is fragile + // around complex parameterized expressions unless inputs are explicitly cast + // and nested selects are kept simple. This query returns the exact raw signals + // the score is built from. The lookback window is applied with a // caller-computed cutoff timestamp. GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error) GetTemplateUsageStats(ctx context.Context, arg GetTemplateUsageStatsParams) ([]TemplateUsageStat, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 9bb42acae01..369ec1fc703 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -35120,9 +35120,10 @@ type GetTemplateRankingSignalsByOwnerIDRow struct { // deleted workspaces the owner used within the lookback window, when the // template was last used, and how many distinct developers in the organization // currently have a non-deleted workspace on it. The affinity score itself is -// computed in Go (see listtemplates.go); the parameterized recency-decay math -// cannot be expressed through sqlc reliably, so this query returns the exact -// raw signals the score is built from. The lookback window is applied with a +// computed in Go (see listtemplates.go) because sqlc type inference is fragile +// around complex parameterized expressions unless inputs are explicitly cast +// and nested selects are kept simple. This query returns the exact raw signals +// the score is built from. The lookback window is applied with a // caller-computed cutoff timestamp. func (q *sqlQuerier) GetTemplateRankingSignalsByOwnerID(ctx context.Context, arg GetTemplateRankingSignalsByOwnerIDParams) ([]GetTemplateRankingSignalsByOwnerIDRow, error) { rows, err := q.db.QueryContext(ctx, getTemplateRankingSignalsByOwnerID, diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql index 225e6cacdca..f615b27cd6a 100644 --- a/coderd/database/queries/workspaces.sql +++ b/coderd/database/queries/workspaces.sql @@ -503,9 +503,10 @@ GROUP BY templates.id; -- deleted workspaces the owner used within the lookback window, when the -- template was last used, and how many distinct developers in the organization -- currently have a non-deleted workspace on it. The affinity score itself is --- computed in Go (see listtemplates.go); the parameterized recency-decay math --- cannot be expressed through sqlc reliably, so this query returns the exact --- raw signals the score is built from. The lookback window is applied with a +-- computed in Go (see listtemplates.go) because sqlc type inference is fragile +-- around complex parameterized expressions unless inputs are explicitly cast +-- and nested selects are kept simple. This query returns the exact raw signals +-- the score is built from. The lookback window is applied with a -- caller-computed cutoff timestamp. WITH org_usage AS ( -- org_usage measures how many distinct developers currently have a diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 480bd237c72..0babf600fca 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -38,11 +38,12 @@ const ( // organization weight matters. They are deliberately explicit so the // ranking can be calibrated as ranking-quality signal accrues. // - // The score is computed in Go (computeAffinityScore) rather than SQL - // because sqlc cannot reliably compile the parameterized decay expression; - // see GetTemplateRankingSignalsByOwnerID. Keeping the score and the - // confidence thresholds in the same place also avoids Postgres-versus-Go - // floating-point differences at confidence boundaries. + // The score is computed in Go (computeAffinityScore) rather than SQL because + // sqlc type inference is fragile around complex parameterized expressions + // unless inputs are explicitly cast and nested selects are kept simple; see + // GetTemplateRankingSignalsByOwnerID. Keeping the score and the confidence + // thresholds in the same place also avoids Postgres-versus-Go floating-point + // differences at confidence boundaries. listTemplatesLookbackDays = 60 listTemplatesHalfLife = 14 * 24 * time.Hour listTemplatesPersonalWeight = 10.0 From 95f7beb9eeb907af2d1ad80d3d5ee03e1f806af6 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 10 Jun 2026 07:27:58 +0000 Subject: [PATCH 12/24] test(coderd/database/dbauthz): cover template ranking signals deny path --- coderd/database/dbauthz/dbauthz_test.go | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 0c5b6f5be9a..d66732e0fae 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3418,6 +3418,25 @@ func (s *MethodTestSuite) TestWorkspace() { dbm.EXPECT().GetTemplateRankingSignalsByOwnerID(gomock.Any(), arg).Return([]database.GetTemplateRankingSignalsByOwnerIDRow{}, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).AnyOrganization(), policy.ActionRead) })) + s.Run("GetTemplateRankingSignalsByOwnerID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + // Deny path: when any requested template is missing from the + // authorized set, the call is rejected and the underlying query is + // never reached. No expectation is registered for + // GetTemplateRankingSignalsByOwnerID, so reaching it fails the test. + arg := database.GetTemplateRankingSignalsByOwnerIDParams{ + OwnerID: uuid.New(), + OrganizationID: uuid.New(), + TemplateIDs: []uuid.UUID{uuid.New(), uuid.New()}, + } + dbm.EXPECT().GetAuthorizedTemplates(gomock.Any(), database.GetTemplatesWithFilterParams{ + Deleted: false, + OrganizationID: arg.OrganizationID, + IDs: arg.TemplateIDs, + }, gomock.Any()).Return([]database.Template{{ID: arg.TemplateIDs[0]}}, nil).AnyTimes() + check.Args(arg). + Asserts(rbac.ResourceWorkspace.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionRead). + Errors(dbauthz.NotAuthorizedError{Err: xerrors.Errorf("not authorized to read template %s", arg.TemplateIDs[1])}) + })) s.Run("GetWorkspaceACLByID", s.Mocked(func(dbM *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ws := testutil.Fake(s.T(), faker, database.Workspace{}) dbM.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes() From 71eca18be56d67b19d22bd836fbfe17eca0482fe Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 10 Jun 2026 07:38:10 +0000 Subject: [PATCH 13/24] refactor(coderd/x/chatd/chattool): simplify templateQueryScore tier matching --- coderd/x/chatd/chattool/listtemplates.go | 56 ++++++++++++------------ 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 0babf600fca..e159b69b694 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -471,40 +471,17 @@ func relevanceSignals(t rankedTemplate) string { func templateQueryScore(t database.Template, query string) int { query = normalizeTemplateSearch(query) - if query == "" { - return 0 - } - queryCompact := compactTemplateSearch(query) - if queryCompact == "" { + if query == "" || queryCompact == "" { return 0 } + + best := 0 for _, field := range []string{t.Name, t.DisplayName} { - field = normalizeTemplateSearch(field) - if field == "" { - continue - } - if field == query || compactTemplateSearch(field) == queryCompact { - return queryScoreExactName - } + best = max(best, nameQueryScore(field, query, queryCompact)) } - for _, field := range []string{t.Name, t.DisplayName} { - field = normalizeTemplateSearch(field) - if field == "" { - continue - } - if strings.HasPrefix(field, query) || strings.HasPrefix(compactTemplateSearch(field), queryCompact) { - return queryScoreNamePrefix - } - } - for _, field := range []string{t.Name, t.DisplayName} { - field = normalizeTemplateSearch(field) - if field == "" { - continue - } - if strings.Contains(field, query) || strings.Contains(compactTemplateSearch(field), queryCompact) { - return queryScoreNameContains - } + if best > 0 { + return best } desc := normalizeTemplateSearch(t.Description) if strings.Contains(desc, query) || strings.Contains(compactTemplateSearch(desc), queryCompact) { @@ -513,6 +490,27 @@ func templateQueryScore(t database.Template, query string) int { return 0 } +// nameQueryScore returns the relevance tier of a single name-like field: +// exact match outranks prefix match, which outranks substring match, on +// either the normalized or compact form. Returns 0 when the field does not +// 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)) } From 206c7383803f5f4f3deae393b45100ab362eba71 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 17 Jul 2026 01:31:07 +0000 Subject: [PATCH 14/24] feat(coderd): add hourly hb_agent_runtime_v1 usage events for Coder Agent runtime Adds a new heartbeat usage event type, hb_agent_runtime_v1, that measures the total agent-loop runtime (in milliseconds) of Coder Agents (chats) per UTC hour, summed from chat_messages.runtime_ms. Events are produced by a new lock-free hourly reconciler (enterprise/coderd/usage.Generator) that scans a trailing 7-day window for missing hourly buckets and fills them with deterministic IDs (hb_agent_runtime_v1:) and created_at set to the bucket start. Idle hours are zero-filled, hours missed during downtime are backfilled, and ON CONFLICT (id) DO NOTHING makes concurrent replicas safe without locking. Events flow to Tallyman through the existing publisher; the publish_usage_data license flag gates egress only. The aggregate_usage_event() trigger sums runtime_ms per day into usage_events_daily for the new type. Also fixes two latent bugs found along the way: - The ai-seats usage cron job never registered because the CronJob was missing its EventType and the Register error was discarded. - subjectUsagePublisher lacked usage_event create permission, so any heartbeat insert under AsUsagePublisher would fail authz once the cron actually ran. InsertHeartbeatUsageEvent now takes an explicit createdAt so generators can backfill historical buckets; the cron passes clock.Now() to preserve its existing behavior. Do not include this in a release until Tallyman accepts hb_agent_runtime_v1: permanent rejections are marked done-forever locally and those buckets would never be re-sent. --- coderd/coderdtest/usage.go | 3 +- coderd/database/dbauthz/dbauthz.go | 24 +- coderd/database/dbauthz/dbauthz_test.go | 21 ++ coderd/database/dbmetrics/querymetrics.go | 16 + coderd/database/dbmock/dbmock.go | 30 ++ coderd/database/dump.sql | 13 +- ...000517_usage_events_agent_runtime.down.sql | 50 +++ .../000517_usage_events_agent_runtime.up.sql | 58 +++ .../000517_usage_events_agent_runtime.up.sql | 20 + coderd/database/querier.go | 6 + coderd/database/querier_test.go | 163 ++++++++ coderd/database/queries.sql.go | 59 +++ coderd/database/queries/chats.sql | 9 + coderd/database/queries/usageevents.sql | 8 + coderd/pproflabel/pproflabel.go | 3 + coderd/usage/inserter.go | 9 +- coderd/usage/usagetypes/events.go | 41 ++ coderd/usage/usagetypes/events_test.go | 24 ++ enterprise/cli/server.go | 25 +- enterprise/coderd/usage/cron.go | 2 +- enterprise/coderd/usage/generator.go | 208 +++++++++++ enterprise/coderd/usage/generator_test.go | 350 ++++++++++++++++++ enterprise/coderd/usage/inserter.go | 5 +- enterprise/coderd/usage/inserter_test.go | 39 ++ 24 files changed, 1170 insertions(+), 16 deletions(-) create mode 100644 coderd/database/migrations/000517_usage_events_agent_runtime.down.sql create mode 100644 coderd/database/migrations/000517_usage_events_agent_runtime.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000517_usage_events_agent_runtime.up.sql create mode 100644 enterprise/coderd/usage/generator.go create mode 100644 enterprise/coderd/usage/generator_test.go diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go index c7139128670..d4d249dc863 100644 --- a/coderd/coderdtest/usage.go +++ b/coderd/coderdtest/usage.go @@ -3,6 +3,7 @@ package coderdtest import ( "context" "sync" + "time" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/usage" @@ -33,7 +34,7 @@ func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.S return nil } -func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error { +func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, _ time.Time, event usagetypes.HeartbeatEvent) error { u.Lock() defer u.Unlock() if _, seen := u.seenHeartbeats[id]; seen { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 6fe38428877..75a47bc9607 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -598,9 +598,9 @@ var ( Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceLicense.Type: {policy.ActionRead}, rbac.ResourceAiSeat.Type: {policy.ActionRead}, // Required for GetActiveAISeatCount. - // The usage publisher doesn't create events, just - // reads/processes them. - rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate}, + // Create is required by the usage cron and generator, + // which insert heartbeat events under this subject. + rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate}, }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, @@ -4626,6 +4626,17 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe return q.db.GetAuthorizedTemplates(ctx, arg, prep) } +func (q *querier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + // This query exists solely to compute hb_agent_runtime_v1 usage event + // payloads. It returns a bare runtime sum across all chats with no chat + // content, so it's gated on the ability to create usage events rather + // than on reading chats. + if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil { + return 0, err + } + return q.db.GetTotalChatMessageRuntimeMsInRange(ctx, arg) +} + func (q *querier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil { return 0, err @@ -6378,6 +6389,13 @@ func (q *querier) ListTasks(ctx context.Context, arg database.ListTasksParams) ( return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListTasks)(ctx, arg) } +func (q *querier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil { + return nil, err + } + return q.db.ListUsageEventCreatedAtsByTypeSince(ctx, arg) +} + func (q *querier) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { u, err := q.db.GetUserByID(ctx, userID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index d66732e0fae..eef5f8f90aa 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6234,6 +6234,27 @@ func (s *MethodTestSuite) TestUsageEvents() { EndDate: time.Time{}, }).Asserts(rbac.ResourceUsageEvent, policy.ActionRead) })) + + s.Run("ListUsageEventCreatedAtsByTypeSince", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: "hb_agent_runtime_v1", + Since: dbtime.Now(), + } + db.EXPECT().ListUsageEventCreatedAtsByTypeSince(gomock.Any(), params).Return([]time.Time{}, nil) + check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionRead) + })) + + // GetTotalChatMessageRuntimeMsInRange exists solely to compute usage + // event payloads, so it asserts usage event creation rather than chat + // read permissions. + s.Run("GetTotalChatMessageRuntimeMsInRange", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: time.Time{}, + EndTime: time.Time{}, + } + db.EXPECT().GetTotalChatMessageRuntimeMsInRange(gomock.Any(), params).Return(int64(0), nil) + check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate) + })) } // Ensures that the prebuilds actor may never insert an api key. diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 6f9d0e12d7c..15780947370 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2921,6 +2921,14 @@ func (m queryMetricsStore) GetTemplatesWithFilter(ctx context.Context, arg datab return r0, r1 } +func (m queryMetricsStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.GetTotalChatMessageRuntimeMsInRange(ctx, arg) + m.queryLatencies.WithLabelValues("GetTotalChatMessageRuntimeMsInRange").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTotalChatMessageRuntimeMsInRange").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { start := time.Now() r0, r1 := m.s.GetTotalUsageDCManagedAgentsV1(ctx, arg) @@ -4505,6 +4513,14 @@ func (m queryMetricsStore) ListTasks(ctx context.Context, arg database.ListTasks return r0, r1 } +func (m queryMetricsStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + start := time.Now() + r0, r1 := m.s.ListUsageEventCreatedAtsByTypeSince(ctx, arg) + m.queryLatencies.WithLabelValues("ListUsageEventCreatedAtsByTypeSince").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListUsageEventCreatedAtsByTypeSince").Inc() + return r0, r1 +} + func (m queryMetricsStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { start := time.Now() r0, r1 := m.s.ListUserChatCompactionThresholds(ctx, userID) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 4a3c88f3de8..45a32d9a9bb 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5475,6 +5475,21 @@ func (mr *MockStoreMockRecorder) GetTemplatesWithFilter(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatesWithFilter", reflect.TypeOf((*MockStore)(nil).GetTemplatesWithFilter), ctx, arg) } +// GetTotalChatMessageRuntimeMsInRange mocks base method. +func (m *MockStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTotalChatMessageRuntimeMsInRange", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetTotalChatMessageRuntimeMsInRange indicates an expected call of GetTotalChatMessageRuntimeMsInRange. +func (mr *MockStoreMockRecorder) GetTotalChatMessageRuntimeMsInRange(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTotalChatMessageRuntimeMsInRange", reflect.TypeOf((*MockStore)(nil).GetTotalChatMessageRuntimeMsInRange), ctx, arg) +} + // GetTotalUsageDCManagedAgentsV1 mocks base method. func (m *MockStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { m.ctrl.T.Helper() @@ -8519,6 +8534,21 @@ func (mr *MockStoreMockRecorder) ListTasks(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTasks", reflect.TypeOf((*MockStore)(nil).ListTasks), ctx, arg) } +// ListUsageEventCreatedAtsByTypeSince mocks base method. +func (m *MockStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListUsageEventCreatedAtsByTypeSince", ctx, arg) + ret0, _ := ret[0].([]time.Time) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListUsageEventCreatedAtsByTypeSince indicates an expected call of ListUsageEventCreatedAtsByTypeSince. +func (mr *MockStoreMockRecorder) ListUsageEventCreatedAtsByTypeSince(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListUsageEventCreatedAtsByTypeSince", reflect.TypeOf((*MockStore)(nil).ListUsageEventCreatedAtsByTypeSince), ctx, arg) +} + // ListUserChatCompactionThresholds mocks base method. func (m *MockStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 9d2b8e3fc56..ebc1ef70e19 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -681,7 +681,7 @@ CREATE FUNCTION aggregate_usage_event() RETURNS trigger AS $$ BEGIN -- Check for supported event types and throw error for unknown types. - IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; END IF; @@ -709,6 +709,13 @@ BEGIN COALESCE((NEW.event_data->>'count')::bigint, 0) ) ) + -- Hourly runtime heartbeats: sum the runtime per day. + WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN + jsonb_build_object( + 'runtime_ms', + COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) + + COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0) + ) END; RETURN NEW; @@ -3173,7 +3180,7 @@ CREATE TABLE usage_events ( publish_started_at timestamp with time zone, published_at timestamp with time zone, failure_message text, - CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text]))) + CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text, 'hb_agent_runtime_v1'::text]))) ); COMMENT ON TABLE usage_events IS 'usage_events contains usage data that is collected from the product and potentially shipped to the usage collector service.'; @@ -4375,6 +4382,8 @@ CREATE INDEX idx_template_versions_has_ai_task ON template_versions USING btree CREATE UNIQUE INDEX idx_unique_preset_name ON template_version_presets USING btree (name, template_version_id); +CREATE INDEX idx_usage_events_agent_runtime ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_agent_runtime_v1'::text); + CREATE INDEX idx_usage_events_ai_seats ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_ai_seats_v1'::text); CREATE INDEX idx_usage_events_select_for_publishing ON usage_events USING btree (published_at, publish_started_at, created_at); diff --git a/coderd/database/migrations/000517_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000517_usage_events_agent_runtime.down.sql new file mode 100644 index 00000000000..a890669ab01 --- /dev/null +++ b/coderd/database/migrations/000517_usage_events_agent_runtime.down.sql @@ -0,0 +1,50 @@ +DROP INDEX IF EXISTS idx_usage_events_agent_runtime; + +-- Remove hb_agent_runtime_v1 rows so the previous constraint can be restored. +DELETE FROM usage_events WHERE event_type = 'hb_agent_runtime_v1'; +DELETE FROM usage_events_daily WHERE event_type = 'hb_agent_runtime_v1'; + +-- Restore previous constraint. +ALTER TABLE usage_events + DROP CONSTRAINT usage_event_type_check, + ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1')); + +-- Restore the previous aggregate function without hb_agent_runtime_v1 +-- support. +CREATE OR REPLACE FUNCTION aggregate_usage_event() +RETURNS TRIGGER AS $$ +BEGIN + -- Check for supported event types and throw error for unknown types. + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN + RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; + END IF; + + INSERT INTO usage_events_daily (day, event_type, usage_data) + VALUES ( + date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date, + NEW.event_type, + NEW.event_data + ) + ON CONFLICT (day, event_type) DO UPDATE SET + usage_data = CASE + -- Handle simple counter events by summing the count. + WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN + jsonb_build_object( + 'count', + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) + + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + -- Heartbeat events: keep the max value seen that day + WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN + jsonb_build_object( + 'count', + GREATEST( + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0), + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + ) + END; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/coderd/database/migrations/000517_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000517_usage_events_agent_runtime.up.sql new file mode 100644 index 00000000000..7647917eee1 --- /dev/null +++ b/coderd/database/migrations/000517_usage_events_agent_runtime.up.sql @@ -0,0 +1,58 @@ +-- Expand the CHECK constraint to allow hb_agent_runtime_v1. +ALTER TABLE usage_events + DROP CONSTRAINT usage_event_type_check, + ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1')); + +-- Partial index for efficient lookups of agent runtime heartbeat events by +-- time. Used by the usage generator to find missing hourly buckets. +CREATE INDEX idx_usage_events_agent_runtime + ON usage_events (event_type, created_at) + WHERE event_type = 'hb_agent_runtime_v1'; + +-- Update the aggregate function to handle hb_agent_runtime_v1 events. +-- Each hourly runtime event contributes to the daily total, so they are +-- summed per day. +CREATE OR REPLACE FUNCTION aggregate_usage_event() +RETURNS TRIGGER AS $$ +BEGIN + -- Check for supported event types and throw error for unknown types. + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN + RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; + END IF; + + INSERT INTO usage_events_daily (day, event_type, usage_data) + VALUES ( + date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date, + NEW.event_type, + NEW.event_data + ) + ON CONFLICT (day, event_type) DO UPDATE SET + usage_data = CASE + -- Handle simple counter events by summing the count. + WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN + jsonb_build_object( + 'count', + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) + + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + -- Heartbeat events: keep the max value seen that day + WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN + jsonb_build_object( + 'count', + GREATEST( + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0), + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + ) + -- Hourly runtime heartbeats: sum the runtime per day. + WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN + jsonb_build_object( + 'runtime_ms', + COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) + + COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0) + ) + END; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/coderd/database/migrations/testdata/fixtures/000517_usage_events_agent_runtime.up.sql b/coderd/database/migrations/testdata/fixtures/000517_usage_events_agent_runtime.up.sql new file mode 100644 index 00000000000..71a42bd1a0a --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000517_usage_events_agent_runtime.up.sql @@ -0,0 +1,20 @@ +INSERT INTO usage_events ( + id, + event_type, + event_data, + created_at, + publish_started_at, + published_at, + failure_message +) +VALUES +-- Unpublished hb_agent_runtime_v1 event. +( + 'hb_agent_runtime_v1:2023-06-01_00:00:00', + 'hb_agent_runtime_v1', + '{"runtime_ms":3600000}', + '2023-06-01 00:00:00+00', + NULL, + NULL, + NULL +); diff --git a/coderd/database/querier.go b/coderd/database/querier.go index f958e008931..172b0b611bd 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -750,6 +750,9 @@ type sqlcQuerier interface { GetTemplateVersionsCreatedAfter(ctx context.Context, createdAt time.Time) ([]TemplateVersion, error) GetTemplates(ctx context.Context) ([]Template, error) GetTemplatesWithFilter(ctx context.Context, arg GetTemplatesWithFilterParams) ([]Template, error) + // Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop + // runtime across ALL chats (including soft-deleted messages) in [start, end). + GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) // Gets the total number of managed agents created between two dates. Uses the // aggregate table to avoid large scans or a complex index on the usage_events // table. @@ -1080,6 +1083,9 @@ type sqlcQuerier interface { ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error) + // Returns created_at of all events of the given type since @since. Used by + // the usage generator to find missing heartbeat buckets. + ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]UserConfig, error) ListUserChatPersonalModelOverrides(ctx context.Context, userID uuid.UUID) ([]ListUserChatPersonalModelOverridesRow, error) // Returns metadata only (no value or value_key_id) for the diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index b7c3490c982..fa34fc63eeb 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -9645,6 +9645,56 @@ func TestUsageEventsTrigger(t *testing.T) { require.Len(t, rows, 3) }) + t.Run("HeartbeatAgentRuntime", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + + insert := func(id, eventType, eventData string, createdAt time.Time) { + t.Helper() + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: id, + EventType: eventType, + EventData: []byte(eventData), + CreatedAt: createdAt, + }) + require.NoError(t, err) + } + requireDaily := func(wantUsageData ...string) { + t.Helper() + rows := getDailyRows(ctx, sqlDB) + require.Len(t, rows, len(wantUsageData)) + for i, want := range wantUsageData { + require.JSONEq(t, want, string(rows[i].UsageData)) + } + } + + day1 := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + day2 := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC) + + // Insert an hourly runtime heartbeat event. + insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1) + requireDaily(`{"runtime_ms": 1000}`) + + // Unlike hb_ai_seats_v1, hourly runtime events are summed per day. + insert("hb_agent_runtime_v1:2025-01-01_12:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 500}`, day1.Add(12*time.Hour)) + requireDaily(`{"runtime_ms": 1500}`) + + // Zero-valued events (idle hours) do not change the sum. + insert("hb_agent_runtime_v1:2025-01-01_18:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 0}`, day1.Add(18*time.Hour)) + requireDaily(`{"runtime_ms": 1500}`) + + // Insert on a different day: separate daily row. + insert("hb_agent_runtime_v1:2025-01-02_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 250}`, day2) + requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) + + // A different event type on the same day gets its own daily row. + insert("hb-seats-1", "hb_ai_seats_v1", `{"count": 3}`, day2) + rows := getDailyRows(ctx, sqlDB) + require.Len(t, rows, 3) + }) + t.Run("UnknownEventType", func(t *testing.T) { t.Parallel() @@ -9678,6 +9728,119 @@ func TestUsageEventsTrigger(t *testing.T) { }) } +func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + + rangeStart := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + rangeEnd := rangeStart.Add(time.Hour) + + // No messages at all: zero. + total, err := db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: rangeStart, + EndTime: rangeEnd, + }) + require.NoError(t, err) + require.EqualValues(t, 0, total) + + 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.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + }) + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "test-model", + ContextLimit: 8192, + }) + chat1 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + chat2 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + + insertMessage := func(chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { + t.Helper() + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: mc.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true}, + }) + _, err := sqlDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID) + require.NoError(t, err) + } + + // Counted: on the inclusive start boundary, in the middle (across two + // chats), soft-deleted, and just before the exclusive end boundary. + insertMessage(chat1.ID, 1, rangeStart, false) + insertMessage(chat2.ID, 2, rangeStart.Add(30*time.Minute), false) + insertMessage(chat1.ID, 4, rangeStart.Add(45*time.Minute), true) + insertMessage(chat1.ID, 8, rangeEnd.Add(-time.Second), false) + // Not counted: before the range, on the exclusive end boundary, and a + // NULL runtime (runtime 0 is stored as NULL). + insertMessage(chat1.ID, 16, rangeStart.Add(-time.Second), false) + insertMessage(chat1.ID, 32, rangeEnd, false) + insertMessage(chat1.ID, 0, rangeStart.Add(10*time.Minute), false) + + total, err = db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: rangeStart, + EndTime: rangeEnd, + }) + require.NoError(t, err) + require.EqualValues(t, 15, total) +} + +func TestListUsageEventCreatedAtsByTypeSince(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _ := dbtestutil.NewDB(t) + + since := time.Date(2025, 3, 10, 0, 0, 0, 0, time.UTC) + + insertEvent := func(id, eventType string, eventData string, createdAt time.Time) { + t.Helper() + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: id, + EventType: eventType, + EventData: []byte(eventData), + CreatedAt: createdAt, + }) + require.NoError(t, err) + } + + // Matching type: one before since (excluded), one exactly at since + // (included), one after (included). + insertEvent("rt-old", "hb_agent_runtime_v1", `{"runtime_ms": 1}`, since.Add(-time.Hour)) + insertEvent("rt-at", "hb_agent_runtime_v1", `{"runtime_ms": 2}`, since) + insertEvent("rt-new", "hb_agent_runtime_v1", `{"runtime_ms": 3}`, since.Add(time.Hour)) + // Different type after since: excluded. + insertEvent("seats-new", "hb_ai_seats_v1", `{"count": 1}`, since.Add(time.Hour)) + + createdAts, err := db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: "hb_agent_runtime_v1", + Since: since, + }) + require.NoError(t, err) + require.Len(t, createdAts, 2) + normalized := make([]time.Time, len(createdAts)) + for i, ts := range createdAts { + normalized[i] = ts.UTC() + } + require.ElementsMatch(t, []time.Time{since, since.Add(time.Hour)}, normalized) +} + func TestListTasks(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 369ec1fc703..d7b46feff74 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8699,6 +8699,28 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time return items, nil } +const getTotalChatMessageRuntimeMsInRange = `-- name: GetTotalChatMessageRuntimeMsInRange :one +SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms +FROM chat_messages cm +WHERE cm.created_at >= $1::timestamptz + AND cm.created_at < $2::timestamptz + AND cm.runtime_ms IS NOT NULL +` + +type GetTotalChatMessageRuntimeMsInRangeParams struct { + StartTime time.Time `db:"start_time" json:"start_time"` + EndTime time.Time `db:"end_time" json:"end_time"` +} + +// Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop +// runtime across ALL chats (including soft-deleted messages) in [start, end). +func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + row := q.db.QueryRowContext(ctx, getTotalChatMessageRuntimeMsInRange, arg.StartTime, arg.EndTime) + var total_runtime_ms int64 + err := row.Scan(&total_runtime_ms) + return total_runtime_ms, err +} + const getUserChatSpendInPeriod = `-- name: GetUserChatSpendInPeriod :one SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros FROM chat_messages cm @@ -26516,6 +26538,43 @@ func (q *sqlQuerier) InsertUsageEvent(ctx context.Context, arg InsertUsageEventP return err } +const listUsageEventCreatedAtsByTypeSince = `-- name: ListUsageEventCreatedAtsByTypeSince :many +SELECT created_at +FROM usage_events +WHERE event_type = $1 + AND created_at >= $2::timestamptz +` + +type ListUsageEventCreatedAtsByTypeSinceParams struct { + EventType string `db:"event_type" json:"event_type"` + Since time.Time `db:"since" json:"since"` +} + +// Returns created_at of all events of the given type since @since. Used by +// the usage generator to find missing heartbeat buckets. +func (q *sqlQuerier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + rows, err := q.db.QueryContext(ctx, listUsageEventCreatedAtsByTypeSince, arg.EventType, arg.Since) + if err != nil { + return nil, err + } + defer rows.Close() + var items []time.Time + for rows.Next() { + var created_at time.Time + if err := rows.Scan(&created_at); err != nil { + return nil, err + } + items = append(items, created_at) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const selectUsageEventsForPublishing = `-- name: SelectUsageEventsForPublishing :many WITH usage_events AS ( UPDATE diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 0d3810acb0f..e1a01126fc1 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -2119,6 +2119,15 @@ WHERE c.owner_id = @user_id::uuid AND cm.created_at < @end_time::timestamptz AND cm.total_cost_micros IS NOT NULL; +-- name: GetTotalChatMessageRuntimeMsInRange :one +-- Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop +-- runtime across ALL chats (including soft-deleted messages) in [start, end). +SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms +FROM chat_messages cm +WHERE cm.created_at >= @start_time::timestamptz + AND cm.created_at < @end_time::timestamptz + AND cm.runtime_ms IS NOT NULL; + -- name: CountEnabledModelsWithoutPricing :one -- Counts enabled, non-deleted model configs that lack both input and -- output pricing in their JSONB options.cost configuration. diff --git a/coderd/database/queries/usageevents.sql b/coderd/database/queries/usageevents.sql index 7ffcb1173b5..9166267a367 100644 --- a/coderd/database/queries/usageevents.sql +++ b/coderd/database/queries/usageevents.sql @@ -20,6 +20,14 @@ SELECT EXISTS( SELECT 1 FROM usage_events WHERE id = @id )::bool; +-- name: ListUsageEventCreatedAtsByTypeSince :many +-- Returns created_at of all events of the given type since @since. Used by +-- the usage generator to find missing heartbeat buckets. +SELECT created_at +FROM usage_events +WHERE event_type = @event_type + AND created_at >= @since::timestamptz; + -- name: SelectUsageEventsForPublishing :many WITH usage_events AS ( UPDATE diff --git a/coderd/pproflabel/pproflabel.go b/coderd/pproflabel/pproflabel.go index f686c1c4288..e5f1f6963f3 100644 --- a/coderd/pproflabel/pproflabel.go +++ b/coderd/pproflabel/pproflabel.go @@ -35,6 +35,9 @@ const ( // ServiceTallymanPublisher publishes usage events to coder/tallyman. ServiceTallymanPublisher = "tallyman-publisher" ServiceUsageEventCron = "usage-event-cron" + // ServiceUsageEventGenerator reconciles hourly heartbeat usage events + // over a trailing window. + ServiceUsageEventGenerator = "usage-event-generator" RequestTypeTag = "coder_request_type" ) diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index 891f5c73873..d9ddb6ae4a4 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -2,6 +2,7 @@ package usage import ( "context" + "time" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/usage/usagetypes" @@ -28,7 +29,11 @@ type Inserter interface { // // Inserts with the same `id` must be idempotent. The database enforces this by // ignoring duplicate records. - InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error + // + // The `createdAt` timestamp is stored on the event row. Generators that + // backfill historical buckets set it to the bucket start rather than the + // insertion time. + InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error } // AGPLInserter is a no-op implementation of Inserter. @@ -48,6 +53,6 @@ func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store // InsertHeartbeatUsageEvent is a no-op implementation of // InsertHeartbeatUsageEvent. -func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ usagetypes.HeartbeatEvent) error { +func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ time.Time, _ usagetypes.HeartbeatEvent) error { return nil } diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 6c8fde416eb..aee59127b69 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -30,6 +30,7 @@ type UsageEventType string const ( UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1" UsageEventTypeHBAISeatsV1 UsageEventType = "hb_ai_seats_v1" + UsageEventTypeHBAgentRuntimeV1 UsageEventType = "hb_agent_runtime_v1" ) func (e UsageEventType) Valid() bool { @@ -38,6 +39,8 @@ func (e UsageEventType) Valid() bool { return true case UsageEventTypeHBAISeatsV1: return true + case UsageEventTypeHBAgentRuntimeV1: + return true default: return false } @@ -105,6 +108,12 @@ func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event, return nil, err } return event, nil + case UsageEventTypeHBAgentRuntimeV1: + var event HBAgentRuntime + if err := ParseEvent(data, &event); err != nil { + return nil, err + } + return event, nil default: return nil, UnknownEventTypeError{EventType: string(eventType)} } @@ -192,3 +201,35 @@ func (e HBAISeats) Fields() map[string]any { "count": e.Count, } } + +// HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs +// is the total agent-loop runtime (LLM streaming, tool execution, and +// retries) in milliseconds consumed by Coder Agents (chats) in one UTC hour. +// +// This measures the new Coder Agents (the `chats` tables), not the deprecated +// Tasks counted by dc_managed_agents_v1. +type HBAgentRuntime struct { + RuntimeMs int64 `json:"runtime_ms"` +} + +var _ HeartbeatEvent = HBAgentRuntime{} + +func (HBAgentRuntime) usageEvent() {} +func (HBAgentRuntime) heartbeatUsageEvent() {} +func (HBAgentRuntime) EventType() UsageEventType { + return UsageEventTypeHBAgentRuntimeV1 +} + +func (e HBAgentRuntime) Valid() error { + if e.RuntimeMs < 0 { + return xerrors.New("runtime_ms cannot be negative") + } + // The runtime can be 0 (idle hour). + return nil +} + +func (e HBAgentRuntime) Fields() map[string]any { + return map[string]any{ + "runtime_ms": e.RuntimeMs, + } +} diff --git a/coderd/usage/usagetypes/events_test.go b/coderd/usage/usagetypes/events_test.go index fcfd076fc0e..f812b26e264 100644 --- a/coderd/usage/usagetypes/events_test.go +++ b/coderd/usage/usagetypes/events_test.go @@ -76,4 +76,28 @@ func TestParseEventWithType(t *testing.T) { require.Equal(t, eventType, event.EventType()) require.Equal(t, map[string]any{"count": int64(1)}, event.Fields()) }) + + t.Run("HBAgentRuntimeV1", func(t *testing.T) { + t.Parallel() + + eventType := usagetypes.UsageEventTypeHBAgentRuntimeV1 + event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1234}`)) + require.NoError(t, err) + require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 1234}, event) + require.Equal(t, eventType, event.EventType()) + require.Equal(t, map[string]any{"runtime_ms": int64(1234)}, event.Fields()) + + // Zero runtime is valid (idle hour). + event, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 0}`)) + require.NoError(t, err) + require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 0}, event) + + // Negative runtime is invalid. + _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": -1}`)) + require.ErrorContains(t, err, "runtime_ms cannot be negative") + + // Unknown fields are rejected. + _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1, "extra": "field"}`)) + require.ErrorContains(t, err, "unmarshal *usagetypes.HBAgentRuntime event") + }) } diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index 37febd028b7..942fd0dd86a 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -17,6 +17,7 @@ import ( agplcoderd "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/usage/usagetypes" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/enterprise/audit" "github.com/coder/coder/v2/enterprise/audit/backends" @@ -153,15 +154,29 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { usageCron := usage.NewCron(quartz.NewReal(), options.Logger.Named("usage-cron"), options.Database, *options.UsageInserter.Load()) // ai-seats heartbeats track the number of users that have used an AI feature. // These users consume a seat for the AI addon to our License. - _ = usageCron.Register(usage.CronJob{ - Name: "ai-seats", - Interval: usage.AISeatsInterval, - Jitter: 10 * time.Minute, - Fn: usage.AISeatsHeartbeat(options.Database), + err = usageCron.Register(usage.CronJob{ + Name: "ai-seats", + Interval: usage.AISeatsInterval, + EventType: usagetypes.UsageEventTypeHBAISeatsV1, + Jitter: 10 * time.Minute, + Fn: usage.AISeatsHeartbeat(options.Database), }) + if err != nil { + _ = closers.Close() + return nil, nil, xerrors.Errorf("register ai-seats usage cron job: %w", err) + } usageCron.Start(ctx) closers.Add(usageCron) + // The usage generator reconciles hourly agent runtime heartbeat + // events from data already in the database, backfilling hours missed + // while the deployment was down. Events are generated unconditionally + // in enterprise builds; the publish_usage_data license flag only + // gates publishing to Tallyman. + usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-generator"), options.Database, *options.UsageInserter.Load()) + usageGenerator.Start(ctx) + closers.Add(usageGenerator) + // In-memory AI Bridge Proxy daemon. The bridge daemon itself is // started unconditionally by AGPL cli/server.go (chatd uses its // in-memory roundtripper regardless of license); only the proxy diff --git a/enterprise/coderd/usage/cron.go b/enterprise/coderd/usage/cron.go index 13ccbb927c4..4c15ba4d7b9 100644 --- a/enterprise/coderd/usage/cron.go +++ b/enterprise/coderd/usage/cron.go @@ -184,7 +184,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) { continue } - if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, event); err != nil { + if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, c.clock.Now(), event); err != nil { c.log.Warn(ctx, "cron heartbeat insert failed", slog.F("job", job.Name), slog.Error(err), diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go new file mode 100644 index 00000000000..40187fc61ce --- /dev/null +++ b/enterprise/coderd/usage/generator.go @@ -0,0 +1,208 @@ +package usage + +import ( + "context" + "math/rand" + "sync" + "time" + + "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/pproflabel" + agplusage "github.com/coder/coder/v2/coderd/usage" + "github.com/coder/coder/v2/coderd/usage/usagetypes" + "github.com/coder/quartz" +) + +const ( + // AgentRuntimeInterval is the bucket size for hb_agent_runtime_v1 usage + // events. Each event covers one UTC hour [H, H+1). + AgentRuntimeInterval = time.Hour + // AgentRuntimeWindow is the trailing window scanned for missing buckets. + // Buckets that are still missing beyond this window (e.g. because the + // deployment was down for longer) are forfeited, which can only ever + // undercount usage. + AgentRuntimeWindow = 7 * 24 * time.Hour + // AgentRuntimeEligibilityLag is how long after a bucket closes before it + // becomes eligible for generation. This gives replicas time to commit + // in-flight chat messages with timestamps inside the bucket. + AgentRuntimeEligibilityLag = 5 * time.Minute + // agentRuntimeJitter is the maximum random delay added after each hour + // boundary. It staggers replicas so one is likely to complete the work + // before others attempt it (inserts are idempotent either way). + agentRuntimeJitter = 4 * time.Minute + // generatorTimerName tags the quartz timer so tests can trap it. + generatorTimerName = "agent-runtime-generator" +) + +// Generator reconciles hb_agent_runtime_v1 heartbeat usage events. Unlike +// Cron jobs, which sample state at the time they fire, the Generator derives +// events from data already persisted in the database (chat_messages), so it +// can deterministically backfill hours that were missed while the deployment +// was down. +// +// Every tick it scans the trailing AgentRuntimeWindow for missing hourly +// buckets and fills each one with the total chat message runtime recorded in +// that hour, inserting zero-valued events for idle hours. Deterministic event +// IDs plus the database's ON CONFLICT (id) DO NOTHING make concurrent +// replicas safe without locking. +// +// Events are generated unconditionally in enterprise builds; the +// publish_usage_data license flag only gates publishing to Tallyman. +type Generator struct { + clock quartz.Clock + log slog.Logger + db database.Store + ins agplusage.Inserter + + // cancel cancels the context of the running goroutine. If the ctx passed + // into Start is canceled, the goroutine also stops. + cancel context.CancelFunc + + // wg ensures the goroutine has exited before Close returns. + wg sync.WaitGroup + + // startOnce ensures Start is idempotent. + startOnce sync.Once +} + +// NewGenerator creates a Generator that reconciles agent runtime heartbeat +// events. The clock controls all timers so that tests can advance time +// deterministically via quartz.Mock. +func NewGenerator(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Generator { + return &Generator{ + clock: clock, + log: log, + db: db, + ins: ins, + } +} + +// Start launches the reconciliation goroutine. Subsequent calls are no-ops. +// On daemon restart a new Generator should be created. +func (g *Generator) Start(ctx context.Context) { + g.startOnce.Do(func() { + ctx, g.cancel = context.WithCancel(ctx) + g.wg.Add(1) + pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceUsageEventGenerator), func(ctx context.Context) { + g.run(ctx) + }) + }) +} + +// Close cancels the goroutine and waits for it to exit. +func (g *Generator) Close() error { + if g.cancel != nil { + g.cancel() + } + g.wg.Wait() + return nil +} + +func (g *Generator) run(ctx context.Context) { + //nolint:gocritic // We are a publisher in this function. + ctx = dbauthz.AsUsagePublisher(ctx) + defer g.wg.Done() + + // The first pass runs shortly after startup to heal any gaps that + // accumulated while the deployment was down. The uniform random delay in + // [1m, 5m) staggers replicas that start simultaneously. + //nolint:gosec // Jitter does not need cryptographic randomness. + delay := time.Minute + time.Duration(rand.Int63n(int64(4*time.Minute))) + for { + // Use a quartz timer so the wait honors ctx cancellation and tests + // can advance time deterministically. + timer := g.clock.NewTimer(delay, generatorTimerName) + + select { + case <-ctx.Done(): + if !timer.Stop() { + // Drain the channel if the timer already fired. + <-timer.C + } + return + case <-timer.C: + } + + err := g.generateAgentRuntimeEvents(ctx) + if ctx.Err() != nil { + return + } + if err != nil { + // The next tick rescans the whole window, so failed ticks heal + // automatically. + g.log.Warn(ctx, "generate agent runtime usage events", slog.Error(err)) + } + + // Wake shortly after the next hour boundary, once the just-closed + // bucket becomes eligible. + _, delay = nextTick(g.clock.Now(), AgentRuntimeInterval, agentRuntimeJitter) + delay += AgentRuntimeEligibilityLag + } +} + +// generateAgentRuntimeEvents scans the trailing window for missing hourly +// buckets and inserts one hb_agent_runtime_v1 event per missing bucket. Any +// error aborts the current tick; remaining buckets are retried on the next +// tick. +func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { + now := g.clock.Now().UTC() + // The most recent bucket that is old enough to generate. Bucket [H, H+1) + // becomes eligible at H + interval + lag. + latestEligible := now.Add(-AgentRuntimeInterval - AgentRuntimeEligibilityLag).Truncate(AgentRuntimeInterval) + // The oldest bucket we are willing to backfill. + earliest := now.Truncate(AgentRuntimeInterval).Add(-AgentRuntimeWindow) + if latestEligible.Before(earliest) { + return nil + } + + existingTimes, err := g.db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: string(usagetypes.UsageEventTypeHBAgentRuntimeV1), + Since: earliest, + }) + if err != nil { + return xerrors.Errorf("list existing agent runtime events: %w", err) + } + existing := make(map[time.Time]struct{}, len(existingTimes)) + for _, ts := range existingTimes { + // Events of this type always have created_at set to the exact bucket + // start; truncation just normalizes timezone and precision. + existing[ts.UTC().Truncate(AgentRuntimeInterval)] = struct{}{} + } + + var filled int + for bucket := earliest; !bucket.After(latestEligible); bucket = bucket.Add(AgentRuntimeInterval) { + if _, ok := existing[bucket]; ok { + continue + } + + runtimeMs, err := g.db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: bucket, + EndTime: bucket.Add(AgentRuntimeInterval), + }) + if err != nil { + return xerrors.Errorf("sum chat message runtime for bucket %s: %w", bucket, err) + } + + // The deterministic ID makes concurrent inserts of the same bucket + // idempotent. created_at is the bucket start (not insertion time) so + // daily rollups attribute backfilled hours to the correct day. + id := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(cronDateFormat) + err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, id, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) + if err != nil { + return xerrors.Errorf("insert agent runtime event for bucket %s: %w", bucket, err) + } + filled++ + } + if filled > 0 { + g.log.Info(ctx, "generated agent runtime usage events", + slog.F("buckets_filled", filled), + slog.F("window_start", earliest), + slog.F("latest_eligible", latestEligible), + ) + } + return nil +} diff --git a/enterprise/coderd/usage/generator_test.go b/enterprise/coderd/usage/generator_test.go new file mode 100644 index 00000000000..e1dc851f938 --- /dev/null +++ b/enterprise/coderd/usage/generator_test.go @@ -0,0 +1,350 @@ +package usage_test + +import ( + "context" + "database/sql" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/usage/usagetypes" + "github.com/coder/coder/v2/enterprise/coderd/usage" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// generatorTimerName must match the tag the Generator passes to +// clock.NewTimer so tests can trap its timers. +const generatorTimerName = "agent-runtime-generator" + +// generatorHarness bundles a real database with seeded chat dependencies. +// The generator itself runs against a dbauthz-wrapped store so the tests +// also verify that the usage publisher subject holds the permissions the +// generator's queries require. +type generatorHarness struct { + db database.Store + authzDB database.Store + rawDB *sql.DB + + user database.User + modelConfig database.ChatModelConfig + chat database.Chat + chat2 database.Chat +} + +func newGeneratorHarness(t *testing.T) *generatorHarness { + t.Helper() + db, _, rawDB := dbtestutil.NewDBWithSQLDB(t) + log := slogtest.Make(t, nil) + authzDB := dbauthz.New(db, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer()) + + 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.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + }) + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "test-model", + ContextLimit: 8192, + }) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + chat2 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + return &generatorHarness{ + db: db, + authzDB: authzDB, + rawDB: rawDB, + user: user, + modelConfig: mc, + chat: chat, + chat2: chat2, + } +} + +// insertRuntimeMessage inserts an assistant message with the given runtime +// and backdates it to createdAt. +func (h *generatorHarness) insertRuntimeMessage(ctx context.Context, t *testing.T, chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { + t.Helper() + msg := dbgen.ChatMessage(t, h.db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: h.user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: h.modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true}, + }) + _, err := h.rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID) + require.NoError(t, err) +} + +// fetchRuntimeEvents returns all hb_agent_runtime_v1 events keyed by their +// created_at (bucket start, UTC), plus a map of created_at to event ID. +func (h *generatorHarness) fetchRuntimeEvents(ctx context.Context, t *testing.T) (map[time.Time]int64, map[time.Time]string) { + t.Helper() + rows, err := h.rawDB.QueryContext(ctx, ` + SELECT id, (event_data->>'runtime_ms')::bigint, created_at + FROM usage_events + WHERE event_type = 'hb_agent_runtime_v1' + `) + require.NoError(t, err) + defer rows.Close() + + runtimes := make(map[time.Time]int64) + ids := make(map[time.Time]string) + for rows.Next() { + var ( + id string + runtimeMs int64 + createdAt time.Time + ) + require.NoError(t, rows.Scan(&id, &runtimeMs, &createdAt)) + bucket := createdAt.UTC() + _, ok := runtimes[bucket] + require.False(t, ok, "duplicate event for bucket %s", bucket) + runtimes[bucket] = runtimeMs + ids[bucket] = id + } + require.NoError(t, rows.Err()) + return runtimes, ids +} + +// expectedBuckets returns a map of every hourly bucket in [first, last] set +// to 0, then applies overrides. +func expectedBuckets(first, last time.Time, overrides map[time.Time]int64) map[time.Time]int64 { + expected := make(map[time.Time]int64) + for bucket := first; !bucket.After(last); bucket = bucket.Add(time.Hour) { + expected[bucket] = 0 + } + for bucket, runtimeMs := range overrides { + expected[bucket] = runtimeMs + } + return expected +} + +func TestGenerator(t *testing.T) { + t.Parallel() + + // startTime is exactly on an hour boundary, so the first tick (which + // fires 1-5 minutes later) always lands before the just-closed bucket + // [13:00, 14:00) becomes eligible at 14:05. + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + // Buckets within the window at the first tick. + bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + bucketB = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + // The most recent closed bucket; not eligible at the first tick. + bucketC = time.Date(2025, 3, 10, 13, 0, 0, 0, time.UTC) + // Window bounds at the first tick. + windowFirst = startTime.Add(-usage.AgentRuntimeWindow) // 2025-03-03 14:00 + windowLast = startTime.Add(-2 * time.Hour) // 2025-03-10 12:00 + ) + + // Bucket A: a message on the bucket start boundary, a message from a + // second chat, and a soft-deleted message. All must be counted. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA, false) + h.insertRuntimeMessage(ctx, t, h.chat2.ID, 2000, bucketA.Add(15*time.Minute), false) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 4000, bucketA.Add(30*time.Minute), true) + // Bucket B: a message exactly on the A/B boundary belongs to B. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 8000, bucketB, false) + // Older than the window: must never be generated. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 16000, windowFirst.Add(-30*time.Minute), false) + // Bucket C: only becomes eligible at 14:05, after the first tick. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 32000, bucketC.Add(30*time.Minute), false) + + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + // The initial pass is delayed by a uniform random duration in [1m, 5m). + call := trap.MustWait(ctx) + call.MustRelease(ctx) + require.GreaterOrEqual(t, call.Duration, time.Minute) + require.Less(t, call.Duration, 5*time.Minute) + clock.Advance(call.Duration).MustWait(ctx) + + // The generator creates the next timer only after the tick completes, + // so trapping it synchronizes with the end of the pass. + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The first pass fills every bucket in [windowFirst, windowLast]: + // bucket C is not yet eligible and the bucket before windowFirst is + // outside the window, message runtimes land in their buckets, and idle + // hours are zero-filled. + runtimes, ids := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast, map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + }), runtimes) + require.Equal(t, "hb_agent_runtime_v1:2025-03-10_10:00:00", ids[bucketA]) + + // The next tick fires shortly after the next hour boundary (15:00), once + // bucket [14:00, 15:00) is eligible. + fireTime := clock.Now().Add(call.Duration) + require.Equal(t, time.Date(2025, 3, 10, 15, 0, 0, 0, time.UTC), fireTime.Truncate(usage.AgentRuntimeInterval)) + require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The second tick fills only the two newly-eligible buckets: C (13:00) + // and the idle 14:00. The window's start advanced by an hour, but the + // bucket at the old windowFirst is kept (rows are never deleted). + runtimes, _ = h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(2*time.Hour), map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + bucketC: 32000, + }), runtimes) + + // A separate generator started later (e.g. another replica restarting) + // finds nothing to do: all buckets in its window already exist. + gen2 := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen2.Start(ctx) + defer gen2.Close() + call = trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + runtimes2, _ := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, runtimes, runtimes2) +} + +// TestGeneratorBackfillAfterDowntime simulates a deployment that was down +// for several hours: a fresh generator's first pass backfills exactly the +// missing buckets. +func TestGeneratorBackfillAfterDowntime(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + + windowFirst := startTime.Add(-usage.AgentRuntimeWindow) + gapBucket := time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 5000, gapBucket.Add(45*time.Minute), false) + + // Simulate events generated before the deployment went down at 09:00: + // every bucket in [windowFirst, 08:00] exists with runtime 1. These are + // inserted through the unwrapped store, which performs no authz. + inserter := usage.NewDBInserter() + for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) { + err := inserter.InsertHeartbeatUsageEvent(ctx, h.db, "hb_agent_runtime_v1:"+bucket.Format("2006-01-02_15:04:05"), bucket, usagetypes.HBAgentRuntime{RuntimeMs: 1}) + require.NoError(t, err) + } + + clock := quartz.NewMock(t) + clock.Set(startTime) + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The pass must fill exactly the gap [09:00, 12:00] and leave the + // pre-existing rows untouched. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + expected := expectedBuckets(windowFirst, startTime.Add(-2*time.Hour), map[time.Time]int64{ + gapBucket: 5000, + }) + for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) { + expected[bucket] = 1 + } + require.Equal(t, expected, runtimes) +} + +// TestGeneratorConcurrentReplicas runs two generators against the same +// database concurrently and verifies exactly one event is produced per +// bucket. Each replica gets its own mock clock (as real replicas have their +// own wall clocks) so their first passes can be fired independently and run +// at the same time. +func TestGeneratorConcurrentReplicas(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + + bucketA := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false) + + // Both replicas' first passes fire between 14:01 and 14:05, so they + // compute identical windows regardless of their random startup jitter. + var traps []*quartz.Trap + for range 2 { + clock := quartz.NewMock(t) + clock.Set(startTime) + trap := clock.Trap().NewTimer(generatorTimerName) + t.Cleanup(trap.Close) + traps = append(traps, trap) + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + // Cleanups run LIFO, so each generator is closed before its trap. + t.Cleanup(func() { _ = gen.Close() }) + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + // Fire the initial timer without waiting for the pass to complete so + // both replicas' passes overlap. + clock.Advance(call.Duration).MustWait(ctx) + } + + // Wait for both passes to complete (each requests its next timer only + // after the pass finishes). + for _, trap := range traps { + call := trap.MustWait(ctx) + call.MustRelease(ctx) + } + + // fetchRuntimeEvents fails on duplicate buckets; the expected map proves + // both replicas raced without erroring or double-inserting. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets( + startTime.Add(-usage.AgentRuntimeWindow), + startTime.Add(-2*time.Hour), + map[time.Time]int64{bucketA: 1000}, + ), runtimes) +} diff --git a/enterprise/coderd/usage/inserter.go b/enterprise/coderd/usage/inserter.go index 90fb6ab4ca8..61286f53eb9 100644 --- a/enterprise/coderd/usage/inserter.go +++ b/enterprise/coderd/usage/inserter.go @@ -3,6 +3,7 @@ package usage import ( "context" "encoding/json" + "time" "github.com/google/uuid" "golang.org/x/xerrors" @@ -68,7 +69,7 @@ func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.S } // InsertHeartbeatUsageEvent implements agplusage.Inserter. -func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error { +func (*dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error { if !event.EventType().IsHeartbeat() { return xerrors.Errorf("event type %q is not a heartbeat event", event.EventType()) } @@ -87,6 +88,6 @@ func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database. ID: id, EventType: string(event.EventType()), EventData: jsonData, - CreatedAt: dbtime.Time(i.clock.Now()), + CreatedAt: dbtime.Time(createdAt), }) } diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index 7ac915be7a5..c0076836173 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -68,6 +68,39 @@ func TestInserter(t *testing.T) { } }) + t.Run("Heartbeat", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + clock := quartz.NewMock(t) + inserter := usage.NewDBInserter(usage.InserterWithClock(clock)) + + // Heartbeat events use the provided ID and createdAt verbatim rather + // than generating them, so backfilled buckets keep their bucket start + // time. + event := usagetypes.HBAgentRuntime{RuntimeMs: 1234} + eventJSON := jsoninate(t, event) + id := "hb_agent_runtime_v1:2025-01-02_03:00:00" + createdAt := time.Date(2025, 1, 2, 3, 0, 0, 0, time.UTC) + // Set the clock to a different time to prove it is not used. + clock.Set(createdAt.Add(30 * time.Hour)) + + db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx interface{}, params database.InsertUsageEventParams) error { + assert.Equal(t, id, params.ID) + assert.Equal(t, event.EventType(), usagetypes.UsageEventType(params.EventType)) + assert.JSONEq(t, eventJSON, string(params.EventData)) + assert.Equal(t, dbtime.Time(createdAt), params.CreatedAt) + return nil + }, + ).Times(1) + + err := inserter.InsertHeartbeatUsageEvent(ctx, db, id, createdAt, event) + require.NoError(t, err) + }) + t.Run("InvalidEvent", func(t *testing.T) { t.Parallel() @@ -81,5 +114,11 @@ func TestInserter(t *testing.T) { Count: 0, // invalid }) assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`) + + // Same for heartbeat events. + err = inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Now(), usagetypes.HBAgentRuntime{ + RuntimeMs: -1, // invalid + }) + assert.ErrorContains(t, err, `invalid "hb_agent_runtime_v1" event: runtime_ms cannot be negative`) }) } From a02ddde40a183154b832c688dac82025f58259a5 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 22 Jul 2026 14:08:36 +0000 Subject: [PATCH 15/24] docs(coderd/database): document usage_events.created_at semantics created_at on usage_events is the event occurrence time rather than the row insertion time: it drives the daily rollup day and is sent to the usage collector as the event timestamp. Backfilled heartbeat events are the first producer where the two diverge, so make the semantics discoverable at the schema level like the table's other columns. --- coderd/database/dump.sql | 2 ++ .../migrations/000546_usage_events_agent_runtime.down.sql | 2 ++ .../migrations/000546_usage_events_agent_runtime.up.sql | 5 +++++ coderd/database/models.go | 3 ++- 4 files changed, 11 insertions(+), 1 deletion(-) diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 64effc353b1..d55a48314d7 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -3538,6 +3538,8 @@ COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version. COMMENT ON COLUMN usage_events.event_data IS 'Event payload. Determined by the matching usage struct for this event type.'; +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; + COMMENT ON COLUMN usage_events.publish_started_at IS 'Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired.'; COMMENT ON COLUMN usage_events.published_at IS 'Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again.'; diff --git a/coderd/database/migrations/000546_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000546_usage_events_agent_runtime.down.sql index a890669ab01..2de42b8a983 100644 --- a/coderd/database/migrations/000546_usage_events_agent_runtime.down.sql +++ b/coderd/database/migrations/000546_usage_events_agent_runtime.down.sql @@ -1,3 +1,5 @@ +COMMENT ON COLUMN usage_events.created_at IS NULL; + DROP INDEX IF EXISTS idx_usage_events_agent_runtime; -- Remove hb_agent_runtime_v1 rows so the previous constraint can be restored. diff --git a/coderd/database/migrations/000546_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000546_usage_events_agent_runtime.up.sql index 7647917eee1..0e54fde40d6 100644 --- a/coderd/database/migrations/000546_usage_events_agent_runtime.up.sql +++ b/coderd/database/migrations/000546_usage_events_agent_runtime.up.sql @@ -3,6 +3,11 @@ ALTER TABLE usage_events DROP CONSTRAINT usage_event_type_check, ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1')); +-- Document that created_at is the event occurrence time, not the row +-- insertion time. Backfilled events are the first case where the two +-- diverge. +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; + -- Partial index for efficient lookups of agent runtime heartbeat events by -- time. Used by the usage generator to find missing hourly buckets. CREATE INDEX idx_usage_events_agent_runtime diff --git a/coderd/database/models.go b/coderd/database/models.go index 73c11d14680..f87dae4a7d9 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -6130,7 +6130,8 @@ type UsageEvent struct { EventType string `db:"event_type" json:"event_type"` // Event payload. Determined by the matching usage struct for this event type. EventData json.RawMessage `db:"event_data" json:"event_data"` - CreatedAt time.Time `db:"created_at" json:"created_at"` + // The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp. + CreatedAt time.Time `db:"created_at" json:"created_at"` // Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired. PublishStartedAt sql.NullTime `db:"publish_started_at" json:"publish_started_at"` // Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again. From f1d30c9f5e28cb69f34dc52e1cd88389025c4980 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 23 Jul 2026 15:19:33 +0000 Subject: [PATCH 16/24] refactor: strip narrative comments from usage event generator change Keeps only comments documenting non-obvious behavior: window forfeiture and eligibility-lag semantics, lock-free replica safety, created_at bucket-start semantics, authz gating rationale, deliberate inclusion of soft-deleted messages, and test determinism/synchronization tricks. --- coderd/database/dbauthz/dbauthz.go | 5 +- ...000551_usage_events_agent_runtime.down.sql | 4 +- .../000551_usage_events_agent_runtime.up.sql | 10 +-- coderd/database/querier.go | 7 +- coderd/database/querier_test.go | 3 - coderd/database/queries.sql.go | 7 +- coderd/database/queries/chats.sql | 4 +- coderd/database/queries/usageevents.sql | 3 +- coderd/pproflabel/pproflabel.go | 4 +- coderd/usage/inserter.go | 5 +- coderd/usage/usagetypes/events_test.go | 3 - enterprise/cli/server.go | 7 +- enterprise/coderd/usage/generator.go | 82 +++++++------------ enterprise/coderd/usage/generator_test.go | 24 ++---- enterprise/coderd/usage/inserter_test.go | 6 +- 15 files changed, 54 insertions(+), 120 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 25a6120cd3f..982ecd0fc44 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4927,9 +4927,8 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe func (q *querier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { // This query exists solely to compute hb_agent_runtime_v1 usage event - // payloads. It returns a bare runtime sum across all chats with no chat - // content, so it's gated on the ability to create usage events rather - // than on reading chats. + // payloads and returns a bare sum with no chat content, so it is gated + // on usage event creation rather than on reading chats. if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil { return 0, err } diff --git a/coderd/database/migrations/000551_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000551_usage_events_agent_runtime.down.sql index 2de42b8a983..8c39020c1b7 100644 --- a/coderd/database/migrations/000551_usage_events_agent_runtime.down.sql +++ b/coderd/database/migrations/000551_usage_events_agent_runtime.down.sql @@ -6,13 +6,11 @@ DROP INDEX IF EXISTS idx_usage_events_agent_runtime; DELETE FROM usage_events WHERE event_type = 'hb_agent_runtime_v1'; DELETE FROM usage_events_daily WHERE event_type = 'hb_agent_runtime_v1'; --- Restore previous constraint. ALTER TABLE usage_events DROP CONSTRAINT usage_event_type_check, ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1')); --- Restore the previous aggregate function without hb_agent_runtime_v1 --- support. +-- Restores the 000444 version of the function. CREATE OR REPLACE FUNCTION aggregate_usage_event() RETURNS TRIGGER AS $$ BEGIN diff --git a/coderd/database/migrations/000551_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000551_usage_events_agent_runtime.up.sql index 0e54fde40d6..ce8993b817a 100644 --- a/coderd/database/migrations/000551_usage_events_agent_runtime.up.sql +++ b/coderd/database/migrations/000551_usage_events_agent_runtime.up.sql @@ -1,22 +1,14 @@ --- Expand the CHECK constraint to allow hb_agent_runtime_v1. ALTER TABLE usage_events DROP CONSTRAINT usage_event_type_check, ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1')); --- Document that created_at is the event occurrence time, not the row --- insertion time. Backfilled events are the first case where the two --- diverge. COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; --- Partial index for efficient lookups of agent runtime heartbeat events by --- time. Used by the usage generator to find missing hourly buckets. +-- Serves the usage generator's scan for missing hourly buckets. CREATE INDEX idx_usage_events_agent_runtime ON usage_events (event_type, created_at) WHERE event_type = 'hb_agent_runtime_v1'; --- Update the aggregate function to handle hb_agent_runtime_v1 events. --- Each hourly runtime event contributes to the daily total, so they are --- summed per day. CREATE OR REPLACE FUNCTION aggregate_usage_event() RETURNS TRIGGER AS $$ BEGIN diff --git a/coderd/database/querier.go b/coderd/database/querier.go index baf8e9318fb..2a114cec5cd 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -829,8 +829,8 @@ type sqlcQuerier interface { GetTemplateVersionsCreatedAfter(ctx context.Context, createdAt time.Time) ([]TemplateVersion, error) GetTemplates(ctx context.Context) ([]Template, error) GetTemplatesWithFilter(ctx context.Context, arg GetTemplatesWithFilterParams) ([]Template, error) - // Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop - // runtime across ALL chats (including soft-deleted messages) in [start, end). + // Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes + // soft-deleted messages and messages from all chats. GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) // Gets the total number of managed agents created between two dates. Uses the // aggregate table to avoid large scans or a complex index on the usage_events @@ -1214,8 +1214,7 @@ type sqlcQuerier interface { ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error) - // Returns created_at of all events of the given type since @since. Used by - // the usage generator to find missing heartbeat buckets. + // Used by the usage generator to find missing heartbeat buckets. ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]UserConfig, error) ListUserChatPersonalModelOverrides(ctx context.Context, userID uuid.UUID) ([]ListUserChatPersonalModelOverridesRow, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index ce4b83848a2..d097df7fbf5 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10591,7 +10591,6 @@ func TestUsageEventsTrigger(t *testing.T) { day1 := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) day2 := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC) - // Insert an hourly runtime heartbeat event. insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1) requireDaily(`{"runtime_ms": 1000}`) @@ -10603,7 +10602,6 @@ func TestUsageEventsTrigger(t *testing.T) { insert("hb_agent_runtime_v1:2025-01-01_18:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 0}`, day1.Add(18*time.Hour)) requireDaily(`{"runtime_ms": 1500}`) - // Insert on a different day: separate daily row. insert("hb_agent_runtime_v1:2025-01-02_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 250}`, day2) requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) @@ -10655,7 +10653,6 @@ func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { rangeStart := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) rangeEnd := rangeStart.Add(time.Hour) - // No messages at all: zero. total, err := db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ StartTime: rangeStart, EndTime: rangeEnd, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 76679e73c38..73eaaa2a66a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -9679,8 +9679,8 @@ type GetTotalChatMessageRuntimeMsInRangeParams struct { EndTime time.Time `db:"end_time" json:"end_time"` } -// Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop -// runtime across ALL chats (including soft-deleted messages) in [start, end). +// Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes +// soft-deleted messages and messages from all chats. func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { row := q.db.QueryRowContext(ctx, getTotalChatMessageRuntimeMsInRange, arg.StartTime, arg.EndTime) var total_runtime_ms int64 @@ -28523,8 +28523,7 @@ type ListUsageEventCreatedAtsByTypeSinceParams struct { Since time.Time `db:"since" json:"since"` } -// Returns created_at of all events of the given type since @since. Used by -// the usage generator to find missing heartbeat buckets. +// Used by the usage generator to find missing heartbeat buckets. func (q *sqlQuerier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { rows, err := q.db.QueryContext(ctx, listUsageEventCreatedAtsByTypeSince, arg.EventType, arg.Since) if err != nil { diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index dde8807a111..7f226f5a5d9 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -2374,8 +2374,8 @@ WHERE c.owner_id = @user_id::uuid AND cm.total_cost_micros IS NOT NULL; -- name: GetTotalChatMessageRuntimeMsInRange :one --- Used solely to compute hb_agent_runtime_v1 usage events. Sums agent-loop --- runtime across ALL chats (including soft-deleted messages) in [start, end). +-- Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes +-- soft-deleted messages and messages from all chats. SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms FROM chat_messages cm WHERE cm.created_at >= @start_time::timestamptz diff --git a/coderd/database/queries/usageevents.sql b/coderd/database/queries/usageevents.sql index 9166267a367..8ba706b0c8a 100644 --- a/coderd/database/queries/usageevents.sql +++ b/coderd/database/queries/usageevents.sql @@ -21,8 +21,7 @@ SELECT EXISTS( )::bool; -- name: ListUsageEventCreatedAtsByTypeSince :many --- Returns created_at of all events of the given type since @since. Used by --- the usage generator to find missing heartbeat buckets. +-- Used by the usage generator to find missing heartbeat buckets. SELECT created_at FROM usage_events WHERE event_type = @event_type diff --git a/coderd/pproflabel/pproflabel.go b/coderd/pproflabel/pproflabel.go index 875f7519a87..dcc17c5c0cc 100644 --- a/coderd/pproflabel/pproflabel.go +++ b/coderd/pproflabel/pproflabel.go @@ -38,9 +38,7 @@ const ( // ServiceWorkspaceBuildOrchestrator fulfills workspace build // orchestrations once their parent build reaches a terminal state. ServiceWorkspaceBuildOrchestrator = "workspace-build-orchestrator" - // ServiceUsageEventGenerator reconciles hourly heartbeat usage events - // over a trailing window. - ServiceUsageEventGenerator = "usage-event-generator" + ServiceUsageEventGenerator = "usage-event-generator" RequestTypeTag = "coder_request_type" ) diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index d9ddb6ae4a4..04c18c24f5c 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -30,9 +30,8 @@ type Inserter interface { // Inserts with the same `id` must be idempotent. The database enforces this by // ignoring duplicate records. // - // The `createdAt` timestamp is stored on the event row. Generators that - // backfill historical buckets set it to the bucket start rather than the - // insertion time. + // Generators that backfill historical buckets pass the bucket start as + // `createdAt` rather than the insertion time. InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error } diff --git a/coderd/usage/usagetypes/events_test.go b/coderd/usage/usagetypes/events_test.go index f812b26e264..ad017b0e073 100644 --- a/coderd/usage/usagetypes/events_test.go +++ b/coderd/usage/usagetypes/events_test.go @@ -87,16 +87,13 @@ func TestParseEventWithType(t *testing.T) { require.Equal(t, eventType, event.EventType()) require.Equal(t, map[string]any{"runtime_ms": int64(1234)}, event.Fields()) - // Zero runtime is valid (idle hour). event, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 0}`)) require.NoError(t, err) require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 0}, event) - // Negative runtime is invalid. _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": -1}`)) require.ErrorContains(t, err, "runtime_ms cannot be negative") - // Unknown fields are rejected. _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1, "extra": "field"}`)) require.ErrorContains(t, err, "unmarshal *usagetypes.HBAgentRuntime event") }) diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index bdfcdc90bed..d17a8447fd5 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -160,11 +160,8 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { usageCron.Start(ctx) closers.Add(usageCron) - // The usage generator reconciles hourly agent runtime heartbeat - // events from data already in the database, backfilling hours missed - // while the deployment was down. Events are generated unconditionally - // in enterprise builds; the publish_usage_data license flag only - // gates publishing to Tallyman. + // Usage generation is deliberately not license-gated; the + // publish_usage_data license flag only gates publishing to Tallyman. usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-generator"), options.Database, *options.UsageInserter.Load()) usageGenerator.Start(ctx) closers.Add(usageGenerator) diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go index 40187fc61ce..b2e2dd2ffda 100644 --- a/enterprise/coderd/usage/generator.go +++ b/enterprise/coderd/usage/generator.go @@ -18,37 +18,30 @@ import ( ) const ( - // AgentRuntimeInterval is the bucket size for hb_agent_runtime_v1 usage - // events. Each event covers one UTC hour [H, H+1). + // AgentRuntimeInterval is the bucket size of hb_agent_runtime_v1 events. AgentRuntimeInterval = time.Hour // AgentRuntimeWindow is the trailing window scanned for missing buckets. - // Buckets that are still missing beyond this window (e.g. because the - // deployment was down for longer) are forfeited, which can only ever - // undercount usage. + // Buckets still missing beyond this window (e.g. because the deployment + // was down for longer) are forfeited, which can only ever undercount + // usage. AgentRuntimeWindow = 7 * 24 * time.Hour // AgentRuntimeEligibilityLag is how long after a bucket closes before it - // becomes eligible for generation. This gives replicas time to commit + // becomes eligible for generation, giving replicas time to commit // in-flight chat messages with timestamps inside the bucket. AgentRuntimeEligibilityLag = 5 * time.Minute - // agentRuntimeJitter is the maximum random delay added after each hour - // boundary. It staggers replicas so one is likely to complete the work - // before others attempt it (inserts are idempotent either way). + // agentRuntimeJitter staggers replicas after each hour boundary so one + // is likely to complete the work before others attempt it. agentRuntimeJitter = 4 * time.Minute // generatorTimerName tags the quartz timer so tests can trap it. generatorTimerName = "agent-runtime-generator" ) // Generator reconciles hb_agent_runtime_v1 heartbeat usage events. Unlike -// Cron jobs, which sample state at the time they fire, the Generator derives -// events from data already persisted in the database (chat_messages), so it -// can deterministically backfill hours that were missed while the deployment -// was down. -// -// Every tick it scans the trailing AgentRuntimeWindow for missing hourly -// buckets and fills each one with the total chat message runtime recorded in -// that hour, inserting zero-valued events for idle hours. Deterministic event -// IDs plus the database's ON CONFLICT (id) DO NOTHING make concurrent -// replicas safe without locking. +// Cron jobs, which sample live state when they fire, the Generator derives +// events from data already persisted in the database, so it can +// deterministically backfill hours missed while the deployment was down, +// zero-filling idle hours. Deterministic event IDs plus the database's +// ON CONFLICT (id) DO NOTHING make concurrent replicas safe without locking. // // Events are generated unconditionally in enterprise builds; the // publish_usage_data license flag only gates publishing to Tallyman. @@ -58,20 +51,12 @@ type Generator struct { db database.Store ins agplusage.Inserter - // cancel cancels the context of the running goroutine. If the ctx passed - // into Start is canceled, the goroutine also stops. - cancel context.CancelFunc - - // wg ensures the goroutine has exited before Close returns. - wg sync.WaitGroup - - // startOnce ensures Start is idempotent. + cancel context.CancelFunc + wg sync.WaitGroup startOnce sync.Once } -// NewGenerator creates a Generator that reconciles agent runtime heartbeat -// events. The clock controls all timers so that tests can advance time -// deterministically via quartz.Mock. +// NewGenerator creates an unstarted Generator. func NewGenerator(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Generator { return &Generator{ clock: clock, @@ -81,8 +66,8 @@ func NewGenerator(clock quartz.Clock, log slog.Logger, db database.Store, ins ag } } -// Start launches the reconciliation goroutine. Subsequent calls are no-ops. -// On daemon restart a new Generator should be created. +// Start launches the reconciliation goroutine. Subsequent calls are no-ops; +// a closed Generator cannot be restarted. func (g *Generator) Start(ctx context.Context) { g.startOnce.Do(func() { ctx, g.cancel = context.WithCancel(ctx) @@ -93,7 +78,7 @@ func (g *Generator) Start(ctx context.Context) { }) } -// Close cancels the goroutine and waits for it to exit. +// Close stops the Generator and waits for its goroutine to exit. func (g *Generator) Close() error { if g.cancel != nil { g.cancel() @@ -107,20 +92,15 @@ func (g *Generator) run(ctx context.Context) { ctx = dbauthz.AsUsagePublisher(ctx) defer g.wg.Done() - // The first pass runs shortly after startup to heal any gaps that - // accumulated while the deployment was down. The uniform random delay in - // [1m, 5m) staggers replicas that start simultaneously. + // The random initial delay staggers replicas that start simultaneously. //nolint:gosec // Jitter does not need cryptographic randomness. delay := time.Minute + time.Duration(rand.Int63n(int64(4*time.Minute))) for { - // Use a quartz timer so the wait honors ctx cancellation and tests - // can advance time deterministically. timer := g.clock.NewTimer(delay, generatorTimerName) select { case <-ctx.Done(): if !timer.Stop() { - // Drain the channel if the timer already fired. <-timer.C } return @@ -132,28 +112,21 @@ func (g *Generator) run(ctx context.Context) { return } if err != nil { - // The next tick rescans the whole window, so failed ticks heal - // automatically. g.log.Warn(ctx, "generate agent runtime usage events", slog.Error(err)) } - // Wake shortly after the next hour boundary, once the just-closed - // bucket becomes eligible. _, delay = nextTick(g.clock.Now(), AgentRuntimeInterval, agentRuntimeJitter) delay += AgentRuntimeEligibilityLag } } -// generateAgentRuntimeEvents scans the trailing window for missing hourly -// buckets and inserts one hb_agent_runtime_v1 event per missing bucket. Any -// error aborts the current tick; remaining buckets are retried on the next -// tick. +// generateAgentRuntimeEvents inserts one hb_agent_runtime_v1 event per +// missing hourly bucket in the trailing window. Any error aborts the current +// tick; the next tick rescans the whole window, so failures self-heal. func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { now := g.clock.Now().UTC() - // The most recent bucket that is old enough to generate. Bucket [H, H+1) - // becomes eligible at H + interval + lag. + // Bucket [H, H+1) becomes eligible at H + interval + lag. latestEligible := now.Add(-AgentRuntimeInterval - AgentRuntimeEligibilityLag).Truncate(AgentRuntimeInterval) - // The oldest bucket we are willing to backfill. earliest := now.Truncate(AgentRuntimeInterval).Add(-AgentRuntimeWindow) if latestEligible.Before(earliest) { return nil @@ -168,8 +141,8 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { } existing := make(map[time.Time]struct{}, len(existingTimes)) for _, ts := range existingTimes { - // Events of this type always have created_at set to the exact bucket - // start; truncation just normalizes timezone and precision. + // created_at is always the exact bucket start for this event type; + // truncation just normalizes timezone and precision. existing[ts.UTC().Truncate(AgentRuntimeInterval)] = struct{}{} } @@ -188,8 +161,9 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { } // The deterministic ID makes concurrent inserts of the same bucket - // idempotent. created_at is the bucket start (not insertion time) so - // daily rollups attribute backfilled hours to the correct day. + // idempotent, and created_at is the bucket start (not the insertion + // time) so daily rollups attribute backfilled hours to the correct + // day. id := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(cronDateFormat) err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, id, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) if err != nil { diff --git a/enterprise/coderd/usage/generator_test.go b/enterprise/coderd/usage/generator_test.go index 3a09b78c96c..3617ba81370 100644 --- a/enterprise/coderd/usage/generator_test.go +++ b/enterprise/coderd/usage/generator_test.go @@ -27,10 +27,9 @@ import ( // clock.NewTimer so tests can trap its timers. const generatorTimerName = "agent-runtime-generator" -// generatorHarness bundles a real database with seeded chat dependencies. -// The generator itself runs against a dbauthz-wrapped store so the tests -// also verify that the usage publisher subject holds the permissions the -// generator's queries require. +// generatorHarness runs the generator against a dbauthz-wrapped store so the +// tests also verify that the usage publisher subject holds the permissions +// the generator's queries require. type generatorHarness struct { db database.Store authzDB database.Store @@ -80,8 +79,6 @@ func newGeneratorHarness(t *testing.T) *generatorHarness { } } -// insertRuntimeMessage inserts an assistant message with the given runtime -// and backdates it to createdAt. func (h *generatorHarness) insertRuntimeMessage(ctx context.Context, t *testing.T, chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { t.Helper() msg := dbgen.ChatMessage(t, h.db, database.ChatMessage{ @@ -95,8 +92,7 @@ func (h *generatorHarness) insertRuntimeMessage(ctx context.Context, t *testing. require.NoError(t, err) } -// fetchRuntimeEvents returns all hb_agent_runtime_v1 events keyed by their -// created_at (bucket start, UTC), plus a map of created_at to event ID. +// fetchRuntimeEvents fails the test if any bucket has more than one event. func (h *generatorHarness) fetchRuntimeEvents(ctx context.Context, t *testing.T) (map[time.Time]int64, map[time.Time]string) { t.Helper() rows, err := h.rawDB.QueryContext(ctx, ` @@ -126,8 +122,6 @@ func (h *generatorHarness) fetchRuntimeEvents(ctx context.Context, t *testing.T) return runtimes, ids } -// expectedBuckets returns a map of every hourly bucket in [first, last] set -// to 0, then applies overrides. func expectedBuckets(first, last time.Time, overrides map[time.Time]int64) map[time.Time]int64 { expected := make(map[time.Time]int64) for bucket := first; !bucket.After(last); bucket = bucket.Add(time.Hour) { @@ -154,7 +148,6 @@ func TestGenerator(t *testing.T) { clock.Set(startTime) var ( - // Buckets within the window at the first tick. bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) bucketB = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) // The most recent closed bucket; not eligible at the first tick. @@ -183,7 +176,6 @@ func TestGenerator(t *testing.T) { gen.Start(ctx) defer gen.Close() - // The initial pass is delayed by a uniform random duration in [1m, 5m). call := trap.MustWait(ctx) call.MustRelease(ctx) require.GreaterOrEqual(t, call.Duration, time.Minute) @@ -195,9 +187,8 @@ func TestGenerator(t *testing.T) { call = trap.MustWait(ctx) call.MustRelease(ctx) - // The first pass fills every bucket in [windowFirst, windowLast]: - // bucket C is not yet eligible and the bucket before windowFirst is - // outside the window, message runtimes land in their buckets, and idle + // The first pass fills every bucket in [windowFirst, windowLast]: bucket + // C is not yet eligible, the pre-window message is excluded, and idle // hours are zero-filled. runtimes, ids := h.fetchRuntimeEvents(ctx, t) require.Equal(t, expectedBuckets(windowFirst, windowLast, map[time.Time]int64{ @@ -257,8 +248,7 @@ func TestGeneratorBackfillAfterDowntime(t *testing.T) { h.insertRuntimeMessage(ctx, t, h.chat.ID, 5000, gapBucket.Add(45*time.Minute), false) // Simulate events generated before the deployment went down at 09:00: - // every bucket in [windowFirst, 08:00] exists with runtime 1. These are - // inserted through the unwrapped store, which performs no authz. + // every bucket in [windowFirst, 08:00] exists with runtime 1. inserter := usage.NewDBInserter() for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) { err := inserter.InsertHeartbeatUsageEvent(ctx, h.db, "hb_agent_runtime_v1:"+bucket.Format("2006-01-02_15:04:05"), bucket, usagetypes.HBAgentRuntime{RuntimeMs: 1}) diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index c0076836173..4dd50b99051 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -77,14 +77,11 @@ func TestInserter(t *testing.T) { clock := quartz.NewMock(t) inserter := usage.NewDBInserter(usage.InserterWithClock(clock)) - // Heartbeat events use the provided ID and createdAt verbatim rather - // than generating them, so backfilled buckets keep their bucket start - // time. event := usagetypes.HBAgentRuntime{RuntimeMs: 1234} eventJSON := jsoninate(t, event) id := "hb_agent_runtime_v1:2025-01-02_03:00:00" createdAt := time.Date(2025, 1, 2, 3, 0, 0, 0, time.UTC) - // Set the clock to a different time to prove it is not used. + // Prove the provided createdAt is stored rather than clock.Now(). clock.Set(createdAt.Add(30 * time.Hour)) db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn( @@ -115,7 +112,6 @@ func TestInserter(t *testing.T) { }) assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`) - // Same for heartbeat events. err = inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Now(), usagetypes.HBAgentRuntime{ RuntimeMs: -1, // invalid }) From b3eacc76c335993954e3f5330955f8453042df29 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 27 Jul 2026 08:27:14 +0000 Subject: [PATCH 17/24] fix: register ai-seats usage cron and grant usage-publisher create permission The ai-seats usage cron never ran: its CronJob was registered without the required EventType, so Register rejected it, and the error was discarded. ai-seats heartbeats drive AI addon seat consumption, so no seat usage events were being generated. Registering the job also surfaces a second latent bug: the usage-publisher subject lacked usage_event create permission, so every heartbeat insert performed under AsUsagePublisher would fail authz. Grant the permission and add a strict-authorizer regression test for the insert path. --- coderd/database/dbauthz/dbauthz.go | 6 +++--- enterprise/cli/server.go | 16 +++++++++++----- enterprise/coderd/usage/cron_test.go | 24 ++++++++++++++++++++++++ 3 files changed, 38 insertions(+), 8 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index b35529b0286..e6d76667264 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -666,9 +666,9 @@ var ( Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceLicense.Type: {policy.ActionRead}, rbac.ResourceAiSeat.Type: {policy.ActionRead}, // Required for GetActiveAISeatCount. - // The usage publisher doesn't create events, just - // reads/processes them. - rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate}, + // Create is required by the usage cron, which inserts + // heartbeat events under this subject. + rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate}, }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index ec3bc52d393..a39de5f8cde 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -16,6 +16,7 @@ import ( agplcoderd "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/usage/usagetypes" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/enterprise/audit" "github.com/coder/coder/v2/enterprise/audit/backends" @@ -145,12 +146,17 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { usageCron := usage.NewCron(quartz.NewReal(), options.Logger.Named("usage-cron"), options.Database, *options.UsageInserter.Load()) // ai-seats heartbeats track the number of users that have used an AI feature. // These users consume a seat for the AI addon to our License. - _ = usageCron.Register(usage.CronJob{ - Name: "ai-seats", - Interval: usage.AISeatsInterval, - Jitter: 10 * time.Minute, - Fn: usage.AISeatsHeartbeat(options.Database), + err = usageCron.Register(usage.CronJob{ + Name: "ai-seats", + Interval: usage.AISeatsInterval, + EventType: usagetypes.UsageEventTypeHBAISeatsV1, + Jitter: 10 * time.Minute, + Fn: usage.AISeatsHeartbeat(options.Database), }) + if err != nil { + _ = closers.Close() + return nil, nil, xerrors.Errorf("register ai-seats usage cron job: %w", err) + } usageCron.Start(ctx) closers.Add(usageCron) diff --git a/enterprise/coderd/usage/cron_test.go b/enterprise/coderd/usage/cron_test.go index 8381e6e77ff..215d2d5e1ce 100644 --- a/enterprise/coderd/usage/cron_test.go +++ b/enterprise/coderd/usage/cron_test.go @@ -80,6 +80,30 @@ func TestCron(t *testing.T) { }) } +// TestUsagePublisherCanInsertHeartbeats wraps a mock database with +// dbauthz to verify that the AsUsagePublisher subject can insert +// heartbeat usage events; the cron performs its inserts under this +// subject, so a missing usage_event create permission would make every +// heartbeat insert fail authz at runtime. +func TestUsagePublisherCanInsertHeartbeats(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + db.EXPECT().Wrappers().Return([]string{}).AnyTimes() + db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).Return(nil) + + authz := rbac.NewStrictAuthorizer(prometheus.NewRegistry()) + authzDB := dbauthz.New(db, authz, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) + + //nolint:gocritic // Testing the usage publisher subject is the point. + ctx := dbauthz.AsUsagePublisher(testutil.Context(t, testutil.WaitLong)) + inserter := usage.NewDBInserter() + err := inserter.InsertHeartbeatUsageEvent(ctx, authzDB, "hb_ai_seats_v1:2025-01-01_00:00:00", usagetypes.HBAISeats{Count: 1}) + require.NoError(t, err) +} + // TestAISeatsHeartbeat checks that AISeatsHeartbeat returns the // correct event type and count. It wraps a mock database with dbauthz // to verify that the AsUsagePublisher subject has the required From 446d77cacdf5f30cd217804d72ae7fe3758374e3 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 27 Jul 2026 09:06:04 +0000 Subject: [PATCH 18/24] refactor(enterprise): drop ai-seats cron fix in favor of standalone PR The ai-seats cron registration fix now ships independently in #27508 so it is not coupled to this feature's Tallyman release gate. The usage-publisher create permission stays here because the generator's heartbeat inserts also require it; both PRs carry the identical hunk. --- coderd/database/dbauthz/dbauthz.go | 4 ++-- enterprise/cli/server.go | 16 +++++----------- enterprise/coderd/usage/cron_test.go | 24 ------------------------ 3 files changed, 7 insertions(+), 37 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index c74aad9141d..6fc2e05737b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -666,8 +666,8 @@ var ( Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceLicense.Type: {policy.ActionRead}, rbac.ResourceAiSeat.Type: {policy.ActionRead}, // Required for GetActiveAISeatCount. - // Create is required by the usage cron and generator, - // which insert heartbeat events under this subject. + // Create is required to insert heartbeat usage events + // under this subject. rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate}, }), User: []rbac.Permission{}, diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index d17a8447fd5..430cfa65558 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -16,7 +16,6 @@ import ( agplcoderd "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/usage/usagetypes" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/enterprise/audit" "github.com/coder/coder/v2/enterprise/audit/backends" @@ -146,17 +145,12 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { usageCron := usage.NewCron(quartz.NewReal(), options.Logger.Named("usage-cron"), options.Database, *options.UsageInserter.Load()) // ai-seats heartbeats track the number of users that have used an AI feature. // These users consume a seat for the AI addon to our License. - err = usageCron.Register(usage.CronJob{ - Name: "ai-seats", - Interval: usage.AISeatsInterval, - EventType: usagetypes.UsageEventTypeHBAISeatsV1, - Jitter: 10 * time.Minute, - Fn: usage.AISeatsHeartbeat(options.Database), + _ = usageCron.Register(usage.CronJob{ + Name: "ai-seats", + Interval: usage.AISeatsInterval, + Jitter: 10 * time.Minute, + Fn: usage.AISeatsHeartbeat(options.Database), }) - if err != nil { - _ = closers.Close() - return nil, nil, xerrors.Errorf("register ai-seats usage cron job: %w", err) - } usageCron.Start(ctx) closers.Add(usageCron) diff --git a/enterprise/coderd/usage/cron_test.go b/enterprise/coderd/usage/cron_test.go index 36a52f33a47..8381e6e77ff 100644 --- a/enterprise/coderd/usage/cron_test.go +++ b/enterprise/coderd/usage/cron_test.go @@ -80,30 +80,6 @@ func TestCron(t *testing.T) { }) } -// TestUsagePublisherCanInsertHeartbeats wraps a mock database with -// dbauthz to verify that the AsUsagePublisher subject can insert -// heartbeat usage events; the cron performs its inserts under this -// subject, so a missing usage_event create permission would make every -// heartbeat insert fail authz at runtime. -func TestUsagePublisherCanInsertHeartbeats(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - db.EXPECT().Wrappers().Return([]string{}).AnyTimes() - db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).Return(nil) - - authz := rbac.NewStrictAuthorizer(prometheus.NewRegistry()) - authzDB := dbauthz.New(db, authz, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) - - //nolint:gocritic // Testing the usage publisher subject is the point. - ctx := dbauthz.AsUsagePublisher(testutil.Context(t, testutil.WaitLong)) - inserter := usage.NewDBInserter() - err := inserter.InsertHeartbeatUsageEvent(ctx, authzDB, "hb_ai_seats_v1:2025-01-01_00:00:00", time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), usagetypes.HBAISeats{Count: 1}) - require.NoError(t, err) -} - // TestAISeatsHeartbeat checks that AISeatsHeartbeat returns the // correct event type and count. It wraps a mock database with dbauthz // to verify that the AsUsagePublisher subject has the required From 7ad3f847b7a76a3f9b3a64190b2922fb399da9bc Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 27 Jul 2026 09:26:59 +0000 Subject: [PATCH 19/24] fix(enterprise/coderd/usage): address generator review findings - Schedule ticks against the lag-shifted clock so a bucket whose eligibility instant is still pending in the current hour is picked up minutes after startup instead of an hour later. - Skip only a failed bucket instead of aborting the whole pass, so one poison bucket (e.g. a negative runtime sum) cannot stall every later bucket until it ages out of the window. - Document that an existing row marks its bucket complete regardless of publish outcome, and the commit-latency assumption behind the eligibility lag. - Reference the jitter constant in the startup delay, align the logger name with the pprof service label, and reword the created_at column comment (every hb_agent_runtime_v1 event carries the bucket start, not just backfilled ones). - Tests: fail the concurrent-replica race on any Warn-level log so a lost ON CONFLICT dedup is detected, cover the poison-bucket skip, and assert owner/member roles cannot call the runtime-sum query. --- coderd/database/dbauthz/dbauthz_test.go | 23 ++++ coderd/database/dump.sql | 2 +- .../000552_usage_events_agent_runtime.up.sql | 2 +- coderd/database/models.go | 2 +- enterprise/cli/server.go | 2 +- enterprise/coderd/usage/generator.go | 79 +++++++++---- enterprise/coderd/usage/generator_test.go | 105 ++++++++++++++++-- enterprise/coderd/usage/inserter_test.go | 7 +- 8 files changed, 184 insertions(+), 38 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index e7d2dd5e271..a4189facfdc 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6736,6 +6736,29 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) { require.True(t, dbauthz.IsNotAuthorizedError(err)) } +// TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied mechanically +// checks the invariant the query's authz gate relies on: it exposes a +// deployment-wide aggregate behind usage_event create, which roles.go +// excludes from every human-assignable role (including owner) via +// allPermsExcept. If that grant ever widens, this test fails. +func TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied(t *testing.T) { + t.Parallel() + + for _, role := range []rbac.RoleIdentifier{rbac.RoleOwner(), rbac.RoleMember()} { + subj := rbac.Subject{ + ID: uuid.NewString(), + Roles: rbac.RoleIdentifiers{role}, + Scope: rbac.ScopeAll, + } + ctx := dbauthz.As(testutil.Context(t, testutil.WaitShort), subj) + mDB := dbmock.NewMockStore(gomock.NewController(t)) + mDB.EXPECT().Wrappers().Times(1).Return([]string{}) + dbz := dbauthz.New(mDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) + _, err := dbz.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{}) + require.True(t, dbauthz.IsNotAuthorizedError(err), "role %s must be denied", role) + } +} + func (s *MethodTestSuite) TestAIBridge() { s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { initID := uuid.UUID{3} diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index d725799aca8..a9c6f92ac34 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -3545,7 +3545,7 @@ COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version. COMMENT ON COLUMN usage_events.event_data IS 'Event payload. Determined by the matching usage struct for this event type.'; -COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; COMMENT ON COLUMN usage_events.publish_started_at IS 'Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired.'; diff --git a/coderd/database/migrations/000552_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000552_usage_events_agent_runtime.up.sql index ce8993b817a..1adb5d821a3 100644 --- a/coderd/database/migrations/000552_usage_events_agent_runtime.up.sql +++ b/coderd/database/migrations/000552_usage_events_agent_runtime.up.sql @@ -2,7 +2,7 @@ ALTER TABLE usage_events DROP CONSTRAINT usage_event_type_check, ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1')); -COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; -- Serves the usage generator's scan for missing hourly buckets. CREATE INDEX idx_usage_events_agent_runtime diff --git a/coderd/database/models.go b/coderd/database/models.go index 1f40872f4fd..f89ab6435ab 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -6136,7 +6136,7 @@ type UsageEvent struct { EventType string `db:"event_type" json:"event_type"` // Event payload. Determined by the matching usage struct for this event type. EventData json.RawMessage `db:"event_data" json:"event_data"` - // The time the usage occurred, which is not necessarily the time the row was inserted. Backfilled heartbeat events (e.g. hb_agent_runtime_v1) set this to the start of the measured time bucket rather than the insertion time. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp. + // The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp. CreatedAt time.Time `db:"created_at" json:"created_at"` // Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired. PublishStartedAt sql.NullTime `db:"publish_started_at" json:"publish_started_at"` diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index 430cfa65558..6a334d1b8e4 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -156,7 +156,7 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { // Usage generation is deliberately not license-gated; the // publish_usage_data license flag only gates publishing to Tallyman. - usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-generator"), options.Database, *options.UsageInserter.Load()) + usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-event-generator"), options.Database, *options.UsageInserter.Load()) usageGenerator.Start(ctx) closers.Add(usageGenerator) diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go index b2e2dd2ffda..c29d47f37d8 100644 --- a/enterprise/coderd/usage/generator.go +++ b/enterprise/coderd/usage/generator.go @@ -27,7 +27,11 @@ const ( AgentRuntimeWindow = 7 * 24 * time.Hour // AgentRuntimeEligibilityLag is how long after a bucket closes before it // becomes eligible for generation, giving replicas time to commit - // in-flight chat messages with timestamps inside the bucket. + // in-flight chat messages with timestamps inside the bucket. This + // assumes chat message inserts commit within the lag of their + // statement-time created_at; a message committing later than the lag + // after its bucket closes lands in an already-sealed bucket and its + // runtime is dropped (undercount-only). AgentRuntimeEligibilityLag = 5 * time.Minute // agentRuntimeJitter staggers replicas after each hour boundary so one // is likely to complete the work before others attempt it. @@ -94,7 +98,7 @@ func (g *Generator) run(ctx context.Context) { // The random initial delay staggers replicas that start simultaneously. //nolint:gosec // Jitter does not need cryptographic randomness. - delay := time.Minute + time.Duration(rand.Int63n(int64(4*time.Minute))) + delay := time.Minute + time.Duration(rand.Int63n(int64(agentRuntimeJitter))) for { timer := g.clock.NewTimer(delay, generatorTimerName) @@ -115,14 +119,19 @@ func (g *Generator) run(ctx context.Context) { g.log.Warn(ctx, "generate agent runtime usage events", slog.Error(err)) } - _, delay = nextTick(g.clock.Now(), AgentRuntimeInterval, agentRuntimeJitter) - delay += AgentRuntimeEligibilityLag + // Wake at the next eligibility instant (hour boundary + lag), not + // the next hour boundary. Computing the tick against the + // lag-shifted clock keeps a bucket whose eligibility is still + // pending in this hour (e.g. a pass that ran just before HH:05) + // from waiting a whole extra hour. + _, delay = nextTick(g.clock.Now().Add(-AgentRuntimeEligibilityLag), AgentRuntimeInterval, agentRuntimeJitter) } } // generateAgentRuntimeEvents inserts one hb_agent_runtime_v1 event per -// missing hourly bucket in the trailing window. Any error aborts the current -// tick; the next tick rescans the whole window, so failures self-heal. +// missing hourly bucket in the trailing window. Per-bucket errors skip only +// that bucket; the next tick rescans the whole window, so transient failures +// self-heal. func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { now := g.clock.Now().UTC() // Bucket [H, H+1) becomes eligible at H + interval + lag. @@ -139,6 +148,12 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { if err != nil { return xerrors.Errorf("list existing agent runtime events: %w", err) } + // A row marks its bucket complete regardless of publish outcome. If + // Tallyman permanently rejects an event, its bucket is never + // regenerated (a re-insert under the deterministic ID would be a no-op + // anyway); recovery is manual. The release gate (Tallyman must accept + // this event type before coderd ships it) is what keeps permanent + // rejections exceptional. existing := make(map[time.Time]struct{}, len(existingTimes)) for _, ts := range existingTimes { // created_at is always the exact bucket start for this event type; @@ -146,37 +161,57 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { existing[ts.UTC().Truncate(AgentRuntimeInterval)] = struct{}{} } - var filled int + var filled, failed int for bucket := earliest; !bucket.After(latestEligible); bucket = bucket.Add(AgentRuntimeInterval) { if _, ok := existing[bucket]; ok { continue } - - runtimeMs, err := g.db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ - StartTime: bucket, - EndTime: bucket.Add(AgentRuntimeInterval), - }) - if err != nil { - return xerrors.Errorf("sum chat message runtime for bucket %s: %w", bucket, err) + if ctx.Err() != nil { + return ctx.Err() } - // The deterministic ID makes concurrent inserts of the same bucket - // idempotent, and created_at is the bucket start (not the insertion - // time) so daily rollups attribute backfilled hours to the correct - // day. - id := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(cronDateFormat) - err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, id, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) + err := g.generateBucket(ctx, bucket) if err != nil { - return xerrors.Errorf("insert agent runtime event for bucket %s: %w", bucket, err) + // Skip only the failed bucket so one bad bucket (e.g. an + // invalid runtime sum) cannot stall every later bucket until + // it ages out of the window. + g.log.Warn(ctx, "generate agent runtime usage event for bucket", + slog.F("bucket", bucket), + slog.Error(err), + ) + failed++ + continue } filled++ } - if filled > 0 { + if filled > 0 || failed > 0 { g.log.Info(ctx, "generated agent runtime usage events", slog.F("buckets_filled", filled), + slog.F("buckets_failed", failed), slog.F("window_start", earliest), slog.F("latest_eligible", latestEligible), ) } return nil } + +// generateBucket computes and inserts the event for a single hourly bucket. +func (g *Generator) generateBucket(ctx context.Context, bucket time.Time) error { + runtimeMs, err := g.db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: bucket, + EndTime: bucket.Add(AgentRuntimeInterval), + }) + if err != nil { + return xerrors.Errorf("sum chat message runtime: %w", err) + } + + // The deterministic ID makes concurrent inserts of the same bucket + // idempotent, and created_at is the bucket start (not the insertion + // time) so daily rollups attribute backfilled hours to the correct day. + id := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(cronDateFormat) + err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, id, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) + if err != nil { + return xerrors.Errorf("insert usage event: %w", err) + } + return nil +} diff --git a/enterprise/coderd/usage/generator_test.go b/enterprise/coderd/usage/generator_test.go index 3617ba81370..ae75c1e74b5 100644 --- a/enterprise/coderd/usage/generator_test.go +++ b/enterprise/coderd/usage/generator_test.go @@ -3,6 +3,7 @@ package usage_test import ( "context" "database/sql" + "sync/atomic" "testing" "time" @@ -10,6 +11,7 @@ import ( "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" + "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" @@ -27,6 +29,19 @@ import ( // clock.NewTimer so tests can trap its timers. const generatorTimerName = "agent-runtime-generator" +// warnSink counts log entries at Warn or above. The generator downgrades +// per-bucket failures to Warn logs (which slogtest tolerates), so tests that +// must prove the error paths stayed quiet assert on this counter instead. +type warnSink struct{ count atomic.Int64 } + +func (s *warnSink) LogEntry(_ context.Context, e slog.SinkEntry) { + if e.Level >= slog.LevelWarn { + s.count.Add(1) + } +} + +func (*warnSink) Sync() {} + // generatorHarness runs the generator against a dbauthz-wrapped store so the // tests also verify that the usage publisher subject holds the permissions // the generator's queries require. @@ -197,18 +212,35 @@ func TestGenerator(t *testing.T) { }), runtimes) require.Equal(t, "hb_agent_runtime_v1:2025-03-10_10:00:00", ids[bucketA]) - // The next tick fires shortly after the next hour boundary (15:00), once - // bucket [14:00, 15:00) is eligible. + // The next tick fires at bucket C's eligibility instant (14:05 plus + // jitter), in the same hour the first pass ran, rather than waiting for + // the next hour boundary. fireTime := clock.Now().Add(call.Duration) - require.Equal(t, time.Date(2025, 3, 10, 15, 0, 0, 0, time.UTC), fireTime.Truncate(usage.AgentRuntimeInterval)) + require.Equal(t, startTime, fireTime.Truncate(usage.AgentRuntimeInterval)) + require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The second tick fills only the newly-eligible bucket C (13:00). The + // window's start advanced by an hour, but the bucket at the old + // windowFirst is kept (rows are never deleted). + runtimes, _ = h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(time.Hour), map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + bucketC: 32000, + }), runtimes) + + // The third tick fires after the next hour boundary (15:05 plus jitter) + // and fills the idle 14:00 bucket. + fireTime = clock.Now().Add(call.Duration) + require.Equal(t, startTime.Add(time.Hour), fireTime.Truncate(usage.AgentRuntimeInterval)) require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag) clock.Advance(call.Duration).MustWait(ctx) call = trap.MustWait(ctx) call.MustRelease(ctx) - // The second tick fills only the two newly-eligible buckets: C (13:00) - // and the idle 14:00. The window's start advanced by an hour, but the - // bucket at the old windowFirst is kept (rows are never deleted). runtimes, _ = h.fetchRuntimeEvents(ctx, t) require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(2*time.Hour), map[time.Time]int64{ bucketA: 7000, @@ -293,7 +325,11 @@ func TestGeneratorConcurrentReplicas(t *testing.T) { startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) ctx := testutil.Context(t, testutil.WaitLong) - log := slogtest.Make(t, nil) + // If insert idempotency regressed (e.g. ON CONFLICT DO NOTHING was + // dropped), the losing replica's inserts would error and surface as + // Warn logs, which slogtest alone would not catch. + sink := &warnSink{} + log := slogtest.Make(t, nil).AppendSinks(sink) h := newGeneratorHarness(t) bucketA := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) @@ -329,11 +365,64 @@ func TestGeneratorConcurrentReplicas(t *testing.T) { } // fetchRuntimeEvents fails on duplicate buckets; the expected map proves - // both replicas raced without erroring or double-inserting. + // both replicas raced without double-inserting, and the warn counter + // proves the duplicate inserts were deduplicated rather than rejected + // with errors. runtimes, _ := h.fetchRuntimeEvents(ctx, t) require.Equal(t, expectedBuckets( startTime.Add(-usage.AgentRuntimeWindow), startTime.Add(-2*time.Hour), map[time.Time]int64{bucketA: 1000}, ), runtimes) + require.Zero(t, sink.count.Load(), "no replica may log Warn or above during the race") +} + +// TestGeneratorPoisonBucket verifies that a bucket which fails +// deterministically on every tick does not stall generation of later +// buckets. +func TestGeneratorPoisonBucket(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + sink := &warnSink{} + log := slogtest.Make(t, nil).AppendSinks(sink) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + poisonBucket = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + laterBucket = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + ) + // Nothing prevents a negative runtime_ms row in the database, and a + // negative sum fails event validation on every tick. + h.insertRuntimeMessage(ctx, t, h.chat.ID, -5, poisonBucket.Add(10*time.Minute), false) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, laterBucket.Add(10*time.Minute), false) + + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // Every bucket except the poison bucket is generated, including buckets + // after it, and the failure is logged. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + expected := expectedBuckets( + startTime.Add(-usage.AgentRuntimeWindow), + startTime.Add(-2*time.Hour), + map[time.Time]int64{laterBucket: 1000}, + ) + delete(expected, poisonBucket) + require.Equal(t, expected, runtimes) + require.NotZero(t, sink.count.Load(), "the poison bucket failure must be logged") } diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index 4dd50b99051..f4558931816 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -74,15 +74,14 @@ func TestInserter(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) - clock := quartz.NewMock(t) - inserter := usage.NewDBInserter(usage.InserterWithClock(clock)) + inserter := usage.NewDBInserter() + // Heartbeat inserts must store the provided id and createdAt + // verbatim. event := usagetypes.HBAgentRuntime{RuntimeMs: 1234} eventJSON := jsoninate(t, event) id := "hb_agent_runtime_v1:2025-01-02_03:00:00" createdAt := time.Date(2025, 1, 2, 3, 0, 0, 0, time.UTC) - // Prove the provided createdAt is stored rather than clock.Now(). - clock.Set(createdAt.Add(30 * time.Hour)) db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn( func(ctx interface{}, params database.InsertUsageEventParams) error { From 8cc70e692eda47c8477edceea330aba10ff2161c Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 28 Jul 2026 10:41:01 +0000 Subject: [PATCH 20/24] chore: updates for PR review --- coderd/coderdtest/usage.go | 24 +++++++--- coderd/usage/inserter.go | 3 +- enterprise/coderd/usage/generator.go | 2 + enterprise/coderd/usage/generator_test.go | 53 +++++++++++++++++++++++ enterprise/coderd/usage/inserter.go | 6 +++ enterprise/coderd/usage/inserter_test.go | 16 +++++++ 6 files changed, 96 insertions(+), 8 deletions(-) diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go index d4d249dc863..7af644fe015 100644 --- a/coderd/coderdtest/usage.go +++ b/coderd/coderdtest/usage.go @@ -12,10 +12,16 @@ import ( var _ usage.Inserter = (*UsageInserter)(nil) +type HeartbeatEvent struct { + ID string + CreatedAt time.Time + Event usagetypes.HeartbeatEvent +} + type UsageInserter struct { sync.Mutex discreteEvents []usagetypes.DiscreteEvent - heartbeatEvents []usagetypes.HeartbeatEvent + heartbeatEvents []HeartbeatEvent seenHeartbeats map[string]struct{} } @@ -23,7 +29,7 @@ func NewUsageInserter() *UsageInserter { return &UsageInserter{ discreteEvents: []usagetypes.DiscreteEvent{}, seenHeartbeats: map[string]struct{}{}, - heartbeatEvents: []usagetypes.HeartbeatEvent{}, + heartbeatEvents: []HeartbeatEvent{}, } } @@ -34,7 +40,7 @@ func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.S return nil } -func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, _ time.Time, event usagetypes.HeartbeatEvent) error { +func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error { u.Lock() defer u.Unlock() if _, seen := u.seenHeartbeats[id]; seen { @@ -42,14 +48,18 @@ func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database. } u.seenHeartbeats[id] = struct{}{} - u.heartbeatEvents = append(u.heartbeatEvents, event) + u.heartbeatEvents = append(u.heartbeatEvents, HeartbeatEvent{ + ID: id, + CreatedAt: createdAt, + Event: event, + }) return nil } -func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent { +func (u *UsageInserter) GetHeartbeatEvents() []HeartbeatEvent { u.Lock() defer u.Unlock() - eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents)) + eventsCopy := make([]HeartbeatEvent, len(u.heartbeatEvents)) copy(eventsCopy, u.heartbeatEvents) return eventsCopy } @@ -73,5 +83,5 @@ func (u *UsageInserter) Reset() { defer u.Unlock() u.seenHeartbeats = map[string]struct{}{} u.discreteEvents = []usagetypes.DiscreteEvent{} - u.heartbeatEvents = []usagetypes.HeartbeatEvent{} + u.heartbeatEvents = []HeartbeatEvent{} } diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index 04c18c24f5c..acaaf4bbeb7 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -31,7 +31,8 @@ type Inserter interface { // ignoring duplicate records. // // Generators that backfill historical buckets pass the bucket start as - // `createdAt` rather than the insertion time. + // `createdAt` rather than the insertion time. `createdAt` must be + // non-zero; implementations reject the zero time rather than storing it. InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error } diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go index c29d47f37d8..728ca8cdac4 100644 --- a/enterprise/coderd/usage/generator.go +++ b/enterprise/coderd/usage/generator.go @@ -83,6 +83,8 @@ func (g *Generator) Start(ctx context.Context) { } // Close stops the Generator and waits for its goroutine to exit. +// It always returns nil; the error return exists to satisfy io.Closer, as the +// Generator is registered with the server's closer list. func (g *Generator) Close() error { if g.cancel != nil { g.cancel() diff --git a/enterprise/coderd/usage/generator_test.go b/enterprise/coderd/usage/generator_test.go index ae75c1e74b5..aeff3c9c309 100644 --- a/enterprise/coderd/usage/generator_test.go +++ b/enterprise/coderd/usage/generator_test.go @@ -314,6 +314,59 @@ func TestGeneratorBackfillAfterDowntime(t *testing.T) { require.Equal(t, expected, runtimes) } +func TestGeneratorInserterArguments(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + windowFirst = startTime.Add(-usage.AgentRuntimeWindow) + // The first pass stops at 12:00; the second pass fires after 14:05 + // and adds the newly eligible 13:00 bucket. + windowLast = startTime.Add(-time.Hour) + ) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false) + + ins := coderdtest.NewUsageInserter() + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, ins) + gen.Start(ctx) + defer gen.Close() + + // Each pass requests its next timer only after finishing, so trapping + // that request synchronizes with the end of the pass. + for range 2 { + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + } + call := trap.MustWait(ctx) + call.MustRelease(ctx) + + var expected []coderdtest.HeartbeatEvent + for bucket := windowFirst; !bucket.After(windowLast); bucket = bucket.Add(usage.AgentRuntimeInterval) { + var runtimeMs int64 + if bucket.Equal(bucketA) { + runtimeMs = 1000 + } + expected = append(expected, coderdtest.HeartbeatEvent{ + ID: "hb_agent_runtime_v1:" + bucket.Format("2006-01-02_15:04:05"), + CreatedAt: bucket, + Event: usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}, + }) + } + require.Equal(t, expected, ins.GetHeartbeatEvents()) +} + // TestGeneratorConcurrentReplicas runs two generators against the same // database concurrently and verifies exactly one event is produced per // bucket. Each replica gets its own mock clock (as real replicas have their diff --git a/enterprise/coderd/usage/inserter.go b/enterprise/coderd/usage/inserter.go index 61286f53eb9..80b209fd906 100644 --- a/enterprise/coderd/usage/inserter.go +++ b/enterprise/coderd/usage/inserter.go @@ -73,6 +73,12 @@ func (*dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.St if !event.EventType().IsHeartbeat() { return xerrors.Errorf("event type %q is not a heartbeat event", event.EventType()) } + // A zero createdAt stores the row at year 1, where bucket reconciliation + // can never match it again while the deterministic id turns every retry + // into a no-op, silently forfeiting the bucket's usage. + if createdAt.IsZero() { + return xerrors.Errorf("createdAt must be set for %q event", event.EventType()) + } if err := event.Valid(); err != nil { return xerrors.Errorf("invalid %q event: %w", event.EventType(), err) } diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index f4558931816..af48341f816 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -116,4 +116,20 @@ func TestInserter(t *testing.T) { }) assert.ErrorContains(t, err, `invalid "hb_agent_runtime_v1" event: runtime_ms cannot be negative`) }) + + t.Run("ZeroCreatedAt", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + // The mock store fails the test on any unexpected call, so no insert + // may reach the database. + db := dbmock.NewMockStore(ctrl) + + inserter := usage.NewDBInserter() + err := inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Time{}, usagetypes.HBAgentRuntime{ + RuntimeMs: 1, + }) + assert.ErrorContains(t, err, `createdAt must be set for "hb_agent_runtime_v1" event`) + }) } From abd74ba54b17b939ddaeadd7d6042d0d2797066c Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 28 Jul 2026 10:56:14 +0000 Subject: [PATCH 21/24] chore: address review findings for hb_agent_runtime_v1 usage events - Suppress the per-bucket generation warning when the pass is aborted by context cancellation, so shutdown mid-insert is not logged as a bucket failure (CRF-17). - Correct the HBAgentRuntime doc: the measured step ends when the model stream finishes, so local tool execution between steps is excluded (CRF-16). - Rename cronDateFormat to usageEventIDTimeFormat, name the startup delay floor, and use stableID in the generator to match the cron (CRF-15). - Widen the authz denial test to every built-in site and organization role, and describe why org-scoped usage_event grants cannot satisfy the site-scoped check (CRF-14). - Assert the daily rollup is unchanged when a bucket ID is inserted twice (CRF-13). - Document the publish re-arm SQL and its 30-day created_at bound at the dedup site (CRF-11). - List hb_agent_runtime_v1 in the usage data reporting doc. --- coderd/database/dbauthz/dbauthz_test.go | 20 ++++++++++--- coderd/database/querier_test.go | 8 +++++ coderd/usage/usagetypes/events.go | 7 +++-- docs/ai-coder/usage-data-reporting.md | 9 ++++++ enterprise/coderd/usage/cron.go | 7 +++-- enterprise/coderd/usage/generator.go | 39 ++++++++++++++++++++----- 6 files changed, 74 insertions(+), 16 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index a4189facfdc..ac3ad354156 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6738,13 +6738,25 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) { // TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied mechanically // checks the invariant the query's authz gate relies on: it exposes a -// deployment-wide aggregate behind usage_event create, which roles.go -// excludes from every human-assignable role (including owner) via -// allPermsExcept. If that grant ever widens, this test fails. +// deployment-wide aggregate behind usage_event create at site scope, which no +// human-assignable role holds. Owner is excluded from usage_event via +// allPermsExcept in roles.go; org roles such as org-admin do carry +// usage_event permissions, but only at org scope, which cannot satisfy a +// site-scoped check. If either of those ever changes, this test fails. func TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied(t *testing.T) { t.Parallel() - for _, role := range []rbac.RoleIdentifier{rbac.RoleOwner(), rbac.RoleMember()} { + orgID := uuid.New() + var roles []rbac.RoleIdentifier + for _, role := range rbac.SiteBuiltInRoles() { + roles = append(roles, role.Identifier) + } + for _, role := range rbac.OrganizationRoles(orgID) { + roles = append(roles, role.Identifier) + } + require.NotEmpty(t, roles) + + for _, role := range roles { subj := rbac.Subject{ ID: uuid.NewString(), Roles: rbac.RoleIdentifiers{role}, diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 4cee00df6bb..63c588400e1 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10605,6 +10605,14 @@ func TestUsageEventsTrigger(t *testing.T) { insert("hb_agent_runtime_v1:2025-01-02_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 250}`, day2) requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) + // Re-inserting a bucket must not double-count it. The daily rollup + // sums runtime_ms, so idempotency rests on the aggregate trigger + // being AFTER INSERT: Postgres does not fire it for rows suppressed + // by ON CONFLICT (id) DO NOTHING. Concurrent replicas and backfill + // re-runs both take this path. + insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1) + requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) + // A different event type on the same day gets its own daily row. insert("hb-seats-1", "hb_ai_seats_v1", `{"count": 3}`, day2) rows := getDailyRows(ctx, sqlDB) diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index aee59127b69..9004fb2514c 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -203,8 +203,11 @@ func (e HBAISeats) Fields() map[string]any { } // HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs -// is the total agent-loop runtime (LLM streaming, tool execution, and -// retries) in milliseconds consumed by Coder Agents (chats) in one UTC hour. +// is the total agent-loop runtime in milliseconds consumed by Coder Agents +// (chats) in one UTC hour. Each measured step spans model streaming (including +// provider-executed tools) and stream retries, and ends when the model stream +// finishes. Time spent executing local tools between steps, including +// sub-agents that bill their own model calls, is excluded. // // This measures the new Coder Agents (the `chats` tables), not the deprecated // Tasks counted by dc_managed_agents_v1. diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 21c1e42d47b..707787b01dd 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -4,6 +4,7 @@ The [AI Governance Add-On](./ai-governance.md) requires reporting usage data to - number of agent workspace builds consumed - number of AI Governance seats consumed +- total Coder Agent runtime consumed per hour, in milliseconds No user-identifiable information or additional metrics are sent to Tallyman. This information is also shared with [Metronome](https://metronome.com), a Stripe product and Coder partner for usage-based billing and reporting. @@ -33,6 +34,14 @@ Coder-Deployment-ID: 8a4e92f1-3b7c-4d5e-9f12-abc123def456 # your deployment ID "count": 1 }, "created_at": "2025-01-15T14:30:00Z" + }, + { + "id": "hb_agent_runtime_v1:2025-01-15_14:00:00", # unique event ID generated by Coder + "event_type": "hb_agent_runtime_v1", # aka. Coder Agent runtime + "event_data": { + "runtime_ms": 3600000 + }, + "created_at": "2025-01-15T14:00:00Z" # start of the hour the runtime was consumed in } ] } diff --git a/enterprise/coderd/usage/cron.go b/enterprise/coderd/usage/cron.go index 4c15ba4d7b9..47b85f6d377 100644 --- a/enterprise/coderd/usage/cron.go +++ b/enterprise/coderd/usage/cron.go @@ -23,7 +23,10 @@ import ( var epoch = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) const ( - cronDateFormat = "2006-01-02_15:04:05" + // usageEventIDTimeFormat is the timestamp layout used in every + // deterministic usage event ID, both the cron's boundary IDs and the + // generator's bucket IDs. + usageEventIDTimeFormat = "2006-01-02_15:04:05" ) // HeartbeatFunc generates a heartbeat event and its stable ID. @@ -145,7 +148,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) { // Use the boundary (not wall-clock "now") for the stable ID // so all replicas targeting the same boundary produce the // same key. - stableID := string(job.EventType) + ":" + boundary.UTC().Format(cronDateFormat) + stableID := string(job.EventType) + ":" + boundary.UTC().Format(usageEventIDTimeFormat) // Skip if this bucket was already recorded — avoids running // the potentially expensive heartbeat function for a diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go index 728ca8cdac4..91d911b57a7 100644 --- a/enterprise/coderd/usage/generator.go +++ b/enterprise/coderd/usage/generator.go @@ -36,6 +36,10 @@ const ( // agentRuntimeJitter staggers replicas after each hour boundary so one // is likely to complete the work before others attempt it. agentRuntimeJitter = 4 * time.Minute + // agentRuntimeStartupDelay is the floor on the first pass after start, + // giving the deployment time to finish booting before the generator + // competes for database work. Jitter is added on top of it. + agentRuntimeStartupDelay = time.Minute // generatorTimerName tags the quartz timer so tests can trap it. generatorTimerName = "agent-runtime-generator" ) @@ -100,7 +104,7 @@ func (g *Generator) run(ctx context.Context) { // The random initial delay staggers replicas that start simultaneously. //nolint:gosec // Jitter does not need cryptographic randomness. - delay := time.Minute + time.Duration(rand.Int63n(int64(agentRuntimeJitter))) + delay := agentRuntimeStartupDelay + time.Duration(rand.Int63n(int64(agentRuntimeJitter))) for { timer := g.clock.NewTimer(delay, generatorTimerName) @@ -150,11 +154,24 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { if err != nil { return xerrors.Errorf("list existing agent runtime events: %w", err) } - // A row marks its bucket complete regardless of publish outcome. If - // Tallyman permanently rejects an event, its bucket is never - // regenerated (a re-insert under the deterministic ID would be a no-op - // anyway); recovery is manual. The release gate (Tallyman must accept - // this event type before coderd ships it) is what keeps permanent + // A row marks its bucket complete regardless of publish outcome, so a + // bucket whose event Tallyman permanently rejected is never + // regenerated (re-inserting under the deterministic ID is a no-op via + // ON CONFLICT (id) DO NOTHING). + // + // The runtime is not lost locally: the row still holds it, and the + // event can be re-queued for publishing with + // + // UPDATE usage_events + // SET published_at = NULL, publish_started_at = NULL, failure_message = NULL + // WHERE id = 'hb_agent_runtime_v1:'; + // + // That re-arm only has an effect while the bucket is inside the + // publisher's 30-day cutoff: SelectUsageEventsForPublishing also + // filters created_at > now - INTERVAL '30 days', and created_at is the + // bucket start, so past that the UPDATE reports success but the row is + // never picked up again. The release gate (Tallyman must accept this + // event type before coderd ships it) is what keeps permanent // rejections exceptional. existing := make(map[time.Time]struct{}, len(existingTimes)) for _, ts := range existingTimes { @@ -174,6 +191,12 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { err := g.generateBucket(ctx, bucket) if err != nil { + if ctx.Err() != nil { + // A cancel landing mid-query surfaces as a bucket error. + // Report the cancellation instead of logging shutdown as + // a bucket failure. + return ctx.Err() + } // Skip only the failed bucket so one bad bucket (e.g. an // invalid runtime sum) cannot stall every later bucket until // it ages out of the window. @@ -210,8 +233,8 @@ func (g *Generator) generateBucket(ctx context.Context, bucket time.Time) error // The deterministic ID makes concurrent inserts of the same bucket // idempotent, and created_at is the bucket start (not the insertion // time) so daily rollups attribute backfilled hours to the correct day. - id := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(cronDateFormat) - err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, id, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) + stableID := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(usageEventIDTimeFormat) + err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, stableID, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) if err != nil { return xerrors.Errorf("insert usage event: %w", err) } From 18e82e9787fb522a2e57d7b5643cbb08cedff6ef Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 28 Jul 2026 12:59:20 +0000 Subject: [PATCH 22/24] chore(coderd/database/migrations): renumber usage events agent runtime migration to 000556 main added 000552 through 000555 after this branch was cut, so the merged tree had two 000552 migrations and every migration-dependent CI job failed with "duplicate migration file". --- ...untime.down.sql => 000556_usage_events_agent_runtime.down.sql} | 0 ...nt_runtime.up.sql => 000556_usage_events_agent_runtime.up.sql} | 0 ...nt_runtime.up.sql => 000556_usage_events_agent_runtime.up.sql} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename coderd/database/migrations/{000552_usage_events_agent_runtime.down.sql => 000556_usage_events_agent_runtime.down.sql} (100%) rename coderd/database/migrations/{000552_usage_events_agent_runtime.up.sql => 000556_usage_events_agent_runtime.up.sql} (100%) rename coderd/database/migrations/testdata/fixtures/{000552_usage_events_agent_runtime.up.sql => 000556_usage_events_agent_runtime.up.sql} (100%) diff --git a/coderd/database/migrations/000552_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000556_usage_events_agent_runtime.down.sql similarity index 100% rename from coderd/database/migrations/000552_usage_events_agent_runtime.down.sql rename to coderd/database/migrations/000556_usage_events_agent_runtime.down.sql diff --git a/coderd/database/migrations/000552_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000556_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/000552_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/000556_usage_events_agent_runtime.up.sql diff --git a/coderd/database/migrations/testdata/fixtures/000552_usage_events_agent_runtime.up.sql b/coderd/database/migrations/testdata/fixtures/000556_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/testdata/fixtures/000552_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/testdata/fixtures/000556_usage_events_agent_runtime.up.sql From 206f0349b2413d967d31546493f2d18c6f8a0bb4 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 29 Jul 2026 09:59:20 +0000 Subject: [PATCH 23/24] fix(coderd/database/migrations): renumber agent runtime migration to 000560 The agent runtime usage events migration was numbered 000556, which collides with 000556_user_secrets_enabled from main. golang-migrate rejects duplicate versions, so every test that opens a database failed with "duplicate migration file: 000556_user_secrets_enabled.down.sql". Renumber the migration and its fixture to 000560, the next free version after main's 000559. --- ...untime.down.sql => 000560_usage_events_agent_runtime.down.sql} | 0 ...nt_runtime.up.sql => 000560_usage_events_agent_runtime.up.sql} | 0 ...nt_runtime.up.sql => 000560_usage_events_agent_runtime.up.sql} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename coderd/database/migrations/{000556_usage_events_agent_runtime.down.sql => 000560_usage_events_agent_runtime.down.sql} (100%) rename coderd/database/migrations/{000556_usage_events_agent_runtime.up.sql => 000560_usage_events_agent_runtime.up.sql} (100%) rename coderd/database/migrations/testdata/fixtures/{000556_usage_events_agent_runtime.up.sql => 000560_usage_events_agent_runtime.up.sql} (100%) diff --git a/coderd/database/migrations/000556_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000560_usage_events_agent_runtime.down.sql similarity index 100% rename from coderd/database/migrations/000556_usage_events_agent_runtime.down.sql rename to coderd/database/migrations/000560_usage_events_agent_runtime.down.sql diff --git a/coderd/database/migrations/000556_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000560_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/000556_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/000560_usage_events_agent_runtime.up.sql diff --git a/coderd/database/migrations/testdata/fixtures/000556_usage_events_agent_runtime.up.sql b/coderd/database/migrations/testdata/fixtures/000560_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/testdata/fixtures/000556_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/testdata/fixtures/000560_usage_events_agent_runtime.up.sql From b72d53cbe1d872eea4e5771da373d55982d37417 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 30 Jul 2026 05:58:47 +0000 Subject: [PATCH 24/24] fix(coderd/database/migrations): renumber agent runtime migration to 000561 Main added 000560_chat_summary_boundary_index_by_id, which collides with this branch's 000560_usage_events_agent_runtime. golang-migrate rejects duplicate version numbers, so every job that runs migrations panics with "duplicate migration file". Renumber to 000561 to restore a unique sequence. --- ...untime.down.sql => 000561_usage_events_agent_runtime.down.sql} | 0 ...nt_runtime.up.sql => 000561_usage_events_agent_runtime.up.sql} | 0 ...nt_runtime.up.sql => 000561_usage_events_agent_runtime.up.sql} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename coderd/database/migrations/{000560_usage_events_agent_runtime.down.sql => 000561_usage_events_agent_runtime.down.sql} (100%) rename coderd/database/migrations/{000560_usage_events_agent_runtime.up.sql => 000561_usage_events_agent_runtime.up.sql} (100%) rename coderd/database/migrations/testdata/fixtures/{000560_usage_events_agent_runtime.up.sql => 000561_usage_events_agent_runtime.up.sql} (100%) diff --git a/coderd/database/migrations/000560_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000561_usage_events_agent_runtime.down.sql similarity index 100% rename from coderd/database/migrations/000560_usage_events_agent_runtime.down.sql rename to coderd/database/migrations/000561_usage_events_agent_runtime.down.sql diff --git a/coderd/database/migrations/000560_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000561_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/000560_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/000561_usage_events_agent_runtime.up.sql diff --git a/coderd/database/migrations/testdata/fixtures/000560_usage_events_agent_runtime.up.sql b/coderd/database/migrations/testdata/fixtures/000561_usage_events_agent_runtime.up.sql similarity index 100% rename from coderd/database/migrations/testdata/fixtures/000560_usage_events_agent_runtime.up.sql rename to coderd/database/migrations/testdata/fixtures/000561_usage_events_agent_runtime.up.sql