1use crate::backend::{
18 AgentBackend, AgentEvent, AgentSession, PromptMode, SessionExit, SessionSpec,
19};
20use crate::error::{EngineError, Result};
21use crate::stream_bounds::TailWindow;
22use crate::types::TokenUsage;
23use serde_json::{json, Value};
24use std::collections::VecDeque;
25use std::time::Duration;
26
27const BODY_TAIL_CHARS: usize = 500;
29
30const REQUEST_TIMEOUT: Duration = Duration::from_secs(600);
36
37const RESPONSE_BODY_CAP: usize = 8 * 1024 * 1024;
42
43const CHARS_PER_TOKEN: usize = 4;
49
50fn last_chars(text: &str, max: usize) -> String {
52 let chars: Vec<char> = text.chars().collect();
53 let start = chars.len().saturating_sub(max);
54 chars[start..].iter().collect()
55}
56
57fn prompt_text(spec: &SessionSpec) -> &str {
58 match &spec.prompt {
59 PromptMode::SingleShot(text) => text.as_str(),
60 PromptMode::Streaming(text) => text.as_str(),
61 }
62}
63
64fn build_messages(spec: &SessionSpec) -> Vec<Value> {
67 let mut messages = Vec::new();
68 if let Some(system) = &spec.append_system_prompt {
69 if !system.is_empty() {
70 messages.push(json!({"role": "system", "content": system}));
71 }
72 }
73 messages.push(json!({"role": "user", "content": prompt_text(spec)}));
74 messages
75}
76
77fn estimate_tokens(messages: &[Value]) -> u32 {
79 let total_chars: usize = messages
80 .iter()
81 .filter_map(|m| m.get("content").and_then(Value::as_str))
82 .map(|s| s.chars().count())
83 .sum();
84 total_chars.div_ceil(CHARS_PER_TOKEN) as u32
85}
86
87fn extract_completion(body: &Value) -> std::result::Result<(String, TokenUsage), String> {
90 let content = body
91 .get("choices")
92 .and_then(Value::as_array)
93 .and_then(|choices| choices.first())
94 .and_then(|choice| choice.get("message"))
95 .and_then(|message| message.get("content"))
96 .and_then(Value::as_str)
97 .ok_or_else(|| {
98 format!(
99 "response missing choices[0].message.content; body tail: {}",
100 last_chars(&body.to_string(), BODY_TAIL_CHARS)
101 )
102 })?
103 .to_string();
104 let usage = body.get("usage");
105 let input = usage
106 .and_then(|u| u.get("prompt_tokens"))
107 .and_then(Value::as_u64)
108 .unwrap_or(0);
109 let output = usage
110 .and_then(|u| u.get("completion_tokens"))
111 .and_then(Value::as_u64)
112 .unwrap_or(0);
113 Ok((
114 content,
115 TokenUsage {
116 input,
117 output,
118 cache_read: 0,
119 cache_write: 0,
120 },
121 ))
122}
123
124async fn read_body_tail(mut response: reqwest::Response, cap: usize) -> String {
130 let mut window = TailWindow::new(cap);
131 while let Ok(Some(chunk)) = response.chunk().await {
134 window.push(&chunk);
135 }
136 window.render()
137}
138
139#[derive(Debug, Clone)]
147pub struct LocalBackend {
148 base_url: String,
149 temperature: Option<f64>,
150 context_budget: u32,
151 request_timeout: Duration,
152 body_cap: usize,
153 client: reqwest::Client,
154}
155
156impl LocalBackend {
157 pub fn new(base_url: String, temperature: Option<f64>, context_budget: u32) -> Self {
158 LocalBackend {
159 base_url,
160 temperature,
161 context_budget,
162 request_timeout: REQUEST_TIMEOUT,
163 body_cap: RESPONSE_BODY_CAP,
164 client: reqwest::Client::new(),
165 }
166 }
167}
168
169#[async_trait::async_trait]
170impl AgentBackend for LocalBackend {
171 async fn start(&self, spec: SessionSpec) -> Result<Box<dyn AgentSession>> {
172 if spec.resume.is_some() {
173 return Err(EngineError::Backend(
174 "local backend is single-shot only; resume is unsupported".to_string(),
175 ));
176 }
177
178 let session_id = spec.session_id.clone();
179 let model = spec.model.clone();
180 let messages = build_messages(&spec);
181
182 let estimated_tokens = estimate_tokens(&messages);
183 if estimated_tokens > self.context_budget {
184 return Ok(Box::new(LocalSession::context_budget_exceeded(
185 session_id,
186 model,
187 estimated_tokens,
188 self.context_budget,
189 )));
190 }
191
192 let mut request_body = json!({
193 "model": model,
194 "messages": messages,
195 });
196 if let Some(temperature) = self.temperature {
197 request_body["temperature"] = json!(temperature);
198 }
199
200 let url = format!(
201 "{}/v1/chat/completions",
202 self.base_url.trim_end_matches('/')
203 );
204 let outcome = match self
205 .client
206 .post(&url)
207 .json(&request_body)
208 .timeout(self.request_timeout)
209 .send()
210 .await
211 {
212 Ok(response) => {
213 let status = response.status();
214 let body_text = read_body_tail(response, self.body_cap).await;
215 if status.is_success() {
216 match serde_json::from_str::<Value>(&body_text) {
217 Ok(parsed) => Ok(parsed),
218 Err(e) => Err(format!(
219 "failed to parse local backend response as JSON: {e}; body tail: {}",
220 last_chars(&body_text, BODY_TAIL_CHARS)
221 )),
222 }
223 } else {
224 Err(format!(
225 "local backend request failed with HTTP {status}; body tail: {}",
226 last_chars(&body_text, BODY_TAIL_CHARS)
227 ))
228 }
229 }
230 Err(e) if e.is_timeout() => Err(format!(
231 "local backend request timed out after {:?}",
232 self.request_timeout
233 )),
234 Err(e) => Err(format!("local backend request failed: {e}")),
235 };
236
237 Ok(Box::new(LocalSession::from_response(
238 session_id, model, outcome,
239 )))
240 }
241}
242
243pub struct LocalSession {
252 session_id: String,
253 queue: VecDeque<AgentEvent>,
254 pending_exit: Option<SessionExit>,
255 exit: Option<SessionExit>,
256}
257
258impl LocalSession {
259 fn from_response(
264 session_id: String,
265 model: String,
266 outcome: std::result::Result<Value, String>,
267 ) -> Self {
268 match outcome {
269 Ok(body) => match extract_completion(&body) {
270 Ok((content, usage)) => {
271 let mut queue = VecDeque::new();
272 queue.push_back(AgentEvent::Init {
273 session_id: session_id.clone(),
274 model,
275 raw: body.clone(),
276 });
277 queue.push_back(AgentEvent::Text {
278 text: content.clone(),
279 raw: body.clone(),
280 });
281 queue.push_back(AgentEvent::Result {
282 text: content,
283 is_error: false,
284 usage,
285 cost_usd: Some(0.0),
286 num_turns: Some(1),
287 raw: body,
288 });
289 LocalSession {
290 session_id,
291 queue,
292 pending_exit: Some(SessionExit::Completed),
293 exit: None,
294 }
295 }
296 Err(message) => LocalSession {
297 session_id,
298 queue: VecDeque::new(),
299 pending_exit: Some(SessionExit::Failed(message)),
300 exit: None,
301 },
302 },
303 Err(message) => LocalSession {
304 session_id,
305 queue: VecDeque::new(),
306 pending_exit: Some(SessionExit::Failed(message)),
307 exit: None,
308 },
309 }
310 }
311
312 fn context_budget_exceeded(
316 session_id: String,
317 model: String,
318 estimated_tokens: u32,
319 context_budget: u32,
320 ) -> Self {
321 let mut queue = VecDeque::new();
322 queue.push_back(AgentEvent::Init {
323 session_id: session_id.clone(),
324 model,
325 raw: json!({}),
326 });
327 LocalSession {
328 session_id,
329 queue,
330 pending_exit: Some(SessionExit::Failed(format!(
331 "prompt estimated at {estimated_tokens} tokens exceeds context budget of \
332 {context_budget} tokens; no request was sent"
333 ))),
334 exit: None,
335 }
336 }
337}
338
339#[async_trait::async_trait]
340impl AgentSession for LocalSession {
341 fn session_id(&self) -> String {
342 self.session_id.clone()
343 }
344
345 async fn next_event(&mut self) -> Result<Option<AgentEvent>> {
346 if let Some(event) = self.queue.pop_front() {
347 return Ok(Some(event));
348 }
349 if self.exit.is_none() {
350 self.exit = self.pending_exit.take();
351 }
352 Ok(None)
353 }
354
355 async fn send_user_message(&mut self, _text: &str) -> Result<()> {
356 Err(EngineError::Backend(
357 "local backend is single-shot only; send_user_message is unsupported".to_string(),
358 ))
359 }
360
361 async fn abort(&mut self) -> Result<()> {
362 self.queue.clear();
363 if self.exit.is_none() {
364 self.exit = Some(SessionExit::Aborted);
365 }
366 self.pending_exit = None;
367 Ok(())
368 }
369
370 fn exit_status(&self) -> Option<SessionExit> {
371 self.exit.clone()
372 }
373}
374
375#[cfg(test)]
376pub(crate) mod tests {
380 use super::*;
381 use std::collections::HashMap;
382 use std::net::SocketAddr;
383 use std::path::PathBuf;
384 use std::sync::atomic::{AtomicUsize, Ordering};
385 use std::sync::{Arc, Mutex};
386 use tokio::io::{AsyncReadExt, AsyncWriteExt};
387 use tokio::net::{TcpListener, TcpStream};
388
389 const TEST_MODEL: &str = "local-test-model";
390
391 fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
392 haystack
393 .windows(needle.len())
394 .position(|window| window == needle)
395 }
396
397 async fn read_http_request(socket: &mut TcpStream) -> Vec<u8> {
401 let mut buf = Vec::new();
402 let mut chunk = [0u8; 4096];
403 loop {
404 let header_end = find_subslice(&buf, b"\r\n\r\n");
405 if let Some(header_end) = header_end {
406 let headers = String::from_utf8_lossy(&buf[..header_end]).to_string();
407 let content_length: usize = headers
408 .lines()
409 .find_map(|line| {
410 let (name, value) = line.split_once(':')?;
411 if name.eq_ignore_ascii_case("content-length") {
412 value.trim().parse().ok()
413 } else {
414 None
415 }
416 })
417 .unwrap_or(0);
418 let body_start = header_end + 4;
419 if buf.len() >= body_start + content_length {
420 break;
421 }
422 }
423 match socket.read(&mut chunk).await {
424 Ok(0) => break,
425 Ok(n) => buf.extend_from_slice(&chunk[..n]),
426 Err(_) => break,
427 }
428 }
429 buf
430 }
431
432 pub(crate) async fn spawn_stub(
441 status_line: &'static str,
442 body: String,
443 ) -> (String, Arc<AtomicUsize>, Arc<Mutex<Vec<u8>>>) {
444 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind stub");
445 let addr: SocketAddr = listener.local_addr().expect("stub addr");
446 let count = Arc::new(AtomicUsize::new(0));
447 let count_for_task = Arc::clone(&count);
448 let received = Arc::new(Mutex::new(Vec::new()));
449 let received_for_task = Arc::clone(&received);
450 tokio::spawn(async move {
451 loop {
452 let (mut socket, _) = match listener.accept().await {
453 Ok(v) => v,
454 Err(_) => break,
455 };
456 count_for_task.fetch_add(1, Ordering::SeqCst);
457 let request_bytes = read_http_request(&mut socket).await;
458 *received_for_task.lock().expect("stub request lock") = request_bytes;
459 let response = format!(
460 "{status_line}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
461 body.len(),
462 body
463 );
464 let _ = socket.write_all(response.as_bytes()).await;
465 let _ = socket.shutdown().await;
466 }
467 });
468 (format!("http://{addr}"), count, received)
469 }
470
471 fn base_spec(session_id: &str, prompt: &str, context_budget_prompt: bool) -> SessionSpec {
472 let _ = context_budget_prompt;
473 SessionSpec {
474 cwd: PathBuf::from("."),
475 prompt: PromptMode::SingleShot(prompt.to_string()),
476 append_system_prompt: Some("be terse".to_string()),
477 model: TEST_MODEL.to_string(),
478 effort: String::new(),
479 session_id: session_id.to_string(),
480 resume: None,
481 permission_mode: None,
482 allowed_tools: vec![],
483 disallowed_tools: vec![],
484 tools: vec![],
485 writable: true,
486 settings_json: None,
487 json_schema: None,
488 max_budget_usd: None,
489 max_turns: None,
490 env: HashMap::new(),
491 sandbox: None,
492 hook_status: None,
493 }
494 }
495
496 async fn drain(session: &mut dyn AgentSession) -> Vec<AgentEvent> {
497 let mut events = Vec::new();
498 while let Some(event) = session.next_event().await.expect("next_event") {
499 events.push(event);
500 }
501 events
502 }
503
504 #[tokio::test]
505 async fn local_http_roundtrip_yields_init_text_result_with_usage_and_zero_cost() {
506 let stub_body = json!({
507 "id": "chatcmpl-1",
508 "choices": [{"message": {"role": "assistant", "content": "hello from stub"}}],
509 "usage": {"prompt_tokens": 12, "completion_tokens": 34, "total_tokens": 46}
510 })
511 .to_string();
512 let (base_url, requests, received) = spawn_stub("HTTP/1.1 200 OK", stub_body).await;
513
514 let backend = LocalBackend::new(base_url, Some(0.2), 100_000);
515 let spec = base_spec("sess-1", "do the thing", false);
516 let mut session = backend.start(spec).await.expect("start");
517
518 let events = drain(session.as_mut()).await;
519 assert_eq!(requests.load(Ordering::SeqCst), 1);
520
521 let raw_request = received.lock().expect("stub request lock").clone();
522 let request_text = String::from_utf8_lossy(&raw_request).to_string();
523 let request_line = request_text.lines().next().expect("request line");
524 assert!(
525 request_line.starts_with("POST "),
526 "expected a POST request, got: {request_line}"
527 );
528 assert!(
529 request_line
530 .split_whitespace()
531 .nth(1)
532 .expect("request target")
533 .ends_with("/v1/chat/completions"),
534 "expected the request target to end with /v1/chat/completions, got: {request_line}"
535 );
536 let header_end = find_subslice(&raw_request, b"\r\n\r\n").expect("request headers");
537 let request_body: Value = serde_json::from_slice(&raw_request[header_end + 4..])
538 .expect("request body should be JSON");
539 assert_eq!(request_body["model"], json!(TEST_MODEL));
540 assert_eq!(request_body["temperature"], json!(0.2));
541 let messages = request_body["messages"].as_array().expect("messages array");
542 assert!(
543 messages
544 .iter()
545 .any(|m| m["role"] == "system" && m["content"] == "be terse"),
546 "expected the system message from append_system_prompt, got: {messages:?}"
547 );
548 assert_eq!(
549 messages.last().expect("at least one message"),
550 &json!({"role": "user", "content": "do the thing"})
551 );
552
553 assert!(
554 matches!(&events[0], AgentEvent::Init { session_id, model, .. }
555 if session_id == "sess-1" && model == TEST_MODEL)
556 );
557 assert!(matches!(&events[1], AgentEvent::Text { text, .. } if text == "hello from stub"));
558 match &events[2] {
559 AgentEvent::Result {
560 text,
561 is_error,
562 usage,
563 cost_usd,
564 num_turns,
565 ..
566 } => {
567 assert_eq!(text, "hello from stub");
568 assert!(!is_error);
569 assert_eq!(usage.input, 12);
570 assert_eq!(usage.output, 34);
571 assert_eq!(*cost_usd, Some(0.0));
572 assert_eq!(*num_turns, Some(1));
573 }
574 other => panic!("expected terminal Result, got {other:?}"),
575 }
576 assert_eq!(events.len(), 3);
577 assert_eq!(session.exit_status(), Some(SessionExit::Completed));
578 }
579
580 #[tokio::test]
581 async fn local_http_rejects_resumed_spec() {
582 let backend = LocalBackend::new("http://127.0.0.1:1".to_string(), None, 100_000);
583 let mut spec = base_spec("sess-1", "do the thing", false);
584 spec.resume = Some("sess-0".to_string());
585
586 let result = backend.start(spec).await;
587 assert!(result.is_err(), "expected resume to be rejected");
588 }
589
590 #[tokio::test]
591 async fn local_http_send_user_message_errors() {
592 let stub_body = json!({
593 "choices": [{"message": {"role": "assistant", "content": "hi"}}],
594 "usage": {"prompt_tokens": 1, "completion_tokens": 1}
595 })
596 .to_string();
597 let (base_url, _requests, _received) = spawn_stub("HTTP/1.1 200 OK", stub_body).await;
598
599 let backend = LocalBackend::new(base_url, None, 100_000);
600 let spec = base_spec("sess-1", "do the thing", false);
601 let mut session = backend.start(spec).await.expect("start");
602
603 let result = session.send_user_message("nope").await;
604 assert!(result.is_err(), "expected send_user_message to be rejected");
605 }
606
607 #[tokio::test]
608 async fn local_http_500_fails_cleanly() {
609 let (base_url, requests, _received) =
610 spawn_stub("HTTP/1.1 500 Internal Server Error", "boom".to_string()).await;
611
612 let backend = LocalBackend::new(base_url, None, 100_000);
613 let spec = base_spec("sess-1", "do the thing", false);
614 let mut session = backend.start(spec).await.expect("start");
615
616 let events = drain(session.as_mut()).await;
617 assert!(events.is_empty(), "expected no events on a failed response");
618 assert_eq!(requests.load(Ordering::SeqCst), 1);
619
620 match session.exit_status() {
621 Some(SessionExit::Failed(message)) => {
622 assert!(
623 message.contains("500"),
624 "expected the failure message to include the HTTP status, got: {message}"
625 );
626 }
627 other => panic!("expected SessionExit::Failed, got {other:?}"),
628 }
629 }
630
631 #[tokio::test]
632 async fn local_http_context_budget_exceeds_fails_cleanly() {
633 let (base_url, requests, _received) = spawn_stub("HTTP/1.1 200 OK", "{}".to_string()).await;
634
635 let backend = LocalBackend::new(base_url, None, 1);
637 let spec = base_spec(
638 "sess-1",
639 "this prompt is far too long for a one-token budget",
640 false,
641 );
642 let mut session = backend.start(spec).await.expect("start");
643
644 let events = drain(session.as_mut()).await;
645 assert_eq!(
646 requests.load(Ordering::SeqCst),
647 0,
648 "must not send an HTTP request when the context budget is exceeded"
649 );
650 assert_eq!(events.len(), 1, "expected only a synthesized Init event");
651 assert!(matches!(&events[0], AgentEvent::Init { .. }));
652
653 match session.exit_status() {
654 Some(SessionExit::Failed(message)) => {
655 assert!(
656 message.contains("context budget"),
657 "expected the failure message to name the context budget, got: {message}"
658 );
659 }
660 other => panic!("expected SessionExit::Failed, got {other:?}"),
661 }
662 }
663
664 async fn spawn_hung_stub() -> String {
669 let listener = TcpListener::bind("127.0.0.1:0")
670 .await
671 .expect("bind hung stub");
672 let addr: SocketAddr = listener.local_addr().expect("hung stub addr");
673 tokio::spawn(async move {
674 let mut held = Vec::new();
675 while let Ok((socket, _)) = listener.accept().await {
676 held.push(socket);
677 }
678 });
679 format!("http://{addr}")
680 }
681
682 #[tokio::test]
686 async fn local_http_hung_endpoint_times_out_instead_of_stalling() {
687 let base_url = spawn_hung_stub().await;
688 let mut backend = LocalBackend::new(base_url, None, 100_000);
689 backend.request_timeout = Duration::from_millis(200);
690
691 let start = std::time::Instant::now();
692 let spec = base_spec("sess-hung", "do the thing", false);
693 let mut session = backend.start(spec).await.expect("start");
694 let events = drain(session.as_mut()).await;
695
696 assert!(events.is_empty(), "a timed-out request yields no events");
697 assert!(
698 start.elapsed() < Duration::from_secs(10),
699 "the request returned near the 200ms timeout, not after a stall"
700 );
701 match session.exit_status() {
702 Some(SessionExit::Failed(message)) => {
703 assert!(
704 message.contains("timed out"),
705 "expected the failure to name the timeout, got: {message}"
706 );
707 }
708 other => panic!("expected SessionExit::Failed, got {other:?}"),
709 }
710 }
711
712 #[tokio::test]
715 async fn local_http_error_body_over_the_cap_is_tailed_with_a_marker() {
716 let body = format!("{}{}", "x".repeat(4096), "BODY-END");
717 let (base_url, _requests, _received) =
718 spawn_stub("HTTP/1.1 500 Internal Server Error", body).await;
719 let mut backend = LocalBackend::new(base_url, None, 100_000);
720 backend.body_cap = 128;
721
722 let spec = base_spec("sess-cap", "do the thing", false);
723 let mut session = backend.start(spec).await.expect("start");
724 let events = drain(session.as_mut()).await;
725 assert!(events.is_empty(), "expected no events on a failed response");
726
727 match session.exit_status() {
728 Some(SessionExit::Failed(message)) => {
729 assert!(
730 message.contains(crate::stream_bounds::TRUNCATION_MARKER),
731 "expected the truncation marker, got: {message}"
732 );
733 assert!(
734 message.contains("BODY-END"),
735 "expected the END of the body to be kept, got: {message}"
736 );
737 assert!(
738 message.len() < 1024,
739 "the surfaced body stayed bounded, got {} bytes",
740 message.len()
741 );
742 }
743 other => panic!("expected SessionExit::Failed, got {other:?}"),
744 }
745 }
746
747 #[tokio::test]
751 async fn local_http_success_body_over_the_cap_fails_honestly_with_a_marker() {
752 let body = format!(
753 "{}{}",
754 json!({"choices": [{"message": {"role": "assistant", "content": "hi"}}],
755 "usage": {"prompt_tokens": 1, "completion_tokens": 1}}),
756 " ".repeat(4096)
757 );
758 let (base_url, _requests, _received) = spawn_stub("HTTP/1.1 200 OK", body).await;
759 let mut backend = LocalBackend::new(base_url, None, 100_000);
760 backend.body_cap = 128;
761
762 let spec = base_spec("sess-cap200", "do the thing", false);
763 let mut session = backend.start(spec).await.expect("start");
764 let events = drain(session.as_mut()).await;
765 assert!(
766 events.is_empty(),
767 "an unparseable over-cap body yields no events"
768 );
769
770 match session.exit_status() {
771 Some(SessionExit::Failed(message)) => {
772 assert!(
773 message.contains("failed to parse"),
774 "expected an honest parse failure, got: {message}"
775 );
776 assert!(
777 message.contains(crate::stream_bounds::TRUNCATION_MARKER),
778 "expected the truncation marker, got: {message}"
779 );
780 }
781 other => panic!("expected SessionExit::Failed, got {other:?}"),
782 }
783 }
784}