Skip to main content

khive_runtime/
keyed_memory.rs

1//! Memory identity is published only by the final DML of its atomic create.
2
3use khive_storage::note::Note;
4use khive_storage::{SqlStatement, SqlValue};
5use khive_types::{Details, KhiveError};
6use serde_json::Value;
7use uuid::Uuid;
8
9use crate::atomic_message::{AtomicNoteOptions, AtomicNoteSpec};
10use crate::atomic_runner::{run_atomic_unit, AtomicOpFailure, AtomicRunOutcome};
11use crate::note_create::{prepare_note_create, KeyPublication, KEY_CLAIM};
12use crate::{KhiveRuntime, NamespaceToken, RuntimeError, RuntimeResult};
13
14pub struct KeyedMemorySpec<'a> {
15    pub content: &'a str,
16    pub key: &'a str,
17    pub salience: f64,
18    pub decay_factor: f64,
19    pub properties: Value,
20    pub source_id: Option<Uuid>,
21    pub embedding_model: Option<&'a str>,
22}
23
24pub fn validate_memory_key(key: &str) -> RuntimeResult<()> {
25    crate::note_validation::validate_key(
26        key,
27        "key must be at most 512 UTF-8 bytes and must not contain U+0000",
28    )
29}
30
31fn idempotency_conflict(key: &str, existing: &Note) -> RuntimeError {
32    KhiveError::conflict(format!(
33        "idempotency_key_conflict: key {key:?} already exists; stored content differs (existing memory {})",
34        existing.id
35    ))
36        .with_details(Details::new_owned([
37            ("reason", "idempotency_key_conflict".into()),
38            ("key", key.to_owned()),
39            ("existing_id", existing.id.to_string()),
40        ]))
41        .into()
42}
43
44async fn resolve_holder(
45    runtime: &KhiveRuntime,
46    token: &NamespaceToken,
47    key: &str,
48) -> RuntimeResult<Option<Note>> {
49    let mut matches = runtime
50        .notes(token)?
51        .get_live_notes_by_key(token.namespace().as_str(), key, Some("memory"))
52        .await?;
53    match matches.len() {
54        0 => Ok(None),
55        1 => Ok(matches.pop()),
56        _ => Err(RuntimeError::Internal(
57            "memory key lookup returned multiple live holders".into(),
58        )),
59    }
60}
61
62struct VisibilityReceiptState {
63    note_present: bool,
64    epoch: Option<String>,
65    receipt_present: bool,
66    fences: Option<Vec<(String, u64)>>,
67}
68
69/// Join identity, independent provenance and original receipt in one snapshot.
70/// Joins use note identity first so a conflicting namespace cannot disappear
71/// behind a filter and masquerade as a genuinely absent receipt or marker.
72async fn visibility_receipt_state(
73    runtime: &KhiveRuntime,
74    token: &NamespaceToken,
75    note_id: Uuid,
76) -> RuntimeResult<VisibilityReceiptState> {
77    let unavailable = || {
78        crate::visibility_receipts::receipt_failure(
79            "receipt_store_unavailable",
80            Some(note_id),
81            false,
82        )
83    };
84    runtime
85        .require_visibility_cutover()
86        .map_err(|_| unavailable())?;
87    let mut reader = runtime.sql().reader().await.map_err(|_| unavailable())?;
88    let rows = reader
89        .query_all(SqlStatement {
90            sql: "SELECT n.kind AS note_kind, n.namespace AS note_namespace, \
91              e.namespace AS epoch_namespace, e.epoch, \
92              r.note_id AS receipt_id, r.namespace AS receipt_namespace, r.model_count, \
93              f.namespace AS fence_namespace, f.model, f.ann_write_log_seq \
94              FROM notes n LEFT JOIN memory_visibility_epochs e ON e.note_id = n.id \
95              LEFT JOIN memory_visibility_receipts r ON r.note_id = n.id \
96              LEFT JOIN memory_visibility_fences f ON f.note_id = n.id \
97              WHERE n.id = ?1 ORDER BY r.namespace, f.namespace, f.model"
98                .into(),
99            params: vec![SqlValue::Text(note_id.to_string())],
100            label: Some("memory-visibility-receipt-read".into()),
101        })
102        .await
103        .map_err(|_| unavailable())?;
104    let mut state = VisibilityReceiptState {
105        note_present: !rows.is_empty(),
106        epoch: None,
107        receipt_present: false,
108        fences: None,
109    };
110    let Some(first) = rows.first() else {
111        return Ok(state);
112    };
113    let namespace = token.namespace().as_str();
114    let text = |value: Option<&SqlValue>, expected: &str| match value {
115        Some(SqlValue::Text(actual)) => actual == expected,
116        _ => false,
117    };
118    let identity_matches = text(first.get("note_kind"), "memory")
119        && text(first.get("note_namespace"), namespace)
120        && text(first.get("epoch_namespace"), namespace);
121    if identity_matches {
122        if let Some(SqlValue::Text(epoch)) = first.get("epoch") {
123            state.epoch = Some(epoch.clone());
124        }
125    }
126    state.receipt_present = rows
127        .iter()
128        .any(|row| matches!(row.get("receipt_id"), Some(SqlValue::Text(_))));
129    let expected_count = match first.get("model_count") {
130        Some(SqlValue::Integer(count)) if *count >= 0 => usize::try_from(*count).ok(),
131        _ => None,
132    };
133    let mut valid = text(first.get("note_kind"), "memory")
134        && text(first.get("note_namespace"), namespace)
135        && state.receipt_present
136        && expected_count.is_some();
137    let mut fences = std::collections::BTreeMap::new();
138    for row in &rows {
139        if (matches!(row.get("receipt_id"), Some(SqlValue::Text(_)))
140            && !text(row.get("receipt_namespace"), namespace))
141            || matches!(row.get("fence_namespace"), Some(SqlValue::Text(ns))
142                if ns != namespace || !text(row.get("receipt_namespace"), ns))
143        {
144            state.epoch = None;
145        }
146        valid &= text(row.get("receipt_namespace"), namespace)
147            && matches!(row.get("model_count"), Some(SqlValue::Integer(count))
148                if usize::try_from(*count).ok() == expected_count);
149        match (
150            row.get("fence_namespace"),
151            row.get("model"),
152            row.get("ann_write_log_seq"),
153        ) {
154            (
155                Some(SqlValue::Text(ns)),
156                Some(SqlValue::Text(model)),
157                Some(SqlValue::Integer(seq)),
158            ) if ns == namespace && !model.is_empty() && *seq > 0 => {
159                valid &= fences.insert(model.clone(), *seq as u64).is_none();
160            }
161            (
162                Some(SqlValue::Null) | None,
163                Some(SqlValue::Null) | None,
164                Some(SqlValue::Null) | None,
165            ) => {}
166            _ => valid = false,
167        }
168    }
169    if valid && expected_count == Some(fences.len()) {
170        state.fences = Some(fences.into_iter().collect());
171    }
172    Ok(state)
173}
174
175/// Read original fences without issuing a token. Callers needing a replay token
176/// must additionally apply the independent provenance classification below.
177pub async fn memory_visibility_receipt(
178    runtime: &KhiveRuntime,
179    token: &NamespaceToken,
180    note_id: Uuid,
181) -> RuntimeResult<Option<Vec<(String, u64)>>> {
182    Ok(visibility_receipt_state(runtime, token, note_id)
183        .await?
184        .fences)
185}
186
187async fn classified_visibility_receipt(
188    runtime: &KhiveRuntime,
189    token: &NamespaceToken,
190    note_id: Uuid,
191    replay: bool,
192) -> RuntimeResult<Vec<(String, u64)>> {
193    let state = visibility_receipt_state(runtime, token, note_id).await?;
194    if replay && !state.note_present {
195        // The attempt's own unit rolled back on the key claim, so this request
196        // wrote nothing; the reason carries that to the error projection.
197        return Err(KhiveError::not_found("memory", "keyed replay holder")
198            .with_details(Details::new_owned([(
199                "reason",
200                "keyed_replay_holder_missing".into(),
201            )]))
202            .into());
203    }
204    let reason = match state.epoch.as_deref() {
205        Some("modern") => {
206            if let Some(fences) = state.fences {
207                return Ok(fences);
208            }
209            "receipt_temporarily_unavailable"
210        }
211        Some("legacy") if !state.receipt_present => "legacy_receipt_absent",
212        _ => "receipt_epoch_unknown",
213    };
214    Err(crate::visibility_receipts::receipt_failure(
215        reason,
216        Some(note_id),
217        replay,
218    ))
219}
220
221pub async fn create_keyed_memory(
222    runtime: &KhiveRuntime,
223    token: &NamespaceToken,
224    spec: KeyedMemorySpec<'_>,
225) -> RuntimeResult<(Note, Option<Uuid>, bool)> {
226    let (note, edge_id, replayed, _) =
227        create_keyed_memory_with_report(runtime, token, spec).await?;
228    Ok((note, edge_id, replayed))
229}
230
231/// Same as [`create_keyed_memory`], also returning the embedding-input
232/// truncation report computed for this call. A replay stores nothing new but
233/// retains the report from preparing this call's identical content.
234pub async fn create_keyed_memory_with_report(
235    runtime: &KhiveRuntime,
236    token: &NamespaceToken,
237    spec: KeyedMemorySpec<'_>,
238) -> RuntimeResult<(
239    Note,
240    Option<Uuid>,
241    bool,
242    crate::retrieval::EmbeddingTruncationReport,
243)> {
244    let (note, edge_id, replayed, _, report) =
245        create_keyed_memory_with_receipt_and_report(runtime, token, spec).await?;
246    Ok((note, edge_id, replayed, report))
247}
248
249/// Receipt-bearing keyed create. The original per-model receipt is retained
250/// on replay, including an explicit receipt for a write with zero models.
251pub async fn create_keyed_memory_with_receipt(
252    runtime: &KhiveRuntime,
253    token: &NamespaceToken,
254    spec: KeyedMemorySpec<'_>,
255) -> RuntimeResult<(Note, Option<Uuid>, bool, Vec<(String, u64)>)> {
256    let (note, edge_id, replayed, fences, _) =
257        create_keyed_memory_with_receipt_and_report(runtime, token, spec).await?;
258    Ok((note, edge_id, replayed, fences))
259}
260
261/// Return both the original visibility receipt and the embedding-input report.
262/// Replays retain the stored fences and the report computed while preparing
263/// this call's embedding input; the report does not describe the original write.
264pub async fn create_keyed_memory_with_receipt_and_report(
265    runtime: &KhiveRuntime,
266    token: &NamespaceToken,
267    spec: KeyedMemorySpec<'_>,
268) -> RuntimeResult<(
269    Note,
270    Option<Uuid>,
271    bool,
272    Vec<(String, u64)>,
273    crate::retrieval::EmbeddingTruncationReport,
274)> {
275    validate_memory_key(spec.key)?;
276    if spec.content.trim().is_empty() {
277        return Err(RuntimeError::InvalidInput(
278            "content must not be empty".into(),
279        ));
280    }
281    let (mut prepared, annotation_ids) = prepare_note_create(
282        runtime,
283        AtomicNoteSpec {
284            token,
285            id: None,
286            kind: "memory",
287            name: None,
288            content: spec.content,
289            properties: Some(spec.properties),
290        },
291        AtomicNoteOptions {
292            salience: Some(spec.salience),
293            decay_factor: Some(spec.decay_factor),
294            embedding_model: spec.embedding_model,
295            key: Some(spec.key),
296            memory_visibility_receipt: true,
297            ..Default::default()
298        },
299        &spec.source_id.into_iter().collect::<Vec<_>>(),
300        KeyPublication::AfterDependents,
301    )
302    .await?;
303    let mut note = prepared.notes.remove(0);
304    let edge_id = annotation_ids.first().copied();
305
306    for _attempt in 0..2 {
307        #[cfg(test)]
308        crate::keyed_memory_tests::checkpoint(token.namespace().as_str(), _attempt, false).await;
309        match run_atomic_unit(runtime.sql().as_ref(), prepared.plans.clone()).await {
310            Ok(AtomicRunOutcome::Committed { .. }) => {
311                note.key = Some(spec.key.to_owned());
312                note.version = 2;
313                let fences = classified_visibility_receipt(runtime, token, note.id, false).await?;
314                return Ok((note, edge_id, false, fences, prepared.embedding_truncation));
315            }
316            Ok(AtomicRunOutcome::RolledBack {
317                failure:
318                    AtomicOpFailure::GuardFailed {
319                        statement_label,
320                        observed: 0,
321                        ..
322                    },
323                ..
324            }) if statement_label.as_deref() == Some(KEY_CLAIM) => {
325                #[cfg(test)]
326                crate::keyed_memory_tests::checkpoint(token.namespace().as_str(), _attempt, true)
327                    .await;
328                if let Some(holder) = resolve_holder(runtime, token, spec.key).await? {
329                    if holder.content == spec.content {
330                        let fences =
331                            classified_visibility_receipt(runtime, token, holder.id, true).await?;
332                        return Ok((holder, None, true, fences, prepared.embedding_truncation));
333                    }
334                    return Err(idempotency_conflict(spec.key, &holder));
335                }
336            }
337            Ok(AtomicRunOutcome::RolledBack {
338                failed_op_index,
339                failure,
340            }) => {
341                // Admission is cached after the first successful cutover check,
342                // so a receipt store that broke since then surfaces here.
343                runtime.recheck_visibility_cutover()?;
344                return Err(RuntimeError::Internal(format!(
345                    "atomic memory write rolled back at op {failed_op_index}: {failure:?}"
346                )));
347            }
348            Err(error) => return Err(RuntimeError::Storage(error.0)),
349        }
350    }
351    Err(
352        KhiveError::unavailable("memory key holder disappeared during reconciliation")
353            .with_details(Details::new_owned([
354                ("reason", "key_holder_unresolved".into()),
355                ("key", spec.key.to_owned()),
356            ]))
357            .into(),
358    )
359}
360
361#[cfg(test)]
362mod receipt_read_tests {
363    use super::*;
364    use crate::DomainDisposition;
365
366    const KEY: &str = "receipt-race-key";
367    const CONTENT: &str = "private receipt race content";
368
369    async fn captured_holder() -> (KhiveRuntime, NamespaceToken, Note) {
370        let runtime = KhiveRuntime::memory().unwrap();
371        runtime.install_kind_registry(vec![], vec!["memory".into()]);
372        let token = runtime
373            .authorize(khive_types::Namespace::parse("receipt-read-race").unwrap())
374            .unwrap();
375        let (note, _, replayed, fences) = create_keyed_memory_with_receipt(
376            &runtime,
377            &token,
378            KeyedMemorySpec {
379                content: CONTENT,
380                key: KEY,
381                salience: 0.7,
382                decay_factor: 0.95,
383                properties: serde_json::json!({}),
384                source_id: None,
385                embedding_model: None,
386            },
387        )
388        .await
389        .unwrap();
390        assert!(!replayed);
391        assert!(fences.is_empty());
392        let holder = resolve_holder(&runtime, &token, KEY)
393            .await
394            .unwrap()
395            .unwrap();
396        assert_eq!(holder, note);
397        assert!(
398            classified_visibility_receipt(&runtime, &token, holder.id, true)
399                .await
400                .unwrap()
401                .is_empty()
402        );
403        (runtime, token, holder)
404    }
405
406    #[tokio::test]
407    async fn replay_receipt_reports_missing_after_captured_holder_is_hard_deleted() {
408        let (runtime, token, holder) = captured_holder().await;
409        // This is the exact boundary in the replay path: holder lookup has
410        // completed, but its joined receipt/epoch read has not started.
411        assert!(runtime.delete_note(&token, holder.id, true).await.unwrap());
412        assert!(runtime
413            .notes(&token)
414            .unwrap()
415            .get_note_including_deleted(holder.id)
416            .await
417            .unwrap()
418            .is_none());
419        let error = classified_visibility_receipt(&runtime, &token, holder.id, true)
420            .await
421            .unwrap_err();
422        assert!(
423            matches!(&error, RuntimeError::Khive(e) if e.kind() == khive_types::ErrorKind::NotFound)
424        );
425        let value = crate::error_projection::runtime_error_value(error, DomainDisposition::Unknown);
426        assert_eq!(value["kind"], "not_found");
427        assert_eq!(value["details"]["reason"], "keyed_replay_holder_missing");
428        assert_eq!(value["domain_disposition"], "not_committed");
429        assert!(value.get("retryable").is_none());
430        let encoded = value.to_string();
431        let id = holder.id.to_string();
432        for private in [KEY, CONTENT, id.as_str()] {
433            assert!(!encoded.contains(private));
434        }
435        assert!(memory_visibility_receipt(&runtime, &token, holder.id)
436            .await
437            .unwrap()
438            .is_none());
439        assert!(resolve_holder(&runtime, &token, KEY)
440            .await
441            .unwrap()
442            .is_none());
443
444        // The post-commit caller retains its existing uncertainty semantics.
445        let error = classified_visibility_receipt(&runtime, &token, holder.id, false)
446            .await
447            .unwrap_err();
448        let value = crate::error_projection::runtime_error_value(error, DomainDisposition::Unknown);
449        assert_eq!(value["details"]["reason"], "receipt_epoch_unknown");
450        assert_eq!(value["domain_disposition"], "unknown");
451    }
452
453    #[tokio::test]
454    async fn present_holder_with_missing_or_invalid_epoch_is_not_missing() {
455        for mutation in [
456            "DELETE FROM memory_visibility_epochs WHERE note_id = ?1",
457            "UPDATE memory_visibility_epochs SET epoch = 'unknown' WHERE note_id = ?1",
458            "UPDATE memory_visibility_epochs SET epoch = 'malformed' WHERE note_id = ?1",
459            "UPDATE memory_visibility_epochs SET epoch = 'legacy' WHERE note_id = ?1",
460            "UPDATE memory_visibility_epochs SET namespace = 'foreign-private-namespace' WHERE note_id = ?1",
461            "UPDATE notes SET namespace = 'foreign-private-namespace' WHERE id = ?1",
462        ] {
463            let (runtime, token, holder) = captured_holder().await;
464            {
465                // Model a damaged epoch value as well as valid but incomplete
466                // provenance. Restore constraint checks before reading it.
467                let writer = runtime.backend().pool().try_writer().unwrap();
468                writer
469                    .conn()
470                    .pragma_update(None, "ignore_check_constraints", true)
471                    .unwrap();
472                writer.conn().execute(mutation, [holder.id.to_string()]).unwrap();
473                writer
474                    .conn()
475                    .pragma_update(None, "ignore_check_constraints", false)
476                    .unwrap();
477            }
478            assert!(runtime
479                .notes(&token)
480                .unwrap()
481                .get_note_including_deleted(holder.id)
482                .await
483                .unwrap()
484                .is_some());
485            let error = classified_visibility_receipt(&runtime, &token, holder.id, true)
486                .await
487                .unwrap_err();
488            let value = crate::error_projection::runtime_error_value(error, DomainDisposition::Unknown);
489            assert_eq!(value["details"]["reason"], "receipt_epoch_unknown");
490            assert_eq!(value["details"]["memory_id"], holder.id.to_string());
491            assert_eq!(value["domain_disposition"], "not_committed");
492            assert_eq!(value["retryable"], false);
493            for private in [KEY, CONTENT, "foreign-private-namespace"] {
494                assert!(!value.to_string().contains(private));
495            }
496        }
497    }
498
499    #[tokio::test]
500    async fn unreadable_note_store_is_unavailable_not_missing() {
501        let (runtime, token, holder) = captured_holder().await;
502        runtime
503            .sql()
504            .writer()
505            .await
506            .unwrap()
507            .execute(SqlStatement {
508                sql: "ALTER TABLE notes RENAME TO private_unreadable_notes".into(),
509                params: vec![],
510                label: Some("receipt-unreadable-note-control".into()),
511            })
512            .await
513            .unwrap();
514        let error = classified_visibility_receipt(&runtime, &token, holder.id, true)
515            .await
516            .unwrap_err();
517        let value = crate::error_projection::runtime_error_value(error, DomainDisposition::Unknown);
518        assert_eq!(value["details"]["reason"], "receipt_store_unavailable");
519        assert_eq!(value["details"]["memory_id"], holder.id.to_string());
520        assert_eq!(value["domain_disposition"], "unknown");
521        assert_eq!(value["retryable"], true);
522        for private in [KEY, CONTENT, "private_unreadable_notes"] {
523            assert!(!value.to_string().contains(private));
524        }
525    }
526}