mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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{}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
Generated
+30
@@ -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()
|
||||
|
||||
Generated
+13
-2
@@ -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;
|
||||
+20
@@ -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
|
||||
);
|
||||
Generated
+2
-1
@@ -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.
|
||||
|
||||
Generated
+5
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Generated
+58
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user