rig-ecs 0.44.0

//! Steering as systems (CONTRACT §9): the components a user writes and
//! what the library does with them, in this crate's own words — the
//! corpus's hook cells are the full proof (`corpus/world_hooks.rs`).
//!
//! | claim | test |
//! |---|---|
//! | `Cancelled` on a run before its first turn: no record, the reason is the failure's | `cancelled_at_start_dispatches_nothing_and_names_the_reason` |
//! | `Cancelled` in `Patch`: the folded completion is despawned before the bus | `cancelled_in_patch_leaves_no_record` |
//! | `RequestPatch` on the fresh turn is folded in: the preamble, a document, the history with the prompt kept | `a_request_patch_is_folded_into_the_turn` |
//! | `Retry { feedback }` on a text turn makes the turn and the feedback history, then another turn | `a_retry_with_feedback_asks_again` |
//! | `Retry { feedback }` on an empty turn makes the feedback history, not the turn, then another turn; without a retry the empty turn settles | `a_retry_written_on_an_empty_turn_asks_again` |
//! | `UsesModel` written on the run before `Select` routes the turn; `Route` is in the required row | `uses_model_written_before_select_routes_the_turn` |
//! | a decision written and not yet read survives a checkpoint (§10 of the dissolves doc) | `a_retry_written_before_a_save_is_read_after_the_load` |
//! | `RequestPatch` on the turn and `Resolution` on the invalid call survive a checkpoint the same way | `a_patch_and_a_resolution_written_before_a_save_are_read_after_the_load` |
//! | `Cancelled` is read at once; the checkpoint carries it and the ending it made, and the loaded run stays ended when the re-issued effect answers | `a_cancel_written_before_a_save_is_the_ending_after_the_load` |

use crate::run_support;

use std::sync::{Arc, Mutex};

use bevy_ecs::prelude::*;
use rig_cassette::ecs::EffectLogResource;
use rig_cassette::ecs::identity::required_row;
use rig_cassette::effect_log::EffectLogRecorder;
use rig_core::{
    effect::{EffectFamily, HandlerKey},
    error::ErrorKind,
    message::{AssistantContent, Message},
};
use rig_ecs::{
    agent::{
        Cancelled, Cursor, DocumentId, DocumentText, Failed, Failure, InvalidCall, RequestPatch,
        Resolution, Retry, Route, RunResult, Settled, UsesModel,
    },
    bus::{Handlers, PendingEffect, RigSchedule},
    checkpoint::{Checkpoint, RestoreMode, load_world, save_world},
    systems::{Fresh, RigSet, RunCommands},
};
use run_support::*;

const MODEL: &str = "t/model:default";
const FAST: &str = "t/model:fast";

fn add_system<M>(
    app: &mut bevy_app::App,
    system: impl IntoScheduleConfigs<bevy_ecs::system::ScheduleSystem, M>,
) {
    app.world_mut()
        .resource_mut::<Schedules>()
        .add_systems(RigSchedule, system);
}

fn records(app: &bevy_app::App) -> Vec<String> {
    app.world()
        .resource::<EffectLogResource>()
        .log()
        .records
        .iter()
        .map(|record| record.key.to_string())
        .collect()
}

#[test]
fn cancelled_at_start_dispatches_nothing_and_names_the_reason() {
    let mut app = app();
    EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
    let (agent, _) = capturing_agent(&mut app, MODEL, MODEL, "never");
    let run = app.world_mut().spawn_run(agent, &[], "go", false, None);
    app.world_mut()
        .entity_mut(run)
        .insert(Cancelled("stopped at run start".to_owned()));
    ended(&mut app, run, "cancelled");
    let Some(Failed(Failure::Cancelled(report))) = app.world().get::<Failed>(run) else {
        panic!("cancelled");
    };
    assert_eq!(report.kind, ErrorKind::Cancelled);
    assert_eq!(report.message, "stopped at run start");
    app.update();
    assert!(records(&app).is_empty(), "nothing was dispatched");
}

fn stop_in_patch(
    folded: Query<&ChildOf, Added<PendingEffect>>,
    turns: Query<&ChildOf>,
    mut commands: Commands,
) {
    for turn_of in &folded {
        if let Ok(run_of) = turns.get(turn_of.parent()) {
            commands
                .entity(run_of.parent())
                .insert(Cancelled("stopped before the completion call".to_owned()));
        }
    }
}

