talos_core/
provider.rs

1//! Provider trait and error types for LLM backends.
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use tokio::sync::mpsc;
6
7use crate::message::{AgentEvent, Message};
8
9pub type Receiver<T> = mpsc::Receiver<T>;
10
11/// Non-secret, request-local progress reported by a language-model provider.
12///
13/// Retry ordinals are zero-based: `0` is the initial dispatch and positive values are the exact
14/// ordinals returned by the provider's retry policy. Progress is transient and must not be
15/// persisted as conversation history.
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
17#[serde(tag = "stage", rename_all = "snake_case")]
18#[non_exhaustive]
19pub enum ProviderProgress {
20    /// The initial request is being dispatched.
21    InitialDispatch {
22        /// Zero-based retry ordinal. This is always `0` for the initial dispatch.
23        attempt: u32,
24        /// Configured maximum retry ordinal.
25        max_attempts: u32,
26    },
27    /// A retry dispatch is being attempted after its scheduled backoff.
28    RetryDispatch {
29        /// Zero-based retry ordinal returned by the provider retry decision.
30        attempt: u32,
31        /// Configured maximum retry ordinal.
32        max_attempts: u32,
33    },
34    /// A bounded retry backoff has been scheduled.
35    ScheduledBackoff {
36        /// Retry ordinal that will be dispatched after the backoff.
37        attempt: u32,
38        /// Configured maximum retry ordinal.
39        max_attempts: u32,
40        /// Actual bounded delay selected by the provider retry policy.
41        delay_ms: u64,
42    },
43    /// Response headers arrived and the provider is waiting for the first stream packet.
44    FirstPacketWait {
45        /// Zero-based retry ordinal whose response is being streamed.
46        attempt: u32,
47        /// Configured maximum retry ordinal.
48        max_attempts: u32,
49    },
50}
51
52#[derive(Debug, thiserror::Error)]
53pub enum ProviderError {
54    #[error("authentication failed: {0}")]
55    AuthenticationFailed(String),
56
57    #[error("rate limited: {0}")]
58    RateLimited(String),
59
60    #[error("server error: {0}")]
61    ServerError(String),
62
63    #[error("network error: {0}")]
64    NetworkError(String),
65
66    #[error("invalid response: {0}")]
67    InvalidResponse(String),
68}
69
70pub type ProviderResult<T> = Result<T, ProviderError>;
71
72/// Provider-enforced limits for an isolated, tool-free decision request.
73#[derive(Debug, Clone, Copy, Serialize, Deserialize, schemars::JsonSchema)]
74pub struct DecisionRequestLimits {
75    /// Positive maximum generated tokens, including any provider reasoning.
76    pub max_output_tokens: u32,
77    /// Maximum additional transport dispatches; zero disables retries.
78    pub max_retries: u32,
79}
80
81#[derive(Debug, Clone, PartialEq)]
82pub struct ToolDefinition {
83    pub name: String,
84    pub description: String,
85    pub parameters: Value,
86}
87
88impl ToolDefinition {
89    /// Creates a new tool definition.
90    #[must_use]
91    pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
92        Self {
93            name: name.into(),
94            description: description.into(),
95            parameters,
96        }
97    }
98
99    /// Formats this tool definition as a text block suitable for inclusion
100    /// in the system prompt.
101    #[must_use]
102    pub fn to_prompt_text(&self) -> String {
103        format!(
104            "## {}\n{}\nParameters: {}",
105            self.name,
106            self.description,
107            serde_json::to_string_pretty(&self.parameters).unwrap_or_default()
108        )
109    }
110}
111
112#[async_trait::async_trait]
113pub trait LanguageModel: Send + Sync {
114    /// Dispatches an isolated text decision with no tools or inherited reasoning settings.
115    ///
116    /// Implementations must enforce the supplied token and retry limits. Unsupported
117    /// providers fail before dispatch; falling back to unrestricted `stream` is unsafe.
118    async fn stream_decision(
119        &self,
120        _messages: &[Message],
121        _limits: DecisionRequestLimits,
122    ) -> ProviderResult<Receiver<AgentEvent>> {
123        Err(ProviderError::InvalidResponse(
124            "provider does not support bounded decisions".into(),
125        ))
126    }
127    /// Dispatches an isolated Auto permission review with provider-supported low-latency reasoning.
128    ///
129    /// The default retains bounded-decision isolation. Adapters may override reasoning
130    /// only for models with a known request contract, without mutating conversation settings.
131    async fn stream_auto_review(
132        &self,
133        messages: &[Message],
134        limits: DecisionRequestLimits,
135    ) -> ProviderResult<Receiver<AgentEvent>> {
136        self.stream_decision(messages, limits).await
137    }
138
139    /// Returns a stable, non-secret scope for capability evidence caching.
140    /// `None` disables caching when the provider cannot describe its endpoint/model safely.
141    fn protocol_capability_scope(&self) -> Option<String> {
142        None
143    }
144
145    /// Reports non-secret evidence for automatic tool-protocol selection.
146    ///
147    /// Implementations must return [`crate::tool::CapabilityProbe::Unknown`] unless
148    /// the evidence is tied to the configured endpoint and model. Unknown evidence
149    /// is handled conservatively by selecting the validated compatibility path.
150    fn protocol_capabilities(&self) -> crate::tool::CapabilityProbe {
151        crate::tool::CapabilityProbe::Unknown
152    }
153
154    /// Performs a request-scoped native capability probe.
155    ///
156    /// Providers must return `Unknown` unless the response was validated as native
157    /// wire evidence for the configured endpoint and model. The conservative default
158    /// preserves compatibility for existing third-party implementations.
159    async fn probe_protocol_capabilities(
160        &self,
161        _tools: &[ToolDefinition],
162    ) -> crate::tool::CapabilityProbe {
163        crate::tool::CapabilityProbe::Unknown
164    }
165
166    async fn stream(&self, messages: &[Message]) -> ProviderResult<Receiver<AgentEvent>>;
167
168    async fn stream_with_tools(
169        &self,
170        messages: &[Message],
171        tools: &[ToolDefinition],
172    ) -> ProviderResult<Receiver<AgentEvent>> {
173        let _ = tools;
174        self.stream(messages).await
175    }
176
177    /// Streams a response while optionally reporting typed request-local provider progress.
178    ///
179    /// The default preserves source compatibility for third-party providers by delegating to
180    /// [`LanguageModel::stream_with_tools`] without emitting progress.
181    async fn stream_with_tools_and_progress(
182        &self,
183        messages: &[Message],
184        tools: &[ToolDefinition],
185        progress_tx: mpsc::UnboundedSender<ProviderProgress>,
186    ) -> ProviderResult<Receiver<AgentEvent>> {
187        drop(progress_tx);
188        self.stream_with_tools(messages, tools).await
189    }
190
191    /// Streams a response with an explicit tool protocol selected by the caller.
192    ///
193    /// Legacy providers retain their Native behavior. Compatibility modes require an
194    /// override implementing both request projection and validated response parsing;
195    /// the default rejects them before dispatch rather than silently sending native tools.
196    /// Auto must be resolved by the caller before dispatch.
197    async fn stream_with_protocol(
198        &self,
199        messages: &[Message],
200        tools: &[ToolDefinition],
201        protocol: crate::tool::ToolProtocol,
202        progress_tx: mpsc::UnboundedSender<ProviderProgress>,
203    ) -> ProviderResult<Receiver<AgentEvent>> {
204        if protocol != crate::tool::ToolProtocol::Native {
205            return Err(ProviderError::InvalidResponse(
206                "provider has no adapter for the selected tool protocol".into(),
207            ));
208        }
209        self.stream_with_tools_and_progress(messages, tools, progress_tx)
210            .await
211    }
212
213    fn request_preview(&self, _messages: &[Message]) -> Option<Value> {
214        None
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    struct LegacyModel;
223
224    struct CountingLegacyModel(std::sync::atomic::AtomicUsize);
225
226    #[async_trait::async_trait]
227    impl LanguageModel for CountingLegacyModel {
228        async fn stream(&self, _: &[Message]) -> ProviderResult<Receiver<AgentEvent>> {
229            self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
230            let (_, rx) = mpsc::channel(1);
231            Ok(rx)
232        }
233    }
234
235    #[tokio::test]
236    async fn legacy_protocol_adapter_rejects_unsupported_modes_before_dispatch() {
237        use crate::tool::ToolProtocol;
238        let model = CountingLegacyModel(std::sync::atomic::AtomicUsize::new(0));
239        for mode in [
240            ToolProtocol::Compat,
241            ToolProtocol::TalosStrict,
242            ToolProtocol::Auto,
243        ] {
244            let (tx, mut rx) = mpsc::unbounded_channel();
245            assert!(
246                model
247                    .stream_with_protocol(&[], &[], mode, tx)
248                    .await
249                    .is_err()
250            );
251            assert_eq!(rx.recv().await, None);
252            assert_eq!(model.0.load(std::sync::atomic::Ordering::SeqCst), 0);
253        }
254        let (tx, _) = mpsc::unbounded_channel();
255        assert!(
256            model
257                .stream_with_protocol(&[], &[], ToolProtocol::Native, tx)
258                .await
259                .is_ok()
260        );
261        assert_eq!(model.0.load(std::sync::atomic::Ordering::SeqCst), 1);
262    }
263
264    #[async_trait::async_trait]
265    impl LanguageModel for LegacyModel {
266        async fn stream(&self, _messages: &[Message]) -> ProviderResult<Receiver<AgentEvent>> {
267            let (_tx, rx) = mpsc::channel(1);
268            Ok(rx)
269        }
270    }
271
272    #[tokio::test]
273    async fn legacy_provider_uses_default_progress_aware_entrypoint() {
274        let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
275        let result = LegacyModel
276            .stream_with_tools_and_progress(&[], &[], progress_tx)
277            .await;
278
279        assert!(result.is_ok());
280        assert_eq!(progress_rx.recv().await, None);
281    }
282
283    #[test]
284    fn provider_progress_roundtrips_without_unbounded_diagnostics() {
285        let progress = ProviderProgress::ScheduledBackoff {
286            attempt: 2,
287            max_attempts: 3,
288            delay_ms: 750,
289        };
290        let encoded = serde_json::to_string(&progress).expect("serialize progress");
291        assert_eq!(
292            encoded,
293            r#"{"stage":"scheduled_backoff","attempt":2,"max_attempts":3,"delay_ms":750}"#
294        );
295        let decoded: ProviderProgress =
296            serde_json::from_str(&encoded).expect("deserialize progress");
297        assert_eq!(decoded, progress);
298    }
299}
talos_core/provider.rs

talos_core/
provider.rs
