diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3c601d3e..8ae572cc 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -178,7 +178,7 @@ reads `response.create` messages off the socket and drives the *same* `stream_id` values run concurrently, while requests in the same lane remain FIFO; requests without a `stream_id` share a default FIFO lane. The session admits at most 64 active or queued requests and 12 MiB of aggregate request data. WebSocket sessions -always force `stream: true, store: true`. Because axum's built-in graceful shutdown +force `stream: true` and honor the requested `store` value. Because axum's built-in graceful shutdown doesn't wait for upgraded connections, `AppState` carries a separate `WebSocketTracker` so shutdown can drain in-flight sessions. @@ -197,6 +197,59 @@ Errors are modeled by a dedicated `WsError` enum (`handler/websocket/error.rs`) than reusing the HTTP JSON-error path, since some failure modes (a dead socket) must not attempt to write a response. +### Opt-in core continuation sessions (`executor/session.rs`) + +Core callers can use `ExecuteRequest::with_session` or `rehydrate_in_session` to +retain response state without durable storage. The WebSocket multiplexer owns one +`ResponseSessionGroup` per connection and keeps one session per lane, including idle +lanes. No-session HTTP and split execution keep their existing behavior. + +The connection retains at most 128 lanes (including the default lane), 32,768 items +and 16 MiB per checkpoint, and 32 MiB of aggregate serialized checkpoints. These are +retention ceilings, not measured process-memory bounds. New lanes receive an +immediate 429 once the lifetime lane limit is reached. Request-count and request-byte +overloads also return immediate 429 responses without mutating retained state, so +rejected work cannot invalidate an earlier accepted queued continuation. Other +validation errors execute in lane order and evict only a matching referenced parent +when routing is valid. Parent lookup happens when execution begins; accepted queues +do not reserve parent snapshots. Existing fork and execution-failure eviction rules +still apply. Disconnect aborts and joins request tasks, waits for active leases to +release pinned state, and drops the entire connection group before the close handshake. + +A `ResponseSession` owns one latest canonical checkpoint and one execution slot. +The executor pins a parent before inference and publishes completed or incomplete +state before exposing terminal completion. Failed continuations discard only a +referenced checkpoint owned by that session. Dropping the owner closes the session +and rejects late publication. Callers still cancel and join active work explicitly; +`wait_until_idle` waits for the execution lease to end, but does not cancel it. + +`ResponseSessionGroup` allows independent serial members to find and pin each +other's latest checkpoints. Failed forks cannot evict the source member's state. +The group bounds lifetime member count, each checkpoint's items and serialized +bytes, and aggregate retained bytes. Shared parent references count once; replaced +parents still pinned by active work and prepared checkpoints awaiting persistence +remain charged until their last reference is released. Reservation happens before +durable writes, and failure or cancellation returns unused capacity. Replacement +requires room for both old and new snapshots until publication. These retention +budgets are not a bound on temporary allocations, execution copies or process +memory. Scheduling, FIFO queues and active-work limits remain caller concerns. + +Response-scoped session history records reasoning, messages and calls in inference-round order, +then built-in call outputs. Public output is accumulated separately and is not +appended to retained history twice. Compaction replaces the canonical window while +preserving MCP discovery records needed for orchestration. Durable restoration +and replayed compaction input select the effective compacted window before validating current calls; +obsolete stored rows are not deleted or charged to that retained window. + +For response-scoped session execution, `store: false` has no durable writes or +database fallback. A stored child of a transient parent persists the complete +canonical window without creating a row or dangling database reference for that +parent. Explicit `conversation_id` requests keep the existing durable Conversations +policy and append output only through the conversation handler. Their session lease +still serializes execution without recording a second copy. Core commit supports +prewarming without inference; WebSocket `generate: false` uses that same session +rehydration and commit path. + ### `handler/common.rs` Transport helpers shared by the HTTP and WS handlers: body reading with a shared size diff --git a/TERMINOLOGY.md b/TERMINOLOGY.md index c48c10c9..5c1be85b 100644 --- a/TERMINOLOGY.md +++ b/TERMINOLOGY.md @@ -92,7 +92,8 @@ round. ### conversation state The prior items and metadata made available to a later turn. State may be managed with a conversation, chained with -`previous_response_id`, or replayed manually. +`previous_response_id`, or replayed manually. WebSocket response state may also be +retained transiently by the active connection; it is distinct from durable conversation storage. ### stored response @@ -106,7 +107,8 @@ mechanism could be either `previous_response_id` or a conversation. ### previous response ID -The response identifier passed in the `previous_response_id` field to continue from a prior stored response. In prose, +The response identifier passed in the `previous_response_id` field to continue from prior response state. That state +may be durable or cached on the active WebSocket connection. In prose, write **previous response ID**; in code and wire-format discussion, use `previous_response_id`. ### rehydration @@ -127,7 +129,8 @@ Describes a flow in which the service retains or resolves prior state, such as R ### stateless Describes a flow in which the request supplies all required context and the service does not rely on retained response -or conversation state. `store: false` disables stored-response state, although callers may still replay prior items. +or conversation state. `store: false` disables durable response storage, although callers may still replay prior items +or continue from the active WebSocket connection's transient checkpoint. An explicit conversation remains durable. ### compaction diff --git a/crates/agentic-llm-d/src/context.rs b/crates/agentic-llm-d/src/context.rs index 5403ad5e..afb46f62 100644 --- a/crates/agentic-llm-d/src/context.rs +++ b/crates/agentic-llm-d/src/context.rs @@ -65,6 +65,7 @@ impl From for RequestContext { conversation_id: wire.conversation_id, // Conversation mode is rejected, so there is no version to resume. conversation_version: None, + continuation: None, } } } diff --git a/crates/agentic-llm-d/tests/split_execution_integration.rs b/crates/agentic-llm-d/tests/split_execution_integration.rs index 67c557b7..1681ac21 100644 --- a/crates/agentic-llm-d/tests/split_execution_integration.rs +++ b/crates/agentic-llm-d/tests/split_execution_integration.rs @@ -287,12 +287,12 @@ async fn a_function_call_stream_passes_strict_validation() { #[tokio::test] async fn relayed_json_tool_call_ids_are_validated_before_persistence() { - assert_relayed_tool_call_ids_are_validated(false).await; + Box::pin(assert_relayed_tool_call_ids_are_validated(false)).await; } #[tokio::test] async fn relayed_sse_tool_call_ids_are_validated_before_persistence() { - assert_relayed_tool_call_ids_are_validated(true).await; + Box::pin(assert_relayed_tool_call_ids_are_validated(true)).await; } async fn assert_relayed_tool_call_ids_are_validated(stream: bool) { @@ -402,12 +402,12 @@ async fn assert_relayed_tool_call_ids_are_validated(stream: bool) { #[tokio::test] async fn relayed_json_call_id_cannot_reuse_continued_history() { - assert_relayed_call_id_cannot_reuse_continued_history(false).await; + Box::pin(assert_relayed_call_id_cannot_reuse_continued_history(false)).await; } #[tokio::test] async fn relayed_sse_call_id_cannot_reuse_continued_history() { - assert_relayed_call_id_cannot_reuse_continued_history(true).await; + Box::pin(assert_relayed_call_id_cannot_reuse_continued_history(true)).await; } async fn assert_relayed_call_id_cannot_reuse_continued_history(stream: bool) { diff --git a/crates/agentic-server-core/src/executor/compaction.rs b/crates/agentic-server-core/src/executor/compaction.rs index 68aae007..a2045e77 100644 --- a/crates/agentic-server-core/src/executor/compaction.rs +++ b/crates/agentic-server-core/src/executor/compaction.rs @@ -417,6 +417,7 @@ pub(crate) async fn compact_items( response_id: uuid7_str("resp_"), conversation_id: None, conversation_version: None, + continuation: None, }; let mut agent = agent_pipeline(ctx, None, None); let response = @@ -469,6 +470,9 @@ pub(crate) async fn maybe_compact_context( let (compacted, usage) = compact_items(&model, input, instructions.as_deref(), exec_ctx, auth).await?; ctx.enriched_request.input = ResponsesInput::Items(compacted.clone()); ctx.new_input_items = compacted; + if let Some(continuation) = &mut ctx.continuation { + continuation.mark_history_replaced(); + } Ok(Some(usage)) } @@ -605,6 +609,7 @@ mod tests { response_id: "resp_test".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } @@ -622,7 +627,7 @@ mod tests { async fn mock_execution_context(response_store: ResponseStore) -> (ExecutionContext, tokio::task::JoinHandle<()>) { let app = Router::new().route( "/v1/responses", - post(|| async { + post(|_body: axum::body::Bytes| async { axum::Json(serde_json::json!({ "id": "resp_upstream", "object": "response", diff --git a/crates/agentic-server-core/src/executor/engine.rs b/crates/agentic-server-core/src/executor/engine.rs index 529bdbad..8fbaa2aa 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::executor::inference::DONE_MARKER; use crate::executor::persist::persist_if_needed; use crate::executor::pipeline::{AgentPipeline, emit_deferred_stream_events}; use crate::executor::prepare::prepare_request_tools; -use crate::executor::rehydrate::{prepare_reasoning_for_vllm, rehydrate_conversation, validate_reasoning_for_vllm}; +use crate::executor::rehydrate::{prepare_reasoning_for_vllm, validate_reasoning_for_vllm}; use crate::executor::request::{ExecutionContext, RequestContext}; use crate::executor::response_budget::ExecutorResponseBudget; #[cfg(test)] @@ -199,6 +199,34 @@ fn prepare_initial_reasoning_for_vllm(input: &mut ResponsesInput, round: usize, Ok(()) } +fn record_round_history( + ctx: &mut RequestContext, + output_items: &[OutputItem], + registry: &ToolRegistry, + public_output_count: usize, +) { + // Explicit conversations append public output through their durable handler; + // the session lease still serializes execution but must not record it twice. + if let Some(continuation) = ctx + .continuation + .as_mut() + .filter(|_| ctx.original_request.conversation_id.is_none()) + { + // The canonical sequence includes reasoning and intermediate messages in + // their original positions, followed by this round's tool call outputs. + // Discovery records are appended separately from the public response. + ctx.new_input_items.extend( + output_items + .iter() + .filter(|item| !matches!(item, OutputItem::McpListTools(_))) + .filter_map(OutputItem::to_input_item), + ); + continuation.mark_outputs_recorded(public_output_count); + } else { + append_gateway_calls_to_new_input(ctx, output_items, registry); + } +} + /// Request-scoped owner of registry-backed tool orchestration and its shared byte budget. struct EngineOrchestration<'a> { agent: &'a mut AgentPipeline, @@ -282,11 +310,17 @@ impl<'a> EngineOrchestration<'a> { .await?; let public_output = public_output_items(¤t_output, &self.registry, &gateway_results)?; combined_output.extend(public_output); + record_round_history( + &mut self.agent.request, + ¤t_output, + &self.registry, + combined_output.len(), + ); // A terminal incomplete response may still contain completed gateway // calls. Record those results, but never start another inference round. if payload.status == "incomplete" { - self.record_round_input(¤t_output, gateway_results); + self.record_gateway_results(gateway_results); finalize_loop(&mut payload, combined_output, combined_usage, &self.agent.request); let tool_search_metadata = self.agent.take_tool_search_metadata(); return Ok((payload, tool_search_metadata)); @@ -297,7 +331,7 @@ impl<'a> EngineOrchestration<'a> { // are handed back to the caller. Gateway calls in the same round are // still recorded so the returned conversation is complete. LoopDecision::RequiresClientAction => { - self.record_round_input(¤t_output, gateway_results); + self.record_gateway_results(gateway_results); finalize_loop(&mut payload, combined_output, combined_usage, &self.agent.request); let tool_search_metadata = self.agent.take_tool_search_metadata(); return Ok((payload, tool_search_metadata)); @@ -314,7 +348,7 @@ impl<'a> EngineOrchestration<'a> { // The final round's gateway calls and outputs are recorded so a // continuation is not fed a dangling tool call. LoopDecision::Incomplete(reason) => { - self.record_round_input(¤t_output, gateway_results); + self.record_gateway_results(gateway_results); finalize_loop(&mut payload, combined_output, combined_usage, &self.agent.request); "incomplete".clone_into(&mut payload.status); payload.incomplete_details = Some(IncompleteDetails { reason: Some(reason) }); @@ -325,7 +359,7 @@ impl<'a> EngineOrchestration<'a> { LoopDecision::Continue => { self.agent.request.enriched_request.tool_choice = Some(ToolChoice::Auto); append_output_items_to_input(&mut self.agent.request.enriched_request.input, ¤t_output); - self.record_round_input(¤t_output, gateway_results); + self.record_gateway_results(gateway_results); } } } @@ -333,8 +367,7 @@ impl<'a> EngineOrchestration<'a> { unreachable!("the final round returns Done, RequiresClientAction, or Incomplete"); } - fn record_round_input(&mut self, output: &[OutputItem], results: Vec) { - append_gateway_calls_to_new_input(&mut self.agent.request, output, &self.registry); + fn record_gateway_results(&mut self, results: Vec) { append_tool_outputs( &mut self.agent.request, results.into_iter().map(|result| result.input_item).collect(), @@ -450,6 +483,9 @@ async fn run_compaction_trigger( unreachable!("compact_items always appends a compaction item"); }; ctx.new_input_items = compacted; + if let Some(continuation) = &mut ctx.continuation { + continuation.mark_history_replaced(); + } let mut payload = ResponsePayload { id: ctx.response_id.clone(), object: "response".to_owned(), @@ -683,6 +719,7 @@ pub struct ExecuteRequest { payload: RequestPayload, exec_ctx: Arc, client_auth: Option, + continuation: Option, } impl ExecuteRequest { @@ -692,6 +729,7 @@ impl ExecuteRequest { payload, exec_ctx, client_auth: None, + continuation: None, } } @@ -702,6 +740,15 @@ impl ExecuteRequest { self } + /// Retain this turn's continuation state in the supplied serial session. + /// + /// # Errors + /// Returns an error when the session is busy or closed. + pub fn with_session(mut self, session: &super::ResponseSession) -> ExecutorResult { + self.continuation = Some(session.begin(self.payload.previous_response_id.as_deref())?); + Ok(self) + } + /// Execute one stateful conversation turn. /// /// Returns `Either::Left(ResponsePayload)` for non-streaming requests, or @@ -720,7 +767,8 @@ impl ExecuteRequest { tools = self.payload.tools.as_ref().map_or(0, Vec::len), "executor received responses request" ); - let ctx = rehydrate_conversation(self.payload, &self.exec_ctx).await?; + let ctx = + super::rehydrate::rehydrate_with_continuation(self.payload, &self.exec_ctx, self.continuation).await?; if !ctx.enriched_request.input.has_compaction_trigger() { validate_reasoning_for_vllm(&ctx.enriched_request.input)?; } @@ -838,10 +886,18 @@ mod tests { "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_upstream\",\"status\":\"completed\",\"usage\":null}}\n\n", "data: [DONE]\n\n", ); - let app = axum::Router::new().route( - "/v1/responses", - axum::routing::post(|| async { ([(axum::http::header::CONTENT_TYPE, "text/event-stream")], UPSTREAM_SSE) }), - ); + let app = axum::Router::new() + .route( + "/v1/responses", + axum::routing::post(|_body: axum::body::Bytes| async { + ([(axum::http::header::CONTENT_TYPE, "text/event-stream")], UPSTREAM_SSE) + }), + ) + // Read the oversized test request before replying, allowing JSON framing + // above the response budget without relying on an early HTTP response. + .layer(axum::extract::DefaultBodyLimit::max( + MAX_EXECUTOR_RESPONSE_BYTES + 64 * 1024, + )); let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("bind streaming mock inference server"); @@ -874,6 +930,7 @@ mod tests { response_id: "resp_test".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; let mut exec_ctx = ExecutionContext::new( ConversationHandler::new(ConversationStore::disabled()), @@ -934,6 +991,7 @@ mod tests { response_id: "resp_mcp".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; let plain_payload: RequestPayload = serde_json::from_value(serde_json::json!({ "model": "test-model", @@ -948,6 +1006,7 @@ mod tests { response_id: "resp_plain".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; let mut exec_ctx = ExecutionContext::new( ConversationHandler::new(ConversationStore::disabled()), diff --git a/crates/agentic-server-core/src/executor/error.rs b/crates/agentic-server-core/src/executor/error.rs index 7e1e3888..1011669c 100644 --- a/crates/agentic-server-core/src/executor/error.rs +++ b/crates/agentic-server-core/src/executor/error.rs @@ -78,6 +78,10 @@ pub enum ExecutorError { #[error("{entity} not found: {id}")] NotFound { entity: String, id: String }, + /// A response session cannot resolve the requested continuation checkpoint. + #[error("Previous response with id '{id}' not found.")] + PreviousResponseNotFound { id: String }, + #[error("invalid request: {0}")] InvalidRequest(String), @@ -128,6 +132,7 @@ impl ExecutorError { Self::ConversationLocked { .. } | Self::Tool(ToolError::Config(_) | ToolError::MissingOutput { .. }) | Self::InvalidRequest(_) + | Self::PreviousResponseNotFound { .. } | Self::JsonError(_) => StatusCode::BAD_REQUEST, Self::Tool( ToolError::Execution(_) @@ -149,6 +154,7 @@ impl ExecutorError { Self::ConversationLocked { .. } | Self::Tool(ToolError::Config(_) | ToolError::MissingOutput { .. }) | Self::InvalidRequest(_) + | Self::PreviousResponseNotFound { .. } | Self::ParseError(_) | Self::JsonError(_) | Self::PayloadTooLarge(_) => "invalid_request_error", @@ -169,6 +175,7 @@ impl ExecutorError { pub fn error_code(&self) -> &'static str { match self.client_visible_error() { Self::ConversationLocked { .. } => "conversation_locked", + Self::PreviousResponseNotFound { .. } => "previous_response_not_found", Self::Conflict(_) => "response_already_stored", Self::PayloadTooLarge(_) => "body_too_large", other => other.error_type(), @@ -180,6 +187,7 @@ impl ExecutorError { pub fn error_param(&self) -> Option<&'static str> { match self.client_visible_error() { Self::ConversationLocked { .. } => Some("conversation"), + Self::PreviousResponseNotFound { .. } => Some("previous_response_id"), Self::Tool(ToolError::MissingOutput { .. }) => Some("input"), _ => None, } @@ -260,6 +268,22 @@ mod tests { assert!(exec_err.to_string().contains("storage error")); } + #[test] + fn previous_response_error_has_the_continuation_envelope() { + let error = ExecutorError::PreviousResponseNotFound { + id: "resp_missing".to_owned(), + }; + assert_eq!(error.http_status(), StatusCode::BAD_REQUEST); + assert_eq!( + serde_json::from_slice::(&error.into_response_body()).unwrap(), + serde_json::json!({"error": { + "message": "Previous response with id 'resp_missing' not found.", + "type": "invalid_request_error", "code": "previous_response_not_found", + "param": "previous_response_id" + }}) + ); + } + #[test] fn tool_search_configuration_errors_are_bad_requests() { let error = ExecutorError::from(ToolError::Config("invalid tool_search request".to_owned())); diff --git a/crates/agentic-server-core/src/executor/mod.rs b/crates/agentic-server-core/src/executor/mod.rs index 7dcde006..3880c934 100644 --- a/crates/agentic-server-core/src/executor/mod.rs +++ b/crates/agentic-server-core/src/executor/mod.rs @@ -14,6 +14,7 @@ pub mod persist; mod prepare; pub mod rehydrate; pub mod request; +pub mod session; pub mod translate; mod gateway; @@ -33,7 +34,8 @@ pub use messages_request::normalize_native_web_search_for_upstream; pub use messages_stream::run_messages_stream; pub use modes::{ConversationHandler, ResponseHandler}; pub use persist::{commit, persist_response, persist_turn}; -pub use rehydrate::rehydrate_conversation; +pub use rehydrate::{rehydrate_conversation, rehydrate_in_session}; pub use request::ExecutionContext; pub use request::RequestContext; +pub use session::{ResponseSession, ResponseSessionGroup}; pub use upstream::{UpstreamBody, decode_upstream, upstream_request}; diff --git a/crates/agentic-server-core/src/executor/modes/conversation.rs b/crates/agentic-server-core/src/executor/modes/conversation.rs index b511e65f..d79b3672 100644 --- a/crates/agentic-server-core/src/executor/modes/conversation.rs +++ b/crates/agentic-server-core/src/executor/modes/conversation.rs @@ -210,6 +210,7 @@ mod tests { response_id: "resp_test".into(), conversation_id: conversation_id.map(str::to_string), conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/src/executor/modes/response.rs b/crates/agentic-server-core/src/executor/modes/response.rs index 718222da..ee3629da 100644 --- a/crates/agentic-server-core/src/executor/modes/response.rs +++ b/crates/agentic-server-core/src/executor/modes/response.rs @@ -60,11 +60,12 @@ impl ResponseHandler { self.store.rehydrate(prev_id).await.map_err(ExecutorError::Storage) } - /// Persists a response record — only the new items from this turn. + /// Completes response-scoped persistence and optional transient checkpoint retention. /// /// Takes `ctx` and `output_items` by value so fields can be moved directly - /// into [`crate::storage::ResponseMetadata`]. Prior history must not be - /// re-inserted; the response store records item IDs for this response only. + /// into [`ResponseMetadata`]. Normally only this turn's items are inserted; + /// promoting a transient parent to a stored child requires its full canonical + /// history because there are no durable parent item IDs to reference. /// /// # Errors /// Returns `ExecutorError` if the store is disabled or the database operation fails. @@ -84,31 +85,80 @@ impl ResponseHandler { /// Persists a response using metadata prepared by request-scoped tool behavior. pub(crate) async fn execute_turn_with_metadata( &self, - ctx: RequestContext, + mut ctx: RequestContext, output_items: Vec, metadata: ResponseMetadata, ) -> ExecutorResult<()> { + let continuation = ctx.continuation.take(); + let write_durable = continuation.is_none() || ctx.original_request.store; let mut new_items = Vec::with_capacity(ctx.new_input_items.len() + output_items.len()); new_items.extend(ctx.new_input_items.into_iter().map(InOutItem::Input)); - new_items.extend(output_items.into_iter().map(InOutItem::Output)); - - let result = self - .store - .persist_with_conversation_id( - &ctx.response_id, - ctx.conversation_id.as_deref(), - metadata.previous_response_id.as_deref(), - new_items, - &metadata, - ) - .await; + new_items.extend( + output_items + .into_iter() + .enumerate() + .filter(|(index, item)| { + continuation + .as_ref() + .is_none_or(|lease| lease.retains_output(*index, item)) + }) + .map(|(_, item)| InOutItem::Output(item)), + ); + + let checkpoint = continuation + .as_ref() + .map(|lease| { + lease.checkpoint( + ctx.response_id.clone(), + ctx.conversation_id.clone(), + &metadata, + &new_items, + write_durable, + ) + }) + .transpose()?; + + // A stored child of a transient parent needs its complete canonical + // checkpoint, not a database reference to a response that was never stored. + let transient_parent = continuation + .as_ref() + .filter(|lease| lease.parent.as_ref().is_some_and(|parent| !parent.durable)); + let previous_durable_id = if let Some(lease) = transient_parent { + if write_durable { + let mut full_items = lease.parent_items().cloned().map(InOutItem::Input).collect::>(); + full_items.append(&mut new_items); + new_items = full_items; + } + None + } else { + metadata.previous_response_id.as_deref() + }; + + let result = if write_durable { + self.store + .persist_with_conversation_id( + &ctx.response_id, + ctx.conversation_id.as_deref(), + previous_durable_id, + new_items, + &metadata, + ) + .await + } else { + Ok(()) + }; match result { Err(error) if error.is_unique_violation() => Err(ExecutorError::Conflict(format!( "a turn is already stored under '{}'", ctx.response_id ))), Err(error) => Err(ExecutorError::Storage(error)), - Ok(()) => Ok(()), + Ok(()) => { + if let Some((lease, checkpoint)) = continuation.zip(checkpoint) { + lease.publish(checkpoint)?; + } + Ok(()) + } } } } @@ -154,6 +204,7 @@ mod tests { response_id: "resp_test".into(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/src/executor/persist.rs b/crates/agentic-server-core/src/executor/persist.rs index f6b91405..13df68ac 100644 --- a/crates/agentic-server-core/src/executor/persist.rs +++ b/crates/agentic-server-core/src/executor/persist.rs @@ -18,7 +18,8 @@ use tracing::error; #[must_use] pub(crate) fn should_persist(ctx: &RequestContext) -> bool { - ctx.original_request.store + ctx.continuation.is_some() + || ctx.original_request.store || ctx.original_request.previous_response_id.is_some() || ctx.original_request.conversation_id.is_some() } @@ -32,7 +33,7 @@ pub(crate) async fn persist_if_needed( ) -> ExecutorResult<()> { if should_persist(&ctx) { match persist_prepared_response(payload, ctx, tool_search_metadata, conv_handler, resp_handler).await { - Err(error @ ExecutorError::Conflict(_)) => Err(error), + Err(error @ (ExecutorError::Conflict(_) | ExecutorError::PayloadTooLarge(_))) => Err(error), Err(source) => { error!(error = ?source, "failed to persist response"); Err(ExecutorError::Persistence(Box::new(source))) @@ -204,13 +205,13 @@ async fn validate_output_call_ids( return Ok(()); } - for (history_index, item) in resp_handler.rehydrate(ctx).await?.iter().enumerate() { - if let Some((output_index, item_type)) = stored_call_id(item).and_then(|call_id| call_ids.get(call_id)) { - return Err(ExecutorError::InvalidRequest(format!( - "upstream response output[{output_index}] {item_type} repeats 'call_id' from continued history item[{history_index}]" - ))); - } + if ctx.continuation.is_some() { + // Rehydration already combined the pinned parent with normalized input. + // Validate that effective window, excluding calls superseded by compaction. + return validate_history_call_ids(ctx.enriched_request.input.model_items().map(input_call_id), &call_ids); } + let history = resp_handler.rehydrate(ctx).await?; + validate_history_call_ids(history.iter().map(stored_call_id), &call_ids)?; for (input_index, item) in ctx.new_input_items.iter().enumerate() { if let Some((output_index, item_type)) = input_call_id(item).and_then(|call_id| call_ids.get(call_id)) { return Err(ExecutorError::InvalidRequest(format!( @@ -221,6 +222,20 @@ async fn validate_output_call_ids( Ok(()) } +fn validate_history_call_ids<'a>( + history: impl Iterator>, + call_ids: &HashMap<&str, (usize, &str)>, +) -> ExecutorResult<()> { + for (history_index, call_id) in history.enumerate() { + if let Some((output_index, item_type)) = call_id.and_then(|id| call_ids.get(id)) { + return Err(ExecutorError::InvalidRequest(format!( + "upstream response output[{output_index}] {item_type} repeats 'call_id' from continued history item[{history_index}]" + ))); + } + } + Ok(()) +} + fn stored_call_id(item: &InOutItem) -> Option<&str> { match item { InOutItem::Input(item) => input_call_id(item), diff --git a/crates/agentic-server-core/src/executor/prepare.rs b/crates/agentic-server-core/src/executor/prepare.rs index c32be5cb..0c8db43d 100644 --- a/crates/agentic-server-core/src/executor/prepare.rs +++ b/crates/agentic-server-core/src/executor/prepare.rs @@ -34,6 +34,11 @@ async fn restored_loaded_tools( } if ctx.original_request.previous_response_id.is_some() { + // Session parents may never have been written to durable storage. + // Their checkpoint retains the same public metadata as a stored response. + if let Some(parent) = ctx.continuation.as_ref().and_then(|lease| lease.parent.as_ref()) { + return Ok(parent.metadata.tool_search_loaded_tools.clone().unwrap_or_default()); + } return Ok(resp_handler .get(ctx) .await? @@ -55,3 +60,52 @@ async fn restored_loaded_tools( } Ok(Vec::new()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::executor::ResponseSession; + use crate::storage::{ConversationStore, ResponseMetadata, ResponseStore}; + use crate::types::request_response::RequestPayload; + use std::num::NonZeroUsize; + + #[tokio::test] + async fn compacted_session_restores_loaded_tools_without_durable_storage() { + let session = ResponseSession::new(NonZeroUsize::new(100).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let loaded: ResponsesTool = serde_json::from_value(serde_json::json!({ + "type": "function", "name": "weather", "parameters": {"type": "object"} + })) + .unwrap(); + let metadata = ResponseMetadata { + tool_search_loaded_tools: Some(vec![loaded.clone()]), + ..ResponseMetadata::default() + }; + let lease = session.begin(None).unwrap(); + let checkpoint = lease + .checkpoint("resp_parent".to_owned(), None, &metadata, &[], false) + .unwrap(); + lease.publish(checkpoint).unwrap(); + let request: RequestPayload = serde_json::from_value(serde_json::json!({ + "model": "test", "store": false, "previous_response_id": "resp_parent", + "input": [{"type": "compaction", "id": "cmp_1", "encrypted_content": "summary"}] + })) + .unwrap(); + let mut ctx = RequestContext { + original_request: request.clone(), + enriched_request: request, + new_input_items: Vec::new(), + response_id: "resp_child".to_owned(), + conversation_id: None, + conversation_version: None, + continuation: Some(session.begin(Some("resp_parent")).unwrap()), + }; + let restored = restored_loaded_tools( + &mut ctx, + &ConversationHandler::new(ConversationStore::disabled()), + &ResponseHandler::new(ResponseStore::disabled()), + ) + .await + .unwrap(); + assert_eq!(serde_json::to_value(restored).unwrap(), serde_json::json!([loaded])); + } +} diff --git a/crates/agentic-server-core/src/executor/rehydrate.rs b/crates/agentic-server-core/src/executor/rehydrate.rs index 77556334..f4ebc215 100644 --- a/crates/agentic-server-core/src/executor/rehydrate.rs +++ b/crates/agentic-server-core/src/executor/rehydrate.rs @@ -3,6 +3,7 @@ //! Builds a [`RequestContext`] by loading prior turns from storage and //! injecting them into the enriched request before it is forwarded to the LLM. +use super::session::{ResponseCheckpoint, ResponseContinuation, ResponseSession, canonical_session_history}; use crate::executor::error::{ExecutorError, ExecutorResult}; use crate::executor::pending_calls::pending_calls; use crate::executor::request::{ExecutionContext, RequestContext}; @@ -122,6 +123,27 @@ pub(super) fn prepare_reasoning_for_vllm(input: &mut ResponsesInput) -> Executor pub async fn rehydrate_conversation( request: RequestPayload, exec_ctx: &ExecutionContext, +) -> ExecutorResult { + rehydrate_with_continuation(request, exec_ctx, None).await +} + +/// Resolve one turn in a transient session without changing the durable Conversations API. +/// +/// # Errors +/// Returns an error when the session is busy/closed or the requested history is unavailable. +pub async fn rehydrate_in_session( + request: RequestPayload, + exec_ctx: &ExecutionContext, + session: &ResponseSession, +) -> ExecutorResult { + let continuation = session.begin(request.previous_response_id.as_deref())?; + rehydrate_with_continuation(request, exec_ctx, Some(continuation)).await +} + +pub(crate) async fn rehydrate_with_continuation( + request: RequestPayload, + exec_ctx: &ExecutionContext, + continuation: Option, ) -> ExecutorResult { // Fail before storage work for new files; check again once history is resolved. validate_message_files(&request.input)?; @@ -142,6 +164,7 @@ pub async fn rehydrate_conversation( response_id, conversation_id: None, conversation_version: None, + continuation, }; if ctx.original_request.conversation_id.is_some() && ctx.original_request.previous_response_id.is_some() { @@ -168,6 +191,9 @@ pub async fn rehydrate_conversation( /// tools and tool choice from the stored metadata, and prepends the history to /// the enriched request input. async fn from_response(ctx: &mut RequestContext, exec_ctx: &ExecutionContext) -> ExecutorResult<()> { + if ctx.continuation.is_some() { + return from_session_response(ctx, exec_ctx).await; + } let stored = exec_ctx.resp_handler.get(ctx).await?; let history = exec_ctx.resp_handler.rehydrate(ctx).await?; @@ -187,6 +213,69 @@ async fn from_response(ctx: &mut RequestContext, exec_ctx: &ExecutionContext) -> Ok(()) } +async fn from_session_response(ctx: &mut RequestContext, exec_ctx: &ExecutionContext) -> ExecutorResult<()> { + let cached = ctx + .continuation + .as_ref() + .and_then(|continuation| continuation.parent.clone()); + let parent = if let Some(parent) = cached { + parent + } else { + let previous_id = ctx.original_request.previous_response_id.as_deref().unwrap_or_default(); + if !ctx.original_request.store { + return Err(ExecutorError::PreviousResponseNotFound { + id: previous_id.to_owned(), + }); + } + let map_missing = |error| match error { + ExecutorError::Storage(ref source) if source.is_not_found() => ExecutorError::PreviousResponseNotFound { + id: previous_id.to_owned(), + }, + other => other, + }; + let stored = exec_ctx.resp_handler.get(ctx).await.map_err(map_missing)?; + let history = exec_ctx.resp_handler.rehydrate(ctx).await.map_err(map_missing)?; + let checkpoint = ResponseCheckpoint { + response_id: stored.response_id, + conversation_id: stored.conversation_id, + history: canonical_session_history(InOutItem::into_input_items(history)), + metadata: stored.metadata, + durable: true, + }; + let continuation = ctx.continuation.as_ref().ok_or_else(|| { + ExecutorError::InvalidRequest("session response requires a continuation lease".to_owned()) + })?; + std::sync::Arc::new(continuation.retain_parent(checkpoint)?) + }; + let mut items = parent.history.clone(); + // Match durable rehydration: lower custom/shell call outputs only in the + // inference copy. new_input_items retains their public wire types. + items.extend(Vec::from(&ctx.original_request.input)); + let items = canonical_session_history(items); + if let Some(pending) = pending_calls(&items)?.into_iter().next() { + return Err(ExecutorError::Tool(ToolError::MissingOutput { + call_id: pending.call_id, + })); + } + ctx.enriched_request.previous_response_id = None; + ctx.enriched_request.input = ResponsesInput::Items(items); + ctx.enriched_request.tools = resolve_tools( + ctx.original_request.tools.as_deref(), + parent.metadata.effective_tools.as_deref(), + ctx.original_request.tools.is_some(), + ); + ctx.enriched_request.tool_choice = Some(resolve_tool_choice( + ctx.original_request.tool_choice.as_ref(), + &parent.metadata.effective_tool_choice, + ctx.original_request.tool_choice.is_some(), + )); + ctx.conversation_id.clone_from(&parent.conversation_id); + if let Some(continuation) = ctx.continuation.as_mut() { + continuation.parent = Some(parent); + } + Ok(()) +} + /// Hydrates `ctx` from the conversation store. /// /// Gets or creates the conversation (depending on `store`) and rehydrates its @@ -245,6 +334,72 @@ mod tests { use crate::types::io::output::{McpListTools, OutputItem}; use crate::types::request_response::RequestPayload; + #[test] + fn session_history_without_compaction_is_unchanged() { + let history: Vec = serde_json::from_value(serde_json::json!([ + {"role":"user", "content":"question"}, + {"type":"mcp_list_tools", "id":"mcpl_1", "server_label":"counter", "tools":[]}, + {"type":"function_call", "call_id":"pending", "name":"lookup", "arguments":"{}"} + ])) + .unwrap(); + let expected = serde_json::to_value(&history).unwrap(); + assert_eq!( + serde_json::to_value(canonical_session_history(history)).unwrap(), + expected + ); + } + + #[test] + fn restored_session_keeps_canonical_compaction_and_mcp_discovery() { + let history: Vec = serde_json::from_value(serde_json::json!([ + {"type":"mcp_list_tools", "id":"mcpl_1", "server_label":"counter", "tools":[]}, + {"role":"assistant", "content":"obsolete detail"}, + {"type":"function_call", "call_id":"obsolete", "name":"lookup", "arguments":"{}"}, + {"role":"user", "content":"noncanonical old user"}, + {"role":"user", "id":"msg_keep", "status":"completed", "content":"canonical user"}, + {"type":"compaction", "id":"cmp_latest", "encrypted_content":"latest summary"}, + {"type":"function_call", "call_id":"live", "name":"lookup", "arguments":"{}"}, + {"type":"function_call_output", "call_id":"live", "output":"live result"} + ])) + .unwrap(); + let history = canonical_session_history(history); + assert_eq!(history.len(), 5); + assert!(matches!(history[0], InputItem::McpListTools(_))); + assert!(matches!(&history[2], InputItem::Compaction(item) + if item.id.as_deref() == Some("cmp_latest") && item.encrypted_content == "latest summary")); + assert!(pending_calls(&history).unwrap().is_empty()); + let serialized = serde_json::to_value(&history).unwrap(); + assert_eq!(serialized[1]["id"], "msg_keep"); + assert_eq!(serialized[3]["call_id"], "live"); + assert_eq!(serialized[4]["call_id"], "live"); + assert!(!serialized.to_string().contains("obsolete")); + } + + #[test] + fn restored_session_uses_the_latest_window_and_preserves_live_pending_calls() { + let history: Vec = serde_json::from_value(serde_json::json!([ + {"role":"user", "id":"msg_old", "status":"completed", "content":"old user"}, + {"type":"mcp_list_tools", "id":"mcpl_old", "server_label":"old", "tools":[]}, + {"type":"compaction", "id":"cmp_old", "encrypted_content":"old summary"}, + {"role":"user", "id":"msg_keep", "status":"completed", "content":"retained user"}, + {"role":"user", "id":"msg_incomplete", "status":"in_progress", "content":"superseded"}, + {"type":"mcp_list_tools", "id":"mcpl_new", "server_label":"new", "tools":[]}, + {"type":"compaction", "id":"cmp_new", "encrypted_content":"latest summary"}, + {"type":"function_call", "call_id":"live", "name":"lookup", "arguments":"{}"} + ])) + .unwrap(); + let history = canonical_session_history(history); + let serialized = serde_json::to_value(&history).unwrap(); + assert_eq!(history.len(), 5); + assert_eq!(serialized[0]["id"], "mcpl_old"); + assert_eq!(serialized[1]["id"], "msg_keep"); + assert_eq!(serialized[2]["id"], "mcpl_new"); + assert_eq!(serialized[3]["id"], "cmp_new"); + let pending = pending_calls(&history).unwrap(); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].call_id, "live"); + } + fn reasoning_item(content: &[&str], encrypted_content: Option) -> InputItem { InputItem::Reasoning(ReasoningOutput { id: "rs_prior".to_owned(), diff --git a/crates/agentic-server-core/src/executor/request.rs b/crates/agentic-server-core/src/executor/request.rs index 5ffa9ed7..9b77325a 100644 --- a/crates/agentic-server-core/src/executor/request.rs +++ b/crates/agentic-server-core/src/executor/request.rs @@ -27,7 +27,9 @@ pub struct RequestContext { /// Enriched request with rehydrated conversation history injected into `.input`. /// This is the request forwarded to the LLM. pub enriched_request: RequestPayload, - /// Only the new input items submitted by the client this turn (used for persistence). + /// This turn's canonical input items for persistence, including built-in calls + /// and call outputs. Session execution also records each round's other output + /// items here in inference order; compaction can replace them with a new window. pub new_input_items: Vec, /// Our generated response ID (uuid7 with "resp_" prefix). pub response_id: String, @@ -36,6 +38,8 @@ pub struct RequestContext { /// Conversation version captured with rehydrated history. /// `None` for non-conversation and `previous_response_id` execution. pub conversation_version: Option, + /// Optional transient-session lease and canonical parent snapshot. Never serialized. + pub continuation: Option, } impl RequestContext { diff --git a/crates/agentic-server-core/src/executor/session.rs b/crates/agentic-server-core/src/executor/session.rs new file mode 100644 index 00000000..ada87f20 --- /dev/null +++ b/crates/agentic-server-core/src/executor/session.rs @@ -0,0 +1,873 @@ +//! Transient continuation state owned by one serial response session. + +#[cfg(test)] +#[path = "session_budget_tests.rs"] +mod budget_tests; + +use std::fmt; +use std::num::NonZeroUsize; +use std::ops::Deref; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard, Weak}; + +use serde::Serialize; +use tokio::sync::Notify; + +use super::{ExecutorError, ExecutorResult}; +use crate::storage::{InOutItem, ResponseMetadata}; +use crate::types::io::input::latest_compaction_window; +use crate::types::io::{InputItem, OutputItem}; +use crate::utils::common::serialized_size_up_to; + +/// One latest response checkpoint with a lifetime independent of durable storage. +/// +/// A session admits one active turn. The owner must outlive execution; dropping it +/// clears cached state and prevents a late completion from publishing another +/// checkpoint. Active turns must still be cancelled/joined by their caller. +/// Budgets constrain retained item count and serialized size, not total heap use. +#[derive(Debug)] +pub struct ResponseSession { + state: Arc>, + idle: Arc, + group: Option>>, + budget: Option>, +} + +#[derive(Debug)] +struct SessionState { + latest: Option>, + active: bool, + closed: bool, + max_items: NonZeroUsize, + max_bytes: NonZeroUsize, +} + +/// A bounded set of independent serial sessions sharing cached parent lookup. +/// +/// The caller owns routing and scheduling. Create one member per logical session +/// and keep it while idle; queue disposal must not discard its retained state. +/// The group caps sessions created over its entire lifetime, not only active +/// sessions. Every member inherits the same item and serialized-byte budgets. +/// An aggregate serialized-byte budget covers cached checkpoints, pinned parents +/// and prepared replacements (including those awaiting durable persistence). +/// Sharing the same immutable parent counts it once. Replacement needs headroom +/// for old and new state until publication succeeds; no credit is granted for a +/// still-live parent. Executor input copies and temporary construction allocations +/// are outside this retention budget. Callers must separately bound active work; +/// this is not a heap-memory limit. +/// +/// Dropping the group invalidates every member even if its handle survives. +/// Callers must still cancel and join active execution tasks. +#[derive(Debug)] +pub struct ResponseSessionGroup { + state: Arc>, + max_sessions: NonZeroUsize, + max_items: NonZeroUsize, + max_bytes: NonZeroUsize, + budget: Arc, +} + +#[derive(Debug)] +struct SessionGroupState { + closed: bool, + members: Vec>>, +} + +impl ResponseSessionGroup { + #[must_use] + pub fn new( + max_sessions: NonZeroUsize, + max_items: NonZeroUsize, + max_bytes: NonZeroUsize, + max_retained_bytes: NonZeroUsize, + ) -> Self { + Self { + state: Arc::new(Mutex::new(SessionGroupState { + closed: false, + members: Vec::new(), + })), + max_sessions, + max_items, + max_bytes, + budget: Arc::new(CheckpointBudget { + used: AtomicUsize::new(0), + limit: max_retained_bytes, + }), + } + } + + /// Create another independent member without retaining its handle here. + /// + /// # Errors + /// + /// Returns an invalid-request error at the lifetime session cap or after + /// closure, and a stream error if the group state lock is poisoned. + pub fn new_session(&self) -> ExecutorResult { + let mut state = lock_group(&self.state)?; + if state.closed || state.members.len() >= self.max_sessions.get() { + return Err(ExecutorError::InvalidRequest( + "response session group is closed or has reached its session limit".to_owned(), + )); + } + let mut session = ResponseSession::new(self.max_items, self.max_bytes); + session.group = Some(Arc::clone(&self.state)); + session.budget = Some(Arc::clone(&self.budget)); + state.members.push(Arc::downgrade(&session.state)); + Ok(session) + } +} + +impl Drop for ResponseSessionGroup { + fn drop(&mut self) { + let members = { + let mut state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + std::mem::take(&mut state.members) + }; + // Never hold the group lock while closing members. A begin operation + // takes the group lock before looking at any member, with no await. + for member in members { + if let Some(member) = member.upgrade() { + let mut state = member.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + state.latest = None; + } + } + } +} + +impl ResponseSession { + #[must_use] + pub fn new(max_items: NonZeroUsize, max_bytes: NonZeroUsize) -> Self { + Self { + state: Arc::new(Mutex::new(SessionState { + latest: None, + active: false, + closed: false, + max_items, + max_bytes, + })), + idle: Arc::new(Notify::new()), + group: None, + budget: None, + } + } + + /// Wait until the active turn has published or dropped its continuation. + /// + /// This is not cancellation: callers must first drop the execution stream + /// or otherwise stop its worker. It prevents a subsequent serial request + /// racing the asynchronous disposal triggered by dropping that stream. + /// The ending lease releases its parent reference before becoming idle; + /// other sessions may still keep that shared checkpoint alive and charged. + /// + /// # Errors + /// + /// Returns a stream error if the session state lock is poisoned. + pub async fn wait_until_idle(&self) -> ExecutorResult<()> { + loop { + let notified = self.idle.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if !lock_state(&self.state)?.active { + return Ok(()); + } + notified.await; + } + } + + /// Discard a referenced checkpoint when a request fails before execution. + /// + /// An unrelated checkpoint is preserved. This does not cancel an active turn + /// or release its execution slot; the caller still owns that turn's lifetime. + /// + /// # Errors + /// + /// Returns a stream error if the session state lock is poisoned. + pub fn discard_cached_response(&self, response_id: &str) -> ExecutorResult<()> { + let mut state = lock_state(&self.state)?; + if state + .latest + .as_ref() + .is_some_and(|latest| latest.response_id == response_id) + { + state.latest = None; + } + Ok(()) + } + + pub(crate) fn begin(&self, previous_id: Option<&str>) -> ExecutorResult { + let group = self.group.as_ref().map(|group| lock_group(group)).transpose()?; + if group.as_ref().is_some_and(|group| group.closed) { + return Err(ExecutorError::InvalidRequest( + "response session group has closed".to_owned(), + )); + } + let mut state = lock_state(&self.state)?; + if state.closed || state.active { + return Err(ExecutorError::InvalidRequest( + "response session is closed or already executing a turn".to_owned(), + )); + } + let parent = state + .latest + .as_ref() + .filter(|entry| Some(entry.response_id.as_str()) == previous_id) + .cloned(); + state.active = true; + let mut continuation = ResponseContinuation { + state: Arc::clone(&self.state), + idle: Arc::clone(&self.idle), + parent, + budget: self.budget.clone(), + history_replaced: false, + recorded_output_count: 0, + finished: false, + }; + drop(state); + // Resolve a fork only when execution begins, not when its request queues. + // Never hold two member locks. Once cloned, the immutable parent remains + // pinned even if its source session advances or fails. + if let (None, Some(previous_id), Some(group)) = (continuation.parent.as_ref(), previous_id, group.as_ref()) { + for member in group.members.iter().filter_map(Weak::upgrade) { + if Arc::ptr_eq(&member, &self.state) { + continue; + } + let source = lock_state(&member)?; + if source.closed { + continue; + } + if let Some(parent) = source + .latest + .as_ref() + .filter(|parent| parent.response_id == previous_id) + { + continuation.parent = Some(Arc::clone(parent)); + break; + } + } + } + Ok(continuation) + } +} + +impl Drop for ResponseSession { + fn drop(&mut self) { + // Recover a poisoned guard only for disposal, never for reuse. + let mut state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + state.closed = true; + state.latest = None; + } +} + +/// An active turn's coherent parent snapshot and completion lease. +/// +/// Constructed by the executor, not serialized into split-execution contexts. +/// Dropping an unfinished lease evicts its referenced cached parent and releases +/// the serial execution slot, including on rehydration, tool or stream errors. +#[derive(Debug)] +pub struct ResponseContinuation { + state: Arc>, + idle: Arc, + pub(crate) parent: Option>, + budget: Option>, + history_replaced: bool, + recorded_output_count: usize, + finished: bool, +} + +#[derive(Serialize)] +pub(crate) struct ResponseCheckpoint { + pub response_id: String, + pub conversation_id: Option, + pub history: Vec, + pub metadata: ResponseMetadata, + pub durable: bool, +} + +impl fmt::Debug for ResponseCheckpoint { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ResponseCheckpoint") + .field("items", &self.history.len()) + .finish_non_exhaustive() + } +} + +/// An immutable checkpoint paired with its lifetime charge. Keeping the fields +/// private and omitting `DerefMut` prevents post-reservation size changes. +#[derive(Debug, Serialize)] +#[serde(transparent)] +pub(crate) struct RetainedCheckpoint { + checkpoint: ResponseCheckpoint, + #[serde(skip)] + _reservation: Option, +} + +impl Deref for RetainedCheckpoint { + type Target = ResponseCheckpoint; + + fn deref(&self) -> &Self::Target { + &self.checkpoint + } +} + +#[derive(Debug)] +struct CheckpointBudget { + used: AtomicUsize, + limit: NonZeroUsize, +} + +impl CheckpointBudget { + fn reserve(self: &Arc, bytes: usize) -> ExecutorResult { + self.used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |used| { + used.checked_add(bytes).filter(|total| *total <= self.limit.get()) + }) + .map_err(|_| aggregate_budget_error())?; + Ok(CheckpointReservation { + budget: Arc::clone(self), + bytes, + }) + } +} + +#[derive(Debug)] +struct CheckpointReservation { + budget: Arc, + bytes: usize, +} + +impl Drop for CheckpointReservation { + fn drop(&mut self) { + let previous = self.budget.used.fetch_sub(self.bytes, Ordering::AcqRel); + debug_assert!(previous >= self.bytes, "checkpoint charge released exactly once"); + } +} + +fn aggregate_budget_error() -> ExecutorError { + ExecutorError::PayloadTooLarge( + "response continuation exceeds the aggregate retained checkpoint budget; release inactive state or replay a compacted input window" + .to_owned(), + ) +} + +/// Restore only the current canonical window, while preserving orchestration +/// records that the model-facing projection would strip. Durable response chains +/// can still reference superseded parent rows; those must not consume the session +/// checkpoint budget or introduce obsolete pending calls after reconnecting. +pub(super) fn canonical_session_history(history: Vec) -> Vec { + let Some(window) = latest_compaction_window(&history) else { + return history; + }; + history + .into_iter() + .enumerate() + .filter(|(index, item)| { + *index >= window.latest_index() + || window.retains_user_item(*index, item) + || matches!(item, InputItem::McpListTools(_)) + }) + .map(|(_, item)| item) + .collect() +} + +impl ResponseContinuation { + /// The executor has replaced this turn's input with a canonical compacted window. + pub(crate) fn mark_history_replaced(&mut self) { + self.history_replaced = true; + } + + /// The loop has recorded these outputs in canonical inference-round order. + /// Public output remains complete, but persistence must not append it again. + pub(crate) fn mark_outputs_recorded(&mut self, output_count: usize) { + self.recorded_output_count = output_count; + } + + pub(crate) fn retains_output(&self, index: usize, item: &OutputItem) -> bool { + index >= self.recorded_output_count || matches!(item, OutputItem::McpListTools(_)) + } + + /// Retain orchestration records across compaction without restoring the old + /// model context. MCP discovery still needs to know which servers were listed. + pub(crate) fn parent_items(&self) -> impl Iterator { + self.parent + .iter() + .flat_map(|parent| parent.history.iter()) + .filter(|item| !self.history_replaced || matches!(item, InputItem::McpListTools(_))) + } + + pub(crate) fn checkpoint( + &self, + response_id: String, + conversation_id: Option, + metadata: &ResponseMetadata, + new_items: &[InOutItem], + durable: bool, + ) -> ExecutorResult { + let mut history = self.parent_items().cloned().collect::>(); + history.extend(InOutItem::into_input_items(new_items.to_vec())); + let mut metadata = metadata.clone(); + if let Some(tools) = metadata.effective_tools.as_mut() { + for tool in tools { + tool.sanitize_for_persistence(); + } + } + let checkpoint = ResponseCheckpoint { + response_id, + conversation_id, + history: canonical_session_history(history), + metadata, + durable, + }; + let (max_items, max_bytes) = { + let state = lock_state(&self.state)?; + (state.max_items.get(), state.max_bytes.get()) + }; + let bytes = if checkpoint.history.len() <= max_items { + serialized_size_up_to(&checkpoint, max_bytes)? + } else { + None + }; + let Some(bytes) = bytes else { + return Err(ExecutorError::PayloadTooLarge( + "response continuation exceeds the session checkpoint budget; replay a compacted input window" + .to_owned(), + )); + }; + self.retain(checkpoint, bytes) + } + + /// A durable fallback becomes a live pinned parent before inference. It must + /// share the aggregate budget instead of bypassing it through storage. + pub(crate) fn retain_parent(&self, checkpoint: ResponseCheckpoint) -> ExecutorResult { + let bytes = if let Some(budget) = &self.budget { + serialized_size_up_to(&checkpoint, budget.limit.get())?.ok_or_else(aggregate_budget_error)? + } else { + 0 // Standalone sessions retain their existing per-completion policy. + }; + self.retain(checkpoint, bytes) + } + + fn retain(&self, checkpoint: ResponseCheckpoint, bytes: usize) -> ExecutorResult { + let reservation = self.budget.as_ref().map(|budget| budget.reserve(bytes)).transpose()?; + Ok(RetainedCheckpoint { + checkpoint, + _reservation: reservation, + }) + } + + pub(crate) fn publish(mut self, checkpoint: RetainedCheckpoint) -> ExecutorResult<()> { + { + let mut state = lock_state(&self.state)?; + if state.closed { + return Err(ExecutorError::InvalidRequest("response session has closed".to_owned())); + } + state.latest = Some(Arc::new(checkpoint)); + // Return this lease's parent charge before another turn can observe + // an idle slot, including a waiter resumed by notify_waiters below. + self.parent = None; + state.active = false; + } + self.finished = true; + self.idle.notify_waiters(); + Ok(()) + } +} + +impl Drop for ResponseContinuation { + fn drop(&mut self) { + if self.finished { + return; + } + let mut state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner); + if self.parent.as_ref().is_some_and(|parent| { + state + .latest + .as_ref() + .is_some_and(|latest| latest.response_id == parent.response_id) + }) { + state.latest = None; + } + // Field destructors run after Drop returns, too late for a waiter that + // can immediately reuse the slot and reserve the released capacity. + self.parent = None; + state.active = false; + drop(state); + self.idle.notify_waiters(); + } +} + +fn lock_state(state: &Mutex) -> ExecutorResult> { + state + .lock() + .map_err(|_| ExecutorError::StreamError("response session state is unavailable".to_owned())) +} + +fn lock_group(state: &Mutex) -> ExecutorResult> { + state + .lock() + .map_err(|_| ExecutorError::StreamError("response session group state is unavailable".to_owned())) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::io::ResponsesInput; + use serde_json::json; + + fn session(items: usize, bytes: usize) -> ResponseSession { + ResponseSession::new(NonZeroUsize::new(items).unwrap(), NonZeroUsize::new(bytes).unwrap()) + } + + fn new_checkpoint(lease: &ResponseContinuation, id: &str) -> RetainedCheckpoint { + let input = ResponsesInput::Text("private prompt".to_owned()); + let items = Vec::from(&input).into_iter().map(InOutItem::Input).collect::>(); + lease + .checkpoint(id.to_owned(), None, &ResponseMetadata::default(), &items, false) + .unwrap() + } + + fn complete(session: &ResponseSession, id: &str) { + let lease = session.begin(None).unwrap(); + let checkpoint = new_checkpoint(&lease, id); + lease.publish(checkpoint).unwrap(); + } + + fn group(count: usize) -> ResponseSessionGroup { + ResponseSessionGroup::new( + NonZeroUsize::new(count).unwrap(), + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(10_000).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + ) + } + + #[test] + fn group_session_cap_counts_idle_and_dropped_members() { + let group = group(2); + let first = group.new_session().unwrap(); + let second = group.new_session().unwrap(); + complete(&first, "resp_1"); + assert!(group.new_session().is_err()); + drop(second); + assert!(group.new_session().is_err(), "the cap covers the entire group lifetime"); + assert!(first.begin(Some("resp_1")).unwrap().parent.is_some()); + } + + #[test] + fn grouped_members_inherit_item_and_byte_budgets() { + for (items, bytes) in [(1, 10_000), (10, 1)] { + let group = ResponseSessionGroup::new( + NonZeroUsize::new(1).unwrap(), + NonZeroUsize::new(items).unwrap(), + NonZeroUsize::new(bytes).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + ); + let member = group.new_session().unwrap(); + let lease = member.begin(None).unwrap(); + let input = Vec::from(&ResponsesInput::Text("private prompt".to_owned())); + let input = input + .iter() + .cycle() + .take(2) + .cloned() + .map(InOutItem::Input) + .collect::>(); + assert!(matches!( + lease.checkpoint("resp_1".to_owned(), None, &ResponseMetadata::default(), &input, false), + Err(ExecutorError::PayloadTooLarge(_)) + )); + } + } + + #[test] + fn groups_do_not_share_parent_lookup() { + let first_group = group(1); + let second_group = group(1); + let source = first_group.new_session().unwrap(); + let other = second_group.new_session().unwrap(); + complete(&source, "resp_1"); + assert!(other.begin(Some("resp_1")).unwrap().parent.is_none()); + } + + #[test] + fn failed_fork_preserves_source_and_unrelated_destination_checkpoint() { + let group = group(2); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_source"); + complete(&target, "resp_target"); + let fork = target.begin(Some("resp_source")).unwrap(); + assert!(fork.parent.is_some()); + drop(fork); + assert!(source.begin(Some("resp_source")).unwrap().parent.is_some()); + assert!(target.begin(Some("resp_target")).unwrap().parent.is_some()); + } + + #[test] + fn fork_cannot_find_parent_evicted_before_execution_starts() { + let group = group(2); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_source"); + drop(source.begin(Some("resp_source")).unwrap()); + assert!(target.begin(Some("resp_source")).unwrap().parent.is_none()); + } + + #[test] + fn dropped_group_invalidates_live_members_and_rejects_late_fork_publication() { + let group = group(2); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_source"); + let fork = target.begin(Some("resp_source")).unwrap(); + let weak = Arc::downgrade(fork.parent.as_ref().unwrap()); + let checkpoint = new_checkpoint(&fork, "resp_fork"); + drop(group); + assert!(source.begin(None).is_err()); + assert!(weak.upgrade().is_some(), "the active fork still pins its parent"); + assert!(fork.publish(checkpoint).is_err()); + assert!(weak.upgrade().is_none()); + assert!(target.begin(None).is_err()); + } + + #[test] + fn group_does_not_keep_dropped_members_or_their_checkpoints_alive() { + let group = group(2); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_source"); + let member = Arc::downgrade(&source.state); + let checkpoint = Arc::downgrade(source.state.lock().unwrap().latest.as_ref().unwrap()); + drop(source); + assert!(member.upgrade().is_none()); + assert!(checkpoint.upgrade().is_none()); + assert!(target.begin(Some("resp_source")).unwrap().parent.is_none()); + } + + #[test] + fn session_rejects_concurrent_turns_and_releases_cancelled_lease() { + let session = session(10, 10_000); + let lease = session.begin(None).unwrap(); + assert!(session.begin(None).is_err()); + drop(lease); + assert!(session.begin(None).is_ok()); + } + + #[tokio::test] + async fn idle_wait_is_not_cancellation_and_is_notified_by_lease_disposal() { + let session = session(10, 10_000); + session.wait_until_idle().await.unwrap(); + let lease = session.begin(None).unwrap(); + let waiting = session.wait_until_idle(); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + assert!(session.begin(None).is_err()); + drop(lease); + waiting.await.unwrap(); + assert!(session.begin(None).is_ok()); + } + + #[tokio::test] + async fn successful_publication_wakes_all_idle_waiters() { + let session = session(10, 10_000); + let lease = session.begin(None).unwrap(); + let checkpoint = new_checkpoint(&lease, "resp_1"); + let first = session.wait_until_idle(); + let second = session.wait_until_idle(); + tokio::pin!(first, second); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + lease.publish(checkpoint).unwrap(); + first.await.unwrap(); + second.await.unwrap(); + assert!(session.begin(Some("resp_1")).unwrap().parent.is_some()); + } + + #[test] + fn owner_drop_releases_checkpoint_and_rejects_late_publication() { + let session = session(10, 10_000); + complete(&session, "resp_1"); + let lease = session.begin(Some("resp_1")).unwrap(); + let weak = Arc::downgrade(lease.parent.as_ref().unwrap()); + let checkpoint = new_checkpoint(&lease, "resp_2"); + drop(session); + assert!(lease.publish(checkpoint).is_err()); + assert!(weak.upgrade().is_none()); + } + + #[test] + fn successful_replacement_does_not_keep_ancestor_objects() { + let session = session(10, 10_000); + complete(&session, "resp_1"); + let lease = session.begin(Some("resp_1")).unwrap(); + let weak = Arc::downgrade(lease.parent.as_ref().unwrap()); + let checkpoint = new_checkpoint(&lease, "resp_2"); + assert_eq!(checkpoint.history.len(), 2); + lease.publish(checkpoint).unwrap(); + assert!(weak.upgrade().is_none()); + let lease = session.begin(Some("resp_2")).unwrap(); + assert_eq!(lease.parent.as_ref().unwrap().history.len(), 2); + } + + #[test] + fn failed_continuation_evicts_parent_but_unknown_id_does_not() { + let session = session(10, 10_000); + complete(&session, "resp_1"); + drop(session.begin(Some("unknown")).unwrap()); + let lease = session.begin(Some("resp_1")).unwrap(); + assert!(lease.parent.is_some()); + drop(lease); + assert!(session.begin(Some("resp_1")).unwrap().parent.is_none()); + } + + #[test] + fn separate_sessions_have_no_shared_lookup() { + let first = session(10, 10_000); + let second = session(10, 10_000); + complete(&first, "resp_1"); + assert!(second.begin(Some("resp_1")).unwrap().parent.is_none()); + } + + #[test] + fn explicit_discard_preserves_the_active_snapshot_and_execution_slot() { + let session = session(10, 10_000); + complete(&session, "resp_1"); + session.discard_cached_response("unrelated").unwrap(); + let lease = session.begin(Some("resp_1")).unwrap(); + assert!(lease.parent.is_some()); + session.discard_cached_response("resp_1").unwrap(); + assert!(lease.parent.is_some(), "the active snapshot remains pinned"); + assert!(session.begin(None).is_err(), "discard is not cancellation"); + drop(lease); + assert!(session.begin(Some("resp_1")).unwrap().parent.is_none()); + } + + #[test] + fn item_budget_is_enforced_on_the_complete_history() { + let session = session(1, 10_000); + complete(&session, "resp_1"); + let lease = session.begin(Some("resp_1")).unwrap(); + let input = serde_json::from_value(json!({"type":"message", "role":"user", "content":"next"})).unwrap(); + assert!(matches!( + lease.checkpoint( + "resp_2".to_owned(), + None, + &ResponseMetadata::default(), + &[InOutItem::Input(input)], + false + ), + Err(ExecutorError::PayloadTooLarge(_)) + )); + } + + #[test] + fn byte_budget_includes_metadata_and_exact_boundary() { + let source = session(10, 10_000); + let lease = source.begin(None).unwrap(); + let checkpoint = new_checkpoint(&lease, "resp_1"); + let size = serialized_size_up_to(&checkpoint, usize::MAX).unwrap().unwrap(); + let exact = session(10, size); + complete(&exact, "resp_1"); + let too_small = session(10, size - 1); + let lease = too_small.begin(None).unwrap(); + let metadata = ResponseMetadata { + effective_instructions: Some("large metadata".repeat(100)), + ..ResponseMetadata::default() + }; + assert!(matches!( + lease.checkpoint("resp_1".to_owned(), None, &metadata, &[], false), + Err(ExecutorError::PayloadTooLarge(_)) + )); + } + + #[test] + fn canonical_items_and_sanitized_mcp_metadata_are_retained() { + let session = session(10, 10_000); + let lease = session.begin(None).unwrap(); + let metadata = ResponseMetadata { + effective_tools: Some( + serde_json::from_value(json!([{ + "type":"mcp", "server_label":"counter", "server_url":"https://example.com/mcp", + "headers":{"X-API-Key":"secret-header"}, "authorization":"secret-token", "require_approval":"never" + }])) + .unwrap(), + ), + ..ResponseMetadata::default() + }; + let input = serde_json::from_value(json!({"type":"reasoning", "id":"rs_1", + "content":[{"type":"reasoning_text", "text":"part one"},{"type":"reasoning_text", "text":"part two"}], + "summary":[{"type":"summary_text", "text":"canonical summary"}], "encrypted_content":"opaque" + })) + .unwrap(); + let checkpoint = lease + .checkpoint("resp_1".to_owned(), None, &metadata, &[InOutItem::Input(input)], false) + .unwrap(); + let serialized = serde_json::to_value(&checkpoint).unwrap(); + assert_eq!(serialized["history"][0]["content"].as_array().unwrap().len(), 2); + assert_eq!(serialized["history"][0]["encrypted_content"], "opaque"); + assert!(!serialized.to_string().contains("secret-header")); + assert!(!serialized.to_string().contains("secret-token")); + assert!(!format!("{checkpoint:?}").contains("canonical summary")); + } + + #[test] + fn compaction_replacement_preserves_mcp_discovery_without_old_context() { + let session = session(3, 10_000); + let lease = session.begin(None).unwrap(); + let items = serde_json::from_value::>(json!([ + {"type":"message", "role":"user", "content":"private prompt"}, + {"type":"mcp_list_tools", "id":"mcp_1", "server_label":"counter", "tools":[]} + ])) + .unwrap() + .into_iter() + .map(InOutItem::Input) + .collect::>(); + let parent = lease + .checkpoint("resp_1".to_owned(), None, &ResponseMetadata::default(), &items, false) + .unwrap(); + lease.publish(parent).unwrap(); + let mut lease = session.begin(Some("resp_1")).unwrap(); + lease.mark_history_replaced(); + let compacted: Vec = serde_json::from_value(json!([ + {"type":"message", "role":"user", "id":"msg_kept", "status":"completed", "content":"retained user"}, + {"type":"compaction", "id":"cmp_1", "encrypted_content":"canonical summary"} + ])) + .unwrap(); + let items = compacted.into_iter().map(InOutItem::Input).collect::>(); + let checkpoint = lease + .checkpoint("resp_2".to_owned(), None, &ResponseMetadata::default(), &items, false) + .unwrap(); + assert_eq!(checkpoint.history.len(), 3); + assert!(matches!(checkpoint.history[0], InputItem::McpListTools(_))); + assert!(matches!(checkpoint.history[2], InputItem::Compaction(_))); + assert!( + !serde_json::to_string(&checkpoint.history) + .unwrap() + .contains("private prompt") + ); + } + + #[test] + fn recorded_rounds_are_not_duplicated_and_mcp_discovery_is_retained() { + let session = session(10, 10_000); + let mut lease = session.begin(None).unwrap(); + let message: OutputItem = serde_json::from_value(json!({ + "type":"message", "id":"msg_1", "role":"assistant", "status":"completed", "content":[] + })) + .unwrap(); + let discovery: OutputItem = serde_json::from_value(json!({ + "type":"mcp_list_tools", "id":"mcp_1", "server_label":"counter", "tools":[] + })) + .unwrap(); + assert!(lease.retains_output(0, &message)); + lease.mark_outputs_recorded(2); + assert!(!lease.retains_output(1, &message)); + assert!(lease.retains_output(2, &message)); + lease.mark_history_replaced(); + lease.mark_outputs_recorded(4); + assert!(!lease.retains_output(2, &message)); + assert!(lease.retains_output(4, &message)); + assert!(lease.retains_output(0, &discovery)); + } +} diff --git a/crates/agentic-server-core/src/executor/session_budget_tests.rs b/crates/agentic-server-core/src/executor/session_budget_tests.rs new file mode 100644 index 00000000..ed672cb8 --- /dev/null +++ b/crates/agentic-server-core/src/executor/session_budget_tests.rs @@ -0,0 +1,349 @@ +//! Aggregate retention is charged per immutable object, not per Arc reference. + +use super::*; +use crate::types::io::ResponsesInput; +use std::sync::atomic::AtomicBool; +use std::task::{Context, Poll, Wake, Waker}; + +/// Observe the handoff synchronously, as a newly scheduled waiter could on a +/// different executor thread. No sleeps or probabilistic task ordering are used. +struct HandoffObserver { + budget: Arc, + reservation_bytes: usize, + used_on_wake: AtomicUsize, + reserved_on_wake: AtomicBool, + wake_count: AtomicUsize, +} + +impl HandoffObserver { + fn new(group: &ResponseSessionGroup, reservation_bytes: usize) -> Arc { + Arc::new(Self { + budget: Arc::clone(&group.budget), + reservation_bytes, + used_on_wake: AtomicUsize::new(usize::MAX), + reserved_on_wake: AtomicBool::new(false), + wake_count: AtomicUsize::new(0), + }) + } +} + +impl Wake for HandoffObserver { + fn wake(self: Arc) { + self.used_on_wake + .store(self.budget.used.load(Ordering::Acquire), Ordering::Release); + let reservation = self.budget.reserve(self.reservation_bytes); + self.reserved_on_wake.store(reservation.is_ok(), Ordering::Release); + self.wake_count.fetch_add(1, Ordering::AcqRel); + } +} + +fn group(limit: usize) -> ResponseSessionGroup { + ResponseSessionGroup::new( + NonZeroUsize::new(32).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(limit).unwrap(), + ) +} + +fn capture(lease: &ResponseContinuation, id: &str) -> ExecutorResult { + let items = Vec::from(&ResponsesInput::Text("retained prompt".to_owned())) + .into_iter() + .map(InOutItem::Input) + .collect::>(); + lease.checkpoint(id.to_owned(), None, &ResponseMetadata::default(), &items, false) +} + +fn size(checkpoint: &RetainedCheckpoint) -> usize { + serialized_size_up_to(checkpoint, usize::MAX).unwrap().unwrap() +} + +fn root_size() -> usize { + let session = ResponseSession::new(NonZeroUsize::new(100).unwrap(), NonZeroUsize::new(100_000).unwrap()); + size(&capture(&session.begin(None).unwrap(), "resp_1").unwrap()) +} + +fn used(group: &ResponseSessionGroup) -> usize { + group.budget.used.load(Ordering::Acquire) +} + +fn complete(session: &ResponseSession, id: &str) { + let lease = session.begin(None).unwrap(); + let checkpoint = capture(&lease, id).unwrap(); + lease.publish(checkpoint).unwrap(); +} + +#[test] +fn aggregate_budget_handoff_cancellation_releases_parent_before_waking_waiters() { + let bytes = root_size(); + let group = group(bytes); + let session = group.new_session().unwrap(); + complete(&session, "resp_1"); + let lease = session.begin(Some("resp_1")).unwrap(); + let observer = HandoffObserver::new(&group, bytes); + let waker = Waker::from(Arc::clone(&observer)); + let mut context = Context::from_waker(&waker); + let mut idle = Box::pin(session.wait_until_idle()); + assert!(idle.as_mut().poll(&mut context).is_pending()); + + drop(lease); + + assert_eq!(observer.wake_count.load(Ordering::Acquire), 1); + assert_eq!(observer.used_on_wake.load(Ordering::Acquire), 0); + assert!(observer.reserved_on_wake.load(Ordering::Acquire)); + assert!(matches!(idle.as_mut().poll(&mut context), Poll::Ready(Ok(())))); + complete(&session, "resp_2"); + assert_eq!(used(&group), bytes); +} + +#[test] +fn aggregate_budget_handoff_publication_releases_replaced_parent_before_waking_waiters() { + let bytes = root_size(); + let group = group(3 * bytes); + let session = group.new_session().unwrap(); + complete(&session, "resp_1"); + let lease = session.begin(Some("resp_1")).unwrap(); + let checkpoint = capture(&lease, "resp_2").unwrap(); + let checkpoint_bytes = size(&checkpoint); + let observer = HandoffObserver::new(&group, 3 * bytes - checkpoint_bytes); + let waker = Waker::from(Arc::clone(&observer)); + let mut context = Context::from_waker(&waker); + let mut idle = Box::pin(session.wait_until_idle()); + assert!(idle.as_mut().poll(&mut context).is_pending()); + + lease.publish(checkpoint).unwrap(); + + assert_eq!(observer.wake_count.load(Ordering::Acquire), 1); + assert_eq!(observer.used_on_wake.load(Ordering::Acquire), checkpoint_bytes); + assert!(observer.reserved_on_wake.load(Ordering::Acquire)); + assert!(matches!(idle.as_mut().poll(&mut context), Poll::Ready(Ok(())))); + assert!(session.begin(Some("resp_2")).unwrap().parent.is_some()); +} + +#[test] +fn aggregate_budget_handoff_keeps_a_shared_parent_charged() { + let bytes = root_size(); + let group = group(bytes); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_1"); + let fork = target.begin(Some("resp_1")).unwrap(); + let observer = HandoffObserver::new(&group, 1); + let waker = Waker::from(Arc::clone(&observer)); + let mut context = Context::from_waker(&waker); + let mut idle = Box::pin(target.wait_until_idle()); + assert!(idle.as_mut().poll(&mut context).is_pending()); + + drop(fork); + + assert_eq!(observer.wake_count.load(Ordering::Acquire), 1); + assert_eq!(observer.used_on_wake.load(Ordering::Acquire), bytes); + assert!(!observer.reserved_on_wake.load(Ordering::Acquire)); + assert!(matches!(idle.as_mut().poll(&mut context), Poll::Ready(Ok(())))); + assert!(source.begin(Some("resp_1")).unwrap().parent.is_some()); +} + +#[test] +fn aggregate_budget_reserves_before_publication_and_releases_abandoned_candidates() { + let bytes = root_size(); + let group = group(bytes); + let first = group.new_session().unwrap(); + let second = group.new_session().unwrap(); + let first_lease = first.begin(None).unwrap(); + let checkpoint = capture(&first_lease, "resp_1").unwrap(); + assert_eq!(used(&group), bytes); + assert!(first.state.lock().unwrap().latest.is_none(), "not yet published"); + let second_lease = second.begin(None).unwrap(); + assert!(matches!( + capture(&second_lease, "resp_2"), + Err(ExecutorError::PayloadTooLarge(_)) + )); + assert_eq!(used(&group), bytes, "failed reservation must not change the counter"); + drop(checkpoint); + assert_eq!(used(&group), 0); + let checkpoint = capture(&second_lease, "resp_2").unwrap(); + second_lease.publish(checkpoint).unwrap(); + assert_eq!(used(&group), bytes); + drop(second); + assert_eq!(used(&group), 0); +} + +#[test] +fn aggregate_budget_counts_shared_parent_once_and_replaced_pinned_parent_until_last_release() { + let bytes = root_size(); + let group = group(3 * bytes); + let source = group.new_session().unwrap(); + let first = group.new_session().unwrap(); + let second = group.new_session().unwrap(); + complete(&source, "resp_1"); + let fork = first.begin(Some("resp_1")).unwrap(); + let other_fork = second.begin(Some("resp_1")).unwrap(); + assert!(Arc::ptr_eq( + fork.parent.as_ref().unwrap(), + other_fork.parent.as_ref().unwrap() + )); + assert_eq!(used(&group), bytes, "fork handles share the original charge"); + complete(&source, "resp_2"); + assert_eq!( + used(&group), + 2 * bytes, + "source replacement cannot release a pinned parent" + ); + drop(source); + assert_eq!(used(&group), bytes); + drop(fork); + assert_eq!(used(&group), bytes, "the second fork is still using the old parent"); + drop(other_fork); + assert_eq!(used(&group), 0); +} + +#[test] +fn aggregate_budget_replacement_requires_headroom_without_premature_eviction() { + let bytes = root_size(); + let group = group(bytes); + let session = group.new_session().unwrap(); + complete(&session, "resp_1"); + let lease = session.begin(None).unwrap(); + assert!(matches!( + capture(&lease, "resp_2"), + Err(ExecutorError::PayloadTooLarge(_)) + )); + assert_eq!(used(&group), bytes); + assert_eq!( + session.state.lock().unwrap().latest.as_ref().unwrap().response_id, + "resp_1" + ); + drop(lease); // Failed fresh root does not reference/evict the old checkpoint. + session.discard_cached_response("resp_1").unwrap(); + assert_eq!(used(&group), 0); + complete(&session, "resp_2"); + assert_eq!(used(&group), bytes); +} + +#[test] +fn aggregate_budget_successful_continuation_releases_replaced_parent() { + let group = group(100_000); + let session = group.new_session().unwrap(); + complete(&session, "resp_1"); + let before = used(&group); + let lease = session.begin(Some("resp_1")).unwrap(); + let checkpoint = capture(&lease, "resp_2").unwrap(); + let after = size(&checkpoint); + assert_eq!(used(&group), before + after, "both survive until publication"); + lease.publish(checkpoint).unwrap(); + assert_eq!(used(&group), after); + session.discard_cached_response("resp_2").unwrap(); + assert_eq!(used(&group), 0); +} + +#[test] +fn aggregate_budget_group_drop_keeps_only_live_parent_and_pending_candidate_charges() { + let group = group(100_000); + let budget = Arc::clone(&group.budget); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + complete(&source, "resp_1"); + let parent_size = used(&group); + let fork = target.begin(Some("resp_1")).unwrap(); + let checkpoint = capture(&fork, "resp_2").unwrap(); + let total = parent_size + size(&checkpoint); + drop(group); + assert_eq!(budget.used.load(Ordering::Acquire), total); + assert!(fork.publish(checkpoint).is_err()); + assert_eq!(budget.used.load(Ordering::Acquire), 0); +} + +#[test] +fn aggregate_budget_charges_canonical_durable_fallback_before_inference() { + let group = group(100_000); + let session = group.new_session().unwrap(); + let lease = session.begin(None).unwrap(); + let checkpoint = ResponseCheckpoint { + response_id: "durable".to_owned(), + conversation_id: None, + history: Vec::from(&ResponsesInput::Text("durable history".to_owned())), + metadata: ResponseMetadata { + effective_instructions: Some("metadata".repeat(10)), + ..ResponseMetadata::default() + }, + durable: true, + }; + let expected = serialized_size_up_to(&checkpoint, usize::MAX).unwrap().unwrap(); + let retained = lease.retain_parent(checkpoint).unwrap(); + assert_eq!(used(&group), expected); + assert_eq!( + size(&retained), + expected, + "budget metadata is not part of serialized history" + ); + drop(retained); + assert_eq!(used(&group), 0); +} + +#[test] +fn aggregate_budget_simultaneous_candidates_cannot_oversubscribe_or_leak() { + let bytes = root_size(); + let group = group(4 * bytes); + let sessions = (0..16).map(|_| group.new_session().unwrap()).collect::>(); + let barrier = std::sync::Barrier::new(sessions.len()); + let candidates = std::thread::scope(|scope| { + let tasks = sessions + .iter() + .map(|session| { + let barrier = &barrier; + scope.spawn(move || { + let lease = session.begin(None).unwrap(); + barrier.wait(); + capture(&lease, "resp_1") + }) + }) + .collect::>(); + tasks.into_iter().map(|task| task.join().unwrap()).collect::>() + }); + assert_eq!(candidates.iter().filter(|candidate| candidate.is_ok()).count(), 4); + assert!( + candidates + .iter() + .filter_map(|candidate| candidate.as_ref().err()) + .all(|error| matches!(error, ExecutorError::PayloadTooLarge(_))) + ); + assert_eq!(used(&group), 4 * bytes); + drop(candidates); + assert_eq!(used(&group), 0); + complete(&sessions[0], "resp_1"); + assert_eq!(used(&group), bytes); +} + +#[test] +fn aggregate_budget_integer_overflow_is_rejected_without_changing_usage() { + let group = group(usize::MAX); + let reservation = group.budget.reserve(usize::MAX).unwrap(); + assert!(group.budget.reserve(1).is_err()); + assert_eq!(used(&group), usize::MAX); + drop(reservation); + assert_eq!(used(&group), 0); +} + +#[tokio::test] +async fn aggregate_budget_cancellation_drops_pending_candidate_and_failed_parent() { + let group = group(100_000); + let session = Arc::new(group.new_session().unwrap()); + complete(&session, "resp_1"); + let worker_session = Arc::clone(&session); + let (prepared_tx, prepared_rx) = tokio::sync::oneshot::channel(); + let worker = tokio::spawn(async move { + let lease = worker_session.begin(Some("resp_1")).unwrap(); + let checkpoint = capture(&lease, "resp_2").unwrap(); + prepared_tx.send(()).unwrap(); + std::future::pending::<()>().await; + lease.publish(checkpoint).unwrap(); + }); + prepared_rx.await.unwrap(); + assert!(used(&group) > root_size()); + worker.abort(); + assert!(worker.await.unwrap_err().is_cancelled()); + session.wait_until_idle().await.unwrap(); + assert_eq!(used(&group), 0); + complete(&session, "resp_3"); + assert_eq!(used(&group), root_size()); +} diff --git a/crates/agentic-server-core/src/executor/upstream.rs b/crates/agentic-server-core/src/executor/upstream.rs index ada00418..7e903580 100644 --- a/crates/agentic-server-core/src/executor/upstream.rs +++ b/crates/agentic-server-core/src/executor/upstream.rs @@ -318,6 +318,7 @@ pub(super) mod tests { response_id: "resp_test".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/src/utils/common.rs b/crates/agentic-server-core/src/utils/common.rs index 42fd402c..8127183c 100644 --- a/crates/agentic-server-core/src/utils/common.rs +++ b/crates/agentic-server-core/src/utils/common.rs @@ -23,6 +23,44 @@ pub fn serialize_to_string(value: &T) -> Result(value: &T, limit: usize) -> Result, serde_json::Error> { + struct Counter { + used: usize, + limit: usize, + exceeded: bool, + } + impl std::io::Write for Counter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if bytes.len() > self.limit.saturating_sub(self.used) { + self.exceeded = true; + return Err(std::io::Error::other("serialized size limit exceeded")); + } + self.used += bytes.len(); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + let mut counter = Counter { + used: 0, + limit, + exceeded: false, + }; + let result = serde_json::to_writer(&mut counter, value); + if counter.exceeded { + return Ok(None); + } + result?; + Ok(Some(counter.used)) +} + /// Serialize any type to a `serde_json::Value`. /// /// # Errors diff --git a/crates/agentic-server-core/tests/relayed_stream_validation_test.rs b/crates/agentic-server-core/tests/relayed_stream_validation_test.rs index b29afc18..55cf890f 100644 --- a/crates/agentic-server-core/tests/relayed_stream_validation_test.rs +++ b/crates/agentic-server-core/tests/relayed_stream_validation_test.rs @@ -23,6 +23,7 @@ fn request_context() -> RequestContext { response_id: "resp_reserved".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/tests/session_continuation_test.rs b/crates/agentic-server-core/tests/session_continuation_test.rs new file mode 100644 index 00000000..ead92d56 --- /dev/null +++ b/crates/agentic-server-core/tests/session_continuation_test.rs @@ -0,0 +1,1516 @@ +//! Public core continuation/commit and lifecycle seams. +//! Tests use disabled storage or isolated SQLite and local protocol stubs, never a live model. +//! Set `AGENTIC_SESSION_TEST_POSTGRES_URL` to an isolated PostgreSQL database with +//! CREATE DATABASE privileges to rerun persistence cases using fresh databases. + +use std::collections::VecDeque; +use std::future::Future; +use std::num::NonZeroUsize; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use agentic_core::executor::{ + ConversationHandler, ExecuteRequest, ExecutionContext, ExecutorError, RequestContext, ResponseHandler, + ResponseSession, ResponseSessionGroup, commit, rehydrate_in_session, +}; +use agentic_core::storage::{ + ConversationStore, InOutItem, ResponseMetadata, ResponseStore, StorageError, create_pool_with_schema, +}; +use agentic_core::tool::{GatewayExecutor, ToolError, ToolHandler, ToolOutput, ToolType}; +use agentic_core::types::io::FunctionTool; +use agentic_core::types::io::output::{FunctionToolCall, GatewayCallStatus, OutputItem, WebSearchCall}; +use agentic_core::types::request_response::{RequestPayload, ResponsePayload}; +use agentic_core::types::tools::WebSearchToolParam; +use axum::{Json, Router, routing::post}; +use serde_json::{Value, json}; +use tokio::sync::Mutex; + +#[path = "support/media_session_checks.rs"] +mod media_session_checks; + +struct LocalSearch { + calls: Arc, +} + +impl ToolHandler for LocalSearch { + type ToolParams = WebSearchToolParam; + + fn tool_type(&self) -> ToolType { + ToolType::WebSearch + } + + fn validate(&self, _params: &WebSearchToolParam) -> Result<(), ToolError> { + Ok(()) + } + + fn normalize(&self, _params: &WebSearchToolParam) -> Vec { + vec![FunctionTool { + type_: "function".to_owned(), + name: "web_search".to_owned(), + description: None, + parameters: Some(json!({"type":"object", "properties":{"query":{"type":"string"}}})), + strict: Some(false), + }] + } +} + +impl GatewayExecutor for LocalSearch { + type ExecutionParams = WebSearchToolParam; + + fn execute( + &self, + call_id: &str, + _tool_name: &str, + _arguments: &str, + _params: &WebSearchToolParam, + ) -> Pin> + Send + '_>> { + self.calls.fetch_add(1, Ordering::Relaxed); + let call_id = call_id.to_owned(); + Box::pin(async move { + Ok(ToolOutput { + call_id, + output: "local search result".to_owned(), + }) + }) + } + + fn public_output( + &self, + call: &FunctionToolCall, + _output: &ToolOutput, + status: GatewayCallStatus, + _params: &WebSearchToolParam, + ) -> Option { + Some(OutputItem::WebSearchCall( + WebSearchCall::try_new(&call.id, status, vec!["local".to_owned()], Vec::new()).unwrap(), + )) + } +} + +struct LocalModel { + exec: ExecutionContext, + requests: Arc>>, + handle: Option>, +} + +impl LocalModel { + async fn start(responses: Vec) -> Self { + let responses = Arc::new(Mutex::new(VecDeque::from(responses))); + let requests = Arc::new(Mutex::new(Vec::new())); + let captured = Arc::clone(&requests); + let app = Router::new().route( + "/v1/responses", + post(move |Json(request): Json| { + let responses = Arc::clone(&responses); + let captured = Arc::clone(&captured); + async move { + captured.lock().await.push(request); + Json( + responses + .lock() + .await + .pop_front() + .expect("planned local model response"), + ) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let exec = ExecutionContext::new( + ConversationHandler::new(ConversationStore::disabled()), + ResponseHandler::new(ResponseStore::disabled()), + Arc::new(reqwest::Client::new()), + format!("http://{address}"), + ); + Self { + exec, + requests, + handle: Some(handle), + } + } + + async fn close(mut self) { + let handle = self.handle.take().unwrap(); + handle.abort(); + assert!(handle.await.unwrap_err().is_cancelled()); + } +} + +impl Drop for LocalModel { + fn drop(&mut self) { + if let Some(handle) = &self.handle { + handle.abort(); + } + } +} + +fn model_message(id: &str, text: &str) -> Value { + json!({"id": id, "object": "response", "created_at": 0, "model": "test-model", "status": "completed", + "output": [{"type":"message", "id":format!("msg_{id}"), "role":"assistant", "status":"completed", + "content":[{"type":"output_text", "text":text, "annotations":[]}]}]}) +} + +fn search_round(id: &str, text: &str) -> Value { + let mut response = model_message(id, text); + response["output"].as_array_mut().unwrap().push(json!({ + "type":"function_call", "id":format!("fc_{id}"), "call_id":format!("call_{id}"), + "name":"web_search", "arguments":"{\"query\":\"local\"}", "status":"completed" + })); + response +} + +fn with_local_search(model: &mut LocalModel) -> Arc { + let calls = Arc::new(AtomicUsize::new(0)); + model.exec = model.exec.clone().with_gateway_executor(Arc::new(LocalSearch { + calls: Arc::clone(&calls), + })); + calls +} + +fn search_request() -> RequestPayload { + let mut payload = request(None, json!("short")); + payload.tools = Some(serde_json::from_value(json!([{"type":"web_search_preview"}])).unwrap()); + payload +} + +async fn execute_local(model: &LocalModel, session: &ResponseSession, payload: RequestPayload) -> ResponsePayload { + let result = ExecuteRequest::new(payload, Arc::new(model.exec.clone())) + .with_session(session) + .unwrap() + .run() + .await + .unwrap(); + let either::Either::Left(payload) = result else { + panic!("expected blocking response") + }; + payload +} + +async fn storage_pool() -> Arc { + let base_url = match std::env::var("AGENTIC_SESSION_TEST_POSTGRES_URL") { + Ok(url) => url, + Err(std::env::VarError::NotPresent) => { + return create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); + } + Err(error) => panic!("invalid PostgreSQL test configuration: {error}"), + }; + let mut database_url = reqwest::Url::parse(&base_url).expect("valid PostgreSQL fixture URL"); + assert!(matches!(database_url.scheme(), "postgres" | "postgresql")); + let admin = agentic_core::storage::create_pool(Some(&base_url)).await.unwrap(); + // UUID-only identifiers are safe here; each case owns a new database and + // never truncates or drops any database supplied by the caller. + let database = format!("session_{}", uuid::Uuid::now_v7().simple()); + sqlx::query(&format!("CREATE DATABASE {database}")) + .execute(admin.as_ref()) + .await + .unwrap(); + admin.close().await; + database_url.set_path(&database); + let pool = agentic_core::storage::create_pool_with_schema_and_configs( + Some(database_url.as_str()), + agentic_core::config::SqliteConfig::default(), + agentic_core::config::PostgresConfig { + max_connections: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + eprintln!("PostgreSQL session fixture ready: {database}"); + pool +} + +fn execution() -> ExecutionContext { + ExecutionContext::new( + ConversationHandler::new(ConversationStore::disabled()), + ResponseHandler::new(ResponseStore::disabled()), + Arc::new(reqwest::Client::new()), + "http://127.0.0.1:1".to_owned(), + ) +} + +fn session() -> ResponseSession { + ResponseSession::new(NonZeroUsize::new(100).unwrap(), NonZeroUsize::new(100_000).unwrap()) +} + +fn request(previous: Option<&str>, input: Value) -> RequestPayload { + let mut value = json!({"model": "test-model", "store": false, "previous_response_id": previous}); + value["input"] = input; + serde_json::from_value(value).unwrap() +} + +fn response(ctx: &RequestContext, output: Value) -> ResponsePayload { + let mut value = json!({ + "id": ctx.response_id, "object": "response", "created_at": 0, + "model": "test-model", "status": "completed" + }); + value["output"] = output; + serde_json::from_value(value).unwrap() +} + +fn function_call(id: &str) -> Value { + json!({"type": "function_call", "id": format!("fc_{id}"), "call_id": id, + "name": "weather", "arguments": "{}", "status": "completed"}) +} + +#[tokio::test] +async fn session_explicit_conversation_generated_turns_match_no_session_history() { + for use_session in [false, true] { + for store in [false, true] { + let mut model = LocalModel::start(vec![ + model_message("resp_first", "first answer"), + model_message("resp_second", "second answer"), + ]) + .await; + let pool = storage_pool().await; + let conversations = ConversationStore::new(pool.clone()); + model.exec.conv_handler = ConversationHandler::new(conversations.clone()); + model.exec.resp_handler = ResponseHandler::new(ResponseStore::new(Arc::clone(&pool))); + let conversation = model.exec.conv_handler.create().await.unwrap(); + let session = session(); + for input in ["first question", "second question"] { + let mut payload = request(None, json!(input)); + payload.conversation_id = Some(conversation.conversation_id.clone()); + payload.store = store; + let execution = ExecuteRequest::new(payload, Arc::new(model.exec.clone())); + let execution = if use_session { + execution.with_session(&session).unwrap() + } else { + execution + }; + let result = execution.run().await.unwrap(); + assert!(matches!(result, either::Either::Left(_))); + session.wait_until_idle().await.unwrap(); + } + let requests = model.requests.lock().await; + let input = requests[1]["input"].as_array().unwrap(); + assert_eq!(input.len(), 3, "session={use_session}, store={store}: {input:?}"); + assert!(input[0].to_string().contains("first question")); + assert!(input[1].to_string().contains("first answer")); + assert!(input[2].to_string().contains("second question")); + drop(requests); + let history = conversations.rehydrate(&conversation.conversation_id).await.unwrap(); + assert_eq!(history.len(), 4, "session={use_session}, store={store}"); + model.close().await; + pool.close().await; + } + } +} + +async fn explicit_conversation_tool_history(use_session: bool, gateway_tool: bool, store: bool) -> (Value, Value) { + let first_response = if gateway_tool { + search_round("resp_search", "before search") + } else { + json!({"id":"resp_call", "object":"response", "created_at":0, "model":"test-model", + "status":"completed", "output":[function_call("client")]}) + }; + let mut responses = vec![first_response]; + if gateway_tool { + responses.push(model_message("resp_first", "first answer")); + } + responses.push(model_message("resp_second", "second answer")); + let mut model = LocalModel::start(responses).await; + let calls = with_local_search(&mut model); + let pool = storage_pool().await; + let conversations = ConversationStore::new(Arc::clone(&pool)); + model.exec.conv_handler = ConversationHandler::new(conversations.clone()); + model.exec.resp_handler = ResponseHandler::new(ResponseStore::new(Arc::clone(&pool))); + let conversation = model.exec.conv_handler.create().await.unwrap(); + let session = session(); + let next_input = if gateway_tool { + json!("second question") + } else { + json!([{"type":"function_call_output", "call_id":"client", "output":"sunny"}]) + }; + for (index, input) in [json!("first question"), next_input].into_iter().enumerate() { + let mut payload = request(None, input); + payload.conversation_id = Some(conversation.conversation_id.clone()); + payload.store = store; + if index == 0 { + payload.tools = Some(if gateway_tool { + serde_json::from_value(json!([{"type":"web_search_preview"}])).unwrap() + } else { + serde_json::from_value(json!([{ + "type":"function", "name":"weather", "parameters":{"type":"object", "properties":{}} + }])) + .unwrap() + }); + } + let execution = ExecuteRequest::new(payload, Arc::new(model.exec.clone())); + let execution = if use_session { + execution.with_session(&session).unwrap() + } else { + execution + }; + let result = execution.run().await.unwrap(); + assert!(matches!(result, either::Either::Left(_))); + session.wait_until_idle().await.unwrap(); + } + assert_eq!(calls.load(Ordering::Relaxed), usize::from(gateway_tool)); + let requests = model.requests.lock().await; + assert_eq!(requests.len(), if gateway_tool { 3 } else { 2 }); + let model_history = requests.last().unwrap()["input"].clone(); + drop(requests); + let history = conversations.rehydrate(&conversation.conversation_id).await.unwrap(); + let stored_history = serde_json::to_value(InOutItem::into_input_items(history)).unwrap(); + model.close().await; + pool.close().await; + (model_history, stored_history) +} + +#[tokio::test] +async fn session_explicit_conversation_tool_turns_match_no_session_history() { + for gateway_tool in [false, true] { + for store in [false, true] { + let baseline = explicit_conversation_tool_history(false, gateway_tool, store).await; + let session = explicit_conversation_tool_history(true, gateway_tool, store).await; + assert_eq!(session, baseline, "gateway_tool={gateway_tool}, store={store}"); + } + } +} + +#[tokio::test] +async fn grouped_session_fork_pins_parent_while_source_advances() { + let exec = execution(); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(1_000_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let fork = group.new_session().unwrap(); + let ctx = rehydrate_in_session(request(None, json!("source prompt")), &exec, &source) + .await + .unwrap(); + let parent_id = ctx.response_id.clone(); + let output = response(&ctx, json!([])); + commit(ctx, output, &exec).await.unwrap(); + + let fork_ctx = rehydrate_in_session(request(Some(&parent_id), json!("fork prompt")), &exec, &fork) + .await + .expect("a grouped session can fork the cached source parent"); + let fork_id = fork_ctx.response_id.clone(); + let source_ctx = rehydrate_in_session(request(Some(&parent_id), json!("source advanced")), &exec, &source) + .await + .expect("source and fork can execute concurrently"); + let source_id = source_ctx.response_id.clone(); + let output = response(&source_ctx, json!([])); + commit(source_ctx, output, &exec).await.unwrap(); + let output = response(&fork_ctx, json!([])); + commit(fork_ctx, output, &exec).await.unwrap(); + + for (session, id, own_text, other_text) in [ + (&source, &source_id, "source advanced", "fork prompt"), + (&fork, &fork_id, "fork prompt", "source advanced"), + ] { + let ctx = rehydrate_in_session(request(Some(id), json!("next")), &exec, session) + .await + .unwrap(); + let history = serde_json::to_value(&ctx.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 3); + assert!(history.to_string().contains("source prompt")); + assert!(history.to_string().contains(own_text)); + assert!(!history.to_string().contains(other_text)); + } +} + +#[tokio::test] +async fn grouped_session_failed_fork_preserves_source_tool_continuation() { + let exec = execution(); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(1_000_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let fork = group.new_session().unwrap(); + let parent = parent_with_call(&exec, &source).await; + let error = rehydrate_in_session(request(Some(&parent), json!("missing output")), &exec, &fork) + .await + .unwrap_err(); + assert!(matches!(error, ExecutorError::Tool(_))); + let ctx = resolved_continuation(&exec, &source, &parent).await; + let output = response(&ctx, json!([])); + commit(ctx, output, &exec).await.unwrap(); +} + +#[tokio::test] +async fn grouped_session_late_fork_does_not_restore_an_evicted_transient_parent() { + let exec = execution(); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(1_000_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let fork = group.new_session().unwrap(); + let parent = parent_with_call(&exec, &source).await; + rehydrate_in_session(request(Some(&parent), json!("missing output")), &exec, &source) + .await + .unwrap_err(); + let error = rehydrate_in_session( + request( + Some(&parent), + json!([{"type":"function_call_output", "call_id":"call_first", "output":"sunny"}]), + ), + &exec, + &fork, + ) + .await + .unwrap_err(); + assert!(matches!(error, ExecutorError::PreviousResponseNotFound { .. })); +} + +#[tokio::test] +async fn grouped_session_budget_failure_does_not_evict_the_source_parent() { + let exec = execution(); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(3).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(1_000_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let fork = group.new_session().unwrap(); + let parent = parent_with_call(&exec, &source).await; + let ctx = resolved_continuation(&exec, &fork, &parent).await; + let output = response( + &ctx, + json!([{ + "type":"message", "id":"msg_fork", "role":"assistant", "status":"completed", + "content":[{"type":"output_text", "text":"fork answer", "annotations":[]}] + }]), + ); + let error = commit(ctx, output, &exec).await.unwrap_err(); + assert!(matches!(error, ExecutorError::PayloadTooLarge(_))); + let ctx = resolved_continuation(&exec, &source, &parent).await; + let output = response(&ctx, json!([])); + commit(ctx, output, &exec).await.unwrap(); + rehydrate_in_session(request(None, json!("fresh fork")), &exec, &fork) + .await + .expect("failed publication must release the fork's execution slot"); +} + +#[tokio::test] +async fn session_cache_miss_does_not_mask_disabled_durable_storage() { + let exec = execution(); + let session = session(); + let mut payload = request(Some("resp_missing"), json!("next")); + let missing = rehydrate_in_session(payload.clone(), &exec, &session) + .await + .unwrap_err(); + assert!(matches!(missing, ExecutorError::PreviousResponseNotFound { .. })); + payload.store = true; + let unavailable = rehydrate_in_session(payload, &exec, &session).await.unwrap_err(); + assert!(matches!( + unavailable, + ExecutorError::Storage(StorageError::NotConfigured) + )); +} + +async fn parent_with_call(exec: &ExecutionContext, session: &ResponseSession) -> String { + let ctx = rehydrate_in_session(request(None, json!("first")), exec, session) + .await + .unwrap(); + let payload = response(&ctx, json!([function_call("call_first")])); + let id = ctx.response_id.clone(); + commit(ctx, payload, exec) + .await + .expect("initial unstored call checkpoint"); + id +} + +async fn resolved_continuation(exec: &ExecutionContext, session: &ResponseSession, parent: &str) -> RequestContext { + rehydrate_in_session( + request( + Some(parent), + json!([ + {"type": "function_call_output", "call_id": "call_first", "output": "sunny"} + ]), + ), + exec, + session, + ) + .await + .expect("cached call output resolves its parent") +} + +#[tokio::test] +async fn session_commit_accepts_a_new_call_without_reading_disabled_storage() { + let exec = execution(); + let session = session(); + let parent = parent_with_call(&exec, &session).await; + let ctx = resolved_continuation(&exec, &session, &parent).await; + let payload = response(&ctx, json!([function_call("call_second")])); + let result = commit(ctx, payload, &exec).await; + assert!( + result.is_ok(), + "a cached history must not require durable storage: {result:?}" + ); + let second = result.unwrap().id; + let next = rehydrate_in_session( + request( + Some(&second), + json!([ + {"type": "function_call_output", "call_id": "call_second", "output": "cloudy"} + ]), + ), + &exec, + &session, + ) + .await + .expect("second call is also retained in the checkpoint"); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 5); +} + +#[tokio::test] +async fn session_commit_rejects_reused_call_id_against_pinned_history() { + let exec = execution(); + let session = session(); + let parent = parent_with_call(&exec, &session).await; + let ctx = resolved_continuation(&exec, &session, &parent).await; + let payload = response(&ctx, json!([function_call("call_first")])); + let error = commit(ctx, payload, &exec).await.unwrap_err(); + assert!( + matches!(&error, ExecutorError::InvalidRequest(message) if message.contains("continued history")), + "must reject the repeated call, not fail on a database lookup: {error:?}" + ); +} + +#[tokio::test] +async fn session_commit_custom_call_uses_the_same_cached_history_validation() { + let exec = execution(); + let session = session(); + let parent = parent_with_call(&exec, &session).await; + let ctx = resolved_continuation(&exec, &session, &parent).await; + let payload = response( + &ctx, + json!([{"type":"custom_tool_call", "id":"ctc_second", + "call_id":"call_second", "name":"apply_patch", "input":"patch text", "status":"completed"}]), + ); + commit(ctx, payload, &exec) + .await + .expect("custom call validation uses the pinned parent too"); +} + +#[tokio::test] +async fn session_commit_message_only_continuation_is_a_positive_control() { + let exec = execution(); + let session = session(); + let parent = parent_with_call(&exec, &session).await; + let ctx = resolved_continuation(&exec, &session, &parent).await; + let payload = response( + &ctx, + json!([{"type":"message", "id":"msg_done", "role":"assistant", + "status":"completed", "content":[{"type":"output_text", "text":"done", "annotations":[]}]}]), + ); + commit(ctx, payload, &exec) + .await + .expect("message-only completion does not inspect output call IDs"); +} + +#[tokio::test] +async fn session_automatic_compaction_fits_the_replaced_history_budget() { + let model = LocalModel::start(vec![ + model_message("resp_old", "obsolete assistant detail"), + model_message("resp_summary", "compact summary"), + model_message("resp_compacted", "new answer"), + ]) + .await; + let session = ResponseSession::new(NonZeroUsize::new(4).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let first = execute_local(&model, &session, request(None, json!("first user"))).await; + let mut payload = request(Some(&first.id), json!("second user")); + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":1}])).unwrap()); + let second = execute_local(&model, &session, payload).await; + let next = rehydrate_in_session(request(Some(&second.id), json!("third user")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!( + history.as_array().unwrap().len(), + 5, + "four retained items plus the new user item" + ); + assert!(!history.to_string().contains("obsolete assistant detail")); + assert_eq!( + history[2]["type"], "compaction", + "retain canonical items, not model-normalized summaries" + ); + assert_eq!(model.requests.lock().await.len(), 3); + model.close().await; +} + +#[tokio::test] +async fn session_compaction_trigger_replaces_history_before_retention() { + let model = LocalModel::start(vec![ + model_message("resp_old", "obsolete assistant detail"), + model_message("resp_summary", "compact summary"), + ]) + .await; + let session = ResponseSession::new(NonZeroUsize::new(2).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let first = execute_local(&model, &session, request(None, json!("first user"))).await; + let second = execute_local( + &model, + &session, + request(Some(&first.id), json!([{"type":"compaction_trigger"}])), + ) + .await; + assert_eq!(serde_json::to_value(&second.output).unwrap()[0]["type"], "compaction"); + let next = rehydrate_in_session(request(Some(&second.id), json!("second user")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 3); + assert!(!history.to_string().contains("obsolete assistant detail")); + assert_eq!(model.requests.lock().await.len(), 2); + model.close().await; +} + +#[tokio::test] +async fn session_compacted_transient_parent_promotes_only_the_canonical_window() { + let mut model = LocalModel::start(vec![ + model_message("resp_old", "obsolete assistant detail"), + model_message("resp_summary", "compact summary"), + model_message("resp_compacted", "new answer"), + ]) + .await; + let pool = storage_pool().await; + let store = ResponseStore::new(Arc::clone(&pool)); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let session = ResponseSession::new(NonZeroUsize::new(4).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let first = execute_local(&model, &session, request(None, json!("first user"))).await; + let mut payload = request(Some(&first.id), json!("second user")); + payload.store = true; + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":1}])).unwrap()); + let second = execute_local(&model, &session, payload).await; + assert!(store.get(&first.id).await.unwrap_err().is_not_found()); + let durable = + serde_json::to_value(InOutItem::into_input_items(store.rehydrate(&second.id).await.unwrap())).unwrap(); + assert_eq!(durable.as_array().unwrap().len(), 4); + assert!(!durable.to_string().contains("obsolete assistant detail")); + let fresh = ResponseSession::new(NonZeroUsize::new(4).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let mut continuation = request(Some(&second.id), json!("third user")); + continuation.store = true; + let restored = rehydrate_in_session(continuation, &model.exec, &fresh).await.unwrap(); + assert_eq!( + serde_json::to_value(&restored.enriched_request.input) + .unwrap() + .as_array() + .unwrap() + .len(), + 5 + ); + drop(restored); + model.close().await; + pool.close().await; +} + +#[tokio::test] +async fn session_failed_response_evicts_only_a_referenced_parent() { + for references_parent in [false, true] { + let failed = json!({"id":"resp_failed", "object":"response", "created_at":0, "model":"test-model", + "status":"failed", "output":[], "error":{"code":"upstream_error", "message":"failed"}}); + let model = LocalModel::start(vec![model_message("resp_first", "first answer"), failed]).await; + let session = session(); + let first = execute_local(&model, &session, request(None, json!("first user"))).await; + let parent = references_parent.then_some(first.id.as_str()); + let failed = execute_local(&model, &session, request(parent, json!("failed request"))).await; + // The existing core status enum normalizes upstream "failed" to "error". + assert_eq!(failed.status, "error"); + assert_eq!(failed.error.as_ref().unwrap()["code"], "upstream_error"); + let result = rehydrate_in_session(request(Some(&first.id), json!("retry")), &model.exec, &session).await; + if references_parent { + assert!(matches!(result, Err(ExecutorError::PreviousResponseNotFound { .. }))); + } else { + assert!(result.is_ok(), "an unrelated failed turn must not erase the checkpoint"); + } + drop(result); + assert!(matches!( + rehydrate_in_session( + request(Some(&failed.id), json!("retry failed id")), + &model.exec, + &session + ) + .await, + Err(ExecutorError::PreviousResponseNotFound { .. }) + )); + model.close().await; + } +} + +#[tokio::test] +async fn session_incomplete_response_keeps_partial_output_for_continuation() { + let mut incomplete = model_message("resp_partial", "partial answer"); + incomplete["status"] = json!("incomplete"); + incomplete["incomplete_details"] = json!({"reason":"max_output_tokens"}); + let model = LocalModel::start(vec![incomplete]).await; + let session = session(); + let first = execute_local(&model, &session, request(None, json!("first user"))).await; + assert_eq!(first.status, "incomplete"); + let next = rehydrate_in_session(request(Some(&first.id), json!("continue")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 3); + assert!(history[1].to_string().contains("partial answer")); + model.close().await; +} + +#[tokio::test] +async fn session_tool_round_compaction_does_not_restore_superseded_output() { + for store_durable in [false, true] { + let obsolete = "obsolete intermediate assistant detail ".repeat(20); + let mut first_round = model_message("resp_tool", &obsolete); + first_round["output"].as_array_mut().unwrap().push(json!({ + "type":"function_call", "id":"fc_search", "call_id":"call_search", + "name":"web_search", "arguments":"{\"query\":\"local\"}", "status":"completed" + })); + let mut model = LocalModel::start(vec![ + first_round, + model_message("resp_summary", "compact summary"), + model_message("resp_answer", "final answer"), + ]) + .await; + let calls = Arc::new(AtomicUsize::new(0)); + model.exec = model.exec.clone().with_gateway_executor(Arc::new(LocalSearch { + calls: Arc::clone(&calls), + })); + let pool = storage_pool().await; + let store = ResponseStore::new(pool); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let session = session(); + let mut payload = request(None, json!("short")); + payload.store = store_durable; + payload.tools = Some(serde_json::from_value(json!([{"type":"web_search_preview"}])).unwrap()); + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":50}])).unwrap()); + let result = execute_local(&model, &session, payload).await; + let public_output = serde_json::to_value(&result.output).unwrap(); + assert!(public_output.to_string().contains(&obsolete)); + assert!(public_output.to_string().contains("final answer")); + assert_eq!(calls.load(Ordering::Relaxed), 1); + { + let requests = model.requests.lock().await; + assert_eq!(requests.len(), 3); + assert_eq!(requests[0]["input"].as_array().unwrap().len(), 1); + assert!(requests[1]["input"].to_string().contains(&obsolete)); + assert_eq!(requests[2]["input"].as_array().unwrap().len(), 2); + assert!(!requests[2]["input"].to_string().contains(&obsolete)); + } + let next = rehydrate_in_session(request(Some(&result.id), json!("next")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert!( + !history.to_string().contains(&obsolete), + "pre-compaction output must remain public without returning to retained history: {history}" + ); + assert_eq!(history.as_array().unwrap().len(), 4); + assert_eq!(history[1]["type"], "compaction"); + if store_durable { + let durable = InOutItem::into_input_items(store.rehydrate(&result.id).await.unwrap()); + let durable = serde_json::to_value(durable).unwrap(); + assert_eq!(durable.as_array().unwrap().len(), 3); + assert!(!durable.to_string().contains(&obsolete)); + } else { + assert!(matches!( + store.get(&result.id).await, + Err(StorageError::NotFound { .. }) + )); + } + model.close().await; + } +} + +#[tokio::test] +async fn grouped_aggregate_rejection_precedes_durable_child_writes_and_preserves_other_lanes() { + for durable_parent in [false, true] { + let mut model = LocalModel::start(vec![ + model_message("resp_parent", "parent answer"), + model_message("resp_large", &"large output ".repeat(400)), + model_message("resp_recovered", "recovered"), + ]) + .await; + let pool = storage_pool().await; + model.exec.resp_handler = ResponseHandler::new(ResponseStore::new(Arc::clone(&pool))); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(3_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + let mut payload = request(None, json!("source")); + payload.store = durable_parent; + let parent = execute_local(&model, &source, payload).await; + let target_ctx = rehydrate_in_session(request(None, json!("unrelated target")), &model.exec, &target) + .await + .unwrap(); + let target_id = target_ctx.response_id.clone(); + let output = response(&target_ctx, json!([])); + commit(target_ctx, output, &model.exec).await.unwrap(); + let before: (i64, i64) = + sqlx::query_as("SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items)") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + let mut payload = request(Some(&parent.id), json!("fork")); + payload.store = true; + let result = ExecuteRequest::new(payload, Arc::new(model.exec.clone())) + .with_session(&target) + .unwrap() + .run() + .await; + assert!(matches!(result, Err(ExecutorError::PayloadTooLarge(ref message)) if message.contains("aggregate"))); + let after: (i64, i64) = sqlx::query_as("SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items)") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(after, before, "aggregate admission must precede the durable write"); + let recovered = execute_local(&model, &target, request(Some(&target_id), json!("retry own target"))).await; + assert_eq!(recovered.status, "completed"); + rehydrate_in_session( + request(Some(&parent.id), json!("source still available")), + &model.exec, + &source, + ) + .await + .unwrap(); + assert_eq!(model.requests.lock().await.len(), 3); + model.close().await; + } +} + +#[tokio::test] +async fn grouped_aggregate_reservation_is_rolled_back_when_durable_storage_fails() { + let exec = execution(); // Both stores intentionally disabled. + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(6_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let target = group.new_session().unwrap(); + let source_ctx = rehydrate_in_session(request(None, json!("s".repeat(2_000))), &exec, &source) + .await + .unwrap(); + let source_id = source_ctx.response_id.clone(); + let output = response(&source_ctx, json!([])); + commit(source_ctx, output, &exec).await.unwrap(); + let mut payload = request(None, json!("t".repeat(2_000))); + payload.store = true; + let ctx = rehydrate_in_session(payload, &exec, &target).await.unwrap(); + let output = response(&ctx, json!([])); + let result = commit(ctx, output, &exec).await; + let Err(ExecutorError::Persistence(cause)) = result else { + panic!("expected wrapped durable-storage failure, got {result:?}"); + }; + assert!(matches!(*cause, ExecutorError::Storage(StorageError::NotConfigured))); + // Two such checkpoints fit, but a third does not. Recovery proves that the + // failed persistence did not leak its already-reserved candidate charge. + let ctx = rehydrate_in_session(request(None, json!("t".repeat(2_000))), &exec, &target) + .await + .unwrap(); + let output = response(&ctx, json!([])); + commit(ctx, output, &exec) + .await + .expect("failed durable write must return its reservation"); + rehydrate_in_session(request(Some(&source_id), json!("source survives")), &exec, &source) + .await + .unwrap(); +} + +#[tokio::test] +async fn grouped_aggregate_sql_failure_rolls_back_items_and_releases_capacity() { + for fork in [false, true] { + let pool = storage_pool().await; + let store = ResponseStore::new(Arc::clone(&pool)); + let mut exec = execution(); + exec.resp_handler = ResponseHandler::new(store.clone()); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(3).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(6_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let destination = group.new_session().unwrap(); + let probe = group.new_session().unwrap(); + let source_ctx = rehydrate_in_session(request(None, json!("s".repeat(2_000))), &exec, &source) + .await + .unwrap(); + let source_id = source_ctx.response_id.clone(); + let output = response(&source_ctx, json!([])); + commit(source_ctx, output, &exec).await.unwrap(); + + let target = if fork { &destination } else { &source }; + let mut payload = request(Some(&source_id), json!("continue")); + payload.store = true; + let ctx = rehydrate_in_session(payload, &exec, target).await.unwrap(); + let failed_id = ctx.response_id.clone(); + // Seed only the response ID, not its items. The real persist path inserts + // the candidate's items before hitting this duplicate response ID. + // This tests a SQL failure after writes, not a disabled store/pool wait. + store + .persist(&failed_id, None, vec![], &ResponseMetadata::default()) + .await + .unwrap(); + let existing = store.get(&failed_id).await.unwrap(); + let output = response(&ctx, json!([])); + let result = commit(ctx, output, &exec).await; + // The public handler translates this unique violation to Conflict. + assert!( + matches!(result, Err(ExecutorError::Conflict(ref message)) + if message == &format!("a turn is already stored under '{failed_id}'")), + "must reach the duplicate-response insert: {result:?}" + ); + tokio::time::timeout(std::time::Duration::from_secs(1), target.wait_until_idle()) + .await + .expect("SQL failure must release the execution slot") + .unwrap(); + let counts: (i64, i64) = + sqlx::query_as("SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items)") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!( + counts, + (1, 0), + "candidate items must roll back; only the seeded response remains" + ); + let after = store.get(&failed_id).await.unwrap(); + assert_eq!(after.history_item_ids, existing.history_item_ids); + assert_eq!(after.created_at, existing.created_at); + assert!(matches!( + rehydrate_in_session(request(Some(&failed_id), json!("no failed checkpoint")), &exec, target).await, + Err(ExecutorError::PreviousResponseNotFound { .. }) + )); + + // With a leaked prepared checkpoint neither size fits. A failed fork + // retains its source charge; a failed own continuation evicts it. + let recovery_bytes = if fork { 2_000 } else { 4_000 }; + let recovered = rehydrate_in_session(request(None, json!("r".repeat(recovery_bytes))), &exec, &probe) + .await + .unwrap(); + let output = response(&recovered, json!([])); + commit(recovered, output, &exec) + .await + .expect("SQL failure must return its prepared-checkpoint reservation"); + let source_result = + rehydrate_in_session(request(Some(&source_id), json!("source check")), &exec, &source).await; + if fork { + let source_ctx = source_result.expect("a failed fork must preserve its source checkpoint"); + assert!( + serde_json::to_string(&source_ctx.enriched_request.input) + .unwrap() + .contains(&"s".repeat(2_000)) + ); + } else { + assert!(matches!( + source_result, + Err(ExecutorError::PreviousResponseNotFound { .. }) + )); + } + pool.close().await; + } +} + +#[tokio::test] +async fn grouped_aggregate_cancellation_while_waiting_for_storage_releases_capacity() { + for fork in [false, true] { + let pool = storage_pool().await; + assert_eq!(pool.options().get_max_connections(), 1); + let mut exec = execution(); + exec.resp_handler = ResponseHandler::new(ResponseStore::new(Arc::clone(&pool))); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(3).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(6_000).unwrap(), + ); + let source = group.new_session().unwrap(); + let destination = group.new_session().unwrap(); + let probe = group.new_session().unwrap(); + let source_ctx = rehydrate_in_session(request(None, json!("s".repeat(2_000))), &exec, &source) + .await + .unwrap(); + let source_id = source_ctx.response_id.clone(); + let output = response(&source_ctx, json!([])); + commit(source_ctx, output, &exec).await.unwrap(); + + let target = if fork { &destination } else { &source }; + let mut payload = request(Some(&source_id), json!("continue")); + payload.store = true; + let ctx = rehydrate_in_session(payload, &exec, target).await.unwrap(); + let output = response(&ctx, json!([])); + // Holding the pool's only connection stops persistence before any write. + // One poll reaches that wait after the candidate's budget reservation. + let held_connection = pool.acquire().await.unwrap(); + let mut pending = Box::pin(commit(ctx, output, &exec)); + assert!(futures::poll!(&mut pending).is_pending()); + assert!( + ExecuteRequest::new(request(None, json!("busy")), Arc::new(exec.clone())) + .with_session(target) + .is_err() + ); + let probe_ctx = rehydrate_in_session(request(None, json!("p".repeat(2_000))), &exec, &probe) + .await + .unwrap(); + let output = response(&probe_ctx, json!([])); + let pressured = commit(probe_ctx, output, &exec).await; + assert!( + matches!(pressured, Err(ExecutorError::PayloadTooLarge(ref message)) if message.contains("aggregate")), + "the pending durable candidate must already hold a reservation: fork={fork}" + ); + + // Drop the actual public commit future while the pool remains occupied. + // This is cancellation, not a storage error or a completed transaction. + drop(pending); + tokio::time::timeout(std::time::Duration::from_secs(1), target.wait_until_idle()) + .await + .expect("cancelled persistence must release the execution slot") + .unwrap(); + drop(held_connection); + let counts: (i64, i64) = + sqlx::query_as("SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items)") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(counts, (0, 0), "cancelled persistence must not leave durable rows"); + + // For a fork the source stays charged; for an own continuation it is + // evicted. Either recovery would exceed 6000 if the candidate leaked. + let recovery_bytes = if fork { 2_000 } else { 4_000 }; + let recovered = rehydrate_in_session(request(None, json!("r".repeat(recovery_bytes))), &exec, &probe) + .await + .unwrap(); + let output = response(&recovered, json!([])); + commit(recovered, output, &exec) + .await + .expect("cancelled persistence must return its reserved capacity"); + let source_result = + rehydrate_in_session(request(Some(&source_id), json!("source check")), &exec, &source).await; + if fork { + let source_ctx = source_result.expect("cancelling a fork must preserve its source checkpoint"); + assert!( + serde_json::to_string(&source_ctx.enriched_request.input) + .unwrap() + .contains(&"s".repeat(2_000)) + ); + } else { + assert!(matches!( + source_result, + Err(ExecutorError::PreviousResponseNotFound { .. }) + )); + } + pool.close().await; + } +} + +#[tokio::test] +async fn grouped_aggregate_durable_fallback_is_rejected_before_inference_and_can_recover() { + let mut model = LocalModel::start(vec![ + model_message("resp_parent", "stored"), + model_message("resp_new", "fresh"), + ]) + .await; + let pool = storage_pool().await; + let store = ResponseStore::new(pool); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let original = session(); + let mut payload = request(None, json!("durable input ".repeat(1_000))); + payload.store = true; + let parent = execute_local(&model, &original, payload).await; + drop(original); + let group = ResponseSessionGroup::new( + NonZeroUsize::new(1).unwrap(), + NonZeroUsize::new(100).unwrap(), + NonZeroUsize::new(100_000).unwrap(), + NonZeroUsize::new(1_500).unwrap(), + ); + let restored = group.new_session().unwrap(); + let mut payload = request(Some(&parent.id), json!("continue")); + payload.store = true; + let result = ExecuteRequest::new(payload, Arc::new(model.exec.clone())) + .with_session(&restored) + .unwrap() + .run() + .await; + assert!(matches!(result, Err(ExecutorError::PayloadTooLarge(ref message)) if message.contains("aggregate"))); + assert_eq!( + model.requests.lock().await.len(), + 1, + "over-budget durable parent must not reach inference" + ); + assert!( + store.get(&parent.id).await.is_ok(), + "retention failure must not delete durable history" + ); + let recovered = execute_local(&model, &restored, request(None, json!("fresh root"))).await; + assert_eq!(recovered.status, "completed"); + model.close().await; +} + +#[tokio::test] +async fn session_budget_rejection_does_not_partially_store_a_child() { + for durable_parent in [false, true] { + for byte_budget in [false, true] { + let mut model = LocalModel::start(vec![ + model_message("resp_parent", "parent answer"), + model_message("resp_too_large", &"oversized answer ".repeat(200)), + model_message("resp_recovery", "recovered"), + ]) + .await; + let pool = storage_pool().await; + model.exec.resp_handler = ResponseHandler::new(ResponseStore::new(Arc::clone(&pool))); + let session = ResponseSession::new( + NonZeroUsize::new(if byte_budget { 100 } else { 3 }).unwrap(), + NonZeroUsize::new(if byte_budget { 2_000 } else { 100_000 }).unwrap(), + ); + let mut payload = request(None, json!("parent")); + payload.store = durable_parent; + let parent = execute_local(&model, &session, payload).await; + let before: (i64, i64, i64) = sqlx::query_as( + "SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items), (SELECT COUNT(*) FROM conversations)", + ) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(before, if durable_parent { (1, 2, 0) } else { (0, 0, 0) }); + let mut payload = request(Some(&parent.id), json!("child")); + payload.store = true; + let result = ExecuteRequest::new(payload, Arc::new(model.exec.clone())) + .with_session(&session) + .unwrap() + .run() + .await; + assert!(matches!(result, Err(ExecutorError::PayloadTooLarge(_)))); + let after: (i64, i64, i64) = sqlx::query_as( + "SELECT (SELECT COUNT(*) FROM responses), (SELECT COUNT(*) FROM items), (SELECT COUNT(*) FROM conversations)", + ) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(after, before, "a failed budget check must precede durable writes"); + assert!(matches!( + rehydrate_in_session(request(Some(&parent.id), json!("next")), &model.exec, &session).await, + Err(ExecutorError::PreviousResponseNotFound { .. }) + )); + let recovery = execute_local(&model, &session, request(None, json!("fresh"))).await; + assert_eq!( + recovery.status, "completed", + "budget failure releases the execution slot" + ); + model.close().await; + } + } +} + +#[tokio::test] +async fn session_tool_round_preserves_reasoning_message_call_output_order() { + for store_durable in [false, true] { + let mut first_round = model_message("resp_tool", "intermediate answer"); + let output = first_round["output"].as_array_mut().unwrap(); + output.insert( + 0, + json!({"type":"reasoning", "id":"rs_round", "summary":[], + "content":[{"type":"reasoning_text", "text":"intermediate reasoning"}]}), + ); + output.push( + json!({"type":"function_call", "id":"fc_search", "call_id":"call_search", + "name":"web_search", "arguments":"{\"query\":\"local\"}", "status":"completed"}), + ); + let mut model = LocalModel::start(vec![first_round, model_message("resp_answer", "final answer")]).await; + model.exec = model.exec.clone().with_gateway_executor(Arc::new(LocalSearch { + calls: Arc::new(AtomicUsize::new(0)), + })); + let pool = storage_pool().await; + let store = ResponseStore::new(pool); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let session = session(); + let mut payload = request(None, json!("question")); + payload.store = store_durable; + payload.tools = Some(serde_json::from_value(json!([{"type":"web_search_preview"}])).unwrap()); + let result = execute_local(&model, &session, payload).await; + let next = rehydrate_in_session(request(Some(&result.id), json!("next")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + let kinds = history + .as_array() + .unwrap() + .iter() + .map(|item| item["type"].clone()) + .collect::>(); + assert_eq!( + kinds, + vec![ + json!("message"), + json!("reasoning"), + json!("message"), + json!("function_call"), + json!("function_call_output"), + json!("message"), + json!("message") + ], + "canonical continuation must preserve inference-round order: {history}" + ); + assert!(history[2].to_string().contains("intermediate answer")); + assert!(history[5].to_string().contains("final answer")); + if store_durable { + let durable = InOutItem::into_input_items(store.rehydrate(&result.id).await.unwrap()); + let durable = serde_json::to_value(durable).unwrap(); + assert_eq!(durable, Value::Array(history.as_array().unwrap()[..6].to_vec())); + } + model.close().await; + } +} + +#[tokio::test] +async fn session_mixed_tool_round_records_calls_once_before_client_continuation() { + let mut first_round = search_round("search", "before the calls"); + first_round["output"] + .as_array_mut() + .unwrap() + .push(function_call("client")); + let mut model = LocalModel::start(vec![first_round]).await; + let calls = with_local_search(&mut model); + let session = session(); + let mut payload = search_request(); + payload.tools.as_mut().unwrap().push( + serde_json::from_value(json!({ + "type":"function", "name":"weather", "parameters":{"type":"object", "properties":{}} + })) + .unwrap(), + ); + let result = execute_local(&model, &session, payload).await; + assert_eq!(result.output.len(), 3); + assert_eq!(calls.load(Ordering::Relaxed), 1); + let next = rehydrate_in_session( + request( + Some(&result.id), + json!([{"type":"function_call_output", "call_id":"client", "output":"sunny"}]), + ), + &model.exec, + &session, + ) + .await + .expect("the client call is retained exactly once despite the mixed round"); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 6); + assert!(history[1].to_string().contains("before the calls")); + assert_eq!(history[2]["call_id"], "call_search"); + assert_eq!(history[3]["call_id"], "client"); + assert_eq!(history[4]["type"], "function_call_output"); + assert_eq!(history[4]["call_id"], "call_search"); + assert_eq!(history[5]["call_id"], "client"); + model.close().await; +} + +#[tokio::test] +async fn session_round_limit_retains_the_final_call_output_in_order() { + let mut model = LocalModel::start( + (0..10) + .map(|round| search_round(&format!("round_{round}"), "intermediate")) + .collect(), + ) + .await; + let calls = with_local_search(&mut model); + let session = session(); + let result = execute_local(&model, &session, search_request()).await; + assert_eq!(result.status, "incomplete"); + assert_eq!(calls.load(Ordering::Relaxed), 10); + let next = rehydrate_in_session(request(Some(&result.id), json!("continue")), &model.exec, &session) + .await + .expect("round exhaustion must not leave an unanswered built-in call"); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + let items = history.as_array().unwrap(); + assert_eq!(items.len(), 32); + for (round, items) in items[1..31].chunks_exact(3).enumerate() { + assert_eq!(items[0]["type"], "message"); + assert_eq!(items[1]["type"], "function_call"); + assert_eq!(items[2]["type"], "function_call_output"); + assert_eq!(items[1]["call_id"], format!("call_round_{round}")); + assert_eq!(items[2]["call_id"], items[1]["call_id"]); + } + model.close().await; +} + +#[tokio::test] +async fn session_repeated_in_round_compaction_retains_only_the_latest_window() { + let mut model = LocalModel::start(vec![ + search_round("first", &"obsolete first round ".repeat(30)), + model_message("resp_summary_one", "first summary"), + search_round("second", &"obsolete second round ".repeat(30)), + model_message("resp_summary_two", "second summary"), + model_message("resp_final", "final answer"), + ]) + .await; + let calls = with_local_search(&mut model); + let session = ResponseSession::new(NonZeroUsize::new(3).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let mut payload = search_request(); + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":50}])).unwrap()); + let result = execute_local(&model, &session, payload).await; + assert_eq!( + result.output.len(), + 5, + "both intermediate messages and search calls remain public" + ); + assert_eq!(calls.load(Ordering::Relaxed), 2); + assert_eq!(model.requests.lock().await.len(), 5); + let next = rehydrate_in_session(request(Some(&result.id), json!("next")), &model.exec, &session) + .await + .unwrap(); + let history = serde_json::to_value(&next.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 4); + assert_eq!(history[1]["type"], "compaction"); + assert_eq!(history[1]["encrypted_content"], "second summary"); + assert!(history[2].to_string().contains("final answer")); + assert!(!history.to_string().contains("obsolete")); + assert!(!history.to_string().contains("first summary")); + model.close().await; +} + +#[tokio::test] +async fn session_compacted_durable_parent_restores_within_the_canonical_budget() { + for (max_items, max_bytes, aggregate) in [(6, 100_000, None), (100, 3_000, None), (100, 100_000, Some(6_000))] { + let mut model = LocalModel::start(vec![ + model_message("resp_old", &"obsolete durable assistant detail".repeat(1_000)), + model_message("resp_summary", "compact durable summary"), + model_message("resp_compacted", "compacted answer"), + model_message("resp_child", "new answer"), + ]) + .await; + let pool = storage_pool().await; + let store = ResponseStore::new(pool); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let original_session = session(); + let mut payload = request(None, json!("first user")); + payload.store = true; + let first = execute_local(&model, &original_session, payload).await; + let mut payload = request(Some(&first.id), json!("second user")); + payload.store = true; + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":1}])).unwrap()); + let compacted = execute_local(&model, &original_session, payload).await; + assert_eq!( + store.rehydrate(&compacted.id).await.unwrap().len(), + 6, + "durable history still references its pre-compaction parent rows" + ); + drop(original_session); + let group = aggregate.map(|limit| { + ResponseSessionGroup::new( + NonZeroUsize::new(1).unwrap(), + NonZeroUsize::new(max_items).unwrap(), + NonZeroUsize::new(max_bytes).unwrap(), + NonZeroUsize::new(limit).unwrap(), + ) + }); + let restored = match &group { + Some(group) => group.new_session().unwrap(), + None => ResponseSession::new( + NonZeroUsize::new(max_items).unwrap(), + NonZeroUsize::new(max_bytes).unwrap(), + ), + }; + let mut payload = request(Some(&compacted.id), json!("third user")); + payload.store = true; + let child = execute_local(&model, &restored, payload).await; + let ctx = rehydrate_in_session(request(Some(&child.id), json!("next")), &model.exec, &restored) + .await + .unwrap(); + let history = serde_json::to_value(&ctx.enriched_request.input).unwrap(); + assert_eq!(history.as_array().unwrap().len(), 7); + assert_eq!(history[2]["type"], "compaction"); + assert!(!history.to_string().contains("obsolete durable assistant detail")); + assert!( + store.get(&first.id).await.is_ok(), + "cache canonicalization must not delete durable history" + ); + model.close().await; + } +} + +#[tokio::test] +async fn replayed_compaction_releases_superseded_session_history() { + let exec = execution(); + for pending_call in [false, true] { + let session = ResponseSession::new(NonZeroUsize::new(2).unwrap(), NonZeroUsize::new(100_000).unwrap()); + let ctx = rehydrate_in_session(request(None, json!("obsolete question")), &exec, &session) + .await + .unwrap(); + let parent = ctx.response_id.clone(); + let items = if pending_call { + json!([function_call("obsolete")]) + } else { + model_message("resp_old", "obsolete answer")["output"].clone() + }; + let output = response(&ctx, items); + commit(ctx, output, &exec).await.unwrap(); + let ctx = rehydrate_in_session( + request( + Some(&parent), + json!([ + {"type":"compaction", "id":"cmp_replayed", "encrypted_content":"current summary"} + ]), + ), + &exec, + &session, + ) + .await + .expect("a replayed window supersedes old pending calls"); + let child = ctx.response_id.clone(); + let output = response(&ctx, json!([])); + commit(ctx, output, &exec) + .await + .expect("superseded history must not consume the checkpoint budget"); + let ctx = rehydrate_in_session(request(Some(&child), json!([])), &exec, &session) + .await + .unwrap(); + let items = serde_json::to_value(&ctx.enriched_request.input).unwrap(); + assert_eq!(items.as_array().unwrap().len(), 1); + assert_eq!(items[0]["type"], "compaction"); + } +} + +#[tokio::test] +async fn session_commit_validates_call_ids_against_replayed_compaction_window() { + let exec = execution(); + let session = session(); + let parent = parent_with_call(&exec, &session).await; + let ctx = rehydrate_in_session( + request( + Some(&parent), + json!([ + {"type":"compaction", "encrypted_content":"superseded old tool call"} + ]), + ), + &exec, + &session, + ) + .await + .unwrap(); + let payload = response(&ctx, json!([function_call("call_first")])); + let child = commit(ctx, payload, &exec) + .await + .expect("the old call is outside the effective history"); + let ctx = resolved_continuation(&exec, &session, &child.id).await; + let items = serde_json::to_value(&ctx.enriched_request.input).unwrap(); + assert_eq!(items.as_array().unwrap().len(), 3); + assert_eq!(items[0]["type"], "compaction"); + assert_eq!(items[1]["type"], "function_call"); + assert_eq!(items[2]["type"], "function_call_output"); +} diff --git a/crates/agentic-server-core/tests/shell_tool_test.rs b/crates/agentic-server-core/tests/shell_tool_test.rs index 3d2c13ab..3da56000 100644 --- a/crates/agentic-server-core/tests/shell_tool_test.rs +++ b/crates/agentic-server-core/tests/shell_tool_test.rs @@ -42,6 +42,7 @@ fn context() -> RequestContext { response_id: "resp_reserved".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/tests/stateful_responses_integration.rs b/crates/agentic-server-core/tests/stateful_responses_integration.rs index b2a27f81..62457c7c 100644 --- a/crates/agentic-server-core/tests/stateful_responses_integration.rs +++ b/crates/agentic-server-core/tests/stateful_responses_integration.rs @@ -891,6 +891,7 @@ async fn tool_search_store_false_manual_replay_completes_without_reusable_respon response_id: "resp_lookup".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; let error = fixture .exec_ctx @@ -955,6 +956,7 @@ async fn test_previous_response_id_persists_inherited_tools_and_choice() { response_id: "resp_lookup".into(), conversation_id: None, conversation_version: None, + continuation: None, }; let stored = fixture @@ -1391,6 +1393,7 @@ fn lookup_context(previous_response_id: Option, conversation_id: Option< response_id: "resp_lookup".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/tests/support/media_session_checks.rs b/crates/agentic-server-core/tests/support/media_session_checks.rs new file mode 100644 index 00000000..4b7a8757 --- /dev/null +++ b/crates/agentic-server-core/tests/support/media_session_checks.rs @@ -0,0 +1,206 @@ +//! Session contracts for media-aware compaction and typed message-file rejection. +use super::*; + +fn media_request(previous: Option<&str>, input: Value) -> RequestPayload { + let mut payload = request(previous, input); + payload.context_management = + Some(serde_json::from_value(json!([{"type":"compaction", "compact_threshold":4096}])).unwrap()); + payload +} + +fn media_input(image_bytes: usize) -> Value { + json!([{"role":"user", "content":[ + {"type":"input_text", "text":"describe"}, + {"type":"input_image", "image_url":format!("data:image/png;base64,{}", "A".repeat(image_bytes))} + ]}]) +} + +fn retained_image(input: &Value) -> &Value { + let matches: Vec<_> = input + .as_array() + .unwrap() + .iter() + .filter_map(|item| item.get("content").and_then(Value::as_array)) + .flatten() + .filter(|part| part["type"] == "input_image") + .collect(); + assert_eq!(matches.len(), 1, "exactly one retained image"); + matches[0] +} + +#[tokio::test] +async fn combined_media_session_history_and_forks_do_not_compact_image_bytes() { + let model = LocalModel::start(vec![ + model_message("resp_image", "first answer"), + model_message("resp_fork", "fork answer"), + model_message("resp_source", "source answer"), + ]) + .await; + let group = ResponseSessionGroup::new( + NonZeroUsize::new(2).unwrap(), + NonZeroUsize::new(64).unwrap(), + NonZeroUsize::new(1024 * 1024).unwrap(), + NonZeroUsize::new(4 * 1024 * 1024).unwrap(), + ); + let source = group.new_session().unwrap(); + let fork = group.new_session().unwrap(); + let input = media_input(256 * 1024); + let first = execute_local(&model, &source, media_request(None, input.clone())).await; + assert_eq!( + model.requests.lock().await.len(), + 1, + "image alone must not add a summary" + ); + execute_local(&model, &fork, media_request(Some(&first.id), json!("fork question"))).await; + execute_local( + &model, + &source, + media_request(Some(&first.id), json!("source question")), + ) + .await; + let requests = model.requests.lock().await; + assert_eq!(requests.len(), 3, "one inference per turn, no image-driven summary"); + for request in requests.iter() { + assert_eq!(retained_image(&request["input"]), &input[0]["content"][1]); + } + drop(requests); + source.wait_until_idle().await.unwrap(); + fork.wait_until_idle().await.unwrap(); + model.close().await; +} + +#[tokio::test] +async fn combined_media_text_compaction_promotes_and_restores_canonical_image_history() { + let old_answer = "obsolete assistant details ".repeat(1024); + let mut model = LocalModel::start(vec![ + model_message("resp_image", &old_answer), + model_message("resp_summary", "compact summary"), + model_message("resp_compacted", "second answer"), + model_message("resp_restored", "third answer"), + ]) + .await; + let pool = storage_pool().await; + let store = ResponseStore::new(Arc::clone(&pool)); + model.exec.resp_handler = ResponseHandler::new(store.clone()); + let session = ResponseSession::new( + NonZeroUsize::new(64).unwrap(), + NonZeroUsize::new(2 * 1024 * 1024).unwrap(), + ); + let input = media_input(256 * 1024); + let first = execute_local(&model, &session, media_request(None, input.clone())).await; + assert_eq!(model.requests.lock().await.len(), 1, "image alone must not compact"); + let mut next = media_request(Some(&first.id), json!("second question")); + next.store = true; + let second = execute_local(&model, &session, next).await; + assert_eq!( + model.requests.lock().await.len(), + 3, + "long text must still cause a summary" + ); + assert!(store.get(&first.id).await.unwrap_err().is_not_found()); + let durable = + serde_json::to_value(InOutItem::into_input_items(store.rehydrate(&second.id).await.unwrap())).unwrap(); + assert_eq!(retained_image(&durable), &input[0]["content"][1]); + assert!(!durable.to_string().contains(&old_answer)); + assert_eq!( + durable + .as_array() + .unwrap() + .iter() + .filter(|item| item["type"] == "compaction") + .count(), + 1 + ); + drop(session); + let fresh = ResponseSession::new( + NonZeroUsize::new(64).unwrap(), + NonZeroUsize::new(2 * 1024 * 1024).unwrap(), + ); + let mut restored = media_request(Some(&second.id), json!("third question")); + restored.store = true; + execute_local(&model, &fresh, restored).await; + let requests = model.requests.lock().await; + assert_eq!(requests.len(), 4, "canonical restore must not recompact image bytes"); + assert_eq!(retained_image(&requests[3]["input"]), &input[0]["content"][1]); + assert!(!requests[3]["input"].to_string().contains(&old_answer)); + drop(requests); + model.close().await; + pool.close().await; +} + +#[tokio::test] +async fn combined_media_token_allowance_does_not_bypass_serialized_retention_budget() { + let model = LocalModel::start(vec![ + model_message("resp_too_large", "answer"), + model_message("resp_recovered", "small answer"), + ]) + .await; + let session = ResponseSession::new(NonZeroUsize::new(64).unwrap(), NonZeroUsize::new(16 * 1024).unwrap()); + let error = ExecuteRequest::new( + media_request(None, media_input(128 * 1024)), + Arc::new(model.exec.clone()), + ) + .with_session(&session) + .unwrap() + .run() + .await + .err() + .expect("real image bytes exceed retention"); + assert!(matches!(error, ExecutorError::PayloadTooLarge(_)), "{error}"); + assert!(error.to_string().contains("checkpoint budget"), "{error}"); + assert_eq!( + model.requests.lock().await.len(), + 1, + "retention is distinct from image token cost" + ); + tokio::time::timeout(std::time::Duration::from_secs(1), session.wait_until_idle()) + .await + .unwrap() + .unwrap(); + execute_local(&model, &session, media_request(None, json!("small new question"))).await; + assert_eq!( + model.requests.lock().await.len(), + 2, + "one rejected publication, then a successful turn" + ); + model.close().await; +} + +#[tokio::test] +async fn combined_media_message_file_rejection_precedes_compaction_and_releases_session() { + let model = LocalModel::start(vec![ + model_message("resp_recovered", "answer"), + model_message("resp_unexpected_extra", "extra answer only if compaction regresses"), + ]) + .await; + let session = ResponseSession::new(NonZeroUsize::new(64).unwrap(), NonZeroUsize::new(1024 * 1024).unwrap()); + let mut input = media_input(256 * 1024); + input[0]["content"] + .as_array_mut() + .unwrap() + .push(json!({"type":"input_file", "file_data":"encoded document"})); + let error = ExecuteRequest::new(media_request(None, input), Arc::new(model.exec.clone())) + .with_session(&session) + .unwrap() + .run() + .await + .err() + .expect("unsupported message file"); + assert!(matches!(error, ExecutorError::InvalidRequest(_))); + assert!(error.to_string().contains("input_file"), "{error}"); + assert!( + model.requests.lock().await.is_empty(), + "neither summary nor answer may start" + ); + tokio::time::timeout(std::time::Duration::from_secs(1), session.wait_until_idle()) + .await + .unwrap() + .unwrap(); + execute_local(&model, &session, media_request(None, media_input(256 * 1024))).await; + assert_eq!( + model.requests.lock().await.len(), + 1, + "valid image remains supported after rejection" + ); + model.close().await; +} diff --git a/crates/agentic-server-core/tests/terminal_status_test.rs b/crates/agentic-server-core/tests/terminal_status_test.rs index 5b797d59..42ffae76 100644 --- a/crates/agentic-server-core/tests/terminal_status_test.rs +++ b/crates/agentic-server-core/tests/terminal_status_test.rs @@ -12,6 +12,7 @@ fn context() -> RequestContext { response_id: "resp_local".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, } } diff --git a/crates/agentic-server-core/tests/tool_normalization_test.rs b/crates/agentic-server-core/tests/tool_normalization_test.rs index 2088f418..a7ca91fd 100644 --- a/crates/agentic-server-core/tests/tool_normalization_test.rs +++ b/crates/agentic-server-core/tests/tool_normalization_test.rs @@ -102,6 +102,7 @@ fn upstream_request_value(payload: RequestPayload, stream: bool) -> Value { response_id: "resp_test".to_string(), conversation_id: None, conversation_version: None, + continuation: None, }; let upstream_request = ctx .enriched_request diff --git a/crates/agentic-server/src/handler/websocket/responses.rs b/crates/agentic-server/src/handler/websocket/responses.rs index bbb9ba0f..3b5b7334 100644 --- a/crates/agentic-server/src/handler/websocket/responses.rs +++ b/crates/agentic-server/src/handler/websocket/responses.rs @@ -1,4 +1,5 @@ use std::collections::{HashMap, VecDeque}; +use std::num::NonZeroUsize; use std::sync::Arc; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; @@ -17,7 +18,8 @@ use tracing::{debug, warn}; use agentic_core::ResponseUsage; use agentic_core::executor::{ - BoxStream, ExecuteRequest, ExecutorError, RequestContext, persist_turn, rehydrate_conversation, + BoxStream, ExecuteRequest, ExecutorError, RequestContext, ResponseSession, ResponseSessionGroup, persist_turn, + rehydrate_in_session, }; use agentic_core::types::request_response::RequestPayload; use agentic_core::utils::common::utcnow_str; @@ -34,6 +36,12 @@ const WS_MAX_EVENT_BYTES: usize = 1024 * 1024; const WS_MAX_OUTSTANDING_REQUESTS: usize = 64; const WS_MAX_OUTSTANDING_BYTES: usize = 12 * 1024 * 1024; const WS_MAX_STREAM_ID_CHARS: usize = 256; +// Retained state is bounded separately from queued requests and outbound events. +// Limits cover serialized checkpoints, including pinned parents and replacements. +const WS_MAX_SESSION_LANES: usize = 128; +const WS_MAX_CHECKPOINT_ITEMS: usize = 32_768; +const WS_MAX_CHECKPOINT_BYTES: usize = 16 * 1024 * 1024; +const WS_MAX_RETAINED_BYTES: usize = 32 * 1024 * 1024; /// Serialized and size-checked before entering the bounded outbound queue. struct WsOutboundEvent(String); @@ -92,6 +100,7 @@ struct WsRequest { #[derive(Debug)] struct WsRequestParseError { + previous_response_id: Option, error: WsError, stream_id: Option, } @@ -127,6 +136,7 @@ impl WsWorkItem { } struct WsAdmissionError { + error: WsError, stream_id: Option, } @@ -162,6 +172,9 @@ struct WsMultiplexer { principal: Option>, outbound_tx: mpsc::Sender, lanes: HashMap, VecDeque>, + // Idle lanes retain their latest checkpoint until the connection closes. + sessions: HashMap, Arc>, + session_group: ResponseSessionGroup, request_tasks: JoinSet, queued_requests: usize, byte_budget: WsByteBudget, @@ -182,6 +195,13 @@ impl WsMultiplexer { principal: principal.map(Arc::new), outbound_tx, lanes: HashMap::new(), + sessions: HashMap::new(), + session_group: ResponseSessionGroup::new( + NonZeroUsize::new(WS_MAX_SESSION_LANES).expect("positive session limit"), + NonZeroUsize::new(WS_MAX_CHECKPOINT_ITEMS).expect("positive item limit"), + NonZeroUsize::new(WS_MAX_CHECKPOINT_BYTES).expect("positive checkpoint limit"), + NonZeroUsize::new(WS_MAX_RETAINED_BYTES).expect("positive retention limit"), + ), request_tasks: JoinSet::new(), queued_requests: 0, byte_budget: WsByteBudget::default(), @@ -197,12 +217,26 @@ impl WsMultiplexer { fn schedule(&mut self, work: WsWorkItem) -> Result<(), WsAdmissionError> { if !self.has_capacity_for(work.input_bytes()) { return Err(WsAdmissionError { + error: WsError::TooManyRequests, stream_id: work.stream_id().cloned(), }); } - self.byte_budget.reserve(work.input_bytes()); let lane = work.lane(); + if !self.sessions.contains_key(&lane) { + if self.sessions.len() >= WS_MAX_SESSION_LANES { + return Err(WsAdmissionError { + error: WsError::TooManyRequests, + stream_id: lane, + }); + } + let session = self.session_group.new_session().map_err(|error| WsAdmissionError { + error: WsError::from(error), + stream_id: lane.clone(), + })?; + self.sessions.insert(lane.clone(), Arc::new(session)); + } + self.byte_budget.reserve(work.input_bytes()); if let Some(queue) = self.lanes.get_mut(&lane) { queue.push_back(work); self.queued_requests += 1; @@ -273,6 +307,8 @@ impl WsMultiplexer { let shutdown_token = self.shutdown_token.clone(); let stream_id = work.stream_id().cloned(); let input_bytes = work.input_bytes(); + // schedule creates a session before admitting work; idle sessions survive schedule_next. + let session = Arc::clone(self.sessions.get(&lane).expect("admitted lane has a session")); self.request_tasks.spawn(async move { // Admission may precede dispatch by an entire inference/tool round. // Recheck here for both new lanes and work dequeued by schedule_next. @@ -285,9 +321,24 @@ impl WsMultiplexer { } let result = match work { WsWorkItem::Execute { request, .. } => { - handle_ws_request(*request, &state, auth, &outbound_tx, &shutdown_token).await + handle_ws_request(*request, &state, auth, &outbound_tx, &shutdown_token, &session).await } - WsWorkItem::Reject { error, .. } => Err(error.error), + WsWorkItem::Reject { error, .. } => { + if let Some(parent) = error.previous_response_id.as_deref() { + session + .discard_cached_response(parent) + .map_err(WsError::from) + .and(Err(error.error)) + } else { + Err(error.error) + } + } + }; + // Dropping a failed executor stream aborts its worker asynchronously. + // Do not dispatch the next turn until its lease has been released. + let result = match session.wait_until_idle().await { + Ok(()) => result, + Err(error) => Err(WsError::from(error)), }; let result = match result { Ok(()) => Ok(()), @@ -425,8 +476,15 @@ async fn responses_ws_loop( if client_disconnected { multiplexer.request_tasks.abort_all(); while multiplexer.request_tasks.join_next().await.is_some() {} + // Executor stream disposal aborts its nested inference worker. Wait for + // every lease to release its pinned state before ending this connection. + for session in multiplexer.sessions.values() { + if let Err(error) = session.wait_until_idle().await { + warn!(%error, "failed to await websocket continuation disposal"); + } + } } - drop(multiplexer.outbound_tx); + drop(multiplexer); close_ws(&mut sender, &mut receiver).await; debug!("responses websocket session closed"); } @@ -484,7 +542,7 @@ async fn handle_ws_client_message( } match multiplexer.schedule(work) { Ok(()) => true, - Err(rejected) => handle_ws_error(sender, WsError::TooManyRequests, rejected.stream_id.as_ref()).await, + Err(rejected) => handle_ws_error(sender, rejected.error, rejected.stream_id.as_ref()).await, } } Message::Binary(_) if *draining => true, @@ -547,6 +605,7 @@ where fn parse_ws_request(text: &str) -> Result { let value = serde_json::from_str::(text).map_err(|error| WsRequestParseError { error: WsError::InvalidJson(error), + previous_response_id: None, stream_id: None, })?; let stream_id = value @@ -560,30 +619,36 @@ fn parse_ws_request(text: &str) -> Result { .transpose() .map_err(|error| WsRequestParseError { error: WsError::from(ExecutorError::InvalidRequest(error)), + previous_response_id: None, stream_id: None, })?; if value.get("type").and_then(Value::as_str) != Some("response.create") { return Err(WsRequestParseError { error: WsError::UnexpectedType, + previous_response_id: None, stream_id, }); } + // Only valid routing plus response.create may identify a checkpoint for eviction. + // In particular, an explicit null/invalid stream_id must not target the default lane. + let previous_response_id = value + .get("previous_response_id") + .and_then(Value::as_str) + .map(str::to_owned); let generate = value.get("generate").and_then(Value::as_bool); let mut payload = serde_json::from_value::(value).map_err(|error| WsRequestParseError { error: WsError::from(ExecutorError::from(error)), + previous_response_id, stream_id: stream_id.clone(), })?; let requested_stream = payload.stream; - let requested_store = payload.store; payload.stream = true; - payload.store = true; debug!( requested_stream, - requested_store, forced_stream = payload.stream, - forced_store = payload.store, + store = payload.store, has_previous_response_id = payload.previous_response_id.is_some(), has_conversation_id = payload.conversation_id.is_some(), stream_id = stream_id.as_ref().map(StreamId::as_str), @@ -605,6 +670,7 @@ async fn handle_ws_request( auth: Option, outbound_tx: &mpsc::Sender, shutdown_token: &CancellationToken, + session: &ResponseSession, ) -> Result<(), WsError> { let WsRequest { payload, @@ -614,11 +680,12 @@ async fn handle_ws_request( if generate == Some(false) { debug!("handling non-generating websocket request locally"); - return complete_without_inference(outbound_tx, state, payload, stream_id.as_ref()).await; + return complete_without_inference(outbound_tx, state, payload, stream_id.as_ref(), session).await; } let result = ExecuteRequest::new(payload, Arc::clone(&state.exec_ctx)) .with_auth(auth) + .with_session(session)? .run() .await?; let Some(result) = keep_if_running(shutdown_token, result) else { @@ -639,8 +706,9 @@ async fn complete_without_inference( state: &AppState, payload: RequestPayload, stream_id: Option<&StreamId>, + session: &ResponseSession, ) -> Result<(), WsError> { - let ctx = rehydrate_conversation(payload, &state.exec_ctx).await?; + let ctx = rehydrate_in_session(payload, &state.exec_ctx, session).await?; let created_at = utcnow_str(); let created_event = empty_response_event(&ctx, created_at, "response.created", "in_progress", 0, None); let completed_event = empty_response_event( diff --git a/crates/agentic-server/tests/input_file_repro.rs b/crates/agentic-server/tests/input_file_repro.rs index 88e8ecd2..b479f350 100644 --- a/crates/agentic-server/tests/input_file_repro.rs +++ b/crates/agentic-server/tests/input_file_repro.rs @@ -615,6 +615,7 @@ fn composed_upstream_request_cannot_bypass_file_validation() { response_id: "resp_composed".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; for stream in [false, true] { let error = upstream_request(&ctx, stream).expect_err("direct typed upstream preparation must reject files"); diff --git a/crates/agentic-server/tests/responses_websocket_test.rs b/crates/agentic-server/tests/responses_websocket_test.rs index 999a564f..3760e9b0 100644 --- a/crates/agentic-server/tests/responses_websocket_test.rs +++ b/crates/agentic-server/tests/responses_websocket_test.rs @@ -867,7 +867,7 @@ async fn assert_oversized_local_completion_rejected(only_completion_oversized: b } #[tokio::test] -async fn test_websocket_generate_false_prewarm_persists_context_without_inference() { +async fn test_websocket_generate_false_prewarm_retains_context_without_inference() { let mock = MockResponsesServer::start(vec![sse_response("resp_upstream_1", "msg_upstream_1", "READY")]).await; let fixture = storage_backed_state(&mock.url).await; let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; @@ -962,7 +962,7 @@ async fn test_websocket_generate_false_prewarm_redacts_mcp_runtime_credentials() "require_approval": "never" }], "generate": false, - "store": false, + "store": true, "stream": true }), ) @@ -988,6 +988,7 @@ async fn test_websocket_generate_false_prewarm_redacts_mcp_runtime_credentials() response_id: "resp_lookup".to_owned(), conversation_id: None, conversation_version: None, + continuation: None, }; let stored = fixture .state @@ -2730,8 +2731,8 @@ async fn test_websocket_unknown_previous_response_returns_error_event() { let error = recv_json(&mut ws).await; assert_eq!(error["type"], "error"); - assert_eq!(error["status"], StatusCode::NOT_FOUND.as_u16()); - assert_eq!(error["error"]["code"], "not_found"); + assert_eq!(error["status"], StatusCode::BAD_REQUEST.as_u16()); + assert_eq!(error["error"]["code"], "previous_response_not_found"); assert_eq!(error["stream_id"], "missing-response-lane"); assert!(mock.request_bodies().await.is_empty()); } @@ -3018,3 +3019,423 @@ async fn websocket_client_close_cancels_all_active_stream_lanes() { .expect("upstream drop sender should notify"); } } + +async fn assert_no_response_state(pool: &DbPool) { + for table in ["responses", "items"] { + let count: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table}")) + .fetch_one(pool) + .await + .unwrap(); + assert_eq!(count, 0, "store:false wrote to {table}"); + } +} + +#[tokio::test] +async fn websocket_unstored_continuation_is_connection_local() { + for generate in [false, true] { + for lane in [None, Some("private")] { + let mock = MockResponsesServer::start(vec![ + sse_response("resp_first", "msg_first", "PRIVATE"), + sse_response("resp_second", "msg_second", "CONTINUED"), + ]) + .await; + let fixture = storage_backed_state(&mock.url).await; + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let mut request = json!({"type":"response.create", "model":"test-model", + "input":"private prefix", "store":false, "generate":generate}); + if let Some(lane) = lane { + request["stream_id"] = json!(lane); + } + send_json(&mut ws, request.clone()).await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + let parent = events.last().unwrap()["response"]["id"].as_str().unwrap().to_owned(); + assert_no_response_state(&fixture.pool).await; + request["previous_response_id"] = json!(parent); + request["input"] = json!("next turn"); + request["generate"] = json!(true); + send_json(&mut ws, request.clone()).await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + let child = events.last().unwrap()["response"]["id"].as_str().unwrap().to_owned(); + assert_no_response_state(&fixture.pool).await; + let bodies = mock.request_bodies().await; + let input = bodies.last().unwrap()["input"].to_string(); + assert!(input.contains("private prefix")); + assert!(input.contains("next turn")); + ws.close(None).await.unwrap(); + for store in [false, true] { + let mut fresh = connect_responses_ws(&gateway_url).await; + request["previous_response_id"] = json!(child); + request["store"] = json!(store); + send_json(&mut fresh, request.clone()).await; + let events = recv_until_completed(&mut fresh).await; + assert_eq!(events.last().unwrap()["error"]["code"], "previous_response_not_found"); + fresh.close(None).await.unwrap(); + } + assert_eq!(mock.request_bodies().await.len(), bodies.len()); + } + } +} + +#[tokio::test] +async fn websocket_unstored_prewarm_and_generation_work_without_storage() { + let mock = MockResponsesServer::start(vec![sse_response("resp_answer", "msg_answer", "READY")]).await; + let (gateway_url, _gateway) = spawn_gateway(persistence_disabled_state(&mock.url)).await; + let mut ws = connect_responses_ws(&gateway_url).await; + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", + "input":"warmup prefix", "generate":false, "store":false}), + ) + .await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + let parent = events.last().unwrap()["response"]["id"].as_str().unwrap(); + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", + "previous_response_id":parent, "input":"next", "store":false}), + ) + .await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + assert!( + mock.request_bodies().await[0]["input"] + .to_string() + .contains("warmup prefix") + ); +} + +async fn prewarm_private( + ws: &mut WebSocketStream>, + lane: Option<&str>, + text: &str, +) -> String { + let mut request = json!({"type":"response.create", "model":"test-model", "store":false, + "generate":false, "input":text}); + if let Some(lane) = lane { + request["stream_id"] = json!(lane); + } + send_json(ws, request).await; + let events = recv_until_completed(ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + events.last().unwrap()["response"]["id"].as_str().unwrap().to_owned() +} + +#[tokio::test] +async fn websocket_validation_failure_evicts_only_its_valid_lane_parent() { + for (routing, should_evict) in [(json!("source"), true), (json!("fork"), false), (Value::Null, false)] { + let mock = MockResponsesServer::start(vec![]).await; + let (gateway_url, _gateway) = spawn_gateway(persistence_disabled_state(&mock.url)).await; + let mut ws = connect_responses_ws(&gateway_url).await; + // Null routing must never be interpreted as the default lane for eviction. + let lane = if routing.is_null() { None } else { Some("source") }; + let parent = prewarm_private(&mut ws, lane, "private").await; + send_json( + &mut ws, + json!({"type":"response.create", "stream_id":routing, + "previous_response_id":parent, "model":7, "store":false, "input":[]}), + ) + .await; + assert_eq!(recv_json(&mut ws).await["type"], "error"); + let mut request = json!({"type":"response.create", "model":"test-model", + "previous_response_id":parent, "input":[], "store":false, "generate":false}); + if let Some(lane) = lane { + request["stream_id"] = json!(lane); + } + send_json(&mut ws, request).await; + let events = recv_until_completed(&mut ws).await; + if should_evict { + assert_eq!(events.last().unwrap()["error"]["code"], "previous_response_not_found"); + } else { + assert_eq!(events.last().unwrap()["type"], "response.completed"); + } + assert!(mock.request_bodies().await.is_empty()); + } +} + +#[tokio::test] +async fn websocket_stored_child_of_unstored_parent_survives_reconnect() { + let mock = MockResponsesServer::start(vec![sse_response("resp_final", "msg_final", "READY")]).await; + let fixture = storage_backed_state(&mock.url).await; + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let parent = prewarm_private(&mut ws, Some("source"), "private prefix").await; + send_json( + &mut ws, + json!({"type":"response.create", "stream_id":"source", "model":"test-model", + "previous_response_id":parent, "input":"stored child", "store":true, "generate":false}), + ) + .await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + let child = events.last().unwrap()["response"]["id"].as_str().unwrap().to_owned(); + let ids: Vec = sqlx::query_scalar("SELECT id FROM responses") + .fetch_all(fixture.pool.as_ref()) + .await + .unwrap(); + assert_eq!( + ids, + vec![child.clone()], + "only the explicitly stored child gets a response row" + ); + ws.close(None).await.unwrap(); + let mut fresh = connect_responses_ws(&gateway_url).await; + send_json( + &mut fresh, + json!({"type":"response.create", "model":"test-model", + "previous_response_id":child, "input":"after reconnect", "store":true}), + ) + .await; + let events = recv_until_completed(&mut fresh).await; + assert_eq!(events.last().unwrap()["type"], "response.completed"); + let input = mock.request_bodies().await[0]["input"].to_string(); + assert!(input.contains("private prefix")); + assert!(input.contains("stored child")); + assert!(input.contains("after reconnect")); +} + +#[tokio::test] +async fn websocket_unstored_custom_tool_round_trip() { + let mock = MockResponsesServer::start(vec![ + sse_custom_tool_call_response(), + sse_response("resp_after_custom", "msg_after_custom", "DONE"), + ]) + .await; + let (gateway_url, _gateway) = spawn_gateway(persistence_disabled_state(&mock.url)).await; + let mut ws = connect_responses_ws(&gateway_url).await; + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", "input":"apply patch", + "tools":[{"type":"custom", "name":"apply_patch"}], "store":false}), + ) + .await; + let events = recv_until_completed(&mut ws).await; + assert_eq!( + events.last().unwrap()["response"]["output"][0]["type"], + "custom_tool_call" + ); + let parent = events.last().unwrap()["response"]["id"].as_str().unwrap(); + send_json(&mut ws, json!({"type":"response.create", "model":"test-model", "store":false, + "previous_response_id":parent, "input":[{"type":"custom_tool_call_output", "call_id":"call_custom_1", "output":"Done!"}]})).await; + let events = recv_until_completed(&mut ws).await; + assert_eq!( + events.last().unwrap()["response"]["output"][0]["content"][0]["text"], + "DONE" + ); + let requests = mock.request_bodies().await; + let input = requests[1]["input"].as_array().unwrap(); + assert_eq!(input.iter().filter(|item| item["type"] == "function_call").count(), 1); + assert_eq!( + input + .iter() + .filter(|item| item["type"] == "function_call_output") + .count(), + 1 + ); +} + +#[tokio::test] +async fn websocket_overload_preserves_parent_for_accepted_queued_continuation() { + for byte_limit in [false, true] { + for lane in [None, Some("source")] { + let (mock, arrived, release) = MockResponsesServer::start_gated(sse_failed_response()).await; + let (gateway_url, _gateway) = spawn_gateway(persistence_disabled_state(&mock.url)).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let parent = prewarm_private(&mut ws, lane, "private parent").await; + let mut active = + json!({"type":"response.create", "model":"test-model", "store":false, "input":"independent turn"}); + if let Some(lane) = lane { + active["stream_id"] = json!(lane); + } + send_json(&mut ws, active.clone()).await; + tokio::time::timeout(std::time::Duration::from_secs(2), arrived) + .await + .unwrap() + .unwrap(); + let mut queued = active.clone(); + queued["previous_response_id"] = json!(parent); + queued["generate"] = json!(false); + queued["input"] = json!("accepted continuation"); + if byte_limit { + queued["metadata"] = json!({"padding":"x".repeat(7 * 1024 * 1024)}); + } + send_json(&mut ws, queued.clone()).await; + if !byte_limit { + let mut filler = active.clone(); + filler["model"] = json!(7); + for _ in 0..62 { + send_json(&mut ws, filler.clone()).await; + } + } + send_json(&mut ws, queued).await; + let rejected = recv_until_completed(&mut ws).await; + assert_eq!( + rejected.last().unwrap()["status"], + 429, + "overload must reply without waiting for active work" + ); + release.send(()).unwrap(); + let failed = recv_until_completed(&mut ws).await; + assert_eq!(failed.last().unwrap()["type"], "response.failed"); + let accepted = recv_until_completed(&mut ws).await; + assert_eq!( + accepted.last().unwrap()["type"], + "response.completed", + "rejection must not evict accepted work's parent" + ); + if !byte_limit { + for _ in 0..62 { + assert_eq!(recv_json(&mut ws).await["type"], "error"); + } + } + assert_eq!(mock.request_bodies().await.len(), 1, "queued prewarm stays local"); + } + } +} + +#[tokio::test] +async fn websocket_bounds_idle_lane_retention_and_preserves_existing_lanes() { + let mock = MockResponsesServer::start(vec![]).await; + let (gateway_url, _gateway) = spawn_gateway(persistence_disabled_state(&mock.url)).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let parent = prewarm_private(&mut ws, None, "default parent").await; + for lane in 1..128 { + prewarm_private(&mut ws, Some(&format!("lane-{lane}")), "private").await; + } + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", "stream_id":"overflow", + "input":[], "store":false, "generate":false}), + ) + .await; + assert_eq!(recv_json(&mut ws).await["status"], 429); + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", + "previous_response_id":parent, "input":[], "store":false, "generate":false}), + ) + .await; + assert_eq!( + recv_until_completed(&mut ws).await.last().unwrap()["type"], + "response.completed" + ); + assert!(mock.request_bodies().await.is_empty()); +} + +#[tokio::test] +async fn websocket_retention_budget_rejects_before_durable_write() { + let mock = MockResponsesServer::start(vec![]).await; + let fixture = storage_backed_state(&mock.url).await; + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let text = "x".repeat(9 * 1024 * 1024); + for lane in 0..3 { + prewarm_private(&mut ws, Some(&format!("lane-{lane}")), &text).await; + } + send_json( + &mut ws, + json!({"type":"response.create", "model":"test-model", "stream_id":"overflow", + "input":text, "store":true, "generate":false}), + ) + .await; + let events = recv_until_completed(&mut ws).await; + assert_eq!(events.len(), 1, "budget rejection precedes response.created"); + assert_eq!(events[0]["status"], 413); + assert_no_response_state(&fixture.pool).await; + assert!(mock.request_bodies().await.is_empty()); +} + +#[tokio::test] +async fn websocket_unstored_disconnect_cancels_active_and_queued_work() { + let (first_chunk, _) = phased_sse_response("resp_hanging", "msg_hanging", "never complete"); + let (mock, upstream_dropped) = MockResponsesServer::start_hanging(first_chunk).await; + let fixture = storage_backed_state(&mock.url).await; + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let parent = prewarm_private(&mut ws, None, "private").await; + let request = json!({"type":"response.create", "model":"test-model", "store":false, + "previous_response_id":parent, "input":"continue"}); + send_json(&mut ws, request.clone()).await; + assert_eq!(recv_json(&mut ws).await["type"], "response.created"); + send_json(&mut ws, request).await; + ws.close(None).await.unwrap(); + drop(ws); + tokio::time::timeout(std::time::Duration::from_secs(2), upstream_dropped) + .await + .unwrap() + .unwrap(); + tokio::time::timeout( + std::time::Duration::from_secs(2), + fixture.state.websocket_tracker.wait_until_idle(), + ) + .await + .unwrap(); + assert_no_response_state(&fixture.pool).await; + assert_eq!(mock.request_bodies().await.len(), 1); + let mut fresh = connect_responses_ws(&gateway_url).await; + send_json( + &mut fresh, + json!({"type":"response.create", "model":"test-model", "store":false, + "previous_response_id":parent, "input":[], "generate":false}), + ) + .await; + assert_eq!( + recv_json(&mut fresh).await["error"]["code"], + "previous_response_not_found" + ); +} + +#[tokio::test] +async fn websocket_unstored_fork_keeps_pinned_parent_after_source_failure() { + let (arrived_tx, arrived_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let mock = MockResponsesServer::start_with_responses(vec![ + MockResponse::Gated { + response: sse_response("resp_fork", "msg_fork", "FORK"), + arrived: arrived_tx, + release: release_rx, + }, + MockResponse::Static(sse_failed_response()), + ]) + .await; + let fixture = storage_backed_state(&mock.url).await; + let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await; + let mut ws = connect_responses_ws(&gateway_url).await; + let parent = prewarm_private(&mut ws, Some("source"), "shared private parent").await; + let mut request = json!({"type":"response.create", "model":"test-model", "store":false, + "stream_id":"fork", "previous_response_id":parent, "input":"fork input"}); + send_json(&mut ws, request.clone()).await; + tokio::time::timeout(std::time::Duration::from_secs(2), arrived_rx) + .await + .unwrap() + .unwrap(); + request["stream_id"] = json!("source"); + request["input"] = json!("source fails"); + send_json(&mut ws, request.clone()).await; + let failed = recv_until_completed(&mut ws).await; + assert_eq!(failed.last().unwrap()["type"], "response.failed"); + assert_eq!(failed.last().unwrap()["stream_id"], "source"); + request["generate"] = json!(false); + send_json(&mut ws, request.clone()).await; + assert_eq!(recv_json(&mut ws).await["error"]["code"], "previous_response_not_found"); + release_tx.send(()).unwrap(); + let fork = recv_until_completed(&mut ws).await; + assert_eq!(fork.last().unwrap()["type"], "response.completed"); + assert_eq!(fork.last().unwrap()["stream_id"], "fork"); + request["stream_id"] = json!("fork"); + request["previous_response_id"] = fork.last().unwrap()["response"]["id"].clone(); + request["input"] = json!([]); + send_json(&mut ws, request).await; + assert_eq!( + recv_until_completed(&mut ws).await.last().unwrap()["type"], + "response.completed" + ); + let requests = mock.request_bodies().await; + assert_eq!(requests.len(), 2); + assert!(requests[0]["input"].to_string().contains("shared private parent")); + assert!(!requests[0]["input"].to_string().contains("source fails")); + assert_no_response_state(&fixture.pool).await; +} diff --git a/docs/api/index.md b/docs/api/index.md index 400b1afc..99659169 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -84,11 +84,23 @@ continuations. Send one JSON text frame per turn: ``` The server normalizes the frame into the internal Responses request model and -uses the same response-store continuation path as HTTP. WebSocket replies are +uses the same core executor as HTTP, with connection-local continuation state. WebSocket replies are JSON Responses stream events, including `response.created`, `response.output_item.added`, `response.output_text.delta`, and `response.completed`. +`store: false` keeps response state only in memory on the active connection. Each +lane retains its latest response, including `generate: false` prewarm responses, +so you can continue with `previous_response_id` without a database. After reconnecting, +replay the full item history or a compacted window; an unstored response ID returns +`400 previous_response_not_found`. With `store: true`, an uncached response can be +loaded from durable storage. Explicit `conversation_id` requests retain the durable +Conversations API behavior. + +A failed same-lane continuation evicts its referenced cached parent. Failed forks +preserve the source lane's parent. Admission rejections (429) preserve existing +checkpoints and accepted queued work. Parent lookup happens when execution begins. + Set `stream_id` to a string containing 1 to 256 characters to multiplex responses over one connection. Requests with different `stream_id` values can run concurrently, while requests with the same value run first in, first out. @@ -97,7 +109,12 @@ its `stream_id`. Requests that omit `stream_id` share a default first-in, first-out lane for backward compatibility. A connection accepts at most 64 outstanding requests and 12 MiB of aggregate request data; additional requests receive a `429` error event -until capacity is available. Upstream SSE lines are limited to 256 KiB, normalized +until capacity is available. A connection retains at most 128 lanes, including the +default lane; reconnect to start new lanes after that limit. Each retained checkpoint +is limited to 32,768 items and 16 MiB of serialized state, with a 32 MiB connection +budget covering cached, active-parent and prepared replacement checkpoints. Exceeding +a checkpoint budget returns 413 before writing response state. These limits do not +measure total process memory. Upstream SSE lines are limited to 256 KiB, normalized executor events are limited to 1 MiB, and each request shares a 1 MiB response budget across MCP discovery, upstream rounds, and normalized gateway tool output. Every outbound WebSocket event, including `stream_id`, is limited to 1 MiB of @@ -112,11 +129,12 @@ Invalid requests are returned as JSON WebSocket error events: { "type": "error", "stream_id": "turn-1", - "status": 404, + "status": 400, "error": { - "message": "human-readable error details", - "type": "not_found", - "code": "not_found" + "message": "Previous response with id 'resp_missing' not found.", + "type": "invalid_request_error", + "code": "previous_response_not_found", + "param": "previous_response_id" } } ```