johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::collections::HashMap;
use std::sync::Arc;

use futures::StreamExt;
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};

use crate::actions::{
    AddMessageAction, AgiAction, CheckThreadAction, IngestFileAction, NewAgentAction,
    NewThreadAction, RunThreadAction,
};
use crate::backend::{AgentBackend, AgentSpec, BackendError, BackendEvent, RunOutcome, RunRequest};
use crate::concurrent::WorkSet;
use crate::error::AgiError;
use crate::events::AgiEvent;
use crate::events::{
    AgentEventData, FileIngestedEvent, NewThreadCreatedEvent, NewThreadMessageEvent,
    NewThreadRunCreatedEvent, ThreadMessageAddedEvent, ThreadRunEventData, ThreadRunFailedEvent,
};
use crate::ids::{AgentId, IdGen, RunId, ThreadId};
use crate::message::{AgentMessage, RunStatus};

struct AgentRuntimeThread {
    agent_id: AgentId,
    messages: Vec<AgentMessage>,
    most_recent_run_id: Option<RunId>,
    most_recent_run_status: Option<RunStatus>,
    running: Option<RunningRun>,
    queued_run: bool,
}

impl AgentRuntimeThread {
    fn new(agent_id: AgentId) -> Self {
        Self {
            agent_id,
            messages: Vec::new(),
            most_recent_run_id: None,
            most_recent_run_status: None,
            running: None,
            queued_run: false,
        }
    }
}

struct RunningRun {
    run_id: RunId,
    outcome: Option<Result<RunOutcome, BackendError>>,
    messages_done: bool,
    message_count: usize,
}

enum RuntimeItem {
    Action {
        receiver: UnboundedReceiver<AgiAction>,
        action: Option<AgiAction>,
    },
    RunFinished {
        thread_id: ThreadId,
        run_id: RunId,
        result: Result<RunOutcome, BackendError>,
    },
    Message {
        thread_id: ThreadId,
        run_id: RunId,
        receiver: UnboundedReceiver<AgentMessage>,
        message: Option<AgentMessage>,
    },
}

pub struct AgentRuntime {
    backend: Arc<dyn AgentBackend>,
    action_rx: Option<UnboundedReceiver<AgiAction>>,
    event_tx: UnboundedSender<AgiEvent>,
    ids: IdGen,
    agents: HashMap<AgentId, AgentSpec>,
    threads: HashMap<ThreadId, AgentRuntimeThread>,
    work: WorkSet<(), RuntimeItem>,
    shutdown_requested: bool,
}

impl AgentRuntime {
    pub fn new(
        backend: Arc<dyn AgentBackend>,
        action_rx: UnboundedReceiver<AgiAction>,
        event_tx: UnboundedSender<AgiEvent>,
        ids: IdGen,
    ) -> Self {
        Self {
            backend,
            action_rx: Some(action_rx),
            event_tx,
            ids,
            agents: HashMap::new(),
            threads: HashMap::new(),
            work: WorkSet::new(),
            shutdown_requested: false,
        }
    }

    pub async fn run(mut self) -> Result<(), AgiError> {
        let action_rx = self.action_rx.take().expect("action receiver present");
        self.work.spawn((), recv_action(action_rx));

        while let Some(item) = self.work.next().await {
            let ((), item) = item.map_err(AgiError::Join)?;
            match item {
                RuntimeItem::Action { receiver, action } => match action {
                    Some(action) => {
                        let shutdown = matches!(action, AgiAction::Shutdown);
                        if !shutdown {
                            self.work.spawn((), recv_action(receiver));
                        }
                        self.handle_action(action)?;
                    }
                    None => {
                        self.shutdown_requested = true;
                    }
                },
                RuntimeItem::RunFinished {
                    thread_id,
                    run_id,
                    result,
                } => {
                    self.handle_run_finished(thread_id, run_id, result)?;
                }
                RuntimeItem::Message {
                    thread_id,
                    run_id,
                    receiver,
                    message,
                } => match message {
                    Some(message) => {
                        self.work
                            .spawn((), recv_message(thread_id.clone(), run_id, receiver));
                        self.emit_thread_message(&thread_id, message)?;
                    }
                    None => {
                        self.mark_messages_done(&thread_id, &run_id)?;
                    }
                },
            }

            if self.try_shutdown()? {
                return Ok(());
            }
        }

        Ok(())
    }

    fn try_shutdown(&mut self) -> Result<bool, AgiError> {
        if !self.shutdown_requested {
            return Ok(false);
        }
        if self.threads.values().any(|thread| thread.running.is_some()) {
            return Ok(false);
        }

        self.emit(AgiEvent::EndEvents)?;
        Ok(true)
    }

