use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicU8, Ordering};
use serde_json::json;
use tokio::sync::Mutex;
use oxicode_agent::tools::{MemoryBackend, MemoryItem, ToolError};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BrainHealth {
Connected,
Degraded,
Unavailable,
}
impl BrainHealth {
pub fn info(self) -> &'static str {
match self {
BrainHealth::Connected => "ok: oxibrain daemon connected",
BrainHealth::Degraded => "degraded: oxibrain daemon unreachable",
BrainHealth::Unavailable => "degraded: oxibrain daemon unreachable",
}
}
}
const HEALTH_CONNECTED: u8 = 0;
const HEALTH_DEGRADED: u8 = 1;
const HEALTH_UNAVAILABLE: u8 = 2;
fn encode_health(h: BrainHealth) -> u8 {
match h {
BrainHealth::Connected => HEALTH_CONNECTED,
BrainHealth::Degraded => HEALTH_DEGRADED,
BrainHealth::Unavailable => HEALTH_UNAVAILABLE,
}
}
fn decode_health(b: u8) -> BrainHealth {
match b {
HEALTH_CONNECTED => BrainHealth::Connected,
HEALTH_DEGRADED => BrainHealth::Degraded,
_ => BrainHealth::Unavailable,
}
}
#[derive(Debug, Clone)]
pub enum MigrationError {
BackendOffline,
Backend(String),
Runtime(String),
Checkpoint(String),
}
impl std::fmt::Display for MigrationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MigrationError::BackendOffline => f.write_str("brain daemon offline"),
MigrationError::Backend(e) => write!(f, "brain write failed: {e}"),
MigrationError::Runtime(e) => write!(f, "migration runtime: {e}"),
MigrationError::Checkpoint(e) => write!(f, "checkpoint write failed: {e}"),
}
}
}
impl std::error::Error for MigrationError {}
pub const DEFAULT_BRAIN_SCOPE: &str = "default";
pub struct BrainMemoryBackend {
socket_path: std::path::PathBuf,
client: Arc<Mutex<Option<oxibrain_client::BrainClient>>>,
health: Arc<AtomicU8>,
scope: String,
}
impl std::fmt::Debug for BrainMemoryBackend {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BrainMemoryBackend")
.field("socket_path", &self.socket_path)
.field("scope", &self.scope)
.field(
"health",
&decode_health(self.health.load(Ordering::SeqCst)).info(),
)
.finish()
}
}
impl BrainMemoryBackend {
pub fn new(socket_path: impl Into<std::path::PathBuf>) -> Self {
Self {
socket_path: socket_path.into(),
client: Arc::new(Mutex::new(None)),
health: Arc::new(AtomicU8::new(HEALTH_UNAVAILABLE)),
scope: DEFAULT_BRAIN_SCOPE.to_string(),
}
}
pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
self.scope = scope.into();
self
}
pub fn health(&self) -> BrainHealth {
decode_health(self.health.load(Ordering::SeqCst))
}
pub async fn connected(&self) -> bool {
self.client.lock().await.is_some()
}
pub fn put_sync(&self, content: &str, kind: &str, subject: &str) -> Result<String, ToolError> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_io()
.build()
.map_err(|e| format!("backend unavailable: tokio runtime: {e}"))?;
let future = <Self as MemoryBackend>::put(self, content, kind, subject);
rt.block_on(future)
}
async fn ensure_connected(&self) -> Result<(), ToolError> {
let mut guard = self.client.lock().await;
if guard.is_some() {
return Ok(());
}
match oxibrain_client::BrainClient::connect(&self.socket_path).await {
Ok(client) => {
*guard = Some(client);
self.health.store(HEALTH_CONNECTED, Ordering::SeqCst);
Ok(())
}
Err(e) => {
self.health.store(HEALTH_UNAVAILABLE, Ordering::SeqCst);
Err(format!(
"backend unavailable: oxibrain daemon unreachable at {}: {e}",
self.socket_path.display()
))
}
}
}
async fn with_client<R>(
&self,
f: impl FnOnce(
&mut oxibrain_client::BrainClient,
) -> futures::future::BoxFuture<'_, anyhow::Result<R>>,
) -> Result<R, ToolError> {
self.ensure_connected().await?;
let mut guard = self.client.lock().await;
let client = guard.as_mut().ok_or_else(|| {
"backend unavailable: oxibrain client missing after handshake".to_string()
})?;
match f(client).await {
Ok(v) => {
self.health.store(HEALTH_CONNECTED, Ordering::SeqCst);
Ok(v)
}
Err(e) => {
self.health.store(HEALTH_DEGRADED, Ordering::SeqCst);
*guard = None;
Err(format!("backend unavailable: oxibrain call failed: {e}"))
}
}
}
}
impl MemoryBackend for BrainMemoryBackend {
fn put<'a>(
&'a self,
content: &'a str,
kind: &'a str,
subject: &'a str,
) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
Box::pin(async move {
let content = content.to_string();
let kind = kind.to_string();
let subject = subject.to_string();
let scope = self.scope.clone();
let args = json!({
"content": content,
"kind": kind,
"subject": subject,
"scope": scope,
});
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("memory.put", args).await }))
.await?;
let id = serde_json::from_str::<serde_json::Value>(&raw)
.ok()
.and_then(|v| {
v.get("id")
.and_then(|i| i.as_str().map(|s| s.to_string()))
.or_else(|| v.get("id").and_then(|i| i.as_u64().map(|n| n.to_string())))
})
.unwrap_or_else(|| raw.trim().to_string());
Ok(id)
})
}
fn search<'a>(
&'a self,
query: &'a str,
k: usize,
) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
Box::pin(async move {
let query = query.to_string();
let scope = self.scope.clone();
let args = json!({
"query": query,
"k": k,
"scope": scope,
});
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("memory.search", args).await }))
.await?;
parse_memory_items(&raw)
})
}
fn list<'a>(
&'a self,
subject: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
Box::pin(async move {
let subject = subject.to_string();
let scope = self.scope.clone();
let args = json!({
"subject": subject,
"scope": scope,
});
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("memory.list", args).await }))
.await?;
parse_memory_items(&raw)
})
}
fn delete<'a>(
&'a self,
id: &'a str,
) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
Box::pin(async move {
let id = id.to_string();
let args = json!({ "id": id });
let _ = self
.with_client(|c| Box::pin(async move { c.call_tool("memory.delete", args).await }))
.await?;
Ok(())
})
}
fn memory_info(&self) -> Option<String> {
Some(self.health().info().to_string())
}
}
fn parse_memory_items(raw: &str) -> Result<Vec<MemoryItem>, ToolError> {
let value: serde_json::Value = serde_json::from_str(raw)
.map_err(|e| format!("backend unavailable: malformed memory response: {e}"))?;
let items = value
.get("items")
.and_then(|i| i.as_array())
.or_else(|| value.as_array())
.ok_or_else(|| "backend unavailable: memory response missing 'items' array".to_string())?;
let mut out = Vec::with_capacity(items.len());
for item in items {
let id = match item.get("id") {
Some(serde_json::Value::String(s)) => s.clone(),
Some(serde_json::Value::Number(n)) => n.to_string(),
_ => String::new(),
};
let kind = item
.get("kind")
.and_then(|v| v.as_str())
.unwrap_or("fact")
.to_string();
let content = item
.get("content")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let subject = item
.get("subject")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
out.push(MemoryItem {
id,
kind,
content,
subject,
});
}
Ok(out)
}
pub fn default_socket_path() -> std::path::PathBuf {
if let Ok(p) = std::env::var("OXIBRAIN_SOCKET") {
return std::path::PathBuf::from(p);
}
if let Some(runtime) = std::env::var_os("XDG_RUNTIME_DIR") {
return std::path::PathBuf::from(runtime).join("oxibrain.sock");
}
if let Some(home) = dirs::home_dir() {
return home.join(".oxi").join("run").join("oxibrain.sock");
}
std::path::PathBuf::from("/tmp/oxibrain.sock")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn health_info_strings_match_spec() {
assert_eq!(
BrainHealth::Connected.info(),
"ok: oxibrain daemon connected"
);
assert_eq!(
BrainHealth::Degraded.info(),
"degraded: oxibrain daemon unreachable"
);
assert_eq!(
BrainHealth::Unavailable.info(),
"degraded: oxibrain daemon unreachable"
);
}
#[test]
fn health_round_trip() {
for h in [
BrainHealth::Connected,
BrainHealth::Degraded,
BrainHealth::Unavailable,
] {
assert_eq!(decode_health(encode_health(h)), h);
}
}
#[test]
fn parse_memory_items_accepts_bare_array() {
let raw = r#"[{"id":"a","kind":"fact","content":"hello","subject":"proj"}]"#;
let items = parse_memory_items(raw).unwrap();
assert_eq!(items.len(), 1);
assert_eq!(items[0].id, "a");
assert_eq!(items[0].kind, "fact");
assert_eq!(items[0].content, "hello");
assert_eq!(items[0].subject, "proj");
}
#[test]
fn parse_memory_items_accepts_wrapped_array() {
let raw = r#"{"items":[{"id":"a","kind":"fact","content":"hello","subject":"proj"}]}"#;
let items = parse_memory_items(raw).unwrap();
assert_eq!(items.len(), 1);
assert_eq!(items[0].id, "a");
}
#[test]
fn parse_memory_items_rejects_missing_items() {
let raw = r#"{"unexpected":"shape"}"#;
let err = parse_memory_items(raw).unwrap_err();
assert!(err.starts_with("backend unavailable"));
}
#[test]
fn backend_starts_unavailable() {
let backend = BrainMemoryBackend::new("/tmp/does-not-exist.sock");
assert_eq!(backend.health(), BrainHealth::Unavailable);
assert_eq!(
backend.memory_info().as_deref(),
Some("degraded: oxibrain daemon unreachable")
);
}
#[test]
fn backend_with_scope_keeps_scope() {
let backend = BrainMemoryBackend::new("/tmp/x.sock").with_scope("oxicode/main");
assert_eq!(backend.scope, "oxicode/main");
}
}