Skip to main content

remem/memory/preference/
render.rs

1use anyhow::Result;
2use rusqlite::{params, Connection, OptionalExtension};
3
4use crate::memory::poisoning::{scan_instruction_pattern, InstructionPatternMatch};
5use crate::memory::Memory;
6
7use super::{
8    consolidation::{classify_preference_texts, PreferenceConsolidationKind},
9    query_global_preferences, query_project_preferences,
10};
11
12#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13pub struct PreferenceRenderSummary {
14    pub rendered: usize,
15    pub project_rendered: usize,
16    pub global_rendered: usize,
17}
18
19#[derive(Debug, Clone, Default)]
20pub(crate) struct PreferenceRenderDetails {
21    pub summary: PreferenceRenderSummary,
22    pub rendered_ids: Vec<i64>,
23    /// Canonical rows from the same query snapshot that produced the text.
24    /// Bundle compilation must not query these identities a second time:
25    /// another process could supersede or suppress a row between rendering
26    /// and compilation.
27    pub rendered_memories: Vec<Memory>,
28    /// Unsafe rows rejected before rendering. Bundle consumers retain only a
29    /// redacted identity for their selection/drop audit.
30    pub poisoning_drops: Vec<Memory>,
31    /// Safe canonical rows intentionally omitted by preference selection.
32    /// Context Bundle consumers use these identities and stable reasons to
33    /// account for every fetched row without re-running the selector.
34    pub selection_drops: Vec<PreferenceSelectionDrop>,
35    /// Character offsets, relative to the start of this preference render,
36    /// immediately after each rendered list item.
37    pub item_end_chars: Vec<usize>,
38}
39
40#[derive(Debug, Clone)]
41pub(crate) struct PreferenceSelectionDrop {
42    pub(crate) memory: Memory,
43    pub(crate) reason: &'static str,
44}
45
46pub(crate) const DROP_CLAUDE_MD_DEDUP: &str = "claude_md_dedup";
47pub(crate) const DROP_PREFERENCE_SIMILARITY_DEDUP: &str = "preference_similarity_dedup";
48pub(crate) const DROP_PREFERENCE_CHAR_LIMIT: &str = "preference_char_limit";
49pub(crate) const DROP_PROJECT_TOPIC_OVERRIDE: &str = "project_topic_override";
50
51impl PreferenceRenderDetails {
52    pub(crate) fn absolute_item_end_chars(&self, output_start: usize) -> Vec<usize> {
53        self.item_end_chars
54            .iter()
55            .map(|end| output_start + end)
56            .collect()
57    }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61enum PreferenceSource {
62    Project,
63    Global,
64}
65
66pub fn dedup_with_claude_md(prefs: &[Memory], cwd: &str) -> Vec<usize> {
67    let claude_md_path = std::path::Path::new(cwd).join("CLAUDE.md");
68    let claude_md_content = std::fs::read_to_string(&claude_md_path).unwrap_or_default();
69
70    if claude_md_content.is_empty() {
71        return (0..prefs.len()).collect();
72    }
73
74    let claude_lower = claude_md_content.to_lowercase();
75    (0..prefs.len())
76        .filter(|&i| {
77            let title_lower = prefs[i].title.to_lowercase();
78            let search_term = title_lower
79                .strip_prefix("preference: ")
80                .unwrap_or(&title_lower);
81            !claude_lower.contains(search_term)
82        })
83        .collect()
84}
85
86pub fn render_preferences(
87    output: &mut String,
88    conn: &Connection,
89    project: &str,
90    cwd: &str,
91) -> Result<()> {
92    render_preferences_with_limits(output, conn, project, cwd, 20, 0, 1500).map(|_| ())
93}
94
95pub fn render_preferences_with_limits(
96    output: &mut String,
97    conn: &Connection,
98    project: &str,
99    cwd: &str,
100    project_limit: usize,
101    global_limit: usize,
102    char_limit: usize,
103) -> Result<usize> {
104    render_preferences_with_context_details(
105        output,
106        conn,
107        project,
108        cwd,
109        project_limit,
110        global_limit,
111        char_limit,
112    )
113    .map(|details| details.summary.rendered)
114}
115
116pub fn render_preferences_with_limits_detailed(
117    output: &mut String,
118    conn: &Connection,
119    project: &str,
120    cwd: &str,
121    project_limit: usize,
122    global_limit: usize,
123    char_limit: usize,
124) -> Result<PreferenceRenderSummary> {
125    render_preferences_with_context_details(
126        output,
127        conn,
128        project,
129        cwd,
130        project_limit,
131        global_limit,
132        char_limit,
133    )
134    .map(|details| details.summary)
135}
136
137pub(crate) fn render_preferences_with_context_details(
138    output: &mut String,
139    conn: &Connection,
140    project: &str,
141    cwd: &str,
142    project_limit: usize,
143    global_limit: usize,
144    char_limit: usize,
145) -> Result<PreferenceRenderDetails> {
146    let project_prefs = query_project_preferences(conn, project, project_limit)?;
147    let global_prefs = query_global_preferences(conn, global_limit)?;
148
149    let mut all_prefs: Vec<(Memory, PreferenceSource)> = project_prefs
150        .into_iter()
151        .map(|memory| (memory, PreferenceSource::Project))
152        .collect();
153    let project_topics: std::collections::HashSet<String> = all_prefs
154        .iter()
155        .filter_map(|(memory, _)| memory.topic_key.clone())
156        .collect();
157    let mut selection_drops = Vec::new();
158    for global_pref in global_prefs {
159        if let Some(ref topic_key) = global_pref.topic_key {
160            if project_topics.contains(topic_key) {
161                selection_drops.push(PreferenceSelectionDrop {
162                    memory: global_pref,
163                    reason: DROP_PROJECT_TOPIC_OVERRIDE,
164                });
165                continue;
166            }
167        }
168        all_prefs.push((global_pref, PreferenceSource::Global));
169    }
170    all_prefs = filter_non_current_preferences(conn, all_prefs, &mut selection_drops)?;
171    let (filtered_prefs, poisoning_drops) =
172        filter_unacknowledged_poisoned_preferences(conn, all_prefs)?;
173    all_prefs = filtered_prefs;
174
175    if all_prefs.is_empty() {
176        return Ok(PreferenceRenderDetails {
177            poisoning_drops,
178            selection_drops,
179            ..PreferenceRenderDetails::default()
180        });
181    }
182
183    let memories = all_prefs
184        .iter()
185        .map(|(memory, _)| memory.clone())
186        .collect::<Vec<_>>();
187    let keep_indices = dedup_with_claude_md(&memories, cwd);
188    record_selection_drops(
189        &all_prefs,
190        &(0..all_prefs.len()).collect::<Vec<_>>(),
191        &keep_indices,
192        DROP_CLAUDE_MD_DEDUP,
193        &mut selection_drops,
194    );
195    if keep_indices.is_empty() {
196        return Ok(PreferenceRenderDetails {
197            poisoning_drops,
198            selection_drops,
199            ..PreferenceRenderDetails::default()
200        });
201    }
202    let similarity_keep_indices = dedup_with_preference_similarity(&memories, &keep_indices);
203    record_selection_drops(
204        &all_prefs,
205        &keep_indices,
206        &similarity_keep_indices,
207        DROP_PREFERENCE_SIMILARITY_DEDUP,
208        &mut selection_drops,
209    );
210    if similarity_keep_indices.is_empty() {
211        return Ok(PreferenceRenderDetails {
212            poisoning_drops,
213            selection_drops,
214            ..PreferenceRenderDetails::default()
215        });
216    }
217
218    output.push_str("## Your Preferences (always apply these)\n");
219    let mut total_chars = 0usize;
220    let mut summary = PreferenceRenderSummary::default();
221    let mut rendered_ids = Vec::new();
222    let mut rendered_memories = Vec::new();
223    let mut item_end_chars = Vec::new();
224    let header_chars = "## Your Preferences (always apply these)\n".chars().count();
225    for (position, &idx) in similarity_keep_indices.iter().enumerate() {
226        let (pref, source) = &all_prefs[idx];
227        let text = normalize_rendered_preference_text(&pref.text);
228        let preview: String = text.chars().take(120).collect();
229        let line = if preview.chars().count() < text.chars().count() {
230            format!("- {}...\n", preview)
231        } else {
232            format!("- {text}\n")
233        };
234        let line_chars = line.chars().count();
235        if total_chars + line_chars > char_limit && total_chars > 0 {
236            selection_drops.extend(similarity_keep_indices[position..].iter().map(|&drop_idx| {
237                PreferenceSelectionDrop {
238                    memory: all_prefs[drop_idx].0.clone(),
239                    reason: DROP_PREFERENCE_CHAR_LIMIT,
240                }
241            }));
242            break;
243        }
244        output.push_str(&line);
245        total_chars += line_chars;
246        summary.rendered += 1;
247        rendered_ids.push(pref.id);
248        rendered_memories.push(pref.clone());
249        item_end_chars.push(header_chars + total_chars);
250        match source {
251            PreferenceSource::Project => summary.project_rendered += 1,
252            PreferenceSource::Global => summary.global_rendered += 1,
253        }
254    }
255    output.push('\n');
256
257    Ok(PreferenceRenderDetails {
258        summary,
259        rendered_ids,
260        rendered_memories,
261        poisoning_drops,
262        selection_drops,
263        item_end_chars,
264    })
265}
266
267fn filter_non_current_preferences(
268    conn: &Connection,
269    preferences: Vec<(Memory, PreferenceSource)>,
270    drops: &mut Vec<PreferenceSelectionDrop>,
271) -> Result<Vec<(Memory, PreferenceSource)>> {
272    let as_of_epoch = chrono::Utc::now().timestamp();
273    let ids = preferences
274        .iter()
275        .map(|(memory, _)| memory.id)
276        .collect::<Vec<_>>();
277    let visibility = crate::truth::admit_many_for_current_context(conn, &ids, as_of_epoch)?;
278    let mut current = Vec::with_capacity(preferences.len());
279    for (memory, source) in preferences {
280        let visibility = visibility.get(&memory.id).copied().ok_or_else(|| {
281            anyhow::anyhow!("missing visibility for preference memory {}", memory.id)
282        })?;
283        if visibility.current_context_eligible {
284            current.push((memory, source));
285        } else {
286            drops.push(PreferenceSelectionDrop {
287                memory,
288                reason: visibility.reason.as_str(),
289            });
290        }
291    }
292    Ok(current)
293}
294
295fn record_selection_drops(
296    prefs: &[(Memory, PreferenceSource)],
297    before: &[usize],
298    after: &[usize],
299    reason: &'static str,
300    drops: &mut Vec<PreferenceSelectionDrop>,
301) {
302    drops.extend(
303        before
304            .iter()
305            .filter(|idx| !after.contains(idx))
306            .map(|&idx| PreferenceSelectionDrop {
307                memory: prefs[idx].0.clone(),
308                reason,
309            }),
310    );
311}
312
313#[derive(Debug, Default)]
314struct PreferencePoisoningState {
315    acknowledged_pattern_id: Option<String>,
316    acknowledged_pattern_version: Option<i64>,
317    source_trust_class: String,
318    source_project: Option<String>,
319}
320
321fn filter_unacknowledged_poisoned_preferences(
322    conn: &Connection,
323    prefs: Vec<(Memory, PreferenceSource)>,
324) -> Result<(Vec<(Memory, PreferenceSource)>, Vec<Memory>)> {
325    let mut kept = Vec::with_capacity(prefs.len());
326    let mut dropped = Vec::new();
327    for (memory, source) in prefs {
328        let Some(pattern_match) =
329            scan_instruction_pattern(&format!("{}\n{}", memory.title, memory.text))
330        else {
331            kept.push((memory, source));
332            continue;
333        };
334        let state = load_preference_poisoning_state(conn, memory.id)?;
335        if state.acknowledged_pattern_id.as_deref() == Some(pattern_match.pattern_id)
336            && state.acknowledged_pattern_version == Some(pattern_match.pattern_set_version)
337        {
338            kept.push((memory, source));
339            continue;
340        }
341        crate::log::error(
342            "context-poisoning",
343            &format!(
344                "dropping unacknowledged poisoned preference memory id={} pattern={}@v{}",
345                memory.id, pattern_match.pattern_id, pattern_match.pattern_set_version
346            ),
347        );
348        record_preference_injection_drop(conn, &memory, &state, pattern_match)?;
349        dropped.push(memory);
350    }
351    Ok((kept, dropped))
352}
353
354fn load_preference_poisoning_state(
355    conn: &Connection,
356    memory_id: i64,
357) -> Result<PreferencePoisoningState> {
358    Ok(conn
359        .query_row(
360            "SELECT acknowledged_pattern_id, acknowledged_pattern_version,
361                    source_trust_class, source_project
362             FROM memories WHERE id = ?1",
363            params![memory_id],
364            |row| {
365                Ok(PreferencePoisoningState {
366                    acknowledged_pattern_id: row.get(0)?,
367                    acknowledged_pattern_version: row.get(1)?,
368                    source_trust_class: row.get(2)?,
369                    source_project: row.get(3)?,
370                })
371            },
372        )
373        .optional()?
374        .unwrap_or_else(|| PreferencePoisoningState {
375            acknowledged_pattern_id: None,
376            acknowledged_pattern_version: None,
377            source_trust_class: "external_content".to_string(),
378            source_project: None,
379        }))
380}
381
382fn record_preference_injection_drop(
383    conn: &Connection,
384    memory: &Memory,
385    state: &PreferencePoisoningState,
386    pattern_match: InstructionPatternMatch,
387) -> Result<()> {
388    conn.execute(
389        "INSERT INTO memory_poisoning_injection_drops
390         (memory_id, pattern_id, pattern_version, source_trust_class, source_project,
391          title, created_at_epoch)
392         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
393        params![
394            memory.id,
395            pattern_match.pattern_id,
396            pattern_match.pattern_set_version,
397            state.source_trust_class.as_str(),
398            state.source_project.as_deref(),
399            memory.title.as_str(),
400            chrono::Utc::now().timestamp(),
401        ],
402    )?;
403    Ok(())
404}
405
406fn normalize_rendered_preference_text(text: &str) -> String {
407    text.trim()
408        .lines()
409        .filter(|line| !line.trim().is_empty())
410        .collect::<Vec<_>>()
411        .join(" ")
412}
413
414fn dedup_with_preference_similarity(prefs: &[Memory], indices: &[usize]) -> Vec<usize> {
415    let mut kept: Vec<usize> = Vec::new();
416    for &idx in indices {
417        let incoming = &prefs[idx];
418        let already_represented = kept.iter().any(|&kept_idx| {
419            let existing = &prefs[kept_idx];
420            classify_preference_texts(existing.id, &existing.text, &incoming.text).is_some_and(
421                |matched| {
422                    matches!(
423                        matched.kind,
424                        PreferenceConsolidationKind::SamePreference
425                            | PreferenceConsolidationKind::Refinement
426                    )
427                },
428            )
429        });
430        if !already_represented {
431            kept.push(idx);
432        }
433    }
434    kept
435}