#[test]
fn cancelled_in_patch_leaves_no_record() {
    let mut app = app();
    EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
    let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "never");
    add_system(&mut app, stop_in_patch.in_set(RigSet::Patch));
    let run = app.world_mut().spawn_run(agent, &[], "go", false, None);
    ended(&mut app, run, "cancelled");
    app.update();
    assert!(
        records(&app).is_empty(),
        "the folded effect never reached the bus"
    );
    assert!(requests.lock().unwrap().is_empty());
    assert!(matches!(
        app.world().get::<Failed>(run),
        Some(Failed(Failure::Cancelled(report))) if report.message == "stopped before the completion call"
    ));
}

fn patch_the_turn(fresh: Query<Entity, Added<Fresh>>, mut commands: Commands) {
    for turn in &fresh {
        commands.entity(turn).insert(RequestPatch {
            preamble: Some("You are a pirate.".to_owned()),
            extra_context: vec![rig_core::completion::Document {
                id: "extra".to_owned(),
                text: "a glarb-glarb".to_owned(),
                additional_props: Default::default(),
            }],
            history: Some(
                [
                    Message::user("My name is Ada."),
                    Message::assistant("Hello, Ada."),
                ]
                .iter()
                .filter_map(rig_ecs::agent::MessageParts::from_message)
                .collect(),
            ),
            ..RequestPatch::default()
        });
    }
}

#[test]
fn a_request_patch_is_folded_into_the_turn() {
    let mut app = app();
    let (agent, requests) = capturing_agent(&mut app, MODEL, MODEL, "Ada");
    let document = app
        .world_mut()
        .spawn((
            DocumentId("static".to_owned()),
            DocumentText("always".to_owned()),
        ))
        .id();
    app.world_mut()
        .spawn((rig_ecs::agent::Context(document), ChildOf(agent)));
    add_system(
        &mut app,
        patch_the_turn
            .after(RigSet::Advance)
            .before(RigSet::Assemble),
    );
    let run = app
        .world_mut()
        .spawn_run(agent, &[], "What is my name?", false, None);
    ended(&mut app, run, "answered");
    let requests = requests.lock().unwrap();
    assert_eq!(
        texts(&requests[0]),
        vec![
            "system:You are a pirate.",
            "user:My name is Ada.",
            "assistant:Hello, Ada.",
            "user:What is my name?"
        ],
        "the preamble replaced, the history replaced, the prompt kept"
    );
    let documents: Vec<&str> = requests[0]
        .documents
        .iter()
        .map(|d| d.id.as_str())
        .collect();
    assert_eq!(
        documents,
        ["static", "extra"],
        "the patch's document after the turn's"
    );
}

fn demand_done(
    turns: Query<
        (Entity, &rig_ecs::agent::Outputs),
        (Without<rig_ecs::systems::Materialised>, Without<Retry>),
    >,
    mut commands: Commands,
) {
    for (turn, outs) in &turns {
        if outs.done
            && !rig_core::transcript::assistant_text_from_choice(&outs.content).contains("DONE")
        {
            commands.entity(turn).insert(Retry {
                feedback: Some("End with DONE.".to_owned()),
            });
        }
    }
}

#[test]
fn a_retry_with_feedback_asks_again() {
    let mut app = app();
    let (agent, requests) = scripted_agent(
        &mut app,
        MODEL,
        vec![
            vec![AssistantContent::text("first")],
            vec![AssistantContent::text("second DONE")],
        ],
    );
    app.world_mut()
        .entity_mut(agent)
        .insert(rig_ecs::agent::MaxTurns(3));
    add_system(&mut app, demand_done.in_set(RigSet::Judge));
    let run = app.world_mut().spawn_run(agent, &[], "say it", false, None);
    ended(&mut app, run, "answered");
    assert_eq!(
        app.world().get::<RunResult>(run).map(|r| r.0.as_str()),
        Some("second DONE")
    );
    let requests = requests.lock().unwrap();
    assert_eq!(requests.len(), 2);
    assert_eq!(
        texts(&requests[1]),
        vec![
            "system:You are terse.",
            "user:say it",
            "assistant:first",
            "user:End with DONE."
        ]
    );
}

#[test]
fn a_retry_written_on_an_empty_turn_asks_again() {
    // A model that answers nothing (a provider can return only a thought
    // signature, which the fold drops to an empty turn) would settle the
    // run on "". A judge that writes `Retry` on that turn is honoured
    // before the empty settlement: the feedback is history, the empty
    // turn is not, and the next turn answers.
    let mut app = app();
    let (agent, requests) = scripted_agent(
        &mut app,
        MODEL,
        vec![Vec::new(), vec![AssistantContent::text("second DONE")]],
    );
    app.world_mut()
        .entity_mut(agent)
        .insert(rig_ecs::agent::MaxTurns(3));
    add_system(&mut app, demand_done.in_set(RigSet::Judge));
    let run = app.world_mut().spawn_run(agent, &[], "say it", false, None);
    ended(&mut app, run, "answered after an empty turn");
    assert_eq!(
        app.world().get::<RunResult>(run).map(|r| r.0.as_str()),
        Some("second DONE"),
        "{:?}",
        app.world().get::<Failed>(run)
    );
    let requests = requests.lock().unwrap();
    assert_eq!(requests.len(), 2);
    assert_eq!(
        texts(&requests[1]),
        vec![
            "system:You are terse.",
            "user:say it",
            "user:End with DONE."
        ],
        "the empty turn is not history; the feedback is"
    );
}

