[Refactor] Decouple TTS usage metrics from OpenAI-compatible `/v1/audio/speech` response

[zhaochenyang20]

Background

In PR #173 we added X-Prompt-Tokens, X-Completion-Tokens, and X-Engine-Time as custom HTTP response headers on the /v1/audio/speech endpoint to expose per-request token usage and engine timing for benchmarking.

While this works, it does not conform to the OpenAI /v1/audio/speech API specification. The OpenAI speech endpoint returns a raw audio body with no usage metadata — there is no standard way to carry token counts or timing info in the response.

Embedding custom headers is a pragmatic short-term solution, but it diverges from the API contract we claim to be compatible with, and may confuse clients that expect strict OpenAI compatibility.

Problem

• Custom X-* headers on a compatibility endpoint break the "drop-in replacement" promise.
• There is no OpenAI-standard mechanism to return usage info from the speech endpoint.
• As we add more models and metrics, stuffing everything into headers does not scale.

Possible Directions (open for discussion)

1. Separate /v1/audio/speech/usage or query param — return usage in a sidecar endpoint or via ?include_usage=true that wraps the response in JSON.
2. Trailing headers (HTTP chunked) — send audio as chunked body, append usage as trailing headers. Requires client support.
3. Keep headers but behind an opt-in flag — only emit X-* headers when the client sends a specific request header (e.g., X-Include-Usage: true), so default behavior stays OpenAI-compatible.
4. SSE / streaming mode with structured events — similar to chat completions streaming, emit audio chunks + a final usage event.

None of these is clearly superior. Community input is welcome.

Current Behavior (status quo)

The /v1/audio/speech endpoint returns:

• Body: raw WAV/MP3 audio bytes
• Headers: X-Prompt-Tokens, X-Completion-Tokens, X-Engine-Time (non-standard)

Relevant Code

• sglang_omni/serve/openai_api.py — header injection in _register_speech
• sglang_omni/models/fishaudio_s2_pro/pipeline/stages.py — usage dict construction in vocoder stage
• sglang_omni/client/types.py — UsageInfo.engine_time_s field
• benchmarks/benchmark_tts_speed.py — reads X-* headers for metrics

[Ratish1]

Hey @zhaochenyang20 , I think there are really two valid options here, and they solve the issue in two different ways:

• Option A is the lowest-risk compatibility fix for the current issue.
• Option B is the cleaner long-term architecture if we want to solve this permanently with a structured API.

I would suggest both and we can discuss further:

• if the goal is a minimal, safe fix now, choose Option A
• if the goal is the durable end-state and we are willing to expand scope, choose Option B

Current architecture

The internal stack is already mostly decoupled:

• sglang_omni/models/fishaudio_s2_pro/pipeline/engine_io.py computes prompt_tokens and completion_tokens
• sglang_omni/executors/engine_executor.py injects engine_time_s
• sglang_omni/models/fishaudio_s2_pro/pipeline/stages.py assembles those fields into payload.data["usage"]
• sglang_omni/client/client.py preserves that as SpeechResult.usage
• sglang_omni/serve/openai_api.py is the only layer that turns usage into non-standard X-* response headers

So the compatibility break is at the HTTP adapter boundary, not in the TTS pipeline itself.

Option A: Lowest-Risk Compatibility Fix

• Keep /v1/audio/speech returning raw audio bytes only by default.
• Remove unconditional X-Prompt-Tokens, X-Completion-Tokens, and X-Engine-Time from the default response.
• Add an explicit non-standard opt-in request signal for benchmark/debug use, and only emit the X-* headers when that opt-in is present.
• Centralize speech-header emission behind one helper in sglang_omni/serve/openai_api.py.
• No changes needed for TTS pipeline, UsageInfo, or SpeechResult.

flowchart LR
    A["POST /v1/audio/speech"] --> B["build internal GenerateRequest"]
    B --> C["Client.speech()"]
    C --> D["SpeechResult.audio_bytes + SpeechResult.usage"]
    D --> E["speech response policy"]
    E --> F["default response: raw audio only"]
    E --> G["opt-in response: raw audio + X-* usage headers"]

Loading

Why Option A is the lowest-risk path

• the current compatibility break is localized to sglang_omni/serve/openai_api.py
• the benchmark already consumes headers, so opt-in migration is small
• no public speech request-ID contract is required
• no retention store is required
• no new endpoint is required

