use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc;
use crate::domain::Msg;
use crate::domain::tasks::{
ApplyReport, EvidenceEntry, Stamp, TaskEdit, TaskItem, TaskOrigin, TaskSpec, TaskStatus,
TaskStore, UserTaskEdit,
};
#[derive(Clone)]
pub struct TaskBroker {
store: Arc<Mutex<TaskStore>>,
tokens: Arc<AtomicU64>,
msg_tx: mpsc::Sender<Msg>,
}
impl TaskBroker {
pub fn new(msg_tx: mpsc::Sender<Msg>) -> Self {
Self {
store: Arc::new(Mutex::new(TaskStore::default())),
tokens: Arc::new(AtomicU64::new(0)),
msg_tx,
}
}
pub fn seed(&self, store: TaskStore) {
*self.lock() = store;
}
pub fn add_tokens(&self, completion_tokens: u64) {
self.tokens.fetch_add(completion_tokens, Ordering::Relaxed);
}
pub async fn create(
&self,
specs: Vec<TaskSpec>,
origin: TaskOrigin,
) -> (Vec<TaskItem>, TaskStore) {
let (created, snapshot) = {
let mut store = self.lock();
let ids = store.create(specs, origin, self.stamp());
let created = store
.tasks
.iter()
.filter(|t| ids.contains(&t.id))
.cloned()
.collect();
(created, store.clone())
};
self.publish(snapshot.clone()).await;
(created, snapshot)
}
pub async fn update(&self, edits: Vec<TaskEdit>) -> (ApplyReport, TaskStore) {
let (report, snapshot) = {
let mut store = self.lock();
let report = store.apply(&edits, self.stamp());
(report, store.clone())
};
self.publish(snapshot.clone()).await;
(report, snapshot)
}
pub async fn user_edit(&self, edit: UserTaskEdit) -> (String, TaskStore) {
let (line, snapshot) = {
let mut store = self.lock();
let subject_of = |store: &TaskStore, id: u32| {
store
.tasks
.iter()
.find(|t| t.id == id)
.map(|t| t.subject.clone())
.unwrap_or_default()
};
let line = match edit {
UserTaskEdit::Add { subject } => {
let ids = store.create(
vec![TaskSpec {
active_form: subject.clone(),
subject: subject.clone(),
description: None,
in_progress: false,
}],
TaskOrigin::User,
self.stamp(),
);
format!("Added task #{} '{subject}'", ids[0])
},
UserTaskEdit::Remove { id } => {
let subject = subject_of(&store, id);
let report = store.apply(
&[TaskEdit {
id,
status: Some(TaskStatus::Deleted),
..TaskEdit::default()
}],
self.stamp(),
);
match report.errors.first() {
Some(err) => err.clone(),
None => format!("Removed task #{id} '{subject}'"),
}
},
UserTaskEdit::Done { id } => {
let subject = subject_of(&store, id);
let report = store.apply(
&[TaskEdit {
id,
status: Some(TaskStatus::Completed),
..TaskEdit::default()
}],
self.stamp(),
);
match report.errors.first() {
Some(err) => err.clone(),
None => format!("Marked task #{id} '{subject}' completed"),
}
},
UserTaskEdit::Clear => {
*store = TaskStore::default();
"Cleared the task list".to_string()
},
};
(line, store.clone())
};
self.publish(snapshot.clone()).await;
(line, snapshot)
}
pub async fn record_evidence(&self, entry: EvidenceEntry) {
let snapshot = {
let mut store = self.lock();
store.record_evidence(entry).then(|| store.clone())
};
if let Some(snapshot) = snapshot {
self.publish(snapshot).await;
}
}
pub fn snapshot(&self) -> TaskStore {
self.lock().clone()
}
fn stamp(&self) -> Stamp {
Stamp {
now_epoch: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0),
run_tokens: self.tokens.load(Ordering::Relaxed),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, TaskStore> {
self.store
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
async fn publish(&self, store: TaskStore) {
let _ = self.msg_tx.send(Msg::TasksUpdated { store }).await;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn spec(subject: &str, in_progress: bool) -> TaskSpec {
TaskSpec {
subject: subject.into(),
active_form: format!("{subject}ing"),
description: None,
in_progress,
}
}
async fn recv_store(rx: &mut mpsc::Receiver<Msg>) -> TaskStore {
match rx.recv().await {
Some(Msg::TasksUpdated { store }) => store,
other => panic!("expected TasksUpdated, got {other:?}"),
}
}
#[tokio::test]
async fn create_and_update_publish_snapshots() {
let (tx, mut rx) = mpsc::channel(8);
let broker = TaskBroker::new(tx);
let (created, _) = broker
.create(vec![spec("a", true), spec("b", false)], TaskOrigin::Model)
.await;
assert_eq!(created.len(), 2);
assert_eq!(recv_store(&mut rx).await.counts(), (0, 2));
let (report, _) = broker
.update(vec![TaskEdit {
id: created[0].id,
status: Some(TaskStatus::Completed),
..TaskEdit::default()
}])
.await;
assert!(report.errors.is_empty());
let published = recv_store(&mut rx).await;
assert_eq!(published.counts(), (1, 2));
}
#[tokio::test]
async fn token_readings_feed_cost_stamps() {
let (tx, _rx) = mpsc::channel(8);
let broker = TaskBroker::new(tx);
broker.add_tokens(1_000);
let (created, _) = broker
.create(vec![spec("a", true)], TaskOrigin::Model)
.await;
broker.add_tokens(8_400);
let (_, snapshot) = broker
.update(vec![TaskEdit {
id: created[0].id,
status: Some(TaskStatus::Completed),
..TaskEdit::default()
}])
.await;
assert_eq!(snapshot.tasks[0].tokens_spent, Some(8_400));
}
#[tokio::test]
async fn seed_overwrites_without_publishing() {
let (tx, mut rx) = mpsc::channel(8);
let broker = TaskBroker::new(tx);
let mut store = TaskStore::default();
store.create(
vec![spec("seeded", false)],
TaskOrigin::Model,
Stamp::default(),
);
broker.seed(store);
assert_eq!(broker.snapshot().tasks.len(), 1);
assert!(rx.try_recv().is_err(), "seed must not publish");
}
#[tokio::test]
async fn user_edits_apply_and_report() {
let (tx, mut rx) = mpsc::channel(8);
let broker = TaskBroker::new(tx);
let (line, _) = broker
.user_edit(UserTaskEdit::Add {
subject: "review the docs".into(),
})
.await;
assert_eq!(line, "Added task #1 'review the docs'");
assert_eq!(recv_store(&mut rx).await.tasks[0].origin, TaskOrigin::User);
let (line, snapshot) = broker.user_edit(UserTaskEdit::Remove { id: 9 }).await;
assert_eq!(line, "#9: no such task");
assert_eq!(snapshot.visible().count(), 1);
}
#[tokio::test]
async fn evidence_publishes_only_when_recorded() {
let (tx, mut rx) = mpsc::channel(8);
let broker = TaskBroker::new(tx);
broker
.record_evidence(EvidenceEntry {
tool: "edit_file".into(),
target: "a.rs".into(),
status: "ok".into(),
})
.await;
assert!(rx.try_recv().is_err());
broker
.create(vec![spec("a", true)], TaskOrigin::Model)
.await;
let _ = recv_store(&mut rx).await;
broker
.record_evidence(EvidenceEntry {
tool: "edit_file".into(),
target: "a.rs".into(),
status: "ok".into(),
})
.await;
let published = recv_store(&mut rx).await;
assert_eq!(published.tasks[0].evidence.len(), 1);
}
}