kcode-k1-chat-thread-durable-state 0.4.17

use super::*;
use kcode_k1_access::K1Access;
use kcode_k1_chat_codex_state::Call;
use kcode_k1_chat_persistence::{K1ChatPersistence, Session};
use kcode_k1_chat_state::{AGENT_RESPONSE_TYPE, TOOL_CALL_TYPE, TOOL_RESULT_TYPE};
use kcode_k1_groups::K1Groups;
use kcode_k1_kmap::K1Kmap;
use kcode_k1_peering::K1Peering;
use kcode_k1_txn_ordering::K1TxnOrdering;
use tempfile::TempDir;

fn fixture() -> (TempDir, Session, Arc<K1AccessKmap>) {
    let root = TempDir::new().unwrap();
    let path = |name| root.path().join(name);
    let ordering = Arc::new(K1TxnOrdering::open(&path("ordering")).unwrap());
    let peering = Arc::new(K1Peering::open(&path("peering"), ordering.clone()).unwrap());
    let groups =
        Arc::new(K1Groups::open(&path("groups"), ordering.clone(), peering.clone()).unwrap());
    let access = Arc::new(
        K1Access::open(&path("access"), ordering.clone(), peering.clone(), groups).unwrap(),
    );
    let kmap = Arc::new(K1Kmap::open(&path("kmap"), ordering.clone(), peering.clone()).unwrap());
    let access_kmap = Arc::new(K1AccessKmap::open(access, kmap).unwrap());
    let persistence = K1ChatPersistence::open(&path("persistence"), ordering, peering).unwrap();
    let (session, original) = persistence.session([11; 12]).unwrap();
    assert!(original.records.is_empty());
    (root, session, access_kmap)
}

fn accept_user_box(thread: &mut DurableThread, contents: &str) {
    thread
        .accept_box(
            USER_MESSAGE_TYPE.into(),
            contents.into(),
            String::new(),
            String::new(),
        )
        .unwrap();
}

#[test]
fn preflight_support_is_bound_to_the_concrete_executor() {
    let (_root, session, access_kmap) = fixture();
    let thread = DurableThread::recover(session, access_kmap).unwrap();
    assert_eq!(
        thread.preflight_executor().supports("CurrentTime"),
        Ok(true)
    );
    assert_eq!(thread.preflight_executor().supports("WebSearch"), Ok(false));
    assert!(thread.boxes().is_empty());
}

#[test]
fn send_message_is_durable_ordered_and_recovered_once() {
    let (_root, session, access_kmap) = fixture();
    let mut thread = DurableThread::recover(session.clone(), Arc::clone(&access_kmap)).unwrap();
    accept_user_box(&mut thread, "hello");
    let (job, _) = thread.begin_input().unwrap().unwrap();
    let arguments = r#"{"message":"WORKING_MESSAGE"}"#;
    let calls = thread
        .prepare_stage(
            job,
            String::new(),
            vec![BoxValue::Call(Ok(Call {
                name: "SendMessage".into(),
                arguments: arguments.into(),
            }))],
        )
        .unwrap();
    assert_eq!(calls.len(), 1);
    let result = thread.launch_action("SendMessage", arguments);
    assert_eq!(result, Ok("success".into()));
    thread
        .accept_tool_return(calls[0].tool_call_id, result)
        .unwrap();

    let prepared = thread.prepare_mailbox_flush(job).unwrap().unwrap();
    let input = thread.prepared_input(&prepared).unwrap();
    let call_at = input.find("| Tool Call]").unwrap();
    let message_at = input.find("| Agent Message]").unwrap();
    let result_at = input.find("| Tool Result]").unwrap();
    assert!(call_at < message_at && message_at < result_at);
    assert!(input.ends_with("| Agent Response]\n"));
    assert_eq!(
        thread
            .boxes()
            .iter()
            .map(ChatBox::box_type)
            .collect::<Vec<_>>(),
        [
            USER_MESSAGE_TYPE,
            AGENT_RESPONSE_TYPE,
            TOOL_CALL_TYPE,
            AGENT_MESSAGE_TYPE,
            TOOL_RESULT_TYPE,
        ]
    );
    let message = thread
        .boxes()
        .iter()
        .find(|value| value.box_type() == AGENT_MESSAGE_TYPE)
        .unwrap();
    assert_eq!(message.contents(), "WORKING_MESSAGE");
    assert_eq!((message.hidden_type(), message.hidden_contents()), ("", ""));

    thread.commit_mailbox_flush(prepared).unwrap();
    assert!(
        !thread
            .complete(job, ShimOutput { items: Vec::new() })
            .unwrap()
    );
    drop(thread);

    let mut recovered = DurableThread::recover(session, access_kmap).unwrap();
    assert_eq!(
        recovered
            .boxes()
            .iter()
            .filter(|value| value.box_type() == AGENT_MESSAGE_TYPE)
            .count(),
        1
    );
    let before = recovered.boxes().len();
    assert!(recovered.launch_action("CurrentTime", "{}").is_ok());
    assert_eq!(
        recovered.launch_action("NotRegistered", "{}"),
        Err("unknown Ktool".into())
    );
    assert_eq!(
        recovered.launch_action("SetLaunchNode", "{}"),
        Err("unknown Ktool".into())
    );
    let social = recovered.launch_action("ListContacts", "{}");
    assert_eq!(social, Err("social Ktools are unavailable".into()));
    assert_eq!(recovered.boxes().len(), before);
}

#[test]
fn complete_returns_terminal_response_before_queued_arrivals() {
    let (_root, session, access_kmap) = fixture();
    let mut thread = DurableThread::recover(session, access_kmap).unwrap();
    accept_user_box(&mut thread, "first");
    let (job, _) = thread.begin_input().unwrap().unwrap();
    let terminal_index = thread.boxes().len();
    accept_user_box(&mut thread, "queued");

    let (resume, terminal_id) = thread
        .complete_with_terminal_response(job, ShimOutput { items: Vec::new() })
        .unwrap();
    assert!(resume);
    assert_eq!(terminal_id, thread.boxes()[terminal_index].id().get());
    assert_eq!(
        thread.boxes()[terminal_index].box_type(),
        AGENT_RESPONSE_TYPE
    );
    assert_eq!(
        thread.boxes()[terminal_index + 1].box_type(),
        USER_MESSAGE_TYPE
    );
}