//! Production Daemon implementation of PeerHost.
//!
//! Bridges peer requests from SessionRuntime to SessionRegistry and the daemon runner,
//! ensuring peer turns are dispatched with isolated history and workspace,
//! idempotent request admission, and correlated results.
use crate::session_registry::SessionRegistry;
use klbr_core::models::LlmClient;
use klbr_ipc::v2::V2ServerMsg;
use klbr_runtime::{
behavior::Release,
peer::LocalPeerHost,
protocol::{
HostFuture, PeerCancelReq, PeerGetReq, PeerHost, PeerOpenReq, PeerResultReq, PeerSendReq,
PeerStopReq, PeerSubmitResultReq, SessionId,
},
store::Store,
};
use std::sync::{Arc, Weak};
use tokio::sync::{broadcast, RwLock};
pub struct DaemonPeerHost {
local: LocalPeerHost,
registry: Weak<SessionRegistry>,
llm: Arc<RwLock<Option<LlmClient>>>,
initial_release: Arc<RwLock<Option<Release>>>,
v2_event_tx: Arc<RwLock<Option<broadcast::Sender<V2ServerMsg>>>>,
}
impl DaemonPeerHost {
pub fn new(registry: Weak<SessionRegistry>) -> Self {
Self {
local: LocalPeerHost,
registry,
llm: Arc::new(RwLock::new(None)),
initial_release: Arc::new(RwLock::new(None)),
v2_event_tx: Arc::new(RwLock::new(None)),
}
}
pub fn with_components(
registry: Weak<SessionRegistry>,
llm: Option<LlmClient>,
initial_release: Option<Release>,
v2_event_tx: Option<broadcast::Sender<V2ServerMsg>>,
) -> Self {
Self {
local: LocalPeerHost,
registry,
llm: Arc::new(RwLock::new(llm)),
initial_release: Arc::new(RwLock::new(initial_release)),
v2_event_tx: Arc::new(RwLock::new(v2_event_tx)),
}
}
pub async fn set_llm(&self, llm: LlmClient) {
*self.llm.write().await = Some(llm);
}
pub async fn set_release(&self, release: Release) {
*self.initial_release.write().await = Some(release);
}
pub async fn set_v2_event_tx(&self, tx: broadcast::Sender<V2ServerMsg>) {
*self.v2_event_tx.write().await = Some(tx);
}
}
impl PeerHost for DaemonPeerHost {
fn open_peer<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerOpenReq,
) -> HostFuture<'a> {
self.local.open_peer(store, caller_session_id, req)
}
fn get_peer<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerGetReq,
) -> HostFuture<'a> {
self.local.get_peer(store, caller_session_id, req)
}
fn list_peers<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
) -> HostFuture<'a> {
self.local.list_peers(store, caller_session_id)
}
fn send_peer_message<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerSendReq,
) -> HostFuture<'a> {
self.local.send_peer_message(store, caller_session_id, req)
}
fn get_peer_result<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerResultReq,
) -> HostFuture<'a> {
self.local.get_peer_result(store, caller_session_id, req)
}
fn submit_peer_result<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerSubmitResultReq,
) -> HostFuture<'a> {
self.local.submit_peer_result(store, caller_session_id, req)
}
fn cancel_peer_request<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerCancelReq,
) -> HostFuture<'a> {
self.local
.cancel_peer_request(store, caller_session_id, req)
}
fn stop_peer<'a>(
&'a self,
store: &'a Store,
caller_session_id: &'a SessionId,
req: PeerStopReq,
) -> HostFuture<'a> {
Box::pin(async move {
let peer_sid = SessionId::try_from(req.peer_id.clone())
.map_err(|e| anyhow::anyhow!("invalid peer_id: {e}"))?;
let drain = req.drain.unwrap_or(false);
if let Some(reg) = self.registry.upgrade() {
if drain {
let _ = reg.drain(&peer_sid).await;
} else {
let _ = reg.stop(&peer_sid).await;
}
}
self.local.stop_peer(store, caller_session_id, req).await
})
}
}