gproxy_protocol/protocol/claude/message/
stream.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::protocol::claude::common::{
6    AssistantRole, Citation, ClaudeModel, Container, ContentBlock, ContextManagementResponse,
7    InputTransformation, JsonObject, MessageObjectType, StopDetails, StopReason, TypedObject,
8    Usage,
9};
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12#[serde(untagged)]
13#[non_exhaustive]
14pub enum StreamEvent {
15    Known(Box<KnownStreamEvent>),
16    Unknown(TypedObject),
17}
18
19impl StreamEvent {
20    /// SSE event name: the wire `type` of this event, if any.
21    pub fn event_name(&self) -> Option<&str> {
22        match self {
23            Self::Known(event) => Some(event.event_name()),
24            Self::Unknown(object) => Some(object.type_.as_str()),
25        }
26    }
27}
28
29#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30#[serde(tag = "type")]
31#[non_exhaustive]
32pub enum KnownStreamEvent {
33    #[serde(rename = "message_start")]
34    MessageStart {
35        message: Box<CreateMessageStartBody>,
36        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
37        extra: JsonObject,
38    },
39    #[serde(rename = "content_block_start")]
40    ContentBlockStart {
41        index: u64,
42        content_block: Box<ContentBlock>,
43        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
44        extra: JsonObject,
45    },
46    #[serde(rename = "content_block_delta")]
47    ContentBlockDelta {
48        index: u64,
49        delta: Box<EventDelta>,
50        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
51        extra: JsonObject,
52    },
53    #[serde(rename = "content_block_stop")]
54    ContentBlockStop {
55        index: u64,
56        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
57        extra: JsonObject,
58    },
59    #[serde(rename = "message_delta")]
60    MessageDelta {
61        #[serde(skip_serializing_if = "Option::is_none")]
62        context_management: Option<Box<ContextManagementResponse>>,
63        delta: Box<MessageDelta>,
64        #[serde(skip_serializing_if = "Option::is_none")]
65        input_transformations: Option<Vec<InputTransformation>>,
66        #[serde(skip_serializing_if = "Option::is_none")]
67        usage: Option<Box<Usage>>,
68        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
69        extra: JsonObject,
70    },
71    #[serde(rename = "message_stop")]
72    MessageStop {
73        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
74        extra: JsonObject,
75    },
76    #[serde(rename = "ping")]
77    Ping {
78        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
79        extra: JsonObject,
80    },
81    #[serde(rename = "error")]
82    Error {
83        error: StreamError,
84        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
85        extra: JsonObject,
86    },
87}
88
89impl KnownStreamEvent {
90    /// SSE event name: the exact serde rename of this variant.
91    pub fn event_name(&self) -> &'static str {
92        match self {
93            Self::MessageStart { .. } => "message_start",
94            Self::ContentBlockStart { .. } => "content_block_start",
95            Self::ContentBlockDelta { .. } => "content_block_delta",
96            Self::ContentBlockStop { .. } => "content_block_stop",
97            Self::MessageDelta { .. } => "message_delta",
98            Self::MessageStop { .. } => "message_stop",
99            Self::Ping { .. } => "ping",
100            Self::Error { .. } => "error",
101        }
102    }
103}
104
105#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
106#[serde(untagged)]
107#[non_exhaustive]
108pub enum EventDelta {
109    Known(Box<KnownEventDelta>),
110    Unknown(TypedObject),
111}
112
113#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
114#[serde(tag = "type")]
115#[non_exhaustive]
116pub enum KnownEventDelta {
117    #[serde(rename = "text_delta")]
118    Text {
119        text: String,
120        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
121        extra: JsonObject,
122    },
123    #[serde(rename = "input_json_delta")]
124    InputJson {
125        partial_json: String,
126        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
127        extra: JsonObject,
128    },
129    #[serde(rename = "citations_delta")]
130    Citations {
131        citation: Box<Citation>,
132        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
133        extra: JsonObject,
134    },
135    #[serde(rename = "thinking_delta")]
136    Thinking {
137        #[serde(default, skip_serializing_if = "Option::is_none")]
138        estimated_tokens: Option<u64>,
139        thinking: String,
140        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
141        extra: JsonObject,
142    },
143    #[serde(rename = "signature_delta")]
144    Signature {
145        signature: String,
146        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
147        extra: JsonObject,
148    },
149    #[serde(rename = "compaction_delta")]
150    Compaction {
151        content: String,
152        encrypted_content: String,
153        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
154        extra: JsonObject,
155    },
156}
157
158#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
159#[non_exhaustive]
160pub struct CreateMessageStartBody {
161    pub id: String,
162    #[serde(rename = "type")]
163    pub type_: MessageObjectType,
164    pub role: AssistantRole,
165    pub content: Vec<ContentBlock>,
166    pub model: ClaudeModel,
167    pub stop_reason: Option<StopReason>,
168    pub stop_sequence: Option<String>,
169    pub usage: Usage,
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub input_transformations: Option<Vec<InputTransformation>>,
172    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
173    pub extra: JsonObject,
174}
175
176#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
177#[non_exhaustive]
178pub struct MessageDelta {
179    #[serde(skip_serializing_if = "Option::is_none")]
180    pub container: Option<Container>,
181    #[serde(skip_serializing_if = "Option::is_none")]
182    pub stop_reason: Option<StopReason>,
183    #[serde(skip_serializing_if = "Option::is_none")]
184    pub stop_sequence: Option<String>,
185    #[serde(skip_serializing_if = "Option::is_none")]
186    pub stop_details: Option<StopDetails>,
187    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
188    pub extra: JsonObject,
189}
190
191#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
192#[non_exhaustive]
193pub struct StreamError {
194    #[serde(rename = "type")]
195    pub type_: String,
196    pub message: String,
197    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
198    pub extra: JsonObject,
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    /// Regression: `event_name()` must equal the serialized `type` tag.
206    #[test]
207    fn event_name_matches_serialized_type_tag() {
208        let events: Vec<StreamEvent> = [
209            r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}"#,
210            r#"{"type":"message_stop"}"#,
211            r#"{"type":"some_future_event","x":1}"#,
212        ]
213        .iter()
214        .map(|raw| serde_json::from_str(raw).unwrap())
215        .collect();
216        for event in events {
217            let value = serde_json::to_value(&event).unwrap();
218            assert_eq!(
219                event.event_name(),
220                value.get("type").and_then(serde_json::Value::as_str),
221                "{value}"
222            );
223        }
224    }
225}
gproxy_protocol/protocol/claude/message/stream.rs

gproxy_protocol/protocol/claude/message/
stream.rs
