1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
use anyhow::Result;
use tokio::sync::mpsc;
use super::types::Usage;
/// Events emitted by the SSE stream.
#[derive(Debug, Clone)]
pub enum ApiEvent {
/// Streaming text from assistant
Text(String),
/// Tool use request
ToolUse {
id: String,
name: String,
input: serde_json::Value,
},
/// Usage information
Usage(Usage),
/// Stream complete
Done,
/// Error from API
Error(String),
}
/// Read an SSE response and send parsed events to the channel.
pub async fn read_sse_stream(
response: reqwest::Response,
tx: mpsc::Sender<ApiEvent>,
) -> Result<()> {
use futures_util::StreamExt as _;
let mut stream = response.bytes_stream();
let mut buffer = String::new();
// Tool use accumulation state
let mut current_tool_id = String::new();
let mut current_tool_name = String::new();
let mut current_tool_input = String::new();
let mut input_tokens: u32 = 0;
let mut output_tokens: u32 = 0;
let mut cache_read_tokens: u32 = 0;
let mut cache_creation_tokens: u32 = 0;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
// Process complete lines
while let Some(newline_pos) = buffer.find('\n') {
let line = buffer[..newline_pos].to_string();
buffer = buffer[newline_pos + 1..].to_string();
let line = line.trim();
if line.is_empty() {
continue;
}
let Some(data) = line.strip_prefix("data: ") else {
continue;
};
if data == "[DONE]" {
let _ = tx
.send(ApiEvent::Usage(Usage {
input_tokens,
output_tokens,
cache_read_tokens,
cache_creation_tokens,
}))
.await;
let _ = tx.send(ApiEvent::Done).await;
return Ok(());
}
let Ok(event) = serde_json::from_str::<serde_json::Value>(data) else {
continue;
};
let event_type = event["type"].as_str().unwrap_or("");
match event_type {
"message_start" => {
if let Some(usage) = event.pointer("/message/usage") {
input_tokens = usage["input_tokens"].as_u64().unwrap_or(0) as u32;
output_tokens = usage["output_tokens"].as_u64().unwrap_or(0) as u32;
cache_read_tokens =
usage["cache_read_input_tokens"].as_u64().unwrap_or(0) as u32;
cache_creation_tokens =
usage["cache_creation_input_tokens"].as_u64().unwrap_or(0) as u32;
}
}
"message_delta" => {
if let Some(usage) = event.get("usage") {
output_tokens = usage["output_tokens"]
.as_u64()
.unwrap_or(output_tokens as u64)
as u32;
}
}
"content_block_start" => {
if let Some(cb) = event.get("content_block") {
if cb["type"].as_str() == Some("tool_use") {
current_tool_id = cb["id"].as_str().unwrap_or("").to_string();
current_tool_name = cb["name"].as_str().unwrap_or("").to_string();
current_tool_input.clear();
}
}
}
"content_block_delta" => {
if let Some(delta) = event.get("delta") {
match delta["type"].as_str().unwrap_or("") {
"text_delta" => {
if let Some(text) = delta["text"].as_str() {
let _ = tx.send(ApiEvent::Text(text.to_string())).await;
}
}
"input_json_delta" => {
if let Some(json) = delta["partial_json"].as_str() {
current_tool_input.push_str(json);
}
}
_ => {}
}
}
}
"content_block_stop" => {
if !current_tool_name.is_empty() && !current_tool_input.is_empty() {
if let Ok(input) = serde_json::from_str(¤t_tool_input) {
let _ = tx
.send(ApiEvent::ToolUse {
id: current_tool_id.clone(),
name: current_tool_name.clone(),
input,
})
.await;
}
current_tool_name.clear();
current_tool_input.clear();
current_tool_id.clear();
}
}
"message_stop" => {
let _ = tx
.send(ApiEvent::Usage(Usage {
input_tokens,
output_tokens,
cache_read_tokens,
cache_creation_tokens,
}))
.await;
let _ = tx.send(ApiEvent::Done).await;
return Ok(());
}
"error" => {
let msg = event["error"]["message"]
.as_str()
.unwrap_or("unknown error");
let _ = tx.send(ApiEvent::Error(msg.to_string())).await;
return Ok(());
}
_ => {}
}
}
}
// Stream ended without message_stop
let _ = tx.send(ApiEvent::Done).await;
Ok(())
}