Skip to main content

claude_code_rs/types/
messages.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3use tokio_stream::{Stream, StreamExt};
4
5use super::content::ContentBlock;
6use crate::error::Result;
7
8/// A message from the Claude CLI streaming protocol.
9///
10/// The CLI emits newline-delimited JSON objects with a top-level `type` field.
11/// Each variant corresponds to one of these message types.
12#[derive(Debug, Clone)]
13#[non_exhaustive]
14pub enum Message {
15    /// System-level message (init acknowledgment, etc.)
16    System {
17        subtype: String,
18        data: Value,
19    },
20
21    /// Assistant (Claude) response message.
22    Assistant {
23        message: AssistantMessage,
24    },
25
26    /// User message echo.
27    User {
28        message: UserMessage,
29    },
30
31    /// Result/completion message - signals end of a turn.
32    Result {
33        result: ResultMessage,
34    },
35
36    /// An unknown message type we don't recognize but preserve.
37    Unknown {
38        message_type: String,
39        raw: Value,
40    },
41}
42
43/// An assistant response with content blocks.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct AssistantMessage {
46    #[serde(default)]
47    pub id: Option<String>,
48    #[serde(default)]
49    pub model: Option<String>,
50    #[serde(default)]
51    pub content: Vec<ContentBlock>,
52    #[serde(default)]
53    pub stop_reason: Option<String>,
54    #[serde(default)]
55    pub usage: Option<Usage>,
56    /// Raw extra fields we don't explicitly model.
57    #[serde(flatten)]
58    pub extra: Value,
59}
60
61/// A user message as echoed back by the CLI.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct UserMessage {
64    #[serde(default)]
65    pub id: Option<String>,
66    #[serde(default)]
67    pub content: UserContent,
68    #[serde(flatten)]
69    pub extra: Value,
70}
71
72/// User message content can be a string or structured blocks.
73#[derive(Debug, Clone, Serialize, Deserialize, Default)]
74#[serde(untagged)]
75pub enum UserContent {
76    Text(String),
77    Blocks(Vec<ContentBlock>),
78    #[default]
79    Empty,
80}
81
82/// Result message indicating the end of a query turn.
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct ResultMessage {
85    #[serde(default)]
86    pub subtype: Option<String>,
87    #[serde(default)]
88    pub is_error: bool,
89    #[serde(default)]
90    pub error: Option<String>,
91    #[serde(default)]
92    pub duration_ms: Option<f64>,
93    #[serde(default)]
94    pub duration_api_ms: Option<f64>,
95    #[serde(default)]
96    pub num_turns: Option<u32>,
97    #[serde(default)]
98    pub session_id: Option<String>,
99    #[serde(default)]
100    pub cost_usd: Option<f64>,
101    #[serde(default)]
102    pub total_cost_usd: Option<f64>,
103    #[serde(default)]
104    pub usage: Option<Usage>,
105    #[serde(flatten)]
106    pub extra: Value,
107}
108
109/// Token usage information.
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct Usage {
112    #[serde(default)]
113    pub input_tokens: Option<u64>,
114    #[serde(default)]
115    pub output_tokens: Option<u64>,
116    #[serde(default)]
117    pub cache_creation_input_tokens: Option<u64>,
118    #[serde(default)]
119    pub cache_read_input_tokens: Option<u64>,
120    #[serde(flatten)]
121    pub extra: Value,
122}
123
124impl Message {
125    /// Returns true if this is a Result message (end of turn).
126    pub fn is_result(&self) -> bool {
127        matches!(self, Message::Result { .. })
128    }
129
130    /// Returns true if this is a Result message with an error.
131    pub fn is_error(&self) -> bool {
132        matches!(self, Message::Result { result } if result.is_error)
133    }
134
135    /// Extract all text content from an Assistant message.
136    pub fn text(&self) -> Option<String> {
137        match self {
138            Message::Assistant { message } => {
139                let mut result = String::new();
140                for block in &message.content {
141                    if let Some(text) = block.as_text() {
142                        result.push_str(text);
143                    }
144                }
145                if result.is_empty() { None } else { Some(result) }
146            }
147            _ => None,
148        }
149    }
150
151    /// Get the session ID from a Result message.
152    pub fn session_id(&self) -> Option<&str> {
153        match self {
154            Message::Result { result } => result.session_id.as_deref(),
155            _ => None,
156        }
157    }
158}
159
160/// Collect messages from a stream until a Result message is received.
161pub(crate) async fn collect_until_result(stream: &mut (impl Stream<Item = Result<Message>> + Unpin)) -> Result<Vec<Message>> {
162    let mut messages = Vec::new();
163    while let Some(msg) = stream.next().await {
164        let msg = msg?;
165        let is_result = msg.is_result();
166        messages.push(msg);
167        if is_result {
168            break;
169        }
170    }
171    Ok(messages)
172}