heartbit_core/llm/
mod.rs

1//! LLM provider abstractions — `LlmProvider` trait, Anthropic/Gemini/OpenRouter/OpenAI-compat backends, retry, cascade, and circuit-breaker wrappers.
2
3pub mod anthropic;
4pub mod circuit;
5
6/// Maximum bytes read from an upstream LLM error body.
7///
8/// SECURITY (F-LLM-5, F-LLM-6): a hostile or compromised provider can stream
9/// gigabytes of body in response to a 4xx/5xx and OOM the agent. We also
10/// truncate to keep accidentally-included secrets / internal IPs from
11/// flooding logs.
12const ERROR_BODY_MAX_BYTES: usize = 8 << 10; // 8 KiB
13
14/// SECURITY (F-LLM-4): hard cap on accumulated streaming text bytes per
15/// response. A drip-fed `text_delta` event sequence would otherwise grow
16/// unbounded.
17pub(crate) const STREAM_MAX_TEXT_BYTES: usize = 16 << 20; // 16 MiB
18
19/// SECURITY (F-LLM-4): hard cap on accumulated tool-call arguments JSON per
20/// individual tool call.
21pub(crate) const STREAM_MAX_TOOL_ARGS_BYTES: usize = 1 << 20; // 1 MiB
22
23/// SECURITY (F-LLM-4): hard cap on the number of tool calls a single
24/// streaming response may emit. Protects against a hostile `tool_calls[].index`
25/// of `u32::MAX` triggering a multi-billion-entry Vec allocation.
26pub(crate) const STREAM_MAX_TOOL_CALLS: usize = 256;
27
28/// SECURITY (F-LLM-5): hard cap on a NON-streaming success body. A hostile or
29/// compromised provider can stream gigabytes to a `200 OK` just as easily as
30/// to an error, OOMing the agent before `serde_json` ever runs. 16 MiB is far
31/// above any real completion JSON (max-tokens × bytes/token) yet bounds the
32/// allocation.
33pub(crate) const SUCCESS_BODY_MAX_BYTES: usize = 16 << 20; // 16 MiB
34
35/// Deserialize a successful response body as JSON, capped at
36/// [`SUCCESS_BODY_MAX_BYTES`]. Use instead of `response.json().await?`, which
37/// buffers the entire (potentially unbounded) body first.
38pub(crate) async fn read_json_capped<T: serde::de::DeserializeOwned>(
39    response: reqwest::Response,
40) -> Result<T, Error> {
41    let (bytes, truncated) =
42        crate::http::read_body_capped(response, SUCCESS_BODY_MAX_BYTES).await?;
43    if truncated {
44        return Err(Error::Api {
45            status: 200,
46            message: format!("response body exceeded {SUCCESS_BODY_MAX_BYTES}-byte cap"),
47        });
48    }
49    serde_json::from_slice(&bytes).map_err(|e| Error::Api {
50        status: 200,
51        message: format!("failed to parse response JSON: {e}"),
52    })
53}
54
55/// Build an `Error::Api` from a failed HTTP response, sanitizing auth errors.
56///
57/// For 401/403 responses, the body is NOT read to avoid leaking API key
58/// fragments in logs. For all other statuses the response body is included
59/// but capped at `ERROR_BODY_MAX_BYTES` and stripped of control characters
60/// (newlines, ANSI escapes) so a hostile body cannot poison structured logs
61/// (F-LLM-6).
62pub(crate) async fn api_error_from_response(response: reqwest::Response) -> Error {
63    use futures::TryStreamExt;
64    let status = response.status().as_u16();
65    let message = if status == 401 || status == 403 {
66        format!("authentication failed (HTTP {status})")
67    } else {
68        let mut buf: Vec<u8> = Vec::with_capacity(2048);
69        let mut stream = response.bytes_stream();
70        let mut overflowed = false;
71        loop {
72            match stream.try_next().await {
73                Ok(Some(chunk)) => {
74                    let remaining = ERROR_BODY_MAX_BYTES.saturating_sub(buf.len());
75                    if remaining == 0 {
76                        overflowed = true;
77                        break;
78                    }
79                    let take = chunk.len().min(remaining);
80                    buf.extend_from_slice(&chunk[..take]);
81                    if take < chunk.len() {
82                        overflowed = true;
83                        break;
84                    }
85                }
86                Ok(None) => break,
87                Err(e) => {
88                    return Error::Api {
89                        status,
90                        message: format!("<body read error: {e}>"),
91                    };
92                }
93            }
94        }
95        let mut text = String::from_utf8_lossy(&buf).to_string();
96        // Strip control characters (CR/LF/ANSI ESC). Keep tabs and printable.
97        text.retain(|c| c == '\t' || (!c.is_control() && c != '\u{1b}'));
98        if overflowed {
99            text.push_str("…[truncated]");
100        }
101        text
102    };
103    Error::Api { status, message }
104}
105pub mod cascade;
106pub mod error_class;
107pub mod gemini;
108pub mod openai_compat;
109pub mod openrouter;
110pub mod pricing;
111pub mod registry;
112pub mod retry;
113pub mod types;
114
115use std::future::Future;
116use std::pin::Pin;
117use std::sync::Arc;
118
119use crate::error::Error;
120use crate::llm::types::{CompletionRequest, CompletionResponse};
121
122/// Callback invoked with each text delta during streaming.
123pub type OnText = dyn Fn(&str) + Send + Sync;
124
125/// Callback invoked with each reasoning (chain-of-thought) delta during
126/// streaming, for reasoning models (qwen3-thinking, deepseek-r1). Separate from
127/// [`OnText`] so a UI can render thinking distinctly from the answer.
128pub type OnReasoning = dyn Fn(&str) + Send + Sync;
129
130/// Decision returned by the `OnApproval` callback.
131///
132/// `Allow` and `Deny` behave like the previous `true`/`false` return.
133/// `AlwaysAllow` and `AlwaysDeny` additionally persist the decision as a
134/// learned permission rule so it survives across sessions.
135#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub enum ApprovalDecision {
137    /// Allow this time.
138    Allow,
139    /// Deny this time.
140    Deny,
141    /// Allow and persist as a permission rule.
142    AlwaysAllow,
143    /// Deny and persist as a permission rule.
144    AlwaysDeny,
145}
146
147impl ApprovalDecision {
148    /// Returns `true` when the decision allows execution.
149    pub fn is_allowed(self) -> bool {
150        matches!(self, Self::Allow | Self::AlwaysAllow)
151    }
152
153    /// Returns `true` when the decision should be persisted.
154    pub fn is_persistent(self) -> bool {
155        matches!(self, Self::AlwaysAllow | Self::AlwaysDeny)
156    }
157}
158
159impl From<bool> for ApprovalDecision {
160    fn from(allowed: bool) -> Self {
161        if allowed { Self::Allow } else { Self::Deny }
162    }
163}
164
165/// Callback invoked before tool execution for human-in-the-loop approval.
166///
167/// Receives the list of tool calls the LLM wants to execute.
168/// Returns an [`ApprovalDecision`] indicating whether to proceed.
169/// `AlwaysAllow`/`AlwaysDeny` additionally persist the decision as a
170/// learned permission rule.
171pub type OnApproval = dyn Fn(&[crate::llm::types::ToolCall]) -> ApprovalDecision + Send + Sync;
172
173/// Trait for LLM providers.
174///
175/// Uses RPITIT (`impl Future`) which means this trait is NOT dyn-compatible.
176/// All consumers are generic over `P: LlmProvider`. This is intentional:
177/// one provider per process, no need for trait objects.
178///
179/// For dynamic dispatch, use [`BoxedProvider`] which wraps any `LlmProvider`
180/// behind [`DynLlmProvider`].
181pub trait LlmProvider: Send + Sync {
182    /// Send a completion request and wait for the full response.
183    fn complete(
184        &self,
185        request: CompletionRequest,
186    ) -> impl Future<Output = Result<CompletionResponse, Error>> + Send;
187
188    /// Stream a completion, calling `on_text` for each text delta as it arrives.
189    ///
190    /// The returned `CompletionResponse` contains the full accumulated response
191    /// (same as `complete()`), but text was emitted incrementally via the callback.
192    ///
193    /// Default: falls back to `complete()` (no incremental streaming).
194    fn stream_complete(
195        &self,
196        request: CompletionRequest,
197        on_text: &OnText,
198    ) -> impl Future<Output = Result<CompletionResponse, Error>> + Send {
199        let _ = on_text;
200        self.complete(request)
201    }
202
203    /// Stream a completion, additionally calling `on_reasoning` for each
204    /// chain-of-thought delta (reasoning models). The answer still arrives via
205    /// `on_text`; reasoning is a separate channel.
206    ///
207    /// Default: delegates to [`stream_complete`](Self::stream_complete),
208    /// ignoring `on_reasoning` (providers that don't stream reasoning, and the
209    /// non-streaming fallback, are unaffected — `CompletionResponse.reasoning`
210    /// still carries the post-hoc accumulation).
211    fn stream_complete_with_reasoning(
212        &self,
213        request: CompletionRequest,
214        on_text: &OnText,
215        on_reasoning: &OnReasoning,
216    ) -> impl Future<Output = Result<CompletionResponse, Error>> + Send {
217        let _ = on_reasoning;
218        self.stream_complete(request, on_text)
219    }
220
221    /// Return the model identifier, if known.
222    ///
223    /// Used for audit trail events. Default returns `None`.
224    fn model_name(&self) -> Option<&str> {
225        None
226    }
227}
228
229// ---------------------------------------------------------------------------
230// DynLlmProvider — object-safe adapter for LlmProvider (RPITIT → dyn)
231// ---------------------------------------------------------------------------
232
233/// Object-safe version of [`LlmProvider`] for dynamic dispatch.
234///
235/// `LlmProvider` uses RPITIT (not dyn-compatible). This trait wraps it via
236/// `Pin<Box<dyn Future>>` so providers can be stored as `Arc<dyn DynLlmProvider>`.
237///
238/// A blanket impl covers all `LlmProvider` types automatically.
239///
240/// Used by the Restate service layer (`AgentServiceImpl`) and by
241/// [`BoxedProvider`] for type-erased standalone use.
242pub trait DynLlmProvider: Send + Sync {
243    /// Boxed-future version of [`LlmProvider::complete`] for object-safe dispatch.
244    fn complete<'a>(
245        &'a self,
246        request: CompletionRequest,
247    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>>;
248
249    /// Boxed-future version of [`LlmProvider::stream_complete`] for object-safe dispatch.
250    fn stream_complete<'a>(
251        &'a self,
252        request: CompletionRequest,
253        on_text: &'a OnText,
254    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>>;
255
256    /// Boxed-future version of [`LlmProvider::stream_complete_with_reasoning`].
257    ///
258    /// Default delegates to [`stream_complete`](Self::stream_complete), ignoring
259    /// reasoning — so direct `DynLlmProvider` impls (e.g. test mocks) need not
260    /// implement it. The blanket impl overrides this to forward reasoning.
261    fn stream_complete_with_reasoning<'a>(
262        &'a self,
263        request: CompletionRequest,
264        on_text: &'a OnText,
265        on_reasoning: &'a OnReasoning,
266    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>> {
267        let _ = on_reasoning;
268        self.stream_complete(request, on_text)
269    }
270
271    /// Return the model identifier, if known.
272    fn model_name(&self) -> Option<&str>;
273}
274
275impl<P: LlmProvider> DynLlmProvider for P {
276    fn complete<'a>(
277        &'a self,
278        request: CompletionRequest,
279    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>> {
280        Box::pin(LlmProvider::complete(self, request))
281    }
282
283    fn stream_complete<'a>(
284        &'a self,
285        request: CompletionRequest,
286        on_text: &'a OnText,
287    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>> {
288        Box::pin(LlmProvider::stream_complete(self, request, on_text))
289    }
290
291    fn stream_complete_with_reasoning<'a>(
292        &'a self,
293        request: CompletionRequest,
294        on_text: &'a OnText,
295        on_reasoning: &'a OnReasoning,
296    ) -> Pin<Box<dyn Future<Output = Result<CompletionResponse, Error>> + Send + 'a>> {
297        Box::pin(LlmProvider::stream_complete_with_reasoning(
298            self,
299            request,
300            on_text,
301            on_reasoning,
302        ))
303    }
304
305    fn model_name(&self) -> Option<&str> {
306        LlmProvider::model_name(self)
307    }
308}
309
310// ---------------------------------------------------------------------------
311// BoxedProvider — type-erased LlmProvider via DynLlmProvider
312// ---------------------------------------------------------------------------
313
314/// Type-erased LLM provider for use when dynamic dispatch is needed.
315///
316/// Wraps any [`LlmProvider`] behind `Box<dyn DynLlmProvider>`. Implements
317/// `LlmProvider` itself, so it can be used with `AgentRunner<BoxedProvider>`
318/// and `Orchestrator<BoxedProvider>`, eliminating the need for generic code
319/// at the call site.
320///
321/// # Example
322///
323/// ```ignore
324/// let provider = BoxedProvider::new(AnthropicProvider::new(key, model));
325/// let runner = AgentRunner::builder(Arc::new(provider))
326///     .name("agent")
327///     .build()?;
328/// ```
329pub struct BoxedProvider(Box<dyn DynLlmProvider>);
330
331impl BoxedProvider {
332    /// Create a type-erased provider from any concrete `LlmProvider`.
333    pub fn new<P: LlmProvider + 'static>(provider: P) -> Self {
334        Self(Box::new(provider))
335    }
336
337    /// Create a type-erased provider from an `Arc<P>`.
338    ///
339    /// Useful when the provider is already behind an `Arc` (e.g., shared between
340    /// the orchestrator and sub-agents) and needs to be converted to `BoxedProvider`
341    /// for type erasure without consuming the original.
342    pub fn from_arc<P: LlmProvider + 'static>(provider: Arc<P>) -> Self {
343        /// Internal adapter: delegates to the `Arc<P>` inner provider.
344        struct ArcAdapter<P>(Arc<P>);
345
346        impl<P: LlmProvider> LlmProvider for ArcAdapter<P> {
347            async fn complete(
348                &self,
349                request: CompletionRequest,
350            ) -> Result<CompletionResponse, Error> {
351                self.0.complete(request).await
352            }
353
354            async fn stream_complete(
355                &self,
356                request: CompletionRequest,
357                on_text: &OnText,
358            ) -> Result<CompletionResponse, Error> {
359                self.0.stream_complete(request, on_text).await
360            }
361
362            async fn stream_complete_with_reasoning(
363                &self,
364                request: CompletionRequest,
365                on_text: &OnText,
366                on_reasoning: &OnReasoning,
367            ) -> Result<CompletionResponse, Error> {
368                self.0
369                    .stream_complete_with_reasoning(request, on_text, on_reasoning)
370                    .await
371            }
372
373            fn model_name(&self) -> Option<&str> {
374                self.0.model_name()
375            }
376        }
377
378        Self(Box::new(ArcAdapter(provider)))
379    }
380}
381
382impl LlmProvider for BoxedProvider {
383    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, Error> {
384        self.0.complete(request).await
385    }
386
387    async fn stream_complete(
388        &self,
389        request: CompletionRequest,
390        on_text: &OnText,
391    ) -> Result<CompletionResponse, Error> {
392        self.0.stream_complete(request, on_text).await
393    }
394
395    async fn stream_complete_with_reasoning(
396        &self,
397        request: CompletionRequest,
398        on_text: &OnText,
399        on_reasoning: &OnReasoning,
400    ) -> Result<CompletionResponse, Error> {
401        self.0
402            .stream_complete_with_reasoning(request, on_text, on_reasoning)
403            .await
404    }
405
406    fn model_name(&self) -> Option<&str> {
407        self.0.model_name()
408    }
409}
410
411#[cfg(test)]
412mod tests {
413    use super::*;
414    use crate::llm::types::{ContentBlock, Message, StopReason, TokenUsage};
415    use std::sync::{Arc, Mutex};
416
417    struct FakeProvider;
418
419    impl LlmProvider for FakeProvider {
420        async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
421            Ok(CompletionResponse {
422                content: vec![ContentBlock::Text {
423                    text: "fake".into(),
424                }],
425                stop_reason: StopReason::EndTurn,
426                reasoning: None,
427                usage: TokenUsage::default(),
428                model: None,
429            })
430        }
431    }
432
433    struct StreamingFakeProvider;
434
435    impl LlmProvider for StreamingFakeProvider {
436        async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
437            panic!("should call stream_complete, not complete");
438        }
439
440        async fn stream_complete(
441            &self,
442            _request: CompletionRequest,
443            on_text: &OnText,
444        ) -> Result<CompletionResponse, Error> {
445            on_text("hello");
446            on_text(" world");
447            Ok(CompletionResponse {
448                content: vec![ContentBlock::Text {
449                    text: "hello world".into(),
450                }],
451                stop_reason: StopReason::EndTurn,
452                reasoning: None,
453                usage: TokenUsage::default(),
454                model: None,
455            })
456        }
457    }
458
459    fn test_request() -> CompletionRequest {
460        CompletionRequest {
461            system: String::new(),
462            messages: vec![Message::user("test")],
463            tools: vec![],
464            max_tokens: 100,
465            tool_choice: None,
466            reasoning_effort: None,
467        }
468    }
469
470    #[test]
471    fn dyn_llm_provider_wraps_provider() {
472        let provider = FakeProvider;
473        let dyn_provider: &dyn DynLlmProvider = &provider;
474        let _ = dyn_provider;
475    }
476
477    #[tokio::test]
478    async fn boxed_provider_delegates_complete() {
479        let provider = BoxedProvider::new(FakeProvider);
480        // Disambiguate: BoxedProvider implements both LlmProvider and DynLlmProvider
481        let response = LlmProvider::complete(&provider, test_request())
482            .await
483            .unwrap();
484        assert_eq!(response.text(), "fake");
485    }
486
487    #[tokio::test]
488    async fn boxed_provider_delegates_stream_complete() {
489        let provider = BoxedProvider::new(StreamingFakeProvider);
490        let received = Arc::new(Mutex::new(Vec::<String>::new()));
491        let received_clone = received.clone();
492        let on_text: &OnText = &move |text: &str| {
493            received_clone
494                .lock()
495                .expect("test lock")
496                .push(text.to_string());
497        };
498
499        let response = LlmProvider::stream_complete(&provider, test_request(), on_text)
500            .await
501            .unwrap();
502        assert_eq!(response.text(), "hello world");
503
504        let texts = received.lock().expect("test lock");
505        assert_eq!(*texts, vec!["hello", " world"]);
506    }
507
508    #[test]
509    fn boxed_provider_is_send_sync() {
510        fn assert_send_sync<T: Send + Sync>() {}
511        assert_send_sync::<BoxedProvider>();
512    }
513
514    #[tokio::test]
515    async fn boxed_provider_default_stream_falls_back_to_complete() {
516        // FakeProvider only implements complete; stream_complete should fall back
517        let provider = BoxedProvider::new(FakeProvider);
518        let on_text: &OnText = &|_| {};
519        let response = LlmProvider::stream_complete(&provider, test_request(), on_text)
520            .await
521            .unwrap();
522        assert_eq!(response.text(), "fake");
523    }
524
525    #[tokio::test]
526    async fn boxed_provider_from_arc_delegates_complete() {
527        let provider = Arc::new(FakeProvider);
528        let boxed = BoxedProvider::from_arc(provider);
529        let response = LlmProvider::complete(&boxed, test_request()).await.unwrap();
530        assert_eq!(response.text(), "fake");
531    }
532
533    #[tokio::test]
534    async fn boxed_provider_from_arc_delegates_stream_complete() {
535        let provider = Arc::new(StreamingFakeProvider);
536        let boxed = BoxedProvider::from_arc(provider);
537        let received = Arc::new(Mutex::new(Vec::<String>::new()));
538        let received_clone = received.clone();
539        let on_text: &OnText = &move |text: &str| {
540            received_clone
541                .lock()
542                .expect("test lock")
543                .push(text.to_string());
544        };
545        let response = LlmProvider::stream_complete(&boxed, test_request(), on_text)
546            .await
547            .unwrap();
548        assert_eq!(response.text(), "hello world");
549        let texts = received.lock().expect("test lock");
550        assert_eq!(*texts, vec!["hello", " world"]);
551    }
552
553    #[test]
554    fn model_name_default_is_none() {
555        let provider = FakeProvider;
556        assert!(LlmProvider::model_name(&provider).is_none());
557    }
558
559    #[test]
560    fn boxed_provider_preserves_model_name() {
561        struct NamedProvider;
562        impl LlmProvider for NamedProvider {
563            async fn complete(
564                &self,
565                _request: CompletionRequest,
566            ) -> Result<CompletionResponse, Error> {
567                unimplemented!()
568            }
569            fn model_name(&self) -> Option<&str> {
570                Some("test-model")
571            }
572        }
573        let boxed = BoxedProvider::new(NamedProvider);
574        assert_eq!(LlmProvider::model_name(&boxed), Some("test-model"));
575    }
576
577    #[test]
578    fn boxed_provider_from_arc_preserves_model_name() {
579        struct NamedProvider;
580        impl LlmProvider for NamedProvider {
581            async fn complete(
582                &self,
583                _request: CompletionRequest,
584            ) -> Result<CompletionResponse, Error> {
585                unimplemented!()
586            }
587            fn model_name(&self) -> Option<&str> {
588                Some("arc-model")
589            }
590        }
591        let boxed = BoxedProvider::from_arc(Arc::new(NamedProvider));
592        assert_eq!(LlmProvider::model_name(&boxed), Some("arc-model"));
593    }
594
595    #[tokio::test]
596    async fn boxed_provider_from_arc_shares_underlying_provider() {
597        // Verify from_arc shares the underlying provider via Arc (not a copy)
598        let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
599        struct CountingProvider(Arc<std::sync::atomic::AtomicUsize>);
600        impl LlmProvider for CountingProvider {
601            async fn complete(
602                &self,
603                _request: CompletionRequest,
604            ) -> Result<CompletionResponse, crate::error::Error> {
605                self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
606                Ok(CompletionResponse {
607                    content: vec![ContentBlock::Text {
608                        text: "counted".into(),
609                    }],
610                    stop_reason: StopReason::EndTurn,
611                    reasoning: None,
612                    usage: TokenUsage::default(),
613                    model: None,
614                })
615            }
616        }
617
618        let inner = Arc::new(CountingProvider(call_count.clone()));
619        let boxed1 = BoxedProvider::from_arc(inner.clone());
620        let boxed2 = BoxedProvider::from_arc(inner);
621
622        LlmProvider::complete(&boxed1, test_request())
623            .await
624            .unwrap();
625        LlmProvider::complete(&boxed2, test_request())
626            .await
627            .unwrap();
628
629        assert_eq!(
630            call_count.load(std::sync::atomic::Ordering::Relaxed),
631            2,
632            "both boxed providers should share the same underlying provider"
633        );
634    }
635
636    // --- ApprovalDecision ---
637
638    #[test]
639    fn approval_decision_from_true() {
640        let decision = ApprovalDecision::from(true);
641        assert_eq!(decision, ApprovalDecision::Allow);
642        assert!(decision.is_allowed());
643        assert!(!decision.is_persistent());
644    }
645
646    #[test]
647    fn approval_decision_from_false() {
648        let decision = ApprovalDecision::from(false);
649        assert_eq!(decision, ApprovalDecision::Deny);
650        assert!(!decision.is_allowed());
651        assert!(!decision.is_persistent());
652    }
653
654    #[test]
655    fn approval_decision_always_allow() {
656        let decision = ApprovalDecision::AlwaysAllow;
657        assert!(decision.is_allowed());
658        assert!(decision.is_persistent());
659    }
660
661    #[test]
662    fn approval_decision_always_deny() {
663        let decision = ApprovalDecision::AlwaysDeny;
664        assert!(!decision.is_allowed());
665        assert!(decision.is_persistent());
666    }
667}
heartbit_core/llm/mod.rs

heartbit_core/llm/
mod.rs
