Skip to main content

gproxy_protocol/adapt/guardian/
invoke_sse.rs

1use super::{
2    GuardianLimits, GuardianStreamContext, GuardianStreamInvocation, GuardianStreamLimits,
3};
4use crate::{
5    capability::Upstream,
6    transform::{TransformError, guardian::GuardianPreparedRequest, identity::IdentityFlow},
7};
8
9/// Run one prepared review request and return the client's native Responses SSE
10/// body. Return controls are taken from the prepared client request.
11pub async fn review_sse<U: Upstream>(
12    upstream: &U,
13    target: &U::Target,
14    request: GuardianPreparedRequest,
15    context: GuardianStreamContext,
16    flow: &mut IdentityFlow,
17    limits: GuardianLimits,
18    stream_limits: GuardianStreamLimits,
19) -> Result<GuardianStreamInvocation, TransformError> {
20    let context = source_context(&request, context)?;
21    super::review(upstream, target, request, limits)
22        .await?
23        .into_stream(context, flow, stream_limits)
24}
25
26pub async fn classify_sse<U: Upstream>(
27    upstream: &U,
28    target: &U::Target,
29    request: GuardianPreparedRequest,
30    context: GuardianStreamContext,
31    flow: &mut IdentityFlow,
32    limits: GuardianLimits,
33    stream_limits: GuardianStreamLimits,
34) -> Result<GuardianStreamInvocation, TransformError> {
35    let context = source_context(&request, context)?;
36    super::classify(upstream, target, request, limits)
37        .await?
38        .into_stream(context, flow, stream_limits)
39}
40
41fn source_context(
42    request: &GuardianPreparedRequest,
43    context: GuardianStreamContext,
44) -> Result<GuardianStreamContext, TransformError> {
45    let source = request.response_request()?;
46    Ok(match context {
47        GuardianStreamContext::Responses => GuardianStreamContext::Responses,
48        GuardianStreamContext::Chat(mut context) => {
49            context.request = source;
50            GuardianStreamContext::Chat(context)
51        }
52        GuardianStreamContext::Claude(mut context) => {
53            context.request = source;
54            GuardianStreamContext::Claude(context)
55        }
56        GuardianStreamContext::Gemini(mut context) => {
57            context.request = source;
58            GuardianStreamContext::Gemini(context)
59        }
60    })
61}