zeph_memory/store/
compression_guidelines.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! SQLite-backed store for ACON compression guidelines and failure pairs.
5
6use std::borrow::Cow;
7use std::sync::LazyLock;
8#[allow(unused_imports)]
9use zeph_db::sql;
10
11use regex::Regex;
12use zeph_common::secrets::{BEARER_TOKEN_PATTERN, JWT_PATTERN, PATH_PREFIXES, SECRET_PREFIXES};
13use zeph_common::text::truncate_to_bytes_ref;
14
15use crate::error::MemoryError;
16use crate::store::SqliteStore;
17use crate::types::ConversationId;
18
19// Prefixes come from zeph_common::secrets::SECRET_PREFIXES (the canonical, unescaped
20// list shared with zeph-core::redact — see #5917) and are regex-escaped here since e.g.
21// `ya29.` contains a literal dot.
22static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| {
23    let alt = SECRET_PREFIXES
24        .iter()
25        .map(|p| regex::escape(p))
26        .collect::<Vec<_>>()
27        .join("|");
28    Regex::new(&format!(r#"(?:{alt})[^\s"'`,;\{{\}}\[\]]*"#)).expect("secret regex")
29});
30
31static PATH_RE: LazyLock<Regex> = LazyLock::new(|| {
32    let alt = PATH_PREFIXES.join("|");
33    Regex::new(&format!(r#"(?:{alt})[^\s"'`,;\{{\}}\[\]]*"#)).expect("path regex")
34});
35
36/// Matches `Authorization: Bearer <token>` headers; captures the token value for redaction.
37static BEARER_RE: LazyLock<Regex> =
38    LazyLock::new(|| Regex::new(BEARER_TOKEN_PATTERN).expect("bearer regex"));
39
40/// Matches standalone JWT tokens (three Base64url-encoded parts separated by dots).
41/// The signature segment uses `*` to handle `alg=none` JWTs with an empty signature.
42static JWT_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(JWT_PATTERN).expect("jwt regex"));
43
44/// Redact secrets and filesystem paths from text before persistent storage.
45///
46/// Returns `Cow::Borrowed` when no sensitive content is found (zero-alloc fast path).
47pub(crate) fn redact_sensitive(text: &str) -> Cow<'_, str> {
48    // Each replace_all may return Cow::Borrowed (no match) or Cow::Owned (replaced).
49    // We materialise intermediate Owned values into String so that subsequent steps
50    // do not hold a borrow of a local.
51    let s0: Cow<'_, str> = SECRET_RE.replace_all(text, "[REDACTED]");
52    let s1: Cow<'_, str> = match PATH_RE.replace_all(s0.as_ref(), "[PATH]") {
53        Cow::Borrowed(_) => s0,
54        Cow::Owned(o) => Cow::Owned(o),
55    };
56    // Replace only the token value in Bearer headers, keeping the header name intact.
57    let s2: Cow<'_, str> = match BEARER_RE.replace_all(s1.as_ref(), "${1}[REDACTED]") {
58        Cow::Borrowed(_) => s1,
59        Cow::Owned(o) => Cow::Owned(o),
60    };
61    match JWT_RE.replace_all(s2.as_ref(), "[REDACTED_JWT]") {
62        Cow::Borrowed(_) => s2,
63        Cow::Owned(o) => Cow::Owned(o),
64    }
65}
66
67/// A recorded compression failure pair: the compressed context and the response
68/// that indicated context was lost.
69#[derive(Debug, Clone)]
70pub struct CompressionFailurePair {
71    pub id: i64,
72    pub conversation_id: ConversationId,
73    pub compressed_context: String,
74    pub failure_reason: String,
75    pub category: String,
76    pub created_at: String,
77}
78
79/// Maximum characters stored per `compressed_context` or `failure_reason` field.
80const MAX_FIELD_CHARS: usize = 4096;
81
82fn truncate_field(s: &str) -> &str {
83    truncate_to_bytes_ref(s, MAX_FIELD_CHARS)
84}
85
86impl SqliteStore {
87    /// Load the latest active compression guidelines.
88    ///
89    /// When `conversation_id` is `Some`, returns conversation-specific guidelines
90    /// preferred over global (NULL) ones. When `None`, returns only global guidelines.
91    ///
92    /// Returns `(version, guidelines_text)`. Returns `(0, "")` if no guidelines exist yet.
93    ///
94    /// # Errors
95    ///
96    /// Returns an error if the database query fails.
97    pub async fn load_compression_guidelines(
98        &self,
99        conversation_id: Option<ConversationId>,
100    ) -> Result<(i64, String), MemoryError> {
101        // `version` is `INTEGER` (`INT4`) on Postgres, so it decodes as `i32`, not `i64`;
102        // widened back to `i64` below to keep this function's public return type unchanged.
103        let row = zeph_db::query_as::<_, (i32, String)>(sql!(
104            // When conversation_id is Some(cid): `conversation_id = cid` matches
105            // conversation-specific rows; `conversation_id IS NULL` matches global rows.
106            // The CASE ensures conversation-specific rows sort before global ones.
107            // When conversation_id is None: `conversation_id = NULL` is always false in SQL,
108            // so only `conversation_id IS NULL` rows match — correct global-only behavior.
109            "SELECT version, guidelines FROM compression_guidelines \
110             WHERE conversation_id = ? OR conversation_id IS NULL \
111             ORDER BY CASE WHEN conversation_id IS NOT NULL THEN 0 ELSE 1 END, \
112                      version DESC \
113             LIMIT 1"
114        ))
115        .bind(conversation_id.map(|c| c.0))
116        .fetch_optional(&self.pool)
117        .await?;
118
119        Ok(row.map_or((0, String::new()), |(version, guidelines)| {
120            (i64::from(version), guidelines)
121        }))
122    }
123
124    /// Load only the version and creation timestamp of the latest active compression guidelines.
125    ///
126    /// Same scoping rules as [`Self::load_compression_guidelines`]: conversation-specific rows are
127    /// preferred over global ones.  Returns `(0, "")` if no guidelines exist yet.
128    ///
129    /// Use this in hot paths where the full text is not needed (e.g. metrics sync).
130    ///
131    /// # Errors
132    ///
133    /// Returns an error if the database query fails.
134    pub async fn load_compression_guidelines_meta(
135        &self,
136        conversation_id: Option<ConversationId>,
137    ) -> Result<(i64, String), MemoryError> {
138        // `created_at` is `TIMESTAMPTZ` on Postgres (`TEXT` on SQLite); project through
139        // `Dialect::select_as_text` so it decodes into the `String` tuple field below.
140        // `version` is `INTEGER` (`INT4`) on Postgres, so it decodes as `i32`, not `i64`;
141        // widened back to `i64` below to keep this function's public return type unchanged.
142        let created_at_sel =
143            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
144        let raw = format!(
145            "SELECT version, {created_at_sel} FROM compression_guidelines \
146             WHERE conversation_id = ? OR conversation_id IS NULL \
147             ORDER BY CASE WHEN conversation_id IS NOT NULL THEN 0 ELSE 1 END, \
148                      version DESC \
149             LIMIT 1"
150        );
151        let query_sql = zeph_db::rewrite_placeholders(&raw);
152        let row = zeph_db::query_as::<_, (i32, String)>(sqlx::AssertSqlSafe(query_sql))
153            .bind(conversation_id.map(|c| c.0)) // lgtm[rust/cleartext-logging]
154            .fetch_optional(&self.pool)
155            .await?;
156
157        Ok(row.map_or((0, String::new()), |(version, created_at)| {
158            (i64::from(version), created_at)
159        }))
160    }
161
162    /// Save a new version of the compression guidelines.
163    ///
164    /// When `conversation_id` is `Some`, the guidelines are scoped to that conversation.
165    /// When `None`, the guidelines are global (apply as fallback for all conversations).
166    ///
167    /// Inserts a new row; older versions are retained for audit.
168    /// Returns the new version number.
169    ///
170    /// Note: version numbers are globally sequential across all conversation scopes —
171    /// they are not per-conversation counters. The UNIQUE(version) constraint from
172    /// migration 033 is preserved.
173    ///
174    /// # Errors
175    ///
176    /// Returns an error if the database insert fails (including FK violation if
177    /// `conversation_id` does not reference a valid conversation row).
178    pub async fn save_compression_guidelines(
179        &self,
180        guidelines: &str,
181        token_count: i64,
182        conversation_id: Option<ConversationId>,
183    ) -> Result<i64, MemoryError> {
184        // The INSERT...SELECT computes MAX(version)+1 across all rows (global + per-conversation)
185        // and inserts it in a single statement. SQLite's single-writer WAL guarantee makes this
186        // atomic — no concurrent writer can observe the same MAX and produce a duplicate version.
187        // `version` is `INTEGER` (`INT4`) on Postgres, so `RETURNING version` decodes as `i32`,
188        // not `i64`; widened back to `i64` below to keep this function's public return type
189        // unchanged.
190        let new_version: i32 = zeph_db::query_scalar(
191            sql!("INSERT INTO compression_guidelines (version, guidelines, token_count, conversation_id) \
192             SELECT COALESCE(MAX(version), 0) + 1, ?, ?, ? \
193             FROM compression_guidelines \
194             RETURNING version"),
195        )
196        .bind(guidelines)
197        .bind(token_count)
198        .bind(conversation_id.map(|c| c.0))
199        .fetch_one(&self.pool)
200        .await?;
201        Ok(i64::from(new_version))
202    }
203
204    /// Log a compression failure pair.
205    ///
206    /// Both `compressed_context` and `failure_reason` are truncated to 4096 chars.
207    /// `category` should be one of: `tool_output`, `assistant_reasoning`, `user_context`, `unknown`.
208    /// Returns the inserted row id.
209    ///
210    /// # Errors
211    ///
212    /// Returns an error if the database insert fails.
213    pub async fn log_compression_failure(
214        &self,
215        conversation_id: ConversationId,
216        compressed_context: &str,
217        failure_reason: &str,
218        category: &str,
219    ) -> Result<i64, MemoryError> {
220        let ctx = redact_sensitive(compressed_context);
221        let ctx = truncate_field(&ctx);
222        let reason = redact_sensitive(failure_reason);
223        let reason = truncate_field(&reason);
224        let id = zeph_db::query_scalar(sql!(
225            "INSERT INTO compression_failure_pairs \
226             (conversation_id, compressed_context, failure_reason, category) \
227             VALUES (?, ?, ?, ?) RETURNING id"
228        ))
229        .bind(conversation_id.0)
230        .bind(ctx)
231        .bind(reason)
232        .bind(category)
233        .fetch_one(&self.pool)
234        .await?;
235        Ok(id)
236    }
237
238    /// Get unused failure pairs (oldest first), up to `limit`.
239    ///
240    /// # Errors
241    ///
242    /// Returns an error if the database query fails.
243    pub async fn get_unused_failure_pairs(
244        &self,
245        limit: usize,
246    ) -> Result<Vec<CompressionFailurePair>, MemoryError> {
247        let limit = i64::try_from(limit).unwrap_or(i64::MAX);
248        // `created_at` is `TIMESTAMPTZ` on Postgres (`TEXT` on SQLite); project through
249        // `Dialect::select_as_text` so it decodes into the `String` tuple field below.
250        // `ORDER BY` is table-qualified so it sorts on the native timestamp, not the cast text.
251        let created_at_sel =
252            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
253        let raw = format!(
254            "SELECT id, conversation_id, compressed_context, failure_reason, category, {created_at_sel} \
255             FROM compression_failure_pairs \
256             WHERE used_in_update = FALSE \
257             ORDER BY compression_failure_pairs.created_at ASC \
258             LIMIT ?"
259        );
260        let query_sql = zeph_db::rewrite_placeholders(&raw);
261        let rows = zeph_db::query_as::<_, (i64, i64, String, String, String, String)>(
262            sqlx::AssertSqlSafe(query_sql),
263        )
264        .bind(limit)
265        .fetch_all(&self.pool)
266        .await?;
267
268        Ok(rows
269            .into_iter()
270            .map(
271                |(id, cid, ctx, reason, category, created_at)| CompressionFailurePair {
272                    id,
273                    conversation_id: ConversationId(cid),
274                    compressed_context: ctx,
275                    failure_reason: reason,
276                    category,
277                    created_at,
278                },
279            )
280            .collect())
281    }
282
283    /// Get unused failure pairs for a specific category (oldest first), up to `limit`.
284    ///
285    /// Used by the categorized ACON updater to run per-category update cycles.
286    ///
287    /// # Errors
288    ///
289    /// Returns an error if the database query fails.
290    pub async fn get_unused_failure_pairs_by_category(
291        &self,
292        category: &str,
293        limit: usize,
294    ) -> Result<Vec<CompressionFailurePair>, MemoryError> {
295        let limit = i64::try_from(limit).unwrap_or(i64::MAX);
296        // `created_at` is `TIMESTAMPTZ` on Postgres — see `get_unused_failure_pairs`.
297        let created_at_sel =
298            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
299        let raw = format!(
300            "SELECT id, conversation_id, compressed_context, failure_reason, category, {created_at_sel} \
301             FROM compression_failure_pairs \
302             WHERE used_in_update = FALSE AND category = ? \
303             ORDER BY compression_failure_pairs.created_at ASC \
304             LIMIT ?"
305        );
306        let query_sql = zeph_db::rewrite_placeholders(&raw);
307        let rows = zeph_db::query_as::<_, (i64, i64, String, String, String, String)>(
308            sqlx::AssertSqlSafe(query_sql),
309        )
310        .bind(category)
311        .bind(limit)
312        .fetch_all(&self.pool)
313        .await?;
314
315        Ok(rows
316            .into_iter()
317            .map(
318                |(id, cid, ctx, reason, cat, created_at)| CompressionFailurePair {
319                    id,
320                    conversation_id: ConversationId(cid),
321                    compressed_context: ctx,
322                    failure_reason: reason,
323                    category: cat,
324                    created_at,
325                },
326            )
327            .collect())
328    }
329
330    /// Count unused failure pairs for a specific category.
331    ///
332    /// # Errors
333    ///
334    /// Returns an error if the database query fails.
335    pub async fn count_unused_failure_pairs_by_category(
336        &self,
337        category: &str,
338    ) -> Result<i64, MemoryError> {
339        let count = zeph_db::query_scalar(sql!(
340            "SELECT COUNT(*) FROM compression_failure_pairs \
341             WHERE used_in_update = FALSE AND category = ?"
342        ))
343        .bind(category)
344        .fetch_one(&self.pool)
345        .await?;
346        Ok(count)
347    }
348
349    /// Load the latest compression guidelines for a specific category.
350    ///
351    /// When `conversation_id` is `Some`, prefers conversation-specific rows.
352    /// Returns `(0, "")` if no guidelines exist for this category.
353    ///
354    /// # Errors
355    ///
356    /// Returns an error if the database query fails.
357    pub async fn load_compression_guidelines_by_category(
358        &self,
359        category: &str,
360        conversation_id: Option<ConversationId>,
361    ) -> Result<(i64, String), MemoryError> {
362        // `version` is `INTEGER` (`INT4`) on Postgres, so it decodes as `i32`, not `i64`;
363        // widened back to `i64` below to keep this function's public return type unchanged.
364        let row = zeph_db::query_as::<_, (i32, String)>(sql!(
365            "SELECT version, guidelines FROM compression_guidelines \
366             WHERE category = ? \
367             AND (conversation_id = ? OR conversation_id IS NULL) \
368             ORDER BY CASE WHEN conversation_id IS NOT NULL THEN 0 ELSE 1 END, \
369                      version DESC \
370             LIMIT 1"
371        ))
372        .bind(category)
373        .bind(conversation_id.map(|c| c.0))
374        .fetch_optional(&self.pool)
375        .await?;
376
377        Ok(row.map_or((0, String::new()), |(version, guidelines)| {
378            (i64::from(version), guidelines)
379        }))
380    }
381
382    /// Save a new version of compression guidelines for a specific category.
383    ///
384    /// # Errors
385    ///
386    /// Returns an error if the database insert fails.
387    pub async fn save_compression_guidelines_with_category(
388        &self,
389        guidelines: &str,
390        token_count: i64,
391        category: &str,
392        conversation_id: Option<ConversationId>,
393    ) -> Result<i64, MemoryError> {
394        // `version` is `INTEGER` (`INT4`) on Postgres, so `RETURNING version` decodes as `i32`,
395        // not `i64`; widened back to `i64` below to keep this function's public return type
396        // unchanged.
397        let new_version: i32 = zeph_db::query_scalar(sql!(
398            "INSERT INTO compression_guidelines \
399             (version, category, guidelines, token_count, conversation_id) \
400             SELECT COALESCE(MAX(version), 0) + 1, ?, ?, ?, ? \
401             FROM compression_guidelines \
402             RETURNING version"
403        ))
404        .bind(category)
405        .bind(guidelines)
406        .bind(token_count)
407        .bind(conversation_id.map(|c| c.0))
408        .fetch_one(&self.pool)
409        .await?;
410        Ok(i64::from(new_version))
411    }
412
413    /// Mark failure pairs as consumed by the updater.
414    ///
415    /// # Errors
416    ///
417    /// Returns an error if the database update fails.
418    pub async fn mark_failure_pairs_used(&self, ids: &[i64]) -> Result<(), MemoryError> {
419        if ids.is_empty() {
420            return Ok(());
421        }
422        let placeholders = zeph_db::placeholder_list(1, ids.len());
423        let query = format!(
424            "UPDATE compression_failure_pairs SET used_in_update = TRUE WHERE id IN ({placeholders})"
425        );
426        let mut q = zeph_db::query(sqlx::AssertSqlSafe(query));
427        for id in ids {
428            q = q.bind(id);
429        }
430        q.execute(&self.pool).await?;
431        Ok(())
432    }
433
434    /// Count unused failure pairs.
435    ///
436    /// # Errors
437    ///
438    /// Returns an error if the database query fails.
439    pub async fn count_unused_failure_pairs(&self) -> Result<i64, MemoryError> {
440        let count = zeph_db::query_scalar(sql!(
441            "SELECT COUNT(*) FROM compression_failure_pairs WHERE used_in_update = FALSE"
442        ))
443        .fetch_one(&self.pool)
444        .await?;
445        Ok(count)
446    }
447
448    /// Delete old used failure pairs, keeping the most recent `keep_recent` unused pairs.
449    ///
450    /// Removes all rows where `used_in_update = TRUE`. Unused rows are managed by the
451    /// `max_stored_pairs` enforcement below: if there are more than `keep_recent` unused pairs,
452    /// the oldest excess rows are deleted.
453    ///
454    /// # Errors
455    ///
456    /// Returns an error if the database query fails.
457    pub async fn cleanup_old_failure_pairs(&self, keep_recent: usize) -> Result<(), MemoryError> {
458        // Delete all used pairs (they've already been processed).
459        zeph_db::query(sql!(
460            "DELETE FROM compression_failure_pairs WHERE used_in_update = TRUE"
461        ))
462        .execute(&self.pool)
463        .await?;
464
465        // Keep only the most recent `keep_recent` unused pairs.
466        let keep = i64::try_from(keep_recent).unwrap_or(i64::MAX);
467        zeph_db::query(sql!(
468            "DELETE FROM compression_failure_pairs \
469             WHERE used_in_update = FALSE \
470             AND id NOT IN ( \
471                 SELECT id FROM compression_failure_pairs \
472                 WHERE used_in_update = FALSE \
473                 ORDER BY created_at DESC \
474                 LIMIT ? \
475             )"
476        ))
477        .bind(keep)
478        .execute(&self.pool)
479        .await?;
480
481        Ok(())
482    }
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488
489    // pool_size=1 is required: SQLite :memory: creates an isolated database per
490    // connection, so multiple connections would each see an empty schema.
491    async fn make_store() -> SqliteStore {
492        SqliteStore::with_pool_size(":memory:", 1)
493            .await
494            .expect("in-memory SqliteStore")
495    }
496
497    #[tokio::test]
498    async fn load_guidelines_meta_returns_defaults_when_empty() {
499        let store = make_store().await;
500        let (version, created_at) = store.load_compression_guidelines_meta(None).await.unwrap();
501        assert_eq!(version, 0);
502        assert!(created_at.is_empty());
503    }
504
505    #[tokio::test]
506    async fn load_guidelines_meta_returns_version_and_created_at() {
507        let store = make_store().await;
508        store
509            .save_compression_guidelines("keep file paths", 4, None)
510            .await
511            .unwrap();
512        let (version, created_at) = store.load_compression_guidelines_meta(None).await.unwrap();
513        assert_eq!(version, 1);
514        assert!(!created_at.is_empty(), "created_at should be populated");
515    }
516
517    #[tokio::test]
518    async fn load_guidelines_returns_defaults_when_empty() {
519        let store = make_store().await;
520        let (version, text) = store.load_compression_guidelines(None).await.unwrap();
521        assert_eq!(version, 0);
522        assert!(text.is_empty());
523    }
524
525    #[tokio::test]
526    async fn save_and_load_guidelines() {
527        let store = make_store().await;
528        let v1 = store
529            .save_compression_guidelines("always preserve file paths", 4, None)
530            .await
531            .unwrap();
532        assert_eq!(v1, 1);
533        let v2 = store
534            .save_compression_guidelines(
535                "always preserve file paths\nalways preserve errors",
536                8,
537                None,
538            )
539            .await
540            .unwrap();
541        assert_eq!(v2, 2);
542        // Loading should return the latest version.
543        let (v, text) = store.load_compression_guidelines(None).await.unwrap();
544        assert_eq!(v, 2);
545        assert!(text.contains("errors"));
546    }
547
548    #[tokio::test]
549    async fn load_guidelines_prefers_conversation_specific() {
550        let store = make_store().await;
551        let cid = ConversationId(store.create_conversation().await.unwrap().0);
552        store
553            .save_compression_guidelines("global rule", 2, None)
554            .await
555            .unwrap();
556        store
557            .save_compression_guidelines("conversation rule", 2, Some(cid))
558            .await
559            .unwrap();
560        let (_, text) = store.load_compression_guidelines(Some(cid)).await.unwrap();
561        assert_eq!(text, "conversation rule");
562    }
563
564    #[tokio::test]
565    async fn load_guidelines_falls_back_to_global() {
566        let store = make_store().await;
567        let cid = ConversationId(store.create_conversation().await.unwrap().0);
568        store
569            .save_compression_guidelines("global rule", 2, None)
570            .await
571            .unwrap();
572        // No conversation-specific guidelines; should fall back to global.
573        let (_, text) = store.load_compression_guidelines(Some(cid)).await.unwrap();
574        assert_eq!(text, "global rule");
575    }
576
577    #[tokio::test]
578    async fn load_guidelines_none_returns_global_only() {
579        let store = make_store().await;
580        let cid = ConversationId(store.create_conversation().await.unwrap().0);
581        store
582            .save_compression_guidelines("conversation rule", 2, Some(cid))
583            .await
584            .unwrap();
585        // None should not return conversation-scoped guidelines.
586        let (version, text) = store.load_compression_guidelines(None).await.unwrap();
587        assert_eq!(version, 0);
588        assert!(text.is_empty());
589    }
590
591    #[tokio::test]
592    async fn load_guidelines_scope_isolation() {
593        let store = make_store().await;
594        let cid_a = ConversationId(store.create_conversation().await.unwrap().0);
595        let cid_b = ConversationId(store.create_conversation().await.unwrap().0);
596
597        // Global guideline (conversation_id = None) — visible to all conversations.
598        store
599            .save_compression_guidelines("Use bullet points", 1, None)
600            .await
601            .unwrap();
602        // Conversation-A-specific guideline — must NOT be visible to B.
603        store
604            .save_compression_guidelines("Be concise", 2, Some(cid_a))
605            .await
606            .unwrap();
607
608        // Conversation B: gets only the global guideline, not A's.
609        let (_, text_b) = store
610            .load_compression_guidelines(Some(cid_b))
611            .await
612            .unwrap();
613        assert_eq!(
614            text_b, "Use bullet points",
615            "conversation B must see global guideline"
616        );
617
618        // Conversation A: gets its own guideline (preferred over global).
619        let (_, text_a) = store
620            .load_compression_guidelines(Some(cid_a))
621            .await
622            .unwrap();
623        assert_eq!(
624            text_a, "Be concise",
625            "conversation A must prefer its own guideline over global"
626        );
627
628        // None scope: gets only the global guideline.
629        let (_, text_global) = store.load_compression_guidelines(None).await.unwrap();
630        assert_eq!(
631            text_global, "Use bullet points",
632            "None scope must see only the global guideline"
633        );
634    }
635
636    #[tokio::test]
637    async fn save_with_nonexistent_conversation_id_fails() {
638        let store = make_store().await;
639        let nonexistent = ConversationId(99999);
640        let result = store
641            .save_compression_guidelines("rule", 1, Some(nonexistent))
642            .await;
643        assert!(
644            result.is_err(),
645            "FK violation expected for nonexistent conversation_id"
646        );
647    }
648
649    #[tokio::test]
650    async fn cascade_delete_removes_conversation_guidelines() {
651        let store = make_store().await;
652        let cid = ConversationId(store.create_conversation().await.unwrap().0);
653        store
654            .save_compression_guidelines("rule", 1, Some(cid))
655            .await
656            .unwrap();
657        // Delete the conversation row directly — should cascade-delete the guideline.
658        zeph_db::query(sql!("DELETE FROM conversations WHERE id = ?"))
659            .bind(cid.0)
660            .execute(store.pool())
661            .await
662            .unwrap();
663        let (version, text) = store.load_compression_guidelines(Some(cid)).await.unwrap();
664        assert_eq!(version, 0);
665        assert!(text.is_empty());
666    }
667
668    #[tokio::test]
669    async fn log_and_count_failure_pairs() {
670        let store = make_store().await;
671        let cid = ConversationId(store.create_conversation().await.unwrap().0);
672        store
673            .log_compression_failure(cid, "compressed ctx", "i don't recall that", "unknown")
674            .await
675            .unwrap();
676        let count = store.count_unused_failure_pairs().await.unwrap();
677        assert_eq!(count, 1);
678    }
679
680    #[tokio::test]
681    async fn get_unused_pairs_sorted_oldest_first() {
682        let store = make_store().await;
683        let cid = ConversationId(store.create_conversation().await.unwrap().0);
684        store
685            .log_compression_failure(cid, "ctx A", "reason A", "unknown")
686            .await
687            .unwrap();
688        store
689            .log_compression_failure(cid, "ctx B", "reason B", "unknown")
690            .await
691            .unwrap();
692        let pairs = store.get_unused_failure_pairs(10).await.unwrap();
693        assert_eq!(pairs.len(), 2);
694        assert_eq!(pairs[0].compressed_context, "ctx A");
695    }
696
697    #[tokio::test]
698    async fn mark_pairs_used_reduces_count() {
699        let store = make_store().await;
700        let cid = ConversationId(store.create_conversation().await.unwrap().0);
701        let id = store
702            .log_compression_failure(cid, "ctx", "reason", "unknown")
703            .await
704            .unwrap();
705        store.mark_failure_pairs_used(&[id]).await.unwrap();
706        let count = store.count_unused_failure_pairs().await.unwrap();
707        assert_eq!(count, 0);
708    }
709
710    #[tokio::test]
711    async fn cleanup_deletes_used_and_trims_unused() {
712        let store = make_store().await;
713        let cid = ConversationId(store.create_conversation().await.unwrap().0);
714        // Add 3 pairs and mark 1 used.
715        let id1 = store
716            .log_compression_failure(cid, "ctx1", "r1", "tool_output")
717            .await
718            .unwrap();
719        store
720            .log_compression_failure(cid, "ctx2", "r2", "tool_output")
721            .await
722            .unwrap();
723        store
724            .log_compression_failure(cid, "ctx3", "r3", "unknown")
725            .await
726            .unwrap();
727        store.mark_failure_pairs_used(&[id1]).await.unwrap();
728        // Cleanup: keep at most 1 unused.
729        store.cleanup_old_failure_pairs(1).await.unwrap();
730        let count = store.count_unused_failure_pairs().await.unwrap();
731        assert_eq!(count, 1, "only 1 unused pair should remain");
732    }
733
734    #[test]
735    fn redact_sensitive_api_key_is_redacted() {
736        let result = redact_sensitive("token sk-abc123def456 used for auth");
737        assert!(result.contains("[REDACTED]"), "API key must be redacted");
738        assert!(
739            !result.contains("sk-abc123"),
740            "original key must not appear"
741        );
742    }
743
744    #[test]
745    fn redact_sensitive_plain_text_borrows() {
746        let text = "safe text, no secrets here";
747        let result = redact_sensitive(text);
748        assert!(
749            matches!(result, Cow::Borrowed(_)),
750            "plain text must return Cow::Borrowed (zero-alloc)"
751        );
752    }
753
754    #[test]
755    fn redact_sensitive_filesystem_path_is_redacted() {
756        let result = redact_sensitive("config loaded from /Users/dev/project/config.toml");
757        assert!(
758            result.contains("[PATH]"),
759            "filesystem path must be redacted"
760        );
761        assert!(
762            !result.contains("/Users/dev/"),
763            "original path must not appear"
764        );
765    }
766
767    #[test]
768    fn redact_sensitive_combined_secret_and_path() {
769        let result = redact_sensitive("key sk-abc at /home/user/file");
770        assert!(result.contains("[REDACTED]"), "secret must be redacted");
771        assert!(result.contains("[PATH]"), "path must be redacted");
772    }
773
774    #[tokio::test]
775    async fn log_compression_failure_redacts_secrets() {
776        let store = make_store().await;
777        let cid = ConversationId(store.create_conversation().await.unwrap().0);
778        store
779            .log_compression_failure(
780                cid,
781                "token sk-abc123def456 used for auth",
782                "context lost",
783                "unknown",
784            )
785            .await
786            .unwrap();
787        let pairs = store.get_unused_failure_pairs(10).await.unwrap();
788        assert_eq!(pairs.len(), 1);
789        assert!(
790            pairs[0].compressed_context.contains("[REDACTED]"),
791            "stored context must have redacted secret"
792        );
793        assert!(
794            !pairs[0].compressed_context.contains("sk-abc123"),
795            "stored context must not contain raw secret"
796        );
797    }
798
799    #[tokio::test]
800    async fn log_compression_failure_redacts_paths() {
801        let store = make_store().await;
802        let cid = ConversationId(store.create_conversation().await.unwrap().0);
803        store
804            .log_compression_failure(
805                cid,
806                "/Users/dev/project/config.toml was loaded",
807                "lost",
808                "unknown",
809            )
810            .await
811            .unwrap();
812        let pairs = store.get_unused_failure_pairs(10).await.unwrap();
813        assert!(
814            pairs[0].compressed_context.contains("[PATH]"),
815            "stored context must have redacted path"
816        );
817        assert!(
818            !pairs[0].compressed_context.contains("/Users/dev/"),
819            "stored context must not contain raw path"
820        );
821    }
822
823    #[tokio::test]
824    async fn log_compression_failure_reason_also_redacted() {
825        let store = make_store().await;
826        let cid = ConversationId(store.create_conversation().await.unwrap().0);
827        store
828            .log_compression_failure(
829                cid,
830                "some context",
831                "secret ghp_abc123xyz was leaked",
832                "unknown",
833            )
834            .await
835            .unwrap();
836        let pairs = store.get_unused_failure_pairs(10).await.unwrap();
837        assert!(
838            pairs[0].failure_reason.contains("[REDACTED]"),
839            "failure_reason must also be redacted"
840        );
841        assert!(
842            !pairs[0].failure_reason.contains("ghp_abc123xyz"),
843            "raw secret must not appear in failure_reason"
844        );
845    }
846
847    #[tokio::test]
848    async fn truncate_field_respects_char_boundary() {
849        let s = "а".repeat(5000); // Cyrillic 'а', 2 bytes each
850        let truncated = truncate_field(&s);
851        assert!(truncated.len() <= MAX_FIELD_CHARS);
852        assert!(s.is_char_boundary(truncated.len()));
853    }
854
855    #[tokio::test]
856    async fn unique_constraint_prevents_duplicate_version() {
857        let store = make_store().await;
858        // Insert version 1 via the public API.
859        store
860            .save_compression_guidelines("first", 1, None)
861            .await
862            .unwrap();
863        // store.pool() access is intentional: we need direct pool access to bypass
864        // the public API and test the UNIQUE constraint at the SQL level.
865        let result = zeph_db::query(
866            sql!("INSERT INTO compression_guidelines (version, guidelines, token_count) VALUES (1, 'dup', 0)"),
867        )
868        .execute(store.pool())
869        .await;
870        assert!(
871            result.is_err(),
872            "duplicate version insert should violate UNIQUE constraint"
873        );
874    }
875
876    #[test]
877    fn redact_sensitive_bearer_token_is_redacted() {
878        let result =
879            redact_sensitive("Authorization: Bearer eyJhbGciOiJSUzI1NiJ9.payload.signature");
880        assert!(
881            result.contains("[REDACTED]"),
882            "Bearer token must be redacted: {result}"
883        );
884        assert!(
885            !result.contains("eyJhbGciOiJSUzI1NiJ9"),
886            "raw JWT header must not appear: {result}"
887        );
888        assert!(
889            result.contains("Authorization:"),
890            "header name must be preserved: {result}"
891        );
892    }
893
894    #[test]
895    fn redact_sensitive_bearer_token_case_insensitive() {
896        let result =
897            redact_sensitive("authorization: bearer eyJhbGciOiJSUzI1NiJ9.payload.signature");
898        assert!(
899            result.contains("[REDACTED]"),
900            "Bearer header match must be case-insensitive: {result}"
901        );
902    }
903
904    #[test]
905    fn redact_sensitive_standalone_jwt_is_redacted() {
906        let jwt = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.SflKxwRJSMeKKF2";
907        let input = format!("token value: {jwt} was found in logs");
908        let result = redact_sensitive(&input);
909        assert!(
910            result.contains("[REDACTED_JWT]"),
911            "standalone JWT must be replaced with [REDACTED_JWT]: {result}"
912        );
913        assert!(
914            !result.contains("eyJhbGci"),
915            "raw JWT must not appear: {result}"
916        );
917    }
918
919    #[test]
920    fn redact_sensitive_mixed_content_all_redacted() {
921        let input =
922            "key sk-abc123 at /home/user/f with Authorization: Bearer eyJhbG.pay.sig and eyJx.b.c";
923        let result = redact_sensitive(input);
924        assert!(result.contains("[REDACTED]"), "API key must be redacted");
925        assert!(result.contains("[PATH]"), "path must be redacted");
926        assert!(!result.contains("sk-abc123"), "raw API key must not appear");
927        assert!(!result.contains("eyJhbG"), "raw JWT must not appear");
928    }
929
930    #[test]
931    fn redact_sensitive_partial_jwt_not_redacted() {
932        // A string starting with eyJ but missing the third segment is not a valid JWT.
933        let input = "eyJhbGciOiJSUzI1NiJ9.onlytwoparts";
934        let result = redact_sensitive(input);
935        // Should not be replaced by the JWT regex (only two dot-separated parts).
936        assert!(
937            !result.contains("[REDACTED_JWT]"),
938            "two-part eyJ string must not be treated as JWT: {result}"
939        );
940        // No substitution occurred — must be zero-alloc Cow::Borrowed.
941        assert!(
942            matches!(result, Cow::Borrowed(_)),
943            "no-match input must return Cow::Borrowed: {result}"
944        );
945    }
946
947    #[test]
948    fn redact_sensitive_alg_none_jwt_empty_signature_redacted() {
949        // alg=none JWTs have an empty third segment: <header>.<payload>.
950        let input =
951            "token: eyJhbGciOiJub25lIn0.eyJzdWIiOiJ1c2VyIn0. was submitted without signature";
952        let result = redact_sensitive(input);
953        assert!(
954            result.contains("[REDACTED_JWT]"),
955            "alg=none JWT with empty signature must be redacted: {result}"
956        );
957        assert!(
958            !result.contains("eyJhbGciOiJub25lIn0"),
959            "raw alg=none JWT header must not appear: {result}"
960        );
961    }
962
963    // ── Category-aware store methods (MF-4) ──────────────────────────────────
964
965    #[tokio::test]
966    async fn get_unused_pairs_by_category_filters_correctly() {
967        let store = make_store().await;
968        let cid = ConversationId(store.create_conversation().await.unwrap().0);
969        store
970            .log_compression_failure(cid, "tool ctx", "lost tool output", "tool_output")
971            .await
972            .unwrap();
973        store
974            .log_compression_failure(cid, "user ctx", "lost user context", "user_context")
975            .await
976            .unwrap();
977
978        let tool_pairs = store
979            .get_unused_failure_pairs_by_category("tool_output", 10)
980            .await
981            .unwrap();
982        assert_eq!(tool_pairs.len(), 1);
983        assert_eq!(tool_pairs[0].category, "tool_output");
984        assert_eq!(tool_pairs[0].compressed_context, "tool ctx");
985
986        let user_pairs = store
987            .get_unused_failure_pairs_by_category("user_context", 10)
988            .await
989            .unwrap();
990        assert_eq!(user_pairs.len(), 1);
991        assert_eq!(user_pairs[0].category, "user_context");
992
993        // Unknown category returns nothing.
994        let unknown_pairs = store
995            .get_unused_failure_pairs_by_category("assistant_reasoning", 10)
996            .await
997            .unwrap();
998        assert!(unknown_pairs.is_empty());
999    }
1000
1001    #[tokio::test]
1002    async fn count_unused_pairs_by_category_returns_correct_count() {
1003        let store = make_store().await;
1004        let cid = ConversationId(store.create_conversation().await.unwrap().0);
1005        store
1006            .log_compression_failure(cid, "ctx A", "reason", "tool_output")
1007            .await
1008            .unwrap();
1009        store
1010            .log_compression_failure(cid, "ctx B", "reason", "tool_output")
1011            .await
1012            .unwrap();
1013        store
1014            .log_compression_failure(cid, "ctx C", "reason", "user_context")
1015            .await
1016            .unwrap();
1017
1018        let tool_count = store
1019            .count_unused_failure_pairs_by_category("tool_output")
1020            .await
1021            .unwrap();
1022        assert_eq!(tool_count, 2);
1023
1024        let user_count = store
1025            .count_unused_failure_pairs_by_category("user_context")
1026            .await
1027            .unwrap();
1028        assert_eq!(user_count, 1);
1029
1030        let unknown_count = store
1031            .count_unused_failure_pairs_by_category("assistant_reasoning")
1032            .await
1033            .unwrap();
1034        assert_eq!(unknown_count, 0);
1035    }
1036
1037    #[tokio::test]
1038    async fn save_and_load_guidelines_by_category() {
1039        let store = make_store().await;
1040        store
1041            .save_compression_guidelines_with_category(
1042                "preserve tool names",
1043                3,
1044                "tool_output",
1045                None,
1046            )
1047            .await
1048            .unwrap();
1049        store
1050            .save_compression_guidelines_with_category("keep user intent", 3, "user_context", None)
1051            .await
1052            .unwrap();
1053
1054        let (_, tool_text) = store
1055            .load_compression_guidelines_by_category("tool_output", None)
1056            .await
1057            .unwrap();
1058        assert_eq!(tool_text, "preserve tool names");
1059
1060        let (_, user_text) = store
1061            .load_compression_guidelines_by_category("user_context", None)
1062            .await
1063            .unwrap();
1064        assert_eq!(user_text, "keep user intent");
1065    }
1066
1067    #[tokio::test]
1068    async fn load_guidelines_by_category_returns_defaults_when_empty() {
1069        let store = make_store().await;
1070        // No guidelines saved for this category.
1071        let (version, text) = store
1072            .load_compression_guidelines_by_category("tool_output", None)
1073            .await
1074            .unwrap();
1075        assert_eq!(version, 0, "version must be 0 when no entries exist");
1076        assert!(text.is_empty(), "text must be empty when no entries exist");
1077    }
1078
1079    /// Concurrent saves must produce strictly unique versions with no collisions.
1080    ///
1081    /// Uses a file-backed database because `SQLite` `:memory:` creates an isolated
1082    /// database per connection — a multi-connection pool over `:memory:` would give
1083    /// each writer its own empty schema and cannot test shared-state atomicity.
1084    #[tokio::test]
1085    async fn concurrent_saves_produce_unique_versions() {
1086        use std::collections::HashSet;
1087        use std::sync::Arc;
1088
1089        let dir = tempfile::tempdir().expect("tempdir");
1090        let db_path = dir.path().join("test.db");
1091        let store = Arc::new(
1092            SqliteStore::with_pool_size(db_path.to_str().expect("utf8 path"), 4)
1093                .await
1094                .expect("file-backed SqliteStore"),
1095        );
1096
1097        let tasks: Vec<_> = (0..8_i64)
1098            .map(|i| {
1099                let s = Arc::clone(&store);
1100                let fut = async move {
1101                    s.save_compression_guidelines(&format!("guideline {i}"), i, None)
1102                        .await
1103                        .expect("concurrent save must succeed")
1104                };
1105                tokio::spawn(fut) // EXEMPT: test-only concurrent writers for UNIQUE version constraint test
1106            })
1107            .collect();
1108
1109        let mut versions = HashSet::new();
1110        for task in tasks {
1111            let v = task.await.expect("task must not panic");
1112            assert!(versions.insert(v), "version {v} appeared more than once");
1113        }
1114        assert_eq!(
1115            versions.len(),
1116            8,
1117            "all 8 saves must produce distinct versions"
1118        );
1119    }
1120}
zeph_memory/store/compression_guidelines.rs

zeph_memory/store/
compression_guidelines.rs
