use dashmap::DashMap;
use dk_core::RepoId;
use tokio::sync::broadcast;
use uuid::Uuid;
#[derive(Debug, Clone)]
pub enum SessionEvent {
SessionCreated {
session_id: Uuid,
agent_id: String,
intent: String,
},
FileModified {
session_id: Uuid,
file_path: String,
},
ChangesetSubmitted {
session_id: Uuid,
files_modified: usize,
},
ChangesetMerged {
session_id: Uuid,
commit_hash: String,
},
SessionDisconnected {
session_id: Uuid,
},
}
const DEFAULT_CHANNEL_CAPACITY: usize = 256;
pub struct RepoEventBus {
channels: DashMap<RepoId, broadcast::Sender<SessionEvent>>,
}
impl RepoEventBus {
pub fn new() -> Self {
Self {
channels: DashMap::new(),
}
}
pub fn publish(&self, repo_id: RepoId, event: SessionEvent) {
let sender = self.get_or_create_sender(repo_id);
let _ = sender.send(event);
}
pub fn subscribe(&self, repo_id: RepoId) -> broadcast::Receiver<SessionEvent> {
let sender = self.get_or_create_sender(repo_id);
sender.subscribe()
}
pub fn active_repos(&self) -> usize {
self.channels.len()
}
pub fn subscriber_count(&self, repo_id: RepoId) -> usize {
self.channels
.get(&repo_id)
.map(|s| s.receiver_count())
.unwrap_or(0)
}
pub fn prune_dead_channels(&self) {
self.channels.retain(|_repo_id, sender| sender.receiver_count() > 0);
}
fn get_or_create_sender(&self, repo_id: RepoId) -> broadcast::Sender<SessionEvent> {
self.channels
.entry(repo_id)
.or_insert_with(|| broadcast::channel(DEFAULT_CHANNEL_CAPACITY).0)
.value()
.clone()
}
}
impl Default for RepoEventBus {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn publish_and_receive() {
let bus = RepoEventBus::new();
let repo = Uuid::new_v4();
let mut rx = bus.subscribe(repo);
bus.publish(
repo,
SessionEvent::SessionCreated {
session_id: Uuid::new_v4(),
agent_id: "agent-1".into(),
intent: "fix bug".into(),
},
);
let event = rx.recv().await.expect("should receive event");
match event {
SessionEvent::SessionCreated { agent_id, .. } => {
assert_eq!(agent_id, "agent-1");
}
other => panic!("unexpected event: {other:?}"),
}
}
#[tokio::test]
async fn no_subscriber_does_not_panic() {
let bus = RepoEventBus::new();
let repo = Uuid::new_v4();
bus.publish(
repo,
SessionEvent::SessionDisconnected {
session_id: Uuid::new_v4(),
},
);
}
#[test]
fn subscriber_count() {
let bus = RepoEventBus::new();
let repo = Uuid::new_v4();
assert_eq!(bus.subscriber_count(repo), 0);
let _rx1 = bus.subscribe(repo);
assert_eq!(bus.subscriber_count(repo), 1);
let _rx2 = bus.subscribe(repo);
assert_eq!(bus.subscriber_count(repo), 2);
}
#[test]
fn active_repos_count() {
let bus = RepoEventBus::new();
assert_eq!(bus.active_repos(), 0);
let _rx = bus.subscribe(Uuid::new_v4());
assert_eq!(bus.active_repos(), 1);
let _rx2 = bus.subscribe(Uuid::new_v4());
assert_eq!(bus.active_repos(), 2);
}
#[tokio::test]
async fn multiple_subscribers_receive_same_event() {
let bus = RepoEventBus::new();
let repo = Uuid::new_v4();
let mut rx1 = bus.subscribe(repo);
let mut rx2 = bus.subscribe(repo);
bus.publish(
repo,
SessionEvent::FileModified {
session_id: Uuid::new_v4(),
file_path: "src/main.rs".into(),
},
);
let e1 = rx1.recv().await.expect("rx1 should receive");
let e2 = rx2.recv().await.expect("rx2 should receive");
match (e1, e2) {
(
SessionEvent::FileModified { file_path: p1, .. },
SessionEvent::FileModified { file_path: p2, .. },
) => {
assert_eq!(p1, "src/main.rs");
assert_eq!(p2, "src/main.rs");
}
_ => panic!("both should receive FileModified"),
}
}
}