gproxy-protocol 4.1.9

#[path = "adapt_generate_image_resources/host.rs"]
mod host;
#[allow(dead_code)]
#[path = "adapt_generate_stream/host.rs"]
mod http_host;
#[path = "adapt_generate_image_resources/read_cancel.rs"]
mod read_cancel;
use gproxy_protocol::{
    Dialect, HttpBody, WireResponse,
    adapt::generate::{
        gemini_responses::*,
        image_resources::*,
        stream::{event::EventLimits, reader::SourceFraming, *},
        *,
    },
    codec::CodecLimits,
    transform::{
        generate::gemini_responses as pair,
        identity::{IdNamespace, IdentityTarget},
    },
    wire::{gemini as g, openai::responses as r},
};
use host::*;
use http_host::{Feed, Host, Store, ready};
use serde_json::json;
use std::{
    future::Future,
    sync::{Arc, atomic::Ordering},
    task::{Context, Waker},
    time::{Duration, UNIX_EPOCH},
};
fn limits() -> CodecLimits {
    CodecLimits {
        max_buffer_bytes: 65536,
        max_value_bytes: 65536,
        max_body_bytes: 262144,
        max_line_bytes: 65536,
        max_part_bytes: 65536,
        max_parts: 64,
    }
}
fn settings() -> StreamSettings {
    StreamSettings {
        codec: limits(),
        source_framing: SourceFraming::Sse,
        client_framing: SourceFraming::Sse,
        events: EventLimits {
            max_bytes: 262144,
            max_pending_bytes: 65536,
        },
    }
}
fn state(store: &Store, dialect: Dialect) -> GenerationStateAccess<'_, Store> {
    GenerationStateAccess {
        store,
        scope: &(),
        target: IdentityTarget::new("selected", dialect)
            .unwrap()
            .with_origin("origin")
            .unwrap(),
        conversation_key: "image-conversation".into(),
        now: UNIX_EPOCH,
        expires_at: UNIX_EPOCH + Duration::from_secs(1000),
    }
}
fn resources(host: &Resources) -> GenerationResources<'_, Resources> {
    GenerationResources {
        access: host,
        scope: &1,
        limits: limits(),
        max_references: 8,
        now: UNIX_EPOCH,
        target: Dialect::OpenAi,
    }
}
fn ids(client: Dialect, native: Dialect) -> GenerationIdentity {
    GenerationIdentity::new(
        IdNamespace::with_bytes([121; 16]),
        IdNamespace::with_bytes([122; 16]),
        client,
        native,
    )
    .unwrap()
}
fn rrequest() -> r::GenerateContentRequestBody {
    serde_json::from_value(json!({"model":"client","input":"draw","tools":[{"type":"image_generation"}],"tool_choice":"auto"})).unwrap()
}
fn grequest() -> g::GenerateContentRequestBody {
    serde_json::from_value(json!({"contents":[{"role":"user","parts":[{"text":"draw"}]}],
        "generationConfig":{"responseModalities":["TEXT","IMAGE"],"responseFormat":{"image":{"delivery":"URI"}}}})).unwrap()
}
fn gresponse() -> g::GenerateContentResponseBody {
    serde_json::from_value(json!({"responseId":"g-image","modelVersion":"selected","candidates":[{"index":0,
        "content":{"role":"model","parts":[{"fileData":{"fileUri":"gs://private/generated.png","mimeType":"image/png"},"thoughtSignature":"original-signed-file"}]},"finishReason":"STOP"}],
        "usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2,"totalTokenCount":5,"cachedContentTokenCount":0,"thoughtsTokenCount":0,"toolUsePromptTokenCount":0}})).unwrap()
}
fn rresponse() -> r::GenerateContentResponseBody {
    serde_json::from_value(json!({"id":"r-image","object":"response","created_at":7,"model":"selected","status":"completed","error":null,"incomplete_details":null,
        "instructions":null,"metadata":null,"parallel_tool_calls":false,"temperature":null,"top_p":null,"tool_choice":"auto","tools":[{"type":"image_generation"}],
        "output":[{"type":"image_generation_call","id":"ig_source","result":PNG,"status":"completed"}],
        "usage":{"input_tokens":3,"input_tokens_details":{"cached_tokens":0,"cache_write_tokens":0},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":5}})).unwrap()
}
fn facts() -> GeminiReturnFacts {
    GeminiReturnFacts {
        parallel_tool_calls: false,
        tool_choice: r::ToolChoice::Mode(r::ToolChoiceMode::Auto),
        prompt_cache_options: None,
        usage: pair::GeminiUsageFacts {
            cache_write_tokens: Some(0),
            cached_tokens: None,
        },
        created_at: 7,
    }
}
fn host<N: serde::Serialize>(store: Arc<Store>, native: &N) -> Host {
    let host = Host::stream(store, Feed::default());
    *host.response.lock().unwrap() = Some(WireResponse {
        status: http::StatusCode::OK,
        headers: http::HeaderMap::new(),
        body: HttpBody::Bytes(serde_json::to_vec(native).unwrap().into()),
    });
    host
}
#[test]
fn signed_file_output_carries_its_signature_and_replays_the_client_bytes() {
    let store = Arc::new(Store::default());
    let state = state(&store, Dialect::Gemini);
    let resource_host = Resources::default();
    let resources = resources(&resource_host);
    let native = gresponse();
    let host = host(store.clone(), &native);
    let mut call = ResponsesViaGemini::prepare(
        rrequest(),
        "selected",
        Endpoint::new("/generateContent").unwrap(),
        ids(Dialect::OpenAi, Dialect::Gemini),
        Default::default(),
    )
    .unwrap();
    let mut progress = ImageResourceProgress::default();
    let GenerationOutcome::Success { response, .. } = ready(call.invoke_with_image_resources(
        &host,
        &(),
        limits(),
        &state,
        &resources,
        &mut progress,
        |_| Ok(facts()),
    ))
    .unwrap() else {
        panic!()
    };
    // The signature travels on the reasoning item before the image, and
    // the image carries the bytes the gateway read.
    let r::ResponseOutputItem::Reasoning(carrier) = &response.body.output[0] else {
        panic!()
    };
    assert_eq!(
        carrier.encrypted_content,
        Some(Some("gemini-next:original-signed-file".into()))
    );
    let r::ResponseOutputItem::ImageGenerationCall(image) = &response.body.output[1] else {
        panic!()
    };
    assert_eq!(image.result.as_deref(), Some(PNG));
    assert_eq!(resource_host.reads.load(Ordering::SeqCst), 1);
    assert!(
        store
            .entries
            .lock()
            .unwrap()
            .keys()
            .all(|key| !key.starts_with("generate:"))
    );
    let mut input = rrequest();
    input.input = Some(r::Input::Items(vec![
        r::InputItem::Reasoning(carrier.clone()),
        r::InputItem::ImageGenerationCall(image.clone()),
    ]));
    let empty = Store::default();
    let prepared = ready(ResponsesViaGemini::prepare_with_state(
        input.clone(),
        Endpoint::new("/generateContent").unwrap(),
        ids(Dialect::OpenAi, Dialect::Gemini),
        &self::state(&empty, Dialect::Gemini),
        Default::default(),
    ))
    .unwrap();
    // The client's bytes replay inline, with the carried signature, in place
    // of the fileData URI the client never saw.
    let restored = &prepared.target_request().contents[0]
        .parts
        .as_ref()
        .unwrap()[0];
    assert_eq!(
        restored.thought_signature.as_deref(),
        Some("original-signed-file")
    );
    assert_eq!(restored.inline_data.as_ref().unwrap().data, PNG);
    assert!(restored.file_data.is_none());
    ready(
        call.recover_with_image_resources(limits(), &state, &resources, &mut progress, |_| {
            Ok(facts())
        }),
    )
    .unwrap();
    assert_eq!(resource_host.reads.load(Ordering::SeqCst), 1);
    assert_eq!(host.sent.lock().unwrap().len(), 1);
}
#[test]
fn uri_publication_cancellation_queries_receipt_and_never_repeats_post_or_publish() {
    let store = Arc::new(Store::default());
    let state = state(&store, Dialect::OpenAi);
    let resource_host = Resources::default();
    resource_host.hang_publish.store(true, Ordering::SeqCst);
    let resources = resources(&resource_host);
    let native = rresponse();
    let host = host(store.clone(), &native);
    let mut call = ready(GeminiViaResponses::prepare_with_image_resources(
        grequest(),
        Endpoint::new("/responses").unwrap(),
        ids(Dialect::Gemini, Dialect::OpenAi),
        &state,
        &resources,
    ))
    .unwrap();
    assert!(call.convert_response(native, Default::default()).is_err());
    let mut progress = ImageResourceProgress::default();
    {
        let mut pending = Box::pin(call.invoke_with_image_resources(
            &host,
            &(),
            limits(),
            &state,
            &resources,
            &mut progress,
            |_| Ok(Default::default()),
        ));
        assert!(
            pending
                .as_mut()
                .poll(&mut Context::from_waker(Waker::noop()))
                .is_pending()
        );
        assert!(resource_host.hung.load(Ordering::SeqCst));
    }
    assert_eq!(progress.publications.operation_ids().count(), 1);
    let GenerationOutcome::Success { response, .. } = ready(call.recover_with_image_resources(
        limits(),
        &state,
        &resources,
        &mut progress,
        |_| Ok(Default::default()),
    ))
    .unwrap() else {
        panic!()
    };
    let image = &response.body.candidates.as_ref().unwrap()[0]
        .content
        .as_ref()
        .unwrap()
        .parts
        .as_ref()
        .unwrap()[0];
    assert!(image.inline_data.is_none());
    assert!(
        image
            .file_data
            .as_ref()
            .unwrap()
            .file_uri
            .starts_with("https://published.example/")
    );
    assert_eq!(resource_host.publishes.load(Ordering::SeqCst), 1);
    assert_eq!(resource_host.statuses.load(Ordering::SeqCst), 1);
    assert_eq!(host.sent.lock().unwrap().len(), 1);
    let wrong = GenerationResources {
        scope: &2,
        ..resources
    };
    assert!(
        ready(
            call.recover_with_image_resources(limits(), &state, &wrong, &mut progress, |_| Ok(
                Default::default()
            ))
        )
        .is_err()
    );
    assert_eq!(resource_host.publishes.load(Ordering::SeqCst), 1);
    resource_host.contradict_publication_metadata();
    let statuses = resource_host.statuses.load(Ordering::SeqCst);
    let receipts = progress.publications.receipts().len();
    for _ in 0..3 {
        assert!(
            ready(call.recover_with_image_resources(
                limits(),
                &state,
                &resources,
                &mut progress,
                |_| Ok(Default::default())
            ))
            .is_err()
        );
    }
    assert_eq!(resource_host.statuses.load(Ordering::SeqCst), statuses + 1);
    assert_eq!(progress.publications.receipts().len(), receipts + 1);
}
#[test]
fn file_image_stream_exposes_validated_bytes_before_eof_and_saves_nothing() {
    let store = Arc::new(Store::default());
    let state = state(&store, Dialect::Gemini);
    let resource_host = Resources::default();
    let resources = resources(&resource_host);
    let request = rrequest();
    let f = facts();
    let mut call = ready(ResponsesViaGemini::prepare_stream(
        request.clone(),
        StreamTarget {
            endpoint: Endpoint::new("/streamGenerateContent").unwrap(),
            identities: ids(Dialect::OpenAi, Dialect::Gemini),
        },
        ResponsesViaGeminiStreamFacts {
            request: Default::default(),
            response: pair::stream::GeminiToResponsesContext {
                response: pair::GeminiResponseContext {
                    request,
                    effective_parallel_tool_calls: f.parallel_tool_calls,
                    effective_tool_choice: f.tool_choice,
                    effective_prompt_cache_options: f.prompt_cache_options,
                    usage: f.usage,
                    created_at: f.created_at,
                },
                actual_model: Some("selected".into()),
                final_thinking_tokens: None,
            },
        },
        settings(),
        &state,
    ))
    .unwrap();
    let feed = Feed::default();
    let mut prefix = gresponse();
    prefix.candidates.as_mut().unwrap()[0].finish_reason = None;
    prefix.usage_metadata = None;
    feed.push(format!(
        "data: {}\n\n",
        serde_json::to_string(&prefix).unwrap()
    ));
    let host = Host::stream(store.clone(), feed.clone());
    ready(call.start(&host, &(), &state)).unwrap();
    let mut progress = ImageStreamProgress::default();
    #[cfg(not(target_arch = "wasm32"))]
    {
        fn assert_send<T: Send>(_: &T) {}
        let normal = call.next(&state);
        assert_send(&normal);
        drop(normal);
        let mapped = call.next_with_image_resources(&state, &resources, &mut progress);
        assert_send(&mapped);
        drop(mapped);
    }
    loop {
        let chunk = ready(call.next_with_image_resources(&state, &resources, &mut progress))
            .unwrap()
            .unwrap();
        if String::from_utf8_lossy(&chunk.bytes).contains(PNG) {
            break;
        }
    }
    assert!(call.client_result().is_none());
    assert_eq!(resource_host.reads.load(Ordering::SeqCst), 1);
    let mut terminal = gresponse();
    terminal.candidates.as_mut().unwrap()[0].content = None;
    feed.push(format!(
        "data: {}\n\n",
        serde_json::to_string(&terminal).unwrap()
    ));
    feed.close();
    while let Some(chunk) =
        ready(call.next_with_image_resources(&state, &resources, &mut progress)).unwrap()
    {
        if matches!(chunk.event, Some(r::stream::StreamEvent::Completed(_))) {
            // The signature went to the client; nothing is kept for it.
            assert!(
                store
                    .entries
                    .lock()
                    .unwrap()
                    .keys()
                    .all(|key| !key.starts_with("generate:"))
            );
        }
    }
    let original = &call.native_result().unwrap().candidates.as_ref().unwrap()[0]
        .content
        .as_ref()
        .unwrap()
        .parts
        .as_ref()
        .unwrap()[0];
    assert!(original.file_data.is_some());
    assert!(original.inline_data.is_none());
    assert_eq!(
        original.thought_signature.as_deref(),
        Some("original-signed-file")
    );
    assert_eq!(host.sent.lock().unwrap().len(), 1);
}
#[test]
fn uri_stream_retains_pending_client_event_across_uncertain_publication() {
    let store = Arc::new(Store::default());
    let state = state(&store, Dialect::OpenAi);
    let resource_host = Resources::default();
    resource_host.hang_publish.store(true, Ordering::SeqCst);
    let resources = resources(&resource_host);
    let mut call = ready(GeminiViaResponses::prepare_stream_with_capabilities(
        grequest(),
        StreamTarget {
            endpoint: Endpoint::new("/responses").unwrap(),
            identities: ids(Dialect::Gemini, Dialect::OpenAi),
        },
        Default::default(),
        settings(),
        &state,
        &resources,
    ))
    .unwrap();
    let feed = Feed::default();
    let events =
        gproxy_protocol::transform::generate::stream::responses::synthesize_responses_stream(
            rresponse(),
            &mut gproxy_protocol::transform::identity::IdentityFlow::new(IdNamespace::with_bytes(
                [125; 16],
            )),
            Default::default(),
        )
        .unwrap()
        .value;
    for event in events {
        feed.push(format!(
            "data: {}\n\n",
            serde_json::to_string(&event).unwrap()
        ));
    }
    let host = Host::stream(store.clone(), feed.clone());
    ready(call.start(&host, &(), &state)).unwrap();
    assert!(
        ready(call.next(&state)).is_err(),
        "ordinary next must not bypass URI publication"
    );
    let mut progress = ImageStreamProgress::default();
    for _ in 0..10000 {
        let mut pending =
            Box::pin(call.next_with_image_resources(&state, &resources, &mut progress));
        match pending
            .as_mut()
            .poll(&mut Context::from_waker(Waker::noop()))
        {
            std::task::Poll::Ready(Ok(Some(chunk))) => {
                assert!(!String::from_utf8_lossy(&chunk.bytes).contains(PNG))
            }
            std::task::Poll::Pending => {}
            _ => panic!("unexpected stream outcome before image publication"),
        }
        if resource_host.hung.load(Ordering::SeqCst) {
            break;
        }
    }
    assert!(resource_host.hung.load(Ordering::SeqCst));
    let mut fresh = ImageStreamProgress::default();
    assert!(ready(call.next_with_image_resources(&state, &resources, &mut fresh)).is_err());
    assert_eq!(resource_host.publishes.load(Ordering::SeqCst), 1);
    let chunk = ready(call.next_with_image_resources(&state, &resources, &mut progress))
        .unwrap()
        .unwrap();
    assert!(String::from_utf8_lossy(&chunk.bytes).contains("https://published.example/"));
    feed.close();
    while ready(call.next_with_image_resources(&state, &resources, &mut progress))
        .unwrap()
        .is_some()
    {}
    assert!(call.client_result().is_some());
    assert_eq!(resource_host.publishes.load(Ordering::SeqCst), 1);
    assert_eq!(resource_host.statuses.load(Ordering::SeqCst), 1);
    assert_eq!(host.sent.lock().unwrap().len(), 1);
}