Skip to content

Commit 4d943e3

Browse files
fix: extract tokens from SSE stream (message_start + message_delta events)
1 parent 6daae64 commit 4d943e3

2 files changed

Lines changed: 37 additions & 2 deletions

File tree

aperture.exe~

1 KB
Binary file not shown.

internal/server/server.go

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -444,16 +444,51 @@ func (s *Server) calculateModelCost(model string, tokensIn, tokensOut int) float
444444
}
445445

446446
func (s *Server) extractAnthropicTokens(data []byte) (int, int) {
447+
// Try non-stream format first (single JSON object)
447448
var resp struct {
448449
Usage struct {
449450
InputTokens int `json:"input_tokens"`
450451
OutputTokens int `json:"output_tokens"`
451452
} `json:"usage"`
452453
}
453-
if json.Unmarshal(data, &resp) == nil {
454+
if json.Unmarshal(data, &resp) == nil && (resp.Usage.InputTokens > 0 || resp.Usage.OutputTokens > 0) {
454455
return resp.Usage.InputTokens, resp.Usage.OutputTokens
455456
}
456-
return 0, 0
457+
458+
// SSE stream: scan each line for message_start (input) and message_delta (output)
459+
var inputTokens, outputTokens int
460+
for _, line := range bytes.Split(data, []byte{'\n'}) {
461+
if !bytes.HasPrefix(line, []byte("data: ")) {
462+
continue
463+
}
464+
payload := bytes.TrimPrefix(line, []byte("data: "))
465+
466+
// message_start: {"type":"message_start","message":{"usage":{"input_tokens":N}}}
467+
var start struct {
468+
Type string `json:"type"`
469+
Message struct {
470+
Usage struct {
471+
InputTokens int `json:"input_tokens"`
472+
} `json:"usage"`
473+
} `json:"message"`
474+
}
475+
if json.Unmarshal(payload, &start) == nil && start.Type == "message_start" {
476+
inputTokens = start.Message.Usage.InputTokens
477+
}
478+
479+
// message_delta: {"type":"message_delta","usage":{"output_tokens":N}}
480+
var delta struct {
481+
Type string `json:"type"`
482+
Usage struct {
483+
OutputTokens int `json:"output_tokens"`
484+
} `json:"usage"`
485+
}
486+
if json.Unmarshal(payload, &delta) == nil && delta.Type == "message_delta" {
487+
outputTokens = delta.Usage.OutputTokens
488+
}
489+
}
490+
491+
return inputTokens, outputTokens
457492
}
458493

459494
func (s *Server) getProviderAnthropicURL(model string) string {

0 commit comments

Comments
 (0)