1use std::collections::HashMap;
4use std::time::Duration;
5
6use ferrin_spec::ToolName;
7use serde::Deserialize;
8use serde::Serialize;
9
10#[derive(Debug, Clone, Default, PartialEq, Eq)]
15pub struct Timeout {
16 pub total: Option<Duration>,
18 pub step: Option<Duration>,
20 pub first_chunk: Option<Duration>,
22 pub chunk: Option<Duration>,
24 pub tool: Option<Duration>,
26 pub per_tool: HashMap<ToolName, Duration>,
28}
29
30impl Timeout {
31 #[must_use]
33 pub fn none() -> Self {
34 Self::default()
35 }
36
37 #[must_use]
39 pub fn with_total(mut self, total: Duration) -> Self {
40 self.total = Some(total);
41 self
42 }
43
44 #[must_use]
46 pub fn with_step(mut self, step: Duration) -> Self {
47 self.step = Some(step);
48 self
49 }
50
51 #[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 #[must_use]
60 pub fn with_chunk(mut self, chunk: Duration) -> Self {
61 self.chunk = Some(chunk);
62 self
63 }
64
65 #[must_use]
67 pub fn with_tool(mut self, tool: Duration) -> Self {
68 self.tool = Some(tool);
69 self
70 }
71
72 #[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 #[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 fn from(total: Duration) -> Self {
90 Self::default().with_total(total)
91 }
92}
93
94#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
96#[serde(tag = "scope", rename_all = "kebab-case")]
97#[non_exhaustive]
98pub enum TimeoutScope {
99 Total,
101 Step,
103 FirstChunk,
105 Chunk,
107 Tool {
109 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}