Skip to main content

agentic_core/executor/
inference.rs

1//! HTTP transport layer for LLM backend communication.
2//!
3//! Handles sending requests, reading streaming chunks, and mapping network
4//! and HTTP errors to [`ExecutorError`].
5
6use std::sync::Arc;
7use std::time::Duration;
8
9use async_stream::stream;
10use futures::{Stream, StreamExt};
11
12use crate::executor::error::{ExecutorError, ExecutorResult};
13use crate::proxy::processed_response_headers;
14
15/// SSE stream of raw lines sent to the client (`data: …\n\n` per event).
16pub type BoxStream = std::pin::Pin<Box<dyn Stream<Item = String> + Send>>;
17
18/// Wire-format marker signalling end-of-stream to the client.
19pub(super) const DONE_MARKER: &str = "data: [DONE]\n\n";
20
21/// Fetch the next raw bytes chunk from a streaming response.
22///
23/// Returns `Ok(Some(bytes))` on data, `Ok(None)` when the stream ends cleanly,
24/// and `Err` on a network failure or chunk timeout.
25pub(super) async fn next_chunk<S>(stream: &mut S, timeout: Duration) -> ExecutorResult<Option<bytes::Bytes>>
26where
27    S: futures::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Unpin,
28{
29    let item = if timeout.is_zero() {
30        stream.next().await
31    } else {
32        tokio::time::timeout(timeout, stream.next()).await.map_err(|_| {
33            ExecutorError::StreamError("chunk timeout: no data received within the configured window".into())
34        })?
35    };
36    item.transpose().map_err(ExecutorError::NetworkError)
37}
38
39fn drain_complete_utf8_lines(buffer: &mut Vec<u8>) -> Vec<String> {
40    let mut lines = Vec::new();
41    while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
42        let line = buffer.drain(..=pos).collect::<Vec<_>>();
43        let line_end = if pos > 0 && line.get(pos - 1) == Some(&b'\r') {
44            pos - 1
45        } else {
46            pos
47        };
48        if let Ok(line) = std::str::from_utf8(&line[..line_end]) {
49            lines.push(line.to_string());
50        }
51    }
52    lines
53}
54
55/// Build, send, and validate an HTTP POST to the LLM backend.
56///
57/// Shared by both the blocking path (caller reads `.text()`) and the streaming
58/// path (caller reads `.bytes_stream()`). Maps connect/timeout failures and
59/// non-2xx status codes to [`ExecutorError::LLMRequest`] and connection
60/// failures to [`ExecutorError::LLMTransport`].
61pub(super) async fn send_request(
62    client: &reqwest::Client,
63    url: &str,
64    body: String,
65    auth: Option<&str>,
66    forwarded_headers: Option<&reqwest::header::HeaderMap>,
67) -> ExecutorResult<reqwest::Response> {
68    let mut headers = forwarded_headers.cloned().unwrap_or_default();
69    headers
70        .entry(reqwest::header::CONTENT_TYPE)
71        .or_insert(reqwest::header::HeaderValue::from_static("application/json"));
72    let mut req = client.post(url).headers(headers).body(body);
73    if let Some(key) = auth {
74        req = req.bearer_auth(key);
75    }
76
77    let resp = req.send().await.map_err(|e| ExecutorError::LLMTransport {
78        status: if e.is_timeout() {
79            http::StatusCode::GATEWAY_TIMEOUT
80        } else {
81            http::StatusCode::BAD_GATEWAY
82        },
83        message: if e.is_timeout() {
84            "LLM timeout"
85        } else {
86            "LLM unavailable"
87        },
88    })?;
89
90    if !resp.status().is_success() {
91        let status = resp.status().as_u16();
92        let headers = processed_response_headers(resp.headers());
93        // Log and discard any error reading the error body — the status code
94        // is the primary signal; an empty body is acceptable here.
95        let body = resp
96            .text()
97            .await
98            .inspect_err(|e| tracing::debug!("failed to read error response body: {e}"))
99            .unwrap_or_default();
100        return Err(ExecutorError::LLMRequest {
101            status: http::StatusCode::from_u16(status).unwrap_or(http::StatusCode::INTERNAL_SERVER_ERROR),
102            body,
103            headers,
104        });
105    }
106
107    Ok(resp)
108}
109
110/// Makes a non-streaming HTTP POST to the LLM backend and returns the full JSON body.
111///
112/// Used by `run_blocking` so it can pass the result to [`ResponseAccumulator::from_json`](crate::executor::accumulator::ResponseAccumulator::from_json).
113pub(super) async fn fetch_response_json(
114    upstream_json: String,
115    url: &str,
116    client: &reqwest::Client,
117    auth: Option<&str>,
118) -> ExecutorResult<String> {
119    let resp = send_request(client, url, upstream_json, auth, None).await?;
120    // Preserve the reqwest::Error as the typed source (NetworkError).
121    resp.text().await.map_err(ExecutorError::NetworkError)
122}
123
124/// Makes a non-streaming HTTP POST with caller-supplied upstream headers.
125pub(super) async fn fetch_response_json_with_headers(
126    upstream_json: String,
127    url: &str,
128    client: &reqwest::Client,
129    headers: &reqwest::header::HeaderMap,
130) -> ExecutorResult<(String, http::HeaderMap)> {
131    let resp = send_request(client, url, upstream_json, None, Some(headers)).await?;
132    let response_headers = processed_response_headers(resp.headers());
133    let body = resp.text().await.map_err(ExecutorError::NetworkError)?;
134    Ok((body, response_headers))
135}
136
137/// Step 2 — Call the LLM inference backend; yields raw SSE lines (`data: …`).
138///
139/// Always requests `stream=true` upstream. Stops on `[DONE]`.
140///
141/// # Errors
142/// Each stream item is `Result<String, ExecutorError>`. The stream yields `Err` on:
143/// - [`ExecutorError::LLMTransport`] — connect timeout (504) or connection failure (502)
144/// - [`ExecutorError::LLMRequest`] — non-2xx HTTP status from the backend
145/// - [`ExecutorError::NetworkError`] — network failure while reading the response body
146pub fn call_inference(
147    upstream_json: String,
148    url: String,
149    client: Arc<reqwest::Client>,
150    auth: Option<String>,
151    chunk_timeout: Duration,
152) -> impl Stream<Item = Result<String, ExecutorError>> + Send + 'static {
153    stream! {
154        let resp = match send_request(&client, &url, upstream_json, auth.as_deref(), None).await {
155            Ok(r) => r,
156            Err(e) => { yield Err(e); return; }
157        };
158
159        let mut lines = Box::pin(response_lines(resp, chunk_timeout));
160        while let Some(line) = lines.next().await {
161            yield line;
162        }
163    }
164}
165
166/// Convert a successful upstream response body into normalized SSE data lines.
167pub(super) fn response_lines(
168    resp: reqwest::Response,
169    chunk_timeout: Duration,
170) -> impl Stream<Item = Result<String, ExecutorError>> + Send + 'static {
171    stream! {
172        let mut bytes = resp.bytes_stream();
173        let mut buf = Vec::with_capacity(8192);
174
175        loop {
176            let chunk = match next_chunk(&mut bytes, chunk_timeout).await {
177                Ok(Some(c)) => c,
178                Ok(None) => break,
179                Err(e) => { yield Err(e); return; }
180            };
181
182            buf.extend_from_slice(&chunk);
183
184            for line in drain_complete_utf8_lines(&mut buf) {
185                match line.as_str() {
186                    "data: [DONE]" => return,
187                    l if l.starts_with("data: ") => yield Ok(line),
188                    _ => {}
189                }
190            }
191        }
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn utf8_line_reader_preserves_split_multibyte_characters() {
201        let snowman = "\u{2603}";
202        let line = format!(r#"data: {{"delta":"snow {snowman}"}}"#);
203        let bytes = format!("{line}\n").into_bytes();
204        let split_at = bytes
205            .windows(snowman.len())
206            .position(|window| window == snowman.as_bytes())
207            .expect("snowman bytes present")
208            + 1;
209        let mut buffer = bytes[..split_at].to_vec();
210
211        assert!(drain_complete_utf8_lines(&mut buffer).is_empty());
212
213        buffer.extend_from_slice(&bytes[split_at..]);
214        let lines = drain_complete_utf8_lines(&mut buffer);
215
216        assert!(buffer.is_empty());
217        assert_eq!(lines, vec![line]);
218        assert!(!lines[0].contains('\u{FFFD}'));
219    }
220}