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(()) +}