diff --git a/Cargo.lock b/Cargo.lock index b08223d..b78f197 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,6 +26,16 @@ version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -70,6 +80,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -131,6 +150,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -165,14 +190,52 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "digest" version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", ] [[package]] @@ -186,6 +249,12 @@ dependencies = [ "syn", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "equivalent" version = "1.0.2" @@ -226,6 +295,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.2.0" @@ -241,6 +316,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -248,6 +338,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -256,6 +347,40 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + [[package]] name = "futures-task" version = "0.3.32" @@ -268,8 +393,13 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] @@ -323,6 +453,25 @@ dependencies = [ "log", ] +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -337,16 +486,25 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash", +] [[package]] name = "hashlink" -version = "0.11.1" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "824e001ac4f3012dd16a264bec811403a67ca9deb6c102fc5049b32c4574b35f" +checksum = "32069d97bb81e38fa67eab65e3393bf804bb85969f2bc06bf13f64aef5aba248" dependencies = [ - "hashbrown 0.16.1", + "hashbrown 0.17.1", ] +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "http" version = "1.4.2" @@ -386,6 +544,21 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -396,9 +569,11 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -615,6 +790,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -635,9 +816,9 @@ dependencies = [ [[package]] name = "libsqlite3-sys" -version = "0.37.0" +version = "0.38.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" +checksum = "f6c19a05435c21ac299d71b6a9c13db3e3f47c520517d58990a462a1397a61db" dependencies = [ "cc", "pkg-config", @@ -715,6 +896,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -886,6 +1077,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "regex" version = "1.13.0" @@ -924,6 +1135,7 @@ dependencies = [ "base64", "bytes", "futures-core", + "futures-util", "http", "http-body", "http-body-util", @@ -943,12 +1155,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots", ] @@ -979,9 +1193,9 @@ dependencies = [ [[package]] name = "rusqlite" -version = "0.39.0" +version = "0.40.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0d2b0146dd9661bf67bb107c0bb2a55064d556eeb3fc314151b957f313bcd4e" +checksum = "11438310b19e3109b6446c33d1ed5e889428cf2e278407bc7896bc4aaea43323" dependencies = [ "bitflags", "fallible-iterator", @@ -1067,6 +1281,31 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1103,6 +1342,17 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_json" version = "1.0.150" @@ -1145,7 +1395,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -1264,6 +1525,24 @@ dependencies = [ "syn", ] +[[package]] +name = "tinyagents" +version = "1.9.0" +source = "git+https://github.com/senamakel/tinyagents?rev=8f75355#8f753557278e448c6c5bcf48ab8503505a47c663" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "reqwest", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror", + "tokio", + "tracing", +] + [[package]] name = "tinycortex" version = "0.1.1" @@ -1271,21 +1550,27 @@ dependencies = [ "anyhow", "async-trait", "chrono", + "futures", "git2", + "log", "parking_lot", "rand", "regex", "reqwest", "rusqlite", + "schemars", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "tempfile", "thiserror", + "tinyagents", "tokio", "toml", + "tracing", "uuid", "walkdir", + "wiremock", ] [[package]] @@ -1349,6 +1634,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "toml" version = "0.8.23" @@ -1442,9 +1740,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -1600,6 +1910,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.103" @@ -1788,6 +2111,29 @@ dependencies = [ "memchr", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64", + "deadpool", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/Cargo.toml b/Cargo.toml index af289ea..764944f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,10 @@ git-diff = ["dep:git2"] # gates the dependency. providers-http = ["dep:reqwest", "tokio"] +# Live source synchronization (Composio HTTP pipelines and workspace scans). +# Host schedulers, credentials, RPC, and event buses remain outside the crate. +sync = ["dep:reqwest", "dep:tracing", "tokio"] + # serde schema / envelope surface for the RPC boundary (`memory::rpc`). Reserved # for goal C5; gates the wire-facing surface without adding heavy dependencies # (serde is already a core dependency). @@ -46,16 +50,20 @@ rpc = [] [dependencies] anyhow = "1" +log = "0.4" +futures = "0.3" async-trait = "0.1" chrono = { version = "0.4", features = ["serde"] } parking_lot = "0.12" rand = "0.10" regex = "1" -rusqlite = { version = "0.39", features = ["bundled"] } +rusqlite = { version = "0.40", features = ["bundled"] } serde = { version = "1", features = ["derive"] } serde_json = "1" +schemars = "1" sha2 = "0.10" thiserror = "2" +tinyagents = { git = "https://github.com/senamakel/tinyagents", rev = "8f75355" } toml = "0.8" uuid = { version = "1", features = ["serde", "v4"] } walkdir = "2" @@ -73,6 +81,7 @@ reqwest = { version = "0.12", default-features = false, features = [ "json", "rustls-tls", ], optional = true } +tracing = { version = "0.1", optional = true } # tokio powers the optional background worker loops (`tokio` feature). It is # also listed under dev-dependencies so the async test suite always compiles. tokio = { version = "1", features = [ @@ -86,3 +95,12 @@ tokio = { version = "1", features = [ [dev-dependencies] tempfile = "3" tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] } +wiremock = "0.6" + +[[test]] +name = "composio_sync_mock" +required-features = ["sync"] + +[[test]] +name = "composio_sync_live" +required-features = ["sync"] diff --git a/src/memory/chunks/connection.rs b/src/memory/chunks/connection.rs index d5dab39..a7f06fb 100644 --- a/src/memory/chunks/connection.rs +++ b/src/memory/chunks/connection.rs @@ -219,7 +219,7 @@ fn apply_schema(conn: &Connection) -> Result<()> { // currently a no-op — a refusal is silently accepted here, and the // synchronous=FULL crash-safety assumption in `init_db` is only actually // valid when the mode really did become TRUNCATE. See audit finding SC-9. - if !journal_mode.eq_ignore_ascii_case("truncate") {} + let _ = journal_mode.eq_ignore_ascii_case("truncate"); conn.execute_batch(SCHEMA) .context("Failed to initialize chunk DB schema")?; // Additive, idempotent migrations. @@ -370,7 +370,7 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result { @@ -384,7 +384,7 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result( f(&guard) } +/// Return the initialized connection shared by chunk, tree, and auxiliary +/// stores for this workspace. +/// +/// Embedding applications may pass this handle to shared-connection stores +/// such as `KvStore` and `EntityIndex`. Callers must not change connection +/// pragmas or hold the mutex across an await point. +pub fn shared_connection(config: &MemoryConfig) -> Result>> { + get_or_init_connection(config) +} + #[cfg(test)] #[path = "connection_tests.rs"] mod tests; diff --git a/src/memory/chunks/embeddings.rs b/src/memory/chunks/embeddings.rs index e4b0517..24575b5 100644 --- a/src/memory/chunks/embeddings.rs +++ b/src/memory/chunks/embeddings.rs @@ -120,7 +120,7 @@ pub fn set_chunk_embedding_for_signature( /// /// # Errors /// See [`upsert_chunk_embedding_conn`]. -pub(crate) fn set_chunk_embedding_for_signature_tx( +pub fn set_chunk_embedding_for_signature_tx( tx: &rusqlite::Transaction<'_>, chunk_id: &str, model_signature: &str, @@ -133,7 +133,7 @@ pub(crate) fn set_chunk_embedding_for_signature_tx( /// /// # Errors /// See [`upsert_summary_embedding_conn`]. -pub(crate) fn set_summary_embedding_for_signature_tx( +pub fn set_summary_embedding_for_signature_tx( tx: &rusqlite::Transaction<'_>, summary_id: &str, model_signature: &str, @@ -199,6 +199,49 @@ pub fn clear_chunk_reembed_skipped( }) } +pub fn mark_summary_reembed_skipped( + config: &MemoryConfig, + summary_id: &str, + model_signature: &str, + reason: &str, +) -> Result<()> { + let summary_id = validate_reembed_skip_key("summary_id", summary_id)?; + let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; + with_connection(config, |conn| { + conn.execute( + "INSERT INTO mem_tree_summary_reembed_skipped + (summary_id, model_signature, reason, skipped_at_ms) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT(summary_id, model_signature) DO UPDATE SET + reason = excluded.reason, skipped_at_ms = excluded.skipped_at_ms", + rusqlite::params![ + summary_id, + model_signature, + reason, + Utc::now().timestamp_millis() + ], + )?; + Ok(()) + }) +} + +pub fn clear_summary_reembed_skipped( + config: &MemoryConfig, + summary_id: &str, + model_signature: &str, +) -> Result<()> { + let summary_id = validate_reembed_skip_key("summary_id", summary_id)?; + let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; + with_connection(config, |conn| { + conn.execute( + "DELETE FROM mem_tree_summary_reembed_skipped + WHERE summary_id = ?1 AND model_signature = ?2", + rusqlite::params![summary_id, model_signature], + )?; + Ok(()) + }) +} + /// Clear all chunk and summary tombstones for a model signature. Returns the /// total number of rows removed across both tombstone tables. Idempotent — /// calling again with nothing left to clear returns `Ok(0)`. @@ -226,7 +269,7 @@ pub fn clear_reembed_skipped_for_signature( /// Bounds attacker-controlled ids/signatures passed to reembed-skipped admin /// helpers. Rejects NUL bytes so SQLite bindings cannot be truncated. -pub(crate) const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048; +pub const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048; /// Validate and trim a `chunk_id` / `model_signature` value before it reaches /// a reembed-skipped table query. `label` is only used to make the error @@ -290,7 +333,7 @@ pub fn get_chunk_embedding(config: &MemoryConfig, chunk_id: &str) -> Result Vec { +pub fn embedding_to_blob(embedding: &[f32]) -> Vec { embedding.iter().flat_map(|f| f.to_le_bytes()).collect() } @@ -331,6 +374,30 @@ fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result rusqlite::Result { + conn.query_row( + "SELECT EXISTS( + SELECT 1 FROM mem_tree_chunks c + WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e + WHERE e.chunk_id = c.id AND e.model_signature = ?1) + AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk + WHERE sk.chunk_id = c.id AND sk.model_signature = ?1)) + OR EXISTS( + SELECT 1 FROM mem_tree_summaries s + WHERE s.deleted = 0 + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e + WHERE e.summary_id = s.id AND e.model_signature = ?1) + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk + WHERE sk.summary_id = s.id AND sk.model_signature = ?1))", + rusqlite::params![model_signature], + |row| row.get(0), + ) +} + /// Defensive cap for batched `IN (?,?,…)` reads, well below SQLite's /// `SQLITE_MAX_VARIABLE_NUMBER` (32 766). const MAX_EMBEDDING_BATCH: usize = 500; diff --git a/src/memory/chunks/mod.rs b/src/memory/chunks/mod.rs index 21a8eac..e8851e1 100644 --- a/src/memory/chunks/mod.rs +++ b/src/memory/chunks/mod.rs @@ -4,15 +4,14 @@ //! OpenHuman's `memory_store/chunks`: //! //! - [`types`] — [`Chunk`], [`Metadata`], [`SourceKind`], [`DataSource`], -//! [`SourceRef`], [`StagedChunk`] and the deterministic -//! [`chunk_id`] / token-estimate functions. The persisted -//! shape. +//! [`SourceRef`], [`StagedChunk`] and the deterministic +//! [`chunk_id`] / token-estimate functions. The persisted shape. //! - [`produce`] — source-kind-dispatch chunker (chat / email / document). -//! Splits canonical Markdown into bounded chunks with stable -//! per-source sequence numbers. +//! Splits canonical Markdown into bounded chunks with stable per-source +//! sequence numbers. //! - [`semantic`] — heading- and paragraph-aware chunker used to split large -//! documents into LLM-context-sized pieces while preserving -//! heading context. Exported as [`chunk_semantic`]. +//! documents into LLM-context-sized pieces while preserving heading context. +//! Exported as [`chunk_semantic`]. //! - `store` / `connection` / `migrations` / `raw_refs` / `embeddings` — //! the SQLite-backed chunk store (the `mem_tree_chunks` table plus its //! per-model embedding sidecars and source ingest gates). @@ -80,29 +79,39 @@ pub use types::{ Chunk, DataSource, Metadata, SourceKind, SourceRef, StagedChunk, }; -pub use connection::with_connection; +pub use connection::{shared_connection, with_connection}; pub use embeddings::{ - clear_chunk_reembed_skipped, clear_reembed_skipped_for_signature, get_chunk_embedding, + clear_chunk_reembed_skipped, clear_reembed_skipped_for_signature, + clear_summary_reembed_skipped, embedding_to_blob, get_chunk_embedding, get_chunk_embedding_for_signature, get_chunk_embeddings_batch, - get_chunk_embeddings_for_signature_batch, mark_chunk_reembed_skipped, set_chunk_embedding, - set_chunk_embedding_for_signature, tree_active_signature, + get_chunk_embeddings_for_signature_batch, has_uncovered_reembed_work, + mark_chunk_reembed_skipped, mark_summary_reembed_skipped, set_chunk_embedding, + set_chunk_embedding_for_signature, set_chunk_embedding_for_signature_tx, + set_summary_embedding_for_signature_tx, tree_active_signature, REEMBED_SKIP_KEY_MAX_LEN, }; pub use raw_refs::{ get_chunk_content_path, get_chunk_content_pointers, get_chunk_raw_refs, get_summary_content_pointers, list_chunk_raw_ref_paths_with_prefix, list_summaries_with_content_path, set_chunk_raw_refs, set_chunk_raw_refs_tx, RawRef, }; +pub use recovery::{ + is_io_open_error, is_transient_cold_start, recover_corrupt_db, try_cleanup_stale_files, +}; pub use store::{ claim_source_ingest_tx, count_chunks, count_chunks_by_lifecycle_status, count_raw_paths_ingested_with_prefix, delete_source_ingest, extraction_coverage, filter_raw_paths_not_ingested, get_chunk, get_chunk_lifecycle_status, get_chunks_batch, - is_source_ingested, list_chunks, mark_raw_paths_ingested, set_chunk_lifecycle_status, - upsert_chunks, ListChunksQuery, CHUNK_STATUS_ADMITTED, CHUNK_STATUS_BUFFERED, - CHUNK_STATUS_DROPPED, CHUNK_STATUS_PENDING_EXTRACTION, CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, + is_source_ingested, list_chunks, list_source_ids_with_prefix, mark_raw_paths_ingested, + set_chunk_lifecycle_status, update_chunk_content_sha256, update_summary_content_sha256, + upsert_chunks, upsert_chunks_tx, upsert_staged_chunks_tx, ListChunksQuery, + CHUNK_STATUS_ADMITTED, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, + CHUNK_STATUS_PENDING_EXTRACTION, CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, }; pub use store_delete::{ delete_chunks_by_owner, delete_chunks_by_source, delete_chunks_by_source_prefix, + delete_orphaned_source_tree, }; +pub use store_sources::{get_chunk_lifecycle_status_tx, set_chunk_lifecycle_status_tx}; // ── Shared internal constants / helpers ───────────────────────────────────── @@ -128,7 +137,10 @@ pub(crate) fn db_path_for(config: &MemoryConfig) -> PathBuf { /// Root directory for on-disk chunk/summary content files associated with the /// chunk DB. Bodies that are too large for the SQLite preview column live here. pub(crate) fn content_root(config: &MemoryConfig) -> PathBuf { - config.workspace.join(DB_DIR).join("content") + config + .content_root + .clone() + .unwrap_or_else(|| config.workspace.join(DB_DIR).join("content")) } /// Redact a PII-bearing string for log output by hashing it to 8 hex chars. diff --git a/src/memory/chunks/produce.rs b/src/memory/chunks/produce.rs index 109d691..0c84985 100644 --- a/src/memory/chunks/produce.rs +++ b/src/memory/chunks/produce.rs @@ -269,13 +269,13 @@ fn split_email_messages(md: &str) -> Vec { if line == "---" { // Check if one of the next 8 lines starts with `From:` let window_end = (i + 9).min(n); - for j in (i + 1)..window_end { - if lines[j].starts_with("From:") { + for candidate in lines.iter().take(window_end).skip(i + 1) { + if candidate.starts_with("From:") { split_positions.push(i); break; } // Skip blank lines between `---` and `From:` - if !lines[j].trim().is_empty() { + if !candidate.trim().is_empty() { break; } } @@ -298,7 +298,7 @@ fn split_email_messages(md: &str) -> Vec { } else { n }; - let piece_lines: Vec<&str> = lines[start..end].iter().copied().collect(); + let piece_lines: Vec<&str> = lines[start..end].to_vec(); let piece = piece_lines.join("\n").trim_end().to_string(); if !piece.is_empty() { pieces.push(piece); diff --git a/src/memory/chunks/produce_split.rs b/src/memory/chunks/produce_split.rs index e568789..07914f1 100644 --- a/src/memory/chunks/produce_split.rs +++ b/src/memory/chunks/produce_split.rs @@ -128,12 +128,9 @@ fn split_on_sentences<'a>(text: &'a str, budget: u32, out: &mut Vec<&'a str>) -> let mut start = 0usize; for i in 0..bytes.len() { let c = bytes[i]; - let boundary_end = if c == b'\n' { - Some(i + 1) - } else if (c == b'.' || c == b'!' || c == b'?') - && i + 1 < bytes.len() - && bytes[i + 1] == b' ' - { + let sentence_end = + matches!(c, b'.' | b'!' | b'?') && i + 1 < bytes.len() && bytes[i + 1] == b' '; + let boundary_end = if c == b'\n' || sentence_end { Some(i + 1) } else { None diff --git a/src/memory/chunks/raw_refs.rs b/src/memory/chunks/raw_refs.rs index c184911..50f32dd 100644 --- a/src/memory/chunks/raw_refs.rs +++ b/src/memory/chunks/raw_refs.rs @@ -127,17 +127,12 @@ pub fn list_chunk_raw_ref_paths_with_prefix( let mut out: std::collections::HashSet = std::collections::HashSet::new(); for row in rows { let json = row?; - match serde_json::from_str::>(&json) { - Ok(refs) => { - for raw_ref in refs { - if raw_ref.path.starts_with(rel_prefix) { - out.insert(raw_ref.path); - } + if let Ok(refs) = serde_json::from_str::>(&json) { + for raw_ref in refs { + if raw_ref.path.starts_with(rel_prefix) { + out.insert(raw_ref.path); } } - // Tolerate individually-corrupt rows: skip rather than failing - // the whole coverage scan. - Err(_) => {} } } Ok(out) diff --git a/src/memory/chunks/recovery.rs b/src/memory/chunks/recovery.rs index 0fb895c..023f220 100644 --- a/src/memory/chunks/recovery.rs +++ b/src/memory/chunks/recovery.rs @@ -58,7 +58,7 @@ const SQLITE_IOERR_IN_PAGE: i32 = 8714; /// it is currently only reachable from this crate's test modules, not from /// any production call site. #[allow(dead_code)] -pub(crate) fn is_transient_cold_start(err: &anyhow::Error) -> bool { +pub fn is_transient_cold_start(err: &anyhow::Error) -> bool { fn is_transient_sqlite(e: &(dyn std::error::Error + 'static)) -> bool { if let Some(rusqlite::Error::SqliteFailure(ffi, _)) = e.downcast_ref::() { return matches!( @@ -95,7 +95,7 @@ pub(crate) fn is_transient_cold_start(err: &anyhow::Error) -> bool { /// error (`{err:#}`) for error shapes that don't downcast cleanly — this /// catches messages surfaced through `anyhow::Context` wrapping that loses /// the original `rusqlite::Error` type. -pub(crate) fn is_io_open_error(err: &anyhow::Error) -> bool { +pub fn is_io_open_error(err: &anyhow::Error) -> bool { if let Some(rusqlite::Error::SqliteFailure(f, _)) = err.downcast_ref::() { return matches!( f.extended_code, @@ -134,7 +134,7 @@ pub(crate) fn is_io_open_error(err: &anyhow::Error) -> bool { /// drives the cold-start `IOERR_SHM*` failures (SQLite rebuilds it on open). /// /// Returns `true` if anything was checkpointed, quarantined, or removed. -pub(crate) fn try_cleanup_stale_files(db_path: &Path) -> bool { +pub fn try_cleanup_stale_files(db_path: &Path) -> bool { let mut cleaned = false; let wal = with_name_suffix(db_path, "-wal"); let shm = with_name_suffix(db_path, "-shm"); @@ -258,7 +258,7 @@ fn quick_check_ok(db_path: &Path) -> Result { /// succeeding on platforms with mandatory file locking), or if rebuilding the /// schema via [`get_or_init_connection`] fails. #[allow(dead_code)] -pub(crate) fn recover_corrupt_db(config: &MemoryConfig) -> Result { +pub fn recover_corrupt_db(config: &MemoryConfig) -> Result { let db_path = db_path_for(config); // 1. Drop any cached (corrupt) connection + breaker so the OS file handle diff --git a/src/memory/chunks/store.rs b/src/memory/chunks/store.rs index ee4b33b..f9a8713 100644 --- a/src/memory/chunks/store.rs +++ b/src/memory/chunks/store.rs @@ -62,15 +62,21 @@ pub fn upsert_chunks(config: &MemoryConfig, chunks: &[Chunk]) -> Result { } with_connection(config, |conn| { let tx = conn.unchecked_transaction()?; - { - let mut stmt = tx.prepare(UPSERT_SQL)?; - upsert_chunks_with_statement(&mut stmt, chunks)?; - } + let count = upsert_chunks_tx(&tx, chunks)?; tx.commit()?; - Ok(chunks.len()) + Ok(count) }) } +pub fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result { + if chunks.is_empty() { + return Ok(0); + } + let mut stmt = tx.prepare(UPSERT_SQL)?; + upsert_chunks_with_statement(&mut stmt, chunks)?; + Ok(chunks.len()) +} + /// `INSERT ... ON CONFLICT(id) DO UPDATE` for the plain (non-staged) chunk /// columns. Deliberately omits `content_path` / `content_sha256` / /// `lifecycle_status` — see the gotcha on [`upsert_chunks`]. @@ -148,10 +154,7 @@ fn upsert_chunks_with_statement( /// chunk's tags to JSON fails, or if any `stmt.execute` call fails. Does not /// commit `tx` itself — the caller owns the transaction lifecycle. #[allow(dead_code)] -pub(crate) fn upsert_staged_chunks_tx( - tx: &Transaction<'_>, - staged: &[StagedChunk], -) -> Result { +pub fn upsert_staged_chunks_tx(tx: &Transaction<'_>, staged: &[StagedChunk]) -> Result { if staged.is_empty() { return Ok(0); } @@ -208,6 +211,52 @@ pub(crate) fn upsert_staged_chunks_tx( /// positions here are load-bearing: `row_to_chunk` reads columns by /// numeric index, so this list and that function's `row.get(N)` calls must /// stay in lockstep. +pub fn update_chunk_content_sha256( + config: &MemoryConfig, + chunk_id: &str, + new_sha256: &str, +) -> Result<()> { + with_connection(config, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET content_sha256 = ?1 WHERE id = ?2", + params![new_sha256, chunk_id], + )?; + Ok(()) + }) +} + +pub fn update_summary_content_sha256( + config: &MemoryConfig, + summary_id: &str, + new_sha256: &str, +) -> Result<()> { + with_connection(config, |conn| { + conn.execute( + "UPDATE mem_tree_summaries SET content_sha256 = ?1 WHERE id = ?2", + params![new_sha256, summary_id], + )?; + Ok(()) + }) +} + +pub fn list_source_ids_with_prefix( + config: &MemoryConfig, + source_kind: SourceKind, + prefix: &str, +) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT DISTINCT source_id FROM mem_tree_chunks + WHERE source_kind = ?1 ORDER BY source_id ASC", + )?; + let rows = stmt.query_map(params![source_kind.as_str()], |row| row.get::<_, String>(0))?; + Ok(rows + .filter_map(std::result::Result::ok) + .filter(|id| id.starts_with(prefix)) + .collect()) + }) +} + const SELECT_COLUMNS: &str = "id, source_kind, source_id, path_scope, source_ref, owner, timestamp_ms, time_range_start_ms, time_range_end_ms, tags_json, content, token_count, seq_in_source, created_at_ms"; diff --git a/src/memory/chunks/store_conn_tests.rs b/src/memory/chunks/store_conn_tests.rs index 741d964..e06bf70 100644 --- a/src/memory/chunks/store_conn_tests.rs +++ b/src/memory/chunks/store_conn_tests.rs @@ -252,7 +252,7 @@ fn connection_cache_uses_separate_connections_for_different_workspaces() { ); let c = sample_chunk("s", 0, 1_700_000_000_000); - upsert_chunks(&cfg1, &[c.clone()]).unwrap(); + upsert_chunks(&cfg1, std::slice::from_ref(&c)).unwrap(); assert_eq!(count_chunks(&cfg1).unwrap(), 1); assert_eq!(count_chunks(&cfg2).unwrap(), 0); } @@ -372,7 +372,7 @@ fn existing_wal_db_migrates_to_truncate() { fn delete_chunks_by_source_removes_symlink_entry_not_target_file() { let (_tmp, cfg) = test_config(); let linked_chunk = sample_chunk("slack:c-1", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[linked_chunk.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&linked_chunk)).unwrap(); let root = content_root(&cfg); let target_path = root.join("chunks/target.md"); diff --git a/src/memory/chunks/store_delete.rs b/src/memory/chunks/store_delete.rs index f24309b..959071e 100644 --- a/src/memory/chunks/store_delete.rs +++ b/src/memory/chunks/store_delete.rs @@ -130,7 +130,7 @@ fn delete_chunks_by_source_filter( let chunks = { let mut stmt = tx.prepare( - "SELECT id, source_id, owner, content_path + "SELECT id, source_id, owner, content_path, path_scope FROM mem_tree_chunks WHERE source_kind = ?1", )?; @@ -140,11 +140,14 @@ fn delete_chunks_by_source_filter( row.get::<_, String>(1)?, row.get::<_, String>(2)?, row.get::<_, Option>(3)?, + row.get::<_, Option>(4)?, )) })?; rows.filter_map(|row| match row { - Ok((id, source_id, owner, content_path)) if matches_chunk(&source_id, &owner) => { - Some(Ok((id, source_id, content_path))) + Ok((id, source_id, owner, content_path, path_scope)) + if matches_chunk(&source_id, &owner) => + { + Some(Ok((id, source_id, content_path, path_scope))) } Ok(_) => None, Err(error) => Some(Err(error)), @@ -155,10 +158,16 @@ fn delete_chunks_by_source_filter( let deleted_source_ids: HashSet = chunks .iter() - .map(|(_, source_id, _)| source_id.clone()) + .map(|(_, source_id, _, _)| source_id.clone()) + .collect(); + let deleted_tree_scopes: HashSet = chunks + .iter() + .map(|(_, source_id, _, path_scope)| { + path_scope.clone().unwrap_or_else(|| source_id.clone()) + }) .collect(); - for (chunk_id, _source_id, content_path) in &chunks { + for (chunk_id, _source_id, content_path, _path_scope) in &chunks { tx.execute( "DELETE FROM mem_tree_score WHERE chunk_id = ?1", params![chunk_id], @@ -229,6 +238,26 @@ fn delete_chunks_by_source_filter( )?; } + for scope in &deleted_tree_scopes { + let remaining: i64 = tx.query_row( + "SELECT COUNT(*) FROM mem_tree_chunks + WHERE source_kind = ?1 AND COALESCE(path_scope, source_id) = ?2", + params![source_kind.as_str(), scope], + |row| row.get(0), + )?; + if remaining != 0 { + continue; + } + if let Some(tree) = crate::memory::tree::store::get_tree_by_scope_conn( + &tx, + crate::memory::tree::store::TreeKind::Source, + scope, + )? { + let cascade = crate::memory::tree::store::delete_tree_cascade_tx(&tx, &tree.id)?; + content_paths.extend(cascade.content_paths); + } + } + let deleted = chunks.len(); tx.commit()?; Ok(deleted) @@ -238,6 +267,60 @@ fn delete_chunks_by_source_filter( Ok(deleted) } +/// Remove stale gates and a source-scoped tree once no chunks remain. +pub fn delete_orphaned_source_tree( + config: &MemoryConfig, + source_kind: SourceKind, + source_id: &str, +) -> Result { + let mut content_paths = Vec::new(); + let cascaded = with_connection(config, |conn| { + let tx = conn.unchecked_transaction()?; + let remaining: i64 = tx.query_row( + "SELECT COUNT(*) FROM mem_tree_chunks WHERE source_kind=?1 AND source_id=?2", + params![source_kind.as_str(), source_id], + |row| row.get(0), + )?; + if remaining > 0 { + return Ok(false); + } + let versioned_prefix = format!("{source_id}@"); + let gate_ids = { + let mut stmt = + tx.prepare("SELECT source_id FROM mem_tree_ingested_sources WHERE source_kind=?1")?; + let rows = + stmt.query_map(params![source_kind.as_str()], |row| row.get::<_, String>(0))?; + rows.filter_map(|row| match row { + Ok(id) if id == source_id || id.starts_with(&versioned_prefix) => Some(Ok(id)), + Ok(_) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()? + }; + for gate_id in gate_ids { + tx.execute( + "DELETE FROM mem_tree_ingested_sources WHERE source_kind=?1 AND source_id=?2", + params![source_kind.as_str(), gate_id], + )?; + } + let cascaded = if let Some(tree) = crate::memory::tree::store::get_tree_by_scope_conn( + &tx, + crate::memory::tree::store::TreeKind::Source, + source_id, + )? { + let deleted = crate::memory::tree::store::delete_tree_cascade_tx(&tx, &tree.id)?; + content_paths.extend(deleted.content_paths); + true + } else { + false + }; + tx.commit()?; + Ok(cascaded) + })?; + remove_chunk_content_files(config, &content_paths); + Ok(cascaded) +} + /// Best-effort removal of on-disk chunk content files, with strict sandboxing: /// a `content_path` that escapes the content root (via `..`, an absolute path, /// or a symlink pointing outside) is refused rather than followed. diff --git a/src/memory/chunks/store_embed_tests.rs b/src/memory/chunks/store_embed_tests.rs index 4360632..3f6e00f 100644 --- a/src/memory/chunks/store_embed_tests.rs +++ b/src/memory/chunks/store_embed_tests.rs @@ -17,14 +17,17 @@ use super::recovery::{is_transient_cold_start, try_cleanup_stale_files}; use super::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; use super::{ claim_source_ingest_tx, clear_chunk_reembed_skipped, clear_reembed_skipped_for_signature, - content_root, count_chunks, db_path_for, delete_chunks_by_owner, delete_chunks_by_source, - extraction_coverage, get_chunk, get_chunk_embedding, get_chunk_embedding_for_signature, - get_chunk_embeddings_for_signature_batch, get_chunks_batch, is_source_ingested, list_chunks, - mark_chunk_reembed_skipped, set_chunk_embedding, set_chunk_embedding_for_signature, - tree_active_signature, upsert_chunks, ListChunksQuery, DB_DIR, - GLOBAL_TOPIC_PURGE_MIGRATION_VERSION, + clear_summary_reembed_skipped, content_root, count_chunks, db_path_for, delete_chunks_by_owner, + delete_chunks_by_source, extraction_coverage, get_chunk, get_chunk_embedding, + get_chunk_embedding_for_signature, get_chunk_embeddings_for_signature_batch, get_chunks_batch, + is_source_ingested, list_chunks, mark_chunk_reembed_skipped, mark_summary_reembed_skipped, + set_chunk_embedding, set_chunk_embedding_for_signature, tree_active_signature, upsert_chunks, + ListChunksQuery, DB_DIR, GLOBAL_TOPIC_PURGE_MIGRATION_VERSION, }; use crate::memory::config::MemoryConfig; +use crate::memory::tree::store::{ + insert_summary_tx, insert_tree, SummaryNode, Tree, TreeKind, TreeStatus, +}; use chrono::{TimeZone, Utc}; use rusqlite::params; use std::sync::Arc; @@ -62,7 +65,7 @@ fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk { fn clear_chunk_reembed_skipped_is_idempotent() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); let sig = tree_active_signature(&cfg); mark_chunk_reembed_skipped(&cfg, &c.id, &sig, "test orphan").unwrap(); clear_chunk_reembed_skipped(&cfg, &c.id, &sig).unwrap(); @@ -79,6 +82,70 @@ fn clear_chunk_reembed_skipped_is_idempotent() { assert_eq!(count, 0); } +#[test] +fn summary_reembed_tombstone_roundtrips_and_clears() { + let (_tmp, cfg) = test_config(); + let now = Utc::now(); + insert_tree( + &cfg, + &Tree { + id: "tree-1".into(), + kind: TreeKind::Source, + scope: "source-1".into(), + root_id: None, + max_level: 0, + status: TreeStatus::Active, + created_at: now, + last_sealed_at: None, + }, + ) + .unwrap(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_summary_tx( + &tx, + &SummaryNode { + id: "summary-1".into(), + tree_id: "tree-1".into(), + tree_kind: TreeKind::Source, + level: 1, + parent_id: None, + child_ids: vec![], + content: "summary".into(), + token_count: 2, + entities: vec![], + topics: vec![], + time_range_start: now, + time_range_end: now, + score: 1.0, + sealed_at: now, + deleted: false, + embedding: None, + doc_id: None, + version_ms: None, + }, + "model@3", + )?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + mark_summary_reembed_skipped(&cfg, "summary-1", "model@3", "oversize").unwrap(); + with_connection(&cfg, |conn| { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_summary_reembed_skipped + WHERE summary_id='summary-1' AND model_signature='model@3'", + [], + |row| row.get(0), + )?; + assert_eq!(count, 1); + Ok(()) + }) + .unwrap(); + clear_summary_reembed_skipped(&cfg, "summary-1", "model@3").unwrap(); + clear_summary_reembed_skipped(&cfg, "summary-1", "model@3").unwrap(); +} + #[test] fn clear_reembed_skipped_for_signature_removes_all_tombstones_for_sig() { let (_tmp, cfg) = test_config(); @@ -187,7 +254,7 @@ fn get_chunks_batch_empty_input_and_missing_ids() { assert!(empty.is_empty()); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); let ids = vec![ c.id.clone(), "ghost:no-such-1".into(), @@ -196,8 +263,8 @@ fn get_chunks_batch_empty_input_and_missing_ids() { let map = get_chunks_batch(&cfg, &ids).unwrap(); assert_eq!(map.len(), 1); assert_eq!(map.get(&c.id), Some(&c)); - assert!(map.get("ghost:no-such-1").is_none()); - assert!(map.get("ghost:no-such-2").is_none()); + assert!(!map.contains_key("ghost:no-such-1")); + assert!(!map.contains_key("ghost:no-such-2")); } #[test] @@ -219,7 +286,7 @@ fn batch_embedding_lookup_returns_only_signature_scoped_rows() { assert_eq!(map_a.len(), 2, "only c1 and c2 are under sig_a"); assert_eq!(map_a.get(&c1.id).cloned(), Some(vec![0.1, 0.2])); assert_eq!(map_a.get(&c2.id).cloned(), Some(vec![0.3, 0.4])); - assert!(map_a.get(&c3.id).is_none(), "c3 has only sig_b"); + assert!(!map_a.contains_key(&c3.id), "c3 has only sig_b"); let map_b = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); assert_eq!(map_b.len(), 1); @@ -237,7 +304,7 @@ fn batch_embedding_lookup_empty_input_returns_empty_map() { fn batch_embedding_lookup_unknown_ids_absent_from_map() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); let sig = "openai/text-embedding-3-small@1536"; set_chunk_embedding_for_signature(&cfg, &c.id, sig, &[0.1]).unwrap(); diff --git a/src/memory/chunks/store_sources.rs b/src/memory/chunks/store_sources.rs index be872a6..a0f7f0e 100644 --- a/src/memory/chunks/store_sources.rs +++ b/src/memory/chunks/store_sources.rs @@ -37,7 +37,7 @@ pub fn set_chunk_lifecycle_status( /// # Errors /// See [`set_chunk_lifecycle_status`]. #[allow(dead_code)] -pub(crate) fn set_chunk_lifecycle_status_tx( +pub fn set_chunk_lifecycle_status_tx( tx: &Transaction<'_>, chunk_id: &str, status: &str, @@ -88,6 +88,19 @@ pub fn get_chunk_lifecycle_status(config: &MemoryConfig, chunk_id: &str) -> Resu }) } +pub fn get_chunk_lifecycle_status_tx( + tx: &Transaction<'_>, + chunk_id: &str, +) -> Result> { + Ok(tx + .query_row( + "SELECT lifecycle_status FROM mem_tree_chunks WHERE id = ?1", + params![chunk_id], + |row| row.get(0), + ) + .optional()?) +} + /// Count chunks currently sitting at a given lifecycle status. Matches /// exactly (no prefix/wildcard semantics) — an unknown status string simply /// counts as `0`. diff --git a/src/memory/chunks/store_tests.rs b/src/memory/chunks/store_tests.rs index 9df4606..1f30fde 100644 --- a/src/memory/chunks/store_tests.rs +++ b/src/memory/chunks/store_tests.rs @@ -61,7 +61,7 @@ fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk { fn upsert_then_get() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - assert_eq!(upsert_chunks(&cfg, &[c.clone()]).unwrap(), 1); + assert_eq!(upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(), 1); let got = get_chunk(&cfg, &c.id).unwrap().expect("chunk stored"); assert_eq!(got, c); } @@ -73,7 +73,7 @@ fn upsert_persists_path_scope() { c.metadata.source_kind = SourceKind::Document; c.metadata.path_scope = Some("notion:conn-1".to_string()); - assert_eq!(upsert_chunks(&cfg, &[c.clone()]).unwrap(), 1); + assert_eq!(upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(), 1); let got = get_chunk(&cfg, &c.id).unwrap().expect("chunk stored"); assert_eq!(got.metadata.source_id, "notion:conn-1:page-abc"); @@ -119,8 +119,8 @@ fn list_chunks_source_scope_filters_before_limit() { fn upsert_is_idempotent() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); assert_eq!(count_chunks(&cfg).unwrap(), 1); } @@ -128,12 +128,12 @@ fn upsert_is_idempotent() { fn reingest_preserves_existing_embedding() { let (_tmp, cfg) = test_config(); let mut c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); set_chunk_embedding(&cfg, &c.id, &[0.1, 0.2, 0.3]).unwrap(); c.content = "updated content".into(); c.token_count = 99; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); let embedding = get_chunk_embedding(&cfg, &c.id).unwrap().unwrap(); assert_eq!(embedding, vec![0.1, 0.2, 0.3]); @@ -146,7 +146,7 @@ fn reingest_preserves_existing_embedding() { fn chunk_embeddings_are_scoped_by_model_signature() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&c)).unwrap(); set_chunk_embedding_for_signature( &cfg, @@ -353,6 +353,33 @@ fn delete_chunks_by_source_removes_chunks_side_rows_and_ingest_gate() { .unwrap(); } +#[test] +fn delete_chunks_by_source_cascades_path_scoped_tree() { + let (_tmp, cfg) = test_config(); + let mut chunk = sample_chunk("notion:item", 0, 1_700_000_000_000); + chunk.metadata.source_kind = SourceKind::Document; + chunk.metadata.path_scope = Some("notion:workspace".into()); + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); + crate::memory::tree::registry::get_or_create_tree( + &cfg, + crate::memory::tree::store::TreeKind::Source, + "notion:workspace", + ) + .unwrap(); + + assert_eq!( + delete_chunks_by_source(&cfg, SourceKind::Document, "notion:item").unwrap(), + 1 + ); + assert!(crate::memory::tree::store::get_tree_by_scope( + &cfg, + crate::memory::tree::store::TreeKind::Source, + "notion:workspace", + ) + .unwrap() + .is_none()); +} + #[test] fn delete_chunks_by_owner_preserves_other_owners_for_same_source() { let (_tmp, cfg) = test_config(); @@ -495,7 +522,7 @@ fn raw_refs_round_trip_and_prefix_listing_tolerates_corrupt_rows() { fn content_pointer_accessors_return_only_complete_non_deleted_rows() { let (_tmp, cfg) = test_config(); let chunk = sample_chunk("notion:page-1", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); with_connection(&cfg, |conn| { conn.execute( "UPDATE mem_tree_chunks SET content_path = ?1, content_sha256 = ?2 WHERE id = ?3", diff --git a/src/memory/chunks/types.rs b/src/memory/chunks/types.rs index e0c935f..f8bbe8f 100644 --- a/src/memory/chunks/types.rs +++ b/src/memory/chunks/types.rs @@ -336,7 +336,7 @@ pub fn conservative_token_estimate(text: &str) -> u32 { .chars() .map(|c| u64::from(char_token_quarters(c))) .sum(); - let tokens = (quarters + 3) / 4; // ceil(quarters / 4) + let tokens = quarters.div_ceil(4); // ceil(quarters / 4) tokens.min(u64::from(u32::MAX)) as u32 } diff --git a/src/memory/config.rs b/src/memory/config.rs index 879a4d6..a4fcb90 100644 --- a/src/memory/config.rs +++ b/src/memory/config.rs @@ -37,6 +37,10 @@ pub struct MemoryConfig { /// exist yet at construction time — callers create it (or fail informatively) /// when they first open the workspace. pub workspace: PathBuf, + /// Optional override for the chunk/summary content vault. `None` uses + /// `/memory_tree/content`. + #[serde(default)] + pub content_root: Option, /// Embedding configuration. #[serde(default)] pub embedding: EmbeddingConfig, @@ -49,6 +53,9 @@ pub struct MemoryConfig { /// Per-source sync budget ceilings (enforced when a host invokes ingest). #[serde(default)] pub sync_budget: SyncBudgetConfig, + /// Live synchronization configuration. + #[serde(default)] + pub sync: SyncConfig, } impl MemoryConfig { @@ -67,10 +74,12 @@ impl MemoryConfig { pub fn new(workspace: impl Into) -> Self { Self { workspace: workspace.into(), + content_root: None, embedding: EmbeddingConfig::default(), tree: TreeConfig::default(), retrieval: RetrievalConfig::default(), sync_budget: SyncBudgetConfig::default(), + sync: SyncConfig::default(), } } } @@ -132,6 +141,7 @@ impl Default for TreeConfig { /// that the four weights sum to `1.0` — the built-in profiles are chosen that /// way by convention so scores land in a familiar `[0.0, 1.0]`-ish range when /// every signal is itself in `[0.0, 1.0]`, but a custom profile with a +/// /// different total simply rescales the final score. #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] pub struct WeightProfile { @@ -221,6 +231,8 @@ impl Default for RetrievalConfig { /// than an already-enforced global cap. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct SyncBudgetConfig { + /// Maximum records persisted by one sync tick. + pub max_items: Option, /// Token ceiling per ingest run; `None` leaves token spend unbounded. pub max_tokens_per_sync: Option, /// USD cost ceiling per ingest run; `None` leaves cost unbounded. @@ -229,6 +241,78 @@ pub struct SyncBudgetConfig { pub sync_depth_days: Option, } +/// Live synchronization configuration. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SyncConfig { + /// Request and ingest ceilings. + #[serde(default)] + pub budget: SyncBudgetConfig, + /// Composio transport configuration; absent disables Composio sync. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub composio: Option, + /// Global periodic cadence. `Some(0)` means manual-only; shorter non-zero + /// values are clamped to the 24-hour default. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub interval_secs: Option, +} + +/// Composio HTTP transport mode. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ComposioMode { + /// Call backend.composio.dev using a BYO API key. + Direct, + /// Call a host proxy using a bearer token. + #[default] + Proxied, +} + +/// Redacted secret value. Debug and Display never reveal the payload. +#[derive(Clone, Default, PartialEq, Eq)] +pub struct SecretString(String); + +impl SecretString { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn expose(&self) -> &str { + &self.0 + } + + pub fn is_empty(&self) -> bool { + self.0.trim().is_empty() + } +} + +impl std::fmt::Debug for SecretString { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("SecretString([REDACTED])") + } +} + +impl std::fmt::Display for SecretString { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("[REDACTED]") + } +} + +/// Composio connection settings injected by the host. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ComposioSyncConfig { + #[serde(default)] + pub mode: ComposioMode, + pub base_url: String, + /// Direct-mode key. Never serialized or printed. + #[serde(skip)] + pub api_key: Option, + /// Proxied-mode bearer. Never serialized or printed. + #[serde(skip)] + pub bearer_token: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub entity_id: Option, +} + #[cfg(test)] #[path = "config_tests.rs"] mod tests; diff --git a/src/memory/config_tests.rs b/src/memory/config_tests.rs index fb8defd..85f78e1 100644 --- a/src/memory/config_tests.rs +++ b/src/memory/config_tests.rs @@ -2,6 +2,24 @@ use super::*; +#[test] +fn sync_secrets_are_redacted_and_not_serialized() { + let secret = SecretString::new("super-secret"); + assert_eq!(secret.to_string(), "[REDACTED]"); + assert!(!format!("{secret:?}").contains("super-secret")); + + let config = ComposioSyncConfig { + mode: ComposioMode::Direct, + base_url: "https://backend.composio.dev".into(), + api_key: Some(secret), + bearer_token: Some(SecretString::new("bearer-secret")), + entity_id: Some("entity".into()), + }; + let json = serde_json::to_string(&config).unwrap(); + assert!(!json.contains("super-secret")); + assert!(!json.contains("bearer-secret")); +} + #[test] fn default_config_uses_openhuman_constants() { let cfg = MemoryConfig::new("/tmp/ws"); diff --git a/src/memory/conversations/inverted_index.rs b/src/memory/conversations/inverted_index.rs index 8c8c7ab..6f12055 100644 --- a/src/memory/conversations/inverted_index.rs +++ b/src/memory/conversations/inverted_index.rs @@ -260,7 +260,7 @@ impl InvertedIndex { // Phase 2: verify each candidate by exact substring match. // Count distinct terms per doc for the score. let mut hit_counts: HashMap = HashMap::new(); - for (term, candidates) in terms.iter().zip(per_term.into_iter()) { + for (term, candidates) in terms.iter().zip(per_term) { for doc_id in candidates { let Some(entry) = self.docs[doc_id as usize].as_ref() else { continue; diff --git a/src/memory/goals/types.rs b/src/memory/goals/types.rs index 178a384..9c9c380 100644 --- a/src/memory/goals/types.rs +++ b/src/memory/goals/types.rs @@ -14,7 +14,13 @@ use serde::{Deserialize, Serialize}; use crate::memory::error::MemoryError; -use crate::memory::store::safety::{has_likely_pii, has_likely_secret}; +use crate::memory::store::safety::{ + has_likely_email, has_likely_pii, has_likely_secret, sanitize_text, +}; + +fn has_goal_pii(text: &str) -> bool { + has_likely_email(text) || has_likely_pii(text) || sanitize_text(text).report.pii_redactions > 0 +} /// Markdown header rendered at the top of `MEMORY_GOALS.md`. pub(crate) const HEADER: &str = "# Long-term Goals"; @@ -140,7 +146,7 @@ impl GoalsDoc { "goal text must be a single line".to_string(), )); } - if has_likely_secret(text) || has_likely_pii(text) { + if has_likely_secret(text) || has_goal_pii(text) { return Err(MemoryError::Invalid( "goal text must not contain secrets or PII".to_string(), )); diff --git a/src/memory/graph/bfs.rs b/src/memory/graph/bfs.rs new file mode 100644 index 0000000..2bc0f66 --- /dev/null +++ b/src/memory/graph/bfs.rs @@ -0,0 +1,118 @@ +//! Bounded shortest paths over persisted co-occurrence edges. + +use std::collections::{HashMap, HashSet, VecDeque}; + +use anyhow::Result; + +use crate::memory::config::MemoryConfig; +use crate::memory::graph::edge_store::edge_neighbors; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PairDistance { + pub a: String, + pub b: String, + pub dist: u32, +} + +pub fn pair_distances( + config: &MemoryConfig, + entity_ids: &[String], + max_h: u32, +) -> Result> { + if max_h == 0 || entity_ids.len() < 2 { + return Ok(Vec::new()); + } + let mut unique = entity_ids.to_vec(); + unique.sort(); + unique.dedup(); + let targets: HashSet<_> = unique.iter().cloned().collect(); + let mut cache: HashMap> = HashMap::new(); + let mut emitted = HashSet::new(); + let mut output = Vec::new(); + + for source in &unique { + let mut remaining = targets.clone(); + remaining.remove(source); + let mut visited = HashSet::from([source.clone()]); + let mut queue = VecDeque::from([(source.clone(), 0)]); + while let Some((node, distance)) = queue.pop_front() { + if distance >= max_h || remaining.is_empty() { + continue; + } + if !cache.contains_key(&node) { + cache.insert( + node.clone(), + edge_neighbors(config, &node)? + .into_iter() + .map(|(neighbor, _)| neighbor) + .collect(), + ); + } + for neighbor in cache.get(&node).cloned().unwrap_or_default() { + if !visited.insert(neighbor.clone()) { + continue; + } + let next_distance = distance + 1; + if remaining.remove(&neighbor) { + let pair = if source < &neighbor { + (source.clone(), neighbor.clone()) + } else { + (neighbor.clone(), source.clone()) + }; + if emitted.insert(pair.clone()) { + output.push(PairDistance { + a: pair.0, + b: pair.1, + dist: next_distance, + }); + } + } + queue.push_back((neighbor, next_distance)); + } + } + } + output.sort_by(|a, b| { + a.dist + .cmp(&b.dist) + .then_with(|| a.a.cmp(&b.a)) + .then_with(|| a.b.cmp(&b.b)) + }); + Ok(output) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::graph::{pairs_from_entities, upsert_edges}; + + #[test] + fn bounded_bfs_finds_two_hop_pair() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + upsert_edges( + &config, + &pairs_from_entities(&["alice".into(), "bob".into()]), + 1, + ) + .unwrap(); + upsert_edges( + &config, + &pairs_from_entities(&["bob".into(), "carol".into()]), + 1, + ) + .unwrap(); + assert!( + pair_distances(&config, &["alice".into(), "carol".into()], 1) + .unwrap() + .is_empty() + ); + assert_eq!( + pair_distances(&config, &["alice".into(), "carol".into()], 2).unwrap(), + vec![PairDistance { + a: "alice".into(), + b: "carol".into(), + dist: 2, + }] + ); + } +} diff --git a/src/memory/graph/edge_store.rs b/src/memory/graph/edge_store.rs new file mode 100644 index 0000000..4c7e927 --- /dev/null +++ b/src/memory/graph/edge_store.rs @@ -0,0 +1,130 @@ +//! Persisted undirected entity co-occurrence edges. + +use std::collections::BTreeSet; + +use anyhow::Result; +use rusqlite::{params, Transaction}; + +use crate::memory::chunks::with_connection; +use crate::memory::config::MemoryConfig; + +fn order_pair<'a>(a: &'a str, b: &'a str) -> Option<(&'a str, &'a str)> { + match a.cmp(b) { + std::cmp::Ordering::Less => Some((a, b)), + std::cmp::Ordering::Greater => Some((b, a)), + std::cmp::Ordering::Equal => None, + } +} + +pub fn pairs_from_entities(entity_ids: &[String]) -> Vec<(String, String)> { + let unique: Vec<_> = entity_ids + .iter() + .collect::>() + .into_iter() + .collect(); + let mut pairs = Vec::new(); + for left in 0..unique.len() { + for right in left + 1..unique.len() { + if let Some((a, b)) = order_pair(unique[left], unique[right]) { + pairs.push((a.to_string(), b.to_string())); + } + } + } + pairs +} + +pub fn upsert_edges_tx( + tx: &Transaction<'_>, + pairs: &[(String, String)], + timestamp_ms: i64, +) -> Result { + let canonical: BTreeSet<_> = pairs + .iter() + .filter_map(|(a, b)| order_pair(a, b).map(|(a, b)| (a.to_string(), b.to_string()))) + .collect(); + let mut statement = tx.prepare( + "INSERT INTO mem_tree_entity_edges (entity_a, entity_b, weight, updated_ms) + VALUES (?1, ?2, 1, ?3) + ON CONFLICT(entity_a, entity_b) + DO UPDATE SET weight = weight + 1, updated_ms = ?3", + )?; + for (a, b) in &canonical { + statement.execute(params![a, b, timestamp_ms])?; + } + Ok(canonical.len()) +} + +pub fn upsert_edges( + config: &MemoryConfig, + pairs: &[(String, String)], + timestamp_ms: i64, +) -> Result { + if pairs.is_empty() { + return Ok(0); + } + with_connection(config, |connection| { + let transaction = connection.unchecked_transaction()?; + let count = upsert_edges_tx(&transaction, pairs, timestamp_ms)?; + transaction.commit()?; + Ok(count) + }) +} + +pub fn edge_neighbors(config: &MemoryConfig, entity_id: &str) -> Result> { + with_connection(config, |connection| { + let mut statement = connection.prepare( + "SELECT entity_b, weight FROM mem_tree_entity_edges WHERE entity_a = ?1 + UNION ALL + SELECT entity_a, weight FROM mem_tree_entity_edges WHERE entity_b = ?1 + ORDER BY weight DESC", + )?; + let rows = statement + .query_map(params![entity_id], |row| Ok((row.get(0)?, row.get(1)?)))? + .collect::>>()?; + Ok(rows) + }) +} + +pub fn clear_edges_for_entities_tx(tx: &Transaction<'_>, entity_ids: &[String]) -> Result { + let mut removed = 0; + let mut statement = + tx.prepare("DELETE FROM mem_tree_entity_edges WHERE entity_a = ?1 OR entity_b = ?1")?; + for id in entity_ids { + removed += statement.execute(params![id])?; + } + Ok(removed) +} + +pub fn count_edges(config: &MemoryConfig) -> Result { + with_connection(config, |connection| { + let count: i64 = + connection.query_row("SELECT COUNT(*) FROM mem_tree_entity_edges", [], |row| { + row.get(0) + })?; + Ok(count.max(0) as u64) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn persisted_edges_are_canonical_weighted_and_symmetric() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + let pairs = pairs_from_entities(&[ + "person:bob".into(), + "person:alice".into(), + "person:alice".into(), + ]); + assert_eq!(pairs, vec![("person:alice".into(), "person:bob".into())]); + upsert_edges(&config, &pairs, 1).unwrap(); + upsert_edges(&config, &pairs, 2).unwrap(); + assert_eq!( + edge_neighbors(&config, "person:bob").unwrap(), + vec![("person:alice".into(), 2)] + ); + assert_eq!(count_edges(&config).unwrap(), 1); + } +} diff --git a/src/memory/graph/mod.rs b/src/memory/graph/mod.rs index 5aa2c14..1846bff 100644 --- a/src/memory/graph/mod.rs +++ b/src/memory/graph/mod.rs @@ -42,8 +42,16 @@ //! not a fixed query budget. Callers on hot paths (e.g. per-request retrieval) //! should cap `limit` and be mindful of the underlying index's per-call cost. +mod bfs; +mod edge_store; pub mod query; pub mod types; +pub use bfs::{pair_distances, PairDistance}; +pub use edge_store::{ + clear_edges_for_entities_tx, count_edges, edge_neighbors, pairs_from_entities, upsert_edges, + upsert_edges_tx, +}; + pub use query::{co_occurring_entities, group_by_weight, neighbors}; pub use types::{EntityOccurrenceIndex, GraphEdge}; diff --git a/src/memory/ingest/extract/mod.rs b/src/memory/ingest/extract/mod.rs index ab7644a..17417cd 100644 --- a/src/memory/ingest/extract/mod.rs +++ b/src/memory/ingest/extract/mod.rs @@ -42,6 +42,33 @@ pub use types::{ MemoryIngestionRequest, MemoryIngestionResult, DEFAULT_MEMORY_EXTRACTION_MODEL, }; +use crate::memory::types::NamespaceDocumentInput; + +/// Extract structured knowledge and merge deterministic ingestion metadata +/// into a document ready for a host-owned namespace store. +pub fn extract_enriched_document( + input: &NamespaceDocumentInput, + config: &MemoryIngestionConfig, +) -> (NamespaceDocumentInput, MemoryIngestionResult) { + let parsed = parse::parse_document(&input.content, &input.title, config); + let (enriched, tags) = header::enrich_document_metadata(input, &parsed, config); + let result = MemoryIngestionResult { + document_id: String::new(), + namespace: String::new(), + model_name: config.model_name.clone(), + extraction_mode: config.extraction_mode.as_str().to_string(), + chunk_count: parsed.chunk_count, + entity_count: parsed.entities.len(), + relation_count: parsed.relations.len(), + preference_count: parsed.preference_count, + decision_count: parsed.decision_count, + tags, + entities: parsed.entities, + relations: parsed.relations, + }; + (enriched, result) +} + #[cfg(test)] #[path = "extract_tests.rs"] mod extract_tests; diff --git a/src/memory/ingest/extract/types.rs b/src/memory/ingest/extract/types.rs index 94324a4..24fb6f5 100644 --- a/src/memory/ingest/extract/types.rs +++ b/src/memory/ingest/extract/types.rs @@ -26,7 +26,7 @@ pub enum ExtractionMode { impl ExtractionMode { /// Returns the string representation of the extraction mode. - pub(super) fn as_str(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::Sentence => "sentence", Self::Chunk => "chunk", diff --git a/src/memory/ingest/mod.rs b/src/memory/ingest/mod.rs index ff5fc17..101350f 100644 --- a/src/memory/ingest/mod.rs +++ b/src/memory/ingest/mod.rs @@ -40,8 +40,9 @@ pub use canonicalize::email::{EmailMessage, EmailThread}; pub use canonicalize::{CanonicaliseRequest, CanonicalisedSource}; pub use extract::{ - extract_document, ExtractedEntity, ExtractedRelation, ExtractionMode, MemoryIngestionConfig, - MemoryIngestionRequest, MemoryIngestionResult, DEFAULT_MEMORY_EXTRACTION_MODEL, + extract_document, extract_enriched_document, ExtractedEntity, ExtractedRelation, + ExtractionMode, MemoryIngestionConfig, MemoryIngestionRequest, MemoryIngestionResult, + DEFAULT_MEMORY_EXTRACTION_MODEL, }; pub use pipeline::{ diff --git a/src/memory/ingest/types.rs b/src/memory/ingest/types.rs index 849b4c9..11f018e 100644 --- a/src/memory/ingest/types.rs +++ b/src/memory/ingest/types.rs @@ -12,7 +12,7 @@ use crate::memory::chunks::RawRef; /// is handed to a [`TreeJobSink`] for downstream extraction/admission/tree /// append/seal. A host wires the real queue behind this trait; tests use an /// in-memory recorder. -pub trait TreeJobSink { +pub trait TreeJobSink: Send + Sync { /// Enqueue an `extract_chunk` job for `chunk_id`. The downstream worker runs /// full extraction, the admission gate, buffer append, and sealing — none of /// which happen on this hot path. @@ -48,7 +48,7 @@ pub struct IngestOptions { /// downstream tree append; the actual buffer append and summary seal happen in /// the (separately ported) async worker, so this summary reports work *pending*, /// not completed. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct IngestSummary { /// Logical source id this call ingested. pub source_id: String, diff --git a/src/memory/mod.rs b/src/memory/mod.rs index 8e9c6ac..2892210 100644 --- a/src/memory/mod.rs +++ b/src/memory/mod.rs @@ -75,6 +75,9 @@ pub mod queue; pub mod retrieval; pub mod score; pub mod sources; +/// Live synchronization engine. Network code is opt-in through `sync`. +#[cfg(feature = "sync")] +pub mod sync; pub mod tool_memory; pub mod tree; diff --git a/src/memory/retrieval/cover.rs b/src/memory/retrieval/cover.rs index af0e752..1f6b937 100644 --- a/src/memory/retrieval/cover.rs +++ b/src/memory/retrieval/cover.rs @@ -26,7 +26,7 @@ use crate::memory::config::MemoryConfig; use crate::memory::tree::store::{get_tree_by_scope, list_summaries_in_window}; use crate::memory::tree::{SummaryNode, TreeKind}; -use super::types::{hit_from_chunk, hit_from_summary, QueryResponse, RetrievalHit}; +use super::types::{hydrated_chunk_hit, hydrated_summary_hit, QueryResponse, RetrievalHit}; /// Default cap on returned cover items when the caller passes `limit = 0`. const DEFAULT_LIMIT: usize = 200; @@ -77,6 +77,29 @@ pub fn cover_window( source_id: Option<&str>, source_kind: Option, limit: usize, +) -> Result { + cover_window_scoped( + config, + since_ms, + until_ms, + source_id, + source_kind, + None, + limit, + ) +} + +/// Compute a cover while restricting memory-source chunks to `source_scope`. +/// The filter runs before the scan cap and result limit so disallowed sources +/// cannot crowd permitted sources out of the response. +pub fn cover_window_scoped( + config: &MemoryConfig, + since_ms: i64, + until_ms: i64, + source_id: Option<&str>, + source_kind: Option, + source_scope: Option>, + limit: usize, ) -> Result { let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; if until_ms < since_ms { @@ -85,7 +108,14 @@ pub fn cover_window( )); } - let mut hits = collect_cover(config, since_ms, until_ms, source_id, source_kind)?; + let mut hits = collect_cover( + config, + since_ms, + until_ms, + source_id, + source_kind, + source_scope, + )?; // Group by source, then chronological ascending within each source. hits.sort_by(|a, b| { @@ -107,7 +137,9 @@ fn collect_cover( until_ms: i64, source_id: Option<&str>, source_kind: Option, + source_scope: Option>, ) -> Result> { + let restrict_summaries = source_id.is_some() || source_scope.is_some(); let chunks = list_chunks( config, &ListChunksQuery { @@ -117,6 +149,7 @@ fn collect_cover( until_ms: Some(until_ms), limit: Some(MAX_WINDOW_CHUNKS), exclude_dropped: true, + source_scope, ..Default::default() }, )?; @@ -130,7 +163,6 @@ fn collect_cover( // An exact `source_id` filter means `chunks` is a strict subset of its // (possibly shared) tree, so shared-tree summaries must be restricted to the // requested leaves. - let exact_source = source_id.is_some(); let mut hits: Vec = Vec::new(); for (source, src_chunks) in by_source { cover_one_source( @@ -139,7 +171,7 @@ fn collect_cover( since_ms, until_ms, src_chunks, - exact_source, + restrict_summaries, &mut hits, )?; } @@ -173,7 +205,7 @@ fn cover_one_source( let by_id: HashMap<&str, &SummaryNode> = eligible.iter().map(|s| (s.id.as_str(), s)).collect(); for id in &plan.maximal_ids { if let Some(node) = by_id.get(id.as_str()) { - out.push(hit_from_summary(node, source)); + out.push(hydrated_summary_hit(config, node, source)); } } @@ -181,7 +213,7 @@ fn cover_one_source( if plan.covered_chunk_ids.contains(&chunk.id) || suppressed_chunk_ids.contains(&chunk.id) { continue; } - out.push(hit_from_chunk(chunk, &tree_id, source, 0.0)); + out.push(hydrated_chunk_hit(config, chunk, &tree_id, source, 0.0)); } Ok(()) } diff --git a/src/memory/retrieval/cover_tests.rs b/src/memory/retrieval/cover_tests.rs index 832a89f..8b0c021 100644 --- a/src/memory/retrieval/cover_tests.rs +++ b/src/memory/retrieval/cover_tests.rs @@ -144,3 +144,29 @@ fn cover_window_emits_raw_chunks_when_no_tree() { .iter() .all(|h| h.node_kind == crate::memory::retrieval::NodeKind::Leaf)); } + +#[test] +fn scoped_cover_filters_memory_sources_before_result_limit() { + let (_tmp, cfg) = test_config(); + let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let mut denied = sample_chunk_at("slack:#denied", 0, "denied", ts); + denied.metadata.tags.push("memory_sources".to_string()); + let mut allowed = sample_chunk_at("slack:#allowed", 0, "allowed", ts); + allowed.metadata.tags.push("memory_sources".to_string()); + insert_chunks(&cfg, &[denied, allowed]); + + let scope = Some(["slack:#allowed".to_string()].into_iter().collect()); + let response = cover_window_scoped( + &cfg, + ts.timestamp_millis() - 1, + ts.timestamp_millis() + 1, + None, + None, + scope, + 1, + ) + .unwrap(); + + assert_eq!(response.hits.len(), 1); + assert_eq!(response.hits[0].tree_scope, "slack:#allowed"); +} diff --git a/src/memory/retrieval/drill_down.rs b/src/memory/retrieval/drill_down.rs index f49c079..4c4f133 100644 --- a/src/memory/retrieval/drill_down.rs +++ b/src/memory/retrieval/drill_down.rs @@ -35,7 +35,7 @@ use crate::memory::score::embed::Embedder; use crate::memory::tree::store::{get_summaries_batch, get_summary, get_tree, get_trees_batch}; use super::rerank::rerank_by_semantic_similarity; -use super::types::{hit_from_chunk, hit_from_summary, RetrievalHit}; +use super::types::{hydrated_chunk_hit, hydrated_summary_hit, RetrievalHit}; /// Pre-size hint for the next-level BFS frontier. const EXPECTED_CHILD_FANOUT: usize = 10; @@ -195,7 +195,7 @@ fn walk_with_embeddings( .unwrap_or_else(|| root_tree_scope.clone()); embeddings.push(summary.embedding.clone()); let child_ids = summary.child_ids.clone(); - out.push(hit_from_summary(&summary, &scope)); + out.push(hydrated_summary_hit(config, &summary, &scope)); if depth < max_depth { next_level.extend(child_ids); } @@ -204,7 +204,13 @@ fn walk_with_embeddings( if let Some(chunk) = chunk_by_id.remove(id) { let emb = emb_by_id.get(id).cloned(); embeddings.push(emb); - out.push(hit_from_chunk(&chunk, "", &chunk.metadata.source_id, 0.0)); + out.push(hydrated_chunk_hit( + config, + &chunk, + "", + &chunk.metadata.source_id, + 0.0, + )); continue; } // Child points at nothing (missing row) — skip it. diff --git a/src/memory/retrieval/fast.rs b/src/memory/retrieval/fast.rs new file mode 100644 index 0000000..70a09f6 --- /dev/null +++ b/src/memory/retrieval/fast.rs @@ -0,0 +1,360 @@ +//! Deterministic graph-routed retrieval over explicit query entity ids. + +use std::cmp::Reverse; +use std::collections::{HashMap, HashSet}; + +use anyhow::Result; + +use crate::memory::config::MemoryConfig; +use crate::memory::graph::{pair_distances, PairDistance}; +use crate::memory::score::embed::Embedder; +use crate::memory::score::store::lookup_entity; +use crate::memory::tree::store::{get_summaries_batch, get_tree}; + +use super::fetch::{fetch_leaves, MAX_BATCH}; +use super::source::query_source; +use super::types::{hydrated_summary_hit, QueryResponse}; + +const LOOKUP_LIMIT: usize = 500; +const DEFAULT_LIMIT: usize = 10; +const MAX_RETRIEVE_LIMIT: usize = 100; +const DEFAULT_MAX_HOPS: u32 = 2; +const MAX_GRAPH_HOPS: u32 = 4; + +#[derive(Clone, Debug)] +pub struct FastRetrieveOptions { + pub limit: usize, + pub max_hops: u32, + pub time_window_days: Option, +} + +impl Default for FastRetrieveOptions { + fn default() -> Self { + Self { + limit: DEFAULT_LIMIT, + max_hops: DEFAULT_MAX_HOPS, + time_window_days: None, + } + } +} + +pub async fn fast_retrieve( + config: &MemoryConfig, + query: &str, + query_entity_ids: &[String], + embedder: &dyn Embedder, + source_scope: Option<&HashSet>, + options: FastRetrieveOptions, +) -> Result { + let limit = if options.limit == 0 { + DEFAULT_LIMIT + } else { + options.limit.min(MAX_RETRIEVE_LIMIT) + }; + let max_hops = if options.max_hops == 0 { + DEFAULT_MAX_HOPS + } else { + options.max_hops.min(MAX_GRAPH_HOPS) + }; + let query = query.trim(); + if query.is_empty() { + return Ok(QueryResponse::empty()); + } + let entity_ids = dedup_ids(query_entity_ids.iter().cloned()); + log::debug!( + "[memory_retrieval:fast] entities={} limit={} hops={}", + entity_ids.len(), + limit, + max_hops + ); + if entity_ids.is_empty() { + return dense( + config, + query, + embedder, + source_scope, + limit, + options.time_window_days, + ) + .await; + } + let pairs = pair_distances(config, &entity_ids, max_hops)?; + if pairs.is_empty() { + return global_occurrence( + config, + query, + &entity_ids, + embedder, + source_scope, + limit, + options.time_window_days, + ) + .await; + } + let mut hops = max_hops; + let mut candidates = local_candidates(config, &pairs)?; + let mut last_non_empty = candidates.clone(); + while candidates.len() > limit && hops > 1 { + hops -= 1; + let next = local_candidates(config, &pair_distances(config, &entity_ids, hops)?)?; + if next.is_empty() { + break; + } + last_non_empty = next.clone(); + candidates = next; + } + if candidates.is_empty() { + candidates = last_non_empty; + } + if candidates.is_empty() { + return global_occurrence( + config, + query, + &entity_ids, + embedder, + source_scope, + limit, + options.time_window_days, + ) + .await; + } + resolve_local(config, candidates, source_scope, limit) +} + +#[derive(Clone, Debug)] +struct Candidate { + node_kind: String, + matched: HashSet, + latest_ts: i64, +} + +fn local_candidates( + config: &MemoryConfig, + pairs: &[PairDistance], +) -> Result> { + let mut output = HashMap::new(); + for pair in pairs { + let left = lookup_entity(config, &pair.a, Some(LOOKUP_LIMIT))?; + let right = lookup_entity(config, &pair.b, Some(LOOKUP_LIMIT))?; + let right_nodes: HashMap<_, _> = right + .iter() + .map(|hit| (hit.node_id.as_str(), hit.timestamp_ms)) + .collect(); + for hit in left { + let Some(right_ts) = right_nodes.get(hit.node_id.as_str()) else { + continue; + }; + let candidate = output.entry(hit.node_id).or_insert_with(|| Candidate { + node_kind: hit.node_kind, + matched: HashSet::new(), + latest_ts: 0, + }); + candidate.matched.insert(pair.a.clone()); + candidate.matched.insert(pair.b.clone()); + candidate.latest_ts = candidate.latest_ts.max(hit.timestamp_ms).max(*right_ts); + } + } + Ok(output) +} + +fn resolve_local( + config: &MemoryConfig, + candidates: HashMap, + source_scope: Option<&HashSet>, + limit: usize, +) -> Result { + let mut ordered: Vec<_> = candidates.into_iter().collect(); + ordered.sort_by(|a, b| { + b.1.matched + .len() + .cmp(&a.1.matched.len()) + .then_with(|| b.1.latest_ts.cmp(&a.1.latest_ts)) + .then_with(|| a.0.cmp(&b.0)) + }); + let coverage: HashMap<_, _> = ordered + .iter() + .map(|(id, candidate)| (id.clone(), candidate.matched.len() as f32)) + .collect(); + let leaf_ids: Vec<_> = ordered + .iter() + .filter(|(_, candidate)| candidate.node_kind == "leaf") + .map(|(id, _)| id.clone()) + .collect(); + let summary_ids: Vec<_> = ordered + .iter() + .filter(|(_, candidate)| candidate.node_kind != "leaf") + .map(|(id, _)| id.clone()) + .collect(); + let mut by_id = HashMap::new(); + for ids in leaf_ids.chunks(MAX_BATCH) { + for hit in fetch_leaves(config, ids)? { + by_id.insert(hit.node_id.clone(), hit); + } + } + let summaries = get_summaries_batch(config, &summary_ids)?; + let mut scope_cache: HashMap = HashMap::new(); + for (id, node) in summaries { + let scope = match scope_cache.get(&node.tree_id) { + Some(scope) => scope.clone(), + None => { + let scope = get_tree(config, &node.tree_id)? + .map(|tree| tree.scope) + .unwrap_or_default(); + scope_cache.insert(node.tree_id.clone(), scope.clone()); + scope + } + }; + by_id.insert(id, hydrated_summary_hit(config, &node, &scope)); + } + let mut hits = Vec::new(); + for (id, _) in ordered { + if let Some(mut hit) = by_id.remove(&id) { + if source_scope.is_some_and(|scope| !scope.contains(&hit.tree_scope)) { + continue; + } + hit.score = coverage.get(&id).copied().unwrap_or_default(); + hits.push(hit); + } + } + let total = hits.len(); + hits.truncate(limit); + Ok(QueryResponse::new(hits, total)) +} + +async fn dense( + config: &MemoryConfig, + query: &str, + embedder: &dyn Embedder, + source_scope: Option<&HashSet>, + limit: usize, + window: Option, +) -> Result { + let mut response = query_source( + config, + None, + None, + window, + Some(query), + embedder, + usize::MAX, + ) + .await?; + if let Some(scope) = source_scope { + response.hits.retain(|hit| scope.contains(&hit.tree_scope)); + } + let total = response.hits.len(); + response.hits.truncate(limit); + Ok(QueryResponse::new(response.hits, total)) +} + +async fn global_occurrence( + config: &MemoryConfig, + query: &str, + entity_ids: &[String], + embedder: &dyn Embedder, + source_scope: Option<&HashSet>, + limit: usize, + window: Option, +) -> Result { + let mut response = dense( + config, + query, + embedder, + source_scope, + limit.saturating_mul(2), + window, + ) + .await?; + let ids: HashSet<_> = entity_ids.iter().map(String::as_str).collect(); + response.hits.sort_by_key(|hit| { + Reverse( + hit.entities + .iter() + .filter(|entity| ids.contains(entity.as_str())) + .count(), + ) + }); + response.hits.truncate(limit); + Ok(QueryResponse::new(response.hits, response.total)) +} + +fn dedup_ids(ids: impl Iterator) -> Vec { + let mut seen = HashSet::new(); + ids.filter(|id| seen.insert(id.clone())).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::graph::{pairs_from_entities, upsert_edges}; + use crate::memory::retrieval::test_support::{ + fixed_ts, index_entity_occurrence, insert_summary, insert_tree_row, source_tree, + summary_node, test_config, + }; + use crate::memory::score::embed::InertEmbedder; + use crate::memory::score::extract::EntityKind; + + #[test] + fn options_and_ids_are_bounded_deterministically() { + assert_eq!(FastRetrieveOptions::default().limit, 10); + assert_eq!( + dedup_ids(vec!["a".into(), "a".into(), "b".into()].into_iter()), + vec!["a", "b"] + ); + } + + #[tokio::test] + async fn local_branch_intersects_entities_and_applies_scope_before_limit() { + let (_temp, config) = test_config(); + let allowed_tree = source_tree("allowed-tree", "slack:#allowed", Some("allowed"), 1); + let denied_tree = source_tree("denied-tree", "slack:#denied", Some("denied"), 1); + insert_tree_row(&config, &allowed_tree); + insert_tree_row(&config, &denied_tree); + for (id, tree_id) in [("allowed", "allowed-tree"), ("denied", "denied-tree")] { + insert_summary( + &config, + &summary_node(id, tree_id, 1, None, &[], id, fixed_ts()), + ); + for (entity, kind) in [ + ("person:alice", EntityKind::Person), + ("topic:launch", EntityKind::Topic), + ] { + index_entity_occurrence( + &config, + entity, + kind, + entity, + id, + "summary", + fixed_ts().timestamp_millis(), + Some(tree_id), + ); + } + } + upsert_edges( + &config, + &pairs_from_entities(&["person:alice".into(), "topic:launch".into()]), + fixed_ts().timestamp_millis(), + ) + .unwrap(); + let scope = HashSet::from(["slack:#allowed".to_string()]); + + let response = fast_retrieve( + &config, + "alice launch", + &["person:alice".into(), "topic:launch".into()], + &InertEmbedder, + Some(&scope), + FastRetrieveOptions { + limit: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(response.total, 1); + assert_eq!(response.hits[0].node_id, "allowed"); + assert_eq!(response.hits[0].score, 2.0); + } +} diff --git a/src/memory/retrieval/fetch.rs b/src/memory/retrieval/fetch.rs index f0ed8cc..e92fd2b 100644 --- a/src/memory/retrieval/fetch.rs +++ b/src/memory/retrieval/fetch.rs @@ -17,7 +17,7 @@ use crate::memory::chunks::get_chunks_batch; use crate::memory::config::MemoryConfig; use crate::memory::score::store::get_scores_batch; -use super::types::{hit_from_chunk, RetrievalHit}; +use super::types::{hydrated_chunk_hit, RetrievalHit}; /// Max batch size. Callers that pass more than this get truncated (no error so /// the caller still sees a partial result). @@ -49,7 +49,7 @@ pub fn fetch_leaves(config: &MemoryConfig, chunk_ids: &[String]) -> Result = Vec::new(); for i in 0..(MAX_BATCH + 5) as u32 { let c = sample_chunk("slack:#eng", i, &format!("content-{i}")); - insert_chunks(&cfg, &[c.clone()]); + insert_chunks(&cfg, std::slice::from_ref(&c)); ids.push(c.id); } let out = fetch_leaves(&cfg, &ids).unwrap(); @@ -55,8 +55,8 @@ fn over_cap_is_truncated() { fn leaf_hit_carries_source_ref_and_scope() { let (_tmp, cfg) = test_config(); let c = sample_chunk("slack:#eng", 0, "content-0"); - insert_chunks(&cfg, &[c.clone()]); - let out = fetch_leaves(&cfg, &[c.id.clone()]).unwrap(); + insert_chunks(&cfg, std::slice::from_ref(&c)); + let out = fetch_leaves(&cfg, std::slice::from_ref(&c.id)).unwrap(); assert_eq!(out.len(), 1); assert_eq!(out[0].source_ref.as_deref(), Some("slack://slack:#eng/0")); assert_eq!(out[0].tree_scope, "slack:#eng"); @@ -96,3 +96,26 @@ fn fetch_leaves_preserves_input_order_and_propagates_scores() { ); assert!((out[2].score - 0.1).abs() < 1e-6, "c1 score"); } + +#[test] +fn fetch_leaves_hydrates_full_staged_body_instead_of_sql_preview() { + let (_tmp, cfg) = test_config(); + let full_body = "full-body ".repeat(120); + let chunk = sample_chunk("slack:#eng", 0, &full_body); + let staged = crate::memory::store::content::stage_chunks( + &crate::memory::chunks::content_root(&cfg), + std::slice::from_ref(&chunk), + ) + .unwrap(); + crate::memory::chunks::with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::memory::chunks::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + + let out = fetch_leaves(&cfg, &[chunk.id]).unwrap(); + assert_eq!(out[0].content, full_body); + assert!(out[0].content.len() > 500); +} diff --git a/src/memory/retrieval/global_tests.rs b/src/memory/retrieval/global_tests.rs index 31b3fef..91db0c7 100644 --- a/src/memory/retrieval/global_tests.rs +++ b/src/memory/retrieval/global_tests.rs @@ -125,7 +125,7 @@ async fn query_topic_resolves_indexed_leaf() { let (_tmp, cfg) = test_config(); let ts = fixed_ts(); let chunk = sample_chunk("slack:#eng", 0, "phoenix launch notes"); - insert_chunks(&cfg, &[chunk.clone()]); + insert_chunks(&cfg, std::slice::from_ref(&chunk)); index_entity_occurrence( &cfg, "topic:phoenix", diff --git a/src/memory/retrieval/mod.rs b/src/memory/retrieval/mod.rs index 3a3a586..be860c8 100644 --- a/src/memory/retrieval/mod.rs +++ b/src/memory/retrieval/mod.rs @@ -54,6 +54,7 @@ pub mod cover; pub mod drill_down; +pub mod fast; pub mod fetch; pub mod global; pub mod graph_adapter; @@ -72,8 +73,9 @@ pub(crate) mod test_support; // ── Public re-exports ─────────────────────────────────────────────────────── -pub use cover::cover_window; +pub use cover::{cover_window, cover_window_scoped}; pub use drill_down::drill_down; +pub use fast::{fast_retrieve, FastRetrieveOptions}; pub use fetch::{fetch_leaves, MAX_BATCH}; pub use global::{query_global, query_topic}; pub use graph_adapter::ConfigEntityIndex; diff --git a/src/memory/retrieval/source.rs b/src/memory/retrieval/source.rs index d142a8a..eb87301 100644 --- a/src/memory/retrieval/source.rs +++ b/src/memory/retrieval/source.rs @@ -37,7 +37,7 @@ use crate::memory::tree::store::{ use crate::memory::tree::{Tree, TreeKind}; use super::rerank::rerank_by_semantic_similarity; -use super::types::{hit_from_summary, QueryResponse, RetrievalHit}; +use super::types::{hydrated_summary_hit, QueryResponse, RetrievalHit}; /// Default result cap applied when the caller passes `limit = 0`. const DEFAULT_LIMIT: usize = 10; @@ -128,7 +128,7 @@ pub(crate) fn collect_source_hits( } node_ids.push(node.id.clone()); embeddings.push(node.embedding.clone()); - hits.push(hit_from_summary(&node, &tree.scope)); + hits.push(hydrated_summary_hit(config, &node, &tree.scope)); } } } diff --git a/src/memory/retrieval/test_support.rs b/src/memory/retrieval/test_support.rs index 2633245..4826131 100644 --- a/src/memory/retrieval/test_support.rs +++ b/src/memory/retrieval/test_support.rs @@ -109,6 +109,7 @@ pub fn set_summary_embedding(cfg: &MemoryConfig, summary_id: &str, vec: &[f32]) } /// Index one entity occurrence against a node. +#[allow(clippy::too_many_arguments)] pub fn index_entity_occurrence( cfg: &MemoryConfig, canonical_id: &str, diff --git a/src/memory/retrieval/types.rs b/src/memory/retrieval/types.rs index 4d1a294..df19138 100644 --- a/src/memory/retrieval/types.rs +++ b/src/memory/retrieval/types.rs @@ -179,6 +179,40 @@ pub fn hit_from_chunk(chunk: &Chunk, tree_id: &str, tree_scope: &str, score: f32 } } +pub(crate) fn hydrated_summary_hit( + config: &crate::memory::MemoryConfig, + node: &SummaryNode, + tree_scope: &str, +) -> RetrievalHit { + let mut hit = hit_from_summary(node, tree_scope); + match crate::memory::store::content::read_summary_body(config, &node.id) { + Ok(body) => hit.content = body, + Err(error) => log::warn!( + "[memory:retrieval] full summary body unavailable; serving preview node_id={}: {error}", + node.id + ), + } + hit +} + +pub(crate) fn hydrated_chunk_hit( + config: &crate::memory::MemoryConfig, + chunk: &Chunk, + tree_id: &str, + tree_scope: &str, + score: f32, +) -> RetrievalHit { + let mut hit = hit_from_chunk(chunk, tree_id, tree_scope, score); + match crate::memory::store::content::read_chunk_body(config, &chunk.id) { + Ok(body) => hit.content = body, + Err(error) => log::warn!( + "[memory:retrieval] full chunk body unavailable; serving preview chunk_id={}: {error}", + chunk.id + ), + } + hit +} + /// Decide the placeholder [`TreeKind`] to report on a leaf hit. Leaves live /// under source trees regardless of the underlying [`SourceKind`], so we always /// return [`TreeKind::Source`]. Accepting the `SourceKind` argument keeps the diff --git a/src/memory/score/mod.rs b/src/memory/score/mod.rs index 1ae0917..2a2ec86 100644 --- a/src/memory/score/mod.rs +++ b/src/memory/score/mod.rs @@ -288,8 +288,6 @@ pub async fn score_chunk(chunk: &Chunk, cfg: &ScoringConfig) -> Result = entities + .iter() + .map(|entity| entity.canonical_id.clone()) + .collect(); + upsert_edges_tx(&tx, &pairs_from_entities(&entity_ids), timestamp_ms)?; tx.commit()?; Ok(entities.len()) }) @@ -364,6 +370,8 @@ pub fn index_summary_entity_ids_tx( canonical_id_is_user(canonical_id) as i32, ])?; } + drop(stmt); + upsert_edges_tx(tx, &pairs_from_entities(entity_ids), timestamp_ms)?; Ok(entity_ids.len()) } @@ -395,6 +403,12 @@ pub fn index_entities_tx( entity_is_user(e) as i32, ])?; } + drop(stmt); + let entity_ids: Vec<_> = entities + .iter() + .map(|entity| entity.canonical_id.clone()) + .collect(); + upsert_edges_tx(tx, &pairs_from_entities(&entity_ids), timestamp_ms)?; Ok(entities.len()) } diff --git a/src/memory/score/store_tests.rs b/src/memory/score/store_tests.rs index b77033f..2d40c70 100644 --- a/src/memory/score/store_tests.rs +++ b/src/memory/score/store_tests.rs @@ -103,6 +103,7 @@ fn index_batch() { let n = index_entities(&cfg, &entities, "chunk-1", "leaf", 1000, None).unwrap(); assert_eq!(n, 3); assert_eq!(count_entity_index(&cfg).unwrap(), 3); + assert_eq!(crate::memory::graph::count_edges(&cfg).unwrap(), 3); } #[test] @@ -234,5 +235,5 @@ fn get_scores_batch_empty_input_and_missing_chunk_ids() { let map = get_scores_batch(&cfg, &ids).unwrap(); assert_eq!(map.len(), 1); assert!((map.get("c1").copied().unwrap() - 0.7).abs() < 1e-6); - assert!(map.get("ghost:no-such").is_none()); + assert!(!map.contains_key("ghost:no-such")); } diff --git a/src/memory/sources/mod.rs b/src/memory/sources/mod.rs index a4dc670..95450c9 100644 --- a/src/memory/sources/mod.rs +++ b/src/memory/sources/mod.rs @@ -29,6 +29,8 @@ pub mod types; pub mod validation; pub use readers::{is_locally_readable, reader_for, SourceReader}; -pub use registry::{memory_sync_defaults_for_toolkit, MemorySourcePatch, SourceRegistry}; +pub use registry::{ + memory_sync_defaults_for_toolkit, ComposioUpsertTarget, MemorySourcePatch, SourceRegistry, +}; pub use types::{ContentType, MemorySourceEntry, SourceContent, SourceItem, SourceKind}; pub use validation::{ensure_within_base, validate_entry}; diff --git a/src/memory/sources/registry.rs b/src/memory/sources/registry.rs index 7e389e3..d564b43 100644 --- a/src/memory/sources/registry.rs +++ b/src/memory/sources/registry.rs @@ -249,6 +249,19 @@ impl SourceRegistry { Ok(entry) } + /// Batch-upsert Composio sources with one load and one atomic save. + pub fn upsert_composio_sources_batch(&self, targets: &[ComposioUpsertTarget]) -> Result { + if targets.is_empty() { + return Ok(0); + } + let mut sources = self.list()?; + for (toolkit, connection_id, label) in targets { + upsert_composio_entry_in_place(&mut sources, toolkit, connection_id, label); + } + self.write_all(&sources)?; + Ok(targets.len().min(u32::MAX as usize) as u32) + } + /// Enable every source and clear all per-source caps ("All In" mode). pub fn apply_all_in(&self) -> Result> { let mut sources = self.list()?; @@ -268,6 +281,8 @@ impl SourceRegistry { } } +pub type ComposioUpsertTarget = (String, String, String); + /// Apply a single composio upsert to an in-memory source list. /// /// Pure (no I/O) so the registry path and unit tests share one find-or-push diff --git a/src/memory/sources/types.rs b/src/memory/sources/types.rs index d03c6ed..5f3b756 100644 --- a/src/memory/sources/types.rs +++ b/src/memory/sources/types.rs @@ -14,6 +14,7 @@ //! Wire strings are snake_case and are part of the persisted contract — do not //! rename them when porting from OpenHuman. +use schemars::JsonSchema; use serde::{Deserialize, Serialize}; pub(crate) fn default_true() -> bool { @@ -24,7 +25,7 @@ pub(crate) fn default_true() -> bool { /// /// The wire representation is snake_case (`github_repo`, `rss_feed`, …) and is /// persisted in `config.toml`; it must stay stable across versions. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum SourceKind { /// A Composio OAuth connector (Gmail, Slack, Notion, …). Network-backed; @@ -65,7 +66,7 @@ impl SourceKind { /// [`kind`](MemorySourceEntry::kind) discriminator determines which fields are /// required; validation is enforced at add/update time via /// [`MemorySourceEntry::validate`]. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] pub struct MemorySourceEntry { /// Stable unique id (e.g. `src_`). pub id: String, @@ -159,7 +160,7 @@ impl MemorySourceEntry { /// /// `id` is reader-scoped (e.g. a folder-relative path or a thread id) and is /// stable enough to pass back into [`crate::memory::sources::readers::SourceReader::read_item`]. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] pub struct SourceItem { /// Reader-scoped item id. pub id: String, @@ -171,7 +172,7 @@ pub struct SourceItem { } /// The rendered content type of a [`SourceContent`] body. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum ContentType { /// Markdown body. @@ -183,7 +184,7 @@ pub enum ContentType { } /// Content read from a single source item. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] pub struct SourceContent { /// Reader-scoped item id (matches the [`SourceItem::id`] it was read from). pub id: String, diff --git a/src/memory/store/content/atomic.rs b/src/memory/store/content/atomic.rs index 9bbe469..00237ca 100644 --- a/src/memory/store/content/atomic.rs +++ b/src/memory/store/content/atomic.rs @@ -64,6 +64,32 @@ pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result { } } +/// Ensure `abs_path` contains `full_bytes` with the expected body digest. +/// +/// Matching files are left untouched. Missing, malformed, or stale files are +/// replaced atomically so callers can safely persist `body_sha256` alongside +/// the returned content path. +pub fn write_or_replace_body( + abs_path: &Path, + full_bytes: &[u8], + body_sha256: &str, +) -> anyhow::Result<()> { + if abs_path.exists() { + let disk_sha = read_body_sha256(abs_path).unwrap_or_default(); + if disk_sha == body_sha256 { + return Ok(()); + } + } + + if let Some(parent) = abs_path.parent() { + std::fs::create_dir_all(parent) + .map_err(|e| anyhow::anyhow!("create content parent {:?}: {e}", parent))?; + } + + crate::memory::fsutil::atomic_write(abs_path, full_bytes) + .map_err(|e| anyhow::anyhow!("atomic content write {:?}: {e}", abs_path)) +} + /// A summary that has been written to disk and is ready for SQLite upsert. #[derive(Debug, Clone)] pub struct StagedSummary { @@ -115,24 +141,7 @@ pub fn stage_summary_with_layout( let full_bytes = composed.full.as_bytes(); - // Idempotent re-stage: matching on-disk body sha → return unchanged. - // Mismatch → the file is stale; overwrite it via temp-file + rename so the - // destination is only ever the old or the new file, never missing between a - // `remove_file` and a re-write if the process crashes. - if abs_path.exists() { - let disk_sha = read_body_sha256(&abs_path).unwrap_or_default(); - if disk_sha == sha256 { - return Ok(StagedSummary { - summary_id: input.summary_id.to_string(), - content_path: rel_path, - content_sha256: sha256, - }); - } - crate::memory::fsutil::atomic_write(&abs_path, full_bytes) - .map_err(|e| anyhow::anyhow!("atomic overwrite {:?}: {e}", abs_path))?; - } else { - write_if_new(&abs_path, full_bytes)?; - } + write_or_replace_body(&abs_path, full_bytes, &sha256)?; Ok(StagedSummary { summary_id: input.summary_id.to_string(), diff --git a/src/memory/store/content/content_tests.rs b/src/memory/store/content/content_tests.rs index 55f9827..58eed69 100644 --- a/src/memory/store/content/content_tests.rs +++ b/src/memory/store/content/content_tests.rs @@ -59,6 +59,29 @@ fn stage_chunks_is_idempotent() { assert_eq!(first[0].content_path, second[0].content_path); } +#[test] +fn stage_chunks_replaces_stale_on_disk_body() { + let dir = TempDir::new().unwrap(); + let chunk = sample_chunk(0); + let abs = paths::chunk_abs_path( + dir.path(), + chunk.metadata.source_kind.as_str(), + &chunk.metadata.source_id, + &chunk.id, + ); + std::fs::create_dir_all(abs.parent().unwrap()).unwrap(); + std::fs::write(&abs, b"---\nstale: true\n---\nSTALE BODY").unwrap(); + + let staged = stage_chunks(dir.path(), std::slice::from_ref(&chunk)).unwrap(); + let on_disk = std::fs::read_to_string(&abs).unwrap(); + assert!(on_disk.ends_with(&chunk.content)); + let (_, body) = compose::split_front_matter(&on_disk).unwrap(); + assert_eq!( + staged[0].content_sha256, + atomic::sha256_hex(body.as_bytes()) + ); +} + #[test] fn stage_chunks_email_skips_disk_write() { let dir = TempDir::new().unwrap(); diff --git a/src/memory/store/content/mod.rs b/src/memory/store/content/mod.rs index a927713..5905894 100644 --- a/src/memory/store/content/mod.rs +++ b/src/memory/store/content/mod.rs @@ -43,12 +43,13 @@ pub use paths::{ SummaryDiskLayout, SummaryTreeKind, }; pub use raw::{ - raw_kind_dir, raw_rel_path, raw_source_dir, slug_account_email, write_raw_items, RawItem, - RawKind, + raw_kind_dir, raw_rel_path, raw_source_dir, sanitize_uid, slug_account_email, write_raw_items, + RawItem, RawKind, }; pub use read::{ - read_chunk_file, read_summary_file, resolve_within_content_root, verify_chunk_file, - verify_summary_file, ChunkFileContents, VerifyResult, + read_chunk_body, read_chunk_file, read_summary_body, read_summary_file, + resolve_within_content_root, verify_chunk_file, verify_summary_file, ChunkFileContents, + VerifyResult, }; pub use tags::{entity_tag, slugify_tag_kind, slugify_tag_value, update_chunk_tags}; @@ -88,7 +89,7 @@ pub fn stage_chunks(content_root: &Path, chunks: &[Chunk]) -> anyhow::Result String { /// Replace path-illegal characters in the upstream uid before splicing it into /// a filename. -pub(crate) fn sanitize_uid(uid: &str) -> String { +/// Sanitize an upstream item identifier for use as a raw-archive filename. +pub fn sanitize_uid(uid: &str) -> String { let cleaned: String = uid .chars() .map(|c| match c { @@ -183,8 +184,8 @@ pub fn slug_account_email(email: &str) -> String { let lower = email.trim().to_lowercase(); let mut out = String::with_capacity(lower.len() + 8); let mut last_dash = true; - let mut chars = lower.chars().peekable(); - while let Some(ch) = chars.next() { + let chars = lower.chars().peekable(); + for ch in chars { match ch { '@' => { if !last_dash { diff --git a/src/memory/store/content/read.rs b/src/memory/store/content/read.rs index c863364..fa3babf 100644 --- a/src/memory/store/content/read.rs +++ b/src/memory/store/content/read.rs @@ -1,16 +1,17 @@ //! Read and verify chunk and summary `.md` files from the content store. //! -//! This port covers the self-contained read + integrity surface: front-matter -//! splitting, body SHA-256 verification, and untrusted-path resolution. The -//! Config/SQLite-aware high-level body readers (`read_chunk_body` / -//! `read_summary_body`, which resolve a chunk id → on-disk path through the -//! chunk store) are **deferred** along with the rest of the SQLite chunk-store -//! integration. +//! Includes high-level id-based readers that resolve SQLite pointers, hydrate +//! raw-backed chunks, and repair stale checksum tokens after external edits. use std::path::{Component, Path, PathBuf}; use super::atomic::sha256_hex; use super::compose::split_front_matter; +use crate::memory::chunks::{ + content_root, get_chunk_content_pointers, get_chunk_raw_refs, get_summary_content_pointers, + update_chunk_content_sha256, update_summary_content_sha256, RawRef, +}; +use crate::memory::config::MemoryConfig; /// Resolve a DB-stored relative forward-slash path against `content_root`, /// rejecting any traversal (`..`), absolute, or non-normal component. @@ -125,6 +126,85 @@ pub fn verify_summary_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Re } } +/// Read a full chunk body by id. Raw archive references take precedence over +/// staged Markdown pointers. A drifted staged-file checksum is repaired +/// best-effort while the authoritative on-disk body is still returned. +pub fn read_chunk_body(config: &MemoryConfig, chunk_id: &str) -> anyhow::Result { + if let Some(refs) = get_chunk_raw_refs(config, chunk_id)? { + if !refs.is_empty() { + return read_chunk_body_from_raw(config, &refs); + } + } + + let (rel_path, expected_sha256) = get_chunk_content_pointers(config, chunk_id)? + .ok_or_else(|| anyhow::anyhow!("no content pointer or raw refs for chunk {chunk_id}"))?; + if rel_path.is_empty() { + anyhow::bail!("empty content pointer and no raw refs for chunk {chunk_id}"); + } + let abs_path = resolve_within_content_root(&content_root(config), &rel_path)?; + let result = read_chunk_file(&abs_path)?; + if result.sha256 != expected_sha256 { + log::warn!("[memory:content] repairing stale chunk checksum token chunk_id={chunk_id}"); + if let Err(error) = update_chunk_content_sha256(config, chunk_id, &result.sha256) { + log::warn!( + "[memory:content] chunk checksum repair failed chunk_id={chunk_id}: {error}" + ); + } + } + Ok(result.body) +} + +fn read_chunk_body_from_raw(config: &MemoryConfig, refs: &[RawRef]) -> anyhow::Result { + let root = content_root(config); + let mut parts = Vec::with_capacity(refs.len()); + for reference in refs { + let path = match resolve_within_content_root(&root, &reference.path) { + Ok(path) => path, + Err(error) => { + log::warn!("[memory:content] rejected unsafe raw reference: {error}"); + continue; + } + }; + let bytes = match std::fs::read(path) { + Ok(bytes) => bytes, + Err(error) => { + log::warn!("[memory:content] raw reference read failed: {error}"); + continue; + } + }; + let start = reference.start.min(bytes.len()); + let end = reference.end.unwrap_or(bytes.len()).min(bytes.len()); + if end <= start { + continue; + } + match std::str::from_utf8(&bytes[start..end]) { + Ok(value) => parts.push(value.to_owned()), + Err(error) => log::warn!("[memory:content] raw reference is not UTF-8: {error}"), + } + } + Ok(parts.join("\n\n")) +} + +/// Read a full summary body by id and repair a stale checksum token +/// best-effort when an external editor changed the staged file. +pub fn read_summary_body(config: &MemoryConfig, summary_id: &str) -> anyhow::Result { + let (rel_path, expected_sha256) = get_summary_content_pointers(config, summary_id)? + .ok_or_else(|| anyhow::anyhow!("no content pointer for summary {summary_id}"))?; + let abs_path = resolve_within_content_root(&content_root(config), &rel_path)?; + let result = read_summary_file(&abs_path)?; + if result.sha256 != expected_sha256 { + log::warn!( + "[memory:content] repairing stale summary checksum token summary_id={summary_id}" + ); + if let Err(error) = update_summary_content_sha256(config, summary_id, &result.sha256) { + log::warn!( + "[memory:content] summary checksum repair failed summary_id={summary_id}: {error}" + ); + } + } + Ok(result.body) +} + #[cfg(test)] #[path = "read_tests.rs"] mod tests; diff --git a/src/memory/store/content/read_tests.rs b/src/memory/store/content/read_tests.rs index cdd55ed..b0b850d 100644 --- a/src/memory/store/content/read_tests.rs +++ b/src/memory/store/content/read_tests.rs @@ -172,3 +172,63 @@ fn resolve_within_content_root_rejects_traversal_and_absolute() { let ok = resolve_within_content_root(root, "sub/dir/file.md").unwrap(); assert_eq!(ok, root.join("sub").join("dir").join("file.md")); } + +#[test] +fn high_level_chunk_reader_uses_custom_root_and_repairs_checksum() { + let dir = TempDir::new().unwrap(); + let custom_root = dir.path().join("custom-content"); + let mut config = crate::memory::MemoryConfig::new(dir.path()); + config.content_root = Some(custom_root.clone()); + let chunk = sample_chunk(); + let staged = + crate::memory::store::content::stage_chunks(&custom_root, std::slice::from_ref(&chunk)) + .unwrap(); + crate::memory::chunks::with_connection(&config, |conn| { + let tx = conn.unchecked_transaction()?; + crate::memory::chunks::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + crate::memory::chunks::update_chunk_content_sha256(&config, &chunk.id, "stale").unwrap(); + + assert_eq!(read_chunk_body(&config, &chunk.id).unwrap(), chunk.content); + let (_, repaired) = crate::memory::chunks::get_chunk_content_pointers(&config, &chunk.id) + .unwrap() + .unwrap(); + assert_eq!(repaired, staged[0].content_sha256); +} + +#[test] +fn high_level_chunk_reader_joins_clamped_raw_references() { + let dir = TempDir::new().unwrap(); + let mut config = crate::memory::MemoryConfig::new(dir.path()); + let root = dir.path().join("raw-root"); + config.content_root = Some(root.clone()); + std::fs::create_dir_all(root.join("raw/source")).unwrap(); + std::fs::write(root.join("raw/source/item.md"), "alpha beta").unwrap(); + let chunk = sample_chunk(); + crate::memory::chunks::upsert_chunks(&config, std::slice::from_ref(&chunk)).unwrap(); + crate::memory::chunks::set_chunk_raw_refs( + &config, + &chunk.id, + &[ + crate::memory::chunks::RawRef { + path: "raw/source/item.md".into(), + start: 0, + end: Some(5), + }, + crate::memory::chunks::RawRef { + path: "raw/source/item.md".into(), + start: 6, + end: Some(100), + }, + ], + ) + .unwrap(); + + assert_eq!( + read_chunk_body(&config, &chunk.id).unwrap(), + "alpha\n\nbeta" + ); +} diff --git a/src/memory/store/entity_index/mod.rs b/src/memory/store/entity_index/mod.rs index 8bb11e9..5518b06 100644 --- a/src/memory/store/entity_index/mod.rs +++ b/src/memory/store/entity_index/mod.rs @@ -25,5 +25,8 @@ pub mod store; pub mod types; -pub use store::{index_entities_tx, EntityIndex, NoSelfIdentity, SelfIdentity}; +pub use store::{ + clear_entity_index_for_node_tx, index_entities_tx, index_entities_tx_with_identity, + index_summary_entity_ids_tx_with_identity, EntityIndex, NoSelfIdentity, SelfIdentity, +}; pub use types::{CanonicalEntity, EntityHit, EntityKind}; diff --git a/src/memory/store/entity_index/store.rs b/src/memory/store/entity_index/store.rs index 6a63c89..8e59a3f 100644 --- a/src/memory/store/entity_index/store.rs +++ b/src/memory/store/entity_index/store.rs @@ -41,10 +41,7 @@ impl SelfIdentity for NoSelfIdentity { /// /// Primary key `(entity_id, node_id)` makes re-indexing the same association a /// no-op update (idempotent). Indexes back both query directions. -const INIT_SQL: &str = " - PRAGMA journal_mode = WAL; - PRAGMA synchronous = NORMAL; - +const SCHEMA_SQL: &str = " CREATE TABLE IF NOT EXISTS mem_tree_entity_index ( entity_id TEXT NOT NULL, node_id TEXT NOT NULL, @@ -61,6 +58,11 @@ const INIT_SQL: &str = " CREATE INDEX IF NOT EXISTS idx_entity_index_node ON mem_tree_entity_index(node_id); "; +const OWNED_CONNECTION_PRAGMAS: &str = " + PRAGMA journal_mode = WAL; + PRAGMA synchronous = NORMAL; +"; + const UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_entity_index ( entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms, tree_id, is_user @@ -89,7 +91,7 @@ impl EntityIndex { std::fs::create_dir_all(parent)?; } let conn = Connection::open(db_path)?; - conn.execute_batch(INIT_SQL)?; + Self::init_owned_connection(&conn)?; Ok(Self { conn: Arc::new(Mutex::new(conn)), identity, @@ -104,13 +106,31 @@ impl EntityIndex { /// Open an in-memory index with a custom identity resolver. pub fn open_in_memory_with_identity(identity: Arc) -> anyhow::Result { let conn = Connection::open_in_memory()?; - conn.execute_batch(INIT_SQL)?; + Self::init_owned_connection(&conn)?; Ok(Self { conn: Arc::new(Mutex::new(conn)), identity, }) } + /// Use a connection owned and configured by the embedding application. + /// + /// Only the entity-index schema is installed. Connection pragmas and + /// transaction policy remain owned by the caller. + pub fn from_shared_connection( + conn: Arc>, + identity: Arc, + ) -> anyhow::Result { + conn.lock().execute_batch(SCHEMA_SQL)?; + Ok(Self { conn, identity }) + } + + fn init_owned_connection(conn: &Connection) -> anyhow::Result<()> { + conn.execute_batch(OWNED_CONNECTION_PRAGMAS)?; + conn.execute_batch(SCHEMA_SQL)?; + Ok(()) + } + /// Resolve `is_user` for a single-occurrence entity via the configured /// [`SelfIdentity`] resolver. fn is_user(&self, entity: &CanonicalEntity) -> bool { @@ -355,6 +375,27 @@ pub fn index_entities_tx( node_kind: &str, timestamp_ms: i64, tree_id: Option<&str>, +) -> anyhow::Result { + index_entities_tx_with_identity( + tx, + entities, + node_id, + node_kind, + timestamp_ms, + tree_id, + &NoSelfIdentity, + ) +} + +/// Identity-aware transaction-scoped batch index. +pub fn index_entities_tx_with_identity( + tx: &Transaction<'_>, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, + identity: &dyn SelfIdentity, ) -> anyhow::Result { if entities.is_empty() { return Ok(0); @@ -370,12 +411,56 @@ pub fn index_entities_tx( e.score, timestamp_ms, tree_id, - 0_i32, + identity.is_self(e.kind, &e.surface) as i32, ])?; } Ok(entities.len()) } +/// Remove a node's entity rows inside a caller-owned transaction. +pub fn clear_entity_index_for_node_tx( + tx: &Transaction<'_>, + node_id: &str, +) -> anyhow::Result { + Ok(tx.execute( + "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", + params![node_id], + )?) +} + +/// Index summary entity ids inside a caller-owned transaction. +pub fn index_summary_entity_ids_tx_with_identity( + tx: &Transaction<'_>, + entity_ids: &[String], + node_id: &str, + score: f32, + timestamp_ms: i64, + tree_id: Option<&str>, + identity: &dyn SelfIdentity, +) -> anyhow::Result { + if entity_ids.is_empty() { + return Ok(0); + } + let mut stmt = tx.prepare(UPSERT_SQL)?; + for canonical_id in entity_ids { + let entity_kind = canonical_id + .split_once(':') + .map_or(canonical_id.as_str(), |(kind, _)| kind); + stmt.execute(params![ + canonical_id, + node_id, + "summary", + entity_kind, + canonical_id, + score, + timestamp_ms, + tree_id, + canonical_id_is_user(identity, canonical_id) as i32, + ])?; + } + Ok(entity_ids.len()) +} + #[cfg(test)] #[path = "store_tests.rs"] mod tests; diff --git a/src/memory/store/entity_index/store_tests.rs b/src/memory/store/entity_index/store_tests.rs index efd3500..a5bf7ac 100644 --- a/src/memory/store/entity_index/store_tests.rs +++ b/src/memory/store/entity_index/store_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::memory::store::entity_index::types::{CanonicalEntity, EntityKind}; +use std::sync::Arc; use tempfile::TempDir; fn index() -> EntityIndex { @@ -17,6 +18,75 @@ fn sample_entity(id: &str) -> CanonicalEntity { } } +#[test] +fn shared_connection_preserves_owner_pragmas_and_visibility() { + let conn = Arc::new(Mutex::new(Connection::open_in_memory().unwrap())); + conn.lock() + .execute_batch("PRAGMA synchronous = OFF;") + .unwrap(); + + let index = + EntityIndex::from_shared_connection(Arc::clone(&conn), Arc::new(NoSelfIdentity)).unwrap(); + index + .index_entity(&sample_entity("alice"), "chunk-1", "leaf", 1000, None) + .unwrap(); + + let synchronous: i64 = conn + .lock() + .query_row("PRAGMA synchronous", [], |row| row.get(0)) + .unwrap(); + let count: i64 = conn + .lock() + .query_row("SELECT COUNT(*) FROM mem_tree_entity_index", [], |row| { + row.get(0) + }) + .unwrap(); + assert_eq!(synchronous, 0); + assert_eq!(count, 1); +} + +#[derive(Debug)] +struct EmailSelf; + +impl SelfIdentity for EmailSelf { + fn is_self(&self, kind: EntityKind, surface: &str) -> bool { + kind == EntityKind::Email && surface == "alice@example.com" + } +} + +#[test] +fn identity_aware_transaction_helpers_preserve_atomic_indexing() { + let idx = EntityIndex::open_in_memory().unwrap(); + let entity = sample_entity("alice"); + idx.with_transaction(|tx| { + index_entities_tx_with_identity( + tx, + std::slice::from_ref(&entity), + "chunk-1", + "leaf", + 1000, + None, + &EmailSelf, + )?; + index_summary_entity_ids_tx_with_identity( + tx, + &["email:alice@example.com".into()], + "summary-1", + 1.0, + 2000, + None, + &EmailSelf, + )?; + Ok(()) + }) + .unwrap(); + + let hits = idx.lookup_entity("email:alice", None).unwrap(); + assert!(hits[0].is_user); + let summary_hits = idx.lookup_entity("email:alice@example.com", None).unwrap(); + assert!(summary_hits[0].is_user); +} + #[test] fn entity_kind_wire_strings_round_trip_and_reject_unknowns() { for (kind, wire, mechanical) in [ diff --git a/src/memory/store/entity_index/types.rs b/src/memory/store/entity_index/types.rs index ca89984..dff814d 100644 --- a/src/memory/store/entity_index/types.rs +++ b/src/memory/store/entity_index/types.rs @@ -104,7 +104,7 @@ impl EntityKind { /// /// Same surface form (after normalisation) → same `canonical_id` regardless /// of how many times it appears. One row per occurrence preserves source spans. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct CanonicalEntity { /// Stable id following the `":"` convention. pub canonical_id: String, diff --git a/src/memory/store/kv.rs b/src/memory/store/kv.rs index 7218413..1c359ca 100644 --- a/src/memory/store/kv.rs +++ b/src/memory/store/kv.rs @@ -18,10 +18,7 @@ use std::sync::Arc; use crate::memory::store::safety; use crate::memory::types::MemoryKvRecord; -const INIT_SQL: &str = " - PRAGMA journal_mode = WAL; - PRAGMA synchronous = NORMAL; - +const SCHEMA_SQL: &str = " CREATE TABLE IF NOT EXISTS kv_global ( key TEXT PRIMARY KEY, value_json TEXT NOT NULL, @@ -38,6 +35,11 @@ const INIT_SQL: &str = " CREATE INDEX IF NOT EXISTS idx_kv_ns ON kv_namespace(namespace); "; +const OWNED_CONNECTION_PRAGMAS: &str = " + PRAGMA journal_mode = WAL; + PRAGMA synchronous = NORMAL; +"; + /// SQLite-backed global + namespace JSON key-value store. /// /// Thread-safe: the connection is behind a `parking_lot::Mutex`. @@ -52,7 +54,7 @@ impl KvStore { std::fs::create_dir_all(parent)?; } let conn = Connection::open(db_path)?; - conn.execute_batch(INIT_SQL)?; + Self::init_owned_connection(&conn)?; Ok(Self { conn: Arc::new(Mutex::new(conn)), }) @@ -61,12 +63,27 @@ impl KvStore { /// Open an in-memory KV store (useful for tests). pub fn open_in_memory() -> anyhow::Result { let conn = Connection::open_in_memory()?; - conn.execute_batch(INIT_SQL)?; + Self::init_owned_connection(&conn)?; Ok(Self { conn: Arc::new(Mutex::new(conn)), }) } + /// Use a connection owned and configured by the embedding application. + /// + /// Only the KV schema is installed. Journal mode, synchronous mode, busy + /// timeout, and transaction policy remain owned by the caller. + pub fn from_shared_connection(conn: Arc>) -> anyhow::Result { + conn.lock().execute_batch(SCHEMA_SQL)?; + Ok(Self { conn }) + } + + fn init_owned_connection(conn: &Connection) -> anyhow::Result<()> { + conn.execute_batch(OWNED_CONNECTION_PRAGMAS)?; + conn.execute_batch(SCHEMA_SQL)?; + Ok(()) + } + /// Seconds since the Unix epoch, as a float (matches OpenHuman's `updated_at`). fn now_ts() -> f64 { use std::time::{SystemTime, UNIX_EPOCH}; @@ -98,7 +115,7 @@ impl KvStore { if safety::has_likely_secret(key) { return Err("kv key cannot contain secrets".to_string()); } - if safety::has_likely_pii(key) { + if safety::has_likely_email(key) || safety::has_likely_pii(key) { return Err("kv key cannot contain personal identifiers".to_string()); } @@ -138,7 +155,11 @@ impl KvStore { if safety::has_likely_secret(namespace) || safety::has_likely_secret(key) { return Err("kv namespace/key cannot contain secrets".to_string()); } - if safety::has_likely_pii(namespace) || safety::has_likely_pii(key) { + if safety::has_likely_email(namespace) + || safety::has_likely_email(key) + || safety::has_likely_pii(namespace) + || safety::has_likely_pii(key) + { return Err("kv namespace/key cannot contain personal identifiers".to_string()); } diff --git a/src/memory/store/kv_tests.rs b/src/memory/store/kv_tests.rs index c998eb9..57e4ff1 100644 --- a/src/memory/store/kv_tests.rs +++ b/src/memory/store/kv_tests.rs @@ -1,10 +1,37 @@ use super::*; use serde_json::json; +use std::sync::Arc; fn store() -> KvStore { KvStore::open_in_memory().unwrap() } +#[test] +fn shared_connection_preserves_owner_pragmas_and_visibility() { + let conn = Arc::new(Mutex::new(Connection::open_in_memory().unwrap())); + conn.lock() + .execute_batch("PRAGMA synchronous = OFF;") + .unwrap(); + + let kv = KvStore::from_shared_connection(Arc::clone(&conn)).unwrap(); + kv.set_global("theme", &json!("dark")).unwrap(); + + let synchronous: i64 = conn + .lock() + .query_row("PRAGMA synchronous", [], |row| row.get(0)) + .unwrap(); + let stored: String = conn + .lock() + .query_row( + "SELECT value_json FROM kv_global WHERE key='theme'", + [], + |row| row.get(0), + ) + .unwrap(); + assert_eq!(synchronous, 0); + assert_eq!(stored, "\"dark\""); +} + #[test] fn global_kv_roundtrips_and_deletes() { let kv = store(); diff --git a/src/memory/store/mod.rs b/src/memory/store/mod.rs index 26a1fec..42ad17b 100644 --- a/src/memory/store/mod.rs +++ b/src/memory/store/mod.rs @@ -33,6 +33,7 @@ pub mod kv; /// Secret / PII detection guard applied before KV writes are persisted. pub mod safety; /// Starter in-memory reference backend used by the smoke test. +#[allow(clippy::module_inception)] pub mod store; /// Core memory contract types ([`MemoryError`], [`MemoryRecord`], etc.). pub mod types; diff --git a/src/memory/store/safety.rs b/src/memory/store/safety/mod.rs similarity index 88% rename from src/memory/store/safety.rs rename to src/memory/store/safety/mod.rs index d07f3b2..c466fc5 100644 --- a/src/memory/store/safety.rs +++ b/src/memory/store/safety/mod.rs @@ -4,20 +4,25 @@ //! prefers false positives over leaking credentials into long-lived stores. //! //! The exhaustive multilingual national-ID PII module (`safety::pii`, ~1k lines -//! of checksum logic) is **deferred**: this port keeps the full secret-pattern -//! surface that the KV write guard depends on, plus a lightweight PII screen -//! (email / phone / US SSN / Brazilian CPF) sufficient for the KV "reject -//! PII-like key" contract. Hosts that need full national-ID redaction can layer -//! it back on top. +//! of checksum logic) is ported from OpenHuman and runs as part of +//! [`sanitize_text`]. The write-rejection boundary stays stricter than content +//! scrubbing: formatted national IDs are rejected, while phone/email-like text is +//! scrubbed from content without rejecting every write that mentions them. use std::sync::LazyLock; use regex::Regex; use serde_json::Value; +/// Exhaustive checksum-gated multilingual national-ID PII module (ported from +/// OpenHuman). Content scrubbing runs from [`sanitize_text`]; the boundary +/// check is re-exported as [`has_likely_pii`]. +pub mod pii; + +pub use pii::{has_likely_email, has_likely_pii}; + const REDACTED_SECRET: &str = "[REDACTED_SECRET]"; const REDACTED_PRIVATE_KEY: &str = "[REDACTED_PRIVATE_KEY]"; -const REDACTED_PII: &str = "[REDACTED_PII]"; const MAX_JSON_SANITIZE_DEPTH: usize = 128; /// Tally of what a sanitization pass changed. @@ -175,40 +180,12 @@ static REDACTION_PATTERNS: LazyLock> = LazyLock::new( ] }); -// Lightweight PII screen (the heavy multilingual national-ID module is deferred). -static PII_PATTERNS: LazyLock> = LazyLock::new(|| { - vec![ - ( - Regex::new(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b") - .expect("valid email pii"), - REDACTED_PII, - ), - ( - Regex::new(r"\b\d{3}-\d{2}-\d{4}\b").expect("valid ssn pii"), - REDACTED_PII, - ), - ( - Regex::new(r"\b\d{3}\.\d{3}\.\d{3}-\d{2}\b").expect("valid cpf pii"), - REDACTED_PII, - ), - ( - Regex::new(r"\+\d{7,15}\b").expect("valid e164 pii"), - REDACTED_PII, - ), - ] -}); - /// True when `value` looks like it contains a credential. pub fn has_likely_secret(value: &str) -> bool { BLOCK_PATTERNS.iter().any(|p| p.is_match(value)) || REDACTION_PATTERNS.iter().any(|(p, _)| p.is_match(value)) } -/// True when `value` looks like it contains a personal identifier. -pub fn has_likely_pii(value: &str) -> bool { - PII_PATTERNS.iter().any(|(p, _)| p.is_match(value)) -} - /// Scrub secrets and PII from free text, returning the cleaned text plus a /// [`SanitizationReport`]. pub fn sanitize_text(value: &str) -> Sanitized { @@ -231,13 +208,12 @@ pub fn sanitize_text(value: &str) -> Sanitized { } } - for (pattern, replacement) in PII_PATTERNS.iter() { - let hits = pattern.find_iter(&out).count(); - if hits > 0 { - report.pii_redactions += hits; - out = pattern.replace_all(&out, *replacement).into_owned(); - } - } + // Full multilingual national-ID PII scrub (checksum-gated, normalization + // pre-pass) — runs after secret redaction so every call site that scrubs + // secrets also scrubs PII. + let pii = pii::redact_pii(&out); + report = report.merge(pii.report); + out = pii.value; Sanitized { value: out, report } } diff --git a/src/memory/store/safety/pii.rs b/src/memory/store/safety/pii.rs new file mode 100644 index 0000000..5f9b756 --- /dev/null +++ b/src/memory/store/safety/pii.rs @@ -0,0 +1,1113 @@ +//! Multilingual personal-PII redaction (national IDs, financial identifiers, +//! international phone) — on-device, regex + checksum only, zero network. +//! +//! ## Design — security first +//! +//! 1. **Checksum gating where possible.** CPF, CNPJ, CUIT, credit-card (Luhn), +//! IBAN (mod-97), Aadhaar (Verhoeff), Spanish DNI/NIE (check letter), and +//! US SSN reserved-range filters all reject look-alikes that aren't real +//! identifiers. The false-positive rate from format alone is too high; the +//! checksums bring it back to acceptable. +//! +//! 2. **Bypass-resistant.** Inputs are run through [`normalize`] before +//! matching, which: +//! - strips zero-width characters (U+200B/200C/200D/FEFF/2060/180E), +//! - folds fullwidth digits (`０-９` → `0-9`) and fullwidth `．－／：` +//! to their ASCII counterparts, +//! - folds Arabic-Indic and Eastern Arabic-Indic digits to ASCII. +//! Match offsets are mapped back to the original text so we only redact +//! the bytes that actually carry PII; surrounding text is untouched. +//! +//! 3. **Overlap-safe.** Patterns are run in priority order; later matches +//! that overlap an earlier redaction are dropped, so a credit-card span +//! can't also be partially matched as a phone number. +//! +//! 4. **Out of scope.** Contextual PII (`"call me at the usual number"`), +//! compound PII (`name + employer + city`), arbitrary names, and freeform +//! dates-of-birth all require NER/LLM and are NOT addressed here. This +//! module is honest about its scope. + +use regex::{Regex, RegexSet}; +use std::sync::LazyLock; + +use super::{SanitizationReport, Sanitized}; + +// ---------- Replacement tokens ---------- + +const PII_RFC: &str = "[REDACTED_PII_RFC]"; +const PII_CPF: &str = "[REDACTED_PII_CPF]"; +const PII_CNPJ: &str = "[REDACTED_PII_CNPJ]"; +const PII_CUIT: &str = "[REDACTED_PII_CUIT]"; +const PII_MYNUM: &str = "[REDACTED_PII_MYNUMBER]"; +const PII_PHONE: &str = "[REDACTED_PII_PHONE]"; +const PII_SSN: &str = "[REDACTED_PII_SSN]"; +const PII_CC: &str = "[REDACTED_PII_CREDIT_CARD]"; +const PII_IBAN: &str = "[REDACTED_PII_IBAN]"; +const PII_AADHAAR: &str = "[REDACTED_PII_AADHAAR]"; +const PII_PAN_IN: &str = "[REDACTED_PII_PAN_IN]"; +const PII_NINO: &str = "[REDACTED_PII_NINO]"; +const PII_DNI: &str = "[REDACTED_PII_DNI]"; +const PII_RRN: &str = "[REDACTED_PII_RRN]"; + +// ---------- Patterns ---------- + +// Brazilian CPF, formatted: NNN.NNN.NNN-NN +static CPF_FMT_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{3}\.\d{3}\.\d{3}-\d{2}\b").expect("cpf fmt")); +// Brazilian CPF, bare: 11 consecutive digits. Checksum-gated; ~1% raw FP. +static CPF_BARE_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{11}\b").expect("cpf bare")); + +// Brazilian CNPJ, formatted: NN.NNN.NNN/NNNN-NN +static CNPJ_FMT_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{2}\.\d{3}\.\d{3}/\d{4}-\d{2}\b").expect("cnpj fmt")); +// Brazilian CNPJ, bare: 14 consecutive digits. +static CNPJ_BARE_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{14}\b").expect("cnpj bare")); + +// Argentine CUIT/CUIL: NN-NNNNNNNN-N (formatted only — bare 11-digit with +// single check digit has ~9% FP on random IDs, too noisy without context). +static CUIT_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{2}-\d{8}-\d\b").expect("cuit")); + +// Mexican RFC: 3-4 letters (incl. Ñ &) + 6 digits + 3 alphanumeric homoclave. +static RFC_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?i)\b[A-ZÑ&]{3,4}\d{6}[A-Z0-9]{3}\b").expect("rfc")); + +// Japan My Number (12 digits) gated by a Japanese or English keyword within +// ~30 chars. Bare 12-digit runs without keyword are too noisy. +static MYNUM_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?:マイナンバー|個人番号|My\s?Number)[\s:はがを、.\-]{0,12}(\d{12})\b") + .expect("my number") +}); + +// E.164 phone: + followed by 7-15 digits, no separators. +static PHONE_E164_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\+\d{7,15}\b").expect("e164")); + +// NANP (US/Canada) formatted phone. Area code must start 2-9; first digit of +// central-office code also 2-9 (real NANP rule). +static PHONE_NANP_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"\b(?:\+?1[\s.\-]?)?\(?([2-9]\d{2})\)?[\s.\-]?([2-9]\d{2})[\s.\-]?(\d{4})\b") + .expect("nanp phone") +}); + +// US SSN: NNN-NN-NNNN. Range filter applied below. +static SSN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{3}-\d{2}-\d{4}\b").expect("ssn")); + +// Credit card: 13-19 digits with optional spaces/dashes every 4. Luhn-gated. +static CC_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b(?:\d[\s\-]?){13,19}\b").expect("credit card")); + +// IBAN: 2 letter country code + 2 check digits + 11-30 alphanumeric. +// Allow optional spaces every 4 chars (common human format). +static IBAN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b[A-Z]{2}\d{2}(?:[\s]?[A-Z0-9]){11,30}\b").expect("iban")); + +// India Aadhaar: 4-4-4 digit groups (space or hyphen) OR contiguous 12 digits +// gated by keyword. Verhoeff-checksum-gated when grouped, keyword-gated when +// bare (Verhoeff alone has ~10% raw FP rate on random 12-digit runs). +static AADHAAR_FMT_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{4}[\s\-]\d{4}[\s\-]\d{4}\b").expect("aadhaar formatted")); +static AADHAAR_KW_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?i)(?:aadhaar|aadhar|आधार|uidai|uid)[\s:#\-no.]{0,10}(\d{12})\b") + .expect("aadhaar keyword") +}); + +// India PAN: 5 letters, 4 digits, 1 letter. Very high signal — no checksum. +static PAN_IN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?i)\b[A-Z]{5}\d{4}[A-Z]\b").expect("pan-in")); + +// UK NINO: 2 letters + 6 digits + suffix A/B/C/D. +static NINO_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?i)\b[A-Z]{2}\d{6}[A-D]\b").expect("nino")); + +// Spain DNI: 8 digits + check letter. NIE: starts X/Y/Z, then 7 digits + letter. +static DNI_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?i)\b\d{8}[A-Z]\b").expect("dni")); +static NIE_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?i)\b[XYZ]\d{7}[A-Z]\b").expect("nie")); + +// South Korea RRN: NNNNNN-CXXXXXX where C is gender/century digit (1-4). +static RRN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\b\d{6}-[1-4]\d{6}\b").expect("rrn")); +static EMAIL_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?i)\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b").expect("email")); + +// Cheap whole-text pre-filter so we skip the per-pattern scans entirely on +// PII-free text. Each entry roughly corresponds to one of the patterns above. +static SCREEN: LazyLock = LazyLock::new(|| { + RegexSet::new([ + r"\d{11,}", // any long digit run → CPF/CNPJ/CC/Aadhaar/IBAN + r"\d{3}\.\d{3}\.\d{3}-\d{2}", // CPF + r"\d{2}\.\d{3}\.\d{3}/\d{4}-\d{2}", // CNPJ + r"\d{2}-\d{8}-\d", // CUIT + r"(?i)[A-Z]{3,4}\d{6}", // RFC / general alphanumeric ID + r"(?:マイナンバー|個人番号|My\s?Number)", // JP keyword + r"\+\d{7}", // E.164 + r"\(?[2-9]\d{2}\)?[\s.\-]\d{3}[\s.\-]\d{4}", // NANP (parens optional) + r"\d{3}-\d{2}-\d{4}", // SSN + r"\b[A-Z]{2}\d{2}[A-Z0-9]", // IBAN prefix + r"\d{4}[\s\-]\d{4}[\s\-]\d{4}", // Aadhaar formatted + r"(?i)aadhaar|aadhar|आधार|uidai", // Aadhaar keyword + r"(?i)[A-Z]{5}\d{4}[A-Z]", // PAN-IN + r"(?i)[A-Z]{2}\d{6}[A-D]", // NINO + r"\b\d{8}[A-Z]\b", // DNI + r"(?i)[XYZ]\d{7}[A-Z]", // NIE + r"\d{6}-[1-4]\d{6}", // RRN + ]) + .expect("screen regex set") +}); + +// ---------- Public API ---------- + +/// Redact format-based multilingual PII from `text`. +/// +/// Runs a Unicode normalization pre-pass to defeat fullwidth-digit and +/// zero-width-char bypasses. Match indices from the normalized form are +/// translated back to original byte offsets so only the PII bytes are +/// replaced — surrounding text (including any preserved fullwidth glyphs) +/// is untouched. +pub fn redact_pii(text: &str) -> Sanitized { + let mut report = SanitizationReport::default(); + + // Fast path: no candidate at all. + if !SCREEN.is_match(text) { + // Even the screen might miss normalized inputs; check normalized too. + let nview = NormalizedView::build(text); + if !SCREEN.is_match(&nview.normalized) { + return Sanitized { + value: text.to_string(), + report, + }; + } + return splice_redactions( + text, + &nview, + collect_redactions(&nview.normalized), + &mut report, + ); + } + + let nview = NormalizedView::build(text); + let redactions = collect_redactions(&nview.normalized); + splice_redactions(text, &nview, redactions, &mut report) +} + +/// True if `value` looks like it carries any PII. Used to *reject* +/// namespace/key inputs at boundary checks (analogous to +/// [`super::has_likely_secret`]). +/// +/// Uses the **strict** match set — only formatted / keyword-gated patterns. +/// Bare-numeric patterns whose only signal is a digit run (credit card via +/// Luhn, bare CPF, bare CNPJ) or a phone-shaped digit run (NANP without +/// separators, E.164 leading `+`) are excluded here because their false- +/// positive rate against scanner-built namespace/key identifiers (WhatsApp +/// JIDs like `12025551234-1543890267@g.us`, telegram numeric peer IDs, +/// millisecond timestamps, padded counters) is too high to use as a hard +/// rejection signal. Content scrubbing via [`redact_pii`] still applies +/// those patterns — false positives are tolerable there because they only +/// replace bytes inside a string, not reject the whole write. +pub fn has_likely_pii(value: &str) -> bool { + let nview = NormalizedView::build(value); + SCREEN.is_match(&nview.normalized) && !collect_strict_redactions(&nview.normalized).is_empty() +} + +/// True when `value` contains an ordinary email address. Kept separate from +/// [`has_likely_pii`] because scanner-built identifiers may legitimately +/// contain email-like `@` segments. +pub fn has_likely_email(value: &str) -> bool { + EMAIL_RE.is_match(value) +} + +// ---------- Match collection ---------- + +#[derive(Debug)] +struct Hit { + start: usize, // byte offset in NORMALIZED text + end: usize, + token: &'static str, +} + +fn collect_redactions(norm: &str) -> Vec { + collect_redactions_inner(norm, true) +} + +/// Variant of [`collect_redactions`] that omits bare-numeric patterns +/// whose only signal is a digit-run shape: credit card via Luhn, bare +/// CPF, bare CNPJ, NANP phones (separators optional, so any 10-11 digit +/// run starting `[2-9]`/`1[2-9]` matches), and E.164 phones (literal `+` +/// the only signal). Used for boundary checks like [`has_likely_pii`] +/// where rejection on such a hit alone would have too many false +/// positives on scanner-built identifiers (WhatsApp group JIDs +/// `-@g.us`, timestamps, padded counters). +fn collect_strict_redactions(norm: &str) -> Vec { + collect_redactions_inner(norm, false) +} + +fn collect_redactions_inner(norm: &str, include_bare_numeric: bool) -> Vec { + let mut hits: Vec = Vec::new(); + + // Priority order: most specific / highest-confidence first. + push_checksum(&mut hits, norm, &CPF_FMT_RE, PII_CPF, |s| { + valid_cpf(digits(s).as_slice()) + }); + push_checksum(&mut hits, norm, &CNPJ_FMT_RE, PII_CNPJ, |s| { + valid_cnpj(digits(s).as_slice()) + }); + push_checksum(&mut hits, norm, &CUIT_RE, PII_CUIT, |s| { + valid_cuit(digits(s).as_slice()) + }); + + // IBAN before credit card: CC can match an IBAN tail of all digits. + push_checksum(&mut hits, norm, &IBAN_RE, PII_IBAN, valid_iban); + + if include_bare_numeric { + // Credit card before bare CPF/CNPJ to avoid catching a 13-19 digit run as CPF/CNPJ. + push_checksum(&mut hits, norm, &CC_RE, PII_CC, valid_luhn); + + push_checksum(&mut hits, norm, &CNPJ_BARE_RE, PII_CNPJ, |s| { + valid_cnpj(digits(s).as_slice()) + }); + push_checksum(&mut hits, norm, &CPF_BARE_RE, PII_CPF, |s| { + valid_cpf(digits(s).as_slice()) + }); + } + + push_checksum(&mut hits, norm, &AADHAAR_FMT_RE, PII_AADHAAR, |s| { + valid_verhoeff(digits(s).as_slice()) + }); + // Keyword-gated Aadhaar redacts only the captured 12-digit group. + push_captured(&mut hits, norm, &AADHAAR_KW_RE, PII_AADHAAR, |digits_str| { + valid_verhoeff(digits(digits_str).as_slice()) + }); + + push_checksum(&mut hits, norm, &DNI_RE, PII_DNI, valid_dni_es); + push_checksum(&mut hits, norm, &NIE_RE, PII_DNI, valid_nie_es); + push_checksum(&mut hits, norm, &NINO_RE, PII_NINO, valid_nino); + push_checksum(&mut hits, norm, &SSN_RE, PII_SSN, valid_ssn); + push_simple(&mut hits, norm, &RRN_RE, PII_RRN); + push_simple(&mut hits, norm, &RFC_RE, PII_RFC); + push_simple(&mut hits, norm, &PAN_IN_RE, PII_PAN_IN); + + if include_bare_numeric { + // Phones: E.164 first (more specific), then NANP. Both are bare-numeric + // shapes — NANP allows optional separators (`\b\d{10,11}\b` matches as + // `XXX-XXX-XXXX`), and E.164 keys on a literal `+` with no further gate. + // Strict callers (boundary checks like `has_likely_pii`) exclude these + // so scanner-built namespace/key values (WhatsApp JIDs + // `-@g.us`, telegram numeric peer IDs) don't get rejected. + push_simple(&mut hits, norm, &PHONE_E164_RE, PII_PHONE); + push_simple(&mut hits, norm, &PHONE_NANP_RE, PII_PHONE); + } + + // My Number — captured digit group only, keyword remains visible. + push_captured(&mut hits, norm, &MYNUM_RE, PII_MYNUM, |_| true); + + dedupe_overlaps(&mut hits); + hits +} + +fn push_simple(hits: &mut Vec, norm: &str, re: &Regex, token: &'static str) { + for m in re.find_iter(norm) { + hits.push(Hit { + start: m.start(), + end: m.end(), + token, + }); + } +} + +fn push_checksum( + hits: &mut Vec, + norm: &str, + re: &Regex, + token: &'static str, + ok: impl Fn(&str) -> bool, +) { + for m in re.find_iter(norm) { + if ok(m.as_str()) { + hits.push(Hit { + start: m.start(), + end: m.end(), + token, + }); + } + } +} + +fn push_captured( + hits: &mut Vec, + norm: &str, + re: &Regex, + token: &'static str, + ok: impl Fn(&str) -> bool, +) { + for caps in re.captures_iter(norm) { + let Some(group) = caps.get(1) else { continue }; + if ok(group.as_str()) { + hits.push(Hit { + start: group.start(), + end: group.end(), + token, + }); + } + } +} + +// Sort by start asc, length desc. Then walk in order, dropping any hit whose +// range overlaps a kept hit. Result: earlier + longer wins; no double-redact. +fn dedupe_overlaps(hits: &mut Vec) { + hits.sort_by(|a, b| { + a.start + .cmp(&b.start) + .then((b.end - b.start).cmp(&(a.end - a.start))) + }); + let mut kept: Vec = Vec::with_capacity(hits.len()); + for h in hits.drain(..) { + let overlaps = kept.last().is_some_and(|k| h.start < k.end); + if !overlaps { + kept.push(h); + } + } + *hits = kept; +} + +// Splice redactions (whose indices reference NORMALIZED text) back into the +// ORIGINAL text via NormalizedView's byte-offset mapping. This preserves +// non-PII original bytes verbatim (including fullwidth glyphs the user +// intentionally typed) while still scrubbing detected PII. +fn splice_redactions( + original: &str, + nview: &NormalizedView, + hits: Vec, + report: &mut SanitizationReport, +) -> Sanitized { + if hits.is_empty() { + return Sanitized { + value: original.to_string(), + report: *report, + }; + } + let mut out = String::with_capacity(original.len()); + let mut cursor = 0; + for h in &hits { + let start_orig = nview.norm_to_orig(h.start); + let end_orig = nview.norm_to_orig(h.end); + if start_orig < cursor || start_orig > original.len() || end_orig > original.len() { + continue; + } + out.push_str(&original[cursor..start_orig]); + out.push_str(h.token); + cursor = end_orig; + } + out.push_str(&original[cursor..]); + report.pii_redactions += hits.len(); + Sanitized { + value: out, + report: *report, + } +} + +// ---------- Unicode normalization for matching ---------- + +struct NormalizedView { + normalized: String, + // For each byte offset i in `normalized`, `byte_map[i]` is the byte offset + // in the original string where the corresponding char *starts*. + // The last entry maps the normalized length to the original length, so + // `norm_to_orig(normalized.len())` is well-defined. + byte_map: Vec, +} + +impl NormalizedView { + fn build(original: &str) -> Self { + let mut normalized = String::with_capacity(original.len()); + let mut byte_map: Vec = Vec::with_capacity(original.len() + 1); + for (idx, ch) in original.char_indices() { + if is_zero_width(ch) { + continue; + } + let mapped = fold_char(ch); + let start = normalized.len(); + normalized.push(mapped); + // One byte_map entry per byte of the normalized char. + let added = normalized.len() - start; + for _ in 0..added { + byte_map.push(idx); + } + } + byte_map.push(original.len()); + Self { + normalized, + byte_map, + } + } + + fn norm_to_orig(&self, norm_byte: usize) -> usize { + if norm_byte >= self.byte_map.len() { + return *self.byte_map.last().unwrap_or(&0); + } + self.byte_map[norm_byte] + } +} + +fn is_zero_width(c: char) -> bool { + matches!( + c, + '\u{200B}' + | '\u{200C}' + | '\u{200D}' + | '\u{200E}' + | '\u{200F}' + | '\u{2060}' + | '\u{180E}' + | '\u{FEFF}' + ) +} + +fn fold_char(c: char) -> char { + match c { + // Fullwidth digits 0-9 + '\u{FF10}'..='\u{FF19}' => char::from_u32(c as u32 - 0xFF10 + 0x30).unwrap_or(c), + // Arabic-Indic digits ٠-٩ + '\u{0660}'..='\u{0669}' => char::from_u32(c as u32 - 0x0660 + 0x30).unwrap_or(c), + // Eastern Arabic-Indic digits ۰-۹ + '\u{06F0}'..='\u{06F9}' => char::from_u32(c as u32 - 0x06F0 + 0x30).unwrap_or(c), + // Common fullwidth punctuation we care about for PII formats + '\u{FF0D}' => '-', + '\u{FF0E}' => '.', + '\u{FF0F}' => '/', + '\u{FF1A}' => ':', + '\u{2010}'..='\u{2015}' => '-', // various unicode hyphens/dashes + '\u{2212}' => '-', // minus sign + other => other, + } +} + +// ---------- Checksum helpers ---------- + +fn digits(s: &str) -> Vec { + s.chars() + .filter(|c| c.is_ascii_digit()) + .map(|c| c.to_digit(10).expect("ascii digit")) + .collect() +} + +fn valid_cpf(d: &[u32]) -> bool { + if d.len() != 11 || d.iter().all(|x| *x == d[0]) { + return false; + } + let s1: u32 = (0..9).map(|i| d[i] * (10 - i as u32)).sum(); + let dv1 = (s1 * 10) % 11 % 10; + if dv1 != d[9] { + return false; + } + let s2: u32 = (0..10).map(|i| d[i] * (11 - i as u32)).sum(); + let dv2 = (s2 * 10) % 11 % 10; + dv2 == d[10] +} + +fn valid_cnpj(d: &[u32]) -> bool { + if d.len() != 14 || d.iter().all(|x| *x == d[0]) { + return false; + } + let w1: [u32; 12] = [5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; + let s1: u32 = (0..12).map(|i| d[i] * w1[i]).sum(); + let r1 = s1 % 11; + let dv1 = if r1 < 2 { 0 } else { 11 - r1 }; + if dv1 != d[12] { + return false; + } + let w2: [u32; 13] = [6, 5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; + let s2: u32 = (0..13).map(|i| d[i] * w2[i]).sum(); + let r2 = s2 % 11; + let dv2 = if r2 < 2 { 0 } else { 11 - r2 }; + dv2 == d[13] +} + +fn valid_cuit(d: &[u32]) -> bool { + if d.len() != 11 { + return false; + } + let w: [u32; 10] = [5, 4, 3, 2, 7, 6, 5, 4, 3, 2]; + let s: u32 = (0..10).map(|i| d[i] * w[i]).sum(); + let r = s % 11; + let dv = match r { + 0 => 0, + 1 => return false, + _ => 11 - r, + }; + dv == d[10] +} + +// Luhn — used for credit-card validation. +fn valid_luhn(s: &str) -> bool { + let d = digits(s); + if d.len() < 13 || d.len() > 19 { + return false; + } + let mut sum = 0u32; + let mut alt = false; + for x in d.iter().rev() { + let v = if alt { + let doubled = x * 2; + if doubled > 9 { + doubled - 9 + } else { + doubled + } + } else { + *x + }; + sum += v; + alt = !alt; + } + sum.is_multiple_of(10) +} + +// IBAN mod-97. Steps: strip spaces, move first 4 chars to end, expand letters +// (A=10..Z=35), divide as a big-integer mod 97, require remainder == 1. +fn valid_iban(s: &str) -> bool { + let cleaned: String = s.chars().filter(|c| !c.is_whitespace()).collect(); + if cleaned.len() < 15 || cleaned.len() > 34 { + return false; + } + if !cleaned.chars().take(2).all(|c| c.is_ascii_alphabetic()) { + return false; + } + if !cleaned[2..4].chars().all(|c| c.is_ascii_digit()) { + return false; + } + let rotated: String = cleaned[4..].chars().chain(cleaned[..4].chars()).collect(); + let mut remainder: u64 = 0; + for c in rotated.chars() { + let chunk = if let Some(d) = c.to_digit(10) { + d as u64 + } else if c.is_ascii_alphabetic() { + (c.to_ascii_uppercase() as u64) - ('A' as u64) + 10 + } else { + return false; + }; + // Expand into the running remainder digit-by-digit so we never need + // u128. Each letter contributes 2 decimal digits. + if chunk >= 10 { + remainder = (remainder * 100 + chunk) % 97; + } else { + remainder = (remainder * 10 + chunk) % 97; + } + } + remainder == 1 +} + +// Verhoeff — used for Aadhaar. +const VERHOEFF_D: [[u8; 10]; 10] = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 0, 6, 7, 8, 9, 5], + [2, 3, 4, 0, 1, 7, 8, 9, 5, 6], + [3, 4, 0, 1, 2, 8, 9, 5, 6, 7], + [4, 0, 1, 2, 3, 9, 5, 6, 7, 8], + [5, 9, 8, 7, 6, 0, 4, 3, 2, 1], + [6, 5, 9, 8, 7, 1, 0, 4, 3, 2], + [7, 6, 5, 9, 8, 2, 1, 0, 4, 3], + [8, 7, 6, 5, 9, 3, 2, 1, 0, 4], + [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], +]; +const VERHOEFF_P: [[u8; 10]; 8] = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 5, 7, 6, 2, 8, 3, 0, 9, 4], + [5, 8, 0, 3, 7, 9, 6, 1, 4, 2], + [8, 9, 1, 6, 0, 4, 3, 5, 2, 7], + [9, 4, 5, 3, 1, 2, 6, 8, 7, 0], + [4, 2, 8, 6, 5, 7, 3, 9, 0, 1], + [2, 7, 9, 3, 8, 0, 6, 4, 1, 5], + [7, 0, 4, 6, 9, 1, 3, 2, 5, 8], +]; + +fn valid_verhoeff(d: &[u32]) -> bool { + if d.len() != 12 { + return false; + } + // Aadhaar can't start with 0 or 1. + if d[0] < 2 { + return false; + } + let mut c: u8 = 0; + for (i, digit) in d.iter().rev().enumerate() { + c = VERHOEFF_D[c as usize][VERHOEFF_P[i % 8][*digit as usize] as usize]; + } + c == 0 +} + +// US SSN reserved/invalid ranges per SSA. +fn valid_ssn(s: &str) -> bool { + let d = digits(s); + if d.len() != 9 { + return false; + } + let area = d[0] * 100 + d[1] * 10 + d[2]; + let group = d[3] * 10 + d[4]; + let serial = d[5] * 1000 + d[6] * 100 + d[7] * 10 + d[8]; + if area == 0 || area == 666 || area >= 900 { + return false; + } + if group == 0 || serial == 0 { + return false; + } + true +} + +// Spain DNI check letter — 8 digits mod 23 indexes into a fixed letter table. +const DNI_LETTERS: &[u8; 23] = b"TRWAGMYFPDXBNJZSQVHLCKE"; + +fn valid_dni_es(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + let num_str = &upper[..8]; + let letter = bytes[8]; + let Ok(num) = num_str.parse::() else { + return false; + }; + DNI_LETTERS[(num % 23) as usize] == letter +} + +fn valid_nie_es(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + let prefix = match bytes[0] { + b'X' => 0u32, + b'Y' => 1, + b'Z' => 2, + _ => return false, + }; + let Ok(rest) = std::str::from_utf8(&bytes[1..8]) else { + return false; + }; + let Ok(num) = rest.parse::() else { + return false; + }; + let composed = prefix * 10_000_000 + num; + DNI_LETTERS[(composed % 23) as usize] == bytes[8] +} + +// UK NINO reserved-prefix blacklist. +fn valid_nino(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + // First char cannot be D F I Q U V; second cannot be D F I O Q U V. + let bad_first = b"DFIQUV"; + let bad_second = b"DFIOQUV"; + if bad_first.contains(&bytes[0]) || bad_second.contains(&bytes[1]) { + return false; + } + // Reserved two-letter prefixes. + let reserved = ["BG", "GB", "KN", "NK", "NT", "TN", "ZZ"]; + let prefix = &upper[..2]; + if reserved.contains(&prefix) { + return false; + } + true +} + +// ---------- Tests ---------- + +#[cfg(test)] +mod tests { + use super::*; + + fn redacts(input: &str, token: &str) { + let out = redact_pii(input); + assert!( + out.value.contains(token), + "expected {token} in output. input={input:?} output={out:?}" + ); + } + + fn unchanged(input: &str) { + let out = redact_pii(input); + assert_eq!( + out.value, input, + "expected no change; report={:?}", + out.report + ); + assert_eq!(out.report.pii_redactions, 0); + } + + // --- CPF --- + #[test] + fn cpf_formatted_valid_redacted() { + redacts("CPF: 111.444.777-35.", PII_CPF); + } + #[test] + fn cpf_formatted_invalid_kept() { + unchanged("CPF 111.444.777-99 nope"); + } + #[test] + fn cpf_all_same_digits_rejected() { + unchanged("Test 111.111.111-11"); + } + #[test] + fn cpf_bare_valid_redacted() { + redacts("Sem mascara 11144477735 ok", PII_CPF); + } + + // --- CNPJ --- + #[test] + fn cnpj_formatted_valid_redacted() { + redacts("CNPJ 11.222.333/0001-81", PII_CNPJ); + } + #[test] + fn cnpj_bare_valid_redacted() { + redacts("contract 11222333000181 yes", PII_CNPJ); + } + + // --- CUIT --- + #[test] + fn cuit_valid_redacted() { + redacts("CUIT 20-11111111-2", PII_CUIT); + } + #[test] + fn cuit_invalid_kept() { + unchanged("noise 20-12345678-0 noise"); + } + + // --- RFC --- + #[test] + fn rfc_redacted() { + redacts("Mi RFC VECJ880326XK4 .", PII_RFC); + } + #[test] + fn rfc_lowercase_redacted() { + redacts("rfc vecj880326xk4", PII_RFC); + } + + // --- My Number --- + #[test] + fn my_number_redacted_with_keyword() { + redacts("マイナンバー: 123456789012", PII_MYNUM); + } + #[test] + fn bare_12_digits_without_keyword_kept() { + unchanged("Order 123456789012 shipped today."); + } + + // --- E.164 + NANP phone --- + #[test] + fn e164_redacted() { + redacts("phone +15551234567", PII_PHONE); + } + #[test] + fn nanp_formatted_redacted() { + redacts("call 415-555-0123 thanks", PII_PHONE); + } + #[test] + fn nanp_with_country_code_redacted() { + redacts("+1 (212) 555-7890", PII_PHONE); + } + #[test] + fn nanp_invalid_area_code_kept() { + unchanged("score 115-555-0123 ish"); + } + + // --- SSN --- + #[test] + fn ssn_valid_redacted() { + redacts("ssn 123-45-6789", PII_SSN); + } + #[test] + fn ssn_reserved_area_kept() { + unchanged("test 666-12-3456"); + } + #[test] + fn ssn_zero_serial_kept() { + unchanged("test 123-45-0000"); + } + + // --- Credit card / Luhn --- + #[test] + fn credit_card_visa_redacted() { + // Visa test number with valid Luhn. + redacts("card 4111 1111 1111 1111 thanks", PII_CC); + } + #[test] + fn credit_card_amex_redacted() { + redacts("card 378282246310005 used", PII_CC); + } + #[test] + fn credit_card_invalid_luhn_kept() { + unchanged("invoice 4111 1111 1111 1112"); + } + + // --- IBAN --- + #[test] + fn iban_de_redacted() { + // Known test IBAN with valid mod-97. + redacts("IBAN DE89370400440532013000 ok", PII_IBAN); + } + #[test] + fn iban_invalid_kept() { + unchanged("noise DE89370400440532013001 noise"); + } + + // --- Aadhaar --- + #[test] + fn aadhaar_formatted_verhoeff_valid_redacted() { + // 234123412346 is a known Verhoeff-valid Aadhaar test number. + redacts("Aadhaar 2341 2341 2346", PII_AADHAAR); + } + #[test] + fn aadhaar_keyword_bare_redacted() { + redacts("Aadhaar: 234123412346", PII_AADHAAR); + } + #[test] + fn aadhaar_invalid_verhoeff_kept() { + unchanged("Random 2341 2341 2345 nope"); + } + + // --- PAN-IN --- + #[test] + fn pan_in_redacted() { + redacts("PAN: ABCDE1234F", PII_PAN_IN); + } + + // --- NINO --- + #[test] + fn nino_redacted() { + redacts("NI no AB123456C", PII_NINO); + } + #[test] + fn nino_reserved_prefix_kept() { + unchanged("BG123456A"); + } + + // --- DNI / NIE --- + #[test] + fn dni_es_redacted() { + redacts("DNI 12345678Z", PII_DNI); + } + #[test] + fn dni_es_bad_letter_kept() { + unchanged("ID 12345678A code"); + } + #[test] + fn nie_es_redacted() { + redacts("NIE X1234567L", PII_DNI); + } + + // --- RRN Korea --- + #[test] + fn rrn_kr_redacted() { + redacts("주민번호 900101-1234567", PII_RRN); + } + #[test] + fn rrn_kr_bad_gender_digit_kept() { + unchanged("ref 900101-5234567 nope"); + } + + // --- Bypass resistance --- + #[test] + fn fullwidth_digits_cannot_bypass_cpf() { + // 111.444.777-35 with fullwidth digits and punctuation. + let input = "CPF: １１１．４４４．７７７－３５ done"; + let out = redact_pii(input); + assert!(out.value.contains(PII_CPF), "got {out:?}"); + } + + #[test] + fn zero_width_chars_cannot_bypass_ssn() { + // U+200B inserted between digits. + let input = "ssn 1\u{200B}23-4\u{200B}5-6789 done"; + let out = redact_pii(input); + assert!(out.value.contains(PII_SSN), "got {out:?}"); + } + + #[test] + fn arabic_indic_digits_normalize_for_phone() { + let input = "phone +١٥٥٥١٢٣٤٥٦٧"; + let out = redact_pii(input); + assert!(out.value.contains(PII_PHONE), "got {out:?}"); + } + + // --- Aggressive mix end-to-end --- + #[test] + fn aggressive_mixed_document() { + let input = "\ +Cliente RFC VECJ880326XK4. \ +Empresa CNPJ 11.222.333/0001-81. \ +Argentino CUIT 20-11111111-2. \ +Brasileiro CPF 111.444.777-35. \ +マイナンバー: 123456789012. \ +SSN 123-45-6789. \ +Card 4111 1111 1111 1111. \ +IBAN DE89370400440532013000. \ +PAN ABCDE1234F. \ +NI AB123456C. \ +DNI 12345678Z. \ +RRN 900101-1234567. \ +Phone +15551234567."; + let out = redact_pii(input); + for token in [ + PII_RFC, PII_CNPJ, PII_CUIT, PII_CPF, PII_MYNUM, PII_SSN, PII_CC, PII_IBAN, PII_PAN_IN, + PII_NINO, PII_DNI, PII_RRN, PII_PHONE, + ] { + assert!( + out.value.contains(token), + "missing {token} in: {}", + out.value + ); + } + assert!(out.report.pii_redactions >= 13); + } + + // --- has_likely_pii --- + #[test] + fn has_likely_pii_detects_cpf() { + assert!(has_likely_pii("user/111.444.777-35")); + } + + #[test] + fn has_likely_email_detects_email_without_changing_boundary_pii() { + assert!(has_likely_email("user/alice@example.com")); + assert!(!has_likely_pii("user/alice@example.com")); + } + #[test] + fn has_likely_pii_quiet_on_normal_text() { + assert!(!has_likely_pii("memory/global/preferences")); + } + + /// Regression: zero-padded millisecond-timestamp keys must NOT be + /// flagged as PII even when the digit run happens to satisfy Luhn. + /// `redact_pii` content scrubbing may still flag the same string — + /// `has_likely_pii` (used for boundary rejection of internal keys) + /// must stay strict to formatted/keyword PII only. + #[test] + fn has_likely_pii_ignores_bare_luhn_timestamp_keys() { + // 18-digit padded timestamps where the digit total mod 10 == 0 + // (the Luhn-passing case that previously rejected autocomplete + // KV writes and screen-intelligence document writes). + for key in [ + "accepted:000001747729035001", + "completion:000001747729035011", + "screen_intelligence_vision-1747729035001-VSCode", + ] { + assert!( + !has_likely_pii(key), + "internal key {key:?} must not be rejected as PII" + ); + } + } + + /// Strict boundary check should still reject formatted PII even though + /// it skips bare-numeric checksum patterns. + #[test] + fn has_likely_pii_still_blocks_formatted_secrets() { + assert!(has_likely_pii("ssn-123-45-6789")); + assert!(has_likely_pii("cliente-RFC-VECJ880326XK4")); + assert!(has_likely_pii("cuit-20-11111111-2")); + } + + /// Regression for Sentry TAURI-RUST-54T / GH #2848: scanner-built + /// `namespace` and `key` values containing bare-numeric phone-shaped + /// digit runs (WhatsApp group JID `-@g.us`, WhatsApp + /// broadcast `@broadcast`, US-prefixed WhatsApp 1:1 JID, + /// telegram numeric peer ID) must NOT be rejected by the boundary + /// PII check. NANP matches `\d{10,11}` with optional separators — + /// strict mode must skip it. Content scrubbing via `redact_pii` + /// continues to redact these substrings (see + /// `redact_pii_still_blurs_bare_phone_in_content` below). + #[test] + fn has_likely_pii_ignores_scanner_bare_phone_keys() { + for key in [ + // WhatsApp group JID — chat_id = "-@g.us" + "12025551234-1543890267@g.us:2026-05-30", + // WhatsApp broadcast list + "12025551234@broadcast:2026-05-30", + // WhatsApp 1:1 JID, country-coded US number (`1` + 10 digits) + "12025551234@c.us:2026-05-30", + // Same shape carried in the namespace + "whatsapp-web:12025551234@c.us", + "whatsapp-web:12025551234-1543890267@g.us", + // Telegram numeric peer_id key + "4123456789:2026-05-30", + ] { + assert!( + !has_likely_pii(key), + "scanner-built key {key:?} must not be rejected as PII" + ); + } + } + + /// Same regression but for the E.164 (`+`-prefixed) shape — iMessage + /// posts `key = format!("{chat_id}:{day}")` where `chat_id` can be + /// `+12025551234`. Strict mode must skip; content redaction stays. + #[test] + fn has_likely_pii_ignores_bare_e164_phone_keys() { + for key in [ + "+12025551234:2026-05-30", + "imessage:+12025551234", + "imessage:+12025551234:2026-05-30", + ] { + assert!( + !has_likely_pii(key), + "E.164-shaped key {key:?} must not be rejected as PII" + ); + } + } + + /// `redact_pii` (content scrubbing path — NOT the boundary check) + /// must still redact formatted NANP and E.164 phone numbers found + /// inside document bodies. False positives in the content path only + /// blur substring bytes; they do not reject the write — which is the + /// asymmetry this PR preserves vs. the boundary check. + /// + /// Note: bare 10-digit NANP runs (`2025551234` with no separators) + /// are NOT reached by `redact_pii` at all — the SCREEN fast-path + /// requires either `\d{11,}`, a separator, or `+`, so a bare 10-digit + /// run short-circuits as "no candidate". That pre-existed this PR; a + /// pinning sentinel for it lives below. + #[test] + fn redact_pii_still_blurs_formatted_and_e164_phone_in_content() { + let out = redact_pii("call me at 202-555-1234 or +12025551234"); + let n_phone = out.value.matches(PII_PHONE).count(); + assert!( + n_phone >= 2, + "redact_pii must still blur both formatted NANP and E.164 phones in content, \ + got {n_phone} PII_PHONE token(s) in: {}", + out.value + ); + assert!(out.report.pii_redactions >= 2); + } + + /// Sentinel pinning a pre-existing SCREEN limitation: a bare 10-digit + /// NANP run (`2025551234` with no separators) is short-circuited by + /// the `SCREEN` fast-path because no `SCREEN` regex matches a 10-digit + /// bare run (`\d{11,}` is the closest, but it needs 11+). This is the + /// status quo on `main` — this PR does not change it. The test exists + /// so any future widening of `SCREEN` (e.g. to catch bare NANP) trips + /// here as a deliberate review checkpoint, NOT a regression. + #[test] + fn redact_pii_does_not_reach_bare_10_digit_nanp_today() { + let out = redact_pii("call me at 2025551234 thanks"); + assert!( + !out.value.contains(PII_PHONE), + "SCREEN fast-path historically skips bare 10-digit NANP — \ + if this test fails, SCREEN was widened; revisit the boundary-check \ + behavior in `has_likely_pii` before adjusting. Got: {}", + out.value + ); + } + + #[test] + fn empty_text_is_noop() { + unchanged(""); + } +} diff --git a/src/memory/store/safety_tests.rs b/src/memory/store/safety/safety_tests.rs similarity index 80% rename from src/memory/store/safety_tests.rs rename to src/memory/store/safety/safety_tests.rs index f1fcb69..8f2c593 100644 --- a/src/memory/store/safety_tests.rs +++ b/src/memory/store/safety/safety_tests.rs @@ -68,19 +68,25 @@ fn has_likely_secret_detects_common_patterns() { } #[test] -fn has_likely_pii_detects_email_and_phone_and_ssn() { - assert!(has_likely_pii("contact alice@example.com")); - assert!(has_likely_pii("call +15551234567")); +fn has_likely_pii_strict_boundary_flags_formatted_national_ids() { + // The write-rejection boundary is the *strict* set: formatted national IDs + // only. Bare-numeric / phone-shaped runs and email are excluded (too many + // false positives against scanner-built identifiers); they are still + // scrubbed by content redaction. Exhaustive coverage lives in `pii`'s tests. assert!(has_likely_pii("ssn 123-45-6789")); + assert!(has_likely_pii("CPF 111.444.777-35")); + assert!(has_likely_pii("cliente RFC VECJ880326XK4")); + assert!(!has_likely_pii("call +15551234567")); // phone: content-scrub only + assert!(!has_likely_pii("contact alice@example.com")); // email: out of scope assert!(!has_likely_pii("just a normal note")); } #[test] fn sanitize_text_scrubs_pii_after_secrets() { - let input = "Token sk-abcdefghijklmnopqrstuvwxyz; email a@b.com; phone +15551234567"; + let input = "Token sk-abcdefghijklmnopqrstuvwxyz; CPF 111.444.777-35; phone +15551234567"; let sanitized = sanitize_text(input); assert!(!sanitized.value.contains("sk-abcdefghijklmnopqrstuvwxyz")); - assert!(!sanitized.value.contains("a@b.com")); + assert!(!sanitized.value.contains("111.444.777-35")); assert!(!sanitized.value.contains("+15551234567")); assert!(sanitized.report.pii_redactions >= 2); } diff --git a/src/memory/store/store.rs b/src/memory/store/store.rs index 5f5f77c..dd83adb 100644 --- a/src/memory/store/store.rs +++ b/src/memory/store/store.rs @@ -93,7 +93,7 @@ impl MemoryStore for InMemoryMemoryStore { query .namespace .as_deref() - .map_or(true, |namespace| record.namespace == namespace) + .is_none_or(|namespace| record.namespace == namespace) }) .filter_map(|record| { let score = match needle.as_deref() { diff --git a/src/memory/store/vectors/embedding.rs b/src/memory/store/vectors/embedding.rs index 7030599..ddc1498 100644 --- a/src/memory/store/vectors/embedding.rs +++ b/src/memory/store/vectors/embedding.rs @@ -20,7 +20,7 @@ use async_trait::async_trait; /// instantiated backend. Drift between the two would silently split one /// embedding space into two. pub fn format_embedding_signature(name: &str, model_id: &str, dims: usize) -> String { - format!("provider={name};model={model_id};dims={dims}") + tinyagents::harness::embeddings::format_embedding_signature(name, model_id, dims) } /// Interface for embedding backends that convert text into numerical vectors. @@ -60,6 +60,38 @@ pub trait EmbeddingBackend: Send + Sync { } } +#[async_trait] +impl EmbeddingBackend for T +where + T: tinyagents::harness::embeddings::EmbeddingModel + ?Sized, +{ + fn name(&self) -> &str { + tinyagents::harness::embeddings::EmbeddingModel::name(self) + } + + fn model_id(&self) -> &str { + tinyagents::harness::embeddings::EmbeddingModel::model_id(self) + } + + fn dimensions(&self) -> usize { + tinyagents::harness::embeddings::EmbeddingModel::dimensions(self) + } + + fn signature(&self) -> String { + tinyagents::harness::embeddings::EmbeddingModel::signature(self) + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let owned = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + tinyagents::harness::embeddings::EmbeddingModel::embed(self, &owned) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} + /// An inert backend that returns deterministic all-zero vectors of a fixed /// dimension. Used by tests and by hosts that want keyword-only behaviour /// without wiring a model. Dimension defaults to diff --git a/src/memory/store/vectors/store_tests.rs b/src/memory/store/vectors/store_tests.rs index f1df118..5c3416b 100644 --- a/src/memory/store/vectors/store_tests.rs +++ b/src/memory/store/vectors/store_tests.rs @@ -66,7 +66,7 @@ fn fake_store(dims: usize) -> VectorStore { #[test] fn roundtrip_vec_bytes() { - let original = vec![1.0_f32, -2.5, 3.14, 0.0, f32::MAX, f32::MIN]; + let original = vec![1.0_f32, -2.5, 3.125, 0.0, f32::MAX, f32::MIN]; let bytes = vec_to_bytes(&original); assert_eq!(bytes.len(), original.len() * 4); assert_eq!(original, bytes_to_vec(&bytes)); diff --git a/src/memory/sync/audit.rs b/src/memory/sync/audit.rs new file mode 100644 index 0000000..3a4711b --- /dev/null +++ b/src/memory/sync/audit.rs @@ -0,0 +1,229 @@ +//! Append-only synchronization cost and outcome audit log. + +use std::io::Write; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +use crate::memory::config::MemoryConfig; + +const AUDIT_FILENAME: &str = "sync_audit.jsonl"; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SyncAuditEntry { + pub timestamp: DateTime, + pub source_id: String, + pub source_kind: String, + pub scope: String, + pub items_fetched: u32, + pub batches: u32, + pub input_tokens: u64, + pub output_tokens: u64, + pub estimated_cost_usd: f64, + #[serde(default)] + pub composio_actions_called: u32, + #[serde(default)] + pub composio_cost_usd: f64, + #[serde(default)] + pub actual_charged_usd: Option, + pub duration_ms: u64, + pub success: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl SyncAuditEntry { + pub fn combined_cost_usd(&self) -> f64 { + self.effective_cost_usd() + } + + pub fn effective_cost_usd(&self) -> f64 { + self.actual_charged_usd.unwrap_or(self.estimated_cost_usd) + self.composio_cost_usd + } +} + +pub fn append_audit_entry(config: &MemoryConfig, entry: &SyncAuditEntry) -> anyhow::Result<()> { + let directory = config.workspace.join("memory_tree"); + std::fs::create_dir_all(&directory)?; + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(directory.join(AUDIT_FILENAME))?; + serde_json::to_writer(&mut file, entry)?; + writeln!(file)?; + tracing::debug!(source_id = %entry.source_id, success = entry.success, "[memory_sync:audit] entry appended"); + Ok(()) +} + +pub fn read_audit_log(config: &MemoryConfig) -> anyhow::Result> { + let path = config.workspace.join("memory_tree").join(AUDIT_FILENAME); + let content = match std::fs::read_to_string(path) { + Ok(content) => content, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(error.into()), + }; + let mut entries: Vec<_> = content + .lines() + .filter(|line| !line.trim().is_empty()) + .filter_map(|line| match serde_json::from_str(line) { + Ok(entry) => Some(entry), + Err(error) => { + tracing::warn!(%error, "[memory_sync:audit] malformed audit line skipped"); + None + } + }) + .collect(); + entries.reverse(); + Ok(entries) +} + +pub fn estimate_cost_usd(input_tokens: u64, output_tokens: u64) -> f64 { + input_tokens as f64 * 0.07 / 1_000_000.0 + output_tokens as f64 * 0.28 / 1_000_000.0 +} + +#[derive(Debug, Default, Clone)] +pub struct RealCostAccumulator { + total_batches: u32, + batches_with_usage: u32, + batches_with_charge: u32, + est_input_tokens: u64, + est_output_tokens: u64, + real_input_tokens: u64, + real_output_tokens: u64, + real_charged_usd: f64, +} + +impl RealCostAccumulator { + pub fn new() -> Self { + Self::default() + } + + pub fn add_batch( + &mut self, + est_input: u64, + est_output: u64, + real_input: u64, + real_output: u64, + charge: Option, + ) { + self.total_batches = self.total_batches.saturating_add(1); + self.est_input_tokens = self.est_input_tokens.saturating_add(est_input); + self.est_output_tokens = self.est_output_tokens.saturating_add(est_output); + if real_input > 0 || real_output > 0 { + self.batches_with_usage = self.batches_with_usage.saturating_add(1); + self.real_input_tokens = self.real_input_tokens.saturating_add(real_input); + self.real_output_tokens = self.real_output_tokens.saturating_add(real_output); + } + if let Some(charge) = charge { + self.batches_with_charge = self.batches_with_charge.saturating_add(1); + self.real_charged_usd += charge; + } + } + + fn usage_is_complete(&self) -> bool { + self.total_batches > 0 && self.batches_with_usage == self.total_batches + } + + fn charge_is_complete(&self) -> bool { + self.total_batches > 0 && self.batches_with_charge == self.total_batches + } + + pub fn audit_input_tokens(&self) -> u64 { + if self.usage_is_complete() { + self.real_input_tokens + } else { + self.est_input_tokens + } + } + + pub fn audit_output_tokens(&self) -> u64 { + if self.usage_is_complete() { + self.real_output_tokens + } else { + self.est_output_tokens + } + } + + pub fn estimated_cost(&self) -> f64 { + estimate_cost_usd(self.est_input_tokens, self.est_output_tokens) + } + + pub fn actual_charged_usd(&self) -> Option { + self.charge_is_complete().then_some(self.real_charged_usd) + } + + pub fn usage_is_real(&self) -> bool { + self.usage_is_complete() + } + + pub fn effective_cost_usd(&self) -> f64 { + self.actual_charged_usd() + .unwrap_or_else(|| self.estimated_cost()) + } + + pub fn cost_is_actual(&self) -> bool { + self.charge_is_complete() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn entry(id: &str, timestamp: DateTime) -> SyncAuditEntry { + SyncAuditEntry { + timestamp, + source_id: id.into(), + source_kind: "test".into(), + scope: "all".into(), + items_fetched: 1, + batches: 1, + input_tokens: 10, + output_tokens: 2, + estimated_cost_usd: 0.1, + composio_actions_called: 1, + composio_cost_usd: 0.02, + actual_charged_usd: None, + duration_ms: 5, + success: true, + error: None, + } + } + + #[test] + fn audit_round_trip_is_newest_first_and_skips_malformed_lines() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + let first = entry("first", Utc::now()); + let second = entry("second", Utc::now()); + append_audit_entry(&config, &first).unwrap(); + std::fs::OpenOptions::new() + .append(true) + .open(temp.path().join("memory_tree/sync_audit.jsonl")) + .unwrap() + .write_all(b"not-json\n") + .unwrap(); + append_audit_entry(&config, &second).unwrap(); + let entries = read_audit_log(&config).unwrap(); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0].source_id, "second"); + assert_eq!(entries[1].source_id, "first"); + } + + #[test] + fn accumulator_uses_real_values_only_when_every_batch_reports_them() { + let mut complete = RealCostAccumulator::new(); + complete.add_batch(100, 10, 80, 8, Some(0.01)); + complete.add_batch(100, 10, 90, 9, Some(0.02)); + assert_eq!(complete.audit_input_tokens(), 170); + assert_eq!(complete.actual_charged_usd(), Some(0.03)); + assert!(complete.cost_is_actual()); + + let mut partial = RealCostAccumulator::new(); + partial.add_batch(100, 10, 80, 8, Some(0.01)); + partial.add_batch(200, 20, 0, 0, None); + assert_eq!(partial.audit_input_tokens(), 300); + assert_eq!(partial.actual_charged_usd(), None); + assert!(!partial.usage_is_real()); + } +} diff --git a/src/memory/sync/composio/client.rs b/src/memory/sync/composio/client.rs new file mode 100644 index 0000000..52f0be1 --- /dev/null +++ b/src/memory/sync/composio/client.rs @@ -0,0 +1,261 @@ +//! Minimal direct/proxied Composio action client. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::memory::config::{ComposioMode, ComposioSyncConfig}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecuteResponse { + #[serde(default)] + pub data: serde_json::Value, + #[serde(default)] + pub successful: bool, + #[serde(default)] + pub error: Option, + #[serde(rename = "costUsd", default)] + pub cost_usd: f64, + #[serde(rename = "markdownFormatted", default)] + pub markdown_formatted: Option, + #[serde(skip, default = "one_attempt")] + pub attempts: u32, +} + +fn one_attempt() -> u32 { + 1 +} + +#[derive(Debug, thiserror::Error)] +#[error("{message}")] +pub struct ExecuteError { + pub attempts: u32, + message: String, +} + +#[derive(Clone)] +pub struct ComposioClient { + http: reqwest::Client, + config: ComposioSyncConfig, +} + +#[async_trait] +pub trait ActionExecutor: Send + Sync { + async fn execute( + &self, + action: &str, + arguments: serde_json::Value, + connection_id: Option<&str>, + ) -> anyhow::Result; +} + +#[async_trait] +impl ActionExecutor for ComposioClient { + async fn execute( + &self, + action: &str, + arguments: serde_json::Value, + connection_id: Option<&str>, + ) -> anyhow::Result { + ComposioClient::execute(self, action, arguments, connection_id).await + } +} + +impl ComposioClient { + pub fn new(config: ComposioSyncConfig) -> Self { + Self { + http: reqwest::Client::new(), + config, + } + } + + pub fn with_http_client(mut self, http: reqwest::Client) -> Self { + self.http = http; + self + } + + pub async fn execute( + &self, + action: &str, + arguments: serde_json::Value, + connection_id: Option<&str>, + ) -> anyhow::Result { + let action = action.trim(); + anyhow::ensure!(!action.is_empty(), "Composio action must not be empty"); + const MAX_ATTEMPTS: u32 = 3; + for attempt in 1..=MAX_ATTEMPTS { + let result = match self.config.mode { + ComposioMode::Direct => { + self.execute_direct(action, arguments.clone(), connection_id) + .await + } + ComposioMode::Proxied => self.execute_proxied(action, arguments.clone()).await, + }; + match result { + Ok(mut response) + if response.successful + || !retryable_provider_error(response.error.as_deref()) + || attempt == MAX_ATTEMPTS => + { + response.attempts = attempt; + return Ok(response); + } + Ok(_) => tracing::warn!( + action, + attempt, + "[sync:composio] retrying provider rate limit" + ), + Err(error) if retryable_transport_error(&error) && attempt < MAX_ATTEMPTS => { + tracing::warn!(action, attempt, %error, "[sync:composio] retrying transient transport failure"); + } + Err(error) => { + return Err(ExecuteError { + attempts: attempt, + message: error.to_string(), + } + .into()) + } + } + tokio::time::sleep(std::time::Duration::from_millis( + 250 * 2u64.pow(attempt - 1), + )) + .await; + } + unreachable!("retry loop always returns") + } + + async fn execute_direct( + &self, + action: &str, + arguments: serde_json::Value, + connection_id: Option<&str>, + ) -> anyhow::Result { + let key = self + .config + .api_key + .as_ref() + .filter(|key| !key.is_empty()) + .map(|key| key.expose().to_owned()) + .or_else(|| std::env::var("COMPOSIO_API_KEY").ok()) + .filter(|key| !key.trim().is_empty()) + .ok_or_else(|| anyhow::anyhow!("Composio direct API key is not configured"))?; + let url = format!( + "{}/tools/execute/{action}", + self.config.base_url.trim_end_matches('/') + ); + let mut body = serde_json::json!({ "arguments": arguments }); + if let Some(entity_id) = self + .config + .entity_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + body["user_id"] = serde_json::json!(entity_id); + } + if let Some(connection_id) = connection_id + .map(str::trim) + .filter(|value| !value.is_empty()) + { + body["connected_account_id"] = serde_json::json!(connection_id); + } + + let response = self + .http + .post(url) + .header("x-api-key", key) + .json(&body) + .send() + .await + .map_err(|error| anyhow::anyhow!("Composio direct request failed: {error}"))?; + let status = response.status(); + if !status.is_success() { + let _ = response.bytes().await; + anyhow::bail!("Composio direct request failed with HTTP {status}"); + } + let raw: serde_json::Value = decode_response(response, "direct").await?; + let successful = raw + .get("successful") + .and_then(serde_json::Value::as_bool) + .or_else(|| raw.get("success").and_then(serde_json::Value::as_bool)) + .unwrap_or(true); + let error = raw + .get("error") + .and_then(serde_json::Value::as_str) + .map(str::to_owned); + let data = raw.get("data").cloned().unwrap_or(raw); + Ok(ExecuteResponse { + data, + successful, + error, + cost_usd: 0.0, + markdown_formatted: None, + attempts: 1, + }) + } + + async fn execute_proxied( + &self, + action: &str, + arguments: serde_json::Value, + ) -> anyhow::Result { + let bearer = self + .config + .bearer_token + .as_ref() + .filter(|token| !token.is_empty()) + .ok_or_else(|| anyhow::anyhow!("Composio proxy bearer token is not configured"))?; + let url = format!( + "{}/agent-integrations/composio/execute", + self.config.base_url.trim_end_matches('/') + ); + let response = self + .http + .post(url) + .bearer_auth(bearer.expose()) + .json(&serde_json::json!({ "tool": action, "arguments": arguments })) + .send() + .await + .map_err(|error| anyhow::anyhow!("Composio proxy request failed: {error}"))?; + let status = response.status(); + if !status.is_success() { + let _ = response.bytes().await; + anyhow::bail!("Composio proxy request failed with HTTP {status}"); + } + response + .json() + .await + .map_err(|error| anyhow::anyhow!("Composio proxy response decode failed: {error}")) + } +} + +fn retryable_provider_error(error: Option<&str>) -> bool { + error.is_some_and(|error| { + let lower = error.to_ascii_lowercase(); + lower.contains("ratelimit") + || lower.contains("rate limit") + || lower.contains("too many requests") + }) +} + +fn retryable_transport_error(error: &anyhow::Error) -> bool { + let message = error.to_string(); + [ + "HTTP 429", + "HTTP 502", + "HTTP 503", + "HTTP 504", + "request failed", + ] + .iter() + .any(|needle| message.contains(needle)) +} + +async fn decode_response( + response: reqwest::Response, + mode: &str, +) -> anyhow::Result { + response + .json() + .await + .map_err(|error| anyhow::anyhow!("Composio {mode} response decode failed: {error}")) +} diff --git a/src/memory/sync/composio/gmail.rs b/src/memory/sync/composio/gmail.rs new file mode 100644 index 0000000..247fa75 --- /dev/null +++ b/src/memory/sync/composio/gmail.rs @@ -0,0 +1,222 @@ +//! Incremental Gmail synchronization through Composio. + +use async_trait::async_trait; +use serde_json::Value; + +use super::client::{ActionExecutor, ComposioClient}; +use super::orchestrator::{ + run_incremental_sync, IncrementalSource, PageFetch, SyncItem, SyncScope, +}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_FETCH_EMAILS: &str = "GMAIL_FETCH_EMAILS"; + +pub struct GmailSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, + query_override: Option, +} + +impl GmailSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 10, + page_size: 50, + query_override: None, + } + } + + pub fn with_limits(mut self, max_pages: usize, page_size: usize) -> Self { + self.max_pages = max_pages.max(1); + self.page_size = page_size.max(1); + self + } + + pub fn with_query(mut self, query: impl Into) -> Self { + self.query_override = Some(query.into()); + self + } +} + +#[async_trait] +impl SyncPipeline for GmailSyncPipeline { + fn id(&self) -> &str { + "composio:gmail" + } + + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + + async fn init(&self, _config: &MemoryConfig, _context: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + + async fn tick( + &self, + _config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, _config, context).await + } +} + +#[async_trait] +impl IncrementalSource for GmailSyncPipeline { + fn toolkit(&self) -> &'static str { + "gmail" + } + + fn action(&self) -> &'static str { + ACTION_FETCH_EMAILS + } + + fn max_pages(&self) -> usize { + self.max_pages + } + + fn server_side_depth(&self) -> bool { + true + } + + fn arguments( + &self, + _scope: &SyncScope, + config: &MemoryConfig, + state: &SyncState, + page: Option<&str>, + ) -> Value { + let mut arguments = serde_json::json!({ + "max_results": self.page_size, + "include_payload": true, + }); + if let Some(token) = page { + arguments["page_token"] = serde_json::json!(token); + } + if let Some(query) = self.query_override.as_deref() { + arguments["query"] = Value::String(query.into()); + } else if let Some(cursor) = state.cursor.as_deref() { + arguments["query"] = serde_json::json!(format!( + "after:{}", + cursor_to_seconds(cursor).unwrap_or_default() + )); + } else if let Some(days) = config.sync.budget.sync_depth_days { + arguments["query"] = serde_json::json!(format!( + "after:{}", + (chrono::Utc::now() - chrono::Duration::days(days as i64)).timestamp() + )); + } + arguments + } + + fn extract_page(&self, data: &Value, _page: Option<&str>) -> PageFetch { + PageFetch { + items: extract_messages(data), + next: extract_page_token(data), + } + } + + fn dedup_key(&self, item: &Value) -> Option { + item_id(item) + } + + fn sort_cursor(&self, item: &Value) -> Option { + item_cursor(item) + } + + async fn document( + &self, + _scope: &SyncScope, + connection_id: &str, + item: SyncItem, + _executor: &dyn ActionExecutor, + _state: &mut SyncState, + ) -> anyhow::Result { + let id = item_id(&item.raw).unwrap_or_else(|| item.dedup_key.clone()); + Ok(SkillDocument { + namespace_skill_id: "gmail".into(), + connection_id: connection_id.into(), + document_id: format!("gmail:{id}"), + title: message_title(&item.raw), + content: serde_json::to_string_pretty(&item.raw)?, + toolkit: "gmail".into(), + metadata: serde_json::json!({ + "source": "composio-provider-incremental", + "taint": "external_sync", + "message_id": id, + }), + }) + } +} + +fn extract_messages(data: &Value) -> Vec { + [ + "/data/messages", + "/messages", + "/data/data/messages", + "/data/items", + "/items", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_array)) + .cloned() + .unwrap_or_default() +} + +fn extract_page_token(data: &Value) -> Option { + [ + "/data/nextPageToken", + "/nextPageToken", + "/data/data/nextPageToken", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_str)) + .map(str::trim) + .filter(|token| !token.is_empty()) + .map(str::to_owned) +} + +fn item_id(message: &Value) -> Option { + ["id", "messageId", "message_id"] + .iter() + .find_map(|key| message.get(key).and_then(Value::as_str)) + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(str::to_owned) +} + +fn item_cursor(message: &Value) -> Option { + ["internalDate", "internal_date", "date"] + .iter() + .find_map(|key| message.get(key).and_then(Value::as_str)) + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .map(str::to_owned) +} + +fn message_title(message: &Value) -> String { + ["subject", "title"] + .iter() + .find_map(|key| message.get(key).and_then(Value::as_str)) + .map(str::trim) + .filter(|title| !title.is_empty()) + .unwrap_or("Gmail message") + .to_owned() +} + +fn cursor_to_seconds(cursor: &str) -> Option { + if let Ok(milliseconds) = cursor.trim().parse::() { + return Some(milliseconds / 1000); + } + chrono::DateTime::parse_from_rfc3339(cursor) + .ok() + .map(|date| date.timestamp()) +} diff --git a/src/memory/sync/composio/mod.rs b/src/memory/sync/composio/mod.rs new file mode 100644 index 0000000..c1ed56b --- /dev/null +++ b/src/memory/sync/composio/mod.rs @@ -0,0 +1,14 @@ +//! Composio-backed synchronization. + +pub mod client; +pub mod gmail; +pub mod orchestrator; +pub mod providers; + +pub use client::{ActionExecutor, ComposioClient, ExecuteError, ExecuteResponse}; +pub use gmail::GmailSyncPipeline; +pub use orchestrator::{run_incremental_sync, IncrementalSource, PageFetch, SyncItem, SyncScope}; +pub use providers::{ + ClickUpSyncPipeline, GitHubSyncPipeline, LinearSyncPipeline, NotionSyncPipeline, + SlackSearchBackfillPipeline, SlackSyncPipeline, +}; diff --git a/src/memory/sync/composio/orchestrator.rs b/src/memory/sync/composio/orchestrator.rs new file mode 100644 index 0000000..dc64fba --- /dev/null +++ b/src/memory/sync/composio/orchestrator.rs @@ -0,0 +1,435 @@ +//! Shared bounded incremental synchronization control flow. + +use async_trait::async_trait; +use serde_json::Value; + +use super::client::ActionExecutor; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncEvent, SyncOutcome, SyncRunError, SyncStage, +}; + +#[derive(Debug)] +pub struct PageFetch { + pub items: Vec, + pub next: Option, +} + +#[derive(Debug)] +pub struct SyncItem { + pub dedup_key: String, + pub sort_cursor: Option, + pub raw: Value, +} + +#[derive(Clone, Debug, Default)] +pub struct SyncScope { + pub id: String, + pub label: String, + pub metadata: Value, +} + +impl SyncScope { + pub fn flat() -> Self { + Self::default() + } + + pub fn named(id: impl Into, label: impl Into) -> Self { + Self { + id: id.into(), + label: label.into(), + metadata: Value::Null, + } + } + + pub fn with_metadata(mut self, metadata: Value) -> Self { + self.metadata = metadata; + self + } +} + +#[async_trait] +pub trait IncrementalSource: Send + Sync { + fn toolkit(&self) -> &'static str; + fn action(&self) -> &'static str; + fn max_pages(&self) -> usize { + 10 + } + fn per_scope_cursors(&self) -> bool { + false + } + fn tolerate_scope_errors(&self) -> bool { + false + } + fn retain_dedup_keys(&self) -> bool { + true + } + fn server_side_depth(&self) -> bool { + false + } + fn depth_floor(&self, config: &MemoryConfig, state: &SyncState) -> Option { + if state.cursor.is_some() { + return None; + } + config.sync.budget.sync_depth_days.map(|days| { + (chrono::Utc::now() - chrono::Duration::days(days as i64)) + .format("%Y-%m-%dT%H:%M:%SZ") + .to_string() + }) + } + fn advance_scope_cursor(&self, _state: &mut SyncState, _scope: &SyncScope, _cursor: &str) {} + async fn scopes( + &self, + _executor: &dyn ActionExecutor, + _connection_id: &str, + _state: &mut SyncState, + ) -> anyhow::Result> { + Ok(vec![SyncScope::flat()]) + } + fn arguments( + &self, + scope: &SyncScope, + config: &MemoryConfig, + state: &SyncState, + page: Option<&str>, + ) -> Value; + fn extract_page(&self, data: &Value, page: Option<&str>) -> PageFetch; + fn dedup_key(&self, item: &Value) -> Option; + fn sort_cursor(&self, item: &Value) -> Option; + async fn document( + &self, + scope: &SyncScope, + connection_id: &str, + item: SyncItem, + executor: &dyn ActionExecutor, + state: &mut SyncState, + ) -> anyhow::Result; +} + +pub async fn run_incremental_sync( + source: &dyn IncrementalSource, + executor: &dyn ActionExecutor, + connection_id: &str, + config: &MemoryConfig, + context: &SyncContext, +) -> anyhow::Result { + let toolkit = source.toolkit(); + emit(context, toolkit, connection_id, SyncStage::Fetching, None).await; + tracing::debug!(toolkit, connection_id, "[sync:orchestrator] sync starting"); + + let mut state = SyncState::load(context.state.as_ref(), toolkit, connection_id).await?; + if state.budget_exhausted() { + tracing::debug!( + toolkit, + connection_id, + "[sync:orchestrator] daily budget exhausted" + ); + return Ok(SyncOutcome { + note: Some("daily request budget exhausted".into()), + ..SyncOutcome::default() + }); + } + + let result = match source.scopes(executor, connection_id, &mut state).await { + Ok(scopes) => { + run_pages( + source, + executor, + connection_id, + config, + context, + &mut state, + &scopes, + ) + .await + } + Err(error) => Err(error), + }; + state.last_sync_at_ms = Some(now_ms()); + if let Err(error) = state.save(context.state.as_ref()).await { + emit( + context, + toolkit, + connection_id, + SyncStage::Failed, + Some("sync state persistence failed".into()), + ) + .await; + return Err(error); + } + + match result { + Ok(outcome) => { + emit( + context, + toolkit, + connection_id, + SyncStage::Stored, + Some(format!("{} records", outcome.records_ingested)), + ) + .await; + emit(context, toolkit, connection_id, SyncStage::Completed, None).await; + tracing::debug!( + toolkit, + connection_id, + records = outcome.records_ingested, + more_pending = outcome.more_pending, + "[sync:orchestrator] sync completed" + ); + Ok(outcome) + } + Err(error) => { + tracing::warn!(toolkit, connection_id, %error, "[sync:orchestrator] sync failed"); + emit( + context, + toolkit, + connection_id, + SyncStage::Failed, + Some(error.to_string()), + ) + .await; + Err(SyncRunError::new( + error.to_string(), + state.run_requests, + state.run_provider_cost_usd, + ) + .into()) + } + } +} + +async fn run_pages( + source: &dyn IncrementalSource, + executor: &dyn ActionExecutor, + connection_id: &str, + config: &MemoryConfig, + context: &SyncContext, + state: &mut SyncState, + scopes: &[SyncScope], +) -> anyhow::Result { + let mut newest_cursor = state.cursor.clone(); + let mut ingested = 0u32; + let mut more_pending = false; + let depth_floor = (!source.server_side_depth()) + .then(|| source.depth_floor(config, state)) + .flatten(); + + 'scopes: for scope in scopes { + let mut page_token = None; + let mut scope_newest_cursor: Option = None; + let mut scope_failed = false; + for page_index in 0..source.max_pages().max(1) { + if state.budget_exhausted() { + more_pending = true; + break 'scopes; + } + let response = match executor + .execute( + source.action(), + source.arguments(scope, config, state, page_token.as_deref()), + Some(connection_id), + ) + .await + { + Ok(response) => response, + Err(error) if source.tolerate_scope_errors() => { + if let Some(execute_error) = error.downcast_ref::() + { + state.record_requests(execute_error.attempts); + } + tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] scope fetch failed; continuing"); + scope_failed = true; + break; + } + Err(error) => { + if let Some(execute_error) = error.downcast_ref::() + { + state.record_requests(execute_error.attempts); + } + return Err(error); + } + }; + // A completed provider round-trip is billable even when its envelope + // reports failure. Transport failures return before this point. + state.record_action(response.attempts, response.cost_usd); + if !response.successful { + let error = anyhow::anyhow!( + "{} provider failure: {}", + source.toolkit(), + response + .error + .unwrap_or_else(|| "unknown provider error".into()) + ); + if source.tolerate_scope_errors() { + tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] provider rejected scope; continuing"); + scope_failed = true; + break; + } + return Err(error); + } + + let fetched = source.extract_page(&response.data, page_token.as_deref()); + let mut reached_cursor_boundary = false; + for raw in fetched.items { + if config + .sync + .budget + .max_items + .is_some_and(|limit| ingested >= limit) + { + more_pending = true; + break 'scopes; + } + if state.budget_exhausted() { + more_pending = true; + break 'scopes; + } + let Some(dedup_key) = source.dedup_key(&raw) else { + continue; + }; + if state.is_synced(&dedup_key) { + continue; + } + let sort_cursor = source.sort_cursor(&raw); + if sort_cursor + .as_deref() + .zip(depth_floor.as_deref()) + .is_some_and(|(item_cursor, floor)| item_cursor < floor) + { + reached_cursor_boundary = true; + break; + } + if !source.per_scope_cursors() + && sort_cursor + .as_deref() + .zip(state.cursor.as_deref()) + .is_some_and(|(item_cursor, persisted_cursor)| { + item_cursor <= persisted_cursor + }) + { + tracing::debug!( + toolkit = source.toolkit(), + connection_id, + scope = %scope.label, + "[sync:orchestrator] reached persisted cursor boundary" + ); + reached_cursor_boundary = true; + break; + } + let document = match source + .document( + scope, + connection_id, + SyncItem { + dedup_key: dedup_key.clone(), + sort_cursor: sort_cursor.clone(), + raw, + }, + executor, + state, + ) + .await + { + Ok(document) => document, + Err(error) if source.tolerate_scope_errors() => { + tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] scope document conversion failed; continuing"); + scope_failed = true; + break; + } + Err(error) => return Err(error), + }; + if let Err(error) = context.documents.store(document).await { + if source.tolerate_scope_errors() { + tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] scope document store failed; continuing"); + scope_failed = true; + break; + } + return Err(error); + } + if source.retain_dedup_keys() { + state.mark_synced(dedup_key); + } + if let Some(cursor) = sort_cursor { + let target = if source.per_scope_cursors() { + &mut scope_newest_cursor + } else { + &mut newest_cursor + }; + if target + .as_deref() + .is_none_or(|current| cursor.as_str() > current) + { + *target = Some(cursor); + } + } + ingested = ingested.saturating_add(1); + if config + .sync + .budget + .max_items + .is_some_and(|limit| ingested >= limit) + { + more_pending = true; + break 'scopes; + } + } + + page_token = fetched.next; + if reached_cursor_boundary { + break; + } + if page_token.is_none() { + break; + } + if page_index + 1 == source.max_pages().max(1) { + more_pending = true; + } + } + if source.per_scope_cursors() && !scope_failed && !more_pending { + if let Some(cursor) = scope_newest_cursor.as_deref() { + source.advance_scope_cursor(state, scope, cursor); + state.save(context.state.as_ref()).await?; + } + } + } + + if !source.per_scope_cursors() && !more_pending { + if let Some(cursor) = newest_cursor { + state.advance_cursor(cursor); + } + } + Ok(SyncOutcome { + records_ingested: ingested, + more_pending, + actions_called: state.run_requests, + provider_cost_usd: state.run_provider_cost_usd, + note: None, + }) +} + +async fn emit( + context: &SyncContext, + toolkit: &str, + connection_id: &str, + stage: SyncStage, + message: Option, +) { + let _ = context + .events + .emit(SyncEvent { + source_id: format!("composio:{toolkit}:{connection_id}"), + toolkit: toolkit.into(), + connection_id: Some(connection_id.into()), + stage, + message, + }) + .await; +} + +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} diff --git a/src/memory/sync/composio/providers/clickup.rs b/src/memory/sync/composio/providers/clickup.rs new file mode 100644 index 0000000..5917ad5 --- /dev/null +++ b/src/memory/sync/composio/providers/clickup.rs @@ -0,0 +1,196 @@ +use async_trait::async_trait; +use serde_json::Value; + +use super::common::{checked_execute, document, first_array, pick_str}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::composio::{ + run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, + SyncScope, +}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_USER: &str = "CLICKUP_GET_AUTHORIZED_USER"; +const ACTION_WORKSPACES: &str = "CLICKUP_GET_AUTHORIZED_TEAMS_WORKSPACES"; +const ACTION_TASKS: &str = "CLICKUP_GET_FILTERED_TEAM_TASKS"; + +pub struct ClickUpSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, +} + +impl ClickUpSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 20, + page_size: 50, + } + } +} + +#[async_trait] +impl SyncPipeline for ClickUpSyncPipeline { + fn id(&self) -> &str { + "composio:clickup" + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, config, context).await + } +} + +#[async_trait] +impl IncrementalSource for ClickUpSyncPipeline { + fn toolkit(&self) -> &'static str { + "clickup" + } + fn action(&self) -> &'static str { + ACTION_TASKS + } + fn max_pages(&self) -> usize { + self.max_pages + } + fn depth_floor(&self, config: &MemoryConfig, state: &SyncState) -> Option { + if state.cursor.is_some() { + return None; + } + config.sync.budget.sync_depth_days.map(|days| { + (chrono::Utc::now() - chrono::Duration::days(days as i64)) + .timestamp_millis() + .to_string() + }) + } + async fn scopes( + &self, + executor: &dyn ActionExecutor, + connection_id: &str, + state: &mut SyncState, + ) -> anyhow::Result> { + let user_response = checked_execute( + executor, + ACTION_USER, + serde_json::json!({}), + connection_id, + state, + ) + .await?; + let user_id = ["/user/id", "/data/user/id", "/id", "/data/id"] + .iter() + .find_map(|path| user_response.data.pointer(path)) + .and_then(value_string) + .ok_or_else(|| anyhow::anyhow!("{ACTION_USER} returned no user id"))?; + if state.budget_exhausted() { + return Ok(Vec::new()); + } + let workspace_response = checked_execute( + executor, + ACTION_WORKSPACES, + serde_json::json!({}), + connection_id, + state, + ) + .await?; + let workspaces = first_array( + &workspace_response.data, + &["/teams", "/data/teams", "/workspaces", "/data/workspaces"], + ); + Ok(workspaces + .into_iter() + .filter_map(|workspace| pick_str(&workspace, &["id", "team_id", "workspace_id"])) + .map(|id| { + SyncScope::named(id.clone(), format!("workspace:{id}")) + .with_metadata(serde_json::json!({"user_id": user_id})) + }) + .collect()) + } + fn arguments( + &self, + scope: &SyncScope, + _: &MemoryConfig, + _: &SyncState, + page: Option<&str>, + ) -> Value { + serde_json::json!({"team_id": scope.id, "assignees": [scope.metadata.get("user_id").and_then(Value::as_str).unwrap_or_default()], "order_by": "updated", "reverse": true, "page": page.and_then(|value| value.parse::().ok()).unwrap_or(0), "page_size": self.page_size, "subtasks": true}) + } + fn extract_page(&self, data: &Value, page: Option<&str>) -> PageFetch { + let items = first_array( + data, + &[ + "/data/tasks", + "/tasks", + "/data/data/tasks", + "/data/results", + "/results", + "/data/items", + "/items", + ], + ); + let page_number = page + .and_then(|value| value.parse::().ok()) + .unwrap_or(0); + let next = (items.len() == self.page_size).then(|| (page_number + 1).to_string()); + PageFetch { items, next } + } + fn dedup_key(&self, item: &Value) -> Option { + let id = pick_str(item, &["id", "data.id", "task_id", "data.task_id"])?; + Some(match self.sort_cursor(item) { + Some(updated) => format!("{id}@{updated}"), + None => id, + }) + } + fn sort_cursor(&self, item: &Value) -> Option { + pick_str( + item, + &[ + "date_updated", + "data.date_updated", + "updated_at", + "data.updated_at", + "dateUpdated", + "data.dateUpdated", + ], + ) + } + async fn document( + &self, + scope: &SyncScope, + connection_id: &str, + item: SyncItem, + _: &dyn ActionExecutor, + _: &mut SyncState, + ) -> anyhow::Result { + let id = pick_str(&item.raw, &["id", "data.id", "task_id", "data.task_id"]) + .unwrap_or_else(|| item.dedup_key.clone()); + let title = pick_str(&item.raw, &["name", "data.name", "title", "data.title"]) + .unwrap_or_else(|| format!("ClickUp task {id}")); + let content = serde_json::to_string_pretty(&item.raw)?; + let mut result = document("clickup", connection_id, &id, title, content, item.raw); + result.metadata["workspace_id"] = Value::String(scope.id.clone()); + Ok(result) + } +} + +fn value_string(value: &Value) -> Option { + match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + } + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} diff --git a/src/memory/sync/composio/providers/common.rs b/src/memory/sync/composio/providers/common.rs new file mode 100644 index 0000000..a5ec838 --- /dev/null +++ b/src/memory/sync/composio/providers/common.rs @@ -0,0 +1,81 @@ +use serde_json::Value; + +use crate::memory::sync::traits::SkillDocument; + +pub fn pick_str(value: &Value, paths: &[&str]) -> Option { + paths.iter().find_map(|path| { + let pointer = format!("/{}", path.replace('.', "/")); + value + .pointer(&pointer) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }) + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) + }) +} + +pub fn first_array(data: &Value, pointers: &[&str]) -> Vec { + pointers + .iter() + .find_map(|pointer| data.pointer(pointer).and_then(Value::as_array)) + .cloned() + .unwrap_or_default() +} + +pub fn document( + toolkit: &str, + connection_id: &str, + id: &str, + title: String, + content: String, + raw: Value, +) -> SkillDocument { + SkillDocument { + namespace_skill_id: toolkit.into(), + connection_id: connection_id.into(), + document_id: format!("{toolkit}:{id}"), + title, + content, + toolkit: toolkit.into(), + metadata: serde_json::json!({ + "source": "composio-provider-incremental", + "taint": "external_sync", + "provider_id": id, + "raw": raw, + }), + } +} + +pub async fn checked_execute( + executor: &dyn super::super::client::ActionExecutor, + action: &str, + arguments: Value, + connection_id: &str, + state: &mut crate::memory::sync::state::SyncState, +) -> anyhow::Result { + let response = match executor + .execute(action, arguments, Some(connection_id)) + .await + { + Ok(response) => response, + Err(error) => { + if let Some(error) = error.downcast_ref::() { + state.record_requests(error.attempts); + } + return Err(error); + } + }; + state.record_action(response.attempts, response.cost_usd); + anyhow::ensure!( + response.successful, + "{action} provider failure: {}", + response + .error + .as_deref() + .unwrap_or("unknown provider error") + ); + Ok(response) +} diff --git a/src/memory/sync/composio/providers/github.rs b/src/memory/sync/composio/providers/github.rs new file mode 100644 index 0000000..68d3b83 --- /dev/null +++ b/src/memory/sync/composio/providers/github.rs @@ -0,0 +1,178 @@ +use async_trait::async_trait; +use serde_json::Value; + +use super::common::{checked_execute, document, first_array, pick_str}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::composio::{ + run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, + SyncScope, +}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_USER: &str = "GITHUB_GET_THE_AUTHENTICATED_USER"; +const ACTION_SEARCH: &str = "GITHUB_SEARCH_ISSUES_AND_PULL_REQUESTS"; + +pub struct GitHubSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, +} + +impl GitHubSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 20, + page_size: 50, + } + } +} + +#[async_trait] +impl SyncPipeline for GitHubSyncPipeline { + fn id(&self) -> &str { + "composio:github" + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, config, context).await + } +} + +#[async_trait] +impl IncrementalSource for GitHubSyncPipeline { + fn toolkit(&self) -> &'static str { + "github" + } + fn action(&self) -> &'static str { + ACTION_SEARCH + } + fn max_pages(&self) -> usize { + self.max_pages + } + fn server_side_depth(&self) -> bool { + true + } + async fn scopes( + &self, + executor: &dyn ActionExecutor, + connection_id: &str, + state: &mut SyncState, + ) -> anyhow::Result> { + let response = checked_execute( + executor, + ACTION_USER, + serde_json::json!({}), + connection_id, + state, + ) + .await?; + let login = pick_str(&response.data, &["login", "data.login"]) + .ok_or_else(|| anyhow::anyhow!("{ACTION_USER} returned no login"))?; + Ok(vec![SyncScope::named( + login.clone(), + format!("involves:{login}"), + )]) + } + fn arguments( + &self, + scope: &SyncScope, + config: &MemoryConfig, + state: &SyncState, + page: Option<&str>, + ) -> Value { + let mut query = format!("involves:{}", scope.id); + if let Some(cursor) = state.cursor.as_deref() { + query.push_str(&format!(" updated:>{cursor}")); + } else if let Some(days) = config.sync.budget.sync_depth_days { + let floor = chrono::Utc::now() - chrono::Duration::days(days as i64); + query.push_str(&format!(" updated:>{}", floor.format("%Y-%m-%dT%H:%M:%SZ"))); + } + serde_json::json!({ "q": query, "sort": "updated", "order": "desc", "per_page": self.page_size, "page": page.and_then(|value| value.parse::().ok()).unwrap_or(1) }) + } + fn extract_page(&self, data: &Value, page: Option<&str>) -> PageFetch { + let items = first_array( + data, + &[ + "/data/items", + "/items", + "/data/data/items", + "/data/results", + "/results", + ], + ); + let page_number = page + .and_then(|value| value.parse::().ok()) + .unwrap_or(1); + let next = (items.len() == self.page_size).then(|| (page_number + 1).to_string()); + PageFetch { items, next } + } + fn dedup_key(&self, item: &Value) -> Option { + let id = issue_id(item)?; + Some(match self.sort_cursor(item) { + Some(updated) => format!("{id}@{updated}"), + None => id, + }) + } + fn sort_cursor(&self, item: &Value) -> Option { + pick_str( + item, + &[ + "updated_at", + "data.updated_at", + "updatedAt", + "data.updatedAt", + ], + ) + } + async fn document( + &self, + _: &SyncScope, + connection_id: &str, + item: SyncItem, + _: &dyn ActionExecutor, + _: &mut SyncState, + ) -> anyhow::Result { + let id = issue_id(&item.raw).unwrap_or_else(|| item.dedup_key.clone()); + let title = pick_str(&item.raw, &["title", "data.title"]) + .unwrap_or_else(|| format!("GitHub issue {id}")); + let content = serde_json::to_string_pretty(&item.raw)?; + Ok(document( + "github", + connection_id, + &id, + title, + content, + item.raw, + )) + } +} + +fn issue_id(item: &Value) -> Option { + pick_str(item, &["id", "data.id"]).or_else(|| { + let url = pick_str(item, &["html_url", "data.html_url", "url", "data.url"])?; + let parts: Vec<_> = url.trim_end_matches('/').split('/').collect(); + (parts.len() >= 7).then(|| { + format!( + "{}/{}#{}", + parts[parts.len() - 4], + parts[parts.len() - 3], + parts[parts.len() - 1] + ) + }) + }) +} diff --git a/src/memory/sync/composio/providers/linear.rs b/src/memory/sync/composio/providers/linear.rs new file mode 100644 index 0000000..2c302d6 --- /dev/null +++ b/src/memory/sync/composio/providers/linear.rs @@ -0,0 +1,193 @@ +use async_trait::async_trait; +use serde_json::Value; + +use super::common::{checked_execute, document, first_array, pick_str}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::composio::{ + run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, + SyncScope, +}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_USERS: &str = "LINEAR_LIST_LINEAR_USERS"; +const ACTION_ISSUES: &str = "LINEAR_LIST_LINEAR_ISSUES"; + +pub struct LinearSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, +} + +impl LinearSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 20, + page_size: 50, + } + } +} + +#[async_trait] +impl SyncPipeline for LinearSyncPipeline { + fn id(&self) -> &str { + "composio:linear" + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, config, context).await + } +} + +#[async_trait] +impl IncrementalSource for LinearSyncPipeline { + fn toolkit(&self) -> &'static str { + "linear" + } + fn action(&self) -> &'static str { + ACTION_ISSUES + } + fn max_pages(&self) -> usize { + self.max_pages + } + async fn scopes( + &self, + executor: &dyn ActionExecutor, + connection_id: &str, + state: &mut SyncState, + ) -> anyhow::Result> { + let response = checked_execute( + executor, + ACTION_USERS, + serde_json::json!({"isMe": true}), + connection_id, + state, + ) + .await?; + let users = first_array( + &response.data, + &[ + "/data/nodes", + "/nodes", + "/data/data/nodes", + "/data/users/nodes", + ], + ); + let viewer = users.first().unwrap_or(&response.data); + let id = pick_str(viewer, &["id", "data.id"]) + .ok_or_else(|| anyhow::anyhow!("{ACTION_USERS} returned no viewer id"))?; + Ok(vec![SyncScope::named(id, "assignee:me")]) + } + fn arguments( + &self, + scope: &SyncScope, + _: &MemoryConfig, + _: &SyncState, + page: Option<&str>, + ) -> Value { + let mut args = serde_json::json!({"assigneeId": scope.id, "first": self.page_size, "orderBy": "updatedAt"}); + if let Some(page) = page { + args["after"] = serde_json::json!(page); + } + args + } + fn extract_page(&self, data: &Value, _: Option<&str>) -> PageFetch { + let items = first_array( + data, + &[ + "/data/nodes", + "/nodes", + "/data/data/nodes", + "/data/issues/nodes", + "/data/results", + "/results", + "/data/items", + "/items", + ], + ); + let page_info = [ + "/data/pageInfo", + "/pageInfo", + "/data/data/pageInfo", + "/data/issues/pageInfo", + ] + .iter() + .find_map(|path| data.pointer(path)); + let next = page_info + .filter(|info| { + info.get("hasNextPage") + .and_then(Value::as_bool) + .unwrap_or(false) + }) + .and_then(|info| info.get("endCursor").and_then(Value::as_str)) + .map(str::to_owned); + PageFetch { items, next } + } + fn dedup_key(&self, item: &Value) -> Option { + let id = pick_str(item, &["id", "data.id", "identifier", "data.identifier"])?; + Some(match self.sort_cursor(item) { + Some(updated) => format!("{id}@{updated}"), + None => id, + }) + } + fn sort_cursor(&self, item: &Value) -> Option { + pick_str( + item, + &[ + "updatedAt", + "data.updatedAt", + "updated_at", + "data.updated_at", + ], + ) + } + async fn document( + &self, + _: &SyncScope, + connection_id: &str, + item: SyncItem, + _: &dyn ActionExecutor, + _: &mut SyncState, + ) -> anyhow::Result { + let id = pick_str( + &item.raw, + &["id", "data.id", "identifier", "data.identifier"], + ) + .unwrap_or_else(|| item.dedup_key.clone()); + let title = pick_str( + &item.raw, + &[ + "title", + "data.title", + "name", + "data.name", + "identifier", + "data.identifier", + ], + ) + .unwrap_or_else(|| format!("Linear issue {id}")); + let content = serde_json::to_string_pretty(&item.raw)?; + Ok(document( + "linear", + connection_id, + &id, + title, + content, + item.raw, + )) + } +} diff --git a/src/memory/sync/composio/providers/mod.rs b/src/memory/sync/composio/providers/mod.rs new file mode 100644 index 0000000..3cad8d1 --- /dev/null +++ b/src/memory/sync/composio/providers/mod.rs @@ -0,0 +1,14 @@ +//! Incremental Composio provider pipelines. + +mod clickup; +mod common; +mod github; +mod linear; +mod notion; +mod slack; + +pub use clickup::ClickUpSyncPipeline; +pub use github::GitHubSyncPipeline; +pub use linear::LinearSyncPipeline; +pub use notion::NotionSyncPipeline; +pub use slack::{SlackSearchBackfillPipeline, SlackSyncPipeline}; diff --git a/src/memory/sync/composio/providers/notion.rs b/src/memory/sync/composio/providers/notion.rs new file mode 100644 index 0000000..33c4038 --- /dev/null +++ b/src/memory/sync/composio/providers/notion.rs @@ -0,0 +1,193 @@ +use async_trait::async_trait; +use serde_json::Value; + +use super::common::{checked_execute, document, first_array, pick_str}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::composio::{ + run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, + SyncScope, +}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_FETCH: &str = "NOTION_FETCH_DATA"; +const ACTION_MARKDOWN: &str = "NOTION_GET_PAGE_MARKDOWN"; + +pub struct NotionSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, +} + +impl NotionSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 20, + page_size: 25, + } + } +} + +#[async_trait] +impl SyncPipeline for NotionSyncPipeline { + fn id(&self) -> &str { + "composio:notion" + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, config, context).await + } +} + +#[async_trait] +impl IncrementalSource for NotionSyncPipeline { + fn toolkit(&self) -> &'static str { + "notion" + } + fn action(&self) -> &'static str { + ACTION_FETCH + } + fn max_pages(&self) -> usize { + self.max_pages + } + fn arguments( + &self, + _: &SyncScope, + _: &MemoryConfig, + _: &SyncState, + page: Option<&str>, + ) -> Value { + let mut args = serde_json::json!({"page_size": self.page_size, "filter": {"value": "page", "property": "object"}, "sort": {"direction": "descending", "timestamp": "last_edited_time"}}); + if let Some(page) = page { + args["start_cursor"] = serde_json::json!(page); + } + args + } + fn extract_page(&self, data: &Value, _: Option<&str>) -> PageFetch { + PageFetch { + items: first_array( + data, + &[ + "/data/results", + "/results", + "/data/data/results", + "/data/items", + "/items", + ], + ), + next: [ + "/data/next_cursor", + "/next_cursor", + "/data/data/next_cursor", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_str)) + .map(str::to_owned), + } + } + fn dedup_key(&self, item: &Value) -> Option { + let id = pick_str(item, &["id", "data.id", "pageId", "data.pageId"])?; + Some(match self.sort_cursor(item) { + Some(edited) => format!("{id}@{edited}"), + None => id, + }) + } + fn sort_cursor(&self, item: &Value) -> Option { + pick_str( + item, + &[ + "last_edited_time", + "data.last_edited_time", + "lastEditedTime", + "data.lastEditedTime", + ], + ) + } + async fn document( + &self, + _: &SyncScope, + connection_id: &str, + item: SyncItem, + executor: &dyn ActionExecutor, + state: &mut SyncState, + ) -> anyhow::Result { + let id = pick_str(&item.raw, &["id", "data.id", "pageId", "data.pageId"]) + .unwrap_or_else(|| item.dedup_key.clone()); + let title = notion_title(&item.raw).unwrap_or_else(|| format!("Notion page {id}")); + let response = checked_execute( + executor, + ACTION_MARKDOWN, + serde_json::json!({"page_id": id}), + connection_id, + state, + ) + .await?; + let content = [ + "/markdown", + "/data/markdown", + "/data/response_data/markdown", + "/response_data/markdown", + "/data/content", + "/content", + "/text", + "/data/text", + ] + .iter() + .find_map(|path| response.data.pointer(path).and_then(Value::as_str)) + .filter(|value| !value.trim().is_empty()) + .map(str::to_owned) + .unwrap_or(serde_json::to_string_pretty(&item.raw)?); + Ok(document( + "notion", + connection_id, + &id, + title, + content, + item.raw, + )) + } +} + +fn notion_title(page: &Value) -> Option { + let properties = page + .get("properties") + .or_else(|| page.pointer("/data/properties")); + properties + .and_then(Value::as_object) + .and_then(|props| { + props.values().find_map(|property| { + (property.get("type").and_then(Value::as_str) == Some("title")) + .then(|| { + property + .get("title") + .and_then(Value::as_array) + .map(|parts| { + parts + .iter() + .filter_map(|part| { + part.get("plain_text").and_then(Value::as_str) + }) + .collect::>() + .join("") + }) + }) + .flatten() + .filter(|title| !title.is_empty()) + }) + }) + .or_else(|| pick_str(page, &["title", "data.title", "name", "data.name"])) +} diff --git a/src/memory/sync/composio/providers/slack.rs b/src/memory/sync/composio/providers/slack.rs new file mode 100644 index 0000000..5dbea46 --- /dev/null +++ b/src/memory/sync/composio/providers/slack.rs @@ -0,0 +1,520 @@ +use std::collections::{BTreeMap, HashMap}; +use std::sync::OnceLock; + +use async_trait::async_trait; +use chrono::Utc; +use serde_json::Value; + +use super::common::{checked_execute, document, first_array, pick_str}; +use crate::memory::config::MemoryConfig; +use crate::memory::sync::composio::{ + run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, + SyncScope, +}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, +}; + +const ACTION_CHANNELS: &str = "SLACK_LIST_CONVERSATIONS"; +const ACTION_HISTORY: &str = "SLACK_FETCH_CONVERSATION_HISTORY"; +const ACTION_SEARCH: &str = "SLACK_SEARCH_MESSAGES"; + +pub struct SlackSyncPipeline { + client: ComposioClient, + connection_id: String, + max_pages: usize, + page_size: usize, + backfill_days: i64, +} + +pub struct SlackSearchBackfillPipeline { + client: ComposioClient, + connection_id: String, + backfill_days: i64, + max_pages: u32, +} + +impl SlackSearchBackfillPipeline { + pub fn new( + client: ComposioClient, + connection_id: impl Into, + backfill_days: i64, + ) -> Self { + Self { + client, + connection_id: connection_id.into(), + backfill_days: backfill_days.max(1), + max_pages: 50, + } + } +} + +#[async_trait] +impl SyncPipeline for SlackSearchBackfillPipeline { + fn id(&self) -> &str { + "composio:slack:search-backfill" + } + + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + + async fn tick(&self, _: &MemoryConfig, context: &SyncContext) -> anyhow::Result { + let mut state = + SyncState::load(context.state.as_ref(), "slack", &self.connection_id).await?; + if state.budget_exhausted() { + return Ok(SyncOutcome { + note: Some("slack search-backfill skipped: daily budget exhausted".into()), + ..SyncOutcome::default() + }); + } + + let directory = SlackSyncPipeline::new(self.client.clone(), self.connection_id.clone()); + let scopes = directory + .scopes(&self.client, &self.connection_id, &mut state) + .await?; + let channels: HashMap<_, _> = scopes + .into_iter() + .map(|scope| (scope.id.clone(), scope)) + .collect(); + let users = channels + .values() + .find_map(|scope| scope.metadata.get("users").and_then(Value::as_object)); + let after = (Utc::now() - chrono::Duration::days(self.backfill_days)) + .format("%Y-%m-%d") + .to_string(); + let mut page = 1u32; + let mut total_pages = 1u32; + let mut stored = 0u32; + + loop { + if state.budget_exhausted() || page > self.max_pages { + break; + } + let response = checked_execute( + &self.client, + ACTION_SEARCH, + serde_json::json!({ + "query": format!("after:{after}"), + "count": 100, + "sort": "timestamp", + "sort_dir": "asc", + "page": page, + }), + &self.connection_id, + &mut state, + ) + .await?; + if page == 1 { + total_pages = search_total_pages(&response.data).min(self.max_pages); + } + let matches = search_matches(&response.data); + let fetched = matches.len(); + for raw in matches { + let Some(ts) = pick_str(&raw, &["ts"]) else { + continue; + }; + if parse_ts(&ts).is_none() { + continue; + } + let Some(text) = pick_str(&raw, &["text"]).filter(|text| !text.trim().is_empty()) + else { + continue; + }; + let Some(channel_id) = pick_str(&raw, &["channel.id", "channel_id"]) else { + continue; + }; + let Some(scope) = channels.get(&channel_id) else { + tracing::warn!(channel_id, "[sync:slack-search] unknown channel skipped"); + continue; + }; + let author_id = + pick_str(&raw, &["user", "bot_id"]).unwrap_or_else(|| "unknown".into()); + let author = users + .and_then(|users| users.get(&author_id)) + .and_then(Value::as_str) + .unwrap_or(&author_id); + let text = replace_mentions(&text, users); + let mut doc = document( + "slack", + &self.connection_id, + &format!("{channel_id}:{ts}"), + format!("Slack {} from {author}", scope.label), + format!("[{ts}] {author}: {text}"), + raw, + ); + doc.metadata["channel_id"] = Value::String(channel_id); + doc.metadata["channel_label"] = Value::String(scope.label.clone()); + context.documents.store(doc).await?; + stored = stored.saturating_add(1); + } + if fetched == 0 || page >= total_pages { + break; + } + page = page.saturating_add(1); + } + + state.last_sync_at_ms = Some(Utc::now().timestamp_millis() as u64); + state.save(context.state.as_ref()).await?; + Ok(SyncOutcome { + records_ingested: stored, + more_pending: page < total_pages, + actions_called: state.run_requests, + provider_cost_usd: state.run_provider_cost_usd, + note: Some(format!( + "slack search-backfill: pages={page} records={stored}" + )), + }) + } +} + +impl SlackSyncPipeline { + pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self { + client, + connection_id: connection_id.into(), + max_pages: 20, + page_size: 200, + backfill_days: 30, + } + } +} + +#[async_trait] +impl SyncPipeline for SlackSyncPipeline { + fn id(&self) -> &str { + "composio:slack" + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Composio + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + run_incremental_sync(self, &self.client, &self.connection_id, config, context).await + } +} + +#[async_trait] +impl IncrementalSource for SlackSyncPipeline { + fn toolkit(&self) -> &'static str { + "slack" + } + fn action(&self) -> &'static str { + ACTION_HISTORY + } + fn max_pages(&self) -> usize { + self.max_pages + } + fn per_scope_cursors(&self) -> bool { + true + } + fn server_side_depth(&self) -> bool { + true + } + fn tolerate_scope_errors(&self) -> bool { + true + } + fn retain_dedup_keys(&self) -> bool { + true + } + + fn advance_scope_cursor(&self, state: &mut SyncState, scope: &SyncScope, cursor: &str) { + let mut cursors = decode_cursors(state.cursor.as_deref()); + cursors.insert(scope.id.clone(), cursor.into()); + state.cursor = serde_json::to_string(&cursors).ok(); + } + + async fn scopes( + &self, + executor: &dyn ActionExecutor, + connection_id: &str, + state: &mut SyncState, + ) -> anyhow::Result> { + let users = fetch_users(executor, connection_id, state).await; + let mut cursor: Option = None; + let mut channels = Vec::new(); + for _ in 0..20 { + if state.budget_exhausted() { + break; + } + let mut args = serde_json::json!({"limit": 200, "types": "public_channel,private_channel", "exclude_archived": true}); + if let Some(cursor) = cursor.as_deref() { + args["cursor"] = Value::String(cursor.into()); + } + let response = + checked_execute(executor, ACTION_CHANNELS, args, connection_id, state).await?; + channels.extend(first_array( + &response.data, + &["/data/channels", "/channels", "/data/data/channels"], + )); + cursor = next_cursor(&response.data); + if cursor.is_none() { + break; + } + } + Ok(channels + .into_iter() + .filter_map(|channel| { + let id = pick_str(&channel, &["id", "data.id"])?; + let name = pick_str(&channel, &["name", "data.name"]).unwrap_or_else(|| id.clone()); + let private = channel + .get("is_private") + .and_then(Value::as_bool) + .unwrap_or(false); + let label = if private { + format!("private:{name}") + } else { + format!("#{name}") + }; + Some( + SyncScope::named(id, label).with_metadata(serde_json::json!({ + "channel": channel, + "users": users, + })), + ) + }) + .collect()) + } + + fn arguments( + &self, + scope: &SyncScope, + config: &MemoryConfig, + state: &SyncState, + page: Option<&str>, + ) -> Value { + let cursors = decode_cursors(state.cursor.as_deref()); + let oldest = cursors.get(&scope.id).cloned().unwrap_or_else(|| { + format!( + "{}.000000", + (Utc::now() + - chrono::Duration::days( + config + .sync + .budget + .sync_depth_days + .map(i64::from) + .unwrap_or(self.backfill_days) + )) + .timestamp() + ) + }); + let mut args = serde_json::json!({"channel": scope.id, "oldest": oldest, "inclusive": false, "limit": self.page_size}); + if let Some(page) = page { + args["cursor"] = Value::String(page.into()); + } + args + } + + fn extract_page(&self, data: &Value, _: Option<&str>) -> PageFetch { + PageFetch { + items: first_array( + data, + &["/data/messages", "/messages", "/data/data/messages"], + ), + next: next_cursor(data), + } + } + + fn dedup_key(&self, item: &Value) -> Option { + let ts = pick_str(item, &["ts", "data.ts"])?; + parse_ts(&ts)?; + let text = pick_str(item, &["text", "data.text"])?; + (!text.trim().is_empty()).then_some(ts) + } + + fn sort_cursor(&self, item: &Value) -> Option { + pick_str(item, &["ts", "data.ts"]) + } + + async fn document( + &self, + scope: &SyncScope, + connection_id: &str, + item: SyncItem, + _: &dyn ActionExecutor, + _: &mut SyncState, + ) -> anyhow::Result { + let ts = pick_str(&item.raw, &["ts", "data.ts"]).unwrap_or(item.dedup_key); + let raw_text = pick_str(&item.raw, &["text", "data.text"]).unwrap_or_default(); + let author_id = pick_str( + &item.raw, + &["user", "data.user", "username", "data.username"], + ) + .unwrap_or_else(|| "unknown".into()); + let users = scope.metadata.get("users").and_then(Value::as_object); + let author = users + .and_then(|users| users.get(&author_id)) + .and_then(Value::as_str) + .unwrap_or(&author_id) + .to_owned(); + let text = replace_mentions(&raw_text, users); + let title = format!("Slack {} from {}", scope.label, author); + let content = format!("[{ts}] {author}: {text}"); + let mut result = document( + "slack", + connection_id, + &format!("{}:{ts}", scope.id), + title, + content, + item.raw, + ); + result.metadata["channel_id"] = Value::String(scope.id.clone()); + result.metadata["channel_label"] = Value::String(scope.label.clone()); + Ok(result) + } +} + +async fn fetch_users( + executor: &dyn ActionExecutor, + connection_id: &str, + state: &mut SyncState, +) -> HashMap { + let mut users = HashMap::new(); + let mut cursor: Option = None; + for page in 0..20 { + if state.budget_exhausted() { + break; + } + let mut arguments = serde_json::json!({"limit": 200}); + if let Some(cursor) = cursor.as_deref() { + arguments["cursor"] = Value::String(cursor.into()); + } + let response = match executor + .execute("SLACK_LIST_ALL_USERS", arguments, Some(connection_id)) + .await + { + Ok(response) => response, + Err(error) => { + if let Some(error) = error.downcast_ref::() { + state.record_requests(error.attempts); + } + tracing::warn!(page, %error, "[sync:slack] user directory fetch failed; using collected users"); + break; + } + }; + state.record_action(response.attempts, response.cost_usd); + if !response.successful { + tracing::warn!( + page, + error = response.error.as_deref().unwrap_or("provider failure"), + "[sync:slack] user directory rejected; using collected users" + ); + break; + } + let members = first_array( + &response.data, + &[ + "/data/members", + "/members", + "/data/users", + "/users", + "/data/data/members", + ], + ); + for member in members { + let Some(id) = pick_str(&member, &["id"]) else { + continue; + }; + if let Some(name) = [ + "profile.display_name", + "profile.real_name", + "real_name", + "name", + "profile.display_name_normalized", + "profile.real_name_normalized", + ] + .iter() + .find_map(|path| pick_str(&member, &[*path])) + { + users.insert(id, name); + } + } + cursor = next_cursor(&response.data); + if cursor.is_none() { + break; + } + } + users +} + +fn mention_regex() -> &'static regex::Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| regex::Regex::new(r"<@(U[A-Z0-9]+)>").expect("Slack mention regex")) +} + +fn replace_mentions(text: &str, users: Option<&serde_json::Map>) -> String { + mention_regex() + .replace_all(text, |captures: ®ex::Captures<'_>| { + let id = &captures[1]; + let resolved = users + .and_then(|users| users.get(id)) + .and_then(Value::as_str) + .unwrap_or(id); + format!("@{resolved}") + }) + .into_owned() +} + +fn next_cursor(data: &Value) -> Option { + [ + "/data/response_metadata/next_cursor", + "/response_metadata/next_cursor", + "/data/next_cursor", + "/next_cursor", + "/data/data/response_metadata/next_cursor", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_str)) + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .map(str::to_owned) +} + +fn search_matches(data: &Value) -> Vec { + first_array( + data, + &[ + "/data/messages/matches", + "/messages/matches", + "/data/data/messages/matches", + "/messages", + ], + ) +} + +fn search_total_pages(data: &Value) -> u32 { + [ + "/data/messages/paging/pages", + "/messages/paging/pages", + "/data/data/messages/paging/pages", + "/pages", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_u64)) + .unwrap_or(1) as u32 +} + +fn decode_cursors(raw: Option<&str>) -> BTreeMap { + raw.and_then(|raw| serde_json::from_str(raw).ok()) + .unwrap_or_default() +} + +fn parse_ts(ts: &str) -> Option<(i64, u64)> { + let mut parts = ts.splitn(2, '.'); + Some(( + parts.next()?.parse().ok()?, + parts.next().unwrap_or("0").parse().ok()?, + )) +} diff --git a/src/memory/sync/dispatcher.rs b/src/memory/sync/dispatcher.rs new file mode 100644 index 0000000..8f8e1f6 --- /dev/null +++ b/src/memory/sync/dispatcher.rs @@ -0,0 +1,252 @@ +//! Pipeline registry and fault-isolated synchronization dispatcher. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +use crate::memory::config::MemoryConfig; +use crate::memory::sync::traits::{SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SyncRunResult { + pub pipeline_id: String, + pub kind: SyncPipelineKind, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub outcome: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Default)] +pub struct SyncDispatcher { + pipelines: BTreeMap>, +} + +impl SyncDispatcher { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, pipeline: Arc) -> anyhow::Result<()> { + let id = pipeline.id().trim(); + anyhow::ensure!(!id.is_empty(), "sync pipeline id must not be empty"); + anyhow::ensure!( + !self.pipelines.contains_key(id), + "sync pipeline already registered: {id}" + ); + tracing::debug!( + pipeline_id = id, + kind = pipeline.kind().as_str(), + "[memory_sync:dispatcher] registering pipeline" + ); + self.pipelines.insert(id.to_owned(), pipeline); + Ok(()) + } + + pub fn ids(&self) -> Vec<&str> { + self.pipelines.keys().map(String::as_str).collect() + } + + pub async fn init_all( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> Vec { + let mut results = Vec::with_capacity(self.pipelines.len()); + for (id, pipeline) in &self.pipelines { + tracing::debug!( + pipeline_id = id, + "[memory_sync:dispatcher] initializing pipeline" + ); + let result = pipeline.init(config, context).await; + results.push(SyncRunResult { + pipeline_id: id.clone(), + kind: pipeline.kind(), + outcome: result.as_ref().ok().map(|_| SyncOutcome::default()), + error: result.err().map(|error| error.to_string()), + }); + } + results + } + + pub async fn tick( + &self, + pipeline_id: &str, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + let pipeline = self + .pipelines + .get(pipeline_id) + .ok_or_else(|| anyhow::anyhow!("unknown sync pipeline: {pipeline_id}"))?; + tracing::debug!( + pipeline_id, + "[memory_sync:dispatcher] pipeline tick starting" + ); + let outcome = pipeline.tick(config, context).await; + match &outcome { + Ok(outcome) => tracing::debug!( + pipeline_id, + records = outcome.records_ingested, + more_pending = outcome.more_pending, + "[memory_sync:dispatcher] pipeline tick completed" + ), + Err(error) => { + tracing::warn!(pipeline_id, %error, "[memory_sync:dispatcher] pipeline tick failed") + } + } + outcome + } + + pub async fn tick_all( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> Vec { + let mut results = Vec::with_capacity(self.pipelines.len()); + for (id, pipeline) in &self.pipelines { + let result = pipeline.tick(config, context).await; + results.push(SyncRunResult { + pipeline_id: id.clone(), + kind: pipeline.kind(), + outcome: result.as_ref().ok().cloned(), + error: result.err().map(|error| error.to_string()), + }); + } + results + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Mutex; + + use async_trait::async_trait; + + use super::*; + use crate::memory::sync::state::SyncStateStore; + use crate::memory::sync::traits::{SkillDocSink, SkillDocument, SyncEvent, SyncEventSink}; + + struct FakePipeline { + id: &'static str, + fail: bool, + } + + #[async_trait] + impl SyncPipeline for FakePipeline { + fn id(&self) -> &str { + self.id + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Workspace + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + async fn tick(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result { + if self.fail { + anyhow::bail!("expected failure") + } + Ok(SyncOutcome { + records_ingested: 3, + more_pending: false, + actions_called: 0, + provider_cost_usd: 0.0, + note: None, + }) + } + } + + #[derive(Default)] + struct NoopHost(Mutex>); + #[async_trait] + impl SkillDocSink for NoopHost { + async fn store(&self, _: SkillDocument) -> anyhow::Result<()> { + Ok(()) + } + async fn delete(&self, _: &str, _: &str) -> anyhow::Result<()> { + Ok(()) + } + } + #[async_trait] + impl SyncEventSink for NoopHost { + async fn emit(&self, _: SyncEvent) -> anyhow::Result<()> { + Ok(()) + } + } + #[async_trait] + impl SyncStateStore for NoopHost { + async fn get( + &self, + namespace: &str, + key: &str, + ) -> anyhow::Result> { + Ok(self + .0 + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.0 + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } + } + + fn context() -> SyncContext { + let host = Arc::new(NoopHost::default()); + SyncContext { + events: host.clone(), + documents: host.clone(), + state: host, + local_documents: None, + external_sources: None, + summariser: None, + } + } + + #[tokio::test] + async fn tick_all_is_deterministic_and_isolates_failures() { + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(Arc::new(FakePipeline { + id: "z-fail", + fail: true, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "a-ok", + fail: false, + })) + .unwrap(); + assert_eq!(dispatcher.ids(), vec!["a-ok", "z-fail"]); + assert!(dispatcher + .register(Arc::new(FakePipeline { + id: "a-ok", + fail: false + })) + .is_err()); + let results = dispatcher + .tick_all(&MemoryConfig::new("/tmp/unused"), &context()) + .await; + assert_eq!(results.len(), 2); + assert_eq!(results[0].outcome.as_ref().unwrap().records_ingested, 3); + assert!(results[1] + .error + .as_deref() + .unwrap() + .contains("expected failure")); + } +} diff --git a/src/memory/sync/github.rs b/src/memory/sync/github.rs new file mode 100644 index 0000000..8b5c2d7 --- /dev/null +++ b/src/memory/sync/github.rs @@ -0,0 +1,199 @@ +//! GitHub repository synchronization through a host-provided network reader. + +use async_trait::async_trait; + +use crate::memory::config::MemoryConfig; +use crate::memory::sources::MemorySourceEntry; +use crate::memory::store::content::{write_raw_items, RawItem, RawKind}; + +use super::rebuild::rebuild_tree_from_raw_with_audit; +use super::traits::{ + SyncContext, SyncEvent, SyncOutcome, SyncPipeline, SyncPipelineKind, SyncStage, +}; + +pub struct GithubRepoSyncPipeline { + id: String, + source: MemorySourceEntry, +} + +impl GithubRepoSyncPipeline { + pub fn new(source: MemorySourceEntry) -> anyhow::Result { + source.validate().map_err(anyhow::Error::msg)?; + if source.kind != crate::memory::sources::SourceKind::GithubRepo { + anyhow::bail!("GitHub repo pipeline requires github_repo source"); + } + Ok(Self { + id: format!("workspace:github_repo:{}", source.id), + source, + }) + } + + async fn event(&self, context: &SyncContext, stage: SyncStage, message: Option) { + let _ = context + .events + .emit(SyncEvent { + source_id: self.id.clone(), + toolkit: "github_repo".into(), + connection_id: Some(self.source.id.clone()), + stage, + message, + }) + .await; + } +} + +#[async_trait] +impl SyncPipeline for GithubRepoSyncPipeline { + fn id(&self) -> &str { + &self.id + } + + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Workspace + } + + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + let reader = context + .external_sources + .as_ref() + .ok_or_else(|| anyhow::anyhow!("GitHub repo pipeline requires an external reader"))?; + let summariser = context + .summariser + .as_ref() + .ok_or_else(|| anyhow::anyhow!("GitHub repo pipeline requires a summariser"))?; + let url = self + .source + .url + .as_deref() + .ok_or_else(|| anyhow::anyhow!("GitHub repo source missing url"))?; + let (owner, repo) = parse_repo(url)?; + let tree_scope = format!("github:{owner}/{repo}"); + let archive_source_id = format!("github.com/{owner}/{repo}"); + + self.event(context, SyncStage::Fetching, None).await; + let items = reader.list_items(&self.source).await?; + let content_root = crate::memory::chunks::content_root(config); + let mut archived = 0u32; + for item in &items { + let Some((kind, uid)) = raw_coordinates(&item.id) else { + tracing::warn!(item_id = %item.id, "[memory_sync:github] unsupported item id skipped"); + continue; + }; + let content = match reader.read_item(&self.source, &item.id).await { + Ok(content) => content, + Err(error) => { + tracing::warn!(item_id = %item.id, %error, "[memory_sync:github] item read failed"); + continue; + } + }; + write_raw_items( + &content_root, + &archive_source_id, + &[RawItem { + uid: &uid, + created_at_ms: item.updated_at_ms.unwrap_or(0), + markdown: &content.body, + kind, + }], + )?; + archived = archived.saturating_add(1); + } + + self.event( + context, + SyncStage::Ingesting, + Some(format!("{archived} items archived")), + ) + .await; + let rebuilt = rebuild_tree_from_raw_with_audit( + config, + &tree_scope, + &archive_source_id, + summariser.as_ref(), + &self.source.id, + self.source.kind.as_str(), + ) + .await?; + self.event(context, SyncStage::Completed, None).await; + Ok(SyncOutcome { + records_ingested: archived, + more_pending: false, + note: Some(format!( + "{archived} items archived, {} summaries rebuilt", + rebuilt.batches + )), + ..SyncOutcome::default() + }) + } +} + +fn parse_repo(url: &str) -> anyhow::Result<(String, String)> { + let path = url + .trim() + .trim_end_matches('/') + .strip_suffix(".git") + .unwrap_or(url.trim().trim_end_matches('/')); + let path = path + .strip_prefix("https://github.com/") + .or_else(|| path.strip_prefix("http://github.com/")) + .or_else(|| path.strip_prefix("git@github.com:")) + .ok_or_else(|| anyhow::anyhow!("unsupported GitHub repository URL"))?; + let mut parts = path.split('/'); + let owner = parts.next().filter(|part| !part.is_empty()); + let repo = parts.next().filter(|part| !part.is_empty()); + if parts.next().is_some() { + anyhow::bail!("GitHub repository URL must identify one repository"); + } + Ok(( + owner + .ok_or_else(|| anyhow::anyhow!("GitHub URL missing owner"))? + .into(), + repo.ok_or_else(|| anyhow::anyhow!("GitHub URL missing repository"))? + .into(), + )) +} + +fn raw_coordinates(item_id: &str) -> Option<(RawKind, String)> { + let (prefix, uid) = item_id.split_once(':')?; + let kind = match prefix { + "commit" => RawKind::Commit, + "issue" => RawKind::Issue, + "pr" => RawKind::PullRequest, + _ => return None, + }; + (!uid.is_empty()).then(|| (kind, uid.into())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_supported_repo_urls() { + assert_eq!( + parse_repo("https://github.com/tinyhumansai/openhuman.git").unwrap(), + ("tinyhumansai".into(), "openhuman".into()) + ); + assert_eq!( + parse_repo("git@github.com:tinyhumansai/openhuman").unwrap(), + ("tinyhumansai".into(), "openhuman".into()) + ); + } + + #[test] + fn maps_raw_item_coordinates() { + assert_eq!( + raw_coordinates("issue:42"), + Some((RawKind::Issue, "42".into())) + ); + assert_eq!(raw_coordinates("unknown:42"), None); + } +} diff --git a/src/memory/sync/mod.rs b/src/memory/sync/mod.rs new file mode 100644 index 0000000..c27f3ad --- /dev/null +++ b/src/memory/sync/mod.rs @@ -0,0 +1,35 @@ +//! Live source synchronization engine. + +pub mod audit; +pub mod composio; +pub mod dispatcher; +pub mod github; +pub mod periodic; +pub mod rebuild; +pub mod state; +pub mod status; +pub mod traits; +pub mod workspace; + +pub use audit::{ + append_audit_entry, estimate_cost_usd, read_audit_log, RealCostAccumulator, SyncAuditEntry, +}; +pub use composio::{ + ClickUpSyncPipeline, ComposioClient, GitHubSyncPipeline, GmailSyncPipeline, LinearSyncPipeline, + NotionSyncPipeline, SlackSearchBackfillPipeline, SlackSyncPipeline, +}; +pub use dispatcher::{SyncDispatcher, SyncRunResult}; +pub use github::GithubRepoSyncPipeline; +pub use periodic::{due_workspace_sources, effective_interval_secs, DEFAULT_SYNC_INTERVAL_SECS}; +pub use rebuild::{ + needs_rebuild, raw_coverage, rebuild_tree_from_raw, rebuild_tree_from_raw_with_audit, + RawCoverage, RawFileRef, RebuildOutcome, +}; +pub use state::{DailyBudget, SyncState, SyncStateStore}; +pub use status::{list_sync_statuses, FreshnessLabel, MemorySyncStatus, StatusListResponse}; +pub use traits::{ + ExternalSourceReader, LocalDocument, LocalDocumentSink, SkillDocSink, SkillDocument, + SyncContext, SyncEvent, SyncEventSink, SyncOutcome, SyncPipeline, SyncPipelineKind, + SyncRunError, SyncStage, +}; +pub use workspace::WorkspaceSourcePipeline; diff --git a/src/memory/sync/periodic.rs b/src/memory/sync/periodic.rs new file mode 100644 index 0000000..23dd899 --- /dev/null +++ b/src/memory/sync/periodic.rs @@ -0,0 +1,150 @@ +//! Host-neutral periodic synchronization cadence policy. + +use std::collections::HashMap; +use std::time::Duration; + +use chrono::{DateTime, Utc}; + +use crate::memory::sources::{MemorySourceEntry, SourceKind}; +use crate::memory::sync::audit::SyncAuditEntry; + +pub const DEFAULT_SYNC_INTERVAL_SECS: u64 = 24 * 60 * 60; + +pub fn effective_interval_secs(configured: Option) -> Option { + match configured { + Some(0) => None, + Some(seconds) => Some(seconds.max(DEFAULT_SYNC_INTERVAL_SECS)), + None => Some(DEFAULT_SYNC_INTERVAL_SECS), + } +} + +pub fn due_workspace_sources( + sources: &[MemorySourceEntry], + audit: &[SyncAuditEntry], + configured_interval_secs: Option, + now: DateTime, +) -> Vec { + let Some(interval) = effective_interval_secs(configured_interval_secs) else { + return Vec::new(); + }; + let successes = last_success_by_source(audit); + sources + .iter() + .filter(|source| source.enabled && is_periodic_workspace_kind(&source.kind)) + .filter(|source| { + successes + .get(&source.id) + .is_none_or(|last| elapsed_since(*last, now) >= Duration::from_secs(interval)) + }) + .cloned() + .collect() +} + +fn is_periodic_workspace_kind(kind: &SourceKind) -> bool { + matches!( + kind, + SourceKind::GithubRepo | SourceKind::Folder | SourceKind::RssFeed | SourceKind::WebPage + ) +} + +fn last_success_by_source(audit: &[SyncAuditEntry]) -> HashMap> { + let mut result = HashMap::new(); + for entry in audit.iter().filter(|entry| entry.success) { + if !matches!( + entry.source_kind.as_str(), + "github_repo" | "folder" | "rss_feed" | "web_page" + ) { + continue; + } + result + .entry(entry.source_id.clone()) + .and_modify(|current: &mut DateTime| *current = (*current).max(entry.timestamp)) + .or_insert(entry.timestamp); + } + result +} + +fn elapsed_since(timestamp: DateTime, now: DateTime) -> Duration { + Duration::from_secs((now - timestamp).num_seconds().max(0) as u64) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn source(id: &str, kind: SourceKind, enabled: bool) -> MemorySourceEntry { + MemorySourceEntry { + id: id.into(), + kind, + label: id.into(), + enabled, + toolkit: None, + connection_id: None, + path: Some("/tmp".into()), + glob: None, + url: Some("https://example.com".into()), + branch: None, + paths: Vec::new(), + max_commits: None, + max_issues: None, + max_prs: None, + query: None, + since_days: None, + max_items: None, + selector: None, + max_tokens_per_sync: None, + max_cost_per_sync_usd: None, + sync_depth_days: None, + } + } + + fn audit(id: &str, timestamp: DateTime, success: bool) -> SyncAuditEntry { + SyncAuditEntry { + timestamp, + source_id: id.into(), + source_kind: "folder".into(), + scope: id.into(), + items_fetched: 0, + batches: 0, + input_tokens: 0, + output_tokens: 0, + estimated_cost_usd: 0.0, + composio_actions_called: 0, + composio_cost_usd: 0.0, + actual_charged_usd: None, + duration_ms: 0, + success, + error: None, + } + } + + #[test] + fn cadence_handles_manual_minimum_and_persisted_success() { + assert_eq!(effective_interval_secs(Some(0)), None); + assert_eq!( + effective_interval_secs(Some(60)), + Some(DEFAULT_SYNC_INTERVAL_SECS) + ); + let now = Utc::now(); + let sources = vec![ + source("new", SourceKind::Folder, true), + source("recent", SourceKind::Folder, true), + source("old", SourceKind::Folder, true), + source("disabled", SourceKind::Folder, false), + source("conversation", SourceKind::Conversation, true), + ]; + let history = vec![ + audit("recent", now - chrono::Duration::hours(1), true), + audit("old", now - chrono::Duration::hours(25), true), + audit("recent", now - chrono::Duration::hours(30), true), + ]; + let due = due_workspace_sources(&sources, &history, None, now); + assert_eq!( + due.iter() + .map(|source| source.id.as_str()) + .collect::>(), + vec!["new", "old"] + ); + assert!(due_workspace_sources(&sources, &history, Some(0), now).is_empty()); + } +} diff --git a/src/memory/sync/rebuild.rs b/src/memory/sync/rebuild.rs new file mode 100644 index 0000000..39e0573 --- /dev/null +++ b/src/memory/sync/rebuild.rs @@ -0,0 +1,488 @@ +//! Raw archive coverage detection and incremental rebuild inputs. + +use std::collections::HashSet; +use std::path::{Path, PathBuf}; + +use crate::memory::chunks::{ + count_raw_paths_ingested_with_prefix, filter_raw_paths_not_ingested, + list_chunk_raw_ref_paths_with_prefix, mark_raw_paths_ingested, +}; +use crate::memory::config::MemoryConfig; +use crate::memory::store::content::{raw_source_dir, sanitize_uid, slugify_source_id}; +use crate::memory::tree::{ + fallback_summary, ingest_summary, Summariser, SummaryContext, SummaryIngestInput, SummaryInput, + TreeKind, +}; +use crate::memory::tree::{store::list_summaries_at_level, TreeFactory}; + +use super::audit::{append_audit_entry, RealCostAccumulator, SyncAuditEntry}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RawFileRef { + pub abs: PathBuf, + pub rel: String, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RawCoverage { + pub total: usize, + pub covered: usize, + pub pending: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct RebuildOutcome { + pub files_read: usize, + pub batches: usize, + pub input_tokens: u64, + pub output_tokens: u64, + pub estimated_cost_usd: f64, + pub actual_charged_usd: Option, +} + +pub fn raw_coverage( + config: &MemoryConfig, + tree_scope: &str, + archive_source_id: &str, +) -> anyhow::Result { + let content_root = crate::memory::chunks::content_root(config); + let source_directory = raw_source_dir(&content_root, archive_source_id); + if !source_directory.exists() { + return Ok(RawCoverage::default()); + } + + let mut files = Vec::new(); + collect_raw_files(&source_directory, &mut files)?; + files.sort(); + let references: Vec<_> = files + .into_iter() + .filter_map(|abs| { + let rel = abs + .strip_prefix(&content_root) + .ok()? + .to_str()? + .replace(std::path::MAIN_SEPARATOR, "/"); + Some(RawFileRef { abs, rel }) + }) + .collect(); + let total = references.len(); + if total == 0 { + return Ok(RawCoverage::default()); + } + + let relative_prefix = format!("raw/{}/", slugify_source_id(archive_source_id)); + if count_raw_paths_ingested_with_prefix(config, &relative_prefix)? == 0 { + backfill_coverage_from_summaries(config, tree_scope, &references)?; + } + + let relative_paths: Vec<_> = references + .iter() + .map(|reference| reference.rel.clone()) + .collect(); + let mut pending_paths: HashSet<_> = filter_raw_paths_not_ingested(config, &relative_paths)? + .into_iter() + .collect(); + let chunk_covered = list_chunk_raw_ref_paths_with_prefix(config, &relative_prefix)?; + pending_paths.retain(|path| !chunk_covered.contains(path)); + let pending: Vec<_> = references + .into_iter() + .filter(|reference| pending_paths.contains(&reference.rel)) + .collect(); + tracing::debug!( + tree_scope, + archive_source_id, + total, + pending = pending.len(), + "[memory_sync:rebuild] raw coverage computed" + ); + Ok(RawCoverage { + total, + covered: total.saturating_sub(pending.len()), + pending, + }) +} + +pub fn needs_rebuild(config: &MemoryConfig, tree_scope: &str, archive_source_id: &str) -> bool { + match raw_coverage(config, tree_scope, archive_source_id) { + Ok(coverage) => !coverage.pending.is_empty(), + Err(error) => { + tracing::warn!(tree_scope, archive_source_id, %error, "[memory_sync:rebuild] coverage check failed"); + false + } + } +} + +pub async fn rebuild_tree_from_raw( + config: &MemoryConfig, + tree_scope: &str, + archive_source_id: &str, + summariser: &dyn Summariser, +) -> anyhow::Result { + rebuild_tree_from_raw_with_audit( + config, + tree_scope, + archive_source_id, + summariser, + &format!("rebuild:{tree_scope}"), + "rebuild", + ) + .await +} + +pub async fn rebuild_tree_from_raw_with_audit( + config: &MemoryConfig, + tree_scope: &str, + archive_source_id: &str, + summariser: &dyn Summariser, + audit_source_id: &str, + audit_source_kind: &str, +) -> anyhow::Result { + let started = std::time::Instant::now(); + let coverage = raw_coverage(config, tree_scope, archive_source_id)?; + if coverage.pending.is_empty() { + return Ok(RebuildOutcome::default()); + } + + let mut inputs = Vec::new(); + for file in &coverage.pending { + let body = match std::fs::read_to_string(&file.abs) { + Ok(body) => body, + Err(error) => { + tracing::warn!(path = %file.abs.display(), %error, "[memory_sync:rebuild] unreadable raw file skipped"); + continue; + } + }; + let stem = file + .abs + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("unknown") + .to_owned(); + let timestamp_ms = stem + .split('_') + .next() + .and_then(|value| value.parse().ok()) + .unwrap_or(0); + let timestamp = + chrono::DateTime::from_timestamp_millis(timestamp_ms).unwrap_or_else(chrono::Utc::now); + inputs.push(RebuildInput { + summary: SummaryInput { + id: stem.clone(), + token_count: crate::memory::chunks::approx_token_count(&body).max(1), + content: body, + entities: Vec::new(), + topics: Vec::new(), + time_range_start: timestamp, + time_range_end: timestamp, + score: 0.5, + }, + label: stem, + relative_path: file.rel.clone(), + }); + } + if inputs.is_empty() { + return Ok(RebuildOutcome { + files_read: coverage.pending.len(), + ..RebuildOutcome::default() + }); + } + + let tree = TreeFactory::source(tree_scope).get_or_create(config)?; + let batches = batch_inputs(inputs, config.tree.input_token_budget); + let batch_count = batches.len(); + let files_read = batches.iter().map(Vec::len).sum(); + let mut cost = RealCostAccumulator::new(); + + for (batch_index, batch) in batches.into_iter().enumerate() { + let summary_inputs: Vec<_> = batch.iter().map(|item| item.summary.clone()).collect(); + let context = SummaryContext { + tree_id: &tree.id, + tree_kind: TreeKind::Source, + target_level: 1, + token_budget: config.tree.output_token_budget, + }; + let call = match summariser + .summarise_with_usage(&summary_inputs, &context) + .await + { + Ok(call) if !call.output.content.trim().is_empty() => call, + Ok(_) => { + tracing::warn!( + batch_index, + "[memory_sync:rebuild] blank summary; using fallback" + ); + crate::memory::tree::SummaryCall { + output: fallback_summary(&summary_inputs, context.token_budget), + ..Default::default() + } + } + Err(error) => { + tracing::warn!(batch_index, %error, "[memory_sync:rebuild] summariser failed; using fallback"); + crate::memory::tree::SummaryCall { + output: fallback_summary(&summary_inputs, context.token_budget), + ..Default::default() + } + } + }; + let output = call.output; + let batch_input_tokens: u64 = summary_inputs + .iter() + .map(|input| input.token_count as u64) + .sum(); + cost.add_batch( + batch_input_tokens, + output.token_count as u64, + call.input_tokens, + call.output_tokens, + call.charged_amount_usd, + ); + let time_range_start = summary_inputs + .iter() + .map(|input| input.time_range_start) + .min() + .unwrap_or_else(chrono::Utc::now); + let time_range_end = summary_inputs + .iter() + .map(|input| input.time_range_end) + .max() + .unwrap_or_else(chrono::Utc::now); + ingest_summary( + config, + &tree, + SummaryIngestInput { + content: output.content, + token_count: output.token_count, + entities: output.entities, + topics: output.topics, + time_range_start, + time_range_end, + score: 0.5, + child_labels: batch.iter().map(|input| input.label.clone()).collect(), + child_basenames: Vec::new(), + }, + summariser, + ) + .await?; + let covered: Vec<_> = batch + .iter() + .map(|input| input.relative_path.clone()) + .collect(); + mark_raw_paths_ingested(config, &covered)?; + tracing::info!( + tree_scope, + batch_index, + files = covered.len(), + "[memory_sync:rebuild] batch ingested and covered" + ); + } + + let input_tokens = cost.audit_input_tokens(); + let output_tokens = cost.audit_output_tokens(); + let estimated_cost_usd = cost.estimated_cost(); + let actual_charged_usd = cost.actual_charged_usd(); + append_audit_entry( + config, + &SyncAuditEntry { + timestamp: chrono::Utc::now(), + source_id: audit_source_id.into(), + source_kind: audit_source_kind.into(), + scope: tree_scope.into(), + items_fetched: files_read as u32, + batches: batch_count as u32, + input_tokens, + output_tokens, + estimated_cost_usd, + composio_actions_called: 0, + composio_cost_usd: 0.0, + actual_charged_usd, + duration_ms: started.elapsed().as_millis() as u64, + success: true, + error: None, + }, + )?; + Ok(RebuildOutcome { + files_read, + batches: batch_count, + input_tokens, + output_tokens, + estimated_cost_usd, + actual_charged_usd, + }) +} + +#[derive(Clone)] +struct RebuildInput { + summary: SummaryInput, + label: String, + relative_path: String, +} + +fn batch_inputs(inputs: Vec, token_budget: u32) -> Vec> { + let budget = token_budget.max(1) as u64; + let mut batches = Vec::new(); + let mut current = Vec::new(); + let mut current_tokens = 0u64; + for input in inputs { + let tokens = input.summary.token_count as u64; + if !current.is_empty() && current_tokens.saturating_add(tokens) > budget { + batches.push(std::mem::take(&mut current)); + current_tokens = 0; + } + current_tokens = current_tokens.saturating_add(tokens); + current.push(input); + } + if !current.is_empty() { + batches.push(current); + } + batches +} + +fn collect_raw_files(directory: &Path, output: &mut Vec) -> anyhow::Result<()> { + for entry in std::fs::read_dir(directory)? { + let path = entry?.path(); + if path.is_dir() { + collect_raw_files(&path, output)?; + } else if path.extension().and_then(|extension| extension.to_str()) == Some("md") + && path.file_name().and_then(|name| name.to_str()) != Some("_source.md") + { + output.push(path); + } + } + Ok(()) +} + +fn backfill_coverage_from_summaries( + config: &MemoryConfig, + tree_scope: &str, + files: &[RawFileRef], +) -> anyhow::Result { + let tree = TreeFactory::source(tree_scope).get_or_create(config)?; + if tree.max_level == 0 { + return Ok(0); + } + let summaries = list_summaries_at_level(config, &tree.id, 1)?; + let mut kind_uids: HashSet<(&'static str, String)> = HashSet::new(); + let mut stems = HashSet::new(); + for summary in summaries { + for label in summary.child_ids { + if let Some(uid) = label.strip_prefix("commit:") { + kind_uids.insert(("commits", sanitize_uid(uid))); + } else if let Some(uid) = label.strip_prefix("issue:") { + kind_uids.insert(("issues", sanitize_uid(uid))); + } else if let Some(uid) = label.strip_prefix("pr:") { + kind_uids.insert(("prs", sanitize_uid(uid))); + } else { + stems.insert(label); + } + } + } + let covered: Vec<_> = files + .iter() + .filter(|reference| { + let path = Path::new(&reference.rel); + let kind = path + .parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()); + let stem = path + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or_default(); + stems.contains(stem) + || kind_uids.iter().any(|(expected_kind, uid)| { + kind == Some(*expected_kind) + && stem + .split_once('_') + .is_some_and(|(_, file_uid)| file_uid == uid) + }) + }) + .map(|reference| reference.rel.clone()) + .collect(); + if covered.is_empty() { + Ok(0) + } else { + mark_raw_paths_ingested(config, &covered) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn coverage_is_incremental_and_ignores_source_metadata() { + let temp = tempfile::tempdir().unwrap(); + let mut config = MemoryConfig::new(temp.path()); + let custom_root = temp.path().join("custom-content"); + config.content_root = Some(custom_root.clone()); + let root = custom_root.join("raw/github-com-org-repo/issues"); + std::fs::create_dir_all(&root).unwrap(); + std::fs::write(root.join("100_one.md"), "one").unwrap(); + std::fs::write(root.join("200_two.md"), "two").unwrap(); + std::fs::write(root.parent().unwrap().join("_source.md"), "metadata").unwrap(); + + let first = raw_coverage(&config, "github:org/repo", "github.com/org/repo").unwrap(); + assert_eq!(first.total, 2); + assert_eq!(first.pending.len(), 2); + mark_raw_paths_ingested(&config, &[first.pending[0].rel.clone()]).unwrap(); + let second = raw_coverage(&config, "github:org/repo", "github.com/org/repo").unwrap(); + assert_eq!(second.covered, 1); + assert_eq!(second.pending.len(), 1); + assert!(needs_rebuild( + &config, + "github:org/repo", + "github.com/org/repo" + )); + } + + #[tokio::test] + async fn rebuild_ingests_l1_summaries_marks_coverage_and_is_idempotent() { + let temp = tempfile::tempdir().unwrap(); + let mut config = MemoryConfig::new(temp.path()); + config.tree.input_token_budget = 2; + let root = config + .workspace + .join("memory_tree/content/raw/github-com-org-repo/issues"); + std::fs::create_dir_all(&root).unwrap(); + std::fs::write(root.join("100_one.md"), "first body").unwrap(); + std::fs::write(root.join("200_two.md"), "second body").unwrap(); + + let first = rebuild_tree_from_raw( + &config, + "github:org/repo", + "github.com/org/repo", + &crate::memory::tree::ConcatSummariser, + ) + .await + .unwrap(); + assert_eq!(first.files_read, 2); + assert_eq!(first.batches, 2); + assert!(!needs_rebuild( + &config, + "github:org/repo", + "github.com/org/repo" + )); + assert_eq!(crate::memory::chunks::count_chunks(&config).unwrap(), 0); + let tree = TreeFactory::source("github:org/repo") + .get_or_create(&config) + .unwrap(); + assert_eq!( + list_summaries_at_level(&config, &tree.id, 1).unwrap().len(), + 2 + ); + + let second = rebuild_tree_from_raw( + &config, + "github:org/repo", + "github.com/org/repo", + &crate::memory::tree::ConcatSummariser, + ) + .await + .unwrap(); + assert_eq!(second, RebuildOutcome::default()); + assert_eq!( + list_summaries_at_level(&config, &tree.id, 1).unwrap().len(), + 2 + ); + } +} diff --git a/src/memory/sync/state.rs b/src/memory/sync/state.rs new file mode 100644 index 0000000..9a11227 --- /dev/null +++ b/src/memory/sync/state.rs @@ -0,0 +1,250 @@ +//! Persistence-neutral cursor, deduplication, and daily-budget state. + +use std::collections::{HashMap, HashSet}; + +use async_trait::async_trait; +use chrono::Utc; +use serde::{Deserialize, Serialize}; + +pub const DEFAULT_DAILY_REQUEST_LIMIT: u32 = 500; +pub const STATE_NAMESPACE: &str = "composio-sync-state"; + +#[async_trait] +pub trait SyncStateStore: Send + Sync { + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result>; + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()>; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DailyBudget { + pub date: String, + pub requests_used: u32, + pub limit: u32, +} + +impl Default for DailyBudget { + fn default() -> Self { + Self { + date: today(), + requests_used: 0, + limit: DEFAULT_DAILY_REQUEST_LIMIT, + } + } +} + +impl DailyBudget { + pub fn remaining(&self) -> u32 { + if self.date != today() { + self.limit + } else { + self.limit.saturating_sub(self.requests_used) + } + } + + pub fn is_exhausted(&self) -> bool { + self.remaining() == 0 + } + + pub fn record_requests(&mut self, count: u32) { + let today = today(); + if self.date != today { + self.date = today; + self.requests_used = 0; + } + self.requests_used = self.requests_used.saturating_add(count); + } + + pub fn record_request(&mut self) { + self.record_requests(1); + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SyncState { + pub toolkit: String, + pub connection_id: String, + #[serde(default)] + pub cursor: Option, + #[serde(default)] + pub synced_ids: HashSet, + #[serde(default)] + pub item_versions: HashMap, + #[serde(default)] + pub daily_budget: DailyBudget, + #[serde(default)] + pub last_seen_id: Option, + #[serde(default)] + pub last_sync_at_ms: Option, + #[serde(skip)] + pub run_requests: u32, + #[serde(skip)] + pub run_provider_cost_usd: f64, +} + +impl SyncState { + pub fn new(toolkit: impl Into, connection_id: impl Into) -> Self { + Self { + toolkit: toolkit.into(), + connection_id: connection_id.into(), + cursor: None, + synced_ids: HashSet::new(), + item_versions: HashMap::new(), + daily_budget: DailyBudget::default(), + last_seen_id: None, + last_sync_at_ms: None, + run_requests: 0, + run_provider_cost_usd: 0.0, + } + } + + pub fn key(toolkit: &str, connection_id: &str) -> String { + format!("{toolkit}:{connection_id}") + } + + pub fn is_synced(&self, id: &str) -> bool { + self.synced_ids.contains(id) + } + + pub fn mark_synced(&mut self, id: impl Into) { + self.synced_ids.insert(id.into()); + } + + pub fn advance_cursor(&mut self, cursor: impl Into) { + self.cursor = Some(cursor.into()); + } + + pub fn set_last_seen_id(&mut self, id: impl Into) { + self.last_seen_id = Some(id.into()); + } + + pub fn set_last_sync_at_ms(&mut self, timestamp_ms: u64) { + self.last_sync_at_ms = Some(timestamp_ms); + } + + pub fn budget_exhausted(&self) -> bool { + self.daily_budget.is_exhausted() + } + + pub fn budget_remaining(&self) -> u32 { + self.daily_budget.remaining() + } + + pub fn record_requests(&mut self, count: u32) { + self.daily_budget.record_requests(count); + self.run_requests = self.run_requests.saturating_add(count); + } + + pub fn record_action(&mut self, attempts: u32, cost_usd: f64) { + self.record_requests(attempts.max(1)); + if cost_usd.is_finite() && cost_usd > 0.0 { + self.run_provider_cost_usd += cost_usd; + } + } + + pub async fn load( + store: &dyn SyncStateStore, + toolkit: &str, + connection_id: &str, + ) -> anyhow::Result { + let key = Self::key(toolkit, connection_id); + match store.get(STATE_NAMESPACE, &key).await? { + Some(value) => { + let mut state: Self = serde_json::from_value(value)?; + if state.daily_budget.date != today() { + state.daily_budget.date = today(); + state.daily_budget.requests_used = 0; + } + Ok(state) + } + None => Ok(Self::new(toolkit, connection_id)), + } + } + + pub async fn save(&self, store: &dyn SyncStateStore) -> anyhow::Result<()> { + let value = serde_json::to_value(self)?; + store + .set( + STATE_NAMESPACE, + &Self::key(&self.toolkit, &self.connection_id), + &value, + ) + .await + } +} + +fn today() -> String { + Utc::now().format("%Y-%m-%d").to_string() +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Mutex; + + use super::*; + + #[derive(Default)] + struct MemoryStateStore(Mutex>); + + #[async_trait] + impl SyncStateStore for MemoryStateStore { + async fn get( + &self, + namespace: &str, + key: &str, + ) -> anyhow::Result> { + Ok(self + .0 + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.0 + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } + } + + #[tokio::test] + async fn state_round_trips_cursor_dedup_and_budget() { + let store = MemoryStateStore::default(); + let mut state = SyncState::new("gmail", "conn-1"); + state.advance_cursor("cursor-2"); + state.mark_synced("message-1"); + state.record_requests(3); + state.save(&store).await.unwrap(); + + let loaded = SyncState::load(&store, "gmail", "conn-1").await.unwrap(); + assert_eq!(loaded.cursor.as_deref(), Some("cursor-2")); + assert!(loaded.is_synced("message-1")); + assert_eq!(loaded.daily_budget.requests_used, 3); + } + + #[test] + fn stale_budget_reports_full_and_resets_on_record() { + let mut budget = DailyBudget { + date: "2000-01-01".into(), + requests_used: 499, + limit: 500, + }; + assert_eq!(budget.remaining(), 500); + budget.record_requests(1); + assert_eq!(budget.requests_used, 1); + assert_eq!(budget.remaining(), 499); + } +} diff --git a/src/memory/sync/status.rs b/src/memory/sync/status.rs new file mode 100644 index 0000000..11a9e99 --- /dev/null +++ b/src/memory/sync/status.rs @@ -0,0 +1,175 @@ +//! Pull-based synchronization status derived from the authoritative chunk store. + +use serde::{Deserialize, Serialize}; + +use crate::memory::chunks::with_connection; +use crate::memory::config::MemoryConfig; + +const WAVE_WINDOW_MS: i64 = 10 * 60 * 1000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FreshnessLabel { + Active, + Recent, + Idle, +} + +impl FreshnessLabel { + pub fn from_age_ms(last_chunk_at_ms: Option, now_ms: i64) -> Self { + match last_chunk_at_ms { + None => Self::Idle, + Some(timestamp) => match now_ms.saturating_sub(timestamp) { + age if age <= 30_000 => Self::Active, + age if age <= 5 * 60_000 => Self::Recent, + _ => Self::Idle, + }, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct MemorySyncStatus { + pub provider: String, + pub chunks_synced: u64, + pub chunks_pending: u64, + pub batch_total: u64, + pub batch_processed: u64, + pub last_chunk_at_ms: Option, + pub freshness: FreshnessLabel, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct StatusListResponse { + pub statuses: Vec, +} + +pub fn list_sync_statuses(config: &MemoryConfig) -> anyhow::Result> { + list_sync_statuses_at(config, chrono::Utc::now().timestamp_millis()) +} + +fn list_sync_statuses_at( + config: &MemoryConfig, + now_ms: i64, +) -> anyhow::Result> { + with_connection(config, |connection| { + let mut statement = connection.prepare( + "WITH provider_chunks AS ( \ + SELECT CASE WHEN INSTR(source_id, ':') > 0 \ + THEN SUBSTR(source_id, 1, INSTR(source_id, ':') - 1) \ + ELSE source_kind END AS provider, \ + created_at_ms, \ + CASE WHEN EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e WHERE e.chunk_id = c.id) \ + OR c.lifecycle_status = 'dropped' \ + OR EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped s WHERE s.chunk_id = c.id) \ + THEN 1 ELSE 0 END AS resolved, timestamp_ms \ + FROM mem_tree_chunks c \ + ), provider_max AS ( \ + SELECT provider, MAX(created_at_ms) AS max_created FROM provider_chunks GROUP BY provider \ + ), provider_pending AS ( \ + SELECT p.provider, SUM(CASE WHEN p.resolved = 0 AND p.created_at_ms >= m.max_created - ?1 THEN 1 ELSE 0 END) AS pending \ + FROM provider_chunks p JOIN provider_max m ON p.provider = m.provider GROUP BY p.provider \ + ), wave_anchors AS ( \ + SELECT p.provider, MIN(p.created_at_ms) AS anchor \ + FROM provider_chunks p JOIN provider_max m ON p.provider = m.provider \ + JOIN provider_pending pp ON p.provider = pp.provider \ + WHERE pp.pending > 0 AND p.created_at_ms >= m.max_created - ?1 GROUP BY p.provider \ + ) SELECT p.provider, COUNT(*) AS chunks_synced, \ + SUM(CASE WHEN p.resolved = 0 THEN 1 ELSE 0 END) AS chunks_pending, \ + SUM(CASE WHEN w.anchor IS NOT NULL AND p.created_at_ms >= w.anchor THEN 1 ELSE 0 END) AS batch_total, \ + SUM(CASE WHEN w.anchor IS NOT NULL AND p.created_at_ms >= w.anchor AND p.resolved = 1 THEN 1 ELSE 0 END) AS batch_processed, \ + MAX(p.timestamp_ms) AS last_chunk_at_ms \ + FROM provider_chunks p LEFT JOIN wave_anchors w ON p.provider = w.provider \ + GROUP BY p.provider ORDER BY last_chunk_at_ms DESC", + )?; + let rows = statement.query_map([WAVE_WINDOW_MS], |row| { + let last_chunk_at_ms = row.get(5)?; + Ok(MemorySyncStatus { + provider: row.get(0)?, + chunks_synced: nonnegative(row.get(1)?), + chunks_pending: nonnegative(row.get(2)?), + batch_total: nonnegative(row.get(3)?), + batch_processed: nonnegative(row.get(4)?), + last_chunk_at_ms, + freshness: FreshnessLabel::from_age_ms(last_chunk_at_ms, now_ms), + }) + })?; + Ok(rows.collect::>>()?) + }) +} + +fn nonnegative(value: i64) -> u64 { + value.max(0) as u64 +} + +#[cfg(test)] +mod tests { + use chrono::{TimeZone, Utc}; + + use super::*; + use crate::memory::chunks::{upsert_chunks, Chunk, Metadata, SourceKind}; + + fn chunk(id: &str, source_id: &str, created_ms: i64) -> Chunk { + let timestamp = Utc.timestamp_millis_opt(created_ms).unwrap(); + Chunk { + id: id.into(), + content: "content".into(), + token_count: 1, + seq_in_source: 0, + created_at: timestamp, + partial_message: false, + metadata: Metadata { + source_kind: SourceKind::Document, + source_id: source_id.into(), + path_scope: None, + source_ref: None, + owner: "test".into(), + timestamp, + time_range: (timestamp, timestamp), + tags: Vec::new(), + }, + } + } + + #[test] + fn status_groups_provider_and_tracks_active_wave_resolution() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + let now = 1_777_000_000_000i64; + upsert_chunks( + &config, + &[ + chunk("a", "gmail:conn", now - 2_000), + chunk("b", "gmail:conn", now - 1_000), + chunk("c", "slack:conn", now - 600_000), + ], + ) + .unwrap(); + with_connection(&config, |connection| { + connection.execute( + "INSERT INTO mem_tree_chunk_embeddings (chunk_id, model_signature, vector, dim, created_at) VALUES ('a', 'test', X'00', 1, 0)", + [], + )?; + connection.execute("UPDATE mem_tree_chunks SET lifecycle_status = 'dropped' WHERE id = 'c'", [])?; + Ok(()) + }).unwrap(); + + let statuses = list_sync_statuses_at(&config, now).unwrap(); + let gmail = statuses + .iter() + .find(|status| status.provider == "gmail") + .unwrap(); + assert_eq!(gmail.chunks_synced, 2); + assert_eq!(gmail.chunks_pending, 1); + assert_eq!(gmail.batch_total, 2); + assert_eq!(gmail.batch_processed, 1); + assert_eq!(gmail.freshness, FreshnessLabel::Active); + let slack = statuses + .iter() + .find(|status| status.provider == "slack") + .unwrap(); + assert_eq!(slack.chunks_pending, 0); + assert_eq!(slack.batch_total, 0); + assert_eq!(slack.freshness, FreshnessLabel::Idle); + } +} diff --git a/src/memory/sync/traits.rs b/src/memory/sync/traits.rs new file mode 100644 index 0000000..dfc6996 --- /dev/null +++ b/src/memory/sync/traits.rs @@ -0,0 +1,155 @@ +//! Host seams and pipeline contracts for live synchronization. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::memory::config::MemoryConfig; +use crate::memory::sources::{MemorySourceEntry, SourceContent, SourceItem}; +use crate::memory::tree::Summariser; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SyncPipelineKind { + Composio, + Workspace, + Mcp, +} + +impl SyncPipelineKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Composio => "composio", + Self::Workspace => "workspace", + Self::Mcp => "mcp", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SyncStage { + Requested, + Fetching, + Stored, + Ingesting, + Completed, + Failed, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SyncEvent { + pub source_id: String, + pub toolkit: String, + pub connection_id: Option, + pub stage: SyncStage, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub message: Option, +} + +#[async_trait] +pub trait SyncEventSink: Send + Sync { + async fn emit(&self, event: SyncEvent) -> anyhow::Result<()>; +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SkillDocument { + pub namespace_skill_id: String, + pub connection_id: String, + pub document_id: String, + pub title: String, + pub content: String, + pub toolkit: String, + #[serde(default)] + pub metadata: serde_json::Value, +} + +#[async_trait] +pub trait SkillDocSink: Send + Sync { + async fn store(&self, document: SkillDocument) -> anyhow::Result<()>; + async fn delete(&self, namespace_skill_id: &str, document_id: &str) -> anyhow::Result<()>; +} + +#[derive(Clone, Debug)] +pub struct LocalDocument { + pub source_id: String, + pub path_scope: Option, + pub owner: String, + pub tags: Vec, + pub title: String, + pub body: String, + pub modified_at: chrono::DateTime, + pub source_ref: Option, +} + +#[async_trait] +pub trait LocalDocumentSink: Send + Sync { + async fn upsert(&self, document: LocalDocument) -> anyhow::Result<()>; + async fn delete(&self, source_id: &str) -> anyhow::Result<()>; +} + +/// Host adapter for source kinds whose fetch transport is product-owned. +/// Tinycortex still owns lifecycle, deduplication, persistence, and events. +#[async_trait] +pub trait ExternalSourceReader: Send + Sync { + async fn list_items(&self, source: &MemorySourceEntry) -> anyhow::Result>; + async fn read_item( + &self, + source: &MemorySourceEntry, + item_id: &str, + ) -> anyhow::Result; +} + +/// Host capabilities required by sync pipelines. +#[derive(Clone)] +pub struct SyncContext { + pub events: Arc, + pub documents: Arc, + pub state: Arc, + pub local_documents: Option>, + pub external_sources: Option>, + pub summariser: Option>, +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct SyncOutcome { + pub records_ingested: u32, + pub more_pending: bool, + #[serde(default)] + pub actions_called: u32, + #[serde(default)] + pub provider_cost_usd: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub note: Option, +} + +#[derive(Debug, thiserror::Error)] +#[error("{message}")] +pub struct SyncRunError { + pub actions_called: u32, + pub provider_cost_usd: f64, + message: String, +} + +impl SyncRunError { + pub fn new(message: impl Into, actions_called: u32, provider_cost_usd: f64) -> Self { + Self { + actions_called, + provider_cost_usd, + message: message.into(), + } + } +} + +#[async_trait] +pub trait SyncPipeline: Send + Sync { + fn id(&self) -> &str; + fn kind(&self) -> SyncPipelineKind; + async fn init(&self, config: &MemoryConfig, context: &SyncContext) -> anyhow::Result<()>; + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result; +} diff --git a/src/memory/sync/workspace.rs b/src/memory/sync/workspace.rs new file mode 100644 index 0000000..5010435 --- /dev/null +++ b/src/memory/sync/workspace.rs @@ -0,0 +1,479 @@ +//! Local workspace source synchronization through crate-owned readers. + +use std::collections::{HashMap, HashSet}; + +use async_trait::async_trait; + +use crate::memory::config::MemoryConfig; +use crate::memory::sources::{reader_for, MemorySourceEntry, SourceReader}; +use crate::memory::sync::state::SyncState; +use crate::memory::sync::traits::{ + LocalDocument, SyncContext, SyncEvent, SyncOutcome, SyncPipeline, SyncPipelineKind, SyncStage, +}; + +pub struct WorkspaceSourcePipeline { + id: String, + source: MemorySourceEntry, + reader: Option>, +} + +impl WorkspaceSourcePipeline { + pub fn new(source: MemorySourceEntry) -> anyhow::Result { + source.validate().map_err(anyhow::Error::msg)?; + let reader = reader_for(&source.kind); + Ok(Self { + id: format!("workspace:{}:{}", source.kind.as_str(), source.id), + source, + reader, + }) + } + + fn source_id(&self, item_id: &str) -> String { + format!("mem_src:{}:{item_id}", self.source.id) + } + + async fn event(&self, context: &SyncContext, stage: SyncStage, message: Option) { + let _ = context + .events + .emit(SyncEvent { + source_id: self.id.clone(), + toolkit: self.source.kind.as_str().into(), + connection_id: Some(self.source.id.clone()), + stage, + message, + }) + .await; + } +} + +#[async_trait] +impl SyncPipeline for WorkspaceSourcePipeline { + fn id(&self) -> &str { + &self.id + } + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Workspace + } + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + Ok(()) + } + + async fn tick( + &self, + config: &MemoryConfig, + context: &SyncContext, + ) -> anyhow::Result { + if !self.source.enabled { + return Ok(SyncOutcome { + note: Some("source disabled".into()), + ..SyncOutcome::default() + }); + } + self.event(context, SyncStage::Fetching, None).await; + let state_toolkit = format!("workspace:{}", self.source.kind.as_str()); + let mut state = + SyncState::load(context.state.as_ref(), &state_toolkit, &self.source.id).await?; + let local_documents = context + .local_documents + .as_ref() + .ok_or_else(|| anyhow::anyhow!("workspace pipeline requires a local document sink"))?; + let items = match &self.reader { + Some(reader) => reader + .list_items(&self.source, config) + .await + .map_err(anyhow::Error::msg)?, + None => { + context + .external_sources + .as_ref() + .ok_or_else(|| { + anyhow::anyhow!( + "source kind requires an external reader: {}", + self.source.kind.as_str() + ) + })? + .list_items(&self.source) + .await? + } + }; + let current_ids: HashSet<_> = items.iter().map(|item| item.id.clone()).collect(); + let mut versions = HashMap::with_capacity(items.len()); + let mut ingested = 0u32; + + for item in items { + let version = item + .updated_at_ms + .map(|value| value.to_string()) + .unwrap_or_else(|| "unknown".into()); + versions.insert(item.id.clone(), version.clone()); + if item.updated_at_ms.is_some() && state.item_versions.get(&item.id) == Some(&version) { + continue; + } + let source_id = self.source_id(&item.id); + let content = match &self.reader { + Some(reader) => reader + .read_item(&self.source, &item.id, config) + .await + .map_err(anyhow::Error::msg)?, + None => { + context + .external_sources + .as_ref() + .expect("external reader checked before item loop") + .read_item(&self.source, &item.id) + .await? + } + }; + local_documents + .upsert(LocalDocument { + source_id, + path_scope: None, + owner: "user".into(), + tags: vec!["memory_sources".into(), self.source.kind.as_str().into()], + title: content.title, + body: content.body, + modified_at: item + .updated_at_ms + .and_then(chrono::DateTime::from_timestamp_millis) + .unwrap_or_else(chrono::Utc::now), + source_ref: Some(format!("{}:{}", self.source.id, item.id)), + }) + .await?; + ingested = ingested.saturating_add(1); + } + + let removed: Vec<_> = if self.source.kind == crate::memory::sources::SourceKind::Folder { + state + .item_versions + .keys() + .filter(|id| !current_ids.contains(*id)) + .cloned() + .collect() + } else { + Vec::new() + }; + for item_id in &removed { + local_documents.delete(&self.source_id(item_id)).await?; + } + state.item_versions = versions; + state.last_sync_at_ms = Some(chrono::Utc::now().timestamp_millis() as u64); + state.save(context.state.as_ref()).await?; + self.event( + context, + SyncStage::Stored, + Some(format!("{ingested} stored, {} removed", removed.len())), + ) + .await; + self.event(context, SyncStage::Completed, None).await; + Ok(SyncOutcome { + records_ingested: ingested, + more_pending: false, + actions_called: 0, + provider_cost_usd: 0.0, + note: (!removed.is_empty()).then(|| format!("{} removed", removed.len())), + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use super::*; + use crate::memory::sources::SourceKind; + use crate::memory::sync::state::SyncStateStore; + use crate::memory::sync::traits::{ + ExternalSourceReader, LocalDocumentSink, SkillDocSink, SkillDocument, SyncEventSink, + }; + + #[derive(Default)] + struct Host { + documents: Mutex>, + state: Mutex>, + events: Mutex>, + deletes: Mutex>, + external_items: Mutex>, + external_bodies: Mutex>, + } + + #[async_trait] + impl SkillDocSink for Host { + async fn store(&self, _: SkillDocument) -> anyhow::Result<()> { + Ok(()) + } + async fn delete(&self, _: &str, _: &str) -> anyhow::Result<()> { + Ok(()) + } + } + + #[async_trait] + impl LocalDocumentSink for Host { + async fn upsert(&self, document: LocalDocument) -> anyhow::Result<()> { + self.documents + .lock() + .unwrap() + .insert(document.source_id.clone(), document); + Ok(()) + } + + async fn delete(&self, source_id: &str) -> anyhow::Result<()> { + self.documents.lock().unwrap().remove(source_id); + self.deletes.lock().unwrap().push(source_id.into()); + Ok(()) + } + } + + #[async_trait] + impl SyncEventSink for Host { + async fn emit(&self, event: SyncEvent) -> anyhow::Result<()> { + self.events.lock().unwrap().push(event); + Ok(()) + } + } + + #[async_trait] + impl SyncStateStore for Host { + async fn get( + &self, + namespace: &str, + key: &str, + ) -> anyhow::Result> { + Ok(self + .state + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.state + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } + } + + #[async_trait] + impl ExternalSourceReader for Host { + async fn list_items( + &self, + _: &MemorySourceEntry, + ) -> anyhow::Result> { + Ok(self.external_items.lock().unwrap().clone()) + } + + async fn read_item( + &self, + _: &MemorySourceEntry, + item_id: &str, + ) -> anyhow::Result { + self.external_bodies + .lock() + .unwrap() + .get(item_id) + .cloned() + .ok_or_else(|| anyhow::anyhow!("missing external item {item_id}")) + } + } + + fn folder_source(path: &std::path::Path) -> MemorySourceEntry { + MemorySourceEntry { + id: "folder-1".into(), + kind: SourceKind::Folder, + label: "Notes".into(), + enabled: true, + toolkit: None, + connection_id: None, + path: Some(path.to_string_lossy().into_owned()), + glob: Some("**/*.md".into()), + url: None, + branch: None, + paths: Vec::new(), + max_commits: None, + max_issues: None, + max_prs: None, + query: None, + since_days: None, + max_items: None, + selector: None, + max_tokens_per_sync: None, + max_cost_per_sync_usd: None, + sync_depth_days: None, + } + } + + #[tokio::test] + async fn folder_pipeline_tracks_create_update_noop_and_remove() { + let temp = tempfile::tempdir().unwrap(); + let notes = temp.path().join("notes"); + std::fs::create_dir_all(¬es).unwrap(); + let file = notes.join("daily.md"); + std::fs::write(&file, "first").unwrap(); + let config = MemoryConfig::new(temp.path().join("workspace")); + let pipeline = WorkspaceSourcePipeline::new(folder_source(¬es)).unwrap(); + let host = Arc::new(Host::default()); + let context = SyncContext { + events: host.clone(), + documents: host.clone(), + state: host.clone(), + local_documents: Some(host.clone()), + external_sources: None, + summariser: None, + }; + + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 1 + ); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 0 + ); + std::thread::sleep(std::time::Duration::from_millis(5)); + std::fs::write(&file, "second").unwrap(); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 1 + ); + assert_eq!( + host.documents.lock().unwrap()["mem_src:folder-1:daily.md"].body, + "second" + ); + std::fs::remove_file(&file).unwrap(); + let removed = pipeline.tick(&config, &context).await.unwrap(); + assert_eq!(removed.records_ingested, 0); + assert_eq!(removed.note.as_deref(), Some("1 removed")); + assert!(host.documents.lock().unwrap().is_empty()); + assert_eq!( + host.deletes.lock().unwrap().as_slice(), + ["mem_src:folder-1:daily.md"] + ); + } + + #[tokio::test] + async fn external_pipeline_tracks_versions_without_destructive_absence() { + use crate::memory::sources::{ContentType, SourceContent, SourceItem}; + + let config = MemoryConfig::new(tempfile::tempdir().unwrap().path().join("workspace")); + let mut source = folder_source(std::path::Path::new("unused")); + source.id = "rss-1".into(); + source.kind = SourceKind::RssFeed; + source.path = None; + source.glob = None; + source.url = Some("https://example.test/feed.xml".into()); + let pipeline = WorkspaceSourcePipeline::new(source).unwrap(); + let host = Arc::new(Host::default()); + host.external_items.lock().unwrap().push(SourceItem { + id: "post-1".into(), + title: "Post".into(), + updated_at_ms: Some(1), + }); + host.external_bodies.lock().unwrap().insert( + "post-1".into(), + SourceContent { + id: "post-1".into(), + title: "Post".into(), + body: "first".into(), + content_type: ContentType::Plaintext, + metadata: serde_json::Value::Null, + }, + ); + let context = SyncContext { + events: host.clone(), + documents: host.clone(), + state: host.clone(), + local_documents: Some(host.clone()), + external_sources: Some(host.clone()), + summariser: None, + }; + + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 1 + ); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 0 + ); + host.external_items.lock().unwrap()[0].updated_at_ms = Some(2); + host.external_bodies.lock().unwrap().remove("post-1"); + assert!(pipeline.tick(&config, &context).await.is_err()); + assert_eq!( + host.documents.lock().unwrap()["mem_src:rss-1:post-1"].body, + "first" + ); + host.external_bodies.lock().unwrap().insert( + "post-1".into(), + SourceContent { + id: "post-1".into(), + title: "Post".into(), + body: "second".into(), + content_type: ContentType::Plaintext, + metadata: serde_json::Value::Null, + }, + ); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 1 + ); + host.external_items.lock().unwrap()[0].updated_at_ms = None; + host.external_bodies + .lock() + .unwrap() + .get_mut("post-1") + .unwrap() + .body = "timestamp-less refresh".into(); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 1 + ); + host.external_items.lock().unwrap().clear(); + assert_eq!( + pipeline + .tick(&config, &context) + .await + .unwrap() + .records_ingested, + 0 + ); + assert!(host + .documents + .lock() + .unwrap() + .contains_key("mem_src:rss-1:post-1")); + } +} diff --git a/src/memory/tool_memory/store.rs b/src/memory/tool_memory/store.rs index d711600..d7b65fe 100644 --- a/src/memory/tool_memory/store.rs +++ b/src/memory/tool_memory/store.rs @@ -147,12 +147,7 @@ impl ToolMemoryStore { let mut rules: Vec = entries .into_iter() .filter(|entry| entry.key.starts_with("rule/")) - .filter_map( - |entry| match serde_json::from_str::(&entry.content) { - Ok(rule) => Some(rule), - Err(_) => None, - }, - ) + .filter_map(|entry| serde_json::from_str::(&entry.content).ok()) .collect(); rules.sort_by(|a, b| { diff --git a/src/memory/tool_memory/test_helpers.rs b/src/memory/tool_memory/test_helpers.rs index 8db6a98..1564cd1 100644 --- a/src/memory/tool_memory/test_helpers.rs +++ b/src/memory/tool_memory/test_helpers.rs @@ -84,7 +84,7 @@ impl Memory for MockMemory { .filter(|((n, _), _)| n == ns) .map(|(_, v)| v.clone()) .collect(), - None => lock.iter().map(|(_, v)| v.clone()).collect(), + None => lock.values().cloned().collect(), }) } @@ -98,7 +98,7 @@ impl Memory for MockMemory { async fn namespace_summaries(&self) -> anyhow::Result> { let mut counts: HashMap = HashMap::new(); - for ((ns, _), _) in self.entries.lock().iter() { + for (ns, _) in self.entries.lock().keys() { *counts.entry(ns.clone()).or_default() += 1; } Ok(counts diff --git a/src/memory/tool_memory/types.rs b/src/memory/tool_memory/types.rs index 9e4dcd5..0079e2d 100644 --- a/src/memory/tool_memory/types.rs +++ b/src/memory/tool_memory/types.rs @@ -29,8 +29,10 @@ use serde::{Deserialize, Serialize}; /// and retrieval (to sort high-priority guidance ahead of advisory notes). #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] +#[derive(Default)] pub enum ToolMemoryPriority { /// Soft suggestion — surfaced on demand, not eagerly injected. + #[default] Normal, /// Important guidance — eagerly injected at tool-selection time. High, @@ -39,12 +41,6 @@ pub enum ToolMemoryPriority { Critical, } -impl Default for ToolMemoryPriority { - fn default() -> Self { - Self::Normal - } -} - impl ToolMemoryPriority { /// True for priorities that must be eagerly surfaced to the agent /// (Critical/High rules are both pinned into the system prompt and @@ -60,6 +56,7 @@ impl ToolMemoryPriority { /// edicts apart from auto-captured observations. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] +#[derive(Default)] pub enum ToolMemorySource { /// User explicitly asked the agent to remember this rule. UserExplicit, @@ -67,15 +64,10 @@ pub enum ToolMemorySource { /// repeated correction, etc.). PostTurn, /// Written by another subsystem (e.g. an integration provisioner). + #[default] Programmatic, } -impl Default for ToolMemorySource { - fn default() -> Self { - Self::Programmatic - } -} - /// A single tool-scoped memory rule. /// /// Stored under the `tool-{tool_name}` namespace as an entry keyed by diff --git a/src/memory/tree/bucket_seal.rs b/src/memory/tree/bucket_seal.rs index 453bcf2..119c56f 100644 --- a/src/memory/tree/bucket_seal.rs +++ b/src/memory/tree/bucket_seal.rs @@ -26,12 +26,18 @@ use std::sync::Arc; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; +use futures::stream::{StreamExt, TryStreamExt}; use crate::memory::chunks::with_connection; use crate::memory::config::MemoryConfig; +use crate::memory::score::embed::Embedder; use crate::memory::score::extract::EntityExtractor; use crate::memory::score::resolver::canonicalise; use crate::memory::score::store::index_summary_entity_ids_tx; +use crate::memory::store::content::{ + slugify_source_id, stage_summary_with_layout, SummaryComposeInput, SummaryDiskLayout, + SummaryTreeKind, +}; use crate::memory::tree::hydrate::hydrate_inputs; use crate::memory::tree::registry::new_summary_id; use crate::memory::tree::store::{self, Buffer, SummaryNode, Tree}; @@ -41,6 +47,32 @@ use crate::memory::tree::summarise::{fallback_summary, Summariser, SummaryContex /// ever slips. const MAX_CASCADE_DEPTH: u32 = 32; +/// Product callbacks around a seal. Engine state is already durable when +/// `summary_committed` runs; hosts use it for mirrors such as wiki-git. +pub trait SealObserver: Send + Sync { + fn progress(&self, _tree: &Tree, _step: &str, _level: u32, _item_count: Option) {} + fn summary_committed( + &self, + _tree: &Tree, + _node: &SummaryNode, + _content_path: &str, + _reason: &str, + ) -> Result<()> { + Ok(()) + } +} + +pub(crate) struct NoopSealObserver; +impl SealObserver for NoopSealObserver {} + +/// Injected compute and product notifications used by the crate-owned seal +/// pipeline. `embedder = None` deliberately persists a re-embeddable summary. +pub struct SealServices<'a> { + pub summariser: &'a dyn Summariser, + pub embedder: Option<&'a dyn Embedder>, + pub observer: &'a dyn SealObserver, +} + /// How a sealed summary node's `entities` and `topics` fields get populated. #[derive(Clone)] pub enum LabelStrategy { @@ -205,12 +237,36 @@ pub async fn cascade_all_from( force: bool, summariser: &dyn Summariser, strategy: &LabelStrategy, +) -> Result> { + cascade_all_from_with_services( + config, + tree, + start_level, + force, + &SealServices { + summariser, + embedder: None, + observer: &NoopSealObserver, + }, + strategy, + false, + ) + .await +} + +pub async fn cascade_all_from_with_services( + config: &MemoryConfig, + tree: &Tree, + start_level: u32, + force: bool, + services: &SealServices<'_>, + strategy: &LabelStrategy, + enqueue_follow_ups: bool, ) -> Result> { let mut sealed_ids: Vec = Vec::new(); - let mut level: u32 = start_level; let mut first_iteration = true; - for _ in 0..MAX_CASCADE_DEPTH { + for level in (start_level..).take(MAX_CASCADE_DEPTH as usize) { let buf = store::get_buffer(config, &tree.id, level)?; let forced = first_iteration && force; first_iteration = false; @@ -222,16 +278,23 @@ pub async fn cascade_all_from( break; } - let summary_id = seal_one_level(config, tree, &buf, summariser, strategy).await?; + let summary_id = seal_one_level_with_services( + config, + tree, + &buf, + services, + strategy, + enqueue_follow_ups, + ) + .await?; sealed_ids.push(summary_id); - level += 1; } Ok(sealed_ids) } /// Level-aware seal gate. L0 gates on `token_sum`; L≥1 gates on sibling count. /// Budgets are read from [`MemoryConfig::tree`], not hardcoded. -pub(crate) fn should_seal(config: &MemoryConfig, buf: &Buffer) -> bool { +pub fn should_seal(config: &MemoryConfig, buf: &Buffer) -> bool { if buf.is_empty() { return false; } @@ -254,12 +317,13 @@ pub(crate) fn should_seal(config: &MemoryConfig, buf: &Buffer) -> bool { /// transaction therefore re-reads and consumes only the snapshotted prefix; /// concurrent appends remain buffered, while a competing seal causes this /// transaction to abort cleanly. -pub(crate) async fn seal_one_level( +pub async fn seal_one_level_with_services( config: &MemoryConfig, tree: &Tree, buf: &Buffer, - summariser: &dyn Summariser, + services: &SealServices<'_>, strategy: &LabelStrategy, + enqueue_follow_ups: bool, ) -> Result { let level = buf.level; let target_level = level + 1; @@ -298,13 +362,30 @@ pub(crate) async fn seal_one_level( }; // Treat a blank summary the same as a hard error — fall back to the // deterministic concat so we never persist `content = ""`. - let output = match summariser.summarise(&inputs, &ctx).await { + services + .observer + .progress(tree, "summarising", level, Some(inputs.len() as u32)); + let output = match services.summariser.summarise(&inputs, &ctx).await { Ok(o) if !o.content.trim().is_empty() => o, _ => fallback_summary(&inputs, budget), }; let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?; + services.observer.progress(tree, "embedding", level, None); + let embedding = match services.embedder { + Some(embedder) if !output.content.trim().is_empty() => { + let input: String = output.content.chars().take(4_000).collect(); + Some( + embedder + .embed(&input) + .await + .context("embed sealed summary")?, + ) + } + _ => None, + }; + let now = Utc::now(); let summary_id = new_summary_id(target_level); let node = SummaryNode { @@ -325,14 +406,70 @@ pub(crate) async fn seal_one_level( score, sealed_at: now, deleted: false, - embedding: None, + embedding, doc_id: None, version_ms: None, }; + services + .observer + .progress(tree, "persisting", target_level, None); + let summary_kind = match tree.kind { + crate::memory::tree::TreeKind::Source => SummaryTreeKind::Source, + crate::memory::tree::TreeKind::Topic => SummaryTreeKind::Topic, + crate::memory::tree::TreeKind::Global => SummaryTreeKind::Global, + }; + let child_basenames = if target_level == 1 { + Some( + node.child_ids + .iter() + .map(|id| { + crate::memory::chunks::get_chunk_raw_refs(config, id) + .ok() + .flatten() + .and_then(|refs| refs.into_iter().next()) + .map(|raw| { + raw.path + .strip_suffix(".md") + .unwrap_or(&raw.path) + .to_string() + }) + }) + .collect::>(), + ) + } else { + None + }; + let layout = if target_level >= 1_000 { + SummaryDiskLayout::Merge + } else { + SummaryDiskLayout::Standard + }; + let staged = stage_summary_with_layout( + &crate::memory::chunks::content_root(config), + &SummaryComposeInput { + summary_id: &node.id, + tree_kind: summary_kind, + tree_id: &node.tree_id, + tree_scope: &tree.scope, + level: node.level, + child_ids: &node.child_ids, + child_basenames: child_basenames.as_deref(), + child_count: node.child_ids.len(), + time_range_start: node.time_range_start, + time_range_end: node.time_range_end, + sealed_at: node.sealed_at, + body: &node.content, + }, + &slugify_source_id(&tree.scope), + layout, + )?; + let signature = crate::memory::chunks::tree_active_signature(config); let tree_id = tree.id.clone(); let summary_id_for_tx = summary_id.clone(); + let node_for_tx = node.clone(); + let staged_for_tx = staged.clone(); with_connection(config, move |conn| { let tx = conn.unchecked_transaction()?; @@ -345,17 +482,17 @@ pub(crate) async fn seal_one_level( .map(|n| n.max(0) as u32) .context("Failed to read current max_level for tree")?; - store::insert_summary_tx(&tx, &node, &signature)?; + store::insert_staged_summary_tx(&tx, &node_for_tx, Some(&staged_for_tx), &signature)?; index_summary_entity_ids_tx( &tx, - &node.entities, - &node.id, - node.score, + &node_for_tx.entities, + &node_for_tx.id, + node_for_tx.score, now.timestamp_millis(), Some(&tree_id), )?; // Backlink children → new parent for single-row traversal. - for child_id in &node.child_ids { + for child_id in &node_for_tx.child_ids { if level == 0 { tx.execute( "UPDATE mem_tree_chunks SET parent_summary_id = ?1 @@ -377,13 +514,24 @@ pub(crate) async fn seal_one_level( // Append the new summary to the parent buffer. let mut parent = store::get_buffer_conn(&tx, &tree_id, target_level)?; parent.item_ids.push(summary_id_for_tx.clone()); - parent.token_sum = parent.token_sum.saturating_add(node.token_count as i64); + parent.token_sum = parent + .token_sum + .saturating_add(node_for_tx.token_count as i64); parent.oldest_at = match parent.oldest_at { Some(existing) => Some(existing.min(time_range_start)), None => Some(time_range_start), }; store::upsert_buffer_tx(&tx, &parent)?; + if enqueue_follow_ups && should_seal(config, &parent) { + let payload = crate::memory::queue::SealPayload { + tree_id: tree_id.clone(), + level: target_level, + force_now_ms: None, + }; + crate::memory::queue::enqueue_tx(&tx, &crate::memory::queue::NewJob::seal(&payload)?)?; + } + if target_level > current_max { store::update_tree_after_seal_tx(&tx, &tree_id, &summary_id_for_tx, target_level, now)?; } else { @@ -394,9 +542,293 @@ pub(crate) async fn seal_one_level( Ok(()) })?; + services + .observer + .summary_committed(tree, &node, &staged.content_path, "bucket_seal")?; + Ok(summary_id) } +/// Level offset reserved for cross-document merge nodes. +pub const MERGE_LEVEL_BASE: u32 = 1_000; +const DOC_SUBTREE_MAX_FANIN: usize = 32; +const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8; + +/// Build one immutable document-version subtree and feed its root into the +/// shared cross-document merge tier. +pub async fn seal_document_subtree_with_services( + config: &MemoryConfig, + tree: &Tree, + doc_id: &str, + version_ms: Option, + chunk_ids: &[String], + services: &SealServices<'_>, + strategy: &LabelStrategy, +) -> Result { + if chunk_ids.is_empty() { + anyhow::bail!("seal_document_subtree: empty chunk set"); + } + log::debug!( + "[memory_tree:seal_document] enter chunks={} has_version={}", + chunk_ids.len(), + version_ms.is_some() + ); + let mut level = 0; + let mut current_ids = chunk_ids.to_vec(); + let doc_root = loop { + let batches = if level == 0 { + batch_leaves_by_token_budget(config, ¤t_ids)? + } else { + batch_by_count(¤t_ids, DOC_SUBTREE_MAX_FANIN) + }; + let batch_futures: Vec<_> = batches + .iter() + .map(|batch| { + seal_explicit_children( + config, tree, level, batch, doc_id, version_ms, services, strategy, + ) + }) + .collect(); + let nodes: Vec = futures::stream::iter(batch_futures) + .buffered(DOC_SUBTREE_SEAL_CONCURRENCY) + .try_collect() + .await?; + current_ids = nodes.iter().map(|node| node.id.clone()).collect(); + level += 1; + if current_ids.len() <= 1 { + break nodes + .into_iter() + .next() + .context("document seal produced no root")?; + } + }; + + append_to_buffer( + config, + &tree.id, + MERGE_LEVEL_BASE, + &doc_root.id, + doc_root.token_count as i64, + doc_root.time_range_start, + )?; + let root_id = doc_root.id.clone(); + let root_level = doc_root.level; + with_connection(config, move |connection| { + let transaction = connection.unchecked_transaction()?; + let (current_root, current_max): (Option, u32) = transaction.query_row( + "SELECT root_id, max_level FROM mem_tree_trees WHERE id = ?1", + [&tree.id], + |row| Ok((row.get(0)?, row.get(1)?)), + )?; + if current_root.as_deref() != Some(root_id.as_str()) || root_level > current_max { + store::update_tree_after_seal_tx( + &transaction, + &tree.id, + &root_id, + root_level, + Utc::now(), + )?; + } + transaction.commit()?; + Ok(()) + })?; + cascade_all_from_with_services( + config, + tree, + MERGE_LEVEL_BASE, + false, + services, + strategy, + false, + ) + .await?; + log::debug!("[memory_tree:seal_document] complete levels={level}"); + Ok(doc_root.id) +} + +fn batch_leaves_by_token_budget( + config: &MemoryConfig, + chunk_ids: &[String], +) -> Result>> { + let chunks = crate::memory::chunks::get_chunks_batch(config, chunk_ids)?; + let mut batches = Vec::new(); + let mut current = Vec::new(); + let mut tokens = 0_i64; + for id in chunk_ids { + let Some(chunk) = chunks.get(id) else { + continue; + }; + let next = chunk.token_count as i64; + if !current.is_empty() + && (tokens + next > config.tree.input_token_budget as i64 + || current.len() >= DOC_SUBTREE_MAX_FANIN) + { + batches.push(std::mem::take(&mut current)); + tokens = 0; + } + current.push(id.clone()); + tokens += next; + } + if !current.is_empty() { + batches.push(current); + } + if batches.is_empty() { + anyhow::bail!("seal_document_subtree: no resolvable chunks"); + } + Ok(batches) +} + +fn batch_by_count(ids: &[String], max: usize) -> Vec> { + ids.chunks(max.max(1)).map(<[String]>::to_vec).collect() +} + +#[allow(clippy::too_many_arguments)] +async fn seal_explicit_children( + config: &MemoryConfig, + tree: &Tree, + level: u32, + child_ids: &[String], + doc_id: &str, + version_ms: Option, + services: &SealServices<'_>, + strategy: &LabelStrategy, +) -> Result { + let target_level = level + 1; + let inputs = hydrate_inputs(config, level, child_ids)?; + if inputs.is_empty() { + anyhow::bail!("document seal has no hydrated inputs at level {level}"); + } + let time_range_start = inputs.iter().map(|i| i.time_range_start).min().unwrap(); + let time_range_end = inputs.iter().map(|i| i.time_range_end).max().unwrap(); + let score = inputs + .iter() + .map(|i| i.score) + .fold(f32::NEG_INFINITY, f32::max) + .max(0.0); + let output = if inputs.len() == 1 && inputs[0].token_count <= config.tree.output_token_budget { + crate::memory::tree::SummaryOutput { + content: inputs[0].content.clone(), + token_count: inputs[0].token_count, + ..Default::default() + } + } else { + let context = SummaryContext { + tree_id: &tree.id, + tree_kind: tree.kind, + target_level, + token_budget: config.tree.output_token_budget, + }; + match services.summariser.summarise(&inputs, &context).await { + Ok(output) if !output.content.trim().is_empty() => output, + _ => fallback_summary(&inputs, context.token_budget), + } + }; + let (entities, topics) = resolve_labels(strategy, &inputs, &output.content).await?; + let embedding = match services.embedder { + Some(embedder) if !output.content.trim().is_empty() => { + let input: String = output.content.chars().take(4_000).collect(); + Some( + embedder + .embed(&input) + .await + .context("embed document summary")?, + ) + } + _ => None, + }; + let now = Utc::now(); + let node = SummaryNode { + id: new_summary_id(target_level), + tree_id: tree.id.clone(), + tree_kind: tree.kind, + level: target_level, + parent_id: None, + child_ids: child_ids.to_vec(), + content: output.content, + token_count: output.token_count, + entities, + topics, + time_range_start, + time_range_end, + score, + sealed_at: now, + deleted: false, + embedding, + doc_id: Some(doc_id.to_string()), + version_ms, + }; + let summary_kind = match tree.kind { + crate::memory::tree::TreeKind::Source => SummaryTreeKind::Source, + crate::memory::tree::TreeKind::Topic => SummaryTreeKind::Topic, + crate::memory::tree::TreeKind::Global => SummaryTreeKind::Global, + }; + let doc_slug = slugify_source_id(doc_id); + let staged = stage_summary_with_layout( + &crate::memory::chunks::content_root(config), + &SummaryComposeInput { + summary_id: &node.id, + tree_kind: summary_kind, + tree_id: &node.tree_id, + tree_scope: &tree.scope, + level: node.level, + child_ids: &node.child_ids, + child_basenames: None, + child_count: node.child_ids.len(), + time_range_start: node.time_range_start, + time_range_end: node.time_range_end, + sealed_at: node.sealed_at, + body: &node.content, + }, + &slugify_source_id(&tree.scope), + SummaryDiskLayout::DocSubtree { + doc_slug: &doc_slug, + version_ms, + }, + )?; + let signature = crate::memory::chunks::tree_active_signature(config); + let node_for_tx = node.clone(); + let staged_for_tx = staged.clone(); + with_connection(config, move |connection| { + let transaction = connection.unchecked_transaction()?; + store::insert_staged_summary_tx( + &transaction, + &node_for_tx, + Some(&staged_for_tx), + &signature, + )?; + index_summary_entity_ids_tx( + &transaction, + &node_for_tx.entities, + &node_for_tx.id, + node_for_tx.score, + now.timestamp_millis(), + Some(&node_for_tx.tree_id), + )?; + for child_id in &node_for_tx.child_ids { + if level == 0 { + transaction.execute( + "UPDATE mem_tree_chunks SET parent_summary_id = ?1 WHERE id = ?2", + rusqlite::params![&node_for_tx.id, child_id], + )?; + } else { + transaction.execute( + "UPDATE mem_tree_summaries SET parent_id = ?1 WHERE id = ?2 AND parent_id IS NULL", + rusqlite::params![&node_for_tx.id, child_id], + )?; + } + } + transaction.commit()?; + Ok(()) + })?; + services.observer.summary_committed( + tree, + &node, + &staged.content_path, + "document_subtree_seal", + )?; + Ok(node) +} + #[cfg(test)] #[path = "bucket_seal_label_tests.rs"] mod label_tests; diff --git a/src/memory/tree/bucket_seal_label_tests.rs b/src/memory/tree/bucket_seal_label_tests.rs index daeb6c7..ac1fbaf 100644 --- a/src/memory/tree/bucket_seal_label_tests.rs +++ b/src/memory/tree/bucket_seal_label_tests.rs @@ -72,7 +72,7 @@ fn seed_chunk( created_at: ts, partial_message: false, }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); + upsert_chunks(cfg, std::slice::from_ref(&c)).unwrap(); c } @@ -347,11 +347,31 @@ async fn hydrate_summary_inputs_preserves_order_and_skips_missing() { version_ms: None, }; let sum_a = mk("sum-a", "BODY-A", 11, 0.11, "entity:alice"); - let sum_b = mk("sum-b", "BODY-B", 22, 0.22, "entity:bob"); + let full_body_b = "B".repeat(700); + let sum_b = mk("sum-b", &full_body_b, 22, 0.22, "entity:bob"); + let staged_b = crate::memory::store::content::stage_summary( + &crate::memory::chunks::content_root(&cfg), + &crate::memory::store::content::SummaryComposeInput { + summary_id: &sum_b.id, + tree_kind: crate::memory::store::content::SummaryTreeKind::Source, + tree_id: &tree.id, + tree_scope: &tree.scope, + level: sum_b.level, + child_ids: &sum_b.child_ids, + child_basenames: None, + child_count: sum_b.child_ids.len(), + time_range_start: ts, + time_range_end: ts, + sealed_at: ts, + body: &full_body_b, + }, + "slack-eng", + ) + .unwrap(); crate::memory::chunks::with_connection(&cfg, |conn| { let tx = conn.unchecked_transaction()?; store::insert_summary_tx(&tx, &sum_a, "test")?; - store::insert_summary_tx(&tx, &sum_b, "test")?; + store::insert_staged_summary_tx(&tx, &sum_b, Some(&staged_b), "test")?; tx.commit()?; Ok(()) }) @@ -363,7 +383,7 @@ async fn hydrate_summary_inputs_preserves_order_and_skips_missing() { assert_eq!(out[0].id, "sum-b"); assert_eq!(out[1].id, "sum-a"); assert_eq!(out[0].token_count, 22); - assert_eq!(out[0].content, "BODY-B"); + assert_eq!(out[0].content, full_body_b); assert_eq!(out[1].content, "BODY-A"); assert_eq!(out[0].entities, vec!["entity:bob".to_string()]); } diff --git a/src/memory/tree/bucket_seal_tests.rs b/src/memory/tree/bucket_seal_tests.rs index ec31f5c..1ecf98f 100644 --- a/src/memory/tree/bucket_seal_tests.rs +++ b/src/memory/tree/bucket_seal_tests.rs @@ -82,7 +82,7 @@ fn seed_chunk( created_at: ts, partial_message: false, }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); + upsert_chunks(cfg, std::slice::from_ref(&c)).unwrap(); c } @@ -103,6 +103,88 @@ async fn append_below_budget_does_not_seal() { assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); } +#[tokio::test] +async fn service_seal_stages_body_and_enqueues_parent_atomically() { + let (_tmp, mut cfg) = test_config(); + cfg.tree.summary_fanout = 1; + let tree = get_or_create_tree(&cfg, TreeKind::Source, "slack:#eng").unwrap(); + let chunk = seed_chunk(&cfg, 0, "full staged body", 10, Vec::new()); + append_to_buffer( + &cfg, + &tree.id, + 0, + &chunk.id, + chunk.token_count as i64, + chunk.metadata.timestamp, + ) + .unwrap(); + let buffer = store::get_buffer(&cfg, &tree.id, 0).unwrap(); + let summariser = ConcatSummariser::new(); + + let summary_id = seal_one_level_with_services( + &cfg, + &tree, + &buffer, + &SealServices { + summariser: &summariser, + embedder: None, + observer: &NoopSealObserver, + }, + &LabelStrategy::Empty, + true, + ) + .await + .unwrap(); + + assert!( + crate::memory::store::content::read_summary_body(&cfg, &summary_id) + .unwrap() + .contains("full staged body") + ); + assert_eq!(crate::memory::queue::count_total(&cfg).unwrap(), 1); +} + +#[tokio::test] +async fn document_subtree_stages_versioned_passthrough_root() { + let (_tmp, cfg) = test_config(); + let tree = get_or_create_tree(&cfg, TreeKind::Source, "notion:connection").unwrap(); + let chunk = seed_chunk(&cfg, 0, "document body", 10, Vec::new()); + let summariser = ConcatSummariser::new(); + let root_id = seal_document_subtree_with_services( + &cfg, + &tree, + "notion:connection:page", + Some(42), + &[chunk.id], + &SealServices { + summariser: &summariser, + embedder: None, + observer: &NoopSealObserver, + }, + &LabelStrategy::Empty, + ) + .await + .unwrap(); + + let root = store::get_summary(&cfg, &root_id).unwrap().unwrap(); + assert_eq!(root.content, "document body"); + assert_eq!(root.doc_id.as_deref(), Some("notion:connection:page")); + assert_eq!(root.version_ms, Some(42)); + assert_eq!( + crate::memory::store::content::read_summary_body(&cfg, &root_id).unwrap(), + "document body" + ); + assert_eq!( + store::get_buffer(&cfg, &tree.id, MERGE_LEVEL_BASE) + .unwrap() + .item_ids, + vec![root_id] + ); + let published = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); + assert_eq!(published.root_id.as_deref(), Some(root.id.as_str())); + assert_eq!(published.max_level, root.level); +} + #[tokio::test] async fn crossing_budget_triggers_seal() { let (_tmp, cfg) = test_config(); @@ -186,7 +268,20 @@ async fn seal_preserves_items_appended_while_summariser_is_running() { started: started.clone(), release: release.clone(), }; - let seal = seal_one_level(&cfg, &tree, &snapshot, &summariser, &LabelStrategy::Empty); + let observer = NoopSealObserver; + let services = SealServices { + summariser: &summariser, + embedder: None, + observer: &observer, + }; + let seal = seal_one_level_with_services( + &cfg, + &tree, + &snapshot, + &services, + &LabelStrategy::Empty, + false, + ); let append = async { started.notified().await; append_to_buffer(&cfg, &tree.id, 0, &c2.id, 10, ts).unwrap(); diff --git a/src/memory/tree/direct_ingest.rs b/src/memory/tree/direct_ingest.rs new file mode 100644 index 0000000..dedaeb4 --- /dev/null +++ b/src/memory/tree/direct_ingest.rs @@ -0,0 +1,185 @@ +//! Direct ingestion of pre-built L1 summaries. + +use anyhow::Result; +use chrono::{DateTime, Utc}; + +use crate::memory::chunks::{tree_active_signature, with_connection}; +use crate::memory::config::MemoryConfig; +use crate::memory::score::store::index_summary_entity_ids_tx; +use crate::memory::store::content::{ + slugify_source_id, stage_summary, SummaryComposeInput, SummaryTreeKind, +}; +use crate::memory::tree::bucket_seal::cascade_all_from; +use crate::memory::tree::registry::new_summary_id; +use crate::memory::tree::store::{self, SummaryNode, Tree}; +use crate::memory::tree::summarise::Summariser; +use crate::memory::tree::TreeFactory; + +#[derive(Clone, Debug)] +pub struct SummaryIngestInput { + pub content: String, + pub token_count: u32, + pub entities: Vec, + pub topics: Vec, + pub time_range_start: DateTime, + pub time_range_end: DateTime, + pub score: f32, + pub child_labels: Vec, + pub child_basenames: Vec>, +} + +#[derive(Clone, Debug)] +pub struct SummaryIngestOutcome { + pub summary_id: String, + pub content_path: String, + pub sealed_ids: Vec, +} + +pub async fn ingest_summary( + config: &MemoryConfig, + tree: &Tree, + input: SummaryIngestInput, + summariser: &dyn Summariser, +) -> Result { + const TARGET_LEVEL: u32 = 1; + let summary_id = new_summary_id(TARGET_LEVEL); + let sealed_at = Utc::now(); + let node = SummaryNode { + id: summary_id.clone(), + tree_id: tree.id.clone(), + tree_kind: tree.kind, + level: TARGET_LEVEL, + parent_id: None, + child_ids: input.child_labels, + content: input.content, + token_count: input.token_count, + entities: input.entities, + topics: input.topics, + time_range_start: input.time_range_start, + time_range_end: input.time_range_end, + score: input.score, + sealed_at, + deleted: false, + embedding: None, + doc_id: None, + version_ms: None, + }; + let content_root = crate::memory::chunks::content_root(config); + let summary_tree_kind = match tree.kind { + crate::memory::tree::TreeKind::Source => SummaryTreeKind::Source, + crate::memory::tree::TreeKind::Topic => SummaryTreeKind::Topic, + crate::memory::tree::TreeKind::Global => SummaryTreeKind::Global, + }; + let staged = stage_summary( + &content_root, + &SummaryComposeInput { + summary_id: &summary_id, + tree_kind: summary_tree_kind, + tree_id: &tree.id, + tree_scope: &tree.scope, + level: TARGET_LEVEL, + child_ids: &node.child_ids, + child_basenames: if input.child_basenames.is_empty() { + None + } else { + Some(&input.child_basenames) + }, + child_count: node.child_ids.len(), + time_range_start: node.time_range_start, + time_range_end: node.time_range_end, + sealed_at, + body: &node.content, + }, + &slugify_source_id(&tree.scope), + )?; + let signature = tree_active_signature(config); + with_connection(config, |connection| { + let transaction = connection.unchecked_transaction()?; + let current_max: i64 = transaction.query_row( + "SELECT max_level FROM mem_tree_trees WHERE id = ?1", + [&tree.id], + |row| row.get(0), + )?; + store::insert_staged_summary_tx(&transaction, &node, Some(&staged), &signature)?; + index_summary_entity_ids_tx( + &transaction, + &node.entities, + &node.id, + node.score, + sealed_at.timestamp_millis(), + Some(&tree.id), + )?; + let mut buffer = store::get_buffer_conn(&transaction, &tree.id, TARGET_LEVEL)?; + if !buffer.item_ids.contains(&summary_id) { + buffer.item_ids.push(summary_id.clone()); + buffer.token_sum = buffer.token_sum.saturating_add(node.token_count as i64); + buffer.oldest_at = Some(buffer.oldest_at.map_or(node.time_range_start, |existing| { + existing.min(node.time_range_start) + })); + store::upsert_buffer_tx(&transaction, &buffer)?; + } + if TARGET_LEVEL > current_max.max(0) as u32 { + store::update_tree_after_seal_tx( + &transaction, + &tree.id, + &summary_id, + TARGET_LEVEL, + sealed_at, + )?; + } + transaction.commit()?; + Ok(()) + })?; + + let strategy = TreeFactory::from_tree(tree).label_strategy(); + let sealed_ids = + cascade_all_from(config, tree, TARGET_LEVEL, false, summariser, &strategy).await?; + #[cfg(feature = "sync")] + tracing::info!(summary_id, tree_id = %tree.id, cascaded = sealed_ids.len(), "[memory_tree:direct_ingest] summary ingested"); + Ok(SummaryIngestOutcome { + summary_id, + content_path: staged.content_path, + sealed_ids, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::tree::{store::get_buffer, ConcatSummariser}; + + #[tokio::test] + async fn direct_summary_lands_at_l1_without_creating_chunks() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + let tree = TreeFactory::source("github:org/repo") + .get_or_create(&config) + .unwrap(); + let now = Utc::now(); + let outcome = ingest_summary( + &config, + &tree, + SummaryIngestInput { + content: "summary".into(), + token_count: 2, + entities: Vec::new(), + topics: Vec::new(), + time_range_start: now, + time_range_end: now, + score: 0.5, + child_labels: vec!["100_issue-1".into()], + child_basenames: Vec::new(), + }, + &ConcatSummariser, + ) + .await + .unwrap(); + assert_eq!(crate::memory::chunks::count_chunks(&config).unwrap(), 0); + let summary = store::get_summary(&config, &outcome.summary_id) + .unwrap() + .unwrap(); + assert_eq!(summary.level, 1); + assert_eq!(summary.child_ids, vec!["100_issue-1"]); + assert_eq!(get_buffer(&config, &tree.id, 1).unwrap().item_ids.len(), 1); + } +} diff --git a/src/memory/tree/factory_tests.rs b/src/memory/tree/factory_tests.rs index 9c507a2..e41f4e0 100644 --- a/src/memory/tree/factory_tests.rs +++ b/src/memory/tree/factory_tests.rs @@ -38,7 +38,7 @@ fn seed_chunk(cfg: &MemoryConfig, seq: u32, content: &str) -> Chunk { created_at: ts, partial_message: false, }; - upsert_chunks(cfg, &[chunk.clone()]).unwrap(); + upsert_chunks(cfg, std::slice::from_ref(&chunk)).unwrap(); chunk } diff --git a/src/memory/tree/flush.rs b/src/memory/tree/flush.rs index e994724..02b6c6e 100644 --- a/src/memory/tree/flush.rs +++ b/src/memory/tree/flush.rs @@ -11,7 +11,9 @@ use anyhow::Result; use chrono::{DateTime, Duration, Utc}; use crate::memory::config::MemoryConfig; -use crate::memory::tree::bucket_seal::{cascade_all_from, LabelStrategy}; +use crate::memory::tree::bucket_seal::{ + cascade_all_from, cascade_all_from_with_services, LabelStrategy, SealServices, +}; use crate::memory::tree::store::{self, DEFAULT_FLUSH_AGE_SECS}; use crate::memory::tree::summarise::Summariser; @@ -22,6 +24,25 @@ pub async fn flush_stale_buffers( max_age: Duration, summariser: &dyn Summariser, strategy: &LabelStrategy, +) -> Result { + flush_stale_buffers_with_services( + config, + max_age, + &SealServices { + summariser, + embedder: None, + observer: &super::bucket_seal::NoopSealObserver, + }, + strategy, + ) + .await +} + +pub async fn flush_stale_buffers_with_services( + config: &MemoryConfig, + max_age: Duration, + services: &SealServices<'_>, + strategy: &LabelStrategy, ) -> Result { let now = Utc::now(); let cutoff = now - max_age; @@ -45,7 +66,10 @@ pub async fn flush_stale_buffers( let Some(tree) = tree_by_id.get(&buf.tree_id) else { continue; // orphan buffer — tree row gone }; - let sealed = cascade_all_from(config, tree, buf.level, true, summariser, strategy).await?; + let sealed = cascade_all_from_with_services( + config, tree, buf.level, true, services, strategy, false, + ) + .await?; seals += sealed.len(); } Ok(seals) diff --git a/src/memory/tree/flush_tests.rs b/src/memory/tree/flush_tests.rs index 7e248d5..47eabb0 100644 --- a/src/memory/tree/flush_tests.rs +++ b/src/memory/tree/flush_tests.rs @@ -38,7 +38,7 @@ fn seed_chunk(cfg: &MemoryConfig, src: &str, seq: u32, content: &str, ts: DateTi created_at: ts, partial_message: false, }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); + upsert_chunks(cfg, std::slice::from_ref(&c)).unwrap(); c } diff --git a/src/memory/tree/hydrate.rs b/src/memory/tree/hydrate.rs index 2f7989f..3608733 100644 --- a/src/memory/tree/hydrate.rs +++ b/src/memory/tree/hydrate.rs @@ -64,9 +64,17 @@ pub(crate) fn hydrate_summary_inputs( let Some(node) = node_by_id.get(id) else { continue; }; + let content = crate::memory::store::content::read_summary_body(config, &node.id) + .unwrap_or_else(|error| { + log::warn!( + "[memory_tree:hydrate] staged summary read failed id={}: {error}; using SQL content", + node.id + ); + node.content.clone() + }); out.push(SummaryInput { id: node.id.clone(), - content: node.content.clone(), + content, token_count: node.token_count, entities: node.entities.clone(), topics: node.topics.clone(), diff --git a/src/memory/tree/mod.rs b/src/memory/tree/mod.rs index dcffebb..0b94f3b 100644 --- a/src/memory/tree/mod.rs +++ b/src/memory/tree/mod.rs @@ -25,6 +25,7 @@ pub mod store; pub mod bucket_seal; +mod direct_ingest; pub mod factory; pub mod flush; mod hydrate; @@ -37,10 +38,17 @@ pub mod summarise; // ── Public API surface ────────────────────────────────────────────────────── pub use bucket_seal::{ - append_leaf, append_leaf_deferred, append_to_buffer, cascade_all_from, LabelStrategy, LeafRef, + append_leaf, append_leaf_deferred, append_to_buffer, cascade_all_from, + cascade_all_from_with_services, seal_document_subtree_with_services, + seal_one_level_with_services, should_seal, LabelStrategy, LeafRef, SealObserver, SealServices, + MERGE_LEVEL_BASE, }; +pub use direct_ingest::{ingest_summary, SummaryIngestInput, SummaryIngestOutcome}; pub use factory::{TreeFactory, TreeProfile, GLOBAL_SCOPE}; -pub use flush::{flush_stale_buffers, flush_stale_buffers_default, force_flush_tree}; +pub use flush::{ + flush_stale_buffers, flush_stale_buffers_default, flush_stale_buffers_with_services, + force_flush_tree, +}; pub use hydrate::fetch_leaves; pub use io::{ TreeLabelStrategy, TreeLeafPayload, TreeReadHit, TreeReadRequest, TreeReadResult, @@ -53,5 +61,6 @@ pub use store::{ SUMMARY_FANOUT, }; pub use summarise::{ - fallback_summary, ConcatSummariser, Summariser, SummaryContext, SummaryInput, SummaryOutput, + fallback_summary, finish_provider_summary, prepare_summary_prompt, ConcatSummariser, + PreparedSummaryPrompt, Summariser, SummaryCall, SummaryContext, SummaryInput, SummaryOutput, }; diff --git a/src/memory/tree/runtime/engine.rs b/src/memory/tree/runtime/engine.rs index a999868..7204fe2 100644 --- a/src/memory/tree/runtime/engine.rs +++ b/src/memory/tree/runtime/engine.rs @@ -31,6 +31,25 @@ pub trait Summariser: Send + Sync { async fn summarise(&self, system: Option<&str>, content: &str) -> Result; } +/// Product hook for runtime progress notifications. Implementations must not +/// retain PII-bearing summary content; only structural identifiers and counts +/// are exposed here. +pub trait RuntimeObserver: Send + Sync { + fn hour_completed(&self, _namespace: &str, _node_id: &str, _token_count: u32) {} + fn node_propagated( + &self, + _namespace: &str, + _node_id: &str, + _level: NodeLevel, + _token_count: u32, + ) { + } + fn rebuild_completed(&self, _namespace: &str, _total_nodes: u64) {} +} + +struct NoopObserver; +impl RuntimeObserver for NoopObserver {} + /// Run the summarisation job for a namespace: drain the buffer, group entries /// by hour, summarise each hour into its leaf, then propagate upward. Returns /// the last hour leaf created, or `None` if the buffer was empty. @@ -54,6 +73,16 @@ pub async fn run_summarization( summariser: &dyn Summariser, namespace: &str, _ts: DateTime, +) -> Result> { + run_summarization_observed(config, summariser, namespace, _ts, &NoopObserver).await +} + +pub async fn run_summarization_observed( + config: &MemoryConfig, + summariser: &dyn Summariser, + namespace: &str, + _ts: DateTime, + observer: &dyn RuntimeObserver, ) -> Result> { let buffered = store::buffer_read(config, namespace)?; if buffered.is_empty() { @@ -100,6 +129,7 @@ pub async fn run_summarization( metadata: None, }; store::write_node(config, &hour_node)?; + observer.hour_completed(namespace, hour_id, hour_node.token_count); let (_, day_id, month_id, year_id, root_id) = derive_node_ids_from_hour_id(hour_id); all_propagation_ids.push((day_id, NodeLevel::Day)); @@ -120,7 +150,9 @@ pub async fn run_summarization( ] { for (node_id, node_level) in &all_propagation_ids { if *node_level == level && seen.insert(node_id.clone()) { - if let Err(_e) = propagate_node(config, summariser, namespace, node_id, level).await + if let Err(_e) = + propagate_node_observed(config, summariser, namespace, node_id, level, observer) + .await { failed.push(node_id.clone()); } @@ -159,6 +191,15 @@ pub async fn rebuild_tree( config: &MemoryConfig, summariser: &dyn Summariser, namespace: &str, +) -> Result { + rebuild_tree_observed(config, summariser, namespace, &NoopObserver).await +} + +pub async fn rebuild_tree_observed( + config: &MemoryConfig, + summariser: &dyn Summariser, + namespace: &str, + observer: &dyn RuntimeObserver, ) -> Result { let status = store::get_tree_status(config, namespace)?; if status.total_nodes == 0 { @@ -217,26 +258,72 @@ pub async fn rebuild_tree( // Partial-success propagation: a single node's failure does not abort the // rebuild (the hour leaves are already re-written above). for day_id in &day_ids { - let _ = propagate_node(config, summariser, namespace, day_id, NodeLevel::Day).await; + let _ = propagate_node_observed( + config, + summariser, + namespace, + day_id, + NodeLevel::Day, + observer, + ) + .await; } for month_id in &month_ids { - let _ = propagate_node(config, summariser, namespace, month_id, NodeLevel::Month).await; + let _ = propagate_node_observed( + config, + summariser, + namespace, + month_id, + NodeLevel::Month, + observer, + ) + .await; } for year_id in &year_ids { - let _ = propagate_node(config, summariser, namespace, year_id, NodeLevel::Year).await; + let _ = propagate_node_observed( + config, + summariser, + namespace, + year_id, + NodeLevel::Year, + observer, + ) + .await; } - let _ = propagate_node(config, summariser, namespace, "root", NodeLevel::Root).await; + let _ = propagate_node_observed( + config, + summariser, + namespace, + "root", + NodeLevel::Root, + observer, + ) + .await; - store::get_tree_status(config, namespace) + let status = store::get_tree_status(config, namespace)?; + observer.rebuild_completed(namespace, status.total_nodes); + Ok(status) } /// Re-summarise a single non-leaf node from its children. +#[cfg(test)] pub(crate) async fn propagate_node( config: &MemoryConfig, summariser: &dyn Summariser, namespace: &str, node_id: &str, level: NodeLevel, +) -> Result<()> { + propagate_node_observed(config, summariser, namespace, node_id, level, &NoopObserver).await +} + +async fn propagate_node_observed( + config: &MemoryConfig, + summariser: &dyn Summariser, + namespace: &str, + node_id: &str, + level: NodeLevel, + observer: &dyn RuntimeObserver, ) -> Result<()> { let children = store::read_children(config, namespace, node_id)?; if children.is_empty() { @@ -274,6 +361,7 @@ pub(crate) async fn propagate_node( metadata: None, }; store::write_node(config, &node)?; + observer.node_propagated(namespace, node_id, level, node.token_count); Ok(()) } diff --git a/src/memory/tree/runtime/engine_tests.rs b/src/memory/tree/runtime/engine_tests.rs index 14098e4..41b78fe 100644 --- a/src/memory/tree/runtime/engine_tests.rs +++ b/src/memory/tree/runtime/engine_tests.rs @@ -5,6 +5,7 @@ use super::*; use async_trait::async_trait; use chrono::{TimeZone, Utc}; +use std::sync::Mutex; use tempfile::TempDir; use crate::memory::config::MemoryConfig; @@ -18,6 +19,33 @@ fn test_config(tmp: &TempDir) -> MemoryConfig { struct StubSummariser { reply: String, } + +#[derive(Default)] +struct RecordingObserver { + events: Mutex>, +} + +impl RuntimeObserver for RecordingObserver { + fn hour_completed(&self, namespace: &str, node_id: &str, _token_count: u32) { + self.events + .lock() + .unwrap() + .push(format!("hour:{namespace}:{node_id}")); + } + + fn node_propagated( + &self, + namespace: &str, + node_id: &str, + _level: NodeLevel, + _token_count: u32, + ) { + self.events + .lock() + .unwrap() + .push(format!("propagated:{namespace}:{node_id}")); + } +} impl StubSummariser { fn with_reply(reply: impl Into) -> Self { Self { @@ -201,6 +229,34 @@ async fn run_summarization_builds_ancestor_chain() { assert!(store::buffer_read(&cfg, ns).unwrap().is_empty()); } +#[tokio::test] +async fn observed_run_reports_hour_and_propagation_progress() { + let tmp = TempDir::new().unwrap(); + let cfg = test_config(&tmp); + let ns = "observed"; + let ts = Utc.with_ymd_and_hms(2024, 4, 1, 8, 0, 0).unwrap(); + store::buffer_write(&cfg, ns, "raw", &ts, None).unwrap(); + let observer = RecordingObserver::default(); + + run_summarization_observed( + &cfg, + &StubSummariser::with_reply("summary"), + ns, + ts, + &observer, + ) + .await + .unwrap(); + + let events = observer.events.lock().unwrap(); + assert!(events + .iter() + .any(|event| event.starts_with("hour:observed:"))); + assert!(events + .iter() + .any(|event| event == "propagated:observed:root")); +} + #[tokio::test] async fn run_summarization_multi_hour_groups_produce_multiple_leaves() { let tmp = TempDir::new().unwrap(); diff --git a/src/memory/tree/runtime/mod.rs b/src/memory/tree/runtime/mod.rs index 1d5070f..95c93ea 100644 --- a/src/memory/tree/runtime/mod.rs +++ b/src/memory/tree/runtime/mod.rs @@ -12,7 +12,10 @@ pub mod engine; pub mod store; pub mod types; -pub use engine::{discover_active_namespaces, rebuild_tree, run_summarization, Summariser}; +pub use engine::{ + discover_active_namespaces, rebuild_tree, rebuild_tree_observed, run_summarization, + run_summarization_observed, RuntimeObserver, Summariser, +}; pub use types::{ derive_node_ids, derive_parent_id, estimate_tokens, level_from_node_id, node_id_to_path, IngestRequest, NodeLevel, QueryResult, TreeNode, TreeStatus, diff --git a/src/memory/tree/store/buffers.rs b/src/memory/tree/store/buffers.rs index 80d7c6f..56f083e 100644 --- a/src/memory/tree/store/buffers.rs +++ b/src/memory/tree/store/buffers.rs @@ -14,7 +14,7 @@ pub fn get_buffer(config: &MemoryConfig, tree_id: &str, level: u32) -> Result Result { +pub fn get_buffer_conn(conn: &Connection, tree_id: &str, level: u32) -> Result { let mut stmt = conn.prepare( "SELECT tree_id, level, item_ids_json, token_sum, oldest_at_ms FROM mem_tree_buffers WHERE tree_id = ?1 AND level = ?2", @@ -27,7 +27,7 @@ pub(crate) fn get_buffer_conn(conn: &Connection, tree_id: &str, level: u32) -> R } /// Upsert a buffer row. -pub(crate) fn upsert_buffer_tx(tx: &Transaction<'_>, buf: &Buffer) -> Result<()> { +pub fn upsert_buffer_tx(tx: &Transaction<'_>, buf: &Buffer) -> Result<()> { let now_ms = Utc::now().timestamp_millis(); tx.execute( "INSERT INTO mem_tree_buffers ( @@ -56,6 +56,13 @@ pub(crate) fn upsert_buffer_tx(tx: &Transaction<'_>, buf: &Buffer) -> Result<()> Ok(()) } +/// Reset a buffer at `(tree_id, level)` to empty. +/// +/// Seal transactions should use the snapshot-aware internal consumer instead. +pub fn clear_buffer_tx(tx: &Transaction<'_>, tree_id: &str, level: u32) -> Result<()> { + upsert_buffer_tx(tx, &Buffer::empty(tree_id, level)) +} + /// Consume exactly a previously-read buffer snapshot during a seal. /// /// The current row must still begin with the snapshot ids. Items appended diff --git a/src/memory/tree/store/mod.rs b/src/memory/tree/store/mod.rs index 1922410..ae9d885 100644 --- a/src/memory/tree/store/mod.rs +++ b/src/memory/tree/store/mod.rs @@ -26,22 +26,25 @@ mod tests; // ── Tree rows ─────────────────────────────────────────────────────────────── pub use trees::{ - archive_tree, get_tree, get_tree_by_scope, get_trees_batch, insert_tree, list_trees_by_kind, + archive_tree, delete_tree_cascade_tx, get_tree, get_tree_by_scope, get_tree_by_scope_conn, + get_trees_batch, insert_tree, insert_tree_conn, list_trees_by_kind, refresh_last_sealed_tx, + update_tree_after_seal_tx, TreeCascadeDeletion, }; -pub(crate) use trees::{refresh_last_sealed_tx, update_tree_after_seal_tx}; // ── Summary rows + embeddings ─────────────────────────────────────────────── pub use summaries::{ count_summaries, get_summaries_batch, get_summary, get_summary_embedding, get_summary_embedding_for_signature, get_summary_embeddings_batch, - get_summary_embeddings_for_signature_batch, insert_summary_tx, list_children_of_summary, - list_summaries_at_level, list_summaries_in_window, set_summary_embedding, - set_summary_embedding_for_signature, + get_summary_embeddings_for_signature_batch, insert_staged_summary_tx, insert_summary_tx, + list_children_of_summary, list_summaries_at_level, list_summaries_in_window, + set_summary_embedding, set_summary_embedding_for_signature, }; // ── Buffers ───────────────────────────────────────────────────────────────── -pub(crate) use buffers::{consume_snapshot_tx, get_buffer_conn, upsert_buffer_tx}; -pub use buffers::{get_buffer, list_stale_buffers}; +pub(crate) use buffers::consume_snapshot_tx; +pub use buffers::{ + clear_buffer_tx, get_buffer, get_buffer_conn, list_stale_buffers, upsert_buffer_tx, +}; // ── Type + constant re-exports ────────────────────────────────────────────── pub use types::{ diff --git a/src/memory/tree/store/store_tests.rs b/src/memory/tree/store/store_tests.rs index 72d911d..c2b844b 100644 --- a/src/memory/tree/store/store_tests.rs +++ b/src/memory/tree/store/store_tests.rs @@ -11,6 +11,7 @@ use tempfile::TempDir; use crate::memory::chunks::with_connection; use crate::memory::config::MemoryConfig; +use crate::memory::store::content::StagedSummary; fn test_config() -> (TempDir, MemoryConfig) { let tmp = TempDir::new().unwrap(); @@ -93,6 +94,58 @@ fn summary_insert_and_fetch() { assert_eq!(count_summaries(&cfg, "tree-1").unwrap(), 1); } +#[test] +fn staged_summary_persists_preview_and_content_pointer_atomically() { + let (_tmp, cfg) = test_config(); + insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap(); + let mut node = sample_summary("sum-staged", "tree-1", 1); + node.content = "x".repeat(700); + let staged = StagedSummary { + summary_id: node.id.clone(), + content_path: "summaries/tree-1/sum-staged.md".into(), + content_sha256: "abc123".into(), + }; + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_staged_summary_tx(&tx, &node, Some(&staged), "test")?; + tx.commit()?; + let row: (String, String, String) = conn.query_row( + "SELECT content, content_path, content_sha256 FROM mem_tree_summaries WHERE id=?1", + [&node.id], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + )?; + assert_eq!(row.0.len(), 500); + assert_eq!(row.1, staged.content_path); + assert_eq!(row.2, staged.content_sha256); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn cascade_delete_returns_staged_paths_and_removes_tree_rows() { + let (_tmp, cfg) = test_config(); + insert_tree(&cfg, &sample_tree("tree-delete", "slack:#delete")).unwrap(); + let node = sample_summary("sum-delete", "tree-delete", 1); + let staged = StagedSummary { + summary_id: node.id.clone(), + content_path: "summaries/tree-delete/sum-delete.md".into(), + content_sha256: "abc123".into(), + }; + let deleted = with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_staged_summary_tx(&tx, &node, Some(&staged), "test")?; + let deleted = delete_tree_cascade_tx(&tx, "tree-delete")?; + tx.commit()?; + Ok(deleted) + }) + .unwrap(); + assert_eq!(deleted.removed_summaries, 1); + assert_eq!(deleted.content_paths, vec![staged.content_path]); + assert!(get_tree(&cfg, "tree-delete").unwrap().is_none()); + assert!(get_summary(&cfg, "sum-delete").unwrap().is_none()); +} + #[test] fn list_summaries_in_window_keeps_only_fully_contained() { let (_tmp, cfg) = test_config(); @@ -376,7 +429,7 @@ fn get_trees_batch_returns_present_ids_and_skips_missing() { assert_eq!(map.len(), 2); assert_eq!(map.get("tree-a").unwrap(), &a); assert_eq!(map.get("tree-b").unwrap(), &b); - assert!(map.get("ghost").is_none()); + assert!(!map.contains_key("ghost")); } #[test] @@ -429,7 +482,7 @@ fn summary_batch_embedding_lookup_returns_only_signature_scoped_rows() { let map_a = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); assert_eq!(map_a.len(), 2); assert_eq!(map_a.get("sum-1").cloned(), Some(vec![0.1, 0.2])); - assert!(map_a.get("sum-3").is_none()); + assert!(!map_a.contains_key("sum-3")); let map_b = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); assert_eq!(map_b.len(), 1); diff --git a/src/memory/tree/store/summaries.rs b/src/memory/tree/store/summaries.rs index 8019a80..85d4869 100644 --- a/src/memory/tree/store/summaries.rs +++ b/src/memory/tree/store/summaries.rs @@ -16,6 +16,7 @@ use super::types::{SummaryNode, TreeKind}; use crate::memory::chunks::{tree_active_signature, with_connection}; use crate::memory::config::MemoryConfig; use crate::memory::score::embed::decode_optional_blob; +use crate::memory::store::content::StagedSummary; /// Insert a sealed summary. Immutable; the caller mints a fresh id per seal. /// Idempotent on the primary key (`INSERT OR IGNORE`). @@ -27,8 +28,28 @@ pub fn insert_summary_tx( tx: &Transaction<'_>, node: &SummaryNode, model_signature: &str, +) -> Result<()> { + insert_staged_summary_tx(tx, node, None, model_signature) +} + +/// Insert a summary and optional staged Markdown pointer in one transaction. +pub fn insert_staged_summary_tx( + tx: &Transaction<'_>, + node: &SummaryNode, + staged: Option<&StagedSummary>, + model_signature: &str, ) -> Result<()> { let embedding_blob: Option> = None; + let (content, content_path, content_sha256) = staged.map_or_else( + || (node.content.clone(), None, None), + |staged| { + ( + node.content.chars().take(500).collect(), + Some(staged.content_path.clone()), + Some(staged.content_sha256.clone()), + ) + }, + ); tx.execute( "INSERT OR IGNORE INTO mem_tree_summaries ( id, tree_id, tree_kind, level, parent_id, @@ -36,8 +57,8 @@ pub fn insert_summary_tx( entities_json, topics_json, time_range_start_ms, time_range_end_ms, score, sealed_at_ms, deleted, embedding, - doc_id, version_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)", + doc_id, version_ms, content_path, content_sha256 + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20)", params![ node.id, node.tree_id, @@ -45,7 +66,7 @@ pub fn insert_summary_tx( node.level, node.parent_id, serde_json::to_string(&node.child_ids)?, - node.content, + content, node.token_count, serde_json::to_string(&node.entities)?, serde_json::to_string(&node.topics)?, @@ -57,6 +78,8 @@ pub fn insert_summary_tx( embedding_blob, node.doc_id, node.version_ms, + content_path, + content_sha256, ], ) .with_context(|| format!("Failed to insert summary id={}", node.id))?; diff --git a/src/memory/tree/store/trees.rs b/src/memory/tree/store/trees.rs index 470855d..aa39b7f 100644 --- a/src/memory/tree/store/trees.rs +++ b/src/memory/tree/store/trees.rs @@ -17,7 +17,7 @@ pub fn insert_tree(config: &MemoryConfig, tree: &Tree) -> Result<()> { with_connection(config, |conn| insert_tree_conn(conn, tree)) } -pub(crate) fn insert_tree_conn(conn: &Connection, tree: &Tree) -> Result<()> { +pub fn insert_tree_conn(conn: &Connection, tree: &Tree) -> Result<()> { conn.execute( "INSERT INTO mem_tree_trees ( id, kind, scope, root_id, max_level, status, @@ -38,6 +38,48 @@ pub(crate) fn insert_tree_conn(conn: &Connection, tree: &Tree) -> Result<()> { Ok(()) } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TreeCascadeDeletion { + pub removed_summaries: usize, + pub content_paths: Vec, +} + +/// Delete a tree and all dependent rows inside a caller-owned transaction. +pub fn delete_tree_cascade_tx(tx: &Transaction<'_>, tree_id: &str) -> Result { + let content_paths = { + let mut stmt = tx.prepare( + "SELECT content_path FROM mem_tree_summaries + WHERE tree_id = ?1 AND content_path IS NOT NULL AND content_path <> ''", + )?; + let paths = stmt + .query_map(params![tree_id], |row| row.get::<_, String>(0))? + .collect::>>()?; + paths + }; + for sql in [ + "DELETE FROM mem_tree_summary_embeddings WHERE summary_id IN + (SELECT id FROM mem_tree_summaries WHERE tree_id = ?1)", + "DELETE FROM mem_tree_summary_reembed_skipped WHERE summary_id IN + (SELECT id FROM mem_tree_summaries WHERE tree_id = ?1)", + "DELETE FROM mem_tree_entity_index WHERE tree_id = ?1", + ] { + tx.execute(sql, params![tree_id])?; + } + let removed_summaries = tx.execute( + "DELETE FROM mem_tree_summaries WHERE tree_id = ?1", + params![tree_id], + )?; + tx.execute( + "DELETE FROM mem_tree_buffers WHERE tree_id = ?1", + params![tree_id], + )?; + tx.execute("DELETE FROM mem_tree_trees WHERE id = ?1", params![tree_id])?; + Ok(TreeCascadeDeletion { + removed_summaries, + content_paths, + }) +} + /// Fetch a tree by `(kind, scope)`. Returns `None` if no such tree exists. pub fn get_tree_by_scope( config: &MemoryConfig, @@ -47,7 +89,7 @@ pub fn get_tree_by_scope( with_connection(config, |conn| get_tree_by_scope_conn(conn, kind, scope)) } -pub(crate) fn get_tree_by_scope_conn( +pub fn get_tree_by_scope_conn( conn: &Connection, kind: TreeKind, scope: &str, @@ -141,7 +183,7 @@ pub fn list_trees_by_kind(config: &MemoryConfig, kind: TreeKind) -> Result, tree_id: &str, root_id: &str, @@ -161,7 +203,7 @@ pub(crate) fn update_tree_after_seal_tx( } /// Refresh `last_sealed_at` without changing the root (same-level seal). -pub(crate) fn refresh_last_sealed_tx( +pub fn refresh_last_sealed_tx( tx: &Transaction<'_>, tree_id: &str, sealed_at: DateTime, diff --git a/src/memory/tree/store/types_tests.rs b/src/memory/tree/store/types_tests.rs index 9ac6099..f9dd234 100644 --- a/src/memory/tree/store/types_tests.rs +++ b/src/memory/tree/store/types_tests.rs @@ -45,8 +45,8 @@ fn budgets_match_config_defaults() { assert_eq!(INPUT_TOKEN_BUDGET, 50_000); assert_eq!(OUTPUT_TOKEN_BUDGET, 5_000); assert_eq!(SUMMARY_FANOUT, 10); - assert!(TOPIC_CREATION_THRESHOLD > TOPIC_ARCHIVE_THRESHOLD); - assert!(TOPIC_RECHECK_EVERY > 0); + const { assert!(TOPIC_CREATION_THRESHOLD > TOPIC_ARCHIVE_THRESHOLD) }; + const { assert!(TOPIC_RECHECK_EVERY > 0) }; } #[test] diff --git a/src/memory/tree/summarise.rs b/src/memory/tree/summarise.rs index dab75d7..3aa2bee 100644 --- a/src/memory/tree/summarise.rs +++ b/src/memory/tree/summarise.rs @@ -14,6 +14,10 @@ use chrono::{DateTime, Utc}; use crate::memory::chunks::approx_token_count; use crate::memory::tree::store::TreeKind; +const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 5_000; +const NUM_CTX_TOKENS: u32 = 60_000; +const OVERHEAD_RESERVE_TOKENS: u32 = 2_048; + /// One contribution being folded — a raw leaf at L0→L1, or a lower-level /// summary at L_n→L_{n+1}. #[derive(Clone, Debug)] @@ -64,6 +68,82 @@ pub struct SummaryOutput { pub topics: Vec, } +#[derive(Clone, Debug, Default)] +pub struct SummaryCall { + pub output: SummaryOutput, + pub input_tokens: u64, + pub output_tokens: u64, + pub charged_amount_usd: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PreparedSummaryPrompt { + pub system: String, + pub user: String, + pub effective_budget: u32, +} + +/// Build the canonical provider prompt for a summary fold. Returns `None` +/// when every input is blank. +pub fn prepare_summary_prompt( + inputs: &[SummaryInput], + ctx: &SummaryContext<'_>, + output_language: Option<&str>, +) -> Option { + let effective_budget = ctx.token_budget.min(MAX_SUMMARY_OUTPUT_TOKENS); + let per_input_cap = if inputs.is_empty() { + 0 + } else { + NUM_CTX_TOKENS + .saturating_sub(effective_budget) + .saturating_sub(OVERHEAD_RESERVE_TOKENS) + / inputs.len() as u32 + }; + let mut ordered: Vec<_> = inputs.iter().collect(); + ordered.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + let user = ordered + .into_iter() + .filter_map(|input| { + let content = input.content.trim(); + (!content.is_empty()).then(|| { + let (content, _) = clamp_to_budget(content, per_input_cap); + format!("[{}]\n{content}", input.id) + }) + }) + .collect::>() + .join("\n\n"); + if user.trim().is_empty() { + return None; + } + let language = output_language + .filter(|language| !language.trim().is_empty()) + .map(|language| format!("\nWrite the summary in {language}.")) + .unwrap_or_default(); + Some(PreparedSummaryPrompt { + system: format!( + "You are folding multiple notes into one compact summary.\n\ + Aim for ~{effective_budget} tokens or fewer. Capture key facts, decisions, and entities.\n\ + Output only the summary prose — no preamble, no JSON, no markdown headings.{language}" + ), + user, + effective_budget, + }) +} + +pub fn finish_provider_summary(text: &str, budget: u32) -> SummaryOutput { + let (content, token_count) = clamp_to_budget(text.trim(), budget); + SummaryOutput { + content, + token_count, + entities: Vec::new(), + topics: Vec::new(), + } +} + /// Backend that folds inputs into one summary. Abstracted so the crate never /// calls a real LLM; the default is the deterministic [`ConcatSummariser`]. #[async_trait] @@ -81,6 +161,17 @@ pub trait Summariser: Send + Sync { inputs: &[SummaryInput], ctx: &SummaryContext<'_>, ) -> Result; + + async fn summarise_with_usage( + &self, + inputs: &[SummaryInput], + ctx: &SummaryContext<'_>, + ) -> Result { + Ok(SummaryCall { + output: self.summarise(inputs, ctx).await?, + ..SummaryCall::default() + }) + } } /// Deterministic, dependency-free summariser: concatenate inputs (priority-first diff --git a/src/memory/tree/summarise_tests.rs b/src/memory/tree/summarise_tests.rs index b57f428..6907b1e 100644 --- a/src/memory/tree/summarise_tests.rs +++ b/src/memory/tree/summarise_tests.rs @@ -59,3 +59,22 @@ async fn concat_summariser_matches_fallback() { assert!(out.content.contains("second")); assert_eq!(out.content, fallback_summary(&inputs, 10_000).content); } + +#[test] +fn provider_prompt_is_priority_ordered_language_aware_and_budgeted() { + let mut low = sample_input("low", "low priority"); + low.score = 0.1; + let mut high = sample_input("high", "high priority"); + high.score = 0.9; + let context = SummaryContext { + tree_id: "tree:test", + tree_kind: TreeKind::Source, + target_level: 1, + token_budget: 9_000, + }; + let prompt = prepare_summary_prompt(&[low, high], &context, Some("French")).unwrap(); + assert!(prompt.user.starts_with("[high]")); + assert!(prompt.system.contains("Write the summary in French")); + assert_eq!(prompt.effective_budget, 5_000); + assert!(prepare_summary_prompt(&[], &context, None).is_none()); +} diff --git a/tests/composio_sync_live.rs b/tests/composio_sync_live.rs new file mode 100644 index 0000000..2ac9c66 --- /dev/null +++ b/tests/composio_sync_live.rs @@ -0,0 +1,26 @@ +use tinycortex::memory::config::{ComposioMode, ComposioSyncConfig, SecretString}; +use tinycortex::memory::sync::ComposioClient; + +#[tokio::test] +#[ignore = "requires COMPOSIO_API_KEY and a connected Gmail account"] +async fn direct_gmail_sync_live_smoke() { + let key = std::env::var("COMPOSIO_API_KEY").expect("COMPOSIO_API_KEY"); + let connection_id = std::env::var("COMPOSIO_GMAIL_CONNECTION_ID").ok(); + let client = ComposioClient::new(ComposioSyncConfig { + mode: ComposioMode::Direct, + base_url: "https://backend.composio.dev/api/v3".into(), + api_key: Some(SecretString::new(key)), + bearer_token: None, + entity_id: std::env::var("COMPOSIO_ENTITY_ID").ok(), + }); + let response = client + .execute( + "GMAIL_FETCH_EMAILS", + serde_json::json!({"max_results": 5, "include_payload": true}), + connection_id.as_deref(), + ) + .await + .unwrap(); + assert!(response.successful, "{:?}", response.error); + assert!(response.data.is_object()); +} diff --git a/tests/composio_sync_mock.rs b/tests/composio_sync_mock.rs new file mode 100644 index 0000000..ec3362b --- /dev/null +++ b/tests/composio_sync_mock.rs @@ -0,0 +1,631 @@ +use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Mutex; + +use async_trait::async_trait; +use serde_json::Value; +use tinycortex::memory::config::{ComposioMode, ComposioSyncConfig, MemoryConfig, SecretString}; +use tinycortex::memory::sync::{ + ClickUpSyncPipeline, ComposioClient, GitHubSyncPipeline, GmailSyncPipeline, LinearSyncPipeline, + NotionSyncPipeline, SkillDocSink, SkillDocument, SlackSearchBackfillPipeline, + SlackSyncPipeline, SyncContext, SyncEvent, SyncEventSink, SyncPipeline, SyncStage, SyncState, + SyncStateStore, +}; +use wiremock::matchers::{body_partial_json, header, method, path, path_regex}; +use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; + +#[derive(Default)] +struct Captures { + documents: Mutex>, + events: Mutex>, + state: Mutex>, +} + +#[async_trait] +impl SkillDocSink for Captures { + async fn store(&self, document: SkillDocument) -> anyhow::Result<()> { + self.documents.lock().unwrap().push(document); + Ok(()) + } + + async fn delete(&self, namespace_skill_id: &str, document_id: &str) -> anyhow::Result<()> { + self.documents.lock().unwrap().retain(|document| { + document.namespace_skill_id != namespace_skill_id || document.document_id != document_id + }); + Ok(()) + } +} + +#[async_trait] +impl SyncEventSink for Captures { + async fn emit(&self, event: SyncEvent) -> anyhow::Result<()> { + self.events.lock().unwrap().push(event); + Ok(()) + } +} + +#[async_trait] +impl SyncStateStore for Captures { + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { + Ok(self + .state + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.state + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } +} + +struct GmailPages; + +impl Respond for GmailPages { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let token = body + .pointer("/arguments/page_token") + .and_then(|v| v.as_str()); + let data = if token == Some("page-2") { + serde_json::json!({ + "successful": true, + "data": {"messages": [ + {"id": "m2", "internalDate": "1700000002000", "subject": "Second"} + ]} + }) + } else { + serde_json::json!({ + "successful": true, + "data": { + "messages": [ + {"id": "m1", "internalDate": "1700000001000", "subject": "First"} + ], + "nextPageToken": "page-2" + } + }) + }; + ResponseTemplate::new(200).set_body_json(data) + } +} + +fn direct_config(base_url: String, key: &str) -> ComposioSyncConfig { + ComposioSyncConfig { + mode: ComposioMode::Direct, + base_url, + api_key: Some(SecretString::new(key)), + bearer_token: None, + entity_id: Some("entity-1".into()), + } +} + +fn test_context() -> (std::sync::Arc, SyncContext) { + let captures = std::sync::Arc::new(Captures::default()); + let context = SyncContext { + events: captures.clone(), + documents: captures.clone(), + state: captures.clone(), + local_documents: None, + external_sources: None, + summariser: None, + }; + (captures, context) +} + +fn test_config() -> MemoryConfig { + MemoryConfig::new("/tmp/tinycortex-sync-test") +} + +#[tokio::test] +async fn gmail_sync_paginates_persists_cursor_taint_and_is_idempotent() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"/tools/execute/GMAIL_FETCH_EMAILS$")) + .and(header("x-api-key", "test-secret")) + .respond_with(GmailPages) + .mount(&server) + .await; + + let captures = std::sync::Arc::new(Captures::default()); + let context = SyncContext { + events: captures.clone(), + documents: captures.clone(), + state: captures.clone(), + local_documents: None, + external_sources: None, + summariser: None, + }; + let pipeline = GmailSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "test-secret")), + "conn-1", + ) + .with_limits(3, 10); + let config = MemoryConfig::new("/tmp/tinycortex-sync-test"); + + let first = pipeline.tick(&config, &context).await.unwrap(); + assert_eq!(first.records_ingested, 2); + assert_eq!(first.actions_called, 2); + let docs = captures.documents.lock().unwrap(); + assert_eq!(docs.len(), 2); + assert!(docs + .iter() + .all(|doc| doc.metadata["taint"] == "external_sync")); + drop(docs); + + let state = SyncState::load(captures.as_ref(), "gmail", "conn-1") + .await + .unwrap(); + assert_eq!(state.cursor.as_deref(), Some("1700000002000")); + assert!(state.is_synced("m1")); + assert!(state.is_synced("m2")); + + let second = pipeline.tick(&config, &context).await.unwrap(); + assert_eq!(second.records_ingested, 0); + assert_eq!(captures.documents.lock().unwrap().len(), 2); + let events = captures.events.lock().unwrap(); + assert!(events + .iter() + .any(|event| event.stage == SyncStage::Fetching)); + assert!(events + .iter() + .any(|event| event.stage == SyncStage::Completed)); +} + +#[tokio::test] +async fn max_items_stops_before_fetching_another_page() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/GMAIL_FETCH_EMAILS")) + .respond_with(GmailPages) + .expect(1) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = GmailSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "capped-conn", + ); + let mut config = test_config(); + config.sync.budget.max_items = Some(1); + let outcome = pipeline.tick(&config, &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 1); + assert_eq!(outcome.actions_called, 1); + assert!(outcome.more_pending); + let state = SyncState::load(captures.as_ref(), "gmail", "capped-conn") + .await + .unwrap(); + assert_eq!( + state.cursor, None, + "capped runs must not advance the cursor" + ); +} + +#[tokio::test] +async fn direct_401_never_leaks_api_key() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(401).set_body_string("unauthorized test-secret")) + .mount(&server) + .await; + let client = ComposioClient::new(direct_config(server.uri(), "test-secret")); + let error = client + .execute("GMAIL_FETCH_EMAILS", serde_json::json!({}), Some("conn")) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("401")); + assert!(!error.contains("test-secret")); +} + +#[tokio::test] +async fn proxied_mode_uses_bearer_and_backend_execute_envelope() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/agent-integrations/composio/execute")) + .and(header("authorization", "Bearer proxy-secret")) + .and(body_partial_json(serde_json::json!({ + "tool": "GMAIL_FETCH_EMAILS" + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"messages": []}, + "costUsd": 0.01 + }))) + .mount(&server) + .await; + let client = ComposioClient::new(ComposioSyncConfig { + mode: ComposioMode::Proxied, + base_url: server.uri(), + api_key: None, + bearer_token: Some(SecretString::new("proxy-secret")), + entity_id: None, + }); + let response = client + .execute("GMAIL_FETCH_EMAILS", serde_json::json!({}), None) + .await + .unwrap(); + assert!(response.successful); + assert_eq!(response.cost_usd, 0.01); +} + +#[tokio::test] +async fn github_resolves_identity_and_stores_timestamped_issue() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/tools/execute/GITHUB_GET_THE_AUTHENTICATED_USER")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, "data": {"login": "alice"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/tools/execute/GITHUB_SEARCH_ISSUES_AND_PULL_REQUESTS")) + .and(body_partial_json(serde_json::json!({"arguments": {"q": "involves:alice"}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"items": [{"id": 42, "title": "Fix sync", "updated_at": "2026-01-02T03:04:05Z"}]} + }))) + .expect(1) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = GitHubSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "github-conn", + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 1); + assert_eq!(outcome.actions_called, 2); + assert_eq!( + captures.documents.lock().unwrap()[0].document_id, + "github:42" + ); + let state = SyncState::load(captures.as_ref(), "github", "github-conn") + .await + .unwrap(); + assert!(state.is_synced("42@2026-01-02T03:04:05Z")); +} + +#[tokio::test] +async fn slack_search_backfill_enriches_and_paginates_workspace_messages() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/SLACK_LIST_ALL_USERS")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"members": [{"id": "U1", "profile": {"display_name": "Alice"}}]} + }))) + .mount(&server) + .await; + Mock::given(path("/tools/execute/SLACK_LIST_CONVERSATIONS")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"channels": [{"id": "C1", "name": "general", "is_private": false}]} + }))) + .mount(&server) + .await; + Mock::given(path("/tools/execute/SLACK_SEARCH_MESSAGES")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"messages": {"matches": [{ + "ts": "1700000000.000001", + "user": "U1", + "text": "hello <@U1>", + "channel": {"id": "C1"} + }], "paging": {"pages": 1}}} + }))) + .expect(1) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = SlackSearchBackfillPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "slack-search-conn", + 30, + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 1); + assert_eq!(outcome.actions_called, 3); + let documents = captures.documents.lock().unwrap(); + assert_eq!(documents[0].title, "Slack #general from Alice"); + assert!(documents[0].content.contains("Alice: hello @Alice")); +} + +struct LinearPages; + +impl Respond for LinearPages { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let second = body.pointer("/arguments/after").is_some(); + let data = if second { + serde_json::json!({"nodes": [{"id": "L-2", "title": "Second", "updatedAt": "2026-02-02T00:00:00Z"}], "pageInfo": {"hasNextPage": false}}) + } else { + serde_json::json!({"nodes": [{"id": "L-1", "title": "First", "updatedAt": "2026-02-01T00:00:00Z"}], "pageInfo": {"hasNextPage": true, "endCursor": "linear-page-2"}}) + }; + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"successful": true, "data": data})) + } +} + +#[tokio::test] +async fn linear_resolves_viewer_and_follows_graphql_cursor() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/LINEAR_LIST_LINEAR_USERS")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"successful": true, "data": {"nodes": [{"id": "viewer-1"}]}}), + )) + .mount(&server) + .await; + Mock::given(path("/tools/execute/LINEAR_LIST_LINEAR_ISSUES")) + .respond_with(LinearPages) + .expect(2) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = LinearSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "linear-conn", + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 2); + assert_eq!(captures.documents.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn notion_fetches_markdown_and_counts_both_requests() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/NOTION_FETCH_DATA")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"successful": true, "data": {"results": [{"id": "page-1", "title": "Roadmap", "last_edited_time": "2026-03-01T00:00:00Z"}]}}))) + .mount(&server).await; + Mock::given(path("/tools/execute/NOTION_GET_PAGE_MARKDOWN")) + .and(body_partial_json( + serde_json::json!({"arguments": {"page_id": "page-1"}}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"successful": true, "data": {"markdown": "# Roadmap\n\nBody"}}), + )) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = NotionSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "notion-conn", + ); + pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!( + captures.documents.lock().unwrap()[0].content, + "# Roadmap\n\nBody" + ); + let state = SyncState::load(captures.as_ref(), "notion", "notion-conn") + .await + .unwrap(); + assert_eq!(state.daily_budget.requests_used, 2); +} + +#[tokio::test] +async fn clickup_pages_each_workspace_with_resolved_user() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/CLICKUP_GET_AUTHORIZED_USER")) + .respond_with( + ResponseTemplate::new(200).set_body_json( + serde_json::json!({"successful": true, "data": {"user": {"id": 77}}}), + ), + ) + .mount(&server) + .await; + Mock::given(path("/tools/execute/CLICKUP_GET_AUTHORIZED_TEAMS_WORKSPACES")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"successful": true, "data": {"teams": [{"id": "ws-1"}, {"id": "ws-2"}]}}))) + .mount(&server).await; + Mock::given(path("/tools/execute/CLICKUP_GET_FILTERED_TEAM_TASKS")) + .and(body_partial_json(serde_json::json!({"arguments": {"assignees": ["77"]}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"successful": true, "data": {"tasks": [{"id": "task", "name": "Scoped", "date_updated": "1700000000000"}]}}))) + .expect(2).mount(&server).await; + let (captures, context) = test_context(); + let pipeline = ClickUpSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "clickup-conn", + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!( + outcome.records_ingested, 1, + "global dedup suppresses the duplicate task returned by both workspaces" + ); + assert_eq!( + captures.documents.lock().unwrap()[0].metadata["workspace_id"], + "ws-1" + ); +} + +struct SlackHistory; + +impl Respond for SlackHistory { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + if body.pointer("/arguments/channel").and_then(Value::as_str) == Some("C_BAD") { + ResponseTemplate::new(200).set_body_json( + serde_json::json!({"successful": false, "error": "channel unavailable"}), + ) + } else { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"messages": [{"ts": "1760000000.000123", "user": "U1", "text": "healthy channel"}]} + })) + } + } +} + +#[tokio::test] +async fn slack_holds_failed_channel_cursor_and_advances_healthy_channel() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/SLACK_LIST_ALL_USERS")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"members": [{"id": "U1", "profile": {"display_name": "Alice"}}]} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/tools/execute/SLACK_LIST_CONVERSATIONS")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"channels": [ + {"id": "C_BAD", "name": "broken", "is_private": false}, + {"id": "C_GOOD", "name": "general", "is_private": false} + ]} + }))) + .mount(&server) + .await; + Mock::given(path("/tools/execute/SLACK_FETCH_CONVERSATION_HISTORY")) + .respond_with(SlackHistory) + .expect(2) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = SlackSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "slack-conn", + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 1); + assert_eq!(outcome.actions_called, 4); + let state = SyncState::load(captures.as_ref(), "slack", "slack-conn") + .await + .unwrap(); + let cursors: Value = serde_json::from_str(state.cursor.as_deref().unwrap()).unwrap(); + assert!(cursors.get("C_BAD").is_none()); + assert_eq!(cursors["C_GOOD"], "1760000000.000123"); + assert_eq!(state.synced_ids.len(), 1); + assert_eq!( + captures.documents.lock().unwrap()[0].metadata["channel_id"], + "C_GOOD" + ); + assert!(captures.documents.lock().unwrap()[0] + .content + .contains("Alice")); +} + +#[tokio::test] +async fn configured_depth_is_server_side_for_github_and_client_side_for_linear() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/GITHUB_GET_THE_AUTHENTICATED_USER")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"successful": true, "data": {"login": "alice"}})), + ) + .mount(&server) + .await; + Mock::given(path( + "/tools/execute/GITHUB_SEARCH_ISSUES_AND_PULL_REQUESTS", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"successful": true, "data": {"items": []}})), + ) + .mount(&server) + .await; + Mock::given(path("/tools/execute/LINEAR_LIST_LINEAR_USERS")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"successful": true, "data": {"nodes": [{"id": "viewer"}]}}), + )) + .mount(&server) + .await; + let recent = (chrono::Utc::now() - chrono::Duration::days(1)) + .format("%Y-%m-%dT%H:%M:%SZ") + .to_string(); + let old = (chrono::Utc::now() - chrono::Duration::days(60)) + .format("%Y-%m-%dT%H:%M:%SZ") + .to_string(); + Mock::given(path("/tools/execute/LINEAR_LIST_LINEAR_ISSUES")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"nodes": [ + {"id": "recent", "title": "Recent", "updatedAt": recent}, + {"id": "old", "title": "Old", "updatedAt": old} + ], "pageInfo": {"hasNextPage": false}} + }))) + .mount(&server) + .await; + let mut config = test_config(); + config.sync.budget.sync_depth_days = Some(30); + let (captures, context) = test_context(); + GitHubSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "github-depth", + ) + .tick(&config, &context) + .await + .unwrap(); + let linear = LinearSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "linear-depth", + ) + .tick(&config, &context) + .await + .unwrap(); + assert_eq!(linear.records_ingested, 1); + assert_eq!(captures.documents.lock().unwrap().len(), 1); + + let requests = server.received_requests().await.unwrap(); + let request = requests + .iter() + .find(|request| { + request + .url + .path() + .ends_with("GITHUB_SEARCH_ISSUES_AND_PULL_REQUESTS") + }) + .unwrap(); + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let query = body + .pointer("/arguments/q") + .and_then(Value::as_str) + .unwrap(); + assert!(query.starts_with("involves:alice updated:>")); +} + +struct RetryThenGmail(AtomicUsize); + +impl Respond for RetryThenGmail { + fn respond(&self, _: &Request) -> ResponseTemplate { + let attempt = self.0.fetch_add(1, Ordering::SeqCst); + if attempt < 2 { + ResponseTemplate::new(429) + } else { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "successful": true, + "data": {"messages": [{"id": "retry-message", "internalDate": "1700000000000", "subject": "Recovered"}]} + })) + } + } +} + +#[tokio::test] +async fn transient_retries_are_counted_in_persisted_budget() { + let server = MockServer::start().await; + Mock::given(path("/tools/execute/GMAIL_FETCH_EMAILS")) + .respond_with(RetryThenGmail(AtomicUsize::new(0))) + .expect(3) + .mount(&server) + .await; + let (captures, context) = test_context(); + let pipeline = GmailSyncPipeline::new( + ComposioClient::new(direct_config(server.uri(), "key")), + "retry-conn", + ); + let outcome = pipeline.tick(&test_config(), &context).await.unwrap(); + assert_eq!(outcome.records_ingested, 1); + assert_eq!(outcome.actions_called, 3); + let state = SyncState::load(captures.as_ref(), "gmail", "retry-conn") + .await + .unwrap(); + assert_eq!(state.daily_budget.requests_used, 3); +}