Skip to main content

agentic_server/handler/
common.rs

1use axum::body::Body;
2use axum::http::HeaderMap;
3use axum::response::Response;
4use bytes::Bytes;
5use futures::StreamExt;
6use http::StatusCode;
7use serde::de::DeserializeOwned;
8use tracing::warn;
9
10use agentic_core::executor::{BoxStream, ExecutorError};
11use agentic_core::proxy::{ProxyAuth, ProxyBody, ProxyResponse, error_response_for_auth};
12use agentic_core::types::request_response::RequestPayload;
13
14pub(super) const MAX_BODY_SIZE: usize = 10 * 1024 * 1024;
15
16/// # Panics
17/// Panics if the response builder produces an invalid response (unreachable in practice).
18pub fn convert_response(resp: ProxyResponse) -> Response {
19    let mut builder = Response::builder().status(resp.status);
20    for (name, value) in &resp.headers {
21        builder = builder.header(name, value);
22    }
23    match resp.body {
24        ProxyBody::Full(bytes) => builder.body(Body::from(bytes)).expect("valid response"),
25        ProxyBody::Stream(stream) => builder.body(Body::from_stream(stream)).expect("valid response"),
26    }
27}
28
29/// # Panics
30/// Panics if the response builder produces an invalid response (unreachable in practice).
31pub fn executor_error_response(err: ExecutorError) -> Response {
32    let status = err.http_status();
33    if !matches!(err, ExecutorError::LLMRequest { .. }) {
34        warn!("executor error ({status}): {err}");
35    }
36    Response::builder()
37        .status(status)
38        .header("Content-Type", "application/json")
39        .body(Body::from(err.into_response_body()))
40        .expect("valid error response")
41}
42
43pub(super) async fn read_bytes(body: Body) -> Result<Bytes, Response> {
44    read_bytes_with_auth(body, ProxyAuth::OpenAiBearer).await
45}
46
47pub(super) async fn read_bytes_with_auth(body: Body, auth: ProxyAuth) -> Result<Bytes, Response> {
48    axum::body::to_bytes(body, MAX_BODY_SIZE).await.map_err(|_| {
49        convert_response(error_response_for_auth(
50            StatusCode::PAYLOAD_TOO_LARGE,
51            "body_too_large",
52            "request body too large",
53            auth,
54        ))
55    })
56}
57
58pub(super) async fn read_and_parse(body: Body) -> Result<(Bytes, RequestPayload), Response> {
59    let bytes = read_bytes(body).await?;
60    let payload = serde_json::from_slice::<RequestPayload>(&bytes)
61        .map_err(|e| executor_error_response(ExecutorError::from(e)))?;
62    Ok((bytes, payload))
63}
64
65pub(super) async fn read_json<T: DeserializeOwned>(body: Body) -> Result<T, Response> {
66    let bytes = read_bytes(body).await?;
67    serde_json::from_slice::<T>(&bytes).map_err(|error| executor_error_response(ExecutorError::from(error)))
68}
69
70pub(super) fn extract_store(bytes: &[u8]) -> bool {
71    serde_json::from_slice::<serde_json::Value>(bytes)
72        .ok()
73        .and_then(|j| j.get("store").and_then(serde_json::Value::as_bool))
74        .unwrap_or(true)
75}
76
77pub(super) fn extract_bearer(headers: &HeaderMap, config_key: Option<&str>) -> Option<String> {
78    headers
79        .get("authorization")
80        .and_then(|v| v.to_str().ok())
81        .and_then(|v| v.strip_prefix("Bearer "))
82        .filter(|s| !s.is_empty())
83        .map(str::to_string)
84        .or_else(|| config_key.filter(|s| !s.is_empty()).map(str::to_string))
85}
86
87pub(super) fn sse_response(stream: BoxStream) -> Response {
88    let byte_stream = stream.map(|line| Ok::<Bytes, std::convert::Infallible>(Bytes::from(line)));
89    Response::builder()
90        .status(StatusCode::OK)
91        .header("Content-Type", "text/event-stream; charset=utf-8")
92        .header("Cache-Control", "no-cache")
93        .header("X-Accel-Buffering", "no")
94        .body(Body::from_stream(byte_stream))
95        .expect("valid SSE response")
96}