johnandersen777.bsky.social/asyncai

Switch repository

Loading…

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

use async_trait::async_trait;
use asyncai::actions::{
    AddMessageAction, AgiAction, NewAgentAction, NewThreadAction, RunThreadAction,
};
use asyncai::backend::{
    AgentBackend, BackendError, BackendEvent, BackendStream, MockBackend, RunOutcome, RunRequest,
};
use asyncai::error::AgiError;
use asyncai::kv::MemoryKvStore;
use asyncai::message::{AgentContent, AgentMessage, Role, RunStatus};
use asyncai::orchestrator::{Orchestrator, OrchestratorConfig, OrchestratorHandle};
use asyncai::output::OutputMessage;
use tokio::task::JoinHandle;

const EVENT_TIMEOUT: Duration = Duration::from_secs(2);

struct Harness {
    handle: OrchestratorHandle,
    task: Option<JoinHandle<Result<(), AgiError>>>,
    seen: Vec<OutputMessage>,
}

impl Harness {
    fn with_backend(backend: Arc<dyn AgentBackend>, config: OrchestratorConfig) -> Self {
        let kv = Arc::new(MemoryKvStore::default());
        let (orchestrator, handle) = Orchestrator::new(backend, kv, config);
        Self {
            handle,
            task: Some(tokio::spawn(orchestrator.run())),
            seen: Vec::new(),
        }
    }

    fn start(config: OrchestratorConfig) -> Self {
        Self::with_backend(Arc::new(MockBackend::echo()), config)
    }

    async fn wait_for_event(&mut self, event_type: &str) -> OutputMessage {
        let wait = async {
            loop {
                let message = self
                    .handle
                    .output_rx
                    .recv()
                    .await
                    .expect("output channel closed");
                if event_type_of(&message) == Some(event_type) {
                    return message;
                }
                self.seen.push(message);
            }
        };
        tokio::time::timeout(EVENT_TIMEOUT, wait)
            .await
            .unwrap_or_else(|_| panic!("timed out waiting for {event_type}"))
    }

    async fn drain(&mut self, duration: Duration) -> Vec<OutputMessage> {
        let deadline = tokio::time::Instant::now() + duration;
        let mut messages = Vec::new();

        loop {
            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            match tokio::time::timeout(remaining, self.handle.output_rx.recv()).await {
                Ok(Some(message)) => messages.push(message),
                _ => break,
            }
        }

        messages
    }

    async fn wait_exit(&mut self) -> Result<(), AgiError> {
        let task = self.task.take().expect("orchestrator task already joined");
        match tokio::time::timeout(EVENT_TIMEOUT, task)
            .await
            .expect("orchestrator did not exit")
        {
            Ok(result) => result,
            Err(join_error) => panic!("orchestrator task failed: {join_error}"),
        }
    }
}

impl Drop for Harness {
    fn drop(&mut self) {
        if let Some(task) = &self.task {
            task.abort();
        }
    }
}

fn event_type_of(message: &OutputMessage) -> Option<&str> {
    message
        .result
        .get("event_type")
        .and_then(|value| value.as_str())
}

fn event_data(message: &OutputMessage) -> &serde_json::Value {
    &message.result["event_data"]
}

fn message_text(message: &OutputMessage) -> &str {
    event_data(message)["message_content"]["text"]
        .as_str()
        .expect("text message content")
}

#[tokio::test]
async fn user_input_drives_agent_turn() {
    let mut harness = Harness::start(OrchestratorConfig::default());
    harness.handle.send_input("hello").unwrap();

    let agent = harness.wait_for_event("new_agent_created").await;
    assert_eq!(event_data(&agent)["agent_name"], "alice");

    let thread = harness.wait_for_event("new_thread_created").await;
    let thread_id = event_data(&thread)["thread_id"]
        .as_str()
        .unwrap()
        .to_string();

    let added = harness.wait_for_event("thread_message_added").await;
    assert_eq!(
        event_data(&added)["thread_id"].as_str(),
        Some(thread_id.as_str())
    );
    assert_eq!(event_data(&added)["message_role"], "user");
    assert_eq!(message_text(&added), "hello");

    let run = harness.wait_for_event("new_thread_run_created").await;
    assert_eq!(
        event_data(&run)["thread_id"].as_str(),
        Some(thread_id.as_str())
    );
    assert_eq!(event_data(&run)["run_status"], "in_progress");

    let message = harness.wait_for_event("new_thread_message").await;
    assert_eq!(event_data(&message)["message_role"], "assistant");
    assert_eq!(message_text(&message), "hello");

    let complete = harness.wait_for_event("thread_run_complete").await;
    assert_eq!(event_data(&complete)["run_status"], "completed");
}

