From 911012490ca9b6529dfaf14b0396363dce981169 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 13 Aug 2026 14:14:56 +0000 Subject: [PATCH] Return Luna samples when streamed JSON completes (#38383) ## Why Structured output can be complete before the Responses stream emits its terminal events. Waiting for those events unnecessarily delays the sampler and can leave it waiting after usable output has arrived. ## What changed - Return a Luna sample as soon as accumulated text deltas parse as a complete JSON object. - Continue draining the remaining response events in the background so the authenticated WebSocket stays reusable for subsequent samples. - Enforce the output-size limit while accumulating deltas on the early-return path. ## Testing - Cover returning complete JSON without terminal response events. - Verify that early return still allows the authenticated WebSocket to be reused. GitOrigin-RevId: 80c9963f1540f7862356aaad3503e3071a8fb92a --- codex-rs/ext/guardian-v2/src/sampler.rs | 21 ++++++- codex-rs/ext/guardian-v2/src/sampler_tests.rs | 62 +++++++++++++++++++ 2 files changed, 82 insertions(+), 1 deletion(-) diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index 4d652077a9..b9e55e97c1 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -242,7 +242,26 @@ impl LunaSampler { let mut deltas = String::new(); while let Some(event) = stream.rx_event.recv().await { match event.map_err(LunaSamplerError::Api)? { - ResponseEvent::OutputTextDelta(delta) => deltas.push_str(&delta), + ResponseEvent::OutputTextDelta(delta) => { + deltas.push_str(&delta); + if deltas.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); + } + + if serde_json::from_str::>(&deltas).is_ok() { + let mut remaining_events = stream.rx_event; + tokio::spawn(async move { + // The transport needs a live consumer until completion to keep its + // authenticated WebSocket available for the next sample. + while let Some(event) = remaining_events.recv().await { + if matches!(event, Ok(ResponseEvent::Completed { .. }) | Err(_)) { + break; + } + } + }); + return Ok(deltas); + } + } ResponseEvent::OutputItemDone(ResponseItem::Message { role, content, .. }) if role == "assistant" => { diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index 1cdfdffa58..b2ee19e039 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -9,11 +9,15 @@ use codex_model_provider_info::ModelProviderInfo; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; use core_test_support::responses; +use core_test_support::responses::WebSocketConnectionConfig; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_model_verification_metadata; +use core_test_support::responses::ev_output_text_delta; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use serde_json::json; +use std::time::Duration; use super::LunaSampler; use super::LunaSamplerConfig; @@ -26,6 +30,9 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ let server = responses::start_websocket_server(vec![vec![ vec![ + ev_output_text_delta(r#"{"score":0"#), + ev_output_text_delta(".25}"), + ev_model_verification_metadata("response-1", vec!["trusted_access_for_cyber"]), ev_assistant_message("sample-1", r#"{"score":0.25}"#), ev_completed("response-1"), ], @@ -126,3 +133,58 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_returns_complete_json_before_terminal_response_events() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![vec![ + ev_output_text_delta(r#"{"score":0"#), + ev_output_text_delta(".25}"), + ]], + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: false, + }]) + .await; + let base_url = server.uri().replacen("ws://", "http://", 1); + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(Some(format!("{base_url}/v1"))), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "test-api-key", + ))), + ); + let sampler = LunaSampler::connect(LunaSamplerConfig { + provider, + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly, + session_source: SessionSource::Exec, + session_id: "session-1".to_owned(), + thread_id: "thread-1".to_owned(), + originator: None, + service_tier: None, + }) + .await?; + + let output = tokio::time::timeout( + Duration::from_secs(2), + sampler.sample(LunaSamplingRequest { + instructions: "Return a risk score.".to_owned(), + input: "The user requested a README summary.".to_owned(), + output_schema: json!({ + "type": "object", + "properties": { "score": { "type": "number" } }, + "required": ["score"], + "additionalProperties": false + }), + reasoning_effort: ReasoningEffort::None, + turn_id: "turn-1".to_owned(), + }), + ) + .await??; + + assert_eq!(output, r#"{"score":0.25}"#); + server.shutdown().await; + Ok(()) +}