use std::path::{Path, PathBuf};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
#[cfg(unix)]
use tokio::net::UnixStream as PlatformStream;
#[cfg(windows)]
use tokio::net::windows::named_pipe::NamedPipeClient as PlatformStream;
use tokio_util::bytes::{Bytes, BytesMut};
use tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec};
use tokio_util::sync::CancellationToken;
use super::cursor::Cursor;
use super::ids::{AgentId, ThreadId};
use super::model::{AgentCard, AgentRecord, Thread};
use super::protocol::{
AgentsStatusReport, CleanupReport, CommsNotification, CommsOut, CommsRequest, CommsResponse, PROTO_VER, SeqMeta,
StatusReport,
};
use super::singleton::{self, CommsPaths};
use super::transport::MAX_FRAME_BYTES;
use super::workspace_pool::AccessedWorkspace;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RescanReport {
pub scanned: usize,
pub updated: usize,
pub docs_indexed: usize,
pub removed: usize,
pub elapsed_ms: u64,
}
const READ_CHUNK: usize = 8 * 1024;
#[cfg(windows)]
const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
pub const REQUEST_TIMEOUT_ENV: &str = "BASEMIND_COMMS_REQUEST_TIMEOUT_SECS";
pub const HANDSHAKE_TIMEOUT_ENV: &str = "BASEMIND_COMMS_HANDSHAKE_TIMEOUT_SECS";
const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 10;
const DEFAULT_HANDSHAKE_TIMEOUT_SECS: u64 = 30;
fn timeout_from_env(var: &str, default_secs: u64) -> std::time::Duration {
let secs = std::env::var(var)
.ok()
.and_then(|raw| raw.trim().parse::<u64>().ok())
.filter(|secs| *secs > 0)
.unwrap_or(default_secs);
std::time::Duration::from_secs(secs)
}
fn request_timeout(req: &CommsRequest) -> Option<std::time::Duration> {
match req {
CommsRequest::Rescan { .. }
| CommsRequest::Cleanup { .. }
| CommsRequest::GitHistory { .. }
| CommsRequest::ResolvedRefs { .. }
| CommsRequest::IndexRead { .. }
| CommsRequest::CodeSearchLanes { .. } => None,
#[cfg(feature = "memory")]
CommsRequest::Memory { .. } | CommsRequest::Governance { .. } => None,
_ => Some(timeout_from_env(REQUEST_TIMEOUT_ENV, DEFAULT_REQUEST_TIMEOUT_SECS)),
}
}
const PENDING_NOTIFICATION_CAP: usize = 1_024;
type SpawnFn = Box<dyn Fn(&CommsPaths) -> std::io::Result<()> + Send + Sync>;
#[derive(Debug, thiserror::Error)]
pub enum CommsClientError {
#[error("comms transport error: {0}")]
Io(#[from] std::io::Error),
#[error("encode error: {0}")]
Encode(#[from] rmp_serde::encode::Error),
#[error("decode error: {0}")]
Decode(#[from] rmp_serde::decode::Error),
#[error(transparent)]
Singleton(#[from] super::singleton::SingletonError),
#[error("connection closed before a response was received")]
Closed,
#[error("broker error [{code}]: {message}")]
Broker {
code: String,
message: String,
},
#[error("unexpected response shape for {request}")]
Unexpected {
request: &'static str,
},
#[error("broker unresponsive: no reply to {what} within {secs}s (retryable)")]
Unresponsive {
what: &'static str,
secs: u64,
},
#[error("wait cancelled")]
Cancelled,
#[error("protocol skew: daemon speaks {daemon}, client speaks {client}")]
ProtoSkew {
daemon: u32,
client: u32,
},
}
pub struct CommsClient {
stream: PlatformStream,
codec: LengthDelimitedCodec,
read_buf: BytesMut,
agent: AgentId,
pending_notifications: std::collections::VecDeque<CommsNotification>,
paths: CommsPaths,
remote: Option<String>,
cwd: Option<PathBuf>,
spawn: SpawnFn,
next_id: u64,
write_incomplete: bool,
request_timeout_override: Option<std::time::Duration>,
}
impl CommsClient {
#[allow(clippy::too_many_arguments)]
pub async fn cleanup_agents(
&mut self,
apply: bool,
message_ttl_secs: u64,
thread_idle_ttl_secs: u64,
thread_retention_ttl_secs: u64,
agent_ttl_secs: u64,
claim_ttl_secs: u64,
) -> Result<CleanupReport, CommsClientError> {
match self
.request(CommsRequest::Cleanup {
apply,
message_ttl_secs,
thread_idle_ttl_secs,
thread_retention_ttl_secs,
agent_ttl_secs,
claim_ttl_secs,
})
.await?
{
CommsResponse::Cleanup(report) => Ok(report),
other => Err(self.shape_err(other, "cleanup_agents")),
}
}
pub async fn agents_status(&mut self, agent_ttl_secs: u64) -> Result<AgentsStatusReport, CommsClientError> {
match self.request(CommsRequest::AgentsStatus { agent_ttl_secs }).await? {
CommsResponse::AgentsStatus(report) => Ok(report),
other => Err(self.shape_err(other, "agents_status")),
}
}
pub async fn connect(
paths: &CommsPaths,
agent: AgentId,
remote: Option<String>,
cwd: Option<PathBuf>,
) -> Result<Self, CommsClientError> {
Self::connect_with_respawn(paths, agent, remote, cwd, |paths| {
singleton::spawn_detached_daemon(paths)
})
.await
}
pub async fn connect_with_respawn(
paths: &CommsPaths,
agent: AgentId,
remote: Option<String>,
cwd: Option<PathBuf>,
spawn: impl Fn(&CommsPaths) -> std::io::Result<()> + Send + Sync + 'static,
) -> Result<Self, CommsClientError> {
let handshake_timeout = timeout_from_env(HANDSHAKE_TIMEOUT_ENV, DEFAULT_HANDSHAKE_TIMEOUT_SECS);
Self::connect_bounded(paths, agent, remote, cwd, spawn, handshake_timeout).await
}
async fn connect_bounded(
paths: &CommsPaths,
agent: AgentId,
remote: Option<String>,
cwd: Option<PathBuf>,
spawn: impl Fn(&CommsPaths) -> std::io::Result<()> + Send + Sync + 'static,
handshake_timeout: std::time::Duration,
) -> Result<Self, CommsClientError> {
let (stream, codec) = bounded("connect", Some(handshake_timeout), Self::dial(paths)).await?;
let mut client = Self {
stream,
codec,
read_buf: BytesMut::with_capacity(READ_CHUNK),
agent,
pending_notifications: std::collections::VecDeque::new(),
paths: paths.clone(),
remote,
cwd,
spawn: Box::new(spawn),
next_id: 1,
write_incomplete: false,
request_timeout_override: None,
};
bounded("handshake", Some(handshake_timeout), client.handshake()).await?;
Ok(client)
}
pub async fn ensure_and_connect(
agent: AgentId,
remote: Option<String>,
cwd: Option<PathBuf>,
) -> Result<Self, CommsClientError> {
let paths = singleton::resolve_paths()?;
singleton::ensure_daemon(&paths).await?;
Self::connect(&paths, agent, remote, cwd).await
}
async fn dial(paths: &CommsPaths) -> Result<(PlatformStream, LengthDelimitedCodec), CommsClientError> {
let stream = Self::connect_stream(&paths.socket_path).await?;
let mut codec = LengthDelimitedCodec::new();
codec.set_max_frame_length(MAX_FRAME_BYTES);
Ok((stream, codec))
}
#[cfg(unix)]
async fn connect_stream(socket_path: &Path) -> Result<PlatformStream, CommsClientError> {
PlatformStream::connect(socket_path)
.await
.map_err(|source| daemon_unreachable_error(socket_path, source))
}
#[cfg(windows)]
async fn connect_stream(socket_path: &Path) -> Result<PlatformStream, CommsClientError> {
use tokio::net::windows::named_pipe::ClientOptions;
const ERROR_PIPE_BUSY: i32 = 231;
const RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(50);
let deadline = std::time::Instant::now() + CONNECT_TIMEOUT;
loop {
match ClientOptions::new().open(socket_path) {
Ok(client) => return Ok(client),
Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY) => {
if std::time::Instant::now() >= deadline {
return Err(daemon_unreachable_error(socket_path, e));
}
tokio::time::sleep(RETRY_INTERVAL).await;
}
Err(source) => return Err(daemon_unreachable_error(socket_path, source)),
}
}
}
async fn handshake(&mut self) -> Result<(), CommsClientError> {
let resp = self
.send_and_await(
CommsRequest::Hello {
agent: self.agent.clone(),
proto_ver: PROTO_VER,
remote: self.remote.clone(),
cwd: self.cwd.clone(),
},
None,
)
.await?;
match resp {
CommsResponse::Welcome { proto_ver, .. } if proto_ver == PROTO_VER => Ok(()),
CommsResponse::Welcome { proto_ver, .. } => Err(CommsClientError::ProtoSkew {
daemon: proto_ver,
client: PROTO_VER,
}),
CommsResponse::Error { code, message } => Err(CommsClientError::Broker { code, message }),
_ => Err(CommsClientError::Unexpected { request: "hello" }),
}
}
async fn reconnect(&mut self) -> Result<(), CommsClientError> {
let spawn = &self.spawn;
singleton::ensure_daemon_with(
&self.paths,
|socket| singleton::off_worker(|| singleton::probe_alive(socket)),
|paths| spawn(paths),
)
.await?;
let handshake_timeout = timeout_from_env(HANDSHAKE_TIMEOUT_ENV, DEFAULT_HANDSHAKE_TIMEOUT_SECS);
let (stream, codec) = bounded("connect", Some(handshake_timeout), Self::dial(&self.paths)).await?;
self.stream = stream;
self.codec = codec;
self.read_buf.clear();
self.pending_notifications.clear();
self.write_incomplete = false;
bounded("handshake", Some(handshake_timeout), self.handshake()).await
}
#[must_use]
pub fn with_request_timeout(mut self, timeout: std::time::Duration) -> Self {
self.request_timeout_override = Some(timeout);
self
}
pub fn agent(&self) -> &AgentId {
&self.agent
}
async fn expect_ok(&mut self, req: CommsRequest, label: &'static str) -> Result<(), CommsClientError> {
match self.request(req).await? {
CommsResponse::Ok => Ok(()),
other => Err(self.shape_err(other, label)),
}
}
pub(super) fn shape_err(&self, resp: CommsResponse, request: &'static str) -> CommsClientError {
match resp {
CommsResponse::Error { code, message } => CommsClientError::Broker { code, message },
_ => CommsClientError::Unexpected { request },
}
}
pub(super) async fn request(&mut self, req: CommsRequest) -> Result<CommsResponse, CommsClientError> {
if self.write_incomplete {
self.reconnect().await?;
}
let timeout = self.request_timeout_override.or_else(|| request_timeout(&req));
let what = req.method();
match bounded(what, timeout, self.send_correlated(req.clone())).await {
Ok(resp) => Ok(resp),
Err(err) if is_connection_lost(&err) => {
self.reconnect().await?;
bounded(what, timeout, self.send_correlated(req)).await
}
Err(err) => Err(err),
}
}
async fn send_correlated(&mut self, req: CommsRequest) -> Result<CommsResponse, CommsClientError> {
let id = self.next_id;
self.next_id = self.next_id.wrapping_add(1);
self.send_and_await(
CommsRequest::Call {
id,
request: Box::new(req),
},
Some(id),
)
.await
}
async fn send_and_await(
&mut self,
req: CommsRequest,
expect_id: Option<u64>,
) -> Result<CommsResponse, CommsClientError> {
self.write_request(&req).await?;
loop {
match self.read_frame().await? {
Some(CommsOut::Reply { id, response }) if Some(id) == expect_id => return Ok(response),
Some(CommsOut::Response(response)) if expect_id.is_none() => return Ok(response),
Some(CommsOut::Reply { id, .. }) => {
tracing::debug!(stale_id = id, expected = ?expect_id, "comms: discarding reply to an abandoned request");
}
Some(CommsOut::Response(_)) => {
tracing::debug!(expected = ?expect_id, "comms: discarding uncorrelated response");
}
Some(CommsOut::Notification(n)) => buffer_notification(&mut self.pending_notifications, n),
None => return Err(CommsClientError::Closed),
}
}
}
async fn write_request(&mut self, req: &CommsRequest) -> Result<(), CommsClientError> {
let body = rmp_serde::to_vec_named(req)?;
let mut framed = BytesMut::new();
self.codec.encode(Bytes::from(body), &mut framed)?;
self.write_incomplete = true;
self.stream.write_all(&framed).await?;
self.stream.flush().await?;
self.write_incomplete = false;
Ok(())
}
async fn read_frame(&mut self) -> Result<Option<CommsOut>, CommsClientError> {
loop {
if let Some(frame) = self.codec.decode(&mut self.read_buf)? {
let out: CommsOut = rmp_serde::from_slice(&frame)?;
return Ok(Some(out));
}
let n = self.stream.read_buf(&mut self.read_buf).await?;
if n == 0 {
if self.read_buf.is_empty() {
return Ok(None);
}
return Err(CommsClientError::Io(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"broker closed mid-frame",
)));
}
}
}
}
async fn bounded<T>(
what: &'static str,
limit: Option<std::time::Duration>,
fut: impl std::future::Future<Output = Result<T, CommsClientError>>,
) -> Result<T, CommsClientError> {
let Some(limit) = limit else {
return fut.await;
};
tokio::time::timeout(limit, fut)
.await
.unwrap_or(Err(CommsClientError::Unresponsive {
what,
secs: limit.as_secs(),
}))
}
fn is_connection_lost(err: &CommsClientError) -> bool {
match err {
CommsClientError::Closed => true,
CommsClientError::Io(io) => matches!(
io.kind(),
std::io::ErrorKind::BrokenPipe
| std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::NotConnected
| std::io::ErrorKind::UnexpectedEof
),
_ => false,
}
}
#[path = "client_ops.rs"]
mod ops;
#[cfg(all(test, feature = "comms", unix))]
#[path = "client_tests.rs"]
mod tests;
fn buffer_notification(queue: &mut std::collections::VecDeque<CommsNotification>, n: CommsNotification) {
while queue.len() >= PENDING_NOTIFICATION_CAP {
queue.pop_front();
}
queue.push_back(n);
}
fn daemon_unreachable_error(socket_path: &Path, source: std::io::Error) -> CommsClientError {
match source.kind() {
std::io::ErrorKind::NotFound | std::io::ErrorKind::ConnectionRefused => {
CommsClientError::Io(std::io::Error::new(
source.kind(),
format!(
"comms daemon is not running (no socket at {}); start it with \
`basemind comms start`",
socket_path.display()
),
))
}
_ => CommsClientError::Io(source),
}
}
pub fn scope_context_for(cwd: &Path) -> (Option<String>, Option<PathBuf>) {
let repo = crate::git::Repo::discover(cwd).ok();
let remote = repo.as_ref().and_then(|r| {
let key = crate::git::scope_key(r);
if key.starts_with("path:") { None } else { Some(key) }
});
(remote, Some(cwd.to_path_buf()))
}