fix(gateway): record prompt/completion token split for streaming request logs by nicemit · Pull Request #3543 · future-agi/future-agi · GitHub
Skip to content

fix(gateway): record prompt/completion token split for streaming request logs - #3543

Open
nicemit wants to merge 2 commits into
future-agi:devfrom
nicemit:fix/gateway-streaming-token-split
Open

fix(gateway): record prompt/completion token split for streaming request logs#3543
nicemit wants to merge 2 commits into
future-agi:devfrom
nicemit:fix/gateway-streaming-token-split

Conversation

@nicemit

@nicemit nicemit commented Oct 17, 2026

Copy link
Copy Markdown

Summary

Streaming responses through the gateway saved only total_tokens in request logs; prompt_tokens and completion_tokens stayed empty. The streaming path already reads usage from the final chunk, but models.Usage only decoded OpenAI field names, so providers reporting input_tokens/output_tokens (Anthropic-style) lost the split. This adds a fallback decoder so the split is recorded for those providers, which keeps cost, budget, and analytics numbers correct for streamed requests.

Linked issues

Closes #3446
Linear: n/a

AI use

AI use: AI-assisted coding, reviewed and tested by me

Type of change

  • 🐛 Bug fix (non-breaking change which fixes an issue)
  • ✨ New feature (non-breaking change which adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation only
  • 🧹 Chore / refactor (no user-visible change)
  • 🚀 Performance improvement
  • 🧪 Test-only change

E2E coverage

E2E: exempt (gateway-internal)

1) What changes were done

A. Usage decoding accepts input/output token names

  • User-visible: streamed request logs now show prompt and completion tokens, not just the total.
  • agentcc-gateway/internal/models/response.go: new Usage.UnmarshalJSON. It reads prompt_tokens/completion_tokens as before, falls back to input_tokens/output_tokens when those are absent, and derives total_tokens when it is missing or 0.
  • No API, schema, or storage changes.

B. Regression tests

  • agentcc-gateway/internal/models/usage_split_test.go: covers the alias decoding, unchanged OpenAI behavior, and total derivation.

2) Why the changes were done

  • Product/UX: Cost tracking, budgets, and analytics depend on the prompt/completion split; streamed requests were silently missing it.
  • Technical: Fixing it in the shared Usage decoder repairs every path that decodes provider usage (streaming chunks included) without touching each provider adapter.
  • Architecture: Keeps provider-format normalization in the model layer, where usage is already defined.

3) Tests written + scenarios each covers

agentcc-gateway/internal/models/usage_split_test.go — usage decoding across provider field names

  • TestStreamChunkUsage_InputOutputAliases — a stream chunk with input_tokens/output_tokens fills prompt/completion (failed before the fix: prompt=0 completion=0 total=300)
  • TestUsage_OpenAIFieldsUnchanged — OpenAI-format usage decodes exactly as before
  • TestUsage_TotalDerivedWhenMissing — total is computed from the split when omitted
$ cd agentcc-gateway && go test ./internal/models/ -run Usage -v
=== RUN   TestExtractRunUsage
--- PASS: TestExtractRunUsage (0.00s)
    --- PASS: TestExtractRunUsage/valid_usage (0.00s)
    --- PASS: TestExtractRunUsage/missing_usage (0.00s)
    --- PASS: TestExtractRunUsage/null_usage (0.00s)
    --- PASS: TestExtractRunUsage/malformed_json (0.00s)
=== RUN   TestStreamChunkUsage_InputOutputAliases
--- PASS: TestStreamChunkUsage_InputOutputAliases (0.00s)
=== RUN   TestUsage_OpenAIFieldsUnchanged
--- PASS: TestUsage_OpenAIFieldsUnchanged (0.00s)
=== RUN   TestUsage_TotalDerivedWhenMissing
--- PASS: TestUsage_TotalDerivedWhenMissing (0.00s)
PASS
ok      github.com/futureagi/agentcc-gateway/internal/models    0.009s

4) How to run / test

Gateway tests

cd agentcc-gateway
go test ./internal/models/ -run Usage -v
# full gateway suite:
go test ./...

Backend (bin/test), frontend, and contract checks are not affected; this is a Go gateway-only change.

5) Screenshots / recordings

Test output

See the test output in section 3.

UI testing

N/A (gateway-internal).

