hehe_core/tool/
call.rs

1use crate::types::{Metadata, Timestamp, ToolCallId};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
6#[serde(rename_all = "snake_case")]
7pub enum ToolCallStatus {
8    Pending,
9    Running,
10    Completed,
11    Failed,
12    Cancelled,
13}
14
15impl ToolCallStatus {
16    pub fn is_terminal(&self) -> bool {
17        matches!(
18            self,
19            ToolCallStatus::Completed | ToolCallStatus::Failed | ToolCallStatus::Cancelled
20        )
21    }
22
23    pub fn is_success(&self) -> bool {
24        matches!(self, ToolCallStatus::Completed)
25    }
26}
27
28#[derive(Clone, Debug, Serialize, Deserialize)]
29pub struct ToolCall {
30    pub id: ToolCallId,
31    pub name: String,
32    pub input: Value,
33    pub status: ToolCallStatus,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub output: Option<Value>,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub error: Option<String>,
38    pub created_at: Timestamp,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub started_at: Option<Timestamp>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub completed_at: Option<Timestamp>,
43    #[serde(default, skip_serializing_if = "Metadata::is_empty")]
44    pub metadata: Metadata,
45}
46
47impl ToolCall {
48    pub fn new(name: impl Into<String>, input: Value) -> Self {
49        Self {
50            id: ToolCallId::new(),
51            name: name.into(),
52            input,
53            status: ToolCallStatus::Pending,
54            output: None,
55            error: None,
56            created_at: Timestamp::now(),
57            started_at: None,
58            completed_at: None,
59            metadata: Metadata::new(),
60        }
61    }
62
63    pub fn with_id(mut self, id: ToolCallId) -> Self {
64        self.id = id;
65        self
66    }
67
68    pub fn start(&mut self) {
69        self.status = ToolCallStatus::Running;
70        self.started_at = Some(Timestamp::now());
71    }
72
73    pub fn complete(&mut self, output: Value) {
74        self.status = ToolCallStatus::Completed;
75        self.output = Some(output);
76        self.completed_at = Some(Timestamp::now());
77    }
78
79    pub fn complete_with_text(&mut self, text: impl Into<String>) {
80        self.complete(Value::String(text.into()));
81    }
82
83    pub fn fail(&mut self, error: impl Into<String>) {
84        self.status = ToolCallStatus::Failed;
85        self.error = Some(error.into());
86        self.completed_at = Some(Timestamp::now());
87    }
88
89    pub fn cancel(&mut self) {
90        self.status = ToolCallStatus::Cancelled;
91        self.completed_at = Some(Timestamp::now());
92    }
93
94    pub fn is_pending(&self) -> bool {
95        self.status == ToolCallStatus::Pending
96    }
97
98    pub fn is_running(&self) -> bool {
99        self.status == ToolCallStatus::Running
100    }
101
102    pub fn is_completed(&self) -> bool {
103        self.status == ToolCallStatus::Completed
104    }
105
106    pub fn is_failed(&self) -> bool {
107        self.status == ToolCallStatus::Failed
108    }
109
110    pub fn is_terminal(&self) -> bool {
111        self.status.is_terminal()
112    }
113
114    pub fn duration_ms(&self) -> Option<i64> {
115        match (self.started_at, self.completed_at) {
116            (Some(start), Some(end)) => Some(end.unix_millis() - start.unix_millis()),
117            _ => None,
118        }
119    }
120
121    pub fn output_as_string(&self) -> Option<String> {
122        self.output.as_ref().and_then(|v| match v {
123            Value::String(s) => Some(s.clone()),
124            _ => serde_json::to_string(v).ok(),
125        })
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn test_tool_call_lifecycle() {
135        let mut call = ToolCall::new("test_tool", serde_json::json!({"arg": "value"}));
136
137        assert!(call.is_pending());
138        assert!(!call.is_terminal());
139
140        call.start();
141        assert!(call.is_running());
142        assert!(call.started_at.is_some());
143
144        call.complete(serde_json::json!({"result": "success"}));
145        assert!(call.is_completed());
146        assert!(call.is_terminal());
147        assert!(call.completed_at.is_some());
148        assert!(call.duration_ms().is_some());
149    }
150
151    #[test]
152    fn test_tool_call_failure() {
153        let mut call = ToolCall::new("failing_tool", serde_json::json!({}));
154
155        call.start();
156        call.fail("Something went wrong");
157
158        assert!(call.is_failed());
159        assert!(call.is_terminal());
160        assert_eq!(call.error, Some("Something went wrong".to_string()));
161    }
162}
hehe_core/tool/call.rs

hehe_core/tool/
call.rs
