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
+19 -14
View File
@@ -28,13 +28,15 @@ type AIBridgeInterception struct {
}
type AIBridgeTokenUsage struct {
ID uuid.UUID `json:"id" format:"uuid"`
InterceptionID uuid.UUID `json:"interception_id" format:"uuid"`
ProviderResponseID string `json:"provider_response_id"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"created_at" format:"date-time"`
ID uuid.UUID `json:"id" format:"uuid"`
InterceptionID uuid.UUID `json:"interception_id" format:"uuid"`
ProviderResponseID string `json:"provider_response_id"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
Metadata map[string]any `json:"metadata"`
CreatedAt time.Time `json:"created_at" format:"date-time"`
}
type AIBridgeUserPrompt struct {
@@ -79,8 +81,10 @@ type AIBridgeSession struct {
}
type AIBridgeSessionTokenUsageSummary struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
}
type AIBridgeListSessionsResponse struct {
@@ -107,12 +111,13 @@ type AIBridgeSessionThreadsResponse struct {
}
// AIBridgeSessionThreadsTokenUsage represents aggregated token usage
// with metadata containing provider-specific fields like
// cache_creation_input, cache_read_input, etc.
// with metadata containing provider-specific fields.
type AIBridgeSessionThreadsTokenUsage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
Metadata map[string]any `json:"metadata"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
Metadata map[string]any `json:"metadata"`
}
// AIBridgeThread represents a single thread within a session.