Compare commits

..
Author SHA1 Message Date
grenadeandClaude Opus 5 ea1e7f1997 fix(neuron): cover the CUDA-only streaming paths
CI / Classify changes (push) Successful in 14s
CI / Web (lint + typecheck + i18n + build) (push) Has been skipped
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m32s
CI / Test (push) Successful in 7m21s
CI / CUDA type-check (push) Successful in 19m6s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The CUDA type-check caught a third `InferenceEvent::Finish` construction
in `stream_inference_via_worker` (the single-GPU device-worker stream)
that the local CPU build cannot see, so `cached_tokens` was missing there
and the cuda feature failed to compile.

Same shape as the TP path: the reuse count is bound inside a match arm
and consumed after it, so it is hoisted to the function scope. Vision
requests clear the cache outright and keep 0, which is accurate rather
than merely convenient.

Also completes #268 on the two streaming paths that were still silent —
the single-GPU worker stream and the CPU fallback. "Both streaming paths"
in the previous commit undercounted: there are four, and only the batch
engine and the TP serialized path had been done.

Audited every `InferenceEvent::Finish` construction and pattern across
the crate rather than waiting for CI to enumerate them one run at a time;
the only one without the field is a test helper matching with `..`.

Refs #268, #269

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XGhUi88ANnSfvk86DtYc9g
2026-08-19 17:21:14 +03:00
grenadeandClaude Opus 5 e94633c3a3 feat(usage): report prefix-cache reuse as cached_tokens
CI / Classify changes (push) Successful in 13s
CI / Web (lint + typecheck + i18n + build) (push) Has been skipped
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m26s
CI / Test (push) Successful in 9m30s
CI / CUDA type-check (push) Failing after 21m0s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Prefix KV caching landed with #11 and works — beast logs `reused=2068`
on prefills — but the number never left the daemon. Every client
therefore reported a 0% cache hit rate on a cache that was saving real
prefill: dsh showed `Cache hit 0%` against a 42.1K-token context.

The stale comment said so out loud, on a field that stayed `None`:

    /// OpenAI-standard breakdown of `prompt_tokens`. Populated once
    /// prompt caching lands (#11); `None` until then.

Prompt caching landed; #11 is closed.

The reuse count was already computed at prefill and only logged, so this
is threading rather than measuring: `op_prefill` returns it beside the
logits (each backend arm yields the pair, so neither can be forgotten),
`Slot` and `RouterMsg::Finish` carry it, and `InferenceEvent::Finish`
delivers it to the projections.

Reported in each surface's own spelling:

- chat completions — `usage.prompt_tokens_details.cached_tokens`
- responses        — `usage.input_tokens_details.cached_tokens`

Both are sub-counts of the prompt/input total and are never added into
`total_tokens`; a cache hit must not inflate what a caller is billed.
Both are omitted entirely at zero, so a client sees byte-identical JSON
when nothing was reused and "absent" stays distinguishable from
"measured as none".

cortex exports `cortex_cached_prompt_tokens_total` beside the existing
`cortex_prompt_tokens_total`, so fleet-wide hit rate is a ratio of two
counters instead of a grep for `reused=` — which was per-host,
unaggregated, and gone on log rotation.

This matters beyond bookkeeping. An agentic harness uses cache-hit rate
to decide whether to compact or restructure history, so reporting a flat
0% actively misinforms it: it may rewrite the very prefix that was being
reused, destroying a hit it could not see.

`usage_chunk`'s arguments are grouped into `UsageCounts` rather than
suppressing the arity lint — four adjacent bare `u32`s at a call site is
how the wrong count ends up in the wrong field.

Closes #269. Makes #270 (reuse looks low on long agentic contexts)
measurable from a dashboard rather than a journal.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XGhUi88ANnSfvk86DtYc9g
2026-08-19 16:55:47 +03:00
grenadeandClaude Opus 5 7cc37ebf37 fix(neuron): a streamed turn now says how it ended
CI / CUDA type-check (push) Blocked by required conditions
CI / Classify changes (push) Successful in 13s
CI / Web (lint + typecheck + i18n + build) (push) Has been skipped
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m32s
CI / Test (push) Successful in 6m58s
CI / Build cortex SRPM (push) Has been cancelled
CI / Build neuron SRPM (push) Has been cancelled
CI / Publish cortex to COPR (push) Has been cancelled
CI / Publish neuron to COPR (push) Has been cancelled
CI / Bump version in source (push) Has been cancelled
A non-streaming request closes its trail with counts and a reason:

    INFO tp_chat: chat_completion: done prompt_tokens=15
         completion_tokens=2 finish_reason=stop total_ms=100779

