Skip to main content

claude_codes/io/
claude_output.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use super::content_blocks::{ContentBlock, ToolUseBlock};
5use super::control::{ControlRequest, ControlResponse};
6use super::errors::{AnthropicError, ParseError};
7use super::message_types::{AssistantMessage, SystemMessage, UserMessage};
8use super::rate_limit::RateLimitEvent;
9use super::result::ResultMessage;
10
11/// Top-level enum for all possible Claude output messages
12#[derive(Debug, Clone, Serialize, Deserialize)]
13#[serde(tag = "type", rename_all = "snake_case")]
14pub enum ClaudeOutput {
15    /// System initialization message
16    System(SystemMessage),
17
18    /// User message echoed back
19    User(UserMessage),
20
21    /// Assistant response
22    Assistant(AssistantMessage),
23
24    /// Result message (completion of a query)
25    Result(ResultMessage),
26
27    /// Control request from CLI (tool permissions, hooks, etc.)
28    ControlRequest(ControlRequest),
29
30    /// Control response from CLI (ack for initialization, etc.)
31    ControlResponse(ControlResponse),
32
33    /// API error from Anthropic (500, 529 overloaded, etc.)
34    Error(AnthropicError),
35
36    /// Rate limit status event
37    RateLimitEvent(RateLimitEvent),
38}
39
40impl ClaudeOutput {
41    /// Get the message type as a string
42    pub fn message_type(&self) -> String {
43        match self {
44            ClaudeOutput::System(_) => "system".to_string(),
45            ClaudeOutput::User(_) => "user".to_string(),
46            ClaudeOutput::Assistant(_) => "assistant".to_string(),
47            ClaudeOutput::Result(_) => "result".to_string(),
48            ClaudeOutput::ControlRequest(_) => "control_request".to_string(),
49            ClaudeOutput::ControlResponse(_) => "control_response".to_string(),
50            ClaudeOutput::Error(_) => "error".to_string(),
51            ClaudeOutput::RateLimitEvent(_) => "rate_limit_event".to_string(),
52        }
53    }
54
55    /// Check if this is a control request (tool permission request)
56    pub fn is_control_request(&self) -> bool {
57        matches!(self, ClaudeOutput::ControlRequest(_))
58    }
59
60    /// Check if this is a control response
61    pub fn is_control_response(&self) -> bool {
62        matches!(self, ClaudeOutput::ControlResponse(_))
63    }
64
65    /// Check if this is an Anthropic API error
66    pub fn is_api_error(&self) -> bool {
67        matches!(self, ClaudeOutput::Error(_))
68    }
69
70    /// Get the control request if this is one
71    pub fn as_control_request(&self) -> Option<&ControlRequest> {
72        match self {
73            ClaudeOutput::ControlRequest(req) => Some(req),
74            _ => None,
75        }
76    }
77
78    /// Get the Anthropic error if this is one
79    ///
80    /// # Example
81    /// ```
82    /// use claude_codes::ClaudeOutput;
83    ///
84    /// let json = r#"{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}"#;
85    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
86    ///
87    /// if let Some(err) = output.as_anthropic_error() {
88    ///     if err.is_overloaded() {
89    ///         println!("API is overloaded, retrying...");
90    ///     }
91    /// }
92    /// ```
93    pub fn as_anthropic_error(&self) -> Option<&AnthropicError> {
94        match self {
95            ClaudeOutput::Error(err) => Some(err),
96            _ => None,
97        }
98    }
99
100    /// Check if this is a rate limit event
101    pub fn is_rate_limit_event(&self) -> bool {
102        matches!(self, ClaudeOutput::RateLimitEvent(_))
103    }
104
105    /// Get the rate limit event if this is one
106    pub fn as_rate_limit_event(&self) -> Option<&RateLimitEvent> {
107        match self {
108            ClaudeOutput::RateLimitEvent(evt) => Some(evt),
109            _ => None,
110        }
111    }
112
113    /// Check if this is a result with error
114    pub fn is_error(&self) -> bool {
115        matches!(self, ClaudeOutput::Result(r) if r.is_error)
116    }
117
118    /// Check if this is an assistant message
119    pub fn is_assistant_message(&self) -> bool {
120        matches!(self, ClaudeOutput::Assistant(_))
121    }
122
123    /// Check if this is a system message
124    pub fn is_system_message(&self) -> bool {
125        matches!(self, ClaudeOutput::System(_))
126    }
127
128    /// Check if this is a system init message
129    ///
130    /// # Example
131    /// ```
132    /// use claude_codes::ClaudeOutput;
133    ///
134    /// let json = r#"{"type":"system","subtype":"init","session_id":"abc"}"#;
135    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
136    /// assert!(output.is_system_init());
137    /// ```
138    pub fn is_system_init(&self) -> bool {
139        matches!(self, ClaudeOutput::System(sys) if sys.is_init())
140    }
141
142    /// Get the session ID from any message type that has one.
143    ///
144    /// Returns the session ID from System, Assistant, or Result messages.
145    /// Returns `None` for User, ControlRequest, and ControlResponse messages.
146    ///
147    /// # Example
148    /// ```
149    /// use claude_codes::ClaudeOutput;
150    ///
151    /// let json = r#"{"type":"result","subtype":"success","is_error":false,
152    ///     "duration_ms":100,"duration_api_ms":200,"num_turns":1,
153    ///     "session_id":"my-session","total_cost_usd":0.01}"#;
154    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
155    /// assert_eq!(output.session_id(), Some("my-session"));
156    /// ```
157    pub fn session_id(&self) -> Option<&str> {
158        match self {
159            ClaudeOutput::System(sys) => sys.data.get("session_id").and_then(|v| v.as_str()),
160            ClaudeOutput::Assistant(ass) => Some(&ass.session_id),
161            ClaudeOutput::Result(res) => Some(&res.session_id),
162            ClaudeOutput::User(_) => None,
163            ClaudeOutput::ControlRequest(_) => None,
164            ClaudeOutput::ControlResponse(_) => None,
165            ClaudeOutput::Error(_) => None,
166            ClaudeOutput::RateLimitEvent(evt) => Some(&evt.session_id),
167        }
168    }
169
170    /// Get a specific tool use by name from an assistant message.
171    ///
172    /// Returns the first `ToolUseBlock` with the given name, or `None` if this
173    /// is not an assistant message or doesn't contain the specified tool.
174    ///
175    /// # Example
176    /// ```
177    /// use claude_codes::ClaudeOutput;
178    ///
179    /// let json = r#"{"type":"assistant","message":{"id":"msg_1","role":"assistant",
180    ///     "model":"claude-3","content":[{"type":"tool_use","id":"tu_1",
181    ///     "name":"Bash","input":{"command":"ls"}}]},"session_id":"abc"}"#;
182    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
183    ///
184    /// if let Some(bash) = output.as_tool_use("Bash") {
185    ///     assert_eq!(bash.name, "Bash");
186    /// }
187    /// ```
188    pub fn as_tool_use(&self, tool_name: &str) -> Option<&ToolUseBlock> {
189        match self {
190            ClaudeOutput::Assistant(ass) => {
191                ass.message.content.iter().find_map(|block| match block {
192                    ContentBlock::ToolUse(tu) if tu.name == tool_name => Some(tu),
193                    _ => None,
194                })
195            }
196            _ => None,
197        }
198    }
199
200    /// Get all tool uses from an assistant message.
201    ///
202    /// Returns an iterator over all `ToolUseBlock`s in the message, or an empty
203    /// iterator if this is not an assistant message.
204    ///
205    /// # Example
206    /// ```
207    /// use claude_codes::ClaudeOutput;
208    ///
209    /// let json = r#"{"type":"assistant","message":{"id":"msg_1","role":"assistant",
210    ///     "model":"claude-3","content":[
211    ///         {"type":"tool_use","id":"tu_1","name":"Read","input":{"file_path":"/tmp/a"}},
212    ///         {"type":"tool_use","id":"tu_2","name":"Write","input":{"file_path":"/tmp/b","content":"x"}}
213    ///     ]},"session_id":"abc"}"#;
214    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
215    ///
216    /// let tools: Vec<_> = output.tool_uses().collect();
217    /// assert_eq!(tools.len(), 2);
218    /// ```
219    pub fn tool_uses(&self) -> impl Iterator<Item = &ToolUseBlock> {
220        let content = match self {
221            ClaudeOutput::Assistant(ass) => Some(&ass.message.content),
222            _ => None,
223        };
224
225        content
226            .into_iter()
227            .flat_map(|c| c.iter())
228            .filter_map(|block| match block {
229                ContentBlock::ToolUse(tu) => Some(tu),
230                _ => None,
231            })
232    }
233
234    /// Get text content from an assistant message.
235    ///
236    /// Returns the concatenated text from all text blocks in the message,
237    /// or `None` if this is not an assistant message or has no text content.
238    ///
239    /// # Example
240    /// ```
241    /// use claude_codes::ClaudeOutput;
242    ///
243    /// let json = r#"{"type":"assistant","message":{"id":"msg_1","role":"assistant",
244    ///     "model":"claude-3","content":[{"type":"text","text":"Hello, world!"}]},
245    ///     "session_id":"abc"}"#;
246    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
247    /// assert_eq!(output.text_content(), Some("Hello, world!".to_string()));
248    /// ```
249    pub fn text_content(&self) -> Option<String> {
250        match self {
251            ClaudeOutput::Assistant(ass) => {
252                let texts: Vec<&str> = ass
253                    .message
254                    .content
255                    .iter()
256                    .filter_map(|block| match block {
257                        ContentBlock::Text(t) => Some(t.text.as_str()),
258                        _ => None,
259                    })
260                    .collect();
261
262                if texts.is_empty() {
263                    None
264                } else {
265                    Some(texts.join(""))
266                }
267            }
268            _ => None,
269        }
270    }
271
272    /// Get the assistant message if this is one.
273    ///
274    /// # Example
275    /// ```
276    /// use claude_codes::ClaudeOutput;
277    ///
278    /// let json = r#"{"type":"assistant","message":{"id":"msg_1","role":"assistant",
279    ///     "model":"claude-3","content":[]},"session_id":"abc"}"#;
280    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
281    ///
282    /// if let Some(assistant) = output.as_assistant() {
283    ///     assert_eq!(assistant.message.model, "claude-3");
284    /// }
285    /// ```
286    pub fn as_assistant(&self) -> Option<&AssistantMessage> {
287        match self {
288            ClaudeOutput::Assistant(ass) => Some(ass),
289            _ => None,
290        }
291    }
292
293    /// Get the result message if this is one.
294    ///
295    /// # Example
296    /// ```
297    /// use claude_codes::ClaudeOutput;
298    ///
299    /// let json = r#"{"type":"result","subtype":"success","is_error":false,
300    ///     "duration_ms":100,"duration_api_ms":200,"num_turns":1,
301    ///     "session_id":"abc","total_cost_usd":0.01}"#;
302    /// let output: ClaudeOutput = serde_json::from_str(json).unwrap();
303    ///
304    /// if let Some(result) = output.as_result() {
305    ///     assert!(!result.is_error);
306    /// }
307    /// ```
308    pub fn as_result(&self) -> Option<&ResultMessage> {
309        match self {
310            ClaudeOutput::Result(res) => Some(res),
311            _ => None,
312        }
313    }
314
315    /// Get the system message if this is one.
316    pub fn as_system(&self) -> Option<&SystemMessage> {
317        match self {
318            ClaudeOutput::System(sys) => Some(sys),
319            _ => None,
320        }
321    }
322
323    /// Parse a JSON string, handling potential ANSI escape codes and other prefixes
324    /// This method will:
325    /// 1. First try to parse as-is
326    /// 2. If that fails, trim until it finds a '{' and try again
327    pub fn parse_json_tolerant(s: &str) -> Result<ClaudeOutput, ParseError> {
328        // First try to parse as-is
329        match Self::parse_json(s) {
330            Ok(output) => Ok(output),
331            Err(first_error) => {
332                // If that fails, look for the first '{' character
333                if let Some(json_start) = s.find('{') {
334                    let trimmed = &s[json_start..];
335                    match Self::parse_json(trimmed) {
336                        Ok(output) => Ok(output),
337                        Err(_) => {
338                            // Return the original error if both attempts fail
339                            Err(first_error)
340                        }
341                    }
342                } else {
343                    Err(first_error)
344                }
345            }
346        }
347    }
348
349    /// Parse a JSON string, returning ParseError with raw JSON if it doesn't match our types
350    pub fn parse_json(s: &str) -> Result<ClaudeOutput, ParseError> {
351        // First try to parse as a Value
352        let value: Value = serde_json::from_str(s).map_err(|e| ParseError {
353            raw_line: s.to_string(),
354            raw_json: None,
355            error_message: format!("Invalid JSON: {}", e),
356        })?;
357
358        // Then try to parse that Value as ClaudeOutput
359        serde_json::from_value::<ClaudeOutput>(value.clone()).map_err(|e| ParseError {
360            raw_line: s.to_string(),
361            raw_json: Some(value),
362            error_message: e.to_string(),
363        })
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370
371    #[test]
372    fn test_deserialize_assistant_message() {
373        let json = r#"{
374            "type": "assistant",
375            "message": {
376                "id": "msg_123",
377                "role": "assistant",
378                "model": "claude-3-sonnet",
379                "content": [{"type": "text", "text": "Hello! How can I help you?"}]
380            },
381            "session_id": "123"
382        }"#;
383
384        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
385        assert!(output.is_assistant_message());
386    }
387
388    #[test]
389    fn test_is_system_init() {
390        let init_json = r#"{
391            "type": "system",
392            "subtype": "init",
393            "session_id": "test-session"
394        }"#;
395        let output: ClaudeOutput = serde_json::from_str(init_json).unwrap();
396        assert!(output.is_system_init());
397
398        let status_json = r#"{
399            "type": "system",
400            "subtype": "status",
401            "session_id": "test-session"
402        }"#;
403        let output: ClaudeOutput = serde_json::from_str(status_json).unwrap();
404        assert!(!output.is_system_init());
405    }
406
407    #[test]
408    fn test_session_id() {
409        // Result message
410        let result_json = r#"{
411            "type": "result",
412            "subtype": "success",
413            "is_error": false,
414            "duration_ms": 100,
415            "duration_api_ms": 200,
416            "num_turns": 1,
417            "session_id": "result-session",
418            "total_cost_usd": 0.01
419        }"#;
420        let output: ClaudeOutput = serde_json::from_str(result_json).unwrap();
421        assert_eq!(output.session_id(), Some("result-session"));
422
423        // Assistant message
424        let assistant_json = r#"{
425            "type": "assistant",
426            "message": {
427                "id": "msg_1",
428                "role": "assistant",
429                "model": "claude-3",
430                "content": []
431            },
432            "session_id": "assistant-session"
433        }"#;
434        let output: ClaudeOutput = serde_json::from_str(assistant_json).unwrap();
435        assert_eq!(output.session_id(), Some("assistant-session"));
436
437        // System message
438        let system_json = r#"{
439            "type": "system",
440            "subtype": "init",
441            "session_id": "system-session"
442        }"#;
443        let output: ClaudeOutput = serde_json::from_str(system_json).unwrap();
444        assert_eq!(output.session_id(), Some("system-session"));
445    }
446
447    #[test]
448    fn test_as_tool_use() {
449        let json = r#"{
450            "type": "assistant",
451            "message": {
452                "id": "msg_1",
453                "role": "assistant",
454                "model": "claude-3",
455                "content": [
456                    {"type": "text", "text": "Let me run that command."},
457                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {"command": "ls -la"}},
458                    {"type": "tool_use", "id": "tu_2", "name": "Read", "input": {"file_path": "/tmp/test"}}
459                ]
460            },
461            "session_id": "abc"
462        }"#;
463        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
464
465        // Find Bash tool
466        let bash = output.as_tool_use("Bash");
467        assert!(bash.is_some());
468        assert_eq!(bash.unwrap().id, "tu_1");
469
470        // Find Read tool
471        let read = output.as_tool_use("Read");
472        assert!(read.is_some());
473        assert_eq!(read.unwrap().id, "tu_2");
474
475        // Non-existent tool
476        assert!(output.as_tool_use("Write").is_none());
477
478        // Not an assistant message
479        let result_json = r#"{
480            "type": "result",
481            "subtype": "success",
482            "is_error": false,
483            "duration_ms": 100,
484            "duration_api_ms": 200,
485            "num_turns": 1,
486            "session_id": "abc",
487            "total_cost_usd": 0.01
488        }"#;
489        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
490        assert!(result.as_tool_use("Bash").is_none());
491    }
492
493    #[test]
494    fn test_tool_uses() {
495        let json = r#"{
496            "type": "assistant",
497            "message": {
498                "id": "msg_1",
499                "role": "assistant",
500                "model": "claude-3",
501                "content": [
502                    {"type": "text", "text": "Running commands..."},
503                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {"command": "ls"}},
504                    {"type": "tool_use", "id": "tu_2", "name": "Read", "input": {"file_path": "/tmp/a"}},
505                    {"type": "tool_use", "id": "tu_3", "name": "Write", "input": {"file_path": "/tmp/b", "content": "x"}}
506                ]
507            },
508            "session_id": "abc"
509        }"#;
510        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
511
512        let tools: Vec<_> = output.tool_uses().collect();
513        assert_eq!(tools.len(), 3);
514        assert_eq!(tools[0].name, "Bash");
515        assert_eq!(tools[1].name, "Read");
516        assert_eq!(tools[2].name, "Write");
517    }
518
519    #[test]
520    fn test_text_content() {
521        // Single text block
522        let json = r#"{
523            "type": "assistant",
524            "message": {
525                "id": "msg_1",
526                "role": "assistant",
527                "model": "claude-3",
528                "content": [{"type": "text", "text": "Hello, world!"}]
529            },
530            "session_id": "abc"
531        }"#;
532        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
533        assert_eq!(output.text_content(), Some("Hello, world!".to_string()));
534
535        // Multiple text blocks
536        let json = r#"{
537            "type": "assistant",
538            "message": {
539                "id": "msg_1",
540                "role": "assistant",
541                "model": "claude-3",
542                "content": [
543                    {"type": "text", "text": "Hello, "},
544                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {}},
545                    {"type": "text", "text": "world!"}
546                ]
547            },
548            "session_id": "abc"
549        }"#;
550        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
551        assert_eq!(output.text_content(), Some("Hello, world!".to_string()));
552
553        // No text blocks
554        let json = r#"{
555            "type": "assistant",
556            "message": {
557                "id": "msg_1",
558                "role": "assistant",
559                "model": "claude-3",
560                "content": [{"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {}}]
561            },
562            "session_id": "abc"
563        }"#;
564        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
565        assert_eq!(output.text_content(), None);
566
567        // Not an assistant message
568        let json = r#"{
569            "type": "result",
570            "subtype": "success",
571            "is_error": false,
572            "duration_ms": 100,
573            "duration_api_ms": 200,
574            "num_turns": 1,
575            "session_id": "abc",
576            "total_cost_usd": 0.01
577        }"#;
578        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
579        assert_eq!(output.text_content(), None);
580    }
581
582    #[test]
583    fn test_as_assistant() {
584        let json = r#"{
585            "type": "assistant",
586            "message": {
587                "id": "msg_1",
588                "role": "assistant",
589                "model": "claude-sonnet-4",
590                "content": []
591            },
592            "session_id": "abc"
593        }"#;
594        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
595
596        let assistant = output.as_assistant();
597        assert!(assistant.is_some());
598        assert_eq!(assistant.unwrap().message.model, "claude-sonnet-4");
599
600        // Not an assistant
601        let result_json = r#"{
602            "type": "result",
603            "subtype": "success",
604            "is_error": false,
605            "duration_ms": 100,
606            "duration_api_ms": 200,
607            "num_turns": 1,
608            "session_id": "abc",
609            "total_cost_usd": 0.01
610        }"#;
611        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
612        assert!(result.as_assistant().is_none());
613    }
614
615    #[test]
616    fn test_as_result() {
617        let json = r#"{
618            "type": "result",
619            "subtype": "success",
620            "is_error": false,
621            "duration_ms": 100,
622            "duration_api_ms": 200,
623            "num_turns": 5,
624            "session_id": "abc",
625            "total_cost_usd": 0.05
626        }"#;
627        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
628
629        let result = output.as_result();
630        assert!(result.is_some());
631        assert_eq!(result.unwrap().num_turns, 5);
632        assert_eq!(result.unwrap().total_cost_usd, 0.05);
633
634        // Not a result
635        let assistant_json = r#"{
636            "type": "assistant",
637            "message": {
638                "id": "msg_1",
639                "role": "assistant",
640                "model": "claude-3",
641                "content": []
642            },
643            "session_id": "abc"
644        }"#;
645        let assistant: ClaudeOutput = serde_json::from_str(assistant_json).unwrap();
646        assert!(assistant.as_result().is_none());
647    }
648
649    #[test]
650    fn test_as_system() {
651        let json = r#"{
652            "type": "system",
653            "subtype": "init",
654            "session_id": "abc",
655            "model": "claude-3"
656        }"#;
657        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
658
659        let system = output.as_system();
660        assert!(system.is_some());
661        assert!(system.unwrap().is_init());
662
663        // Not a system message
664        let result_json = r#"{
665            "type": "result",
666            "subtype": "success",
667            "is_error": false,
668            "duration_ms": 100,
669            "duration_api_ms": 200,
670            "num_turns": 1,
671            "session_id": "abc",
672            "total_cost_usd": 0.01
673        }"#;
674        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
675        assert!(result.as_system().is_none());
676    }
677}