kcode-k1-chat-state 0.3.0

Provider-free boxed chat scheduling state
Documentation
use std::sync::{Arc, atomic::Ordering};

use crate::{ActionId, ActorState, BoxContent, DispatchOutcome, StateError};

fn action(value: u64) -> ActionId {
    ActionId::new([value as u8; 12], value)
}

#[test]
fn arrivals_open_one_round_with_the_exact_frontier_and_shared_attempt() {
    let mut state = ActorState::new(false);
    assert!(state.quiet());
    state.accept_user("ask".into()).unwrap();
    let frontier = state.boxes()[0].id();
    let start = state.begin_inference().unwrap().unwrap();
    assert_eq!(start.frontier, Some(frontier));
    assert_eq!(Arc::strong_count(&start.attempt), 2);
    start.attempt.store(4, Ordering::Relaxed);
    assert!(matches!(
        state.boxes()[1].content(),
        BoxContent::Kennedy { text, complete } if text.is_empty() && !complete
    ));
    state.append_kennedy_text(start.job, "answer").unwrap();
    state.complete_provider_output(start.job, vec![]).unwrap();
    assert!(matches!(
        state.boxes()[1].content(),
        BoxContent::Kennedy { text, complete } if text == "answer" && *complete
    ));
    assert!(state.quiet());
}

#[test]
fn provider_and_dispatch_barriers_drain_the_ordered_suffix() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    state
        .collect_provider_call(start.job, "a".into(), "1".into())
        .unwrap();
    state
        .collect_provider_call(start.job, "b".into(), "2".into())
        .unwrap();
    state.accept_system("queued".into()).unwrap();
    state.accept_attachment().unwrap();
    assert_eq!(state.boxes().len(), 1);
    let calls = state
        .complete_provider_output(start.job, vec![action(7), action(8)])
        .unwrap();
    assert_eq!(calls.len(), 2);
    state.accept_user("after dispatch".into()).unwrap();
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_dispatch(vec![
            DispatchOutcome::Terminal(Ok("done".into())),
            DispatchOutcome::Pending,
        ])
        .unwrap();
    let boxes = state.boxes();
    assert_eq!(boxes.len(), 7);
    assert!(matches!(boxes[0].content(), BoxContent::Kennedy { .. }));
    assert!(matches!(boxes[1].content(), BoxContent::KtoolCall { name, .. } if name == "a"));
    assert!(matches!(boxes[2].content(), BoxContent::KtoolCall { name, .. } if name == "b"));
    assert!(
        matches!(boxes[3].content(), BoxContent::KtoolReturn { result, .. } if result == &Ok("done".into()))
    );
    assert!(matches!(boxes[4].content(), BoxContent::System(text) if text == "queued"));
    assert!(matches!(boxes[5].content(), BoxContent::Attachment));
    assert!(matches!(boxes[6].content(), BoxContent::User(text) if text == "after dispatch"));
    let frontier = boxes[6].id();
    let next = state.begin_inference().unwrap().unwrap();
    assert_eq!(next.frontier, Some(frontier));
    state.complete_provider_output(next.job, vec![]).unwrap();
    assert!(state.quiet());
}

#[test]
fn call_only_output_with_pending_dispatch_becomes_quiet_without_arrival() {
    let mut state = ActorState::new(true);
    let first = state.begin_inference().unwrap().unwrap();
    state
        .collect_provider_call(first.job, "later".into(), "{}".into())
        .unwrap();
    let calls = state
        .complete_provider_output(first.job, vec![action(1)])
        .unwrap();
    assert!(matches!(
        state.boxes()[0].content(),
        BoxContent::Kennedy { text, complete } if text.is_empty() && *complete
    ));
    assert_eq!(calls.len(), 1);
    state
        .complete_dispatch(vec![DispatchOutcome::Pending])
        .unwrap();
    assert!(state.begin_inference().unwrap().is_none());
    assert!(state.quiet());
}

