kcode-k1-chat-codex-state 0.3.0

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

fn call() -> BoxValue {
    BoxValue::Call(Ok(Call {
        name: "tool".into(),
        arguments: "{}".into(),
    }))
}

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

fn active(byte: u8) -> (ConversationState, u64, ToolCallId) {
    let mut state = ConversationState::new([byte; 12]);
    state.accept(Arrival::User("ask".into())).unwrap();
    let job = start(&mut state).job;
    let id = state
        .prepare_stage(job, "thinking".into(), vec![call()])
        .unwrap()[0]
        .tool_call_id;
    (state, job, id)
}

#[test]
fn steer_is_one_stable_external_fifo_prefix() {
    let (mut state, job, id) = active(1);
    assert!(state.prepare_steer(job).unwrap().is_none());
    state.accept(Arrival::User("u".into())).unwrap();
    state
        .accept(Arrival::Return {
            tool_call_id: id,
            result: Ok("r".into()),
        })
        .unwrap();
    let first = state.prepare_steer(job).unwrap().unwrap();
    assert_eq!(
        first.values(),
        state.boxes()[3..5].iter().map(project).collect::<Vec<_>>()
    );
    assert!(matches!(state.boxes()[3].content(), BoxContent::User(text) if text == "u"));
    assert!(matches!(
        state.boxes()[4].content(),
        BoxContent::KtoolReturn { .. }
    ));
    assert_eq!(state.boxes()[4].id().get(), state.boxes()[3].id().get() + 1);
    assert!(
        state
            .prepare_stage(job, String::new(), vec![call()])
            .is_err()
    );
    state.accept(Arrival::System("later".into())).unwrap();
    let repeated = state.prepare_steer(job).unwrap().unwrap();
    assert_eq!(repeated.values(), first.values());
    state.validate_steer(&first).unwrap();
    state.commit_steer(repeated).unwrap();
    assert!(state.commit_steer(first).is_err());
    let later = state.prepare_steer(job).unwrap().unwrap();
    assert!(matches!(&later.values()[0], BoxValue::History(text) if text.contains("later")));
    state.commit_steer(later).unwrap();
    assert!(state.prepare_steer(job).unwrap().is_none());
    state
        .complete(
            job,
            ShimOutput {
                items: vec![ShimItem::Text("done".into())],
            },
        )
        .unwrap();
    assert_eq!(state.status(), Status::Quiet);
}

#[test]
fn tokens_are_exact_once_stale_wrong_job_and_unforgeable() {
    let (mut first, job, _) = active(2);
    first.accept(Arrival::User("queued".into())).unwrap();
    assert!(first.prepare_steer(job + 1).is_err());
    let token = first.prepare_steer(job).unwrap().unwrap();
    first.validate_steer(&token).unwrap();
    let (mut other, other_job, _) = active(2);
    other.accept(Arrival::User("queued".into())).unwrap();
    let foreign = other.prepare_steer(other_job).unwrap().unwrap();
    assert!(first.validate_steer(&foreign).is_err());
    first.commit_steer(token.clone()).unwrap();
    first.complete(job, ShimOutput { items: vec![] }).unwrap();
    first.accept(Arrival::User("next".into())).unwrap();
    assert_ne!(start(&mut first).job, job);
    assert!(first.validate_steer(&token).is_err());
}

#[test]
fn failures_recovery_and_incumbent_rules_hold() {
    let (mut state, job, first_id) = active(3);
    state.accept(Arrival::User("flushed".into())).unwrap();
    let token = state.prepare_steer(job).unwrap().unwrap();
    state.fail(job, "steer failed".into(), false);
    assert!(matches!(
        state.status(),
        Status::Stalled {
            restartable: false,
            ..
        }
    ));
    assert!(state.validate_steer(&token).is_err());
    assert_eq!(state.restart(), Err(RestartError::NotRestartable));
    let expected = state.boxes().iter().map(project).collect::<Vec<_>>();
    let mut recovered = ConversationState::recover([3; 12], state.boxes().to_vec(), true).unwrap();
    let retry = start(&mut recovered);
    assert_eq!(retry.boxes, expected);
    assert!(recovered.validate_steer(&token).is_err());
    let next_id = recovered
        .prepare_stage(retry.job, String::new(), vec![call()])
        .unwrap()[0]
        .tool_call_id;
    assert!(next_id.sequence() > first_id.sequence());

    let (mut incomplete, active_job, _) = active(4);
    incomplete.accept(Arrival::User("queued".into())).unwrap();
    incomplete.prepare_steer(active_job).unwrap().unwrap();
    assert!(
        incomplete
            .complete(active_job, ShimOutput { items: vec![] })
            .is_err()
    );
    assert_eq!(incomplete.restart(), Err(RestartError::NotRestartable));

    let mut early = ConversationState::new([5; 12]);
    early.accept(Arrival::User("ask".into())).unwrap();
    let early_job = start(&mut early).job;
    assert!(
        early
            .prepare_stage(
                early_job,
                "bad".into(),
                vec![BoxValue::History("bad".into())]
            )
            .is_err()
    );
    assert_eq!(early.boxes().len(), 1);
    early.fail(early_job, "early".into(), true);
    assert!(matches!(
        early.status(),
        Status::Stalled {
            restartable: true,
            ..
        }
    ));
    early.restart().unwrap();
    assert_eq!(start(&mut early).boxes.len(), 1);
}