Skip to main content

gproxy_protocol/adapt/guardian/
mod.rs

1//! Host-facing Guardian invocation. One prepared target request is sent once;
2//! non-2xx responses retain their complete original body and no retry is done.
3
4mod invoke_sse;
5mod resources;
6mod stream;
7
8use crate::{
9    HttpBody, WireResponse,
10    adapt::{JsonInvocation, invoke_json},
11    capability::Upstream,
12    codec::CodecLimits,
13    transform::{
14        TransformError,
15        guardian::{
16            self, GuardianClassifyResult, GuardianDialectRequest, GuardianExtraction,
17            GuardianOperation, GuardianPreparedRequest, GuardianReviewResult,
18        },
19    },
20};
21
22pub use invoke_sse::{classify_sse, review_sse};
23pub use resources::{GuardianResourceLimits, prepare_with_resources};
24pub use stream::{GuardianStreamContext, GuardianStreamInvocation, GuardianStreamLimits};
25
26#[derive(Debug, Clone, Copy)]
27pub struct GuardianLimits {
28    pub codec: CodecLimits,
29    pub max_request_bytes: u64,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct GuardianPreflight {
34    pub operation: GuardianOperation,
35    pub model: String,
36    pub max_bytes: u64,
37}
38
39#[derive(Debug)]
40#[allow(clippy::large_enum_variant)]
41pub enum GuardianReviewInvocation {
42    Success {
43        result: GuardianReviewResult,
44        native: GuardianNativeResponse,
45    },
46    Rejected(WireResponse<HttpBody>),
47    Invalid {
48        error: TransformError,
49        native: GuardianNativeResponse,
50    },
51}
52
53#[derive(Debug)]
54#[allow(clippy::large_enum_variant)]
55pub enum GuardianClassifyInvocation {
56    Success {
57        result: GuardianClassifyResult,
58        native: GuardianNativeResponse,
59    },
60    Rejected(WireResponse<HttpBody>),
61    Invalid {
62        error: TransformError,
63        native: GuardianNativeResponse,
64    },
65}
66
67/// The validated semantic label is paired with the complete generated target
68/// response. Hosts can use this provenance to map usage/IDs and synthesize the
69/// native Responses SSE lifecycle without treating the label as a bare JSON
70/// replacement for the client response.
71#[derive(Debug)]
72#[allow(clippy::large_enum_variant)]
73pub enum GuardianNativeResponse {
74    Claude(WireResponse<crate::wire::claude::generate_content::GenerateContentResponseBody>),
75    Gemini(WireResponse<crate::wire::gemini::GenerateContentResponseBody>),
76    Chat(WireResponse<crate::wire::openai::chat::response::GenerateContentResponseBody>),
77    Responses(WireResponse<crate::wire::openai::responses::GenerateContentResponseBody>),
78}
79
80pub async fn review<U: Upstream>(
81    upstream: &U,
82    target: &U::Target,
83    request: GuardianPreparedRequest,
84    limits: GuardianLimits,
85) -> Result<GuardianReviewInvocation, TransformError> {
86    if request.operation() != GuardianOperation::Review {
87        return Err(TransformError::new(
88            crate::transform::TransformErrorKind::Conflict,
89            "guardian.operation",
90            "review received a classify request",
91        ));
92    }
93    invoke(
94        upstream,
95        target,
96        request.into_request(),
97        GuardianOperation::Review,
98        limits,
99    )
100    .await
101    .map(map_review)
102}
103
104pub async fn classify<U: Upstream>(
105    upstream: &U,
106    target: &U::Target,
107    request: GuardianPreparedRequest,
108    limits: GuardianLimits,
109) -> Result<GuardianClassifyInvocation, TransformError> {
110    if request.operation() != GuardianOperation::Classify {
111        return Err(TransformError::new(
112            crate::transform::TransformErrorKind::Conflict,
113            "guardian.operation",
114            "classify received a review request",
115        ));
116    }
117    invoke(
118        upstream,
119        target,
120        request.into_request(),
121        GuardianOperation::Classify,
122        limits,
123    )
124    .await
125    .map(map_classify)
126}
127
128async fn invoke<U: Upstream>(
129    upstream: &U,
130    target: &U::Target,
131    request: GuardianDialectRequest,
132    operation: GuardianOperation,
133    limits: GuardianLimits,
134) -> Result<Extracted, TransformError> {
135    if limits.max_request_bytes == 0 {
136        return Err(TransformError::new(
137            crate::transform::TransformErrorKind::Limit,
138            "guardian.request",
139            "preflight byte cap must be positive",
140        ));
141    }
142    guardian::preflight_request(&request, limits.codec, limits.max_request_bytes)?;
143    match request {
144        GuardianDialectRequest::Claude(request) => {
145            match invoke_json(upstream, target, request, limits.codec).await? {
146                JsonInvocation::Success(response) => {
147                    Ok(GuardianExtractionInvocation::Claude(response))
148                }
149                JsonInvocation::Rejected(response) => {
150                    Ok(GuardianExtractionInvocation::Rejected(response))
151                }
152            }
153        }
154        GuardianDialectRequest::Gemini(request) => {
155            match invoke_json(upstream, target, request, limits.codec).await? {
156                JsonInvocation::Success(response) => {
157                    Ok(GuardianExtractionInvocation::Gemini(response))
158                }
159                JsonInvocation::Rejected(response) => {
160                    Ok(GuardianExtractionInvocation::Rejected(response))
161                }
162            }
163        }
164        GuardianDialectRequest::OpenAiChat(request) => {
165            match invoke_json(upstream, target, request, limits.codec).await? {
166                JsonInvocation::Success(response) => {
167                    Ok(GuardianExtractionInvocation::Chat(response))
168                }
169                JsonInvocation::Rejected(response) => {
170                    Ok(GuardianExtractionInvocation::Rejected(response))
171                }
172            }
173        }
174        GuardianDialectRequest::OpenAiResponses(request) => {
175            match invoke_json(upstream, target, request, limits.codec).await? {
176                JsonInvocation::Success(response) => {
177                    Ok(GuardianExtractionInvocation::Responses(response))
178                }
179                JsonInvocation::Rejected(response) => {
180                    Ok(GuardianExtractionInvocation::Rejected(response))
181                }
182            }
183        }
184    }
185    .and_then(|result| result.extract(operation))
186}
187
188#[allow(clippy::large_enum_variant)]
189enum GuardianExtractionInvocation {
190    Claude(WireResponse<crate::wire::claude::generate_content::GenerateContentResponseBody>),
191    Gemini(WireResponse<crate::wire::gemini::GenerateContentResponseBody>),
192    Chat(WireResponse<crate::wire::openai::chat::response::GenerateContentResponseBody>),
193    Responses(WireResponse<crate::wire::openai::responses::GenerateContentResponseBody>),
194    Rejected(WireResponse<HttpBody>),
195}
196
197enum Extracted {
198    Review {
199        result: GuardianReviewResult,
200        native: GuardianNativeResponse,
201    },
202    Classify {
203        result: GuardianClassifyResult,
204        native: GuardianNativeResponse,
205    },
206    Rejected(WireResponse<HttpBody>),
207    Invalid {
208        error: TransformError,
209        native: GuardianNativeResponse,
210    },
211}
212
213impl GuardianExtractionInvocation {
214    fn extract(self, operation: GuardianOperation) -> Result<Extracted, TransformError> {
215        match self {
216            Self::Claude(response) => map_response(
217                response,
218                operation,
219                guardian::extract_claude,
220                GuardianNativeResponse::Claude,
221            ),
222            Self::Gemini(response) => map_response(
223                response,
224                operation,
225                guardian::extract_gemini,
226                GuardianNativeResponse::Gemini,
227            ),
228            Self::Chat(response) => map_response(
229                response,
230                operation,
231                guardian::extract_chat,
232                GuardianNativeResponse::Chat,
233            ),
234            Self::Responses(response) => map_response(
235                response,
236                operation,
237                guardian::extract_responses,
238                GuardianNativeResponse::Responses,
239            ),
240            Self::Rejected(response) => Ok(Extracted::Rejected(response)),
241        }
242    }
243}
244
245fn map_response<T, F>(
246    response: WireResponse<T>,
247    operation: GuardianOperation,
248    extract: F,
249    native: impl FnOnce(WireResponse<T>) -> GuardianNativeResponse,
250) -> Result<Extracted, TransformError>
251where
252    F: FnOnce(T, GuardianOperation) -> Result<GuardianExtraction, TransformError>,
253    T: Clone,
254{
255    let WireResponse {
256        status,
257        headers,
258        body,
259    } = response;
260    let parsed = extract(body.clone(), operation);
261    let native = native(WireResponse {
262        status,
263        headers,
264        body,
265    });
266    match parsed {
267        Err(error) => Ok(Extracted::Invalid { error, native }),
268        Ok(GuardianExtraction::Review(value)) => Ok(Extracted::Review {
269            result: value,
270            native,
271        }),
272        Ok(GuardianExtraction::Classify(value)) => Ok(Extracted::Classify {
273            result: value,
274            native,
275        }),
276    }
277}
278
279fn map_review(value: Extracted) -> GuardianReviewInvocation {
280    match value {
281        Extracted::Review { result, native } => {
282            GuardianReviewInvocation::Success { result, native }
283        }
284        Extracted::Rejected(value) => GuardianReviewInvocation::Rejected(value),
285        Extracted::Invalid { error, native } => GuardianReviewInvocation::Invalid { error, native },
286        Extracted::Classify { .. } => unreachable!("review extraction returned classify result"),
287    }
288}
289
290fn map_classify(value: Extracted) -> GuardianClassifyInvocation {
291    match value {
292        Extracted::Classify { result, native } => {
293            GuardianClassifyInvocation::Success { result, native }
294        }
295        Extracted::Rejected(value) => GuardianClassifyInvocation::Rejected(value),
296        Extracted::Invalid { error, native } => {
297            GuardianClassifyInvocation::Invalid { error, native }
298        }
299        Extracted::Review { .. } => unreachable!("classify extraction returned review result"),
300    }
301}