Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a26bd94f13
|
||
|
|
0fe7aa38b0
|
||
|
|
d1ea8d28d3
|
||
|
|
b5e77148df
|
||
|
|
ed1f535260
|
||
|
|
115076294f
|
||
|
|
4cb383fbac
|
||
|
|
c5b759928b
|
@@ -289,6 +289,16 @@ pub mod events {
|
||||
pub const OUTPUT_TEXT_DONE: &str = "response.output_text.done";
|
||||
pub const CONTENT_PART_DONE: &str = "response.content_part.done";
|
||||
pub const OUTPUT_ITEM_DONE: &str = "response.output_item.done";
|
||||
/// Reasoning-item event family. A reasoning model's think block is
|
||||
/// its own output item, emitted ahead of the message item, with a
|
||||
/// summary part carrying the text. Streaming these is what keeps the
|
||||
/// connection observably alive while a model thinks — clients time
|
||||
/// out on silence, and SSE comment keep-alives don't count because
|
||||
/// idle timers reset on parsed events, not comments.
|
||||
pub const REASONING_SUMMARY_PART_ADDED: &str = "response.reasoning_summary_part.added";
|
||||
pub const REASONING_SUMMARY_TEXT_DELTA: &str = "response.reasoning_summary_text.delta";
|
||||
pub const REASONING_SUMMARY_TEXT_DONE: &str = "response.reasoning_summary_text.done";
|
||||
pub const REASONING_SUMMARY_PART_DONE: &str = "response.reasoning_summary_part.done";
|
||||
pub const FUNCTION_CALL_ARGUMENTS_DELTA: &str = "response.function_call_arguments.delta";
|
||||
pub const FUNCTION_CALL_ARGUMENTS_DONE: &str = "response.function_call_arguments.done";
|
||||
pub const COMPLETED: &str = "response.completed";
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
//! in Prometheus text format.
|
||||
|
||||
use anyhow::Result;
|
||||
use metrics_exporter_prometheus::PrometheusBuilder;
|
||||
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder};
|
||||
use std::net::SocketAddr;
|
||||
|
||||
/// Install the Prometheus metrics recorder and return a handle.
|
||||
@@ -12,8 +12,7 @@ use std::net::SocketAddr;
|
||||
pub fn install(listen: &str) -> Result<()> {
|
||||
let addr: SocketAddr = listen.parse()?;
|
||||
|
||||
PrometheusBuilder::new()
|
||||
.with_http_listener(addr)
|
||||
with_buckets(PrometheusBuilder::new().with_http_listener(addr))?
|
||||
.install()
|
||||
.map_err(|e| anyhow::anyhow!("failed to install Prometheus exporter: {e}"))?;
|
||||
|
||||
@@ -25,13 +24,77 @@ pub fn install(listen: &str) -> Result<()> {
|
||||
/// Install a recorder for testing (no HTTP listener). Returns a handle
|
||||
/// that can render the current metrics as Prometheus text.
|
||||
pub fn install_test_recorder() -> Result<metrics_exporter_prometheus::PrometheusHandle> {
|
||||
let handle = PrometheusBuilder::new()
|
||||
let handle = with_buckets(PrometheusBuilder::new())?
|
||||
.install_recorder()
|
||||
.map_err(|e| anyhow::anyhow!("failed to install test recorder: {e}"))?;
|
||||
describe_metrics();
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
/// Give every histogram explicit buckets, so it is exported as a
|
||||
/// Prometheus **histogram** rather than a summary.
|
||||
///
|
||||
/// Without this, `metrics-exporter-prometheus` renders every
|
||||
/// `histogram!` as a summary: `{quantile="0.95"}` series and no
|
||||
/// `_bucket` series at all. Three things break as a result, and the
|
||||
/// third is what made it visible.
|
||||
///
|
||||
/// 1. `histogram_quantile()` has nothing to read, so any dashboard
|
||||
/// panel written the idiomatic way returns *No data* forever. The
|
||||
/// fleet dashboard's TTFT panel did exactly that while cortex was
|
||||
/// serving thousands of requests.
|
||||
/// 2. Summary quantiles are computed per process over a rolling
|
||||
/// window, so they **cannot be aggregated**. Averaging p95 across
|
||||
/// two gateways is not p95 of anything.
|
||||
/// 3. They decay: with no samples in the window every quantile reads
|
||||
/// `0`, which is indistinguishable from "genuinely instant" on a
|
||||
/// graph.
|
||||
///
|
||||
/// Buckets are per-metric because the quantities differ by orders of
|
||||
/// magnitude — a request lasting minutes and a TTFT of milliseconds
|
||||
/// have no useful shared scale. Ranges are chosen from observed fleet
|
||||
/// behaviour rather than round numbers: decode runs for minutes on a
|
||||
/// long turn, prefill is seconds on a long prompt, and decode
|
||||
/// throughput sits in the tens of tokens/sec.
|
||||
fn with_buckets(builder: PrometheusBuilder) -> Result<PrometheusBuilder> {
|
||||
let seconds_short = &[
|
||||
0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0, 120.0,
|
||||
];
|
||||
let seconds_long = &[
|
||||
0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1200.0,
|
||||
];
|
||||
let tokens_per_second = &[
|
||||
1.0, 2.5, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 120.0, 200.0,
|
||||
];
|
||||
builder
|
||||
// Whole-request latency: a long agentic turn legitimately runs
|
||||
// for many minutes, so the tail has to reach there.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_request_duration_seconds".into()),
|
||||
seconds_long,
|
||||
)?
|
||||
// Prefill: sub-second for a short prompt, seconds for a long
|
||||
// one. Anything past a minute is pathological and belongs in
|
||||
// the overflow bucket.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_time_to_first_token_seconds".into()),
|
||||
seconds_short,
|
||||
)?
|
||||
// Image generation is inherently slower than a text turn and
|
||||
// scales with resolution and step count.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_images_generation_seconds".into()),
|
||||
seconds_long,
|
||||
)?
|
||||
// Not a duration — decode throughput, tens of tokens/sec on
|
||||
// this fleet.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_tokens_per_second".into()),
|
||||
tokens_per_second,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("failed to configure histogram buckets: {e}"))
|
||||
}
|
||||
|
||||
fn describe_metrics() {
|
||||
metrics::describe_histogram!(
|
||||
"cortex_request_duration_seconds",
|
||||
@@ -130,3 +193,43 @@ fn describe_metrics() {
|
||||
"Live prefill throughput per neuron:model, tokens/sec EMA (#137)"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod bucket_tests {
|
||||
/// Histograms must export as Prometheus **histograms**, with
|
||||
/// `_bucket` series — not as summaries.
|
||||
///
|
||||
/// This asserts the exported *shape*, not the recording call,
|
||||
/// because the recording was never the problem: cortex measured
|
||||
/// TTFT correctly for months while the fleet dashboard's panel sat
|
||||
/// on "No data", since `histogram_quantile()` reads `_bucket`
|
||||
/// series and a summary has none. A metric that is collected but
|
||||
/// cannot be queried is indistinguishable from one that was never
|
||||
/// collected — from the graph, and from the operator's chair.
|
||||
#[test]
|
||||
fn histograms_export_buckets_not_summary_quantiles() {
|
||||
let handle = match super::install_test_recorder() {
|
||||
Ok(h) => h,
|
||||
// Another test in this binary owns the global recorder;
|
||||
// it installs the same buckets, so skipping is honest.
|
||||
Err(_) => return,
|
||||
};
|
||||
metrics::histogram!("cortex_time_to_first_token_seconds", "node" => "n", "model" => "m")
|
||||
.record(0.42);
|
||||
let rendered = handle.render();
|
||||
|
||||
assert!(
|
||||
rendered.contains("cortex_time_to_first_token_seconds_bucket"),
|
||||
"no _bucket series — histogram_quantile() cannot read this:\n{rendered}"
|
||||
);
|
||||
assert!(
|
||||
rendered.contains("# TYPE cortex_time_to_first_token_seconds histogram"),
|
||||
"exported as the wrong metric type:\n{rendered}"
|
||||
);
|
||||
assert!(
|
||||
!rendered
|
||||
.contains(r#"cortex_time_to_first_token_seconds{node="n",model="m",quantile="#),
|
||||
"still exporting summary quantiles, which cannot be aggregated across gateways"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,8 +98,29 @@ impl CortexState {
|
||||
neuron_configs: config.neurons.clone(),
|
||||
eviction: config.eviction.clone(),
|
||||
catalogue,
|
||||
// READ timeout, not a total one. reqwest's `.timeout()` is a
|
||||
// total deadline — "from when the request starts connecting
|
||||
// until the response body has finished" — which for a streamed
|
||||
// completion is a hard cap on how long a model may generate,
|
||||
// not a liveness check. At 300s and ~13 tok/s that capped every
|
||||
// stream at roughly 3.9k tokens: cortex severed the connection
|
||||
// mid-answer, the client reported the stream as terminated, and
|
||||
// neuron carried on generating into a socket nobody was reading.
|
||||
// A long agentic turn could therefore never complete through the
|
||||
// gateway however healthy the fleet was, and the failure looked
|
||||
// like a model fault because the model was blameless.
|
||||
//
|
||||
// `.read_timeout()` resets on every successful read, so it times
|
||||
// out a stalled upstream while letting a producing one run as
|
||||
// long as it needs. The same 300s is kept deliberately: it is
|
||||
// the headroom a cold model load needs before its first byte,
|
||||
// which is what the original total timeout was reaching for.
|
||||
//
|
||||
// This pairs with neuron's 1s SSE keep-alive: a healthy stream
|
||||
// puts bytes on the wire every second, so idle here means idle,
|
||||
// and the deadline finally measures what it was meant to.
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(300))
|
||||
.read_timeout(std::time::Duration::from_secs(300))
|
||||
.build()
|
||||
.expect("failed to build HTTP client"),
|
||||
entitlements,
|
||||
|
||||
@@ -139,11 +139,35 @@ pub fn project_chat_stream_with(
|
||||
}
|
||||
InferenceEvent::ReasoningDelta(text) => {
|
||||
if !config.include_thinking {
|
||||
// Default path — reasoning has no slot in
|
||||
// chat completions, so it's dropped. Naïve
|
||||
// clients (Zed commit-message generator,
|
||||
// any vanilla OpenAI client) get clean
|
||||
// output.
|
||||
// Default path. This used to `continue` —
|
||||
// dropping the token and emitting nothing at
|
||||
// all — to keep model-internal scratchpad out
|
||||
// of naïve clients' UIs. That goal was right;
|
||||
// the method silenced the stream.
|
||||
//
|
||||
// A reasoning model can think for minutes on a
|
||||
// hard prompt, and for that whole span not one
|
||||
// SSE event went out. Clients time that out: the
|
||||
// DeepSeek Harness reports `pi-ai stream idle
|
||||
// timeout after 300000ms` and retries, so a task
|
||||
// that reasons for over five minutes could never
|
||||
// complete no matter how healthy the model was.
|
||||
// SSE comment keep-alives do not help — idle
|
||||
// timers reset on events, not comments.
|
||||
//
|
||||
// Emitting `reasoning_content` instead keeps the
|
||||
// original guarantee (it is not `content`, so
|
||||
// unaware clients still show clean output) while
|
||||
// the stream stays visibly alive, and hands
|
||||
// reasoning-aware clients the field they are
|
||||
// already looking for.
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let chunk = reasoning_chunk(&id, created, &model_id, &text);
|
||||
if tx.send(chunk).await.is_err() {
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
let Some(markers) = config.reasoning_markers.as_ref() else {
|
||||
@@ -246,6 +270,35 @@ fn role_chunk(id: &str, created: u64, model_id: &str) -> ChatCompletionChunk {
|
||||
}
|
||||
}
|
||||
|
||||
/// A reasoning delta, in the field reasoning-aware clients read.
|
||||
///
|
||||
/// `reasoning_content` is not in the OpenAI spec, but it is the
|
||||
/// de-facto slot: DeepSeek's API emits it, and so do vLLM and SGLang
|
||||
/// for reasoning models. Our own consumers already expect it —
|
||||
/// `helexa-acp`'s OpenAI provider reads `choice.delta.reasoning_content`,
|
||||
/// and `helexa-bench` times a stream's first reasoning token from it.
|
||||
///
|
||||
/// Crucially it is NOT `content`, so a client that has never heard of
|
||||
/// reasoning ignores the unknown field and sees clean output — the
|
||||
/// property the old drop-it-entirely default was protecting, kept
|
||||
/// without the silence that came with it.
|
||||
fn reasoning_chunk(id: &str, created: u64, model_id: &str, text: &str) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_id.into(),
|
||||
choices: vec![ChunkChoice {
|
||||
index: 0,
|
||||
delta: json!({ "reasoning_content": text }),
|
||||
finish_reason: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}],
|
||||
usage: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
fn content_chunk(id: &str, created: u64, model_id: &str, text: &str) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
@@ -500,10 +553,24 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_delta_is_dropped_in_chat_projection() {
|
||||
/// Reasoning is streamed as `reasoning_content`, never as `content`.
|
||||
///
|
||||
/// Both halves matter. Keeping it out of `content` is what stops
|
||||
/// scratchpad material appearing in the UI of a client that has
|
||||
/// never heard of reasoning — the guarantee the previous
|
||||
/// drop-everything default provided.
|
||||
///
|
||||
/// But emitting *something* is what keeps the connection alive. A
|
||||
/// model can think for minutes; while it did, this projection sent
|
||||
/// no events at all, and clients closed the stream as dead. The
|
||||
/// DeepSeek Harness reports `pi-ai stream idle timeout after
|
||||
/// 300000ms`, so any task whose reasoning outlasts the client's
|
||||
/// idle timeout could never complete. SSE comment keep-alives do
|
||||
/// not help, because idle timers reset on events, not comments.
|
||||
async fn reasoning_delta_is_streamed_as_reasoning_content() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::ReasoningDelta("<think>".into()))
|
||||
tx.send(InferenceEvent::ReasoningDelta("pondering".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("real".into()))
|
||||
@@ -511,8 +578,27 @@ mod tests {
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].choices[0].delta["content"], "real");
|
||||
assert_eq!(out.len(), 2, "reasoning must produce a chunk, not silence");
|
||||
assert_eq!(out[0].choices[0].delta["reasoning_content"], "pondering");
|
||||
assert!(
|
||||
out[0].choices[0].delta.get("content").is_none(),
|
||||
"reasoning must never arrive as content"
|
||||
);
|
||||
assert_eq!(out[1].choices[0].delta["content"], "real");
|
||||
}
|
||||
|
||||
/// An empty reasoning delta is still nothing to say — it must not
|
||||
/// manufacture a chunk, matching the text path.
|
||||
#[tokio::test]
|
||||
async fn empty_reasoning_delta_is_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::ReasoningDelta(String::new()))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
assert!(out.is_empty(), "empty reasoning deltas produce no chunks");
|
||||
}
|
||||
|
||||
fn pair() -> ReasoningTokenPair {
|
||||
@@ -667,10 +753,12 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn include_thinking_off_drops_reasoning_even_with_markers() {
|
||||
// Default behaviour even when markers happen to be
|
||||
// configured. The flag is the gate, not the marker
|
||||
// presence.
|
||||
/// With the flag off, reasoning stays out of `content` whether or
|
||||
/// not markers are configured — the flag is the gate, not the
|
||||
/// marker presence. It now travels as `reasoning_content` instead
|
||||
/// of being discarded, so this asserts the containment guarantee
|
||||
/// only: nothing model-internal reaches a naive client's `content`.
|
||||
async fn include_thinking_off_keeps_reasoning_out_of_content() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream_with(
|
||||
rx,
|
||||
|
||||
@@ -365,6 +365,12 @@ async fn run_projection(
|
||||
// final `response.completed` payload so clients that only read
|
||||
// the terminal response still see the calls.
|
||||
let mut tool_items: Vec<Value> = Vec::new();
|
||||
// Reasoning item state (#267) and whether the message item's
|
||||
// opening frames have gone out. The message is opened lazily,
|
||||
// because its `output_index` depends on whether a reasoning item
|
||||
// claimed index 0 first.
|
||||
let mut reasoning = ReasoningTracker::default();
|
||||
let mut message_open = false;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
@@ -378,12 +384,22 @@ async fn run_projection(
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Thinking is over the moment visible text starts.
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
if !message_open {
|
||||
if !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
message_open = true;
|
||||
}
|
||||
accumulated.push_str(&text);
|
||||
let frame = ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_TEXT_DELTA,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": reasoning.message_index(),
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}),
|
||||
@@ -392,10 +408,45 @@ async fn run_projection(
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::ReasoningDelta(_) => {
|
||||
// No representation in our Responses model yet.
|
||||
// Stage where it'd land: a `response.reasoning_*`
|
||||
// event family alongside `response.output_text.*`.
|
||||
InferenceEvent::ReasoningDelta(text) => {
|
||||
// Stream the think block as its own output item (#267).
|
||||
// This used to be dropped on the floor, which meant a
|
||||
// model that reasoned for minutes emitted no events at
|
||||
// all and clients closed the stream as dead — dsh
|
||||
// reports `pi-ai stream idle timeout after 300000ms`.
|
||||
// Its watchdog resets on any parsed event, so surfacing
|
||||
// reasoning is what keeps a thinking model alive on the
|
||||
// wire; SSE comment keep-alives cannot, because idle
|
||||
// timers count events, not comments.
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if !reasoning.open {
|
||||
if reasoning.item.is_some() {
|
||||
// Reasoning resumed after visible text. Our
|
||||
// model has one reasoning item per response, so
|
||||
// fold it into the existing summary rather than
|
||||
// opening a second item out of order.
|
||||
continue;
|
||||
}
|
||||
if !emit_reasoning_open_frames(&tx, &meta).await {
|
||||
return;
|
||||
}
|
||||
reasoning.open = true;
|
||||
}
|
||||
reasoning.text.push_str(&text);
|
||||
let frame = ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
data: json!({
|
||||
"item_id": reasoning_item_id(&meta),
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"delta": text,
|
||||
}),
|
||||
};
|
||||
if tx.send(frame).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::ToolCall {
|
||||
index,
|
||||
@@ -403,9 +454,19 @@ async fn run_projection(
|
||||
name,
|
||||
arguments,
|
||||
} => {
|
||||
// The message item opened by the start frames holds
|
||||
// output_index 0; function_call items follow it.
|
||||
let output_index = 1 + index as u64;
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
if !message_open {
|
||||
if !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
message_open = true;
|
||||
}
|
||||
// function_call items follow the message item, which
|
||||
// itself follows any reasoning item — so the base is the
|
||||
// message's index, not a hardcoded 0.
|
||||
let output_index = reasoning.message_index() + 1 + index as u64;
|
||||
let item_id = format!(
|
||||
"fc_{}_{index}",
|
||||
meta.response_id.trim_start_matches("resp_")
|
||||
@@ -508,18 +569,158 @@ async fn run_projection(
|
||||
return;
|
||||
}
|
||||
|
||||
// A stream can end while still thinking — a reasoning model that
|
||||
// exhausts max_output_tokens mid-block never emits visible text.
|
||||
// Close the item so the client sees a finished response rather than
|
||||
// a reasoning item left open forever.
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
// The message item is opened lazily, so a response that produced no
|
||||
// text (tool-call-only turns, or the case above) has not announced
|
||||
// it yet. The finish frames reference it, so it must exist first.
|
||||
if !message_open && !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
|
||||
let reason = finish.unwrap_or(FinishReason::Stop);
|
||||
let _ = emit_finish_frames(
|
||||
&tx,
|
||||
&meta,
|
||||
&accumulated,
|
||||
reason,
|
||||
usage.as_ref(),
|
||||
&tool_items,
|
||||
FinishContext {
|
||||
full_text: &accumulated,
|
||||
reason,
|
||||
usage: usage.as_ref(),
|
||||
tool_items: &tool_items,
|
||||
message_index: reasoning.message_index(),
|
||||
reasoning_item: reasoning.item.as_ref(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Streaming state for a reasoning model's think block (#267).
|
||||
///
|
||||
/// The block is its own output item, ahead of the message, so the
|
||||
/// message's `output_index` depends on whether reasoning happened —
|
||||
/// which is only known once the first token arrives. This tracks that
|
||||
/// decision and the accumulated summary text.
|
||||
#[derive(Default)]
|
||||
struct ReasoningTracker {
|
||||
/// Frames opening the reasoning item have been emitted.
|
||||
open: bool,
|
||||
/// Everything the model has thought so far, for the `*.done` frames
|
||||
/// and the terminal `completed` payload.
|
||||
text: String,
|
||||
/// The finished reasoning item, replayed inside `response.completed`
|
||||
/// so a client that only reads the terminal response still sees it.
|
||||
item: Option<Value>,
|
||||
}
|
||||
|
||||
impl ReasoningTracker {
|
||||
/// `output_index` the message item takes: 1 when a reasoning item
|
||||
/// occupies 0, else 0.
|
||||
fn message_index(&self) -> u64 {
|
||||
u64::from(self.open || self.item.is_some())
|
||||
}
|
||||
}
|
||||
|
||||
fn reasoning_item_id(meta: &ResponseMeta) -> String {
|
||||
format!("rs_{}", meta.response_id.trim_start_matches("resp_"))
|
||||
}
|
||||
|
||||
/// Open the reasoning item and its summary part, once.
|
||||
async fn emit_reasoning_open_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
) -> bool {
|
||||
let item_id = reasoning_item_id(meta);
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_ADDED,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": {
|
||||
"type": "reasoning",
|
||||
"id": item_id,
|
||||
"summary": [],
|
||||
"status": "in_progress",
|
||||
},
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_PART_ADDED,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"part": { "type": "summary_text", "text": "" },
|
||||
}),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Close the reasoning item, if one was opened, and record it for the
|
||||
/// terminal payload. Idempotent — every path that emits a non-reasoning
|
||||
/// item calls this first, and only the first call does anything.
|
||||
async fn close_reasoning(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
reasoning: &mut ReasoningTracker,
|
||||
) -> bool {
|
||||
if !reasoning.open {
|
||||
return true;
|
||||
}
|
||||
reasoning.open = false;
|
||||
let item_id = reasoning_item_id(meta);
|
||||
let full_item = json!({
|
||||
"type": "reasoning",
|
||||
"id": item_id,
|
||||
"summary": [{ "type": "summary_text", "text": reasoning.text }],
|
||||
"status": "completed",
|
||||
});
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_TEXT_DONE,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"text": reasoning.text,
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_PART_DONE,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"part": { "type": "summary_text", "text": reasoning.text },
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_DONE,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": full_item.clone(),
|
||||
}),
|
||||
},
|
||||
];
|
||||
reasoning.item = Some(full_item);
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &ResponseMeta) -> bool {
|
||||
let shell = response_shell(meta, "in_progress", &[], None);
|
||||
let frames = [
|
||||
@@ -540,10 +741,31 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
event_name: events::IN_PROGRESS,
|
||||
data: json!({ "response": shell }),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Open the message item, once, at `output_index`.
|
||||
///
|
||||
/// Deferred rather than emitted with the start frames because the
|
||||
/// message is no longer guaranteed to be `output[0]`: a reasoning model
|
||||
/// puts its think block in an item ahead of it, and an item's index has
|
||||
/// to be right the first time it is announced.
|
||||
async fn emit_message_open_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
output_index: u64,
|
||||
) -> bool {
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_ADDED,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"output_index": output_index,
|
||||
"item": empty_message_item(&meta.message_item_id),
|
||||
}),
|
||||
},
|
||||
@@ -551,7 +773,7 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
event_name: events::CONTENT_PART_ADDED,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": output_index,
|
||||
"content_index": 0,
|
||||
"part": { "type": "output_text", "text": "", "annotations": [] },
|
||||
}),
|
||||
@@ -565,14 +787,34 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
true
|
||||
}
|
||||
|
||||
/// Everything the terminal frames need that isn't stream identity.
|
||||
/// Grouped rather than passed loose: the reasoning item (#267) pushed
|
||||
/// the argument list past the point where positional `&str`/`u64`
|
||||
/// parameters are safe to read at a call site.
|
||||
struct FinishContext<'a> {
|
||||
full_text: &'a str,
|
||||
reason: FinishReason,
|
||||
usage: Option<&'a ResponsesUsage>,
|
||||
tool_items: &'a [Value],
|
||||
/// Where the message item sits — 1 when a reasoning item took 0.
|
||||
message_index: u64,
|
||||
/// The completed reasoning item, replayed inside `completed`.
|
||||
reasoning_item: Option<&'a Value>,
|
||||
}
|
||||
|
||||
async fn emit_finish_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
full_text: &str,
|
||||
reason: FinishReason,
|
||||
usage: Option<&ResponsesUsage>,
|
||||
tool_items: &[Value],
|
||||
ctx: FinishContext<'_>,
|
||||
) -> bool {
|
||||
let FinishContext {
|
||||
full_text,
|
||||
reason,
|
||||
usage,
|
||||
tool_items,
|
||||
message_index,
|
||||
reasoning_item,
|
||||
} = ctx;
|
||||
let status = finish_to_status(reason);
|
||||
let full_part = json!({
|
||||
"type": "output_text",
|
||||
@@ -586,10 +828,14 @@ async fn emit_finish_frames(
|
||||
"content": [full_part.clone()],
|
||||
"status": status,
|
||||
});
|
||||
// Terminal output array: the message item (always present — its
|
||||
// open frames were emitted eagerly at stream start) followed by
|
||||
// any completed function_call items.
|
||||
let mut output_items = Vec::with_capacity(1 + tool_items.len());
|
||||
// Terminal output array, in wire order: the reasoning item (when
|
||||
// the model thought), then the message, then any completed
|
||||
// function_call items. Order matters — a client reconstructing the
|
||||
// response from `completed` alone should see what the stream showed.
|
||||
let mut output_items = Vec::with_capacity(2 + tool_items.len());
|
||||
if let Some(item) = reasoning_item {
|
||||
output_items.push(item.clone());
|
||||
}
|
||||
output_items.push(full_item.clone());
|
||||
output_items.extend(tool_items.iter().cloned());
|
||||
let frames = [
|
||||
@@ -597,7 +843,7 @@ async fn emit_finish_frames(
|
||||
event_name: events::OUTPUT_TEXT_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"content_index": 0,
|
||||
"text": full_text,
|
||||
}),
|
||||
@@ -606,7 +852,7 @@ async fn emit_finish_frames(
|
||||
event_name: events::CONTENT_PART_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"content_index": 0,
|
||||
"part": full_part,
|
||||
}),
|
||||
@@ -614,7 +860,7 @@ async fn emit_finish_frames(
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_DONE,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"item": full_item.clone(),
|
||||
}),
|
||||
},
|
||||
@@ -1095,6 +1341,170 @@ mod tests {
|
||||
out
|
||||
}
|
||||
|
||||
/// The regression #267 was filed for: a reasoning model must put
|
||||
/// events on the wire while it thinks.
|
||||
///
|
||||
/// These deltas used to be dropped, so a model that reasoned for
|
||||
/// minutes emitted nothing at all and clients closed the stream as
|
||||
/// dead — dsh reports `pi-ai stream idle timeout after 300000ms`
|
||||
/// and retries, so a task whose thinking outlasts the client's idle
|
||||
/// 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.
|
||||
#[tokio::test]
|
||||
async fn reasoning_projects_a_full_item_lifecycle_before_the_message() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("think".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("ing".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("answer".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 2,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert_eq!(
|
||||
names,
|
||||
vec![
|
||||
events::CREATED,
|
||||
events::IN_PROGRESS,
|
||||
events::OUTPUT_ITEM_ADDED,
|
||||
events::REASONING_SUMMARY_PART_ADDED,
|
||||
events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
events::REASONING_SUMMARY_TEXT_DONE,
|
||||
events::REASONING_SUMMARY_PART_DONE,
|
||||
events::OUTPUT_ITEM_DONE,
|
||||
events::OUTPUT_ITEM_ADDED,
|
||||
events::CONTENT_PART_ADDED,
|
||||
events::OUTPUT_TEXT_DELTA,
|
||||
events::OUTPUT_TEXT_DONE,
|
||||
events::CONTENT_PART_DONE,
|
||||
events::OUTPUT_ITEM_DONE,
|
||||
events::COMPLETED,
|
||||
]
|
||||
);
|
||||
|
||||
// The reasoning item owns output_index 0, so the message shifts
|
||||
// to 1 — an item's index has to be right when it is announced,
|
||||
// which is why the message is opened lazily.
|
||||
let reasoning_added = &frames[2];
|
||||
assert_eq!(reasoning_added.data["output_index"], 0);
|
||||
assert_eq!(reasoning_added.data["item"]["type"], "reasoning");
|
||||
let message_added = &frames[9];
|
||||
assert_eq!(message_added.data["output_index"], 1);
|
||||
assert_eq!(message_added.data["item"]["type"], "message");
|
||||
assert_eq!(frames[11].data["output_index"], 1, "text delta index");
|
||||
|
||||
// Accumulated thinking, not just the last fragment.
|
||||
assert_eq!(frames[6].data["text"], "thinking");
|
||||
|
||||
// The terminal payload replays both items, in wire order, so a
|
||||
// client reading only `completed` sees what the stream showed.
|
||||
let output = frames.last().unwrap().data["response"]["output"]
|
||||
.as_array()
|
||||
.expect("output array");
|
||||
assert_eq!(output.len(), 2);
|
||||
assert_eq!(output[0]["type"], "reasoning");
|
||||
assert_eq!(output[0]["summary"][0]["text"], "thinking");
|
||||
assert_eq!(output[1]["type"], "message");
|
||||
}
|
||||
|
||||
/// A model can exhaust its budget mid-thought and never produce
|
||||
/// visible text. The reasoning item must still be closed and the
|
||||
/// message still opened, or the client waits on a response that
|
||||
/// never resolves.
|
||||
#[tokio::test]
|
||||
async fn reasoning_only_stream_still_closes_cleanly() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("pondering".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Length,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert!(names.contains(&events::REASONING_SUMMARY_PART_DONE));
|
||||
assert!(
|
||||
names.contains(&events::CONTENT_PART_ADDED),
|
||||
"message item must still be opened so the finish frames refer to something"
|
||||
);
|
||||
assert_eq!(names.last(), Some(&events::COMPLETED));
|
||||
let last = frames.last().unwrap();
|
||||
assert_eq!(last.data["response"]["status"], "incomplete");
|
||||
}
|
||||
|
||||
/// Tool calls sit after the message, which sits after any reasoning
|
||||
/// item — so their index is relative to the message, not a
|
||||
/// hardcoded 0.
|
||||
#[tokio::test]
|
||||
async fn tool_call_index_shifts_past_a_reasoning_item() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("plan".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::ToolCall {
|
||||
index: 0,
|
||||
id: "call_1".into(),
|
||||
name: "bash".into(),
|
||||
arguments: "{}".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::ToolCalls,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let call_added = frames
|
||||
.iter()
|
||||
.find(|f| {
|
||||
f.event_name == events::OUTPUT_ITEM_ADDED
|
||||
&& f.data["item"]["type"] == "function_call"
|
||||
})
|
||||
.expect("function_call item announced");
|
||||
// reasoning=0, message=1, first call=2.
|
||||
assert_eq!(call_added.data["output_index"], 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_stream_emits_expected_event_sequence() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
|
||||
Reference in New Issue
Block a user