mermaid_cli/models/
types.rs

1use crate::domain::ActionDisplay;
2use serde::{Deserialize, Serialize};
3
4/// Represents a chat message
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct ChatMessage {
7    pub role: MessageRole,
8    pub content: String,
9    pub timestamp: chrono::DateTime<chrono::Local>,
10    /// Mermaid-owned message classification. Provider adapters ignore
11    /// this; render/persistence use it to distinguish generated
12    /// checkpoints from normal user/assistant turns.
13    #[serde(default)]
14    pub kind: ChatMessageKind,
15    /// Optional Mermaid-owned structured metadata for UI/replay.
16    #[serde(default)]
17    pub metadata: Option<serde_json::Value>,
18    /// Actions performed during this message (for display purposes)
19    #[serde(default)]
20    pub actions: Vec<ActionDisplay>,
21    /// Thinking/reasoning content (for models that expose their thought process)
22    #[serde(default)]
23    pub thinking: Option<String>,
24    /// Base64-encoded images/PDFs for multimodal models
25    #[serde(default)]
26    pub images: Option<Vec<String>>,
27    /// Tool calls from the model (Ollama native function calling)
28    #[serde(default)]
29    pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
30    /// Tool call ID for tool result messages (OpenAI-compatible format)
31    /// This links the tool result back to the original tool_call from the assistant
32    #[serde(default)]
33    pub tool_call_id: Option<String>,
34    /// Tool name for tool result messages (required by Ollama API)
35    /// This tells the model which function's result is being returned
36    #[serde(default)]
37    pub tool_name: Option<String>,
38    /// Anthropic thinking-block signature — encrypted server state that
39    /// MUST round-trip back into the next request when extended thinking
40    /// is enabled, or the API returns 400 `invalid_request_error`. Set
41    /// only by the Anthropic adapter; other adapters leave it `None`
42    /// and other providers ignore it on the wire.
43    #[serde(default)]
44    pub thinking_signature: Option<String>,
45}
46
47impl ChatMessage {
48    /// Create a user message
49    pub fn user(content: impl Into<String>) -> Self {
50        Self::new(MessageRole::User, content.into())
51    }
52
53    /// Create an assistant message
54    pub fn assistant(content: impl Into<String>) -> Self {
55        Self::new(MessageRole::Assistant, content.into())
56    }
57
58    /// Create a system message
59    pub fn system(content: impl Into<String>) -> Self {
60        Self::new(MessageRole::System, content.into())
61    }
62
63    /// Create a tool result message
64    pub fn tool(
65        tool_call_id: impl Into<String>,
66        tool_name: impl Into<String>,
67        content: impl Into<String>,
68    ) -> Self {
69        Self {
70            role: MessageRole::Tool,
71            content: content.into(),
72            timestamp: chrono::Local::now(),
73            kind: ChatMessageKind::Normal,
74            metadata: None,
75            actions: Vec::new(),
76            thinking: None,
77            images: None,
78            tool_calls: None,
79            tool_call_id: Some(tool_call_id.into()),
80            tool_name: Some(tool_name.into()),
81            thinking_signature: None,
82        }
83    }
84
85    /// Base constructor with role and content
86    fn new(role: MessageRole, content: String) -> Self {
87        Self {
88            role,
89            content,
90            timestamp: chrono::Local::now(),
91            kind: ChatMessageKind::Normal,
92            metadata: None,
93            actions: Vec::new(),
94            thinking: None,
95            images: None,
96            tool_calls: None,
97            tool_call_id: None,
98            tool_name: None,
99            thinking_signature: None,
100        }
101    }
102
103    /// Builder: attach images
104    pub fn with_images(mut self, images: Vec<String>) -> Self {
105        self.images = Some(images);
106        self
107    }
108
109    /// Builder: attach tool calls
110    pub fn with_tool_calls(mut self, tool_calls: Vec<crate::models::tool_call::ToolCall>) -> Self {
111        self.tool_calls = if tool_calls.is_empty() {
112            None
113        } else {
114            Some(tool_calls)
115        };
116        self
117    }
118
119    /// Builder: attach an Anthropic thinking signature. Used by the
120    /// Anthropic adapter when committing assistant messages so the
121    /// signature can round-trip on the next request.
122    pub fn with_thinking_signature(mut self, signature: impl Into<String>) -> Self {
123        self.thinking_signature = Some(signature.into());
124        self
125    }
126
127    /// Extract thinking blocks from message content.
128    /// Returns `(thinking_content, answer_content)`.
129    ///
130    /// Performs a single `find` for the start marker; the previous version
131    /// scanned twice (`contains` + `find`) and called `find("Thinking...")`
132    /// again inside the if-let-chain.
133    ///
134    /// Safety: `str::find()` returns byte offsets. The markers `"Thinking..."`
135    /// and `"...done thinking."` are pure ASCII, so adding their `.len()`
136    /// always lands on a valid UTF-8 char boundary.
137    pub fn extract_thinking(text: &str) -> (Option<String>, String) {
138        let Some(thinking_start) = text.find("Thinking...") else {
139            return (None, text.to_string());
140        };
141        let content_start = thinking_start + "Thinking...".len();
142
143        if let Some(thinking_end) = text.find("...done thinking.") {
144            let thinking_text = text[content_start..thinking_end].trim().to_string();
145            let answer_start = thinking_end + "...done thinking.".len();
146            let answer_text = text[answer_start..].trim().to_string();
147            return (Some(thinking_text), answer_text);
148        }
149
150        // Start marker without end marker — thinking is still in progress.
151        let thinking_text = text[content_start..].trim().to_string();
152        (Some(thinking_text), String::new())
153    }
154}
155
156#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
157pub enum MessageRole {
158    User,
159    Assistant,
160    System,
161    /// Tool result message (OpenAI-compatible format for function calling)
162    Tool,
163}
164
165#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(rename_all = "snake_case")]
167pub enum ChatMessageKind {
168    #[default]
169    Normal,
170    ContextCheckpoint,
171}
172
173/// Response from a model
174#[derive(Debug, Clone)]
175pub struct ModelResponse {
176    /// The actual response text
177    pub content: String,
178    /// Usage statistics if available
179    pub usage: Option<TokenUsage>,
180    /// Model that generated the response
181    pub model_name: String,
182    /// Thinking/reasoning content (for models that expose their thought process)
183    pub thinking: Option<String>,
184    /// Tool calls from the model (Ollama native function calling)
185    pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
186    /// Anthropic thinking-block signature (encrypted server state). Set
187    /// only by `AnthropicAdapter`; other adapters leave it `None`. The
188    /// agent loop's commit step copies this onto the resulting assistant
189    /// `ChatMessage::thinking_signature` so it round-trips on the next
190    /// turn — without this, multi-turn Claude conversations with
191    /// extended thinking 400 with `invalid_request_error`.
192    pub thinking_signature: Option<String>,
193}
194
195/// Where a token count came from. Provider-reported counts are the
196/// billing/request truth; estimates are only for preflight context
197/// diagnostics.
198#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
199#[serde(rename_all = "snake_case")]
200pub enum TokenUsageSource {
201    #[default]
202    Provider,
203    Estimate,
204}
205
206/// Token usage statistics normalized across providers.
207///
208/// `prompt_tokens`, `completion_tokens`, and `total_tokens` preserve
209/// the old public surface. Extra fields keep cache/reasoning detail so
210/// UI can stop flattening unlike provider concepts into one number.
211#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
212pub struct TokenUsage {
213    pub prompt_tokens: usize,
214    pub completion_tokens: usize,
215    pub total_tokens: usize,
216    #[serde(default)]
217    pub cached_input_tokens: usize,
218    #[serde(default)]
219    pub cache_creation_input_tokens: usize,
220    #[serde(default)]
221    pub reasoning_output_tokens: usize,
222    #[serde(default)]
223    pub source: TokenUsageSource,
224}
225
226impl TokenUsage {
227    pub fn provider(prompt_tokens: usize, completion_tokens: usize, total_tokens: usize) -> Self {
228        Self {
229            prompt_tokens,
230            completion_tokens,
231            total_tokens,
232            cached_input_tokens: 0,
233            cache_creation_input_tokens: 0,
234            reasoning_output_tokens: 0,
235            source: TokenUsageSource::Provider,
236        }
237    }
238
239    pub fn estimate(prompt_tokens: usize) -> Self {
240        Self {
241            prompt_tokens,
242            completion_tokens: 0,
243            total_tokens: prompt_tokens,
244            cached_input_tokens: 0,
245            cache_creation_input_tokens: 0,
246            reasoning_output_tokens: 0,
247            source: TokenUsageSource::Estimate,
248        }
249    }
250
251    pub fn with_cached_input(mut self, cached_input_tokens: usize) -> Self {
252        self.cached_input_tokens = cached_input_tokens;
253        self
254    }
255
256    pub fn with_cache_creation(mut self, cache_creation_input_tokens: usize) -> Self {
257        self.cache_creation_input_tokens = cache_creation_input_tokens;
258        self
259    }
260
261    pub fn with_reasoning_output(mut self, reasoning_output_tokens: usize) -> Self {
262        self.reasoning_output_tokens = reasoning_output_tokens;
263        self
264    }
265
266    pub fn input_total_tokens(&self) -> usize {
267        self.prompt_tokens
268            .saturating_add(self.cached_input_tokens)
269            .saturating_add(self.cache_creation_input_tokens)
270    }
271
272    pub fn output_total_tokens(&self) -> usize {
273        self.completion_tokens
274            .saturating_add(self.reasoning_output_tokens)
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn test_message_role_equality() {
284        let user1 = MessageRole::User;
285        let user2 = MessageRole::User;
286        let assistant = MessageRole::Assistant;
287
288        assert_eq!(user1, user2, "User roles should be equal");
289        assert_ne!(user1, assistant, "Different roles should not be equal");
290    }
291
292    #[test]
293    fn test_chat_message_constructors() {
294        let user = ChatMessage::user("Hello!");
295        assert_eq!(user.role, MessageRole::User);
296        assert_eq!(user.content, "Hello!");
297        assert!(user.tool_calls.is_none());
298
299        let assistant = ChatMessage::assistant("Hi there");
300        assert_eq!(assistant.role, MessageRole::Assistant);
301
302        let system = ChatMessage::system("You are helpful");
303        assert_eq!(system.role, MessageRole::System);
304
305        let tool = ChatMessage::tool("call_1", "read_file", "file contents");
306        assert_eq!(tool.role, MessageRole::Tool);
307        assert_eq!(tool.tool_call_id, Some("call_1".to_string()));
308        assert_eq!(tool.tool_name, Some("read_file".to_string()));
309    }
310
311    #[test]
312    fn test_chat_message_builders() {
313        let msg = ChatMessage::user("test").with_images(vec!["base64data".to_string()]);
314        assert_eq!(msg.images, Some(vec!["base64data".to_string()]));
315    }
316
317    #[test]
318    fn test_token_usage_structure() {
319        let usage = TokenUsage::provider(100, 50, 150)
320            .with_cached_input(25)
321            .with_reasoning_output(10);
322
323        assert_eq!(usage.prompt_tokens, 100);
324        assert_eq!(usage.completion_tokens, 50);
325        assert_eq!(usage.total_tokens, 150);
326        assert_eq!(usage.cached_input_tokens, 25);
327        assert_eq!(usage.reasoning_output_tokens, 10);
328        assert_eq!(usage.source, TokenUsageSource::Provider);
329    }
330
331    // --- extract_thinking ---
332
333    #[test]
334    fn extract_thinking_no_marker_returns_text_unchanged() {
335        let (thinking, answer) = ChatMessage::extract_thinking("just a plain answer");
336        assert_eq!(thinking, None);
337        assert_eq!(answer, "just a plain answer");
338    }
339
340    #[test]
341    fn extract_thinking_complete_block() {
342        let raw = "Thinking...\n  reasoning here\n...done thinking.\n\nFinal answer";
343        let (thinking, answer) = ChatMessage::extract_thinking(raw);
344        assert_eq!(thinking.as_deref(), Some("reasoning here"));
345        assert_eq!(answer, "Final answer");
346    }
347
348    #[test]
349    fn thinking_signature_round_trips_through_serde() {
350        // Anthropic encrypted server state — must survive
351        // serialize/deserialize so saved conversations resume cleanly.
352        let msg = ChatMessage::assistant("Step 3 lives.")
353            .with_thinking_signature("sig_abc123_encrypted_blob");
354        let json = serde_json::to_string(&msg).expect("serialize");
355        let back: ChatMessage = serde_json::from_str(&json).expect("deserialize");
356        assert_eq!(
357            back.thinking_signature.as_deref(),
358            Some("sig_abc123_encrypted_blob")
359        );
360        assert_eq!(back.content, "Step 3 lives.");
361    }
362
363    #[test]
364    fn thinking_signature_defaults_to_none() {
365        // Backward compat: messages saved before Step 3 won't have the
366        // field. Serde default kicks in — None — and deserialize
367        // succeeds without errors.
368        let pre_step3_json = r#"{
369            "role": "Assistant",
370            "content": "hello",
371            "timestamp": "2026-04-16T12:00:00-04:00"
372        }"#;
373        let msg: ChatMessage = serde_json::from_str(pre_step3_json).expect("backward compat");
374        assert!(msg.thinking_signature.is_none());
375    }
376
377    #[test]
378    fn extract_thinking_in_progress_no_end_marker() {
379        let raw = "Thinking...\n  partial reasoning so far";
380        let (thinking, answer) = ChatMessage::extract_thinking(raw);
381        assert_eq!(thinking.as_deref(), Some("partial reasoning so far"));
382        assert_eq!(answer, "");
383    }
384
385    #[test]
386    fn test_model_response_creation() {
387        let usage = TokenUsage::provider(100, 50, 150);
388
389        let response = ModelResponse {
390            content: "Hello, world!".to_string(),
391            usage: Some(usage),
392            model_name: "ollama/tinyllama".to_string(),
393            thinking: None,
394            tool_calls: None,
395            thinking_signature: None,
396        };
397
398        assert_eq!(response.content, "Hello, world!");
399        assert!(response.usage.is_some());
400        assert_eq!(response.model_name, "ollama/tinyllama");
401        assert_eq!(response.usage.unwrap().total_tokens, 150);
402        assert!(response.tool_calls.is_none());
403    }
404}
mermaid_cli/models/types.rs

mermaid_cli/models/
types.rs
