use futures::TryStreamExt as _;
#[allow(unused_imports)]
use zeph_common;
use zeph_common::ContextFidelity;
use zeph_db::ActiveDialect;
use zeph_db::fts::sanitize_fts_query;
#[allow(unused_imports)]
use zeph_db::{begin_write, placeholder_list, sql};
use zeph_llm::provider::{Message, MessageMetadata, MessagePart, MessageVisibility, Role};
use super::SqliteStore;
use crate::error::MemoryError;
use crate::types::{ConversationId, MessageId};
const MAX_BATCH: usize = 490;
fn compacted_at_bind_expr() -> String {
<ActiveDialect as zeph_db::dialect::Dialect>::timestamptz_from_epoch("?")
}
fn parse_role(s: &str) -> Role {
match s {
"assistant" => Role::Assistant,
"system" => Role::System,
_ => Role::User,
}
}
type HistoryRow = (String, String, String, String, i64, i32, Option<String>);
fn parse_persisted_trust_level(s: Option<&str>) -> Option<u8> {
match s {
None => None,
Some("trusted") => Some(0),
Some("local_untrusted") => Some(1),
Some("external_untrusted") => Some(2),
Some(other) => {
tracing::warn!(
value = other,
"unrecognized persisted trust_level, treating as external_untrusted"
);
Some(2)
}
}
}
fn trust_level_ordinal_to_str(ordinal: Option<u8>) -> Option<&'static str> {
match ordinal {
None => None,
Some(0) => Some("trusted"),
Some(1) => Some("local_untrusted"),
Some(_) => Some("external_untrusted"),
}
}
#[must_use]
pub fn role_str(role: Role) -> &'static str {
match role {
Role::System => "system",
Role::Assistant => "assistant",
Role::User | _ => "user",
}
}
fn legacy_key_to_kind(key: &str) -> Option<&'static str> {
match key {
"Text" => Some("text"),
"ToolOutput" => Some("tool_output"),
"Recall" => Some("recall"),
"CodeContext" => Some("code_context"),
"Summary" => Some("summary"),
"CrossSession" => Some("cross_session"),
"ToolUse" => Some("tool_use"),
"ToolResult" => Some("tool_result"),
"Image" => Some("image"),
"ThinkingBlock" => Some("thinking_block"),
"RedactedThinkingBlock" => Some("redacted_thinking_block"),
"Compaction" => Some("compaction"),
_ => None,
}
}
fn try_parse_legacy_parts(parts_json: &str) -> Option<Vec<MessagePart>> {
let array: Vec<serde_json::Value> = serde_json::from_str(parts_json).ok()?;
let mut result = Vec::with_capacity(array.len());
for element in array {
let obj = element.as_object()?;
if obj.contains_key("kind") {
return None;
}
if obj.len() != 1 {
return None;
}
let (key, inner) = obj.iter().next()?;
let kind = legacy_key_to_kind(key)?;
let mut new_obj = match inner {
serde_json::Value::Object(m) => m.clone(),
other => {
let mut m = serde_json::Map::new();
m.insert("data".to_string(), other.clone());
m
}
};
new_obj.insert(
"kind".to_string(),
serde_json::Value::String(kind.to_string()),
);
let part: MessagePart = serde_json::from_value(serde_json::Value::Object(new_obj)).ok()?;
result.push(part);
}
Some(result)
}
fn parse_parts_json(role_str: &str, parts_json: &str) -> Vec<MessagePart> {
if parts_json == "[]" {
return vec![];
}
match serde_json::from_str(parts_json) {
Ok(p) => p,
Err(e) => {
if let Some(parts) = try_parse_legacy_parts(parts_json) {
let truncated = parts_json.chars().take(120).collect::<String>();
tracing::warn!(
role = %role_str,
parts_json = %truncated,
"loaded legacy-format message parts via compat path"
);
return parts;
}
let truncated = parts_json.chars().take(120).collect::<String>();
tracing::warn!(
role = %role_str,
parts_json = %truncated,
error = %e,
"failed to deserialize message parts, falling back to empty"
);
vec![]
}
}
}
impl SqliteStore {
pub async fn create_conversation(&self) -> Result<ConversationId, MemoryError> {
let row: (ConversationId,) = zeph_db::query_as(sql!(
"INSERT INTO conversations DEFAULT VALUES RETURNING id"
))
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn save_message(
&self,
conversation_id: ConversationId,
role: &str,
content: &str,
) -> Result<MessageId, MemoryError> {
self.save_message_with_parts(conversation_id, role, content, "[]")
.await
}
pub async fn save_message_with_parts(
&self,
conversation_id: ConversationId,
role: &str,
content: &str,
parts_json: &str,
) -> Result<MessageId, MemoryError> {
self.save_message_with_metadata(
conversation_id,
role,
content,
parts_json,
MessageVisibility::Both,
)
.await
}
pub async fn save_message_with_category(
&self,
conversation_id: ConversationId,
role: &str,
content: &str,
category: Option<&str>,
) -> Result<MessageId, MemoryError> {
let importance_score = crate::semantic::importance::compute_importance(content, role);
let row: (MessageId,) = zeph_db::query_as(sql!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, \
importance_score, category) \
VALUES (?, ?, ?, '[]', 'both', ?, ?) RETURNING id"
))
.bind(conversation_id)
.bind(role)
.bind(content)
.bind(importance_score)
.bind(category)
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn save_message_with_metadata(
&self,
conversation_id: ConversationId,
role: &str,
content: &str,
parts_json: &str,
visibility: MessageVisibility,
) -> Result<MessageId, MemoryError> {
self.save_message_with_provenance(
conversation_id,
role,
content,
parts_json,
visibility,
None,
None,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn save_message_with_provenance(
&self,
conversation_id: ConversationId,
role: &str,
content: &str,
parts_json: &str,
visibility: MessageVisibility,
source_kind: Option<&str>,
trust_level: Option<&str>,
) -> Result<MessageId, MemoryError> {
const MAX_BYTES: usize = 100 * 1024;
let content_cow: std::borrow::Cow<'_, str> = if content.len() > MAX_BYTES {
let boundary = content.floor_char_boundary(MAX_BYTES);
tracing::debug!(
original_bytes = content.len(),
"save_message: content exceeds 100KB, truncating"
);
std::borrow::Cow::Owned(format!(
"{}... [truncated, {} bytes total]",
&content[..boundary],
content.len()
))
} else {
std::borrow::Cow::Borrowed(content)
};
let importance_score = crate::semantic::importance::compute_importance(&content_cow, role);
let json_cast = <ActiveDialect as zeph_db::dialect::Dialect>::JSON_CAST;
let insert_sql = zeph_db::rewrite_placeholders(&format!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, importance_score, \
source_kind, trust_level) \
VALUES (?, ?, ?, ?{json_cast}, ?, ?, ?, ?) RETURNING id"
));
let row: (MessageId,) = zeph_db::query_as(sqlx::AssertSqlSafe(insert_sql))
.bind(conversation_id)
.bind(role)
.bind(content_cow.as_ref())
.bind(parts_json)
.bind(visibility.as_db_str())
.bind(importance_score)
.bind(source_kind)
.bind(trust_level)
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn load_history(
&self,
conversation_id: ConversationId,
limit: u32,
) -> Result<Vec<Message>, MemoryError> {
let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
let raw = format!(
"SELECT role, content, {parts_select} AS parts, visibility, id, \
CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level FROM (\
SELECT role, content, parts, visibility, id, fidelity_tag, trust_level \
FROM messages \
WHERE conversation_id = ? AND deleted_at IS NULL \
ORDER BY id DESC \
LIMIT ?\
) ORDER BY id ASC"
);
let sql = zeph_db::rewrite_placeholders(&raw);
let rows: Vec<HistoryRow> = zeph_db::query_as(sqlx::AssertSqlSafe(sql))
.bind(conversation_id)
.bind(i64::from(limit))
.fetch_all(&self.pool)
.await?;
let messages = rows
.into_iter()
.map(
|(
role_str,
content,
parts_json,
visibility_str,
row_id,
fidelity_raw,
trust_level_str,
)| {
let parts = parse_parts_json(&role_str, &parts_json);
Message {
role: parse_role(&role_str),
content,
parts,
metadata: MessageMetadata {
visibility: MessageVisibility::from_db_str(&visibility_str),
compacted_at: None,
deferred_summary: None,
focus_pinned: false,
focus_marker_id: None,
db_id: Some(row_id),
fidelity_tag: if fidelity_raw == 0 {
None
} else {
u8::try_from(fidelity_raw)
.ok()
.map(ContextFidelity::from_u8)
},
embedding: None,
trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
},
}
},
)
.collect();
Ok(messages)
}
pub async fn load_history_filtered(
&self,
conversation_id: ConversationId,
limit: u32,
agent_visible: Option<bool>,
user_visible: Option<bool>,
) -> Result<Vec<Message>, MemoryError> {
let exclude_user_only = agent_visible == Some(true);
let exclude_agent_only = user_visible == Some(true);
let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
let raw = format!(
"WITH recent AS (\
SELECT role, content, parts, visibility, id, fidelity_tag, trust_level \
FROM messages \
WHERE conversation_id = ? \
AND deleted_at IS NULL \
AND (NOT ? OR visibility != 'user_only') \
AND (NOT ? OR visibility != 'agent_only') \
ORDER BY id DESC \
LIMIT ?\
) SELECT role, content, {parts_select} AS parts, visibility, id, \
CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level \
FROM recent ORDER BY id ASC"
);
let sql = zeph_db::rewrite_placeholders(&raw);
let rows: Vec<HistoryRow> = zeph_db::query_as(sqlx::AssertSqlSafe(sql))
.bind(conversation_id)
.bind(exclude_user_only)
.bind(exclude_agent_only)
.bind(i64::from(limit))
.fetch_all(&self.pool)
.await?;
let messages = rows
.into_iter()
.map(
|(
role_str,
content,
parts_json,
visibility_str,
row_id,
fidelity_raw,
trust_level_str,
)| {
let parts = parse_parts_json(&role_str, &parts_json);
Message {
role: parse_role(&role_str),
content,
parts,
metadata: MessageMetadata {
visibility: MessageVisibility::from_db_str(&visibility_str),
compacted_at: None,
deferred_summary: None,
focus_pinned: false,
focus_marker_id: None,
db_id: Some(row_id),
fidelity_tag: if fidelity_raw == 0 {
None
} else {
u8::try_from(fidelity_raw)
.ok()
.map(ContextFidelity::from_u8)
},
embedding: None,
trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
},
}
},
)
.collect();
Ok(messages)
}
pub async fn update_fidelity_tags(
&self,
updates: &[(MessageId, u8)],
) -> Result<(), MemoryError> {
const MAX_FIDELITY_BATCH: usize = 333;
if updates.is_empty() {
return Ok(());
}
let mut tx = self.pool.begin().await?;
for chunk in updates.chunks(MAX_FIDELITY_BATCH) {
let case_arms: String = chunk
.iter()
.map(|_| "WHEN ? THEN ?")
.collect::<Vec<_>>()
.join(" ");
let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages SET fidelity_tag = CASE id {case_arms} END WHERE id IN ({in_list})"
));
let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
for &(id, tag) in chunk {
q = q.bind(id.0).bind(i32::from(tag));
}
for &(id, _) in chunk {
q = q.bind(id.0);
}
q.execute(&mut *tx).await?;
}
tx.commit().await?;
Ok(())
}
pub async fn replace_conversation(
&self,
conversation_id: ConversationId,
compacted_range: std::ops::RangeInclusive<MessageId>,
summary_role: &str,
summary_content: &str,
trust_level: Option<u8>,
) -> Result<MessageId, MemoryError> {
let now = {
let secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
format!("{secs}")
};
let start_id = compacted_range.start().0;
let end_id = compacted_range.end().0;
let mut tx = self.pool.begin().await?;
let compacted_at_expr = compacted_at_bind_expr();
let update_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages SET visibility = 'user_only', compacted_at = {compacted_at_expr} \
WHERE conversation_id = ? AND id >= ? AND id <= ?"
));
zeph_db::query(sqlx::AssertSqlSafe(update_sql))
.bind(&now)
.bind(conversation_id)
.bind(start_id)
.bind(end_id)
.execute(&mut *tx)
.await?;
let row: (MessageId,) = zeph_db::query_as(sql!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, trust_level) \
VALUES (?, ?, ?, '[]', 'agent_only', ?) RETURNING id"
))
.bind(conversation_id)
.bind(summary_role)
.bind(summary_content)
.bind(trust_level_ordinal_to_str(trust_level))
.fetch_one(&mut *tx)
.await?;
tx.commit().await?;
Ok(row.0)
}
pub async fn apply_tool_pair_summaries(
&self,
conversation_id: ConversationId,
hide_ids: &[i64],
summaries: &[String],
trust_levels: &[Option<u8>],
) -> Result<(), MemoryError> {
if hide_ids.is_empty() && summaries.is_empty() {
return Ok(());
}
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
.to_string();
let mut tx = self.pool.begin().await?;
let compacted_at_expr = compacted_at_bind_expr();
for chunk in hide_ids.chunks(MAX_BATCH) {
let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let update_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages SET visibility = 'user_only', compacted_at = {compacted_at_expr} \
WHERE id IN ({in_list})"
));
let mut q = zeph_db::query(sqlx::AssertSqlSafe(update_sql)).bind(&now);
for &id in chunk {
q = q.bind(id);
}
q.execute(&mut *tx).await?;
}
for (i, summary) in summaries.iter().enumerate() {
let content = format!("[tool summary] {summary}");
let parts = serde_json::to_string(&[MessagePart::Summary {
text: summary.clone(),
}])
.unwrap_or_else(|_| "[]".to_string());
let trust_level = trust_levels.get(i).copied().flatten();
let json_cast = <ActiveDialect as zeph_db::dialect::Dialect>::JSON_CAST;
let insert_sql = zeph_db::rewrite_placeholders(&format!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, trust_level) \
VALUES (?, 'assistant', ?, ?{json_cast}, 'agent_only', ?)"
));
zeph_db::query(sqlx::AssertSqlSafe(insert_sql))
.bind(conversation_id)
.bind(&content)
.bind(&parts)
.bind(trust_level_ordinal_to_str(trust_level))
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(())
}
pub async fn oldest_message_ids(
&self,
conversation_id: ConversationId,
n: u32,
) -> Result<Vec<MessageId>, MemoryError> {
let rows: Vec<(MessageId,)> = zeph_db::query_as(
sql!("SELECT id FROM messages WHERE conversation_id = ? AND deleted_at IS NULL ORDER BY id ASC LIMIT ?"),
)
.bind(conversation_id)
.bind(i64::from(n))
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|r| r.0).collect())
}
pub async fn latest_conversation_id(&self) -> Result<Option<ConversationId>, MemoryError> {
let row: Option<(ConversationId,)> = zeph_db::query_as(sql!(
"SELECT id FROM conversations ORDER BY id DESC LIMIT 1"
))
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|r| r.0))
}
pub async fn message_by_id(
&self,
message_id: MessageId,
) -> Result<Option<Message>, MemoryError> {
let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
let sql = zeph_db::rewrite_placeholders(&format!(
"SELECT role, content, {parts_select} AS parts, visibility, \
CAST(fidelity_tag AS INTEGER) AS fidelity_tag, trust_level FROM messages \
WHERE id = ? AND deleted_at IS NULL"
));
let row: Option<(String, String, String, String, i32, Option<String>)> =
zeph_db::query_as(sqlx::AssertSqlSafe(sql))
.bind(message_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(
|(role_str, content, parts_json, visibility_str, fidelity_raw, trust_level_str)| {
let parts = parse_parts_json(&role_str, &parts_json);
Message {
role: parse_role(&role_str),
content,
parts,
metadata: MessageMetadata {
visibility: MessageVisibility::from_db_str(&visibility_str),
compacted_at: None,
deferred_summary: None,
focus_pinned: false,
focus_marker_id: None,
db_id: Some(message_id.0),
fidelity_tag: if fidelity_raw == 0 {
None
} else {
u8::try_from(fidelity_raw)
.ok()
.map(ContextFidelity::from_u8)
},
embedding: None,
trust_level: parse_persisted_trust_level(trust_level_str.as_deref()),
},
}
},
))
}
pub async fn messages_by_ids(
&self,
ids: &[MessageId],
) -> Result<Vec<(MessageId, Message)>, MemoryError> {
if ids.is_empty() {
return Ok(Vec::new());
}
let placeholders = zeph_db::placeholder_list(1, ids.len());
let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
let query = format!(
"SELECT id, role, content, {parts_select} AS parts FROM messages \
WHERE id IN ({placeholders}) AND visibility != 'user_only' AND deleted_at IS NULL"
);
let mut q =
zeph_db::query_as::<_, (MessageId, String, String, String)>(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
let rows = q.fetch_all(&self.pool).await?;
Ok(rows
.into_iter()
.map(|(id, role_str, content, parts_json)| {
let parts = parse_parts_json(&role_str, &parts_json);
(
id,
Message {
role: parse_role(&role_str),
content,
parts,
metadata: MessageMetadata {
db_id: Some(id.0),
..MessageMetadata::default()
},
},
)
})
.collect())
}
pub async fn unembedded_message_ids(
&self,
limit: Option<usize>,
) -> Result<Vec<(MessageId, ConversationId, String, String)>, MemoryError> {
let effective_limit = limit.map_or(i64::MAX, |l| i64::try_from(l).unwrap_or(i64::MAX));
let rows: Vec<(MessageId, ConversationId, String, String)> = zeph_db::query_as(sql!(
"SELECT m.id, m.conversation_id, m.role, m.content \
FROM messages m \
LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
WHERE em.id IS NULL AND m.deleted_at IS NULL \
ORDER BY m.id ASC \
LIMIT ?"
))
.bind(effective_limit)
.fetch_all(&self.pool)
.await?;
Ok(rows)
}
pub async fn count_unembedded_messages(&self) -> Result<usize, MemoryError> {
let row: (i64,) = zeph_db::query_as(sql!(
"SELECT COUNT(*) FROM messages m \
LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
WHERE em.id IS NULL AND m.deleted_at IS NULL"
))
.fetch_one(&self.pool)
.await?;
Ok(usize::try_from(row.0).unwrap_or(usize::MAX))
}
pub fn stream_unembedded_messages(
&self,
limit: i64,
) -> impl futures::Stream<Item = Result<(MessageId, ConversationId, String, String), MemoryError>> + '_
{
zeph_db::query_as(sql!(
"SELECT m.id, m.conversation_id, m.role, m.content \
FROM messages m \
LEFT JOIN embeddings_metadata em ON m.id = em.message_id \
WHERE em.id IS NULL AND m.deleted_at IS NULL \
ORDER BY m.id ASC \
LIMIT ?"
))
.bind(limit)
.fetch(&self.pool)
.map_err(MemoryError::from)
}
pub async fn count_messages(
&self,
conversation_id: ConversationId,
) -> Result<i64, MemoryError> {
let row: (i64,) = zeph_db::query_as(sql!(
"SELECT COUNT(*) FROM messages WHERE conversation_id = ? AND deleted_at IS NULL"
))
.bind(conversation_id)
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn count_messages_after(
&self,
conversation_id: ConversationId,
after_id: MessageId,
) -> Result<i64, MemoryError> {
let row: (i64,) =
zeph_db::query_as(
sql!("SELECT COUNT(*) FROM messages WHERE conversation_id = ? AND id > ? AND deleted_at IS NULL"),
)
.bind(conversation_id)
.bind(after_id)
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn keyword_search(
&self,
query: &str,
limit: usize,
conversation_id: Option<ConversationId>,
) -> Result<Vec<(MessageId, f64)>, MemoryError> {
let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
let safe_query = sanitize_fts_query(query);
if safe_query.is_empty() {
return Ok(Vec::new());
}
let rows: Vec<(MessageId, f64)> = if let Some(cid) = conversation_id {
zeph_db::query_as(
sql!("SELECT m.id, -rank AS score \
FROM messages_fts f \
JOIN messages m ON m.id = f.rowid \
WHERE messages_fts MATCH ? AND m.conversation_id = ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL \
ORDER BY rank \
LIMIT ?"),
)
.bind(&safe_query)
.bind(cid)
.bind(effective_limit)
.fetch_all(&self.pool)
.await?
} else {
zeph_db::query_as(sql!(
"SELECT m.id, -rank AS score \
FROM messages_fts f \
JOIN messages m ON m.id = f.rowid \
WHERE messages_fts MATCH ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL \
ORDER BY rank \
LIMIT ?"
))
.bind(&safe_query)
.bind(effective_limit)
.fetch_all(&self.pool)
.await?
};
Ok(rows)
}
pub async fn keyword_search_with_time_range(
&self,
query: &str,
limit: usize,
conversation_id: Option<ConversationId>,
after: Option<&str>,
before: Option<&str>,
) -> Result<Vec<(MessageId, f64)>, MemoryError> {
let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
let safe_query = sanitize_fts_query(query);
if safe_query.is_empty() {
return Ok(Vec::new());
}
let after_clause = if after.is_some() {
" AND m.created_at > ?"
} else {
""
};
let before_clause = if before.is_some() {
" AND m.created_at < ?"
} else {
""
};
let conv_clause = if conversation_id.is_some() {
" AND m.conversation_id = ?"
} else {
""
};
let sql = zeph_db::rewrite_placeholders(&format!(
"SELECT m.id, -rank AS score \
FROM messages_fts f \
JOIN messages m ON m.id = f.rowid \
WHERE messages_fts MATCH ? AND m.visibility != 'user_only' AND m.deleted_at IS NULL\
{after_clause}{before_clause}{conv_clause} \
ORDER BY rank \
LIMIT ?"
));
let mut q =
zeph_db::query_as::<_, (MessageId, f64)>(sqlx::AssertSqlSafe(sql)).bind(&safe_query);
if let Some(a) = after {
q = q.bind(a);
}
if let Some(b) = before {
q = q.bind(b);
}
if let Some(cid) = conversation_id {
q = q.bind(cid);
}
q = q.bind(effective_limit);
Ok(q.fetch_all(&self.pool).await?)
}
#[tracing::instrument(name = "memory.store.message_timestamps", skip(self, ids), fields(count = ids.len()))]
pub async fn message_timestamps(
&self,
ids: &[MessageId],
) -> Result<std::collections::HashMap<MessageId, i64>, MemoryError> {
if ids.is_empty() {
return Ok(std::collections::HashMap::new());
}
let placeholders: String =
zeph_db::rewrite_placeholders(&ids.iter().map(|_| "?").collect::<Vec<_>>().join(","));
let epoch_expr = <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
let query = format!(
"SELECT id, {epoch_expr} FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
);
let mut q = zeph_db::query_as::<_, (MessageId, i64)>(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
let rows = q.fetch_all(&self.pool).await?;
Ok(rows.into_iter().collect())
}
pub async fn load_messages_range(
&self,
conversation_id: ConversationId,
after_message_id: MessageId,
limit: usize,
) -> Result<Vec<(MessageId, String, String)>, MemoryError> {
let effective_limit = i64::try_from(limit).unwrap_or(i64::MAX);
let rows: Vec<(MessageId, String, String)> = zeph_db::query_as(sql!(
"SELECT id, role, content FROM messages \
WHERE conversation_id = ? AND id > ? AND deleted_at IS NULL \
ORDER BY id ASC LIMIT ?"
))
.bind(conversation_id)
.bind(after_message_id)
.bind(effective_limit)
.fetch_all(&self.pool)
.await?;
Ok(rows)
}
pub async fn get_eviction_candidates(
&self,
) -> Result<Vec<crate::eviction::EvictionEntry>, crate::error::MemoryError> {
let created_at_sel =
<ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
let last_accessed_sel =
<ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("last_accessed");
let raw = format!(
"SELECT id, {created_at_sel}, {last_accessed_sel}, access_count \
FROM messages WHERE deleted_at IS NULL"
);
let query_sql = zeph_db::rewrite_placeholders(&raw);
let rows: Vec<(MessageId, String, Option<String>, i32)> =
zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
.fetch_all(&self.pool)
.await?;
Ok(rows
.into_iter()
.map(
|(id, created_at, last_accessed, access_count)| crate::eviction::EvictionEntry {
id,
created_at,
last_accessed,
access_count: access_count.try_into().unwrap_or(0),
},
)
.collect())
}
pub async fn soft_delete_messages(
&self,
ids: &[MessageId],
) -> Result<(), crate::error::MemoryError> {
if ids.is_empty() {
return Ok(());
}
for chunk in ids.chunks(MAX_BATCH) {
let placeholders = placeholder_list(1, chunk.len());
let sql = format!(
"UPDATE messages SET deleted_at = CURRENT_TIMESTAMP \
WHERE id IN ({placeholders}) AND deleted_at IS NULL"
);
let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
for &id in chunk {
q = q.bind(id);
}
q.execute(&self.pool).await?;
}
Ok(())
}
pub async fn get_soft_deleted_message_ids(
&self,
) -> Result<Vec<MessageId>, crate::error::MemoryError> {
let rows: Vec<(MessageId,)> = zeph_db::query_as(sql!(
"SELECT id FROM messages WHERE deleted_at IS NOT NULL AND qdrant_cleaned = FALSE"
))
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|(id,)| id).collect())
}
pub async fn filter_out_preserved_episode_ids(
&self,
candidate_ids: &[MessageId],
) -> Result<Vec<MessageId>, crate::error::MemoryError> {
if candidate_ids.is_empty() {
return Ok(Vec::new());
}
let mut safe_to_delete: Vec<MessageId> = Vec::with_capacity(candidate_ids.len());
for chunk in candidate_ids.chunks(MAX_BATCH) {
let placeholders = placeholder_list(1, chunk.len());
let sql = format!(
"SELECT m.id \
FROM messages m \
WHERE m.id IN ({placeholders}) \
AND NOT EXISTS ( \
SELECT 1 \
FROM summaries s \
WHERE s.first_message_id IS NOT NULL \
AND s.last_message_id IS NOT NULL \
AND m.id >= s.first_message_id \
AND m.id <= s.last_message_id \
)"
);
let mut q = zeph_db::query_as::<_, (MessageId,)>(sqlx::AssertSqlSafe(sql));
for &id in chunk {
q = q.bind(id);
}
let rows: Vec<(MessageId,)> = q.fetch_all(&self.pool).await?;
safe_to_delete.extend(rows.into_iter().map(|(id,)| id));
}
Ok(safe_to_delete)
}
pub async fn mark_qdrant_cleaned(
&self,
ids: &[MessageId],
) -> Result<(), crate::error::MemoryError> {
if ids.is_empty() {
return Ok(());
}
for chunk in ids.chunks(MAX_BATCH) {
let placeholders = placeholder_list(1, chunk.len());
let sql =
format!("UPDATE messages SET qdrant_cleaned = TRUE WHERE id IN ({placeholders})");
let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
for &id in chunk {
q = q.bind(id);
}
q.execute(&self.pool).await?;
}
Ok(())
}
pub async fn fetch_importance_scores(
&self,
ids: &[MessageId],
) -> Result<std::collections::HashMap<MessageId, f64>, MemoryError> {
if ids.is_empty() {
return Ok(std::collections::HashMap::new());
}
let placeholders = zeph_db::placeholder_list(1, ids.len());
let query = format!(
"SELECT id, importance_score FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
);
let mut q = zeph_db::query_as::<_, (MessageId, f64)>(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
let rows = q.fetch_all(&self.pool).await?;
Ok(rows.into_iter().collect())
}
pub async fn increment_access_counts(&self, ids: &[MessageId]) -> Result<(), MemoryError> {
if ids.is_empty() {
return Ok(());
}
let placeholders = zeph_db::placeholder_list(1, ids.len());
let query = format!(
"UPDATE messages SET access_count = access_count + 1, last_accessed = CURRENT_TIMESTAMP \
WHERE id IN ({placeholders})"
);
let mut q = zeph_db::query(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
q.execute(&self.pool).await?;
Ok(())
}
#[tracing::instrument(name = "memory.store.message_access_counts", skip(self, ids), fields(count = ids.len()))]
pub async fn message_access_counts(
&self,
ids: &[MessageId],
) -> Result<std::collections::HashMap<MessageId, i64>, MemoryError> {
if ids.is_empty() {
return Ok(std::collections::HashMap::new());
}
let placeholders = zeph_db::placeholder_list(1, ids.len());
let query = format!(
"SELECT id, CAST(access_count AS BIGINT) AS access_count FROM messages \
WHERE id IN ({placeholders}) AND deleted_at IS NULL"
);
let mut q = zeph_db::query_as::<_, (MessageId, i64)>(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
let rows = q.fetch_all(&self.pool).await?;
Ok(rows.into_iter().collect())
}
pub async fn find_promotion_candidates(
&self,
min_sessions: u32,
batch_size: usize,
) -> Result<Vec<PromotionCandidate>, MemoryError> {
let limit = i64::try_from(batch_size).unwrap_or(i64::MAX);
let min = i64::from(min_sessions);
let rows: Vec<(MessageId, ConversationId, String, i64, f64)> = zeph_db::query_as(sql!(
"SELECT id, conversation_id, content, \
CAST(session_count AS BIGINT) AS session_count, importance_score \
FROM messages \
WHERE tier = 'episodic' AND session_count >= ? AND deleted_at IS NULL \
ORDER BY session_count DESC, importance_score DESC \
LIMIT ?"
))
.bind(min)
.bind(limit)
.fetch_all(&self.pool)
.await?;
Ok(rows
.into_iter()
.map(
|(id, conversation_id, content, session_count, importance_score)| {
PromotionCandidate {
id,
conversation_id,
content,
session_count: session_count.try_into().unwrap_or(0),
importance_score,
}
},
)
.collect())
}
pub async fn count_messages_by_tier(&self) -> Result<(i64, i64), MemoryError> {
let rows: Vec<(String, i64)> = zeph_db::query_as(sql!(
"SELECT tier, COUNT(*) FROM messages \
WHERE deleted_at IS NULL AND tier IN ('episodic', 'semantic') \
GROUP BY tier"
))
.fetch_all(&self.pool)
.await?;
let mut episodic = 0i64;
let mut semantic = 0i64;
for (tier, count) in rows {
match tier.as_str() {
"episodic" => episodic = count,
"semantic" => semantic = count,
_ => {}
}
}
Ok((episodic, semantic))
}
pub async fn count_semantic_facts(&self) -> Result<i64, MemoryError> {
let row: (i64,) = zeph_db::query_as(sql!(
"SELECT COUNT(*) FROM messages WHERE tier = 'semantic' AND deleted_at IS NULL"
))
.fetch_one(&self.pool)
.await?;
Ok(row.0)
}
pub async fn promote_to_semantic(
&self,
conversation_id: ConversationId,
merged_content: &str,
original_ids: &[MessageId],
) -> Result<MessageId, MemoryError> {
if original_ids.is_empty() {
return Err(MemoryError::InvalidInput(
"promote_to_semantic: original_ids must not be empty".into(),
));
}
let mut tx = begin_write(&self.pool).await?;
let epoch_now = <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
let promote_insert_raw = format!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, \
tier, promotion_timestamp) \
VALUES (?, 'assistant', ?, '[]', 'agent_only', 'semantic', {epoch_now}) \
RETURNING id"
);
let promote_insert_sql = zeph_db::rewrite_placeholders(&promote_insert_raw);
let row: (MessageId,) = zeph_db::query_as(sqlx::AssertSqlSafe(promote_insert_sql))
.bind(conversation_id)
.bind(merged_content)
.fetch_one(&mut *tx)
.await?;
let new_id = row.0;
for &id in original_ids {
zeph_db::query(sql!(
"UPDATE messages \
SET deleted_at = CURRENT_TIMESTAMP, qdrant_cleaned = FALSE \
WHERE id = ? AND deleted_at IS NULL"
))
.bind(id)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(new_id)
}
pub async fn manual_promote(&self, ids: &[MessageId]) -> Result<usize, MemoryError> {
if ids.is_empty() {
return Ok(0);
}
let epoch_now = <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
let mut count = 0usize;
for chunk in ids.chunks(MAX_BATCH) {
let placeholders = placeholder_list(1, chunk.len());
let manual_promote_sql = format!(
"UPDATE messages \
SET tier = 'semantic', promotion_timestamp = {epoch_now} \
WHERE id IN ({placeholders}) AND deleted_at IS NULL AND tier = 'episodic'"
);
let mut q = zeph_db::query(sqlx::AssertSqlSafe(manual_promote_sql));
for &id in chunk {
q = q.bind(id);
}
let result = q.execute(&self.pool).await?;
count += usize::try_from(result.rows_affected()).unwrap_or(0);
}
Ok(count)
}
pub async fn increment_session_counts_for_conversation(
&self,
conversation_id: ConversationId,
) -> Result<(), MemoryError> {
zeph_db::query(sql!(
"UPDATE messages SET session_count = session_count + 1 \
WHERE conversation_id = ? AND tier = 'episodic' AND deleted_at IS NULL"
))
.bind(conversation_id)
.execute(&self.pool)
.await?;
Ok(())
}
#[tracing::instrument(name = "memory.store.fetch_tiers", skip(self, ids), fields(count = ids.len()))]
pub async fn fetch_tiers(
&self,
ids: &[MessageId],
) -> Result<std::collections::HashMap<MessageId, String>, MemoryError> {
if ids.is_empty() {
return Ok(std::collections::HashMap::new());
}
let placeholders = zeph_db::placeholder_list(1, ids.len());
let query = format!(
"SELECT id, tier FROM messages WHERE id IN ({placeholders}) AND deleted_at IS NULL"
);
let mut q = zeph_db::query_as::<_, (MessageId, String)>(sqlx::AssertSqlSafe(query));
for &id in ids {
q = q.bind(id);
}
let rows = q.fetch_all(&self.pool).await?;
Ok(rows.into_iter().collect())
}
pub async fn conversations_with_unconsolidated_messages(
&self,
) -> Result<Vec<ConversationId>, MemoryError> {
let rows: Vec<(ConversationId,)> = zeph_db::query_as(sql!(
"SELECT DISTINCT conversation_id FROM messages \
WHERE consolidated = FALSE AND deleted_at IS NULL"
))
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|(id,)| id).collect())
}
pub async fn find_unconsolidated_messages(
&self,
conversation_id: ConversationId,
limit: usize,
) -> Result<Vec<(MessageId, String)>, MemoryError> {
let limit = i64::try_from(limit).unwrap_or(i64::MAX);
let rows: Vec<(MessageId, String)> = zeph_db::query_as(sql!(
"SELECT id, content FROM messages \
WHERE conversation_id = ? \
AND consolidated = FALSE \
AND deleted_at IS NULL \
ORDER BY id ASC \
LIMIT ?"
))
.bind(conversation_id)
.bind(limit)
.fetch_all(&self.pool)
.await?;
Ok(rows)
}
pub async fn find_consolidated_for_source(
&self,
source_id: MessageId,
) -> Result<Option<MessageId>, MemoryError> {
let row: Option<(MessageId,)> = zeph_db::query_as(sql!(
"SELECT consolidated_id FROM memory_consolidation_sources \
WHERE source_id = ? \
LIMIT 1"
))
.bind(source_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|(id,)| id))
}
pub async fn apply_consolidation_merge(
&self,
conversation_id: ConversationId,
role: &str,
merged_content: &str,
source_ids: &[MessageId],
confidence: f32,
confidence_threshold: f32,
) -> Result<bool, MemoryError> {
if confidence < confidence_threshold {
return Ok(false);
}
if source_ids.is_empty() {
return Ok(false);
}
let mut tx = self.pool.begin().await?;
let importance = crate::semantic::importance::compute_importance(merged_content, role);
let row: (MessageId,) = zeph_db::query_as(sql!(
"INSERT INTO messages \
(conversation_id, role, content, parts, visibility, \
importance_score, consolidated, consolidation_confidence) \
VALUES (?, ?, ?, '[]', 'both', ?, TRUE, ?) \
RETURNING id"
))
.bind(conversation_id)
.bind(role)
.bind(merged_content)
.bind(importance)
.bind(confidence)
.fetch_one(&mut *tx)
.await?;
let consolidated_id = row.0;
for chunk in source_ids.chunks(MAX_BATCH) {
let values_list: String = chunk
.iter()
.map(|_| "(?, ?)")
.collect::<Vec<_>>()
.join(", ");
let consol_sql = zeph_db::rewrite_placeholders(&format!(
"{} INTO memory_consolidation_sources (consolidated_id, source_id) VALUES {values_list}{}",
<ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
<ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
));
let mut insert_q = zeph_db::query(sqlx::AssertSqlSafe(consol_sql));
for &source_id in chunk {
insert_q = insert_q.bind(consolidated_id).bind(source_id);
}
insert_q.execute(&mut *tx).await?;
let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let update_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages SET consolidated = TRUE WHERE id IN ({in_list})"
));
let mut update_q = zeph_db::query(sqlx::AssertSqlSafe(update_sql));
for &source_id in chunk {
update_q = update_q.bind(source_id);
}
update_q.execute(&mut *tx).await?;
}
tx.commit().await?;
Ok(true)
}
pub async fn apply_consolidation_update(
&self,
target_id: MessageId,
new_content: &str,
additional_source_ids: &[MessageId],
confidence: f32,
confidence_threshold: f32,
) -> Result<bool, MemoryError> {
if confidence < confidence_threshold {
return Ok(false);
}
let mut tx = self.pool.begin().await?;
zeph_db::query(sql!(
"UPDATE messages SET content = ?, consolidation_confidence = ?, consolidated = TRUE WHERE id = ?"
))
.bind(new_content)
.bind(confidence)
.bind(target_id)
.execute(&mut *tx)
.await?;
for chunk in additional_source_ids.chunks(MAX_BATCH) {
let values_list: String = chunk
.iter()
.map(|_| "(?, ?)")
.collect::<Vec<_>>()
.join(", ");
let consol_sql = zeph_db::rewrite_placeholders(&format!(
"{} INTO memory_consolidation_sources (consolidated_id, source_id) VALUES {values_list}{}",
<ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
<ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
));
let mut insert_q = zeph_db::query(sqlx::AssertSqlSafe(consol_sql));
for &source_id in chunk {
insert_q = insert_q.bind(target_id).bind(source_id);
}
insert_q.execute(&mut *tx).await?;
let in_list: String = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let update_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages SET consolidated = TRUE WHERE id IN ({in_list})"
));
let mut update_q = zeph_db::query(sqlx::AssertSqlSafe(update_sql));
for &source_id in chunk {
update_q = update_q.bind(source_id);
}
update_q.execute(&mut *tx).await?;
}
tx.commit().await?;
Ok(true)
}
pub async fn set_importance_score(&self, id: MessageId, score: f64) -> Result<(), MemoryError> {
zeph_db::query(sql!(
"UPDATE messages SET importance_score = ? WHERE id = ? AND deleted_at IS NULL"
))
.bind(score)
.bind(id)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_importance_score(&self, id: MessageId) -> Result<Option<f64>, MemoryError> {
let row: Option<(f64,)> = zeph_db::query_as(sql!(
"SELECT importance_score FROM messages WHERE id = ? AND deleted_at IS NULL"
))
.bind(id)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|(s,)| s))
}
pub async fn batch_increment_access_count(&self, ids: &[MessageId]) -> Result<(), MemoryError> {
self.increment_access_counts(ids).await
}
pub async fn mark_messages_consolidated(&self, ids: &[i64]) -> Result<(), MemoryError> {
if ids.is_empty() {
return Ok(());
}
for chunk in ids.chunks(MAX_BATCH) {
let placeholders = placeholder_list(1, chunk.len());
let sql =
format!("UPDATE messages SET consolidated = TRUE WHERE id IN ({placeholders})");
let mut q = zeph_db::query(sqlx::AssertSqlSafe(sql));
for &id in chunk {
q = q.bind(id);
}
q.execute(&self.pool).await?;
}
Ok(())
}
pub async fn run_forgetting_sweep_tx(
&self,
config: &zeph_common::config::memory::ForgettingConfig,
) -> Result<crate::forgetting::ForgettingResult, MemoryError> {
let mut tx = self.pool.begin().await?;
let decay = f64::from(config.decay_rate);
let floor = f64::from(config.forgetting_floor);
let batch = i64::try_from(config.sweep_batch_size).unwrap_or(i64::MAX);
let replay_window_secs = i64::from(config.replay_window_hours) * 3600;
let replay_min_access = i64::from(config.replay_min_access_count);
let protect_window_secs = i64::from(config.protect_recent_hours) * 3600;
let protect_min_access = i64::from(config.protect_min_access_count);
let candidate_ids: Vec<(MessageId,)> = zeph_db::query_as(sql!(
"SELECT id FROM messages \
WHERE deleted_at IS NULL AND consolidated = FALSE \
ORDER BY importance_score ASC \
LIMIT ?"
))
.bind(batch)
.fetch_all(&mut *tx)
.await?;
#[allow(clippy::cast_possible_truncation)]
let downscaled = candidate_ids.len() as u32;
if downscaled > 0 {
let placeholders = zeph_db::placeholder_list(1, candidate_ids.len());
let downscale_sql = format!(
"UPDATE messages SET importance_score = importance_score * (1.0 - {decay}) \
WHERE id IN ({placeholders})"
);
let mut q = zeph_db::query(sqlx::AssertSqlSafe(downscale_sql));
for &(id,) in &candidate_ids {
q = q.bind(id);
}
q.execute(&mut *tx).await?;
}
let replayed = if downscaled > 0 {
let replay_placeholders: String = candidate_ids
.iter()
.map(|_| "?")
.collect::<Vec<_>>()
.join(",");
let epoch_now = <ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
let last_accessed_epoch =
<ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("last_accessed");
let least_fn = <ActiveDialect as zeph_db::dialect::Dialect>::LEAST_FN;
let replay_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages \
SET importance_score = {least_fn}(1.0, importance_score / (1.0 - {decay})) \
WHERE id IN ({replay_placeholders}) \
AND (\
(last_accessed IS NOT NULL \
AND {last_accessed_epoch} >= {epoch_now} - ?) \
OR access_count >= ?\
)"
));
let mut rq = zeph_db::query(sqlx::AssertSqlSafe(replay_sql));
for &(id,) in &candidate_ids {
rq = rq.bind(id);
}
let replay_result = rq
.bind(replay_window_secs)
.bind(replay_min_access)
.execute(&mut *tx)
.await?;
#[allow(clippy::cast_possible_truncation)]
let n = replay_result.rows_affected() as u32;
n
} else {
0
};
let epoch_now = <ActiveDialect as zeph_db::dialect::Dialect>::EPOCH_NOW;
let last_accessed_epoch =
<ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("last_accessed");
let prune_sql = zeph_db::rewrite_placeholders(&format!(
"UPDATE messages \
SET deleted_at = CURRENT_TIMESTAMP \
WHERE deleted_at IS NULL AND consolidated = FALSE \
AND importance_score < {floor} \
AND (\
last_accessed IS NULL \
OR {last_accessed_epoch} < {epoch_now} - ?\
) \
AND access_count < ?"
));
let prune_result = zeph_db::query(sqlx::AssertSqlSafe(prune_sql))
.bind(protect_window_secs)
.bind(protect_min_access)
.execute(&mut *tx)
.await?;
#[allow(clippy::cast_possible_truncation)]
let pruned = prune_result.rows_affected() as u32;
tx.commit().await?;
Ok(crate::forgetting::ForgettingResult {
downscaled,
replayed,
pruned,
})
}
}
#[derive(Debug, Clone)]
pub struct PromotionCandidate {
pub id: MessageId,
pub conversation_id: ConversationId,
pub content: String,
pub session_count: u32,
pub importance_score: f64,
}
#[cfg(test)]
mod tests;