gproxy_protocol/adapt/generate/
transport.rs

1use crate::{
2    Dialect, HttpBody, WireRequest, WireResponse,
3    capability::Upstream,
4    codec::{self, CodecLimits},
5    transform::{
6        Converted, Report, TransformError, TransformErrorKind,
7        identity::{IdNamespace, IdentityFlow, IdentityRole, KnownIdPrefix, TargetIdPolicy},
8    },
9    wire::DeclaredFields,
10};
11use serde::{Serialize, de::DeserializeOwned};
12
13/// Host-selected operation path/query and headers. Authentication stays in Upstream.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct Endpoint {
16    pub path: String,
17    pub query: Option<String>,
18    pub headers: http::HeaderMap,
19}
20
21impl Endpoint {
22    pub fn new(path: impl Into<String>) -> Result<Self, TransformError> {
23        let value = Self {
24            path: path.into(),
25            query: None,
26            headers: Default::default(),
27        };
28
29        Ok(value)
30    }
31}
32
33/// Distinct invocation namespaces prevent request-history identities from colliding
34/// with response output positions. The host reuses these namespaces on recovery.
35#[derive(Debug, Clone)]
36pub struct GenerationIdentity {
37    pub request: IdentityFlow,
38    pub response: IdentityFlow,
39    pub request_policy: TargetIdPolicy,
40    pub response_policy: TargetIdPolicy,
41}
42
43impl GenerationIdentity {
44    pub fn new(
45        request: IdNamespace,
46        response: IdNamespace,
47        client: Dialect,
48        upstream: Dialect,
49    ) -> Result<Self, TransformError> {
50        if request == response {
51            return Err(TransformError::shape(
52                "identity.namespace",
53                "request and response namespaces must differ",
54            ));
55        }
56        Ok(Self {
57            request: IdentityFlow::new(request),
58            response: IdentityFlow::new(response),
59            request_policy: generation_policy(upstream),
60            response_policy: generation_policy(client).with_reversible_tool_calls(),
61        })
62    }
63}
64
65/// Kept by the caller if an await is cancelled or response conversion fails.
66/// A started send is never implicitly retried. Native bytes are retained before decoding.
67#[derive(Debug)]
68pub struct GenerationProgress<N> {
69    pub send_started: bool,
70    pub(super) exposed_response: Option<bytes::Bytes>,
71    pub(super) binding: Option<InvocationBinding>,
72    pub raw_response: Option<WireResponse<bytes::Bytes>>,
73    pub native_response: Option<N>,
74}
75
76impl<N> Default for GenerationProgress<N> {
77    fn default() -> Self {
78        Self {
79            exposed_response: None,
80            binding: None,
81            send_started: false,
82            raw_response: None,
83            native_response: None,
84        }
85    }
86}
87
88#[derive(Debug)]
89pub enum GenerationOutcome<C> {
90    Success {
91        response: WireResponse<C>,
92        report: Report,
93    },
94    Rejected(WireResponse<bytes::Bytes>),
95}
96
97pub(super) async fn send<
98    U: Upstream,
99    I: Serialize + DeclaredFields,
100    N: DeserializeOwned + DeclaredFields + Clone,
101>(
102    upstream: &U,
103    target: &U::Target,
104    endpoint: Endpoint,
105    body: I,
106    limits: CodecLimits,
107    progress: &mut GenerationProgress<N>,
108) -> Result<bool, TransformError> {
109    if progress.send_started
110        || progress.raw_response.is_some()
111        || progress.native_response.is_some()
112    {
113        return Err(TransformError::new(
114            TransformErrorKind::Conflict,
115            "generation.send",
116            "progress already records a send; recover its result instead of replaying",
117        ));
118    }
119
120    let host = upstream.limits();
121    let bytes = codec::encode_json(&body.into_declared(), bound(limits, host.write_bytes))
122        .map_err(|e| {
123            TransformError::with_source(
124                if e.kind() == codec::CodecErrorKind::Limit {
125                    TransformErrorKind::Limit
126                } else {
127                    TransformErrorKind::InvalidInput
128                },
129                "request.body",
130                e.to_string(),
131                e,
132            )
133        })?;
134    let mut headers = endpoint.headers;
135    json_headers(&mut headers);
136    progress.send_started = true;
137    let response = upstream
138        .send(
139            target,
140            WireRequest {
141                method: http::Method::POST,
142                path: endpoint.path,
143                query: endpoint.query,
144                headers,
145                body: HttpBody::Bytes(bytes),
146            },
147        )
148        .await?;
149    let WireResponse {
150        status,
151        headers,
152        body,
153    } = response;
154    let bytes = codec::read_http_body(body, bound(limits, host.read_bytes))
155        .await
156        .map_err(codec_error)?;
157    progress.raw_response = Some(WireResponse {
158        status,
159        headers,
160        body: bytes,
161    });
162    if !status.is_success() {
163        return Ok(false);
164    }
165    let raw = progress.raw_response.as_ref().expect("just retained");
166    let native: N =
167        codec::decode_json(&raw.body, bound(limits, host.read_bytes)).map_err(codec_error)?;
168    progress.native_response = Some(native.into_declared());
169    Ok(true)
170}
171
172pub(super) fn finish<C: Serialize, N>(
173    progress: &mut GenerationProgress<N>,
174    converted: Converted<C>,
175    mut report: Report,
176    limits: CodecLimits,
177) -> Result<GenerationOutcome<C>, TransformError> {
178    let encoded = codec::encode_json(&converted.value, limits).map_err(codec_error)?;
179    if progress
180        .exposed_response
181        .as_ref()
182        .is_some_and(|previous| previous != &encoded)
183    {
184        return Err(TransformError::new(
185            TransformErrorKind::Conflict,
186            "generation.recovery",
187            "client response changed after first exposure",
188        ));
189    }
190    progress.exposed_response = Some(encoded);
191    let raw = progress
192        .raw_response
193        .as_ref()
194        .expect("successful send retained");
195    let mut headers = raw.headers.clone();
196    json_headers(&mut headers);
197    report.diagnostics.extend(converted.report.diagnostics);
198    Ok(GenerationOutcome::Success {
199        response: WireResponse {
200            status: raw.status,
201            headers,
202            body: converted.value,
203        },
204        report,
205    })
206}
207
208pub(super) fn rejected<C, N>(progress: &GenerationProgress<N>) -> GenerationOutcome<C> {
209    let raw = progress
210        .raw_response
211        .as_ref()
212        .expect("rejected send retained");
213    GenerationOutcome::Rejected(WireResponse {
214        status: raw.status,
215        headers: raw.headers.clone(),
216        body: raw.body.clone(),
217    })
218}
219
220fn json_headers(headers: &mut http::HeaderMap) {
221    for name in [
222        http::header::CONTENT_LENGTH,
223        http::header::CONTENT_ENCODING,
224        http::header::TRANSFER_ENCODING,
225    ] {
226        headers.remove(name);
227    }
228    headers.insert(
229        http::header::CONTENT_TYPE,
230        http::HeaderValue::from_static("application/json"),
231    );
232}
233
234fn bound(mut limits: CodecLimits, bytes: u64) -> CodecLimits {
235    limits.max_body_bytes = limits.max_body_bytes.min(bytes);
236    limits.max_buffer_bytes = limits.max_buffer_bytes.min(bytes);
237    limits.max_value_bytes = limits.max_value_bytes.min(bytes);
238    limits
239}
240
241fn codec_error(error: codec::CodecError) -> TransformError {
242    let kind = match error.kind() {
243        codec::CodecErrorKind::Limit => TransformErrorKind::Limit,
244        codec::CodecErrorKind::Transport => TransformErrorKind::Host,
245        _ => TransformErrorKind::InvalidResult,
246    };
247    TransformError::with_source(kind, "upstream.response", error.to_string(), error)
248}
249
250#[derive(Debug, Clone, PartialEq, Eq)]
251pub(super) struct InvocationBinding {
252    endpoint: Endpoint,
253    selected_model: String,
254    request_namespace: IdNamespace,
255    response_namespace: IdNamespace,
256    request_policy: TargetIdPolicy,
257    response_policy: TargetIdPolicy,
258    original: bytes::Bytes,
259    prepared: bytes::Bytes,
260    state_target: crate::transform::identity::IdentityTarget,
261    conversation: String,
262    expiry: std::time::SystemTime,
263}
264
265pub(super) fn bind<
266    S: crate::capability::StateStore,
267    C: Serialize + DeclaredFields + Clone,
268    I: Serialize + DeclaredFields + Clone,
269    N,
270>(
271    endpoint: (&Endpoint, &str),
272    identities: &GenerationIdentity,
273    requests: (&C, &I),
274    state: &super::GenerationStateAccess<'_, S>,
275    limits: CodecLimits,
276    progress: &mut GenerationProgress<N>,
277    recovery: bool,
278) -> Result<(), TransformError> {
279    if recovery && state.now >= state.expires_at {
280        return Err(TransformError::shape(
281            "generation.recovery",
282            "invocation state expired",
283        ));
284    }
285    let (endpoint, model) = endpoint;
286    let (original, prepared) = requests;
287    let value = InvocationBinding {
288        endpoint: endpoint.clone(),
289        selected_model: model.into(),
290        request_namespace: identities.request.namespace(),
291        response_namespace: identities.response.namespace(),
292        request_policy: identities.request_policy.clone(),
293        response_policy: identities.response_policy.clone(),
294        original: codec::encode_json(&original.clone().into_declared(), limits)
295            .map_err(codec_error)?,
296        prepared: codec::encode_json(&prepared.clone().into_declared(), limits)
297            .map_err(codec_error)?,
298        state_target: state.target.clone(),
299        conversation: state.conversation_key.clone(),
300        expiry: state.expires_at,
301    };
302    if let Some(existing) = &progress.binding {
303        if existing != &value {
304            return Err(TransformError::new(
305                TransformErrorKind::Conflict,
306                "generation.recovery",
307                "prepared request or scope binding changed",
308            ));
309        }
310    } else if recovery {
311        return Err(TransformError::new(
312            TransformErrorKind::MissingState,
313            "generation.recovery",
314            "progress is not bound to this invocation",
315        ));
316    } else {
317        progress.binding = Some(value);
318    }
319    Ok(())
320}
321
322pub(super) fn recover_native<N: DeserializeOwned + DeclaredFields + Clone>(
323    progress: &mut GenerationProgress<N>,
324    limits: CodecLimits,
325) -> Result<N, TransformError> {
326    let raw = progress.raw_response.as_ref().ok_or_else(|| {
327        TransformError::new(
328            TransformErrorKind::MissingState,
329            "generation.recovery",
330            "no raw upstream response retained",
331        )
332    })?;
333    if !raw.status.is_success() {
334        return Err(TransformError::invalid_result(
335            "generation.recovery",
336            "success response required",
337        ));
338    }
339    let native: N = codec::decode_json(&raw.body, limits).map_err(codec_error)?;
340    let native = native.into_declared();
341    progress.native_response = Some(native.clone());
342    Ok(native)
343}
344
345pub(super) fn generation_policy(dialect: Dialect) -> TargetIdPolicy {
346    let policy = TargetIdPolicy::new(dialect);
347    if dialect == Dialect::Claude {
348        policy
349            .with_generated_prefix(IdentityRole::Response, KnownIdPrefix::Message)
350            .with_generated_prefix(IdentityRole::ToolCall, KnownIdPrefix::Tool)
351    } else {
352        policy
353    }
354}
gproxy_protocol/adapt/generate/transport.rs

gproxy_protocol/adapt/generate/
transport.rs
