a3s-code-core 8.2.0

use super::{
    fusion::fuse_lexical_semantic, DurableMemoryMode, DurableMemoryRecallChannel,
    DurableMemoryRecallHit, DurableMemoryRecallPreview, DurableMemorySession,
};
use crate::context::{ContextAssembly, ContextItem, ContextResult, ContextType};
use a3s_memory::repository::{
    DurableMemoryKind, MemoryAccessEvent, MemoryNode, MemoryQuery, MemoryRelationKind,
    MemoryRepositoryError, MemoryStatus,
};
use chrono::{DateTime, Utc};
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use sha2::{Digest, Sha256};
use std::collections::HashSet;
use tokio_util::sync::CancellationToken;

/// Legacy schema-3 profile. It separates sessions and process-local run IDs but
/// cannot distinguish a run ID reused after session retention and restart.
pub const DURABLE_MEMORY_CONTEXT_ID_PROFILE_V1: &str =
    "a3s.code.memory.context.session-run-sequence-sha256.v1";

/// Current profile. The invocation incarnation is generated by Code rather
/// than the host environment because host IDs guarantee uniqueness only inside
/// one process and may legitimately restart from the same sequence.
pub const DURABLE_MEMORY_CONTEXT_ID_PROFILE_V2: &str =
    "a3s.code.memory.context.session-run-invocation-sequence-sha256.v2";

const PROVIDER: &str = "durable_memory_v2";
const RELATED_SCORE_FACTOR: f32 = 0.75;

#[derive(Clone)]
pub(super) struct RecallCandidate {
    pub(super) node: MemoryNode,
    pub(super) score: f32,
    pub(super) channel: DurableMemoryRecallChannel,
    pub(super) related_from: Option<String>,
}

