Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions coderd/apidoc/docs.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions coderd/apidoc/swagger.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 10 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
7 changes: 7 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
8 changes: 8 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 7 additions & 0 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

53 changes: 53 additions & 0 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

39 changes: 39 additions & 0 deletions coderd/database/queries/aibridge.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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;
31 changes: 23 additions & 8 deletions coderd/exp_chats.go
Original file line number Diff line number Diff line change
Expand Up @@ -2472,16 +2472,31 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) {
// @Success 200 {object} codersdk.ChatCost
// @Router /api/experimental/chats/{chat}/cost [get]
// @Description Experimental: this endpoint is subject to change.
// @Description
// @Description Cost covers the whole chat tree: the root chat plus every
// @Description subagent chat beneath it. Requesting cost for a subagent chat
// @Description returns that same total.
Comment thread
ibetitsmike marked this conversation as resolved.
// @Description
// @Description Cost is derived from AI Gateway data, which is subject to its
// @Description own retention period, 60 days by default, configured
// @Description independently of chat retention. Spend for requests older than
// @Description that period is no longer reported, so a chat whose requests
// @Description have all been purged reports zero cost.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
chat := httpmw.ChatParam(r)

// The query rolls up the requested chat's subtree, so a root chat
// reports itself plus all subagents while a subagent reports only
// its own spend (plus any nested subagents it spawned).
row, err := api.Database.GetChatModelUsageCostByChatID(ctx, chat.ID)
// AI Gateway attributes a subagent's requests to the chat that spawned
// it, so cost is only meaningful for a whole chat tree. Resolve the root
// chat and report the tree total, including for subagent chats.
rootChatID := chat.ID
if chat.RootChatID.Valid {
rootChatID = chat.RootChatID.UUID
Comment thread
ibetitsmike marked this conversation as resolved.
}

row, err := api.Database.GetAIBridgeChatCost(ctx, rootChatID)
if err != nil {
if httpapi.Is404Error(err) {
httpapi.ResourceNotFound(rw)
Expand All @@ -2495,10 +2510,10 @@ func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) {
}

httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatCost{
ChatID: row.ChatID,
TotalCostMicros: row.TotalCostMicros,
PricedMessageCount: row.PricedMessageCount,
UnpricedMessagesHavingUsageCount: row.UnpricedMessagesHavingUsageCount,
ChatID: chat.ID,
TotalCostMicros: row.TotalCostMicros,
RequestCount: row.RequestCount,
UnpricedRequestCount: row.UnpricedRequestCount,
})
}

Expand Down
29 changes: 15 additions & 14 deletions coderd/exp_chats_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@ import (
"github.com/coder/coder/v2/testutil"
)

// ExtractChatParam authorizes the read, then GetChatModelUsageCostByChatID
// authorizes it again. A denial on the second check means the ACL changed in
// between (a read-authz race). Assert it surfaces as 404, not 500.
// ExtractChatParam authorizes the read, then GetAIBridgeChatCost authorizes it
// again. A denial on the second check means the ACL changed in between (a
// read-authz race). Assert it surfaces as 404, not 500.
func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
t.Parallel()

Expand All @@ -38,8 +38,8 @@ func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
}

dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil)
dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), chat.ID).Return(
database.GetChatModelUsageCostByChatIDRow{},
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), chat.ID).Return(
database.GetAIBridgeChatCostRow{},
dbauthz.NotAuthorizedError{Err: sql.ErrNoRows},
)

Expand All @@ -56,9 +56,9 @@ func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

// A subagent chat's cost is scoped to its own subtree, so the handler
// must query the requested chat ID rather than resolving to the root.
func TestGetChatCostQueriesRequestedChat(t *testing.T) {
// AI Gateway attributes a subagent's requests to the chat that spawned it, so
// a subagent request must be answered with its root chat's tree cost.
func TestGetChatCostQueriesRootChat(t *testing.T) {
t.Parallel()

ctrl := gomock.NewController(t)
Expand All @@ -73,11 +73,11 @@ func TestGetChatCostQueriesRequestedChat(t *testing.T) {
}

dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil)
dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), child.ID).Return(
database.GetChatModelUsageCostByChatIDRow{
ChatID: child.ID,
TotalCostMicros: 250,
PricedMessageCount: 1,
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), rootID).Return(
database.GetAIBridgeChatCostRow{
TotalCostMicros: 250,
RequestCount: 2,
UnpricedRequestCount: 1,
},
nil,
)
Expand All @@ -97,7 +97,8 @@ func TestGetChatCostQueriesRequestedChat(t *testing.T) {
require.NoError(t, json.NewDecoder(resp.Body).Decode(&cost))
require.Equal(t, child.ID, cost.ChatID)
require.Equal(t, int64(250), cost.TotalCostMicros)
require.Equal(t, int64(1), cost.PricedMessageCount)
require.Equal(t, int64(2), cost.RequestCount)
require.Equal(t, int64(1), cost.UnpricedRequestCount)
}

func TestEnrichMissingChatAgentIDs(t *testing.T) {
Expand Down
Loading
Loading