A streamed one logged nothing at all. The journal stopped after prefill,
so a finished turn, a turn that exhausted its output cap, and a turn
abandoned by a disconnected client were indistinguishable afterwards —
on the path essentially every interactive client uses.

That silence has a measured cost. During #267 an absence of log lines was
twice read as an absence of work on a daemon that was healthy and
decoding the whole time, and one of those readings was published as a
root cause before being retracted. Later, with the fixes deployed, a turn
wrote two of three requested files and stopped; whether it hit `max_new`
could not be answered from the server at all, only guessed from file
timestamps.

Both streaming paths now emit the same field set as the non-streaming
one, plus two the older line lacks:

- `reasoning_tokens`, because a reasoning model can spend its entire
  budget thinking and return nothing — observed at exactly
  `reasoning_tokens == output_tokens == 8192` (#223). A turn that
  produced no answer should say where the tokens went.
- `client_gone`, because an abandoned request keeps running to
  completion holding its admission permit and the TP pool lock. That is
  the amplifier behind #267's self-reinforcing retry loop and it was
  entirely invisible.

The batch-engine path logs from `run_router`, which is already
instrumented with the request's span, so req_id and model come along
without threading new arguments. A router that falls out of its loop
without a Finish — engine dropped the slot on a fatal worker error or
shutdown — warns instead, so that ending is distinguishable from a clean
one rather than looking identical to it.

Pinned by a test that captures the subscriber output and asserts the line
and its counts, not the `Finish` struct: an event that reaches the client
while leaving no server-side record is precisely the state being fixed.
Verified it fails without the line.

Closes #268

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XGhUi88ANnSfvk86DtYc9g
2026-08-19 16:43:13 +03:00
grenade 98d63cc3fd Merge branch 'fix/prometheus-histogram-buckets'
build-prerelease / Resolve version stamps + change detection (push) Successful in 13s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build helexa-angels binary (push) Has been skipped
build-prerelease / Build helexa-router binary (push) Has been skipped
build-prerelease / Package helexa-angels RPM (push) Has been skipped
build-prerelease / Package helexa-router RPM (push) Has been skipped
build-prerelease / Build helexa-tools binary (push) Has been skipped
build-prerelease / Package helexa-tools RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Successful in 2m12s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m37s
build-prerelease / Package cortex RPM (push) Successful in 29s
build-prerelease / Test (push) Successful in 6m54s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 13s
2026-08-19 16:35:45 +03:00
9 changed files with 395 additions and 26 deletions
+6 -4
View File
@@ -149,8 +149,9 @@ pub struct Usage {
/// `completion_tokens`, never added into `total_tokens`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub completion_tokens_details: Option<CompletionTokensDetails>,
/// OpenAI-standard breakdown of `prompt_tokens`. Populated once
/// prompt caching lands (#11); `None` until then.
/// OpenAI-standard breakdown of `prompt_tokens`, carrying
/// `cached_tokens` (#269). Omitted when nothing was reused, so
/// "absent" stays distinguishable from "measured as zero".
#[serde(default, skip_serializing_if = "Option::is_none")]
pub prompt_tokens_details: Option<PromptTokensDetails>,
/// helexa extension (non-OpenAI): server-measured prefill/decode
@@ -182,8 +183,9 @@ pub struct CompletionTokensDetails {
/// Sub-counts of `Usage::prompt_tokens`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptTokensDetails {
/// Prompt tokens served from cache (cache-read rate). Populated
/// once prompt caching lands (#11).
/// Prompt tokens served from neuron's prefix KV cache (#11), so a
/// caller can see the saving rather than assume none (#269). A
/// sub-count of `prompt_tokens`, never added into `total_tokens`.
pub cached_tokens: u64,
}
+4 -4
View File
@@ -249,8 +249,8 @@ pub struct ResponsesUsage {
/// sub-count of `output_tokens`, never added into `total_tokens`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_tokens_details: Option<OutputTokensDetails>,
/// OpenAI-standard breakdown of `input_tokens`. Populated once
/// prompt caching lands (#11); `None` until then.
/// OpenAI-standard breakdown of `input_tokens`, carrying
/// `cached_tokens` (#269). Omitted when nothing was reused.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_tokens_details: Option<InputTokensDetails>,
}
@@ -265,8 +265,8 @@ pub struct OutputTokensDetails {
/// Sub-counts of `ResponsesUsage::input_tokens`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InputTokensDetails {
/// Input tokens served from cache (cache-read rate). Populated
/// once prompt caching lands (#11).
/// Input tokens served from neuron's prefix KV cache (#11), so a
/// caller can see the saving rather than assume none (#269).
pub cached_tokens: u64,
}
+4
View File
@@ -121,6 +121,10 @@ fn describe_metrics() {
"cortex_request_errors_total",
"Total number of failed proxy requests"
);
metrics::describe_counter!(
"cortex_cached_prompt_tokens_total",
"Prompt tokens served from neuron's prefix KV cache; over cortex_prompt_tokens_total this is the cache-hit rate (#269)"
);
metrics::describe_counter!("cortex_evictions_total", "Total number of model evictions");
metrics::describe_counter!(
"cortex_cold_starts_total",
+12
View File
@@ -171,6 +171,9 @@ impl IntoResponse for ProxyError {
// over the decode window (first→last chunk); falls back to the
// full request duration for single-chunk (non-streaming) bodies
// cortex_prompt_tokens_total / cortex_completion_tokens_total (counters)
// cortex_cached_prompt_tokens_total (counter) — prompt tokens served
// from neuron's prefix cache; over prompt_tokens_total this is the
// fleet's cache-hit rate (#269)
/// Cap on the retained body tail. The usage object rides on the final
/// chunk, so a generous tail is plenty; the cap bounds memory on huge
@@ -243,6 +246,15 @@ impl ChunkObserver for CortexMetrics {
if let Some(prompt) = prompt {
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
}
// Prefix-cache reuse (#269), read from the same usage object.
// Against `cortex_prompt_tokens_total` this gives fleet-wide
// cache-hit rate as a ratio, which was previously only
// obtainable by grepping `reused=` out of a neuron journal —
// per-host, unaggregated, and gone on log rotation.
if let Some(cached) = last_count_for(self.tail.as_str(), "cached_tokens") {
metrics::counter!("cortex_cached_prompt_tokens_total", &self.labels)
.increment(cached);
}
if let Some(completion) = completion.filter(|c| *c > 0) {
metrics::counter!("cortex_completion_tokens_total", &self.labels)
.increment(completion);
+48
View File
@@ -5160,6 +5160,10 @@ impl CandleHarness {
let mut prefill_ms_measured: u32 = 0;
let mut decode_start: Option<std::time::Instant> = None;
// Hoisted out of `'work` so the completion path below
// can report it (#269) — the reuse count is decided
// inside the block but consumed after it.
let mut reused_for_task: u32 = 0;
'work: {
// Prefix-cache decision (#11): vision requests
// clear as before; text requests restore the
@@ -5181,6 +5185,7 @@ impl CandleHarness {
}
}
};
reused_for_task = reused as u32;
let mut logits_processor = {
let sampling = if temperature <= 0.0 {
@@ -5599,12 +5604,29 @@ impl CandleHarness {
.prefill_rate
.record_decode(all_tokens.len(), d.elapsed());
}
// Close the journal trail (#268) — the serialized
// streaming path, like the engine one, previously
// ended without saying how. Mirrors the field set
// the non-streaming `chat_completion: done` emits.
tracing::info!(
prompt_tokens = prompt_len,
completion_tokens = all_tokens.len(),
reasoning_tokens = reasoning_token_count,
cached_tokens = reused_for_task,
finish_reason = ?finish_reason,
prefill_ms = prefill_ms_measured,
decode_ms = decode_start
.map(|d| d.elapsed().as_millis() as u32)
.unwrap_or(0),
"chat_completion (stream): done"
);
let _ = tx
.send(InferenceEvent::Finish {
reason: finish_reason,
prompt_tokens: prompt_len as u32,
completion_tokens: all_tokens.len() as u32,
reasoning_tokens: reasoning_token_count,
cached_tokens: reused_for_task,
timing: Some(FinishTiming {
prefill_ms: prefill_ms_measured,
decode_ms: decode_start
@@ -7416,6 +7438,9 @@ async fn stream_inference_via_worker(
// effective rate — and the throughput ceiling it feeds — rises.
let prefill_start = std::time::Instant::now();
let prefill_prompt_len = prompt_tokens.len();
// Hoisted out of the match so the Finish event below can report it
// (#269). Vision requests clear the cache outright, so they keep 0.
let mut cached_tokens: u32 = 0;
let logits_vec = match images {
Some((imgs, image_token_id)) => {
worker
@@ -7430,6 +7455,7 @@ async fn stream_inference_via_worker(
None => {
let reused =
restore_or_clear_via_worker(&worker, handle, prefix_cache, &prompt_tokens).await?;
cached_tokens = reused as u32;
// Two-stage prefill around the retokenization-stable
// snapshot boundary — see `run_inference_via_worker`.
let cut = if prefix_cache.is_some() {
@@ -7613,12 +7639,24 @@ async fn stream_inference_via_worker(
}
// Fold this request's decode throughput into the model tracker (#137).
prefill_rate.record_decode(all_tokens.len(), decode_start.elapsed());
// Close the journal trail (#268).
tracing::info!(
prompt_tokens = prompt_tokens.len(),
completion_tokens = all_tokens.len(),
reasoning_tokens = reasoning_token_count,
cached_tokens,
finish_reason = ?finish_reason,
prefill_ms = prefill_elapsed.as_millis(),
decode_ms = decode_start.elapsed().as_millis(),
"chat_completion (stream): done"
);
let _ = tx
.send(InferenceEvent::Finish {
reason: finish_reason,
prompt_tokens: prompt_tokens.len() as u32,
completion_tokens: all_tokens.len() as u32,
reasoning_tokens: reasoning_token_count,
cached_tokens,
timing: Some(FinishTiming {
prefill_ms: prefill_elapsed.as_millis() as u32,
decode_ms: decode_start.elapsed().as_millis() as u32,
@@ -7885,11 +7923,21 @@ fn run_inference_streaming(
if emitted_tool_call && finish_reason == FinishReason::Stop {
finish_reason = FinishReason::ToolCalls;
}
// Close the journal trail (#268).
tracing::info!(
prompt_tokens = prompt_tokens.len(),
completion_tokens = all_tokens.len(),
reasoning_tokens = reasoning_token_count,
cached_tokens = reused,
finish_reason = ?finish_reason,
"chat_completion (stream): done"
);
let _ = tx.blocking_send(InferenceEvent::Finish {
reason: finish_reason,
prompt_tokens: prompt_tokens.len() as u32,
completion_tokens: all_tokens.len() as u32,
reasoning_tokens: reasoning_token_count,
cached_tokens: reused as u32,
timing: Some(FinishTiming {
prefill_ms: prefill_elapsed.as_millis() as u32,
decode_ms: decode_start.elapsed().as_millis() as u32,
+148 -6
View File
@@ -297,14 +297,23 @@ async fn op_drop_snap(cfg: &EngineConfig, session: &mut ActiveSession, id: u64)
/// Prefill one joining request at B=1 through the backend's existing
/// chunked-prefill + prefix-cache paths. Returns the prefill logits.
/// Prefill, returning the logits **and** how many leading tokens the
/// prefix cache let us skip (#269).
///
/// The reuse count is computed here and was previously only logged, so
/// no caller could report it — every client saw a 0% cache hit rate on
/// a cache that was working. Returning it is what lets `usage` carry
/// `cached_tokens`.
async fn op_prefill(
cfg: &EngineConfig,
session: &mut ActiveSession,
prompt_tokens: &[u32],
) -> Result<Vec<f32>> {
) -> Result<(Vec<f32>, usize)> {
let prompt_len = prompt_tokens.len();
let prefill_start = std::time::Instant::now();
let logits = match (&cfg.backend, session) {
// Each arm yields its logits plus how many leading tokens the
// prefix cache served, so neither can be forgotten.
let (logits, cached_tokens) = match (&cfg.backend, session) {
(
BackendConfig::Single {
worker,
@@ -340,7 +349,7 @@ async fn op_prefill(
None => chunked_prefill_via_worker(worker, *handle, prompt_tokens, reused).await?,
};
prefill_rate.record(prompt_len, prefill_start.elapsed());
logits
(logits, reused)
}
#[cfg(feature = "cuda")]
(BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
@@ -384,12 +393,12 @@ async fn op_prefill(
}
};
tp.prefill_rate.record(prompt_len, prefill_start.elapsed());
logits
(logits, reused)
}
#[cfg(feature = "cuda")]
_ => anyhow::bail!("engine backend/session mismatch"),
};
Ok(logits)
Ok((logits, cached_tokens))
}
/// One queued request. Admission has already been passed — the permit
@@ -439,6 +448,8 @@ enum RouterMsg {
reason: FinishReason,
prompt_tokens: u32,
completion_tokens: u32,
/// Leading prompt tokens the prefix cache let us skip (#269).
cached_tokens: u32,
timing: FinishTiming,
},
}
@@ -462,6 +473,9 @@ struct Slot {
finished: Option<FinishReason>,
prefill_ms: u32,
prefill_tokens: u32,
/// Prompt tokens served from the prefix cache (#269), reported to
/// the caller as `cached_tokens` so a working cache is visible.
cached_tokens: u32,
decode_start: std::time::Instant,
_admit: AdmissionPermit,
}
@@ -627,6 +641,7 @@ async fn finish_slot(slot: &mut Slot, reason: FinishReason, rate: &PrefillRateEm
reason,
prompt_tokens: slot.prompt_len as u32,
completion_tokens: slot.generated.len() as u32,
cached_tokens: slot.cached_tokens,
timing: FinishTiming {
prefill_ms: slot.prefill_ms,
decode_ms: slot.decode_start.elapsed().as_millis() as u32,
@@ -788,7 +803,7 @@ async fn prefill_join(
let prompt_len = prompt_tokens.len();
let prefill_start = std::time::Instant::now();
let logits_vec = op_prefill(cfg, session, &prompt_tokens).await?;
let (logits_vec, cached_tokens) = op_prefill(cfg, session, &prompt_tokens).await?;
let prefill_elapsed = prefill_start.elapsed();
// First token from the prefill logits.
@@ -827,6 +842,7 @@ async fn prefill_join(
let mut slot = Slot {
prefix_len: prompt_len,
prompt_len,
cached_tokens: cached_tokens as u32,
generated,
next_token: first,
max_new,
@@ -890,6 +906,9 @@ async fn run_router(
let mut tool_call_idx: usize = 0;
let mut emitted_tool_call = false;
let mut consumer_alive = true;
// Whether a Finish was seen, so the fall-through below can tell a
// clean end from a dropped slot.
let mut finished = false;
while let Some(msg) = rx.recv().await {
match msg {
@@ -978,17 +997,43 @@ async fn run_router(
mut reason,
prompt_tokens,
completion_tokens,
cached_tokens,
timing,
} => {
if emitted_tool_call && reason == FinishReason::Stop {
reason = FinishReason::ToolCalls;
}
// Close the journal trail for this turn (#268). The
// non-streaming path has always logged this; streamed
// turns logged nothing at all, so whether one finished,
// hit its output cap, or was abandoned mid-flight was
// unanswerable afterwards — and an absence of lines was
// repeatedly misread as an absence of work.
//
// `run_router` is instrumented with the request's span,
// so req_id and model come along without being passed.
tracing::info!(
prompt_tokens,
completion_tokens,
reasoning_tokens = reasoning_token_count,
cached_tokens,
finish_reason = ?reason,
prefill_ms = timing.prefill_ms,
decode_ms = timing.decode_ms,
// The turn ran to completion with nobody reading it:
// it held its admission permit and the pool lock the
// whole time. Invisible until now.
client_gone = !consumer_alive,
"chat_completion (stream): done"
);
finished = true;
let _ = tx
.send(InferenceEvent::Finish {
reason,
prompt_tokens,
completion_tokens,
reasoning_tokens: reasoning_token_count,
cached_tokens,
timing: Some(timing),
})
.await;
@@ -996,6 +1041,17 @@ async fn run_router(
}
}
}
// Fell out of the loop without a Finish: the engine dropped this
// slot's channel (fatal worker error, shutdown). The client's stream
// ends either way, but silently — so say so rather than leaving the
// turn's last trace to be a prefill line (#268).
if !finished {
tracing::warn!(
reasoning_tokens = reasoning_token_count,
client_gone = !consumer_alive,
"chat_completion (stream): ended without a finish event — engine dropped the slot"
);
}
}
/// `tracing::Instrument` without importing the trait at every use
@@ -1013,6 +1069,92 @@ mod tests {
use crate::config::AdmissionConfig;
use crate::harness::admission::AdmissionController;
/// A streamed turn must close its own journal trail (#268).
///
/// Streaming used to end silently: no token counts, no
/// `finish_reason`, no duration. That made a finished turn, a turn
/// that hit its output cap, and a turn abandoned by a disconnected
/// client indistinguishable after the fact — and during the #267
/// investigation an absence of log lines was twice mistaken for an
/// absence of work on a daemon that was healthy and decoding.
///
/// Asserts the emitted line rather than the struct, because a
/// `Finish` event that reaches the client but leaves no server-side
/// record is exactly the state this fixes.
#[tokio::test]
async fn a_streamed_turn_logs_how_it_ended() {
use std::sync::{Arc, Mutex};
use tracing_subscriber::prelude::*;
#[derive(Clone, Default)]
struct Capture(Arc<Mutex<Vec<u8>>>);
impl std::io::Write for Capture {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0
.lock()
.expect("capture poisoned")
.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Capture {
type Writer = Self;
fn make_writer(&'a self) -> Self::Writer {
self.clone()
}
}
let capture = Capture::default();
let subscriber = tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer().with_writer(capture.clone()));
let _guard = tracing::subscriber::set_default(subscriber);
let (router_tx, router_rx) = mpsc::channel::<RouterMsg>(8);
let (tx, mut rx) = mpsc::channel::<InferenceEvent>(8);
let handle = tokio::spawn(run_router(
tiny_tokenizer(8),
None,
None,
ToolSchemas::default(),
false,
tx,
Arc::new(AtomicBool::new(false)),
router_rx,
));
router_tx
.send(RouterMsg::Finish {
reason: FinishReason::Length,
prompt_tokens: 7413,
completion_tokens: 24576,
cached_tokens: 2068,
timing: FinishTiming {
prefill_ms: 6000,
decode_ms: 900_000,
prefill_tokens: 7413,
},
})
.await
.expect("router accepts finish");
drop(router_tx);
handle.await.expect("router task joins");
while rx.recv().await.is_some() {}
let logged = String::from_utf8(capture.0.lock().expect("capture poisoned").clone())
.expect("utf8 log output");
assert!(
logged.contains("chat_completion (stream): done"),
"streamed turn left no completion line:\n{logged}"
);
// The counts are the point — "it ended" without them cannot
// distinguish a complete answer from an exhausted budget.
assert!(logged.contains("24576"), "no completion_tokens:\n{logged}");
assert!(logged.contains("Length"), "no finish_reason:\n{logged}");
}
/// A WordLevel tokenizer whose vocab covers the whole fixture
/// vocab (`w0`..`w511`), so every decoded token maps to a unique
/// word and the emitted text uniquely encodes the token sequence.
+7
View File
@@ -84,6 +84,13 @@ pub enum InferenceEvent {
/// `output_tokens_details.reasoning_tokens` (responses).
/// Zero for non-reasoning models.
reasoning_tokens: u32,
/// Leading prompt tokens served from the prefix cache (#269) — a
/// sub-count of `prompt_tokens`, never added to `total_tokens`.
/// Surfaced as `prompt_tokens_details.cached_tokens` (chat) /
/// `input_tokens_details.cached_tokens` (responses). The cache
/// has worked since #11; without this the saving was invisible
/// and every client reported a 0% hit rate.
cached_tokens: u32,
/// Server-measured prefill/decode timing for the request, or
/// `None` on paths that don't measure it (CPU fallback that
/// doesn't instrument, tests). Streaming projectors surface
+111 -8
View File
@@ -27,7 +27,8 @@
//! propagate without us writing any logic.
use cortex_core::openai::{
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, Usage,
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, PromptTokensDetails,
Usage,
};
use serde_json::json;
use tokio::sync::mpsc;
@@ -219,6 +220,7 @@ pub fn project_chat_stream_with(
prompt_tokens,
completion_tokens,
reasoning_tokens,
cached_tokens,
timing,
} => {
// The finish_reason chunk, then an OpenAI-style
@@ -233,9 +235,12 @@ pub fn project_chat_stream_with(
&id,
created,
&model_id,
prompt_tokens,
completion_tokens,
reasoning_tokens,
UsageCounts {
prompt_tokens,
completion_tokens,
reasoning_tokens,
cached_tokens,
},
timing,
),
]
@@ -385,15 +390,32 @@ fn final_chunk(
/// ignore the empty-choices chunk; clients that do (opencode, and
/// cortex's Anthropic translator) get the token counts they need to
/// track context.
/// The token counts a finished turn reports. Grouped rather than passed
/// loose: adding `cached_tokens` (#269) put the list past the point
/// where four adjacent bare `u32`s are safe to read at a call site.
struct UsageCounts {
prompt_tokens: u32,
completion_tokens: u32,
/// Sub-count of `completion_tokens` — tokens spent reasoning.
reasoning_tokens: u32,
/// Sub-count of `prompt_tokens` — tokens served from the prefix
/// cache. Neither sub-count is added into `total_tokens`.
cached_tokens: u32,
}
fn usage_chunk(
id: &str,
created: u64,
model_id: &str,
prompt_tokens: u32,
completion_tokens: u32,
reasoning_tokens: u32,
counts: UsageCounts,
timing: Option<FinishTiming>,
) -> ChatCompletionChunk {
let UsageCounts {
prompt_tokens,
completion_tokens,
reasoning_tokens,
cached_tokens,
} = counts;
ChatCompletionChunk {
id: id.into(),
object: "chat.completion.chunk".into(),
@@ -409,7 +431,12 @@ fn usage_chunk(
completion_tokens_details: (reasoning_tokens > 0).then_some(CompletionTokensDetails {
reasoning_tokens: reasoning_tokens as u64,
}),
prompt_tokens_details: None,
// Prefix-cache reuse (#269). Omitted at zero so a client
// sees unchanged JSON when nothing was cached, and so
// "absent" stays distinguishable from "measured as none".
prompt_tokens_details: (cached_tokens > 0).then_some(PromptTokensDetails {
cached_tokens: cached_tokens as u64,
}),
// helexa extension (#85): server-measured prefill/decode
// timing for the bench harness. Omitted on paths that don't
// measure it so standard clients see unchanged JSON.
@@ -458,6 +485,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -496,6 +524,7 @@ mod tests {
prompt_tokens: 128,
completion_tokens: 64,
reasoning_tokens: 0,
cached_tokens: 0,
timing: Some(FinishTiming {
prefill_ms: 200,
decode_ms: 1500,
@@ -520,6 +549,73 @@ mod tests {
assert_eq!(timing.prefill_tokens, 128);
}
/// A cache hit must reach the caller (#269).
///
/// Prefix caching has worked since #11 — neuron logs `reused=2068`
/// on prefills — but nothing carried the number onto the wire, so
/// every client reported a 0% hit rate on a cache that was saving
/// real prefill. dsh showed `Cache hit 0%` against a 42.1K-token
/// context, and an agentic harness that believes nothing is cached
/// may restructure the very prefix that was being reused.
#[tokio::test]
async fn cached_tokens_reach_the_usage_chunk() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
prompt_tokens: 7413,
completion_tokens: 12,
reasoning_tokens: 0,
cached_tokens: 2068,
timing: None,
})
.await
.unwrap();
drop(tx);
let out = collect(out_rx).await;
let usage = out
.iter()
.find_map(|c| c.usage.as_ref())
.expect("usage chunk present");
let details = usage
.prompt_tokens_details
.as_ref()
.expect("prompt_tokens_details populated when the cache was hit");
assert_eq!(details.cached_tokens, 2068);
// A sub-count of prompt_tokens, never added into the total —
// otherwise a cache hit would inflate what a caller is billed.
assert_eq!(usage.prompt_tokens, 7413);
assert_eq!(usage.total_tokens, 7425);
}
/// No reuse means no `prompt_tokens_details` at all, so a client
/// sees byte-identical JSON to before this landed — and "absent"
/// stays distinguishable from "measured as zero".
#[tokio::test]
async fn no_cache_hit_omits_the_details_object() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
prompt_tokens: 10,
completion_tokens: 1,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
.unwrap();
drop(tx);
let usage = collect(out_rx)
.await
.iter()
.find_map(|c| c.usage.clone())
.expect("usage chunk present");
assert!(usage.prompt_tokens_details.is_none());
}
#[tokio::test]
async fn empty_text_delta_is_dropped() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
@@ -541,6 +637,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -637,6 +734,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -689,6 +787,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -735,6 +834,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -781,6 +881,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -809,6 +910,7 @@ mod tests {
prompt_tokens: 42,
completion_tokens: 5,
reasoning_tokens: 2,
cached_tokens: 0,
timing: None,
})
.await
@@ -843,6 +945,7 @@ mod tests {
prompt_tokens: 10,
completion_tokens: 7,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
+55 -4
View File
@@ -33,9 +33,9 @@
use cortex_core::openai::{ChatCompletionRequest, ChatMessage, MessageContent};
use cortex_core::responses::{
OutputTokensDetails, ResponsesContentPart, ResponsesInput, ResponsesInputElement,
ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent, ResponsesOutputItem,
ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
InputTokensDetails, OutputTokensDetails, ResponsesContentPart, ResponsesInput,
ResponsesInputElement, ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent,
ResponsesOutputItem, ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
};
use serde_json::{Value, json};
use tokio::sync::mpsc;
@@ -540,6 +540,7 @@ async fn run_projection(
prompt_tokens,
completion_tokens,
reasoning_tokens,
cached_tokens,
// Responses-side `helexa_timing` surfacing not wired yet;
// the bench harness reads timing off the chat path (#85).
timing: _,
@@ -556,7 +557,12 @@ async fn run_projection(
output_tokens_details: (reasoning_tokens > 0).then_some(OutputTokensDetails {
reasoning_tokens: reasoning_tokens as u64,
}),
input_tokens_details: None,
// Prefix-cache reuse (#269). Omitted at zero so a
// client sees unchanged JSON when nothing was
// cached.
input_tokens_details: (cached_tokens > 0).then_some(InputTokensDetails {
cached_tokens: cached_tokens as u64,
}),
});
}
}
@@ -1351,6 +1357,41 @@ mod tests {
/// timeout could never finish. Its watchdog resets on any parsed
/// event, which is why emitting these fixes it and why SSE comment
/// keep-alives could not: idle timers count events, not comments.
/// The Responses surface reports cache reuse too (#269), in its own
/// spelling: `input_tokens_details.cached_tokens`.
///
/// This is the surface dsh uses, and the one that showed
/// `Cache hit 0%` against a 42.1K-token context while neuron was
/// reusing thousands of tokens per prefill.
#[tokio::test]
async fn cached_tokens_reach_the_completed_usage() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
let out = project_responses_stream(rx, meta());
tx.send(InferenceEvent::Start).await.unwrap();
tx.send(InferenceEvent::TextDelta("ok".into()))
.await
.unwrap();
tx.send(InferenceEvent::Finish {
reason: FinishReason::Stop,
prompt_tokens: 7413,
completion_tokens: 12,
reasoning_tokens: 0,
cached_tokens: 2068,
timing: None,
})
.await
.unwrap();
drop(tx);
let frames = collect(out).await;
let usage = &frames.last().unwrap().data["response"]["usage"];
assert_eq!(usage["input_tokens_details"]["cached_tokens"], 2068);
assert_eq!(usage["input_tokens"], 7413);
// Sub-count, not additive — a cache hit must not inflate the total.
assert_eq!(usage["total_tokens"], 7425);
}
#[tokio::test]
async fn reasoning_projects_a_full_item_lifecycle_before_the_message() {
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
@@ -1371,6 +1412,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 2,
cached_tokens: 0,
timing: None,
})
.await
@@ -1444,6 +1486,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 1,
cached_tokens: 0,
timing: None,
})
.await
@@ -1487,6 +1530,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 1,
cached_tokens: 0,
timing: None,
})
.await
@@ -1522,6 +1566,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -1583,6 +1628,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -1655,6 +1701,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -1684,6 +1731,7 @@ mod tests {
prompt_tokens: 30,
completion_tokens: 12,
reasoning_tokens: 4,
cached_tokens: 0,
timing: None,
})
.await
@@ -1715,6 +1763,7 @@ mod tests {
prompt_tokens: 8,
completion_tokens: 3,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -1740,6 +1789,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await
@@ -1787,6 +1837,7 @@ mod tests {
prompt_tokens: 0,
completion_tokens: 0,
reasoning_tokens: 0,
cached_tokens: 0,
timing: None,
})
.await