use std::path::PathBuf;
use std::sync::{Arc, Mutex as StdMutex};
use std::time::Duration;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use crate::config;
use crate::config::paths::AppPaths;
use crate::domain::errors::{AgentError, AgentResult, ErrorCode};
use crate::domain::policy::CommandPolicy;
use crate::infrastructure;
use crate::operational::pi_rpc_manager::PiRpcManager;
use crate::operational::session_manager::SessionManager;
use crate::operational::skill_service::SkillService;
use crate::pairing::{
FileClientStore, PairingService, ServerIdentity, SystemClock, SystemRandomSource,
};
use crate::presentation;
use crate::presentation::ui_handshake::UiPairingService;
use crate::runtime::{RuntimeApplication, RuntimeShutdown};
pub(crate) struct UiRuntimeApplication {
application: presentation::ui_server::UiApplication,
application_task: StdMutex<Option<tokio::task::JoinHandle<AgentResult<()>>>>,
application_cancellation: CancellationToken,
}
pub(crate) struct AgentRuntimeShutdown {
manager: Arc<Mutex<SessionManager>>,
pi_manager: Arc<PiRpcManager>,
}
impl AgentRuntimeShutdown {
pub(crate) fn new(manager: Arc<Mutex<SessionManager>>, pi_manager: Arc<PiRpcManager>) -> Self {
Self {
manager,
pi_manager,
}
}
}
#[async_trait::async_trait]
impl RuntimeShutdown for AgentRuntimeShutdown {
async fn shutdown(&self) {
self.pi_manager.shutdown().await;
self.manager.lock().await.shutdown().await;
}
}
impl std::fmt::Debug for UiRuntimeApplication {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("UiRuntimeApplication { redacted }")
}
}
#[async_trait::async_trait]
impl RuntimeApplication for UiRuntimeApplication {
async fn run(&self, cancellation: CancellationToken) -> AgentResult<()> {
let mut application_task = self
.application_task
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.take()
.ok_or_else(|| {
AgentError::new(
ErrorCode::BackendDisconnected,
"integrated UI application is unavailable",
)
})?;
tokio::select! {
_ = cancellation.cancelled() => {
cancellation.cancel();
self.application_cancellation.cancel();
application_task
.await
.map_err(|_| AgentError::new(ErrorCode::BackendDisconnected, "integrated UI application task failed"))?
}
result = &mut application_task => {
cancellation.cancel();
self.application_cancellation.cancel();
result.map_err(|_| AgentError::new(ErrorCode::BackendDisconnected, "integrated UI application task failed"))?
}
}
}
}
impl UiRuntimeApplication {
pub(crate) fn hub(&self) -> Arc<presentation::ui_channel::UiHub> {
self.application.hub()
}
}
impl Drop for UiRuntimeApplication {
fn drop(&mut self) {
self.application_cancellation.cancel();
}
}
pub(crate) async fn integrated_ui_application(
config: &config::Config,
manager: Arc<Mutex<SessionManager>>,
pi_manager: Arc<PiRpcManager>,
skill_service: SkillService,
paths: &AppPaths,
) -> AgentResult<(Arc<UiRuntimeApplication>, Arc<UiPairingService>)> {
let history_operations = Arc::new(Mutex::new(()));
let policy: CommandPolicy = manager.lock().await.policy_snapshot();
let active_manager = pi_manager.clone();
let history_store = infrastructure::pi_history::PiHistoryStore::with_active_path_lookup(
&config.pi_session_dir,
Arc::new(move |path| active_manager.active_session_id_for_history_path(path)),
);
let directory_browser = config.home_dir.as_ref().map(|home| {
infrastructure::directory_browser::DirectoryBrowser::new(home.clone(), policy.clone())
});
let probe_cwd = config
.home_dir
.as_deref()
.map(PathBuf::from)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
let pi_models = infrastructure::pi_models::PiModels::new(
config.pi_command.clone(),
probe_cwd,
Duration::from_secs(config.pi_rpc_timeout_seconds),
pi_manager.clone(),
);
let pairing_service = load_pairing_service(paths).await?;
let application_cancellation = CancellationToken::new();
let application = presentation::ui_server::start_application(
config.output_buffer_bytes,
presentation::ui_server::UiServerResources {
manager,
pi_manager: Some(pi_manager),
history_store: Some(history_store),
directory_browser,
pi_models: Some(pi_models),
skill_service,
history_operations,
pairing_service: pairing_service.clone(),
},
application_cancellation.clone(),
)
.await?;
let server = Arc::new(UiRuntimeApplication {
application: application.application(),
application_task: StdMutex::new(Some(application.task)),
application_cancellation,
});
Ok((server, pairing_service))
}
async fn load_pairing_service(paths: &AppPaths) -> AgentResult<Arc<UiPairingService>> {
let server_key_file = paths.server_key_file.clone();
let clients_file = paths.clients_file.clone();
tokio::task::spawn_blocking(move || {
let identity = ServerIdentity::load_or_create(&server_key_file)?;
let store = FileClientStore::new(clients_file, identity.clone());
Ok::<_, crate::pairing::PairingStoreError>(Arc::new(PairingService::new(
identity,
store,
SystemRandomSource,
SystemClock::new(),
)))
})
.await
.map_err(|_| {
AgentError::new(
ErrorCode::PairingStorageFailed,
"pairing storage initialization task failed",
)
})?
.map_err(|_| {
AgentError::new(
ErrorCode::PairingStorageFailed,
"pairing storage is invalid or unavailable",
)
})
}
pub(crate) async fn run_agent(
config: config::Config,
manager: Arc<Mutex<SessionManager>>,
pi_manager: Arc<PiRpcManager>,
skill_service: SkillService,
) -> AgentResult<()> {
let history_operations = Arc::new(Mutex::new(()));
let policy: CommandPolicy = manager.lock().await.policy_snapshot();
let active_manager = pi_manager.clone();
let history_store = infrastructure::pi_history::PiHistoryStore::with_active_path_lookup(
&config.pi_session_dir,
Arc::new(move |path| active_manager.active_session_id_for_history_path(path)),
);
let probe_cwd = config
.home_dir
.as_deref()
.map(PathBuf::from)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
let pi_models = infrastructure::pi_models::PiModels::new(
config.pi_command.clone(),
probe_cwd,
Duration::from_secs(config.pi_rpc_timeout_seconds),
pi_manager.clone(),
);
presentation::connection::run(
config,
manager,
pi_models,
history_store,
policy,
skill_service,
history_operations,
)
.await
}
pub(crate) async fn shutdown_signal() {
if let Err(err) = wait_for_shutdown(tokio::signal::ctrl_c(), terminate_signal()).await {
tracing::warn!(error = %err, "failed to listen for shutdown signal");
}
}
async fn terminate_signal() -> std::io::Result<()> {
let mut signal = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
signal.recv().await;
Ok(())
}
pub(crate) async fn wait_for_shutdown<C, T>(ctrl_c: C, terminate: T) -> std::io::Result<()>
where
C: std::future::Future<Output = std::io::Result<()>>,
T: std::future::Future<Output = std::io::Result<()>>,
{
tokio::pin!(ctrl_c);
tokio::pin!(terminate);
tokio::select! {
result = &mut ctrl_c => result,
result = &mut terminate => result,
}
}
pub(crate) async fn run_until_shutdown<R, S>(
manager: Arc<Mutex<SessionManager>>,
pi_manager: Arc<PiRpcManager>,
run: R,
shutdown: S,
) -> AgentResult<()>
where
R: std::future::Future<Output = AgentResult<()>>,
S: std::future::Future<Output = ()>,
{
tokio::pin!(run);
tokio::pin!(shutdown);
let result = tokio::select! {
result = &mut run => result,
_ = &mut shutdown => Ok(()),
};
let terminal_shutdown = async {
manager.lock().await.shutdown().await;
};
tokio::join!(terminal_shutdown, pi_manager.shutdown());
result
}