This is the easiest way to restore OpenAI-compatible default behavior without introducing a larger API redesign.

Testing for Option A

• default speech response returns audio bytes and does not include X-Prompt-Tokens
• default speech response does not include X-Completion-Tokens
• default speech response does not include X-Engine-Time
• opt-in speech response includes usage headers when SpeechResult.usage is present
• opt-in speech response still returns the same raw audio body and MIME type
• opt-in speech response emits no usage headers when usage is absent
• existing chat completion tests continue to pass unchanged
• benchmarks/benchmark_tts_speed.py sends the opt-in request signal and still computes token throughput when enabled

Option B: Sidecar Structured Usage API

If the goal is the permanent design, I think the cleaner architecture is:

• keep /v1/audio/speech strictly raw-audio
• expose per-request speech usage through a separate structured lookup

This matches the general shape used by systems that separate inference transport from request identity and observability. It is also closer in spirit to OpenAI, where speech generation stays audio-focused and usage is exposed through separate usage/reporting surfaces rather than embedded into the speech response contract.

flowchart LR
    A["POST /v1/audio/speech"] --> B["resolve request_id"]
    B --> C["build internal GenerateRequest"]
    C --> D["Client.speech()"]
    D --> E["SpeechResult.audio_bytes + SpeechResult.usage"]
    E --> F["store compact usage record by request_id"]
    E --> G["return raw audio response"]
    G --> H["X-Request-Id"]
    I["GET /v1/audio/speech/usage?request_id=..."] --> J["usage lookup"]
    F --> J
    J --> K["structured usage response"]

Loading

Sidecar design

1. Expose a stable speech request identifier.

• Add request_id: str | None = None to CreateSpeechRequest in sglang_omni/serve/protocol.py, matching chat.
• In create_speech() in sglang_omni/serve/openai_api.py, use req.request_id or f"speech-{uuid.uuid4()}".
• Return the final ID to the caller in X-Request-Id, so clients can retrieve usage even when the server generated the ID.

2. Add a dedicated structured usage endpoint.

• Recommended shape:
  • GET /v1/audio/speech/usage?request_id=<id>
• Recommended public response shape:
  • request_id
  • status
  • usage
• expires_at can be included if we want clients to know the lookup TTL.
• created_at does not need to be part of the public API.

3. Add a bounded lightweight usage store.

• No need to useRequestInfo.result in sglang_omni/pipeline/coordinator.py as the permanent backing store.
• Reason: completed request results are retained in memory today, and for speech the terminal result includes audio_data, sample_rate, modality, and usage.
• Reusing that structure for a sidecar would tie usage lookup to retained full audio payloads, which is the wrong memory shape for a permanent metrics API.
• Instead, store a slim usage record keyed by request_id.

Recommended internal record shape:

• request_id
• status
• usage
• created_at
• expires_at

created_at and expires_at are useful for internal retention and cleanup. expires_at is the only one that may be worth exposing publicly.

Why Option B is the better permanent design

• it keeps /v1/audio/speech strictly compatible by default
• it avoids turning custom response headers into a long-term observability contract
• it scales better if we later want richer metrics than prompt/completion tokens and engine time
• it cleanly separates transport, compatibility, and observability concerns

Why Option B is higher risk than Option A

Even though it is the better long-term design, it is a larger change:

• it requires a public speech request-ID contract
• it requires a new endpoint
• it requires a bounded retention store for usage records
• it requires lifecycle decisions for pending, completed, unknown, and expired requests
• it requires more API and state-management testing than the header-gating approach

Testing for Option B

• request with client-supplied request_id returns raw audio and the same request ID in X-Request-Id
• request without supplied request_id returns raw audio and a server-generated X-Request-Id
• GET /v1/audio/speech/usage?request_id=... returns structured usage for completed requests
• usage lookup for unknown request IDs returns not-found
• usage lookup for in-flight requests returns the chosen pending behavior
• usage lookup after TTL expiry returns not-found
• default speech response remains raw audio and OpenAI-compatible
• existing chat request-ID and usage behavior remains unchanged

What we can conclude from above is that we can choose:

• Option A if the goal is to fix things ASAP with the smallest safe change
• Option B if we want to use this issue to introduce the permanent structured design