#[tokio::test]
async fn second_input_reuses_existing_thread() {
    let mut harness = Harness::start(OrchestratorConfig::default());

    harness.handle.send_input("first").unwrap();
    harness.wait_for_event("new_thread_created").await;
    harness.wait_for_event("thread_run_complete").await;

    harness.handle.send_input("second").unwrap();
    let added = harness.wait_for_event("thread_message_added").await;
    assert_eq!(message_text(&added), "second");
    harness.wait_for_event("thread_run_complete").await;

    let quiet = harness.drain(Duration::from_millis(150)).await;
    assert!(
        quiet
            .iter()
            .all(|message| event_type_of(message) != Some("new_thread_created")),
        "a second thread was created: {quiet:?}"
    );
}

#[tokio::test]
async fn user_input_waits_for_agent_creation() {
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::start(config);

    harness.handle.send_input("deferred").unwrap();
    let early = harness.drain(Duration::from_millis(150)).await;

    assert!(early.iter().all(|message| {
        !matches!(
            event_type_of(message),
            Some("new_agent_created" | "new_thread_created")
        )
    }));
    assert!(
        early
            .iter()
            .any(|message| message.work_name == "main.user.unvalidated.input_action_stream"),
        "raw input was not emitted: {early:?}"
    );

    harness
        .handle
        .send_action(AgiAction::NewAgent(NewAgentAction {
            agent_id: None,
            agent_name: "alice".to_string(),
            agent_instructions: "test".to_string(),
        }))
        .unwrap();

    harness.wait_for_event("new_agent_created").await;
    harness.wait_for_event("new_thread_created").await;

    let added = harness.wait_for_event("thread_message_added").await;
    assert_eq!(message_text(&added), "deferred");

    harness.wait_for_event("thread_run_complete").await;
}

#[tokio::test]
async fn run_requested_while_busy_is_queued() {
    let backend = Arc::new(MockBackend::echo().with_delay(Duration::from_millis(120)));
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::with_backend(backend, config);

    harness
        .handle
        .send_action(AgiAction::NewAgent(NewAgentAction {
            agent_id: Some("agent-1".to_string()),
            agent_name: "alice".to_string(),
            agent_instructions: "test".to_string(),
        }))
        .unwrap();
    harness.wait_for_event("new_agent_created").await;

    harness
        .handle
        .send_action(AgiAction::NewThread(NewThreadAction {
            agent_id: "agent-1".to_string(),
        }))
        .unwrap();
    let thread = harness.wait_for_event("new_thread_created").await;
    let thread_id = event_data(&thread)["thread_id"]
        .as_str()
        .unwrap()
        .to_string();

    harness
        .handle
        .send_action(AgiAction::AddMessage(AddMessageAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
            message_role: Role::User,
            message_content: AgentContent::Text("slow".to_string()),
        }))
        .unwrap();
    harness.wait_for_event("thread_message_added").await;

    harness
        .handle
        .send_action(AgiAction::RunThread(RunThreadAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
        }))
        .unwrap();
    harness.wait_for_event("new_thread_run_created").await;

    harness
        .handle
        .send_action(AgiAction::RunThread(RunThreadAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
        }))
        .unwrap();
    harness.wait_for_event("thread_run_in_progress").await;

    harness.wait_for_event("thread_run_complete").await;
    harness.wait_for_event("new_thread_run_created").await;
    harness.wait_for_event("thread_run_complete").await;
}