impl RecallCandidate {
    fn into_preview_hit(self) -> DurableMemoryRecallHit {
        DurableMemoryRecallHit {
            node_id: self.node.id,
            node_revision: self.node.revision,
            kind: self.node.kind,
            content: self.node.content,
            score: self.score,
            channel: self.channel,
            related_from: self.related_from,
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct DurableMemoryRecallIdentity {
    pub(crate) item_id: String,
    source: String,
    pub(crate) node_id: String,
    pub(crate) node_revision: u64,
    content_digest: String,
}

impl DurableMemoryRecallIdentity {
    fn matches(&self, item: &ContextItem) -> bool {
        item.id == self.item_id
            && item.source.as_deref() == Some(self.source.as_str())
            && digest(&item.content) == self.content_digest
    }
}

pub(crate) struct DurableMemoryContextBatch {
    pub(crate) result: ContextResult,
    pub(crate) identities: Vec<DurableMemoryRecallIdentity>,
}

impl DurableMemorySession {
    /// Run the same pure active-only retrieval used before context assembly.
    /// This diagnostic does not record admission or use and does not inject
    /// any returned content into a model prompt.
    pub async fn preview_recall(
        &self,
        text: &str,
    ) -> Result<DurableMemoryRecallPreview, MemoryRepositoryError> {
        Ok(DurableMemoryRecallPreview {
            hits: self
                .query_recall_candidates(text, CancellationToken::new())
                .await?
                .into_iter()
                .map(RecallCandidate::into_preview_hit)
                .collect(),
        })
    }

    #[cfg(test)]
    pub(crate) async fn query_active_context(
        &self,
        text: &str,
    ) -> Result<DurableMemoryContextBatch, MemoryRepositoryError> {
        self.query_active_context_with_cancellation(text, CancellationToken::new())
            .await
    }

    pub(crate) async fn query_active_context_with_cancellation(
        &self,
        text: &str,
        cancellation: CancellationToken,
    ) -> Result<DurableMemoryContextBatch, MemoryRepositoryError> {
        let hits = self.query_recall_candidates(text, cancellation).await?;
        let mut result = ContextResult::new(PROVIDER);
        let mut identities = Vec::new();
        for hit in hits {
            let node = hit.node;
            let encoded_id = utf8_percent_encode(&node.id, NON_ALPHANUMERIC);
            let source = format!("a3s-memory://{encoded_id}?revision={}", node.revision);
            let item_id = format!("a3s-memory-v2:{}:r{}", node.id, node.revision);
            let content_digest = digest(&node.content);
            let token_count = (node.content.len() / 4).max(1);
            let item = ContextItem::new(&item_id, ContextType::Memory, &node.content)
                .with_relevance(hit.score)
                .with_token_count(token_count)
                .with_source(&source)
                .with_metadata("memory_node_id", serde_json::json!(node.id))
                .with_metadata("memory_node_revision", serde_json::json!(node.revision))
                .with_metadata("memory_kind", serde_json::json!(kind_label(node.kind)))
                .with_metadata("evidence_count", serde_json::json!(node.evidence.len()))
                .with_metadata(
                    "retrieval_channel",
                    serde_json::json!(channel_label(hit.channel)),
                )
                .with_provenance(PROVIDER)
                .with_priority(0.4)
                .with_trust(0.8)
                .with_freshness(0.6);
            let item = match hit.related_from {
                Some(source_id) => item.with_metadata("related_from", serde_json::json!(source_id)),
                None => item,
            };
            identities.push(DurableMemoryRecallIdentity {
                item_id,
                source,
                node_id: node.id,
                node_revision: node.revision,
                content_digest,
            });
            result.add_item(item);
        }
        Ok(DurableMemoryContextBatch { result, identities })
    }

    async fn query_recall_candidates(
        &self,
        text: &str,
        cancellation: CancellationToken,
    ) -> Result<Vec<RecallCandidate>, MemoryRepositoryError> {
        let Some(policy) = self.recall_policy() else {
            return Ok(Vec::new());
        };
        if self.mode() != DurableMemoryMode::ActiveRecall
            || !text.chars().any(char::is_alphanumeric)
        {
            return Ok(Vec::new());
        }

        let query = MemoryQuery::new(self.namespace().clone())
            .with_text(text)
            .with_limit(policy.max_results());
        let lexical = self
            .repository()
            .query(query)
            .await?
            .hits
            .into_iter()
            .filter(|hit| hit.score.total >= policy.min_lexical_score())
            .map(|hit| RecallCandidate {
                node: hit.node,
                score: hit.score.total,
                channel: DurableMemoryRecallChannel::Lexical,
                related_from: None,
            })
            .collect::<Vec<_>>();
        let semantic = match self.semantic_recall() {
            Some(semantic) => match semantic
                .query_verified(
                    self.repository().as_ref(),
                    self.namespace(),
                    text,
                    cancellation.clone(),
                )
                .await
            {
                Ok(candidates) => candidates,
                Err(error) => {
                    tracing::warn!(
                        reason = error.redacted_message(),
                        "Semantic durable-memory recall degraded to lexical recall"
                    );
                    Vec::new()
                }
            },
            None => Vec::new(),
        };
        let mut candidates = fuse_lexical_semantic(lexical, semantic);
        if policy.max_related_lookups() == 0 {
            candidates.truncate(policy.max_results());
            return Ok(candidates);
        }
        let mut known_ids = candidates
            .iter()
            .map(|candidate| candidate.node.id.clone())
            .collect::<HashSet<_>>();
        let mut looked_up = HashSet::new();
        let mut lookup_count = 0;
        let recall_seeds = candidates.clone();
        'seeds: for seed in &recall_seeds {
            for relation in &seed.node.relations {
                if cancellation.is_cancelled() {
                    break 'seeds;
                }
                if relation.kind != MemoryRelationKind::RelatedTo
                    || known_ids.contains(&relation.target_id)
                    || !looked_up.insert(relation.target_id.clone())
                {
                    continue;
                }
                if lookup_count >= policy.max_related_lookups() {
                    break 'seeds;
                }
                lookup_count += 1;
                let Some(node) = self
                    .repository()
                    .get(self.namespace(), &relation.target_id)
                    .await?
                else {
                    continue;
                };
                if node.status != MemoryStatus::Active {
                    continue;
                }
                known_ids.insert(node.id.clone());
                candidates.push(RecallCandidate {
                    node,
                    score: (seed.score * RELATED_SCORE_FACTOR).clamp(0.0, 1.0),
                    channel: DurableMemoryRecallChannel::Related,
                    related_from: Some(seed.node.id.clone()),
                });
            }
        }
        candidates.sort_by(|left, right| {
            right
                .score
                .total_cmp(&left.score)
                .then_with(|| channel_rank(left.channel).cmp(&channel_rank(right.channel)))
                .then_with(|| right.node.updated_at.cmp(&left.node.updated_at))
                .then_with(|| left.node.id.cmp(&right.node.id))
        });
        candidates.truncate(policy.max_results());
        Ok(candidates)
    }

    pub(crate) async fn admit_selected_context(
        &self,
        assembly: &mut ContextAssembly,
        identities: &[DurableMemoryRecallIdentity],
        context_id: Option<&str>,
        occurred_at: Option<DateTime<Utc>>,
    ) -> usize {
        if identities.is_empty() {
            return 0;
        }
        let mut admitted = HashSet::new();
        if let (Some(context_id), Some(occurred_at)) = (context_id, occurred_at) {
            for item in &assembly.items {
                let Some(identity) = identities.iter().find(|identity| identity.matches(item))
                else {
                    continue;
                };
                let event_id = admission_id(context_id, &identity.node_id, identity.node_revision);
                let event = MemoryAccessEvent::new(
                    event_id,
                    self.namespace().clone(),
                    &identity.node_id,
                    identity.node_revision,
                    occurred_at,
                )
                .with_context_id(context_id);
                match self.repository().record_admission(event).await {
                    Ok(()) => {
                        admitted.insert(identity.item_id.clone());
                    }
                    Err(error) => {
                        tracing::warn!(
                            %error,
                            memory_id = %identity.node_id,
                            memory_revision = identity.node_revision,
                            "Dropping V2 memory that could not be admitted"
                        );
                    }
                }
            }
        } else if context_id.is_none() {
            tracing::warn!("Dropping V2 memory context because invocation identity is unavailable");
        } else {
            tracing::warn!("Dropping V2 memory context because host time is invalid");
        }

        assembly.items.retain(|item| {
            let recalled = identities.iter().any(|identity| identity.matches(item));
            !recalled || admitted.contains(&item.id)
        });
        assembly.total_tokens = assembly
            .items
            .iter()
            .map(|item| {
                if item.token_count > 0 {
                    item.token_count
                } else {
                    item.content.split_whitespace().count().max(1)
                }
            })
            .sum();
        admitted.len()
    }
}

pub(crate) fn durable_memory_context_id(
    session_id: &str,
    run_id: &str,
    invocation_incarnation: &str,
    context_sequence: u64,
) -> String {
    let mut hasher = Sha256::new();
    hasher.update(DURABLE_MEMORY_CONTEXT_ID_PROFILE_V2.as_bytes());
    hasher.update(b"\0session\0");
    hasher.update(Sha256::digest(session_id.as_bytes()));
    hasher.update(b"\0run\0");
    hasher.update(Sha256::digest(run_id.as_bytes()));
    hasher.update(b"\0invocation\0");
    hasher.update(Sha256::digest(invocation_incarnation.as_bytes()));
    hasher.update(b"\0sequence\0");
    hasher.update(context_sequence.to_le_bytes());
    format!("a3s-code-context-v2-{:x}", hasher.finalize())
}

fn kind_label(kind: DurableMemoryKind) -> &'static str {
    match kind {
        DurableMemoryKind::Episodic => "episodic",
        DurableMemoryKind::Semantic => "semantic",
        DurableMemoryKind::Procedural => "procedural",
    }
}

fn channel_label(channel: DurableMemoryRecallChannel) -> &'static str {
    match channel {
        DurableMemoryRecallChannel::Lexical => "lexical",
        DurableMemoryRecallChannel::Semantic => "semantic",
        DurableMemoryRecallChannel::Hybrid => "hybrid",
        DurableMemoryRecallChannel::Related => "related",
    }
}

