1use futures::TryStreamExt as _;
5#[allow(unused_imports)]
6use zeph_common;
7use zeph_common::ContextFidelity;
8use zeph_db::ActiveDialect;
9use zeph_db::fts::sanitize_fts_query;
10#[allow(unused_imports)]
11use zeph_db::{begin_write, placeholder_list, sql};
12use zeph_llm::provider::{Message, MessageMetadata, MessagePart, MessageVisibility, Role};
13
14use super::SqliteStore;
15use crate::error::MemoryError;
16use crate::types::{ConversationId, MessageId};
17
18const MAX_BATCH: usize = 490;
23
24fn compacted_at_bind_expr() -> String {
32 <ActiveDialect as zeph_db::dialect::Dialect>::timestamptz_from_epoch("?")
33}
34
35fn parse_role(s: &str) -> Role {
36 match s {
37 "assistant" => Role::Assistant,
38 "system" => Role::System,
39 _ => Role::User,
40 }
41}
42
43type HistoryRow = (String, String, String, String, i64, i32, Option<String>);
48
49fn parse_persisted_trust_level(s: Option<&str>) -> Option<u8> {
65 match s {
66 None => None,
67 Some("trusted") => Some(0),
68 Some("local_untrusted") => Some(1),
69 Some("external_untrusted") => Some(2),
70 Some(other) => {
71 tracing::warn!(
72 value = other,
73 "unrecognized persisted trust_level, treating as external_untrusted"
74 );
75 Some(2)
76 }
77 }
78}
79
80fn trust_level_ordinal_to_str(ordinal: Option<u8>) -> Option<&'static str> {
88 match ordinal {
89 None => None,
90 Some(0) => Some("trusted"),
91 Some(1) => Some("local_untrusted"),
92 Some(_) => Some("external_untrusted"),
93 }
94}
95
96#[must_use]
97pub fn role_str(role: Role) -> &'static str {
98 match role {
99 Role::System => "system",
100 Role::Assistant => "assistant",
101 Role::User | _ => "user",
102 }
103}
104
105fn legacy_key_to_kind(key: &str) -> Option<&'static str> {
108 match key {
109 "Text" => Some("text"),
110 "ToolOutput" => Some("tool_output"),
111 "Recall" => Some("recall"),
112 "CodeContext" => Some("code_context"),
113 "Summary" => Some("summary"),
114 "CrossSession" => Some("cross_session"),
115 "ToolUse" => Some("tool_use"),
116 "ToolResult" => Some("tool_result"),
117 "Image" => Some("image"),
118 "ThinkingBlock" => Some("thinking_block"),
119 "RedactedThinkingBlock" => Some("redacted_thinking_block"),
120 "Compaction" => Some("compaction"),
121 _ => None,
122 }
123}
124
125fn try_parse_legacy_parts(parts_json: &str) -> Option<Vec<MessagePart>> {
133 let array: Vec<serde_json::Value> = serde_json::from_str(parts_json).ok()?;
134 let mut result = Vec::with_capacity(array.len());
135 for element in array {
136 let obj = element.as_object()?;
137 if obj.contains_key("kind") {
138 return None;
139 }
140 if obj.len() != 1 {
141 return None;
142 }
143 let (key, inner) = obj.iter().next()?;
144 let kind = legacy_key_to_kind(key)?;
145 let mut new_obj = match inner {
146 serde_json::Value::Object(m) => m.clone(),
147 other => {
149 let mut m = serde_json::Map::new();
150 m.insert("data".to_string(), other.clone());
151 m
152 }
153 };
154 new_obj.insert(
155 "kind".to_string(),
156 serde_json::Value::String(kind.to_string()),
157 );
158 let part: MessagePart = serde_json::from_value(serde_json::Value::Object(new_obj)).ok()?;
159 result.push(part);
160 }
161 Some(result)
162}
163
164fn parse_parts_json(role_str: &str, parts_json: &str) -> Vec<MessagePart> {
169 if parts_json == "[]" {
170 return vec![];
171 }
172 match serde_json::from_str(parts_json) {
173 Ok(p) => p,
174 Err(e) => {
175 if let Some(parts) = try_parse_legacy_parts(parts_json) {
176 let truncated = parts_json.chars().take(120).collect::<String>();
177 tracing::warn!(
178 role = %role_str,
179 parts_json = %truncated,
180 "loaded legacy-format message parts via compat path"
181 );
182 return parts;
183 }
184 let truncated = parts_json.chars().take(120).collect::<String>();
185 tracing::warn!(
186 role = %role_str,
187 parts_json = %truncated,
188 error = %e,
189 "failed to deserialize message parts, falling back to empty"
190 );
191 vec![]
192 }
193 }
194}
195
196impl SqliteStore {
197 pub async fn create_conversation(&self) -> Result<ConversationId, MemoryError> {
203 let row: (ConversationId,) = zeph_db::query_as(sql!(
204 "INSERT INTO conversations DEFAULT VALUES RETURNING id"
205 ))
206 .fetch_one(&self.pool)
207 .await?;
208 Ok(row.0)
209 }
210
211 pub async fn save_message(
217 &self,
218 conversation_id: ConversationId,
219 role: &str,
220 content: &str,
221 ) -> Result<MessageId, MemoryError> {
222 self.save_message_with_parts(conversation_id, role, content, "[]")
223 .await
224 }
225
226 pub async fn save_message_with_parts(
232 &self,
233 conversation_id: ConversationId,
234 role: &str,
235 content: &str,
236 parts_json: &str,
237 ) -> Result<MessageId, MemoryError> {
238 self.save_message_with_metadata(
239 conversation_id,
240 role,
241 content,
242 parts_json,
243 MessageVisibility::Both,
244 )
245 .await
246 }
247
248 pub async fn save_message_with_category(
256 &self,
257 conversation_id: ConversationId,
258 role: &str,
259 content: &str,
260 category: Option<&str>,
261 ) -> Result<MessageId, MemoryError> {
262 let importance_score = crate::semantic::importance::compute_importance(content, role);
263 let row: (MessageId,) = zeph_db::query_as(sql!(
264 "INSERT INTO messages \
265 (conversation_id, role, content, parts, visibility, \
266 importance_score, category) \
267 VALUES (?, ?, ?, '[]', 'both', ?, ?) RETURNING id"
268 ))
269 .bind(conversation_id)
270 .bind(role)
271 .bind(content)
272 .bind(importance_score)
273 .bind(category)
274 .fetch_one(&self.pool)
275 .await?;
276 Ok(row.0)
277 }
278
279 pub async fn save_message_with_metadata(
289 &self,
290 conversation_id: ConversationId,
291 role: &str,
292 content: &str,
293 parts_json: &str,
294 visibility: MessageVisibility,
295 ) -> Result<MessageId, MemoryError> {
296 self.save_message_with_provenance(
297 conversation_id,
298 role,
299 content,
300 parts_json,
301 visibility,
302 None,
303 None,
304 )
305 .await
306 }
307
308 #[allow(clippy::too_many_arguments)]
321 pub async fn save_message_with_provenance(
322 &self,
323 conversation_id: ConversationId,
324 role: &str,
325 content: &str,
326 parts_json: &str,
327 visibility: MessageVisibility,
328 source_kind: Option<&str>,
329 trust_level: Option<&str>,
330 ) -> Result<MessageId, MemoryError> {
331 const MAX_BYTES: usize = 100 * 1024;
332
333 let content_cow: std::borrow::Cow<'_, str> = if content.len() > MAX_BYTES {
336 let boundary = content.floor_char_boundary(MAX_BYTES);
337 tracing::debug!(
338 original_bytes = content.len(),
339 "save_message: content exceeds 100KB, truncating"
340 );
341 std::borrow::Cow::Owned(format!(
342 "{}... [truncated, {} bytes total]",
343 &content[..boundary],
344 content.len()
345 ))
346 } else {
347 std::borrow::Cow::Borrowed(content)
348 };
349
350 let importance_score = crate::semantic::importance::compute_importance(&content_cow, role);
351 let json_cast = <ActiveDialect as zeph_db::dialect::Dialect>::JSON_CAST;
352 let insert_sql = zeph_db::rewrite_placeholders(&format!(
353 "INSERT INTO messages \
354 (conversation_id, role, content, parts, visibility, importance_score, \
355 source_kind, trust_level) \
356 VALUES (?, ?, ?, ?{json_cast}, ?, ?, ?, ?) RETURNING id"
357 ));
358 let row: (MessageId,) = zeph_db::query_as(sqlx::AssertSqlSafe(insert_sql))
359 .bind(conversation_id)
360 .bind(role)
361 .bind(content_cow.as_ref())
362 .bind(parts_json)
363 .bind(visibility.as_db_str())
364 .bind(importance_score)
365 .bind(source_kind)
366 .bind(trust_level)
367 .fetch_one(&self.pool)
368 .await?;
369 Ok(row.0)
370 }
371
372 pub async fn load_history(
382 &self,
383 conversation_id: ConversationId,
384 limit: u32,
385 ) -> Result<Vec<Message>, MemoryError> {
386 let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
387 let raw = format!(
388 "SELECT role, content, {parts_select} AS parts, visibility, id, \
389 CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level FROM (\
390 SELECT role, content, parts, visibility, id, fidelity_tag, trust_level \
391 FROM messages \
392 WHERE conversation_id = ? AND deleted_at IS NULL \
393 ORDER BY id DESC \
394 LIMIT ?\
395 ) ORDER BY id ASC"
396 );
397 let sql = zeph_db::rewrite_placeholders(&raw);
398 let rows: Vec<HistoryRow> = zeph_db::query_as(sqlx::AssertSqlSafe(sql))
399 .bind(conversation_id)
400 .bind(i64::from(limit))
401 .fetch_all(&self.pool)
402 .await?;
403
404 let messages = rows
405 .into_iter()
406 .map(
407 |(
408 role_str,
409 content,
410 parts_json,
411 visibility_str,
412 row_id,
413 fidelity_raw,
414 trust_level_str,
415 )| {
416 let parts = parse_parts_json(&role_str, &parts_json);
417 Message {
418 role: parse_role(&role_str),
419 content,
420 parts,
421 metadata: MessageMetadata {
422 visibility: MessageVisibility::from_db_str(&visibility_str),
423 compacted_at: None,
424 deferred_summary: None,
425 focus_pinned: false,
426 focus_marker_id: None,
427 db_id: Some(row_id),
428 fidelity_tag: if fidelity_raw == 0 {
429 None
430 } else {
431 u8::try_from(fidelity_raw)
432 .ok()
433 .map(ContextFidelity::from_u8)
434 },
435 embedding: None,
436 trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
437 },
438 }
439 },
440 )
441 .collect();
442 Ok(messages)
443 }
444
445 pub async fn load_history_filtered(
457 &self,
458 conversation_id: ConversationId,
459 limit: u32,
460 agent_visible: Option<bool>,
461 user_visible: Option<bool>,
462 ) -> Result<Vec<Message>, MemoryError> {
463 let exclude_user_only = agent_visible == Some(true);
469 let exclude_agent_only = user_visible == Some(true);
470
471 let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
472 let raw = format!(
473 "WITH recent AS (\
474 SELECT role, content, parts, visibility, id, fidelity_tag, trust_level \
475 FROM messages \
476 WHERE conversation_id = ? \
477 AND deleted_at IS NULL \
478 AND (NOT ? OR visibility != 'user_only') \
479 AND (NOT ? OR visibility != 'agent_only') \
480 ORDER BY id DESC \
481 LIMIT ?\
482 ) SELECT role, content, {parts_select} AS parts, visibility, id, \
483 CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level \
484 FROM recent ORDER BY id ASC"
485 );
486 let sql = zeph_db::rewrite_placeholders(&raw);
487 let rows: Vec<HistoryRow> = zeph_db::query_as(sqlx::AssertSqlSafe(sql))
488 .bind(conversation_id)
489 .bind(exclude_user_only)
490 .bind(exclude_agent_only)
491 .bind(i64::from(limit))
492 .fetch_all(&self.pool)
493 .await?;
494
495 let messages = rows
496 .into_iter()
497 .map(
498 |(
499 role_str,
500 content,
501 parts_json,
502 visibility_str,
503 row_id,
504 fidelity_raw,
505 trust_level_str,
506 )| {
507 let parts = parse_parts_json(&role_str, &parts_json);
508 Message {
509 role: parse_role(&role_str),
510 content,
511 parts,
512 metadata: MessageMetadata {
513 visibility: MessageVisibility::from_db_str(&visibility_str),
514 compacted_at: None,
515 deferred_summary: None,
516 focus_pinned: false,
517 focus_marker_id: None,
518 db_id: Some(row_id),
519 fidelity_tag: if fidelity_raw == 0 {
520 None
521 } else {
522 u8::try_from(fidelity_raw)
523 .ok()
524 .map(ContextFidelity::from_u8)
525 },
526 embedding: None,
527 trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
528 },
529 }
530 },
531 )
532 .collect();
533 Ok(messages)
534 }
535
536 pub async fn update_fidelity_tags(
549 &self,
550 updates: &[(MessageId, u8)],
551 ) -> Result<(), MemoryError> {
552 const MAX_FIDELITY_BATCH: usize = 333;
555 if updates.is_empty() {
556 return Ok(());
557 }
558 let mut tx = self.pool.begin().await?;
559 for chunk in updates.chunks(MAX_FIDELITY_BATCH) {
560 let case_arms: String = chunk
561 .iter()
562 .map(|_| "WHEN ? THEN ?")
563 .collect::<Vec<_>>()
564 .join(" ");
565 let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
566 let sql = zeph_db::rewrite_placeholders(&format!(
567 "UPDATE messages SET fidelity_tag = CASE id {case_arms} END WHERE id IN ({in_list})"
568 ));
569 let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
570 for &(id, tag) in chunk {
571 q = q.bind(id.0).bind(i32::from(tag));
572 }
573 for &(id, _) in chunk {
574 q = q.bind(id.0);
575 }
576 q.execute(&mut *tx).await?;
577 }
578 tx.commit().await?;
579 Ok(())
580 }
581
582 pub async fn replace_conversation(
598 &self,
599 conversation_id: ConversationId,
600 compacted_range: std::ops::RangeInclusive<MessageId>,
601 summary_role: &str,
602 summary_content: &str,
603 trust_level: Option<u8>,
604 ) -> Result<MessageId, MemoryError> {
605 let now = {
606 let secs = std::time::SystemTime::now()
607 .duration_since(std::time::UNIX_EPOCH)
608 .unwrap_or_default()
609 .as_secs();
610 format!("{secs}")
611 };
612 let start_id = compacted_range.start().0;
613 let end_id = compacted_range.end().0;
614
615 let mut tx = self.pool.begin().await?;
616
617 let compacted_at_expr = compacted_at_bind_expr();
618 let update_sql = zeph_db::rewrite_placeholders(&format!(
619 "UPDATE messages SET visibility = 'user_only', compacted_at = {compacted_at_expr} \
620 WHERE conversation_id = ? AND id >= ? AND id <= ?"
621 ));
622 zeph_db::query(sqlx::AssertSqlSafe(update_sql))
623 .bind(&now)
624 .bind(conversation_id)
625 .bind(start_id)
626 .bind(end_id)
627 .execute(&mut *tx)
628 .await?;
629
630 let row: (MessageId,) = zeph_db::query_as(sql!(
632 "INSERT INTO messages \
633 (conversation_id, role, content, parts, visibility, trust_level) \
634 VALUES (?, ?, ?, '[]', 'agent_only', ?) RETURNING id"
635 ))
636 .bind(conversation_id)
637 .bind(summary_role)
638 .bind(summary_content)
639 .bind(trust_level_ordinal_to_str(trust_level))
640 .fetch_one(&mut *tx)
641 .await?;
642
643 tx.commit().await?;
644
645 Ok(row.0)
646 }
647
648 pub async fn apply_tool_pair_summaries(
663 &self,
664 conversation_id: ConversationId,
665 hide_ids: &[i64],
666 summaries: &[String],
667 trust_levels: &[Option<u8>],
668 ) -> Result<(), MemoryError> {
669 if hide_ids.is_empty() && summaries.is_empty() {
670 return Ok(());
671 }
672
673 let now = std::time::SystemTime::now()
674 .duration_since(std::time::UNIX_EPOCH)
675 .unwrap_or_default()
676 .as_secs()
677 .to_string();
678
679 let mut tx = self.pool.begin().await?;
680
681 let compacted_at_expr = compacted_at_bind_expr();
682 for chunk in hide_ids.chunks(MAX_BATCH) {
683 let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
684 let update_sql = zeph_db::rewrite_placeholders(&format!(
685 "UPDATE messages SET visibility = 'user_only', compacted_at = {compacted_at_expr} \
686 WHERE id IN ({in_list})"
687 ));
688 let mut q = zeph_db::query(sqlx::AssertSqlSafe(update_sql)).bind(&now);
689 for &id in chunk {
690 q = q.bind(id);
691 }
692 q.execute(&mut *tx).await?;
693 }
694
695 for (i, summary) in summaries.iter().enumerate() {
696 let content = format!("[tool summary] {summary}");
697 let parts = serde_json::to_string(&[MessagePart::Summary {
698 text: summary.clone(),
699 }])
700 .unwrap_or_else(|_| "[]".to_string());
701 let trust_level = trust_levels.get(i).copied().flatten();
702 let json_cast = <ActiveDialect as zeph_db::dialect::Dialect>::JSON_CAST;
703 let insert_sql = zeph_db::rewrite_placeholders(&format!(
704 "INSERT INTO messages \
705 (conversation_id, role, content, parts, visibility, trust_level) \
706 VALUES (?, 'assistant', ?, ?{json_cast}, 'agent_only', ?)"
707 ));
708 zeph_db::query(sqlx::AssertSqlSafe(insert_sql))
709 .bind(conversation_id)
710 .bind(&content)
711 .bind(&parts)
712 .bind(trust_level_ordinal_to_str(trust_level))
713 .execute(&mut *tx)
714 .await?;
715 }
716
717 tx.commit().await?;
718 Ok(())
719 }
720
721 pub async fn oldest_message_ids(
727 &self,
728 conversation_id: ConversationId,
729 n: u32,
730 ) -> Result<Vec<MessageId>, MemoryError> {
731 let rows: Vec<(MessageId,)> = zeph_db::query_as(
732 sql!("SELECT id FROM messages WHERE conversation_id = ? AND deleted_at IS NULL ORDER BY id ASC LIMIT ?"),
733 )
734 .bind(conversation_id)
735 .bind(i64::from(n))
736 .fetch_all(&self.pool)
737 .await?;
738 Ok(rows.into_iter().map(|r| r.0).collect())
739 }
740
741 pub async fn latest_conversation_id(&self) -> Result<Option<ConversationId>, MemoryError> {
747 let row: Option<(ConversationId,)> = zeph_db::query_as(sql!(
748 "SELECT id FROM conversations ORDER BY id DESC LIMIT 1"
749 ))
750 .fetch_optional(&self.pool)
751 .await?;
752 Ok(row.map(|r| r.0))
753 }
754
755 pub async fn message_by_id(
765 &self,
766 message_id: MessageId,
767 ) -> Result<Option<Message>, MemoryError> {
768 let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
769 let sql = zeph_db::rewrite_placeholders(&format!(
770 "SELECT role, content, {parts_select} AS parts, visibility, \
771 CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level FROM messages \
772 WHERE id = ? AND deleted_at IS NULL"
773 ));
774 let row: Option<(String, String, String, String, i32, Option<String>)> =
775 zeph_db::query_as(sqlx::AssertSqlSafe(sql))
776 .bind(message_id)
777 .fetch_optional(&self.pool)
778 .await?;
779
780 Ok(row.map(
781 |(role_str, content, parts_json, visibility_str, fidelity_raw, trust_level_str)| {
782 let parts = parse_parts_json(&role_str, &parts_json);
783 Message {
784 role: parse_role(&role_str),
785 content,
786 parts,
787 metadata: MessageMetadata {
788 visibility: MessageVisibility::from_db_str(&visibility_str),
789 compacted_at: None,
790 deferred_summary: None,
791 focus_pinned: false,
792 focus_marker_id: None,
793 db_id: Some(message_id.0),
794 fidelity_tag: if fidelity_raw == 0 {
795 None
796 } else {
797 u8::try_from(fidelity_raw)
798 .ok()
799 .map(ContextFidelity::from_u8)
800 },
801 embedding: None,
802 trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
803 },
804 }
805 },
806 ))
807 }
808
809 pub async fn messages_by_ids(
815 &self,
816 ids: &[MessageId],
817 ) -> Result<Vec<(MessageId, Message)>, MemoryError> {
818 if ids.is_empty() {
819 return Ok(Vec::new());
820 }
821
822 let placeholders = zeph_db::placeholder_list(1, ids.len());
823 let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
824
825 let query = format!(
826 "SELECT id, role, content, {parts_select} AS parts FROM messages \
827 WHERE id IN ({placeholders}) AND visibility != 'user_only' AND deleted_at IS NULL"
828 );
829 let mut q =
830 zeph_db::query_as::<_, (MessageId, String, String, String)>(sqlx::AssertSqlSafe(query));
831 for &id in ids {
832 q = q.bind(id);
833 }
834
835 let rows = q.fetch_all(&self.pool).await?;
836
837 Ok(rows
838 .into_iter()
839 .map(|(id, role_str, content, parts_json)| {
840 let parts = parse_parts_json(&role_str, &parts_json);
841 (
842 id,
843 Message {
844 role: parse_role(&role_str),
845 content,
846 parts,
847 metadata: MessageMetadata {
848 db_id: Some(id.0),
849 ..MessageMetadata::default()
850 },
851 },
852 )
853 })
854 .collect())
855 }
856
857 pub async fn unembedded_message_ids(
863 &self,
864 limit: Option<usize>,
865 ) -> Result<Vec<(MessageId, ConversationId, String, String)>, MemoryError> {
866 let effective_limit = limit.map_or(i64::MAX, |l| i64::try_from(l).unwrap_or(i64::MAX));
867
868 let rows: Vec<(MessageId, ConversationId, String, String)> = zeph_db::query_as(sql!(
869 "SELECT m.id, m.conversation_id, m.role, m.content \
870 FROM messages m \
871 LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
872 WHERE em.id IS NULL AND m.deleted_at IS NULL \
873 ORDER BY m.id ASC \
874 LIMIT ?"
875 ))
876 .bind(effective_limit)
877 .fetch_all(&self.pool)
878 .await?;
879
880 Ok(rows)
881 }
882
883 pub async fn count_unembedded_messages(&self) -> Result<usize, MemoryError> {
889 let row: (i64,) = zeph_db::query_as(sql!(
890 "SELECT COUNT(*) FROM messages m \
891 LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
892 WHERE em.id IS NULL AND m.deleted_at IS NULL"
893 ))
894 .fetch_one(&self.pool)
895 .await?;
896 Ok(usize::try_from(row.0).unwrap_or(usize::MAX))
897 }
898
899 pub fn stream_unembedded_messages(
910 &self,
911 limit: i64,
912 ) -> impl futures::Stream<Item = Result<(MessageId, ConversationId, String, String), MemoryError>> + '_
913 {
914 zeph_db::query_as(sql!(
915 "SELECT m.id, m.conversation_id, m.role, m.content \
916 FROM messages m \
917 LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
918 WHERE em.id IS NULL AND m.deleted_at IS NULL \
919 ORDER BY m.id ASC \
920 LIMIT ?"
921 ))
922 .bind(limit)
923 .fetch(&self.pool)
924 .map_err(MemoryError::from)
925 }
926
927 pub async fn count_messages(
933 &self,
934 conversation_id: ConversationId,
935 ) -> Result<i64, MemoryError> {
936 let row: (i64,) = zeph_db::query_as(sql!(
937 "SELECT COUNT(*) FROM messages WHERE conversation_id = ? AND deleted_at IS NULL"
938 ))
939 .bind(conversation_id)
940 .fetch_one(&self.pool)
941 .await?;
942 Ok(row.0)
943 }
944
945 pub async fn count_messages_after(
951 &self,
952 conversation_id: ConversationId,
953 after_id: MessageId,
954 ) -> Result<i64, MemoryError> {
955 let row: (i64,) =
956 zeph_db::query_as(
957 sql!("SELECT COUNT(*) FROM messages WHERE conversation_id = ? AND id > ? AND deleted_at IS NULL"),
958 )
959 .bind(conversation_id)
960 .bind(after_id)
961 .fetch_one(&self.pool)
962 .await?;
963 Ok(row.0)
964 }
965
966 pub async fn keyword_search(
975 &self,
976 query: &str,
977 limit: usize,
978 conversation_id: Option<ConversationId>,
979 ) -> Result<Vec<(MessageId, f64)>, MemoryError> {
980 let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
981 let safe_query = sanitize_fts_query(query);
982 if safe_query.is_empty() {
983 return Ok(Vec::new());
984 }
985
986 let rows: Vec<(MessageId, f64)> = if let Some(cid) = conversation_id {
987 zeph_db::query_as(
988 sql!("SELECT m.id, -rank AS score \
989 FROM messages_fts f \
990 JOIN messages m ON m.id = f.rowid \
991 WHERE messages_fts MATCH ? AND m.conversation_id = ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL \
992 ORDER BY rank \
993 LIMIT ?"),
994 )
995 .bind(&safe_query)
996 .bind(cid)
997 .bind(effective_limit)
998 .fetch_all(&self.pool)
999 .await?
1000 } else {
1001 zeph_db::query_as(sql!(
1002 "SELECT m.id, -rank AS score \
1003 FROM messages_fts f \
1004 JOIN messages m ON m.id = f.rowid \
1005 WHERE messages_fts MATCH ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL \
1006 ORDER BY rank \
1007 LIMIT ?"
1008 ))
1009 .bind(&safe_query)
1010 .bind(effective_limit)
1011 .fetch_all(&self.pool)
1012 .await?
1013 };
1014
1015 Ok(rows)
1016 }
1017
1018 pub async fn keyword_search_with_time_range(
1031 &self,
1032 query: &str,
1033 limit: usize,
1034 conversation_id: Option<ConversationId>,
1035 after: Option<&str>,
1036 before: Option<&str>,
1037 ) -> Result<Vec<(MessageId, f64)>, MemoryError> {
1038 let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
1039 let safe_query = sanitize_fts_query(query);
1040 if safe_query.is_empty() {
1041 return Ok(Vec::new());
1042 }
1043
1044 let after_clause = if after.is_some() {
1046 " AND m.created_at > ?"
1047 } else {
1048 ""
1049 };
1050 let before_clause = if before.is_some() {
1051 " AND m.created_at < ?"
1052 } else {
1053 ""
1054 };
1055 let conv_clause = if conversation_id.is_some() {
1056 " AND m.conversation_id = ?"
1057 } else {
1058 ""
1059 };
1060
1061 let sql = zeph_db::rewrite_placeholders(&format!(
1062 "SELECT m.id, -rank AS score \
1063 FROM messages_fts f \
1064 JOIN messages m ON m.id = f.rowid \
1065 WHERE messages_fts MATCH ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL\
1066 {after_clause}{before_clause}{conv_clause} \
1067 ORDER BY rank \
1068 LIMIT ?"
1069 ));
1070
1071 let mut q =
1072 zeph_db::query_as::<_, (MessageId, f64)>(sqlx::AssertSqlSafe(sql)).bind(&safe_query);
1073 if let Some(a) = after {
1074 q = q.bind(a);
1075 }
1076 if let Some(b) = before {
1077 q = q.bind(b);
1078 }
1079 if let Some(cid) = conversation_id {
1080 q = q.bind(cid);
1081 }
1082 q = q.bind(effective_limit);
1083
1084 Ok(q.fetch_all(&self.pool).await?)
1085 }
1086
1087 #[tracing::instrument(name = "memory.store.message_timestamps", skip(self, ids), fields(count = ids.len()))]
1095 pub async fn message_timestamps(
1096 &self,
1097 ids: &[MessageId],
1098 ) -> Result<std::collections::HashMap<MessageId, i64>, MemoryError> {
1099 if ids.is_empty() {
1100 return Ok(std::collections::HashMap::new());
1101 }
1102
1103 let placeholders: String =
1104 zeph_db::rewrite_placeholders(&ids.iter().map(|_| "?").collect::<Vec<_>>().join(","));
1105 let epoch_expr = <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
1106 let query = format!(
1107 "SELECT id, {epoch_expr} FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1108 );
1109 let mut q = zeph_db::query_as::<_, (MessageId, i64)>(sqlx::AssertSqlSafe(query));
1110 for &id in ids {
1111 q = q.bind(id);
1112 }
1113
1114 let rows = q.fetch_all(&self.pool).await?;
1115 Ok(rows.into_iter().collect())
1116 }
1117
1118 pub async fn load_messages_range(
1124 &self,
1125 conversation_id: ConversationId,
1126 after_message_id: MessageId,
1127 limit: usize,
1128 ) -> Result<Vec<(MessageId, String, String)>, MemoryError> {
1129 let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
1130
1131 let rows: Vec<(MessageId, String, String)> = zeph_db::query_as(sql!(
1132 "SELECT id, role, content FROM messages \
1133 WHERE conversation_id = ? AND id > ? AND deleted_at IS NULL \
1134 ORDER BY id ASC LIMIT ?"
1135 ))
1136 .bind(conversation_id)
1137 .bind(after_message_id)
1138 .bind(effective_limit)
1139 .fetch_all(&self.pool)
1140 .await?;
1141
1142 Ok(rows)
1143 }
1144
1145 pub async fn get_eviction_candidates(
1153 &self,
1154 ) -> Result<Vec<crate::eviction::EvictionEntry>, crate::error::MemoryError> {
1155 let created_at_sel =
1160 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
1161 let last_accessed_sel =
1162 <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("last_accessed");
1163 let raw = format!(
1164 "SELECT id, {created_at_sel}, {last_accessed_sel}, access_count \
1165 FROM messages WHERE deleted_at IS NULL"
1166 );
1167 let query_sql = zeph_db::rewrite_placeholders(&raw);
1168 let rows: Vec<(MessageId, String, Option<String>, i32)> =
1169 zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
1170 .fetch_all(&self.pool)
1171 .await?;
1172
1173 Ok(rows
1174 .into_iter()
1175 .map(
1176 |(id, created_at, last_accessed, access_count)| crate::eviction::EvictionEntry {
1177 id,
1178 created_at,
1179 last_accessed,
1180 access_count: access_count.try_into().unwrap_or(0),
1181 },
1182 )
1183 .collect())
1184 }
1185
1186 pub async fn soft_delete_messages(
1194 &self,
1195 ids: &[MessageId],
1196 ) -> Result<(), crate::error::MemoryError> {
1197 if ids.is_empty() {
1198 return Ok(());
1199 }
1200 for chunk in ids.chunks(MAX_BATCH) {
1201 let placeholders = placeholder_list(1, chunk.len());
1202 let sql = format!(
1203 "UPDATE messages SET deleted_at = CURRENT_TIMESTAMP \
1204 WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1205 );
1206 let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
1207 for &id in chunk {
1208 q = q.bind(id);
1209 }
1210 q.execute(&self.pool).await?;
1211 }
1212 Ok(())
1213 }
1214
1215 pub async fn get_soft_deleted_message_ids(
1221 &self,
1222 ) -> Result<Vec<MessageId>, crate::error::MemoryError> {
1223 let rows: Vec<(MessageId,)> = zeph_db::query_as(sql!(
1224 "SELECT id FROM messages WHERE deleted_at IS NOT NULL AND qdrant_cleaned = FALSE"
1225 ))
1226 .fetch_all(&self.pool)
1227 .await?;
1228 Ok(rows.into_iter().map(|(id,)| id).collect())
1229 }
1230
1231 pub async fn filter_out_preserved_episode_ids(
1246 &self,
1247 candidate_ids: &[MessageId],
1248 ) -> Result<Vec<MessageId>, crate::error::MemoryError> {
1249 if candidate_ids.is_empty() {
1250 return Ok(Vec::new());
1251 }
1252 let mut safe_to_delete: Vec<MessageId> = Vec::with_capacity(candidate_ids.len());
1253 for chunk in candidate_ids.chunks(MAX_BATCH) {
1254 let placeholders = placeholder_list(1, chunk.len());
1255 let sql = format!(
1256 "SELECT m.id \
1257 FROM messages m \
1258 WHERE m.id IN ({placeholders}) \
1259 AND NOT EXISTS ( \
1260 SELECT 1 \
1261 FROM summaries s \
1262 WHERE s.first_message_id IS NOT NULL \
1263 AND s.last_message_id IS NOT NULL \
1264 AND m.id >= s.first_message_id \
1265 AND m.id <= s.last_message_id \
1266 )"
1267 );
1268 let mut q = zeph_db::query_as::<_, (MessageId,)>(sqlx::AssertSqlSafe(sql));
1269 for &id in chunk {
1270 q = q.bind(id);
1271 }
1272 let rows: Vec<(MessageId,)> = q.fetch_all(&self.pool).await?;
1273 safe_to_delete.extend(rows.into_iter().map(|(id,)| id));
1274 }
1275 Ok(safe_to_delete)
1276 }
1277
1278 pub async fn mark_qdrant_cleaned(
1284 &self,
1285 ids: &[MessageId],
1286 ) -> Result<(), crate::error::MemoryError> {
1287 if ids.is_empty() {
1288 return Ok(());
1289 }
1290 for chunk in ids.chunks(MAX_BATCH) {
1291 let placeholders = placeholder_list(1, chunk.len());
1292 let sql =
1293 format!("UPDATE messages SET qdrant_cleaned = TRUE WHERE id IN ({placeholders})");
1294 let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
1295 for &id in chunk {
1296 q = q.bind(id);
1297 }
1298 q.execute(&self.pool).await?;
1299 }
1300 Ok(())
1301 }
1302
1303 pub async fn fetch_importance_scores(
1311 &self,
1312 ids: &[MessageId],
1313 ) -> Result<std::collections::HashMap<MessageId, f64>, MemoryError> {
1314 if ids.is_empty() {
1315 return Ok(std::collections::HashMap::new());
1316 }
1317 let placeholders = zeph_db::placeholder_list(1, ids.len());
1318 let query = format!(
1319 "SELECT id, importance_score FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1320 );
1321 let mut q = zeph_db::query_as::<_, (MessageId, f64)>(sqlx::AssertSqlSafe(query));
1322 for &id in ids {
1323 q = q.bind(id);
1324 }
1325 let rows = q.fetch_all(&self.pool).await?;
1326 Ok(rows.into_iter().collect())
1327 }
1328
1329 pub async fn increment_access_counts(&self, ids: &[MessageId]) -> Result<(), MemoryError> {
1337 if ids.is_empty() {
1338 return Ok(());
1339 }
1340 let placeholders = zeph_db::placeholder_list(1, ids.len());
1341 let query = format!(
1342 "UPDATE messages SET access_count = access_count + 1, last_accessed = CURRENT_TIMESTAMP \
1343 WHERE id IN ({placeholders})"
1344 );
1345 let mut q = zeph_db::query(sqlx::AssertSqlSafe(query));
1346 for &id in ids {
1347 q = q.bind(id);
1348 }
1349 q.execute(&self.pool).await?;
1350 Ok(())
1351 }
1352
1353 #[tracing::instrument(name = "memory.store.message_access_counts", skip(self, ids), fields(count = ids.len()))]
1364 pub async fn message_access_counts(
1365 &self,
1366 ids: &[MessageId],
1367 ) -> Result<std::collections::HashMap<MessageId, i64>, MemoryError> {
1368 if ids.is_empty() {
1369 return Ok(std::collections::HashMap::new());
1370 }
1371 let placeholders = zeph_db::placeholder_list(1, ids.len());
1372 let query = format!(
1373 "SELECT id, CAST(access_count AS BIGINT) AS access_count FROM messages \
1374 WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1375 );
1376 let mut q = zeph_db::query_as::<_, (MessageId, i64)>(sqlx::AssertSqlSafe(query));
1377 for &id in ids {
1378 q = q.bind(id);
1379 }
1380 let rows = q.fetch_all(&self.pool).await?;
1381 Ok(rows.into_iter().collect())
1382 }
1383
1384 pub async fn find_promotion_candidates(
1393 &self,
1394 min_sessions: u32,
1395 batch_size: usize,
1396 ) -> Result<Vec<PromotionCandidate>, MemoryError> {
1397 let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
1398 let min = i64::from(min_sessions);
1399 let rows: Vec<(MessageId, ConversationId, String, i64, f64)> = zeph_db::query_as(sql!(
1402 "SELECT id, conversation_id, content, \
1403 CAST(session_count AS BIGINT) AS session_count, importance_score \
1404 FROM messages \
1405 WHERE tier = 'episodic' AND session_count >= ? AND deleted_at IS NULL \
1406 ORDER BY session_count DESC, importance_score DESC \
1407 LIMIT ?"
1408 ))
1409 .bind(min)
1410 .bind(limit)
1411 .fetch_all(&self.pool)
1412 .await?;
1413
1414 Ok(rows
1415 .into_iter()
1416 .map(
1417 |(id, conversation_id, content, session_count, importance_score)| {
1418 PromotionCandidate {
1419 id,
1420 conversation_id,
1421 content,
1422 session_count: session_count.try_into().unwrap_or(0),
1423 importance_score,
1424 }
1425 },
1426 )
1427 .collect())
1428 }
1429
1430 pub async fn count_messages_by_tier(&self) -> Result<(i64, i64), MemoryError> {
1438 let rows: Vec<(String, i64)> = zeph_db::query_as(sql!(
1439 "SELECT tier, COUNT(*) FROM messages \
1440 WHERE deleted_at IS NULL AND tier IN ('episodic', 'semantic') \
1441 GROUP BY tier"
1442 ))
1443 .fetch_all(&self.pool)
1444 .await?;
1445
1446 let mut episodic = 0i64;
1447 let mut semantic = 0i64;
1448 for (tier, count) in rows {
1449 match tier.as_str() {
1450 "episodic" => episodic = count,
1451 "semantic" => semantic = count,
1452 _ => {}
1453 }
1454 }
1455 Ok((episodic, semantic))
1456 }
1457
1458 pub async fn count_semantic_facts(&self) -> Result<i64, MemoryError> {
1464 let row: (i64,) = zeph_db::query_as(sql!(
1465 "SELECT COUNT(*) FROM messages WHERE tier = 'semantic' AND deleted_at IS NULL"
1466 ))
1467 .fetch_one(&self.pool)
1468 .await?;
1469 Ok(row.0)
1470 }
1471
1472 pub async fn promote_to_semantic(
1485 &self,
1486 conversation_id: ConversationId,
1487 merged_content: &str,
1488 original_ids: &[MessageId],
1489 ) -> Result<MessageId, MemoryError> {
1490 if original_ids.is_empty() {
1491 return Err(MemoryError::InvalidInput(
1492 "promote_to_semantic: original_ids must not be empty".into(),
1493 ));
1494 }
1495
1496 let mut tx = begin_write(&self.pool).await?;
1499
1500 let epoch_now = <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
1502 let promote_insert_raw = format!(
1503 "INSERT INTO messages \
1504 (conversation_id, role, content, parts, visibility, \
1505 tier, promotion_timestamp) \
1506 VALUES (?, 'assistant', ?, '[]', 'agent_only', 'semantic', {epoch_now}) \
1507 RETURNING id"
1508 );
1509 let promote_insert_sql = zeph_db::rewrite_placeholders(&promote_insert_raw);
1510 let row: (MessageId,) = zeph_db::query_as(sqlx::AssertSqlSafe(promote_insert_sql))
1511 .bind(conversation_id)
1512 .bind(merged_content)
1513 .fetch_one(&mut *tx)
1514 .await?;
1515
1516 let new_id = row.0;
1517
1518 for &id in original_ids {
1520 zeph_db::query(sql!(
1521 "UPDATE messages \
1522 SET deleted_at = CURRENT_TIMESTAMP, qdrant_cleaned = FALSE \
1523 WHERE id = ? AND deleted_at IS NULL"
1524 ))
1525 .bind(id)
1526 .execute(&mut *tx)
1527 .await?;
1528 }
1529
1530 tx.commit().await?;
1531 Ok(new_id)
1532 }
1533
1534 pub async fn manual_promote(&self, ids: &[MessageId]) -> Result<usize, MemoryError> {
1543 if ids.is_empty() {
1544 return Ok(0);
1545 }
1546 let epoch_now = <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
1547 let mut count = 0usize;
1548 for chunk in ids.chunks(MAX_BATCH) {
1549 let placeholders = placeholder_list(1, chunk.len());
1550 let manual_promote_sql = format!(
1551 "UPDATE messages \
1552 SET tier = 'semantic', promotion_timestamp = {epoch_now} \
1553 WHERE id IN ({placeholders}) AND deleted_at IS NULL AND tier = 'episodic'"
1554 );
1555 let mut q = zeph_db::query(sqlx::AssertSqlSafe(manual_promote_sql));
1556 for &id in chunk {
1557 q = q.bind(id);
1558 }
1559 let result = q.execute(&self.pool).await?;
1560 count += usize::try_from(result.rows_affected()).unwrap_or(0);
1561 }
1562 Ok(count)
1563 }
1564
1565 pub async fn increment_session_counts_for_conversation(
1574 &self,
1575 conversation_id: ConversationId,
1576 ) -> Result<(), MemoryError> {
1577 zeph_db::query(sql!(
1578 "UPDATE messages SET session_count = session_count + 1 \
1579 WHERE conversation_id = ? AND tier = 'episodic' AND deleted_at IS NULL"
1580 ))
1581 .bind(conversation_id)
1582 .execute(&self.pool)
1583 .await?;
1584 Ok(())
1585 }
1586
1587 #[tracing::instrument(name = "memory.store.fetch_tiers", skip(self, ids), fields(count = ids.len()))]
1595 pub async fn fetch_tiers(
1596 &self,
1597 ids: &[MessageId],
1598 ) -> Result<std::collections::HashMap<MessageId, String>, MemoryError> {
1599 if ids.is_empty() {
1600 return Ok(std::collections::HashMap::new());
1601 }
1602 let placeholders = zeph_db::placeholder_list(1, ids.len());
1603 let query = format!(
1604 "SELECT id, tier FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
1605 );
1606 let mut q = zeph_db::query_as::<_, (MessageId, String)>(sqlx::AssertSqlSafe(query));
1607 for &id in ids {
1608 q = q.bind(id);
1609 }
1610 let rows = q.fetch_all(&self.pool).await?;
1611 Ok(rows.into_iter().collect())
1612 }
1613
1614 pub async fn conversations_with_unconsolidated_messages(
1622 &self,
1623 ) -> Result<Vec<ConversationId>, MemoryError> {
1624 let rows: Vec<(ConversationId,)> = zeph_db::query_as(sql!(
1625 "SELECT DISTINCT conversation_id FROM messages \
1626 WHERE consolidated = FALSE AND deleted_at IS NULL"
1627 ))
1628 .fetch_all(&self.pool)
1629 .await?;
1630 Ok(rows.into_iter().map(|(id,)| id).collect())
1631 }
1632
1633 pub async fn find_unconsolidated_messages(
1642 &self,
1643 conversation_id: ConversationId,
1644 limit: usize,
1645 ) -> Result<Vec<(MessageId, String)>, MemoryError> {
1646 let limit = i64::try_from(limit).unwrap_or(i64::MAX);
1647 let rows: Vec<(MessageId, String)> = zeph_db::query_as(sql!(
1648 "SELECT id, content FROM messages \
1649 WHERE conversation_id = ? \
1650 AND consolidated = FALSE \
1651 AND deleted_at IS NULL \
1652 ORDER BY id ASC \
1653 LIMIT ?"
1654 ))
1655 .bind(conversation_id)
1656 .bind(limit)
1657 .fetch_all(&self.pool)
1658 .await?;
1659 Ok(rows)
1660 }
1661
1662 pub async fn find_consolidated_for_source(
1671 &self,
1672 source_id: MessageId,
1673 ) -> Result<Option<MessageId>, MemoryError> {
1674 let row: Option<(MessageId,)> = zeph_db::query_as(sql!(
1675 "SELECT consolidated_id FROM memory_consolidation_sources \
1676 WHERE source_id = ? \
1677 LIMIT 1"
1678 ))
1679 .bind(source_id)
1680 .fetch_optional(&self.pool)
1681 .await?;
1682 Ok(row.map(|(id,)| id))
1683 }
1684
1685 pub async fn apply_consolidation_merge(
1699 &self,
1700 conversation_id: ConversationId,
1701 role: &str,
1702 merged_content: &str,
1703 source_ids: &[MessageId],
1704 confidence: f32,
1705 confidence_threshold: f32,
1706 ) -> Result<bool, MemoryError> {
1707 if confidence < confidence_threshold {
1708 return Ok(false);
1709 }
1710 if source_ids.is_empty() {
1711 return Ok(false);
1712 }
1713
1714 let mut tx = self.pool.begin().await?;
1715
1716 let importance = crate::semantic::importance::compute_importance(merged_content, role);
1717 let row: (MessageId,) = zeph_db::query_as(sql!(
1718 "INSERT INTO messages \
1719 (conversation_id, role, content, parts, visibility, \
1720 importance_score, consolidated, consolidation_confidence) \
1721 VALUES (?, ?, ?, '[]', 'both', ?, TRUE, ?) \
1722 RETURNING id"
1723 ))
1724 .bind(conversation_id)
1725 .bind(role)
1726 .bind(merged_content)
1727 .bind(importance)
1728 .bind(confidence)
1729 .fetch_one(&mut *tx)
1730 .await?;
1731 let consolidated_id = row.0;
1732
1733 for chunk in source_ids.chunks(MAX_BATCH) {
1734 let values_list: String = chunk
1735 .iter()
1736 .map(|_| "(?, ?)")
1737 .collect::<Vec<_>>()
1738 .join(", ");
1739 let consol_sql = zeph_db::rewrite_placeholders(&format!(
1740 "{} INTO memory_consolidation_sources (consolidated_id, source_id) VALUES {values_list}{}",
1741 <ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
1742 <ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
1743 ));
1744 let mut insert_q = zeph_db::query(sqlx::AssertSqlSafe(consol_sql));
1745 for &source_id in chunk {
1746 insert_q = insert_q.bind(consolidated_id).bind(source_id);
1747 }
1748 insert_q.execute(&mut *tx).await?;
1749
1750 let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
1752 let update_sql = zeph_db::rewrite_placeholders(&format!(
1753 "UPDATE messages SET consolidated = TRUE WHERE id IN ({in_list})"
1754 ));
1755 let mut update_q = zeph_db::query(sqlx::AssertSqlSafe(update_sql));
1756 for &source_id in chunk {
1757 update_q = update_q.bind(source_id);
1758 }
1759 update_q.execute(&mut *tx).await?;
1760 }
1761
1762 tx.commit().await?;
1763 Ok(true)
1764 }
1765
1766 pub async fn apply_consolidation_update(
1779 &self,
1780 target_id: MessageId,
1781 new_content: &str,
1782 additional_source_ids: &[MessageId],
1783 confidence: f32,
1784 confidence_threshold: f32,
1785 ) -> Result<bool, MemoryError> {
1786 if confidence < confidence_threshold {
1787 return Ok(false);
1788 }
1789
1790 let mut tx = self.pool.begin().await?;
1791
1792 zeph_db::query(sql!(
1793 "UPDATE messages SET content = ?, consolidation_confidence = ?, consolidated = TRUE WHERE id = ?"
1794 ))
1795 .bind(new_content)
1796 .bind(confidence)
1797 .bind(target_id)
1798 .execute(&mut *tx)
1799 .await?;
1800
1801 for chunk in additional_source_ids.chunks(MAX_BATCH) {
1802 let values_list: String = chunk
1803 .iter()
1804 .map(|_| "(?, ?)")
1805 .collect::<Vec<_>>()
1806 .join(", ");
1807 let consol_sql = zeph_db::rewrite_placeholders(&format!(
1808 "{} INTO memory_consolidation_sources (consolidated_id, source_id) VALUES {values_list}{}",
1809 <ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
1810 <ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
1811 ));
1812 let mut insert_q = zeph_db::query(sqlx::AssertSqlSafe(consol_sql));
1813 for &source_id in chunk {
1814 insert_q = insert_q.bind(target_id).bind(source_id);
1815 }
1816 insert_q.execute(&mut *tx).await?;
1817
1818 let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
1819 let update_sql = zeph_db::rewrite_placeholders(&format!(
1820 "UPDATE messages SET consolidated = TRUE WHERE id IN ({in_list})"
1821 ));
1822 let mut update_q = zeph_db::query(sqlx::AssertSqlSafe(update_sql));
1823 for &source_id in chunk {
1824 update_q = update_q.bind(source_id);
1825 }
1826 update_q.execute(&mut *tx).await?;
1827 }
1828
1829 tx.commit().await?;
1830 Ok(true)
1831 }
1832
1833 pub async fn set_importance_score(&self, id: MessageId, score: f64) -> Result<(), MemoryError> {
1843 zeph_db::query(sql!(
1844 "UPDATE messages SET importance_score = ? WHERE id = ? AND deleted_at IS NULL"
1845 ))
1846 .bind(score)
1847 .bind(id)
1848 .execute(&self.pool)
1849 .await?;
1850 Ok(())
1851 }
1852
1853 pub async fn get_importance_score(&self, id: MessageId) -> Result<Option<f64>, MemoryError> {
1861 let row: Option<(f64,)> = zeph_db::query_as(sql!(
1862 "SELECT importance_score FROM messages WHERE id = ? AND deleted_at IS NULL"
1863 ))
1864 .bind(id)
1865 .fetch_optional(&self.pool)
1866 .await?;
1867 Ok(row.map(|(s,)| s))
1868 }
1869
1870 pub async fn batch_increment_access_count(&self, ids: &[MessageId]) -> Result<(), MemoryError> {
1878 self.increment_access_counts(ids).await
1879 }
1880
1881 pub async fn mark_messages_consolidated(&self, ids: &[i64]) -> Result<(), MemoryError> {
1889 if ids.is_empty() {
1890 return Ok(());
1891 }
1892 for chunk in ids.chunks(MAX_BATCH) {
1893 let placeholders = placeholder_list(1, chunk.len());
1894 let sql =
1895 format!("UPDATE messages SET consolidated = TRUE WHERE id IN ({placeholders})");
1896 let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
1897 for &id in chunk {
1898 q = q.bind(id);
1899 }
1900 q.execute(&self.pool).await?;
1901 }
1902 Ok(())
1903 }
1904
1905 pub async fn run_forgetting_sweep_tx(
1921 &self,
1922 config: &zeph_common::config::memory::ForgettingConfig,
1923 ) -> Result<crate::forgetting::ForgettingResult, MemoryError> {
1924 let mut tx = self.pool.begin().await?;
1925
1926 let decay = f64::from(config.decay_rate);
1927 let floor = f64::from(config.forgetting_floor);
1928 let batch = i64::try_from(config.sweep_batch_size).unwrap_or(i64::MAX);
1929 let replay_window_secs = i64::from(config.replay_window_hours) * 3600;
1930 let replay_min_access = i64::from(config.replay_min_access_count);
1931 let protect_window_secs = i64::from(config.protect_recent_hours) * 3600;
1932 let protect_min_access = i64::from(config.protect_min_access_count);
1933
1934 let candidate_ids: Vec<(MessageId,)> = zeph_db::query_as(sql!(
1937 "SELECT id FROM messages \
1938 WHERE deleted_at IS NULL AND consolidated = FALSE \
1939 ORDER BY importance_score ASC \
1940 LIMIT ?"
1941 ))
1942 .bind(batch)
1943 .fetch_all(&mut *tx)
1944 .await?;
1945
1946 #[allow(clippy::cast_possible_truncation)]
1947 let downscaled = candidate_ids.len() as u32;
1948
1949 if downscaled > 0 {
1950 let placeholders = zeph_db::placeholder_list(1, candidate_ids.len());
1951 let downscale_sql = format!(
1952 "UPDATE messages SET importance_score = importance_score * (1.0 - {decay}) \
1953 WHERE id IN ({placeholders})"
1954 );
1955 let mut q = zeph_db::query(sqlx::AssertSqlSafe(downscale_sql));
1956 for &(id,) in &candidate_ids {
1957 q = q.bind(id);
1958 }
1959 q.execute(&mut *tx).await?;
1960 }
1961
1962 let replayed = if downscaled > 0 {
1968 let replay_placeholders: String = candidate_ids
1969 .iter()
1970 .map(|_| "?")
1971 .collect::<Vec<_>>()
1972 .join(",");
1973 let epoch_now = <ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
1974 let last_accessed_epoch =
1975 <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("last_accessed");
1976 let least_fn = <ActiveDialect as zeph_db::dialect::Dialect>::LEAST_FN;
1977 let replay_sql = zeph_db::rewrite_placeholders(&format!(
1978 "UPDATE messages \
1979 SET importance_score = {least_fn}(1.0, importance_score / (1.0 - {decay})) \
1980 WHERE id IN ({replay_placeholders}) \
1981 AND (\
1982 (last_accessed IS NOT NULL \
1983 AND {last_accessed_epoch} >= {epoch_now} - ?) \
1984 OR access_count >= ?\
1985 )"
1986 ));
1987 let mut rq = zeph_db::query(sqlx::AssertSqlSafe(replay_sql));
1988 for &(id,) in &candidate_ids {
1989 rq = rq.bind(id);
1990 }
1991 let replay_result = rq
1992 .bind(replay_window_secs)
1993 .bind(replay_min_access)
1994 .execute(&mut *tx)
1995 .await?;
1996 #[allow(clippy::cast_possible_truncation)]
1997 let n = replay_result.rows_affected() as u32;
1998 n
1999 } else {
2000 0
2001 };
2002
2003 let epoch_now = <ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
2005 let last_accessed_epoch =
2006 <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("last_accessed");
2007 let prune_sql = zeph_db::rewrite_placeholders(&format!(
2008 "UPDATE messages \
2009 SET deleted_at = CURRENT_TIMESTAMP \
2010 WHERE deleted_at IS NULL AND consolidated = FALSE \
2011 AND importance_score < {floor} \
2012 AND (\
2013 last_accessed IS NULL \
2014 OR {last_accessed_epoch} < {epoch_now} - ?\
2015 ) \
2016 AND access_count < ?"
2017 ));
2018 let prune_result = zeph_db::query(sqlx::AssertSqlSafe(prune_sql))
2019 .bind(protect_window_secs)
2020 .bind(protect_min_access)
2021 .execute(&mut *tx)
2022 .await?;
2023 #[allow(clippy::cast_possible_truncation)]
2024 let pruned = prune_result.rows_affected() as u32;
2025
2026 tx.commit().await?;
2027
2028 Ok(crate::forgetting::ForgettingResult {
2029 downscaled,
2030 replayed,
2031 pruned,
2032 })
2033 }
2034}
2035
2036#[derive(Debug, Clone)]
2038pub struct PromotionCandidate {
2039 pub id: MessageId,
2040 pub conversation_id: ConversationId,
2041 pub content: String,
2042 pub session_count: u32,
2043 pub importance_score: f64,
2044}
2045
2046#[cfg(test)]
2047mod tests;