gproxy_protocol/transform/generate/claude_responses/stream/
common.rs

1use crate::{
2    transform::{
3        Report, TransformError, TransformErrorKind,
4        identity::{IdentityFlow, IdentityRole, KnownIdPrefix, SourceIdentity, TargetIdPolicy},
5    },
6    wire::{DeclaredFields, claude::generate_content as c, openai::responses as r},
7};
8
9#[derive(Debug, Clone, Copy)]
10pub struct StreamLimits {
11    pub max_bytes: usize,
12    pub max_pending: usize,
13}
14
15impl Default for StreamLimits {
16    fn default() -> Self {
17        Self {
18            max_bytes: 16 * 1024 * 1024,
19            max_pending: 16 * 1024 * 1024,
20        }
21    }
22}
23
24#[derive(Debug)]
25pub struct StreamEnd<T> {
26    pub chunks: Vec<T>,
27    pub identities: IdentityFlow,
28    pub report: Report,
29}
30
31pub(super) fn invalid(message: impl Into<String>) -> TransformError {
32    TransformError::invalid_result("claude_responses.stream", message)
33}
34
35pub(super) fn limit() -> TransformError {
36    TransformError::new(
37        TransformErrorKind::Limit,
38        "claude_responses.stream",
39        "stream limit exceeded",
40    )
41}
42
43pub(super) struct Budget {
44    limits: StreamLimits,
45
46    input_bytes: usize,
47
48    output_bytes: usize,
49}
50
51impl Budget {
52    pub fn new(limits: StreamLimits) -> Self {
53        Self {
54            limits,
55
56            input_bytes: 0,
57
58            output_bytes: 0,
59        }
60    }
61    pub fn input<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
62        self.input_bytes += measure(
63            value,
64            self.limits.max_bytes.saturating_sub(self.input_bytes),
65        )?;
66        Ok(())
67    }
68    pub fn output<T: serde::Serialize>(&mut self, value: &T) -> Result<usize, TransformError> {
69        let n = measure(
70            value,
71            self.limits.max_bytes.saturating_sub(self.output_bytes),
72        )?;
73        self.output_bytes += n;
74        Ok(n)
75    }
76}
77
78pub(super) fn measure<T: serde::Serialize>(value: &T, cap: usize) -> Result<usize, TransformError> {
79    let n = cap as u64;
80    crate::codec::encode_json(
81        value,
82        crate::codec::CodecLimits {
83            max_value_bytes: n,
84            max_body_bytes: n,
85            max_buffer_bytes: n,
86            max_line_bytes: n,
87            max_part_bytes: n,
88            max_parts: usize::MAX,
89        },
90    )
91    .map(|v| v.len())
92    .map_err(|e| {
93        if e.kind() == crate::codec::CodecErrorKind::Limit {
94            limit()
95        } else {
96            invalid(e.to_string())
97        }
98    })
99}
100
101pub(super) fn declared<T: DeclaredFields>(value: T) -> T {
102    value.into_declared()
103}
104
105pub(super) fn claude_policy() -> TargetIdPolicy {
106    TargetIdPolicy::new(crate::Dialect::Claude)
107        .with_generated_prefix(IdentityRole::Response, KnownIdPrefix::Message)
108        .with_generated_prefix(IdentityRole::ToolCall, KnownIdPrefix::Tool)
109}
110
111pub(super) fn id(
112    flow: &mut IdentityFlow,
113    policy: &TargetIdPolicy,
114    dialect: crate::Dialect,
115    from: IdentityRole,
116    to: IdentityRole,
117    source: Option<String>,
118    index: u64,
119) -> Result<String, TransformError> {
120    if source.as_ref().is_some_and(String::is_empty) {
121        return Err(invalid("empty native identity"));
122    }
123    flow.resolve_as(
124        from,
125        to,
126        SourceIdentity::new(dialect, source, index),
127        policy,
128    )
129    .map(|v| v.emitted_id)
130    .map_err(|e| invalid(e.to_string()))
131}
132
133pub(super) fn item_id(item: &r::ResponseOutputItem) -> Option<&str> {
134    match item {
135        r::ResponseOutputItem::Message(v) => Some(&v.id),
136        r::ResponseOutputItem::Reasoning(v) => Some(&v.id),
137        r::ResponseOutputItem::FunctionCall(v) => v.id.as_deref(),
138        r::ResponseOutputItem::McpCall(v) => Some(&v.id),
139        r::ResponseOutputItem::ShellCall(v) => Some(&v.id),
140        r::ResponseOutputItem::ApplyPatchCall(v) => Some(&v.id),
141        r::ResponseOutputItem::ToolSearchCall(v) => Some(&v.id),
142        _ => None,
143    }
144}
145
146pub(super) fn normalize_arguments(
147    body: &mut r::GenerateContentResponseBody,
148) -> Result<(), TransformError> {
149    for item in &mut body.output {
150        let value = match item {
151            r::ResponseOutputItem::FunctionCall(v) => &mut v.arguments,
152            r::ResponseOutputItem::McpCall(v) => &mut v.arguments,
153            _ => continue,
154        };
155        let object: serde_json::Map<String, serde_json::Value> = serde_json::from_str(value)
156            .map_err(|e| invalid(format!("tool input requires an object: {e}")))?;
157        *value = serde_json::to_string(&object)?;
158    }
159    Ok(())
160}
161
162pub(super) fn initial_usage(mut usage: c::Usage) -> Result<c::Usage, TransformError> {
163    usage = declared(usage);
164    if usage.input_tokens < 0 || usage.output_tokens < 0 {
165        return Err(invalid("negative initial usage"));
166    }
167    for count in [
168        usage.cache_read_input_tokens.flatten(),
169        usage.cache_creation_input_tokens.flatten(),
170    ]
171    .into_iter()
172    .flatten()
173    {
174        if count < 0 {
175            return Err(invalid("negative initial cache count"));
176        }
177    }
178    if let Some(cache) = usage.cache_creation.as_ref().and_then(Option::as_ref) {
179        let total = cache
180            .ephemeral_1h_input_tokens
181            .checked_add(cache.ephemeral_5m_input_tokens)
182            .filter(|_| {
183                cache.ephemeral_1h_input_tokens >= 0 && cache.ephemeral_5m_input_tokens >= 0
184            })
185            .ok_or_else(|| invalid("invalid cache breakdown"))?;
186        if usage
187            .cache_creation_input_tokens
188            .flatten()
189            .is_some_and(|n| n != total)
190        {
191            return Err(invalid("cache breakdown contradicts aggregate"));
192        }
193        usage.cache_creation_input_tokens = Some(Some(total));
194    }
195    if usage
196        .output_tokens_details
197        .as_ref()
198        .and_then(Option::as_ref)
199        .is_some_and(|v| v.thinking_tokens < 0 || v.thinking_tokens > usage.output_tokens)
200    {
201        return Err(invalid("invalid initial thinking count"));
202    }
203    Ok(usage)
204}
gproxy_protocol/transform/generate/claude_responses/stream/common.rs

gproxy_protocol/transform/generate/claude_responses/stream/
common.rs