pub(super) fn channel_rank(channel: DurableMemoryRecallChannel) -> u8 {
    match channel {
        DurableMemoryRecallChannel::Hybrid => 0,
        DurableMemoryRecallChannel::Lexical => 1,
        DurableMemoryRecallChannel::Semantic => 2,
        DurableMemoryRecallChannel::Related => 3,
    }
}

fn digest(content: &str) -> String {
    format!("sha256:{:x}", Sha256::digest(content.as_bytes()))
}

fn admission_id(context_id: &str, node_id: &str, node_revision: u64) -> String {
    let mut hasher = Sha256::new();
    hasher.update(b"a3s.code.memory.admission.v1\0");
    hasher.update(context_id.as_bytes());
    hasher.update(b"\0");
    hasher.update(node_id.as_bytes());
    hasher.update(b"\0");
    hasher.update(node_revision.to_le_bytes());
    format!("a3s-code-admission-{:x}", hasher.finalize())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn context_identity_is_deterministic_for_an_exact_tuple() {
        let identity = durable_memory_context_id("session-a", "run-1", "invocation-a", 1);
        assert_eq!(
            identity,
            durable_memory_context_id("session-a", "run-1", "invocation-a", 1)
        );
        assert_eq!(
            identity,
            "a3s-code-context-v2-92a5ba699eac46c932d2b4bc8e02a15603380336c7f2719ace2204b5e1301f5c"
        );
    }

    #[test]
    fn context_identity_separates_sessions_with_colliding_local_run_ids() {
        assert_ne!(
            durable_memory_context_id("session-a", "run-local-1", "invocation-a", 1),
            durable_memory_context_id("session-b", "run-local-1", "invocation-a", 1)
        );
    }

    #[test]
    fn context_identity_separates_runs_and_is_repository_safe() {
        let first = durable_memory_context_id("session-a", "run-1", "invocation-a", 1);
        let second = durable_memory_context_id("session-a", "run-2", "invocation-a", 1);
        assert_ne!(first, second);
        assert!(first.starts_with("a3s-code-context-v2-"));
        assert!(first.len() <= a3s_memory::repository::MAX_IDENTIFIER_BYTES);
    }

    #[test]
    fn context_identity_separates_multiple_contexts_in_one_run() {
        assert_ne!(
            durable_memory_context_id("session-a", "run-1", "invocation-a", 1),
            durable_memory_context_id("session-a", "run-1", "invocation-a", 2)
        );
    }

    #[test]
    fn context_identity_separates_reconstructed_invocations_with_reused_run_ids() {
        assert_ne!(
            durable_memory_context_id("session-a", "run-reused", "invocation-a", 1),
            durable_memory_context_id("session-a", "run-reused", "invocation-b", 1)
        );
    }
}