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

@@ -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::<serde_json::Map<String, Value>>(&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" =>
                 {

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