gproxy_protocol/transform/generate/stream/claude/
collector.rs

1use crate::{
2    codec::{self, CodecErrorKind, CodecLimits},
3    transform::{Converted, TransformError, TransformErrorKind},
4    wire::{
5        DeclaredFields,
6        claude::{
7            generate_content as c,
8            stream::{self, StreamEvent},
9        },
10    },
11};
12use std::collections::{HashMap, HashSet};
13
14#[derive(Debug, Clone, Copy)]
15pub struct ClaudeStreamLimits {
16    pub max_text_bytes: usize,
17    pub max_json_bytes: usize,
18}
19
20impl Default for ClaudeStreamLimits {
21    fn default() -> Self {
22        Self {
23            max_text_bytes: 16 * 1024 * 1024,
24            max_json_bytes: 16 * 1024 * 1024,
25        }
26    }
27}
28
29/// Native Claude lifecycle collector. Every failed push poisons the stream;
30/// EOF never substitutes for message_stop or an unclosed content block.
31pub struct ClaudeStreamCollector {
32    message: Option<stream::StreamMessage>,
33
34    bytes: usize,
35    text_bytes: usize,
36    limits: ClaudeStreamLimits,
37    stopped: bool,
38    failed: bool,
39    message_delta: bool,
40    closed: HashSet<usize>,
41    call_ids: HashSet<String>,
42    json_buffers: HashMap<usize, String>,
43}
44
45impl ClaudeStreamCollector {
46    pub fn new(limits: ClaudeStreamLimits) -> Self {
47        Self {
48            message: None,
49
50            bytes: 0,
51            text_bytes: 0,
52            limits,
53            stopped: false,
54            failed: false,
55            message_delta: false,
56            closed: HashSet::new(),
57            call_ids: HashSet::new(),
58            json_buffers: HashMap::new(),
59        }
60    }
61    pub fn push(&mut self, event: StreamEvent) -> Result<(), TransformError> {
62        if self.failed {
63            return Err(invalid("stream", "collector already failed"));
64        }
65        let result = self.push_declared(event.into_declared());
66        if result.is_err() {
67            self.failed = true;
68        }
69        result
70    }
71    fn push_declared(&mut self, event: StreamEvent) -> Result<(), TransformError> {
72        if self.stopped {
73            return Err(invalid("stream", "event after message_stop"));
74        }
75        let encoded = bounded(
76            &event,
77            self.limits.max_json_bytes.saturating_sub(self.bytes),
78        )?;
79        add(
80            &mut self.bytes,
81            encoded.len(),
82            self.limits.max_json_bytes,
83            "bytes",
84        )?;
85        match event {
86            StreamEvent::MessageStart(event) => {
87                if self.message.is_some() {
88                    return Err(invalid("message_start", "duplicate start"));
89                }
90                if event.message.id.is_empty() || event.message.model.is_empty() {
91                    return Err(invalid("message_start", "empty message identity or model"));
92                }
93                if !event.message.content.is_empty()
94                    || event.message.stop_reason.flatten().is_some()
95                {
96                    return Err(invalid(
97                        "message_start",
98                        "start must precede content and terminal reason",
99                    ));
100                }
101
102                self.message = Some(event.message);
103            }
104            StreamEvent::ContentBlockStart(event) => {
105                if self.message_delta {
106                    return Err(invalid(
107                        "content_block_start",
108                        "content after message_delta",
109                    ));
110                }
111                let message = self
112                    .message
113                    .as_mut()
114                    .ok_or_else(|| invalid("content_block_start", "missing message_start"))?;
115                let index = index(event.index)?;
116                if index != message.content.len() {
117                    return Err(invalid(
118                        "content_block_start.index",
119                        "index must append a new block",
120                    ));
121                }
122
123                if let Some(id) = super::blocks::call_id(&event.content_block)
124                    && (id.is_empty() || !self.call_ids.insert(id.to_owned()))
125                {
126                    return Err(invalid(
127                        "content_block.id",
128                        "empty or duplicate tool call ID",
129                    ));
130                }
131                add(
132                    &mut self.text_bytes,
133                    super::blocks::text_bytes(&event.content_block),
134                    self.limits.max_text_bytes,
135                    "text_bytes",
136                )?;
137                if let c::ResponseContentBlock::Fallback(block) = &event.content_block {
138                    if block.from.model != message.model {
139                        return Err(invalid(
140                            "fallback.from.model",
141                            "fallback source differs from current model",
142                        ));
143                    }
144                    if block.to.model.is_empty() {
145                        return Err(invalid("fallback.to.model", "empty fallback model"));
146                    }
147                    message.model.clone_from(&block.to.model);
148                }
149                message.content.push(event.content_block);
150            }
151            StreamEvent::ContentBlockDelta(event) => {
152                if self.message_delta {
153                    return Err(invalid(
154                        "content_block_delta",
155                        "content after message_delta",
156                    ));
157                }
158                let index = index(event.index)?;
159                if self.closed.contains(&index) {
160                    return Err(invalid(
161                        "content_block_delta.index",
162                        "delta after block stop",
163                    ));
164                }
165                let block = self
166                    .message
167                    .as_mut()
168                    .and_then(|m| m.content.get_mut(index))
169                    .ok_or_else(|| invalid("content_block_delta.index", "unknown block"))?;
170                add(
171                    &mut self.text_bytes,
172                    super::blocks::delta_text_bytes(&event.delta),
173                    self.limits.max_text_bytes,
174                    "text_bytes",
175                )?;
176                super::blocks::apply(
177                    block,
178                    event.delta,
179                    self.json_buffers.entry(index).or_default(),
180                )?;
181            }
182            StreamEvent::ContentBlockStop(event) => {
183                let index = index(event.index)?;
184                let block = self
185                    .message
186                    .as_mut()
187                    .and_then(|m| m.content.get_mut(index))
188                    .ok_or_else(|| invalid("content_block_stop.index", "unknown block"))?;
189                if !self.closed.insert(index) {
190                    return Err(invalid("content_block_stop.index", "duplicate block stop"));
191                }
192                if let Some(json) = self.json_buffers.remove(&index).filter(|s| !s.is_empty())
193                    && let Some(input) = super::blocks::input(block)
194                {
195                    *input = serde_json::from_str(&json)
196                        .map_err(|e| invalid("tool.input", e.to_string()))?;
197                }
198            }
199            StreamEvent::MessageDelta(event) => {
200                let message = self
201                    .message
202                    .as_mut()
203                    .ok_or_else(|| invalid("message_delta", "missing message_start"))?;
204                if self.closed.len() != message.content.len() {
205                    return Err(invalid("message_delta", "content blocks remain open"));
206                }
207                self.message_delta = true;
208                super::usage::merge(message, *event)?;
209            }
210            StreamEvent::MessageStop(_) => {
211                let message = self
212                    .message
213                    .as_ref()
214                    .ok_or_else(|| invalid("message_stop", "missing message_start"))?;
215                if !self.message_delta
216                    || self.closed.len() != message.content.len()
217                    || message.stop_reason.flatten().is_none()
218                {
219                    return Err(invalid(
220                        "message_stop",
221                        "missing message terminal or unclosed content",
222                    ));
223                }
224                self.stopped = true;
225            }
226            StreamEvent::Error(event) => {
227                return Err(TransformError::with_source(
228                    TransformErrorKind::InvalidResult,
229                    "stream.error",
230                    format!("{:?}: {}", event.error.type_, event.error.message),
231                    crate::transform::UpstreamError(event.error),
232                ));
233            }
234            StreamEvent::Ping(_) => {}
235        }
236        Ok(())
237    }
238    pub fn finish(self) -> Result<Converted<c::GenerateContentResponseBody>, TransformError> {
239        if self.failed || !self.stopped {
240            return Err(invalid("stream", "failed stream or premature EOF"));
241        }
242        let m = self
243            .message
244            .ok_or_else(|| invalid("stream", "missing message_start"))?;
245        Ok(Converted {
246            value: c::GenerateContentResponseBody {
247                input_transformations: m.input_transformations,
248                type_: m.type_,
249                id: m.id,
250                container: m.container,
251                content: m.content,
252                context_management: m.context_management,
253                diagnostics: m.diagnostics,
254                model: m.model,
255                role: m.role,
256                stop_details: m.stop_details,
257                stop_reason: m
258                    .stop_reason
259                    .flatten()
260                    .ok_or_else(|| invalid("stop_reason", "missing terminal reason"))?,
261                stop_sequence: m.stop_sequence,
262                usage: m.usage,
263                rest: Default::default(),
264            },
265            report: Default::default(),
266        })
267    }
268}
269
270pub(super) fn bounded<T: serde::Serialize>(
271    value: &T,
272    remaining: usize,
273) -> Result<bytes::Bytes, TransformError> {
274    let n = remaining as u64;
275    codec::encode_json(
276        value,
277        CodecLimits {
278            max_buffer_bytes: n,
279            max_value_bytes: n,
280            max_body_bytes: n,
281            max_line_bytes: n,
282            max_part_bytes: n,
283            max_parts: 0,
284        },
285    )
286    .map_err(|e| {
287        if e.kind() == CodecErrorKind::Limit {
288            limit("bytes")
289        } else {
290            invalid("stream.event", e.to_string())
291        }
292    })
293}
294
295pub(super) fn invalid(field: &str, message: impl Into<String>) -> TransformError {
296    TransformError::invalid_result(field, message)
297}
298
299fn limit(field: &str) -> TransformError {
300    TransformError::new(
301        TransformErrorKind::Limit,
302        format!("stream.{field}"),
303        "Claude stream limit exceeded",
304    )
305}
306
307fn add(total: &mut usize, n: usize, max: usize, field: &str) -> Result<(), TransformError> {
308    *total = total
309        .checked_add(n)
310        .filter(|n| *n <= max)
311        .ok_or_else(|| limit(field))?;
312    Ok(())
313}
314
315fn index(n: i64) -> Result<usize, TransformError> {
316    usize::try_from(n).map_err(|_| invalid("block.index", "negative or unrepresentable index"))
317}
gproxy_protocol/transform/generate/stream/claude/collector.rs

gproxy_protocol/transform/generate/stream/claude/
collector.rs
