kcode-k1-chat-state 0.9.4

use crate::{
    AGENT_MESSAGE_TYPE, AGENT_RESPONSE_TYPE, ActorState, PreflightGenerated, ProviderCall,
    ProviderGenerated, SYSTEM_MESSAGE_TYPE, TOOL_CALL_TYPE, TOOL_MESSAGE_TYPE, TOOL_RESULT_TYPE,
    ToolCallId, USER_MESSAGE_TYPE,
};

fn tool_id(value: u8) -> ToolCallId {
    ToolCallId::new([value; 12], u64::from(value))
}

fn call(value: u8) -> ProviderCall {
    ProviderCall {
        tool_call_id: tool_id(value),
        name: "WebSearch".into(),
        arguments: "{}".into(),
    }
}

fn generated_call(value: u8) -> ProviderGenerated {
    ProviderGenerated::ToolCall(call(value))
}

#[test]
fn preflight_is_ordered_and_does_not_schedule_inference() {
    let mut state = ActorState::new(false);
    let calls = state
        .append_preflight(vec![
            PreflightGenerated::SystemMessage {
                contents: "context".into(),
            },
            PreflightGenerated::ToolCall(call(9)),
        ])
        .unwrap();
    assert_eq!(calls.len(), 1);
    assert_eq!(calls[0].tool_call_id, tool_id(9));
    assert_eq!(calls[0].call_box_id.get(), 2);
    assert_eq!(
        state
            .boxes()
            .iter()
            .map(|value| value.box_type())
            .collect::<Vec<_>>(),
        [SYSTEM_MESSAGE_TYPE, TOOL_CALL_TYPE]
    );
    assert!(state.begin_inference().unwrap().is_none());
}

#[test]
fn tool_messages_do_not_trigger_but_v2_result_does() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(active.job, String::new(), vec![generated_call(1)])
        .unwrap();

    state
        .accept_tool_message(tool_id(1), "first".into())
        .unwrap();
    let flushed = state.flush_active_arrivals(active.job).unwrap();
    assert_eq!(flushed.len(), 1);
    assert_eq!(flushed[0].box_type(), TOOL_MESSAGE_TYPE);

    state
        .accept_tool_message(tool_id(1), "second".into())
        .unwrap();
    state.complete_inference(active.job, String::new()).unwrap();
    assert!(state.quiet());
    assert!(state.begin_inference().unwrap().is_none());

    state
        .accept_async_return_v2(
            tool_id(1),
            Ok("answer".into()),
            "k1.web-search-result/v1".into(),
            "{}".into(),
        )
        .unwrap();
    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 messages_and_result_stay_fifo_and_schedule_one_followup() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(active.job, String::new(), vec![generated_call(2)])
        .unwrap();
    state.accept_tool_message(tool_id(2), "one".into()).unwrap();
    state.accept_tool_message(tool_id(2), "two".into()).unwrap();
    assert_eq!(state.flush_active_arrivals(active.job).unwrap().len(), 2);
    state
        .accept_async_return(tool_id(2), Ok("done".into()))
        .unwrap();
    state.complete_inference(active.job, String::new()).unwrap();

    let kinds = state
        .boxes()
        .iter()
        .map(|value| value.box_type())
        .collect::<Vec<_>>();
    assert_eq!(
        &kinds[kinds.len() - 4..],
        [
            TOOL_MESSAGE_TYPE,
            TOOL_MESSAGE_TYPE,
            AGENT_RESPONSE_TYPE,
            TOOL_RESULT_TYPE,
        ]
    );
    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 retry_uses_a_fresh_job_and_preserves_the_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);
    state.complete_inference(second.job, String::new()).unwrap();
}

#[test]
fn stage_keeps_response_and_provider_order_but_dispatches_only_calls() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    let dispatched = state
        .append_stage(
            active.job,
            String::new(),
            vec![
                ProviderGenerated::AgentMessage {
                    contents: "visible".into(),
                },
                generated_call(3),
            ],
        )
        .unwrap();

    assert_eq!(dispatched.len(), 1);
    assert_eq!(dispatched[0].tool_call_id, tool_id(3));
    let kinds = state
        .boxes()
        .iter()
        .map(|value| value.box_type())
        .collect::<Vec<_>>();
    assert_eq!(
        kinds,
        [AGENT_RESPONSE_TYPE, AGENT_MESSAGE_TYPE, TOOL_CALL_TYPE]
    );
    assert_eq!(state.boxes()[0].contents(), "");
    assert_eq!(state.boxes()[1].contents(), "visible");
}

#[test]
fn done_appends_response_before_queued_arrivals() {
    let mut state = ActorState::new(true);
    let active = state.begin_inference().unwrap().unwrap();
    state
        .append_stage(active.job, "wave".into(), Vec::new())
        .unwrap();
    assert!(state.flush_active_arrivals(active.job).unwrap().is_empty());
    state.accept_user("late".into()).unwrap();

    state.complete_inference(active.job, String::new()).unwrap();

    let boxes = state.boxes();
    assert_eq!(boxes[0].box_type(), AGENT_RESPONSE_TYPE);
    assert_eq!(boxes[0].contents(), "wave");
    assert_eq!(boxes[1].box_type(), AGENT_RESPONSE_TYPE);
    assert_eq!(boxes[1].contents(), "");
    assert_eq!(boxes[2].box_type(), USER_MESSAGE_TYPE);
    assert_eq!(boxes[2].contents(), "late");
}