Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 69 additions & 34 deletions crates/tinymemory-remote/src/cognee.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
//! Self-hosted Cognee REST adapter.

use anyhow::{anyhow, Context};
use anyhow::Context;
use async_trait::async_trait;
use reqwest::{multipart, Method};
use serde_json::{json, Value};
Expand Down Expand Up @@ -265,14 +265,7 @@ impl CogneeDialect {
serde_json::from_str(&raw).context("Cognee record envelope is invalid")?;
entry.remote_id = format!("{}:{id}", dataset.id);
if entry.timestamp.is_empty() {
entry.timestamp = data
.get("updatedAt")
.or_else(|| data.get("updated_at"))
.or_else(|| data.get("createdAt"))
.or_else(|| data.get("created_at"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_owned();
entry.timestamp = Self::listing_timestamp(data);
}
entries.push(entry);
}
Expand All @@ -285,7 +278,21 @@ impl CogneeDialect {
/// key resolves by matching the name — the per-record raw-fetch loop the
/// first cut ran existed only because it read the key out of each
/// envelope body instead.
async fn data_index(&self, dataset: &Dataset) -> anyhow::Result<Vec<(String, String)>> {
/// The listing row's write time. Checked per candidate with `find_map`,
/// NOT an `or_else` chain over `Value::get`: real Cognee serializes
/// `"updatedAt": null` for every never-updated record, and `get` on a
/// present-but-null key answers `Some(Null)` — an `or_else` chain commits
/// to it and never reaches `createdAt`, emptying every timestamp
/// (issue #75).
fn listing_timestamp(data: &Value) -> String {
["updatedAt", "updated_at", "createdAt", "created_at"]
.iter()
.find_map(|key| data.get(key).and_then(Value::as_str))
.unwrap_or_default()
.to_owned()
}

async fn data_index(&self, dataset: &Dataset) -> anyhow::Result<Vec<(String, String, String)>> {
let response: Value = self
.client
.json(
Expand All @@ -303,6 +310,12 @@ impl CogneeDialect {
Some((
data.get("id")?.as_str()?.to_owned(),
data.get("name")?.as_str()?.to_owned(),
// Kept alongside the id: the envelope this adapter
// uploads carries an empty timestamp, so the listing is
// the ONLY source a keyed fetch can backfill from (the
// enumeration path already does — the keyed path must
// agree with it).
Self::listing_timestamp(data),
))
})
.collect())
Expand All @@ -313,16 +326,23 @@ impl CogneeDialect {
/// On a duplicate name (possible only if a historical blind re-add ever
/// raced), newest-listed wins deterministically — the listing is
/// insertion-ordered — rather than an arbitrary pick.
async fn find_data_id(&self, dataset: &Dataset, key: &str) -> anyhow::Result<Option<String>> {
async fn find_data_id(
&self,
dataset: &Dataset,
key: &str,
) -> anyhow::Result<Option<(String, String)>> {
let uploaded = Self::filename(key);
let stripped = uploaded.trim_end_matches(".json").to_owned();
let stripped = uploaded
.strip_suffix(".json")
.unwrap_or(&uploaded)
.to_owned();
Ok(self
.data_index(dataset)
.await?
.into_iter()
.rev()
.find(|(_, name)| *name == uploaded || *name == stripped)
.map(|(id, _)| id))
.find(|(_, name, _)| *name == uploaded || *name == stripped)
.map(|(id, _, timestamp)| (id, timestamp)))
}

/// Downloads and decodes ONE envelope by its ids, verifying it is the
Expand All @@ -333,6 +353,7 @@ impl CogneeDialect {
&self,
dataset: &Dataset,
data_id: &str,
listing_timestamp: &str,
namespace: &str,
key: &str,
) -> anyhow::Result<Option<StoredEntry>> {
Expand All @@ -355,6 +376,13 @@ impl CogneeDialect {
);
}
entry.remote_id = format!("{}:{data_id}", dataset.id);
// The uploaded envelope's timestamp is empty by construction
// (StoredEntry::new), so without this backfill every keyed get would
// answer an empty timestamp while the enumeration path answers the
// listing's — the same record disagreeing with itself.
if entry.timestamp.is_empty() {
entry.timestamp = listing_timestamp.to_owned();
}
Ok(Some(entry))
}

Expand Down Expand Up @@ -393,21 +421,29 @@ impl Dialect for CogneeDialect {
let Some(dataset) = self.find_dataset(namespace).await? else {
return Ok(None);
};
let Some(data_id) = self.find_data_id(&dataset, key).await? else {
let Some((data_id, listing_timestamp)) = self.find_data_id(&dataset, key).await? else {
return Ok(None);
};
self.fetch_entry(&dataset, &data_id, namespace, key).await
self.fetch_entry(&dataset, &data_id, &listing_timestamp, namespace, key)
.await
}

/// Replaces an existing envelope and uploads the new exact record.
async fn upsert(&self, entry: StoredEntry) -> anyhow::Result<()> {
// Through the keyed seam: dataset + listing, no raw fan-out. The
// existing record's ids are all the replace path needs.
// existing record's ids are all the replace path needs. Like delete,
// the PATCH trusts the dataset-scoped filename match without an
// envelope read — the name is the key's SHA-256 digest, so a wrong
// target needs a digest collision, and what a collision would cost
// here is an overwrite of the colliding record's content with THIS
// key's envelope (recoverable by that record's next upsert, unlike
// delete's unrecoverable removal — which is the sharper case and got
// this same argument first).
let existing = match self.find_dataset(&entry.namespace).await? {
Some(dataset) => self
.find_data_id(&dataset, &entry.key)
.await?
.map(|data_id| (dataset, data_id)),
.map(|(data_id, _)| (dataset, data_id)),
None => None,
};
let body = serde_json::to_vec(&entry)?;
Expand All @@ -432,19 +468,7 @@ impl Dialect for CogneeDialect {
.text("run_in_background", "false"),
)
};
let response = self
.client
.multipart(method, &path)?
.multipart(form)
.send()
.await?;
if !response.status().is_success() {
return Err(anyhow!(
"memory API {path} returned HTTP {}",
response.status()
));
}
Ok(())
self.client.send_multipart(method, &path, form).await
}

/// Enumerates records across all TinyMemory-owned Cognee datasets.
Expand All @@ -463,7 +487,18 @@ impl Dialect for CogneeDialect {
limit: usize,
opts: RecallOpts<'_>,
) -> anyhow::Result<Vec<StoredEntry>> {
let datasets = opts.namespace.map(Self::dataset_name);
// Resolve the dataset BEFORE asking recall (issue #75): real Cognee
// 404s ("No datasets found") when a named dataset resolves to
// nothing, so the first recall in a fresh namespace — before its
// first store — errored where every sibling op answers empty. One
// extra listing request, the same price entry()/delete() pay.
let datasets = match opts.namespace {
Some(namespace) => match self.find_dataset(namespace).await? {
Some(dataset) => Some(vec![dataset.name]),
None => return Ok(Vec::new()),
},
None => None,
};
let response: Value = self
.client
.json(
Expand All @@ -472,7 +507,7 @@ impl Dialect for CogneeDialect {
Some(&json!({
"query": query,
"search_type": "CHUNKS",
"datasets": datasets.map(|name| vec![name]),
"datasets": datasets,
"top_k": limit,
"only_context": true,
"session_id": opts.session_id
Expand Down Expand Up @@ -519,7 +554,7 @@ impl Dialect for CogneeDialect {
let Some(dataset) = self.find_dataset(namespace).await? else {
return Ok(false);
};
let Some(data_id) = self.find_data_id(&dataset, key).await? else {
let Some((data_id, _)) = self.find_data_id(&dataset, key).await? else {
return Ok(false);
};
self.client
Expand Down
127 changes: 126 additions & 1 deletion crates/tinymemory-remote/src/cognee_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,15 @@ async fn data(State(state): State<AppState>) -> Json<Value> {
let name = state.2.lock().expect("name lock").clone();
let values = if state.0.lock().expect("state lock").is_some() {
let name = name.unwrap_or_else(|| "6b6579.tinymemory".to_owned());
vec![json!({"id": "data-1", "name": name, "created_at": "2026-08-12T00:00:00Z"})]
// `updatedAt` present-but-null is what real Cognee serializes for a
// never-updated record: the backfill must fall through to
// `created_at` instead of committing to the null (issue #75).
vec![json!({
"id": "data-1",
"name": name,
"updatedAt": null,
"created_at": "2026-08-12T00:00:00Z"
})]
} else {
vec![]
};
Expand Down Expand Up @@ -266,6 +274,23 @@ async fn native_cognee_round_trips_the_tinymemory_contract() {
.expect("entry");
assert_eq!(entry.content, "updated knowledge graph");
assert_eq!(entry.taint, MemoryTaint::ExternalSync);
// The uploaded envelope's timestamp is empty by construction, so the
// keyed fetch must backfill from the listing the way the enumeration
// path always did — get and list answering different timestamps for the
// same record was the #71-review regression, and a null `updatedAt`
// halting the fallback chain was issue #75's.
assert_eq!(
entry.timestamp, "2026-08-12T00:00:00Z",
"keyed get backfills the listing timestamp past the null updatedAt"
);
let listed = driver
.list(Some("project"), None, None)
.await
.expect("list");
assert_eq!(
listed[0].timestamp, entry.timestamp,
"keyed get and namespace list agree on the timestamp"
);
assert_eq!(
driver
.recall(
Expand Down Expand Up @@ -366,3 +391,103 @@ async fn keyed_ops_never_fan_out_over_raw_fetches() {
let counts = *state.1.lock().expect("counts");
assert_eq!(counts.raws, 0, "a keyed delete reads no envelopes");
}

/// The envelope-over-filename trust boundary, pinned the way its mem0 twin
/// is. A file whose deterministic name matches the asked key but whose
/// envelope names a DIFFERENT record (hash collision, foreign file wearing
/// our extension) must refuse — never serve someone else's memory.
#[tokio::test]
async fn a_filename_match_with_a_foreign_envelope_is_refused() {
use crate::common::StoredEntry;

let foreign = serde_json::to_vec(&StoredEntry::new(
"someone-elses-ns",
"decision",
"not yours",
MemoryCategory::Conversation,
None,
MemoryTaint::Internal,
))
.expect("envelope");
let name = super::CogneeDialect::filename("decision");
let app = Router::new()
.route(
"/api/v1/datasets/",
get(|| async {
Json(json!([{
"id": "dataset-1",
"name": super::CogneeDialect::dataset_name("project")
}]))
}),
)
.route(
"/api/v1/datasets/{dataset}/data",
get(move || {
let name = name.clone();
async move {
Json(json!([{
"id": "data-1",
"name": name,
"created_at": "2026-08-12T00:00:00Z"
}]))
}
}),
)
.route(
"/api/v1/datasets/{dataset}/data/{data}/raw",
get(move || {
let body = foreign.clone();
async move { (StatusCode::OK, body) }
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let endpoint = format!("http://{}", listener.local_addr().expect("address"));
tokio::spawn(async move {
axum::serve(listener, app).await.expect("serve");
});

let driver =
crate::cognee_provider(super::CogneeMemory::self_hosted(&endpoint, None).expect("client"));
let err = driver.get("project", "decision").await;
let message = format!("{:?}", err.expect_err("mismatched envelope must refuse"));
assert!(
message.contains("envelope names"),
"the refusal names the mismatch: {message}"
);
}

/// Issue #75: recall in a namespace that has never stored anything answers
/// EMPTY, like every sibling op — real Cognee 404s a recall naming a dataset
/// that resolves to nothing. The double serves no /api/v1/recall route at
/// all, so this also proves the guard short-circuits before asking.
#[tokio::test]
async fn recall_in_a_fresh_namespace_is_empty_not_an_error() {
use tinymemory_api::recall::OwnedRecallOpts;

let app = Router::new().route("/api/v1/datasets/", get(|| async { Json(json!([])) }));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let endpoint = format!("http://{}", listener.local_addr().expect("address"));
tokio::spawn(async move {
axum::serve(listener, app).await.expect("serve");
});

let driver =
crate::cognee_provider(super::CogneeMemory::self_hosted(&endpoint, None).expect("client"));
let hits = driver
.recall(
"anything",
3,
&OwnedRecallOpts {
namespace: Some("never-stored".into()),
..OwnedRecallOpts::default()
},
None,
)
.await
.expect("a fresh namespace recalls empty, not an error");
assert!(hits.is_empty());
}
Loading
Loading