    fn handle_action(&mut self, action: AgiAction) -> Result<(), AgiError> {
        match action {
            AgiAction::NewAgent(action) => self.new_agent(action),
            AgiAction::NewThread(action) => self.new_thread(action),
            AgiAction::AddMessage(action) => self.add_message(action),
            AgiAction::RunThread(action) => self.start_run(action),
            AgiAction::CheckThread(action) => self.check_thread(action),
            AgiAction::IngestFile(action) => self.ingest_file(action),
            AgiAction::Shutdown => {
                self.shutdown_requested = true;
                Ok(())
            }
        }
    }

    fn new_agent(&mut self, action: NewAgentAction) -> Result<(), AgiError> {
        let (agent_id, existing) = match action.agent_id {
            Some(agent_id) if self.agents.contains_key(&agent_id) => (agent_id, true),
            Some(agent_id) => (agent_id, false),
            None => (self.ids.next_agent(), false),
        };

        self.agents.insert(
            agent_id.clone(),
            AgentSpec {
                agent_id: agent_id.clone(),
                agent_name: action.agent_name.clone(),
                instructions: action.agent_instructions,
            },
        );

        let data = AgentEventData {
            agent_id,
            agent_name: action.agent_name,
        };
        let event = if existing {
            AgiEvent::ExistingAgentRetrieved(data)
        } else {
            AgiEvent::NewAgentCreated(data)
        };
        self.emit(event)
    }

    fn new_thread(&mut self, action: NewThreadAction) -> Result<(), AgiError> {
        if !self.agents.contains_key(&action.agent_id) {
            return self.emit_error(format!("agent not found: {}", action.agent_id));
        }

        let thread_id = self.ids.next_thread();
        self.threads.insert(
            thread_id.clone(),
            AgentRuntimeThread::new(action.agent_id.clone()),
        );

        self.emit(AgiEvent::NewThreadCreated(NewThreadCreatedEvent {
            agent_id: action.agent_id,
            thread_id,
        }))
    }

    fn add_message(&mut self, action: AddMessageAction) -> Result<(), AgiError> {
        if !self.threads.contains_key(&action.thread_id) {
            return self.emit_error(format!("thread not found: {}", action.thread_id));
        }

        let message_id = {
            let thread = self
                .threads
                .get_mut(&action.thread_id)
                .expect("thread checked");
            thread.messages.push(AgentMessage {
                role: action.message_role.clone(),
                content: action.message_content.clone(),
            });
            thread.messages.len() - 1
        };

        self.emit(AgiEvent::ThreadMessageAdded(ThreadMessageAddedEvent {
            agent_id: action.agent_id,
            thread_id: action.thread_id,
            message_id,
            message_role: action.message_role,
            message_content: action.message_content,
        }))
    }

    fn ingest_file(&mut self, action: IngestFileAction) -> Result<(), AgiError> {
        if !self.agents.contains_key(&action.agent_id) {
            return self.emit_error(format!("agent not found: {}", action.agent_id));
        }

        let file_id = self.ids.next_file();
        self.emit(AgiEvent::FileIngested(FileIngestedEvent {
            agent_id: action.agent_id,
            file_id,
        }))
    }

    fn start_run(&mut self, action: RunThreadAction) -> Result<(), AgiError> {
        if !self.threads.contains_key(&action.thread_id) {
            return self.emit_error(format!("thread not found: {}", action.thread_id));
        }
        if !self.agents.contains_key(&action.agent_id) {
            return self.emit_error(format!("agent not found: {}", action.agent_id));
        }

        let (running_run_id, messages, agent) = {
            let thread = self
                .threads
                .get_mut(&action.thread_id)
                .expect("thread checked");
            let running_run_id = thread.running.as_ref().map(|run| run.run_id.clone());
            let messages = thread.messages.clone();
            let agent = self
                .agents
                .get(&action.agent_id)
                .expect("agent checked")
                .clone();
            (running_run_id, messages, agent)
        };

        if let Some(run_id) = running_run_id {
            let thread = self
                .threads
                .get_mut(&action.thread_id)
                .expect("thread checked");
            thread.queued_run = true;

            return self.emit(AgiEvent::ThreadRunInProgress(ThreadRunEventData {
                agent_id: action.agent_id,
                thread_id: action.thread_id,
                run_id,
                run_status: RunStatus::InProgress,
            }));
        }

        let run_id = self.ids.next_run();
        {
            let thread = self
                .threads
                .get_mut(&action.thread_id)
                .expect("thread checked");
            thread.running = Some(RunningRun {
                run_id: run_id.clone(),
                outcome: None,
                messages_done: false,
                message_count: messages.len(),
            });
            thread.most_recent_run_id = Some(run_id.clone());
            thread.most_recent_run_status = Some(RunStatus::InProgress);
        }

        self.emit(AgiEvent::NewThreadRunCreated(NewThreadRunCreatedEvent {
            agent_id: action.agent_id.clone(),
            thread_id: action.thread_id.clone(),
            run_id: run_id.clone(),
            run_status: RunStatus::InProgress,
        }))?;

        let (message_tx, message_rx) = mpsc::unbounded_channel();
        let backend = Arc::clone(&self.backend);
        let request = RunRequest {
            agent,
            thread_id: Some(action.thread_id.to_string()),
            messages,
        };
        let run_thread_id = action.thread_id.clone();
        let run_run_id = run_id.clone();

        self.work.spawn((), async move {
            let result = run_backend(backend, request, message_tx).await;
            RuntimeItem::RunFinished {
                thread_id: run_thread_id,
                run_id: run_run_id,
                result,
            }
        });

        let message_thread_id = action.thread_id;
        self.work
            .spawn((), recv_message(message_thread_id, run_id, message_rx));

        Ok(())
    }

