|
| 1 | +//! The endpoints and their axum glue. The split routes require a shared token; |
| 2 | +//! probes do not. Keep the listener cluster-internal regardless — the token |
| 3 | +//! authenticates the calling workload, not a tenant. |
| 4 | +
|
| 5 | +use std::time::Duration; |
| 6 | + |
| 7 | +use axum::body::Body; |
| 8 | +use axum::extract::{Request, State}; |
| 9 | +use axum::http::StatusCode; |
| 10 | +use axum::middleware::Next; |
| 11 | +use axum::response::{IntoResponse, Response}; |
| 12 | +use serde::Deserialize; |
| 13 | +use serde::de::DeserializeOwned; |
| 14 | +use serde_json::value::RawValue; |
| 15 | +use tracing::warn; |
| 16 | + |
| 17 | +use agentic_core::executor::request::RequestContext; |
| 18 | + |
| 19 | +use agentic_core::executor::{ |
| 20 | + ExecutorError, UpstreamBody, commit, decode_upstream, rehydrate_conversation, upstream_request, |
| 21 | +}; |
| 22 | +use agentic_core::types::request_response::RequestPayload; |
| 23 | + |
| 24 | +use crate::BackendState; |
| 25 | +use crate::context::{Hydration, ensure_splittable, seal, unseal}; |
| 26 | + |
| 27 | +const MAX_BODY_SIZE: usize = 10 * 1024 * 1024; |
| 28 | +/// The calling workload's shared secret. |
| 29 | +pub const WORKLOAD_TOKEN_HEADER: &str = "x-agentic-workload-token"; |
| 30 | +/// Readiness means storage answers - llm-d owns the model fleet. |
| 31 | +const STORAGE_PROBE_TIMEOUT: Duration = Duration::from_secs(2); |
| 32 | + |
| 33 | +/// Body of `POST /v1alpha/responses/persist`: the context, plus one response form. |
| 34 | +#[derive(Debug, Deserialize)] |
| 35 | +pub struct PersistRequest { |
| 36 | + context: String, |
| 37 | + response: Option<Box<RawValue>>, |
| 38 | + sse: Option<String>, |
| 39 | +} |
| 40 | + |
| 41 | +/// Rejects any split-route call without the shared secret. The probes are |
| 42 | +/// layered separately and stay open. |
| 43 | +pub async fn require_token(State(state): State<BackendState>, request: Request, next: Next) -> Response { |
| 44 | + // Not `Authorization`: that stays free for the end user's token. |
| 45 | + let presented = request |
| 46 | + .headers() |
| 47 | + .get(WORKLOAD_TOKEN_HEADER) |
| 48 | + .and_then(|value| value.to_str().ok()); |
| 49 | + match presented { |
| 50 | + Some(token) if token_matches(token, &state.api_token) => next.run(request).await, |
| 51 | + _ => json( |
| 52 | + StatusCode::UNAUTHORIZED, |
| 53 | + br#"{"error":{"type":"invalid_request_error","message":"missing or invalid bearer token"}}"#.to_vec(), |
| 54 | + ), |
| 55 | + } |
| 56 | +} |
| 57 | + |
| 58 | +/// No early return, so a wrong token takes the same time whatever byte differs. |
| 59 | +fn token_matches(presented: &str, expected: &str) -> bool { |
| 60 | + presented.len() == expected.len() |
| 61 | + && presented |
| 62 | + .bytes() |
| 63 | + .zip(expected.bytes()) |
| 64 | + .fold(0_u8, |differences, (a, b)| differences | (a ^ b)) |
| 65 | + == 0 |
| 66 | +} |
| 67 | + |
| 68 | +pub async fn health() -> StatusCode { |
| 69 | + StatusCode::OK |
| 70 | +} |
| 71 | + |
| 72 | +pub async fn ready(State(state): State<BackendState>) -> StatusCode { |
| 73 | + if state.exec_ctx.storage_ready(STORAGE_PROBE_TIMEOUT).await { |
| 74 | + StatusCode::OK |
| 75 | + } else { |
| 76 | + StatusCode::SERVICE_UNAVAILABLE |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +pub async fn hydrate(State(state): State<BackendState>, req: Request) -> Response { |
| 81 | + let payload: RequestPayload = match read_json(req.into_body()).await { |
| 82 | + Ok(payload) => payload, |
| 83 | + Err(response) => return response, |
| 84 | + }; |
| 85 | + match build_hydration(payload, &state).await { |
| 86 | + Ok(hydration) => axum::Json(hydration).into_response(), |
| 87 | + Err(error) => error_response(error), |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +/// Rehydrates the turn and builds the request the caller forwards to a model. |
| 92 | +#[allow(clippy::result_large_err)] // `ExecutorError` is core's; boxing it is not ours to decide |
| 93 | +async fn build_hydration( |
| 94 | + request: RequestPayload, |
| 95 | + state: &BackendState, |
| 96 | +) -> agentic_core::executor::ExecutorResult<Hydration> { |
| 97 | + ensure_splittable(&request)?; |
| 98 | + let ctx = rehydrate_conversation(request, state.exec_ctx.as_ref()).await?; |
| 99 | + // Rehydration can restore a gateway-owned tool from the stored turn, so |
| 100 | + // check what will actually run. |
| 101 | + ensure_splittable(&ctx.enriched_request)?; |
| 102 | + let stream = ctx.original_request.stream; |
| 103 | + let request = RawValue::from_string(upstream_request(&ctx, stream)?).map_err(ExecutorError::JsonError)?; |
| 104 | + let context = seal(ctx.into(), &state.signing_key)?; |
| 105 | + Ok(Hydration { request, context }) |
| 106 | +} |
| 107 | + |
| 108 | +pub async fn persist(State(state): State<BackendState>, req: Request) -> Response { |
| 109 | + let PersistRequest { context, response, sse } = match read_json(req.into_body()).await { |
| 110 | + Ok(request) => request, |
| 111 | + Err(response) => return response, |
| 112 | + }; |
| 113 | + // serde rejects `RawValue` in `untagged`, so "exactly one of" is checked here. |
| 114 | + let upstream = match (response.as_deref(), sse.as_deref()) { |
| 115 | + (Some(json), None) => UpstreamBody::Json(json.get()), |
| 116 | + (None, Some(sse)) => UpstreamBody::Sse(sse), |
| 117 | + _ => { |
| 118 | + let message = "exactly one of `response` or `sse` is required".to_owned(); |
| 119 | + return error_response(ExecutorError::InvalidRequest(message)); |
| 120 | + } |
| 121 | + }; |
| 122 | + let context = match unseal(&context, &state.signing_key) { |
| 123 | + Ok(context) => context, |
| 124 | + Err(error) => return error_response(error), |
| 125 | + }; |
| 126 | + let ctx = RequestContext::from(context); |
| 127 | + let stored = match decode_upstream(&ctx, upstream) { |
| 128 | + Ok(payload) => commit(ctx, payload, state.exec_ctx.as_ref()).await, |
| 129 | + Err(error) => Err(error), |
| 130 | + }; |
| 131 | + match stored { |
| 132 | + Ok(payload) => axum::Json(payload).into_response(), |
| 133 | + Err(error) => error_response(error), |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +/// Renders an error with the status and envelope core defines. |
| 138 | +fn error_response(error: ExecutorError) -> Response { |
| 139 | + let status = error.http_status(); |
| 140 | + warn!("backend error ({status}): {error}"); |
| 141 | + json(status, error.into_response_body()) |
| 142 | +} |
| 143 | + |
| 144 | +#[allow(clippy::result_large_err)] // an axum `Response` is the idiomatic error here |
| 145 | +async fn read_json<T: DeserializeOwned>(body: Body) -> Result<T, Response> { |
| 146 | + let too_large = br#"{"error":{"type":"invalid_request_error","message":"request body too large"}}"#; |
| 147 | + let bytes = axum::body::to_bytes(body, MAX_BODY_SIZE) |
| 148 | + .await |
| 149 | + .map_err(|_| json(StatusCode::PAYLOAD_TOO_LARGE, too_large.to_vec()))?; |
| 150 | + serde_json::from_slice(&bytes).map_err(|error| error_response(ExecutorError::from(error))) |
| 151 | +} |
| 152 | + |
| 153 | +fn json(status: StatusCode, body: Vec<u8>) -> Response { |
| 154 | + Response::builder() |
| 155 | + .status(status) |
| 156 | + .header("Content-Type", "application/json") |
| 157 | + .body(Body::from(body)) |
| 158 | + .expect("valid response") |
| 159 | +} |
0 commit comments