Skip to main content

ferrin_core/
timeout.rs

1//! Timeout configuration.
2
3use std::collections::HashMap;
4use std::time::Duration;
5
6use ferrin_spec::ToolName;
7use serde::Deserialize;
8use serde::Serialize;
9
10/// Timeouts applied to a call.
11///
12/// `first_chunk` and `chunk` only apply to streaming calls; `total` covers
13/// the whole call including tool execution.
14#[derive(Debug, Clone, Default, PartialEq, Eq)]
15pub struct Timeout {
16    /// Whole call.
17    pub total: Option<Duration>,
18    /// One step (model call plus its tool executions).
19    pub step: Option<Duration>,
20    /// Streaming: from the request until the first content part.
21    pub first_chunk: Option<Duration>,
22    /// Streaming: between two consecutive content parts.
23    pub chunk: Option<Duration>,
24    /// Default per tool execution.
25    pub tool: Option<Duration>,
26    /// Per-tool overrides of `tool`.
27    pub per_tool: HashMap<ToolName, Duration>,
28}
29
30impl Timeout {
31    /// No timeouts.
32    #[must_use]
33    pub fn none() -> Self {
34        Self::default()
35    }
36
37    /// Sets the total timeout.
38    #[must_use]
39    pub fn with_total(mut self, total: Duration) -> Self {
40        self.total = Some(total);
41        self
42    }
43
44    /// Sets the step timeout.
45    #[must_use]
46    pub fn with_step(mut self, step: Duration) -> Self {
47        self.step = Some(step);
48        self
49    }
50
51    /// Sets the first-chunk timeout (streaming only).
52    #[must_use]
53    pub fn with_first_chunk(mut self, first_chunk: Duration) -> Self {
54        self.first_chunk = Some(first_chunk);
55        self
56    }
57
58    /// Sets the inter-chunk timeout (streaming only).
59    #[must_use]
60    pub fn with_chunk(mut self, chunk: Duration) -> Self {
61        self.chunk = Some(chunk);
62        self
63    }
64
65    /// Sets the default tool timeout.
66    #[must_use]
67    pub fn with_tool(mut self, tool: Duration) -> Self {
68        self.tool = Some(tool);
69        self
70    }
71
72    /// Sets the timeout of one tool.
73    #[must_use]
74    pub fn with_tool_for(mut self, name: impl Into<ToolName>, timeout: Duration) -> Self {
75        self.per_tool.insert(name.into(), timeout);
76        self
77    }
78
79    /// Returns the timeout for `tool`: the per-tool override, else the
80    /// default tool timeout.
81    #[must_use]
82    pub fn tool_timeout(&self, tool: &ToolName) -> Option<Duration> {
83        self.per_tool.get(tool).copied().or(self.tool)
84    }
85}
86
87impl From<Duration> for Timeout {
88    /// A bare duration is the total timeout.
89    fn from(total: Duration) -> Self {
90        Self::default().with_total(total)
91    }
92}
93
94/// Which timeout fired.
95#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
96#[serde(tag = "scope", rename_all = "kebab-case")]
97#[non_exhaustive]
98pub enum TimeoutScope {
99    /// [`Timeout::total`].
100    Total,
101    /// [`Timeout::step`].
102    Step,
103    /// [`Timeout::first_chunk`].
104    FirstChunk,
105    /// [`Timeout::chunk`].
106    Chunk,
107    /// [`Timeout::tool`] or a per-tool override.
108    Tool {
109        /// The tool that timed out.
110        tool_name: ToolName,
111    },
112}
113
114impl std::fmt::Display for TimeoutScope {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        match self {
117            Self::Total => f.write_str("total"),
118            Self::Step => f.write_str("step"),
119            Self::FirstChunk => f.write_str("first chunk"),
120            Self::Chunk => f.write_str("chunk"),
121            Self::Tool { tool_name } => write!(f, "tool `{tool_name}`"),
122        }
123    }
124}