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/// Why a model stopped generating, normalized across providers.
174///
175/// Providers report this as Anthropic `stop_reason`, OpenAI `finish_reason`,
176/// or Gemini `finishReason`. It used to be parsed and discarded, so a
177/// `max_tokens` truncation or a `content_filter`/safety block looked identical
178/// to a clean finish. The agent loop now inspects it: `Length` surfaces a
179/// truncation notice, and an empty `ContentFilter` finish becomes an error.
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
181#[serde(rename_all = "snake_case")]
182pub enum FinishReason {
183    /// Normal end of turn (`end_turn`, `stop`, `stop_sequence`, `STOP`).
184    Stop,
185    /// Stopped to call a tool (`tool_use`, `tool_calls`).
186    ToolUse,
187    /// Hit the output token limit — the response is truncated.
188    Length,
189    /// Blocked by a content filter / safety system / recitation check.
190    ContentFilter,
191    /// A reason we don't specifically model; carries the raw provider string.
192    Other(String),
193}
194
195/// Response from a model
196#[derive(Debug, Clone)]
197pub struct ModelResponse {
198    /// The actual response text
199    pub content: String,
200    /// Usage statistics if available
201    pub usage: Option<TokenUsage>,
202    /// Model that generated the response
203    pub model_name: String,
204    /// Thinking/reasoning content (for models that expose their thought process)
205    pub thinking: Option<String>,
206    /// Tool calls from the model (Ollama native function calling)
207    pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
208    /// Why generation stopped, when the provider reported it. `None` if the
209    /// provider didn't say (or the adapter doesn't yet map it).
210    pub stop_reason: Option<FinishReason>,
211    /// Anthropic thinking-block signature (encrypted server state). Set
212    /// only by `AnthropicAdapter`; other adapters leave it `None`. The
213    /// agent loop's commit step copies this onto the resulting assistant
214    /// `ChatMessage::thinking_signature` so it round-trips on the next
215    /// turn — without this, multi-turn Claude conversations with
216    /// extended thinking 400 with `invalid_request_error`.
217    pub thinking_signature: Option<String>,
218}
219
220/// Where a token count came from. Provider-reported counts are the
221/// billing/request truth; estimates are only for preflight context
222/// diagnostics.
223#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
224#[serde(rename_all = "snake_case")]
225pub enum TokenUsageSource {
226    #[default]
227    Provider,
228    Estimate,
229}
230
231/// Token usage statistics normalized across providers.
232///
233/// `prompt_tokens`, `completion_tokens`, and `total_tokens` preserve
234/// the old public surface. Extra fields keep cache/reasoning detail so
235/// UI can stop flattening unlike provider concepts into one number.
236#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
237pub struct TokenUsage {
238    pub prompt_tokens: usize,
239    pub completion_tokens: usize,
240    pub total_tokens: usize,
241    #[serde(default)]
242    pub cached_input_tokens: usize,
243    #[serde(default)]
244    pub cache_creation_input_tokens: usize,
245    #[serde(default)]
246    pub reasoning_output_tokens: usize,
247    #[serde(default)]
248    pub source: TokenUsageSource,
249}
250
251impl TokenUsage {
252    pub fn provider(prompt_tokens: usize, completion_tokens: usize, total_tokens: usize) -> Self {
253        Self {
254            prompt_tokens,
255            completion_tokens,
256            total_tokens,
257            cached_input_tokens: 0,
258            cache_creation_input_tokens: 0,
259            reasoning_output_tokens: 0,
260            source: TokenUsageSource::Provider,
261        }
262    }
263
264    pub fn estimate(prompt_tokens: usize) -> Self {
265        Self {
266            prompt_tokens,
267            completion_tokens: 0,
268            total_tokens: prompt_tokens,
269            cached_input_tokens: 0,
270            cache_creation_input_tokens: 0,
271            reasoning_output_tokens: 0,
272            source: TokenUsageSource::Estimate,
273        }
274    }
275
276    pub fn with_cached_input(mut self, cached_input_tokens: usize) -> Self {
277        self.cached_input_tokens = cached_input_tokens;
278        self
279    }
280
281    pub fn with_cache_creation(mut self, cache_creation_input_tokens: usize) -> Self {
282        self.cache_creation_input_tokens = cache_creation_input_tokens;
283        self
284    }
285
286    pub fn with_reasoning_output(mut self, reasoning_output_tokens: usize) -> Self {
287        self.reasoning_output_tokens = reasoning_output_tokens;
288        self
289    }
290
291    pub fn input_total_tokens(&self) -> usize {
292        self.prompt_tokens
293            .saturating_add(self.cached_input_tokens)
294            .saturating_add(self.cache_creation_input_tokens)
295    }
296
297    pub fn output_total_tokens(&self) -> usize {
298        self.completion_tokens
299            .saturating_add(self.reasoning_output_tokens)
300    }
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306
307    #[test]
308    fn test_message_role_equality() {
309        let user1 = MessageRole::User;
310        let user2 = MessageRole::User;
311        let assistant = MessageRole::Assistant;
312
313        assert_eq!(user1, user2, "User roles should be equal");
314        assert_ne!(user1, assistant, "Different roles should not be equal");
315    }
316
317    #[test]
318    fn test_chat_message_constructors() {
319        let user = ChatMessage::user("Hello!");
320        assert_eq!(user.role, MessageRole::User);
321        assert_eq!(user.content, "Hello!");
322        assert!(user.tool_calls.is_none());
323
324        let assistant = ChatMessage::assistant("Hi there");
325        assert_eq!(assistant.role, MessageRole::Assistant);
326
327        let system = ChatMessage::system("You are helpful");
328        assert_eq!(system.role, MessageRole::System);
329
330        let tool = ChatMessage::tool("call_1", "read_file", "file contents");
331        assert_eq!(tool.role, MessageRole::Tool);
332        assert_eq!(tool.tool_call_id, Some("call_1".to_string()));
333        assert_eq!(tool.tool_name, Some("read_file".to_string()));
334    }
335
336    #[test]
337    fn test_chat_message_builders() {
338        let msg = ChatMessage::user("test").with_images(vec!["base64data".to_string()]);
339        assert_eq!(msg.images, Some(vec!["base64data".to_string()]));
340    }
341
342    #[test]
343    fn test_token_usage_structure() {
344        let usage = TokenUsage::provider(100, 50, 150)
345            .with_cached_input(25)
346            .with_reasoning_output(10);
347
348        assert_eq!(usage.prompt_tokens, 100);
349        assert_eq!(usage.completion_tokens, 50);
350        assert_eq!(usage.total_tokens, 150);
351        assert_eq!(usage.cached_input_tokens, 25);
352        assert_eq!(usage.reasoning_output_tokens, 10);
353        assert_eq!(usage.source, TokenUsageSource::Provider);
354    }
355
356    // --- extract_thinking ---
357
358    #[test]
359    fn extract_thinking_no_marker_returns_text_unchanged() {
360        let (thinking, answer) = ChatMessage::extract_thinking("just a plain answer");
361        assert_eq!(thinking, None);
362        assert_eq!(answer, "just a plain answer");
363    }
364
365    #[test]
366    fn extract_thinking_complete_block() {
367        let raw = "Thinking...\n  reasoning here\n...done thinking.\n\nFinal answer";
368        let (thinking, answer) = ChatMessage::extract_thinking(raw);
369        assert_eq!(thinking.as_deref(), Some("reasoning here"));
370        assert_eq!(answer, "Final answer");
371    }
372
373    #[test]
374    fn thinking_signature_round_trips_through_serde() {
375        // Anthropic encrypted server state — must survive
376        // serialize/deserialize so saved conversations resume cleanly.
377        let msg = ChatMessage::assistant("Step 3 lives.")
378            .with_thinking_signature("sig_abc123_encrypted_blob");
379        let json = serde_json::to_string(&msg).expect("serialize");
380        let back: ChatMessage = serde_json::from_str(&json).expect("deserialize");
381        assert_eq!(
382            back.thinking_signature.as_deref(),
383            Some("sig_abc123_encrypted_blob")
384        );
385        assert_eq!(back.content, "Step 3 lives.");
386    }
387
388    #[test]
389    fn thinking_signature_defaults_to_none() {
390        // Backward compat: messages saved before Step 3 won't have the
391        // field. Serde default kicks in — None — and deserialize
392        // succeeds without errors.
393        let pre_step3_json = r#"{
394            "role": "Assistant",
395            "content": "hello",
396            "timestamp": "2026-04-16T12:00:00-04:00"
397        }"#;
398        let msg: ChatMessage = serde_json::from_str(pre_step3_json).expect("backward compat");
399        assert!(msg.thinking_signature.is_none());
400    }
401
402    #[test]
403    fn extract_thinking_in_progress_no_end_marker() {
404        let raw = "Thinking...\n  partial reasoning so far";
405        let (thinking, answer) = ChatMessage::extract_thinking(raw);
406        assert_eq!(thinking.as_deref(), Some("partial reasoning so far"));
407        assert_eq!(answer, "");
408    }
409
410    #[test]
411    fn test_model_response_creation() {
412        let usage = TokenUsage::provider(100, 50, 150);
413
414        let response = ModelResponse {
415            content: "Hello, world!".to_string(),
416            usage: Some(usage),
417            model_name: "ollama/tinyllama".to_string(),
418            thinking: None,
419            tool_calls: None,
420            stop_reason: None,
421            thinking_signature: None,
422        };
423
424        assert_eq!(response.content, "Hello, world!");
425        assert!(response.usage.is_some());
426        assert_eq!(response.model_name, "ollama/tinyllama");
427        assert_eq!(response.usage.unwrap().total_tokens, 150);
428        assert!(response.tool_calls.is_none());
429    }
430}
mermaid_cli/models/types.rs

mermaid_cli/models/
types.rs
