cel_memory/
write_hook.rs

1//! The [`MemoryWriteHook`] trait — per-write governance for memory backends.
2//!
3//! Memory writes are themselves governable events. Before a provider
4//! persists a chunk, it calls `MemoryWriteHook::before_write(chunk)`. The
5//! hook can:
6//!
7//! - return `Ok(WriteDecision::Allow)` — the provider persists the chunk
8//!   verbatim.
9//! - return `Ok(WriteDecision::Redact { reason })` — the provider drops
10//!   the chunk silently (the `before_write` event still gets matcher
11//!   coverage for audit; the chunk just doesn't land in `memory_chunks`).
12//!   Used to honor `redact_memory` rules.
13//! - return `Err(_)` — the write surfaces as a hard error to the caller.
14//!
15//! A production runtime can wire this to any policy engine: synthesize an event
16//! from the chunk's kind, source, caller, and content, then return the
17//! appropriate decision. Without a hook, providers persist every write.
18
19use async_trait::async_trait;
20use serde::{Deserialize, Serialize};
21
22use crate::chunk::NewMemoryChunk;
23use crate::error::Result;
24
25/// What a [`MemoryWriteHook`] tells the provider to do with an incoming
26/// chunk.
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
28#[serde(rename_all = "snake_case")]
29pub enum WriteDecision {
30    /// Persist the chunk verbatim. The default outcome.
31    Allow,
32    /// Drop the chunk before persistence. The provider returns success
33    /// to the caller (so write paths don't need to special-case redaction)
34    /// but the chunk never lands in `memory_chunks`. The reason string
35    /// surfaces in logs and in any `record_access`-style audit trail.
36    Redact {
37        /// Human-readable cause (typically the matched rule's name).
38        reason: String,
39    },
40}
41
42/// The hook providers consult before persisting a chunk.
43///
44/// Implementations must be cheap: this fires on every write through the
45/// memory subsystem. Keep implementations proportional to the number of rules
46/// or checks they evaluate.
47#[async_trait]
48pub trait MemoryWriteHook: Send + Sync {
49    /// Decide what to do with the chunk. Default impl returns `Allow`,
50    /// which makes wiring a hook optional — most consumers can ignore
51    /// this trait entirely.
52    async fn before_write(&self, _chunk: &NewMemoryChunk) -> Result<WriteDecision> {
53        Ok(WriteDecision::Allow)
54    }
55}
56
57/// Convenience wrapper for callers that want to plug a closure in without
58/// declaring a struct. Useful for tests, examples, and small policy adapters.
59pub struct ClosureHook<F>(pub F)
60where
61    F: Fn(&NewMemoryChunk) -> WriteDecision + Send + Sync + 'static;
62
63#[async_trait]
64impl<F> MemoryWriteHook for ClosureHook<F>
65where
66    F: Fn(&NewMemoryChunk) -> WriteDecision + Send + Sync + 'static,
67{
68    async fn before_write(&self, chunk: &NewMemoryChunk) -> Result<WriteDecision> {
69        Ok((self.0)(chunk))
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::chunk::{ChunkKind, ChunkSource};
77
78    fn nc(content: &str) -> NewMemoryChunk {
79        NewMemoryChunk {
80            kind: ChunkKind::Chat,
81            source: ChunkSource::Embedded,
82            session_id: None,
83            project_root: None,
84            caller_id: "test".into(),
85            content: content.into(),
86            metadata: serde_json::Value::Null,
87            importance: None,
88            shareable: false,
89            pinned: false,
90        }
91    }
92
93    #[tokio::test]
94    async fn closure_hook_redacts_matching_content() {
95        let hook = ClosureHook(|c: &NewMemoryChunk| {
96            if c.content.contains("secret") {
97                WriteDecision::Redact {
98                    reason: "secret keyword".into(),
99                }
100            } else {
101                WriteDecision::Allow
102            }
103        });
104        let chunk = nc("a totally innocent chat");
105        assert_eq!(
106            hook.before_write(&chunk).await.unwrap(),
107            WriteDecision::Allow
108        );
109        let chunk = nc("contains a secret");
110        match hook.before_write(&chunk).await.unwrap() {
111            WriteDecision::Redact { reason } => assert_eq!(reason, "secret keyword"),
112            other => panic!("expected Redact, got {other:?}"),
113        }
114    }
115
116    #[tokio::test]
117    async fn default_hook_allows_everything() {
118        struct Default;
119        #[async_trait]
120        impl MemoryWriteHook for Default {}
121
122        let chunk = nc("anything");
123        assert_eq!(
124            Default.before_write(&chunk).await.unwrap(),
125            WriteDecision::Allow
126        );
127    }
128}
cel_memory/write_hook.rs

cel_memory/
write_hook.rs
