laser-sdk 0.8.2

use crate::harness;
use bytes::Bytes;
use laser_sdk::agent::StepContext;
use laser_sdk::agent::{Clock, SystemClock};
use laser_sdk::prelude::full::*;
use laser_sdk::wire::agent::{AgentCard, CapabilityDescriptor};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

// A worker that tags its reply with its role and echoes the input, so a step's
// output reflects both its own work and what the prior step produced.
struct Tagger {
    tag: String,
}

impl AgentHandler for Tagger {
    async fn handle(&self, message: &AgentMessage, ctx: &AgentCtx<'_>) -> Result<(), LaserError> {
        // `body()` is the task body for both AGDX commands (workflow/contract) and
        // plain sends, so the worker need not know how it was reached.
        let reply = format!("{}:{}", self.tag, String::from_utf8_lossy(message.body()));
        ctx.respond(Bytes::from(reply)).await
    }
}

fn card(skill: &str) -> AgentCard {
    AgentCard {
        name: None,
        version: None,
        capabilities: vec![CapabilityDescriptor {
            skill_id: skill.to_owned(),
            input: None,
            output: None,
            cost_class: None,
            latency_class: None,
            max_concurrency: None,
            health: None,
            load: None,
        }],
        ttl_micros: None,
    }
}

// Each capability-bearing worker opens its own connection: presence is
// connection-scoped and one connection advertises one agent.
async fn worker(laser: &Laser, id: &str, skill: &str, tag: &str) -> AgentHandle {
    let connection = harness::reconnect(laser).await;
    Agent::builder()
        .id(id.parse().expect("worker id is valid"))
        .listen_on(AgentTopic::Sessions)
        .respond_on(AgentTopic::Sessions)
        .capabilities(card(skill).capabilities)
        .handler(Tagger {
            tag: tag.to_owned(),
        })
        .build()
        .spawn(connection)
}

