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 = "personal";
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 ping(&self) -> Result<(), ToolError> {
self.with_client(|c| Box::pin(async move { c.ping().await }))
.await
}
pub async fn stats(&self) -> Result<serde_json::Value, ToolError> {
let space = self.scope.clone();
self.with_client(|c| Box::pin(async move { c.stats(&space).await }))
.await
}
pub fn stats_sync(&self) -> Result<serde_json::Value, ToolError> {
block_on_sync(self.stats())
}
pub fn search_sync(&self, query: &str, k: usize) -> Result<Vec<MemoryItem>, ToolError> {
block_on_sync(<Self as MemoryBackend>::search(self, query, k))
}
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> {
block_on_sync(<Self as MemoryBackend>::put(self, content, kind, subject))
}
pub fn ping_sync(&self) -> Result<(), ToolError> {
block_on_sync(self.ping())
}
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}"))
}
}
}
}
fn block_on_sync<F: std::future::Future>(fut: F) -> F::Output {
match tokio::runtime::Handle::try_current() {
Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
Err(_) => {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_io()
.build()
.expect("build current-thread tokio runtime");
rt.block_on(fut)
}
}
}
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 args = json!({
"content": content,
"space": self.scope,
"source_path": format!("oxicode/{kind}/{subject}"),
"extract": false,
});
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("ingest", args).await }))
.await?;
let id = raw
.split_once("episode:")
.map(|(_, tail)| tail.trim().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 args = json!({
"query": query,
"space": self.scope,
"token_budget": 4000,
});
let _ = k;
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("recall", 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 args = json!({
"query": subject,
"space": self.scope,
"token_budget": 2000,
});
let raw = self
.with_client(|c| Box::pin(async move { c.call_tool("recall", 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 {
if id.trim().is_empty() {
return Err("backend unavailable: brain delete requires an episode id \
(a `put` return value or a recall provenance id)"
.to_string());
}
let args = json!({
"target_kind": "episode",
"target_id": id,
"space": self.scope,
});
let _ = self
.with_client(|c| Box::pin(async move { c.call_tool("redact", 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 layers = value
.get("layers")
.and_then(|v| v.as_array())
.ok_or_else(|| "backend unavailable: unrecognized memory response shape".to_string())?;
let mut out = Vec::new();
for layer in layers {
let kind = layer
.get("kind")
.and_then(|v| v.as_str())
.unwrap_or("layer");
let text = layer.get("text").and_then(|v| v.as_str()).unwrap_or("");
let provenance: Vec<&str> = layer
.get("provenance")
.and_then(|v| v.as_array())
.map(|a| a.iter().filter_map(|p| p.as_str()).collect())
.unwrap_or_default();
let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
for (i, line) in lines.iter().enumerate() {
let id = provenance.get(i).map(|p| p.to_string()).unwrap_or_default();
out.push(MemoryItem {
id,
kind: kind.to_string(),
content: line.trim().to_string(),
subject: String::new(),
});
}
}
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(home) = dirs::home_dir() {
return home.join(".oxi").join("brain").join("oxibrain.sock");
}
std::path::PathBuf::from(".oxi/brain/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_maps_recall_layers() {
let raw = r#"{"layers":[
{"kind":"recent_episodes",
"text":"first note\nsecond note\n",
"provenance":["ep-1","ep-2"]},
{"kind":"statements","text":"oxicode prefers Korean prose"}
]}"#;
let items = parse_memory_items(raw).unwrap();
assert_eq!(items.len(), 3);
assert_eq!(items[0].id, "ep-1");
assert_eq!(items[0].kind, "recent_episodes");
assert_eq!(items[0].content, "first note");
assert_eq!(items[1].id, "ep-2");
assert_eq!(items[1].content, "second note");
assert_eq!(items[2].id, "");
assert_eq!(items[2].kind, "statements");
assert_eq!(items[2].content, "oxicode prefers Korean prose");
}
#[test]
fn parse_memory_items_skips_blank_lines() {
let raw = r#"{"layers":[{"kind":"recent_episodes",
"text":"\n \nonly line\n","provenance":["ep-9"]}]}"#;
let items = parse_memory_items(raw).unwrap();
assert_eq!(items.len(), 1);
assert_eq!(items[0].id, "ep-9");
assert_eq!(items[0].content, "only line");
}
#[test]
fn parse_memory_items_empty_layers_is_ok() {
let items = parse_memory_items(r#"{"layers":[]}"#).unwrap();
assert!(items.is_empty());
}
#[test]
fn parse_memory_items_rejects_unknown_wrapper() {
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");
}
}