johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::sync::Arc;

use serde_json::Value;
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};

use crate::actions::{
    AddMessageAction, AgiAction, NewAgentAction, NewThreadAction, RunThreadAction,
};
use crate::backend::AgentBackend;
use crate::concurrent::WorkSet;
use crate::error::AgiError;
use crate::events::EventType;
use crate::events::{AgiEvent, ErrorEvent};
use crate::ids::IdGen;
use crate::kv::KvStore;
use crate::message::{AgentContent, Role};
use crate::output::OutputMessage;
use crate::runtime::AgentRuntime;
use crate::state::{AgentRecord, OrchestratorState, ThreadRecord};
use crate::waiting::{WaitingEngine, WaitingRule};

#[derive(Clone, Debug)]
pub struct OrchestratorConfig {
    pub user_name: String,
    pub agi_name: String,
    pub agent_instructions: String,
    pub seed_agent: bool,
}

impl Default for OrchestratorConfig {
    fn default() -> Self {
        Self {
            user_name: "user".to_string(),
            agi_name: "alice".to_string(),
            agent_instructions: "You are asyncai, a Rust port of the AGI event loop.".to_string(),
            seed_agent: true,
        }
    }
}

/// Background input for the agent (e.g. an MCP watch notice). Handled like
/// user input: it becomes a user-role message and a run, coalescing with an
/// active run through the runtime's `queued_run`.
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)]
pub struct Notice {
    /// Origin, e.g. `watch:cm`.
    pub source: String,
    /// Message text for the model.
    pub text: String,
}

pub struct OrchestratorHandle {
    pub input_tx: UnboundedSender<String>,
    pub notice_tx: UnboundedSender<Notice>,
    pub action_tx: UnboundedSender<AgiAction>,
    pub waiting_tx: UnboundedSender<WaitingRule>,
    pub output_rx: UnboundedReceiver<OutputMessage>,
}

impl OrchestratorHandle {
    pub fn send_input(&self, line: impl Into<String>) -> Result<(), AgiError> {
        self.input_tx
            .send(line.into())
            .map_err(|_| AgiError::ActionChannelClosed)
    }

    pub fn send_notice(&self, notice: Notice) -> Result<(), AgiError> {
        self.notice_tx
            .send(notice)
            .map_err(|_| AgiError::ActionChannelClosed)
    }

    pub fn send_action(&self, action: AgiAction) -> Result<(), AgiError> {
        self.action_tx
            .send(action)
            .map_err(|_| AgiError::ActionChannelClosed)
    }

    pub fn send_waiting_rule(&self, rule: WaitingRule) -> Result<(), AgiError> {
        self.waiting_tx
            .send(rule)
            .map_err(|_| AgiError::ActionChannelClosed)
    }

    pub async fn recv_output(&mut self) -> Option<OutputMessage> {
        self.output_rx.recv().await
    }
}

pub struct Orchestrator {
    backend: Arc<dyn AgentBackend>,
    kv: Arc<dyn KvStore>,
    config: OrchestratorConfig,
    input_rx: Option<UnboundedReceiver<String>>,
    notice_rx: Option<UnboundedReceiver<Notice>>,
    output_tx: UnboundedSender<OutputMessage>,
    action_tx: UnboundedSender<AgiAction>,
    action_rx: Option<UnboundedReceiver<AgiAction>>,
    waiting_rx: Option<UnboundedReceiver<WaitingRule>>,
    agent_event_tx: Option<UnboundedSender<AgiEvent>>,
    agent_event_rx: Option<UnboundedReceiver<AgiEvent>>,
    state: OrchestratorState,
    waiting: WaitingEngine,
    new_thread_requested: bool,
    runtime_done: bool,
    agent_events_closed: bool,
}

