codex/codex-rs/core/tests/suite/pending_input.rs

use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_message_item_added;
use core_test_support::responses::ev_output_text_delta;
use core_test_support::responses::ev_response_created;
use core_test_support::streaming_sse::StreamingSseChunk;
use core_test_support::streaming_sse::start_streaming_sse_server;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::Value;
use tokio::sync::oneshot;
use tokio::time::Duration;
use tokio::time::timeout;

fn ev_message_item_done(id: &str, text: &str) -> Value {
    serde_json::json!({
        "type": "response.output_item.done",
        "item": {
            "type": "message",
            "role": "assistant",
            "id": id,
            "content": [{"type": "output_text", "text": text}]
        }
    })
}

fn sse_event(event: Value) -> String {
    responses::sse(vec![event])
}

fn message_input_texts(body: &Value, role: &str) -> Vec<String> {
    body.get("input")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .filter(|item| item.get("type").and_then(Value::as_str) == Some("message"))
        .filter(|item| item.get("role").and_then(Value::as_str) == Some(role))
        .filter_map(|item| item.get("content").and_then(Value::as_array))
        .flatten()
        .filter(|span| span.get("type").and_then(Value::as_str) == Some("input_text"))
        .filter_map(|span| span.get("text").and_then(Value::as_str).map(str::to_owned))
        .collect()
}

fn has_user_text(body: &Value, text: &str) -> bool {
    message_input_texts(body, "user")
        .iter()
        .any(|candidate| candidate == text)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn injected_user_input_triggers_follow_up_request_with_deltas() {
    let (gate_completed_tx, gate_completed_rx) = oneshot::channel();

    let first_chunks = vec![
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_response_created("resp-1")),
        },
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_message_item_added("msg-1", "")),
        },
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_output_text_delta("first ")),
        },
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_output_text_delta("turn")),
        },
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_message_item_done("msg-1", "first turn")),
        },
        StreamingSseChunk {
            gate: Some(gate_completed_rx),
            body: sse_event(ev_completed("resp-1")),
        },
    ];

    let second_chunks = vec![
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_response_created("resp-2")),
        },
        StreamingSseChunk {
            gate: None,
            body: sse_event(ev_completed("resp-2")),
        },
    ];

    let (server, completions) = start_streaming_sse_server(vec![first_chunks, second_chunks]).await;
    let mut completions = completions.into_iter();

    let codex = test_codex()
        .with_model("gpt-5.1")
        .build_with_streaming_server(&server)
        .await
        .unwrap()
        .codex;

    codex
        .submit(Op::UserInput {
            items: vec![UserInput::Text {
                text: "first prompt".into(),
                text_elements: Vec::new(),
            }],
            final_output_json_schema: None,
        })
        .await
        .unwrap();

    wait_for_event(&codex, |event| {
        matches!(event, EventMsg::AgentMessageContentDelta(_))
    })
    .await;
    eprintln!("pending input probe observed first assistant delta");

    codex
        .submit(Op::UserInput {
            items: vec![UserInput::Text {
                text: "second prompt".into(),
                text_elements: Vec::new(),
            }],
            final_output_json_schema: None,
        })
        .await
        .unwrap();
    eprintln!("pending input probe injected second user input");
    timeout(Duration::from_secs(5), async {
        while !codex.has_pending_input().await {
            tokio::task::yield_now().await;
        }
    })
    .await
    .expect("timed out waiting for second user input to become pending");
    eprintln!("pending input probe observed second user input queued as pending");

    let _ = gate_completed_tx.send(());
    eprintln!("pending input probe released first response completion gate");

    let first_completion = completions
        .next()
        .expect("missing first response stream completion handle");
    timeout(Duration::from_secs(5), first_completion)
        .await
        .expect("timed out waiting for first response stream completion")
        .expect("first response stream closed before completion");
    eprintln!("pending input probe observed first response stream completion");

    let second_completion = completions
        .next()
        .expect("missing follow-up response stream completion handle");
    timeout(Duration::from_secs(5), second_completion)
        .await
        .expect("timed out waiting for follow-up response stream completion")
        .expect("follow-up response stream closed before completion");
    eprintln!("pending input probe observed follow-up response stream completion");

    let requests = timeout(Duration::from_secs(5), async {
        loop {
            let requests = server.requests().await;
            if requests.len() >= 2 {
                break requests;
            }
            tokio::task::yield_now().await;
        }
    })
    .await
    .expect("timed out waiting for follow-up request capture");
    eprintln!("pending input probe captured {} requests", requests.len());
    assert_eq!(requests.len(), 2);

    let first_body: Value = serde_json::from_slice(&requests[0]).expect("parse first request");
    let second_body: Value = serde_json::from_slice(&requests[1]).expect("parse second request");
    eprintln!("pending input probe request[0] body: {first_body}");
    eprintln!("pending input probe request[1] body: {second_body}");

    let request_bodies = [&first_body, &second_body];
    let initial_request_matches = request_bodies
        .iter()
        .filter(|body| has_user_text(body, "first prompt") && !has_user_text(body, "second prompt"))
        .count();
    let follow_up_request_matches = request_bodies
        .iter()
        .filter(|body| has_user_text(body, "first prompt") && has_user_text(body, "second prompt"))
        .count();

    assert_eq!(
        initial_request_matches, 1,
        "expected exactly one initial request with only the first prompt, bodies: {request_bodies:?}"
    );
    assert_eq!(
        follow_up_request_matches, 1,
        "expected exactly one follow-up request with both prompts, bodies: {request_bodies:?}"
    );

    server.shutdown().await;
}