feat: remove native chat cost tracking in favor of AI Gateway cost data (#27330)

## Stack Context

This stack makes AI Gateway data and budgets the source of truth for AI
spend controls.

1. Re-back the per-chat cost endpoint with AI Gateway data (#27328,
merged).
2. Remove native chat usage limits (#27329, merged).
3. **This PR, now based on `main`:** remove native chat cost tracking
and its dedicated admin UI.

## Summary

Removes native per-message price calculation, model pricing fields, cost
persistence, aggregate cost queries, and admin cost API types. It also
deletes the Analytics and Spend pages plus their legacy redirects. The
AI Gateway-backed per-chat cost row and compact budget indicators
remain.

The spend documentation is renamed to `spend-management.md` and updated
for the remaining surfaces, group budget APIs, CSV export, upgrade
handling for native pricing and cost history, and the absence of a
deployment-wide spend dashboard. The per-chat cost API documents that
data follows AI Gateway retention and reports zero after all matching
requests are purged.

No schema is dropped in this release. `chat_messages.total_cost_micros`
remains nullable and unwritten so replicas from the previous release can
continue inserting messages during rolling upgrades. #27600 tracks
removal after the compatibility window.

> Mux prepared this PR on Mike's behalf.
This commit is contained in:
Michael Suchacz
2026-08-04 12:27:38 +02:00
committed by GitHub
parent 0b8b48913f
commit 6b8f820493
72 changed files with 159 additions and 5303 deletions
+5 -536
View File
@@ -20,7 +20,6 @@ import (
"github.com/google/uuid"
"github.com/mark3labs/mcp-go/mcp"
"github.com/shopspring/decimal"
"github.com/sqlc-dev/pqtype"
"github.com/stretchr/testify/require"
"golang.org/x/sync/errgroup"
@@ -264,12 +263,11 @@ func (s *failNextUpdateChatModelConfigStore) UpdateChatModelConfig(
return s.Store.UpdateChatModelConfig(ctx, arg)
}
func insertAssistantCostMessage(
func insertAssistantMessage(
t *testing.T,
db database.Store,
chatID uuid.UUID,
modelConfigID uuid.UUID,
totalCostMicros int64,
) {
t.Helper()
@@ -279,11 +277,10 @@ func insertAssistantCostMessage(
require.NoError(t, err)
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chatID,
ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true},
Role: database.ChatMessageRoleAssistant,
Content: assistantContent,
TotalCostMicros: sql.NullInt64{Int64: totalCostMicros, Valid: true},
ChatID: chatID,
ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true},
Role: database.ChatMessageRoleAssistant,
Content: assistantContent,
})
}
@@ -3878,41 +3875,6 @@ func TestListChatModelConfigs(t *testing.T) {
require.Equal(t, enabledConfig.ID, memberConfigs[0].ID)
})
t.Run("DeserializesLegacyPricingJSON", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
aiProvider := createAIProviderForTest(t, client, "openai", "test-api-key")
legacyOptions := json.RawMessage(`{"input_price_per_million_tokens":0.15,"output_price_per_million_tokens":0.6,"cache_read_price_per_million_tokens":0.03,"cache_write_price_per_million_tokens":0.3}`)
storedConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true},
Model: "gpt-4o-mini-legacy",
DisplayName: "GPT-4o Mini Legacy",
CreatedBy: uuid.NullUUID{UUID: firstUser.UserID, Valid: true},
UpdatedBy: uuid.NullUUID{UUID: firstUser.UserID, Valid: true},
ContextLimit: 4096,
CompressionThreshold: 80,
Options: legacyOptions,
})
configs, err := client.ListChatModelConfigs(ctx)
require.NoError(t, err)
require.Len(t, configs, 1)
require.Equal(t, storedConfig.ID, configs[0].ID)
requireChatModelPricing(t, configs[0].ModelConfig, &codersdk.ChatModelCallConfig{
Cost: &codersdk.ModelCostConfig{
InputPricePerMillionTokens: decRef("0.15"),
OutputPricePerMillionTokens: decRef("0.6"),
CacheReadPricePerMillionTokens: decRef("0.03"),
CacheWritePricePerMillionTokens: decRef("0.3"),
},
})
})
t.Run("SuccessForOrganizationMember", func(t *testing.T) {
t.Parallel()
@@ -3954,20 +3916,11 @@ func TestCreateChatModelConfig(t *testing.T) {
contextLimit := int64(4096)
isDefault := true
pricing := &codersdk.ChatModelCallConfig{
Cost: &codersdk.ModelCostConfig{
InputPricePerMillionTokens: decRef("0.15"),
OutputPricePerMillionTokens: decRef("0.6"),
CacheReadPricePerMillionTokens: decRef("0.03"),
CacheWritePricePerMillionTokens: decRef("0.3"),
},
}
modelConfig, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
AIProviderID: &aiProvider.ID,
Model: "gpt-4o-mini",
ContextLimit: &contextLimit,
IsDefault: &isDefault,
ModelConfig: pricing,
})
require.NoError(t, err)
require.NotEqual(t, uuid.Nil, modelConfig.ID)
@@ -3975,12 +3928,10 @@ func TestCreateChatModelConfig(t *testing.T) {
require.Equal(t, "gpt-4o-mini", modelConfig.Model)
require.EqualValues(t, 4096, modelConfig.ContextLimit)
require.True(t, modelConfig.IsDefault)
requireChatModelPricing(t, modelConfig.ModelConfig, pricing)
configs, err := client.ListChatModelConfigs(ctx)
require.NoError(t, err)
require.Len(t, configs, 1)
requireChatModelPricing(t, configs[0].ModelConfig, pricing)
})
t.Run("ConcurrentCreatesElectSingleDefault", func(t *testing.T) {
@@ -4040,35 +3991,6 @@ func TestCreateChatModelConfig(t *testing.T) {
require.Equal(t, []uuid.UUID{claimed.ID}, defaults)
})
t.Run("RejectsNegativePricing", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
_ = coderdtest.CreateFirstUser(t, client.Client)
aiProvider := createAIProviderForTest(t, client, "openai", "test-api-key")
contextLimit := int64(4096)
_, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
AIProviderID: &aiProvider.ID,
Model: "gpt-4o-mini",
ContextLimit: &contextLimit,
ModelConfig: &codersdk.ChatModelCallConfig{
Cost: &codersdk.ModelCostConfig{
InputPricePerMillionTokens: decRef("-0.01"),
},
},
})
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "Invalid model config.", sdkErr.Message)
require.Equal(
t,
"cost.input_price_per_million_tokens must be greater than or equal to zero",
sdkErr.Detail,
)
})
t.Run("ReasoningEffortStored", func(t *testing.T) {
t.Parallel()
@@ -4358,29 +4280,18 @@ func TestUpdateChatModelConfig(t *testing.T) {
modelConfig := createChatModelConfig(t, client)
contextLimit := int64(8192)
pricing := &codersdk.ChatModelCallConfig{
Cost: &codersdk.ModelCostConfig{
InputPricePerMillionTokens: decRef("0.2"),
OutputPricePerMillionTokens: decRef("0.8"),
CacheReadPricePerMillionTokens: decRef("0.04"),
CacheWritePricePerMillionTokens: decRef("0.4"),
},
}
updated, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{
DisplayName: "GPT-4o Mini Updated",
ContextLimit: &contextLimit,
ModelConfig: pricing,
})
require.NoError(t, err)
require.Equal(t, modelConfig.ID, updated.ID)
require.Equal(t, "GPT-4o Mini Updated", updated.DisplayName)
require.EqualValues(t, 8192, updated.ContextLimit)
requireChatModelPricing(t, updated.ModelConfig, pricing)
configs, err := client.ListChatModelConfigs(ctx)
require.NoError(t, err)
require.Len(t, configs, 1)
requireChatModelPricing(t, configs[0].ModelConfig, pricing)
})
t.Run("UnchangedProviderWithoutAIProviderID", func(t *testing.T) {
@@ -4594,30 +4505,6 @@ func TestUpdateChatModelConfig(t *testing.T) {
require.True(t, foundForMember)
})
t.Run("RejectsNegativePricing", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client := newChatClient(t)
_ = coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
_, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{
ModelConfig: &codersdk.ChatModelCallConfig{
Cost: &codersdk.ModelCostConfig{
OutputPricePerMillionTokens: decRef("-1.0"),
},
},
})
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
require.Equal(t, "Invalid model config.", sdkErr.Message)
require.Equal(
t,
"cost.output_price_per_million_tokens must be greater than or equal to zero",
sdkErr.Detail,
)
})
t.Run("UpdateAIProviderID", func(t *testing.T) {
t.Parallel()
@@ -5314,7 +5201,6 @@ func TestGetChatUserPrompts(t *testing.T) {
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
@@ -5396,7 +5282,6 @@ func TestGetChatUserPrompts(t *testing.T) {
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
@@ -5423,7 +5308,6 @@ func TestGetChatUserPrompts(t *testing.T) {
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
@@ -5618,7 +5502,6 @@ func TestGetChatUserPrompts(t *testing.T) {
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
@@ -10660,7 +10543,6 @@ func TestPromoteChatQueuedMessage(t *testing.T) {
CacheReadTokens: []int64{0},
ContextLimit: []int64{0},
Compressed: []bool{false},
TotalCostMicros: []int64{0},
RuntimeMs: []int64{0},
})
require.NoError(t, err)
@@ -11264,194 +11146,6 @@ func TestGetChatFile(t *testing.T) {
})
}
type chatCostTestFixture struct {
Client *codersdk.ExperimentalClient
DB database.Store
ModelConfigID uuid.UUID
ChatID uuid.UUID
EarliestCreatedAt time.Time
LatestCreatedAt time.Time
}
// safeOptions returns an explicit time window around the fixture messages to
// avoid app-time/database-time boundary flakes in summary tests.
func (f chatCostTestFixture) safeOptions() codersdk.ChatCostSummaryOptions {
return codersdk.ChatCostSummaryOptions{
StartDate: f.EarliestCreatedAt.Add(-time.Minute),
EndDate: f.LatestCreatedAt.Add(time.Minute),
}
}
func seedChatCostFixture(t *testing.T) chatCostTestFixture {
t.Helper()
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: firstUser.UserID,
LastModelConfigID: modelConfig.ID,
Title: "test chat",
})
msg1 := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
RuntimeMs: sql.NullInt64{Int64: 1500, Valid: true},
})
msg2 := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
RuntimeMs: sql.NullInt64{Int64: 2500, Valid: true},
})
results := []database.ChatMessage{msg1, msg2}
require.Len(t, results, 2)
earliestCreatedAt := results[0].CreatedAt
latestCreatedAt := results[0].CreatedAt
for _, msg := range results {
if msg.CreatedAt.Before(earliestCreatedAt) {
earliestCreatedAt = msg.CreatedAt
}
if msg.CreatedAt.After(latestCreatedAt) {
latestCreatedAt = msg.CreatedAt
}
}
return chatCostTestFixture{
Client: client,
DB: db,
ModelConfigID: modelConfig.ID,
ChatID: chat.ID,
EarliestCreatedAt: earliestCreatedAt,
LatestCreatedAt: latestCreatedAt,
}
}
func assertChatCostSummary(t *testing.T, summary codersdk.ChatCostSummary, modelConfigID, chatID uuid.UUID) {
t.Helper()
require.Equal(t, int64(1000), summary.TotalCostMicros)
require.Equal(t, int64(2), summary.PricedMessageCount)
require.Equal(t, int64(0), summary.UnpricedMessagesHavingUsageCount)
require.Equal(t, int64(200), summary.TotalInputTokens)
require.Equal(t, int64(100), summary.TotalOutputTokens)
require.Equal(t, int64(4000), summary.TotalRuntimeMs)
require.Len(t, summary.ByModel, 1)
require.Equal(t, modelConfigID, summary.ByModel[0].ModelConfigID)
require.Equal(t, int64(1000), summary.ByModel[0].TotalCostMicros)
require.Equal(t, int64(2), summary.ByModel[0].MessageCount)
require.Equal(t, int64(4000), summary.ByModel[0].TotalRuntimeMs)
require.Len(t, summary.ByChat, 1)
require.Equal(t, chatID, summary.ByChat[0].RootChatID)
require.Equal(t, int64(1000), summary.ByChat[0].TotalCostMicros)
require.Equal(t, int64(2), summary.ByChat[0].MessageCount)
require.Equal(t, int64(4000), summary.ByChat[0].TotalRuntimeMs)
}
func TestChatCostSummary(t *testing.T) {
t.Parallel()
t.Run("BasicSummary", func(t *testing.T) {
t.Parallel()
f := seedChatCostFixture(t)
ctx := testutil.Context(t, testutil.WaitLong)
// Use a window derived from DB timestamps to avoid time boundary flakes.
summary, err := f.Client.GetChatCostSummary(ctx, "me", f.safeOptions())
require.NoError(t, err)
assertChatCostSummary(t, summary, f.ModelConfigID, f.ChatID)
})
}
func TestChatCostSummary_AfterModelDeletion(t *testing.T) {
t.Parallel()
f := seedChatCostFixture(t)
ctx := testutil.Context(t, testutil.WaitLong)
options := f.safeOptions()
// Baseline: use DB-derived timestamps to avoid time boundary flakes.
summary, err := f.Client.GetChatCostSummary(ctx, "me", options)
require.NoError(t, err)
assertChatCostSummary(t, summary, f.ModelConfigID, f.ChatID)
// Soft-delete the model config.
err = f.Client.DeleteChatModelConfig(ctx, f.ModelConfigID)
require.NoError(t, err)
// Costs must survive the deletion unchanged within the same safe window.
summary, err = f.Client.GetChatCostSummary(ctx, "me", options)
require.NoError(t, err)
assertChatCostSummary(t, summary, f.ModelConfigID, f.ChatID)
}
func TestChatCostSummary_AdminDrilldown(t *testing.T) {
t.Parallel()
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
memberClient := codersdk.NewExperimentalClient(memberClientRaw)
modelConfig := createChatModelConfig(t, client)
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: member.ID,
LastModelConfigID: modelConfig.ID,
Title: "member chat",
})
message := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 200, Valid: true},
OutputTokens: sql.NullInt64{Int64: 100, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 750, Valid: true},
})
options := codersdk.ChatCostSummaryOptions{
// Pad the DB-assigned timestamp so the query window cannot race it.
StartDate: message.CreatedAt.Add(-time.Minute),
EndDate: message.CreatedAt.Add(time.Minute),
}
t.Run("AdminCanDrilldown", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
summary, err := client.GetChatCostSummary(ctx, member.ID.String(), options)
require.NoError(t, err)
require.Equal(t, int64(750), summary.TotalCostMicros)
require.Equal(t, int64(1), summary.PricedMessageCount)
})
t.Run("MemberCannotDrilldownOtherUser", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
_, err := memberClient.GetChatCostSummary(ctx, firstUser.UserID.String(), options)
require.Error(t, err)
var sdkErr *codersdk.Error
require.ErrorAs(t, err, &sdkErr)
require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
})
}
// seedChatGatewayRequest records one finished Coder Agents gateway request
// under sessionChatID, mirroring aibridged: the session ID is the spawning
// chat, and each usage is one provider response within that one request.
@@ -11888,231 +11582,6 @@ func TestGetChatCost(t *testing.T) {
})
}
func TestChatCostUsers(t *testing.T) {
t.Parallel()
seedCtx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
memberClient := codersdk.NewExperimentalClient(memberClientRaw)
firstUserRecord, err := db.GetUserByID(dbauthz.AsSystemRestricted(seedCtx), firstUser.UserID)
require.NoError(t, err)
modelConfig := createChatModelConfig(t, client)
adminChat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: firstUser.UserID,
LastModelConfigID: modelConfig.ID,
Title: "admin chat",
})
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: adminChat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 300, Valid: true},
})
memberChat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: member.ID,
LastModelConfigID: modelConfig.ID,
Title: "member chat",
})
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: memberChat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 200, Valid: true},
OutputTokens: sql.NullInt64{Int64: 100, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 800, Valid: true},
})
t.Run("AdminCanListUsers", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
resp, err := client.GetChatCostUsers(ctx, codersdk.ChatCostUsersOptions{})
require.NoError(t, err)
require.Equal(t, int64(2), resp.Count)
require.Len(t, resp.Users, 2)
require.Equal(t, member.ID, resp.Users[0].UserID)
require.Equal(t, member.Username, resp.Users[0].Username)
require.Equal(t, int64(800), resp.Users[0].TotalCostMicros)
require.Equal(t, int64(1), resp.Users[0].MessageCount)
require.Equal(t, int64(1), resp.Users[0].ChatCount)
require.Equal(t, firstUser.UserID, resp.Users[1].UserID)
require.Equal(t, firstUserRecord.Username, resp.Users[1].Username)
require.Equal(t, int64(300), resp.Users[1].TotalCostMicros)
})
t.Run("AdminCanFilterAndPaginateUsers", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
resp, err := client.GetChatCostUsers(ctx, codersdk.ChatCostUsersOptions{
Username: member.Username,
Pagination: codersdk.Pagination{
Limit: 1,
Offset: 0,
},
})
require.NoError(t, err)
require.Equal(t, int64(1), resp.Count)
require.Len(t, resp.Users, 1)
require.Equal(t, member.ID, resp.Users[0].UserID)
require.Equal(t, member.Username, resp.Users[0].Username)
})
t.Run("MemberCannotListUsers", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
_, err := memberClient.GetChatCostUsers(ctx, codersdk.ChatCostUsersOptions{})
require.Error(t, err)
var sdkErr *codersdk.Error
require.ErrorAs(t, err, &sdkErr)
require.Equal(t, http.StatusForbidden, sdkErr.StatusCode())
})
}
func TestChatCostSummary_DateRange(t *testing.T) {
t.Parallel()
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: firstUser.UserID,
LastModelConfigID: modelConfig.ID,
Title: "date range test",
})
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
})
now := time.Now()
t.Run("MessageInRange", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
summary, err := client.GetChatCostSummary(ctx, "me", codersdk.ChatCostSummaryOptions{
StartDate: now.Add(-time.Hour),
EndDate: now.Add(time.Hour),
})
require.NoError(t, err)
require.Equal(t, int64(500), summary.TotalCostMicros)
require.Equal(t, int64(1), summary.PricedMessageCount)
})
t.Run("MessageOutOfRange", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
summary, err := client.GetChatCostSummary(ctx, "me", codersdk.ChatCostSummaryOptions{
StartDate: now.Add(time.Hour),
EndDate: now.Add(2 * time.Hour),
})
require.NoError(t, err)
require.Equal(t, int64(0), summary.TotalCostMicros)
require.Equal(t, int64(0), summary.PricedMessageCount)
})
}
func TestChatCostSummary_UnpricedMessages(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := newChatClientWithDatabase(t)
firstUser := coderdtest.CreateFirstUser(t, client.Client)
modelConfig := createChatModelConfig(t, client)
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: firstUser.OrganizationID,
OwnerID: firstUser.UserID,
LastModelConfigID: modelConfig.ID,
Title: "unpriced test",
})
pricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
})
unpricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chat.ID,
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
InputTokens: sql.NullInt64{Int64: 200, Valid: true},
OutputTokens: sql.NullInt64{Int64: 75, Valid: true},
})
earliestCreatedAt := pricedMessage.CreatedAt
latestCreatedAt := pricedMessage.CreatedAt
if unpricedMessage.CreatedAt.Before(earliestCreatedAt) {
earliestCreatedAt = unpricedMessage.CreatedAt
}
if unpricedMessage.CreatedAt.After(latestCreatedAt) {
latestCreatedAt = unpricedMessage.CreatedAt
}
options := codersdk.ChatCostSummaryOptions{
// Pad the DB-assigned timestamps to avoid time boundary flakes.
StartDate: earliestCreatedAt.Add(-time.Minute),
EndDate: latestCreatedAt.Add(time.Minute),
}
summary, err := client.GetChatCostSummary(ctx, "me", options)
require.NoError(t, err)
require.Equal(t, int64(500), summary.TotalCostMicros)
require.Equal(t, int64(1), summary.PricedMessageCount)
require.Equal(t, int64(1), summary.UnpricedMessagesHavingUsageCount)
require.Equal(t, int64(300), summary.TotalInputTokens)
require.Equal(t, int64(125), summary.TotalOutputTokens)
}
func requireChatModelPricing(
t *testing.T,
actual *codersdk.ChatModelCallConfig,
expected *codersdk.ChatModelCallConfig,
) {
t.Helper()
require.NotNil(t, actual)
require.NotNil(t, expected)
require.NotNil(t, actual.Cost)
require.NotNil(t, expected.Cost)
require.NotNil(t, actual.Cost.InputPricePerMillionTokens)
require.NotNil(t, actual.Cost.OutputPricePerMillionTokens)
require.NotNil(t, actual.Cost.CacheReadPricePerMillionTokens)
require.NotNil(t, actual.Cost.CacheWritePricePerMillionTokens)
require.True(t, expected.Cost.InputPricePerMillionTokens.Equal(*actual.Cost.InputPricePerMillionTokens))
require.True(t, expected.Cost.OutputPricePerMillionTokens.Equal(*actual.Cost.OutputPricePerMillionTokens))
require.True(t, expected.Cost.CacheReadPricePerMillionTokens.Equal(*actual.Cost.CacheReadPricePerMillionTokens))
require.True(t, expected.Cost.CacheWritePricePerMillionTokens.Equal(*actual.Cost.CacheWritePricePerMillionTokens))
}
func decRef(value string) *decimal.Decimal {
d := decimal.RequireFromString(value)
return &d
}
func TestWatchChatDesktop(t *testing.T) {
t.Parallel()