feat: add columns for cached tokens from aibridge (#23832)

Two new columns added to aibridge_token_usages:
  - cache_read_input_tokens (BIGINT, default 0)
  - cache_write_input_tokens (BIGINT, default 0)

Migration backfills existing rows by extracting values from the metadata
JSONB column (cache_read_input, input_cached, prompt_cached for reads
(max value selected since only 1 should be set), cache_creation_input
for writes).

All references to data from metadata were updated to reference new
columns. No other changes then changing where data is extracted from.

Requires aibridge library version bump to include:
https://github.com/coder/aibridge/pull/229
Fixes: https://github.com/coder/aibridge/issues/150
This commit is contained in:
Paweł Banaszewski
2026-04-03 16:27:31 +02:00
committed by GitHub
parent da3c46b557
commit 8369fa88fd
28 changed files with 423 additions and 212 deletions
+26 -18
View File
@@ -1037,8 +1037,10 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
StartedAt: row.StartedAt,
Threads: row.Threads,
TokenUsageSummary: codersdk.AIBridgeSessionTokenUsageSummary{
InputTokens: row.InputTokens,
OutputTokens: row.OutputTokens,
InputTokens: row.InputTokens,
OutputTokens: row.OutputTokens,
CacheReadInputTokens: row.CacheReadInputTokens,
CacheWriteInputTokens: row.CacheWriteInputTokens,
},
}
// Ensure non-nil slices for JSON serialization.
@@ -1062,13 +1064,15 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
func AIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage {
return codersdk.AIBridgeTokenUsage{
ID: usage.ID,
InterceptionID: usage.InterceptionID,
ProviderResponseID: usage.ProviderResponseID,
InputTokens: usage.InputTokens,
OutputTokens: usage.OutputTokens,
Metadata: jsonOrEmptyMap(usage.Metadata),
CreatedAt: usage.CreatedAt,
ID: usage.ID,
InterceptionID: usage.InterceptionID,
ProviderResponseID: usage.ProviderResponseID,
InputTokens: usage.InputTokens,
OutputTokens: usage.OutputTokens,
CacheReadInputTokens: usage.CacheReadInputTokens,
CacheWriteInputTokens: usage.CacheWriteInputTokens,
Metadata: jsonOrEmptyMap(usage.Metadata),
CreatedAt: usage.CreatedAt,
}
}
@@ -1179,9 +1183,11 @@ func AIBridgeSessionThreads(
PageStartedAt: pageStartedAt,
PageEndedAt: pageEndedAt,
TokenUsageSummary: codersdk.AIBridgeSessionThreadsTokenUsage{
InputTokens: session.InputTokens,
OutputTokens: session.OutputTokens,
Metadata: sessionTokenMeta,
InputTokens: session.InputTokens,
OutputTokens: session.OutputTokens,
CacheReadInputTokens: session.CacheReadInputTokens,
CacheWriteInputTokens: session.CacheWriteInputTokens,
Metadata: sessionTokenMeta,
},
Threads: threads,
}
@@ -1314,17 +1320,19 @@ func buildAIBridgeThread(
// aggregateTokenUsage sums token usage rows and aggregates metadata.
func aggregateTokenUsage(tokens []database.AIBridgeTokenUsage) codersdk.AIBridgeSessionThreadsTokenUsage {
var inputTokens, outputTokens int64
var inputTokens, outputTokens, cacheRead, cacheWrite int64
for _, tu := range tokens {
inputTokens += tu.InputTokens
outputTokens += tu.OutputTokens
// TODO: once https://github.com/coder/aibridge/issues/150 lands we
// should aggregate the other token types.
cacheRead += tu.CacheReadInputTokens
cacheWrite += tu.CacheWriteInputTokens
}
return codersdk.AIBridgeSessionThreadsTokenUsage{
InputTokens: inputTokens,
OutputTokens: outputTokens,
Metadata: aggregateTokenMetadata(tokens),
InputTokens: inputTokens,
OutputTokens: outputTokens,
CacheReadInputTokens: cacheRead,
CacheWriteInputTokens: cacheWrite,
Metadata: aggregateTokenMetadata(tokens),
}
}
+9 -5
View File
@@ -259,11 +259,13 @@ func TestAIBridgeInterception(t *testing.T) {
},
tokenUsages: []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
InterceptionID: interceptionID,
ProviderResponseID: "resp-123",
InputTokens: 100,
OutputTokens: 200,
ID: uuid.New(),
InterceptionID: interceptionID,
ProviderResponseID: "resp-123",
InputTokens: 100,
OutputTokens: 200,
CacheReadInputTokens: 50,
CacheWriteInputTokens: 10,
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"cache":"hit"}`),
Valid: true,
@@ -413,6 +415,8 @@ func TestAIBridgeInterception(t *testing.T) {
require.Equal(t, tu.ProviderResponseID, result.TokenUsages[i].ProviderResponseID)
require.Equal(t, tu.InputTokens, result.TokenUsages[i].InputTokens)
require.Equal(t, tu.OutputTokens, result.TokenUsages[i].OutputTokens)
require.Equal(t, tu.CacheReadInputTokens, result.TokenUsages[i].CacheReadInputTokens)
require.Equal(t, tu.CacheWriteInputTokens, result.TokenUsages[i].CacheWriteInputTokens)
}
// Verify user prompts are converted correctly.
+9 -7
View File
@@ -1613,13 +1613,15 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
func AIBridgeTokenUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeTokenUsageParams) database.AIBridgeTokenUsage {
usage, err := db.InsertAIBridgeTokenUsage(genCtx, database.InsertAIBridgeTokenUsageParams{
ID: takeFirst(seed.ID, uuid.New()),
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
InputTokens: takeFirst(seed.InputTokens, 100),
OutputTokens: takeFirst(seed.OutputTokens, 100),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
ID: takeFirst(seed.ID, uuid.New()),
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
InputTokens: takeFirst(seed.InputTokens, 100),
OutputTokens: takeFirst(seed.OutputTokens, 100),
CacheReadInputTokens: seed.CacheReadInputTokens,
CacheWriteInputTokens: seed.CacheWriteInputTokens,
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert aibridge token usage")
return usage
+3 -1
View File
@@ -1134,7 +1134,9 @@ CREATE TABLE aibridge_token_usages (
input_tokens bigint NOT NULL,
output_tokens bigint NOT NULL,
metadata jsonb,
created_at timestamp with time zone NOT NULL
created_at timestamp with time zone NOT NULL,
cache_read_input_tokens bigint DEFAULT 0 NOT NULL,
cache_write_input_tokens bigint DEFAULT 0 NOT NULL
);
COMMENT ON TABLE aibridge_token_usages IS 'Audit log of tokens used by intercepted requests in AI Bridge';
@@ -0,0 +1,3 @@
ALTER TABLE aibridge_token_usages
DROP COLUMN cache_read_input_tokens,
DROP COLUMN cache_write_input_tokens;
@@ -0,0 +1,26 @@
ALTER TABLE aibridge_token_usages
ADD COLUMN cache_read_input_tokens BIGINT NOT NULL DEFAULT 0,
ADD COLUMN cache_write_input_tokens BIGINT NOT NULL DEFAULT 0;
-- Backfill from metadata JSONB. Old rows stored cache tokens under
-- provider-specific keys; new rows use the dedicated columns above.
UPDATE aibridge_token_usages
SET
-- Cache-read metadata keys by provider:
-- Anthropic (/v1/messages): "cache_read_input"
-- OpenAI (/v1/responses): "input_cached"
-- OpenAI (/v1/chat/completions): "prompt_cached"
cache_read_input_tokens = GREATEST(
COALESCE((metadata->>'cache_read_input')::bigint, 0),
COALESCE((metadata->>'input_cached')::bigint, 0),
COALESCE((metadata->>'prompt_cached')::bigint, 0)
),
-- Cache-write metadata keys by provider:
-- Anthropic (/v1/messages): "cache_creation_input"
-- OpenAI does not report cache-write tokens.
cache_write_input_tokens = COALESCE((metadata->>'cache_creation_input')::bigint, 0)
WHERE metadata IS NOT NULL
AND cache_read_input_tokens = 0
AND cache_write_input_tokens = 0;
+2
View File
@@ -1029,6 +1029,8 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg Lis
&i.Threads,
&i.InputTokens,
&i.OutputTokens,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
&i.LastPrompt,
); err != nil {
return nil, err
+7 -5
View File
@@ -4055,11 +4055,13 @@ type AIBridgeTokenUsage struct {
ID uuid.UUID `db:"id" json:"id"`
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
// The ID for the response in which the tokens were used, produced by the provider.
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
}
// Audit log of tool calls in intercepted requests in AI Bridge
+48 -43
View File
@@ -148,21 +148,8 @@ token_aggregates AS (
SELECT
COALESCE(SUM(tu.input_tokens), 0) AS token_count_input,
COALESCE(SUM(tu.output_tokens), 0) AS token_count_output,
-- Cached tokens are stored in metadata JSON, extract if available.
-- Read tokens may be stored in:
-- - cache_read_input (Anthropic)
-- - prompt_cached (OpenAI)
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_read_input')::bigint, 0) +
COALESCE((tu.metadata->>'prompt_cached')::bigint, 0)
), 0) AS token_count_cached_read,
-- Written tokens may be stored in:
-- - cache_creation_input (Anthropic)
-- Note that cache_ephemeral_5m_input and cache_ephemeral_1h_input on
-- Anthropic are included in the cache_creation_input field.
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_creation_input')::bigint, 0)
), 0) AS token_count_cached_written,
COALESCE(SUM(tu.cache_read_input_tokens), 0) AS token_count_cached_read,
COALESCE(SUM(tu.cache_write_input_tokens), 0) AS token_count_cached_written,
COUNT(tu.id) AS token_usages_count
FROM
interceptions_in_range i
@@ -559,7 +546,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
const getAIBridgeTokenUsagesByInterceptionID = `-- name: GetAIBridgeTokenUsagesByInterceptionID :many
SELECT
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
FROM
aibridge_token_usages WHERE interception_id = $1::uuid
ORDER BY
@@ -584,6 +571,8 @@ func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context,
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
); err != nil {
return nil, err
}
@@ -781,21 +770,23 @@ func (q *sqlQuerier) InsertAIBridgeModelThought(ctx context.Context, arg InsertA
const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :one
INSERT INTO aibridge_token_usages (
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at
) VALUES (
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7
$1, $2, $3, $4, $5, $6, $7, COALESCE($8::jsonb, '{}'::jsonb), $9
)
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
`
type InsertAIBridgeTokenUsageParams struct {
ID uuid.UUID `db:"id" json:"id"`
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
ID uuid.UUID `db:"id" json:"id"`
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
}
func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) {
@@ -805,6 +796,8 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
arg.ProviderResponseID,
arg.InputTokens,
arg.OutputTokens,
arg.CacheReadInputTokens,
arg.CacheWriteInputTokens,
arg.Metadata,
arg.CreatedAt,
)
@@ -817,6 +810,8 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
)
return i, err
}
@@ -1448,6 +1443,8 @@ SELECT
sp.threads,
COALESCE(st.input_tokens, 0)::bigint AS input_tokens,
COALESCE(st.output_tokens, 0)::bigint AS output_tokens,
COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens,
COALESCE(st.cache_write_input_tokens, 0)::bigint AS cache_write_input_tokens,
COALESCE(slp.prompt, '') AS last_prompt
FROM
session_page sp
@@ -1469,7 +1466,9 @@ LEFT JOIN LATERAL (
-- Aggregate tokens only for this session's interceptions.
SELECT
COALESCE(SUM(tu.input_tokens), 0)::bigint AS input_tokens,
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens,
COALESCE(SUM(tu.cache_read_input_tokens), 0)::bigint AS cache_read_input_tokens,
COALESCE(SUM(tu.cache_write_input_tokens), 0)::bigint AS cache_write_input_tokens
FROM aibridge_token_usages tu
WHERE tu.interception_id = ANY(sr.interception_ids)
) st ON true
@@ -1501,21 +1500,23 @@ type ListAIBridgeSessionsParams struct {
}
type ListAIBridgeSessionsRow struct {
SessionID string `db:"session_id" json:"session_id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
UserUsername string `db:"user_username" json:"user_username"`
UserName string `db:"user_name" json:"user_name"`
UserAvatarUrl string `db:"user_avatar_url" json:"user_avatar_url"`
Providers []string `db:"providers" json:"providers"`
Models []string `db:"models" json:"models"`
Client string `db:"client" json:"client"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
EndedAt time.Time `db:"ended_at" json:"ended_at"`
Threads int64 `db:"threads" json:"threads"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
LastPrompt string `db:"last_prompt" json:"last_prompt"`
SessionID string `db:"session_id" json:"session_id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
UserUsername string `db:"user_username" json:"user_username"`
UserName string `db:"user_name" json:"user_name"`
UserAvatarUrl string `db:"user_avatar_url" json:"user_avatar_url"`
Providers []string `db:"providers" json:"providers"`
Models []string `db:"models" json:"models"`
Client string `db:"client" json:"client"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
EndedAt time.Time `db:"ended_at" json:"ended_at"`
Threads int64 `db:"threads" json:"threads"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
LastPrompt string `db:"last_prompt" json:"last_prompt"`
}
// Returns paginated sessions with aggregated metadata, token counts, and
@@ -1560,6 +1561,8 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
&i.Threads,
&i.InputTokens,
&i.OutputTokens,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
&i.LastPrompt,
); err != nil {
return nil, err
@@ -1577,7 +1580,7 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
const listAIBridgeTokenUsagesByInterceptionIDs = `-- name: ListAIBridgeTokenUsagesByInterceptionIDs :many
SELECT
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
FROM
aibridge_token_usages
WHERE
@@ -1604,6 +1607,8 @@ func (q *sqlQuerier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Contex
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
); err != nil {
return nil, err
}
+9 -18
View File
@@ -31,9 +31,9 @@ WHERE aibridge_interceptions.id = (
-- name: InsertAIBridgeTokenUsage :one
INSERT INTO aibridge_token_usages (
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at
) VALUES (
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, @cache_read_input_tokens, @cache_write_input_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
)
RETURNING *;
@@ -299,21 +299,8 @@ token_aggregates AS (
SELECT
COALESCE(SUM(tu.input_tokens), 0) AS token_count_input,
COALESCE(SUM(tu.output_tokens), 0) AS token_count_output,
-- Cached tokens are stored in metadata JSON, extract if available.
-- Read tokens may be stored in:
-- - cache_read_input (Anthropic)
-- - prompt_cached (OpenAI)
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_read_input')::bigint, 0) +
COALESCE((tu.metadata->>'prompt_cached')::bigint, 0)
), 0) AS token_count_cached_read,
-- Written tokens may be stored in:
-- - cache_creation_input (Anthropic)
-- Note that cache_ephemeral_5m_input and cache_ephemeral_1h_input on
-- Anthropic are included in the cache_creation_input field.
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_creation_input')::bigint, 0)
), 0) AS token_count_cached_written,
COALESCE(SUM(tu.cache_read_input_tokens), 0) AS token_count_cached_read,
COALESCE(SUM(tu.cache_write_input_tokens), 0) AS token_count_cached_written,
COUNT(tu.id) AS token_usages_count
FROM
interceptions_in_range i
@@ -552,6 +539,8 @@ SELECT
sp.threads,
COALESCE(st.input_tokens, 0)::bigint AS input_tokens,
COALESCE(st.output_tokens, 0)::bigint AS output_tokens,
COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens,
COALESCE(st.cache_write_input_tokens, 0)::bigint AS cache_write_input_tokens,
COALESCE(slp.prompt, '') AS last_prompt
FROM
session_page sp
@@ -573,7 +562,9 @@ LEFT JOIN LATERAL (
-- Aggregate tokens only for this session's interceptions.
SELECT
COALESCE(SUM(tu.input_tokens), 0)::bigint AS input_tokens,
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens,
COALESCE(SUM(tu.cache_read_input_tokens), 0)::bigint AS cache_read_input_tokens,
COALESCE(SUM(tu.cache_write_input_tokens), 0)::bigint AS cache_write_input_tokens
FROM aibridge_token_usages tu
WHERE tu.interception_id = ANY(sr.interception_ids)
) st ON true