6) Edge cases & considerations

  • Both naming styles present: OpenAI fields take precedence.
  • Explicit non-zero total_tokens: never overwritten.
  • Usage missing or null: decodes to zero values as before (TestExtractRunUsage still passes).

7) Pre-existing issues (NOT introduced by this PR)

  • go vet reports an existing warning at errors_test.go:109.
  • TestOrgProviderBuildFailureIsA502ServerError/unresolvable_host can fail due to a DNS timeout in some environments; unrelated to this change.

8) Architectural / important decisions

  • Decode-time normalization: handled in Usage.UnmarshalJSON instead of per-provider stream handlers, so every caller gets the fix at once.
  • Fallback, not replacement: OpenAI field names stay primary to avoid changing behavior for existing providers.

Checklist

  • My code follows the style guide
  • I've added tests that prove my fix is effective or that my feature works
  • bin/test passes locally (backend) — N/A, Go gateway only
  • yarn test:run passes locally (frontend) — N/A
  • yarn contracts:check passes if API surface changed — N/A, no API change
  • I've updated the documentation where relevant — N/A
  • No hardcoded secrets, URLs, or PII
  • I've signed the CLA
  • PR title follows Conventional Commits
  • Linear issue is linked and updated with status — no Linear access

@github-actions

github-actions Bot commented Oct 17, 2026

Copy link
Copy Markdown

@github-actions github-actions Bot added the needs-admission Admission gate found template/issue/branch/hygiene problems label Oct 17, 2026
@nicemit nicemit mentioned this pull request Oct 17, 2026
Open
@nicemit
fix(gateway): take streaming input tokens from Anthropic message_delt…
9963b5f 
…a usage

Anthropic-compatible providers can report input_tokens as 0 in
message_start and send the real count in message_delta. The stream
parser only read input_tokens from message_start, so the final usage
chunk, and the request log built from it, recorded prompt_tokens=0
with total_tokens equal to output tokens alone.

Use message_delta's input_tokens when it is non-zero.
@nicemit

nicemit commented Oct 18, 2026

Copy link
Copy Markdown
Author

I traced the streaming usage path end to end, from the provider stream through the gateway's request logging to the backend ingestion. On current dev, prompt and completion tokens are passed through correctly, and the request-log row is built from the same usage data as the final chunk the client receives.

I did find one real gap along the way. The Anthropic stream parser only read input_tokens from the message_start event, so Anthropic-compatible providers that report it in message_delta instead were logged with prompt tokens as 0. I've added a second commit that handles this, with a regression test.

@strzhao, could you share which provider and api_format you're using, along with a raw SSE sample from the upstream? That would help confirm whether this covers your case.

@strzhao

strzhao commented Oct 18, 2026

Copy link
Copy Markdown

Here's our setup and a live capture that confirms your second commit covers our case exactly.

Deployment: a custom provider (named bigmodel-echo) pointing at Zhipu BigModel's Anthropic-compatible endpoint — https://open.bigmodel.cn/api/anthropic, api_format: anthropic, model glm-5.3-flash.

Production evidence (agentcc_request_log, our pilot stack): all 5 streaming rows for this provider have input_tokens = 0 with non-zero output tokens, while the 32 non-streaming rows carry normal input counts. Exactly the shape you traced.

Raw SSE from the upstream (one minimal streaming call just now, usage fields verbatim):

event: message_start
data: {"type":"message_start","message":{...,"usage":{"input_tokens": 0, "output_tokens": 0}}}

event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"max_tokens","stop_sequence":null},"usage":{"input_tokens": 15, "output_tokens": 32, "cache_read_input_tokens": 0, "service_tier":"standard"}}

event: message_stop

So BigModel reports the real input count only in message_delta — message_start carries zeros — which is precisely the parser gap your second commit (message_delta fallback) fixes. Happy to re-verify on our stack once it lands on dev; the same rows should start carrying the split.

@nicemit

nicemit commented Oct 18, 2026

Copy link
Copy Markdown
Author

Thanks @strzhao, that SSE capture confirms it. BigModel only reports input_tokens in message_delta, which is exactly what the second commit handles. I'd appreciate your re-check once it lands on dev.

Since the reporter has confirmed the root cause and the fix, could a maintainer label #3446 as accepted so this can go through review?

@nicemit

nicemit commented Oct 19, 2026

Copy link
Copy Markdown
Author

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

gateway needs-admission Admission gate found template/issue/branch/hygiene problems

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@nicemit @strzhao