#[tokio::test]
async fn backend_failure_emits_failed_event_and_loop_continues() {
    let backend = Arc::new(MockBackend::echo().with_failure("boom"));
    let mut harness = Harness::with_backend(backend, OrchestratorConfig::default());

    harness.handle.send_input("first").unwrap();
    harness.wait_for_event("new_thread_created").await;
    harness.wait_for_event("new_thread_run_created").await;
    let failed = harness.wait_for_event("thread_run_failed").await;
    assert!(event_data(&failed)["last_error"]
        .as_str()
        .unwrap()
        .contains("boom"));

    harness.handle.send_input("second").unwrap();
    harness.wait_for_event("thread_message_added").await;
    harness.wait_for_event("new_thread_run_created").await;
    harness.wait_for_event("thread_run_failed").await;
}

#[tokio::test]
async fn invalid_action_is_fatal() {
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::start(config);

    harness
        .handle
        .send_action(AgiAction::AddMessage(AddMessageAction {
            agent_id: "missing-agent".to_string(),
            thread_id: "missing-thread".to_string(),
            message_role: Role::User,
            message_content: AgentContent::Text("nope".to_string()),
        }))
        .unwrap();

    let result = harness.wait_exit().await;
    assert!(
        matches!(&result, Err(AgiError::Agent(message)) if message.contains("thread not found")),
        "unexpected result: {result:?}"
    );
}

#[tokio::test]
async fn buffered_inputs_share_one_thread() {
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::start(config);

    harness.handle.send_input("one").unwrap();
    harness.handle.send_input("two").unwrap();
    tokio::time::sleep(Duration::from_millis(100)).await;

    harness
        .handle
        .send_action(AgiAction::NewAgent(NewAgentAction {
            agent_id: None,
            agent_name: "alice".to_string(),
            agent_instructions: "test".to_string(),
        }))
        .unwrap();

    harness.wait_for_event("new_agent_created").await;
    harness.wait_for_event("new_thread_created").await;

    let first = harness.wait_for_event("thread_message_added").await;
    let second = harness.wait_for_event("thread_message_added").await;
    assert_eq!(message_text(&first), "one");
    assert_eq!(message_text(&second), "two");
    assert_eq!(
        event_data(&first)["thread_id"].as_str(),
        event_data(&second)["thread_id"].as_str()
    );

    harness.wait_for_event("thread_run_complete").await;
    harness.wait_for_event("thread_run_complete").await;

    let quiet = harness.drain(Duration::from_millis(150)).await;
    let extra_threads = harness
        .seen
        .iter()
        .chain(quiet.iter())
        .filter(|message| event_type_of(message) == Some("new_thread_created"))
        .count();
    assert_eq!(
        extra_threads, 0,
        "buffered inputs created extra threads: {quiet:?}"
    );
}

#[tokio::test]
async fn shutdown_action_ends_runtime() {
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::start(config);

    harness.handle.send_action(AgiAction::Shutdown).unwrap();

    let end = harness.wait_for_event("end_events").await;
    assert!(end.result.get("event_data").is_none());

    let result = harness.wait_exit().await;
    assert!(result.is_ok(), "unexpected shutdown result: {result:?}");
}

#[derive(Default)]
struct RecordingBackend {
    requests: Mutex<Vec<Vec<AgentMessage>>>,
}

#[async_trait]
impl AgentBackend for RecordingBackend {
    async fn start_run(&self, request: RunRequest) -> Result<BackendStream, BackendError> {
        self.requests.lock().unwrap().push(request.messages.clone());
        tokio::time::sleep(Duration::from_millis(100)).await;

        let assistant = AgentMessage {
            role: Role::Assistant,
            content: AgentContent::Text("reply".to_string()),
        };
        let mut messages = request.messages;
        messages.push(assistant.clone());

        let events = vec![
            Ok(BackendEvent::Message(assistant)),
            Ok(BackendEvent::Completed(RunOutcome {
                status: RunStatus::Completed,
                messages,
                final_output: Some(AgentContent::Text("reply".to_string())),
                last_error: None,
            })),
        ];
        Ok(Box::pin(futures::stream::iter(events)))
    }
}

fn text_of(message: &AgentMessage) -> &str {
    match &message.content {
        AgentContent::Text(text) => text,
        AgentContent::Json(_) => "",
    }
}

