/// ImageGeneratorManager — DB-aware registry of image generation providers. /// /// Two kinds of providers coexist: /// - **DB-backed**: rows in `image_generate_models`, built from `llm_providers` credentials. /// Managed via `add_model` / `update_model` / `delete_model`. Loaded on startup /// and after every mutation. /// - **Plugin-registered**: ephemeral providers registered at runtime by plugins. /// Not persisted — they disappear on plugin stop. /// /// `get(id)` resolves by explicit id across both plugin and DB-backed providers. /// When called without an id, plugin providers take precedence over DB-backed ones. use std::sync::Arc; use anyhow::{Result, anyhow}; use async_trait::async_trait; use sqlx::SqlitePool; use tokio::sync::RwLock; use tracing::{info, warn}; use core_api::image_generate::ImageGenerateRegistry; use crate::llm::LlmProviderRecord; use crate::llm::db as llm_db; use crate::provider::ProviderRegistry; use crate::tools::Tool; use super::{ImageGenerate, ImageGenerateInfo, ImageGenerateModelInfo, ImageGenerateModelRecord}; use super::db as image_db; // ── Internal state ──────────────────────────────────────────────────────────── struct ImageGenerateSlot { record: ImageGenerateModelRecord, provider: LlmProviderRecord, generator: Arc, } struct ManagerState { /// DB-backed generators, ordered by priority ASC. Rebuilt on every reload(). db_slots: Vec, /// Plugin-registered providers (ephemeral — not in DB). plugins: Vec>, } // ── ImageGeneratorManager ───────────────────────────────────────────────────── pub struct ImageGeneratorManager { pool: Arc, registry: Arc, state: RwLock, } impl ImageGeneratorManager { pub async fn new( pool: Arc, registry: Arc, ) -> Result> { let mgr = Arc::new(Self { pool, registry, state: RwLock::new(ManagerState { db_slots: Vec::new(), plugins: Vec::new(), }), }); mgr.reload().await?; Ok(mgr) } // ── Plugin registration (ephemeral) ─────────────────────────────────────── pub async fn register(&self, provider: Arc) { let mut state = self.state.write().await; let id = provider.id().to_string(); state.plugins.retain(|p| p.id() != id); state.plugins.push(provider); info!(provider_id = %id, "image generator registered (plugin)"); } pub async fn unregister(&self, id: &str) { let mut state = self.state.write().await; let before = state.plugins.len(); state.plugins.retain(|p| p.id() != id); if state.plugins.len() < before { info!(provider_id = %id, "image generator unregistered (plugin)"); } } // ── Model CRUD (DB-backed) ──────────────────────────────────────────────── pub async fn add_model(&self, record: ImageGenerateModelRecord) -> Result { let id = image_db::insert(&self.pool, &record).await?; self.reload().await?; Ok(id) } pub async fn update_model(&self, id: i64, record: ImageGenerateModelRecord) -> Result<()> { image_db::update(&self.pool, id, &record).await?; self.reload().await } pub async fn delete_model(&self, id: i64) -> Result<()> { image_db::soft_delete(&self.pool, id).await?; self.reload().await } pub async fn get_model(&self, id: i64) -> Option { self.state.read().await .db_slots.iter() .find(|s| s.record.id == id) .map(|s| s.record.clone()) } pub async fn list_models_info(&self) -> Vec { self.state.read().await.db_slots.iter().map(|s| ImageGenerateModelInfo { id: s.record.id, provider_id: s.provider.id, provider_name: s.provider.name.clone(), model_id: s.record.model_id.clone(), name: s.record.name.clone(), priority: s.record.priority, from_plugin: false, description: None, }).collect() } /// Returns all active providers: plugin-registered first, then DB-backed by priority. pub async fn list_all_info(&self) -> Vec { let state = self.state.read().await; let plugins = state.plugins.iter().map(|p| ImageGenerateModelInfo { id: 0, provider_id: 0, provider_name: "Plugin".into(), model_id: p.id().to_string(), name: p.name().to_string(), priority: 0, from_plugin: true, description: p.description().map(str::to_string), }); let db = state.db_slots.iter().map(|s| ImageGenerateModelInfo { id: s.record.id, provider_id: s.provider.id, provider_name: s.provider.name.clone(), model_id: s.record.model_id.clone(), name: s.record.name.clone(), priority: s.record.priority, from_plugin: false, description: None, }); plugins.chain(db).collect() } // ── Provider queries ─────────────────────────────────────────────────────── /// Returns all active providers as lightweight info structs (for LLM tool). pub async fn list(&self) -> Vec { let state = self.state.read().await; state.plugins.iter() .map(|p| ImageGenerateInfo { id: p.id().to_string(), name: p.name().to_string(), description: p.description().map(str::to_string), extra_params_schema: p.extra_params_schema(), }) .chain(state.db_slots.iter().map(|s| ImageGenerateInfo { id: s.record.name.clone(), name: s.record.name.clone(), description: None, extra_params_schema: None, })) .collect() } /// Looks up a provider by id — plugins first, then DB-backed by name. pub async fn get(&self, id: &str) -> Option> { let state = self.state.read().await; if let Some(p) = state.plugins.iter().find(|p| p.id() == id) { return Some(Arc::clone(p)); } state.db_slots.iter() .find(|s| s.record.name == id) .map(|s| Arc::clone(&s.generator)) } // ── Generation ──────────────────────────────────────────────────────────── /// Renders `prompt` with `provider_id` and hands the raw bytes back. /// /// **Placement is the caller's**, deliberately. This used to write the file /// into the server's own `data/images/` and return that host path to /// the model — a path in nobody's vocabulary: it is not the caller's home, /// not their container, and every consumer downstream resolves agent paths /// (§6). Telegram's `send_attachment` therefore looked for /// `data/images/x.png` under the user's home and answered "file not found", /// and `read_file`/`execute_cmd`/the viewer could not reach it either. The /// manager has no `UserFs` and no session, so the one place that does — the /// tool, through its `ToolContext` — owns where the image lands. pub async fn generate_bytes( &self, provider_id: &str, prompt: &str, extra_params: Option<&serde_json::Value>, ) -> Result> { let provider = self.get(provider_id).await .ok_or_else(|| anyhow!("image provider '{}' not found", provider_id))?; let bytes = provider.generate(prompt, extra_params).await?; info!(provider_id, bytes = bytes.len(), "image generated"); Ok(bytes) } // ── Tool injection ───────────────────────────────────────────────────────── /// Returns the two image tools when at least one provider is active. /// Called per-turn by the session handler to conditionally inject tools. pub async fn tools(self: Arc) -> Vec> { let state = self.state.read().await; if state.plugins.is_empty() && state.db_slots.is_empty() { return vec![]; } drop(state); vec![ Arc::new(crate::tools::image_generate::ImageGenerateProvidersList { mgr: Arc::clone(&self) }) as Arc, Arc::new(crate::tools::image_generate::ImageGenerateTool { mgr: Arc::clone(&self) }) as Arc, ] } // ── Private ─────────────────────────────────────────────────────────────── async fn reload(&self) -> Result<()> { let model_records = image_db::load_all(&self.pool).await?; let provider_records: Vec = llm_db::load_all_providers(&self.pool).await?; let providers: std::collections::HashMap = provider_records.into_iter().map(|p| (p.id, p)).collect(); let mut db_slots = Vec::new(); for model in model_records { let provider = match providers.get(&model.provider_id) { Some(p) => p.clone(), None => { warn!( model = %model.name, provider_id = model.provider_id, "orphaned image model — provider not found, skipping", ); continue; } }; let result = self.registry.get(&provider.provider) .and_then(|p| p.build_image_generator(&provider, &model)) .unwrap_or_else(|| anyhow::bail!("provider '{}' does not support image generation", provider.provider)); match result { Ok(generator) => db_slots.push(ImageGenerateSlot { record: model, provider, generator }), Err(e) => warn!(model = %model.name, error = %e, "failed to build image generator, skipping"), } } let slot_count = db_slots.len(); self.state.write().await.db_slots = db_slots; info!(db_backed = slot_count, "image generator manager reloaded"); Ok(()) } } // ── ImageGenerateRegistry impl ──────────────────────────────────────────────── #[async_trait] impl ImageGenerateRegistry for ImageGeneratorManager { async fn register(&self, provider: Arc) { ImageGeneratorManager::register(self, provider).await; } async fn unregister(&self, id: &str) { ImageGeneratorManager::unregister(self, id).await; } }