kcode-k1-chat-state 0.4.0

Provider-free boxed chat scheduling state
Documentation
use crate::{ActorState, BoxContent, ProviderCall, StateError, ToolCallId};

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(),
    }
}

#[test]
fn begin_adds_no_box() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    assert_eq!(start.frontier, None);
    assert!(state.boxes().is_empty());
    state.complete_inference(start.job, String::new()).unwrap();
}

#[test]
fn stages_are_canonical_and_launchable_before_done() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    let first = state
        .append_stage(start.job, "a".into(), vec![call(1), call(2)])
        .unwrap();
    let second = state
        .append_stage(start.job, String::new(), vec![call(3)])
        .unwrap();
    assert_eq!(first.len(), 2);
    assert_eq!(second[0].tool_call_id, id(3));
    assert!(matches!(state.boxes()[0].content(), BoxContent::Kennedy { text } if text == "a"));
    assert!(
        matches!(state.boxes()[1].content(), BoxContent::KtoolCall { tool_call_id, .. } if *tool_call_id == id(1))
    );
    state.complete_inference(start.job, "z".into()).unwrap();
    assert!(
        matches!(state.boxes().last().unwrap().content(), BoxContent::Kennedy { text } if text == "z")
    );
}

#[test]
fn queued_arrivals_are_fifo_and_schedule_one_followup() {
    let mut state = ActorState::new(true);
    let first = state.begin_inference().unwrap().unwrap();
    state.accept_system("s".into()).unwrap();
    state.accept_user("u".into()).unwrap();
    state.accept_attachment().unwrap();
    state.complete_inference(first.job, String::new()).unwrap();
    assert!(matches!(state.boxes()[0].content(), BoxContent::System(text) if text == "s"));
    assert!(matches!(state.boxes()[1].content(), BoxContent::User(text) if text == "u"));
    assert!(matches!(state.boxes()[2].content(), BoxContent::Attachment));
    let followup = state.begin_inference().unwrap().unwrap();
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_inference(followup.job, String::new())
        .unwrap();
    assert!(state.quiet());
}

#[test]
fn call_only_stage_does_not_schedule() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    assert_eq!(
        state
            .append_stage(start.job, String::new(), vec![call(1)])
            .unwrap()
            .len(),
        1
    );
    state.complete_inference(start.job, String::new()).unwrap();
    assert!(state.quiet());
    assert!(state.begin_inference().unwrap().is_none());
}

#[test]
fn async_returns_accept_out_of_order() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(start.job, String::new(), vec![call(1), call(2)])
        .unwrap();
    state.complete_inference(start.job, String::new()).unwrap();
    state.accept_async_return(id(2), Ok("two".into())).unwrap();
    state.accept_async_return(id(1), Ok("one".into())).unwrap();
    let tail = &state.boxes()[2..];
    assert!(
        matches!(tail[0].content(), BoxContent::KtoolReturn { tool_call_id, .. } if *tool_call_id == id(2))
    );
    assert!(
        matches!(tail[1].content(), BoxContent::KtoolReturn { tool_call_id, .. } if *tool_call_id == id(1))
    );
}

#[test]
fn recovery_is_idle_and_keeps_box_ids() {
    let mut original = ActorState::new(false);
    original.accept_user("prior".into()).unwrap();
    let boxes = original.boxes().to_vec();
    let mut state = ActorState::recover(boxes, false).unwrap();
    assert!(state.quiet());
    state.accept_system("next".into()).unwrap();
    assert_eq!(state.boxes().last().unwrap().id().get(), 2);
}

#[test]
fn retry_uses_a_fresh_job_and_the_same_frontier() {
    let mut state = ActorState::new(false);
    state.accept_user("ask".into()).unwrap();
    let first = state.begin_inference().unwrap().unwrap();
    state.stall_inference(first.job, "retry".into()).unwrap();
    assert_eq!(state.take_halt(), Some("retry".into()));
    state.restart().unwrap();
    let second = state.begin_inference().unwrap().unwrap();
    assert!(second.job > first.job);
    assert_eq!(second.frontier, first.frontier);
    assert_eq!(state.boxes().len(), 1);
    state.complete_inference(second.job, String::new()).unwrap();
}

#[test]
fn wrong_job_and_halt_are_transactional() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    assert!(matches!(
        state.append_stage(start.job + 1, String::new(), vec![]),
        Err(StateError::WrongInference { .. })
    ));
    assert!(state.halt("stop".into()));
    assert!(matches!(state.restart(), Err(StateError::Busy)));
    state.complete_inference(start.job, String::new()).unwrap();
    assert_eq!(state.take_halt(), Some("stop".into()));
    assert!(state.quiet());
}