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_json: Value::String(s.to_string()),
354            error_message: format!("Invalid JSON: {}", e),
355        })?;
356
357        // Then try to parse that Value as ClaudeOutput
358        serde_json::from_value::<ClaudeOutput>(value.clone()).map_err(|e| ParseError {
359            raw_json: value,
360            error_message: e.to_string(),
361        })
362    }
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368
369    #[test]
370    fn test_deserialize_assistant_message() {
371        let json = r#"{
372            "type": "assistant",
373            "message": {
374                "id": "msg_123",
375                "role": "assistant",
376                "model": "claude-3-sonnet",
377                "content": [{"type": "text", "text": "Hello! How can I help you?"}]
378            },
379            "session_id": "123"
380        }"#;
381
382        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
383        assert!(output.is_assistant_message());
384    }
385
386    #[test]
387    fn test_is_system_init() {
388        let init_json = r#"{
389            "type": "system",
390            "subtype": "init",
391            "session_id": "test-session"
392        }"#;
393        let output: ClaudeOutput = serde_json::from_str(init_json).unwrap();
394        assert!(output.is_system_init());
395
396        let status_json = r#"{
397            "type": "system",
398            "subtype": "status",
399            "session_id": "test-session"
400        }"#;
401        let output: ClaudeOutput = serde_json::from_str(status_json).unwrap();
402        assert!(!output.is_system_init());
403    }
404
405    #[test]
406    fn test_session_id() {
407        // Result message
408        let result_json = r#"{
409            "type": "result",
410            "subtype": "success",
411            "is_error": false,
412            "duration_ms": 100,
413            "duration_api_ms": 200,
414            "num_turns": 1,
415            "session_id": "result-session",
416            "total_cost_usd": 0.01
417        }"#;
418        let output: ClaudeOutput = serde_json::from_str(result_json).unwrap();
419        assert_eq!(output.session_id(), Some("result-session"));
420
421        // Assistant message
422        let assistant_json = r#"{
423            "type": "assistant",
424            "message": {
425                "id": "msg_1",
426                "role": "assistant",
427                "model": "claude-3",
428                "content": []
429            },
430            "session_id": "assistant-session"
431        }"#;
432        let output: ClaudeOutput = serde_json::from_str(assistant_json).unwrap();
433        assert_eq!(output.session_id(), Some("assistant-session"));
434
435        // System message
436        let system_json = r#"{
437            "type": "system",
438            "subtype": "init",
439            "session_id": "system-session"
440        }"#;
441        let output: ClaudeOutput = serde_json::from_str(system_json).unwrap();
442        assert_eq!(output.session_id(), Some("system-session"));
443    }
444
445    #[test]
446    fn test_as_tool_use() {
447        let json = r#"{
448            "type": "assistant",
449            "message": {
450                "id": "msg_1",
451                "role": "assistant",
452                "model": "claude-3",
453                "content": [
454                    {"type": "text", "text": "Let me run that command."},
455                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {"command": "ls -la"}},
456                    {"type": "tool_use", "id": "tu_2", "name": "Read", "input": {"file_path": "/tmp/test"}}
457                ]
458            },
459            "session_id": "abc"
460        }"#;
461        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
462
463        // Find Bash tool
464        let bash = output.as_tool_use("Bash");
465        assert!(bash.is_some());
466        assert_eq!(bash.unwrap().id, "tu_1");
467
468        // Find Read tool
469        let read = output.as_tool_use("Read");
470        assert!(read.is_some());
471        assert_eq!(read.unwrap().id, "tu_2");
472
473        // Non-existent tool
474        assert!(output.as_tool_use("Write").is_none());
475
476        // Not an assistant message
477        let result_json = r#"{
478            "type": "result",
479            "subtype": "success",
480            "is_error": false,
481            "duration_ms": 100,
482            "duration_api_ms": 200,
483            "num_turns": 1,
484            "session_id": "abc",
485            "total_cost_usd": 0.01
486        }"#;
487        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
488        assert!(result.as_tool_use("Bash").is_none());
489    }
490
491    #[test]
492    fn test_tool_uses() {
493        let json = r#"{
494            "type": "assistant",
495            "message": {
496                "id": "msg_1",
497                "role": "assistant",
498                "model": "claude-3",
499                "content": [
500                    {"type": "text", "text": "Running commands..."},
501                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {"command": "ls"}},
502                    {"type": "tool_use", "id": "tu_2", "name": "Read", "input": {"file_path": "/tmp/a"}},
503                    {"type": "tool_use", "id": "tu_3", "name": "Write", "input": {"file_path": "/tmp/b", "content": "x"}}
504                ]
505            },
506            "session_id": "abc"
507        }"#;
508        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
509
510        let tools: Vec<_> = output.tool_uses().collect();
511        assert_eq!(tools.len(), 3);
512        assert_eq!(tools[0].name, "Bash");
513        assert_eq!(tools[1].name, "Read");
514        assert_eq!(tools[2].name, "Write");
515    }
516
517    #[test]
518    fn test_text_content() {
519        // Single text block
520        let json = r#"{
521            "type": "assistant",
522            "message": {
523                "id": "msg_1",
524                "role": "assistant",
525                "model": "claude-3",
526                "content": [{"type": "text", "text": "Hello, world!"}]
527            },
528            "session_id": "abc"
529        }"#;
530        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
531        assert_eq!(output.text_content(), Some("Hello, world!".to_string()));
532
533        // Multiple text blocks
534        let json = r#"{
535            "type": "assistant",
536            "message": {
537                "id": "msg_1",
538                "role": "assistant",
539                "model": "claude-3",
540                "content": [
541                    {"type": "text", "text": "Hello, "},
542                    {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {}},
543                    {"type": "text", "text": "world!"}
544                ]
545            },
546            "session_id": "abc"
547        }"#;
548        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
549        assert_eq!(output.text_content(), Some("Hello, world!".to_string()));
550
551        // No text blocks
552        let json = r#"{
553            "type": "assistant",
554            "message": {
555                "id": "msg_1",
556                "role": "assistant",
557                "model": "claude-3",
558                "content": [{"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {}}]
559            },
560            "session_id": "abc"
561        }"#;
562        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
563        assert_eq!(output.text_content(), None);
564
565        // Not an assistant message
566        let json = r#"{
567            "type": "result",
568            "subtype": "success",
569            "is_error": false,
570            "duration_ms": 100,
571            "duration_api_ms": 200,
572            "num_turns": 1,
573            "session_id": "abc",
574            "total_cost_usd": 0.01
575        }"#;
576        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
577        assert_eq!(output.text_content(), None);
578    }
579
580    #[test]
581    fn test_as_assistant() {
582        let json = r#"{
583            "type": "assistant",
584            "message": {
585                "id": "msg_1",
586                "role": "assistant",
587                "model": "claude-sonnet-4",
588                "content": []
589            },
590            "session_id": "abc"
591        }"#;
592        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
593
594        let assistant = output.as_assistant();
595        assert!(assistant.is_some());
596        assert_eq!(assistant.unwrap().message.model, "claude-sonnet-4");
597
598        // Not an assistant
599        let result_json = r#"{
600            "type": "result",
601            "subtype": "success",
602            "is_error": false,
603            "duration_ms": 100,
604            "duration_api_ms": 200,
605            "num_turns": 1,
606            "session_id": "abc",
607            "total_cost_usd": 0.01
608        }"#;
609        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
610        assert!(result.as_assistant().is_none());
611    }
612
613    #[test]
614    fn test_as_result() {
615        let json = r#"{
616            "type": "result",
617            "subtype": "success",
618            "is_error": false,
619            "duration_ms": 100,
620            "duration_api_ms": 200,
621            "num_turns": 5,
622            "session_id": "abc",
623            "total_cost_usd": 0.05
624        }"#;
625        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
626
627        let result = output.as_result();
628        assert!(result.is_some());
629        assert_eq!(result.unwrap().num_turns, 5);
630        assert_eq!(result.unwrap().total_cost_usd, 0.05);
631
632        // Not a result
633        let assistant_json = r#"{
634            "type": "assistant",
635            "message": {
636                "id": "msg_1",
637                "role": "assistant",
638                "model": "claude-3",
639                "content": []
640            },
641            "session_id": "abc"
642        }"#;
643        let assistant: ClaudeOutput = serde_json::from_str(assistant_json).unwrap();
644        assert!(assistant.as_result().is_none());
645    }
646
647    #[test]
648    fn test_as_system() {
649        let json = r#"{
650            "type": "system",
651            "subtype": "init",
652            "session_id": "abc",
653            "model": "claude-3"
654        }"#;
655        let output: ClaudeOutput = serde_json::from_str(json).unwrap();
656
657        let system = output.as_system();
658        assert!(system.is_some());
659        assert!(system.unwrap().is_init());
660
661        // Not a system message
662        let result_json = r#"{
663            "type": "result",
664            "subtype": "success",
665            "is_error": false,
666            "duration_ms": 100,
667            "duration_api_ms": 200,
668            "num_turns": 1,
669            "session_id": "abc",
670            "total_cost_usd": 0.01
671        }"#;
672        let result: ClaudeOutput = serde_json::from_str(result_json).unwrap();
673        assert!(result.as_system().is_none());
674    }
675}