opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: add tallyman events for ai seat tracking (#22689)

Browse Source 
AI seat tracking inserted as heartbeat into usage table.
This commit is contained in:
Steven Masley
2026-03-18 09:30:22 -05:00
committed by GitHub
parent b83b93ea5c
commit 84de391f26
25 changed files with 860 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/querier_test.go
+74
View File
@@ -8315,6 +8315,80 @@ func TestUsageEventsTrigger(t *testing.T) {
require.WithinDuration(t, time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC), rows[1].Day, time.Second)
})
t.Run("HeartbeatAISeats", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
// Insert a heartbeat event.
err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: "hb-1",
EventType: "hb_ai_seats_v1",
EventData: []byte(`{"count": 10}`),
CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
})
require.NoError(t, err)
rows := getDailyRows(ctx, sqlDB)
require.Len(t, rows, 1)
require.Equal(t, "hb_ai_seats_v1", rows[0].EventType)
require.JSONEq(t, `{"count": 10}`, string(rows[0].UsageData))
// Insert a higher count on the same day — should take the max.
err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: "hb-2",
EventType: "hb_ai_seats_v1",
EventData: []byte(`{"count": 50}`),
CreatedAt: time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC),
})
require.NoError(t, err)
rows = getDailyRows(ctx, sqlDB)
require.Len(t, rows, 1)
require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
// Insert a lower count on the same day — should keep the max (50).
err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: "hb-3",
EventType: "hb_ai_seats_v1",
EventData: []byte(`{"count": 25}`),
CreatedAt: time.Date(2025, 1, 1, 18, 0, 0, 0, time.UTC),
})
require.NoError(t, err)
rows = getDailyRows(ctx, sqlDB)
require.Len(t, rows, 1)
require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
// Insert on a different day.
err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: "hb-4",
EventType: "hb_ai_seats_v1",
EventData: []byte(`{"count": 5}`),
CreatedAt: time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC),
})
require.NoError(t, err)
rows = getDailyRows(ctx, sqlDB)
require.Len(t, rows, 2)
require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
require.JSONEq(t, `{"count": 5}`, string(rows[1].UsageData))
// Also insert a dc_managed_agents_v1 on the same first day to
// verify different event types get separate daily rows.
err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: "dc-1",
EventType: "dc_managed_agents_v1",
EventData: []byte(`{"count": 7}`),
CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
})
require.NoError(t, err)
rows = getDailyRows(ctx, sqlDB)
require.Len(t, rows, 3)
})
t.Run("UnknownEventType", func(t *testing.T) {
t.Parallel()