#[test]
fn async_return_keeps_its_call_reference_and_schedules_work() {
    let mut state = ActorState::new(true);
    let first = state.begin_inference().unwrap().unwrap();
    state
        .collect_provider_call(first.job, "async".into(), "{}".into())
        .unwrap();
    let call = state
        .complete_provider_output(first.job, vec![action(3)])
        .unwrap()
        .remove(0);
    state
        .complete_dispatch(vec![DispatchOutcome::Pending])
        .unwrap();
    assert!(state.quiet());
    state
        .accept_async_return(call.action_id, Ok("later".into()))
        .unwrap();
    assert!(matches!(
        state.boxes().last().unwrap().content(),
        BoxContent::KtoolReturn {
            action_id,
            originating_call,
            ..
        } if *action_id == call.action_id && *originating_call == call.call_box_id
    ));
    assert!(!state.quiet());
}

#[test]
fn wrong_jobs_and_count_errors_are_transactional() {
    let mut state = ActorState::new(true);
    let start = state.begin_inference().unwrap().unwrap();
    assert!(matches!(
        state.append_kennedy_text(start.job + 1, "bad"),
        Err(StateError::WrongInference { .. })
    ));
    assert!(matches!(
        state.collect_provider_call(start.job + 1, "bad".into(), "bad".into()),
        Err(StateError::WrongInference { .. })
    ));
    state
        .collect_provider_call(start.job, "good".into(), "{}".into())
        .unwrap();
    assert!(matches!(
        state.complete_provider_output(start.job, vec![]),
        Err(StateError::Transition(_))
    ));
    let calls = state
        .complete_provider_output(start.job, vec![action(4)])
        .unwrap();
    assert_eq!(calls.len(), 1);
    assert!(matches!(
        state.complete_dispatch(vec![]),
        Err(StateError::Transition(_))
    ));
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_dispatch(vec![DispatchOutcome::Pending])
        .unwrap();
}

#[test]
fn stall_restart_reuses_the_open_kennedy_and_frontier_with_a_fresh_job() {
    let mut state = ActorState::new(false);
    state.accept_user("ask".into()).unwrap();
    let first = state.begin_inference().unwrap().unwrap();
    let kennedy = state.boxes().last().unwrap().id();
    state.append_kennedy_text(first.job, "partial").unwrap();
    state.stall_inference(first.job, "first".into()).unwrap();
    assert_eq!(state.take_halt(), Some("first".into()));
    assert!(state.halted());
    assert!(!state.halt("second".into()));
    assert!(state.begin_inference().unwrap().is_none());
    assert!(state.restart().is_ok());
    let second = state.begin_inference().unwrap().unwrap();
    assert!(second.job > first.job);
    assert_eq!(second.frontier, first.frontier);
    assert_eq!(state.boxes().last().unwrap().id(), kennedy);
    assert_eq!(Arc::strong_count(&second.attempt), 2);
    state.complete_provider_output(second.job, vec![]).unwrap();
}

#[test]
fn halt_force_quiet_and_exhaustion_are_exact() {
    let mut state = ActorState::new(false);
    assert!(matches!(state.restart(), Err(StateError::NotStalled)));
    assert!(state.halt("first".into()));
    assert!(!state.halt("second".into()));
    assert_eq!(state.take_halt(), Some("first".into()));
    assert!(state.quiet());
    state.halt("again".into());
    state.force_inference();
    assert!(state.restart().is_ok());
    let start = state.begin_inference().unwrap().unwrap();
    assert_eq!(start.job, 1);
    state.complete_provider_output(start.job, vec![]).unwrap();
    state.next_job = u64::MAX;
    state.force_inference();
    let before = state.boxes().len();
    assert!(matches!(
        state.begin_inference(),
        Err(StateError::JobIdExhausted)
    ));
    assert_eq!(state.boxes().len(), before);
}