fix: harden split execution contracts (vllm-project#235)

## Summary

- harden the split execution endpoints added in vllm-project#216 with validated
secret types, separate hydrate, context, response, and persist size
budgets, stable 401, 409, and 413 error envelopes, and propagated
shutdown errors
- make duplicate persistence atomic by mapping the database uniqueness
violation to a stable `409 response_already_stored` response
- strictly validate caller-relayed JSON and SSE response lifecycles,
required fields, item identity and type, terminal status, and permanent
output ID/index uniqueness before persistence
- treat `response.output_item.done` as authoritative so missing deltas
cannot silently truncate stored message content, and reject stream item
types that cannot be represented safely
- document the workload-token contract, deployment trust boundary,
limits, and retry behavior

The workload token authenticates a trusted coordinator, not a tenant.
Per-response tenant ownership remains tracked in vllm-project#107, so these
endpoints should stay on an encrypted, policy-restricted service
network.

## Test Plan

- `cargo fmt --all -- --check`
- `CARGO_INCREMENTAL=0 RUSTFLAGS='-C debuginfo=0' cargo clippy -j 2
--workspace --all-targets -- -D warnings`
- `CARGO_INCREMENTAL=0 RUSTFLAGS='-C debuginfo=0' cargo test --workspace
-j 2`
- `uv run --with-requirements docs/requirements.txt mkdocs build
--strict`

The standard PostgreSQL tests that require `TEST_POSTGRES_URL` remained
ignored.

---------

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>

Author: franciscojavierarceo

CHANGELOG.md

@@ -18,6 +18,8 @@ All notable changes to Agentic API are documented here.
 
 ### Fixed
 
+- Hardened split execution with atomic duplicate persistence, strict relayed-response validation, independent secret
+  validation, bounded hydrate and persist payloads, stable error envelopes, and graceful shutdown error propagation.
 - Forwarded Responses `text` generation settings, including structured output
   formats and verbosity, through typed HTTP, WebSocket, and gateway-tool paths.
 - Replaced `WebSearchActionSearch::new` and `WebSearchCall::new` with fallible

Cargo.lock

@@ -37,6 +37,7 @@ serde_json = { version = "1", features = ["raw_value"] }
 thiserror = "2"
 tokio = { version = "1", features = ["full"] }
 tokio-util = "0.7"
+tower = { version = "0.5", features = ["util"] }
 tower-http = { version = "0.6", features = ["cors"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

Cargo.toml

@@ -28,3 +28,4 @@ agentic-core = { workspace = true, features = [] }
 reqwest = { workspace = true, features = ["json"] }
 serde_json.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+tower.workspace = true

crates/agentic-llm-d/Cargo.toml

@@ -14,6 +14,8 @@ use agentic_core::types::io::{ResponsesInput, ToolChoice};
 use agentic_core::types::request_response::RequestPayload;
 use agentic_core::types::tools::ResponsesTool;
 
+use crate::SigningKey;
+
 /// What `hydrate` returns. Raw JSON: the caller forwards it uninterpreted.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Hydration {
@@ -95,7 +97,7 @@ struct SealedClaims {
 ///
 /// # Errors
 /// [`ExecutorError::InvalidRequest`] if the token cannot be produced.
-pub fn seal(context: SplitContext, key: &[u8]) -> ExecutorResult<String> {
+pub fn seal(context: SplitContext, key: &SigningKey) -> ExecutorResult<String> {
     let expires = SystemTime::now()
         .checked_add(CONTEXT_TTL)
         .and_then(|at| at.duration_since(UNIX_EPOCH).ok())
@@ -105,20 +107,24 @@ pub fn seal(context: SplitContext, key: &[u8]) -> ExecutorResult<String> {
         aud: AUDIENCE.to_owned(),
         ctx: context,
     };
-    encode(&Header::new(Algorithm::HS256), &claims, &EncodingKey::from_secret(key))
-        .map_err(|error| ExecutorError::InvalidRequest(format!("cannot seal context: {error}")))
+    encode(
+        &Header::new(Algorithm::HS256),
+        &claims,
+        &EncodingKey::from_secret(key.as_bytes()),
+    )
+    .map_err(|error| ExecutorError::InvalidRequest(format!("cannot seal context: {error}")))
 }
 
 /// Opens a sealed context, rejecting one that was tampered with or has expired.
 ///
 /// # Errors
 /// [`ExecutorError::InvalidRequest`] for a bad signature, a wrong audience, or
 /// a context past its expiry.
-pub fn unseal(token: &str, key: &[u8]) -> ExecutorResult<SplitContext> {
+pub fn unseal(token: &str, key: &SigningKey) -> ExecutorResult<SplitContext> {
     let mut validation = Validation::new(Algorithm::HS256);
     validation.set_audience(&[AUDIENCE]);
     validation.set_required_spec_claims(&["exp", "aud"]);
-    decode::<SealedClaims>(token, &DecodingKey::from_secret(key), &validation)
+    decode::<SealedClaims>(token, &DecodingKey::from_secret(key.as_bytes()), &validation)
         .map(|data| data.claims.ctx)
         .map_err(|error| ExecutorError::InvalidRequest(format!("context rejected: {error}")))
 }

crates/agentic-llm-d/src/context.rs

@@ -9,8 +9,8 @@ use axum::extract::{Request, State};
 use axum::http::StatusCode;
 use axum::middleware::Next;
 use axum::response::{IntoResponse, Response};
-use serde::Deserialize;
 use serde::de::DeserializeOwned;
+use serde::{Deserialize, Serialize};
 use serde_json::value::RawValue;
 use tracing::warn;
 
@@ -24,7 +24,10 @@ use agentic_core::types::request_response::RequestPayload;
 use crate::BackendState;
 use crate::context::{Hydration, ensure_splittable, seal, unseal};
 
-const MAX_BODY_SIZE: usize = 10 * 1024 * 1024;
+const MAX_HYDRATE_BODY_SIZE: usize = 2 * 1024 * 1024;
+const MAX_PERSIST_BODY_SIZE: usize = 16 * 1024 * 1024;
+const MAX_CONTEXT_SIZE: usize = 6 * 1024 * 1024;
+const MAX_UPSTREAM_BODY_SIZE: usize = 4 * 1024 * 1024;
 /// The calling workload's shared secret.
 pub const WORKLOAD_TOKEN_HEADER: &str = "x-agentic-workload-token";
 /// Readiness means storage answers - llm-d owns the model fleet.
@@ -47,10 +50,12 @@ pub async fn require_token(State(state): State<BackendState>, request: Request,
         .get(WORKLOAD_TOKEN_HEADER)
         .and_then(|value| value.to_str().ok());
     match presented {
-        Some(token) if token_matches(token, &state.api_token) => next.run(request).await,
-        _ => json(
+        Some(token) if token_matches(token, state.secrets.workload_token()) => next.run(request).await,
+        _ => api_error(
             StatusCode::UNAUTHORIZED,
-            br#"{"error":{"type":"invalid_request_error","message":"missing or invalid bearer token"}}"#.to_vec(),
+            "authentication_error",
+            "invalid_workload_token",
+            "missing or invalid workload token",
         ),
     }
 }
@@ -78,7 +83,7 @@ pub async fn ready(State(state): State<BackendState>) -> StatusCode {
 }
 
 pub async fn hydrate(State(state): State<BackendState>, req: Request) -> Response {
-    let payload: RequestPayload = match read_json(req.into_body()).await {
+    let payload: RequestPayload = match read_json(req.into_body(), MAX_HYDRATE_BODY_SIZE).await {
         Ok(payload) => payload,
         Err(response) => return response,
     };
@@ -101,12 +106,17 @@ async fn build_hydration(
     ensure_splittable(&ctx.enriched_request)?;
     let stream = ctx.original_request.stream;
     let request = RawValue::from_string(upstream_request(&ctx, stream)?).map_err(ExecutorError::JsonError)?;
-    let context = seal(ctx.into(), &state.signing_key)?;
+    let context = seal(ctx.into(), state.secrets.signing_key())?;
+    if context.len() > MAX_CONTEXT_SIZE {
+        return Err(ExecutorError::PayloadTooLarge(
+            "hydrated context exceeds the split-execution size budget".to_owned(),
+        ));
+    }
     Ok(Hydration { request, context })
 }
 
 pub async fn persist(State(state): State<BackendState>, req: Request) -> Response {
-    let PersistRequest { context, response, sse } = match read_json(req.into_body()).await {
+    let PersistRequest { context, response, sse } = match read_json(req.into_body(), MAX_PERSIST_BODY_SIZE).await {
         Ok(request) => request,
         Err(response) => return response,
     };
@@ -119,7 +129,20 @@ pub async fn persist(State(state): State<BackendState>, req: Request) -> Respons
             return error_response(ExecutorError::InvalidRequest(message));
         }
     };
-    let context = match unseal(&context, &state.signing_key) {
+    if context.len() > MAX_CONTEXT_SIZE {
+        return error_response(ExecutorError::PayloadTooLarge(
+            "sealed context exceeds the split-execution size budget".to_owned(),
+        ));
+    }
+    let upstream_size = match upstream {
+        UpstreamBody::Json(body) | UpstreamBody::Sse(body) => body.len(),
+    };
+    if upstream_size > MAX_UPSTREAM_BODY_SIZE {
+        return error_response(ExecutorError::PayloadTooLarge(
+            "upstream response exceeds the split-execution size budget".to_owned(),
+        ));
+    }
+    let context = match unseal(&context, state.secrets.signing_key()) {
         Ok(context) => context,
         Err(error) => return error_response(error),
     };
@@ -142,14 +165,38 @@ fn error_response(error: ExecutorError) -> Response {
 }
 
 #[allow(clippy::result_large_err)] // an axum `Response` is the idiomatic error here
-async fn read_json<T: DeserializeOwned>(body: Body) -> Result<T, Response> {
-    let too_large = br#"{"error":{"type":"invalid_request_error","message":"request body too large"}}"#;
-    let bytes = axum::body::to_bytes(body, MAX_BODY_SIZE)
+async fn read_json<T: DeserializeOwned>(body: Body, limit: usize) -> Result<T, Response> {
+    let bytes = axum::body::to_bytes(body, limit)
         .await
-        .map_err(|_| json(StatusCode::PAYLOAD_TOO_LARGE, too_large.to_vec()))?;
+        .map_err(|_| error_response(ExecutorError::PayloadTooLarge("request body too large".to_owned())))?;
     serde_json::from_slice(&bytes).map_err(|error| error_response(ExecutorError::from(error)))
 }
 
+#[derive(Serialize)]
+struct ApiErrorEnvelope<'a> {
+    error: ApiErrorBody<'a>,
+}
+
+#[derive(Serialize)]
+struct ApiErrorBody<'a> {
+    message: &'a str,
+    #[serde(rename = "type")]
+    error_type: &'a str,
+    code: &'a str,
+}
+
+fn api_error(status: StatusCode, error_type: &str, code: &str, message: &str) -> Response {
+    let body = serde_json::to_vec(&ApiErrorEnvelope {
+        error: ApiErrorBody {
+            message,
+            error_type,
+            code,
+        },
+    })
+    .expect("static API error serializes");
+    json(status, body)
+}
+
 fn json(status: StatusCode, body: Vec<u8>) -> Response {
     Response::builder()
         .status(status)

crates/agentic-llm-d/src/handler.rs

@@ -12,14 +12,120 @@ use axum::{Router, middleware};
 
 use agentic_core::executor::ExecutionContext;
 
+const MIN_SECRET_LEN: usize = 32;
+
+/// Invalid backend authentication or signing material.
+#[derive(Debug, thiserror::Error)]
+pub enum SecretError {
+    #[error("{name} must be at least {MIN_SECRET_LEN} bytes of independently generated randomness")]
+    TooShort { name: &'static str },
+    #[error("the signing key and workload token must be generated independently")]
+    Reused,
+}
+
+/// Validated secrets shared by the backend handlers.
+#[derive(Clone)]
+pub struct BackendSecrets {
+    signing_key: SigningKey,
+    workload_token: Arc<str>,
+}
+
+/// Validated key used to sign and verify split-execution contexts.
+#[derive(Clone)]
+pub struct SigningKey(Arc<[u8]>);
+
+impl SigningKey {
+    /// Validates a context-signing key.
+    ///
+    /// # Errors
+    /// Returns [`SecretError`] when the key is shorter than 32 non-whitespace
+    /// bytes.
+    pub fn new(value: Vec<u8>) -> Result<Self, SecretError> {
+        if non_whitespace_ascii_len(&value) < MIN_SECRET_LEN {
+            return Err(SecretError::TooShort { name: "signing key" });
+        }
+        Ok(Self(value.into()))
+    }
+
+    pub(crate) fn as_bytes(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+impl BackendSecrets {
+    /// Validates and constructs the secrets required by the split routes.
+    ///
+    /// # Errors
+    /// Returns [`SecretError`] when either value is too short or both values are
+    /// identical.
+    pub fn new(signing_key: Vec<u8>, workload_token: String) -> Result<Self, SecretError> {
+        let signing_key = SigningKey::new(signing_key)?;
+        if non_whitespace_ascii_len(workload_token.as_bytes()) < MIN_SECRET_LEN {
+            return Err(SecretError::TooShort { name: "workload token" });
+        }
+        if signing_key.as_bytes() == workload_token.as_bytes() {
+            return Err(SecretError::Reused);
+        }
+        Ok(Self {
+            signing_key,
+            workload_token: workload_token.into(),
+        })
+    }
+
+    pub(crate) fn signing_key(&self) -> &SigningKey {
+        &self.signing_key
+    }
+
+    pub(crate) fn workload_token(&self) -> &str {
+        &self.workload_token
+    }
+}
+
+fn non_whitespace_ascii_len(value: &[u8]) -> usize {
+    value.iter().filter(|byte| !byte.is_ascii_whitespace()).count()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn signing_keys_and_workload_tokens_must_be_long_and_independent() {
+        assert!(matches!(
+            BackendSecrets::new(Vec::new(), "b".repeat(MIN_SECRET_LEN)),
+            Err(SecretError::TooShort { name: "signing key" })
+        ));
+        assert!(matches!(
+            BackendSecrets::new(vec![b'a'; MIN_SECRET_LEN], String::new()),
+            Err(SecretError::TooShort { name: "workload token" })
+        ));
+        assert!(matches!(
+            BackendSecrets::new(vec![b'a'; MIN_SECRET_LEN], "a".repeat(MIN_SECRET_LEN)),
+            Err(SecretError::Reused)
+        ));
+        let sparse = format!("a{}b", " ".repeat(MIN_SECRET_LEN - 2));
+        assert!(matches!(
+            SigningKey::new(sparse.as_bytes().to_vec()),
+            Err(SecretError::TooShort { name: "signing key" })
+        ));
+        assert!(matches!(
+            BackendSecrets::new(vec![b'a'; MIN_SECRET_LEN], sparse),
+            Err(SecretError::TooShort { name: "workload token" })
+        ));
+        assert!(matches!(
+            SigningKey::new(vec![b'a'; MIN_SECRET_LEN - 1]),
+            Err(SecretError::TooShort { name: "signing key" })
+        ));
+        BackendSecrets::new(vec![b'a'; MIN_SECRET_LEN], "b".repeat(MIN_SECRET_LEN)).expect("independent secrets");
+    }
+}
+
 /// All the endpoints need — far less than the gateway's `AppState`.
 #[derive(Clone)]
 pub struct BackendState {
     pub exec_ctx: Arc<ExecutionContext>,
-    /// Seals the context between hydrate and persist.
-    pub signing_key: Arc<Vec<u8>>,
-    /// Shared secret every split-route caller must present.
-    pub api_token: Arc<String>,
+    /// Validated signing and workload-authentication material.
+    pub secrets: BackendSecrets,
 }
 
 /// The whole surface: two split-execution endpoints and two probes.

crates/agentic-llm-d/src/lib.rs

@@ -31,9 +31,6 @@ struct Cli {
     api_token: String,
 }
 
-/// Long enough that a guessed value is not worth trying.
-const MIN_SECRET_LEN: usize = 32;
-
 #[tokio::main]
 async fn main() -> Result<(), runner::Error> {
     tracing_subscriber::fmt::init();
@@ -52,25 +49,6 @@ async fn main() -> Result<(), runner::Error> {
         tools: ToolRuntimeConfig::default(),
     };
 
-    // Checked before binding: a short secret is worse than none, it looks configured.
-    for (name, value) in [
-        ("AGENTIC_LLM_D_SIGNING_KEY", &cli.signing_key),
-        ("AGENTIC_LLM_D_API_TOKEN", &cli.api_token),
-    ] {
-        if value.trim().len() < MIN_SECRET_LEN {
-            eprintln!("{name} must be at least {MIN_SECRET_LEN} characters of independently generated randomness");
-            std::process::exit(2);
-        }
-    }
-
-    // Independence is the property we can actually check. The token travels in a
-    // header on every request, so reusing it as the signing key would let anyone
-    // who sees one forge the other.
-    if cli.signing_key == cli.api_token {
-        eprintln!("AGENTIC_LLM_D_SIGNING_KEY and AGENTIC_LLM_D_API_TOKEN must be generated independently");
-        std::process::exit(2);
-    }
-
     let shutdown = CancellationToken::new();
     let on_signal = shutdown.clone();
     tokio::spawn(async move {