clawmates/crates/tc-runtime/tests/run_loop.rs

use std::sync::Arc;

use tc_domain::{
    AccessPolicy, Agent, AgentId, AgentStatus, MessageRole, Role, RunState, User, UserId,
    Workspace, WorkspaceId,
};
use tc_llm::ScriptedProvider;
use tc_runtime::{RunEventBody, Runtime, RuntimeConfig};

const SCENARIOS: &str = r#"
[[scenario]]
marker = "[[scenario:tool-time]]"

[[scenario.turns]]
events = [
    { type = "text", text = "Let me check the clock." },
    { type = "tool_use", name = "clock.now", input = {} },
]

[[scenario.turns]]
events = [
    { type = "text", text = " It is now known." },
]
"#;

async fn seeded(pool: &sqlx::PgPool) -> Agent {
    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    tc_db::repo::workspaces::insert(pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    tc_db::repo::users::insert(pool, &owner).await.unwrap();
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: ws.id,
        name: "Scout".into(),
        job_title: "Analyst".into(),
        system_prompt: "You are concise.".into(),
        avatar: String::new(),
        accent: String::new(),
        wallpaper: String::new(),
        managed_by: owner.id,
        status: AgentStatus::Online,
    };
    tc_db::repo::agents::insert(pool, &agent, &AccessPolicy::default())
        .await
        .unwrap();
    agent
}

fn runtime(pool: sqlx::PgPool) -> Runtime {
    Runtime::new(
        pool,
        Arc::new(ScriptedProvider::from_toml(SCENARIOS).unwrap()),
        RuntimeConfig::basic("scripted", 1024),
    )
}

async fn drain(
    mut rx: tokio::sync::broadcast::Receiver<tc_runtime::RunEventEnvelope>,
) -> Vec<tc_runtime::RunEventEnvelope> {
    let mut events = Vec::new();
    while let Ok(envelope) = rx.recv().await {
        let done = matches!(
            envelope.event,
            RunEventBody::RunCompleted { .. } | RunEventBody::Error { .. }
        );
        events.push(envelope);
        if done {
            break;
        }
    }
    events
}

#[tokio::test]
async fn plain_message_streams_and_persists_a_reply() {
    let pool = tc_testkit::test_pool().await;
    let agent = seeded(&pool).await;
    let session = tc_db::repo::sessions::create(&pool, agent.id, agent.workspace_id, "Chat")
        .await
        .unwrap();
    let rt = runtime(pool.clone());

    let started = rt.send_message(session.id, "hello there").await.unwrap();
    let events = drain(started.events).await;

    // Envelope sequence is strictly increasing from 1.
    let seqs: Vec<i64> = events.iter().map(|e| e.seq).collect();
    assert_eq!(seqs, (1..=seqs.len() as i64).collect::<Vec<_>>());
    assert!(matches!(events[0].event, RunEventBody::RunStarted { .. }));

    let text: String = events
        .iter()
        .filter_map(|e| match &e.event {
            RunEventBody::TextDelta { delta } => Some(delta.as_str()),
            _ => None,
        })
        .collect();
    assert_eq!(text, "I received: hello there");

    // Transcript persisted: user message + agent reply with that text.
    let history = tc_db::repo::messages::history(&pool, session.id)
        .await
        .unwrap();
    assert_eq!(history.len(), 2);
    assert_eq!(history[0].message.role, MessageRole::User);
    assert_eq!(history[1].message.role, MessageRole::Agent);
    assert_eq!(
        history[1].message.content["text"],
        "I received: hello there"
    );

    let run = tc_db::repo::runs::get(&pool, started.run_id).await.unwrap();
    assert_eq!(run.state, RunState::Completed);
}

#[tokio::test]
async fn tool_scenario_executes_real_tool_and_records_steps() {
    let pool = tc_testkit::test_pool().await;
    let agent = seeded(&pool).await;
    let session = tc_db::repo::sessions::create(&pool, agent.id, agent.workspace_id, "Chat")
        .await
        .unwrap();
    let rt = runtime(pool.clone());

    let started = rt
        .send_message(session.id, "time? [[scenario:tool-time]]")
        .await
        .unwrap();
    let events = drain(started.events).await;

    let step_started = events.iter().find_map(|e| match &e.event {
        RunEventBody::StepStarted { tool, .. } => Some(tool.clone()),
        _ => None,
    });
    assert_eq!(step_started.as_deref(), Some("clock.now"));
    let step_finished = events.iter().find_map(|e| match &e.event {
        RunEventBody::StepFinished { status, output, .. } => Some((status.clone(), output.clone())),
        _ => None,
    });
    let (status, output) = step_finished.expect("step finished event");
    assert_eq!(status, "ok");
    assert!(output["now"].is_string(), "clock output: {output}");

    // Both text legs land in one agent message.
    let history = tc_db::repo::messages::history(&pool, session.id)
        .await
        .unwrap();
    let reply = &history[1];
    let text = reply.message.content["text"].as_str().unwrap();
    assert!(text.contains("Let me check the clock."), "got: {text}");
    assert!(text.contains("It is now known."), "got: {text}");

    // The step trace is persisted on the reply (the "N steps" UI source).
    assert_eq!(reply.steps.len(), 1);
    assert_eq!(reply.steps[0].tool_name.as_deref(), Some("clock.now"));
    assert_eq!(reply.steps[0].status, tc_domain::StepStatus::Ok);
}

