From 068c4549258137873904058045bf796d74aae0bd Mon Sep 17 00:00:00 2001 From: Chuyue Wang Date: Sat, 29 Aug 2026 12:08:55 -0400 Subject: [PATCH] fix: replay persisted reasoning safely Signed-off-by: Chuyue Wang --- .../src/executor/engine.rs | 13 +- .../src/executor/rehydrate.rs | 199 +++++++++- .../tests/stateful_responses_integration.rs | 371 +++++++++++++++++- 3 files changed, 580 insertions(+), 3 deletions(-) diff --git a/crates/agentic-server-core/src/executor/engine.rs b/crates/agentic-server-core/src/executor/engine.rs index a17b676e..cb42899a 100644 --- a/crates/agentic-server-core/src/executor/engine.rs +++ b/crates/agentic-server-core/src/executor/engine.rs @@ -25,7 +25,7 @@ use crate::events::EventFrame; use crate::executor::error::ExecutorResult; use crate::executor::inference::DONE_MARKER; use crate::executor::persist::persist_if_needed; -use crate::executor::rehydrate::rehydrate_conversation; +use crate::executor::rehydrate::{prepare_reasoning_for_vllm, rehydrate_conversation, validate_reasoning_for_vllm}; use crate::executor::request::{ExecutionContext, RequestContext}; use crate::executor::upstream::{emit_deferred_stream_events, fetch_blocking_payload, fetch_stream_payload}; use crate::tool::{ToolRegistry, mcp}; @@ -116,6 +116,13 @@ async fn run_until_gateway_tools_complete( run_gateway_tool_loop(ctx, exec_ctx, auth, stream_upstream, stream).await } +fn prepare_initial_reasoning_for_vllm(input: &mut ResponsesInput, round: usize, compacted: bool) -> ExecutorResult<()> { + if round == 0 && !compacted { + return prepare_reasoning_for_vllm(input); + } + Ok(()) +} + async fn run_gateway_tool_loop( mut ctx: RequestContext, exec_ctx: &ExecutionContext, @@ -137,6 +144,7 @@ async fn run_gateway_tool_loop( for round in 0..MAX_GATEWAY_TOOL_ROUNDS { let compaction_usage = maybe_compact_context(&mut ctx, exec_ctx, auth).await?; + prepare_initial_reasoning_for_vllm(&mut ctx.enriched_request.input, round, compaction_usage.is_some())?; accumulate_usage(&mut combined_usage, compaction_usage); let output_offset = combined_output.len(); let (mut payload, deferred_stream_events): (ResponsePayload, Vec<_>) = if stream_upstream { @@ -563,6 +571,9 @@ impl ExecuteRequest { "executor received responses request" ); let ctx = rehydrate_conversation(self.payload, &self.exec_ctx).await?; + if !ctx.enriched_request.input.has_compaction_trigger() { + validate_reasoning_for_vllm(&ctx.enriched_request.input)?; + } if ctx.original_request.stream { Ok(Either::Right(run_stream(ctx, self.exec_ctx, self.client_auth))) } else { diff --git a/crates/agentic-server-core/src/executor/rehydrate.rs b/crates/agentic-server-core/src/executor/rehydrate.rs index 6b9fdf08..37ab41aa 100644 --- a/crates/agentic-server-core/src/executor/rehydrate.rs +++ b/crates/agentic-server-core/src/executor/rehydrate.rs @@ -6,10 +6,79 @@ use crate::executor::error::{ExecutorError, ExecutorResult}; use crate::executor::request::{ExecutionContext, RequestContext}; use crate::storage::InOutItem; -use crate::types::io::{InputItem, ResponsesInput, resolve_tool_choice, resolve_tools}; +use crate::types::io::{ + InputItem, ReasoningOutput, ReasoningTextContent, ResponsesInput, resolve_tool_choice, resolve_tools, +}; use crate::types::request_response::RequestPayload; use crate::utils::uuid7_str; +fn has_plaintext_reasoning(reasoning: &ReasoningOutput) -> bool { + reasoning.content.iter().any(|content| !content.text.is_empty()) +} + +fn has_opaque_reasoning_state(reasoning: &ReasoningOutput) -> bool { + reasoning + .encrypted_content + .as_ref() + .is_some_and(|encrypted| !encrypted.is_null()) +} + +/// Reject opaque reasoning that vLLM cannot replay before any normal inference call. +pub(super) fn validate_reasoning_for_vllm(input: &ResponsesInput) -> ExecutorResult<()> { + let ResponsesInput::Items(items) = input else { + return Ok(()); + }; + + if items.iter().any(|item| { + matches!(item, InputItem::Reasoning(reasoning) if has_opaque_reasoning_state(reasoning) && !has_plaintext_reasoning(reasoning)) + }) { + return Err(ExecutorError::InvalidRequest( + "reasoning item contains encrypted state without plaintext reasoning content and cannot be replayed to vLLM" + .to_owned(), + )); + } + + Ok(()) +} + +/// Prepare reasoning in the vLLM-bound request copy. +/// +/// vLLM can replay plaintext reasoning content but cannot interpret opaque provider state. Its generic Responses +/// conversion reads only the first reasoning content part and falls back to a summary when content is absent, while +/// its Harmony conversion joins all content parts. Normalize usable plaintext into one ordered, newline-delimited part +/// and remove summaries so both paths receive the same continuation state. Summary-only items have no usable vLLM +/// state and are omitted. [`RequestContext`] keeps the original request and new input items separately, so none of +/// these changes mutate persisted state. +pub(super) fn prepare_reasoning_for_vllm(input: &mut ResponsesInput) -> ExecutorResult<()> { + // Validate the complete input before mutation so an error never leaves a partially prepared request behind. + validate_reasoning_for_vllm(input)?; + + let ResponsesInput::Items(items) = input else { + return Ok(()); + }; + + items.retain_mut(|item| { + let InputItem::Reasoning(reasoning) = item else { + return true; + }; + if !has_plaintext_reasoning(reasoning) { + return false; + } + + let plaintext = reasoning + .content + .iter() + .map(|content| content.text.as_str()) + .collect::>() + .join("\n"); + reasoning.content = vec![ReasoningTextContent::new(plaintext)]; + reasoning.summary.clear(); + reasoning.encrypted_content = None; + true + }); + Ok(()) +} + /// Step 1 — Build [`RequestContext`] by rehydrating conversation history. /// /// `request` is moved into the context as `enriched_request`; one clone is taken @@ -127,6 +196,134 @@ mod tests { }; use crate::types::request_response::RequestPayload; + fn reasoning_item(content: &[&str], encrypted_content: Option) -> InputItem { + InputItem::Reasoning(ReasoningOutput { + id: "rs_prior".to_owned(), + content: content.iter().map(|text| ReasoningTextContent::new(*text)).collect(), + summary: vec![serde_json::json!({"type": "summary_text", "text": "public summary"})], + encrypted_content, + status: Some("completed".to_owned()), + }) + } + + #[test] + fn plaintext_reasoning_is_normalized_for_both_vllm_paths() { + let mut input = ResponsesInput::Items(vec![reasoning_item( + &["first continuation part", "second continuation part"], + Some(serde_json::json!({"ciphertext": "opaque-provider-state"})), + )]); + + prepare_reasoning_for_vllm(&mut input).expect("plaintext reasoning is replayable"); + + let ResponsesInput::Items(items) = input else { + panic!("expected structured input"); + }; + let InputItem::Reasoning(reasoning) = &items[0] else { + panic!("expected reasoning item"); + }; + assert_eq!(reasoning.id, "rs_prior"); + assert_eq!(reasoning.content.len(), 1); + assert_eq!( + reasoning.content[0].text, + "first continuation part\nsecond continuation part" + ); + assert!(reasoning.summary.is_empty()); + assert_eq!(reasoning.status.as_deref(), Some("completed")); + assert_eq!(reasoning.encrypted_content, None); + } + + #[test] + fn encrypted_reasoning_requires_nonempty_plaintext_content() { + for content in [Vec::new(), vec![""], vec!["", ""]] { + let mut input = ResponsesInput::Items(vec![reasoning_item( + &content, + Some(serde_json::json!("opaque-provider-state")), + )]); + + let error = + prepare_reasoning_for_vllm(&mut input).expect_err("encrypted-only reasoning must not reach vLLM"); + + assert_eq!(error.http_status(), http::StatusCode::BAD_REQUEST); + assert!( + error + .to_string() + .contains("encrypted state without plaintext reasoning content") + ); + assert!(!error.to_string().contains("opaque-provider-state")); + } + } + + #[test] + fn plaintext_reasoning_with_null_encrypted_state_is_normalized_without_summary() { + let mut item = reasoning_item(&["plaintext continuation"], Some(serde_json::Value::Null)); + let InputItem::Reasoning(reasoning) = &mut item else { + panic!("expected reasoning item"); + }; + reasoning.content[0].type_ = "unexpected_provider_type".to_owned(); + let mut input = ResponsesInput::Items(vec![item]); + + prepare_reasoning_for_vllm(&mut input).expect("null encrypted state is valid"); + + let ResponsesInput::Items(items) = input else { + panic!("expected structured input"); + }; + let InputItem::Reasoning(reasoning) = &items[0] else { + panic!("expected reasoning item"); + }; + assert_eq!(reasoning.content[0].type_, "reasoning_text"); + assert_eq!(reasoning.content[0].text, "plaintext continuation"); + assert!(reasoning.summary.is_empty()); + assert_eq!(reasoning.encrypted_content, None); + } + + #[test] + fn summary_only_reasoning_without_opaque_state_is_removed_from_vllm_copy() { + for encrypted_content in [None, Some(serde_json::Value::Null)] { + let mut input = ResponsesInput::Items(vec![reasoning_item(&[], encrypted_content)]); + + prepare_reasoning_for_vllm(&mut input).expect("summary-only reasoning has no usable vLLM state"); + + let ResponsesInput::Items(items) = input else { + panic!("expected structured input"); + }; + assert!( + items.is_empty(), + "a reasoning summary must never be promoted to reasoning text" + ); + } + } + + #[test] + fn validation_failure_does_not_partially_mutate_input() { + let valid = reasoning_item( + &["plaintext continuation"], + Some(serde_json::json!("first-opaque-state")), + ); + let invalid = reasoning_item(&[], Some(serde_json::json!("second-opaque-state"))); + let mut input = ResponsesInput::Items(vec![valid.clone(), invalid]); + + prepare_reasoning_for_vllm(&mut input).expect_err("the complete input must validate before normalization"); + + let ResponsesInput::Items(items) = input else { + panic!("expected structured input"); + }; + let (InputItem::Reasoning(actual), InputItem::Reasoning(expected)) = (&items[0], &valid) else { + panic!("expected reasoning items"); + }; + assert_eq!(actual.content[0].text, expected.content[0].text); + assert_eq!(actual.summary, expected.summary); + assert_eq!(actual.encrypted_content, expected.encrypted_content); + } + + #[test] + fn text_input_is_unchanged() { + let mut input = ResponsesInput::Text("plain user input".to_owned()); + + prepare_reasoning_for_vllm(&mut input).expect("text input contains no reasoning item"); + + assert!(matches!(input, ResponsesInput::Text(ref text) if text == "plain user input")); + } + fn request(conversation_id: Option<&str>, previous_response_id: Option<&str>) -> RequestPayload { RequestPayload { model: "test".into(), diff --git a/crates/agentic-server-core/tests/stateful_responses_integration.rs b/crates/agentic-server-core/tests/stateful_responses_integration.rs index 3d36cce5..849b1b33 100644 --- a/crates/agentic-server-core/tests/stateful_responses_integration.rs +++ b/crates/agentic-server-core/tests/stateful_responses_integration.rs @@ -7,12 +7,16 @@ mod support; use agentic_core::executor::execute; use agentic_core::executor::request::RequestContext; +use agentic_core::storage::InOutItem; use agentic_core::types::request_response::RequestPayload; use agentic_core::types::tools::{FunctionToolParam, NonEmptyToolName}; -use agentic_core::{FunctionToolResultMessage, InputItem, ResponsesInput, ResponsesTool, ToolChoice}; +use agentic_core::{ + FunctionToolResultMessage, InputItem, OutputItem, ReasoningOutput, ResponsesInput, ResponsesTool, ToolChoice, +}; use either::Either; use futures::StreamExt; use serde_json::Value; +use std::fmt::Write as _; use std::sync::Arc; use support::{ MockResponse, TestFixture, collect_stream, expected_text, load_cassette, make_request, output_text, @@ -441,6 +445,72 @@ async fn test_previous_response_id_rehydrates_function_call_before_tool_output() assert_eq!(input[2]["call_id"], "call_1"); } +#[tokio::test] +async fn test_previous_response_id_replays_plaintext_reasoning_without_opaque_state() { + assert_plaintext_reasoning_replay(false, false, true).await; +} + +#[tokio::test] +async fn test_streaming_previous_response_id_replays_plaintext_reasoning_without_opaque_state() { + assert_plaintext_reasoning_replay(true, false, true).await; +} + +#[tokio::test] +async fn test_conversation_replays_plaintext_reasoning_without_opaque_state() { + assert_plaintext_reasoning_replay(false, true, true).await; +} + +#[tokio::test] +async fn test_streaming_conversation_replays_plaintext_reasoning_with_null_state() { + assert_plaintext_reasoning_replay(true, true, false).await; +} + +#[tokio::test] +async fn test_summary_only_reasoning_is_not_replayed() { + for stream in [false, true] { + for conversation in [false, true] { + assert_summary_only_reasoning_not_replayed(stream, conversation).await; + } + } +} + +#[tokio::test] +async fn test_encrypted_only_persisted_reasoning_fails_before_upstream() { + let fixture = TestFixture::new_with_responses(vec![reasoning_response(false, &[], true)]).await; + let first = unwrap_blocking( + execute( + make_request("historical user", true, false, None, None), + Arc::clone(&fixture.exec_ctx), + ) + .await + .expect("persist encrypted-only reasoning"), + ); + + for stream in [false, true] { + let result = execute( + make_request("continue", true, stream, Some(first.id.clone()), None), + Arc::clone(&fixture.exec_ctx), + ) + .await; + let Err(error) = result else { + panic!("encrypted-only reasoning must be rejected before inference"); + }; + assert_eq!(error.http_status(), http::StatusCode::BAD_REQUEST); + assert!( + error + .to_string() + .contains("encrypted state without plaintext reasoning content") + ); + assert!(!error.to_string().contains("opaque-provider-state")); + } + + assert_eq!( + fixture.request_bodies().await.len(), + 1, + "invalid continuations must not call the upstream" + ); +} + #[tokio::test] async fn test_mcp_namespace_showcase_round_trip_rehydrates_calls_tools_and_outputs() { let tool_json = mcp_showcase_tools_json(); @@ -836,3 +906,302 @@ fn contains_key(value: &Value, key: &str) -> bool { Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => false, } } + +async fn assert_plaintext_reasoning_replay(stream: bool, conversation: bool, opaque_state: bool) { + let follow_up = if stream { + let cassette = load_cassette(&format!("{DIR}/resp-single-gpt-4o-streaming.yaml")); + MockResponse::from_turn(&cassette.turns[0]) + } else { + text_response("follow-up answer") + }; + let fixture = TestFixture::new_with_responses(vec![ + reasoning_response(stream, &["", "plaintext continuation"], opaque_state), + follow_up, + ]) + .await; + let conversation_id = conversation.then(|| "conv_reasoning_replay".to_owned()); + let first = run_response( + make_request("historical user", true, stream, None, conversation_id.clone()), + Arc::clone(&fixture.exec_ctx), + ) + .await; + let previous_response_id = (!conversation).then(|| first.id.clone()); + + let mut second = make_request( + "ignored", + true, + stream, + previous_response_id.clone(), + conversation_id.clone(), + ); + second.input = serde_json::from_value(serde_json::json!([ + { + "type": "function_call_output", + "call_id": "call_prior", + "output": "tool output" + }, + {"role": "user", "content": "new user input"} + ])) + .expect("valid follow-up input"); + let result = execute(second, Arc::clone(&fixture.exec_ctx)) + .await + .expect("execute continuation"); + if stream { + collect_stream(result).await; + } else { + unwrap_blocking(result); + } + + let requests = fixture.request_bodies().await; + assert_eq!(requests.len(), 2); + let input = requests[1]["input"].as_array().expect("rehydrated input array"); + let item_types = input + .iter() + .map(|item| item["type"].as_str().expect("typed input item")) + .collect::>(); + assert_eq!( + item_types, + [ + "message", + "reasoning", + "message", + "function_call", + "function_call_output", + "message" + ] + ); + + let reasoning = &input[1]; + assert_eq!(reasoning["id"], "rs_prior"); + assert_eq!(reasoning["content"].as_array().map(Vec::len), Some(1)); + assert_eq!(reasoning["content"][0]["text"], "\nplaintext continuation"); + assert_eq!(reasoning["summary"], serde_json::json!([])); + assert_eq!(reasoning["status"], "completed"); + assert!( + reasoning["encrypted_content"].is_null(), + "opaque provider state must not be forwarded to vLLM" + ); + assert!(!contains_key(&requests[1], "_agentic_item_kind")); + + let lookup = lookup_context(previous_response_id, conversation_id); + let history = if conversation { + fixture + .exec_ctx + .conv_handler + .rehydrate(&lookup) + .await + .expect("rehydrate conversation") + } else { + fixture + .exec_ctx + .resp_handler + .rehydrate(&lookup) + .await + .expect("rehydrate response") + }; + let stored = persisted_reasoning(&history); + assert_eq!(stored.content.len(), 2); + assert_eq!(stored.content[0].text, ""); + assert_eq!(stored.content[1].text, "plaintext continuation"); + assert_eq!(stored.summary[0]["text"], "public summary"); + let expected_state = opaque_state.then(|| serde_json::json!("opaque-provider-state")); + assert_eq!(stored.encrypted_content, expected_state); + assert_eq!(stored.status.as_deref(), Some("completed")); +} + +async fn assert_summary_only_reasoning_not_replayed(stream: bool, conversation: bool) { + let follow_up = if stream { + let cassette = load_cassette(&format!("{DIR}/resp-single-gpt-4o-streaming.yaml")); + MockResponse::from_turn(&cassette.turns[0]) + } else { + text_response("follow-up answer") + }; + let fixture = TestFixture::new_with_responses(vec![reasoning_response(stream, &[], false), follow_up]).await; + let conversation_id = conversation.then(|| format!("conv_summary_only_{stream}")); + let first = run_response( + make_request("historical user", true, stream, None, conversation_id.clone()), + Arc::clone(&fixture.exec_ctx), + ) + .await; + let previous_response_id = (!conversation).then(|| first.id.clone()); + + let mut second = make_request( + "ignored", + true, + stream, + previous_response_id.clone(), + conversation_id.clone(), + ); + second.input = serde_json::from_value(serde_json::json!([ + { + "type": "function_call_output", + "call_id": "call_prior", + "output": "tool output" + }, + {"role": "user", "content": "new user input"} + ])) + .expect("valid follow-up input"); + run_response(second, Arc::clone(&fixture.exec_ctx)).await; + + let requests = fixture.request_bodies().await; + assert_eq!(requests.len(), 2); + let input = requests[1]["input"].as_array().expect("rehydrated input array"); + let item_types = input + .iter() + .map(|item| item["type"].as_str().expect("typed input item")) + .collect::>(); + assert_eq!( + item_types, + ["message", "message", "function_call", "function_call_output", "message"] + ); + assert!( + input.iter().all(|item| item["type"] != "reasoning"), + "summary-only reasoning must not reach vLLM" + ); + assert!( + !requests[1].to_string().contains("public summary"), + "a reasoning summary must never be promoted into the vLLM-bound copy" + ); + + let lookup = lookup_context(previous_response_id, conversation_id); + let history = if conversation { + fixture + .exec_ctx + .conv_handler + .rehydrate(&lookup) + .await + .expect("rehydrate conversation") + } else { + fixture + .exec_ctx + .resp_handler + .rehydrate(&lookup) + .await + .expect("rehydrate response") + }; + let stored = persisted_reasoning(&history); + assert!(stored.content.is_empty()); + assert_eq!(stored.summary[0]["text"], "public summary"); + assert_eq!(stored.encrypted_content, None); + assert_eq!(stored.status.as_deref(), Some("completed")); +} + +async fn run_response( + request: RequestPayload, + exec_ctx: Arc, +) -> agentic_core::ResponsePayload { + let stream = request.stream; + let result = execute(request, exec_ctx).await.expect("execute response"); + if stream { + collect_stream(result).await + } else { + unwrap_blocking(result) + } +} + +fn lookup_context(previous_response_id: Option, conversation_id: Option) -> RequestContext { + let request = make_request("lookup", true, false, previous_response_id, conversation_id); + RequestContext { + enriched_request: request.clone(), + original_request: request, + new_input_items: Vec::new(), + response_id: "resp_lookup".to_owned(), + conversation_id: None, + conversation_version: None, + } +} + +fn persisted_reasoning(history: &[InOutItem]) -> &ReasoningOutput { + history + .iter() + .find_map(|item| match item { + InOutItem::Output(OutputItem::Reasoning(reasoning)) => Some(reasoning), + InOutItem::Input(_) | InOutItem::Output(_) => None, + }) + .expect("persisted reasoning item") +} + +fn reasoning_response(stream: bool, plaintext: &[&str], opaque_state: bool) -> MockResponse { + let content = plaintext + .iter() + .map(|text| serde_json::json!({"type": "reasoning_text", "text": text})) + .collect::>(); + let reasoning = serde_json::json!({ + "type": "reasoning", + "id": "rs_prior", + "content": content, + "summary": [{"type": "summary_text", "text": "public summary"}], + "encrypted_content": opaque_state.then_some("opaque-provider-state"), + "status": "completed" + }); + let message = serde_json::json!({ + "type": "message", + "id": "msg_prior", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "assistant context", "annotations": []}] + }); + let function_call = serde_json::json!({ + "type": "function_call", + "id": "fc_prior", + "call_id": "call_prior", + "name": "client_tool", + "arguments": "{}", + "status": "completed" + }); + let response = serde_json::json!({ + "id": "resp_reasoning", + "object": "response", + "created_at": 0, + "model": "test-model", + "status": "completed", + "output": [reasoning.clone(), message.clone(), function_call.clone()], + "usage": null, + "incomplete_details": null, + "error": null, + "previous_response_id": null, + "conversation_id": null, + "instructions": null + }); + if !stream { + return MockResponse::Json(response.to_string()); + } + + let events = vec![ + serde_json::json!({ + "type": "response.output_item.done", + "output_index": 0, + "item": reasoning + }), + serde_json::json!({ + "type": "response.output_item.added", + "output_index": 1, + "item": {"type": "message", "id": "msg_prior", "role": "assistant", "status": "in_progress", "content": []} + }), + serde_json::json!({ + "type": "response.output_text.delta", + "item_id": "msg_prior", + "output_index": 1, + "content_index": 0, + "delta": "assistant context" + }), + serde_json::json!({ + "type": "response.output_item.done", + "output_index": 1, + "item": message + }), + serde_json::json!({ + "type": "response.output_item.done", + "output_index": 2, + "item": function_call + }), + serde_json::json!({"type": "response.completed", "response": response}), + ]; + let mut body = String::new(); + for event in events { + let event_type = event["type"].as_str().expect("event type"); + writeln!(body, "event: {event_type}\ndata: {event}\n").expect("write SSE fixture"); + } + body.push_str("data: [DONE]\n\n"); + MockResponse::Sse(body) +}