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
43#[allow(clippy::result_large_err)]
44pub(super) async fn read_bytes(body: Body) -> Result<Bytes, Response> {
45    read_bytes_with_auth(body, ProxyAuth::OpenAiBearer).await
46}
47
48#[allow(clippy::result_large_err)]
49pub(super) async fn read_bytes_with_auth(body: Body, auth: ProxyAuth) -> Result<Bytes, Response> {
50    axum::body::to_bytes(body, MAX_BODY_SIZE).await.map_err(|_| {
51        convert_response(error_response_for_auth(
52            StatusCode::PAYLOAD_TOO_LARGE,
53            "body_too_large",
54            "request body too large",
55            auth,
56        ))
57    })
58}
59
60#[allow(clippy::result_large_err)]
61pub(super) async fn read_and_parse(body: Body) -> Result<(Bytes, RequestPayload), Response> {
62    let bytes = read_bytes(body).await?;
63    let payload = serde_json::from_slice::<RequestPayload>(&bytes)
64        .map_err(|e| executor_error_response(ExecutorError::from(e)))?;
65    Ok((bytes, payload))
66}
67
68#[allow(clippy::result_large_err)]
69pub(super) async fn read_json<T: DeserializeOwned>(body: Body) -> Result<T, Response> {
70    let bytes = read_bytes(body).await?;
71    serde_json::from_slice::<T>(&bytes).map_err(|error| executor_error_response(ExecutorError::from(error)))
72}
73
74pub(super) fn extract_store(bytes: &[u8]) -> bool {
75    serde_json::from_slice::<serde_json::Value>(bytes)
76        .ok()
77        .and_then(|j| j.get("store").and_then(serde_json::Value::as_bool))
78        .unwrap_or(true)
79}
80
81pub(super) fn extract_bearer(headers: &HeaderMap, config_key: Option<&str>) -> Option<String> {
82    headers
83        .get("authorization")
84        .and_then(|v| v.to_str().ok())
85        .and_then(|v| v.strip_prefix("Bearer "))
86        .filter(|s| !s.is_empty())
87        .map(str::to_string)
88        .or_else(|| config_key.filter(|s| !s.is_empty()).map(str::to_string))
89}
90
91pub(super) fn sse_response(stream: BoxStream) -> Response {
92    sse_response_with_headers(stream, HeaderMap::new())
93}
94
95pub(super) fn sse_response_with_headers(stream: BoxStream, mut headers: HeaderMap) -> Response {
96    let byte_stream = stream.map(|line| Ok::<Bytes, std::convert::Infallible>(Bytes::from(line)));
97    headers.insert(
98        http::header::CONTENT_TYPE,
99        http::HeaderValue::from_static("text/event-stream; charset=utf-8"),
100    );
101    headers.insert(http::header::CACHE_CONTROL, http::HeaderValue::from_static("no-cache"));
102    headers.insert("x-accel-buffering", http::HeaderValue::from_static("no"));
103    let mut builder = Response::builder().status(StatusCode::OK);
104    for (name, value) in &headers {
105        builder = builder.header(name, value);
106    }
107    builder
108        .body(Body::from_stream(byte_stream))
109        .expect("valid SSE response")
110}