kcode-k1-chat-chatend 0.4.0

use super::*;

fn id(n: u64) -> ToolCallId {
    ToolCallId::new([n as u8; 12], n)
}

fn call(n: u64) -> ProviderCall {
    ProviderCall {
        tool_call_id: id(n),
        name: format!("t{n}"),
        arguments: "{}".into(),
    }
}

fn boxed(n: u64, content: BoxContent) -> ChatBox {
    ChatBox::new(BoxId::new(n), content)
}

fn call_box(n: u64, tool: u64) -> ChatBox {
    boxed(
        n,
        BoxContent::KtoolCall {
            tool_call_id: id(tool),
            name: "t".into(),
            arguments: "{}".into(),
        },
    )
}

fn return_box(n: u64, tool: u64, origin: u64) -> ChatBox {
    boxed(
        n,
        BoxContent::KtoolReturn {
            tool_call_id: id(tool),
            originating_call: BoxId::new(origin),
            result: Ok("r".into()),
        },
    )
}

fn kennedy_box(n: u64, text: &str) -> ChatBox {
    boxed(n, BoxContent::Kennedy { text: text.into() })
}

fn is_return(value: &ChatBox, tool: u64) -> bool {
    matches!(value.content(), BoxContent::KtoolReturn { tool_call_id, .. } if *tool_call_id == id(tool))
}

#[test]
fn stages_arrivals_and_rejections_are_exact() {
    let mut chat = Chatend::new();
    assert_eq!(
        chat.accept_system("prior".into()).unwrap().unwrap().get(),
        1
    );
    assert_eq!(chat.start_round().unwrap().unwrap().get(), 1);
    let first = chat
        .append_stage("a".into(), vec![call(1), call(2)])
        .unwrap();
    let second = chat.append_stage("b".into(), vec![call(3)]).unwrap();
    assert_eq!(
        first
            .iter()
            .map(|call| call.call_box_id.get())
            .collect::<Vec<_>>(),
        [3, 4]
    );
    assert_eq!(second[0].call_box_id.get(), 6);

    let before = chat.boxes().to_vec();
    for calls in [vec![call(4), call(4)], vec![call(1)]] {
        assert_eq!(
            chat.append_stage("bad".into(), calls).unwrap_err(),
            TransitionError::DuplicateToolCall
        );
        assert_eq!(chat.boxes(), before);
    }

    chat.accept_system("s".into()).unwrap();
    chat.accept_async_return(id(2), Ok("two".into())).unwrap();
    chat.accept_user("u".into()).unwrap();
    chat.accept_attachment().unwrap();
    assert_eq!(
        chat.accept_async_return(id(2), Ok("again".into()))
            .unwrap_err(),
        TransitionError::DuplicateReturn
    );
    chat.done("z".into()).unwrap();
    let tail = &chat.boxes()[6..];
    assert!(matches!(tail[0].content(), BoxContent::Kennedy { text } if text == "z"));
    assert!(matches!(tail[1].content(), BoxContent::System(text) if text == "s"));
    assert!(is_return(&tail[2], 2));
    assert!(matches!(tail[3].content(), BoxContent::User(text) if text == "u"));
    assert!(matches!(tail[4].content(), BoxContent::Attachment));

    for (tool, expected_id) in [(3, 12), (1, 13)] {
        let actual = chat
            .accept_async_return(id(tool), Ok(tool.to_string()))
            .unwrap();
        assert_eq!(actual.unwrap().get(), expected_id);
    }
    assert_eq!(
        chat.accept_async_return(id(9), Ok("x".into())).unwrap_err(),
        TransitionError::UnknownToolCall
    );
    assert_eq!(
        chat.done(String::new()).unwrap_err(),
        TransitionError::InvalidPhase
    );
}

#[test]
fn active_arrival_flush_is_finite_fifo_and_transactional() {
    let mut chat = Chatend::new();
    assert_eq!(
        chat.flush_active_arrivals().unwrap_err(),
        TransitionError::InvalidPhase
    );
    chat.start_round().unwrap();
    assert!(chat.flush_active_arrivals().unwrap().is_empty());
    chat.append_stage(String::new(), vec![call(1)]).unwrap();
    chat.accept_user("first".into()).unwrap();
    let before = chat.boxes().to_vec();
    assert_eq!(
        chat.accept_async_return(id(9), Ok("bad".into()))
            .unwrap_err(),
        TransitionError::UnknownToolCall
    );
    assert_eq!(chat.boxes(), before);
    chat.accept_async_return(id(1), Ok("result".into()))
        .unwrap();

    let flushed = chat.flush_active_arrivals().unwrap();
    assert_eq!(
        flushed.iter().map(|b| b.id().get()).collect::<Vec<_>>(),
        [2, 3]
    );
    assert!(matches!(flushed[0].content(), BoxContent::User(text) if text == "first"));
    assert!(is_return(&flushed[1], 1));
    assert!(chat.flush_active_arrivals().unwrap().is_empty());
    assert!(chat.append_stage(String::new(), vec![]).unwrap().is_empty());

    chat.accept_user("later".into()).unwrap();
    chat.done("final".into()).unwrap();
    assert!(matches!(chat.boxes()[3].content(), BoxContent::Kennedy { text } if text == "final"));
    assert!(matches!(chat.boxes()[4].content(), BoxContent::User(text) if text == "later"));
}

#[test]
fn recovery_validates_and_continues() {
    let valid = vec![kennedy_box(1, "s"), call_box(2, 1), return_box(3, 1, 2)];
    let mut recovered = Chatend::recover(valid.clone()).unwrap();
    assert_eq!(recovered.boxes(), valid);
    assert_eq!(
        recovered.accept_user("next".into()).unwrap().unwrap().get(),
        4
    );

    let invalid = [
        (
            vec![boxed(2, BoxContent::Attachment)],
            RecoveryError::NonContiguousBoxId,
        ),
        (vec![kennedy_box(1, "")], RecoveryError::EmptyKennedy),
        (
            vec![call_box(1, 1), call_box(2, 1)],
            RecoveryError::DuplicateToolCall,
        ),
        (
            vec![return_box(1, 1, 2), call_box(2, 1)],
            RecoveryError::UnknownToolCall,
        ),
        (
            vec![call_box(1, 1), return_box(2, 1, 2)],
            RecoveryError::WrongOriginatingCall,
        ),
        (
            vec![call_box(1, 1), return_box(2, 1, 1), return_box(3, 1, 1)],
            RecoveryError::DuplicateReturn,
        ),
    ];
    for (boxes, expected) in invalid {
        assert_eq!(Chatend::recover(boxes).err(), Some(expected));
    }
}