Skip to main content

remem/memory/store/
write.rs

1use anyhow::{Context, Result};
2use rusqlite::{params, Connection, OptionalExtension};
3
4use crate::memory::search_context::build_search_context;
5use crate::memory::state_key::{self, StateKeyDecision};
6use crate::memory::{
7    lifecycle::MemoryLifecycleOp,
8    operation::{
9        insert_operation_log, with_operation_savepoint, MemoryOperationInput, MemoryOperationPlan,
10    },
11    preference::consolidation::PreferenceConsolidationKind,
12};
13
14pub fn insert_memory(
15    conn: &Connection,
16    session_id: Option<&str>,
17    project: &str,
18    topic_key: Option<&str>,
19    title: &str,
20    content: &str,
21    memory_type: &str,
22    files: Option<&str>,
23) -> Result<i64> {
24    insert_memory_with_branch(
25        conn,
26        session_id,
27        project,
28        topic_key,
29        title,
30        content,
31        memory_type,
32        files,
33        None,
34    )
35}
36
37pub fn insert_memory_with_branch(
38    conn: &Connection,
39    session_id: Option<&str>,
40    project: &str,
41    topic_key: Option<&str>,
42    title: &str,
43    content: &str,
44    memory_type: &str,
45    files: Option<&str>,
46    branch: Option<&str>,
47) -> Result<i64> {
48    insert_memory_full(
49        conn,
50        session_id,
51        project,
52        topic_key,
53        title,
54        content,
55        memory_type,
56        files,
57        branch,
58        "project",
59        None,
60    )
61}
62
63#[allow(clippy::too_many_arguments)]
64pub fn insert_memory_full(
65    conn: &Connection,
66    session_id: Option<&str>,
67    project: &str,
68    topic_key: Option<&str>,
69    title: &str,
70    content: &str,
71    memory_type: &str,
72    files: Option<&str>,
73    branch: Option<&str>,
74    scope: &str,
75    created_at_override: Option<i64>,
76) -> Result<i64> {
77    insert_memory_full_with_reference_time(
78        conn,
79        session_id,
80        project,
81        topic_key,
82        title,
83        content,
84        memory_type,
85        files,
86        branch,
87        scope,
88        created_at_override,
89        created_at_override,
90    )
91}
92
93#[allow(clippy::too_many_arguments)]
94pub fn insert_memory_full_with_reference_time(
95    conn: &Connection,
96    session_id: Option<&str>,
97    project: &str,
98    topic_key: Option<&str>,
99    title: &str,
100    content: &str,
101    memory_type: &str,
102    files: Option<&str>,
103    branch: Option<&str>,
104    scope: &str,
105    created_at_override: Option<i64>,
106    reference_time_override: Option<i64>,
107) -> Result<i64> {
108    let now = chrono::Utc::now().timestamp();
109    let created_at = created_at_override.unwrap_or(now);
110    let reference_time = reference_time_override
111        .or(created_at_override)
112        .unwrap_or(created_at);
113    let (expires_at_epoch, valid_from_epoch) =
114        crate::memory::lifecycle::ttl_metadata(memory_type, topic_key, content, now);
115    let search_context = build_search_context(memory_type, topic_key, content, files);
116    let ownership = default_ownership(project, scope);
117    let state_key = state_key::derive_state_key(memory_type, topic_key, title, content);
118
119    let mut existing_id = None;
120    let mut preference_conflict = false;
121    if let Some(topic_key) = topic_key {
122        if !topic_key.is_empty() {
123            existing_id = conn
124                .query_row(
125                    "SELECT id FROM memories
126                     WHERE project = ?1 AND topic_key = ?2 AND scope = ?3
127                       AND memory_type = ?4
128                     ORDER BY CASE status WHEN 'active' THEN 0 ELSE 1 END,
129                              updated_at_epoch DESC,
130                              id DESC
131                     LIMIT 1",
132                    params![project, topic_key, scope, memory_type],
133                    |row| row.get(0),
134                )
135                .optional()?;
136        }
137    }
138
139    if existing_id.is_none() {
140        if let Some(decision) = &state_key {
141            if decision.allows_direct_upsert() {
142                existing_id = state_key::current_memory_id(
143                    conn,
144                    ownership.owner_scope,
145                    ownership.owner_key,
146                    memory_type,
147                    &decision.state_key,
148                    now,
149                )?;
150            }
151        }
152    }
153    if memory_type == "preference" && existing_id.is_none() {
154        if let Some(preference_match) =
155            crate::memory::preference::consolidation::find_preference_consolidation(
156                conn,
157                ownership.owner_scope,
158                ownership.owner_key,
159                scope,
160                branch,
161                content,
162                now,
163            )?
164        {
165            match preference_match.kind {
166                PreferenceConsolidationKind::SamePreference
167                | PreferenceConsolidationKind::Refinement => {
168                    existing_id = Some(preference_match.memory_id);
169                }
170                PreferenceConsolidationKind::Contradiction => {
171                    preference_conflict = true;
172                }
173            }
174        }
175    }
176
177    if existing_id.is_none() && !preference_conflict {
178        existing_id = crate::memory::semantic_dedup::find_curated_duplicate_id(
179            conn,
180            project,
181            scope,
182            memory_type,
183            title,
184            content,
185            topic_key,
186            branch,
187            now,
188        )?;
189    }
190
191    if let Some(id) = existing_id {
192        return with_memory_savepoint(conn, || {
193            update_existing_memory(
194                conn,
195                id,
196                session_id,
197                topic_key,
198                title,
199                content,
200                memory_type,
201                files,
202                branch,
203                scope,
204                &search_context,
205                expires_at_epoch,
206                valid_from_epoch,
207                &ownership,
208                state_key.as_ref(),
209                now,
210                reference_time,
211            )?;
212            refresh_memory_entities(conn, id, title, content)?;
213            refresh_memory_embedding(conn, id, title, content, memory_type, topic_key)?;
214            Ok(id)
215        });
216    }
217
218    with_memory_savepoint(conn, || {
219        conn.execute(
220            "INSERT INTO memories \
221             (session_id, project, topic_key, title, content, memory_type, files, search_context, \
222              created_at_epoch, updated_at_epoch, reference_time_epoch, status, branch, scope, \
223              source_project, target_project, owner_scope, owner_key, context_class, \
224              expires_at_epoch, valid_from_epoch) \
225             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, 'active', ?12, ?13, \
226                     ?14, ?15, ?16, ?17, ?18, ?19, ?20)",
227            params![
228                session_id,
229                project,
230                topic_key,
231                title,
232                content,
233                memory_type,
234                files,
235                search_context,
236                created_at,
237                now,
238                reference_time,
239                branch,
240                scope,
241                ownership.source_project,
242                ownership.target_project,
243                ownership.owner_scope,
244                ownership.owner_key,
245                ownership.context_class,
246                expires_at_epoch,
247                valid_from_epoch
248            ],
249        )?;
250        let id = conn.last_insert_rowid();
251        attach_state_key(conn, id, memory_type, &ownership, state_key.as_ref(), now)?;
252        refresh_memory_entities(conn, id, title, content)?;
253        refresh_memory_embedding(conn, id, title, content, memory_type, topic_key)?;
254        Ok(id)
255    })
256}
257
258#[allow(clippy::too_many_arguments)]
259pub fn insert_memory_full_with_operation_log(
260    conn: &Connection,
261    session_id: Option<&str>,
262    project: &str,
263    topic_key: Option<&str>,
264    title: &str,
265    content: &str,
266    memory_type: &str,
267    files: Option<&str>,
268    branch: Option<&str>,
269    scope: &str,
270    created_at_override: Option<i64>,
271    reference_time_override: Option<i64>,
272    operation_input: &MemoryOperationInput,
273    operation_plan: &MemoryOperationPlan,
274) -> Result<(i64, MemoryLifecycleOp)> {
275    with_operation_savepoint(conn, || {
276        let id = insert_memory_full_with_reference_time(
277            conn,
278            session_id,
279            project,
280            topic_key,
281            title,
282            content,
283            memory_type,
284            files,
285            branch,
286            scope,
287            created_at_override,
288            reference_time_override,
289        )?;
290        let mut logged_plan = operation_plan.clone();
291        logged_plan.target_memory_id = Some(id);
292        let operation_id = insert_operation_log(conn, operation_input, &logged_plan, Some(id))?;
293        crate::memory::edge::insert_supersedes_edges(
294            conn,
295            &logged_plan.superseded_ids,
296            id,
297            crate::memory::edge::MemoryEdgeWriteContext {
298                source_candidate_id: operation_input.source_candidate_id,
299                source_operation_id: Some(operation_id),
300                confidence: operation_input.confidence,
301                reason: Some(logged_plan.reason.as_str()),
302                ..Default::default()
303            },
304        )?;
305        crate::memory::edge::insert_conflicts_edges(
306            conn,
307            &logged_plan.conflicting_ids,
308            id,
309            crate::memory::edge::MemoryEdgeWriteContext {
310                source_candidate_id: operation_input.source_candidate_id,
311                source_operation_id: Some(operation_id),
312                confidence: operation_input.confidence,
313                reason: Some(logged_plan.reason.as_str()),
314                ..Default::default()
315            },
316        )?;
317        Ok((id, logged_plan.op))
318    })
319}
320
321pub(crate) struct DefaultOwnership<'a> {
322    pub(crate) source_project: &'a str,
323    pub(crate) target_project: Option<&'a str>,
324    pub(crate) owner_scope: &'static str,
325    pub(crate) owner_key: &'a str,
326    pub(crate) context_class: &'static str,
327}
328
329pub(crate) fn default_ownership<'a>(project: &'a str, scope: &str) -> DefaultOwnership<'a> {
330    if scope == "global" {
331        DefaultOwnership {
332            source_project: project,
333            target_project: None,
334            owner_scope: "user",
335            owner_key: "user:default",
336            context_class: "startup_core",
337        }
338    } else {
339        DefaultOwnership {
340            source_project: project,
341            target_project: Some(project),
342            owner_scope: "repo",
343            owner_key: project,
344            context_class: "startup_core",
345        }
346    }
347}
348
349#[allow(clippy::too_many_arguments)]
350fn update_existing_memory(
351    conn: &Connection,
352    id: i64,
353    session_id: Option<&str>,
354    topic_key: Option<&str>,
355    title: &str,
356    content: &str,
357    memory_type: &str,
358    files: Option<&str>,
359    branch: Option<&str>,
360    scope: &str,
361    search_context: &str,
362    expires_at_epoch: Option<i64>,
363    valid_from_epoch: Option<i64>,
364    ownership: &DefaultOwnership<'_>,
365    state_key: Option<&StateKeyDecision>,
366    now: i64,
367    reference_time: i64,
368) -> Result<()> {
369    let state_key_id = attach_state_key(conn, id, memory_type, ownership, state_key, now)?;
370    clear_obsolete_state_key_links(conn, id, state_key_id, now)?;
371    conn.execute(
372        "UPDATE memories SET session_id = ?1, topic_key = ?2, title = ?3, content = ?4, \
373         memory_type = ?5, files = ?6, updated_at_epoch = ?7, branch = ?8, \
374         scope = ?9, search_context = ?10, reference_time_epoch = ?11, \
375         status = 'active', valid_to_epoch = NULL, \
376         expires_at_epoch = ?12, valid_from_epoch = ?13, \
377         state_key_id = ?14, \
378         source_project = COALESCE(source_project, ?15), \
379         target_project = COALESCE(target_project, ?16), \
380         owner_scope = COALESCE(owner_scope, ?17), \
381         owner_key = COALESCE(owner_key, ?18), \
382         context_class = COALESCE(context_class, ?19) \
383         WHERE id = ?20",
384        params![
385            session_id,
386            topic_key,
387            title,
388            content,
389            memory_type,
390            files,
391            now,
392            branch,
393            scope,
394            search_context,
395            reference_time,
396            expires_at_epoch,
397            valid_from_epoch,
398            state_key_id,
399            ownership.source_project,
400            ownership.target_project,
401            ownership.owner_scope,
402            ownership.owner_key,
403            ownership.context_class,
404            id
405        ],
406    )?;
407    Ok(())
408}
409
410pub(crate) fn clear_obsolete_state_key_links(
411    conn: &Connection,
412    id: i64,
413    active_state_key_id: Option<i64>,
414    now: i64,
415) -> Result<()> {
416    update_state_key_links(conn, id, active_state_key_id, active_state_key_id, now)
417}
418
419pub(crate) fn update_state_key_links(
420    conn: &Connection,
421    id: i64,
422    row_state_key_id: Option<i64>,
423    current_state_key_id: Option<i64>,
424    now: i64,
425) -> Result<()> {
426    conn.execute(
427        "UPDATE memories SET state_key_id = ?1 WHERE id = ?2",
428        params![row_state_key_id, id],
429    )?;
430    conn.execute(
431        "UPDATE memory_state_keys
432         SET current_memory_id = NULL, updated_at_epoch = ?3
433         WHERE current_memory_id = ?1
434           AND (?2 IS NULL OR id <> ?2)",
435        params![id, current_state_key_id, now],
436    )?;
437    Ok(())
438}
439
440fn attach_state_key(
441    conn: &Connection,
442    id: i64,
443    memory_type: &str,
444    ownership: &DefaultOwnership<'_>,
445    state_key: Option<&StateKeyDecision>,
446    now: i64,
447) -> Result<Option<i64>> {
448    state_key
449        .map(|decision| {
450            state_key::attach_current_memory(
451                conn,
452                id,
453                ownership.owner_scope,
454                ownership.owner_key,
455                memory_type,
456                decision,
457                now,
458            )
459        })
460        .transpose()
461}
462
463fn with_memory_savepoint<T>(conn: &Connection, f: impl FnOnce() -> Result<T>) -> Result<T> {
464    conn.execute_batch("SAVEPOINT remem_memory_state_write")?;
465    match f() {
466        Ok(value) => {
467            conn.execute_batch("RELEASE SAVEPOINT remem_memory_state_write")?;
468            Ok(value)
469        }
470        Err(error) => {
471            let rollback = conn.execute_batch(
472                "ROLLBACK TO SAVEPOINT remem_memory_state_write;
473                 RELEASE SAVEPOINT remem_memory_state_write;",
474            );
475            if let Err(rollback_error) = rollback {
476                return Err(error.context(format!(
477                    "memory state-key rollback also failed: {rollback_error}"
478                )));
479            }
480            Err(error)
481        }
482    }
483}
484
485fn refresh_memory_entities(conn: &Connection, id: i64, title: &str, content: &str) -> Result<()> {
486    let entities = crate::retrieval::entity::extract_entities(title, content);
487    crate::retrieval::entity::refresh_memory_entities(conn, id, &entities)
488        .with_context(|| format!("entity refresh failed for memory id={id}"))
489}
490
491fn refresh_memory_embedding(
492    conn: &Connection,
493    id: i64,
494    title: &str,
495    content: &str,
496    memory_type: &str,
497    topic_key: Option<&str>,
498) -> Result<()> {
499    crate::retrieval::vector::upsert_memory_embedding(
500        conn,
501        id,
502        title,
503        content,
504        memory_type,
505        topic_key,
506    )
507}
508
509#[cfg(test)]
510mod tests;