    fn check_thread(&mut self, action: CheckThreadAction) -> Result<(), AgiError> {
        if !self.threads.contains_key(&action.thread_id) {
            return self.emit_error(format!("thread not found: {}", action.thread_id));
        }

        let thread = self.threads.get(&action.thread_id).expect("thread checked");

        let event = if let Some(running) = &thread.running {
            AgiEvent::ThreadRunInProgress(ThreadRunEventData {
                agent_id: action.agent_id,
                thread_id: action.thread_id,
                run_id: running.run_id.clone(),
                run_status: RunStatus::InProgress,
            })
        } else {
            match (
                thread.most_recent_run_id.clone(),
                thread.most_recent_run_status,
            ) {
                (Some(run_id), Some(RunStatus::Completed)) => {
                    AgiEvent::ThreadRunComplete(ThreadRunEventData {
                        agent_id: action.agent_id,
                        thread_id: action.thread_id,
                        run_id,
                        run_status: RunStatus::Completed,
                    })
                }
                (Some(run_id), Some(RunStatus::Failed)) => {
                    AgiEvent::ThreadRunFailed(ThreadRunFailedEvent {
                        agent_id: action.agent_id,
                        thread_id: action.thread_id,
                        run_id,
                        run_status: RunStatus::Failed,
                        last_error: "most recent run failed".to_string(),
                    })
                }
                (run_id, Some(run_status)) => {
                    AgiEvent::ThreadRunEventWithUnknownStatus(ThreadRunEventData {
                        agent_id: action.agent_id,
                        thread_id: action.thread_id,
                        run_id: run_id.unwrap_or_default(),
                        run_status,
                    })
                }
                (_, None) => {
                    return self.emit_error(format!("thread has no run: {}", action.thread_id));
                }
            }
        };

        self.emit(event)
    }

    fn handle_run_finished(
        &mut self,
        thread_id: ThreadId,
        run_id: RunId,
        result: Result<RunOutcome, BackendError>,
    ) -> Result<(), AgiError> {
        let matched = match self.threads.get_mut(&thread_id) {
            Some(thread) => match thread.running.as_mut() {
                Some(running) if running.run_id == run_id => {
                    running.outcome = Some(result);
                    true
                }
                _ => false,
            },
            None => false,
        };

        if !matched {
            return self.emit_error(format!(
                "unexpected run completion for thread {thread_id} run {run_id}"
            ));
        }

        self.maybe_finalize(&thread_id)
    }

    fn mark_messages_done(&mut self, thread_id: &ThreadId, run_id: &RunId) -> Result<(), AgiError> {
        if let Some(thread) = self.threads.get_mut(thread_id) {
            if let Some(running) = thread.running.as_mut() {
                if running.run_id == *run_id {
                    running.messages_done = true;
                }
            }
        }

        self.maybe_finalize(thread_id)
    }