#[derive(Resource)]
struct Fast(Entity);

fn route_first_turn(
    fresh: Query<&ChildOf, Added<Fresh>>,
    runs: Query<&Cursor>,
    fast: Res<Fast>,
    mut commands: Commands,
) {
    for turn_of in &fresh {
        let run = turn_of.parent();
        if runs.get(run).is_ok_and(|cursor| cursor.turn == 1) {
            commands.entity(run).insert(UsesModel(fast.0));
        }
    }
}

#[test]
fn uses_model_written_before_select_routes_the_turn() {
    let mut app = app();
    EffectLogResource::install(app.world_mut(), EffectLogRecorder::new());
    let (agent, _) = capturing_agent(&mut app, MODEL, MODEL, "slow");
    let (fast, _) = Capturing::new(FAST, "fast");
    let fast = register(&mut app, FAST, fast);
    app.world_mut().spawn((Route(fast), ChildOf(agent)));
    app.insert_resource(Fast(fast));
    add_system(
        &mut app,
        route_first_turn
            .after(RigSet::Advance)
            .before(RigSet::Select),
    );
    let row = required_row(app.world_mut(), agent);
    assert_eq!(
        row.get(&HandlerKey::from(FAST)),
        Some(&EffectFamily::Completion)
    );
    let run = app.world_mut().spawn_run(agent, &[], "go", false, None);
    ended(&mut app, run, "answered");
    assert_eq!(
        app.world().get::<RunResult>(run).map(|r| r.0.as_str()),
        Some("fast")
    );
    assert_eq!(records(&app), [FAST]);
}

#[test]
fn a_retry_written_before_a_save_is_read_after_the_load() {
    // The first world: the turn answered, a `Retry` written on it, the
    // world saved before `Materialise` read it.
    let mut first = app();
    let (agent, _) = scripted_agent(
        &mut first,
        MODEL,
        vec![vec![AssistantContent::text("first")]],
    );
    first
        .world_mut()
        .entity_mut(agent)
        .insert(rig_ecs::agent::MaxTurns(3));
    let _run = first
        .world_mut()
        .spawn_run(agent, &[], "say it", false, None);
    tick_until(&mut first, "the turn answered", |world| {
        world
            .query_filtered::<&rig_ecs::agent::Outputs, With<rig_ecs::agent::Turn>>()
            .iter(world)
            .any(|outs| outs.done)
    });
    // Whether Materialise has read the turn is not the point — the point
    // is that the decision survives: written on the turn, saved, restored.
    let turn = first
        .world_mut()
        .query_filtered::<Entity, With<rig_ecs::agent::Turn>>()
        .iter(first.world())
        .next()
        .expect("a turn");
    first.world_mut().entity_mut(turn).insert(Retry {
        feedback: Some("End with DONE.".to_owned()),
    });
    let checkpoint = save_world(first.world_mut()).expect("serializes");
    let json = checkpoint.to_json().expect("serde");
    assert!(
        json.contains("End with DONE."),
        "the decision is checkpoint state"
    );
    drop(first);

    let mut app = app();
    let (model, _) = Capturing::new(MODEL, "again");
    register(&mut app, MODEL, model);
    let loaded = load_world(
        &Checkpoint::from_json(&json).expect("serde"),
        app.world_mut(),
        RestoreMode::Strict,
        [],
    )
    .expect("loads");
    let turn = loaded.with::<rig_ecs::agent::Turn>(app.world())[0];
    assert_eq!(
        app.world().get::<Retry>(turn),
        Some(&Retry {
            feedback: Some("End with DONE.".to_owned())
        }),
        "the decision restored intact"
    );
    let _ = Handlers::with(app.world_mut(), |handlers| handlers.keys());
}

