use crate::utils::clock::{SystemTime, UNIX_EPOCH};
use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::sync::Mutex;
use std::time::Duration;
use anyhow::{anyhow, Result};
use serde_json::Value;
use tokio::sync::{oneshot, Notify};
use tracing::debug;
use crate::llm::ollama_client::Message;
#[derive(Clone, Debug)]
pub struct PendingLlmRequest {
pub id: u64,
pub created_unix_ms: u64,
pub model: String,
pub messages: Vec<Message>,
pub tools: Vec<Value>,
pub claimed: bool,
}
struct Entry {
request: PendingLlmRequest,
responder: oneshot::Sender<Vec<Value>>,
}
struct Inner {
order: VecDeque<u64>,
entries: HashMap<u64, Entry>,
next_id: u64,
}
pub struct LlmRequestQueue {
inner: Mutex<Inner>,
notify: Notify,
pipe_path: Option<PathBuf>,
}
impl LlmRequestQueue {
pub fn new(pipe_path: Option<PathBuf>) -> Self {
Self {
inner: Mutex::new(Inner {
order: VecDeque::new(),
entries: HashMap::new(),
next_id: 1,
}),
notify: Notify::new(),
pipe_path,
}
}
pub fn pipe_path(&self) -> Option<&PathBuf> {
self.pipe_path.as_ref()
}
pub fn submit(
&self,
model: String,
messages: Vec<Message>,
tools: Vec<Value>,
) -> (u64, oneshot::Receiver<Vec<Value>>) {
let (tx, rx) = oneshot::channel();
let created_unix_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0);
let id = {
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let id = inner.next_id;
inner.next_id += 1;
let request = PendingLlmRequest {
id,
created_unix_ms,
model,
messages,
tools,
claimed: false,
};
inner.entries.insert(
id,
Entry {
request,
responder: tx,
},
);
inner.order.push_back(id);
id
};
self.write_pipe_notification(id);
self.notify.notify_waiters();
(id, rx)
}
pub fn claim_next(&self) -> Option<PendingLlmRequest> {
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let mut target = None;
for &id in &inner.order {
if inner.entries.get(&id).map_or(false, |e| !e.request.claimed) {
target = Some(id);
break;
}
}
let id = target?;
let entry = inner.entries.get_mut(&id)?;
entry.request.claimed = true;
Some(entry.request.clone())
}
pub async fn wait_and_claim(&self, timeout: Duration) -> Option<PendingLlmRequest> {
if let Some(req) = self.claim_next() {
return Some(req);
}
let deadline = tokio::time::Instant::now() + timeout;
loop {
let notified = self.notify.notified();
if let Some(req) = self.claim_next() {
return Some(req);
}
tokio::select! {
_ = notified => {
if let Some(req) = self.claim_next() {
return Some(req);
}
}
_ = tokio::time::sleep_until(deadline) => {
return None;
}
}
}
}
pub fn answer(&self, id: u64, actions: Vec<Value>) -> Result<()> {
let entry = {
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let entry = inner.entries.remove(&id).ok_or_else(|| {
anyhow!(
"no pending LLM request with id {} (already answered, expired, or never existed)",
id
)
})?;
inner.order.retain(|&x| x != id);
entry
};
entry.responder.send(actions).map_err(|_| {
anyhow!(
"LLM request {} is no longer waiting (its connection timed out or closed)",
id
)
})?;
debug!("agent-queue: request #{} answered", id);
Ok(())
}
pub fn expire(&self, id: u64) {
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner.entries.remove(&id);
inner.order.retain(|&x| x != id);
}
pub fn list(&self) -> Vec<PendingLlmRequest> {
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner
.order
.iter()
.filter_map(|id| inner.entries.get(id).map(|e| e.request.clone()))
.collect()
}
fn write_pipe_notification(&self, id: u64) {
let Some(path) = self.pipe_path.as_ref() else {
return;
};
#[cfg(unix)]
{
use std::io::Write;
use std::os::unix::fs::{FileTypeExt, OpenOptionsExt};
match std::fs::OpenOptions::new()
.write(true)
.custom_flags(libc::O_NONBLOCK | libc::O_NOFOLLOW)
.open(path)
{
Ok(mut f) => {
if !f
.metadata()
.is_ok_and(|metadata| metadata.file_type().is_fifo())
{
debug!(
"agent-queue: notification path is not a FIFO; skipping #{}",
id
);
return;
}
if let Err(e) = writeln!(f, "{}", id) {
debug!("agent-queue: FIFO write for #{} failed: {}", id, e);
}
}
Err(e) => {
debug!(
"agent-queue: FIFO not ready for #{} ({}): no reader attached",
id, e
);
}
}
}
#[cfg(not(unix))]
{
let _ = (path, id);
}
}
}