1use anyhow::{anyhow, bail, Context, Result};
2use rusqlite::{params, Connection, OptionalExtension};
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5use std::collections::HashSet;
6
7use crate::memory::lifecycle::MemoryLifecycleOp;
8use crate::memory::operation::{insert_operation_log, MemoryOperationInput, MemoryOperationPlan};
9
10use super::audit::load_memory_audit_rows;
11use super::mutate::{insert_scope_cleanup_event, load_target, ObjectMutation};
12use super::preference_cluster::preference_clusters;
13use super::ObjectRef;
14
15pub const CLEANUP_PLANNER_VERSION: &str = "memory-cleanup-v1";
16
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
18pub struct MemoryCleanupPlan {
19 pub project: String,
20 pub created_at_epoch: i64,
21 pub planner_version: String,
22 pub groups: Vec<MemoryCleanupGroup>,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
26pub struct MemoryCleanupGroup {
27 pub cluster_key: String,
28 pub owner_scope: Option<String>,
29 pub owner_key: Option<String>,
30 pub memory_type: String,
31 pub state_key: Option<String>,
32 pub current_id: i64,
33 pub stale_ids: Vec<i64>,
34 pub reason: String,
35 pub confidence: f64,
36 pub preview: Vec<String>,
37 pub merged_content: Option<String>,
38 pub row_snapshots: Vec<MemoryCleanupRowSnapshot>,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
42pub struct MemoryCleanupRowSnapshot {
43 pub id: i64,
44 pub project: String,
45 pub scope: Option<String>,
46 pub source_project: Option<String>,
47 pub target_project: Option<String>,
48 pub status: String,
49 pub content_sha256: String,
50 pub updated_at_epoch: i64,
51 pub owner_scope: Option<String>,
52 pub owner_key: Option<String>,
53 pub memory_type: String,
54 pub topic_key: Option<String>,
55 pub state_key_id: Option<i64>,
56 pub state_key: Option<String>,
57 pub current_memory_id: Option<i64>,
58}
59
60#[derive(Debug, Clone, Serialize)]
61pub struct MemoryCleanupApplyResult {
62 pub project: String,
63 pub planner_version: String,
64 pub groups_applied: usize,
65 pub current_ids: Vec<i64>,
66 pub stale_ids: Vec<i64>,
67 pub operation_ids: Vec<i64>,
68 pub edge_count: usize,
69 pub affected: Vec<ObjectMutation>,
70}
71
72pub fn build_preference_cleanup_plan(
73 conn: &Connection,
74 project: &str,
75) -> Result<MemoryCleanupPlan> {
76 let memories = load_memory_audit_rows(conn, project)?;
77 let clusters = preference_clusters(&memories, project);
78 let mut groups = Vec::with_capacity(clusters.len());
79
80 for cluster in clusters {
81 let current_ref = ObjectRef::parse(&cluster.canonical_ref)?;
82 let stale_ids = cluster
83 .refs
84 .iter()
85 .filter(|object_ref| *object_ref != &cluster.canonical_ref)
86 .map(|object_ref| ObjectRef::parse(object_ref).map(|parsed| parsed.id))
87 .collect::<Result<Vec<_>>>()?;
88 if stale_ids.is_empty() {
89 continue;
90 }
91 let mut ids = Vec::with_capacity(stale_ids.len() + 1);
92 ids.push(current_ref.id);
93 ids.extend(stale_ids.iter().copied());
94 let row_snapshots = load_row_snapshots(conn, &ids)?;
95 let current = snapshot_for(&row_snapshots, current_ref.id)?;
96 let preview = row_snapshots
97 .iter()
98 .take(4)
99 .map(|row| format!("memory:{} {}", row.id, row.status))
100 .collect();
101 groups.push(MemoryCleanupGroup {
102 cluster_key: cluster.cluster_key,
103 owner_scope: current.owner_scope.clone(),
104 owner_key: current.owner_key.clone(),
105 memory_type: current.memory_type.clone(),
106 state_key: current.state_key.clone(),
107 current_id: current_ref.id,
108 stale_ids,
109 reason: cluster.reason,
110 confidence: 1.0,
111 preview,
112 merged_content: cluster.merged_content,
113 row_snapshots,
114 });
115 }
116
117 Ok(MemoryCleanupPlan {
118 project: project.to_string(),
119 created_at_epoch: chrono::Utc::now().timestamp(),
120 planner_version: CLEANUP_PLANNER_VERSION.to_string(),
121 groups,
122 })
123}
124
125pub fn apply_memory_cleanup_plan(
126 conn: &Connection,
127 plan: &MemoryCleanupPlan,
128) -> Result<MemoryCleanupApplyResult> {
129 if plan.planner_version != CLEANUP_PLANNER_VERSION {
130 bail!(
131 "unsupported cleanup planner version: {}",
132 plan.planner_version
133 );
134 }
135
136 let tx = conn.unchecked_transaction()?;
137 validate_plan_shape(plan)?;
138 for group in &plan.groups {
139 validate_group(&tx, plan, group)?;
140 }
141
142 let now = chrono::Utc::now().timestamp();
143 let mut affected = Vec::new();
144 let mut current_ids = Vec::new();
145 let mut stale_ids = Vec::new();
146 let mut operation_ids = Vec::new();
147 let mut edge_count = 0usize;
148
149 for group in &plan.groups {
150 let current_ref = ObjectRef::memory(group.current_id);
151 let canonical = load_target(&tx, current_ref)?;
152 let current_snapshot = snapshot_for(&group.row_snapshots, group.current_id)?;
153 let merged = group.merged_content.as_deref();
154 let final_text = if let Some(merged) = merged {
155 merged.to_string()
156 } else {
157 tx.query_row(
158 "SELECT content FROM memories WHERE id = ?1",
159 [group.current_id],
160 |row| row.get::<_, String>(0),
161 )?
162 };
163 let affected_ids = std::iter::once(group.current_id)
164 .chain(group.stale_ids.iter().copied())
165 .collect::<Vec<_>>();
166 crate::memory::preference::compilation::enqueue_for_memory_ids(&tx, &affected_ids)?;
167 crate::memory::preference::reinforcement::reconcile_cleanup_preference(
168 &tx,
169 group.current_id,
170 &group.stale_ids,
171 &final_text,
172 now,
173 )?;
174 let updated = tx.execute(
175 "UPDATE memories
176 SET content = COALESCE(?1, content),
177 status = 'active',
178 updated_at_epoch = ?2
179 WHERE id = ?3",
180 params![merged, now, group.current_id],
181 )?;
182 if updated != 1 {
183 bail!(
184 "failed to update cleanup current memory {}",
185 group.current_id
186 );
187 }
188 current_ids.push(group.current_id);
189 affected.push(ObjectMutation {
190 object_ref: current_ref.to_string(),
191 title: canonical.title.clone(),
192 previous_status: canonical.status.clone(),
193 new_status: "active".to_string(),
194 previous_owner: canonical.owner.clone(),
195 new_owner: canonical.owner.clone(),
196 });
197 insert_scope_cleanup_event(
198 &tx,
199 "memory-cleanup",
200 &canonical,
201 "active",
202 &canonical.owner,
203 Some(group.reason.as_str()),
204 now,
205 )?;
206
207 if let Some(state_key_id) = current_snapshot.state_key_id {
208 tx.execute(
209 "UPDATE memory_state_keys
210 SET current_memory_id = ?1, updated_at_epoch = ?2
211 WHERE id = ?3",
212 params![group.current_id, now, state_key_id],
213 )?;
214 }
215
216 for stale_id in &group.stale_ids {
217 let stale_ref = ObjectRef::memory(*stale_id);
218 let target = load_target(&tx, stale_ref)?;
219 let updated = tx.execute(
220 "UPDATE memories SET status = 'stale', updated_at_epoch = ?1 WHERE id = ?2",
221 params![now, stale_id],
222 )?;
223 if updated != 1 {
224 bail!("failed to stale cleanup memory {stale_id}");
225 }
226 stale_ids.push(*stale_id);
227 affected.push(ObjectMutation {
228 object_ref: stale_ref.to_string(),
229 title: target.title.clone(),
230 previous_status: target.status.clone(),
231 new_status: "stale".to_string(),
232 previous_owner: target.owner.clone(),
233 new_owner: target.owner.clone(),
234 });
235 insert_scope_cleanup_event(
236 &tx,
237 "memory-cleanup",
238 &target,
239 "stale",
240 &target.owner,
241 Some("duplicate preference superseded by cleanup plan"),
242 now,
243 )?;
244 }
245
246 let operation_id = insert_cleanup_operation_log(&tx, plan, group)?;
247 operation_ids.push(operation_id);
248 edge_count += crate::memory::edge::insert_replacement_edges(
249 &tx,
250 crate::memory::edge::MemoryEdgeType::Duplicates,
251 &group.stale_ids,
252 group.current_id,
253 crate::memory::edge::MemoryEdgeWriteContext {
254 state_key_id: current_snapshot.state_key_id,
255 source_operation_id: Some(operation_id),
256 confidence: Some(group.confidence),
257 reason: Some(group.reason.as_str()),
258 ..Default::default()
259 },
260 )?;
261 }
262
263 tx.commit()?;
264 Ok(MemoryCleanupApplyResult {
265 project: plan.project.clone(),
266 planner_version: plan.planner_version.clone(),
267 groups_applied: plan.groups.len(),
268 current_ids,
269 stale_ids,
270 operation_ids,
271 edge_count,
272 affected,
273 })
274}
275
276fn validate_plan_shape(plan: &MemoryCleanupPlan) -> Result<()> {
277 let mut ids = HashSet::new();
278 for group in &plan.groups {
279 for id in std::iter::once(group.current_id).chain(group.stale_ids.iter().copied()) {
280 if !ids.insert(id) {
281 bail!("cleanup plan lists memory:{id} in more than one action");
282 }
283 }
284 }
285 Ok(())
286}
287
288fn validate_group(
289 conn: &Connection,
290 plan: &MemoryCleanupPlan,
291 group: &MemoryCleanupGroup,
292) -> Result<()> {
293 if group.stale_ids.contains(&group.current_id) {
294 bail!(
295 "cleanup group {} lists current id {} as stale",
296 group.cluster_key,
297 group.current_id
298 );
299 }
300 if group.memory_type != "preference" {
301 bail!(
302 "unsupported cleanup group memory type {}",
303 group.memory_type
304 );
305 }
306 let mut expected_ids = group.stale_ids.clone();
307 expected_ids.push(group.current_id);
308 expected_ids.sort_unstable();
309 expected_ids.dedup();
310 let mut snapshot_ids = group
311 .row_snapshots
312 .iter()
313 .map(|snapshot| snapshot.id)
314 .collect::<Vec<_>>();
315 snapshot_ids.sort_unstable();
316 snapshot_ids.dedup();
317 if snapshot_ids != expected_ids {
318 bail!(
319 "cleanup group {} row snapshots do not match current/stale ids",
320 group.cluster_key
321 );
322 }
323
324 let current_snapshot = snapshot_for(&group.row_snapshots, group.current_id)?;
325 if group.owner_scope != current_snapshot.owner_scope
326 || group.owner_key != current_snapshot.owner_key
327 {
328 bail!(
329 "cleanup group {} owner does not match current row owner",
330 group.cluster_key
331 );
332 }
333 if group.state_key != current_snapshot.state_key {
334 bail!(
335 "cleanup group {} state key does not match current row",
336 group.cluster_key
337 );
338 }
339 let current_owner = current_snapshot.owner_namespace(&plan.project);
340 let current_state_key_id = current_snapshot.state_key_id;
341 let current_state_key = current_snapshot.state_key.as_deref();
342 let topic_group = group.cluster_key.starts_with("topic:");
343
344 for snapshot in &group.row_snapshots {
345 let current = load_row_snapshot(conn, snapshot.id)?
346 .ok_or_else(|| anyhow!("cleanup plan row {} no longer exists", snapshot.id))?;
347 if ¤t != snapshot {
348 bail!(
349 "cleanup plan row {} changed since dry-run; refresh the plan before applying",
350 snapshot.id
351 );
352 }
353 if snapshot.status != "active" {
354 bail!("cleanup plan row {} is no longer active", snapshot.id);
355 }
356 if snapshot.memory_type != group.memory_type {
357 bail!(
358 "cleanup plan row {} type {} does not match group type {}",
359 snapshot.id,
360 snapshot.memory_type,
361 group.memory_type
362 );
363 }
364 if !snapshot.belongs_to_project(&plan.project) {
365 bail!(
366 "cleanup plan row {} does not belong to project {}",
367 snapshot.id,
368 plan.project
369 );
370 }
371 if snapshot.owner_namespace(&plan.project) != current_owner {
372 bail!(
373 "cleanup plan row {} owner does not match current row owner",
374 snapshot.id
375 );
376 }
377 match (current_state_key_id, current_state_key) {
378 (Some(state_key_id), _) if snapshot.state_key_id != Some(state_key_id) => {
379 bail!(
380 "cleanup plan row {} state key does not match current row",
381 snapshot.id
382 );
383 }
384 (None, Some(state_key)) if snapshot.state_key.as_deref() != Some(state_key) => {
385 bail!(
386 "cleanup plan row {} state key does not match current row",
387 snapshot.id
388 );
389 }
390 _ => {}
391 }
392 if topic_group && snapshot.topic_key != current_snapshot.topic_key {
393 bail!(
394 "cleanup plan row {} topic key does not match current row",
395 snapshot.id
396 );
397 }
398 }
399 Ok(())
400}
401
402fn insert_cleanup_operation_log(
403 conn: &Connection,
404 plan: &MemoryCleanupPlan,
405 group: &MemoryCleanupGroup,
406) -> Result<i64> {
407 let current = snapshot_for(&group.row_snapshots, group.current_id)?;
408 let mut operation_plan = MemoryOperationPlan::new(
409 MemoryLifecycleOp::Update,
410 group.state_key.clone(),
411 group.reason.clone(),
412 )
413 .with_target_memory_id(Some(group.current_id))
414 .with_superseded_ids(group.stale_ids.clone());
415 operation_plan.planner_version = CLEANUP_PLANNER_VERSION;
416 let input = MemoryOperationInput {
417 source: "memory_cleanup".to_string(),
418 actor: "memory_cleanup".to_string(),
419 source_project: plan.project.clone(),
420 owner_scope: group
421 .owner_scope
422 .clone()
423 .unwrap_or_else(|| "repo".to_string()),
424 owner_key: group
425 .owner_key
426 .clone()
427 .unwrap_or_else(|| plan.project.clone()),
428 memory_type: group.memory_type.clone(),
429 topic_key: current.topic_key.clone(),
430 state_key: group.state_key.clone(),
431 source_candidate_id: None,
432 confidence: Some(group.confidence),
433 };
434 insert_operation_log(conn, &input, &operation_plan, Some(group.current_id))
435}
436
437fn load_row_snapshots(conn: &Connection, ids: &[i64]) -> Result<Vec<MemoryCleanupRowSnapshot>> {
438 ids.iter()
439 .copied()
440 .map(|id| {
441 load_row_snapshot(conn, id)?
442 .ok_or_else(|| anyhow!("cleanup plan target memory:{id} not found"))
443 })
444 .collect()
445}
446
447fn load_row_snapshot(conn: &Connection, id: i64) -> Result<Option<MemoryCleanupRowSnapshot>> {
448 conn.query_row(
449 "SELECT m.id, m.status, m.content, m.updated_at_epoch, m.owner_scope,
450 m.owner_key, m.memory_type, m.topic_key, m.state_key_id, sk.state_key,
451 sk.current_memory_id, m.project, m.scope, m.source_project, m.target_project
452 FROM memories m
453 LEFT JOIN memory_state_keys sk ON sk.id = m.state_key_id
454 WHERE m.id = ?1",
455 params![id],
456 |row| {
457 let content: String = row.get(2)?;
458 Ok(MemoryCleanupRowSnapshot {
459 id: row.get(0)?,
460 status: row.get(1)?,
461 content_sha256: content_sha256(&content),
462 updated_at_epoch: row.get(3)?,
463 owner_scope: row.get(4)?,
464 owner_key: row.get(5)?,
465 memory_type: row.get(6)?,
466 topic_key: row.get(7)?,
467 state_key_id: row.get(8)?,
468 state_key: row.get(9)?,
469 current_memory_id: row.get(10)?,
470 project: row.get(11)?,
471 scope: row.get(12)?,
472 source_project: row.get(13)?,
473 target_project: row.get(14)?,
474 })
475 },
476 )
477 .optional()
478 .with_context(|| format!("load cleanup plan row snapshot for memory:{id}"))
479}
480
481impl MemoryCleanupRowSnapshot {
482 fn owner_namespace(&self, project: &str) -> (String, String) {
483 match (self.owner_scope.as_deref(), self.owner_key.as_deref()) {
484 (Some(scope), Some(key)) => (scope.to_string(), key.to_string()),
485 _ if self.project == project
486 && self.scope.as_deref().unwrap_or("project") != "global" =>
487 {
488 ("legacy_repo".to_string(), project.to_string())
489 }
490 _ => ("legacy_other".to_string(), self.project.clone()),
491 }
492 }
493
494 fn belongs_to_project(&self, project: &str) -> bool {
495 self.source_project.as_deref() == Some(project)
496 || self.target_project.as_deref() == Some(project)
497 || (self.owner_scope.as_deref() == Some("repo")
498 && self.owner_key.as_deref() == Some(project))
499 || (self.owner_scope.is_none()
500 && self.project == project
501 && self.scope.as_deref().unwrap_or("project") != "global")
502 }
503}
504
505fn snapshot_for(
506 snapshots: &[MemoryCleanupRowSnapshot],
507 id: i64,
508) -> Result<&MemoryCleanupRowSnapshot> {
509 snapshots
510 .iter()
511 .find(|snapshot| snapshot.id == id)
512 .ok_or_else(|| anyhow!("cleanup plan missing snapshot for memory:{id}"))
513}
514
515fn content_sha256(content: &str) -> String {
516 let mut hasher = Sha256::new();
517 hasher.update(content.as_bytes());
518 format!("{:x}", hasher.finalize())
519}