impl Orchestrator {
    pub fn new(
        backend: Arc<dyn AgentBackend>,
        kv: Arc<dyn KvStore>,
        config: OrchestratorConfig,
    ) -> (Self, OrchestratorHandle) {
        let (input_tx, input_rx) = mpsc::unbounded_channel();
        let (notice_tx, notice_rx) = mpsc::unbounded_channel();
        let (output_tx, output_rx) = mpsc::unbounded_channel();
        let (action_tx, action_rx) = mpsc::unbounded_channel();
        let (waiting_tx, waiting_rx) = mpsc::unbounded_channel();
        let (agent_event_tx, agent_event_rx) = mpsc::unbounded_channel();

        let handle = OrchestratorHandle {
            input_tx,
            notice_tx,
            action_tx: action_tx.clone(),
            waiting_tx: waiting_tx.clone(),
            output_rx,
        };

        let orchestrator = Self {
            backend,
            kv,
            config,
            input_rx: Some(input_rx),
            notice_rx: Some(notice_rx),
            output_tx,
            action_tx,
            action_rx: Some(action_rx),
            waiting_rx: Some(waiting_rx),
            agent_event_tx: Some(agent_event_tx),
            agent_event_rx: Some(agent_event_rx),
            state: OrchestratorState::default(),
            waiting: WaitingEngine::new(),
            new_thread_requested: false,
            runtime_done: false,
            agent_events_closed: false,
        };

        (orchestrator, handle)
    }

    pub async fn run(mut self) -> Result<(), AgiError> {
        let runtime = AgentRuntime::new(
            Arc::clone(&self.backend),
            self.action_rx.take().expect("action receiver present"),
            self.agent_event_tx
                .take()
                .expect("agent event sender present"),
            IdGen::new("agi"),
        );

        let mut work: WorkSet<MainWork, MainItem> = WorkSet::new();
        work.spawn(MainWork::Runtime, async move {
            MainItem::Runtime(runtime.run().await)
        });
        work.spawn(
            MainWork::Input,
            recv_input(self.input_rx.take().expect("input receiver present")),
        );
        work.spawn(
            MainWork::Notice,
            recv_notice(self.notice_rx.take().expect("notice receiver present")),
        );
        work.spawn(
            MainWork::Waiting,
            recv_waiting(self.waiting_rx.take().expect("waiting receiver present")),
        );
        work.spawn(
            MainWork::AgentEvent,
            recv_agent_event(
                self.agent_event_rx
                    .take()
                    .expect("agent event receiver present"),
            ),
        );

        if self.config.seed_agent {
            let key = format!("agents.{}.id", self.config.agi_name);
            let agent_id = self.kv.get(&key).await;
            self.action_tx
                .send(AgiAction::NewAgent(NewAgentAction {
                    agent_id,
                    agent_name: self.config.agi_name.clone(),
                    agent_instructions: self.config.agent_instructions.clone(),
                }))
                .map_err(|_| AgiError::ActionChannelClosed)?;
        }

        while let Some(item) = work.next().await {
            let (work_kind, item) = item.map_err(AgiError::Join)?;

            match item {
                MainItem::Runtime(result) => {
                    result?;
                    self.runtime_done = true;
                    if self.agent_events_closed {
                        return Ok(());
                    }
                }
                MainItem::Input { receiver, line } => {
                    if let Some(line) = line {
                        self.emit_output(work_kind.name(), Value::String(line.clone()))?;
                        self.handle_user_input(line);
                        work.spawn(MainWork::Input, recv_input(receiver));
                    }
                }
                MainItem::Notice { receiver, notice } => {
                    if let Some(notice) = notice {
                        let value = serde_json::to_value(&notice)
                            .map_err(|error| AgiError::Serialize(error.to_string()))?;
                        self.emit_output(work_kind.name(), value)?;
                        self.handle_user_input(notice.text);
                        work.spawn(MainWork::Notice, recv_notice(receiver));
                    }
                }
                MainItem::Waiting { receiver, rule } => {
                    if let Some(rule) = rule {
                        let spec = rule.spec.clone();
                        self.emit_output(work_kind.name(), serde_json::json!({ "spec": spec }))?;
                        self.waiting.push(rule);
                        work.spawn(MainWork::Waiting, recv_waiting(receiver));
                    }
                }
                MainItem::AgentEvent { receiver, event } => {
                    if let Some(event) = event {
                        let value = serde_json::to_value(&event)
                            .map_err(|error| AgiError::Serialize(error.to_string()))?;
                        self.emit_output(work_kind.name(), value)?;
                        self.waiting.record(event.event_type());
                        self.handle_agent_event(event).await?;
                        work.spawn(MainWork::AgentEvent, recv_agent_event(receiver));
                    } else {
                        self.agent_events_closed = true;
                        if self.runtime_done {
                            return Ok(());
                        }
                    }
                }
            }

            self.dispatch_waiting()?;
        }

        Ok(())
    }

