use std::{collections::HashMap, sync::Arc, time::Duration};
use anyhow::Result;
use tokio::{
sync::mpsc,
task::{self},
};
use tokio_util::task::AbortOnDropHandle;
use uuid::Uuid;
use crate::indexing::Index;
use crate::{
agent::{self, session::RunningSession},
frontend::App,
git,
repository::Repository,
util::accept_non_zero_exit,
};
use super::{
command::{Command, CommandEvent},
responder::{CommandResponse, Responder},
};
pub struct CommandHandler<S: Index> {
rx: Option<mpsc::UnboundedReceiver<CommandEvent>>,
tx: mpsc::UnboundedSender<CommandEvent>,
repository: Arc<Repository>,
agent_sessions: HashMap<Uuid, RunningSession>,
index: S,
}
impl<I: Index + Clone + 'static> CommandHandler<I> {
pub fn from_repository_and_index(repository: impl Into<Repository>, index: I) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
CommandHandler {
rx: Some(rx),
tx,
repository: Arc::new(repository.into()),
agent_sessions: HashMap::new(),
index,
}
}
#[must_use]
pub fn command_tx(&self) -> &mpsc::UnboundedSender<CommandEvent> {
&self.tx
}
pub fn register_ui(&mut self, app: &mut App) {
app.command_tx = Some(self.tx.clone());
}
pub fn start(mut self) -> AbortOnDropHandle<()> {
let repository = Arc::clone(&self.repository);
let index = self.index.clone();
let mut rx = self.rx.take().expect("Expected a receiver");
let this_handler = Arc::new(tokio::sync::Mutex::new(self));
AbortOnDropHandle::new(task::spawn(async move {
let mut joinset = tokio::task::JoinSet::new();
while let Some(event) = rx.recv().await {
if event.command().is_quit() {
tracing::warn!("Backend received quit command, shutting down");
joinset.shutdown().await;
tracing::warn!("Backend shutdown complete");
break;
}
let repository = Arc::clone(&repository);
let storage = index.clone();
let this_handler = Arc::clone(&this_handler);
joinset.spawn(async move {
let result = Box::pin(this_handler.lock().await.handle_command_event(&repository, &storage, &event, &event.command())).await;
event.responder().send(CommandResponse::Completed).await;
if let Err(error) = result {
tracing::error!(?error, cmd = %event.command(), "Failed to handle command {cmd} with error {error:#}", cmd= event.command());
event.responder().system_message(&format!(
"Failed to handle command: {error:#}"
)).await;
}
});
}
tracing::warn!("CommandHandler shutting down");
}))
}
#[tracing::instrument(skip_all, fields(otel.name = %cmd.to_string(), uuid = %event.uuid()), err)]
async fn handle_command_event(
&mut self,
repository: &Repository,
index: &I,
event: &CommandEvent,
cmd: &Command,
) -> Result<()> {
let now = std::time::Instant::now();
tracing::warn!("Handling command {cmd}");
#[allow(clippy::match_wildcard_for_single_variants)]
match cmd {
Command::StopAgent => {
self.stop_agent(event.uuid(), event.clone_responder())
.await?;
}
Command::IndexRepository => {
index
.index_repository(repository, Some(event.clone_responder()))
.await?;
}
Command::ShowConfig => {
event
.responder()
.system_message(&toml::to_string_pretty(repository.config())?)
.await;
}
Command::Chat { message } => {
let message = message.clone();
let session = self
.find_or_start_agent_by_uuid(event.uuid(), &message, event.clone_responder())
.await?;
let token = session.cancel_token().clone();
tokio::select! {
() = token.cancelled() => Ok(()),
result = session.query_agent(&message) => result,
}?;
}
Command::Diff => {
let Some(session) = self.find_agent_by_uuid(event.uuid()) else {
event
.responder()
.system_message("No agent found (yet), is it starting up?")
.await;
return Ok(());
};
let base_sha = &session.agent_environment().start_ref;
let diff = git::util::diff(session.executor(), &base_sha, true).await?;
event.responder().system_message(&diff).await;
}
Command::Exec { cmd } => {
let Some(session) = self.find_agent_by_uuid(event.uuid()) else {
event
.responder()
.system_message("No agent found (yet), is it starting up?")
.await;
return Ok(());
};
let output = accept_non_zero_exit(session.executor().exec_cmd(cmd).await)?.output;
event.responder().system_message(&output).await;
}
Command::RetryChat => {
let Some(session) = self.find_agent_by_uuid(event.uuid()) else {
event
.responder()
.system_message("No agent found (yet), is it starting up?")
.await;
return Ok(());
};
let mut token = session.cancel_token().clone();
if token.is_cancelled() {
if let Some(session) = self.agent_sessions.get_mut(&event.uuid()) {
session.reset_cancel_token();
token = session.cancel_token().clone();
}
}
session.active_agent().agent_context.redrive().await;
tokio::select! {
() = token.cancelled() => Ok(()),
result = session.run_agent() => result,
}?;
}
Command::Quit => unreachable!("Quit should be handled earlier"),
}
tokio::time::sleep(Duration::from_millis(100)).await;
let mut elapsed = now.elapsed();
if cfg!(debug_assertions) {
elapsed = Duration::from_secs(0);
}
event
.responder()
.system_message(&format!(
"Command {cmd} successful in {} seconds",
elapsed.as_secs_f64().round()
))
.await;
Ok(())
}
async fn find_or_start_agent_by_uuid(
&mut self,
uuid: Uuid,
query: &str,
responder: Arc<dyn Responder>,
) -> Result<RunningSession> {
if let Some(session) = self.agent_sessions.get_mut(&uuid) {
session.reset_cancel_token();
return Ok(session.clone());
}
let running_agent =
agent::start_session(uuid, &self.repository, &self.index, query, responder).await?;
let cloned = running_agent.clone();
self.agent_sessions.insert(uuid, running_agent);
Ok(cloned)
}
fn find_agent_by_uuid(&self, uuid: Uuid) -> Option<RunningSession> {
if let Some(session) = self.agent_sessions.get(&uuid) {
return Some(session.clone());
}
None
}
async fn stop_agent(&self, uuid: Uuid, responder: Arc<dyn Responder>) -> Result<()> {
let Some(session) = self.agent_sessions.get(&uuid) else {
responder
.system_message("No agent found (yet), is it starting up?")
.await;
return Ok(());
};
if session.cancel_token().is_cancelled() {
responder.system_message("Agent already stopped").await;
return Ok(());
}
session.stop().await;
responder.system_message("Agent stopped").await;
Ok(())
}
}