Skip to content

Commit 4090ccf

Browse files
mostlydevmostlydev
committed
feat(cost): unpriced request tracking, reported cost passthrough, timezone context
- cost/usage.go: parse usage from raw map to handle provider-specific fields; add ReportedCostUSD field for passthrough of provider-reported cost - cost/accumulator.go: track unpriced_requests count (requests where cost is nil) - logging/logger.go: CostInfo.CostUSD is now *float64 (nil = unpriced, not zero) - proxy/handler.go: propagate unpriced request count through to cost response - proxy/time_context.go: timezone-aware current time injection for agents - ui/handler.go + dashboard.html: expose total_requests and unpriced_requests in costs API response; surface in dashboard UI
1 parent b20e7e1 commit 4090ccf

File tree

13 files changed

+970
-102
lines changed

13 files changed

+970
-102
lines changed

‎internal/cost/accumulator.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ type CostEntry struct {
1414
TotalOutputTokens int
1515
TotalCostUSD float64
1616
RequestCount int
17+
PricedRequests int
18+
UnpricedRequests int
1719
}
1820

1921
type bucketKey struct {
@@ -33,6 +35,10 @@ func NewAccumulator() *Accumulator {
3335
}
3436

3537
func (a *Accumulator) Record(agentID, provider, model string, inputTokens, outputTokens int, costUSD float64) {
38+
a.RecordWithStatus(agentID, provider, model, inputTokens, outputTokens, costUSD, true)
39+
}
40+
41+
func (a *Accumulator) RecordWithStatus(agentID, provider, model string, inputTokens, outputTokens int, costUSD float64, costKnown bool) {
3642
key := bucketKey{AgentID: agentID, Provider: provider, Model: model}
3743
a.mu.Lock()
3844
defer a.mu.Unlock()
@@ -45,6 +51,11 @@ func (a *Accumulator) Record(agentID, provider, model string, inputTokens, outpu
4551
e.TotalOutputTokens += outputTokens
4652
e.TotalCostUSD += costUSD
4753
e.RequestCount++
54+
if costKnown {
55+
e.PricedRequests++
56+
} else {
57+
e.UnpricedRequests++
58+
}
4859
}
4960

5061
// ByAgent returns all cost entries for a given agent, sorted by model.
@@ -89,3 +100,23 @@ func (a *Accumulator) TotalCost() float64 {
89100
}
90101
return total
91102
}
103+
104+
func (a *Accumulator) TotalRequests() int {
105+
a.mu.RLock()
106+
defer a.mu.RUnlock()
107+
total := 0
108+
for _, e := range a.buckets {
109+
total += e.RequestCount
110+
}
111+
return total
112+
}
113+
114+
func (a *Accumulator) TotalUnpricedRequests() int {
115+
a.mu.RLock()
116+
defer a.mu.RUnlock()
117+
total := 0
118+
for _, e := range a.buckets {
119+
total += e.UnpricedRequests
120+
}
121+
return total
122+
}

‎internal/cost/accumulator_test.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ func TestAccumulatorRecordAndQuery(t *testing.T) {
2424
if entry.RequestCount != 2 {
2525
t.Errorf("expected 2 requests, got %d", entry.RequestCount)
2626
}
27+
if entry.PricedRequests != 2 || entry.UnpricedRequests != 0 {
28+
t.Errorf("expected priced=2 unpriced=0, got priced=%d unpriced=%d", entry.PricedRequests, entry.UnpricedRequests)
29+
}
2730
}
2831

2932
func TestAccumulatorAllAgents(t *testing.T) {
@@ -46,3 +49,26 @@ func TestAccumulatorTotalCost(t *testing.T) {
4649
t.Errorf("expected ~0.003, got %f", total)
4750
}
4851
}
52+
53+
func TestAccumulatorTracksUnpricedRequests(t *testing.T) {
54+
a := NewAccumulator()
55+
a.RecordWithStatus("tiverton", "openrouter", "unknown/model", 500, 100, 0, false)
56+
57+
summary := a.ByAgent("tiverton")
58+
if len(summary) != 1 {
59+
t.Fatalf("expected 1 entry, got %d", len(summary))
60+
}
61+
entry := summary[0]
62+
if entry.PricedRequests != 0 {
63+
t.Errorf("expected 0 priced requests, got %d", entry.PricedRequests)
64+
}
65+
if entry.UnpricedRequests != 1 {
66+
t.Errorf("expected 1 unpriced request, got %d", entry.UnpricedRequests)
67+
}
68+
if got := a.TotalUnpricedRequests(); got != 1 {
69+
t.Errorf("expected total unpriced requests=1, got %d", got)
70+
}
71+
if got := a.TotalRequests(); got != 1 {
72+
t.Errorf("expected total requests=1, got %d", got)
73+
}
74+
}

‎internal/cost/usage.go‎

Lines changed: 133 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,33 +3,30 @@ package cost
33
import (
44
"bytes"
55
"encoding/json"
6+
"strconv"
67
)
78

89
// Usage holds token counts from an OpenAI-compatible response.
910
type Usage struct {
10-
PromptTokens int `json:"prompt_tokens"`
11-
CompletionTokens int `json:"completion_tokens"`
12-
TotalTokens int `json:"total_tokens"`
11+
PromptTokens int `json:"prompt_tokens"`
12+
CompletionTokens int `json:"completion_tokens"`
13+
TotalTokens int `json:"total_tokens"`
14+
ReportedCostUSD *float64 `json:"-"`
1315
}
1416

1517
// ExtractUsage parses usage from a non-streamed JSON response body.
1618
func ExtractUsage(body []byte) (Usage, error) {
17-
var resp struct {
18-
Usage *Usage `json:"usage"`
19-
}
20-
if err := json.Unmarshal(body, &resp); err != nil {
19+
var payload map[string]any
20+
if err := json.Unmarshal(body, &payload); err != nil {
2121
return Usage{}, err
2222
}
23-
if resp.Usage == nil {
24-
return Usage{}, nil
25-
}
26-
return *resp.Usage, nil
23+
return extractUsageFromPayload(payload), nil
2724
}
2825

2926
// ExtractUsageFromSSE scans SSE data lines for the last one containing a "usage" field.
3027
// OpenAI streams include usage in the final data chunk before "data: [DONE]".
3128
func ExtractUsageFromSSE(stream []byte) (Usage, error) {
32-
var lastUsage Usage
29+
var observed Usage
3330
for _, line := range bytes.Split(stream, []byte("\n")) {
3431
line = bytes.TrimSpace(line)
3532
if !bytes.HasPrefix(line, []byte("data: ")) {
@@ -39,12 +36,131 @@ func ExtractUsageFromSSE(stream []byte) (Usage, error) {
3936
if bytes.Equal(payload, []byte("[DONE]")) {
4037
continue
4138
}
42-
var chunk struct {
43-
Usage *Usage `json:"usage"`
39+
var chunk map[string]any
40+
if json.Unmarshal(payload, &chunk) == nil {
41+
observed = mergeUsage(observed, extractUsageFromPayload(chunk))
42+
}
43+
}
44+
return observed, nil
45+
}
46+
47+
func extractUsageFromPayload(payload map[string]any) Usage {
48+
var usage Usage
49+
if raw, ok := payload["usage"]; ok {
50+
usage = mergeUsage(usage, parseUsageObject(raw))
51+
}
52+
if rawMessage, ok := payload["message"].(map[string]any); ok {
53+
if rawUsage, ok := rawMessage["usage"]; ok {
54+
usage = mergeUsage(usage, parseUsageObject(rawUsage))
55+
}
56+
}
57+
return usage
58+
}
59+
60+
func parseUsageObject(raw any) Usage {
61+
obj, ok := raw.(map[string]any)
62+
if !ok {
63+
return Usage{}
64+
}
65+
66+
prompt := firstInt(obj, "prompt_tokens", "input_tokens")
67+
completion := firstInt(obj, "completion_tokens", "output_tokens")
68+
total := intFromAny(obj["total_tokens"])
69+
if total == 0 && (prompt > 0 || completion > 0) {
70+
total = prompt + completion
71+
}
72+
73+
return Usage{
74+
PromptTokens: prompt,
75+
CompletionTokens: completion,
76+
TotalTokens: total,
77+
ReportedCostUSD: floatPtrFromAny(obj["cost"]),
78+
}
79+
}
80+
81+
func mergeUsage(base, next Usage) Usage {
82+
if next.PromptTokens > base.PromptTokens {
83+
base.PromptTokens = next.PromptTokens
84+
}
85+
if next.CompletionTokens > base.CompletionTokens {
86+
base.CompletionTokens = next.CompletionTokens
87+
}
88+
if next.TotalTokens > base.TotalTokens {
89+
base.TotalTokens = next.TotalTokens
90+
}
91+
if derivedTotal := base.PromptTokens + base.CompletionTokens; derivedTotal > base.TotalTokens {
92+
base.TotalTokens = derivedTotal
93+
}
94+
if next.ReportedCostUSD != nil {
95+
if base.ReportedCostUSD == nil || *next.ReportedCostUSD > *base.ReportedCostUSD {
96+
cost := *next.ReportedCostUSD
97+
base.ReportedCostUSD = &cost
98+
}
99+
}
100+
return base
101+
}
102+
103+
func firstInt(obj map[string]any, keys ...string) int {
104+
for _, key := range keys {
105+
if v, ok := obj[key]; ok {
106+
if parsed := intFromAny(v); parsed > 0 {
107+
return parsed
108+
}
109+
}
110+
}
111+
return 0
112+
}
113+
114+
func intFromAny(v any) int {
115+
switch t := v.(type) {
116+
case float64:
117+
return int(t)
118+
case float32:
119+
return int(t)
120+
case int:
121+
return t
122+
case int64:
123+
return int(t)
124+
case int32:
125+
return int(t)
126+
case json.Number:
127+
n, err := t.Int64()
128+
if err == nil {
129+
return int(n)
130+
}
131+
case string:
132+
n, err := strconv.Atoi(t)
133+
if err == nil {
134+
return n
135+
}
136+
}
137+
return 0
138+
}
139+
140+
func floatPtrFromAny(v any) *float64 {
141+
switch t := v.(type) {
142+
case float64:
143+
out := t
144+
return &out
145+
case float32:
146+
out := float64(t)
147+
return &out
148+
case int:
149+
out := float64(t)
150+
return &out
151+
case int64:
152+
out := float64(t)
153+
return &out
154+
case json.Number:
155+
n, err := t.Float64()
156+
if err == nil {
157+
return &n
44158
}
45-
if json.Unmarshal(payload, &chunk) == nil && chunk.Usage != nil {
46-
lastUsage = *chunk.Usage
159+
case string:
160+
n, err := strconv.ParseFloat(t, 64)
161+
if err == nil {
162+
return &n
47163
}
48164
}
49-
return lastUsage, nil
165+
return nil
50166
}

‎internal/cost/usage_test.go‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,53 @@ func TestExtractUsageFromJSON(t *testing.T) {
2525
}
2626
}
2727

28+
func TestExtractUsageFromAnthropicJSON(t *testing.T) {
29+
body := []byte(`{
30+
"id": "msg_01",
31+
"type": "message",
32+
"usage": {
33+
"input_tokens": 321,
34+
"output_tokens": 89
35+
}
36+
}`)
37+
38+
u, err := ExtractUsage(body)
39+
if err != nil {
40+
t.Fatal(err)
41+
}
42+
if u.PromptTokens != 321 {
43+
t.Errorf("expected 321 prompt tokens, got %d", u.PromptTokens)
44+
}
45+
if u.CompletionTokens != 89 {
46+
t.Errorf("expected 89 completion tokens, got %d", u.CompletionTokens)
47+
}
48+
if u.TotalTokens != 410 {
49+
t.Errorf("expected total tokens derived from input/output, got %d", u.TotalTokens)
50+
}
51+
}
52+
53+
func TestExtractUsageIncludesReportedCost(t *testing.T) {
54+
body := []byte(`{
55+
"id": "chatcmpl-1",
56+
"usage": {
57+
"prompt_tokens": 120,
58+
"completion_tokens": 30,
59+
"cost": 0.0042
60+
}
61+
}`)
62+
63+
u, err := ExtractUsage(body)
64+
if err != nil {
65+
t.Fatal(err)
66+
}
67+
if u.ReportedCostUSD == nil {
68+
t.Fatal("expected reported cost to be parsed")
69+
}
70+
if *u.ReportedCostUSD != 0.0042 {
71+
t.Fatalf("expected reported cost 0.0042, got %f", *u.ReportedCostUSD)
72+
}
73+
}
74+
2875
func TestExtractUsageMissing(t *testing.T) {
2976
body := []byte(`{"id": "chatcmpl-1", "choices": []}`)
3077
u, err := ExtractUsage(body)
@@ -53,6 +100,29 @@ func TestExtractUsageFromSSE(t *testing.T) {
53100
}
54101
}
55102

103+
func TestExtractUsageFromAnthropicSSE(t *testing.T) {
104+
stream := []byte("event: message_start\n" +
105+
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":150}}}\n\n" +
106+
"event: message_delta\n" +
107+
"data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":37}}\n\n" +
108+
"event: message_stop\n" +
109+
"data: {\"type\":\"message_stop\"}\n\n")
110+
111+
u, err := ExtractUsageFromSSE(stream)
112+
if err != nil {
113+
t.Fatal(err)
114+
}
115+
if u.PromptTokens != 150 {
116+
t.Errorf("expected 150 prompt tokens, got %d", u.PromptTokens)
117+
}
118+
if u.CompletionTokens != 37 {
119+
t.Errorf("expected 37 completion tokens, got %d", u.CompletionTokens)
120+
}
121+
if u.TotalTokens != 187 {
122+
t.Errorf("expected derived total tokens, got %d", u.TotalTokens)
123+
}
124+
}
125+
56126
func TestExtractUsageFromSSENoUsage(t *testing.T) {
57127
stream := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n")
58128
u, err := ExtractUsageFromSSE(stream)

0 commit comments

Comments
 (0)