    fn emit_output(&self, work_name: impl Into<String>, result: Value) -> Result<(), AgiError> {
        self.output_tx
            .send(OutputMessage::new(work_name, result))
            .map_err(|_| AgiError::OutputChannelClosed)
    }

    fn handle_user_input(&mut self, line: String) {
        self.waiting.forget(EventType::ThreadMessageAdded);

        if self.state.active_thread.is_none() && !self.new_thread_requested {
            self.new_thread_requested = true;
            self.waiting.push(WaitingRule::any(
                vec![
                    EventType::NewAgentCreated,
                    EventType::ExistingAgentRetrieved,
                ],
                |state| {
                    let agent_id = state.active_agent.clone()?;
                    Some(AgiAction::NewThread(NewThreadAction { agent_id }))
                },
            ));
        }

        self.waiting.push(WaitingRule::one(
            EventType::NewThreadCreated,
            move |state| {
                let agent_id = state.active_agent.clone()?;
                let thread_id = state.active_thread.clone()?;
                Some(AgiAction::AddMessage(AddMessageAction {
                    agent_id,
                    thread_id,
                    message_role: Role::User,
                    message_content: AgentContent::Text(line),
                }))
            },
        ));

        self.waiting
            .push(WaitingRule::one(EventType::ThreadMessageAdded, |state| {
                let agent_id = state.active_agent.clone()?;
                let thread_id = state.active_thread.clone()?;
                Some(AgiAction::RunThread(RunThreadAction {
                    agent_id,
                    thread_id,
                }))
            }));
    }

    fn dispatch_waiting(&mut self) -> Result<(), AgiError> {
        let actions = self.waiting.drain_ready(&self.state.dispatch_state());
        for action in actions {
            self.action_tx
                .send(action)
                .map_err(|_| AgiError::ActionChannelClosed)?;
        }
        Ok(())
    }

