tirea_protocol_ai_sdk_v6/
history_encoder.rs

1use super::{StreamState, TextUIPart, ToolState, ToolUIPart, UIMessage, UIMessagePart, UIRole};
2use serde_json::Value;
3use tirea_contract::{Message, Role};
4
5pub struct AiSdkV6HistoryEncoder;
6
7impl AiSdkV6HistoryEncoder {
8    pub fn encode_message(msg: &Message) -> UIMessage {
9        let role = match msg.role {
10            Role::System => UIRole::System,
11            Role::User => UIRole::User,
12            Role::Assistant | Role::Tool => UIRole::Assistant,
13        };
14
15        let mut parts = Vec::new();
16
17        if msg.role == Role::Tool {
18            if let Some(tool_call_id) = &msg.tool_call_id {
19                let mut part = ToolUIPart::dynamic_tool(
20                    "tool",
21                    tool_call_id.clone(),
22                    ToolState::OutputAvailable,
23                );
24                part.output = Some(parse_tool_output(&msg.content));
25                parts.push(UIMessagePart::Tool(part));
26            } else if !msg.content.is_empty() {
27                parts.push(UIMessagePart::Text(TextUIPart::new(
28                    msg.content.clone(),
29                    Some(StreamState::Done),
30                )));
31            }
32        } else {
33            if !msg.content.is_empty() {
34                parts.push(UIMessagePart::Text(TextUIPart::new(
35                    msg.content.clone(),
36                    Some(StreamState::Done),
37                )));
38            }
39
40            if let Some(ref calls) = msg.tool_calls {
41                for tc in calls {
42                    let mut part = ToolUIPart::static_tool(
43                        tc.name.clone(),
44                        tc.id.clone(),
45                        ToolState::InputAvailable,
46                    );
47                    part.input = Some(tc.arguments.clone());
48                    parts.push(UIMessagePart::Tool(part));
49                }
50            }
51        }
52
53        let metadata = msg.metadata.as_ref().and_then(|metadata| {
54            let mut object = serde_json::Map::new();
55            if let Some(step_index) = metadata.step_index {
56                object.insert("step_index".to_string(), Value::from(step_index));
57            }
58            if object.is_empty() {
59                None
60            } else {
61                Some(Value::Object(object))
62            }
63        });
64
65        UIMessage {
66            id: msg.id.clone().unwrap_or_default(),
67            role,
68            metadata,
69            parts,
70        }
71    }
72}
73
74fn parse_tool_output(content: &str) -> Value {
75    if content.is_empty() {
76        Value::Null
77    } else {
78        serde_json::from_str(content).unwrap_or_else(|_| Value::String(content.to_string()))
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use tirea_contract::{Message, MessageMetadata, ToolCall, Visibility};
86
87    #[test]
88    fn test_ai_sdk_history_encoder_user_message() {
89        let msg = Message {
90            id: Some("msg_1".to_string()),
91            role: Role::User,
92            content: "hello".to_string(),
93            tool_calls: None,
94            tool_call_id: None,
95            visibility: Visibility::default(),
96            metadata: None,
97        };
98        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
99        assert_eq!(encoded.id, "msg_1");
100        assert_eq!(encoded.role, UIRole::User);
101        assert_eq!(encoded.parts.len(), 1);
102        assert!(matches!(
103            &encoded.parts[0],
104            UIMessagePart::Text(TextUIPart {
105                text,
106                state: Some(StreamState::Done),
107                ..
108            }) if text == "hello"
109        ));
110    }
111
112    #[test]
113    fn test_ai_sdk_history_encoder_assistant_with_tool_calls() {
114        let msg = Message {
115            id: Some("msg_2".to_string()),
116            role: Role::Assistant,
117            content: "Let me search.".to_string(),
118            tool_calls: Some(vec![
119                ToolCall {
120                    id: "call_1".to_string(),
121                    name: "search".to_string(),
122                    arguments: serde_json::json!({"query": "rust"}),
123                },
124                ToolCall {
125                    id: "call_2".to_string(),
126                    name: "fetch".to_string(),
127                    arguments: serde_json::json!({"url": "https://example.com"}),
128                },
129            ]),
130            tool_call_id: None,
131            visibility: Visibility::default(),
132            metadata: None,
133        };
134        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
135        assert_eq!(encoded.role, UIRole::Assistant);
136        assert_eq!(encoded.parts.len(), 3);
137        assert!(matches!(
138            &encoded.parts[0],
139            UIMessagePart::Text(TextUIPart { text, .. }) if text == "Let me search."
140        ));
141        assert!(matches!(
142            &encoded.parts[1],
143            UIMessagePart::Tool(ToolUIPart { part_type, tool_call_id, state: ToolState::InputAvailable, .. })
144            if tool_call_id == "call_1" && part_type == "tool-search"
145        ));
146        assert!(matches!(
147            &encoded.parts[2],
148            UIMessagePart::Tool(ToolUIPart { part_type, tool_call_id, .. })
149            if tool_call_id == "call_2" && part_type == "tool-fetch"
150        ));
151    }
152
153    #[test]
154    fn test_ai_sdk_history_encoder_tool_role_maps_to_assistant_with_tool_output_part() {
155        let msg = Message {
156            id: Some("msg_3".to_string()),
157            role: Role::Tool,
158            content: "{\"result\":42}".to_string(),
159            tool_calls: None,
160            tool_call_id: Some("call_1".to_string()),
161            visibility: Visibility::default(),
162            metadata: None,
163        };
164        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
165        assert_eq!(encoded.role, UIRole::Assistant);
166        assert_eq!(encoded.parts.len(), 1);
167        assert!(matches!(
168            &encoded.parts[0],
169            UIMessagePart::Tool(ToolUIPart { tool_call_id, state: ToolState::OutputAvailable, output: Some(output), .. })
170            if tool_call_id == "call_1" && output["result"] == 42
171        ));
172    }
173
174    #[test]
175    fn test_ai_sdk_history_encoder_tool_role_without_call_id_falls_back_to_text() {
176        let msg = Message {
177            id: Some("msg_3b".to_string()),
178            role: Role::Tool,
179            content: "plain output".to_string(),
180            tool_calls: None,
181            tool_call_id: None,
182            visibility: Visibility::default(),
183            metadata: None,
184        };
185        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
186        assert_eq!(encoded.role, UIRole::Assistant);
187        assert_eq!(encoded.parts.len(), 1);
188        assert!(matches!(
189            &encoded.parts[0],
190            UIMessagePart::Text(TextUIPart { text, .. }) if text == "plain output"
191        ));
192    }
193
194    #[test]
195    fn test_ai_sdk_history_encoder_empty_content_no_text_part() {
196        let msg = Message {
197            id: Some("msg_4".to_string()),
198            role: Role::Assistant,
199            content: String::new(),
200            tool_calls: Some(vec![ToolCall {
201                id: "call_1".to_string(),
202                name: "search".to_string(),
203                arguments: serde_json::json!({}),
204            }]),
205            tool_call_id: None,
206            visibility: Visibility::default(),
207            metadata: None,
208        };
209        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
210        assert_eq!(encoded.parts.len(), 1);
211        assert!(matches!(&encoded.parts[0], UIMessagePart::Tool(_)));
212    }
213
214    #[test]
215    fn test_ai_sdk_history_encoder_no_id_defaults_empty() {
216        let msg = Message {
217            id: None,
218            role: Role::User,
219            content: "hello".to_string(),
220            tool_calls: None,
221            tool_call_id: None,
222            visibility: Visibility::default(),
223            metadata: None,
224        };
225        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
226        assert_eq!(encoded.id, "");
227    }
228
229    #[test]
230    fn test_ai_sdk_history_encoder_with_metadata() {
231        let msg = Message {
232            id: Some("msg_5".to_string()),
233            role: Role::Assistant,
234            content: "response".to_string(),
235            tool_calls: None,
236            tool_call_id: None,
237            visibility: Visibility::default(),
238            metadata: Some(MessageMetadata {
239                run_id: Some("run_1".to_string()),
240                step_index: Some(2),
241            }),
242        };
243        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
244        assert!(encoded.metadata.is_some());
245        let meta = encoded.metadata.unwrap();
246        assert!(meta.get("run_id").is_none(), "run_id must not be exposed");
247        assert_eq!(meta["step_index"], 2);
248    }
249
250    #[test]
251    fn test_ai_sdk_history_encoder_system_message() {
252        let msg = Message {
253            id: Some("msg_sys".to_string()),
254            role: Role::System,
255            content: "You are helpful.".to_string(),
256            tool_calls: None,
257            tool_call_id: None,
258            visibility: Visibility::default(),
259            metadata: None,
260        };
261        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
262        assert_eq!(encoded.role, UIRole::System);
263    }
264
265    #[test]
266    fn test_ai_sdk_encode_messages_batch() {
267        let msgs = [Message::user("hello"), Message::assistant("world")];
268        let encoded: Vec<_> = msgs
269            .iter()
270            .map(AiSdkV6HistoryEncoder::encode_message)
271            .collect();
272        assert_eq!(encoded.len(), 2);
273        assert_eq!(encoded[0].role, UIRole::User);
274        assert_eq!(encoded[1].role, UIRole::Assistant);
275    }
276
277    #[test]
278    fn test_ai_sdk_history_encoder_serialization() {
279        let msg = Message {
280            id: Some("msg_1".to_string()),
281            role: Role::Assistant,
282            content: "Hello".to_string(),
283            tool_calls: Some(vec![ToolCall {
284                id: "call_1".to_string(),
285                name: "search".to_string(),
286                arguments: serde_json::json!({"q": "test"}),
287            }]),
288            tool_call_id: None,
289            visibility: Visibility::default(),
290            metadata: None,
291        };
292        let encoded = AiSdkV6HistoryEncoder::encode_message(&msg);
293        let json = serde_json::to_string(&encoded).unwrap();
294        assert!(json.contains("toolCallId"));
295        assert!(json.contains("\"type\":\"tool-search\""));
296        assert!(!json.contains("tool_call_id"));
297        assert!(!json.contains("tool_name"));
298    }
299}
tirea_protocol_ai_sdk_v6/history_encoder.rs

tirea_protocol_ai_sdk_v6/
history_encoder.rs
