use async_trait::async_trait;
#[cfg(feature = "grpc")]
pub mod grpc;
#[cfg(feature = "grpc")]
pub use grpc::CortexdbGrpcMemory;
use harness_core::{Memory, MemoryEntry, MemoryError, Tool};
use harness_mcp_client::McpClient;
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
pub struct CortexdbMemory {
_client: McpClient,
save: Arc<dyn Tool>,
search: Arc<dyn Tool>,
consolidate: Option<Arc<dyn Tool>>,
scope: String,
namespace: String,
user_id: Option<String>,
seq: AtomicU64,
}
impl CortexdbMemory {
pub async fn connect_stdio(program: &str, args: &[&str]) -> anyhow::Result<Self> {
Self::from_client(McpClient::connect_stdio(program, args).await?)
}
#[cfg(feature = "http")]
pub async fn connect_http(url: &str) -> anyhow::Result<Self> {
Self::from_client(McpClient::connect_http(url).await?)
}
pub fn from_client(client: McpClient) -> anyhow::Result<Self> {
let tools = client.tools();
let find = |name: &str| tools.iter().find(|t| t.name() == name).cloned();
let save = find("memory_save")
.ok_or_else(|| anyhow::anyhow!("CortexDB MCP server exposes no `memory_save` tool"))?;
let search = find("memory_search").ok_or_else(|| {
anyhow::anyhow!("CortexDB MCP server exposes no `memory_search` tool")
})?;
let consolidate = find("knowledge_memory_consolidate");
Ok(Self {
_client: client,
save,
search,
consolidate,
scope: "global".into(),
namespace: "harness".into(),
user_id: None,
seq: AtomicU64::new(0),
})
}
pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
self.scope = scope.into();
self
}
pub fn with_namespace(mut self, ns: impl Into<String>) -> Self {
self.namespace = ns.into();
self
}
pub fn with_user_id(mut self, user: impl Into<String>) -> Self {
self.user_id = Some(user.into());
self
}
pub fn can_consolidate(&self) -> bool {
self.consolidate.is_some()
}
pub async fn consolidate(&self, extra_args: Value) -> Result<Value, MemoryError> {
let tool = self.consolidate.as_ref().ok_or_else(|| {
MemoryError::Backend(
"this CortexDB build exposes no `knowledge_memory_consolidate` tool".into(),
)
})?;
let args = build_consolidate_args(
&self.scope,
&self.namespace,
self.user_id.as_deref(),
extra_args,
);
let mut world = harness_context::default_world(".");
let res = tool
.invoke(args, &mut world)
.await
.map_err(|e| MemoryError::Backend(e.to_string()))?;
if !res.ok {
return Err(MemoryError::Backend(format!(
"knowledge_memory_consolidate: {}",
res.content
)));
}
Ok(res.content)
}
fn next_id(&self) -> String {
let n = self.seq.fetch_add(1, Ordering::SeqCst);
let t = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("harness-{t}-{n}")
}
}
#[async_trait]
impl Memory for CortexdbMemory {
async fn write(&self, entry: MemoryEntry) -> Result<(), MemoryError> {
let id = if entry.id.is_empty() {
self.next_id()
} else {
entry.id.clone()
};
let args = build_save_args(
&id,
&entry,
&self.scope,
&self.namespace,
self.user_id.as_deref(),
);
let mut world = harness_context::default_world(".");
let res = self
.save
.invoke(args, &mut world)
.await
.map_err(|e| MemoryError::Backend(e.to_string()))?;
if !res.ok {
return Err(MemoryError::Backend(format!(
"memory_save: {}",
res.content
)));
}
Ok(())
}
async fn recall(&self, query: &str, k: usize) -> Result<Vec<MemoryEntry>, MemoryError> {
if k == 0 || query.trim().is_empty() {
return Ok(Vec::new());
}
let mut args = json!({
"query": query,
"top_k": k,
"scope": self.scope,
"namespace": self.namespace,
});
if let Some(u) = &self.user_id {
args["user_id"] = json!(u);
}
let mut world = harness_context::default_world(".");
let res = self
.search
.invoke(args, &mut world)
.await
.map_err(|e| MemoryError::Backend(e.to_string()))?;
if !res.ok {
return Err(MemoryError::Backend(format!(
"memory_search: {}",
res.content
)));
}
let results = res
.content
.get("results")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
let mut out = Vec::with_capacity(results.len());
for hit in &results {
let rec = hit.get("memory").unwrap_or(hit);
let content = rec
.get("content")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_string();
if content.trim().is_empty() {
continue;
}
let mut entry = MemoryEntry::new(content);
if let Some(id) = rec.get("id").and_then(|v| v.as_str()) {
entry.id = id.to_string();
}
if let Some(meta) = rec.get("metadata") {
if let Some(tags) = meta.get("tags").and_then(|v| v.as_array()) {
entry.tags = tags
.iter()
.filter_map(|t| t.as_str().map(String::from))
.collect();
}
if let Some(src) = meta.get("source").and_then(|v| v.as_str()) {
entry.source = Some(src.to_string());
}
}
out.push(entry);
}
Ok(out)
}
}
fn build_save_args(
id: &str,
entry: &MemoryEntry,
scope: &str,
namespace: &str,
user_id: Option<&str>,
) -> Value {
let mut role: Option<String> = None;
let mut session: Option<String> = None;
let mut other_tags: Vec<String> = Vec::new();
for t in &entry.tags {
if let Some(r) = t.strip_prefix("role:") {
role = Some(r.to_string());
} else if let Some(s) = t.strip_prefix("session:") {
session = Some(s.to_string());
} else {
other_tags.push(t.clone());
}
}
let mut metadata = serde_json::Map::new();
if !other_tags.is_empty() {
metadata.insert("tags".into(), json!(other_tags));
}
if let Some(s) = &entry.source {
metadata.insert("source".into(), json!(s));
}
let mut args = json!({
"memory_id": id,
"content": entry.content,
"scope": scope,
"namespace": namespace,
"metadata": Value::Object(metadata),
});
if let Some(r) = role {
args["role"] = json!(r);
}
if let Some(s) = session {
args["session_id"] = json!(s);
}
if let Some(u) = user_id {
args["user_id"] = json!(u);
}
if let Some(expires_ms) = entry.expires_ms {
let now_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
let secs = (expires_ms - now_ms).div_euclid(1000).max(1);
args["ttl_seconds"] = json!(secs);
}
args
}
fn build_consolidate_args(
scope: &str,
namespace: &str,
user_id: Option<&str>,
extra_args: Value,
) -> Value {
let mut args = json!({ "scope": scope, "namespace": namespace });
if let Some(u) = user_id {
args["user_id"] = json!(u);
}
if let Value::Object(extra) = extra_args
&& let Value::Object(base) = &mut args
{
base.extend(extra);
}
if args.get("reflect").is_none() {
args["reflect"] = json!({
"recall": {
"query": "key facts, entities, decisions, tasks, and tools from recent conversations",
"max_memory_items": 30
}
});
}
if args.get("promote_to_knowledge").is_none() {
args["promote_to_knowledge"] = json!(true);
}
args
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn consolidate_null_fills_required_reflect_defaults() {
let args = build_consolidate_args("global", "harness", Some("u1"), Value::Null);
assert_eq!(args["scope"], "global");
assert_eq!(args["user_id"], "u1");
assert!(args["reflect"]["recall"]["query"].is_string());
assert_eq!(args["reflect"]["recall"]["max_memory_items"], 30);
assert_eq!(args["promote_to_knowledge"], true);
}
#[test]
fn consolidate_extra_args_override_defaults() {
let extra = json!({
"reflect": { "recall": { "query": "custom", "max_memory_items": 5 } },
"promote_to_knowledge": false
});
let args = build_consolidate_args("session", "app", None, extra);
assert_eq!(args["reflect"]["recall"]["query"], "custom");
assert_eq!(args["reflect"]["recall"]["max_memory_items"], 5);
assert_eq!(args["promote_to_knowledge"], false);
assert!(args.get("user_id").is_none());
}
#[test]
fn role_and_session_promoted_to_native_params() {
let entry = MemoryEntry::new("hi there")
.with_source("transcript")
.with_tags(["role:assistant", "session:sess-1", "topic:deploy"]);
let args = build_save_args("m1", &entry, "session", "myapp", Some("u42"));
assert_eq!(args["role"], "assistant");
assert_eq!(args["session_id"], "sess-1");
assert_eq!(args["user_id"], "u42");
assert_eq!(args["scope"], "session");
assert_eq!(args["namespace"], "myapp");
let tags = args["metadata"]["tags"].as_array().unwrap();
assert_eq!(tags.len(), 1);
assert_eq!(tags[0], "topic:deploy");
assert_eq!(args["metadata"]["source"], "transcript");
}
#[test]
fn ttl_is_carried_to_the_server_as_seconds() {
let entry = MemoryEntry::new("ephemeral").with_ttl_days(7);
let args = build_save_args("m3", &entry, "session", "app", None);
let ttl = args["ttl_seconds"].as_i64().expect("ttl_seconds sent");
assert!(
(7 * 86_400 - 60..=7 * 86_400).contains(&ttl),
"expected ~7d in seconds, got {ttl}"
);
}
#[test]
fn no_ttl_means_no_ttl_field() {
let entry = MemoryEntry::new("durable fact");
let args = build_save_args("m4", &entry, "global", "app", None);
assert!(args.get("ttl_seconds").is_none());
}
#[test]
fn already_expired_entry_gets_the_shortest_ttl_not_a_negative_one() {
let mut entry = MemoryEntry::new("stale");
entry.expires_ms = Some(1); let args = build_save_args("m5", &entry, "session", "app", None);
assert_eq!(args["ttl_seconds"], 1);
}
#[test]
fn no_role_session_tags_omits_native_params() {
let entry = MemoryEntry::new("note").with_tags(["misc"]);
let args = build_save_args("m2", &entry, "global", "harness", None);
assert!(args.get("role").is_none());
assert!(args.get("session_id").is_none());
assert!(args.get("user_id").is_none());
assert_eq!(args["metadata"]["tags"][0], "misc");
}
}