use std::time::Duration;
use tokio::sync::mpsc::unbounded_channel;
use tokio_util::sync::CancellationToken;
use super::{LlamaSupervisor, MANAGED_READY_TIMEOUT, ServerSupervisor, managed_embed_config};
use crate::shared::api::managed::PortLedger;
use crate::shared::api::{OpenAiClient, ServerHandle, wait_until_ready};
use crate::shared::config::{
EmbedSettings, EngineSettings, ManagedEmbedSettings, ManagedSettings, ServerMode,
};
use crate::shared::i18n::Locale;
use crate::shared::server::ServerStatus;
fn ru() -> &'static Locale {
crate::shared::i18n::locale(crate::shared::i18n::Lang::Ru)
}
async fn spawn_stranger() -> (u16, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let task = tokio::spawn(async move {
while let Ok((mut sock, _)) = listener.accept().await {
tokio::spawn(async move {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let mut buf = [0u8; 1024];
let _ = sock.read(&mut buf).await;
let _ = sock
.write_all(
b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
)
.await;
let _ = sock.shutdown().await;
});
}
});
(port, task)
}
async fn next_status(
rx: &mut tokio::sync::mpsc::UnboundedReceiver<ServerStatus>,
) -> Option<ServerStatus> {
tokio::time::timeout(Duration::from_secs(120), rx.recv())
.await
.expect("bounded")
}
fn managed_embed_on(bin: &str, model: &str, port: u16) -> EmbedSettings {
EmbedSettings {
mode: ServerMode::Managed,
managed: ManagedEmbedSettings {
binary: Some(bin.to_string()),
model_path: Some(model.to_string()),
port,
..Default::default()
},
..Default::default()
}
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "requires a local llama-server binary + model (MINDFORK_LLAMA_BIN, MINDFORK_EMBED_MODEL)"]
async fn a_stranger_on_the_port_is_refused_live() {
let (Ok(bin), Ok(model)) = (
std::env::var("MINDFORK_LLAMA_BIN"),
std::env::var("MINDFORK_EMBED_MODEL"),
) else {
eprintln!("skip: MINDFORK_LLAMA_BIN / MINDFORK_EMBED_MODEL not set");
return;
};
let (port, _stranger) = spawn_stranger().await;
let settings = managed_embed_on(&bin, &model, port);
let cfg = managed_embed_config(&settings.managed, bin.clone().into(), &Default::default());
let handle = ServerHandle::launch(&cfg, ru()).expect("the old launch spawns");
let client = OpenAiClient::new(handle.base_url());
let probe = tokio::time::timeout(
Duration::from_secs(60),
wait_until_ready(&client, MANAGED_READY_TIMEOUT, Some(handle.exited()), ru()),
)
.await
.expect("bounded");
println!("control: the probe said {probe:?} with a stranger on the port");
assert!(
probe.is_ok(),
"control: the stranger passes for a ready server"
);
tokio::time::timeout(Duration::from_secs(60), handle.exited().wait())
.await
.expect("control: our child dies on the bind");
let died = handle
.exited()
.message(crate::shared::i18n::locale(crate::shared::i18n::Lang::En));
println!("control: the child said {died}");
drop(handle);
let supervisor = LlamaSupervisor::default();
let (tx, mut rx) = unbounded_channel();
let setup = supervisor.apply_embed(&settings, None, CancellationToken::new(), tx, ru());
println!("refused: {:?}", setup.status);
assert!(setup.handle.is_none(), "nothing is spawned");
match &setup.status {
ServerStatus::Disconnected(msg) => {
assert!(
msg.contains(&format!("порт {port} занят другой программой")),
"{msg}"
)
}
other => panic!("expected the busy-port refusal, got {other:?}"),
}
assert!(rx.recv().await.is_none(), "no probe behind a refusal");
let free = std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port();
let (tx, mut rx) = unbounded_channel();
let setup = supervisor.apply_embed(
&managed_embed_on(&bin, &model, free),
None,
CancellationToken::new(),
tx,
ru(),
);
let _handle = setup.handle.expect("a free port is launched onto");
assert_eq!(
tokio::time::timeout(Duration::from_secs(120), rx.recv())
.await
.expect("bounded"),
Some(ServerStatus::Ready)
);
}
#[tokio::test]
#[ignore = "requires a local llama-server binary + model (MINDFORK_LLAMA_BIN, MINDFORK_EMBED_MODEL)"]
async fn a_restart_onto_our_own_port_goes_through_live() {
let (Ok(bin), Ok(model)) = (
std::env::var("MINDFORK_LLAMA_BIN"),
std::env::var("MINDFORK_EMBED_MODEL"),
) else {
eprintln!("skip: MINDFORK_LLAMA_BIN / MINDFORK_EMBED_MODEL not set");
return;
};
let port = std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port();
let settings = managed_embed_on(&bin, &model, port);
let supervisor = LlamaSupervisor::default();
let (tx, mut rx) = unbounded_channel();
let first = supervisor.apply_embed(&settings, None, CancellationToken::new(), tx, ru());
assert_eq!(
next_status(&mut rx).await,
Some(ServerStatus::Ready),
"first launch"
);
drop(first.handle); let cfg = managed_embed_config(&settings.managed, bin.clone().into(), &Default::default());
let control = PortLedger::default().launch(&cfg, ru());
println!(
"control: {:?}",
control.as_ref().err().map(ToString::to_string)
);
assert!(
control.is_err(),
"control: the old process still holds the port"
);
let (tx, mut rx) = unbounded_channel();
let second = supervisor.apply_embed(&settings, None, CancellationToken::new(), tx, ru());
assert_eq!(
second.status,
ServerStatus::Connecting,
"{:?}",
second.status
);
assert_eq!(
next_status(&mut rx).await,
Some(ServerStatus::Ready),
"the restart"
);
}
#[tokio::test]
#[ignore = "requires a local llama-server binary (MINDFORK_LLAMA_BIN)"]
async fn a_load_failure_is_reported_once_with_its_cause_live() {
let Ok(bin) = std::env::var("MINDFORK_LLAMA_BIN") else {
eprintln!("skip: MINDFORK_LLAMA_BIN not set");
return;
};
let dir = tempfile::tempdir().expect("tempdir");
let fake = dir.path().join("fake.gguf");
std::fs::write(&fake, b"this is not a gguf file").expect("write");
let port = std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port();
let settings = EngineSettings {
mode: ServerMode::Managed,
managed: ManagedSettings {
binary: Some(bin),
model_path: Some(fake.display().to_string()),
port,
..Default::default()
},
..Default::default()
};
let (tx, mut rx) = unbounded_channel();
let setup = LlamaSupervisor::default().apply_chat(
&settings,
None,
CancellationToken::new(),
tx,
crate::shared::i18n::locale(crate::shared::i18n::Lang::En),
);
let _handle = setup.handle.expect("spawned");
let mut seen = Vec::new();
while let Some(status) = tokio::time::timeout(Duration::from_secs(60), rx.recv())
.await
.expect("the probe must end")
{
seen.push(status);
}
println!("reported: {seen:?}");
match seen.as_slice() {
[ServerStatus::Disconnected(msg)] => assert!(
msg.starts_with("llama-server stopped with an error: "),
"the logged cause, not the generic guess: {msg}"
),
other => panic!("expected exactly one Disconnected, got {other:?}"),
}
}