async fn wait_for_capable(laser: &Laser, skill: &str, expected: usize) {
    let skill = skill.to_owned();
    harness::eventually(|| {
        let laser = laser.clone();
        let skill = skill.clone();
        async move {
            let now = SystemClock.now_micros();
            let mut registry = laser.agent_registry().ok()?;
            registry.refresh(now).await.ok()?;
            (registry.resolve(&skill, now).len() == expected).then_some(())
        }
    })
    .await;
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_concurrent_card_publishers_when_registry_refreshes_then_should_observe_every_card() {
    let laser = harness::laser().await;
    let publisher_a = harness::reconnect(&laser).await;
    let publisher_b = harness::reconnect(&laser).await;
    let card_a = card("publish-a");
    let card_b = card("publish-b");

    let (published_a, published_b) = tokio::join!(
        publisher_a.publish_card("publisher-a".parse().unwrap(), &card_a),
        publisher_b.publish_card("publisher-b".parse().unwrap(), &card_b),
    );
    published_a.expect("publisher a writes its card");
    published_b.expect("publisher b writes its card");

    wait_for_capable(&laser, "publish-a", 1).await;
    wait_for_capable(&laser, "publish-b", 1).await;
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_a_two_step_workflow_when_run_then_should_dispatch_in_dependency_order_and_thread_outputs()
 {
    let laser = harness::laser().await;
    let mut triager = worker(&laser, "triager", "triage", "triaged").await;
    let mut diagnoser = worker(&laser, "diagnoser", "diagnose", "diagnosed").await;
    triager.ready().await.expect("triager joins its group");
    diagnoser.ready().await.expect("diagnoser joins its group");
    wait_for_capable(&laser, "triage", 1).await;
    wait_for_capable(&laser, "diagnose", 1).await;

    // Each step is target-filtered to its worker on the shared topic.
    let outcome = laser
        .workflow("incidentflow")
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "triage",
            Router::to_capable("triage", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"incident").to_vec(),
        )
        .step(
            "diagnose",
            Router::to_capable("diagnose", RoutePolicy::Any),
            // Thread the triage output into the diagnose task.
            |ctx: &StepContext<'_>| ctx.outputs.get("triage").cloned().unwrap_or_default(),
        )
        .after("triage")
        .run()
        .await
        .expect("the workflow runs to completion");

    assert_eq!(
        outcome.outputs.get("triage").map(|o| o.as_slice()),
        Some(b"triaged:incident".as_slice()),
    );
    assert_eq!(
        outcome.outputs.get("diagnose").map(|o| o.as_slice()),
        Some(b"diagnosed:triaged:incident".as_slice()),
        "the diagnose step must receive the triage step's output",
    );

    triager.shutdown().await.expect("triager shuts down");
    diagnoser.shutdown().await.expect("diagnoser shuts down");
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_an_all_capable_step_when_run_then_should_scatter_and_fold_every_reply() {
    let laser = harness::laser().await;
    let mut reviewer_a = worker(&laser, "rev-a", "review", "a-reviewed").await;
    let mut reviewer_b = worker(&laser, "rev-b", "review", "b-reviewed").await;
    reviewer_a
        .ready()
        .await
        .expect("reviewer a joins its group");
    reviewer_b
        .ready()
        .await
        .expect("reviewer b joins its group");
    wait_for_capable(&laser, "review", 2).await;

    let outcome = laser
        .workflow("reviewflow")
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "review",
            Router::all_capable("review", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"check").to_vec(),
        )
        .run()
        .await
        .expect("the all-capable workflow runs");

    let folded = String::from_utf8(
        outcome
            .outputs
            .get("review")
            .cloned()
            .expect("the review step produced an output"),
    )
    .expect("the folded output is utf8");
    // Both reviewers' replies are folded (order is nondeterministic across the
    // concurrent scatter).
    assert!(
        folded.contains("a-reviewed:check") && folded.contains("b-reviewed:check"),
        "both reviewers' replies must be folded, got {folded:?}",
    );

    reviewer_a.shutdown().await.expect("reviewer a shuts down");
    reviewer_b.shutdown().await.expect("reviewer b shuts down");
}

// Counts how many times it is dispatched to, so a resumed run can prove a
// journalled step is not re-executed.
struct Counter {
    calls: Arc<AtomicUsize>,
}

impl AgentHandler for Counter {
    async fn handle(&self, _message: &AgentMessage, ctx: &AgentCtx<'_>) -> Result<(), LaserError> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        ctx.respond(Bytes::from_static(b"counted")).await
    }
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_a_completed_step_when_the_workflow_resumes_then_should_not_re_dispatch_it() {
    let laser = harness::laser().await;
    let calls = Arc::new(AtomicUsize::new(0));
    let mut counter = Agent::builder()
        .id("counter".parse().expect("counter id is valid"))
        .listen_on(AgentTopic::Sessions)
        .respond_on(AgentTopic::Sessions)
        .capabilities(card("count").capabilities)
        .handler(Counter {
            calls: calls.clone(),
        })
        .build()
        .spawn(laser.clone());
    counter.ready().await.expect("counter joins its group");

    // First run dispatches the step once.
    let first = laser
        .workflow("countflow")
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "count",
            Router::to_capable("count", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"go").to_vec(),
        )
        .run()
        .await
        .expect("the first run completes");
    assert_eq!(calls.load(Ordering::SeqCst), 1);

    // Resuming the same run id replays the journal: the completed step keeps its
    // recorded output and is not dispatched again.
    let resumed = laser
        .workflow("countflow")
        .run_id(first.run_id)
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "count",
            Router::to_capable("count", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"go").to_vec(),
        )
        .run()
        .await
        .expect("the resumed run completes");

    assert_eq!(
        calls.load(Ordering::SeqCst),
        1,
        "a journalled step must not be re-dispatched on resume",
    );
    assert_eq!(
        resumed.outputs.get("count").map(|o| o.as_slice()),
        Some(b"counted".as_slice()),
        "resume re-derives the recorded output",
    );

    counter.shutdown().await.expect("counter shuts down");
}

struct SagaWorker {
    dispatched: Arc<AtomicUsize>,
    compensated: Arc<AtomicUsize>,
}