    async fn handle_agent_event(&mut self, event: AgiEvent) -> Result<(), AgiError> {
        match event {
            AgiEvent::Error(ErrorEvent { message }) => Err(AgiError::Agent(message)),
            AgiEvent::EndEvents | AgiEvent::InternalReQueue => Ok(()),
            AgiEvent::NewAgentCreated(data) | AgiEvent::ExistingAgentRetrieved(data) => {
                let existing_thread_ids = self
                    .state
                    .agents
                    .get(&data.agent_id)
                    .map(|agent| agent.thread_ids.clone())
                    .unwrap_or_default();

                self.state.agents.insert(
                    data.agent_id.clone(),
                    AgentRecord {
                        agent_id: data.agent_id.clone(),
                        agent_name: data.agent_name.clone(),
                        file_ids: Vec::new(),
                        thread_ids: existing_thread_ids,
                    },
                );
                self.state.active_agent = Some(data.agent_id.clone());

                let key = format!("agents.{}.id", data.agent_name);
                self.kv.set(&key, data.agent_id).await;
                Ok(())
            }
            AgiEvent::FileIngested(data) => {
                if let Some(agent) = self.state.agents.get_mut(&data.agent_id) {
                    agent.file_ids.push(data.file_id);
                }
                Ok(())
            }
            AgiEvent::NewThreadCreated(data) => {
                self.state.threads.insert(
                    data.thread_id.clone(),
                    ThreadRecord {
                        agent_id: data.agent_id.clone(),
                        thread_id: data.thread_id.clone(),
                        most_recent_run_id: None,
                        most_recent_run_status: None,
                    },
                );
                self.state.active_thread = Some(data.thread_id.clone());
                self.new_thread_requested = false;

                if let Some(agent) = self.state.agents.get_mut(&data.agent_id) {
                    agent.thread_ids.push(data.thread_id);
                }
                Ok(())
            }
            AgiEvent::NewThreadRunCreated(data) => {
                if let Some(thread) = self.state.threads.get_mut(&data.thread_id) {
                    thread.most_recent_run_id = Some(data.run_id);
                    thread.most_recent_run_status = Some(data.run_status);
                }
                Ok(())
            }
            AgiEvent::ThreadRunInProgress(data) => {
                if let Some(thread) = self.state.threads.get_mut(&data.thread_id) {
                    thread.most_recent_run_status = Some(data.run_status);
                }
                Ok(())
            }
            AgiEvent::ThreadRunComplete(data) => {
                if let Some(thread) = self.state.threads.get_mut(&data.thread_id) {
                    thread.most_recent_run_id = Some(data.run_id);
                    thread.most_recent_run_status = Some(data.run_status);
                }
                Ok(())
            }
            AgiEvent::ThreadRunFailed(data) => {
                if let Some(thread) = self.state.threads.get_mut(&data.thread_id) {
                    thread.most_recent_run_status = Some(data.run_status);
                }
                Ok(())
            }
            AgiEvent::ThreadRunEventWithUnknownStatus(data) => {
                if let Some(thread) = self.state.threads.get_mut(&data.thread_id) {
                    thread.most_recent_run_status = Some(data.run_status);
                }
                Ok(())
            }
            AgiEvent::ThreadMessageAdded(_) | AgiEvent::NewThreadMessage(_) => Ok(()),
        }
    }
}

#[derive(Clone, Copy, Debug)]
enum MainWork {
    Input,
    Notice,
    Waiting,
    AgentEvent,
    Runtime,
}

impl MainWork {
    fn name(self) -> &'static str {
        match self {
            Self::Input => "main.user.unvalidated.input_action_stream",
            Self::Notice => "main.system.notice_stream",
            Self::Waiting => "main.system.waiting_event_stream",
            Self::AgentEvent => "main.agent.events",
            Self::Runtime => "main.runtime",
        }
    }
}

enum MainItem {
    Input {
        receiver: UnboundedReceiver<String>,
        line: Option<String>,
    },
    Notice {
        receiver: UnboundedReceiver<Notice>,
        notice: Option<Notice>,
    },
    Waiting {
        receiver: UnboundedReceiver<WaitingRule>,
        rule: Option<WaitingRule>,
    },
    AgentEvent {
        receiver: UnboundedReceiver<AgiEvent>,
        event: Option<AgiEvent>,
    },
    Runtime(Result<(), AgiError>),
}

async fn recv_input(mut receiver: UnboundedReceiver<String>) -> MainItem {
    let line = receiver.recv().await;
    MainItem::Input { receiver, line }
}

async fn recv_notice(mut receiver: UnboundedReceiver<Notice>) -> MainItem {
    let notice = receiver.recv().await;
    MainItem::Notice { receiver, notice }
}

async fn recv_waiting(mut receiver: UnboundedReceiver<WaitingRule>) -> MainItem {
    let rule = receiver.recv().await;
    MainItem::Waiting { receiver, rule }
}

async fn recv_agent_event(mut receiver: UnboundedReceiver<AgiEvent>) -> MainItem {
    let event = receiver.recv().await;
    MainItem::AgentEvent { receiver, event }
}