diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 254ce945..2cabb2a3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -469,7 +469,7 @@ scheduler switch. It is forwarded to vLLM for all supported declaration mixtures defaults to `false` when omitted. Whatever calls the model emits are executed under the global sliding window and each handler's same-tool safety policy. -#### `messages_loop.rs` / `messages_request.rs` / `messages_stream.rs` +#### `messages_context.rs` / `messages_loop.rs` / `messages_request.rs` / `messages_stream.rs` A **parallel, independent implementation** of the same shape of loop for the Anthropic Messages API. `messages_stream.rs`'s own header comment describes it as "structurally @@ -482,6 +482,18 @@ constants (`MAX_GATEWAY_TOOL_ROUNDS`, `GATEWAY_TOOL_TIMEOUT`) are duplicated and manually kept in sync with the Responses-side ones rather than shared — a known seam, not an oversight, per the future-consolidation note. +Both loops take a `MessagesRequestContext` (`messages_context.rs`), the per-request +type that replaced a bare `serde_json::Value` at that boundary. It holds two views of +one request: a typed `MessagesRequest` for reading `tools`/`stream`/`model`, and the +raw JSON body that is actually forwarded upstream. The raw body is deliberately *not* +re-serialized from the typed view — `ContentBlock` catches unmodeled block types in +`#[serde(other)] Unknown` and models only the fields the gateway reads, so a typed +round-trip would drop `cache_control` and `is_error` and collapse `image`/ +`redacted_thinking` into `{"type":"unknown"}`. The context owns every mutation the +loops make to that body (`force_stream`, `append_round`) and the native web-search +budget, so the two views cannot drift apart uncontrolled; `messages` and `system` are +reachable only through the raw body, never the typed view. + ### `storage/` — persistence - **`pool.rs`** — `DbPool = sqlx::Pool`, driver-agnostic across SQLite and diff --git a/crates/agentic-server-core/src/executor/messages_context.rs b/crates/agentic-server-core/src/executor/messages_context.rs new file mode 100644 index 00000000..37b15bd0 --- /dev/null +++ b/crates/agentic-server-core/src/executor/messages_context.rs @@ -0,0 +1,266 @@ +//! Per-request context for the Anthropic Messages gateway tool loops. +//! +//! The Messages loops are a pass-through, not a transform: the client's request +//! is forwarded to vLLM `/v1/messages` essentially untouched, and only `tools` +//! and `stream`/`messages` are read or rewritten. That rules out a typed +//! round-trip through [`MessagesRequest`] as the upstream body — `ContentBlock` +//! carries a `#[serde(other)] Unknown` catch-all, and several block types model +//! only the fields the gateway reads, so re-serializing would silently drop +//! `cache_control`, `is_error`, and every unmodeled block (`image`, +//! `redacted_thinking`, future provider extensions). +//! +//! So this context carries **two views of one request**, built once per request: +//! +//! * `raw` — the JSON body actually sent upstream. It +//! is the single source of truth for `messages` and `system`, and the only +//! thing the loops mutate. +//! * `typed` — the client's request as received, for safe field access to +//! `tools`, `stream`, and `model` in routing and the loops. +//! +//! The two are **not** kept byte-identical, and must not be confused: `typed` is +//! what the client sent, `raw` is what the gateway sends upstream. They diverge +//! wherever the gateway rewrites the body for upstream — today +//! `normalize_native_web_search` rewriting a native `web_search_20250305` +//! declaration into the ordinary function-tool shape vLLM accepts. Accordingly, +//! only the fields the loops never mutate are exposed off `typed` +//! ([`tools`](MessagesRequestContext::tools), +//! [`stream`](MessagesRequestContext::stream), +//! [`model`](MessagesRequestContext::model)); `messages` and `system` are +//! deliberately unreachable through it, so a stale typed view can never be read +//! back after a round is appended. + +use serde::Deserialize; +use serde_json::{Map, Value, json}; + +use crate::executor::error::{ExecutorError, ExecutorResult}; +use crate::executor::messages_request::{WebSearchBudget, normalize_native_web_search}; +use crate::types::messages::{MessagesRequest, ToolParam}; +use crate::utils::common::serialize_to_string; + +/// One `/v1/messages` request, in both the typed and raw views the gateway tool +/// loops need. See the module docs for why both exist. +#[derive(Debug)] +pub struct MessagesRequestContext { + /// The client's request as received. Read-only: routing and registry + /// construction only. + typed: MessagesRequest, + /// The upstream body. Mutated by the loops; the source of truth for + /// `messages` and `system`. + raw: Value, + /// Request-wide native web-search budget, derived while normalizing `raw`. + web_search_budget: WebSearchBudget, +} + +impl MessagesRequestContext { + /// Build the context from a request the caller has already parsed for + /// routing, plus the original body bytes it was parsed from. + /// + /// Reusing the caller's `typed` keeps each body parsed exactly once per + /// view: the routing parse is carried into the loop instead of being + /// discarded, and the raw parse only happens on the loop path, so a proxied + /// request never pays for a view it does not use. + /// + /// Native web-search declarations are validated and normalized here, before + /// a streaming handler commits its HTTP status — an invalid declaration must + /// surface as an error response, not as a mid-stream event. + /// + /// # Errors + /// Returns [`ExecutorError::JsonError`] if `body` is not valid JSON, or + /// [`ExecutorError::InvalidRequest`] if it carries an unsupported or invalid + /// native web-search declaration. + pub fn new(typed: MessagesRequest, body: &[u8]) -> ExecutorResult { + let raw = serde_json::from_slice(body).map_err(ExecutorError::JsonError)?; + Self::from_parts(typed, raw) + } + + /// Build the context from a raw JSON body alone, deriving the typed view + /// from it. + /// + /// Prefer [`new`](Self::new) when the caller has already parsed the request + /// for routing; this exists for callers that only hold a [`Value`]. + /// + /// # Errors + /// Returns [`ExecutorError::JsonError`] if `raw` is not a well-formed + /// Messages request, or [`ExecutorError::InvalidRequest`] if it carries an + /// unsupported or invalid native web-search declaration. + pub fn from_value(raw: Value) -> ExecutorResult { + // Deserializing from the parsed tree avoids re-lexing the body text. + let typed = MessagesRequest::deserialize(&raw).map_err(ExecutorError::JsonError)?; + Self::from_parts(typed, raw) + } + + fn from_parts(typed: MessagesRequest, mut raw: Value) -> ExecutorResult { + let web_search_budget = normalize_native_web_search(&mut raw)?; + Ok(Self { + typed, + raw, + web_search_budget, + }) + } + + /// The tools the client declared, for routing and registry construction. + /// + /// These are the client's declarations as received — before the upstream + /// normalization applied to `raw` — which is what the tool seam + /// needs to recognise a native server-tool declaration. + #[must_use] + pub fn tools(&self) -> Option<&Vec> { + self.typed.tools.as_ref() + } + + /// Whether the client asked for a streaming response. + #[must_use] + pub fn stream(&self) -> bool { + self.typed.stream + } + + /// The model the client requested. + #[must_use] + pub fn model(&self) -> &str { + &self.typed.model + } + + /// The body to POST upstream for the next round. + /// + /// # Errors + /// Returns [`ExecutorError::JsonError`] if the body cannot be serialized. + pub(super) fn upstream_body(&self) -> ExecutorResult { + serialize_to_string(&self.raw).map_err(ExecutorError::JsonError) + } + + /// Force the upstream streaming mode, regardless of what the client asked. + /// + /// Each loop drives its own rounds and so pins `stream` to what it can + /// consume; the client-facing mode is [`stream`](Self::stream), decided by + /// the handler before the loop starts. + pub(super) fn force_stream(&mut self, streaming: bool) { + self.raw["stream"] = Value::Bool(streaming); + } + + /// Reserve up to `requested` native web searches, returning how many may run. + pub(super) fn reserve_searches(&mut self, requested: usize) -> usize { + self.web_search_budget.reserve(requested) + } + + /// Append the model's assistant turn (preserving its `thinking`/`text`/ + /// `tool_use` blocks in order — F3) and a following user turn of + /// `tool_result`s, so the next upstream round sees the full conversation + /// state. These stay internal — the client never sees them (hide-the-call). + /// + /// # Errors + /// Returns [`ExecutorError::InvalidRequest`] if the body has no `messages` + /// array to append to. Unreachable for a context built through either + /// constructor, since `MessagesRequest::messages` is a required array — + /// erroring keeps it from silently no-opping into a loop that re-POSTs an + /// unchanged body until the round cap. + pub(super) fn append_round(&mut self, assistant_content: &[Value], tool_results: Vec) -> ExecutorResult<()> { + let messages = self + .raw + .get_mut("messages") + .and_then(Value::as_array_mut) + .ok_or_else(|| ExecutorError::InvalidRequest("request has no messages array".to_owned()))?; + messages.push(json!({ "role": "assistant", "content": assistant_content })); + // Built by hand rather than with `json!` so the tool outputs move in + // instead of being deep-copied — a web-search result runs to kilobytes. + let mut user = Map::new(); + user.insert("role".to_owned(), Value::String("user".to_owned())); + user.insert("content".to_owned(), Value::Array(tool_results)); + messages.push(Value::Object(user)); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn request() -> Value { + json!({ + "model": "qwen3", "max_tokens": 1024, "stream": true, + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "web_search", "type": "web_search_20250305", "max_uses": 2}] + }) + } + + #[test] + fn typed_view_reads_client_fields_and_raw_carries_upstream_normalization() { + let ctx = MessagesRequestContext::from_value(request()).unwrap(); + + assert_eq!(ctx.model(), "qwen3"); + assert!(ctx.stream()); + // The typed view keeps the client's native declaration, which is what + // the tool seam classifies on... + let tools = ctx.tools().expect("tools"); + assert_eq!(tools[0].name, "web_search"); + assert_eq!(tools[0].type_.as_deref(), Some("web_search_20250305")); + // ...while the raw body carries the function-tool shape vLLM accepts. + assert_eq!(ctx.raw["tools"][0]["name"], "web_search"); + assert!(ctx.raw["tools"][0].get("type").is_none()); + assert!(ctx.raw["tools"][0].get("input_schema").is_some()); + } + + #[test] + fn force_stream_overrides_the_client_mode_without_touching_the_typed_view() { + let mut ctx = MessagesRequestContext::from_value(request()).unwrap(); + ctx.force_stream(false); + + assert_eq!(ctx.raw["stream"], json!(false)); + assert!(ctx.stream(), "the client's requested mode is still readable"); + } + + #[test] + fn append_round_extends_the_raw_history_only() { + let mut ctx = MessagesRequestContext::from_value(request()).unwrap(); + let assistant = vec![json!({"type": "tool_use", "id": "t1", "name": "web_search", "input": {}})]; + ctx.append_round(&assistant, vec![json!({"type": "tool_result", "tool_use_id": "t1"})]) + .unwrap(); + + let messages = ctx.raw["messages"].as_array().expect("messages"); + assert_eq!(messages.len(), 3); + assert_eq!(messages[1]["role"], "assistant"); + assert_eq!(messages[1]["content"], json!(assistant)); + assert_eq!(messages[2]["role"], "user"); + assert_eq!(messages[2]["content"][0]["tool_use_id"], "t1"); + } + + #[test] + fn budget_is_shared_across_rounds() { + let mut ctx = MessagesRequestContext::from_value(request()).unwrap(); + assert_eq!(ctx.reserve_searches(1), 1); + assert_eq!(ctx.reserve_searches(3), 1, "max_uses caps the request-wide total"); + assert_eq!(ctx.reserve_searches(1), 0); + } + + #[test] + fn invalid_native_web_search_declaration_is_rejected_at_construction() { + let mut body = request(); + body["tools"][0]["max_uses"] = json!(0); + let error = MessagesRequestContext::from_value(body).unwrap_err(); + assert!(matches!(error, ExecutorError::InvalidRequest(_)), "{error:?}"); + } + + #[test] + fn unmodeled_blocks_and_cache_control_survive_in_the_raw_body() { + // The reason the raw view exists: a typed round-trip would drop these. + let body = json!({ + "model": "m", "max_tokens": 8, + "system": [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}, + {"type": "redacted_thinking", "data": "enc"} + ]}] + }); + let ctx = MessagesRequestContext::from_value(body.clone()).unwrap(); + assert_eq!(ctx.raw, body); + } + + #[test] + fn new_reuses_the_routing_parse() { + let body = serde_json::to_vec(&request()).unwrap(); + let typed: MessagesRequest = serde_json::from_slice(&body).unwrap(); + let ctx = MessagesRequestContext::new(typed, &body).unwrap(); + + assert_eq!(ctx.model(), "qwen3"); + assert_eq!(ctx.raw["messages"][0]["content"], "hi"); + } +} diff --git a/crates/agentic-server-core/src/executor/messages_loop.rs b/crates/agentic-server-core/src/executor/messages_loop.rs index 683d307e..075b1833 100644 --- a/crates/agentic-server-core/src/executor/messages_loop.rs +++ b/crates/agentic-server-core/src/executor/messages_loop.rs @@ -19,11 +19,12 @@ use serde_json::{Value, json}; use crate::executor::error::{ExecutorError, ExecutorResult}; use crate::executor::inference::fetch_response_json_with_headers; -use crate::executor::messages_request::{normalize_native_web_search, web_search_budget_exhausted_result}; +use crate::executor::messages_context::MessagesRequestContext; +use crate::executor::messages_request::web_search_budget_exhausted_result; use crate::executor::request::ExecutionContext; use crate::tool::ToolRegistry; use crate::types::messages::tool_seam; -use crate::utils::common::{deserialize_from_str, serialize_to_string}; +use crate::utils::common::deserialize_from_str; /// Max gateway rounds before the loop gives up. Each round is one upstream /// `/v1/messages` call. Shared with the streaming loop (`messages_stream`). @@ -71,36 +72,28 @@ pub struct MessagesResponse { pub headers: http::HeaderMap, } -/// The `tool_result` block for one executed gateway call, fed back next round. -/// (The model's own `tool_use` block is carried forward via the preserved -/// assistant content, not reconstructed here — see `append_round_to_history`.) -struct ResolvedCall { - tool_result_block: Value, -} - /// Run the Messages-native gateway tool loop and return the final assistant /// message (Anthropic JSON `Value`). /// -/// `request` is the client's parsed request body as JSON — forwarded upstream -/// with `stream:false` forced and its `messages` extended each round. +/// `ctx` carries the client's request in both views: its raw body is forwarded +/// upstream with `stream:false` forced and its `messages` extended each round. /// /// # Errors /// Returns [`ExecutorError`] on upstream failure or unparseable upstream JSON. /// Gateway-tool execution failures do **not** error — they become error /// `tool_result`s fed back to the model. pub async fn run_messages_loop( - mut request: Value, + mut ctx: MessagesRequestContext, registry: &ToolRegistry, exec_ctx: &ExecutionContext, upstream: &MessagesUpstream, ) -> ExecutorResult> { - let mut web_search_budget = normalize_native_web_search(&mut request)?; // The loop drives turns itself; force non-streaming upstream regardless of // what the client asked (the handler routes streaming elsewhere). - request["stream"] = Value::Bool(false); + ctx.force_stream(false); for _round in 0..MAX_GATEWAY_TOOL_ROUNDS { - let body = serialize_to_string(&request).map_err(ExecutorError::JsonError)?; + let body = ctx.upstream_body()?; let (resp_text, response_headers) = fetch_response_json_with_headers(body, &upstream.url, &exec_ctx.client, &upstream.headers).await?; let message: Value = deserialize_from_str(&resp_text).map_err(ExecutorError::JsonError)?; @@ -165,10 +158,9 @@ pub async fn run_messages_loop( // Pure gateway-tool round: execute the calls, then feed the model's FULL // assistant turn (thinking/text/tool_use, order preserved — F3) plus the // tool_results back for the next round. Gateway blocks stay internal. - let assistant_content = content.clone(); - let allowed_searches = web_search_budget.reserve(gateway_calls.len()); - let resolved = execute_gateway_calls(&gateway_calls, registry, gateway_map, allowed_searches).await; - append_round_to_history(&mut request, &assistant_content, &resolved); + let allowed_searches = ctx.reserve_searches(gateway_calls.len()); + let tool_results = execute_gateway_calls(&gateway_calls, registry, gateway_map, allowed_searches).await; + ctx.append_round(content, tool_results)?; } // Round budget exhausted — re-run once more is not attempted; return the @@ -189,20 +181,22 @@ pub async fn run_messages_loop( /// Execute the gateway-owned `tool_use` blocks concurrently, each bounded by the /// per-call timeout. A failure or timeout becomes an error `tool_result` (E5). +/// +/// Returns one `tool_result` block per call, fed back next round. (The model's +/// own `tool_use` block is carried forward via the preserved assistant content, +/// not reconstructed here — see [`MessagesRequestContext::append_round`].) async fn execute_gateway_calls( gateway_calls: &[Value], registry: &ToolRegistry, gateway_map: &tool_seam::GatewayToolMap, allowed_searches: usize, -) -> Vec { +) -> Vec { let futures = gateway_calls.iter().enumerate().map(|(index, block)| async move { let id = block.get("id").and_then(Value::as_str).unwrap_or_default(); let name = block.get("name").and_then(Value::as_str).unwrap_or_default(); if index >= allowed_searches { - return ResolvedCall { - tool_result_block: web_search_budget_exhausted_result(id), - }; + return web_search_budget_exhausted_result(id); } // F4: reject a malformed/absent input rather than dispatching with args @@ -229,25 +223,7 @@ async fn execute_gateway_calls( ) }; - ResolvedCall { - tool_result_block: tool_seam::tool_result_block(id, &output, is_error), - } + tool_seam::tool_result_block(id, &output, is_error) }); join_all(futures).await } - -/// Append the model's assistant turn (preserving its `thinking`/`text`/`tool_use` -/// blocks in order — F3) and a following user turn of `tool_result`s to the -/// request `messages`, so the next upstream round sees the full conversation -/// state. These stay internal — the client never sees them (hide-the-call). -fn append_round_to_history(request: &mut Value, assistant_content: &[Value], resolved: &[ResolvedCall]) { - let assistant = json!({ "role": "assistant", "content": assistant_content }); - let user = json!({ - "role": "user", - "content": resolved.iter().map(|r| r.tool_result_block.clone()).collect::>() - }); - if let Some(messages) = request.get_mut("messages").and_then(Value::as_array_mut) { - messages.push(assistant); - messages.push(user); - } -} diff --git a/crates/agentic-server-core/src/executor/messages_request.rs b/crates/agentic-server-core/src/executor/messages_request.rs index dca43a3d..df3e751f 100644 --- a/crates/agentic-server-core/src/executor/messages_request.rs +++ b/crates/agentic-server-core/src/executor/messages_request.rs @@ -7,6 +7,7 @@ use crate::tool::web_search::web_search_function_tool; use crate::types::messages::tool_seam::{NATIVE_WEB_SEARCH_TYPE, WEB_SEARCH_EXECUTOR}; /// Request-wide native web-search execution budget. +#[derive(Debug)] pub(super) struct WebSearchBudget { remaining: Option, } @@ -159,16 +160,6 @@ fn native_web_search_max_uses(request: &Value) -> ExecutorResult> Ok(max_uses) } -/// Validate native web-search declarations before an HTTP streaming response -/// commits its status and headers. -/// -/// # Errors -/// Returns [`ExecutorError::InvalidRequest`] when a native declaration uses an -/// unsupported version or invalid policy configuration. -pub fn validate_native_web_search_request(request: &Value) -> ExecutorResult<()> { - native_web_search_max_uses(request).map(drop) -} - /// Normalize native web-search declarations for an upstream endpoint that /// validates ordinary function-tool schemas, returning whether the body changed. /// diff --git a/crates/agentic-server-core/src/executor/messages_stream.rs b/crates/agentic-server-core/src/executor/messages_stream.rs index c2a42bd2..385523d7 100644 --- a/crates/agentic-server-core/src/executor/messages_stream.rs +++ b/crates/agentic-server-core/src/executor/messages_stream.rs @@ -23,7 +23,8 @@ use serde_json::{Value, json}; use crate::executor::error::{ExecutorError, ExecutorResult}; use crate::executor::inference::{BoxStream, response_lines, send_request}; -use crate::executor::messages_request::{normalize_native_web_search, web_search_budget_exhausted_result}; +use crate::executor::messages_context::MessagesRequestContext; +use crate::executor::messages_request::web_search_budget_exhausted_result; use crate::executor::request::ExecutionContext; use crate::proxy::processed_response_headers; use crate::tool::ToolRegistry; @@ -43,17 +44,16 @@ use crate::executor::messages_loop::{ /// Returns an executor error when the initial request cannot be serialized or /// when the upstream rejects it before streaming begins. pub async fn run_messages_stream( - mut request: Value, + mut ctx: MessagesRequestContext, registry: Arc, exec_ctx: Arc, upstream: MessagesUpstream, ) -> ExecutorResult> { - let mut web_search_budget = normalize_native_web_search(&mut request)?; - request["stream"] = Value::Bool(true); + ctx.force_stream(true); // Prime the first upstream request before the handler commits an HTTP 200. // This lets initial vLLM errors retain their original status and body. - let first_body = serialize_to_string(&request)?; + let first_body = ctx.upstream_body()?; let first_response = send_request( &exec_ctx.client, upstream.url(), @@ -72,9 +72,9 @@ pub async fn run_messages_stream( let response = if let Some(response) = prepared_response.take() { response } else { - let body = match serialize_to_string(&request) { + let body = match ctx.upstream_body() { Ok(b) => b, - Err(e) => { yield error_sse(&e.to_string()); return; } + Err(e) => { yield executor_error_sse(&e); return; } }; match send_request( &exec_ctx.client, @@ -118,14 +118,17 @@ pub async fn run_messages_stream( // the gateway tool_use (F3, streaming half). The gateway calls are // derived from the same buffered blocks for dispatch. let (assistant_content, calls) = acc.take_round(); - let allowed_searches = web_search_budget.reserve(calls.len()); - let resolved = execute_gateway_calls( + let allowed_searches = ctx.reserve_searches(calls.len()); + let tool_results = execute_gateway_calls( &calls, ®istry, &exec_ctx.messages_gateway_tools, allowed_searches, ).await; - append_round_to_history(&mut request, &assistant_content, &resolved); + if let Err(e) = ctx.append_round(&assistant_content, tool_results) { + yield executor_error_sse(&e); + return; + } } // Round budget exhausted. @@ -446,17 +449,19 @@ fn executor_error_sse(error: &ExecutorError) -> String { /// Execute reconstructed gateway calls (concurrent, per-call timeout). Errors /// become error `tool_result`s (E5). +/// +/// Returns one `tool_result` block per call, fed back next round. (The assistant +/// turn — including each call's `tool_use` block — is reconstructed from the +/// accumulator's buffered blocks in [`MessagesStreamAccumulator::take_round`].) async fn execute_gateway_calls( calls: &[StreamedCall], registry: &ToolRegistry, gateway_map: &tool_seam::GatewayToolMap, allowed_searches: usize, -) -> Vec { +) -> Vec { let futures = calls.iter().enumerate().map(|(index, c)| async move { if index >= allowed_searches { - return ResolvedStreamCall { - tool_result_block: web_search_budget_exhausted_result(&c.id), - }; + return web_search_budget_exhausted_result(&c.id); } // F4: reject a malformed/incomplete reconstructed input rather than // coercing to {} and dispatching the tool with args the model never sent. @@ -477,36 +482,11 @@ async fn execute_gateway_calls( } Err(reason) => (format!("{reason}; tool was not run"), true), }; - ResolvedStreamCall { - tool_result_block: tool_seam::tool_result_block(&c.id, &output, is_error), - } + tool_seam::tool_result_block(&c.id, &output, is_error) }); futures::future::join_all(futures).await } -/// The `tool_result` block for one executed gateway call, fed back next round. -/// (The assistant turn — including this call's `tool_use` block — is reconstructed -/// from the accumulator's buffered blocks in [`MessagesStreamAccumulator::take_round`].) -struct ResolvedStreamCall { - tool_result_block: Value, -} - -/// Append the model's full assistant turn (`thinking`/`text`/`signature` + -/// gateway `tool_use`, order preserved — F3) and a following user turn of `tool_result`s, -/// so the next upstream round sees the complete conversation state. These stay -/// internal — the client never sees the gateway call (hide-the-call). -fn append_round_to_history(request: &mut Value, assistant_content: &[Value], resolved: &[ResolvedStreamCall]) { - let assistant = json!({ "role": "assistant", "content": assistant_content }); - let user = json!({ - "role": "user", - "content": resolved.iter().map(|r| r.tool_result_block.clone()).collect::>() - }); - if let Some(messages) = request.get_mut("messages").and_then(Value::as_array_mut) { - messages.push(assistant); - messages.push(user); - } -} - #[cfg(test)] mod tests { use super::*; @@ -744,7 +724,7 @@ mod tests { calls.len(), ) .await; - let content = resolved[0].tool_result_block["content"].as_str().unwrap_or_default(); + let content = resolved[0]["content"].as_str().unwrap_or_default(); assert!( content.contains("invalid") || content.contains("malformed") || content.contains("could not"), "malformed args must yield an error tool_result, not an empty-arg dispatch: {content:?}" diff --git a/crates/agentic-server-core/src/executor/mod.rs b/crates/agentic-server-core/src/executor/mod.rs index 3f27aa2a..c83fe0ea 100644 --- a/crates/agentic-server-core/src/executor/mod.rs +++ b/crates/agentic-server-core/src/executor/mod.rs @@ -6,6 +6,7 @@ pub mod engine; pub mod error; pub mod function_sse; pub mod inference; +pub mod messages_context; pub mod messages_loop; mod messages_request; pub mod messages_stream; @@ -23,8 +24,9 @@ pub use compaction::compact_response; pub use engine::{BoxStream, ExecuteRequest, create_conversation, execute}; pub use error::{ExecutorError, ExecutorResult}; pub use inference::call_inference; +pub use messages_context::MessagesRequestContext; pub use messages_loop::{MessagesResponse, MessagesUpstream, run_messages_loop}; -pub use messages_request::{normalize_native_web_search_for_upstream, validate_native_web_search_request}; +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}; diff --git a/crates/agentic-server-core/tests/messages_loop_test.rs b/crates/agentic-server-core/tests/messages_loop_test.rs index 21f348f1..78be049e 100644 --- a/crates/agentic-server-core/tests/messages_loop_test.rs +++ b/crates/agentic-server-core/tests/messages_loop_test.rs @@ -14,7 +14,8 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use agentic_core::executor::{ - ConversationHandler, ExecutionContext, ExecutorResult, MessagesUpstream, ResponseHandler, run_messages_loop, + ConversationHandler, ExecutionContext, ExecutorResult, MessagesRequestContext, MessagesUpstream, ResponseHandler, + run_messages_loop, }; use agentic_core::storage::{ConversationStore, ResponseStore}; use agentic_core::tool::{ToolRegistry, WebSearchHandler}; @@ -205,9 +206,14 @@ async fn run_test_messages_loop( exec_ctx: &ExecutionContext, ) -> ExecutorResult { let upstream = MessagesUpstream::new(&exec_ctx.llm_base_url, None, reqwest::header::HeaderMap::new()); - run_messages_loop(request, registry, exec_ctx, &upstream) - .await - .map(|response| response.body) + run_messages_loop( + MessagesRequestContext::from_value(request)?, + registry, + exec_ctx, + &upstream, + ) + .await + .map(|response| response.body) } fn web_search_request() -> Value { @@ -289,7 +295,8 @@ async fn native_web_search_applies_domain_and_location_configuration() { let registry = build_tool_registry(&tools, &exec_ctx).await; let transport = MessagesUpstream::new(&exec_ctx.llm_base_url, None, reqwest::header::HeaderMap::new()); - run_messages_loop(request, ®istry, &exec_ctx, &transport) + let ctx = MessagesRequestContext::from_value(request).expect("request context"); + run_messages_loop(ctx, ®istry, &exec_ctx, &transport) .await .expect("loop runs"); @@ -316,7 +323,8 @@ async fn native_web_search_applies_blocked_domains() { let registry = build_tool_registry(&tools, &exec_ctx).await; let transport = MessagesUpstream::new(&exec_ctx.llm_base_url, None, reqwest::header::HeaderMap::new()); - run_messages_loop(request, ®istry, &exec_ctx, &transport) + let ctx = MessagesRequestContext::from_value(request).expect("request context"); + run_messages_loop(ctx, ®istry, &exec_ctx, &transport) .await .expect("loop runs"); @@ -351,7 +359,8 @@ async fn native_web_search_enforces_max_uses() { let registry = build_tool_registry(&tools, &exec_ctx).await; let transport = MessagesUpstream::new(&exec_ctx.llm_base_url, None, reqwest::header::HeaderMap::new()); - run_messages_loop(request, ®istry, &exec_ctx, &transport) + let ctx = MessagesRequestContext::from_value(request).expect("request context"); + run_messages_loop(ctx, ®istry, &exec_ctx, &transport) .await .expect("loop runs"); @@ -770,3 +779,83 @@ async fn messages_loop_preserves_claude_code_cache_control_across_rounds() { ); assert!(requests[1]["tools"][0].get("cache_control").is_none()); } + +/// The dual-view context exists so the upstream body stays the client's own +/// JSON. A typed round-trip through `MessagesRequest` would not: `ContentBlock` +/// models only the fields the gateway reads and catches everything else in +/// `#[serde(other)] Unknown`, which re-serializes as a literal +/// `{"type":"unknown"}` block. This locks in that the raw body — not the typed +/// view — is what reaches vLLM, across a gateway round. +#[tokio::test] +async fn messages_loop_preserves_unmodeled_blocks_and_tool_result_fields_across_rounds() { + let round0 = serde_json::json!({ + "id": "m", "type": "message", "role": "assistant", "model": "qwen3", + "content": [{"type": "tool_use", "id": "t2", "name": "WebSearch", "input": {"query": "rust"}}], + "stop_reason": "tool_use", "usage": {"input_tokens": 5, "output_tokens": 3} + }); + let round1 = serde_json::json!({ + "id": "m2", "type": "message", "role": "assistant", "model": "qwen3", + "content": [{"type": "text", "text": "Done."}], + "stop_reason": "end_turn", "usage": {"input_tokens": 5, "output_tokens": 3} + }); + let (vllm_url, upstream, _v) = spawn_mock_vllm_messages(vec![round0, round1]).await; + let (search_url, _search_requests, _s) = spawn_mock_search().await; + let mut exec_ctx = build_exec_ctx(&vllm_url, &search_url).await; + exec_ctx.messages_gateway_tools = GatewayToolMap::from_pairs([("WebSearch", "web_search")]); + + // Every block here is either unmodeled by `ContentBlock` (`image`, + // `redacted_thinking`) or carries fields it does not model (`citations`, + // `is_error`, `cache_control`, a non-text `tool_result` part). + let request = serde_json::json!({ + "model": "qwen3", "max_tokens": 1024, "stream": false, + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "What is this?", + "citations": [{"type": "web_search_result_location", "url": "https://example.com"}]}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}} + ]}, + {"role": "assistant", "content": [ + {"type": "redacted_thinking", "data": "encrypted-blob"}, + {"type": "tool_use", "id": "t1", "name": "WebSearch", "input": {"query": "prior"}} + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "t1", "is_error": true, + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + "content": [ + {"type": "text", "text": "search failed"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "BBBB"}} + ]} + ]} + ], + "tools": [{"name": "WebSearch", "description": "Search the web", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}}] + }); + let original_messages = request["messages"].clone(); + let tools: Vec = serde_json::from_value(request["tools"].clone()).unwrap(); + let registry = build_tool_registry(&tools, &exec_ctx).await; + + let result = run_test_messages_loop(request, ®istry, &exec_ctx).await.unwrap(); + assert_eq!(result["stop_reason"], "end_turn"); + + let requests = upstream.requests.lock().await; + assert_eq!(requests.len(), 2, "tool round + final round"); + for upstream_request in requests.iter() { + let messages = upstream_request["messages"].as_array().expect("messages"); + assert_eq!( + &messages[..3], + original_messages.as_array().unwrap().as_slice(), + "the client's own blocks reach upstream untouched" + ); + // The failure mode a typed round-trip would introduce. + let rendered = serde_json::to_string(upstream_request).unwrap(); + assert!( + !rendered.contains(r#""type":"unknown""#), + "no block collapsed to Unknown" + ); + } + // Round 2 carries the appended gateway turn on top of the untouched prefix. + let round_two = requests[1]["messages"].as_array().expect("messages"); + assert_eq!(round_two.len(), 5, "3 client turns + assistant turn + tool_result turn"); + assert_eq!(round_two[3]["content"][0]["id"], "t2"); + assert_eq!(round_two[4]["content"][0]["type"], "tool_result"); +} diff --git a/crates/agentic-server-core/tests/messages_stream_test.rs b/crates/agentic-server-core/tests/messages_stream_test.rs index 71588f10..8a6f80b5 100644 --- a/crates/agentic-server-core/tests/messages_stream_test.rs +++ b/crates/agentic-server-core/tests/messages_stream_test.rs @@ -12,7 +12,8 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use agentic_core::executor::{ - BoxStream, ConversationHandler, ExecutionContext, MessagesUpstream, ResponseHandler, run_messages_stream, + BoxStream, ConversationHandler, ExecutionContext, MessagesRequestContext, MessagesUpstream, ResponseHandler, + run_messages_stream, }; use agentic_core::storage::{ConversationStore, ResponseStore}; use agentic_core::tool::{ToolRegistry, WebSearchHandler}; @@ -180,7 +181,8 @@ async fn run_test_messages_stream( exec_ctx: Arc, ) -> BoxStream { let upstream = MessagesUpstream::new(&exec_ctx.llm_base_url, None, reqwest::header::HeaderMap::new()); - run_messages_stream(request, registry, exec_ctx, upstream) + let ctx = MessagesRequestContext::from_value(request).expect("request context"); + run_messages_stream(ctx, registry, exec_ctx, upstream) .await .map(|response| response.body) .unwrap() diff --git a/crates/agentic-server/src/handler/http/messages.rs b/crates/agentic-server/src/handler/http/messages.rs index 0f0d0c3e..83f0b331 100644 --- a/crates/agentic-server/src/handler/http/messages.rs +++ b/crates/agentic-server/src/handler/http/messages.rs @@ -7,8 +7,8 @@ use http::HeaderMap; use tracing::debug; use agentic_core::executor::{ - ExecutorError, MessagesUpstream, normalize_native_web_search_for_upstream, run_messages_loop, run_messages_stream, - validate_native_web_search_request, + ExecutorError, MessagesRequestContext, MessagesUpstream, normalize_native_web_search_for_upstream, + run_messages_loop, run_messages_stream, }; use agentic_core::proxy::{ ProxyAuth, ProxyBody, ProxyRequest, ProxyResponse, error_response_for_auth, proxy_request_with_path, @@ -73,7 +73,7 @@ async fn execute_messages( state: &AppState, headers: &HeaderMap, query: Option<&str>, - req: &MessagesRequest, + req: MessagesRequest, body: &Bytes, ) -> Response { // Build the request-scoped registry from the declared tools (M6). Gateway @@ -87,28 +87,29 @@ async fn execute_messages( Err(e) => return messages_error_response(ExecutorError::from(e)), }; - // Parse the raw body to a JSON Value the loop forwards upstream untouched — - // preserving every Anthropic field (tool_choice, stop_sequences, …). - let request_json: serde_json::Value = match serde_json::from_slice(body) { - Ok(v) => v, - Err(e) => return messages_error_response(ExecutorError::from(e)), + // One context per request, carrying both views: the typed request for field + // access, and the raw body the loop forwards upstream untouched — preserving + // every Anthropic field (tool_choice, stop_sequences, cache_control, and + // block types the gateway does not model). Reuses the routing parse, so + // neither view is built twice. Native web-search declarations are validated + // here, before a streaming response commits its status and headers. + let ctx = match MessagesRequestContext::new(req, body) { + Ok(ctx) => ctx, + Err(e) => return messages_error_response(e), }; - if let Err(error) = validate_native_web_search_request(&request_json) { - return messages_error_response(error); - } let upstream = MessagesUpstream::new( &state.exec_ctx.llm_base_url, query, upstream_request_headers(headers, &state.proxy_state.config, ProxyAuth::Anthropic), ); - if req.stream { - match run_messages_stream(request_json, Arc::new(registry), Arc::clone(&state.exec_ctx), upstream).await { + if ctx.stream() { + match run_messages_stream(ctx, Arc::new(registry), Arc::clone(&state.exec_ctx), upstream).await { Ok(response) => sse_response_with_headers(response.body, response.headers), Err(e) => messages_error_response(e), } } else { - match run_messages_loop(request_json, ®istry, &state.exec_ctx, &upstream).await { + match run_messages_loop(ctx, ®istry, &state.exec_ctx, &upstream).await { Ok(message) => { let mut response = axum::Json(message.body).into_response(); response.headers_mut().extend(message.headers); @@ -141,7 +142,7 @@ pub async fn messages(State(state): State, request: Request) -> Respon "routing HTTP messages request" ); if route_to_loop { - return execute_messages(&state, &parts.headers, parts.uri.query(), &req, &bytes).await; + return execute_messages(&state, &parts.headers, parts.uri.query(), req, &bytes).await; } }