kcode-k1-chat-codex-state 0.6.0

use kcode_k1_chat_codex_codec::{BoxValue, Call, project};
use kcode_k1_chat_codex_state::{
    ConversationState, RestartError, Start, TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE, ToolCallId,
    USER_MESSAGE_TYPE,
};
use kcode_k1_codex_adapter::{ShimItem, ShimOutput};

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

fn accept(state: &mut ConversationState, kind: &str, contents: &str) {
    state
        .accept(kind.into(), contents.into(), String::new(), String::new())
        .unwrap();
}

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

fn active(byte: u8) -> (ConversationState, u64, ToolCallId) {
    let mut state = ConversationState::new([byte; 12]);
    accept(&mut state, USER_MESSAGE_TYPE, "ask");
    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 flushed_message_waits_for_a_later_result() {
    let (mut state, job, id) = active(1);
    state.accept_tool_message(id, "working".into()).unwrap();
    assert_eq!(
        state.flush_active_arrivals(job).unwrap()[0].box_type(),
        TOOL_MESSAGE_TYPE
    );
    assert!(state.prepare_steer(job).unwrap().is_none());
    state.accept_tool_return(id, Ok("done".into())).unwrap();
    let token = state.prepare_steer(job).unwrap().unwrap();
    let tail = &state.boxes()[state.boxes().len() - 2..];
    assert_eq!(
        [tail[0].box_type(), tail[1].box_type()],
        [TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE]
    );
    assert_eq!(token.values(), tail.iter().map(project).collect::<Vec<_>>());
}

#[test]
fn unflushed_messages_and_v2_result_stay_fifo() {
    let (mut state, job, id) = active(2);
    state.accept_tool_message(id, "one".into()).unwrap();
    state.accept_tool_message(id, "two".into()).unwrap();
    state
        .accept_tool_return_v2(id, Ok("done".into()), "result/meta-v1".into(), "{}".into())
        .unwrap();
    let token = state.prepare_steer(job).unwrap().unwrap();
    let tail = &state.boxes()[state.boxes().len() - 3..];
    assert_eq!(
        tail.iter()
            .map(|value| value.box_type())
            .collect::<Vec<_>>(),
        [TOOL_MESSAGE_TYPE, TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE]
    );
    assert_eq!(token.values(), tail.iter().map(project).collect::<Vec<_>>());
}

#[test]
fn exact_token_keeps_late_fifo_for_the_next_token() {
    let (mut state, job, id) = active(3);
    accept(&mut state, "Future Arrival", "first");
    let first = state.prepare_steer(job).unwrap().unwrap();
    state
        .accept_tool_message(id, "late message".into())
        .unwrap();
    state
        .accept(
            "Future Arrival".into(),
            "second".into(),
            "hidden/v1".into(),
            "secret".into(),
        )
        .unwrap();
    state.flush_active_arrivals(job).unwrap();
    state.validate_steer(&first).unwrap();
    let repeated = state.prepare_steer(job).unwrap().unwrap();
    state.commit_steer(repeated).unwrap();
    assert!(state.commit_steer(first).is_err());
    let next = state.prepare_steer(job).unwrap().unwrap();
    assert_eq!(next.values().len(), 2);
    assert!(!format!("{:?}", next.values()).contains("secret"));
}

#[test]
fn restart_invalidates_tokens_and_terminal_boxes_stall() {
    let mut state = ConversationState::new([4; 12]);
    accept(&mut state, USER_MESSAGE_TYPE, "ask");
    let job = start(&mut state).job;
    accept(&mut state, "Future Arrival", "queued");
    let token = state.prepare_steer(job).unwrap().unwrap();
    state.fail(job, "retry".into(), true);
    state.restart().unwrap();
    let retry = start(&mut state);
    assert!(state.validate_steer(&token).is_err());
    let output = ShimOutput {
        items: vec![ShimItem::Box(BoxValue::History("forbidden".into()))],
    };
    assert!(state.complete(retry.job, output).is_err());
    assert_eq!(state.restart(), Err(RestartError::NotRestartable));
}