#[tokio::test]
async fn journal_matches_what_subscribers_saw() {
    let pool = tc_testkit::test_pool().await;
    let agent = seeded(&pool).await;
    let session = tc_db::repo::sessions::create(&pool, agent.id, agent.workspace_id, "Chat")
        .await
        .unwrap();
    let rt = runtime(pool.clone());

    let started = rt.send_message(session.id, "ping").await.unwrap();
    let live = drain(started.events).await;

    // Replay from the journal reproduces the exact live sequence (the
    // property reconnect/resumeFrom depends on).
    let journal = tc_db::repo::run_events::list_after(&pool, started.run_id, 0)
        .await
        .unwrap();
    assert_eq!(journal.len(), live.len());
    for (persisted, observed) in journal.iter().zip(live.iter()) {
        assert_eq!(persisted.seq, observed.seq);
        let live_json = serde_json::to_value(&observed.event).unwrap();
        assert_eq!(persisted.event_type, live_json["type"].as_str().unwrap());
    }
}

#[tokio::test]
async fn unknown_tool_fails_the_step_but_not_the_run() {
    let pool = tc_testkit::test_pool().await;
    let agent = seeded(&pool).await;
    let session = tc_db::repo::sessions::create(&pool, agent.id, agent.workspace_id, "Chat")
        .await
        .unwrap();
    let scenarios = r#"
[[scenario]]
marker = "[[scenario:bad-tool]]"

[[scenario.turns]]
events = [ { type = "tool_use", name = "no.such.tool", input = {} } ]

[[scenario.turns]]
events = [ { type = "text", text = "I could not use that tool." } ]
"#;
    let rt = Runtime::new(
        pool.clone(),
        Arc::new(ScriptedProvider::from_toml(scenarios).unwrap()),
        RuntimeConfig::basic("scripted", 1024),
    );

    let started = rt
        .send_message(session.id, "x [[scenario:bad-tool]]")
        .await
        .unwrap();
    let events = drain(started.events).await;

    let failed_step = events.iter().any(
        |e| matches!(&e.event, RunEventBody::StepFinished { status, .. } if status == "error"),
    );
    assert!(failed_step, "expected an errored step");
    // The error went back to the model as a tool result and it continued.
    let run = tc_db::repo::runs::get(&pool, started.run_id).await.unwrap();
    assert_eq!(run.state, RunState::Completed);
}

#[tokio::test]
async fn completed_runs_are_metered_and_decrement_credits() {
    let pool = tc_testkit::test_pool().await;
    let agent = seeded(&pool).await;
    tc_db::repo::credits::add_lot(&pool, agent.workspace_id, 10, "seed")
        .await
        .unwrap();
    let rt = runtime(pool.clone());
    let session = tc_db::repo::sessions::create(&pool, agent.id, agent.workspace_id, "Bill")
        .await
        .unwrap();

    let started = rt.send_message(session.id, "ping").await.unwrap();
    drain(started.events).await;

    // Scripted accounting: 1 word in, "I received: ping" = 3 words out →
    // 4 tokens → 1 credit.
    let mut metered = false;
    for _ in 0..100 {
        let row = sqlx::query!(
            r#"SELECT tokens_in, tokens_out, credits FROM usage_events
               WHERE workspace_id = $1"#,
            agent.workspace_id.as_uuid(),
        )
        .fetch_optional(&pool)
        .await
        .unwrap();
        if let Some(usage) = row {
            assert_eq!(usage.tokens_in, 1);
            assert_eq!(usage.tokens_out, 3);
            assert_eq!(
                tc_db::repo::credits::balance(&pool, agent.workspace_id)
                    .await
                    .unwrap(),
                9
            );
            metered = true;
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }
    assert!(metered, "usage event never recorded");
}