pub mod auth; pub mod http; mod router; pub mod uploads; pub mod ws; use axum::{Router, middleware, routing}; use std::net::SocketAddr; use std::sync::Arc; use tokio::net::TcpListener; use crate::bus::{MessageBus, OutboundDispatcher}; use crate::channels::{ChannelManager, CliChatChannel}; use crate::config::{Config, ensure_workspace_dir, expand_path}; use crate::delivery::{ConversationWriteLocks, DeliveryCoordinator, TurnDeliveryService}; use crate::logging; use crate::mcp; use crate::memory::MemoryManager; use crate::scheduler::Scheduler; use crate::session::{SessionManager, SessionManagerServices}; use crate::task_supervisor::TaskSupervisor; pub struct GatewayState { pub config: Config, pub config_path: std::path::PathBuf, pub workspace_dir: std::path::PathBuf, pub session_manager: Arc, pub channel_manager: ChannelManager, pub storage: Arc, pub task_supervisor: TaskSupervisor, pub delivery_coordinator: DeliveryCoordinator, pub connection_shutdown: tokio_util::sync::CancellationToken, pub auth: auth::AuthManager, pub uploads: uploads::UploadRegistry, } impl GatewayState { pub async fn new() -> Result> { let config_path = crate::config::resolve_default_config_path(); let config = Config::load_from(&config_path)?; Self::from_config(config, config_path).await } async fn from_config( config: Config, config_path: std::path::PathBuf, ) -> Result> { let task_supervisor = TaskSupervisor::new(); let delivery_coordinator = DeliveryCoordinator::new(ConversationWriteLocks::default()); let connection_shutdown = tokio_util::sync::CancellationToken::new(); let auth = auth::AuthManager::load( config.gateway.require_pairing, crate::config::get_user_config_dir().join("web_auth.json"), ) .await?; let uploads = uploads::UploadRegistry::new(config.gateway.file_transfer.clone()); // Initialize workspace directory: expand path and ensure it exists let workspace_path = expand_path(&config.workspace_dir); let workspace_path = ensure_workspace_dir(&workspace_path)?; // Switch current working directory to workspace std::env::set_current_dir(&workspace_path).map_err(|e| { format!( "Failed to switch to workspace directory {}: {}", workspace_path.display(), e ) })?; tracing::info!("Using workspace directory: {}", workspace_path.display()); // Release default AGENTS.md and USER.md to ~/.picobot/ if not exist ensure_default_config_files(); // Get provider config for SessionManager let mut provider_config = config.get_provider_config("default")?; // Override workspace_dir with the ensured path provider_config.workspace_dir = workspace_path.clone(); // Initialize Storage let db_path = if let Some(ref path) = config.gateway.session_db_path { std::path::PathBuf::from(path) } else { workspace_path.join("picobot.db") }; let storage = Arc::new( crate::storage::Storage::new(&db_path) .await .map_err(|e| format!("failed to initialize session storage: {}", e))?, ); tracing::info!("Session storage: {}", db_path.display()); // Resolve consolidation provider/model with fallback to main agent config let consolidation_provider = config .memory .resolve_consolidation_provider(&provider_config.name); let consolidation_model = config .memory .resolve_consolidation_model(&provider_config.model_id); let memory_manager = Arc::new(MemoryManager::new( storage.clone(), consolidation_provider, consolidation_model, )); tracing::info!( consolidation_provider = %memory_manager.consolidation_provider, consolidation_model = %memory_manager.consolidation_model, "Memory system initialized" ); // Create MessageBus first (shared by SessionManager and ChannelManager) let bus = MessageBus::new(100); // Channels are resolved by TurnDeliveryService, while Session workers // depend only on that protocol-neutral delivery facade. let cli_chat_channel = Arc::new(CliChatChannel::with_upload_registry(uploads.clone())); let channel_manager = ChannelManager::with_bus(cli_chat_channel, bus.clone()); channel_manager .init(&config, workspace_path.clone()) .await .map_err(|e| format!("Failed to init channels: {}", e))?; let turn_delivery = TurnDeliveryService::new( delivery_coordinator.clone(), channel_manager.clone(), task_supervisor.clone(), ); let browser_config = if config.browser.enabled { Some(config.browser.clone()) } else { None }; // Create SessionManager with bus injection let session_manager = SessionManager::new( provider_config.clone(), storage.clone(), SessionManagerServices::new( bus.clone(), memory_manager, task_supervisor.clone(), turn_delivery, ), browser_config, config.gateway.max_concurrent_background_tasks, )?; let session_manager = Arc::new(session_manager); // Register send_message tool with available channel names let available_channels = channel_manager.list_channel_names().await; let valid_channels = available_channels.clone(); session_manager.register_outbound_tool(available_channels); // Register chat_manager tool session_manager .tools() .register(crate::tools::ChatManagerTool::new( storage.clone(), valid_channels.clone(), )); // Initialize MCP servers — connect and register discovered tools if !config.mcp.servers.is_empty() { let mcp_tools = mcp::connect_all(&config.mcp).await; for tool_info in mcp_tools { let wrapper = mcp::McpToolWrapper::new( &tool_info.server_name, tool_info.tool_name, tool_info.description, tool_info.schema, tool_info.connection, ); session_manager.tools().register(wrapper); } } // Initialize scheduler if enabled in config let scheduler_config = config.gateway.scheduler.clone().unwrap_or_default(); if scheduler_config.enabled { session_manager .tools() .register(crate::tools::RoutineMaintenanceTool::new( storage.clone(), config.memory.timeline_retention_days, )); storage .ensure_default_maintenance_job() .await .map_err(|e| format!("failed to seed default maintenance job: {e}"))?; // Register cron tools session_manager .tools() .register(crate::tools::cron::CronAddTool::new( storage.clone(), valid_channels, )); session_manager .tools() .register(crate::tools::cron::CronListTool::new(storage.clone())); session_manager .tools() .register(crate::tools::cron::CronRemoveTool::new(storage.clone())); session_manager .tools() .register(crate::tools::cron::CronEnableTool::new(storage.clone())); session_manager .tools() .register(crate::tools::cron::CronDisableTool::new(storage.clone())); session_manager .tools() .register(crate::tools::cron::CronUpdateTool::new(storage.clone())); tracing::info!("Cron tools registered"); } Ok(Self { config, config_path, workspace_dir: workspace_path, session_manager: session_manager.clone(), channel_manager, storage, task_supervisor, delivery_coordinator, connection_shutdown, auth, uploads, }) } /// Get a reference to the MessageBus pub fn bus(&self) -> Arc { self.channel_manager.bus() } /// Get CLI chat channel for WebSocket handling pub fn cli_chat_channel(&self) -> Arc { self.channel_manager.cli_chat_channel() } /// Start the message processing loops pub async fn start_message_processing(&self) { let bus = self.bus(); let bus_for_outbound = bus.clone(); let session_manager = self.session_manager.clone(); if self.uploads.enabled() { let uploads = self.uploads.clone(); self.task_supervisor .spawn("pending-upload-cleanup", async move { let mut interval = tokio::time::interval(std::time::Duration::from_secs(300)); interval.tick().await; loop { interval.tick().await; let removed = uploads.cleanup_expired().await; if removed > 0 { tracing::debug!(removed, "Expired pending uploads removed"); } } }); } // Relay structured plan changes to WebSocket clients. This remains // separate from chat messages, so task UI updates never pollute history. let mut plan_events = self.session_manager.work_manager().subscribe(); let cli_chat = self.cli_chat_channel(); self.task_supervisor.spawn("task-plan-events", async move { loop { match plan_events.recv().await { Ok(event) => cli_chat.publish_plan_changed(event).await, Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { tracing::warn!(skipped, "Task plan event relay lagged"); } Err(tokio::sync::broadcast::error::RecvError::Closed) => break, } } }); router::spawn_message_routers(bus.clone(), session_manager, self.task_supervisor.clone()); // Spawn outbound dispatcher let dispatcher = OutboundDispatcher::new( bus_for_outbound, self.channel_manager.clone(), self.task_supervisor.clone(), self.delivery_coordinator.write_locks(), ); self.task_supervisor .spawn("outbound-dispatcher", async move { tracing::info!("Outbound dispatcher started"); dispatcher.run().await; }); // Spawn scheduler background task if enabled let scheduler_config = self.config.gateway.scheduler.clone().unwrap_or_default(); if scheduler_config.enabled { let sched = Arc::new(Scheduler::new( self.storage.clone(), self.session_manager.clone(), scheduler_config, )); self.task_supervisor.spawn("scheduler", async move { sched.run().await; }); tracing::info!("Scheduler background task spawned"); } } } pub async fn run( host: Option, port: Option, ) -> Result<(), Box> { let config_path = crate::config::resolve_default_config_path(); let config = Config::load_from(&config_path)?; // Initialize logging logging::init_logging(); tracing::info!(config_path = %config_path.display(), "Starting PicoBot Gateway"); let state = Arc::new(GatewayState::from_config(config, config_path).await?); // Start all channels (init already done while constructing GatewayState) state.channel_manager.start_all().await?; // Start message processing (inbound processor + control processor + outbound dispatcher) state.start_message_processing().await; // CLI args override config file values let bind_host = host.unwrap_or_else(|| state.config.gateway.host.clone()); let bind_port = port.unwrap_or(state.config.gateway.port); let protected = Router::new() .route("/api/health", routing::get(http::health)) .route( "/api/config", routing::get(http::get_config).put(http::put_config), ) .route( "/api/profiles/{name}", routing::get(http::get_profile).put(http::put_profile), ) .route("/api/logs", routing::get(http::get_logs)) .route("/api/tasks", routing::get(http::get_tasks)) .route("/api/jobs", routing::get(http::get_jobs)) .route("/api/jobs/{id}/runs", routing::get(http::get_job_runs)) .route("/api/memories", routing::get(http::get_memories)) .route( "/api/chat/{client_id}/uploads", routing::post(http::upload_file).layer(axum::extract::DefaultBodyLimit::disable()), ) .route( "/api/chat/{client_id}/sessions/{session_id}/messages/{message_id}/attachments/{index}", routing::get(http::download_attachment), ) .route("/ws", routing::get(ws::ws_handler)) .route_layer(middleware::from_fn_with_state( state.auth.clone(), auth::require_auth, )); let app = Router::new() .route("/", routing::get(http::webui_index)) .route("/app.js", routing::get(http::webui_script)) .route("/styles.css", routing::get(http::webui_styles)) .route("/health", routing::get(http::health)) .route("/api/auth/status", routing::get(auth::status)) .route("/api/auth/pair", routing::post(auth::pair)) .route("/api/auth/code", routing::post(auth::issue_code)) .merge(protected) .with_state(state.clone()); let addr = format!("{}:{}", bind_host, bind_port); let listener = TcpListener::bind(&addr).await?; tracing::info!(address = %addr, "Gateway listening"); let connection_shutdown = state.connection_shutdown.clone(); let serve_result = axum::serve( listener, app.into_make_service_with_connect_info::(), ) .with_graceful_shutdown(async move { wait_for_shutdown_signal().await; tracing::info!("Shutdown signal received"); connection_shutdown.cancel(); }) .await; // Stop external intake before waiting for internal work to finish. if let Err(error) = state.channel_manager.stop_all().await { tracing::error!(error = %error, "Failed to stop channels cleanly"); } state.task_supervisor.cancel(); state .task_supervisor .shutdown(std::time::Duration::from_secs(10)) .await; serve_result?; Ok(()) } async fn wait_for_shutdown_signal() { #[cfg(unix)] { use tokio::signal::unix::{SignalKind, signal}; match signal(SignalKind::terminate()) { Ok(mut terminate) => { tokio::select! { result = tokio::signal::ctrl_c() => { if let Err(error) = result { tracing::error!(error = %error, "Failed to listen for Ctrl-C"); } } _ = terminate.recv() => {} } } Err(error) => { tracing::error!(error = %error, "Failed to listen for SIGTERM"); if let Err(error) = tokio::signal::ctrl_c().await { tracing::error!(error = %error, "Failed to listen for Ctrl-C"); } } } } #[cfg(not(unix))] if let Err(error) = tokio::signal::ctrl_c().await { tracing::error!(error = %error, "Failed to listen for Ctrl-C"); } } /// Release default AGENTS.md and USER.md templates to ~/.picobot/ if not already present. fn ensure_default_config_files() { let picobot_dir = dirs::home_dir().unwrap_or_default().join(".picobot"); if let Err(e) = std::fs::create_dir_all(&picobot_dir) { tracing::warn!(dir = %picobot_dir.display(), error = %e, "Failed to create ~/.picobot directory"); return; } let agents_path = picobot_dir.join("AGENTS.md"); if !agents_path.exists() { let content = include_str!("../../resources/templates/AGENTS.md"); if let Err(e) = std::fs::write(&agents_path, content) { tracing::warn!(path = %agents_path.display(), error = %e, "Failed to write AGENTS.md template"); } else { tracing::info!(path = %agents_path.display(), "Released default AGENTS.md template"); } } let user_path = picobot_dir.join("USER.md"); if !user_path.exists() { let content = include_str!("../../resources/templates/USER.md"); if let Err(e) = std::fs::write(&user_path, content) { tracing::warn!(path = %user_path.display(), error = %e, "Failed to write USER.md template"); } else { tracing::info!(path = %user_path.display(), "Released default USER.md template"); } } let config_example_path = picobot_dir.join("config.example.json"); if !config_example_path.exists() { let content = include_str!("../../resources/templates/config.example.json"); if let Err(e) = std::fs::write(&config_example_path, content) { tracing::warn!(path = %config_example_path.display(), error = %e, "Failed to write config.example.json template"); } else { tracing::info!(path = %config_example_path.display(), "Released config.example.json template"); } } }