use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc::UnboundedSender;
use tokio_util::sync::CancellationToken;
use crate::app::events::{ServerStatus, ServerStatuses};
use crate::app::supervisor::ServerSupervisor;
use crate::shared::api::{Embedder, EngineBackend, ServerHandle};
use crate::shared::config::{
CloudProvider, EmbedSettings, EngineSettings, ImpersonationEngineSettings, ImpersonationMode,
OpenRouterSettings, SecretSlot,
};
use crate::shared::i18n::Locale;
use crate::shared::secrets::{ApiKeyEntry, SecretKey};
fn stored_key(api_keys: &[ApiKeyEntry], key: Option<SecretKey>) -> Option<String> {
crate::shared::secrets::stored_key(api_keys, &key?.storage_name())
}
fn gateway_for(
provider: Option<CloudProvider>,
gateway: &OpenRouterSettings,
) -> Option<OpenRouterSettings> {
(provider == Some(CloudProvider::OpenRouter)).then(|| gateway.clone())
}
pub(super) type EmbedderDress = Box<dyn FnOnce(Arc<dyn Embedder>) -> Arc<dyn Embedder>>;
type Applied<S> = (S, Vec<ApiKeyEntry>, Option<OpenRouterSettings>);
pub(super) const RESTART_BUDGET: usize = 3;
const RESTART_WINDOW: Duration = Duration::from_secs(300);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum Server {
Chat,
Embed,
Impersonation,
}
#[derive(Default)]
struct RestartBudget {
marks: Vec<Instant>,
}
impl RestartBudget {
fn allow(&mut self, now: Instant) -> bool {
self.marks
.retain(|t| now.duration_since(*t) < RESTART_WINDOW);
if self.marks.len() < RESTART_BUDGET {
self.marks.push(now);
true
} else {
false
}
}
fn reset(&mut self) {
self.marks.clear();
}
}
pub(super) struct EngineManager {
supervisor: Arc<dyn ServerSupervisor>,
pub(super) backend: Option<Arc<dyn EngineBackend>>,
chat_handle: Option<ServerHandle>,
embed_handle: Option<ServerHandle>,
imp_backend: Option<Arc<dyn EngineBackend>>,
imp_handle: Option<ServerHandle>,
pub(super) server_status: ServerStatus,
prefill_noted: bool,
window_noted: bool,
cut_noted: bool,
imp_status: ServerStatus,
embed_status: ServerStatus,
applied_chat: Option<Applied<EngineSettings>>,
applied_embed: Option<Applied<EmbedSettings>>,
applied_imp: Option<Applied<ImpersonationEngineSettings>>,
status_tx: UnboundedSender<ServerStatus>,
imp_status_tx: UnboundedSender<ServerStatus>,
embed_status_tx: UnboundedSender<ServerStatus>,
chat_probe_cancel: Option<CancellationToken>,
imp_probe_cancel: Option<CancellationToken>,
embed_probe_cancel: Option<CancellationToken>,
restarts: [RestartBudget; 3],
refused: [bool; 3],
pub(super) embedder: Arc<dyn Embedder>,
}
impl EngineManager {
pub(super) fn new(
supervisor: Arc<dyn ServerSupervisor>,
status_tx: UnboundedSender<ServerStatus>,
imp_status_tx: UnboundedSender<ServerStatus>,
embed_status_tx: UnboundedSender<ServerStatus>,
) -> Self {
Self {
supervisor,
backend: None,
chat_handle: None,
embed_handle: None,
applied_chat: None,
applied_embed: None,
applied_imp: None,
imp_backend: None,
imp_handle: None,
server_status: ServerStatus::NotConfigured,
prefill_noted: false,
window_noted: false,
cut_noted: false,
imp_status: ServerStatus::NotConfigured,
embed_status: ServerStatus::NotConfigured,
status_tx,
imp_status_tx,
embed_status_tx,
chat_probe_cancel: None,
imp_probe_cancel: None,
embed_probe_cancel: None,
restarts: Default::default(),
refused: [false; 3],
embedder: Arc::new(crate::shared::api::UnavailableEmbedder),
}
}
pub(super) fn apply_chat(
&mut self,
settings: &EngineSettings,
api_keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
loc: &'static Locale,
) {
self.supervisor.set_gateway(gateway);
self.chat_handle = None; if let Some(tok) = self.chat_probe_cancel.take() {
tok.cancel();
}
let cancel = CancellationToken::new();
self.chat_probe_cancel = Some(cancel.clone());
let key = stored_key(api_keys, settings.secret_key());
let setup = self.supervisor.apply_chat(
settings,
key.as_deref(),
cancel,
self.status_tx.clone(),
loc,
);
self.backend = setup.backend;
self.chat_handle = setup.handle;
self.refused[Server::Chat as usize] = matches!(setup.status, ServerStatus::Disconnected(_));
self.server_status = setup.status;
self.applied_chat = Some((
settings.clone(),
api_keys.to_vec(),
gateway_for(settings.mode.cloud_provider(), gateway),
));
}
pub(super) fn apply_embed(
&mut self,
settings: &EmbedSettings,
api_keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
loc: &'static Locale,
dress: EmbedderDress,
) {
self.supervisor.set_gateway(gateway);
self.embed_handle = None; if let Some(tok) = self.embed_probe_cancel.take() {
tok.cancel();
}
let cancel = CancellationToken::new();
self.embed_probe_cancel = Some(cancel.clone());
let key = stored_key(api_keys, settings.secret_key());
let setup = self.supervisor.apply_embed(
settings,
key.as_deref(),
cancel,
self.embed_status_tx.clone(),
loc,
);
self.embedder = dress(setup.embedder);
self.embed_handle = setup.handle;
self.refused[Server::Embed as usize] =
matches!(setup.status, ServerStatus::Disconnected(_));
self.embed_status = setup.status;
self.applied_embed = Some((
settings.clone(),
api_keys.to_vec(),
gateway_for(settings.mode.cloud_provider(), gateway),
));
}
pub(super) fn apply_impersonation(
&mut self,
settings: &ImpersonationEngineSettings,
api_keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
loc: &'static Locale,
) {
self.supervisor.set_gateway(gateway);
self.imp_handle = None; if let Some(tok) = self.imp_probe_cancel.take() {
tok.cancel();
}
match settings.mode {
ImpersonationMode::Shared => {
self.imp_backend = None;
self.refused[Server::Impersonation as usize] = false;
self.imp_status = ServerStatus::NotConfigured;
}
_ => {
let cancel = CancellationToken::new();
self.imp_probe_cancel = Some(cancel.clone());
let key = stored_key(api_keys, settings.secret_key());
let setup = self.supervisor.apply_impersonation(
settings,
key.as_deref(),
cancel,
self.imp_status_tx.clone(),
loc,
);
self.imp_backend = setup.backend;
self.imp_handle = setup.handle;
self.refused[Server::Impersonation as usize] =
matches!(setup.status, ServerStatus::Disconnected(_));
self.imp_status = setup.status;
}
}
self.applied_imp = Some((
settings.clone(),
api_keys.to_vec(),
gateway_for(settings.mode.cloud_provider(), gateway),
));
}
pub(super) fn chat_is_current(
&self,
s: &EngineSettings,
keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
) -> bool {
let wanted = gateway_for(s.mode.cloud_provider(), gateway);
self.applied_chat
.as_ref()
.is_some_and(|(a, k, g)| a == s && k == keys && *g == wanted)
}
pub(super) fn embed_is_current(
&self,
s: &EmbedSettings,
keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
) -> bool {
let wanted = gateway_for(s.mode.cloud_provider(), gateway);
self.applied_embed
.as_ref()
.is_some_and(|(a, k, g)| a == s && k == keys && *g == wanted)
}
pub(super) fn impersonation_is_current(
&self,
s: &ImpersonationEngineSettings,
keys: &[ApiKeyEntry],
gateway: &OpenRouterSettings,
) -> bool {
let wanted = gateway_for(s.mode.cloud_provider(), gateway);
self.applied_imp
.as_ref()
.is_some_and(|(a, k, g)| a == s && k == keys && *g == wanted)
}
pub(super) fn find_local_servers(
&self,
) -> futures_util::future::BoxFuture<'static, Vec<crate::shared::api::local_servers::Found>>
{
self.supervisor.find_local_servers()
}
pub(super) fn set_chat_status(&mut self, status: ServerStatus) {
self.note_recovery(Server::Chat, &status);
self.server_status = status;
}
pub(super) fn set_imp_status(&mut self, status: ServerStatus) {
self.note_recovery(Server::Impersonation, &status);
self.imp_status = status;
}
pub(super) fn set_embed_status(&mut self, status: ServerStatus) {
self.note_recovery(Server::Embed, &status);
self.embed_status = status;
}
fn note_recovery(&mut self, server: Server, status: &ServerStatus) {
self.refused[server as usize] = false;
if matches!(status, ServerStatus::Ready) {
self.restarts[server as usize].reset();
if matches!(server, Server::Chat) {
self.prefill_noted = false;
self.window_noted = false;
self.cut_noted = false;
}
}
}
pub(super) fn claim_prefill_note(&mut self) -> bool {
!std::mem::replace(&mut self.prefill_noted, true)
}
pub(super) fn claim_window_note(&mut self) -> bool {
!std::mem::replace(&mut self.window_noted, true)
}
pub(super) fn claim_cut_note(&mut self) -> bool {
!std::mem::replace(&mut self.cut_noted, true)
}
pub(super) fn launch_refused(&self, server: Server) -> bool {
self.refused[server as usize]
}
pub(super) fn allow_relaunch(&mut self, server: Server, now: Instant) -> bool {
self.restarts[server as usize].allow(now)
}
pub(super) fn status_of(&self, server: Server) -> &ServerStatus {
match server {
Server::Chat => &self.server_status,
Server::Embed => &self.embed_status,
Server::Impersonation => &self.imp_status,
}
}
pub(super) fn statuses(&self) -> ServerStatuses {
ServerStatuses {
chat: self.server_status.clone(),
embed: self.embed_status.clone(),
impersonation: self.imp_status.clone(),
}
}
pub(super) fn backend_if_ready(
&self,
loc: &'static Locale,
) -> Result<Arc<dyn EngineBackend>, String> {
match &self.server_status {
ServerStatus::Ready => self
.backend
.clone()
.ok_or_else(|| loc.t("ui.err.server.not_configured").to_string()),
ServerStatus::Connecting => Err(loc.t("ui.err.server.connecting").to_string()),
ServerStatus::NotConfigured => Err(loc.t("ui.err.server.not_configured").to_string()),
ServerStatus::Disconnected(reason) => {
Err(loc.tf("ui.err.server.unavailable", &[("reason", reason)]))
}
}
}
pub(super) fn impersonation_backend_if_ready(
&self,
mode: ImpersonationMode,
loc: &'static Locale,
) -> Result<Arc<dyn EngineBackend>, String> {
if mode == ImpersonationMode::Shared {
return self.backend_if_ready(loc);
}
match &self.imp_status {
ServerStatus::Ready => self
.imp_backend
.clone()
.ok_or_else(|| loc.t("ui.err.server.imp_not_configured").to_string()),
ServerStatus::Connecting => Err(loc.t("ui.err.server.imp_connecting").to_string()),
ServerStatus::NotConfigured => {
Err(loc.t("ui.err.server.imp_not_configured").to_string())
}
ServerStatus::Disconnected(reason) => {
Err(loc.tf("ui.err.server.imp_unavailable", &[("reason", reason)]))
}
}
}
pub(super) fn embedder(&self) -> Arc<dyn Embedder> {
self.embedder.clone()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn budget_allows_up_to_the_cap_then_refuses() {
let mut b = RestartBudget::default();
let now = Instant::now();
for i in 0..RESTART_BUDGET {
assert!(b.allow(now), "relaunch {i} should be allowed");
}
assert!(!b.allow(now), "the cap should stop the crash loop");
}
#[test]
fn budget_forgets_marks_older_than_the_window() {
let mut b = RestartBudget::default();
let now = Instant::now();
for _ in 0..RESTART_BUDGET {
b.allow(now);
}
assert!(!b.allow(now));
assert!(
b.allow(now + RESTART_WINDOW + Duration::from_secs(1)),
"marks older than the window should be pruned"
);
}
#[test]
fn reaching_ready_clears_the_budget() {
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let (tx2, _rx2) = tokio::sync::mpsc::unbounded_channel();
let (tx3, _rx3) = tokio::sync::mpsc::unbounded_channel();
let mut m = EngineManager::new(
Arc::new(crate::app::supervisor::MockSupervisor::with_backend(None)),
tx,
tx2,
tx3,
);
let now = Instant::now();
for _ in 0..RESTART_BUDGET {
assert!(m.allow_relaunch(Server::Chat, now));
}
assert!(!m.allow_relaunch(Server::Chat, now));
m.set_chat_status(ServerStatus::Ready);
assert!(
m.allow_relaunch(Server::Chat, now),
"a recovered server should get a clean budget"
);
}
#[test]
fn budgets_are_independent_per_server() {
let mut m = EngineManager::new(
Arc::new(crate::app::supervisor::MockSupervisor::with_backend(None)),
tokio::sync::mpsc::unbounded_channel().0,
tokio::sync::mpsc::unbounded_channel().0,
tokio::sync::mpsc::unbounded_channel().0,
);
let now = Instant::now();
for _ in 0..RESTART_BUDGET {
m.allow_relaunch(Server::Embed, now);
}
assert!(!m.allow_relaunch(Server::Embed, now));
assert!(m.allow_relaunch(Server::Chat, now));
}
}