mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: back the per-chat cost endpoint with AI Gateway data (#27328)
## Stack Context
This stack removes native chat cost tracking and native chat usage
limits, making the AI Gateway the single source of AI spend data and
budget enforcement.
1. **This PR:** re-back the per-chat cost endpoint with AI Gateway data.
2. Remove native chat usage limits end to end, rewiring the sidebar
indicator to gateway spend.
3. Remove native chat cost tracking end to end, deleting the
Analytics/Spend cost UI.
## What?
`GET /api/experimental/chats/{chat}/cost` summed
`chat_messages.total_cost_micros`, which native chat cost tracking
maintained. It now aggregates AI Gateway interception data instead, and
has no native fallback.
- New `GetAIBridgeChatCost` query, authorized through the root chat so
members can read their own chat's cost without gaining access to raw
interception rows.
- Response fields renamed: `priced_message_count` -> `request_count`,
`unpriced_messages_having_usage_count` -> `unpriced_request_count`.
- The chat summary sidebar keys its cost cache by root chat, and hides
the cost row where the AI Gateway is off or unlicensed. The root cost is
invalidated when a chat leaves an active status and when a generated
title lands, since title generation bills its own gateway request.
`GetChatModelUsageCostByChatID` and the rest of native cost tracking are
untouched here; PR 3 removes them.
## Why?
Native cost tracking duplicates what the AI Gateway already records, and
the two disagree. Repointing the endpoint first means the cost UI keeps
working while the native implementation is deleted later in the stack.
Two behaviour changes follow from gateway semantics and are intentional:
- **Requests, not messages.** The gateway records interceptions, so
counts are requests. Title-generation traffic now counts.
- **Whole-tree totals.** The gateway records the *spawning* chat's ID as
the interception session ID, so a subagent's requests are attributed to
its immediate parent, not always the root. Only a whole chat tree can be
summed, so the query resolves the root and aggregates the tree, and
every chat in a tree reports the same total. Native returned per-subtree
totals.
## Attribution and counting semantics
The aggregate groups token usage per interception before counting, so
the reported numbers are per request even though a request records one
usage row per provider response:
- `RequestCount` counts finished `Coder Agents` interceptions in the
tree, including unpriced ones.
- `UnpricedRequestCount` counts requests with at least one usage row the
gateway could not price. It is a subset of `RequestCount`.
- `TotalCostMicros` omits only unpriced usage, so a partially priced
request still contributes its priced portion. The sidebar therefore says
`Excludes unpriced usage from N request(s)` rather than claiming whole
requests were dropped.
A recorded cost of zero is a free request, not an unpriced one. Usage
without an effective group is excluded, matching what never reached
`ai_user_daily_spend`.
## Authorization
Reads go through `ExtractChatParam` plus `ResourceChat`, with no
cost-specific RBAC widening. `TestGetChatCost/MemberCanReadOwnChat`
covers a scoped `agents-access` member reading their own chat's cost,
and `MemberCannotReadOtherUsersChat` still asserts 404 for a non-owner.
Plain members without `agents-access` cannot create or read chats at
all, so they never reach this endpoint.
## Known limitation
AI Gateway data has its own retention period, 60 days by default and
configured independently of chat retention, so spend for requests older
than that is no longer reported. A chat whose gateway records have all
been purged reports zero cost, which is indistinguishable from genuinely
free usage under this contract. The endpoint documents the caveat;
#27330 documents it on the Spend Management page.
In-flight interceptions are excluded, since cost is only known once the
response is recorded. A chat's cost therefore lags the active turn by
one request.
## Rebase note
Rebased onto `main` after #27579 removed the `ai-gateway-cost-control`
experiment. The per-chat cost row is now gated on the `aibridge` feature
alone, matching how #27579 degated the other cost-control surfaces.
> Mux prepared this PR on Mike's behalf.
This commit is contained in:
@@ -2793,6 +2793,16 @@ func (q *querier) FindMatchingPresetID(ctx context.Context, arg database.FindMat
|
||||
return q.db.FindMatchingPresetID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
// The aggregate covers one chat tree, so it is authorized through the
|
||||
// root chat. Members cannot read interception rows back, but they can
|
||||
// read their own chats.
|
||||
if _, err := q.GetChatByID(ctx, rootChatID); err != nil {
|
||||
return database.GetAIBridgeChatCostRow{}, err
|
||||
}
|
||||
return q.db.GetAIBridgeChatCost(ctx, rootChatID)
|
||||
}
|
||||
|
||||
func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
return fetch(q.log, q.auth, q.db.GetAIBridgeInterceptionByID)(ctx, id)
|
||||
}
|
||||
|
||||
@@ -1069,6 +1069,13 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().GetChatMessagesByChatID(gomock.Any(), arg).Return(msgs, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs)
|
||||
}))
|
||||
s.Run("GetAIBridgeChatCost", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
row := database.GetAIBridgeChatCostRow{TotalCostMicros: 1000, RequestCount: 2}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), chat.ID).Return(row, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(row)
|
||||
}))
|
||||
s.Run("GetChatModelUsageCostByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
row := database.GetChatModelUsageCostByChatIDRow{ChatID: chat.ID, TotalCostMicros: 1000, PricedMessageCount: 2}
|
||||
|
||||
+8
@@ -1081,6 +1081,14 @@ func (m queryMetricsStore) FindMatchingPresetID(ctx context.Context, arg databas
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAIBridgeChatCost(ctx, rootChatID)
|
||||
m.queryLatencies.WithLabelValues("GetAIBridgeChatCost").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBridgeChatCost").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAIBridgeInterceptionByID(ctx, id)
|
||||
|
||||
Generated
+15
@@ -1858,6 +1858,21 @@ func (mr *MockStoreMockRecorder) FindMatchingPresetID(ctx, arg any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindMatchingPresetID", reflect.TypeOf((*MockStore)(nil).FindMatchingPresetID), ctx, arg)
|
||||
}
|
||||
|
||||
// GetAIBridgeChatCost mocks base method.
|
||||
func (m *MockStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAIBridgeChatCost", ctx, rootChatID)
|
||||
ret0, _ := ret[0].(database.GetAIBridgeChatCostRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAIBridgeChatCost indicates an expected call of GetAIBridgeChatCost.
|
||||
func (mr *MockStoreMockRecorder) GetAIBridgeChatCost(ctx, rootChatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeChatCost", reflect.TypeOf((*MockStore)(nil).GetAIBridgeChatCost), ctx, rootChatID)
|
||||
}
|
||||
|
||||
// GetAIBridgeInterceptionByID mocks base method.
|
||||
func (m *MockStore) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+7
@@ -295,6 +295,13 @@ type sqlcQuerier interface {
|
||||
// The query finds presets where all preset parameters are present in the provided parameters,
|
||||
// and returns the preset with the most parameters (largest subset).
|
||||
FindMatchingPresetID(ctx context.Context, arg FindMatchingPresetIDParams) (uuid.UUID, error)
|
||||
// AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
// beneath it. The spawning chat's ID is recorded as the interception session
|
||||
// ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
// to its parent rather than the root, and only whole trees can be summed. The
|
||||
// owner check guards against session-id collisions. Usage without an
|
||||
// effective group never reaches ai_user_daily_spend.
|
||||
GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (GetAIBridgeChatCostRow, error)
|
||||
GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (AIBridgeInterception, error)
|
||||
// Look up the parent interception and the root of the thread by finding
|
||||
// which interception recorded a tool usage with the given tool call ID.
|
||||
|
||||
Generated
+53
@@ -1156,6 +1156,59 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti
|
||||
return total_deleted, err
|
||||
}
|
||||
|
||||
const getAIBridgeChatCost = `-- name: GetAIBridgeChatCost :one
|
||||
WITH per_request AS (
|
||||
-- One row per interception. A request records one token usage per provider
|
||||
-- response, so aggregating here keeps the outer counts per request and
|
||||
-- flags a request whose cost is partial because some usage was unpriced.
|
||||
-- The usage join is a LEFT JOIN so a request that ended without eligible
|
||||
-- usage, such as one that failed upstream, still counts as a request. The
|
||||
-- tu.id guard keeps that row from reading as unpriced usage, since the
|
||||
-- unmatched side is all NULL.
|
||||
SELECT
|
||||
SUM(tu.cost_micros) AS cost_micros,
|
||||
BOOL_OR(tu.id IS NOT NULL AND tu.cost_micros IS NULL) AS has_unpriced_usage
|
||||
FROM aibridge_interceptions i
|
||||
JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id
|
||||
LEFT JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL
|
||||
WHERE (
|
||||
-- Spelled out instead of COALESCE(c.root_chat_id, c.id) so each branch
|
||||
-- stays a plain comparison against an indexed column.
|
||||
c.root_chat_id = $1::uuid
|
||||
OR (c.root_chat_id IS NULL AND c.id = $1::uuid)
|
||||
)
|
||||
-- Restrict to aibridge.ClientCoderAgents so another client's session
|
||||
-- reference cannot match a chat ID.
|
||||
AND i.client = 'Coder Agents'
|
||||
AND i.ended_at IS NOT NULL
|
||||
GROUP BY i.id
|
||||
)
|
||||
SELECT
|
||||
COALESCE(SUM(cost_micros), 0)::bigint AS total_cost_micros,
|
||||
COUNT(*)::bigint AS request_count,
|
||||
COUNT(*) FILTER (WHERE has_unpriced_usage)::bigint AS unpriced_request_count
|
||||
FROM per_request
|
||||
`
|
||||
|
||||
type GetAIBridgeChatCostRow struct {
|
||||
TotalCostMicros int64 `db:"total_cost_micros" json:"total_cost_micros"`
|
||||
RequestCount int64 `db:"request_count" json:"request_count"`
|
||||
UnpricedRequestCount int64 `db:"unpriced_request_count" json:"unpriced_request_count"`
|
||||
}
|
||||
|
||||
// AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
// beneath it. The spawning chat's ID is recorded as the interception session
|
||||
// ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
// to its parent rather than the root, and only whole trees can be summed. The
|
||||
// owner check guards against session-id collisions. Usage without an
|
||||
// effective group never reaches ai_user_daily_spend.
|
||||
func (q *sqlQuerier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (GetAIBridgeChatCostRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, getAIBridgeChatCost, rootChatID)
|
||||
var i GetAIBridgeChatCostRow
|
||||
err := row.Scan(&i.TotalCostMicros, &i.RequestCount, &i.UnpricedRequestCount)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
|
||||
SELECT
|
||||
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, agent_firewall_session_id, agent_firewall_sequence_number, error_type, error_message
|
||||
|
||||
@@ -663,3 +663,42 @@ GROUP BY
|
||||
LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100)
|
||||
OFFSET @offset_
|
||||
;
|
||||
|
||||
-- name: GetAIBridgeChatCost :one
|
||||
-- AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
-- beneath it. The spawning chat's ID is recorded as the interception session
|
||||
-- ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
-- to its parent rather than the root, and only whole trees can be summed. The
|
||||
-- owner check guards against session-id collisions. Usage without an
|
||||
-- effective group never reaches ai_user_daily_spend.
|
||||
WITH per_request AS (
|
||||
-- One row per interception. A request records one token usage per provider
|
||||
-- response, so aggregating here keeps the outer counts per request and
|
||||
-- flags a request whose cost is partial because some usage was unpriced.
|
||||
-- The usage join is a LEFT JOIN so a request that ended without eligible
|
||||
-- usage, such as one that failed upstream, still counts as a request. The
|
||||
-- tu.id guard keeps that row from reading as unpriced usage, since the
|
||||
-- unmatched side is all NULL.
|
||||
SELECT
|
||||
SUM(tu.cost_micros) AS cost_micros,
|
||||
BOOL_OR(tu.id IS NOT NULL AND tu.cost_micros IS NULL) AS has_unpriced_usage
|
||||
FROM aibridge_interceptions i
|
||||
JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id
|
||||
LEFT JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL
|
||||
WHERE (
|
||||
-- Spelled out instead of COALESCE(c.root_chat_id, c.id) so each branch
|
||||
-- stays a plain comparison against an indexed column.
|
||||
c.root_chat_id = @root_chat_id::uuid
|
||||
OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid)
|
||||
)
|
||||
-- Restrict to aibridge.ClientCoderAgents so another client's session
|
||||
-- reference cannot match a chat ID.
|
||||
AND i.client = 'Coder Agents'
|
||||
AND i.ended_at IS NOT NULL
|
||||
GROUP BY i.id
|
||||
)
|
||||
SELECT
|
||||
COALESCE(SUM(cost_micros), 0)::bigint AS total_cost_micros,
|
||||
COUNT(*)::bigint AS request_count,
|
||||
COUNT(*) FILTER (WHERE has_unpriced_usage)::bigint AS unpriced_request_count
|
||||
FROM per_request;
|
||||
|
||||
Reference in New Issue
Block a user