kcode-k1-chat-chatend 0.3.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 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 value = id(7);
    assert_eq!((value.session(), value.sequence()), ([7; 12], 7));
    assert!(value < id(8));

    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);
    assert!(chat.append_stage(String::new(), vec![]).unwrap().is_empty());

    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(|v| v.call_box_id.get())
            .collect::<Vec<_>>(),
        [3, 4]
    );
    assert_eq!(second[0].call_box_id.get(), 6);
    assert!(matches!(chat.boxes()[1].content(), BoxContent::Kennedy { text } if text == "a"));
    assert!(matches!(
        chat.boxes()[2].content(),
        BoxContent::KtoolCall { .. }
    ));

    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);
    }

    assert!(chat.accept_system("s".into()).unwrap().is_none());
    assert!(
        chat.accept_async_return(id(2), Ok("two".into()))
            .unwrap()
            .is_none()
    );
    assert!(chat.accept_user("u".into()).unwrap().is_none());
    assert!(chat.accept_attachment().unwrap().is_none());
    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));

    assert_eq!(
        chat.accept_async_return(id(3), Ok("3".into()))
            .unwrap()
            .unwrap()
            .get(),
        12
    );
    assert_eq!(
        chat.accept_async_return(id(1), Ok("1".into()))
            .unwrap()
            .unwrap()
            .get(),
        13
    );
    assert!(is_return(&chat.boxes()[11], 3) && is_return(&chat.boxes()[12], 1));
    assert_eq!(
        chat.accept_async_return(id(9), Ok("x".into())).unwrap_err(),
        TransitionError::UnknownToolCall
    );

    let finished = chat.boxes().to_vec();
    assert_eq!(
        chat.done(String::new()).unwrap_err(),
        TransitionError::InvalidPhase
    );
    assert_eq!(chat.boxes(), finished);
}

#[test]
fn recovery_validates_and_continues() {
    let valid = vec![
        boxed(1, BoxContent::Kennedy { text: "s".into() }),
        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 cases = [
        (
            vec![boxed(2, BoxContent::Attachment)],
            RecoveryError::NonContiguousBoxId,
        ),
        (
            vec![boxed(
                1,
                BoxContent::Kennedy {
                    text: String::new(),
                },
            )],
            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 cases {
        match Chatend::recover(boxes) {
            Err(error) => assert_eq!(error, expected),
            Ok(_) => panic!("invalid recovery succeeded"),
        }
    }
}