kcode-k1-chat-codex-state 0.2.0

use kcode_k1_chat_codex_codec::{BoxValue, Call};
use kcode_k1_chat_codex_state::{Arrival, ConversationState, RestartError, Start, Status};
use kcode_k1_chat_state::BoxContent;
use kcode_k1_codex_adapter::{ShimItem, ShimOutput};

fn call(name: &str) -> Call {
    Call {
        name: name.into(),
        arguments: "{}".into(),
    }
}

fn start(state: &mut ConversationState) -> Start {
    state.begin().unwrap().unwrap()
}

#[test]
fn stages_are_ordered_and_canonical_before_done() {
    let mut state = ConversationState::new([1; 12]);
    state.accept(Arrival::User("ask".into())).unwrap();
    let round = start(&mut state);
    let prepared = state
        .prepare_stage(
            round.job,
            "thinking".into(),
            vec![
                BoxValue::Call(Ok(call("one"))),
                BoxValue::Call(Ok(call("two"))),
            ],
        )
        .unwrap();
    assert_eq!(
        prepared
            .iter()
            .map(|call| call.name.as_str())
            .collect::<Vec<_>>(),
        ["one", "two"]
    );
    assert!(
        matches!(state.boxes()[1].content(), BoxContent::Kennedy { text } if text == "thinking")
    );
    assert!(matches!(
        state.boxes()[2].content(),
        BoxContent::KtoolCall { .. }
    ));
    state
        .complete(
            round.job,
            ShimOutput {
                items: vec![ShimItem::Text("done".into())],
            },
        )
        .unwrap();
    assert!(matches!(state.boxes()[4].content(), BoxContent::Kennedy { text } if text == "done"));
}

#[test]
fn queued_arrivals_are_fifo_and_calls_do_not_schedule() {
    let mut state = ConversationState::new([2; 12]);
    state.accept(Arrival::User("first".into())).unwrap();
    let round = start(&mut state);
    state
        .prepare_stage(
            round.job,
            String::new(),
            vec![BoxValue::Call(Ok(call("one")))],
        )
        .unwrap();
    state.accept(Arrival::System("s".into())).unwrap();
    state.accept(Arrival::User("u".into())).unwrap();
    state
        .complete(round.job, ShimOutput { items: vec![] })
        .unwrap();
    let next = start(&mut state);
    assert_eq!(next.boxes.len(), 3);
    assert!(state.begin().unwrap().is_none());
}

#[test]
fn malformed_stage_is_transactional() {
    let mut state = ConversationState::new([3; 12]);
    state.accept(Arrival::User("ask".into())).unwrap();
    let round = start(&mut state);
    assert!(
        state
            .prepare_stage(
                round.job,
                "text".into(),
                vec![BoxValue::History("bad".into())]
            )
            .is_err()
    );
    assert_eq!(state.boxes().len(), 1);
    state
        .prepare_stage(
            round.job,
            String::new(),
            vec![BoxValue::Call(Ok(call("one")))],
        )
        .unwrap();
    assert_eq!(
        state
            .boxes()
            .iter()
            .filter(|box_| matches!(box_.content(), BoxContent::KtoolCall { .. }))
            .count(),
        1
    );
}

#[test]
fn retry_is_only_safe_before_calls() {
    let mut state = ConversationState::new([4; 12]);
    state.accept(Arrival::User("ask".into())).unwrap();
    let round = start(&mut state);
    state.fail(round.job, "early".into(), true);
    assert!(matches!(
        state.status(),
        Status::Stalled {
            restartable: true,
            ..
        }
    ));
    state.restart().unwrap();
    assert_eq!(start(&mut state).boxes.len(), 1);

    let mut state = ConversationState::new([5; 12]);
    state.accept(Arrival::User("ask".into())).unwrap();
    let round = start(&mut state);
    state
        .prepare_stage(
            round.job,
            String::new(),
            vec![BoxValue::Call(Ok(call("one")))],
        )
        .unwrap();
    state.fail(round.job, "late".into(), true);
    assert_eq!(state.restart(), Err(RestartError::NotRestartable));
}

#[test]
fn recovery_reseeds_and_reuses_session_sequence() {
    let session = [6; 12];
    let mut state = ConversationState::new(session);
    state.accept(Arrival::User("ask".into())).unwrap();
    let round = start(&mut state);
    let first_call = state
        .prepare_stage(
            round.job,
            String::new(),
            vec![BoxValue::Call(Ok(call("one")))],
        )
        .unwrap()
        .remove(0);
    state
        .complete(round.job, ShimOutput { items: vec![] })
        .unwrap();
    let mut recovered = ConversationState::recover(session, state.boxes().to_vec(), true).unwrap();
    let retry = start(&mut recovered);
    assert!(!retry.boxes.is_empty());
    let next = recovered
        .prepare_stage(
            retry.job,
            String::new(),
            vec![BoxValue::Call(Ok(call("two")))],
        )
        .unwrap()
        .remove(0);
    assert_ne!(first_call.tool_call_id, next.tool_call_id);
}