a3s_code_core/durable_memory/
context.rs

1use super::{
2    fusion::fuse_lexical_semantic, DurableMemoryMode, DurableMemoryRecallChannel,
3    DurableMemoryRecallHit, DurableMemoryRecallPreview, DurableMemorySession,
4};
5use crate::context::{ContextAssembly, ContextItem, ContextResult, ContextType};
6use a3s_memory::repository::{
7    DurableMemoryKind, MemoryAccessEvent, MemoryNode, MemoryQuery, MemoryRelationKind,
8    MemoryRepositoryError, MemoryStatus,
9};
10use chrono::{DateTime, Utc};
11use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
12use sha2::{Digest, Sha256};
13use std::collections::HashSet;
14use tokio_util::sync::CancellationToken;
15
16/// Legacy schema-3 profile. It separates sessions and process-local run IDs but
17/// cannot distinguish a run ID reused after session retention and restart.
18pub const DURABLE_MEMORY_CONTEXT_ID_PROFILE_V1: &str =
19    "a3s.code.memory.context.session-run-sequence-sha256.v1";
20
21/// Current profile. The invocation incarnation is generated by Code rather
22/// than the host environment because host IDs guarantee uniqueness only inside
23/// one process and may legitimately restart from the same sequence.
24pub const DURABLE_MEMORY_CONTEXT_ID_PROFILE_V2: &str =
25    "a3s.code.memory.context.session-run-invocation-sequence-sha256.v2";
26
27const PROVIDER: &str = "durable_memory_v2";
28const RELATED_SCORE_FACTOR: f32 = 0.75;
29
30#[derive(Clone)]
31pub(super) struct RecallCandidate {
32    pub(super) node: MemoryNode,
33    pub(super) score: f32,
34    pub(super) channel: DurableMemoryRecallChannel,
35    pub(super) related_from: Option<String>,
36}
37
38impl RecallCandidate {
39    fn into_preview_hit(self) -> DurableMemoryRecallHit {
40        DurableMemoryRecallHit {
41            node_id: self.node.id,
42            node_revision: self.node.revision,
43            kind: self.node.kind,
44            content: self.node.content,
45            score: self.score,
46            channel: self.channel,
47            related_from: self.related_from,
48        }
49    }
50}
51
52#[derive(Debug, Clone)]
53pub(crate) struct DurableMemoryRecallIdentity {
54    pub(crate) item_id: String,
55    source: String,
56    pub(crate) node_id: String,
57    pub(crate) node_revision: u64,
58    content_digest: String,
59}
60
61impl DurableMemoryRecallIdentity {
62    fn matches(&self, item: &ContextItem) -> bool {
63        item.id == self.item_id
64            && item.source.as_deref() == Some(self.source.as_str())
65            && digest(&item.content) == self.content_digest
66    }
67}
68
69pub(crate) struct DurableMemoryContextBatch {
70    pub(crate) result: ContextResult,
71    pub(crate) identities: Vec<DurableMemoryRecallIdentity>,
72}
73
74impl DurableMemorySession {
75    /// Run the same pure active-only retrieval used before context assembly.
76    /// This diagnostic does not record admission or use and does not inject
77    /// any returned content into a model prompt.
78    pub async fn preview_recall(
79        &self,
80        text: &str,
81    ) -> Result<DurableMemoryRecallPreview, MemoryRepositoryError> {
82        Ok(DurableMemoryRecallPreview {
83            hits: self
84                .query_recall_candidates(text, CancellationToken::new())
85                .await?
86                .into_iter()
87                .map(RecallCandidate::into_preview_hit)
88                .collect(),
89        })
90    }
91
92    #[cfg(test)]
93    pub(crate) async fn query_active_context(
94        &self,
95        text: &str,
96    ) -> Result<DurableMemoryContextBatch, MemoryRepositoryError> {
97        self.query_active_context_with_cancellation(text, CancellationToken::new())
98            .await
99    }
100
101    pub(crate) async fn query_active_context_with_cancellation(
102        &self,
103        text: &str,
104        cancellation: CancellationToken,
105    ) -> Result<DurableMemoryContextBatch, MemoryRepositoryError> {
106        let hits = self.query_recall_candidates(text, cancellation).await?;
107        let mut result = ContextResult::new(PROVIDER);
108        let mut identities = Vec::new();
109        for hit in hits {
110            let node = hit.node;
111            let encoded_id = utf8_percent_encode(&node.id, NON_ALPHANUMERIC);
112            let source = format!("a3s-memory://{encoded_id}?revision={}", node.revision);
113            let item_id = format!("a3s-memory-v2:{}:r{}", node.id, node.revision);
114            let content_digest = digest(&node.content);
115            let token_count = (node.content.len() / 4).max(1);
116            let item = ContextItem::new(&item_id, ContextType::Memory, &node.content)
117                .with_relevance(hit.score)
118                .with_token_count(token_count)
119                .with_source(&source)
120                .with_metadata("memory_node_id", serde_json::json!(node.id))
121                .with_metadata("memory_node_revision", serde_json::json!(node.revision))
122                .with_metadata("memory_kind", serde_json::json!(kind_label(node.kind)))
123                .with_metadata("evidence_count", serde_json::json!(node.evidence.len()))
124                .with_metadata(
125                    "retrieval_channel",
126                    serde_json::json!(channel_label(hit.channel)),
127                )
128                .with_provenance(PROVIDER)
129                .with_priority(0.4)
130                .with_trust(0.8)
131                .with_freshness(0.6);
132            let item = match hit.related_from {
133                Some(source_id) => item.with_metadata("related_from", serde_json::json!(source_id)),
134                None => item,
135            };
136            identities.push(DurableMemoryRecallIdentity {
137                item_id,
138                source,
139                node_id: node.id,
140                node_revision: node.revision,
141                content_digest,
142            });
143            result.add_item(item);
144        }
145        Ok(DurableMemoryContextBatch { result, identities })
146    }
147
148    async fn query_recall_candidates(
149        &self,
150        text: &str,
151        cancellation: CancellationToken,
152    ) -> Result<Vec<RecallCandidate>, MemoryRepositoryError> {
153        let Some(policy) = self.recall_policy() else {
154            return Ok(Vec::new());
155        };
156        if self.mode() != DurableMemoryMode::ActiveRecall
157            || !text.chars().any(char::is_alphanumeric)
158        {
159            return Ok(Vec::new());
160        }
161
162        let query = MemoryQuery::new(self.namespace().clone())
163            .with_text(text)
164            .with_limit(policy.max_results());
165        let lexical = self
166            .repository()
167            .query(query)
168            .await?
169            .hits
170            .into_iter()
171            .filter(|hit| hit.score.total >= policy.min_lexical_score())
172            .map(|hit| RecallCandidate {
173                node: hit.node,
174                score: hit.score.total,
175                channel: DurableMemoryRecallChannel::Lexical,
176                related_from: None,
177            })
178            .collect::<Vec<_>>();
179        let semantic = match self.semantic_recall() {
180            Some(semantic) => match semantic
181                .query_verified(
182                    self.repository().as_ref(),
183                    self.namespace(),
184                    text,
185                    cancellation.clone(),
186                )
187                .await
188            {
189                Ok(candidates) => candidates,
190                Err(error) => {
191                    tracing::warn!(
192                        reason = error.redacted_message(),
193                        "Semantic durable-memory recall degraded to lexical recall"
194                    );
195                    Vec::new()
196                }
197            },
198            None => Vec::new(),
199        };
200        let mut candidates = fuse_lexical_semantic(lexical, semantic);
201        if policy.max_related_lookups() == 0 {
202            candidates.truncate(policy.max_results());
203            return Ok(candidates);
204        }
205        let mut known_ids = candidates
206            .iter()
207            .map(|candidate| candidate.node.id.clone())
208            .collect::<HashSet<_>>();
209        let mut looked_up = HashSet::new();
210        let mut lookup_count = 0;
211        let recall_seeds = candidates.clone();
212        'seeds: for seed in &recall_seeds {
213            for relation in &seed.node.relations {
214                if cancellation.is_cancelled() {
215                    break 'seeds;
216                }
217                if relation.kind != MemoryRelationKind::RelatedTo
218                    || known_ids.contains(&relation.target_id)
219                    || !looked_up.insert(relation.target_id.clone())
220                {
221                    continue;
222                }
223                if lookup_count >= policy.max_related_lookups() {
224                    break 'seeds;
225                }
226                lookup_count += 1;
227                let Some(node) = self
228                    .repository()
229                    .get(self.namespace(), &relation.target_id)
230                    .await?
231                else {
232                    continue;
233                };
234                if node.status != MemoryStatus::Active {
235                    continue;
236                }
237                known_ids.insert(node.id.clone());
238                candidates.push(RecallCandidate {
239                    node,
240                    score: (seed.score * RELATED_SCORE_FACTOR).clamp(0.0, 1.0),
241                    channel: DurableMemoryRecallChannel::Related,
242                    related_from: Some(seed.node.id.clone()),
243                });
244            }
245        }
246        candidates.sort_by(|left, right| {
247            right
248                .score
249                .total_cmp(&left.score)
250                .then_with(|| channel_rank(left.channel).cmp(&channel_rank(right.channel)))
251                .then_with(|| right.node.updated_at.cmp(&left.node.updated_at))
252                .then_with(|| left.node.id.cmp(&right.node.id))
253        });
254        candidates.truncate(policy.max_results());
255        Ok(candidates)
256    }
257
258    pub(crate) async fn admit_selected_context(
259        &self,
260        assembly: &mut ContextAssembly,
261        identities: &[DurableMemoryRecallIdentity],
262        context_id: Option<&str>,
263        occurred_at: Option<DateTime<Utc>>,
264    ) -> usize {
265        if identities.is_empty() {
266            return 0;
267        }
268        let mut admitted = HashSet::new();
269        if let (Some(context_id), Some(occurred_at)) = (context_id, occurred_at) {
270            for item in &assembly.items {
271                let Some(identity) = identities.iter().find(|identity| identity.matches(item))
272                else {
273                    continue;
274                };
275                let event_id = admission_id(context_id, &identity.node_id, identity.node_revision);
276                let event = MemoryAccessEvent::new(
277                    event_id,
278                    self.namespace().clone(),
279                    &identity.node_id,
280                    identity.node_revision,
281                    occurred_at,
282                )
283                .with_context_id(context_id);
284                match self.repository().record_admission(event).await {
285                    Ok(()) => {
286                        admitted.insert(identity.item_id.clone());
287                    }
288                    Err(error) => {
289                        tracing::warn!(
290                            %error,
291                            memory_id = %identity.node_id,
292                            memory_revision = identity.node_revision,
293                            "Dropping V2 memory that could not be admitted"
294                        );
295                    }
296                }
297            }
298        } else if context_id.is_none() {
299            tracing::warn!("Dropping V2 memory context because invocation identity is unavailable");
300        } else {
301            tracing::warn!("Dropping V2 memory context because host time is invalid");
302        }
303
304        assembly.items.retain(|item| {
305            let recalled = identities.iter().any(|identity| identity.matches(item));
306            !recalled || admitted.contains(&item.id)
307        });
308        assembly.total_tokens = assembly
309            .items
310            .iter()
311            .map(|item| {
312                if item.token_count > 0 {
313                    item.token_count
314                } else {
315                    item.content.split_whitespace().count().max(1)
316                }
317            })
318            .sum();
319        admitted.len()
320    }
321}
322
323pub(crate) fn durable_memory_context_id(
324    session_id: &str,
325    run_id: &str,
326    invocation_incarnation: &str,
327    context_sequence: u64,
328) -> String {
329    let mut hasher = Sha256::new();
330    hasher.update(DURABLE_MEMORY_CONTEXT_ID_PROFILE_V2.as_bytes());
331    hasher.update(b"\0session\0");
332    hasher.update(Sha256::digest(session_id.as_bytes()));
333    hasher.update(b"\0run\0");
334    hasher.update(Sha256::digest(run_id.as_bytes()));
335    hasher.update(b"\0invocation\0");
336    hasher.update(Sha256::digest(invocation_incarnation.as_bytes()));
337    hasher.update(b"\0sequence\0");
338    hasher.update(context_sequence.to_le_bytes());
339    format!("a3s-code-context-v2-{:x}", hasher.finalize())
340}
341
342fn kind_label(kind: DurableMemoryKind) -> &'static str {
343    match kind {
344        DurableMemoryKind::Episodic => "episodic",
345        DurableMemoryKind::Semantic => "semantic",
346        DurableMemoryKind::Procedural => "procedural",
347    }
348}
349
350fn channel_label(channel: DurableMemoryRecallChannel) -> &'static str {
351    match channel {
352        DurableMemoryRecallChannel::Lexical => "lexical",
353        DurableMemoryRecallChannel::Semantic => "semantic",
354        DurableMemoryRecallChannel::Hybrid => "hybrid",
355        DurableMemoryRecallChannel::Related => "related",
356    }
357}
358
359pub(super) fn channel_rank(channel: DurableMemoryRecallChannel) -> u8 {
360    match channel {
361        DurableMemoryRecallChannel::Hybrid => 0,
362        DurableMemoryRecallChannel::Lexical => 1,
363        DurableMemoryRecallChannel::Semantic => 2,
364        DurableMemoryRecallChannel::Related => 3,
365    }
366}
367
368fn digest(content: &str) -> String {
369    format!("sha256:{:x}", Sha256::digest(content.as_bytes()))
370}
371
372fn admission_id(context_id: &str, node_id: &str, node_revision: u64) -> String {
373    let mut hasher = Sha256::new();
374    hasher.update(b"a3s.code.memory.admission.v1\0");
375    hasher.update(context_id.as_bytes());
376    hasher.update(b"\0");
377    hasher.update(node_id.as_bytes());
378    hasher.update(b"\0");
379    hasher.update(node_revision.to_le_bytes());
380    format!("a3s-code-admission-{:x}", hasher.finalize())
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386
387    #[test]
388    fn context_identity_is_deterministic_for_an_exact_tuple() {
389        let identity = durable_memory_context_id("session-a", "run-1", "invocation-a", 1);
390        assert_eq!(
391            identity,
392            durable_memory_context_id("session-a", "run-1", "invocation-a", 1)
393        );
394        assert_eq!(
395            identity,
396            "a3s-code-context-v2-92a5ba699eac46c932d2b4bc8e02a15603380336c7f2719ace2204b5e1301f5c"
397        );
398    }
399
400    #[test]
401    fn context_identity_separates_sessions_with_colliding_local_run_ids() {
402        assert_ne!(
403            durable_memory_context_id("session-a", "run-local-1", "invocation-a", 1),
404            durable_memory_context_id("session-b", "run-local-1", "invocation-a", 1)
405        );
406    }
407
408    #[test]
409    fn context_identity_separates_runs_and_is_repository_safe() {
410        let first = durable_memory_context_id("session-a", "run-1", "invocation-a", 1);
411        let second = durable_memory_context_id("session-a", "run-2", "invocation-a", 1);
412        assert_ne!(first, second);
413        assert!(first.starts_with("a3s-code-context-v2-"));
414        assert!(first.len() <= a3s_memory::repository::MAX_IDENTIFIER_BYTES);
415    }
416
417    #[test]
418    fn context_identity_separates_multiple_contexts_in_one_run() {
419        assert_ne!(
420            durable_memory_context_id("session-a", "run-1", "invocation-a", 1),
421            durable_memory_context_id("session-a", "run-1", "invocation-a", 2)
422        );
423    }
424
425    #[test]
426    fn context_identity_separates_reconstructed_invocations_with_reused_run_ids() {
427        assert_ne!(
428            durable_memory_context_id("session-a", "run-reused", "invocation-a", 1),
429            durable_memory_context_id("session-a", "run-reused", "invocation-b", 1)
430        );
431    }
432}
a3s_code_core/durable_memory/context.rs

a3s_code_core/durable_memory/
context.rs
