feat: token counts on span + cancellation WARN logging [DIS-1643]

[nnshah1]

Summary

• ResponseMetricCollector records request summary on the enclosing span in Drop:
  input_tokens, output_tokens, ttft_ms, avg_itl_ms, prefill_worker_id, decode_worker_id
• All fields inherited by InflightGuard "request completed" log — no per-token counting needed
• Stores already-computed TTFT and accumulates ITL from existing calculations
• WARN log at cancellation point in disconnect.rs with request context
• Connection monitor disconnects upgraded from TRACE to WARN

Depends on: #7734

Fields on "request completed" Log

Field	Source	Streaming	Unary
input_tokens	ResponseMetricCollector ISL	✅	⚠️ (DIS-1653)
output_tokens	ResponseMetricCollector OSL	✅	⚠️ (DIS-1653)
ttft_ms	Already computed for histogram	✅	⚠️
avg_itl_ms	Accumulated from per-chunk ITL	✅	N/A
prefill_worker_id	From LLMMetricAnnotation	✅ (when set)	✅ (when set)
decode_worker_id	From LLMMetricAnnotation	✅ (when set)	N/A

Example: Streaming Success (with all fields)

{
  "level": "INFO",
  "message": "request completed",
  "status": "success",
  "request_id": "32691d61-...",
  "model": "qwen/qwen3-0.6b",
  "endpoint": "chat_completions",
  "request_type": "stream",
  "elapsed_ms": "20",
  "input_tokens": "9",
  "output_tokens": "50",
  "ttft_ms": "5.85",
  "avg_itl_ms": "0.29",
  "trace_id": "bca97f5e..."
}

Example: Cancellation

{"level":"ERROR","message":"request completed","status":"error","error_type":"cancelled","error_detail":"client disconnected before completion","elapsed_ms":"230"}
{"level":"WARN","message":"Stream closed unexpectedly; issuing cancellation"}

Example: Worker Crash (partial tokens before crash)

{"level":"ERROR","message":"request completed","status":"error","error_type":"internal","error_detail":"internal server error during processing","elapsed_ms":"556","input_tokens":"9","output_tokens":"4"}

Performance

• ttft_ms: one Option<f64> assignment when TTFT is already computed (same code path)
• itl_sum_secs/itl_count: one f64 add + one u64 increment per chunk (negligible, same path as histogram publish)
• Span recording: 6 span.record() calls in Drop, once per request at cleanup
• No allocations in streaming hot path

Files Changed

File	Change
lib/llm/src/http/service/metrics.rs	ResponseMetricCollector: ttft_ms, itl accumulation, span recording in Drop
lib/llm/src/http/service/disconnect.rs	WARN cancel log, TRACE→WARN for disconnects
lib/runtime/src/logging.rs	Empty fields: ttft_ms, avg_itl_ms, prefill_worker_id, decode_worker_id

Follow-up

• DIS-1653 — Token counts/TTFT missing on unary requests (collector Drop ordering)
• DIS-1652 — Propagate model name to worker via transport headers

🤖 Generated with Claude Code

Linear: DIS-1643

Summary by CodeRabbit

• Chores
  • Enhanced observability with improved logging for connection and stream closure events, now including detailed contextual information.
  • Implemented performance metrics tracking for request-level timings, including time-to-first-token and inter-token latency measurements.
  • Extended distributed tracing with structured fields to support better monitoring and debugging capabilities.

[nnshah1]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Walkthrough

The changes enhance observability by upgrading disconnect logging from trace to warn level with structured context, adding time-to-first-token (TTFT) and inter-token latency (ITL) tracking to metrics collection, and recording request metrics to tracing spans for log correlation across inflight and response handling phases.

Changes

Cohort / File(s)	Summary
Disconnect Monitoring lib/llm/src/http/service/disconnect.rs	Upgraded log severity for unexpected connection and stream closures from trace! to warn!. Added structured logging fields (request_id, model, endpoint, request_type, error_type, elapsed_ms) when context stops, while preserving existing cancellation metrics and engine context termination.
Latency Metrics Tracking lib/llm/src/http/service/metrics.rs, lib/runtime/src/logging.rs	Enhanced ResponseMetricCollector to track per-request TTFT and ITL (inter-token latency sum and sample count). Computes TTFT on first token chunk and accumulates ITL across subsequent chunks. On drop, records ttft_ms, avg_itl_ms, prefill_worker_id, and decode_worker_id to tracing span. Added corresponding empty span fields to make_inference_request_span for field initialization.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the two main changes: token counts on span and cancellation WARN logging with the ticket reference.
Description check	✅ Passed	The description is comprehensive and well-structured, covering overview, detailed changes, file listings, examples, and follow-ups, exceeding the template requirements.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/llm/src/http/service/metrics.rs (1)

