feat: add hourly hb_agent_runtime_v1 usage events for Coder Agent runtime (#27312)

closes CODAGT-839
closes CODAGT-843
closes CODAGT-773

## Summary

Adds a new heartbeat usage event type, `hb_agent_runtime_v1`, measuring
the total agent-loop runtime of Coder Agents (chats) per UTC hour, plus
a reconciler that generates one event per hour with self-healing
backfill over a trailing 7-day window. Events flow to Tallyman through
the existing publisher unchanged. This measures the new Coder Agents
(the `chats` tables), not the deprecated Tasks counted by
`dc_managed_agents_v1`.

Independent of #27508, which fixes the dead ai-seats cron registration.
Both PRs carry the identical `usage_event` create permission hunk for
the usage-publisher subject (this feature's generator and the ai-seats
cron each need it for heartbeat inserts), so they can land in either
order and the overlap merges cleanly.

> [!WARNING]
> **Do not include this in a release until Tallyman accepts
`hb_agent_runtime_v1`.** The publisher marks permanently rejected events
as done-forever, and the generator then sees those buckets as complete
locally, so their usage would be silently and permanently lost.

## Details

Each event's payload is `{"runtime_ms": N}`: the sum of
`chat_messages.runtime_ms` for messages created in the hour bucket `[H,
H+1)`, across all chats (sub-agents, API-created, archived, and
soft-deleted messages included). Events use deterministic IDs
(`hb_agent_runtime_v1:<bucket start>`) with `created_at` set to the
bucket start, so concurrent replicas race safely via `ON CONFLICT (id)
DO NOTHING` without locking, and daily rollups attribute backfilled
hours to the correct day. Idle hours produce zero-valued events. A
bucket becomes eligible 5 minutes after it closes; hours missing for
longer than the 7-day window are forfeited, which can only undercount.

Note that this makes `usage_events.created_at` explicitly the *event
occurrence time* rather than the row insertion time; the two only
diverge for backfilled events. It already behaved as the occurrence
timestamp (it drives the daily rollup day and is shipped to
Tallyman/Metronome as the event timestamp), and the migration now
documents this with a `COMMENT ON COLUMN`, which also surfaces as a Go
doc comment on `UsageEvent.CreatedAt`.

The new `usage.Generator` runs unconditionally in enterprise builds; the
`publish_usage_data` license flag continues to gate egress only, so
air-gapped deployments still fill their local ledger. The
`aggregate_usage_event()` trigger sums `runtime_ms` per day into
`usage_events_daily` (unlike `hb_ai_seats_v1`, which takes the daily
max).

`InsertHeartbeatUsageEvent` now takes an explicit `createdAt` so
generators can backfill historical buckets; the cron passes
`clock.Now()` to preserve its existing behavior.

## Tallyman follow-up

<details>
<summary>Prompt for the Tallyman-repo agent</summary>

> **Task**: Add support for the new Coder usage event type
`hb_agent_runtime_v1` so Tallyman accepts, validates, and forwards it to
Metronome.
>
> **Background**: coder/coder PR (this PR) adds hourly heartbeat events
measuring Coder Agent runtime. Events arrive via the existing
`/api/v1/events/ingest` endpoint with: `event_type:
"hb_agent_runtime_v1"`, `event_data: {"runtime_ms": <int64 >= 0>}`,
deterministic `id` of the form `hb_agent_runtime_v1:2026-07-15_14:00:00`
(UTC hour bucket start), and `created_at` set to the bucket start (may
be up to ~8 days in the past due to backfill; within Metronome's 34-day
dedup window). Zero-value events are normal (idle hours).
>
> **Work**:
> 1. Update Tallyman's vendored/imported `coderd/usage/usagetypes` (or
equivalent) to the coder/coder commit that adds
`UsageEventTypeHBAgentRuntimeV1` and `HBAgentRuntime`.
> 2. Ensure ingestion validation accepts the type (`Valid()` switches)
and rejects negative `runtime_ms`.
> 3. Ensure Metronome forwarding maps the event with transaction ID
derived from the event `id` as for existing types, passing `runtime_ms`
through as the property for a SUM-aggregated billable metric ("Coder
Agent Hours" = `SUM(runtime_ms) / 3,600,000`).
> 4. Do NOT permanently reject unknown-but-well-formed future `hb_*`
types if avoidable; at minimum confirm current behavior for unknown
types (temporary vs permanent rejection) and report it.
> 5. Tests: ingest accept/validate, dedup by ID, Metronome payload
mapping.
>
> **Constraint**: this must be deployed to tallyman-prod **before** any
coder/coder release containing the event generator; coderd treats
permanent rejections as terminal per event.

</details>
This commit is contained in:
Jaayden Halko
2026-07-30 08:37:45 +01:00
committed by GitHub
parent 2b28515d9b
commit 54d5eb7ec2
26 changed files with 1402 additions and 20 deletions
+18 -7
View File
@@ -3,6 +3,7 @@ package coderdtest
import (
"context"
"sync"
"time"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/usage"
@@ -11,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{}
}
@@ -22,7 +29,7 @@ func NewUsageInserter() *UsageInserter {
return &UsageInserter{
discreteEvents: []usagetypes.DiscreteEvent{},
seenHeartbeats: map[string]struct{}{},
heartbeatEvents: []usagetypes.HeartbeatEvent{},
heartbeatEvents: []HeartbeatEvent{},
}
}
@@ -33,7 +40,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, createdAt time.Time, event usagetypes.HeartbeatEvent) error {
u.Lock()
defer u.Unlock()
if _, seen := u.seenHeartbeats[id]; seen {
@@ -41,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
}
@@ -72,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{}
}
+20 -3
View File
@@ -671,9 +671,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 to insert heartbeat usage events
// under this subject.
rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
}),
User: []rbac.Permission{},
ByOrgID: map[string]rbac.OrgPermissions{},
@@ -4992,6 +4992,16 @@ 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 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
}
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
@@ -6901,6 +6911,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 {
+56
View File
@@ -6717,6 +6717,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.
@@ -6735,6 +6756,41 @@ 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 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()
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},
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}
+16
View File
@@ -3177,6 +3177,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)
@@ -4857,6 +4865,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)
+30
View File
@@ -5938,6 +5938,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()
@@ -9146,6 +9161,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()
+13 -2
View File
@@ -728,7 +728,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;
@@ -756,6 +756,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;
@@ -3529,7 +3536,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.';
@@ -3540,6 +3547,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. 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.';
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.';
@@ -4870,6 +4879,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);
@@ -0,0 +1,50 @@
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.
DELETE FROM usage_events WHERE event_type = 'hb_agent_runtime_v1';
DELETE FROM usage_events_daily WHERE event_type = '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'));
-- Restores the 000444 version of the function.
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;
@@ -0,0 +1,55 @@
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. 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
ON usage_events (event_type, created_at)
WHERE event_type = 'hb_agent_runtime_v1';
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;
@@ -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
);
+2 -1
View File
@@ -6142,7 +6142,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. 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"`
// 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.
+5
View File
@@ -867,6 +867,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)
// 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
// table.
@@ -1257,6 +1260,8 @@ 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)
// 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
+167
View File
@@ -10717,6 +10717,62 @@ 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("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("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)
require.Len(t, rows, 3)
})
t.Run("UnknownEventType", func(t *testing.T) {
t.Parallel()
@@ -10750,6 +10806,117 @@ 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)
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{
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()
+58
View File
@@ -10090,6 +10090,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"`
}
// 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
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
@@ -29088,6 +29110,42 @@ 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"`
}
// 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
+9
View File
@@ -2483,6 +2483,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
-- 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
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.
+7
View File
@@ -20,6 +20,13 @@ SELECT EXISTS(
SELECT 1 FROM usage_events WHERE id = @id
)::bool;
-- name: ListUsageEventCreatedAtsByTypeSince :many
-- 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
+1
View File
@@ -38,6 +38,7 @@ const (
// ServiceWorkspaceBuildOrchestrator fulfills workspace build
// orchestrations once their parent build reaches a terminal state.
ServiceWorkspaceBuildOrchestrator = "workspace-build-orchestrator"
ServiceUsageEventGenerator = "usage-event-generator"
RequestTypeTag = "coder_request_type"
)
+7 -2
View File
@@ -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
//
// Generators that backfill historical buckets pass the bucket start as
// `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
}
// 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
}
+44
View File
@@ -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,38 @@ 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 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.
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,
}
}
+21
View File
@@ -76,4 +76,25 @@ 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())
event, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 0}`))
require.NoError(t, err)
require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 0}, event)
_, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": -1}`))
require.ErrorContains(t, err, "runtime_ms cannot be negative")
_, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1, "extra": "field"}`))
require.ErrorContains(t, err, "unmarshal *usagetypes.HBAgentRuntime event")
})
}