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#[derive(Debug, Clone, Serialize, Deserialize)]
13#[serde(tag = "type", rename_all = "snake_case")]
14pub enum ClaudeOutput {
15 System(SystemMessage),
17
18 User(UserMessage),
20
21 Assistant(AssistantMessage),
23
24 Result(ResultMessage),
26
27 ControlRequest(ControlRequest),
29
30 ControlResponse(ControlResponse),
32
33 Error(AnthropicError),
35
36 RateLimitEvent(RateLimitEvent),
38}
39
40impl ClaudeOutput {
41 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 pub fn is_control_request(&self) -> bool {
57 matches!(self, ClaudeOutput::ControlRequest(_))
58 }
59
60 pub fn is_control_response(&self) -> bool {
62 matches!(self, ClaudeOutput::ControlResponse(_))
63 }
64
65 pub fn is_api_error(&self) -> bool {
67 matches!(self, ClaudeOutput::Error(_))
68 }
69
70 pub fn as_control_request(&self) -> Option<&ControlRequest> {
72 match self {
73 ClaudeOutput::ControlRequest(req) => Some(req),
74 _ => None,
75 }
76 }
77
78 pub fn as_anthropic_error(&self) -> Option<&AnthropicError> {
94 match self {
95 ClaudeOutput::Error(err) => Some(err),
96 _ => None,
97 }
98 }
99
100 pub fn is_rate_limit_event(&self) -> bool {
102 matches!(self, ClaudeOutput::RateLimitEvent(_))
103 }
104
105 pub fn as_rate_limit_event(&self) -> Option<&RateLimitEvent> {
107 match self {
108 ClaudeOutput::RateLimitEvent(evt) => Some(evt),
109 _ => None,
110 }
111 }
112
113 pub fn is_error(&self) -> bool {
115 matches!(self, ClaudeOutput::Result(r) if r.is_error)
116 }
117
118 pub fn is_assistant_message(&self) -> bool {
120 matches!(self, ClaudeOutput::Assistant(_))
121 }
122
123 pub fn is_system_message(&self) -> bool {
125 matches!(self, ClaudeOutput::System(_))
126 }
127
128 pub fn is_system_init(&self) -> bool {
139 matches!(self, ClaudeOutput::System(sys) if sys.is_init())
140 }
141
142 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 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 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 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 pub fn as_assistant(&self) -> Option<&AssistantMessage> {
287 match self {
288 ClaudeOutput::Assistant(ass) => Some(ass),
289 _ => None,
290 }
291 }
292
293 pub fn as_result(&self) -> Option<&ResultMessage> {
309 match self {
310 ClaudeOutput::Result(res) => Some(res),
311 _ => None,
312 }
313 }
314
315 pub fn as_system(&self) -> Option<&SystemMessage> {
317 match self {
318 ClaudeOutput::System(sys) => Some(sys),
319 _ => None,
320 }
321 }
322
323 pub fn parse_json_tolerant(s: &str) -> Result<ClaudeOutput, ParseError> {
328 match Self::parse_json(s) {
330 Ok(output) => Ok(output),
331 Err(first_error) => {
332 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 Err(first_error)
340 }
341 }
342 } else {
343 Err(first_error)
344 }
345 }
346 }
347 }
348
349 pub fn parse_json(s: &str) -> Result<ClaudeOutput, ParseError> {
351 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 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 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 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 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 let bash = output.as_tool_use("Bash");
465 assert!(bash.is_some());
466 assert_eq!(bash.unwrap().id, "tu_1");
467
468 let read = output.as_tool_use("Read");
470 assert!(read.is_some());
471 assert_eq!(read.unwrap().id, "tu_2");
472
473 assert!(output.as_tool_use("Write").is_none());
475
476 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 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 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 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 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 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 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 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}