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

zeph_memory/store/
compression_guidelines.rs