impl AgentHandler for SagaWorker {
    async fn handle(&self, message: &AgentMessage, ctx: &AgentCtx<'_>) -> Result<(), LaserError> {
        if message.body() == b"undo" {
            self.compensated.fetch_add(1, Ordering::SeqCst);
            return ctx.respond(Bytes::from_static(b"undone")).await;
        }
        self.dispatched.fetch_add(1, Ordering::SeqCst);
        ctx.respond(Bytes::from_static(b"done")).await
    }
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_a_budget_breach_after_a_completed_step_when_the_workflow_fails_then_should_run_compensation()
 {
    let laser = harness::laser().await;
    let dispatched = Arc::new(AtomicUsize::new(0));
    let compensated = Arc::new(AtomicUsize::new(0));
    let mut worker = Agent::builder()
        .id("saga-worker".parse().expect("worker id is valid"))
        .listen_on(AgentTopic::Sessions)
        .respond_on(AgentTopic::Sessions)
        .capabilities(card("saga").capabilities)
        .handler(SagaWorker {
            dispatched: dispatched.clone(),
            compensated: compensated.clone(),
        })
        .build()
        .spawn(laser.clone());
    worker.ready().await.expect("worker joins its group");

    let result = laser
        .workflow("budgetflow")
        .budget(WorkflowBudget::unlimited().invocations(1))
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "charge",
            Router::to_capable("saga", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"charge").to_vec(),
        )
        .compensate_with(|_ctx: &StepContext<'_>| Bytes::from_static(b"undo").to_vec())
        .step(
            "ship",
            Router::to_capable("saga", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| Bytes::from_static(b"ship").to_vec(),
        )
        .after("charge")
        .run()
        .await;

    assert!(matches!(result, Err(LaserError::BudgetExceeded { .. })));
    assert_eq!(dispatched.load(Ordering::SeqCst), 1);
    assert_eq!(compensated.load(Ordering::SeqCst), 1);

    worker.shutdown().await.expect("worker shuts down");
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_a_workflow_when_run_then_should_record_the_run_as_a_root_session_with_child_steps() {
    use laser_sdk::wire::agent::{SessionStart, TaskState};
    use laser_sdk::wire::framing::decode_named;
    let laser = harness::laser().await;
    let mut triager = worker(&laser, "rooted-triager", "rooted-triage", "triaged").await;
    triager.ready().await.expect("triager joins its group");
    wait_for_capable(&laser, "rooted-triage", 1).await;
    let outcome = laser
        .workflow("rootedflow")
        .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
        .step(
            "triage",
            Router::to_capable("rooted-triage", RoutePolicy::Any),
            |_ctx: &StepContext<'_>| b"incident".to_vec(),
        )
        .run()
        .await
        .expect("the workflow runs to completion");
    let run = outcome.run_id;
    let child = ConversationId::derive(&format!("{run}/triage/1"));
    let states = |session: ConversationId| {
        let laser = laser.clone();
        async move {
            let turns = laser.sessions().open(session).context().await.ok()?;
            let states: Vec<(TaskState, Vec<u8>)> = turns
                .into_iter()
                .filter_map(|turn| {
                    let envelope = turn.message.envelope?;
                    if envelope.operation.as_deref() != Some("session") {
                        return None;
                    }
                    Some((envelope.task_state?, envelope.body))
                })
                .collect();
            Some(states)
        }
    };
    let root = harness::eventually(|| async {
        let found = states(run).await?;
        (found.len() >= 2).then_some(found)
    })
    .await;
    assert_eq!(root[0].0, TaskState::Working);
    assert_eq!(root.last().map(|state| state.0), Some(TaskState::Completed));
    let step = harness::eventually(|| async {
        let found = states(child).await?;
        found
            .iter()
            .any(|state| state.0 == TaskState::Completed)
            .then_some(found)
    })
    .await;
    assert_eq!(step[0].0, TaskState::Submitted);
    let start: SessionStart = decode_named(&step[0].1).expect("the child start decodes");
    assert_eq!(start.parent, Some(run.into()));
    assert_eq!(start.root, Some(run.into()));
    triager.shutdown().await.expect("triager stops");
}

#[tokio::test]
#[serial_test::serial(integration)]
async fn given_a_cancel_request_on_the_control_topic_when_a_workflow_runs_then_should_stop_as_cancelled()
 {
    let laser = harness::laser().await;
    laser
        .topic(AgentTopic::Control.topic_string())
        .ensure(4)
        .await
        .expect("provisioning creates the control topic");
    let run = ConversationId::new();
    laser
        .sessions()
        .control(laser.default_stream().expect("stream"), run)
        .as_operator(
            "operator"
                .parse::<laser_sdk::types::AgentId>()
                .expect("valid agent id"),
        )
        .cancel()
        .await
        .expect("the cancel request is sent");
    let result = harness::eventually(|| async {
        let result = laser
            .workflow("cancelledflow")
            .run_id(run)
            .inbox_route(InboxRoute::Fixed(AgentTopic::Sessions))
            .step(
                "never",
                Router::to("nobody".parse().expect("valid agent id")),
                |_ctx: &StepContext<'_>| b"x".to_vec(),
            )
            .run()
            .await;
        matches!(result, Err(LaserError::Cancelled { .. })).then_some(result)
    })
    .await;
    assert!(matches!(result, Err(LaserError::Cancelled { .. })));
}