#[tokio::test]
async fn finalize_preserves_messages_added_while_running() {
    let backend = Arc::new(RecordingBackend::default());
    let config = OrchestratorConfig {
        seed_agent: false,
        ..Default::default()
    };
    let mut harness = Harness::with_backend(backend.clone(), config);

    harness
        .handle
        .send_action(AgiAction::NewAgent(NewAgentAction {
            agent_id: Some("agent-1".to_string()),
            agent_name: "alice".to_string(),
            agent_instructions: "test".to_string(),
        }))
        .unwrap();
    harness.wait_for_event("new_agent_created").await;

    harness
        .handle
        .send_action(AgiAction::NewThread(NewThreadAction {
            agent_id: "agent-1".to_string(),
        }))
        .unwrap();
    let thread = harness.wait_for_event("new_thread_created").await;
    let thread_id = event_data(&thread)["thread_id"]
        .as_str()
        .unwrap()
        .to_string();

    harness
        .handle
        .send_action(AgiAction::AddMessage(AddMessageAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
            message_role: Role::User,
            message_content: AgentContent::Text("one".to_string()),
        }))
        .unwrap();
    harness.wait_for_event("thread_message_added").await;

    harness
        .handle
        .send_action(AgiAction::RunThread(RunThreadAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
        }))
        .unwrap();
    harness.wait_for_event("new_thread_run_created").await;

    harness
        .handle
        .send_action(AgiAction::AddMessage(AddMessageAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
            message_role: Role::User,
            message_content: AgentContent::Text("two".to_string()),
        }))
        .unwrap();
    harness.wait_for_event("thread_message_added").await;

    harness
        .handle
        .send_action(AgiAction::RunThread(RunThreadAction {
            agent_id: "agent-1".to_string(),
            thread_id: thread_id.clone(),
        }))
        .unwrap();
    harness.wait_for_event("thread_run_in_progress").await;

    harness.wait_for_event("thread_run_complete").await;
    harness.wait_for_event("thread_run_complete").await;

    let requests = backend.requests.lock().unwrap();
    assert_eq!(requests.len(), 2);
    let second = &requests[1];
    assert_eq!(
        second.last().map(|message| message.role.clone()),
        Some(Role::User),
        "queued run must end with the newly added user message: {second:?}"
    );
    assert_eq!(text_of(second.last().unwrap()), "two");
    assert!(second.iter().any(|message| text_of(message) == "one"));
}

#[tokio::test]
async fn notices_during_a_run_coalesce_into_one_follow_up_run() {
    let backend = Arc::new(MockBackend::echo().with_delay(Duration::from_millis(150)));
    let mut harness = Harness::with_backend(backend, OrchestratorConfig::default());

    harness.handle.send_input("watch configmaps").unwrap();
    harness.wait_for_event("new_thread_run_created").await;

    for name in ["a", "b"] {
        harness
            .handle
            .send_notice(asyncai::Notice {
                source: "watch:cm".to_string(),
                text: format!("<watch-event name=\"cm\">added default/{name}</watch-event>"),
            })
            .unwrap();
    }

    harness.wait_for_event("thread_run_complete").await;
    let follow_up = harness.wait_for_event("new_thread_run_created").await;
    assert_eq!(event_data(&follow_up)["run_status"], "in_progress");
    harness.wait_for_event("thread_run_complete").await;

    let rest = harness.drain(Duration::from_millis(400)).await;
    let mut all = std::mem::take(&mut harness.seen);
    all.extend(rest);
    let notices = all
        .iter()
        .filter(|message| message.work_name == "main.system.notice_stream")
        .collect::<Vec<_>>();
    assert_eq!(notices.len(), 2);
    assert_eq!(notices[0].result["source"], "watch:cm");
    let added = all
        .iter()
        .filter(|message| event_type_of(message) == Some("thread_message_added"))
        .map(message_text)
        .collect::<Vec<_>>();
    assert!(
        added.iter().any(|text| text.contains("default/a")),
        "{added:?}"
    );
    assert!(
        added.iter().any(|text| text.contains("default/b")),
        "{added:?}"
    );
    assert!(
        all.iter()
            .all(|message| event_type_of(message) != Some("new_thread_run_created")),
        "notices during one run must coalesce into exactly one follow-up run"
    );
}