kcode-k1-chat-state 0.7.0

use crate::{
    ActorState, BoxId, DispatchedToolCall, ProviderCall, ResultView, StateError, ToolResult,
    ToolResultStatus, TransitionError, USER_ATTACHMENT_TYPE,
};
use serde_json::json;

fn provider(tool: &str) -> ProviderCall {
    ProviderCall {
        tool: tool.into(),
        tool_version: "1.0.0".into(),
        arguments: json!({"value": tool}),
    }
}

fn result(call: &DispatchedToolCall, status: ToolResultStatus) -> ToolResult {
    let view = match status {
        ToolResultStatus::Ok => ResultView::OneLine("complete".into()),
        ToolResultStatus::Error => ResultView::Error("failed".into()),
    };
    ToolResult::new(
        call.call.call_id(),
        call.call_box_id,
        call.call.tool().into(),
        call.call.tool_version().into(),
        status,
        json!({"complete": true}),
        view,
    )
    .unwrap()
}

fn start(state: &mut ActorState) -> crate::InferenceStart {
    state.begin_inference().unwrap().unwrap()
}

fn finish_followup(state: &mut ActorState) {
    let followup = start(state);
    assert!(state.begin_inference().unwrap().is_none());
    state
        .complete_inference(followup.job, String::new())
        .unwrap();
    assert!(state.quiet());
}

#[test]
fn calls_are_transactional_and_recover_the_next_id() {
    let mut state = ActorState::new(true);
    let active = start(&mut state);
    assert_eq!(active.frontier, None);
    assert!(state.boxes().is_empty());

    assert!(matches!(
        state.append_stage(active.job + 1, String::new(), vec![provider("wrong")]),
        Err(StateError::WrongInference { .. })
    ));
    let invalid = ProviderCall {
        tool: "Tool".into(),
        tool_version: "1.0.0".into(),
        arguments: json!({"tool": "forbidden"}),
    };
    assert_eq!(
        state.append_stage(active.job, "not appended".into(), vec![invalid]),
        Err(StateError::Transition(TransitionError::InvalidToolEnvelope))
    );
    assert!(state.boxes().is_empty());

    let calls = state
        .append_stage(
            active.job,
            "thinking".into(),
            vec![provider("first"), provider("second")],
        )
        .unwrap();
    assert_eq!(calls[0].call.call_id().to_string(), "c1");
    assert_eq!(calls[1].call.call_id().to_string(), "c2");
    assert_eq!(calls[0].call_box_id, BoxId::new(2));
    state
        .complete_inference(active.job, "answer".into())
        .unwrap();

    let mut recovered = ActorState::recover(state.boxes().to_vec(), true).unwrap();
    let next = start(&mut recovered);
    let call = recovered
        .append_stage(next.job, String::new(), vec![provider("third")])
        .unwrap();
    assert_eq!(call[0].call.call_id().to_string(), "c3");
}

#[test]
fn open_arrivals_flush_finitely_and_survive_recovery() {
    let mut state = ActorState::new(true);
    let active = start(&mut state);
    state
        .accept_box(
            "Future Kind".into(),
            "visible".into(),
            "future/v9".into(),
            "hidden".into(),
        )
        .unwrap();
    state
        .accept_attachment(
            "Object ID: abc\nName: report.pdf".into(),
            "k1.attachment/v1".into(),
            "{\"mime\":\"application/pdf\"}".into(),
        )
        .unwrap();

    let flushed = state.flush_active_arrivals(active.job).unwrap();
    assert_eq!(flushed[0].box_type(), "Future Kind");
    assert_eq!(flushed[1].box_type(), USER_ATTACHMENT_TYPE);
    assert_eq!(flushed[1].hidden_type(), "k1.attachment/v1");
    assert!(state.flush_active_arrivals(active.job).unwrap().is_empty());

    state
        .accept_box(
            "Later Future".into(),
            "later".into(),
            "future/v10".into(),
            "later hidden".into(),
        )
        .unwrap();
    state.complete_inference(active.job, String::new()).unwrap();
    finish_followup(&mut state);

    let boxes = state.boxes().to_vec();
    let mut recovered = ActorState::recover(boxes.clone(), false).unwrap();
    assert_eq!(recovered.boxes(), boxes);
    assert!(recovered.quiet());
    recovered.accept_system("next".into()).unwrap();
    assert_eq!(recovered.boxes().last().unwrap().id(), BoxId::new(4));
}

#[test]
fn structured_results_have_one_followup_and_duplicate_rules() {
    let mut state = ActorState::new(true);
    let active = start(&mut state);
    let calls = state
        .append_stage(
            active.job,
            String::new(),
            vec![provider("ok"), provider("error"), provider("idle")],
        )
        .unwrap();

    state
        .accept_async_return(result(&calls[0], ToolResultStatus::Ok))
        .unwrap();
    state
        .accept_async_return(result(&calls[1], ToolResultStatus::Error))
        .unwrap();
    assert_eq!(state.flush_active_arrivals(active.job).unwrap().len(), 2);
    assert!(
        state
            .accept_async_return(result(&calls[0], ToolResultStatus::Ok))
            .is_err()
    );

    state.complete_inference(active.job, String::new()).unwrap();
    assert!(state.quiet());
    state
        .accept_async_return(result(&calls[2], ToolResultStatus::Ok))
        .unwrap();
    finish_followup(&mut state);
}

#[test]
fn stall_restart_and_halt_preserve_job_ownership() {
    let mut retrying = ActorState::new(false);
    retrying.accept_user("ask".into()).unwrap();
    let first = start(&mut retrying);
    retrying.stall_inference(first.job, "retry".into()).unwrap();
    assert_eq!(retrying.take_halt(), Some("retry".into()));
    retrying.restart().unwrap();

    let second = start(&mut retrying);
    assert!(second.job > first.job);
    assert_eq!(second.frontier, first.frontier);
    retrying
        .complete_inference(second.job, String::new())
        .unwrap();

    let mut halted = ActorState::new(true);
    let active = start(&mut halted);
    assert!(halted.halt("stop".into()));
    assert_eq!(halted.restart(), Err(StateError::Busy));
    halted
        .complete_inference(active.job, String::new())
        .unwrap();
    assert_eq!(halted.take_halt(), Some("stop".into()));
    assert!(halted.quiet());
}