diff --git a/Cargo.lock b/Cargo.lock index 3bb2425a..b70e9dde 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,23 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "agentic-llm-d" +version = "0.5.0" +dependencies = [ + "agentic-server-core", + "axum", + "clap", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", +] + [[package]] name = "agentic-praxis" version = "0.5.0" diff --git a/Cargo.toml b/Cargo.toml index b4f812a8..e584dc70 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,7 +33,7 @@ reqwest = { version = "0.12", default-features = false } rmcp-reqwest = { package = "reqwest", version = "0.13.2", default-features = false, features = ["json", "stream", "rustls"] } rmcp = { version = "1.8", default-features = false } serde = { version = "1", features = ["derive"] } -serde_json = "1" +serde_json = { version = "1", features = ["raw_value"] } thiserror = "2" tokio = { version = "1", features = ["full"] } tokio-util = "0.7" diff --git a/Dockerfile b/Dockerfile index ff5a4131..0443474f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -34,9 +34,11 @@ COPY crates ./crates RUN cargo clean \ -p agentic-server-core \ -p agentic-server \ - -p agentic-praxis && \ - cargo build --locked --release -p agentic-server && \ - install -Dm755 -s target/release/agentic-server /out/agentic-server + -p agentic-praxis \ + -p agentic-llm-d && \ + cargo build --locked --release -p agentic-server -p agentic-llm-d && \ + install -Dm755 -s target/release/agentic-server /out/agentic-server && \ + install -Dm755 -s target/release/agentic-llm-d /out/agentic-llm-d FROM debian:${DEBIAN_VERSION}-slim@${DEBIAN_IMAGE_DIGEST} AS runtime @@ -51,6 +53,7 @@ RUN apt-get update && \ chmod g=u,g+s /var/lib/agentic-api COPY --from=rust-build /out/agentic-server /usr/local/bin/agentic-server +COPY --from=rust-build /out/agentic-llm-d /usr/local/bin/agentic-llm-d COPY --chmod=0755 docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh ARG OCI_CREATED="" diff --git a/crates/agentic-llm-d/Cargo.toml b/crates/agentic-llm-d/Cargo.toml new file mode 100644 index 00000000..099f7730 --- /dev/null +++ b/crates/agentic-llm-d/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "agentic-llm-d" +description = "Backend mode for agentic-api: split-execution endpoints for the llm-d coordinator" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +publish = false + +[dependencies] +agentic-core.workspace = true +axum.workspace = true +clap.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true + +[dev-dependencies] +agentic-core = { workspace = true, features = [] } +reqwest = { workspace = true, features = ["json"] } +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/crates/agentic-llm-d/src/handler.rs b/crates/agentic-llm-d/src/handler.rs new file mode 100644 index 00000000..0b08f50a --- /dev/null +++ b/crates/agentic-llm-d/src/handler.rs @@ -0,0 +1,101 @@ +//! The endpoints, and the axum glue they need. They trust their caller: +//! `hydrate` returns full conversation history and `persist` writes a turn from +//! a caller-supplied context, so bind this binary cluster-internal only. + +use std::time::Duration; + +use axum::body::Body; +use axum::extract::{Request, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use serde::Deserialize; +use serde::de::DeserializeOwned; +use serde_json::value::RawValue; +use tracing::warn; + +use agentic_core::executor::ExecutorError; +use agentic_core::executor::request::SplitContext; +use agentic_core::executor::split::{self, UpstreamBody}; +use agentic_core::types::request_response::RequestPayload; + +use crate::InternalState; + +const MAX_BODY_SIZE: usize = 10 * 1024 * 1024; +/// Readiness means storage answers — llm-d owns the model fleet. +const STORAGE_PROBE_TIMEOUT: Duration = Duration::from_secs(2); + +/// Body of `POST /internal/persist`: the context, plus exactly one response form. +#[derive(Debug, Deserialize)] +pub struct PersistRequest { + context: SplitContext, + response: Option>, + sse: Option, +} + +pub async fn health() -> StatusCode { + StatusCode::OK +} + +pub async fn ready(State(state): State) -> StatusCode { + if state.exec_ctx.storage_ready(STORAGE_PROBE_TIMEOUT).await { + StatusCode::OK + } else { + StatusCode::SERVICE_UNAVAILABLE + } +} + +pub async fn internal_hydrate(State(state): State, req: Request) -> Response { + let payload: RequestPayload = match read_json(req.into_body()).await { + Ok(payload) => payload, + Err(response) => return response, + }; + match split::hydrate(payload, state.exec_ctx.as_ref()).await { + Ok(hydration) => axum::Json(hydration).into_response(), + Err(error) => error_response(error), + } +} + +pub async fn internal_persist(State(state): State, req: Request) -> Response { + let PersistRequest { context, response, sse } = match read_json(req.into_body()).await { + Ok(request) => request, + Err(response) => return response, + }; + // serde rejects `RawValue` inside `flatten`/`untagged`, so the wire cannot + // type "exactly one of". Narrow here. + let upstream = match (response.as_deref(), sse.as_deref()) { + (Some(json), None) => UpstreamBody::Json(json.get()), + (None, Some(sse)) => UpstreamBody::Sse(sse), + _ => { + let message = "exactly one of `response` or `sse` is required".to_owned(); + return error_response(ExecutorError::InvalidRequest(message)); + } + }; + match split::persist(context, upstream, state.exec_ctx.as_ref()).await { + Ok(payload) => axum::Json(payload).into_response(), + Err(error) => error_response(error), + } +} + +/// Renders an error with the status and envelope core defines. +fn error_response(error: ExecutorError) -> Response { + let status = error.http_status(); + warn!("backend error ({status}): {error}"); + json(status, error.into_response_body()) +} + +#[allow(clippy::result_large_err)] // an axum `Response` is the idiomatic error here +async fn read_json(body: Body) -> Result { + let too_large = br#"{"error":{"type":"invalid_request_error","message":"request body too large"}}"#; + let bytes = axum::body::to_bytes(body, MAX_BODY_SIZE) + .await + .map_err(|_| json(StatusCode::PAYLOAD_TOO_LARGE, too_large.to_vec()))?; + serde_json::from_slice(&bytes).map_err(|error| error_response(ExecutorError::from(error))) +} + +fn json(status: StatusCode, body: Vec) -> Response { + Response::builder() + .status(status) + .header("Content-Type", "application/json") + .body(Body::from(body)) + .expect("valid response") +} diff --git a/crates/agentic-llm-d/src/lib.rs b/crates/agentic-llm-d/src/lib.rs new file mode 100644 index 00000000..0588e8a4 --- /dev/null +++ b/crates/agentic-llm-d/src/lib.rs @@ -0,0 +1,29 @@ +//! Backend mode — agentic-api as state services for an orchestrator that runs +//! inference itself (the llm-d coordinator). Nothing here proxies or calls a model. + +pub mod handler; +pub mod runner; + +use std::sync::Arc; + +use axum::Router; +use axum::routing::{get, post}; + +use agentic_core::executor::ExecutionContext; + +/// All the endpoints need. The gateway's `AppState` carries eight more fields, +/// every one for machinery backend mode does not run. +#[derive(Clone)] +pub struct InternalState { + pub exec_ctx: Arc, +} + +/// The whole surface: two split-execution endpoints and two probes. +pub fn build_internal_router(state: InternalState) -> Router { + Router::new() + .route("/health", get(handler::health)) + .route("/ready", get(handler::ready)) + .route("/internal/hydrate", post(handler::internal_hydrate)) + .route("/internal/persist", post(handler::internal_persist)) + .with_state(state) +} diff --git a/crates/agentic-llm-d/src/main.rs b/crates/agentic-llm-d/src/main.rs new file mode 100644 index 00000000..b6e45200 --- /dev/null +++ b/crates/agentic-llm-d/src/main.rs @@ -0,0 +1,50 @@ +//! `agentic-llm-d` — the coordinator calls `hydrate`, runs inference against its +//! own model fleet, then calls `persist`. This binary does neither. + +use clap::Parser; +use tokio_util::sync::CancellationToken; + +use agentic_core::config::{Config, PostgresConfig, SqliteConfig, ToolRuntimeConfig}; +use agentic_llm_d::runner; + +#[derive(Parser)] +#[command(name = "agentic-llm-d", about = "agentic-api backend mode for the llm-d coordinator")] +struct Cli { + /// Keep this cluster-internal: the endpoints trust their caller. + #[arg(long, env = "AGENTIC_LLM_D_HOST", default_value = "127.0.0.1")] + host: String, + #[arg(long, env = "AGENTIC_LLM_D_PORT", default_value_t = 8081)] + port: u16, + /// Defaults to the local database under the agentic-api home. + #[arg(long, env = "DATABASE_URL")] + db_url: Option, +} + +#[tokio::main] +async fn main() -> Result<(), runner::Error> { + tracing_subscriber::fmt::init(); + let cli = Cli::parse(); + + // Every model-facing field is inert here; llm-d owns the fleet. + let config = Config { + llm_api_base: String::new(), + openai_api_key: None, + llm_ready_timeout_s: 0.0, + llm_ready_interval_s: 0.0, + skip_llm_ready_check: true, + db_url: cli.db_url, + postgres: PostgresConfig::default(), + sqlite: SqliteConfig::default(), + tools: ToolRuntimeConfig::default(), + }; + + let shutdown = CancellationToken::new(); + let on_signal = shutdown.clone(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + on_signal.cancel(); + } + }); + + runner::serve(&config, &cli.host, cli.port, shutdown).await +} diff --git a/crates/agentic-llm-d/src/runner.rs b/crates/agentic-llm-d/src/runner.rs new file mode 100644 index 00000000..550b1db1 --- /dev/null +++ b/crates/agentic-llm-d/src/runner.rs @@ -0,0 +1,34 @@ +//! Build state, bind, serve, drain. + +use std::sync::Arc; + +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; +use tracing::info; + +use agentic_core::config::Config; +use agentic_core::executor::ExecutionContext; + +use crate::{InternalState, build_internal_router}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("failed to build the execution context: {0}")] + Context(#[from] agentic_core::error::Error), + #[error("failed to bind or serve: {0}")] + Io(#[from] std::io::Error), +} + +/// Opens storage from `config` and serves the internal router until `shutdown`. +/// +/// # Errors +/// If the storage pool cannot be opened, or the address cannot be bound. +pub async fn serve(config: &Config, host: &str, port: u16, shutdown: CancellationToken) -> Result<(), Error> { + let exec_ctx = Arc::new(ExecutionContext::from_config(config).await?); + let listener = TcpListener::bind(format!("{host}:{port}")).await?; + info!("agentic-llm-d listening on {host}:{port} — no proxy, no inference"); + axum::serve(listener, build_internal_router(InternalState { exec_ctx })) + .with_graceful_shutdown(async move { shutdown.cancelled().await }) + .await?; + Ok(()) +} diff --git a/crates/agentic-server-core/src/executor/accumulator.rs b/crates/agentic-server-core/src/executor/accumulator.rs index 2c68b542..fbb38a7b 100644 --- a/crates/agentic-server-core/src/executor/accumulator.rs +++ b/crates/agentic-server-core/src/executor/accumulator.rs @@ -344,6 +344,11 @@ impl ResponseAccumulator { } } + /// Ask before `finish_stream`, which forces an unterminated stream to `completed`. + pub(crate) fn saw_terminal_frame(&self) -> bool { + self.status != ResponseStatus::InProgress + } + pub(crate) fn finish_stream(&mut self) { self.finalize_all(); if self.status == ResponseStatus::InProgress { diff --git a/crates/agentic-server-core/src/executor/mod.rs b/crates/agentic-server-core/src/executor/mod.rs index 7f261d72..d20933d1 100644 --- a/crates/agentic-server-core/src/executor/mod.rs +++ b/crates/agentic-server-core/src/executor/mod.rs @@ -13,6 +13,7 @@ pub mod modes; pub mod persist; pub mod rehydrate; pub mod request; +pub mod split; mod gateway; pub mod gateway_accumulator; diff --git a/crates/agentic-server-core/src/executor/request.rs b/crates/agentic-server-core/src/executor/request.rs index 158b8e22..2c919f6f 100644 --- a/crates/agentic-server-core/src/executor/request.rs +++ b/crates/agentic-server-core/src/executor/request.rs @@ -1,6 +1,8 @@ use std::sync::Arc; use std::time::Duration; +use serde::{Deserialize, Serialize}; + use crate::config::{Config, default_database_url}; use crate::error::Error; use crate::executor::modes::{ConversationHandler, ResponseHandler}; @@ -9,9 +11,10 @@ use crate::storage::{ ConversationStore, ConversationVersion, DatabaseBackend, ResponseStore, create_pool_with_schema_and_configs, }; use crate::tool::{GatewayExecutor, GatewayExecutors}; -use crate::types::io::InputItem; +use crate::types::io::{InputItem, ResponsesInput, ToolChoice}; use crate::types::messages::GatewayToolMap; use crate::types::request_response::{RequestPayload, ResponsePayload}; +use crate::types::tools::ResponsesTool; /// Env var configuring client-tool → gateway-executor aliases for `/v1/messages` /// (e.g. `WebSearch=web_search`). Empty/unset means no aliases — client @@ -19,6 +22,8 @@ use crate::types::request_response::{RequestPayload, ResponsePayload}; const GATEWAY_TOOL_ALIASES_ENV: &str = "MESSAGES_GATEWAY_TOOL_ALIASES"; /// Context built by `rehydrate_conversation`, threaded through the execute pipeline. +/// +/// Converts to and from [`SplitContext`] when a turn crosses a process boundary. #[derive(Debug)] pub struct RequestContext { /// Untouched original request from the client. @@ -49,6 +54,54 @@ impl RequestContext { } } +/// Wire form of a [`RequestContext`]: the subset that survives leaving the process. +/// +/// Absent on purpose: `enriched_request` (the conversation, already in flight as +/// the request — rebuilt on return) and `new_input_items` (derived). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SplitContext { + /// Response id reserved for this turn (`resp_` prefix). + pub response_id: String, + /// The client's request as received, continuation ids intact. + pub original_request: RequestPayload, + pub conversation_id: Option, + /// Resolved against the stored chain; the only part of `enriched_request` + /// that `original_request` cannot supply. + pub effective_tools: Option>, + pub effective_tool_choice: Option, +} + +impl From for SplitContext { + fn from(ctx: RequestContext) -> Self { + Self { + response_id: ctx.response_id, + original_request: ctx.original_request, + conversation_id: ctx.conversation_id, + effective_tools: ctx.enriched_request.tools, + effective_tool_choice: ctx.enriched_request.tool_choice, + } + } +} + +impl From for RequestContext { + fn from(wire: SplitContext) -> Self { + let new_input_items = Vec::from(&wire.original_request.input); + let mut enriched_request = wire.original_request.clone(); + enriched_request.previous_response_id = None; + enriched_request.input = ResponsesInput::Items(new_input_items.clone()); + enriched_request.tools = wire.effective_tools; + enriched_request.tool_choice = wire.effective_tool_choice; + Self { + original_request: wire.original_request, + enriched_request, + new_input_items, + response_id: wire.response_id, + conversation_id: wire.conversation_id, + conversation_version: None, + } + } +} + /// Runtime dependencies passed into `execute()`. /// /// Owns the storage handlers, HTTP client, and LLM endpoint configuration. diff --git a/crates/agentic-server-core/src/executor/split.rs b/crates/agentic-server-core/src/executor/split.rs new file mode 100644 index 00000000..e93b5736 --- /dev/null +++ b/crates/agentic-server-core/src/executor/split.rs @@ -0,0 +1,107 @@ +//! Split execution — hydration and persistence as separately callable steps, +//! for an orchestrator (e.g. the llm-d coordinator) that runs inference itself. +//! +//! [`hydrate`] returns the stateless upstream request plus the [`RequestContext`] +//! for the turn; the orchestrator calls the model, then passes the context and +//! the response to [`persist`]. +//! +//! Both halves call the same steps the in-process flow uses rather than +//! duplicating them — [`rehydrate_conversation`], then +//! `payload_from_upstream`, `persist_if_needed`. The context is the same type +//! either way; it serializes through a reduced wire form so only the fields that +//! mean anything off-process make the trip. + +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; + +use crate::executor::error::{ExecutorError, ExecutorResult}; +use crate::executor::persist::persist_if_needed; +use crate::executor::rehydrate::rehydrate_conversation; +use crate::executor::request::{ExecutionContext, RequestContext, SplitContext}; +pub use crate::executor::upstream::UpstreamBody; +use crate::executor::upstream::{payload_from_upstream, upstream_request_json}; +use crate::types::event::ResponseStatus; +use crate::types::request_response::{RequestPayload, ResponsePayload}; + +/// Result of [`hydrate`]. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Hydration { + /// Upstream request body: history inlined, no continuation or storage fields. + /// Raw JSON — the caller forwards it without interpreting it. + pub request: Box, + /// Echo back to [`persist`] unchanged. + pub context: SplitContext, +} + +/// Rehydrates history and builds the upstream request. +/// +/// Rejects requests that cannot be split; see [`ensure_splittable`], which is +/// public so a caller can pre-validate for a cleaner error. +/// +/// # Errors +/// [`ExecutorError::InvalidRequest`] for a request that cannot be split, +/// not-found for an unknown `previous_response_id`, or a storage error. +pub async fn hydrate(request: RequestPayload, exec_ctx: &ExecutionContext) -> ExecutorResult { + ensure_splittable(&request)?; + let ctx = rehydrate_conversation(request, exec_ctx).await?; + let stream = ctx.original_request.stream; + let request = RawValue::from_string(upstream_request_json(&ctx, stream)?).map_err(ExecutorError::JsonError)?; + Ok(Hydration { + request, + context: ctx.into(), + }) +} + +/// Persists the turn and builds the final envelope. +/// +/// For a stream the caller has already relayed the frames, so nothing is +/// re-emitted here. Stored when the request set `store` or continues a chain — +/// `store: false` alone skips storage, but not alongside `previous_response_id`. +/// +/// # Errors +/// [`ExecutorError::InvalidRequest`] for a non-terminal response or a cut-short +/// stream, a parse error for an invalid body, or a storage error. +pub async fn persist( + context: SplitContext, + upstream: UpstreamBody<'_>, + exec_ctx: &ExecutionContext, +) -> ExecutorResult { + let ctx = RequestContext::from(context); + let payload = payload_from_upstream(&ctx, upstream)?; + + // The in-process flow silently skips non-terminal statuses; here that would + // return an envelope whose id can never be continued, so reject it. + if !matches!( + payload.status.parse::().unwrap_or_default(), + ResponseStatus::Completed | ResponseStatus::Incomplete + ) { + return Err(ExecutorError::InvalidRequest(format!( + "upstream response status '{}' cannot be persisted", + payload.status + ))); + } + + persist_if_needed( + payload.clone(), + ctx, + exec_ctx.conv_handler.clone(), + exec_ctx.resp_handler.clone(), + ) + .await?; + Ok(payload) +} + +/// Rejects requests that cannot cross a process boundary — each needs state the +/// in-process flow keeps between steps. Not a limit of hydration itself; callers +/// running the whole turn in-process should not call this. +/// +/// # Errors +/// [`ExecutorError::InvalidRequest`] naming the feature that cannot be split. +pub fn ensure_splittable(request: &RequestPayload) -> ExecutorResult<()> { + if let Some(feature) = request.in_process_feature() { + return Err(ExecutorError::InvalidRequest(format!( + "{feature} is not supported for split execution" + ))); + } + Ok(()) +} diff --git a/crates/agentic-server-core/src/executor/upstream.rs b/crates/agentic-server-core/src/executor/upstream.rs index 90a48c80..534bed79 100644 --- a/crates/agentic-server-core/src/executor/upstream.rs +++ b/crates/agentic-server-core/src/executor/upstream.rs @@ -32,6 +32,15 @@ pub(super) struct StreamPayload { pub(super) deferred_events: Vec, } +/// Builds the JSON body sent upstream, so the request shape is defined once. +/// +/// # Errors +/// A tool-configuration or serialization failure. +pub(super) fn upstream_request_json(ctx: &RequestContext, stream: bool) -> ExecutorResult { + let request = ctx.enriched_request.to_upstream_request(stream)?; + serialize_to_string(&request).map_err(ExecutorError::JsonError) +} + pub(super) async fn fetch_blocking_payload( ctx: &RequestContext, exec_ctx: &ExecutionContext, @@ -39,20 +48,64 @@ pub(super) async fn fetch_blocking_payload( ) -> ExecutorResult { let url = exec_ctx.responses_url(); // Non-streaming request: stream=false -> full JSON body -> from_json. - let upstream_request = ctx.enriched_request.to_upstream_request(false)?; - let upstream_json = serialize_to_string(&upstream_request).map_err(ExecutorError::JsonError)?; + let upstream_json = upstream_request_json(ctx, false)?; let body = fetch_response_json(upstream_json, &url, &exec_ctx.client, auth).await?; - let acc = ResponseAccumulator::from_json(&body, ctx.conversation_id.as_deref())?; + payload_from_upstream(ctx, UpstreamBody::Json(&body)) +} + +/// A complete upstream response, in whichever form the caller received it. +#[derive(Debug, Clone, Copy)] +pub enum UpstreamBody<'a> { + Json(&'a str), + /// Frames of a streamed response, already relayed by the caller. + Sse(&'a str), +} + +fn absorb_line(acc: &mut ResponseAccumulator, ctx: &RequestContext, line: &str) { + if let Some(frame) = acc.process_sse_line(line) { + log_upstream_failure(&frame, &ctx.response_id); + } +} + +/// Assembles the final [`ResponsePayload`], stamping our ids over the upstream's. +/// +/// # Errors +/// A parse error for an invalid JSON body; [`ExecutorError::InvalidRequest`] for +/// an SSE relay cut short, which `finish_stream` would call complete. +pub(super) fn payload_from_upstream( + ctx: &RequestContext, + upstream: UpstreamBody<'_>, +) -> ExecutorResult { + let acc = match upstream { + UpstreamBody::Json(body) => ResponseAccumulator::from_json(body, ctx.conversation_id.as_deref())?, + UpstreamBody::Sse(sse) => { + let mut acc = ResponseAccumulator::new(ctx.response_id.clone(), ctx.conversation_id.clone()); + for line in sse.lines() { + absorb_line(&mut acc, ctx, line); + } + if !acc.saw_terminal_frame() { + return Err(ExecutorError::InvalidRequest( + "upstream stream ended without a terminal event".to_owned(), + )); + } + acc.finish_stream(); + acc + } + }; + Ok(finalize_payload(ctx, acc)) +} + +/// The tail both legs share: request-derived fields in, our ids stamped on. +fn finalize_payload(ctx: &RequestContext, acc: ResponseAccumulator) -> ResponsePayload { let mut payload = acc.finalize( &ctx.enriched_request.model, ctx.original_request.previous_response_id.as_deref(), ctx.original_request.instructions.as_deref(), ); ctx.inject_ids(&mut payload); - - Ok(payload) + payload } pub(super) async fn fetch_stream_payload( @@ -67,8 +120,7 @@ pub(super) async fn fetch_stream_payload( output_offset: usize, ) -> ExecutorResult { let url = exec_ctx.responses_url(); - let upstream_request = ctx.enriched_request.to_upstream_request(true)?; - let upstream_json = serialize_to_string(&upstream_request).map_err(ExecutorError::JsonError)?; + let upstream_json = upstream_request_json(ctx, true)?; let mut line_stream = Box::pin(call_inference( upstream_json, url, @@ -84,9 +136,7 @@ pub(super) async fn fetch_stream_payload( while let Some(line_result) = line_stream.next().await { let line = line_result?; if stream.is_none() { - if let Some(frame) = acc.process_sse_line(&line) { - log_upstream_failure(&frame, &ctx.response_id); - } + absorb_line(&mut acc, ctx, &line); continue; } if let Some(translation) = acc.process_sse_line_with_translator(&line, &mut function_sse)? { @@ -130,12 +180,7 @@ pub(super) async fn fetch_stream_payload( } } acc.finish_stream(); - let mut payload = acc.finalize( - &ctx.enriched_request.model, - ctx.original_request.previous_response_id.as_deref(), - ctx.original_request.instructions.as_deref(), - ); - ctx.inject_ids(&mut payload); + let payload = finalize_payload(ctx, acc); Ok(StreamPayload { payload, deferred_events, diff --git a/crates/agentic-server-core/src/types/request_response.rs b/crates/agentic-server-core/src/types/request_response.rs index 72ae5eee..75cc47e9 100644 --- a/crates/agentic-server-core/src/types/request_response.rs +++ b/crates/agentic-server-core/src/types/request_response.rs @@ -106,6 +106,35 @@ where } impl RequestPayload { + /// Names the feature in this request that only the in-process executor + /// implements, if any. Such a request cannot be served by the passthrough + /// proxy or by split execution; the returned name lets a caller say which + /// feature forced its hand. + #[must_use] + pub fn in_process_feature(&self) -> Option<&'static str> { + if self.conversation_id.is_some() { + return Some("conversation_id"); + } + if self + .tools + .as_ref() + .is_some_and(|tools| tools.iter().any(|tool| !matches!(tool, ResponsesTool::Function(_)))) + { + return Some("gateway-owned tools"); + } + if self.input.contains_compaction() || self.input.has_compaction_trigger() { + return Some("compaction input"); + } + if self + .context_management + .as_ref() + .is_some_and(|entries| !entries.is_empty()) + { + return Some("context_management"); + } + None + } + /// Construct an `UpstreamRequest` suitable for forwarding to vLLM. /// /// Codex `namespace` tools' members are first renamed to their flat, diff --git a/crates/agentic-server-core/tests/split_execution_integration.rs b/crates/agentic-server-core/tests/split_execution_integration.rs new file mode 100644 index 00000000..9c88e889 --- /dev/null +++ b/crates/agentic-server-core/tests/split_execution_integration.rs @@ -0,0 +1,209 @@ +//! Split execution: hydrate, an external inference call, then persist. + +use std::fmt::Write as _; +use std::sync::Arc; + +use serde_json::{Value, json}; + +use agentic_core::executor::request::{RequestContext, SplitContext}; +use agentic_core::executor::split::{Hydration, UpstreamBody}; +use agentic_core::executor::{ConversationHandler, ExecutionContext, ResponseHandler, rehydrate_conversation, split}; +use agentic_core::storage::{ConversationStore, ResponseStore, create_pool_with_schema}; +use agentic_core::types::request_response::{RequestPayload, ResponsePayload}; + +async fn exec_ctx() -> ExecutionContext { + let pool = create_pool_with_schema(Some("sqlite://?mode=memory")) + .await + .expect("pool"); + ExecutionContext::new( + ConversationHandler::new(ConversationStore::new(Arc::clone(&pool))), + ResponseHandler::new(ResponseStore::new(pool)), + Arc::new(reqwest::Client::new()), + "http://localhost:8000".to_owned(), + ) +} + +/// Built the way a request actually arrives — through deserialization. +fn request(input: &str, previous: Option<&str>) -> RequestPayload { + let mut body = json!({"model": "test-model", "input": input, "store": true}); + if let Some(previous) = previous { + body["previous_response_id"] = json!(previous); + } + serde_json::from_value(body).expect("valid request") +} + +fn message(status: &str, content: &Value) -> Value { + json!({"type": "message", "id": "msg_1", "role": "assistant", "status": status, "content": content}) +} + +/// A complete non-streaming upstream body, as the model backend returns it. +fn upstream_json(body: &str) -> String { + let text = json!([{"type": "output_text", "text": body, "annotations": []}]); + json!({ + "id": "resp_upstream", "object": "response", "created_at": 1_700_000_000, + "model": "test-model", "status": "completed", "output": [message("completed", &text)] + }) + .to_string() +} + +/// The same turn as SSE, the way the caller would have relayed it. +fn upstream_sse(body: &str) -> String { + let text = json!([{"type": "output_text", "text": body, "annotations": []}]); + let frames = [ + json!({"type": "response.output_item.added", "output_index": 0, "item": message("in_progress", &json!([]))}), + json!({"type": "response.output_text.delta", "output_index": 0, "item_id": "msg_1", "delta": body}), + json!({"type": "response.output_item.done", "output_index": 0, "item": message("completed", &text)}), + json!({"type": "response.completed", "response": {"id": "resp_upstream", "status": "completed"}}), + ]; + frames.iter().fold(String::new(), |mut sse, frame| { + writeln!(sse, "data: {frame}\n").expect("write to a String"); + sse + }) +} + +async fn hydrate(input: &str, previous: Option<&str>, ctx: &ExecutionContext) -> Hydration { + split::hydrate(request(input, previous), ctx).await.expect("hydrate") +} + +async fn persist(turn: Hydration, upstream: UpstreamBody<'_>, ctx: &ExecutionContext) -> ResponsePayload { + split::persist(turn.context, upstream, ctx).await.expect("persist") +} + +/// The upstream request is raw JSON; parse it to assert on its shape. +fn sent(turn: &Hydration) -> Value { + serde_json::from_str(turn.request.get()).expect("valid request") +} + +/// How many input items the upstream request replays, and their combined text. +fn replayed(turn: &Hydration) -> (usize, String) { + let items = sent(turn)["input"].as_array().expect("input items").clone(); + (items.len(), items.iter().map(ToString::to_string).collect()) +} + +/// `InputItem` and `ResponsesTool` are not `PartialEq`; compare their wire form. +fn json_of(value: &T) -> Value { + serde_json::to_value(value).expect("serializable") +} + +#[tokio::test] +async fn a_second_turn_replays_the_stored_history() { + let ctx = exec_ctx().await; + + let turn = hydrate("What is 2+2?", None, &ctx).await; + assert_eq!(replayed(&turn).0, 1); + assert!( + sent(&turn).get("previous_response_id").is_none(), + "upstream is stateless" + ); + + let first = persist(turn, UpstreamBody::Json(&upstream_json("4")), &ctx).await; + assert!(first.id.starts_with("resp_")); + assert_ne!( + first.id, "resp_upstream", + "the envelope carries our id, not the model's" + ); + + let turn = hydrate("What did I ask?", Some(&first.id), &ctx).await; + let (count, text) = replayed(&turn); + assert_eq!(count, 3, "prior user + assistant turns, then the new input"); + assert!(text.contains("What is 2+2?") && text.contains('4')); + + let second = persist(turn, UpstreamBody::Json(&upstream_json("2+2")), &ctx).await; + assert_eq!(second.previous_response_id.as_deref(), Some(first.id.as_str())); +} + +/// Every way a turn can fail to be stored, and the status the caller sees. +#[tokio::test] +async fn a_turn_that_cannot_be_stored_is_refused() { + let ctx = exec_ctx().await; + + let unknown = split::hydrate(request("hi", Some("resp_missing")), &ctx).await; + assert_eq!(unknown.expect_err("unknown id").http_status().as_u16(), 404); + + let mut in_progress: Value = serde_json::from_str(&upstream_json("partial")).expect("json"); + in_progress["status"] = json!("in_progress"); + let turn = hydrate("hi", None, &ctx).await; + let error = split::persist(turn.context, UpstreamBody::Json(&in_progress.to_string()), &ctx).await; + assert_eq!(error.expect_err("never stored").http_status().as_u16(), 400); + + // A relay that died mid-stream: `finish_stream` would call this complete. + let cut_short = r#"data: {"type":"response.output_text.delta","item_id":"msg_1","delta":"4"}"#; + let turn = hydrate("hi", None, &ctx).await; + let error = split::persist(turn.context, UpstreamBody::Sse(cut_short), &ctx).await; + assert_eq!(error.expect_err("no terminal event").http_status().as_u16(), 400); +} + +#[test] +fn the_boundary_check_names_what_cannot_be_split() { + let gateway_tool: RequestPayload = serde_json::from_value(json!({ + "model": "test-model", "input": "hi", "store": true, "tools": [{"type": "web_search_preview"}] + })) + .expect("valid request"); + let error = split::ensure_splittable(&gateway_tool).expect_err("the loop needs a caller"); + assert!(error.to_string().contains("tools"), "got: {error}"); + + let mut conversational = request("hi", None); + conversational.conversation_id = Some("conv_1".into()); + let error = split::ensure_splittable(&conversational).expect_err("its version cannot cross"); + assert!(error.to_string().contains("conversation_id"), "got: {error}"); + + let mut streaming = request("hi", None); + streaming.stream = true; + split::ensure_splittable(&streaming).expect("the caller relays the frames, then replays them"); + split::ensure_splittable(&request("hi", None)).expect("a plain turn is splittable"); +} + +/// The wire form drops what it can rebuild, and the rebuild has to agree with +/// what hydration produced — persist stores from it. +#[tokio::test] +async fn the_wire_context_round_trips_into_an_equal_context() { + let ctx = exec_ctx().await; + let live = rehydrate_conversation(request("What is 2+2?", None), &ctx) + .await + .expect("rehydrate"); + let (id, items, tools) = ( + live.response_id.clone(), + json_of(&live.new_input_items), + json_of(&live.enriched_request.tools), + ); + + let wire = serde_json::to_string(&SplitContext::from(live)).expect("serialize"); + assert!(!wire.contains("enriched_request"), "already in flight as the request"); + assert!( + !wire.contains("conversation_version"), + "conversation mode does not split" + ); + + let back = RequestContext::from(serde_json::from_str::(&wire).expect("deserialize")); + assert_eq!(back.response_id, id); + assert_eq!(json_of(&back.new_input_items), items, "derived items match"); + assert_eq!(json_of(&back.enriched_request.tools), tools, "resolved tools survive"); + assert!( + back.enriched_request.previous_response_id.is_none(), + "upstream stays stateless" + ); + assert!( + back.conversation_version.is_none(), + "never resumed with a stale version" + ); +} + +#[tokio::test] +async fn a_streamed_turn_persists_from_the_relayed_frames() { + let ctx = exec_ctx().await; + let mut streaming = request("What is 2+2?", None); + streaming.stream = true; + + let turn = split::hydrate(streaming, &ctx).await.expect("hydrate"); + assert_eq!( + sent(&turn)["stream"], + json!(true), + "the client's flag reaches the model" + ); + + let stored = persist(turn, UpstreamBody::Sse(&upstream_sse("4")), &ctx).await; + assert_ne!(stored.id, "resp_upstream"); + + let next = hydrate("What did I ask?", Some(&stored.id), &ctx).await; + assert_eq!(replayed(&next).0, 3, "the streamed turn is continuable"); +} diff --git a/crates/agentic-server/src/handler/http/responses.rs b/crates/agentic-server/src/handler/http/responses.rs index a27daf71..e454348f 100644 --- a/crates/agentic-server/src/handler/http/responses.rs +++ b/crates/agentic-server/src/handler/http/responses.rs @@ -10,7 +10,6 @@ use std::sync::Arc; use agentic_core::executor::{ExecuteRequest, compact_response as execute_compaction}; use agentic_core::proxy::{ProxyRequest, proxy_request}; use agentic_core::types::request_response::{CompactRequest, RequestPayload}; -use agentic_core::types::tools::ResponsesTool; use super::super::common::{ convert_response, executor_error_response, extract_bearer, read_and_parse, read_json, sse_response, @@ -39,13 +38,6 @@ async fn execute_responses(state: &AppState, parts: Parts, payload: RequestPaylo } } -fn has_gateway_tools(payload: &RequestPayload) -> bool { - payload - .tools - .as_ref() - .is_some_and(|tools| tools.iter().any(|tool| !matches!(tool, ResponsesTool::Function(_)))) -} - pub async fn responses(State(state): State, req: Request) -> Response { let (parts, body) = req.into_parts(); let (bytes, payload) = match read_and_parse(body).await { @@ -53,16 +45,8 @@ pub async fn responses(State(state): State, req: Request) -> Response Err(e) => return e, }; - let should_execute = payload.store - || payload.previous_response_id.is_some() - || payload.conversation_id.is_some() - || payload.input.contains_compaction() - || payload.input.has_compaction_trigger() - || payload - .context_management - .as_ref() - .is_some_and(|entries| !entries.is_empty()) - || has_gateway_tools(&payload); + let should_execute = + payload.store || payload.previous_response_id.is_some() || payload.in_process_feature().is_some(); debug!( route = if should_execute { "executor" } else { "proxy" }, store = payload.store, diff --git a/docs/design/agentic-llm-d.md b/docs/design/agentic-llm-d.md new file mode 100644 index 00000000..029873ff --- /dev/null +++ b/docs/design/agentic-llm-d.md @@ -0,0 +1,62 @@ +# agentic-llm-d + +## Scope + +`agentic-llm-d` runs agentic-api as a pair of state services for the llm-d coordinator, which performs the inference +call itself. It decomposes one stateful Responses turn into two separately callable steps, so a caller that already +routes model traffic does not have to proxy that traffic back through the gateway. + +Serving `previous_response_id` normally means reaching whichever engine holds the earlier turn. Keeping that history +behind an API removes the constraint: the request the coordinator forwards carries its own history, so any engine can +serve it. Cache-aware scoring can still prefer the engine that handled the previous turn without being bound to it. + +## The two steps + +`hydrate` takes the client's request, resolves `previous_response_id` against storage, and returns two things: the +upstream request body with the history inlined and every continuation and storage field removed, and a `SplitContext` +describing the turn. The caller forwards the body to a model unchanged and echoes the context back. + +`persist` takes that context together with the response the model produced — either a complete JSON body, or the SSE +frames a streaming caller has already relayed to its own client — assembles the turn, stores it, and returns the +response envelope carrying the reserved `resp_` identifier. Nothing is re-emitted for a streamed turn, since the caller +has already sent the frames on. + +`SplitContext` is the wire form of the in-process `RequestContext`. It omits the enriched request, which is already in +flight as the request body, and the derived input items; both are rebuilt when the context comes back. Callers treat it +as opaque. + +## Composition + +Neither step reimplements the flow. `hydrate` calls `rehydrate_conversation` and `upstream_request_json`, and `persist` +calls `payload_from_upstream` and `persist_if_needed`. All four are the functions the in-process executor already uses, +so a change to how a turn is rehydrated or stored reaches both paths at once. `split.rs` contains no parsing, no +storage access and no request building of its own. + +What it does contain is the boundary itself: the check for what cannot be split, the conversion between the live and +wire context forms, and a terminal-status check. The last exists because an external caller can return a response the +in-process flow could never produce, such as one still in progress. Storing it would hand back an identifier that could +never be continued, so `persist` rejects it. + +`ensure_splittable` reuses `RequestPayload::in_process_feature`, the predicate that already decides whether the gateway +runs the executor or passes a request through to vLLM. The passthrough proxy and the split boundary have the same +limits, so sharing one predicate stops them drifting apart as features are added. + +## Boundary + +`ensure_splittable` names the feature that prevents a request being split: `conversation_id`, gateway-owned tools, +compaction input, or `context_management`. Each needs state that the in-process executor keeps between steps. + +## The crate + +The endpoints are served by a separate crate and binary that depends on `agentic-server-core` and not on the gateway. +It serves `/internal/hydrate`, `/internal/persist`, `/health` and `/ready`, and nothing else: the passthrough proxy, +`/v1`, the WebSocket transport, upstream readiness probing and vLLM subprocess management are all absent, so the +internal endpoints cannot be exposed on a listener that also serves `/v1`. Readiness reports whether storage answers, +since the coordinator owns the model fleet. + +## Discussion points + +The `/internal` endpoints carry no credential and trust their caller, so restricting them is a network-layer concern +today. A retried `persist` conflicts with the response-identifier primary key and returns a server error rather than +the turn it already stored. Request fields that `RequestPayload` does not model are dropped rather than forwarded, +which narrows what reaches vLLM compared with plain passthrough.