917-922: ⚠️ Potential issue | 🔴 Critical

Fix missing request_id argument in audio_speech handler.

The create_inflight_guard call at openai.rs:2306 in the audio_speech handler is missing the required request_id: String argument. The variable is available at that point (extracted and reassigned around lines 2278–2282), so update the call to:

.create_inflight_guard(&model, Endpoint::Audios, streaming, request_id.clone())

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/http/service/metrics.rs` around lines 917 - 922, The call to
create_inflight_guard in the audio_speech handler is missing the required
request_id argument; update the invocation in the audio_speech handler
(openai.rs) to pass the model by reference and the request id (use
request_id.clone()) and the Audios endpoint, e.g. call
create_inflight_guard(&model, Endpoint::Audios, streaming, request_id.clone())
so the signature pub fn create_inflight_guard(self: Arc<Self>, model: &str,
endpoint: Endpoint, streaming: bool, request_id: String) is satisfied.

🧹 Nitpick comments (1)

lib/llm/src/http/service/metrics.rs (1)

1382-1404: Span recording relies on collector being dropped while span is still active.

This implementation calls tracing::Span::current() in Drop, assuming the request span is still the current span. Based on the codebase architecture (TraceLayer applies the span at middleware layer, collector is created inside handlers and dropped when handler returns/stream ends), this assumption holds for normal request flows.

However, if the collector is somehow moved to a different async task or outlives the span scope, these recordings would go to the wrong span or be lost. Consider adding a brief code comment documenting this contract for future maintainers.

📝 Suggested documentation comment

     // Publish final OSL when the collector is dropped
     self.metrics
         .output_sequence_length
         .with_label_values(&[&self.model])
         .observe(self.osl as f64);

-    // Record request summary on the enclosing span.
-    // InflightGuard::Drop and on_response logs will inherit these.
+    // Record request summary on the enclosing span.
+    // IMPORTANT: This assumes the collector is dropped while the request span
+    // (created by TraceLayer middleware) is still active. The collector must not
+    // be moved to a separate task or outlive the handler's span scope.
+    // InflightGuard::Drop and on_response logs will inherit these fields.
     let span = tracing::Span::current();

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/http/service/metrics.rs` around lines 1382 - 1404, The Drop
implementation in InflightGuard (where tracing::Span::current() is used to
record fields like
input_tokens/output_tokens/ttft_ms/avg_itl_ms/prefill_worker_id/decode_worker_id)
assumes the request span is still the current span; add a clear comment above
this block (or above InflightGuard::drop) documenting the lifetime contract:
that the collector must be dropped while the request span remains active and
must not be moved to another async task or outlive the handler scope, and note
potential misrecording if that contract is violated and how callers should
ensure correct ownership/awaiting to preserve span activity.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@lib/llm/src/http/service/metrics.rs`:
- Around line 917-922: The call to create_inflight_guard in the audio_speech
handler is missing the required request_id argument; update the invocation in
the audio_speech handler (openai.rs) to pass the model by reference and the
request id (use request_id.clone()) and the Audios endpoint, e.g. call
create_inflight_guard(&model, Endpoint::Audios, streaming, request_id.clone())
so the signature pub fn create_inflight_guard(self: Arc<Self>, model: &str,
endpoint: Endpoint, streaming: bool, request_id: String) is satisfied.

---

Nitpick comments:
In `@lib/llm/src/http/service/metrics.rs`:
- Around line 1382-1404: The Drop implementation in InflightGuard (where
tracing::Span::current() is used to record fields like
input_tokens/output_tokens/ttft_ms/avg_itl_ms/prefill_worker_id/decode_worker_id)
assumes the request span is still the current span; add a clear comment above
this block (or above InflightGuard::drop) documenting the lifetime contract:
that the collector must be dropped while the request span remains active and
must not be moved to another async task or outlive the handler scope, and note
potential misrecording if that contract is violated and how callers should
ensure correct ownership/awaiting to preserve span activity.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 536cf9a9-f270-4ebd-8924-47f4b13363e5

📥 Commits

Reviewing files that changed from the base of the PR and between 72021e2 and ac29172.

📒 Files selected for processing (3)

• lib/llm/src/http/service/disconnect.rs
• lib/llm/src/http/service/metrics.rs
• lib/runtime/src/logging.rs

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-cdd2fbaf-856a-43c8-8694-11d7c54dfcd5.docs.buildwithfern.com/dynamo/dev