From 83c62adeaddfe9037e852715d7825a331a61da91 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Tue, 25 Nov 2025 18:39:05 +0200 Subject: [PATCH] chore: actually store translated token metadata Signed-off-by: Danny Kopping --- enterprise/aibridged/aibridged_integration_test.go | 6 ++++-- enterprise/aibridged/translator.go | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/enterprise/aibridged/aibridged_integration_test.go b/enterprise/aibridged/aibridged_integration_test.go index 655992cbc0411..cb1a3ed8f1b3b 100644 --- a/enterprise/aibridged/aibridged_integration_test.go +++ b/enterprise/aibridged/aibridged_integration_test.go @@ -12,6 +12,7 @@ import ( "github.com/prometheus/client_golang/prometheus" promtest "github.com/prometheus/client_golang/prometheus/testutil" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" "github.com/coder/aibridge" "github.com/coder/coder/v2/coderd/coderdtest" @@ -107,7 +108,7 @@ func TestIntegration(t *testing.T) { "completion_tokens": 15, "total_tokens": 75, "prompt_tokens_details": { - "cached_tokens": 0, + "cached_tokens": 15, "audio_tokens": 0 }, "completion_tokens_details": { @@ -244,8 +245,9 @@ func TestIntegration(t *testing.T) { tokens, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, interceptions[0].ID) require.NoError(t, err) require.Len(t, tokens, 1) - require.EqualValues(t, tokens[0].InputTokens, 60) + require.EqualValues(t, tokens[0].InputTokens, 45) require.EqualValues(t, tokens[0].OutputTokens, 15) + require.EqualValues(t, gjson.Get(string(tokens[0].Metadata.RawMessage), "prompt_cached").Int(), 15) tools, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, interceptions[0].ID) require.NoError(t, err) diff --git a/enterprise/aibridged/translator.go b/enterprise/aibridged/translator.go index 73ae5e2ac706c..cbede0bc729f5 100644 --- a/enterprise/aibridged/translator.go +++ b/enterprise/aibridged/translator.go @@ -72,7 +72,7 @@ func (t *recorderTranslation) RecordTokenUsage(ctx context.Context, req *aibridg MsgId: req.MsgID, InputTokens: req.Input, OutputTokens: req.Output, - Metadata: marshalForProto(req.Metadata), + Metadata: marshalForProto(merged), CreatedAt: timestamppb.New(req.CreatedAt), }) return err