/// A first world whose model never answers: the run open, its completion
/// in flight, the turn folded — the moment a decision is written on.
fn open_run() -> (bevy_app::App, Entity, Entity) {
    let mut app = app();
    let model = register(
        &mut app,
        MODEL,
        NeverAnswers {
            label: MODEL.to_owned(),
        },
    );
    let agent = spawn_agent(app.world_mut(), "t", model);
    let run = app.world_mut().spawn_run(agent, &[], "say it", false, None);
    tick_until(&mut app, "the completion in flight", |world| {
        world
            .query_filtered::<(), (With<PendingEffect>, With<rig_ecs::bus::InFlight>)>()
            .iter(world)
            .next()
            .is_some()
    });
    assert!(app.world().get::<Settled>(run).is_none());
    let turn = app
        .world_mut()
        .query_filtered::<Entity, With<rig_ecs::agent::Turn>>()
        .iter(app.world())
        .next()
        .expect("a turn");
    (app, run, turn)
}

/// The checkpoint as JSON, loaded into a fresh world with a capturing model:
/// the run, the turn, and the requests the model saw.
fn reload(
    json: &str,
) -> (
    bevy_app::App,
    Entity,
    Entity,
    Arc<Mutex<Vec<rig_core::completion::CompletionRequest>>>,
) {
    let mut app = app();
    let (model, requests) = Capturing::new(MODEL, "again");
    register(&mut app, MODEL, model);
    let loaded = load_world(
        &Checkpoint::from_json(json).expect("serde"),
        app.world_mut(),
        RestoreMode::Strict,
        [],
    )
    .expect("loads");
    let run = loaded.with::<rig_ecs::agent::Run>(app.world())[0];
    let turn = loaded.with::<rig_ecs::agent::Turn>(app.world())[0];
    (app, run, turn, requests)
}

#[test]
fn a_patch_and_a_resolution_written_before_a_save_are_read_after_the_load() {
    // On the entities CONTRACT §9 puts them on: the patch on the turn, the
    // resolution on an invalid call `ChildOf` the turn.
    let (mut first, _run, turn) = open_run();
    let patch = RequestPatch {
        preamble: Some("You are a pirate.".to_owned()),
        max_tokens: Some(64),
        ..Default::default()
    };
    let resolution = Resolution::Retry {
        feedback: "there is no tool named multiply".to_owned(),
    };
    first.world_mut().entity_mut(turn).insert(patch.clone());
    first.world_mut().spawn((
        InvalidCall {
            origin: None,
            id: rig_core::message::CallId::from_wire("c1"),
            name: "multiply".to_owned(),
            arguments: serde_json::json!({"x": 2, "y": 3}),
            prefix: Vec::new(),
            stream_offset: None,
        },
        resolution.clone(),
        ChildOf(turn),
    ));
    let checkpoint = save_world(first.world_mut()).expect("serializes");
    let json = checkpoint.to_json().expect("serde");
    for decision in ["You are a pirate.", "no tool named multiply"] {
        assert!(
            json.contains(decision),
            "{decision}: the decision is checkpoint state"
        );
    }
    drop(first);

    let (app, _run, turn, _) = reload(&json);
    let call = app
        .world()
        .get::<Children>(turn)
        .into_iter()
        .flat_map(|children| children.iter())
        .find(|child| app.world().get::<InvalidCall>(*child).is_some())
        .expect("the invalid call is the turn's");
    assert_eq!(app.world().get::<RequestPatch>(turn), Some(&patch));
    assert_eq!(app.world().get::<Resolution>(call), Some(&resolution));
}

#[test]
fn a_cancel_written_before_a_save_is_the_ending_after_the_load() {
    // A cancel is read at once (an observer): the scene carries the
    // decision and the ending it made, with the reason; the loaded run
    // stays ended.
    let (mut first, run, _turn) = open_run();
    first
        .world_mut()
        .entity_mut(run)
        .insert(Cancelled("stopped before the save".to_owned()));
    let scene = save_world(first.world_mut()).expect("serializes");
    let json = serde_json::to_string(&scene).expect("serde");
    assert!(
        json.contains("stopped before the save"),
        "the decision is scene state"
    );
    drop(first);

    let (mut app, run, _turn, requests) = reload(&json);
    assert_eq!(
        app.world().get::<Cancelled>(run),
        Some(&Cancelled("stopped before the save".to_owned()))
    );
    ended(&mut app, run, "cancelled after the load");
    for _ in 0..3 {
        app.update();
    }
    assert!(matches!(
        app.world().get::<Failed>(run),
        Some(Failed(Failure::Cancelled(report))) if report.message == "stopped before the save"
    ));
    // The completion in flight at the cut was saved as intent and is
    // re-issued to its handler here (CONTRACT §13: an effect in flight is
    // the handler's); its answer does not reopen a run that ended.
    assert_eq!(requests.lock().unwrap().len(), 1);
    assert!(app.world().get::<Settled>(run).is_none());
    assert!(app.world().get::<RunResult>(run).is_none());
}