    fn maybe_finalize(&mut self, thread_id: &ThreadId) -> Result<(), AgiError> {
        let ready = match self.threads.get(thread_id) {
            Some(thread) => match &thread.running {
                Some(running) => running.messages_done && running.outcome.is_some(),
                None => false,
            },
            None => false,
        };
        if !ready {
            return Ok(());
        }

        let running = match self
            .threads
            .get_mut(thread_id)
            .and_then(|thread| thread.running.take())
        {
            Some(running) => running,
            None => return Ok(()),
        };
        let agent_id = match self.threads.get(thread_id) {
            Some(thread) => thread.agent_id.clone(),
            None => return Ok(()),
        };
        let run_id = running.run_id;
        let base_len = running.message_count;
        let outcome = running.outcome.expect("outcome checked");

        match outcome {
            Ok(outcome) => {
                let RunOutcome {
                    status,
                    messages,
                    last_error,
                    ..
                } = outcome;

                if let Some(thread) = self.threads.get_mut(thread_id) {
                    let appended = if thread.messages.len() >= base_len {
                        thread.messages.split_off(base_len)
                    } else {
                        Vec::new()
                    };
                    thread.messages.extend(messages.into_iter().skip(base_len));
                    thread.messages.extend(appended);
                    thread.most_recent_run_id = Some(run_id.clone());
                    thread.most_recent_run_status = Some(status);
                }

                let event = match status {
                    RunStatus::Completed => AgiEvent::ThreadRunComplete(ThreadRunEventData {
                        agent_id: agent_id.clone(),
                        thread_id: thread_id.clone(),
                        run_id: run_id.clone(),
                        run_status: status,
                    }),
                    RunStatus::Failed => AgiEvent::ThreadRunFailed(ThreadRunFailedEvent {
                        agent_id: agent_id.clone(),
                        thread_id: thread_id.clone(),
                        run_id: run_id.clone(),
                        run_status: status,
                        last_error: last_error
                            .unwrap_or_else(|| "backend reported failure".to_string()),
                    }),
                    other => AgiEvent::ThreadRunEventWithUnknownStatus(ThreadRunEventData {
                        agent_id: agent_id.clone(),
                        thread_id: thread_id.clone(),
                        run_id: run_id.clone(),
                        run_status: other,
                    }),
                };

                self.emit(event)?;
            }
            Err(error) => {
                if let Some(thread) = self.threads.get_mut(thread_id) {
                    thread.most_recent_run_status = Some(RunStatus::Failed);
                }

                self.emit(AgiEvent::ThreadRunFailed(ThreadRunFailedEvent {
                    agent_id: agent_id.clone(),
                    thread_id: thread_id.clone(),
                    run_id: run_id.clone(),
                    run_status: RunStatus::Failed,
                    last_error: error.to_string(),
                }))?;
            }
        }

        let queued = self
            .threads
            .get_mut(thread_id)
            .map(|thread| std::mem::take(&mut thread.queued_run))
            .unwrap_or(false);
        if queued {
            self.start_run(RunThreadAction {
                agent_id,
                thread_id: thread_id.clone(),
            })?;
        }

        Ok(())
    }

    fn emit_thread_message(
        &self,
        thread_id: &ThreadId,
        message: AgentMessage,
    ) -> Result<(), AgiError> {
        let agent_id = match self.threads.get(thread_id) {
            Some(thread) => thread.agent_id.clone(),
            None => {
                return self.emit_error(format!("thread not found: {thread_id}"));
            }
        };

        let AgentMessage { role, content } = message;
        let content_type = content.content_type().to_string();

        self.emit(AgiEvent::NewThreadMessage(NewThreadMessageEvent {
            agent_id,
            thread_id: thread_id.clone(),
            message_role: role,
            message_content_type: content_type,
            message_content: content,
        }))
    }

    fn emit(&self, event: AgiEvent) -> Result<(), AgiError> {
        self.event_tx
            .send(event)
            .map_err(|_| AgiError::Agent("agent event channel closed".to_string()))
    }

    fn emit_error(&self, message: impl Into<String>) -> Result<(), AgiError> {
        self.emit(AgiEvent::Error(crate::events::ErrorEvent {
            message: message.into(),
        }))
    }
}

async fn recv_action(mut receiver: UnboundedReceiver<AgiAction>) -> RuntimeItem {
    let action = receiver.recv().await;
    RuntimeItem::Action { receiver, action }
}

async fn recv_message(
    thread_id: ThreadId,
    run_id: RunId,
    mut receiver: UnboundedReceiver<AgentMessage>,
) -> RuntimeItem {
    let message = receiver.recv().await;
    RuntimeItem::Message {
        thread_id,
        run_id,
        receiver,
        message,
    }
}

async fn run_backend(
    backend: Arc<dyn AgentBackend>,
    request: RunRequest,
    message_tx: UnboundedSender<AgentMessage>,
) -> Result<RunOutcome, BackendError> {
    let mut stream = backend.start_run(request).await?;
    let mut completion = None;

    while let Some(item) = stream.next().await {
        match item? {
            BackendEvent::Message(message) => {
                if message_tx.send(message).is_err() {
                    return Err(BackendError::new("thread message receiver dropped"));
                }
            }
            BackendEvent::Completed(outcome) => {
                completion = Some(outcome);
                break;
            }
        }
    }

    completion.ok_or_else(|| BackendError::new("backend stream ended without completion"))
}