agent_base/types/
checkpoint.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use super::message::ChatMessage;
5use super::session::SessionId;
6
7#[derive(Clone, Debug, Serialize, Deserialize)]
8pub struct CheckpointData {
9    pub session_id: SessionId,
10    pub user_input: String,
11    pub step: CheckpointStep,
12    pub turn_count: u32,
13}
14
15#[derive(Clone, Debug, Serialize, Deserialize)]
16pub enum CheckpointStep {
17    AfterUserInput,
18    BeforeLlm {
19        messages: Vec<ChatMessage>,
20        tools: Vec<Value>,
21    },
22    BeforeToolCalls {
23        tool_calls: Vec<(String, String, String)>,
24    },
25    AfterToolCalls {
26        tool_calls: Vec<(String, String, String)>,
27        results: Vec<ToolResultData>,
28    },
29}
30
31#[derive(Clone, Debug, Serialize, Deserialize)]
32pub struct ToolResultData {
33    pub tool_call_id: String,
34    pub tool_name: String,
35    pub summary: String,
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41
42    fn sample_messages() -> Vec<ChatMessage> {
43        vec![
44            ChatMessage::system("You are a helpful agent."),
45            ChatMessage::user("What is the weather in Tokyo?"),
46        ]
47    }
48
49    fn sample_tools() -> Vec<Value> {
50        vec![serde_json::json!({"name": "weather", "args": {"city": "Tokyo"}})]
51    }
52
53    fn sample_tool_calls() -> Vec<(String, String, String)> {
54        vec![
55            (
56                "call_1".to_string(),
57                "weather".to_string(),
58                "{\"city\": \"Tokyo\"}".to_string(),
59            ),
60            (
61                "call_2".to_string(),
62                "time".to_string(),
63                "{\"tz\": \"Asia/Tokyo\"}".to_string(),
64            ),
65        ]
66    }
67
68    fn sample_results() -> Vec<ToolResultData> {
69        vec![
70            ToolResultData {
71                tool_call_id: "call_1".to_string(),
72                tool_name: "weather".to_string(),
73                summary: "20C, clear".to_string(),
74            },
75            ToolResultData {
76                tool_call_id: "call_2".to_string(),
77                tool_name: "time".to_string(),
78                summary: "09:41 JST".to_string(),
79            },
80        ]
81    }
82
83    fn checkpoint_data(step: CheckpointStep) -> CheckpointData {
84        CheckpointData {
85            session_id: SessionId::with_external_id(42, "test-session"),
86            user_input: "What is the weather in Tokyo?".to_string(),
87            step,
88            turn_count: 3,
89        }
90    }
91
92    #[test]
93    fn checkpoint_step_after_user_input_round_trips() {
94        let step = CheckpointStep::AfterUserInput;
95        let json = serde_json::to_string(&step).unwrap();
96        let decoded: CheckpointStep = serde_json::from_str(&json).unwrap();
97        assert!(matches!(decoded, CheckpointStep::AfterUserInput));
98    }
99
100    #[test]
101    fn checkpoint_step_before_llm_round_trips() {
102        let step = CheckpointStep::BeforeLlm {
103            messages: sample_messages(),
104            tools: sample_tools(),
105        };
106        let json = serde_json::to_string(&step).unwrap();
107        let decoded: CheckpointStep = serde_json::from_str(&json).unwrap();
108        match decoded {
109            CheckpointStep::BeforeLlm { messages, tools } => {
110                // ChatMessage has no PartialEq, so compare canonical JSON.
111                assert_eq!(
112                    serde_json::to_string(&messages).unwrap(),
113                    serde_json::to_string(&sample_messages()).unwrap()
114                );
115                assert_eq!(tools, sample_tools());
116            }
117            other => panic!("expected BeforeLlm, got {other:?}"),
118        }
119    }
120
121    #[test]
122    fn checkpoint_step_before_tool_calls_round_trips() {
123        let step = CheckpointStep::BeforeToolCalls {
124            tool_calls: sample_tool_calls(),
125        };
126        let json = serde_json::to_string(&step).unwrap();
127        let decoded: CheckpointStep = serde_json::from_str(&json).unwrap();
128        match decoded {
129            CheckpointStep::BeforeToolCalls { tool_calls } => {
130                assert_eq!(tool_calls, sample_tool_calls());
131            }
132            other => panic!("expected BeforeToolCalls, got {other:?}"),
133        }
134    }
135
136    #[test]
137    fn checkpoint_step_after_tool_calls_round_trips() {
138        let step = CheckpointStep::AfterToolCalls {
139            tool_calls: sample_tool_calls(),
140            results: sample_results(),
141        };
142        let json = serde_json::to_string(&step).unwrap();
143        let decoded: CheckpointStep = serde_json::from_str(&json).unwrap();
144        match decoded {
145            CheckpointStep::AfterToolCalls {
146                tool_calls,
147                results,
148            } => {
149                assert_eq!(tool_calls, sample_tool_calls());
150                assert_eq!(results.len(), 2);
151                assert_eq!(results[0].tool_call_id, "call_1");
152                assert_eq!(results[0].tool_name, "weather");
153                assert_eq!(results[0].summary, "20C, clear");
154                assert_eq!(results[1].tool_call_id, "call_2");
155                assert_eq!(results[1].tool_name, "time");
156                assert_eq!(results[1].summary, "09:41 JST");
157            }
158            other => panic!("expected AfterToolCalls, got {other:?}"),
159        }
160    }
161
162    #[test]
163    fn checkpoint_data_round_trips_for_every_step_variant() {
164        let variants = [
165            CheckpointStep::AfterUserInput,
166            CheckpointStep::BeforeLlm {
167                messages: sample_messages(),
168                tools: sample_tools(),
169            },
170            CheckpointStep::BeforeToolCalls {
171                tool_calls: sample_tool_calls(),
172            },
173            CheckpointStep::AfterToolCalls {
174                tool_calls: sample_tool_calls(),
175                results: sample_results(),
176            },
177        ];
178
179        for step in variants {
180            let data = checkpoint_data(step);
181            let json = serde_json::to_string(&data).unwrap();
182            let decoded: CheckpointData = serde_json::from_str(&json).unwrap();
183            assert_eq!(decoded.session_id, data.session_id);
184            assert_eq!(decoded.user_input, data.user_input);
185            assert_eq!(decoded.turn_count, data.turn_count);
186            // CheckpointStep does not derive PartialEq, so compare variant
187            // discriminants for the step round-trip.
188            assert_eq!(
189                std::mem::discriminant(&data.step),
190                std::mem::discriminant(&decoded.step)
191            );
192        }
193    }
194
195    #[test]
196    fn checkpoint_step_deserialization_fails_gracefully_on_malformed_input() {
197        // Unknown variant tag.
198        assert!(serde_json::from_str::<CheckpointStep>(r#"{"NotAStep":{}}"#).is_err());
199        // Wrong field type on a known variant.
200        assert!(
201            serde_json::from_str::<CheckpointStep>(
202                r#"{"BeforeLlm":{"messages":"not-a-list","tools":[]}}"#,
203            )
204            .is_err()
205        );
206        // Truncated JSON.
207        assert!(serde_json::from_str::<CheckpointStep>(r#"{"AfterUserInput""#).is_err());
208        // Empty input.
209        assert!(serde_json::from_str::<CheckpointStep>("").is_err());
210    }
211
212    #[test]
213    fn checkpoint_data_deserialization_fails_gracefully_on_malformed_input() {
214        // Missing required fields.
215        assert!(
216            serde_json::from_str::<CheckpointData>(r#"{"session_id":{"id":1},"user_input":"hi"}"#,)
217                .is_err()
218        );
219        // Wrong session_id shape.
220        assert!(
221            serde_json::from_str::<CheckpointData>(
222                r#"{"session_id":"oops","user_input":"hi","step":"AfterUserInput","turn_count":1}"#,
223            )
224            .is_err()
225        );
226        // turn_count is not a number.
227        assert!(
228            serde_json::from_str::<CheckpointData>(
229                r#"{"session_id":{"id":1},"user_input":"hi","step":"AfterUserInput","turn_count":"three"}"#,
230            )
231            .is_err()
232        );
233    }
234}
agent_base/types/checkpoint.rs

agent_base/types/
checkpoint.rs
