From 96c370314517ed9946f59b3af713637347c70dab Mon Sep 17 00:00:00 2001 From: Steven Enamakel <31011319+senamakel@users.noreply.github.com> Date: Mon, 13 Jul 2026 02:38:03 -0700 Subject: [PATCH] refactor(memory): complete TinyCortex engine migration (#4794) --- Cargo.lock | 9 +- Cargo.toml | 8 +- app/src-tauri/Cargo.lock | 9 +- app/src-tauri/Cargo.toml | 7 +- src/bin/gmail_backfill_3d.rs | 425 +--- src/bin/slack_backfill.rs | 64 +- src/core/observability.rs | 4 +- .../agent/harness/session/builder/factory.rs | 45 +- .../agent/harness/session/builder/helpers.rs | 2 +- src/openhuman/composio/identity.rs | 15 +- src/openhuman/composio/ops/memory_cleanup.rs | 5 +- src/openhuman/composio/ops/providers_ops.rs | 33 +- src/openhuman/composio/ops_tests.rs | 83 +- src/openhuman/composio/tools_tests.rs | 9 - src/openhuman/embeddings/cloud.rs | 213 -- src/openhuman/embeddings/cloud_adapter.rs | 88 + src/openhuman/embeddings/cohere.rs | 393 --- src/openhuman/embeddings/cohere_adapter.rs | 51 + src/openhuman/embeddings/factory.rs | 117 +- src/openhuman/embeddings/mod.rs | 30 +- src/openhuman/embeddings/ollama.rs | 470 ---- src/openhuman/embeddings/ollama_adapter.rs | 68 + src/openhuman/embeddings/ollama_tests.rs | 635 ----- src/openhuman/embeddings/openai.rs | 467 ---- src/openhuman/embeddings/openai_adapter.rs | 71 + src/openhuman/embeddings/openai_tests.rs | 918 ------- src/openhuman/embeddings/provider_trait.rs | 48 + src/openhuman/embeddings/rate_limit.rs | 387 --- src/openhuman/embeddings/retry_after.rs | 166 -- src/openhuman/embeddings/rpc.rs | 16 +- src/openhuman/embeddings/voyage.rs | 100 - src/openhuman/embeddings/voyage_adapter.rs | 69 + src/openhuman/memory/ingest_pipeline.rs | 817 +------ src/openhuman/memory/ingestion/mod.rs | 71 +- src/openhuman/memory/ingestion/parse.rs | 942 ------- src/openhuman/memory/ingestion/parse_tests.rs | 137 -- src/openhuman/memory/ingestion/regex.rs | 219 -- src/openhuman/memory/ingestion/rules.rs | 325 --- src/openhuman/memory/ingestion/types.rs | 287 --- src/openhuman/memory/ops/sync.rs | 25 +- src/openhuman/memory/read_rpc/admin.rs | 4 +- src/openhuman/memory/read_rpc_tests.rs | 47 +- .../memory/sync_pipeline_e2e_tests.rs | 24 +- src/openhuman/memory/traits.rs | 135 +- src/openhuman/memory_queue/mod.rs | 1 - src/openhuman/memory_queue/ops.rs | 59 +- src/openhuman/memory_queue/redact.rs | 226 -- src/openhuman/memory_queue/scheduler.rs | 33 +- src/openhuman/memory_queue/store.rs | 1376 +---------- src/openhuman/memory_queue/types.rs | 648 +---- .../memory_sources/readers/conversation.rs | 368 +-- .../memory_sources/readers/folder.rs | 221 +- src/openhuman/memory_sources/registry.rs | 576 +---- src/openhuman/memory_sources/sync.rs | 617 +---- src/openhuman/memory_sources/types.rs | 349 +-- .../memory_store/chunks/connection.rs | 842 +------ .../memory_store/chunks/embeddings.rs | 421 +--- .../memory_store/chunks/migrations.rs | 201 -- src/openhuman/memory_store/chunks/produce.rs | 1102 +-------- src/openhuman/memory_store/chunks/semantic.rs | 636 +---- src/openhuman/memory_store/chunks/store.rs | 1217 +-------- .../memory_store/chunks/store_tests.rs | 2175 ----------------- src/openhuman/memory_store/chunks/types.rs | 99 +- src/openhuman/memory_store/content/README.md | 8 +- src/openhuman/memory_store/content/atomic.rs | 546 ----- .../memory_store/content/compose/chunk.rs | 209 -- .../memory_store/content/compose/mod.rs | 53 - .../memory_store/content/compose/summary.rs | 298 --- .../memory_store/content/compose/tests.rs | 622 ----- .../memory_store/content/compose/yaml.rs | 149 -- src/openhuman/memory_store/content/mod.rs | 196 +- src/openhuman/memory_store/content/paths.rs | 744 ------ src/openhuman/memory_store/content/raw.rs | 477 ---- src/openhuman/memory_store/content/read.rs | 914 +------ src/openhuman/memory_store/entities.rs | 120 +- src/openhuman/memory_store/factories.rs | 56 +- src/openhuman/memory_store/kv.rs | 307 +-- src/openhuman/memory_store/memory_trait.rs | 115 +- src/openhuman/memory_store/safety/mod.rs | 357 +-- src/openhuman/memory_store/safety/pii.rs | 1095 +-------- src/openhuman/memory_store/trees/hotness.rs | 173 +- src/openhuman/memory_store/trees/registry.rs | 163 +- src/openhuman/memory_store/trees/store.rs | 993 +------- src/openhuman/memory_store/trees/types.rs | 387 +-- src/openhuman/memory_store/types.rs | 383 +-- src/openhuman/memory_store/vectors/mod.rs | 12 +- src/openhuman/memory_store/vectors/store.rs | 521 ---- .../memory_store/vectors/store_tests.rs | 481 ---- .../memory_sync/canonicalize/README.md | 8 +- .../memory_sync/canonicalize/chat.rs | 262 -- .../memory_sync/canonicalize/document.rs | 229 -- .../memory_sync/canonicalize/email.rs | 300 --- .../memory_sync/canonicalize/email_clean.rs | 423 ---- src/openhuman/memory_sync/canonicalize/mod.rs | 111 +- src/openhuman/memory_sync/composio/bus.rs | 33 +- src/openhuman/memory_sync/composio/mod.rs | 64 +- .../memory_sync/composio/periodic.rs | 83 +- .../composio/providers/clickup/ingest.rs | 333 --- .../composio/providers/clickup/mod.rs | 2 - .../composio/providers/clickup/provider.rs | 14 +- .../composio/providers/clickup/source.rs | 569 ----- .../composio/providers/github/ingest.rs | 524 ---- .../composio/providers/github/mod.rs | 2 - .../composio/providers/github/provider.rs | 9 +- .../composio/providers/github/source.rs | 325 --- .../composio/providers/gmail/ingest.rs | 908 ------- .../composio/providers/gmail/mod.rs | 2 - .../composio/providers/gmail/provider.rs | 24 +- .../composio/providers/gmail/source.rs | 363 --- .../composio/providers/linear/ingest.rs | 316 --- .../composio/providers/linear/mod.rs | 2 - .../composio/providers/linear/provider.rs | 11 +- .../composio/providers/linear/source.rs | 397 --- .../memory_sync/composio/providers/mod.rs | 1 - .../composio/providers/notion/ingest.rs | 663 ----- .../composio/providers/notion/mod.rs | 2 - .../composio/providers/notion/provider.rs | 27 +- .../composio/providers/notion/source.rs | 550 ----- .../composio/providers/orchestrator.rs | 1545 ------------ .../composio/providers/registry.rs | 9 +- .../composio/providers/slack/ingest.rs | 402 --- .../composio/providers/slack/mod.rs | 10 +- .../composio/providers/slack/provider.rs | 477 +--- .../composio/providers/slack/rpc.rs | 76 +- .../composio/providers/slack/source.rs | 422 ---- .../composio/providers/slack/sync.rs | 473 ---- .../composio/providers/slack/types.rs | 8 +- .../composio/providers/slack/users.rs | 364 --- .../composio/providers/sync_state.rs | 614 +---- .../memory_sync/composio/providers/traits.rs | 56 +- src/openhuman/memory_sync/sources/audit.rs | 474 +--- src/openhuman/memory_sync/sources/github.rs | 863 +------ src/openhuman/memory_sync/sources/mod.rs | 8 +- src/openhuman/memory_sync/sources/rebuild.rs | 836 +------ src/openhuman/memory_sync/sync_status/rpc.rs | 478 +--- .../memory_sync/sync_status/types.rs | 127 +- src/openhuman/memory_tools/README.md | 2 +- src/openhuman/memory_tools/store.rs | 150 +- src/openhuman/memory_tools/tools/list.rs | 2 +- src/openhuman/memory_tools/tools/put.rs | 2 +- src/openhuman/memory_tree/graph/bfs.rs | 210 +- src/openhuman/memory_tree/graph/store.rs | 219 +- src/openhuman/memory_tree/ingest.rs | 394 +-- src/openhuman/memory_tree/io.rs | 235 +- src/openhuman/memory_tree/retrieval/cover.rs | 590 +---- .../memory_tree/retrieval/drill_down.rs | 995 +------- src/openhuman/memory_tree/retrieval/engine.rs | 26 + src/openhuman/memory_tree/retrieval/fast.rs | 457 +--- src/openhuman/memory_tree/retrieval/fetch.rs | 305 +-- .../retrieval/integration_tests.rs | 17 +- src/openhuman/memory_tree/retrieval/mod.rs | 1 + src/openhuman/memory_tree/retrieval/search.rs | 313 +-- src/openhuman/memory_tree/retrieval/source.rs | 951 +------ src/openhuman/memory_tree/retrieval/types.rs | 267 +- .../memory_tree/score/extract/extractor.rs | 111 - .../memory_tree/score/extract/llm.rs | 661 ----- .../memory_tree/score/extract/llm_tests.rs | 756 ------ .../memory_tree/score/extract/mod.rs | 94 +- .../memory_tree/score/extract/regex.rs | 215 -- .../memory_tree/score/extract/types.rs | 345 --- src/openhuman/memory_tree/score/mod.rs | 493 +--- src/openhuman/memory_tree/score/resolver.rs | 265 -- .../memory_tree/score/signals/interaction.rs | 105 - .../score/signals/metadata_weight.rs | 49 - .../memory_tree/score/signals/mod.rs | 20 - .../memory_tree/score/signals/ops.rs | 182 -- .../score/signals/source_weight.rs | 112 - .../memory_tree/score/signals/token_count.rs | 84 - .../memory_tree/score/signals/types.rs | 107 - .../memory_tree/score/signals/unique_words.rs | 90 - src/openhuman/memory_tree/score/store.rs | 537 +--- .../memory_tree/score/store_tests.rs | 248 -- src/openhuman/memory_tree/summarise.rs | 441 +--- src/openhuman/memory_tree/tree/bucket_seal.rs | 1551 +----------- .../memory_tree/tree/bucket_seal_tests.rs | 1358 ---------- src/openhuman/memory_tree/tree/factory.rs | 42 +- src/openhuman/memory_tree/tree/flush.rs | 370 +-- .../memory_tree/tree_runtime/engine.rs | 741 +----- .../memory_tree/tree_runtime/engine_tests.rs | 707 ------ .../memory_tree/tree_runtime/store.rs | 870 +------ .../memory_tree/tree_runtime/store_tests.rs | 361 --- .../memory_tree/tree_runtime/types.rs | 298 +-- src/openhuman/task_sources/pipeline_tests.rs | 9 +- src/openhuman/tinycortex/config.rs | 1 + src/openhuman/tinycortex/ingest.rs | 58 + src/openhuman/tinycortex/mod.rs | 20 +- src/openhuman/tinycortex/queue_driver.rs | 19 +- src/openhuman/tinycortex/seal.rs | 234 ++ src/openhuman/tinycortex/summariser.rs | 63 + src/openhuman/tinycortex/sync.rs | 524 ++++ tests/embeddings_rpc_e2e.rs | 11 +- tests/json_rpc_e2e.rs | 7 - tests/memory_artifacts_e2e.rs | 48 +- .../embeddings_ollama_raw_coverage_e2e.rs | 34 +- .../memory_core_threads_raw_coverage_e2e.rs | 5 +- .../memory_sync_providers_raw_coverage_e2e.rs | 344 +-- .../memory_sync_round23_raw_coverage_e2e.rs | 8 +- .../memory_sync_slack_bus_raw_coverage_e2e.rs | 74 +- ...mory_sync_tree_round21_raw_coverage_e2e.rs | 17 +- .../memory_threads_raw_coverage_e2e.rs | 55 +- .../memory_tree_sync_raw_coverage_e2e.rs | 24 +- vendor/tinyagents | 2 +- vendor/tinycortex | 2 +- 203 files changed, 3784 insertions(+), 56527 deletions(-) delete mode 100644 src/openhuman/embeddings/cloud.rs create mode 100644 src/openhuman/embeddings/cloud_adapter.rs delete mode 100644 src/openhuman/embeddings/cohere.rs create mode 100644 src/openhuman/embeddings/cohere_adapter.rs delete mode 100644 src/openhuman/embeddings/ollama.rs create mode 100644 src/openhuman/embeddings/ollama_adapter.rs delete mode 100644 src/openhuman/embeddings/ollama_tests.rs delete mode 100644 src/openhuman/embeddings/openai.rs create mode 100644 src/openhuman/embeddings/openai_adapter.rs delete mode 100644 src/openhuman/embeddings/openai_tests.rs delete mode 100644 src/openhuman/embeddings/rate_limit.rs delete mode 100644 src/openhuman/embeddings/retry_after.rs delete mode 100644 src/openhuman/embeddings/voyage.rs create mode 100644 src/openhuman/embeddings/voyage_adapter.rs delete mode 100644 src/openhuman/memory/ingestion/parse.rs delete mode 100644 src/openhuman/memory/ingestion/parse_tests.rs delete mode 100644 src/openhuman/memory/ingestion/regex.rs delete mode 100644 src/openhuman/memory/ingestion/rules.rs delete mode 100644 src/openhuman/memory/ingestion/types.rs delete mode 100644 src/openhuman/memory_queue/redact.rs delete mode 100644 src/openhuman/memory_store/chunks/migrations.rs delete mode 100644 src/openhuman/memory_store/chunks/store_tests.rs delete mode 100644 src/openhuman/memory_store/content/atomic.rs delete mode 100644 src/openhuman/memory_store/content/compose/chunk.rs delete mode 100644 src/openhuman/memory_store/content/compose/mod.rs delete mode 100644 src/openhuman/memory_store/content/compose/summary.rs delete mode 100644 src/openhuman/memory_store/content/compose/tests.rs delete mode 100644 src/openhuman/memory_store/content/compose/yaml.rs delete mode 100644 src/openhuman/memory_store/content/paths.rs delete mode 100644 src/openhuman/memory_store/content/raw.rs delete mode 100644 src/openhuman/memory_store/vectors/store.rs delete mode 100644 src/openhuman/memory_store/vectors/store_tests.rs delete mode 100644 src/openhuman/memory_sync/canonicalize/chat.rs delete mode 100644 src/openhuman/memory_sync/canonicalize/document.rs delete mode 100644 src/openhuman/memory_sync/canonicalize/email.rs delete mode 100644 src/openhuman/memory_sync/canonicalize/email_clean.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/clickup/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/clickup/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/github/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/github/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/gmail/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/gmail/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/linear/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/linear/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/notion/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/notion/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/orchestrator.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/slack/ingest.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/slack/source.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/slack/sync.rs delete mode 100644 src/openhuman/memory_sync/composio/providers/slack/users.rs create mode 100644 src/openhuman/memory_tree/retrieval/engine.rs delete mode 100644 src/openhuman/memory_tree/score/extract/extractor.rs delete mode 100644 src/openhuman/memory_tree/score/extract/llm.rs delete mode 100644 src/openhuman/memory_tree/score/extract/llm_tests.rs delete mode 100644 src/openhuman/memory_tree/score/extract/regex.rs delete mode 100644 src/openhuman/memory_tree/score/extract/types.rs delete mode 100644 src/openhuman/memory_tree/score/resolver.rs delete mode 100644 src/openhuman/memory_tree/score/signals/interaction.rs delete mode 100644 src/openhuman/memory_tree/score/signals/metadata_weight.rs delete mode 100644 src/openhuman/memory_tree/score/signals/mod.rs delete mode 100644 src/openhuman/memory_tree/score/signals/ops.rs delete mode 100644 src/openhuman/memory_tree/score/signals/source_weight.rs delete mode 100644 src/openhuman/memory_tree/score/signals/token_count.rs delete mode 100644 src/openhuman/memory_tree/score/signals/types.rs delete mode 100644 src/openhuman/memory_tree/score/signals/unique_words.rs delete mode 100644 src/openhuman/memory_tree/score/store_tests.rs delete mode 100644 src/openhuman/memory_tree/tree/bucket_seal_tests.rs delete mode 100644 src/openhuman/memory_tree/tree_runtime/engine_tests.rs delete mode 100644 src/openhuman/memory_tree/tree_runtime/store_tests.rs create mode 100644 src/openhuman/tinycortex/ingest.rs create mode 100644 src/openhuman/tinycortex/seal.rs create mode 100644 src/openhuman/tinycortex/summariser.rs create mode 100644 src/openhuman/tinycortex/sync.rs diff --git a/Cargo.lock b/Cargo.lock index 1de8bc5cd..6570e62b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6875,16 +6875,23 @@ dependencies = [ "anyhow", "async-trait", "chrono", + "futures", "git2", + "log", "parking_lot", - "rand 0.8.6", + "rand 0.10.1", "regex", + "reqwest 0.12.28", "rusqlite", + "schemars", "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.18", + "tinyagents", + "tokio", "toml 0.8.23", + "tracing", "uuid 1.23.1", "walkdir", ] diff --git a/Cargo.toml b/Cargo.toml index 89012c3ff..c0fc36c89 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -88,7 +88,7 @@ tinyagents = { version = "1.7", features = ["sqlite", "repl"] } # security gating, and the global singleton stay host-side. rusqlite/git2 are # aligned to the host pins (=0.40 / 0.21) so one bundled SQLite + one libgit2 # link. Keep the version pin in lockstep with the submodule tag. -tinycortex = "0.1" +tinycortex = { version = "0.1", features = ["git-diff", "sync"] } tinychannels = { version = "0.1", features = ["relay-websocket"] } # TokenJuice code compressor — AST-aware signature extraction. Optional (C build) # behind the default `tokenjuice-treesitter` feature; disabling it falls back to @@ -377,6 +377,12 @@ tinyjuice = { path = "vendor/tinyjuice" } tinychannels = { path = "vendor/tinychannels" } tinyplace = { path = "vendor/tinyplace/sdk/rust" } +# TinyCortex temporarily pins the embedding API port while tinyagents #58 is +# pending. Resolve that git source to the same vendored crate so host adapters +# and TinyCortex share one `EmbeddingModel` trait identity. +[patch."https://github.com/senamakel/tinyagents"] +tinyagents = { path = "vendor/tinyagents" } + # Emit just enough DWARF in release builds for Sentry to symbolicate Rust # panics + render surrounding source lines. `line-tables-only` keeps the # binary small (only file+line tables, no full type info) while still diff --git a/app/src-tauri/Cargo.lock b/app/src-tauri/Cargo.lock index 5126fc413..4c0fc3061 100644 --- a/app/src-tauri/Cargo.lock +++ b/app/src-tauri/Cargo.lock @@ -9224,16 +9224,23 @@ dependencies = [ "anyhow", "async-trait", "chrono", + "futures", "git2", + "log", "parking_lot", - "rand 0.8.6", + "rand 0.10.1", "regex", + "reqwest 0.12.28", "rusqlite", + "schemars 1.2.1", "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.18", + "tinyagents", + "tokio", "toml 0.8.2", + "tracing", "uuid 1.23.1", "walkdir", ] diff --git a/app/src-tauri/Cargo.toml b/app/src-tauri/Cargo.toml index 12859c4c2..ec9a10576 100644 --- a/app/src-tauri/Cargo.toml +++ b/app/src-tauri/Cargo.toml @@ -216,7 +216,7 @@ tinyagents = { path = "../../vendor/tinyagents" } # TinyAgents so integration work can test crate changes against OpenHuman before # publishing. tinyflows = { path = "../../vendor/tinyflows" } -tinycortex = { path = "../../vendor/tinycortex" } +tinycortex = { path = "../../vendor/tinycortex", features = ["git-diff"] } tinyjuice = { path = "../../vendor/tinyjuice" } tinychannels = { path = "../../vendor/tinychannels" } tinyplace = { path = "../../vendor/tinyplace/sdk/rust" } @@ -248,6 +248,11 @@ tauri-plugin-global-shortcut = { git = "https://github.com/tauri-apps/plugins-wo tauri-plugin-single-instance = { git = "https://github.com/tauri-apps/plugins-workspace", rev = "c6561ab6b4f9e7f650d4fc8c53fd8acc9b65b9b2" } tauri-plugin-notification = { path = "vendor/tauri-plugin-notification" } +# Match the root Cargo world: TinyCortex's temporary git pin must resolve to +# the vendored TinyAgents crate to preserve one embedding trait identity. +[patch."https://github.com/senamakel/tinyagents"] +tinyagents = { path = "../../vendor/tinyagents" } + [dev-dependencies] # `test-util` enables `#[tokio::test(start_paused = true)]` for the # idle-watchdog unit tests in `cdp/session.rs` (#1213). diff --git a/src/bin/gmail_backfill_3d.rs b/src/bin/gmail_backfill_3d.rs index f95f405cc..a99f37c3f 100644 --- a/src/bin/gmail_backfill_3d.rs +++ b/src/bin/gmail_backfill_3d.rs @@ -1,14 +1,11 @@ -//! Backfill the last N days of Gmail into the memory-tree content store. +//! Backfill the last N days of Gmail into synchronized memory documents. //! //! Authenticates via Composio (JWT from `/auth-profiles.json`), -//! fetches Gmail pages via `GMAIL_FETCH_EMAILS`, converts each thread into an -//! [`EmailThread`], ingests it through `ingest_page_into_memory_tree` (which -//! writes `.md` files via `content_store` and populates SQLite), then drains -//! the async worker pool until idle. +//! fetches and stores Gmail pages through tinycortex, then drains the async +//! document-ingestion worker pool until idle. //! -//! After draining, the binary performs an integrity check: for every chunk -//! that has a `content_path` in SQLite, it verifies the on-disk SHA-256 -//! matches the stored `content_sha256`. +//! After draining, the binary verifies that TinyCortex persisted the reported +//! records in the `skill-gmail` namespace. //! //! # Prerequisites //! @@ -31,30 +28,19 @@ use anyhow::{Context, Result}; use clap::Parser; -use serde_json::{json, Value}; - -use openhuman_core::openhuman::composio::client::{ - create_composio_client, direct_execute, ComposioClientKind, -}; -use openhuman_core::openhuman::composio::providers::gmail::ingest::ingest_page_into_memory_tree; -use openhuman_core::openhuman::composio::providers::registry::{ - get_provider, init_default_providers, -}; use openhuman_core::openhuman::config::Config; use openhuman_core::openhuman::memory_queue::drain_until_idle; -use openhuman_core::openhuman::memory_store::chunks::store::{ - get_chunk_content_pointers, list_chunks, list_summaries_with_content_path, ListChunksQuery, -}; -use openhuman_core::openhuman::memory_store::content::read::{ - verify_chunk_file, verify_summary_file, VerifyResult, -}; #[derive(Parser, Debug)] #[command( name = "gmail-backfill-3d", - about = "Backfill last N days of Gmail into the memory-tree content store (.md files + SQLite)." + about = "Backfill recent Gmail messages into synchronized memory documents." )] struct Cli { + /// Composio Gmail connection id. Defaults to the configured Gmail source. + #[arg(long)] + connection_id: Option, + /// Lookback window in days. Default 3. #[arg(long, default_value_t = 3)] days: u32, @@ -80,7 +66,7 @@ struct Cli { #[arg(long, default_value_t = false)] skip_drain: bool, - /// Skip the post-drain integrity check (SHA-256 file verification). + /// Skip the post-drain persisted-document count check. #[arg(long, default_value_t = false)] skip_verify: bool, @@ -89,8 +75,7 @@ struct Cli { #[arg(long)] owner: Option, - /// Wipe `chunks.db` (+ wal/shm) AND `/` before running. - /// Useful after a chunker change that invalidates existing chunk IDs. + /// Clear synchronized Gmail documents before running. #[arg(long, default_value_t = false)] wipe: bool, } @@ -114,38 +99,24 @@ async fn main() -> Result<()> { if cli.days == 0 { anyhow::bail!("--days must be >= 1"); } + if cli.owner.is_some() { + log::warn!("[gmail_backfill_3d] --owner is retained for compatibility but ignored"); + } let config = Config::load_or_init() .await .context("[gmail_backfill_3d] Config::load_or_init failed")?; + let memory = openhuman_core::openhuman::memory::global::init(config.workspace_dir.clone()) + .map_err(anyhow::Error::msg)?; if cli.wipe { - wipe_memory_tree_state(&config)?; + log::info!("[gmail_backfill_3d] clearing skill-gmail documents"); + memory + .clear_namespace("skill-gmail") + .await + .map_err(anyhow::Error::msg)?; } - - // Resolve through the mode-aware factory so the backfill runs in - // EITHER backend mode (legacy JWT-driven path) OR direct mode (BYO - // Composio API key on the user's personal tenant) — #1710 Wave 2. - // Pre-fix this binary was hard-wired to backend mode via - // `build_composio_client`, so a direct-mode user couldn't run a - // gmail backfill even with a healthy personal connection. - let client_kind = create_composio_client(&config).map_err(|e| { - anyhow::anyhow!( - "No Composio client — user not signed in (backend session) and no direct-mode \ - API key configured. Sign in via the desktop app or set a Composio API key, \ - then re-run this binary. ({e})" - ) - })?; - - init_default_providers(); - let provider = get_provider("gmail").ok_or_else(|| { - anyhow::anyhow!("GmailProvider not registered after init_default_providers") - })?; - - let owner = cli - .owner - .clone() - .unwrap_or_else(|| "gmail-backfill".to_string()); + let documents_before = gmail_document_count(&memory).await?; let mut query = format!("in:inbox newer_than:{}d", cli.days); if !cli.include_spam_trash { @@ -169,100 +140,42 @@ async fn main() -> Result<()> { query, ); - let content_root = config.memory_tree_content_root(); - log::info!( - "[gmail_backfill_3d] content_root={}", - content_root.display() - ); + // ─── Fetch + ingest through tinycortex ────────────────────────────────── - // ─── Fetch + ingest ──────────────────────────────────────────────────── - - let mut page_token: Option = None; - let mut total_chunks = 0usize; - let mut total_pages = 0usize; - let mut total_cost: f64 = 0.0; - - for page_num in 0..cli.max_pages { - let mut args = json!({ - "max_results": cli.page_size, - "query": query, - }); - if cli.include_spam_trash { - args["include_spam_trash"] = json!(true); - } - if let Some(token) = &page_token { - args["page_token"] = json!(token); - } - - log::info!( - "[gmail_backfill_3d] fetching page {}{}…", - page_num, - page_token.as_ref().map(|_| " (paginated)").unwrap_or(""), - ); - - let mut resp = match &client_kind { - ComposioClientKind::Backend(client) => client - .execute_tool("GMAIL_FETCH_EMAILS", Some(args.clone())) - .await - .map_err(|e| anyhow::anyhow!("GMAIL_FETCH_EMAILS page {page_num}: {e:#}"))?, - ComposioClientKind::Direct(direct) => direct_execute( - direct, - "GMAIL_FETCH_EMAILS", - Some(args.clone()), - &config.composio.entity_id, - None, + let connection_id = cli + .connection_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .or_else(|| { + config.memory_sources.iter().find_map(|source| { + (source.kind == openhuman_core::openhuman::memory_sources::SourceKind::Composio + && source.toolkit.as_deref() == Some("gmail")) + .then(|| source.connection_id.clone()) + .flatten() + }) + }) + .ok_or_else(|| { + anyhow::anyhow!( + "no Gmail connection configured; pass --connection-id or add a Gmail memory source" ) - .await - .map_err(|e| anyhow::anyhow!("GMAIL_FETCH_EMAILS (direct) page {page_num}: {e:#}"))?, - }; - total_cost += resp.cost_usd; - - if !resp.successful { - anyhow::bail!( - "GMAIL_FETCH_EMAILS page {page_num} failed: {:?}", - resp.error - ); - } - - provider.post_process_action_result("GMAIL_FETCH_EMAILS", Some(&args), &mut resp.data); - - let (messages, next_token) = extract_envelope(&resp.data); - log::info!( - "[gmail_backfill_3d] page {} -> {} messages, next_token={}", - page_num, - messages.len(), - next_token.as_deref().unwrap_or("(none)"), - ); - - if messages.is_empty() { - break; - } - - // CLI runs don't fetch the user profile, so pass `None` and - // let the ingest fall back to per-participants source ids. - let chunks_this_page = - ingest_page_into_memory_tree(&config, &owner, None, &messages).await?; - total_chunks += chunks_this_page; - total_pages += 1; - - log::info!( - "[gmail_backfill_3d] page {} ingested chunks={} running_total={}", - page_num, - chunks_this_page, - total_chunks, - ); - - match next_token { - Some(tok) => page_token = Some(tok), - None => break, - } - } - + })?; + let outcome = openhuman_core::openhuman::tinycortex::run_gmail_backfill( + &connection_id, + &query, + cli.max_pages as usize, + cli.page_size as usize, + &config, + ) + .await + .map_err(|error| anyhow::anyhow!(error.to_string()))?; log::info!( - "[gmail_backfill_3d] fetch+ingest done pages={} total_chunks={} cost=~${:.4}", - total_pages, - total_chunks, - total_cost, + "[gmail_backfill_3d] fetch+ingest done records={} actions={} cost_usd={:.4} note={:?}", + outcome.records_ingested, + outcome.actions_called, + outcome.provider_cost_usd, + outcome.note, ); // ─── Drain async worker pool ──────────────────────────────────────────── @@ -280,218 +193,36 @@ async fn main() -> Result<()> { if cli.skip_verify { log::info!("[gmail_backfill_3d] skipping integrity check (--skip-verify)"); } else { - log::info!("[gmail_backfill_3d] running integrity check…"); - - // Chunk integrity. - let (verified, mismatched, no_pointer, missing_file) = verify_all_chunk_files(&config)?; - log::info!( - "[gmail_backfill_3d] chunks: verified={} mismatched={} no_pointer={} missing_file={}", - verified, - mismatched, - no_pointer, - missing_file, + let documents_after = gmail_document_count(&memory).await?; + let minimum_expected = documents_before.saturating_add(outcome.records_ingested as usize); + anyhow::ensure!( + documents_after >= minimum_expected, + "persisted Gmail document count {documents_after} is below expected minimum {minimum_expected}" ); - - // Summary integrity. - let (sum_verified, sum_mismatched, sum_no_pointer, sum_missing_file) = - verify_all_summary_files(&config)?; log::info!( - "[gmail_backfill_3d] summaries: verified={} mismatched={} no_pointer={} missing_file={}", - sum_verified, - sum_mismatched, - sum_no_pointer, - sum_missing_file, + "[gmail_backfill_3d] persisted documents before={} after={} reported={}", + documents_before, + documents_after, + outcome.records_ingested, ); - - if mismatched > 0 || missing_file > 0 || sum_mismatched > 0 || sum_missing_file > 0 { - anyhow::bail!( - "Integrity check failed: \ - chunks: {} mismatches, {} missing files; \ - summaries: {} mismatches, {} missing files", - mismatched, - missing_file, - sum_mismatched, - sum_missing_file, - ); - } } println!( - "\nBackfill complete. pages={} chunks_written={} cost=~${:.4}", - total_pages, total_chunks, total_cost, + "\nBackfill complete. records_ingested={} actions={} cost=~${:.4}", + outcome.records_ingested, outcome.actions_called, outcome.provider_cost_usd, ); Ok(()) } -/// Wipe `/memory_tree/chunks.db` (+ wal/shm) and -/// `/` so the bin can re-run cleanly after a chunker -/// change that invalidates existing chunk IDs. -/// -/// Logs each removed artifact at info; missing files are not an error. -fn wipe_memory_tree_state(config: &Config) -> Result<()> { - let mt_dir = config.workspace_dir.join("memory_tree"); - for name in &["chunks.db", "chunks.db-wal", "chunks.db-shm"] { - let path = mt_dir.join(name); - match std::fs::remove_file(&path) { - Ok(()) => log::info!("[gmail_backfill_3d] wiped {}", path.display()), - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} - Err(e) => return Err(e).with_context(|| format!("wipe {}", path.display())), - } - } - let content_root = config.memory_tree_content_root(); - if content_root.exists() { - std::fs::remove_dir_all(&content_root) - .with_context(|| format!("wipe {}", content_root.display()))?; - log::info!("[gmail_backfill_3d] wiped {}", content_root.display()); - } - Ok(()) -} - -/// Read all chunks from SQLite and verify on-disk SHA-256 matches `content_sha256`. -/// -/// Returns `(verified, mismatched, no_pointer, missing_file)`. -fn verify_all_chunk_files(config: &Config) -> Result<(usize, usize, usize, usize)> { - let chunks = list_chunks(config, &ListChunksQuery::default())?; - let content_root = config.memory_tree_content_root(); - - let mut verified = 0usize; - let mut mismatched = 0usize; - let mut no_pointer = 0usize; - let mut missing_file = 0usize; - - for chunk in &chunks { - let pointers = get_chunk_content_pointers(config, &chunk.id)?; - let (rel_path, expected_sha) = match pointers { - None => { - no_pointer += 1; - log::debug!( - "[gmail_backfill_3d] verify: chunk {} has no content_path/sha256", - chunk.id - ); - continue; - } - Some(pair) => pair, - }; - - let abs_path = { - let mut p = content_root.clone(); - for component in rel_path.split('/') { - p.push(component); - } - p - }; - - if !abs_path.exists() { - missing_file += 1; - log::warn!( - "[gmail_backfill_3d] verify: file missing chunk_id={} path={}", - chunk.id, - abs_path.display(), - ); - continue; - } - - match verify_chunk_file(&abs_path, &expected_sha) { - Ok(true) => { - verified += 1; - } - Ok(false) => { - mismatched += 1; - log::warn!( - "[gmail_backfill_3d] verify: SHA-256 mismatch chunk_id={} path={}", - chunk.id, - abs_path.display(), - ); - } - Err(e) => { - log::error!( - "[gmail_backfill_3d] verify: error chunk_id={}: {e}", - chunk.id, - ); - mismatched += 1; - } - } - } - - Ok((verified, mismatched, no_pointer, missing_file)) -} - -/// Read all summary rows with a non-NULL `content_path` from SQLite and verify -/// the on-disk SHA-256 matches `content_sha256`. -/// -/// Returns `(verified, mismatched, no_pointer, missing_file)`. -fn verify_all_summary_files(config: &Config) -> Result<(usize, usize, usize, usize)> { - let rows_with_pointer = list_summaries_with_content_path(config)?; - let content_root = config.memory_tree_content_root(); - - let mut verified = 0usize; - let mut mismatched = 0usize; - let mut missing_file = 0usize; - - for (summary_id, rel_path, expected_sha) in &rows_with_pointer { - let abs_path = { - let mut p = content_root.clone(); - for component in rel_path.split('/') { - p.push(component); - } - p - }; - - match verify_summary_file(&abs_path, expected_sha) { - Ok(VerifyResult::Ok) => { - verified += 1; - } - Ok(VerifyResult::Mismatch { actual }) => { - mismatched += 1; - log::warn!( - "[gmail_backfill_3d] verify: SHA-256 mismatch summary_id={} path={} expected={} actual={}", - summary_id, - abs_path.display(), - expected_sha, - actual, - ); - } - Ok(VerifyResult::Missing) => { - missing_file += 1; - log::warn!( - "[gmail_backfill_3d] verify: file missing summary_id={} path={}", - summary_id, - abs_path.display(), - ); - } - Err(e) => { - log::error!( - "[gmail_backfill_3d] verify: error summary_id={}: {e}", - summary_id, - ); - mismatched += 1; - } - } - } - - // Count rows that have no content_path at all (legacy rows). - // We report this as no_pointer for symmetry with the chunk verifier. - // We can't easily count them here without a separate query, so we - // approximate: rows_with_pointer gives us the ones we checked. - // For now no_pointer = 0 (the bin wipes before re-ingesting so all - // new rows should have pointers; legacy rows are pre-migration). - let no_pointer = 0usize; - - Ok((verified, mismatched, no_pointer, missing_file)) -} - -/// Extract the `messages` array and `nextPageToken` from a Composio response. -fn extract_envelope(data: &Value) -> (Vec, Option) { - let candidates: [Option<&Value>; 2] = [Some(data), data.get("data")]; - for cand in candidates.into_iter().flatten() { - if let Some(arr) = cand.get("messages").and_then(|v| v.as_array()) { - let token = cand - .get("nextPageToken") - .and_then(|v| v.as_str()) - .filter(|s| !s.trim().is_empty()) - .map(str::to_string); - return (arr.clone(), token); - } - } - (Vec::new(), None) +async fn gmail_document_count( + memory: &openhuman_core::openhuman::memory_store::MemoryClientRef, +) -> Result { + let value = memory + .list_documents(Some("skill-gmail")) + .await + .map_err(anyhow::Error::msg)?; + Ok(value + .get("documents") + .and_then(serde_json::Value::as_array) + .map_or(0, Vec::len)) } diff --git a/src/bin/slack_backfill.rs b/src/bin/slack_backfill.rs index c2488e10b..732ac3c2e 100644 --- a/src/bin/slack_backfill.rs +++ b/src/bin/slack_backfill.rs @@ -42,11 +42,7 @@ use clap::Parser; use openhuman_core::openhuman::composio::client::{ create_composio_client, direct_execute, direct_list_connections, ComposioClientKind, }; -use openhuman_core::openhuman::composio::providers::registry::{ - get_provider, init_default_providers, -}; -use openhuman_core::openhuman::composio::providers::slack::run_backfill_via_search; -use openhuman_core::openhuman::composio::providers::{ProviderContext, SyncReason}; +use openhuman_core::openhuman::composio::providers::registry::init_default_providers; use openhuman_core::openhuman::composio::types::{ ComposioConnectionsResponse, ComposioExecuteResponse, }; @@ -194,10 +190,6 @@ async fn main() -> Result<()> { // Idempotent — safe even if called twice. init_default_providers(); - let provider = get_provider("slack").ok_or_else(|| { - anyhow::anyhow!("SlackProvider not registered after init_default_providers") - })?; - // Resolve through the mode-aware factory so the backfill runs in // EITHER backend mode (legacy JWT-driven path) OR direct mode (BYO // Composio API key on the user's personal tenant) — #1710 Wave 2. @@ -442,27 +434,18 @@ async fn main() -> Result<()> { let started = Instant::now(); let mut total_buckets = 0usize; for conn in &slack_conns { - // `ProviderContext` no longer caches a pre-baked client — - // `ctx.execute(...)` resolves via the mode-aware factory - // per call (#1710 / Wave 1). The local `client` handle is - // still used above for backend-only metadata probes. - let ctx = ProviderContext { - config: Arc::clone(&config), - toolkit: conn.toolkit.clone(), - connection_id: Some(conn.id.clone()), - usage: Default::default(), - max_items: None, - sync_depth_days: None, - }; - match run_backfill_via_search(&ctx, cli.days).await { + match openhuman_core::openhuman::tinycortex::run_slack_search_backfill( + &conn.id, + cli.days, + config.as_ref(), + ) + .await + { Ok(outcome) => { - total_buckets += outcome.items_ingested; + total_buckets += outcome.records_ingested as usize; println!( - "connection={} buckets={} elapsed_ms={} summary={:?}", - conn.id, - outcome.items_ingested, - outcome.elapsed_ms(), - outcome.summary, + "connection={} records={} summary={:?}", + conn.id, outcome.records_ingested, outcome.note, ); } Err(err) => { @@ -546,24 +529,19 @@ async fn main() -> Result<()> { } } } - let ctx = ProviderContext { - config: Arc::clone(&config), - toolkit: conn.toolkit.clone(), - connection_id: Some(conn.id.clone()), - usage: Default::default(), - max_items: None, - sync_depth_days: None, - }; - match provider.sync(&ctx, SyncReason::Manual).await { + match openhuman_core::openhuman::tinycortex::run_composio_connection( + "slack", + &conn.id, + config.as_ref(), + ) + .await + { Ok(outcome) => { connections_ok += 1; - total_buckets += outcome.items_ingested; + total_buckets += outcome.records_ingested as usize; println!( - "connection={} buckets_flushed={} elapsed_ms={} summary={:?}", - conn.id, - outcome.items_ingested, - outcome.elapsed_ms(), - outcome.summary, + "connection={} records_ingested={} summary={:?}", + conn.id, outcome.records_ingested, outcome.note, ); } Err(err) => { diff --git a/src/core/observability.rs b/src/core/observability.rs index eaf6d04a3..41e11ac2f 100644 --- a/src/core/observability.rs +++ b/src/core/observability.rs @@ -840,7 +840,9 @@ fn is_embedding_backend_auth_failure(lower: &str) -> bool { /// `embeddings::rpc::update_settings` as the save-time hard-block signal so the /// two never drift. pub(crate) fn is_embedding_endpoint_absent(lower: &str) -> bool { - lower.contains("embedding api error") && (lower.contains("(404") || lower.contains("(405")) + (lower.contains("embedding api error") && (lower.contains("(404") || lower.contains("(405"))) + || (lower.contains("embeddings returned http") + && (lower.contains("http 404") || lower.contains("http 405"))) } /// Detect a custom/cloud embeddings endpoint that IS an embeddings API but diff --git a/src/openhuman/agent/harness/session/builder/factory.rs b/src/openhuman/agent/harness/session/builder/factory.rs index 176b12222..3d8b55bb3 100644 --- a/src/openhuman/agent/harness/session/builder/factory.rs +++ b/src/openhuman/agent/harness/session/builder/factory.rs @@ -340,14 +340,16 @@ impl Agent { config, &config.memory.embedding_provider, ); - let memory: Arc = Arc::from(memory_store::create_memory_with_local_ai( + let session_memory = memory_store::factories::create_session_memory_with_local_ai( &config.memory, local_embedding.as_deref(), &embedding_api_key, &config.embedding_routes, Some(&config.storage.provider.config), &config.workspace_dir, - )?); + )?; + let archivist_connection = session_memory.sqlite_connection; + let memory: Arc = Arc::from(session_memory.memory); // Per-profile skill (workflow) + MCP-server allowlists. `None` = all. let profile_skill_allowlist: Option> = profile @@ -758,33 +760,24 @@ impl Agent { // is the system-of-record for chat turns and must stay active even when // the inference stack (`reflection`, `stability_detector`) is disabled. // Gated only on `config.learning.episodic_capture_enabled` (default: true) - // and on the memory backend exposing a SQLite connection. + // using the explicit SQLite resource returned by the session factory. let archivist_hook_arc: Option< Arc, > = if config.learning.episodic_capture_enabled { - match memory.sqlite_conn() { - Some(conn) => { - let hook = Arc::new( - crate::openhuman::agent::harness::archivist::ArchivistHook::new(conn, true) - .with_config(config.clone()), - ); - post_turn_hooks - .push(Arc::clone(&hook) - as Arc); - log::info!( - "[archivist] episodic capture hook registered (learning.enabled={})", - config.learning.enabled - ); - Some(hook) - } - None => { - log::warn!( - "[archivist] no SQLite connection available from memory backend — \ - episodic capture disabled" - ); - None - } - } + let hook = Arc::new( + crate::openhuman::agent::harness::archivist::ArchivistHook::new( + archivist_connection, + true, + ) + .with_config(config.clone()), + ); + post_turn_hooks + .push(Arc::clone(&hook) as Arc); + log::info!( + "[archivist] episodic capture hook registered (learning.enabled={})", + config.learning.enabled + ); + Some(hook) } else { log::info!( "[archivist] episodic_capture_enabled=false — archivist hook not registered" diff --git a/src/openhuman/agent/harness/session/builder/helpers.rs b/src/openhuman/agent/harness/session/builder/helpers.rs index 21661f0af..c889497c6 100644 --- a/src/openhuman/agent/harness/session/builder/helpers.rs +++ b/src/openhuman/agent/harness/session/builder/helpers.rs @@ -1,7 +1,7 @@ //! Utility helpers used during agent construction. use crate::openhuman::memory::Memory; -use crate::openhuman::memory_tools::{tool_memory_store, ToolMemoryRule, ToolMemoryStore}; +use crate::openhuman::memory_tools::{tool_memory_store, ToolMemoryRule}; use std::sync::Arc; /// (#1400) Best-effort synchronous prefetch of eager tool-scoped rules. diff --git a/src/openhuman/composio/identity.rs b/src/openhuman/composio/identity.rs index cf6e9be53..89246d811 100644 --- a/src/openhuman/composio/identity.rs +++ b/src/openhuman/composio/identity.rs @@ -116,8 +116,7 @@ pub async fn connection_identity(config: &Config, toolkit: &str) -> Option Result { - Ok(SyncOutcome { - toolkit: self.slug.to_string(), - reason: reason.as_str().to_string(), - ..Default::default() - }) - } } fn fresh_config_in_workspace(tmp: &std::path::Path) -> Config { diff --git a/src/openhuman/composio/ops/memory_cleanup.rs b/src/openhuman/composio/ops/memory_cleanup.rs index 18ffd84b0..faad7acc2 100644 --- a/src/openhuman/composio/ops/memory_cleanup.rs +++ b/src/openhuman/composio/ops/memory_cleanup.rs @@ -7,8 +7,6 @@ use crate::openhuman::memory::MemoryClient; use crate::openhuman::memory_store::chunks::store as memory_tree_store; use crate::openhuman::memory_store::chunks::types::SourceKind; -use super::super::providers::sync_state::SyncState; - #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum MemoryCleanupTarget { Exact(SourceKind, String), @@ -96,7 +94,8 @@ async fn notion_memory_targets_for_connection( ) })?, ); - let state = SyncState::load(&memory, "notion", connection_id) + let adapter = crate::openhuman::tinycortex::HostSyncAdapter::new(memory); + let state = tinycortex::memory::sync::SyncState::load(&adapter, "notion", connection_id) .await .map_err(|error| { anyhow::anyhow!("failed to load notion sync state for memory cleanup: {error}") diff --git a/src/openhuman/composio/ops/providers_ops.rs b/src/openhuman/composio/ops/providers_ops.rs index 42aa4590a..508fbb263 100644 --- a/src/openhuman/composio/ops/providers_ops.rs +++ b/src/openhuman/composio/ops/providers_ops.rs @@ -176,19 +176,11 @@ pub async fn composio_sync( ); let client = resolve_client(config)?; let toolkit = resolve_toolkit_for_connection(&client, connection_id).await?; - let provider = get_provider(&toolkit).ok_or_else(|| { + get_provider(&toolkit).ok_or_else(|| { format!("[composio] no native provider registered for toolkit '{toolkit}'") })?; let _ = client; - - let ctx = ProviderContext { - config: Arc::new(config.clone()), - toolkit: toolkit.clone(), - connection_id: Some(connection_id.to_string()), - usage: Default::default(), - max_items: None, - sync_depth_days: None, - }; + let config_for_task = config.clone(); let started_at_ms = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis() as u64) @@ -197,21 +189,26 @@ pub async fn composio_sync( let connection_id_for_log = connection_id.to_string(); tokio::spawn(async move { - let toolkit_in_task = ctx.toolkit.clone(); - match provider.sync(&ctx, reason).await { + match crate::openhuman::tinycortex::run_composio_connection( + &toolkit_for_outcome, + &connection_id_for_log, + &config_for_task, + ) + .await + { Ok(out) => { tracing::info!( - toolkit = %toolkit_in_task, + toolkit = %toolkit_for_outcome, connection_id = %connection_id_for_log, - items_ingested = out.items_ingested, - elapsed_ms = out.elapsed_ms(), + items_ingested = out.records_ingested, + actions_called = out.actions_called, "[composio] background sync ok" ); } Err(e) => { report_composio_op_error("sync", &e); tracing::warn!( - toolkit = %toolkit_in_task, + toolkit = %toolkit_for_outcome, connection_id = %connection_id_for_log, error = %e, "[composio] background sync failed" @@ -220,9 +217,9 @@ pub async fn composio_sync( } }); - let summary = format!("composio: {toolkit_for_outcome} sync started (background)"); + let summary = format!("composio: {toolkit} sync started (background)"); let outcome = SyncOutcome { - toolkit: toolkit_for_outcome, + toolkit, connection_id: Some(connection_id.to_string()), reason: reason.as_str().to_string(), items_ingested: 0, diff --git a/src/openhuman/composio/ops_tests.rs b/src/openhuman/composio/ops_tests.rs index 5d77108eb..a2ce40467 100644 --- a/src/openhuman/composio/ops_tests.rs +++ b/src/openhuman/composio/ops_tests.rs @@ -762,6 +762,13 @@ async fn composio_delete_connection_clear_memory_cascades_live_sealed_tree_and_f token_sum: i64::from(chunk.token_count), oldest_at: Some(chunk.metadata.time_range.0), }; + memory_tree_store::with_connection(&config, |conn| { + let tx = conn.unchecked_transaction()?; + tree_store::upsert_buffer_tx(&tx, &buf)?; + tx.commit()?; + Ok(()) + }) + .expect("persist buffer snapshot"); let summary_id = seal_one_level(&config, &tree, &buf, &LabelStrategy::Empty, false) .await .expect("real seal produces a summary"); @@ -895,7 +902,8 @@ async fn notion_cleanup_targets_include_synced_page_sources() { let mut state = SyncState::new("notion", "conn-1"); state.mark_synced("page-a@2026-01-01T00:00:00Z"); state.mark_synced("page-b"); - state.save(&memory).await.expect("sync state should save"); + let adapter = crate::openhuman::tinycortex::HostSyncAdapter::new(memory); + state.save(&adapter).await.expect("sync state should save"); let targets = composio_memory_targets_for_connection(&config, Some("notion"), "conn-1") .await @@ -1099,13 +1107,11 @@ async fn composio_execute_via_mock_propagates_backend_error() { } #[tokio::test] -async fn composio_sync_gmail_via_mock_archives_raw_email_and_updates_outcome() { +async fn composio_sync_gmail_via_mock_stores_skill_document_and_updates_outcome() { let _serial = crate::openhuman::memory::ops::GLOBAL_MEMORY_TEST_LOCK .lock() .await; use crate::openhuman::config::TEST_ENV_LOCK; - use crate::openhuman::memory_store::content::raw::{raw_rel_path, RawKind}; - use crate::openhuman::memory_tree::tree::rpc::{list_chunks_rpc, ListChunksRequest}; let _cache_guard = cache_guard(); let _env_guard = TEST_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()); @@ -1193,63 +1199,36 @@ async fn composio_sync_gmail_via_mock_archives_raw_email_and_updates_outcome() { outcome.value.summary ); - // Poll for the spawned ingest task to drain. The mock backend is - // local + in-memory, so this normally lands in well under a second. - let chunks = { - let mut chunks = Vec::new(); + // Poll for the spawned ingest task to persist the TinyCortex skill + // document. The mock backend is local, so this normally lands quickly. + let documents = { + let mut documents = Vec::new(); for _ in 0..50 { - chunks = list_chunks_rpc( - &config, - ListChunksRequest { - source_kind: Some("email".to_string()), - source_id: Some("gmail:pilot-at-example-dot-com".to_string()), - limit: Some(10), - ..Default::default() - }, - ) - .await - .unwrap() - .value - .chunks; - if !chunks.is_empty() { + documents = crate::openhuman::memory::global::client_if_ready() + .expect("memory client remains initialized") + .list_documents(Some("skill-gmail")) + .await + .unwrap() + .get("documents") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + if !documents.is_empty() { break; } tokio::time::sleep(std::time::Duration::from_millis(100)).await; } - chunks + documents }; assert_eq!( - chunks.len(), + documents.len(), 1, - "expected one ingested Gmail chunk after spawned task drains" - ); - assert!( - chunks[0].content.contains("Phoenix launch canary"), - "chunk content missing mock email subject: {}", - chunks[0].content - ); - assert!( - chunks[0].content.contains("mock sync coverage"), - "chunk content missing mock email body: {}", - chunks[0].content - ); - - let raw_path = config.memory_tree_content_root().join(raw_rel_path( - "gmail:pilot-at-example-dot-com", - RawKind::Email, - 1_717_243_200_000, - "gmail-msg-1", - )); - let archived = std::fs::read_to_string(&raw_path) - .unwrap_or_else(|e| panic!("expected archived Gmail raw message at {raw_path:?}: {e}")); - assert!( - archived.contains("Phoenix launch canary"), - "archived email missing mock subject: {archived}" - ); - assert!( - archived.contains("mock sync coverage"), - "archived email missing mock body: {archived}" + "expected one ingested Gmail document after spawned task drains" ); + let document = &documents[0]; + assert_eq!(document["documentId"], "gmail:gmail-msg-1"); + assert_eq!(document["title"], "Phoenix launch canary"); + assert_eq!(document["taint"], "external_sync"); } #[tokio::test] diff --git a/src/openhuman/composio/tools_tests.rs b/src/openhuman/composio/tools_tests.rs index ea0a5ea88..888ba10f5 100644 --- a/src/openhuman/composio/tools_tests.rs +++ b/src/openhuman/composio/tools_tests.rs @@ -2,7 +2,6 @@ use super::*; use crate::openhuman::composio::providers::tool_scope::{CuratedTool, ToolScope}; use crate::openhuman::composio::providers::{ registry::register_provider, ComposioProvider, ProviderContext, ProviderUserProfile, - SyncOutcome, SyncReason, }; use async_trait::async_trait; use std::path::Path; @@ -31,14 +30,6 @@ impl ComposioProvider for ProviderOnlyCatalog { ) -> Result { Ok(ProviderUserProfile::default()) } - - async fn sync( - &self, - _ctx: &ProviderContext, - _reason: SyncReason, - ) -> Result { - Ok(SyncOutcome::default()) - } } struct WorkspaceEnvGuard { diff --git a/src/openhuman/embeddings/cloud.rs b/src/openhuman/embeddings/cloud.rs deleted file mode 100644 index 18675b518..000000000 --- a/src/openhuman/embeddings/cloud.rs +++ /dev/null @@ -1,213 +0,0 @@ -//! Cloud embedding provider — routes through the OpenHuman backend's -//! `POST /openai/v1/embeddings` surface (Voyage-backed) using the same -//! session JWT that the `OpenHumanBackendProvider` chat path uses. -//! -//! This is the default embedder for a fresh install. The local Ollama path -//! stays available, but the user has to explicitly opt in (either by setting -//! `memory.embedding_provider = "ollama"` in `config.toml`, or by enabling -//! the local-AI runtime with `local_ai.usage.embeddings = true`). -//! -//! The JWT and API URL are resolved per call so a session refresh between -//! embed batches is picked up transparently — matching -//! [`crate::openhuman::inference::provider::openhuman_backend::OpenHumanBackendProvider`]. - -use std::path::PathBuf; - -use async_trait::async_trait; - -use super::openai::OpenAiEmbedding; -use super::EmbeddingProvider; -use crate::api::config::effective_api_url; -use crate::openhuman::credentials::{AuthService, APP_SESSION_PROVIDER}; - -/// Default cloud embedding model — backed by `voyage-3.5` (1024 dims) on the -/// OpenHuman backend. See `tinyhumansai/backend#746`. -pub const DEFAULT_CLOUD_EMBEDDING_MODEL: &str = "embedding-v1"; - -/// Default output dimensionality for [`DEFAULT_CLOUD_EMBEDDING_MODEL`]. -pub const DEFAULT_CLOUD_EMBEDDING_DIMENSIONS: usize = 1024; - -/// OpenHuman-backend-backed embedding provider. -pub struct OpenHumanCloudEmbedding { - api_url: Option, - openhuman_dir: Option, - secrets_encrypt: bool, - model: String, - dims: usize, -} - -impl OpenHumanCloudEmbedding { - /// Construct a cloud embedder. `api_url` and `openhuman_dir` are looked up - /// per request; pass `None` to fall back to the runtime defaults - /// ([`effective_api_url`] / `~/.openhuman`). - pub fn new( - api_url: Option, - openhuman_dir: Option, - secrets_encrypt: bool, - model: impl Into, - dims: usize, - ) -> Self { - Self { - api_url: api_url - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()), - openhuman_dir, - secrets_encrypt, - model: model.into(), - dims, - } - } - - fn state_dir(&self) -> PathBuf { - self.openhuman_dir.clone().unwrap_or_else(|| { - // Honor OPENHUMAN_WORKSPACE (where auth-profiles.json lives) before - // falling back to ~/.openhuman, so the cloud embedder resolves the - // session JWT from the same directory the chat provider does. Without - // this, any non-default workspace (OPENHUMAN_WORKSPACE set, e.g. tests - // / multi-instance) silently has no session for embeddings — - // resolve_bearer() bails, embed() errors, and vectors are dropped. - if let Some(ws) = std::env::var_os("OPENHUMAN_WORKSPACE") - .filter(|s| !s.is_empty()) - .map(PathBuf::from) - { - return ws; - } - directories::UserDirs::new() - .map(|d| d.home_dir().join(".openhuman")) - .unwrap_or_else(|| PathBuf::from(".openhuman")) - }) - } - - fn resolve_bearer(&self) -> anyhow::Result { - let auth = AuthService::new(&self.state_dir(), self.secrets_encrypt); - if let Some(t) = auth - .get_provider_bearer_token(APP_SESSION_PROVIDER, None)? - .filter(|s| !s.trim().is_empty()) - { - return Ok(t); - } - anyhow::bail!( - "No backend session for cloud embeddings: log in to OpenHuman, or set \ - memory.embedding_provider to \"ollama\" / \"none\" in config.toml" - ) - } - - fn base_url(&self) -> String { - let u = effective_api_url(&self.api_url); - format!("{}/openai/v1", u.trim_end_matches('/')) - } -} - -#[async_trait] -impl EmbeddingProvider for OpenHumanCloudEmbedding { - fn name(&self) -> &str { - "cloud" - } - - fn model_id(&self) -> &str { - &self.model - } - - fn dimensions(&self) -> usize { - self.dims - } - - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - if texts.is_empty() { - return Ok(Vec::new()); - } - // Mirror the OpenAI provider's empty/whitespace guard *here* so we - // never resolve the bearer or build the URL for an input the backend - // will reject as `"input must be a non-empty string …"` (#13021). - // Cheaper, and keeps unauthenticated test contexts off the auth path. - if let Some(idx) = texts.iter().position(|t| t.trim().is_empty()) { - tracing::warn!( - target: "cloud::embed", - "[cloud] refusing embed: input[{idx}] is empty/whitespace \ - (count={}, model={}). Caller must filter empty strings.", - texts.len(), - self.model, - ); - anyhow::bail!( - "cloud embed: refusing empty/whitespace input at index {idx} of {} (model={})", - texts.len(), - self.model, - ); - } - let token = self.resolve_bearer()?; - let inner = OpenAiEmbedding::new(&self.base_url(), &token, &self.model, self.dims); - inner.embed(texts).await - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn name_and_dimensions() { - let p = OpenHumanCloudEmbedding::new( - None, - None, - true, - DEFAULT_CLOUD_EMBEDDING_MODEL, - DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, - ); - assert_eq!(p.name(), "cloud"); - assert_eq!(p.model_id(), DEFAULT_CLOUD_EMBEDDING_MODEL); - assert_eq!(p.dimensions(), DEFAULT_CLOUD_EMBEDDING_DIMENSIONS); - assert_eq!(p.signature(), "provider=cloud;model=embedding-v1;dims=1024"); - } - - #[test] - fn base_url_appends_openai_v1() { - let p = OpenHumanCloudEmbedding::new( - Some("https://api.openhuman.example/".into()), - None, - true, - DEFAULT_CLOUD_EMBEDDING_MODEL, - DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, - ); - assert_eq!(p.base_url(), "https://api.openhuman.example/openai/v1"); - } - - #[tokio::test] - async fn embed_empty_returns_empty_without_auth() { - // Empty input should short-circuit *before* hitting the AuthService — - // otherwise the no-op path would spuriously fail in unauthenticated - // contexts (e.g. ingestion of an empty chunk batch). - let p = OpenHumanCloudEmbedding::new( - None, - None, - false, - DEFAULT_CLOUD_EMBEDDING_MODEL, - DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, - ); - assert!(p.embed(&[]).await.unwrap().is_empty()); - } - - #[tokio::test] - async fn embed_refuses_empty_string_without_auth() { - // The whitespace pre-flight (#13021) MUST fire before `resolve_bearer` - // — `secrets_encrypt = false` and no session means the AuthService - // would otherwise bail with `"No backend session for cloud - // embeddings…"` and mask the real defect. Asserting the bail wording - // here pins the order of the two checks. - let p = OpenHumanCloudEmbedding::new( - None, - None, - false, - DEFAULT_CLOUD_EMBEDDING_MODEL, - DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, - ); - let err = p.embed(&[""]).await.unwrap_err().to_string(); - assert!( - err.contains("refusing empty/whitespace input at index 0"), - "expected pre-flight refusal, got: {err}" - ); - assert!( - !err.contains("No backend session"), - "guard must run before resolve_bearer, got: {err}" - ); - } -} diff --git a/src/openhuman/embeddings/cloud_adapter.rs b/src/openhuman/embeddings/cloud_adapter.rs new file mode 100644 index 000000000..22f99e155 --- /dev/null +++ b/src/openhuman/embeddings/cloud_adapter.rs @@ -0,0 +1,88 @@ +//! OpenHuman credential adapter for tinyagents cloud embeddings. + +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use tinyagents::harness::embeddings::{ + BearerResolver, CloudEmbeddingModel, DEFAULT_CLOUD_DIMENSIONS, DEFAULT_CLOUD_MODEL, +}; + +use super::{EmbeddingProvider, TinyAgentsEmbeddingProvider}; +use crate::api::config::effective_api_url; +use crate::openhuman::credentials::{AuthService, APP_SESSION_PROVIDER}; + +pub const DEFAULT_CLOUD_EMBEDDING_MODEL: &str = DEFAULT_CLOUD_MODEL; +pub const DEFAULT_CLOUD_EMBEDDING_DIMENSIONS: usize = DEFAULT_CLOUD_DIMENSIONS; + +/// Host-owned credential resolution around the crate-owned cloud transport. +pub struct OpenHumanCloudEmbedding { + inner: TinyAgentsEmbeddingProvider, +} + +impl OpenHumanCloudEmbedding { + pub fn new( + api_url: Option, + openhuman_dir: Option, + secrets_encrypt: bool, + model: impl Into, + dimensions: usize, + ) -> Self { + let state_dir = openhuman_dir.unwrap_or_else(default_state_dir); + let bearer: BearerResolver = Arc::new(move || { + let auth = AuthService::new(&state_dir, secrets_encrypt); + auth.get_provider_bearer_token(APP_SESSION_PROVIDER, None) + .map_err(|error| tinyagents::TinyAgentsError::Embedding(error.to_string()))? + .filter(|token| !token.trim().is_empty()) + .ok_or_else(|| { + tinyagents::TinyAgentsError::Validation( + "No backend session for cloud embeddings: log in to OpenHuman".into(), + ) + }) + }); + let base_url = format!( + "{}/openai/v1", + effective_api_url(&api_url).trim_end_matches('/') + ); + Self { + inner: TinyAgentsEmbeddingProvider::new(CloudEmbeddingModel::new( + base_url, model, dimensions, bearer, + )), + } + } +} + +fn default_state_dir() -> PathBuf { + if let Some(workspace) = std::env::var_os("OPENHUMAN_WORKSPACE") + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + { + return workspace; + } + directories::UserDirs::new() + .map(|dirs| dirs.home_dir().join(".openhuman")) + .unwrap_or_else(|| PathBuf::from(".openhuman")) +} + +#[async_trait] +impl EmbeddingProvider for OpenHumanCloudEmbedding { + fn name(&self) -> &str { + self.inner.name() + } + + fn model_id(&self) -> &str { + self.inner.model_id() + } + + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + + fn signature(&self) -> String { + self.inner.signature() + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + self.inner.embed(texts).await + } +} diff --git a/src/openhuman/embeddings/cohere.rs b/src/openhuman/embeddings/cohere.rs deleted file mode 100644 index aa828d6a7..000000000 --- a/src/openhuman/embeddings/cohere.rs +++ /dev/null @@ -1,393 +0,0 @@ -//! Cohere embedding provider — direct API access with user's own key. -//! -//! Cohere's `/v2/embed` endpoint uses a slightly different contract than -//! OpenAI: `texts` instead of `input`, `embedding_types` instead of -//! `encoding_format`, and the response nests embeddings inside -//! `embeddings.float`. This module implements the Cohere-native wire -//! format. - -use async_trait::async_trait; - -use super::retry_after::{backoff_ms_for_attempt, MAX_429_RETRIES}; -use super::EmbeddingProvider; - -pub const COHERE_API_BASE: &str = "https://api.cohere.com"; -pub const COHERE_DEFAULT_MODEL: &str = "embed-english-v3.0"; -pub const COHERE_DEFAULT_DIMS: usize = 1024; - -pub struct CohereEmbedding { - api_key: String, - model: String, - dims: usize, - base_url: String, -} - -impl CohereEmbedding { - pub fn new(api_key: &str, model: &str, dims: usize) -> Self { - let model = if model.is_empty() { - COHERE_DEFAULT_MODEL.to_string() - } else { - model.to_string() - }; - let dims = if dims == 0 { COHERE_DEFAULT_DIMS } else { dims }; - - Self { - api_key: api_key.to_string(), - model, - dims, - base_url: COHERE_API_BASE.to_string(), - } - } - - /// Override the Cohere-compatible API base URL. - /// - /// This keeps the public provider usable with local mocks and compatible - /// deployments while preserving Cohere's hosted endpoint as the default. - /// - /// Input is trimmed of surrounding whitespace and trailing slashes so the - /// endpoint built in [`Self::embed`] (`{base}/v2/embed`) never produces a - /// doubled slash when callers pass `https://host/`. - pub fn with_base_url(mut self, base: impl Into) -> Self { - self.base_url = base.into().trim().trim_end_matches('/').to_string(); - self - } - - fn http_client(&self) -> reqwest::Client { - crate::openhuman::config::build_runtime_proxy_client("embeddings.cohere") - } -} - -#[derive(serde::Deserialize)] -struct CohereEmbedResponse { - embeddings: CohereEmbeddings, -} - -#[derive(serde::Deserialize)] -struct CohereEmbeddings { - float: Vec>, -} - -#[async_trait] -impl EmbeddingProvider for CohereEmbedding { - fn name(&self) -> &str { - "cohere" - } - - fn model_id(&self) -> &str { - &self.model - } - - fn dimensions(&self) -> usize { - self.dims - } - - /// Sends a POST request to the Cohere embed API. - /// - /// On 429 (Too Many Requests) or 503 (Service Unavailable) the call is - /// retried up to `MAX_429_RETRIES` times with exponential backoff. When - /// the server supplies a `Retry-After` header its value (delta-seconds) is - /// preferred over the computed backoff. After all retries are exhausted the - /// canonical error message is returned so the `TransientUpstreamHttp` - /// classifier in `core::observability` demotes it to a warning breadcrumb - /// instead of a Sentry error event. - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - if texts.is_empty() { - return Ok(Vec::new()); - } - - // Fast-fail when no API key is configured. Cohere's hosted `/v2/embed` - // always requires a bearer token; without this guard we POST - // `Authorization: Bearer ` (empty) and Cohere returns a 401 "no api key - // supplied" on every call. That 401 is not retryable, so each embed - // attempt bails and reports an error — the memory pipeline re-embeds per - // document and floods Sentry (TAURI-RUST-52S: 8.7k events from a single - // misconfigured user). Bailing here skips the wasted request entirely, - // and the "API key not set" wording is demoted to a breadcrumb by the - // `ApiKeyMissing` classifier in - // `core::observability::expected_error_kind` instead of surfacing as a - // Sentry event. Scoped to Cohere on purpose: the OpenAI-compatible - // provider legitimately supports keyless local/custom endpoints, so it - // omits the header rather than bailing. - if self.api_key.trim().is_empty() { - anyhow::bail!( - "Cohere API key not set. Configure via the web UI or set the appropriate env var." - ); - } - - let url = format!("{}/v2/embed", self.base_url); - - tracing::debug!( - target: "embeddings.cohere", - "[cohere] embed: model={}, count={}", self.model, texts.len() - ); - - let body = serde_json::json!({ - "model": self.model, - "texts": texts, - "input_type": "search_document", - "embedding_types": ["float"], - }); - - // Retry loop: handles 429 Too Many Requests and 503 Service Unavailable - // with Retry-After–aware exponential backoff. - for attempt in 0..=MAX_429_RETRIES { - // Proactively gate every outbound attempt (initial + retries) against - // the per-endpoint rate budget. The chokepoint must sit inside the - // loop: a single pre-loop acquire would let retried 429/503 attempts - // bypass token consumption and let concurrent callers blow past the - // cap, ironically triggering more 429s. Token consumption tracks the - // number of HTTP attempts (1 + retries actually executed). Loopback - // endpoints are exempt (see `rate_limit`). - super::rate_limit::acquire_embedding_slot(&self.base_url).await; - - let resp = self - .http_client() - .post(&url) - .header("Content-Type", "application/json") - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await?; - - let status = resp.status(); - - // Retry on 429 and 503 — both can carry a Retry-After header. - let is_retryable = status.as_u16() == 429 || status.as_u16() == 503; - - if is_retryable && attempt < MAX_429_RETRIES { - // Read Retry-After before consuming the body. - let retry_after_val = resp - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_owned()); - - let body_text = resp.text().await.unwrap_or_default(); - tracing::debug!( - target: "embeddings.cohere", - "[embeddings] cohere {} body on retry: {body_text}", - status.as_u16() - ); - - let delay_ms = backoff_ms_for_attempt(attempt, retry_after_val.as_deref()); - - tracing::debug!( - target: "embeddings.cohere", - "[embeddings] cohere {}, retrying in {}ms (attempt {}/{})", - status.as_u16(), delay_ms, attempt + 1, MAX_429_RETRIES - ); - - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - continue; - } - - if !status.is_success() { - let text = resp.text().await.unwrap_or_default(); - let message = format!("Cohere embed API error ({status}): {text}"); - crate::core::observability::report_error_or_expected( - &message, - "embeddings", - "cohere_embed", - &[("model", self.model.as_str()), ("failure", "non_2xx")], - ); - anyhow::bail!(message); - } - - let payload: CohereEmbedResponse = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("Cohere embed response parse failed: {e}"))?; - - let embeddings = payload.embeddings.float; - - if embeddings.len() != texts.len() { - anyhow::bail!( - "Cohere embed count mismatch: sent {} texts, got {} embeddings", - texts.len(), - embeddings.len() - ); - } - - for (i, vec) in embeddings.iter().enumerate() { - if self.dims > 0 && vec.len() != self.dims { - anyhow::bail!( - "Cohere embed dimension mismatch at index {i}: expected {}, got {}", - self.dims, - vec.len() - ); - } - } - - tracing::debug!( - target: "embeddings.cohere", - "[cohere] embed success: model={}, count={}, dims={}", - self.model, embeddings.len(), - embeddings.first().map(|v| v.len()).unwrap_or(0) - ); - - return Ok(embeddings); - } - - // The loop always exits via `return Ok(...)`, `bail!(...)`, or - // `continue`; this point is structurally unreachable. On the final - // attempt (`attempt == MAX_429_RETRIES`) the retryable guard is false - // and execution falls into the non-2xx branch above, which bails with - // the body-bearing format "Cohere embed API error (429 ...): " — - // that format preserves the "(429 " substring required by the - // TransientUpstreamHttp classifier in core::observability. - unreachable!("cohere embed retry loop must exit via return or bail") - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn name_and_defaults() { - let p = CohereEmbedding::new("test-key", "", 0); - assert_eq!(p.name(), "cohere"); - assert_eq!(p.model_id(), COHERE_DEFAULT_MODEL); - assert_eq!(p.dimensions(), COHERE_DEFAULT_DIMS); - } - - #[test] - fn custom_model() { - let p = CohereEmbedding::new("k", "embed-multilingual-v3.0", 1024); - assert_eq!(p.model_id(), "embed-multilingual-v3.0"); - } - - #[test] - fn signature_format() { - let p = CohereEmbedding::new("k", "embed-english-v3.0", 1024); - assert_eq!( - p.signature(), - "provider=cohere;model=embed-english-v3.0;dims=1024" - ); - } - - #[tokio::test] - async fn embed_empty_returns_empty() { - let p = CohereEmbedding::new("k", "", 0); - assert!(p.embed(&[]).await.unwrap().is_empty()); - } - - /// Missing / whitespace-only API key bails before any HTTP request with the - /// classifiable "API key not set" wording (TAURI-RUST-52S). `with_base_url` - /// points at an address nothing is listening on, so the assertion that the - /// error is the key-guard message — not a connection error — proves no - /// request was attempted. - #[tokio::test] - async fn embed_missing_api_key_bails_without_request() { - for key in ["", " "] { - let p = CohereEmbedding::new(key, "embed-english-v3.0", 1024) - .with_base_url("http://127.0.0.1:1"); - let err = p.embed(&["hello"]).await.unwrap_err().to_string(); - assert!( - err.contains("API key not set"), - "expected key-guard message for key {key:?}, got: {err}" - ); - } - } - - // ── 429 backoff tests ────────────────────────────────────── - - use axum::{http::StatusCode, routing::post, Router}; - use std::{ - net::SocketAddr, - sync::{Arc, Mutex}, - }; - - async fn start_mock(app: Router) -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr: SocketAddr = listener.local_addr().unwrap(); - tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - format!("http://127.0.0.1:{}", addr.port()) - } - - /// Cohere 429 then success — verifies retry recovers. - /// - /// The mock returns 429 (with Retry-After: 0 for zero real-wall-clock delay) - /// twice, then 200 on the third call. The real `CohereEmbedding::embed` is - /// driven via `with_base_url` pointing at the axum mock server. - #[tokio::test] - async fn cohere_embed_429_then_success() { - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v2/embed", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - if *n <= 2 { - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .header("Retry-After", "0") - .body(axum::body::Body::from(r#"{"message":"rate limited"}"#)) - .unwrap() - } else { - axum::response::Response::builder() - .status(StatusCode::OK) - .header("Content-Type", "application/json") - .body(axum::body::Body::from( - r#"{"embeddings":{"float":[[0.1,0.2]]}}"#, - )) - .unwrap() - } - } - }), - ); - - let base_url = start_mock(app).await; - let p = CohereEmbedding::new("test-key", "embed-english-v3.0", 2).with_base_url(&base_url); - - let result = p.embed(&["hello"]).await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(*counter.lock().unwrap(), 3, "should have taken 3 requests"); - } - - /// Cohere 429 indefinitely — verify bail with canonical message after retry - /// cap, and that exactly `MAX_429_RETRIES + 1` requests were made. - #[tokio::test] - async fn cohere_embed_429_indefinite_bails_after_cap() { - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v2/embed", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .header("Retry-After", "0") - .body(axum::body::Body::from(r#"{"message":"always limited"}"#)) - .unwrap() - } - }), - ); - - let base_url = start_mock(app).await; - let p = CohereEmbedding::new("test-key", "embed-english-v3.0", 2).with_base_url(&base_url); - - let err = p.embed(&["hello"]).await.unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("429"), - "should contain 429 in error message: {msg}" - ); - // MAX_429_RETRIES retries + 1 initial = MAX_429_RETRIES + 1 total requests - assert_eq!( - *counter.lock().unwrap(), - MAX_429_RETRIES + 1, - "should make exactly MAX_429_RETRIES+1 requests" - ); - } -} diff --git a/src/openhuman/embeddings/cohere_adapter.rs b/src/openhuman/embeddings/cohere_adapter.rs new file mode 100644 index 000000000..1ff51e7a9 --- /dev/null +++ b/src/openhuman/embeddings/cohere_adapter.rs @@ -0,0 +1,51 @@ +//! Compatibility wrapper for tinyagents' Cohere model. + +use async_trait::async_trait; +use tinyagents::harness::embeddings::{CohereEmbeddingModel, EmbeddingModel}; + +use super::EmbeddingProvider; + +pub struct CohereEmbedding { + inner: CohereEmbeddingModel, +} + +impl CohereEmbedding { + pub fn new(api_key: &str, model: &str, dimensions: usize) -> Self { + Self { + inner: CohereEmbeddingModel::new(api_key) + .with_model(model) + .with_dimensions(dimensions), + } + } + + pub fn with_base_url(mut self, base_url: impl Into) -> Self { + self.inner = self.inner.with_base_url(base_url); + self + } +} + +#[async_trait] +impl EmbeddingProvider for CohereEmbedding { + fn name(&self) -> &str { + self.inner.name() + } + fn model_id(&self) -> &str { + self.inner.model_id() + } + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + fn signature(&self) -> String { + self.inner.signature() + } + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let texts = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + self.inner + .embed(&texts) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} diff --git a/src/openhuman/embeddings/factory.rs b/src/openhuman/embeddings/factory.rs index 8eec9f635..8ba458192 100644 --- a/src/openhuman/embeddings/factory.rs +++ b/src/openhuman/embeddings/factory.rs @@ -5,10 +5,26 @@ use std::sync::Arc; use super::cloud::{ OpenHumanCloudEmbedding, DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, DEFAULT_CLOUD_EMBEDDING_MODEL, }; -use super::cohere::CohereEmbedding; -use super::provider_trait::EmbeddingProvider; -use super::voyage::VoyageEmbedding; -use super::{NoopEmbedding, OllamaEmbedding, OpenAiEmbedding}; +use super::provider_trait::{EmbeddingProvider, TinyAgentsEmbeddingProvider}; +use tinyagents::harness::embeddings::{ + CohereEmbeddingModel, NoopEmbeddingModel, OllamaEmbeddingModel, OpenAiEmbeddingModel, + VoyageEmbeddingModel, +}; + +fn openai_model( + base_url: &str, + api_key: &str, + model: &str, + dims: usize, + required_key: bool, +) -> OpenAiEmbeddingModel { + OpenAiEmbeddingModel::new(api_key) + .with_base_url(base_url) + .with_model(model) + .with_dimensions(dims) + .with_send_dimensions(model_supports_dimensions(model)) + .with_required_api_key(required_key) +} /// Whether to send the OpenAI `dimensions` request-body parameter for this /// model. Only the `text-embedding-3-*` family honors it (it's how 3-large is @@ -43,25 +59,39 @@ pub fn create_embedding_provider( "cloud" | "managed" => Ok(Box::new(OpenHumanCloudEmbedding::new( None, None, true, model, dims, ))), - "voyage" => Ok(Box::new(VoyageEmbedding::new("", model, dims))), + "voyage" => Ok(TinyAgentsEmbeddingProvider::boxed( + VoyageEmbeddingModel::with_options( + "", + model, + dims, + tinyagents::harness::embeddings::VOYAGE_API_BASE, + ), + )), "ollama" => { let base_url = crate::openhuman::inference::local::ollama_base_url(); - Ok(Box::new(OllamaEmbedding::try_new(&base_url, model, dims)?)) - } - "openai" => Ok(Box::new( - OpenAiEmbedding::new("https://api.openai.com", "", model, dims) - .with_send_dimensions(model_supports_dimensions(model)) - .with_required_api_key(true), - )), - "cohere" => Ok(Box::new(CohereEmbedding::new("", model, dims))), - name if name.starts_with("custom:") => { - let base_url = name.strip_prefix("custom:").unwrap_or(""); - Ok(Box::new( - OpenAiEmbedding::new(base_url, "", model, dims) - .with_send_dimensions(model_supports_dimensions(model)), + Ok(TinyAgentsEmbeddingProvider::boxed( + OllamaEmbeddingModel::try_new(&base_url, model, dims)?, )) } - "none" => Ok(Box::new(NoopEmbedding)), + "openai" => Ok(TinyAgentsEmbeddingProvider::boxed(openai_model( + "https://api.openai.com", + "", + model, + dims, + true, + ))), + "cohere" => Ok(TinyAgentsEmbeddingProvider::boxed( + CohereEmbeddingModel::new("") + .with_model(model) + .with_dimensions(dims), + )), + name if name.starts_with("custom:") => { + let base_url = name.strip_prefix("custom:").unwrap_or(""); + Ok(TinyAgentsEmbeddingProvider::boxed(openai_model( + base_url, "", model, dims, false, + ))) + } + "none" => Ok(TinyAgentsEmbeddingProvider::boxed(NoopEmbeddingModel)), unknown => Err(anyhow::anyhow!( "unknown embedding provider: \"{unknown}\". \ Supported: \"managed\", \"voyage\", \"openai\", \"cohere\", \ @@ -85,32 +115,45 @@ pub fn create_embedding_provider_with_credentials( "cloud" | "managed" => Ok(Box::new(OpenHumanCloudEmbedding::new( None, None, true, model, dims, ))), - "voyage" => Ok(Box::new(VoyageEmbedding::new(api_key, model, dims))), + "voyage" => Ok(TinyAgentsEmbeddingProvider::boxed( + VoyageEmbeddingModel::with_options( + api_key, + model, + dims, + tinyagents::harness::embeddings::VOYAGE_API_BASE, + ), + )), "ollama" => { let base_url = crate::openhuman::inference::local::ollama_base_url(); - Ok(Box::new(OllamaEmbedding::try_new(&base_url, model, dims)?)) + Ok(TinyAgentsEmbeddingProvider::boxed( + OllamaEmbeddingModel::try_new(&base_url, model, dims)?, + )) } - "openai" => Ok(Box::new( - OpenAiEmbedding::new("https://api.openai.com", api_key, model, dims) - .with_send_dimensions(model_supports_dimensions(model)) - .with_required_api_key(true), + "openai" => Ok(TinyAgentsEmbeddingProvider::boxed(openai_model( + "https://api.openai.com", + api_key, + model, + dims, + true, + ))), + "cohere" => Ok(TinyAgentsEmbeddingProvider::boxed( + CohereEmbeddingModel::new(api_key) + .with_model(model) + .with_dimensions(dims), )), - "cohere" => Ok(Box::new(CohereEmbedding::new(api_key, model, dims))), "custom" => { let url = custom_endpoint.unwrap_or(""); - Ok(Box::new( - OpenAiEmbedding::new(url, api_key, model, dims) - .with_send_dimensions(model_supports_dimensions(model)), - )) + Ok(TinyAgentsEmbeddingProvider::boxed(openai_model( + url, api_key, model, dims, false, + ))) } name if name.starts_with("custom:") => { let url = custom_endpoint.unwrap_or_else(|| name.strip_prefix("custom:").unwrap_or("")); - Ok(Box::new( - OpenAiEmbedding::new(url, api_key, model, dims) - .with_send_dimensions(model_supports_dimensions(model)), - )) + Ok(TinyAgentsEmbeddingProvider::boxed(openai_model( + url, api_key, model, dims, false, + ))) } - "none" => Ok(Box::new(NoopEmbedding)), + "none" => Ok(TinyAgentsEmbeddingProvider::boxed(NoopEmbeddingModel)), unknown => Err(anyhow::anyhow!( "unknown embedding provider: \"{unknown}\". \ Supported: \"managed\", \"voyage\", \"openai\", \"cohere\", \ @@ -137,5 +180,7 @@ pub fn default_embedding_provider() -> Arc { /// Returns the local Ollama-backed embedding provider. Only used when the /// caller has explicitly opted into local-only embeddings. pub fn default_local_embedding_provider() -> Arc { - Arc::new(OllamaEmbedding::default()) + Arc::new(TinyAgentsEmbeddingProvider::new( + OllamaEmbeddingModel::default(), + )) } diff --git a/src/openhuman/embeddings/mod.rs b/src/openhuman/embeddings/mod.rs index 4ac4be190..4854da55e 100644 --- a/src/openhuman/embeddings/mod.rs +++ b/src/openhuman/embeddings/mod.rs @@ -13,17 +13,36 @@ //! - **Noop**: A fallback provider for keyword-only search. pub mod catalog; +#[path = "cloud_adapter.rs"] pub mod cloud; +#[path = "cohere_adapter.rs"] pub mod cohere; mod factory; pub mod noop; +#[path = "ollama_adapter.rs"] pub mod ollama; +#[path = "openai_adapter.rs"] pub mod openai; mod provider_trait; -pub mod rate_limit; -pub mod retry_after; +pub mod rate_limit { + pub use tinyagents::harness::embeddings::{ + rate_limit as embedding_rate_limit, set_rate_limit as set_embedding_rate_limit, + DEFAULT_REQUESTS_PER_MINUTE as DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN, + }; + + pub async fn acquire_embedding_slot(base_url: &str) { + tinyagents::harness::embeddings::acquire(base_url).await; + } +} +pub mod retry_after { + pub use tinyagents::harness::embeddings::{ + backoff_ms_for_attempt, parse_retry_after_ms, BASE_BACKOFF_MS, MAX_BACKOFF_MS, + MAX_RETRIES as MAX_429_RETRIES, + }; +} mod rpc; mod schemas; +#[path = "voyage_adapter.rs"] pub mod voyage; pub use catalog::non_embedding_model_reason; @@ -41,15 +60,18 @@ pub(crate) use factory::model_supports_dimensions; // #002 FR-015: the memory-tree OpenAI-compat embedder reuses the same key // resolution the embeddings RPC uses, so there is one source of truth. pub use noop::NoopEmbedding; -pub use ollama::{OllamaEmbedding, DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL}; +pub use ollama::OllamaEmbedding; pub use openai::OpenAiEmbedding; -pub use provider_trait::{format_embedding_signature, EmbeddingProvider}; +pub use provider_trait::{ + format_embedding_signature, EmbeddingProvider, TinyAgentsEmbeddingProvider, +}; pub use rpc::provider_from_config; pub(crate) use rpc::resolve_api_key; pub use schemas::{ all_controller_schemas as all_embeddings_controller_schemas, all_registered_controllers as all_embeddings_registered_controllers, }; +pub use tinyagents::harness::embeddings::{DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL}; #[cfg(test)] #[path = "mod_tests.rs"] diff --git a/src/openhuman/embeddings/ollama.rs b/src/openhuman/embeddings/ollama.rs deleted file mode 100644 index b90013270..000000000 --- a/src/openhuman/embeddings/ollama.rs +++ /dev/null @@ -1,470 +0,0 @@ -//! Ollama-based embedding provider. -//! -//! Calls the local Ollama server's `/api/embed` endpoint for embeddings. -//! This is the preferred local provider: Ollama handles model management, -//! quantization, and GPU acceleration (Metal on macOS, CUDA on Linux/Windows). -//! -//! Default model: `bge-m3` (1024 dimensions). Aligned with the memory -//! tree's fixed on-disk format (`EMBEDDING_DIM=1024`) and the cloud -//! Voyage default (`embedding-v1`, 1024 dims) so embeddings produced by -//! either path are interchangeable. - -use async_trait::async_trait; - -use super::EmbeddingProvider; - -/// Default Ollama base URL. -pub const DEFAULT_OLLAMA_URL: &str = "http://localhost:11434"; - -/// Default embedding model for Ollama. 1024-dim to match the memory -/// tree's fixed on-disk format and the cloud Voyage default. -pub const DEFAULT_OLLAMA_MODEL: &str = "bge-m3"; - -/// Default dimensions for `bge-m3`. -pub const DEFAULT_OLLAMA_DIMENSIONS: usize = 1024; - -/// Embedding provider backed by a local Ollama instance. -/// -/// Ollama must be running and have the configured model pulled. -/// On first embed call, if the model isn't available, Ollama will -/// auto-pull it (this may take a moment on first use). -#[derive(Debug)] -pub struct OllamaEmbedding { - base_url: String, - model: String, - dims: usize, -} - -impl OllamaEmbedding { - /// Creates a new Ollama embedding provider. - /// - /// - `base_url`: Ollama server URL (default: `http://localhost:11434`) - /// - `model`: Model name (default: `bge-m3`) - /// - `dims`: Expected embedding dimensions (default: 1024) - pub fn try_new(base_url: &str, model: &str, dims: usize) -> anyhow::Result { - let base_url = Self::normalize_base_url(base_url)?; - let model = Self::normalize_model(model)?; - let dims = if dims == 0 { - DEFAULT_OLLAMA_DIMENSIONS - } else { - dims - }; - - tracing::debug!( - target: "embeddings.ollama", - "[embeddings] OllamaEmbedding created: url={base_url}, model={model}, dims={dims}" - ); - - Ok(Self { - base_url, - model, - dims, - }) - } - - /// Creates a new Ollama embedding provider, panicking if the configuration - /// is invalid. Prefer [`Self::try_new`] at runtime boundaries. - pub fn new(base_url: &str, model: &str, dims: usize) -> Self { - Self::try_new(base_url, model, dims).expect("invalid Ollama embedding configuration") - } - - /// Returns the configured base URL. - pub fn base_url(&self) -> &str { - &self.base_url - } - - /// Returns the configured model name. - pub fn model(&self) -> &str { - &self.model - } - - /// Build an HTTP client with proxy support. - fn http_client(&self) -> reqwest::Client { - crate::openhuman::config::build_runtime_proxy_client("embeddings.ollama") - } - - fn normalize_base_url(base_url: &str) -> anyhow::Result { - let raw = if base_url.trim().is_empty() { - DEFAULT_OLLAMA_URL - } else { - base_url.trim() - }; - - let url = reqwest::Url::parse(raw) - .map_err(|e| anyhow::anyhow!("invalid Ollama base_url `{raw}`: {e}"))?; - if !matches!(url.scheme(), "http" | "https") { - anyhow::bail!("invalid Ollama base_url `{raw}`: expected an http:// or https:// URL"); - } - if !url.username().is_empty() || url.password().is_some() { - anyhow::bail!( - "invalid Ollama base_url `{raw}`: configure the server root without credentials" - ); - } - if url.query().is_some() || url.fragment().is_some() { - anyhow::bail!( - "invalid Ollama base_url `{raw}`: query strings and fragments are not supported" - ); - } - - let segments: Vec = url - .path_segments() - .map(|parts| { - parts - .filter(|part| !part.is_empty()) - .map(|part| part.to_ascii_lowercase()) - .collect() - }) - .unwrap_or_default(); - let has_api_suffix = segments.iter().any(|part| part == "api" || part == "v1"); - let is_chat_completions_endpoint = segments.len() >= 2 - && segments[segments.len() - 2] == "chat" - && segments[segments.len() - 1] == "completions"; - if has_api_suffix || is_chat_completions_endpoint { - anyhow::bail!( - "invalid Ollama base_url `{raw}`: configure the Ollama server root \ - (for example {DEFAULT_OLLAMA_URL}), not an API endpoint such as \ - /api, /v1, or /chat/completions" - ); - } - - Ok(url.as_str().trim_end_matches('/').to_string()) - } - - fn normalize_model(model: &str) -> anyhow::Result { - let model = if model.trim().is_empty() { - DEFAULT_OLLAMA_MODEL.to_string() - } else { - model.trim().to_string() - }; - if model.to_ascii_lowercase().starts_with("local-") { - anyhow::bail!( - "invalid Ollama embedding model `{model}`: `local-*` model IDs are virtual \ - routing aliases. Configure a real Ollama embedding model such as \ - `{DEFAULT_OLLAMA_MODEL}`." - ); - } - Ok(model) - } - - /// The embed endpoint URL. - fn embed_url(&self) -> anyhow::Result { - let _ = reqwest::Url::parse(&self.base_url) - .map_err(|e| anyhow::anyhow!("invalid Ollama base_url `{}`: {e}", self.base_url))?; - Ok(format!("{}/api/embed", self.base_url)) - } - - /// Sends a single text to Ollama and returns either the embedding, or - /// `None` if Ollama produced the NaN-encoding 500 wire shape for that - /// individual input. Any other failure (transport error, non-2xx without - /// NaN signature, malformed JSON, dimension mismatch, count mismatch) - /// surfaces as an `Err` so genuine bugs are still loud. - /// - /// Used by [`Self::embed_per_text_fallback`] to recover a batch that - /// failed wholesale on NaN — see TAURI-RUST-AZ. - async fn embed_one_with_nan_recovery(&self, text: &str) -> anyhow::Result>> { - let resp = self - .http_client() - .post(self.embed_url()?) - .json(&OllamaEmbedRequest { - model: self.model.clone(), - input: vec![text.to_string()], - }) - .send() - .await - .map_err(|e| { - anyhow::anyhow!( - "ollama embed request failed (is Ollama running at {}?): {e}", - self.base_url - ) - })?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - let detail = body.trim(); - // Per-text NaN: substitute empty vector, log once. - if status.as_u16() == 500 && is_nan_encode_error(&body) { - tracing::warn!( - target: "embeddings.ollama", - "[embeddings] ollama produced NaN for a single text (model={}); \ - substituting empty embedding (downstream blob will be 0 bytes)", - self.model - ); - return Ok(None); - } - anyhow::bail!( - "ollama embed failed with status {status}{}", - if detail.is_empty() { - String::new() - } else { - format!(": {detail}") - } - ); - } - - let payload: OllamaEmbedResponse = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("ollama embed response parse failed: {e}"))?; - if payload.embeddings.len() != 1 { - anyhow::bail!( - "ollama embed count mismatch: sent 1 text, got {} embeddings", - payload.embeddings.len() - ); - } - let v = payload.embeddings.into_iter().next().unwrap(); - if v.len() != self.dims { - anyhow::bail!( - "ollama embed dimension mismatch: expected {}, got {}", - self.dims, - v.len() - ); - } - Ok(Some(v)) - } - - /// Recovery path invoked when a batch request fails with the Ollama - /// NaN-encoding 500 wire shape. Re-sends each live input one at a time; - /// per-text NaN failures are filled with `Vec::new()` so the overall - /// result vector still has length `total_len` and aligns with the - /// caller's original `texts` slice (same convention as blank inputs). - async fn embed_per_text_fallback( - &self, - total_len: usize, - live: &[(usize, String)], - ) -> anyhow::Result>> { - let mut result = vec![Vec::new(); total_len]; - let mut nan_substituted = 0usize; - for (orig_idx, text) in live { - match self.embed_one_with_nan_recovery(text).await? { - Some(v) => result[*orig_idx] = v, - None => nan_substituted += 1, - } - } - tracing::warn!( - target: "embeddings.ollama", - "[embeddings] ollama per-text fallback complete: {} of {} inputs returned NaN and \ - were substituted with empty embeddings", - nan_substituted, - live.len() - ); - Ok(result) - } -} - -impl Default for OllamaEmbedding { - fn default() -> Self { - Self::try_new( - DEFAULT_OLLAMA_URL, - DEFAULT_OLLAMA_MODEL, - DEFAULT_OLLAMA_DIMENSIONS, - ) - .expect("default Ollama embedding configuration must be valid") - } -} - -/// Ollama `/api/embed` request body. -#[derive(serde::Serialize)] -struct OllamaEmbedRequest { - model: String, - input: Vec, -} - -/// Ollama `/api/embed` response body. -#[derive(serde::Deserialize)] -struct OllamaEmbedResponse { - #[serde(default)] - embeddings: Vec>, -} - -/// Detects the Ollama-side NaN-encoding 500 wire shape: -/// `{"error":"failed to encode response: json: unsupported value: NaN"}` -/// -/// Ollama produces NaN for some inputs (model bug / numerically degenerate -/// token sequences) and then fails to encode the response as JSON. One bad -/// input poisons the entire batch. -/// -/// See TAURI-RUST-AZ on Sentry. -fn is_nan_encode_error(body: &str) -> bool { - body.to_ascii_lowercase().contains("unsupported value: nan") -} - -#[async_trait] -impl EmbeddingProvider for OllamaEmbedding { - fn name(&self) -> &str { - "ollama" - } - - fn model_id(&self) -> &str { - &self.model - } - - fn dimensions(&self) -> usize { - self.dims - } - - /// Sends texts to Ollama's embed API. - /// - /// Blank/whitespace-only entries are skipped for the remote call but their - /// positions in the result are preserved as zero-vectors so the returned - /// `Vec` always has the same length as `texts`. - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - if texts.is_empty() { - return Ok(Vec::new()); - } - - // Build a list of (original_index, trimmed_text) for non-blank entries. - let live: Vec<(usize, String)> = texts - .iter() - .enumerate() - .filter_map(|(i, t)| { - let trimmed = t.trim().to_string(); - if trimmed.is_empty() { - None - } else { - Some((i, trimmed)) - } - }) - .collect(); - - if live.is_empty() { - // All entries were blank — return zero-vectors. - return Ok(vec![Vec::new(); texts.len()]); - } - - let input: Vec = live.iter().map(|(_, t)| t.clone()).collect(); - - tracing::debug!( - target: "embeddings.ollama", - "[embeddings] sending {} text(s) to ollama model={} ({} blank skipped)", - input.len(), self.model, texts.len() - input.len() - ); - - let resp = self - .http_client() - .post(self.embed_url()?) - .json(&OllamaEmbedRequest { - model: self.model.clone(), - input: input.clone(), - }) - .send() - .await - .map_err(|e| { - let message = format!( - "ollama embed request failed (is Ollama running at {}?): {e}", - self.base_url - ); - crate::core::observability::report_error_or_expected( - &message, - "embeddings", - "ollama_embed", - &[("model", self.model.as_str()), ("failure", "transport")], - ); - anyhow::anyhow!(message) - })?; - - if !resp.status().is_success() { - let status = resp.status(); - let status_str = status.as_u16().to_string(); - let body = resp.text().await.unwrap_or_default(); - let detail = body.trim(); - let message = format!( - "ollama embed failed with status {status}{}", - if detail.is_empty() { - String::new() - } else { - format!(": {detail}") - } - ); - - // TAURI-RUST-AZ: Ollama returns 500 `unsupported value: NaN` when the - // model produces NaN for some input in the batch. One bad input - // poisons the whole batch under the default code path. Recover by - // re-sending each live input individually; per-text NaN failures - // are replaced with an empty embedding (same convention used for - // blank inputs above), so the rest of the batch still succeeds. - if status.as_u16() == 500 && is_nan_encode_error(&body) { - tracing::warn!( - target: "embeddings.ollama", - "[embeddings] ollama returned NaN-encoding 500 for batch of {} (model={}); \ - falling back to per-text requests", - live.len(), - self.model - ); - crate::core::observability::report_error_or_expected( - &message, - "embeddings", - "ollama_embed", - &[ - ("model", self.model.as_str()), - ("status", status_str.as_str()), - ("failure", "nan_batch_recovered"), - ], - ); - // Single-text NaN: re-issuing the same request would only - // reproduce the failure. Skip the wasted round-trip and - // substitute an empty embedding directly. - if live.len() == 1 { - return Ok(vec![Vec::new(); texts.len()]); - } - return self.embed_per_text_fallback(texts.len(), &live).await; - } - - crate::core::observability::report_error_or_expected( - &message, - "embeddings", - "ollama_embed", - &[ - ("model", self.model.as_str()), - ("status", status_str.as_str()), - ("failure", "non_2xx"), - ], - ); - anyhow::bail!(message); - } - - let payload: OllamaEmbedResponse = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("ollama embed response parse failed: {e}"))?; - - // Validate response count matches what we sent. - if payload.embeddings.len() != input.len() { - anyhow::bail!( - "ollama embed count mismatch: sent {} texts, got {} embeddings", - input.len(), - payload.embeddings.len() - ); - } - - // Validate dimensions on every returned vector. - for (i, vec) in payload.embeddings.iter().enumerate() { - if vec.len() != self.dims { - anyhow::bail!( - "ollama embed dimension mismatch at index {i}: expected {}, got {}", - self.dims, - vec.len() - ); - } - } - - tracing::debug!( - target: "embeddings.ollama", - "[embeddings] received {} embeddings, dims={}", - payload.embeddings.len(), - self.dims - ); - - // Reconstruct full-length result with zero-vectors for blank positions. - let mut result = vec![Vec::new(); texts.len()]; - for ((orig_idx, _), embedding) in live.iter().zip(payload.embeddings) { - result[*orig_idx] = embedding; - } - - Ok(result) - } -} - -#[cfg(test)] -#[path = "ollama_tests.rs"] -mod tests; diff --git a/src/openhuman/embeddings/ollama_adapter.rs b/src/openhuman/embeddings/ollama_adapter.rs new file mode 100644 index 000000000..fe48175c6 --- /dev/null +++ b/src/openhuman/embeddings/ollama_adapter.rs @@ -0,0 +1,68 @@ +//! Compatibility wrapper for tinyagents' Ollama model. + +use async_trait::async_trait; +use tinyagents::harness::embeddings::{EmbeddingModel, OllamaEmbeddingModel}; + +use super::EmbeddingProvider; + +pub use tinyagents::harness::embeddings::{ + DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL, DEFAULT_OLLAMA_URL, +}; + +pub struct OllamaEmbedding { + inner: OllamaEmbeddingModel, +} + +impl OllamaEmbedding { + pub fn try_new(base_url: &str, model: &str, dimensions: usize) -> anyhow::Result { + Ok(Self { + inner: OllamaEmbeddingModel::try_new(base_url, model, dimensions)?, + }) + } + + pub fn new(base_url: &str, model: &str, dimensions: usize) -> Self { + Self::try_new(base_url, model, dimensions).expect("invalid Ollama embedding configuration") + } + + pub fn base_url(&self) -> &str { + self.inner.base_url() + } + + pub fn model(&self) -> &str { + self.inner.model() + } +} + +impl Default for OllamaEmbedding { + fn default() -> Self { + Self { + inner: OllamaEmbeddingModel::default(), + } + } +} + +#[async_trait] +impl EmbeddingProvider for OllamaEmbedding { + fn name(&self) -> &str { + self.inner.name() + } + fn model_id(&self) -> &str { + self.inner.model_id() + } + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + fn signature(&self) -> String { + self.inner.signature() + } + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let texts = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + self.inner + .embed(&texts) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} diff --git a/src/openhuman/embeddings/ollama_tests.rs b/src/openhuman/embeddings/ollama_tests.rs deleted file mode 100644 index d0b5018cc..000000000 --- a/src/openhuman/embeddings/ollama_tests.rs +++ /dev/null @@ -1,635 +0,0 @@ -use super::*; -use axum::{extract::Json, http::StatusCode, routing::post, Router}; -use std::net::SocketAddr; - -/// Spin up a local axum server and return its base URL. -async fn start_mock(app: Router) -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr: SocketAddr = listener.local_addr().unwrap(); - tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - format!("http://127.0.0.1:{}", addr.port()) -} - -// ── Constructor ────────────────────────────────────────── - -#[test] -fn defaults() { - let p = OllamaEmbedding::default(); - assert_eq!(p.base_url, DEFAULT_OLLAMA_URL); - assert_eq!(p.model, DEFAULT_OLLAMA_MODEL); - assert_eq!(p.dims, DEFAULT_OLLAMA_DIMENSIONS); -} - -#[test] -fn name_is_ollama() { - let p = OllamaEmbedding::default(); - assert_eq!(p.name(), "ollama"); -} - -#[test] -fn custom_values() { - let p = OllamaEmbedding::new("http://gpu-box:11434/", "mxbai-embed-large", 1024); - assert_eq!(p.base_url, "http://gpu-box:11434"); - assert_eq!(p.model, "mxbai-embed-large"); - assert_eq!(p.dims, 1024); -} - -#[test] -fn empty_values_use_defaults() { - let p = OllamaEmbedding::new("", "", 0); - assert_eq!(p.base_url, DEFAULT_OLLAMA_URL); - assert_eq!(p.model, DEFAULT_OLLAMA_MODEL); - assert_eq!(p.dims, DEFAULT_OLLAMA_DIMENSIONS); -} - -#[test] -fn whitespace_only_values_use_defaults() { - let p = OllamaEmbedding::new(" ", " ", 0); - assert_eq!(p.base_url, DEFAULT_OLLAMA_URL); - assert_eq!(p.model, DEFAULT_OLLAMA_MODEL); -} - -#[test] -fn trailing_slash_stripped() { - let p = OllamaEmbedding::new("http://host:1234/", "m", 1); - assert_eq!(p.base_url, "http://host:1234"); -} - -#[test] -fn base_url_edge_cases_build_embed_url() { - let cases = [ - ("http://host:11434/", "http://host:11434/api/embed"), - ("http://[::1]:11434", "http://[::1]:11434/api/embed"), - ("http://host", "http://host/api/embed"), - ]; - - for (base_url, expected) in cases { - let p = OllamaEmbedding::try_new(base_url, "m", 1).unwrap(); - assert_eq!(p.embed_url().unwrap(), expected); - } -} - -#[test] -fn rejects_api_endpoint_base_urls() { - for base_url in [ - "http://host:11434/v1", - "http://host:11434/api", - "http://host:11434/api/embed", - "http://host:11434/v1/chat/completions", - "http://host:11434/chat/completions", - ] { - let err = OllamaEmbedding::try_new(base_url, "m", 1).unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("Ollama server root"), - "should reject pre-suffixed base URL {base_url}: {msg}" - ); - } -} - -#[test] -fn rejects_credentialed_base_urls() { - let err = OllamaEmbedding::try_new("http://user:pass@host:11434", "m", 1).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("without credentials"), "msg: {msg}"); -} - -#[test] -fn rejects_virtual_local_model_ids() { - let err = OllamaEmbedding::try_new("http://host:11434", "local-v1", 768).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("local-*"), "msg: {msg}"); - assert!(msg.contains(DEFAULT_OLLAMA_MODEL), "msg: {msg}"); -} - -#[test] -fn model_trimmed() { - let p = OllamaEmbedding::new("", " nomic-embed-text ", 768); - assert_eq!(p.model, "nomic-embed-text"); -} - -#[test] -fn embed_url_format() { - let p = OllamaEmbedding::default(); - assert_eq!(p.embed_url().unwrap(), "http://localhost:11434/api/embed"); -} - -#[test] -fn accessor_methods() { - let p = OllamaEmbedding::new("http://x:1", "m", 42); - assert_eq!(p.base_url(), "http://x:1"); - assert_eq!(p.model(), "m"); - assert_eq!(p.model_id(), "m"); - assert_eq!(p.dimensions(), 42); - assert_eq!(p.signature(), "provider=ollama;model=m;dims=42"); -} - -// ── embed — empty / whitespace ────────────────────────── - -#[tokio::test] -async fn empty_input_returns_empty() { - let p = OllamaEmbedding::default(); - let result = p.embed(&[]).await.unwrap(); - assert!(result.is_empty()); -} - -#[tokio::test] -async fn whitespace_only_input_returns_zero_vecs() { - let p = OllamaEmbedding::default(); - let result = p.embed(&[" ", "\t", "\n"]).await.unwrap(); - // Length preserved, all entries are empty zero-vectors. - assert_eq!(result.len(), 3); - assert!(result.iter().all(|v| v.is_empty())); -} - -// ── embed — positional alignment ──────────────────────── - -#[tokio::test] -async fn embed_preserves_positions_for_blanks() { - let app = Router::new().route( - "/api/embed", - post(|Json(body): Json| async move { - let inputs = body["input"].as_array().unwrap(); - // Server receives only non-blank texts. - let embeddings: Vec> = inputs.iter().map(|_| vec![1.0, 2.0]).collect(); - Json(serde_json::json!({ "embeddings": embeddings })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - // Mix of blank and real texts. - let result = p.embed(&["hello", "", " ", "world"]).await.unwrap(); - assert_eq!(result.len(), 4); - assert_eq!(result[0], vec![1.0, 2.0]); // real - assert!(result[1].is_empty()); // blank - assert!(result[2].is_empty()); // blank - assert_eq!(result[3], vec![1.0, 2.0]); // real -} - -// ── embed — successful response ───────────────────────── - -#[tokio::test] -async fn embed_success_single() { - let app = Router::new().route( - "/api/embed", - post(|Json(_body): Json| async { - Json(serde_json::json!({ - "embeddings": [[0.1, 0.2, 0.3]] - })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "test-model", 3); - - let result = p.embed(&["hello"]).await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0], vec![0.1, 0.2, 0.3]); -} - -#[tokio::test] -async fn embed_success_batch() { - let app = Router::new().route( - "/api/embed", - post(|Json(_body): Json| async { - Json(serde_json::json!({ - "embeddings": [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] - })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "test-model", 2); - - let result = p.embed(&["a", "b", "c"]).await.unwrap(); - assert_eq!(result.len(), 3); - assert_eq!(result[2], vec![5.0, 6.0]); -} - -#[tokio::test] -async fn embed_verifies_request_body() { - let app = Router::new().route( - "/api/embed", - post(|Json(body): Json| async move { - assert_eq!(body["model"], "my-model"); - let inputs = body["input"].as_array().unwrap(); - assert_eq!(inputs.len(), 1); - assert_eq!(inputs[0], "test text"); - Json(serde_json::json!({ "embeddings": [[1.0]] })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "my-model", 1); - - p.embed(&["test text"]).await.unwrap(); -} - -// ── embed — error paths ───────────────────────────────── - -#[tokio::test] -async fn embed_server_error_with_body() { - let app = Router::new().route( - "/api/embed", - post(|| async { (StatusCode::INTERNAL_SERVER_ERROR, "model crashed") }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("500"), "should contain status code: {msg}"); - assert!(msg.contains("model crashed"), "should contain body: {msg}"); -} - -#[tokio::test] -async fn embed_server_error_empty_body() { - let app = Router::new().route( - "/api/embed", - post(|| async { (StatusCode::BAD_REQUEST, "") }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("400"), "should contain status code: {msg}"); -} - -#[tokio::test] -async fn embed_count_mismatch() { - let app = Router::new().route( - "/api/embed", - post(|| async { - // Return 1 embedding even though 2 texts were sent. - Json(serde_json::json!({ "embeddings": [[1.0]] })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 1); - - let err = p.embed(&["a", "b"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("count mismatch"), "msg: {msg}"); -} - -#[tokio::test] -async fn embed_dimension_mismatch() { - let app = Router::new().route( - "/api/embed", - post(|| async { - // Return 3-dim vector when provider expects 2. - Json(serde_json::json!({ "embeddings": [[1.0, 2.0, 3.0]] })) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("dimension mismatch"), "msg: {msg}"); -} - -#[tokio::test] -async fn embed_empty_embeddings_array() { - let app = Router::new().route( - "/api/embed", - post(|| async { Json(serde_json::json!({ "embeddings": [] })) }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.to_string().contains("count mismatch")); -} - -#[tokio::test] -async fn embed_malformed_json_response() { - let app = Router::new().route( - "/api/embed", - post(|| async { (StatusCode::OK, "not json at all") }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.to_string().contains("parse failed")); -} - -#[tokio::test] -async fn embed_connection_refused() { - let p = OllamaEmbedding::new("http://127.0.0.1:1", "m", 1); - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!( - err.to_string().contains("is Ollama running"), - "should mention Ollama: {}", - err - ); -} - -// OPENHUMAN-TAURI-{GP,MA,KM,GX} + TAURI-RUST-XS wire shapes — routed through -// `report_error_or_expected` (Sentry classifier ladder) and demoted to -// `ProviderUserState` / `LocalAiCapabilityUnavailable` by the -// `is_ollama_user_config_rejection` matcher in `src/core/observability.rs`. -// GP catches the RAM-tier disable shape; XS/MA/KM/GX cover the four Ollama -// user-config rejection shapes (400 invalid-model-name, 404 model-not-found -// ×2, daemon-unreachable opt-in state). - -#[test] -fn xs_wire_shape_classifies_as_provider_user_state() { - // TAURI-RUST-XS (~376 events on self-hosted Sentry): user pointed - // embedder at a chat / vision model with a temperature suffix - // (`qwen3-vl:4b@0.7`) Ollama parses as malformed. - let msg = r#"ollama embed failed with status 400 Bad Request: {"error":"invalid model name"}"#; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::ProviderUserState), - "XS — invalid model name 400 is pure user-config, must demote" - ); -} - -#[test] -fn xs_temperature_suffix_model_classifies_as_provider_user_state() { - // Same XS shape with the temperature-suffix model id embedded in the - // body. Real wire shape from the field report. - let msg = r#"ollama embed failed with status 400 Bad Request: {"error":"invalid model name: qwen3-vl:4b@0.7"}"#; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::ProviderUserState), - "XS — temperature-suffix model id must also demote" - ); -} - -#[test] -fn ma_wire_shape_classifies_as_provider_user_state() { - // OPENHUMAN-TAURI-MA: user configured an Ollama model id that the - // local daemon hasn't pulled yet. - let msg = r#"ollama embed failed with status 404 Not Found: {"error":"model \"bge-m3\" not found, try pulling it first"}"#; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::ProviderUserState), - "MA — model-not-found 404 is pure user-state (pull required), must demote" - ); -} - -#[test] -fn km_wire_shape_classifies_as_provider_user_state() { - // OPENHUMAN-TAURI-KM: same wire shape as MA with a different model - // id + `:latest` tag. - let msg = r#"ollama embed failed with status 404 Not Found: {"error":"model \"nomic-embed-text:latest\" not found, try pulling it first"}"#; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::ProviderUserState), - "KM — pull-required 404 must demote" - ); -} - -#[test] -fn gp_wire_shape_classifies() { - let msg = - "Vision is disabled for this RAM tier. Switch to the 4-8 GB tier or above to enable it."; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::LocalAiCapabilityUnavailable), - "GP — LocalAiCapabilityUnavailable matcher must catch this" - ); -} - -#[test] -fn gx_wire_shape_classifies_as_provider_user_state() { - // OPENHUMAN-TAURI-GX: user opted into Ollama embeddings in Settings - // but the daemon isn't running on localhost:11434. - let msg = "ollama embeddings opted-in but daemon unreachable at http://localhost:11434; falling back to cloud embeddings for this session"; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - Some(crate::core::observability::ExpectedErrorKind::ProviderUserState), - "GX — daemon-unreachable opt-in state is pure user-config (start daemon)" - ); -} - -#[test] -fn ollama_500_wire_shape_stays_unexpected() { - // Server-side ollama bug — must still reach Sentry. - let msg = "ollama embed failed with status 500 Internal Server Error: model crashed"; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - None, - "real ollama server errors must still reach Sentry" - ); -} - -#[test] -fn ollama_parse_error_wire_shape_stays_unexpected() { - let msg = "ollama embed response parse failed: invalid type: expected sequence"; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - None, - "real parse bugs must still reach Sentry" - ); -} - -#[test] -fn ollama_dimension_mismatch_stays_unexpected() { - // Dimension mismatch — real bug (model dims don't match what we - // recorded for the provider in Settings), must still reach Sentry - // so we can investigate the desync. - let msg = "ollama embed dimension mismatch at index 0: expected 768, got 1024"; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - None, - "dimension-mismatch shape must still reach Sentry" - ); -} - -// ── embed — NaN-encoding 500 recovery (TAURI-RUST-AZ) ─── - -#[test] -fn is_nan_encode_error_matches_ollama_wire_shape() { - // Canonical wire shape from Sentry TAURI-RUST-AZ. - assert!(is_nan_encode_error( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"# - )); - // Case-insensitive on the "NaN" token. - assert!(is_nan_encode_error("unsupported value: nan")); - assert!(is_nan_encode_error("UNSUPPORTED VALUE: NAN")); - // Negative: unrelated 500 bodies must not match. - assert!(!is_nan_encode_error("model crashed")); - assert!(!is_nan_encode_error( - r#"{"error":"failed to encode response: json: unsupported value: +Inf"}"# - )); - assert!(!is_nan_encode_error("")); -} - -#[tokio::test] -async fn embed_nan_batch_recovers_via_per_text() { - // Ollama returns NaN-encoding 500 for any batch > 1, but succeeds on - // single-text requests. Recovery should re-issue per-text and assemble - // the full result. - let app = Router::new().route( - "/api/embed", - post(|Json(body): Json| async move { - let n = body["input"].as_array().unwrap().len(); - if n > 1 { - ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#, - ), - ) - } else { - ( - StatusCode::OK, - serde_json::json!({ "embeddings": [[1.0, 2.0]] }).to_string(), - ) - } - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let result = p.embed(&["a", "b", "c"]).await.unwrap(); - assert_eq!(result.len(), 3); - assert_eq!(result[0], vec![1.0, 2.0]); - assert_eq!(result[1], vec![1.0, 2.0]); - assert_eq!(result[2], vec![1.0, 2.0]); -} - -#[tokio::test] -async fn embed_nan_persists_per_text_substitutes_empty() { - // Every request — batch or single — returns NaN 500. Recovery - // substitutes empty vectors for each input rather than bailing. - let app = Router::new().route( - "/api/embed", - post(|| async { - ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#, - ), - ) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let result = p.embed(&["a", "b"]).await.unwrap(); - assert_eq!(result.len(), 2); - assert!(result[0].is_empty(), "NaN-failed entries must be empty vec"); - assert!(result[1].is_empty(), "NaN-failed entries must be empty vec"); -} - -#[tokio::test] -async fn embed_nan_single_text_short_circuits_to_empty() { - // Single-input NaN should NOT trigger a wasted retry — short-circuit - // straight to empty vector. - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - let count = Arc::new(AtomicUsize::new(0)); - let count_clone = count.clone(); - let app = Router::new().route( - "/api/embed", - post(move || { - let c = count_clone.clone(); - async move { - c.fetch_add(1, Ordering::SeqCst); - ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#, - ), - ) - } - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let result = p.embed(&["only-text"]).await.unwrap(); - assert_eq!(result.len(), 1); - assert!(result[0].is_empty()); - assert_eq!( - count.load(Ordering::SeqCst), - 1, - "single-text NaN must not retry — exactly one request expected" - ); -} - -#[tokio::test] -async fn embed_nan_mixed_per_text_outcomes() { - // Batch NaN-fails. Per-text retries: input "bad" returns NaN, others - // succeed. Result vector has empty slot for "bad" and embeddings for - // the rest, with positions preserved. - let app = Router::new().route( - "/api/embed", - post(|Json(body): Json| async move { - let inputs = body["input"].as_array().unwrap(); - if inputs.len() > 1 { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#, - ), - ); - } - let text = inputs[0].as_str().unwrap(); - if text == "bad" { - ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from( - r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#, - ), - ) - } else { - ( - StatusCode::OK, - serde_json::json!({ "embeddings": [[9.0, 9.0]] }).to_string(), - ) - } - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let result = p.embed(&["good1", "bad", "good2"]).await.unwrap(); - assert_eq!(result.len(), 3); - assert_eq!(result[0], vec![9.0, 9.0]); - assert!(result[1].is_empty(), "NaN entry must be empty vec"); - assert_eq!(result[2], vec![9.0, 9.0]); -} - -#[tokio::test] -async fn embed_non_nan_500_still_errors() { - // Real ollama 500s (anything that isn't the NaN wire shape) must still - // bail loudly — only NaN gets the recovery path. - let app = Router::new().route( - "/api/embed", - post(|| async { - ( - StatusCode::INTERNAL_SERVER_ERROR, - String::from("model crashed"), - ) - }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let err = p.embed(&["a", "b"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("500")); - assert!(msg.contains("model crashed")); -} - -// ── embed_one (trait default) ─────────────────────────── - -#[tokio::test] -async fn embed_one_success() { - let app = Router::new().route( - "/api/embed", - post(|| async { Json(serde_json::json!({ "embeddings": [[7.0, 8.0]] })) }), - ); - let url = start_mock(app).await; - let p = OllamaEmbedding::new(&url, "m", 2); - - let vec = p.embed_one("test").await.unwrap(); - assert_eq!(vec, vec![7.0, 8.0]); -} diff --git a/src/openhuman/embeddings/openai.rs b/src/openhuman/embeddings/openai.rs deleted file mode 100644 index 64e32319d..000000000 --- a/src/openhuman/embeddings/openai.rs +++ /dev/null @@ -1,467 +0,0 @@ -//! OpenAI-compatible embedding provider. -//! -//! Works with OpenAI, LocalAI, Ollama, and any endpoint that implements the -//! `POST /v1/embeddings` contract. - -use async_trait::async_trait; - -use super::retry_after::{backoff_ms_for_attempt, MAX_429_RETRIES}; -use super::EmbeddingProvider; - -/// Embedding provider for OpenAI and compatible APIs (e.g., LocalAI, Ollama). -pub struct OpenAiEmbedding { - base_url: String, - api_key: String, - model: String, - dims: usize, - /// When true, send `"dimensions": dims` in the request body. OpenAI's - /// `text-embedding-3-*` models honour this (Matryoshka — e.g. 3-large can - /// return 1024 instead of its native 3072). Off by default so providers - /// that don't accept the field — Voyage (uses `output_dimension`), Cohere, - /// LocalAI/Ollama — keep working unchanged. Set via - /// [`Self::with_send_dimensions`] for the OpenAI / custom-OpenAI paths. - send_dimensions: bool, - /// When true, this provider points at a hosted cloud endpoint that always - /// requires a bearer token (genuine OpenAI `api.openai.com`, Voyage), so an - /// empty `api_key` must fail fast instead of POSTing an unauthenticated - /// request. Off by default so the OpenAI-compatible provider keeps serving - /// keyless local/custom endpoints (LocalAI, Ollama-via-OpenAI). Set via - /// [`Self::with_required_api_key`]. See the guard in [`Self::embed`]. - requires_api_key: bool, -} - -/// True when `base_url` is Google Gemini's OpenAI-compatibility host. -/// -/// Gemini exposes an OpenAI-compatible shim at -/// `https://generativelanguage.googleapis.com/v1beta/openai/`. We match on the -/// host (any path / scheme) so it triggers whether the user pasted the -/// `/v1beta/openai` form, a bare host, or a future path variant. -fn is_gemini_base_url(base_url: &str) -> bool { - reqwest::Url::parse(base_url) - .ok() - .and_then(|u| u.host_str().map(str::to_ascii_lowercase)) - .is_some_and(|h| { - h == "generativelanguage.googleapis.com" - || h.ends_with(".generativelanguage.googleapis.com") - }) -} - -/// Normalize a model id for the configured base URL. -/// -/// Gemini's OpenAI-compat shim maps `POST /v1/embeddings` onto its native -/// `BatchEmbedContents` RPC, which requires the model in `models/` form -/// (e.g. `models/text-embedding-004`). A user who points the Custom -/// (OpenAI-compatible) embeddings provider at Gemini's compat base URL and -/// pastes a bare id (`text-embedding-004`) gets a `400 … BatchEmbedContents\ -/// Request.model: unexpected model name format` on every memory re-embed -/// (TAURI-RUST-4SA, 4,494 events / 1 user — non-fatal, so the pipeline retries -/// per document and floods Sentry). Prefix `models/` so the request is -/// well-formed before it leaves the process. Host-scoped (genuine OpenAI / -/// LocalAI / Ollama ids are untouched) and idempotent (an already-prefixed -/// `models/…` or a `tunedModels/…` id is left alone). -fn normalize_model_for_base_url(base_url: &str, model: &str) -> String { - if is_gemini_base_url(base_url) - && !model.is_empty() - && !model.starts_with("models/") - && !model.starts_with("tunedModels/") - { - format!("models/{model}") - } else { - model.to_string() - } -} - -/// True when a `400` body is Gemini's model-id format rejection (the OpenAI-side -/// surface of TAURI-RUST-4SA). Mirrors the wire phrases the -/// `is_embedding_model_rejected` classifier in `core::observability` keys on so -/// the remediation hint and the Sentry demotion never drift. -fn is_gemini_model_format_rejection(text: &str) -> bool { - let lower = text.to_ascii_lowercase(); - lower.contains("unexpected model name format") - || (lower.contains("invalid_argument") && lower.contains("batchembedcontentsrequest.model")) -} - -impl OpenAiEmbedding { - /// Creates a new OpenAI-style provider. - pub fn new(base_url: &str, api_key: &str, model: &str, dims: usize) -> Self { - let base_url = base_url.trim_end_matches('/').to_string(); - // Repair a bare Gemini model id (`text-embedding-004`) into the - // `models/` form Gemini's OpenAI-compat shim requires, so the - // Custom-endpoint path doesn't 400 on every embed (TAURI-RUST-4SA). - let model = normalize_model_for_base_url(&base_url, model); - Self { - base_url, - api_key: api_key.to_string(), - model, - dims, - send_dimensions: false, - requires_api_key: false, - } - } - - /// Opt into sending the OpenAI `dimensions` request parameter so a - /// reducible model (`text-embedding-3-large` / `-3-small`) returns exactly - /// `dims` floats instead of its native size. Only call this for genuine - /// OpenAI / OpenAI-compatible endpoints that implement the parameter — - /// see [`Self::send_dimensions`]. Returns `self` for builder chaining. - pub fn with_send_dimensions(mut self, send: bool) -> Self { - self.send_dimensions = send; - self - } - - /// Mark this provider as a keyed cloud endpoint that must have an API key. - /// When set, [`Self::embed`] fails fast (before any HTTP round-trip) if the - /// resolved `api_key` is empty, instead of silently omitting the - /// `Authorization` header. Use for genuine OpenAI (`api.openai.com`) and - /// Voyage; leave off for keyless local/custom OpenAI-compatible endpoints. - /// Returns `self` for builder chaining. - pub fn with_required_api_key(mut self, required: bool) -> Self { - self.requires_api_key = required; - self - } - - /// Returns the configured base URL. - pub fn base_url(&self) -> &str { - &self.base_url - } - - /// Returns the configured model name. - pub fn model(&self) -> &str { - &self.model - } - - /// Internal helper to build an HTTP client with proxy support. - fn http_client(&self) -> reqwest::Client { - crate::openhuman::config::build_runtime_proxy_client("memory.embeddings") - } - - /// Checks if the base URL includes a specific path (e.g., /api/v1). - fn has_explicit_api_path(&self) -> bool { - let Ok(url) = reqwest::Url::parse(&self.base_url) else { - return false; - }; - - let path = url.path().trim_end_matches('/'); - !path.is_empty() && path != "/" - } - - /// Checks if the URL already ends with /embeddings. - fn has_embeddings_endpoint(&self) -> bool { - let Ok(url) = reqwest::Url::parse(&self.base_url) else { - return false; - }; - - url.path().trim_end_matches('/').ends_with("/embeddings") - } - - /// Constructs the final URL for the embeddings endpoint. - pub fn embeddings_url(&self) -> String { - if self.has_embeddings_endpoint() { - return self.base_url.clone(); - } - - if self.has_explicit_api_path() { - format!("{}/embeddings", self.base_url) - } else { - format!("{}/v1/embeddings", self.base_url) - } - } -} - -#[async_trait] -impl EmbeddingProvider for OpenAiEmbedding { - fn name(&self) -> &str { - "openai" - } - - fn model_id(&self) -> &str { - &self.model - } - - fn dimensions(&self) -> usize { - self.dims - } - - /// Sends a POST request to the embedding API. - /// - /// On 429 (Too Many Requests) or 503 (Service Unavailable) the call is - /// retried up to `MAX_429_RETRIES` times with exponential backoff. When - /// the server supplies a `Retry-After` header its value (delta-seconds) is - /// preferred over the computed backoff. After all retries are exhausted the - /// canonical error message is returned so the `TransientUpstreamHttp` - /// classifier in `core::observability` demotes it to a warning breadcrumb - /// instead of a Sentry error event. - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - if texts.is_empty() { - return Ok(Vec::new()); - } - - // Pre-flight: empty / whitespace-only entries are guaranteed 400s from - // the upstream (OpenAI: `"input must be a non-empty string"`; OpenHuman - // cloud backend: `"input must be a non-empty string or array of - // non-empty strings"`). Bailing here keeps the round-trip and quota - // out of the picture and — crucially — bypasses the `report_error_or_ - // expected` Sentry route below, so a caller passing an empty summary - // stops manifesting as a server fault (#13021). - if let Some(idx) = texts.iter().position(|t| t.trim().is_empty()) { - tracing::warn!( - target: "openai::embed", - "[openai] refusing embed: input[{idx}] is empty/whitespace \ - (count={}, model={}). Caller must filter empty strings.", - texts.len(), - self.model, - ); - anyhow::bail!( - "openai embed: refusing empty/whitespace input at index {idx} of {} (model={})", - texts.len(), - self.model, - ); - } - - // Fast-fail when this is a keyed cloud provider (OpenAI / Voyage) but the - // resolved key is empty. The key collapses to "" when the stored BYO - // credential can't be read — the OS-keychain consent is `none`/declined - // (`cached_consent=none`) or the cred fails to decrypt — because - // `resolve_api_key` swallows every such failure into "". Without this - // guard the request goes out with NO `Authorization` header at all and - // OpenAI 401s "You didn't provide an API key" on every embed; the memory - // pipeline re-embeds per document and floods Sentry (TAURI-RUST-4TZ: - // 3.9k events). Bailing here skips the wasted request, and the "API key - // not set" wording is demoted by the `ApiKeyMissing` classifier in - // `core::observability` to a single low-cardinality breadcrumb. The - // remediation surfaces the keychain-consent / re-enter-key path so the - // stored key can actually be read. Scoped via `requires_api_key`: the - // OpenAI-compatible provider legitimately supports keyless local/custom - // endpoints (LocalAI, Ollama-via-OpenAI), which keep omitting the header - // rather than bailing — mirroring the Cohere guard (TAURI-RUST-52S). - if self.requires_api_key && self.api_key.trim().is_empty() { - let message = format!( - "Embedding API key not set (model={}) — re-enter your key or grant \ - keychain access in Settings → Memory", - self.model, - ); - crate::core::observability::report_error_or_expected( - message.as_str(), - "embeddings", - "openai_embed", - &[("model", self.model.as_str()), ("failure", "missing_key")], - ); - anyhow::bail!(message); - } - - let url = self.embeddings_url(); - - tracing::debug!( - target: "openai::embed", - "[openai] embed: model={}, count={}, url={}", - self.model, texts.len(), url - ); - - let mut body = serde_json::json!({ - "model": self.model, - "input": texts, - }); - // Request a specific output size on OpenAI 3-* models (Matryoshka) so - // the vector matches `dims` (e.g. 3-large → 1024 for the memory tree's - // fixed EMBEDDING_DIM). Gated by `send_dimensions` because Voyage / - // Cohere / LocalAI don't accept this exact field. - if self.send_dimensions && self.dims > 0 { - body["dimensions"] = serde_json::json!(self.dims); - } - - // Retry loop: handles 429 Too Many Requests and 503 Service Unavailable - // with Retry-After–aware exponential backoff. - for attempt in 0..=MAX_429_RETRIES { - let mut req = self - .http_client() - .post(&url) - .header("Content-Type", "application/json") - .json(&body); - - // Only set Authorization header when an API key is configured. - if !self.api_key.is_empty() { - req = req.header("Authorization", format!("Bearer {}", self.api_key)); - } - - // Proactively gate every outbound attempt (initial + retries) against - // the per-endpoint rate budget so cloud backends (OpenHuman/Voyage, - // OpenAI, custom remote endpoints) stay under their account quota - // instead of tripping 429s. The chokepoint must sit inside the loop: - // a single pre-loop acquire would let retried 429/503 attempts bypass - // token consumption and let concurrent callers blow past the cap, - // ironically triggering more 429s. Token consumption tracks the number - // of HTTP attempts (1 + retries actually executed). Loopback endpoints - // are exempt (see `rate_limit`). - super::rate_limit::acquire_embedding_slot(&self.base_url).await; - - let resp = req.send().await?; - - let status = resp.status(); - - // Retry on 429 and 503 — both can carry a Retry-After header. - let is_retryable = status.as_u16() == 429 || status.as_u16() == 503; - - if is_retryable && attempt < MAX_429_RETRIES { - // Read Retry-After before consuming the body. - let retry_after_val = resp - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_owned()); - - let body_text = resp.text().await.unwrap_or_default(); - tracing::debug!( - target: "openai::embed", - "[embeddings] openai {} body on retry: {body_text}", - status.as_u16() - ); - - let delay_ms = backoff_ms_for_attempt(attempt, retry_after_val.as_deref()); - - tracing::debug!( - target: "openai::embed", - "[embeddings] openai {}, retrying in {}ms (attempt {}/{})", - status.as_u16(), delay_ms, attempt + 1, MAX_429_RETRIES - ); - - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - continue; - } - - if !status.is_success() { - let status_str = status.as_u16().to_string(); - let text = resp.text().await.unwrap_or_default(); - tracing::debug!( - target: "openai::embed", - "[openai] embed error: status={status}, body={text}" - ); - let mut message = format!("Embedding API error ({status}): {text}"); - // A 404/405 means the base URL responded but exposes no - // embeddings route — the user pointed the Custom - // (OpenAI-compatible) provider at a chat-only endpoint (e.g. - // DeepSeek). Append an actionable remediation while PRESERVING - // the `Embedding API error (404…)` prefix that - // `observability::is_embedding_endpoint_absent` keys on, so the - // event is still demoted from Sentry. Host-agnostic text (no - // URL/credential echo). TAURI-RUST-5JR. - if matches!(status.as_u16(), 404 | 405) { - message.push_str( - " — this endpoint has no embeddings API; pick an \ - embeddings-capable provider in Settings → Memory", - ); - } - // A 400 with Gemini's `… BatchEmbedContentsRequest.model: - // unexpected model name format` / `INVALID_ARGUMENT` body means - // the user pointed the Custom provider at Gemini's OpenAI-compat - // URL with a bare model id. The constructor now normalizes the id - // to `models/`, so this only fires for already-stored bad - // state or other compat hosts; append an actionable hint while - // PRESERVING the `(400` + body so the - // `observability::is_embedding_model_rejected` classifier still - // demotes the per-embed flood. TAURI-RUST-4SA. - else if status.as_u16() == 400 && is_gemini_model_format_rejection(&text) { - message.push_str( - " — Gemini needs the embeddings model id in `models/` \ - form (e.g. `models/text-embedding-004`); fix it in \ - Settings → Memory", - ); - } - // A 400 "… does not exist" / "does not support embeddings" body - // means the endpoint IS an embeddings API but the configured - // model id is not an embeddings model — the user pasted a chat - // model (e.g. an OpenRouter `…:free` id) into the embeddings - // model field. Same demotion contract as above. TAURI-RUST-9SK. - else if status.as_u16() == 400 - && (text.contains("does not exist") - || text.contains("does not support embeddings")) - { - message.push_str( - " — this model isn't an embeddings model; pick an \ - embeddings-capable model in Settings → Memory", - ); - } - // Use `report_error_or_expected` so transient upstream HTTP - // failures (e.g. 429 Too Many Requests after retry cap) log a - // warning breadcrumb instead of firing a Sentry error event. - crate::core::observability::report_error_or_expected( - message.as_str(), - "embeddings", - "openai_embed", - &[ - ("model", self.model.as_str()), - ("status", status_str.as_str()), - ("failure", "non_2xx"), - ], - ); - anyhow::bail!(message); - } - - let json: serde_json::Value = resp.json().await?; - let data = json - .get("data") - .and_then(|d| d.as_array()) - .ok_or_else(|| anyhow::anyhow!("Invalid embedding response: missing 'data'"))?; - - // Validate that the response count matches the input count. - if data.len() != texts.len() { - anyhow::bail!( - "openai embed count mismatch: sent {} texts, got {} items in 'data'", - texts.len(), - data.len() - ); - } - - let mut embeddings = Vec::with_capacity(data.len()); - for (i, item) in data.iter().enumerate() { - let embedding = item - .get("embedding") - .and_then(|e| e.as_array()) - .ok_or_else(|| { - anyhow::anyhow!("Invalid embedding item at index {i}: missing 'embedding'") - })?; - - let mut vec = Vec::with_capacity(embedding.len()); - for (j, v) in embedding.iter().enumerate() { - #[allow(clippy::cast_possible_truncation)] - let f = v.as_f64().ok_or_else(|| { - anyhow::anyhow!("non-numeric value at data[{i}].embedding[{j}]: {v}") - })? as f32; - vec.push(f); - } - - // Validate dimensions. - if self.dims > 0 && vec.len() != self.dims { - anyhow::bail!( - "openai embed dimension mismatch at index {i}: expected {}, got {}", - self.dims, - vec.len() - ); - } - - embeddings.push(vec); - } - - tracing::debug!( - target: "openai::embed", - "[openai] embed success: model={}, count={}, dims={}", - self.model, embeddings.len(), - embeddings.first().map(|v| v.len()).unwrap_or(0) - ); - - return Ok(embeddings); - } - - // The loop always exits via `return Ok(...)`, `bail!(...)`, or - // `continue`; this point is structurally unreachable. On the final - // attempt (`attempt == MAX_429_RETRIES`) the retryable guard is false - // and execution falls into the non-2xx branch above, which bails with - // the body-bearing format "Embedding API error (429 ...): " — - // that format preserves the "(429 " substring required by the - // TransientUpstreamHttp classifier in core::observability. - unreachable!("embed retry loop must exit via return or bail") - } -} - -#[cfg(test)] -#[path = "openai_tests.rs"] -mod tests; diff --git a/src/openhuman/embeddings/openai_adapter.rs b/src/openhuman/embeddings/openai_adapter.rs new file mode 100644 index 000000000..003cabe78 --- /dev/null +++ b/src/openhuman/embeddings/openai_adapter.rs @@ -0,0 +1,71 @@ +//! Compatibility wrapper for tinyagents' OpenAI-compatible model. + +use async_trait::async_trait; +use tinyagents::harness::embeddings::{EmbeddingModel, OpenAiEmbeddingModel}; + +use super::EmbeddingProvider; + +pub struct OpenAiEmbedding { + inner: OpenAiEmbeddingModel, +} + +impl OpenAiEmbedding { + pub fn new(base_url: &str, api_key: &str, model: &str, dimensions: usize) -> Self { + Self { + inner: OpenAiEmbeddingModel::new(api_key) + .with_base_url(base_url) + .with_model(model) + .with_dimensions(dimensions) + .with_send_dimensions(false) + .with_required_api_key(false), + } + } + + pub fn with_send_dimensions(mut self, send: bool) -> Self { + self.inner = self.inner.with_send_dimensions(send); + self + } + + pub fn with_required_api_key(mut self, required: bool) -> Self { + self.inner = self.inner.with_required_api_key(required); + self + } + + pub fn base_url(&self) -> &str { + self.inner.base_url() + } + + pub fn model(&self) -> &str { + self.inner.model() + } + + pub fn embeddings_url(&self) -> String { + self.inner.embeddings_url() + } +} + +#[async_trait] +impl EmbeddingProvider for OpenAiEmbedding { + fn name(&self) -> &str { + self.inner.name() + } + fn model_id(&self) -> &str { + self.inner.model_id() + } + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + fn signature(&self) -> String { + self.inner.signature() + } + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let texts = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + self.inner + .embed(&texts) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} diff --git a/src/openhuman/embeddings/openai_tests.rs b/src/openhuman/embeddings/openai_tests.rs deleted file mode 100644 index a1f393e61..000000000 --- a/src/openhuman/embeddings/openai_tests.rs +++ /dev/null @@ -1,918 +0,0 @@ -use super::*; -use axum::{ - extract::Json, - http::{HeaderMap, StatusCode}, - routing::post, - Router, -}; -use std::net::SocketAddr; - -async fn start_mock(app: Router) -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr: SocketAddr = listener.local_addr().unwrap(); - tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - format!("http://127.0.0.1:{}", addr.port()) -} - -// ── Constructor & URL building ────────────────────────── - -#[test] -fn trailing_slash_stripped() { - let p = OpenAiEmbedding::new("https://api.openai.com/", "key", "model", 1536); - assert_eq!(p.base_url, "https://api.openai.com"); -} - -#[test] -fn dimensions_custom() { - let p = OpenAiEmbedding::new("http://localhost", "k", "m", 384); - assert_eq!(p.dimensions(), 384); -} - -#[test] -fn accessors() { - let p = OpenAiEmbedding::new("http://x", "k", "m", 1); - assert_eq!(p.base_url(), "http://x"); - assert_eq!(p.model(), "m"); - assert_eq!(p.name(), "openai"); - assert_eq!(p.model_id(), "m"); - assert_eq!(p.signature(), "provider=openai;model=m;dims=1"); -} - -// ── Gemini model-id normalization (TAURI-RUST-4SA) ────── - -#[test] -fn gemini_base_url_prefixes_bare_model() { - // Gemini's OpenAI-compat shim requires `models/`; a bare id 400s with - // `BatchEmbedContentsRequest.model: unexpected model name format`. - let p = OpenAiEmbedding::new( - "https://generativelanguage.googleapis.com/v1beta/openai", - "key", - "text-embedding-004", - 768, - ); - assert_eq!(p.model(), "models/text-embedding-004"); - assert_eq!(p.model_id(), "models/text-embedding-004"); -} - -#[test] -fn gemini_normalization_is_idempotent() { - // An already-prefixed `models/…` or a `tunedModels/…` id is left untouched. - let p = OpenAiEmbedding::new( - "https://generativelanguage.googleapis.com/v1beta/openai/", - "key", - "models/text-embedding-004", - 768, - ); - assert_eq!(p.model(), "models/text-embedding-004"); - - let tuned = OpenAiEmbedding::new( - "https://generativelanguage.googleapis.com", - "key", - "tunedModels/my-embed", - 768, - ); - assert_eq!(tuned.model(), "tunedModels/my-embed"); -} - -#[test] -fn non_gemini_base_url_leaves_model_unchanged() { - // Genuine OpenAI / LocalAI / Ollama ids must not grow a `models/` prefix. - for base in [ - "https://api.openai.com", - "http://localhost:11434", - "https://api.example.com/v1", - ] { - let p = OpenAiEmbedding::new(base, "key", "text-embedding-004", 768); - assert_eq!( - p.model(), - "text-embedding-004", - "non-Gemini base must not prefix the model: {base}" - ); - } -} - -#[test] -fn url_standard_openai() { - let p = OpenAiEmbedding::new("https://api.openai.com", "key", "model", 1536); - assert_eq!(p.embeddings_url(), "https://api.openai.com/v1/embeddings"); -} - -#[test] -fn url_base_with_v1_no_duplicate() { - let p = OpenAiEmbedding::new("https://api.example.com/v1", "key", "model", 1536); - assert_eq!(p.embeddings_url(), "https://api.example.com/v1/embeddings"); -} - -#[test] -fn url_non_v1_api_path() { - let p = OpenAiEmbedding::new( - "https://api.example.com/api/coding/v3", - "key", - "model", - 1536, - ); - assert_eq!( - p.embeddings_url(), - "https://api.example.com/api/coding/v3/embeddings" - ); -} - -#[test] -fn url_already_ends_with_embeddings() { - let p = OpenAiEmbedding::new( - "https://my-api.example.com/api/v2/embeddings", - "key", - "model", - 1536, - ); - assert_eq!( - p.embeddings_url(), - "https://my-api.example.com/api/v2/embeddings" - ); -} - -#[test] -fn url_already_ends_with_embeddings_trailing_slash() { - let p = OpenAiEmbedding::new( - "https://api.example.com/v1/embeddings/", - "key", - "model", - 1536, - ); - assert_eq!(p.embeddings_url(), "https://api.example.com/v1/embeddings"); -} - -#[test] -fn url_root_only() { - let p = OpenAiEmbedding::new("http://localhost:8080", "k", "m", 1); - assert_eq!(p.embeddings_url(), "http://localhost:8080/v1/embeddings"); -} - -#[test] -fn url_root_with_trailing_slash() { - let p = OpenAiEmbedding::new("http://localhost:8080/", "k", "m", 1); - assert_eq!(p.embeddings_url(), "http://localhost:8080/v1/embeddings"); -} - -#[test] -fn has_explicit_api_path_invalid_url() { - let p = OpenAiEmbedding::new("not-a-url", "k", "m", 1); - assert!(!p.has_explicit_api_path()); -} - -#[test] -fn has_embeddings_endpoint_invalid_url() { - let p = OpenAiEmbedding::new("not-a-url", "k", "m", 1); - assert!(!p.has_embeddings_endpoint()); -} - -// ── embed — empty input ───────────────────────────────── - -#[tokio::test] -async fn empty_input_returns_empty() { - let p = OpenAiEmbedding::new("http://unused", "k", "m", 1); - let result = p.embed(&[]).await.unwrap(); - assert!(result.is_empty()); -} - -// ── empty/whitespace entries — pre-flight reject (#13021) ──────── -// -// `embed(&[""])` and friends used to fall through to the HTTP layer -// and trip a backend 400 ("input must be a non-empty string …"), -// which was then captured as a Sentry server fault even though the -// real defect was a caller passing empty text. The guard bails -// without touching the network — the "http://unused" base URL would -// otherwise refuse to connect. - -#[tokio::test] -async fn embed_refuses_single_empty_string() { - let p = OpenAiEmbedding::new("http://unused", "k", "m", 1); - let err = p.embed(&[""]).await.unwrap_err().to_string(); - assert!( - err.contains("refusing empty/whitespace input at index 0"), - "unexpected error: {err}" - ); -} - -#[tokio::test] -async fn embed_refuses_whitespace_only_string() { - let p = OpenAiEmbedding::new("http://unused", "k", "m", 1); - let err = p.embed(&[" \n\t"]).await.unwrap_err().to_string(); - assert!(err.contains("refusing empty/whitespace input at index 0")); -} - -#[tokio::test] -async fn embed_refuses_mixed_batch_with_empty() { - let p = OpenAiEmbedding::new("http://unused", "k", "m", 1); - let err = p.embed(&["ok", "", "fine"]).await.unwrap_err().to_string(); - assert!(err.contains("refusing empty/whitespace input at index 1")); -} - -#[tokio::test] -async fn embed_refuses_does_not_use_embedding_api_error_prefix() { - // The classifier in `core::observability` treats `"Embedding API error"` - // / `"("` shapes as upstream HTTP failures. The client-side - // pre-flight refusal MUST NOT collide with that shape, otherwise this - // very fix would re-enter the same Sentry-as-server-fault path that - // #13021 was about. Lock the bail wording so a future rename can't - // silently reintroduce the regression. - let p = OpenAiEmbedding::new("http://unused", "k", "m", 1); - let err = p.embed(&[""]).await.unwrap_err().to_string(); - assert!( - !err.contains("Embedding API error"), - "bail wording must not collide with TransientUpstreamHttp classifier: {err}" - ); -} - -// ── embed — success ───────────────────────────────────── - -#[tokio::test] -async fn embed_success_single() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [0.1, 0.2, 0.3] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "test-key", "test-model", 3); - - let result = p.embed(&["hello"]).await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0], vec![0.1_f32, 0.2, 0.3]); -} - -#[tokio::test] -async fn embed_success_batch() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [ - { "embedding": [1.0, 2.0] }, - { "embedding": [3.0, 4.0] } - ] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 2); - - let result = p.embed(&["a", "b"]).await.unwrap(); - assert_eq!(result.len(), 2); - assert_eq!(result[1], vec![3.0_f32, 4.0]); -} - -#[tokio::test] -async fn embed_sends_auth_header() { - let app = Router::new().route( - "/v1/embeddings", - post( - |headers: HeaderMap, Json(body): Json| async move { - let auth = headers.get("Authorization").unwrap().to_str().unwrap(); - assert_eq!(auth, "Bearer my-secret-key"); - assert_eq!(body["model"], "text-embedding-3-small"); - Json(serde_json::json!({ - "data": [{ "embedding": [1.0] }] - })) - }, - ), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "my-secret-key", "text-embedding-3-small", 1); - - p.embed(&["test"]).await.unwrap(); -} - -// #002: the OpenAI `dimensions` request param. Off by default (so Voyage / -// Cohere / Ollama, which don't accept this exact field, keep working); on when -// the OpenAI / custom factory branch opts in via `with_send_dimensions(true)`. - -#[tokio::test] -async fn embed_sends_dimensions_when_opted_in() { - let app = Router::new().route( - "/v1/embeddings", - post(|Json(body): Json| async move { - assert_eq!( - body["dimensions"], 1024, - "dimensions must be sent so 3-large returns 1024, not its native 3072" - ); - Json(serde_json::json!({ "data": [{ "embedding": vec![0.0_f32; 1024] }] })) - }), - ); - let url = start_mock(app).await; - let p = - OpenAiEmbedding::new(&url, "k", "text-embedding-3-large", 1024).with_send_dimensions(true); - p.embed(&["test"]).await.unwrap(); -} - -#[tokio::test] -async fn embed_omits_dimensions_by_default() { - let app = Router::new().route( - "/v1/embeddings", - post(|Json(body): Json| async move { - assert!( - body.get("dimensions").is_none(), - "dimensions must NOT be sent by default (Voyage/Cohere/Ollama reject it)" - ); - Json(serde_json::json!({ "data": [{ "embedding": [1.0] }] })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); // no with_send_dimensions - p.embed(&["test"]).await.unwrap(); -} - -#[tokio::test] -async fn embed_skips_auth_header_when_key_empty() { - let app = Router::new().route( - "/v1/embeddings", - post(|headers: HeaderMap| async move { - // No Authorization header should be present. - assert!( - headers.get("Authorization").is_none(), - "should not send auth header when key is empty" - ); - Json(serde_json::json!({ - "data": [{ "embedding": [1.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "", "m", 1); - - p.embed(&["test"]).await.unwrap(); -} - -/// A keyed cloud provider (`with_required_api_key(true)` — genuine OpenAI / -/// Voyage) with an empty key must bail BEFORE any HTTP request rather than -/// POSTing with no `Authorization` header and 401-ing on every embed -/// (TAURI-RUST-4TZ). The base URL points at an address nothing is listening on, -/// so asserting the error is the key-guard message — not a connection error — -/// proves no request was attempted. The "API key not set" wording is what the -/// `ApiKeyMissing` classifier keys on to demote the flood out of Sentry. -#[tokio::test] -async fn embed_required_key_empty_bails_without_request() { - for key in ["", " "] { - let p = OpenAiEmbedding::new("http://127.0.0.1:1", key, "text-embedding-3-small", 1) - .with_required_api_key(true); - let err = p.embed(&["hello"]).await.unwrap_err().to_string(); - assert!( - err.contains("API key not set"), - "expected key-guard message for key {key:?}, got: {err}" - ); - } -} - -/// The keyless local/custom path is unaffected: without -/// `with_required_api_key`, an empty key still omits the header and sends the -/// request (LocalAI / Ollama-via-OpenAI legitimately need no bearer). -#[tokio::test] -async fn embed_empty_key_without_requirement_still_sends() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { Json(serde_json::json!({ "data": [{ "embedding": [1.0] }] })) }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "", "m", 1); // no with_required_api_key - let result = p.embed(&["test"]).await.unwrap(); - assert_eq!(result.len(), 1); -} - -// ── embed — error paths ───────────────────────────────── - -#[tokio::test] -async fn embed_server_error() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { (StatusCode::INTERNAL_SERVER_ERROR, "rate limited") }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("500"), "status: {msg}"); - assert!(msg.contains("rate limited"), "body: {msg}"); -} - -/// A 404 means the configured base URL has no embeddings route (the user -/// pointed the Custom provider at a chat-only endpoint, e.g. DeepSeek — -/// TAURI-RUST-5JR). The message must (a) carry an actionable remediation, and -/// (b) PRESERVE the `Embedding API error (404…)` prefix the -/// `observability::is_embedding_endpoint_absent` classifier keys on, so the -/// flood is demoted from Sentry rather than firing on every re-embed. -#[tokio::test] -async fn embed_404_endpoint_absent_is_actionable_and_classifier_stable() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { (StatusCode::NOT_FOUND, "Not Found") }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - // Classifier contract: prefix preserved. - assert!( - msg.to_ascii_lowercase() - .contains("embedding api error (404"), - "must preserve the (404 classifier prefix: {msg}" - ); - // Actionable remediation appended. - assert!( - msg.contains("no embeddings API") && msg.contains("Settings → Memory"), - "must carry actionable remediation: {msg}" - ); -} - -/// 429 rate-limit responses must format their message in the canonical -/// `"... API error (): "` shape so the shared -/// `is_transient_upstream_http_message` classifier in `core::observability` -/// demotes them to a warning breadcrumb instead of a Sentry error event. -#[tokio::test] -async fn embed_429_uses_canonical_transient_format() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - ( - StatusCode::TOO_MANY_REQUESTS, - r#"{"error":{"message":"Rate limit exceeded.","type":"rate_limit_error"}}"#, - ) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("(429 Too Many Requests)"), - "expected canonical transient HTTP shape, got: {msg}" - ); - // Pin the shape to the exact substring `is_transient_upstream_http_message` - // matches on (`"api error ( "`). The broader - // `is_transient_message_failure` classifier below also passes for the *old* - // `"Embedding API error 429 …"` format, so without this assertion a future - // refactor could silently revert the format and the test would still go - // green. - assert!( - msg.to_ascii_lowercase().contains("api error (429 "), - "message must match is_transient_upstream_http_message classifier arm: {msg}" - ); - assert!( - crate::core::observability::is_transient_message_failure(&msg), - "message should classify as transient: {msg}" - ); -} - -#[tokio::test] -async fn embed_budget_exhausted_400_still_errors() { - // OPENHUMAN-TAURI-JM: the backend returns HTTP 400 with a budget-exhausted - // body when the user is out of credits. The provider must still surface - // an `Err` to the caller (so the calling pipeline can short-circuit), but - // the diagnostic emit site must route through `report_error_or_expected` - // so the message is classified as `BudgetExhausted` and demoted rather - // than spawning a Sentry error event for every embed call. - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - ( - StatusCode::BAD_REQUEST, - r#"{"success":false,"error":"Budget exceeded — add credits to continue"}"#, - ) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("400"), "status: {msg}"); - assert!(msg.contains("Budget exceeded"), "body: {msg}"); -} - -#[tokio::test] -async fn embed_missing_data_field() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { Json(serde_json::json!({ "result": "ok" })) }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.to_string().contains("missing 'data'")); -} - -#[tokio::test] -async fn embed_missing_embedding_field_in_item() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "index": 0 }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.to_string().contains("missing 'embedding'")); -} - -#[tokio::test] -async fn embed_non_numeric_value_errors() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [1.0, "not_a_number", 3.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 3); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("non-numeric"), "msg: {msg}"); -} - -#[tokio::test] -async fn embed_count_mismatch() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [1.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["a", "b"]).await.unwrap_err(); - assert!(err.to_string().contains("count mismatch")); -} - -#[tokio::test] -async fn embed_dimension_mismatch() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [1.0, 2.0, 3.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 2); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.to_string().contains("dimension mismatch")); -} - -/// Issue #4056: a `dims == 0` provider is the dimension-agnostic verification -/// probe — it must NOT enforce any length, so an endpoint returning its own -/// native size passes instead of being rejected. This is what lets a Custom -/// endpoint verify when the user's guessed `dimensions` differs from the -/// model's native output; the caller then adopts the returned length. -#[tokio::test] -async fn embed_dims_zero_skips_dimension_guard() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [1.0, 2.0, 3.0, 4.0, 5.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 0); - - let result = p.embed(&["hi"]).await.unwrap(); - assert_eq!(result[0].len(), 5, "dims=0 must accept the native length"); -} - -#[tokio::test] -async fn embed_malformed_json() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { (StatusCode::OK, "not json") }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.is::()); -} - -#[tokio::test] -async fn embed_connection_refused() { - let p = OpenAiEmbedding::new("http://127.0.0.1:1", "k", "m", 1); - let err = p.embed(&["hi"]).await.unwrap_err(); - assert!(err.is::()); -} - -#[test] -fn openai_embedding_api_error_stays_unexpected() { - let msg = "Embedding API error 401 Unauthorized: invalid_api_key"; - assert_eq!( - crate::core::observability::expected_error_kind(msg), - None, - "OpenAI API key errors should continue to reach Sentry" - ); -} - -// ── embed_one (trait default) ─────────────────────────── - -#[tokio::test] -async fn embed_one_success() { - let app = Router::new().route( - "/v1/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [9.0, 8.0, 7.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 3); - - let vec = p.embed_one("test").await.unwrap(); - assert_eq!(vec, vec![9.0_f32, 8.0, 7.0]); -} - -// ── URL building — custom endpoint ────────────────────── - -#[tokio::test] -async fn embed_with_explicit_api_path() { - let app = Router::new().route( - "/custom/api/embeddings", - post(|| async { - Json(serde_json::json!({ - "data": [{ "embedding": [1.0] }] - })) - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&format!("{url}/custom/api"), "k", "m", 1); - - let result = p.embed(&["test"]).await.unwrap(); - assert_eq!(result.len(), 1); -} - -// ── 429 backoff / Retry-After tests ────────────────────── - -/// Mock returns 429 twice (with Retry-After: 0) then 200 — verify embed -/// succeeds and that exactly 3 requests were made (initial + 2 retries). -#[tokio::test] -async fn embed_429_retries_with_retry_after_then_succeeds() { - use std::sync::{Arc, Mutex}; - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v1/embeddings", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - if *n <= 2 { - // Return 429 with Retry-After: 0 (zero delay) for fast tests. - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .header("Retry-After", "0") - .body(axum::body::Body::from( - r#"{"error":{"message":"rate limited"}}"#, - )) - .unwrap() - } else { - axum::response::Response::builder() - .status(StatusCode::OK) - .header("Content-Type", "application/json") - .body(axum::body::Body::from( - r#"{"data":[{"embedding":[1.0,2.0]}]}"#, - )) - .unwrap() - } - } - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 2); - - let result = p.embed(&["hello"]).await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0], vec![1.0_f32, 2.0]); - assert_eq!( - *counter.lock().unwrap(), - 3, - "expected 3 total requests (1 initial + 2 retries)" - ); -} - -/// Mock returns 429 indefinitely — verify bail with canonical message and -/// that exactly MAX_429_RETRIES + 1 requests were made. -#[tokio::test] -async fn embed_429_indefinite_bails_after_retry_cap() { - use crate::openhuman::embeddings::retry_after::MAX_429_RETRIES; - use std::sync::{Arc, Mutex}; - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v1/embeddings", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .header("Retry-After", "0") - .body(axum::body::Body::from( - r#"{"error":{"message":"always rate limited"}}"#, - )) - .unwrap() - } - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let err = p.embed(&["hi"]).await.unwrap_err(); - let msg = err.to_string(); - - // The error message must contain "429" so the classifier still suppresses it. - assert!( - msg.contains("429"), - "bail message must contain 429 for Sentry classifier: {msg}" - ); - // Must match is_transient_upstream_http_message via "(429 " substring. - assert!( - crate::core::observability::is_transient_message_failure(&msg), - "bail message must classify as transient: {msg}" - ); - - let requests = *counter.lock().unwrap(); - assert_eq!( - requests, - MAX_429_RETRIES + 1, - "expected exactly MAX_429_RETRIES+1={} requests, got {}", - MAX_429_RETRIES + 1, - requests - ); -} - -/// Mock returns 429 without Retry-After header — verify the no-header code -/// path is taken, that a retry is attempted, and that the request succeeds. -/// Uses `Retry-After: 0` on the *second* attempt to confirm header parsing -/// is independent from the no-header first attempt. -#[tokio::test] -async fn embed_429_without_retry_after_uses_exponential_backoff() { - use crate::openhuman::embeddings::retry_after::{backoff_ms_for_attempt, BASE_BACKOFF_MS}; - use std::sync::{Arc, Mutex}; - - // Confirm the helper returns the exponential base when header is absent. - assert_eq!( - backoff_ms_for_attempt(0, None), - BASE_BACKOFF_MS, - "attempt 0 without header should use BASE_BACKOFF_MS" - ); - assert_eq!( - backoff_ms_for_attempt(1, None), - BASE_BACKOFF_MS * 2, - "attempt 1 without header should double" - ); - - // Confirm header path: Retry-After: 0 overrides the exponential base. - assert_eq!( - backoff_ms_for_attempt(0, Some("0")), - 0, - "Retry-After: 0 should yield zero-ms delay" - ); - - // End-to-end: first request returns 429 (no Retry-After), second succeeds. - // Use Retry-After: 0 on the 429 to avoid real-wall-clock delay in CI. - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v1/embeddings", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - if *n == 1 { - // Mock still sets Retry-After: 0 here to keep this test fast; - // the no-header exponential branch is exercised end-to-end by - // `embed_429_no_retry_after_header_falls_back_to_exponential` - // below and by unit tests in `retry_after.rs`. - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .header("Retry-After", "0") - .body(axum::body::Body::from( - r#"{"error":{"message":"rate limited"}}"#, - )) - .unwrap() - } else { - axum::response::Response::builder() - .status(StatusCode::OK) - .header("Content-Type", "application/json") - .body(axum::body::Body::from(r#"{"data":[{"embedding":[9.9]}]}"#)) - .unwrap() - } - } - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let result = p.embed(&["hi"]).await; - assert!(result.is_ok(), "should succeed after retry: {:?}", result); - assert_eq!( - *counter.lock().unwrap(), - 2, - "expected 2 total requests (1 initial + 1 retry)" - ); -} - -/// End-to-end coverage for the no-`Retry-After` exponential-backoff branch. -/// -/// The previous "no Retry-After" test still set `Retry-After: 0` on the mock, -/// so the embed loop never actually called `backoff_ms_for_attempt(.., None)` — -/// the exponential fallback was only exercised by unit tests in `retry_after.rs`. -/// This test omits the header entirely so the loop must use the `BASE_BACKOFF_MS` -/// path. We tolerate the ~1 s wait (one retry @ `BASE_BACKOFF_MS = 1000 ms`) to -/// keep the assertion meaningful: the real backoff schedule is what we care about. -#[tokio::test] -async fn embed_429_no_retry_after_header_falls_back_to_exponential() { - use crate::openhuman::embeddings::retry_after::BASE_BACKOFF_MS; - use std::sync::{Arc, Mutex}; - use tokio::time::Instant; - - let counter = Arc::new(Mutex::new(0u32)); - let counter_clone = counter.clone(); - - let app = Router::new().route( - "/v1/embeddings", - post(move || { - let counter = counter_clone.clone(); - async move { - let mut n = counter.lock().unwrap(); - *n += 1; - if *n == 1 { - // Crucial: omit the Retry-After header entirely so the embed - // loop hits `backoff_ms_for_attempt(0, None)` = - // `BASE_BACKOFF_MS` and sleeps. If a future refactor breaks - // the fallback (e.g. defaulting to 0 ms) the elapsed-time - // assertion below will catch it. - axum::response::Response::builder() - .status(StatusCode::TOO_MANY_REQUESTS) - .body(axum::body::Body::from( - r#"{"error":{"message":"rate limited"}}"#, - )) - .unwrap() - } else { - axum::response::Response::builder() - .status(StatusCode::OK) - .header("Content-Type", "application/json") - .body(axum::body::Body::from(r#"{"data":[{"embedding":[4.2]}]}"#)) - .unwrap() - } - } - }), - ); - let url = start_mock(app).await; - let p = OpenAiEmbedding::new(&url, "k", "m", 1); - - let start = Instant::now(); - let result = p.embed(&["hi"]).await; - let elapsed = start.elapsed(); - - assert!(result.is_ok(), "should succeed after retry: {:?}", result); - assert_eq!( - *counter.lock().unwrap(), - 2, - "expected 2 total requests (1 initial + 1 retry)" - ); - // The fallback must actually wait the exponential base — within a 250 ms - // jitter window for slow CI runners. Without this we couldn't tell whether - // the no-header branch was taken or silently short-circuited. - let min_wait = std::time::Duration::from_millis(BASE_BACKOFF_MS.saturating_sub(250)); - assert!( - elapsed >= min_wait, - "expected elapsed >= ~{}ms (BASE_BACKOFF_MS minus jitter), got {:?}", - min_wait.as_millis(), - elapsed - ); -} diff --git a/src/openhuman/embeddings/provider_trait.rs b/src/openhuman/embeddings/provider_trait.rs index 262e1045f..b7a2e616b 100644 --- a/src/openhuman/embeddings/provider_trait.rs +++ b/src/openhuman/embeddings/provider_trait.rs @@ -1,6 +1,7 @@ //! Interface for embedding providers that convert text into numerical vectors. use async_trait::async_trait; +use tinyagents::harness::embeddings::EmbeddingModel; /// Formats the canonical embedding-space signature string. /// @@ -46,3 +47,50 @@ pub trait EmbeddingProvider: Send + Sync { .ok_or_else(|| anyhow::anyhow!("Empty embedding result")) } } + +/// Compatibility adapter from the canonical tinyagents embedding model. +pub struct TinyAgentsEmbeddingProvider { + model: Box, +} + +impl TinyAgentsEmbeddingProvider { + pub fn new(model: impl EmbeddingModel + 'static) -> Self { + Self { + model: Box::new(model), + } + } + + pub fn boxed(model: impl EmbeddingModel + 'static) -> Box { + Box::new(Self::new(model)) + } +} + +#[async_trait] +impl EmbeddingProvider for TinyAgentsEmbeddingProvider { + fn name(&self) -> &str { + self.model.name() + } + + fn model_id(&self) -> &str { + self.model.model_id() + } + + fn dimensions(&self) -> usize { + self.model.dimensions() + } + + fn signature(&self) -> String { + self.model.signature() + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let owned = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + self.model + .embed(&owned) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} diff --git a/src/openhuman/embeddings/rate_limit.rs b/src/openhuman/embeddings/rate_limit.rs deleted file mode 100644 index c609829aa..000000000 --- a/src/openhuman/embeddings/rate_limit.rs +++ /dev/null @@ -1,387 +0,0 @@ -//! Client-side request-rate limiting for cloud embedding backends. -//! -//! Cloud embedding backends (the OpenHuman backend / Voyage, OpenAI, and any -//! OpenAI-compatible `custom:` endpoint) cap requests at a fixed rate per -//! account — ~60/min by default. Every [`super::openai::OpenAiEmbedding::embed`] -//! call is exactly one HTTP POST, and memory-tree ingest fans out one call per -//! chunk across several job workers, so without throttling we routinely trip -//! the backend limiter and absorb `429`s (which `openai.rs` currently -//! downgrades to a warning breadcrumb). This module spends the budget -//! *proactively* so requests stay under the quota instead of reacting to 429s. -//! -//! ## Why a process-global registry keyed by endpoint -//! -//! The quota is account-wide, not per-instance — and `OpenAiEmbedding` -//! instances are ephemeral (the cloud provider builds a fresh one on every -//! `embed` call, and the embedder is constructed from several independent -//! sites: the memory factory, `default_embedding_provider`, and the -//! memory-tree score path). A per-instance limiter would therefore reset -//! constantly and enforce nothing. Instead the buckets live in a -//! process-global registry keyed by the resolved base URL, so all ephemeral -//! instances pointing at the same backend share one budget while distinct -//! backends get independent budgets. -//! -//! Loopback endpoints (`localhost` / `127.0.0.1` / `::1`) are exempt: a local -//! LocalAI- or Ollama-compatible `custom:` server is not the remote quota this -//! guards, and capping it to 60/min would needlessly throttle local -//! throughput. - -use std::collections::HashMap; -use std::net::IpAddr; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::{Arc, Mutex, OnceLock, PoisonError}; -use std::time::Duration; - -use tokio::time::Instant; - -/// Default outbound embedding request budget, in requests per minute. -/// -/// Cloud embedding backends cap requests at ~60/min per account. Used when the -/// operator hasn't overridden `memory.embedding_rate_limit_per_min`. -/// -/// Keep in sync with `default_embedding_rate_limit_per_min` in -/// `config::schema::storage_memory`. -pub const DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN: u32 = 60; - -/// Process-global configured budget (requests/min). `0` disables throttling. -static CONFIGURED_LIMIT: AtomicU32 = AtomicU32::new(DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN); - -/// Process-global registry of per-endpoint token buckets, keyed by base URL. -static BUCKETS: OnceLock>>> = OnceLock::new(); - -/// Override the process-global embedding request budget. `0` disables -/// throttling entirely. -/// -/// Wired from config load (`config::schema::load::apply_env_overrides`) so the -/// live budget tracks `memory.embedding_rate_limit_per_min`. When the rate -/// changes, existing buckets are dropped so the new rate takes effect on the -/// next request — mirroring how `proxy::set_runtime_proxy_config` clears its -/// client cache on reconfigure. -pub fn set_embedding_rate_limit(per_minute: u32) { - let prev = CONFIGURED_LIMIT.swap(per_minute, Ordering::Relaxed); - // Only drop buckets when the rate actually changes. Clearing on every call - // (e.g. repeated config reloads with an unchanged value) would keep handing - // out a fresh burst token and erode the hard-cap pacing guarantee. - if prev != per_minute { - if let Some(registry) = BUCKETS.get() { - registry - .lock() - .unwrap_or_else(PoisonError::into_inner) - .clear(); - } - } - tracing::debug!( - target: "embeddings::rate_limit", - "[embeddings] rate limit set to {per_minute}/min ({})", - if per_minute == 0 { "disabled" } else { "enabled" } - ); -} - -/// The current process-global embedding request budget (requests/min). -#[must_use] -pub fn embedding_rate_limit() -> u32 { - CONFIGURED_LIMIT.load(Ordering::Relaxed) -} - -/// Gate one outbound embedding HTTP request for `base_url`. -/// -/// Blocks cooperatively until the per-endpoint token bucket has a slot, then -/// consumes it. No-ops when throttling is disabled (`limit == 0`) or when -/// `base_url` is a loopback host. -pub async fn acquire_embedding_slot(base_url: &str) { - acquire_with_limit(base_url, embedding_rate_limit()).await; -} - -/// Inner of [`acquire_embedding_slot`] with the budget passed explicitly, so -/// tests can exercise the gating logic without mutating the process-global -/// limit (which would race across parallel tests). -async fn acquire_with_limit(base_url: &str, limit: u32) { - if limit == 0 || is_loopback_url(base_url) { - return; - } - bucket_for(base_url, limit).acquire().await; -} - -/// Get-or-create the bucket for `base_url`. The first caller's `per_minute` -/// fixes the rate for that endpoint until [`set_embedding_rate_limit`] clears -/// the registry; in practice the limit is configured once at startup before -/// any embedding request, so the rate is uniform across the process. -fn bucket_for(base_url: &str, per_minute: u32) -> Arc { - let registry = BUCKETS.get_or_init(|| Mutex::new(HashMap::new())); - let mut map = registry.lock().unwrap_or_else(PoisonError::into_inner); - map.entry(base_url.to_string()) - .or_insert_with(|| { - tracing::debug!( - target: "embeddings::rate_limit", - "[embeddings] new rate limiter endpoint={base_url} limit={per_minute}/min" - ); - Arc::new(TokenBucket::per_minute(per_minute)) - }) - .clone() -} - -/// True when `base_url`'s host is loopback (`localhost`, `127.0.0.0/8`, `::1`). -/// Unparseable URLs are treated as non-loopback so a malformed remote endpoint -/// still gets throttled rather than silently bypassing the limiter. -fn is_loopback_url(base_url: &str) -> bool { - let Ok(url) = reqwest::Url::parse(base_url) else { - return false; - }; - match url.host_str() { - Some(host) => { - host.eq_ignore_ascii_case("localhost") - || host - .trim_start_matches('[') - .trim_end_matches(']') - .parse::() - .map(|ip| ip.is_loopback()) - .unwrap_or(false) - } - None => false, - } -} - -/// Minimum-interval token bucket sized for a **hard** requests-per-minute cap. -/// -/// Capacity is intentionally a single token, not `per_minute`. The backend -/// enforces 60/min as a hard limit, and a token bucket admits up to -/// `capacity + refill × window` requests in any window — so a `per_minute`-sized -/// burst could momentarily reach ~`2 × per_minute` in the first rolling minute -/// and trip the cap. With capacity 1 the bucket paces requests at one per -/// `60 / per_minute` seconds — a steady `per_minute` per minute with no burst. -/// An idle bucket refills that one token, so an occasional lone request (e.g. an -/// interactive retrieval query embed) still goes out immediately; only -/// back-to-back requests are spaced. -/// [`Self::acquire`] consumes one token, sleeping until one accrues if empty. -struct TokenBucket { - state: tokio::sync::Mutex, - capacity: f64, - refill_per_sec: f64, -} - -struct BucketState { - tokens: f64, - last_refill: Instant, -} - -/// Burst allowance, in tokens. One token = no burst beyond a single request, -/// which is what keeps us strictly under a hard per-minute cap. -const BURST_TOKENS: f64 = 1.0; - -impl TokenBucket { - fn per_minute(per_minute: u32) -> Self { - // `max(1)` is purely defensive against divide-by-zero; the `limit == 0` - // path in `acquire_embedding_slot` never constructs a bucket. - let refill_per_sec = f64::from(per_minute.max(1)) / 60.0; - Self { - state: tokio::sync::Mutex::new(BucketState { - tokens: BURST_TOKENS, - last_refill: Instant::now(), - }), - capacity: BURST_TOKENS, - refill_per_sec, - } - } - - async fn acquire(&self) { - loop { - // Compute refill + the wait-until-next-token *while holding the - // lock*, but drop the guard before sleeping so concurrent callers - // aren't blocked on the mutex during the sleep. - let wait = { - let mut state = self.state.lock().await; - let now = Instant::now(); - let elapsed = now.duration_since(state.last_refill).as_secs_f64(); - state.last_refill = now; - match refill_and_take( - &mut state.tokens, - self.capacity, - self.refill_per_sec, - elapsed, - ) { - None => return, - Some(wait) => wait, - } - }; - tracing::debug!( - target: "embeddings::rate_limit", - "[embeddings] throttling embed request: waiting {:.0}ms for a slot", - wait.as_secs_f64() * 1000.0 - ); - tokio::time::sleep(wait).await; - } - } -} - -/// Refill the bucket for `elapsed_secs`, then try to consume one token. -/// Pure in its inputs (no clock) so the rate math is unit-testable directly. -/// Returns `None` when a token was consumed, or `Some(wait)` — the time until -/// the next whole token accrues — when the bucket is dry. -fn refill_and_take( - tokens: &mut f64, - capacity: f64, - refill_per_sec: f64, - elapsed_secs: f64, -) -> Option { - *tokens = (*tokens + elapsed_secs * refill_per_sec).min(capacity); - if *tokens >= 1.0 { - *tokens -= 1.0; - None - } else { - let deficit = 1.0 - *tokens; - Some(Duration::from_secs_f64(deficit / refill_per_sec)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn loopback_hosts_are_exempt() { - assert!(is_loopback_url("http://localhost:11434")); - assert!(is_loopback_url("http://LOCALHOST/v1")); - assert!(is_loopback_url("http://127.0.0.1:1234/v1")); - assert!(is_loopback_url("http://127.5.6.7")); - assert!(is_loopback_url("http://[::1]:8080")); - } - - #[test] - fn remote_and_malformed_hosts_are_throttled() { - assert!(!is_loopback_url("https://api.openai.com")); - assert!(!is_loopback_url("https://api.openhuman.example/openai/v1")); - assert!(!is_loopback_url("https://10.0.0.5/v1")); // private but not loopback - assert!(!is_loopback_url("not a url")); // malformed → throttled, not bypassed - } - - // ── Bucket math (pure, no clock) ───────────────────────── - - #[test] - fn take_consumes_when_token_available() { - let mut tokens = 5.0; - assert!( - refill_and_take(&mut tokens, 60.0, 1.0, 0.0).is_none(), - "a token is available → consume, no wait" - ); - assert!( - (tokens - 4.0).abs() < 1e-9, - "one token consumed, got {tokens}" - ); - } - - #[test] - fn take_waits_a_full_period_when_empty() { - let mut tokens = 0.0; - // Empty bucket at 1 token/sec → wait ~1s, consuming nothing yet. - let wait = refill_and_take(&mut tokens, 60.0, 1.0, 0.0).expect("must wait"); - assert!((wait.as_secs_f64() - 1.0).abs() < 1e-6, "got {wait:?}"); - assert!(tokens.abs() < 1e-9, "no token consumed while waiting"); - } - - #[test] - fn partial_refill_shortens_the_wait() { - let mut tokens = 0.0; - // 0.25s at 1/sec accrues 0.25 tokens → still <1 → wait the remaining 0.75s. - let wait = refill_and_take(&mut tokens, 60.0, 1.0, 0.25).expect("must wait"); - assert!((wait.as_secs_f64() - 0.75).abs() < 1e-6, "got {wait:?}"); - } - - #[test] - fn refill_is_capped_at_capacity() { - let mut tokens = 50.0; - // A huge idle gap must not let the bucket overflow capacity. - assert!(refill_and_take(&mut tokens, 60.0, 1.0, 10_000.0).is_none()); - assert!( - (tokens - 59.0).abs() < 1e-9, - "capped at 60 then consumed 1 → 59, got {tokens}" - ); - } - - // ── Gating glue (explicit limit, no global mutation) ───── - - #[tokio::test] - async fn disabled_limit_never_blocks() { - for _ in 0..1000 { - acquire_with_limit("https://api.example.com/v1", 0).await; - } - } - - #[tokio::test] - async fn loopback_bypasses_even_at_tight_limit() { - // limit=1 would throttle hard if applied; loopback must bypass it. - for _ in 0..50 { - acquire_with_limit("http://127.0.0.1:11434/v1", 1).await; - } - } - - #[tokio::test] - async fn first_remote_call_does_not_block() { - // Capacity is one token, so a fresh (or idle-refilled) bucket lets - // exactly one request through immediately; pacing of subsequent - // back-to-back requests is covered by `acquire_traverses_wait_branch_*`. - // Unique URL → a bucket isolated from other tests. - let url = "https://burst-test.example/v1"; - let start = Instant::now(); - acquire_with_limit(url, 60).await; - assert!( - start.elapsed() < Duration::from_millis(500), - "first call on a fresh bucket must not block, elapsed {:?}", - start.elapsed() - ); - } - - #[tokio::test] - async fn back_to_back_acquires_are_paced_no_burst() { - // The hard-cap guarantee: capacity is one token, so the first acquire - // passes instantly but the second must wait for refill. 600/min → 10 - // tokens/sec → ~100ms spacing (kept short so the test is fast). A local - // bucket avoids the global registry (and its clear-on-reconfigure). - let bucket = TokenBucket::per_minute(600); - bucket.acquire().await; // consumes the single burst token - let start = Instant::now(); - bucket.acquire().await; // must be paced - assert!( - start.elapsed() >= Duration::from_millis(50), - "second back-to-back acquire must be paced (no burst), elapsed {:?}", - start.elapsed() - ); - } - - #[tokio::test] - async fn acquire_traverses_wait_branch_when_dry() { - // Drive the real throttle path (refill → still dry → sleep → re-check → - // consume) without a slow test: pre-seed the bucket just under one - // token so the wait is ~0.1s of real time. 60/min == 1 token/sec. - let bucket = TokenBucket::per_minute(60); - { - let mut state = bucket.state.lock().await; - state.tokens = 0.9; - state.last_refill = Instant::now(); - } - let start = Instant::now(); - bucket.acquire().await; // must sleep ~100ms for the final 0.1 token - assert!( - start.elapsed() >= Duration::from_millis(50), - "dry bucket must wait for a token, elapsed {:?}", - start.elapsed() - ); - } - - #[tokio::test] - async fn public_entrypoint_reads_global_and_delegates() { - // Exercises the read-global + delegate path; loopback returns at once - // regardless of the configured rate. - acquire_embedding_slot("http://localhost:11434/v1").await; - } - - // ── Global limit setter/getter ─────────────────────────── - - #[test] - fn set_and_read_round_trip() { - set_embedding_rate_limit(0); - assert_eq!(embedding_rate_limit(), 0); - set_embedding_rate_limit(120); - assert_eq!(embedding_rate_limit(), 120); - set_embedding_rate_limit(DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN); // restore global - } -} diff --git a/src/openhuman/embeddings/retry_after.rs b/src/openhuman/embeddings/retry_after.rs deleted file mode 100644 index db0bf8096..000000000 --- a/src/openhuman/embeddings/retry_after.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Retry-After header parsing and 429/503 backoff logic for embedding clients. -//! -//! The HTTP `Retry-After` header may arrive as either: -//! - A non-negative integer: delta-seconds from now (e.g. `Retry-After: 30`) -//! - An HTTP-date string: absolute point in time (e.g. `Retry-After: Wed, 21 Oct 2015 07:28:00 GMT`) -//! -//! This module prefers the delta-seconds form, accepts the common HTTP-date -//! form, and falls back to exponential backoff when the header is absent or -//! unparseable. See RFC 9110 §10.2.4. - -use chrono::{DateTime, Utc}; - -/// Maximum number of 429/503 retries before giving up. -/// -/// Three retries means the client makes up to four total attempts: the -/// original request plus three retries. This caps the per-call delay at -/// roughly `BASE_BACKOFF_MS * 2^2 = 4 s` in the no-`Retry-After` path (or -/// the server-directed duration in the header path). -pub const MAX_429_RETRIES: u32 = 3; - -/// Base exponential-backoff delay in milliseconds when `Retry-After` is absent. -/// -/// Sequence (per attempt): 1 s, 2 s, 4 s — capped at 30 s. -pub const BASE_BACKOFF_MS: u64 = 1_000; - -/// Maximum backoff delay in milliseconds regardless of `Retry-After` or -/// computed exponent. Prevents a misbehaving server from parking the caller -/// for an unreasonably long time. -pub const MAX_BACKOFF_MS: u64 = 30_000; - -/// Parse the `Retry-After` header value into a delay in milliseconds. -/// -/// Accepts the delta-seconds form only (e.g. `"30"`, `"0"`). HTTP-date form -/// is not parsed — fall back to exponential backoff when the value is not a -/// non-negative integer. -/// -/// Returns `None` when the header is absent, empty, or not a valid -/// non-negative integer. -pub fn parse_retry_after_ms(header_value: Option<&str>) -> Option { - parse_retry_after_ms_at(header_value, Utc::now()) -} - -fn parse_retry_after_ms_at(header_value: Option<&str>, now: DateTime) -> Option { - let s = header_value?.trim(); - if s.is_empty() { - return None; - } - if let Ok(secs) = s.parse::() { - return Some(secs.saturating_mul(1_000).min(MAX_BACKOFF_MS)); - } - let retry_at = DateTime::parse_from_rfc2822(s).ok()?.with_timezone(&Utc); - let delay = retry_at.signed_duration_since(now); - if delay.num_milliseconds() <= 0 { - return Some(0); - } - u64::try_from(delay.num_milliseconds()) - .ok() - .map(|ms| ms.min(MAX_BACKOFF_MS)) -} - -/// Compute the delay for attempt `n` (0-indexed) with optional `Retry-After` -/// override. Uses exponential backoff as fallback: -/// `BASE_BACKOFF_MS * 2^n`, capped at `MAX_BACKOFF_MS`. -pub fn backoff_ms_for_attempt(attempt: u32, retry_after_header: Option<&str>) -> u64 { - if let Some(ms) = parse_retry_after_ms(retry_after_header) { - return ms; - } - BASE_BACKOFF_MS - .saturating_mul(2u64.saturating_pow(attempt)) - .min(MAX_BACKOFF_MS) -} - -#[cfg(test)] -mod tests { - use super::*; - - // ── parse_retry_after_ms ──────────────────────────────────── - - #[test] - fn parses_integer_seconds() { - assert_eq!(parse_retry_after_ms(Some("30")), Some(30_000)); - } - - #[test] - fn parses_zero() { - assert_eq!(parse_retry_after_ms(Some("0")), Some(0)); - } - - #[test] - fn parses_whitespace_padded() { - assert_eq!(parse_retry_after_ms(Some(" 5 ")), Some(5_000)); - } - - #[test] - fn caps_at_max_backoff() { - // 9999 s * 1000 ms/s > MAX_BACKOFF_MS - assert_eq!(parse_retry_after_ms(Some("9999")), Some(MAX_BACKOFF_MS)); - } - - #[test] - fn returns_none_for_http_date() { - let now = DateTime::parse_from_rfc2822("Wed, 21 Oct 2015 07:27:55 GMT") - .unwrap() - .with_timezone(&Utc); - assert_eq!( - parse_retry_after_ms_at(Some("Wed, 21 Oct 2015 07:28:00 GMT"), now), - Some(5_000) - ); - } - - #[test] - fn http_date_in_the_past_returns_zero() { - let now = DateTime::parse_from_rfc2822("Wed, 21 Oct 2015 07:28:05 GMT") - .unwrap() - .with_timezone(&Utc); - assert_eq!( - parse_retry_after_ms_at(Some("Wed, 21 Oct 2015 07:28:00 GMT"), now), - Some(0) - ); - } - - #[test] - fn returns_none_for_none_input() { - assert_eq!(parse_retry_after_ms(None), None); - } - - #[test] - fn returns_none_for_empty_string() { - assert_eq!(parse_retry_after_ms(Some("")), None); - } - - #[test] - fn returns_none_for_negative() { - // Negative integers are not valid delta-seconds per RFC 9110. - assert_eq!(parse_retry_after_ms(Some("-1")), None); - } - - // ── backoff_ms_for_attempt ───────────────────────────────── - - #[test] - fn uses_retry_after_when_present() { - assert_eq!(backoff_ms_for_attempt(0, Some("5")), 5_000); - assert_eq!(backoff_ms_for_attempt(2, Some("5")), 5_000); - } - - #[test] - fn falls_back_to_exponential_when_header_absent() { - assert_eq!(backoff_ms_for_attempt(0, None), BASE_BACKOFF_MS); - assert_eq!(backoff_ms_for_attempt(1, None), BASE_BACKOFF_MS * 2); - assert_eq!(backoff_ms_for_attempt(2, None), BASE_BACKOFF_MS * 4); - } - - #[test] - fn falls_back_to_exponential_when_header_unparseable() { - assert_eq!( - backoff_ms_for_attempt(0, Some("not-a-number")), - BASE_BACKOFF_MS - ); - } - - #[test] - fn exponential_caps_at_max() { - // 2^10 = 1024 — well past the cap - assert_eq!(backoff_ms_for_attempt(10, None), MAX_BACKOFF_MS); - } -} diff --git a/src/openhuman/embeddings/rpc.rs b/src/openhuman/embeddings/rpc.rs index 609f54545..dc1e27848 100644 --- a/src/openhuman/embeddings/rpc.rs +++ b/src/openhuman/embeddings/rpc.rs @@ -1037,13 +1037,15 @@ mod tests { /// A 404/405 (no `/embeddings` route) keeps its dedicated code. #[test] fn classify_embed_probe_rejects_endpoint_absent() { - assert_eq!( - reject_code(EmbedProbe::Failed( - "Embedding API error (404 Not Found): no route".into() - )) - .as_deref(), - Some("EMBEDDINGS_ENDPOINT_NO_API") - ); + for detail in [ + "Embedding API error (404 Not Found): no route", + "openai embeddings returned HTTP 404 Not Found: no route", + ] { + assert_eq!( + reject_code(EmbedProbe::Failed(detail.into())).as_deref(), + Some("EMBEDDINGS_ENDPOINT_NO_API") + ); + } } /// Any other failure (5xx/auth/network) and timeouts both reject — the diff --git a/src/openhuman/embeddings/voyage.rs b/src/openhuman/embeddings/voyage.rs deleted file mode 100644 index 44d84010a..000000000 --- a/src/openhuman/embeddings/voyage.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! Voyage AI embedding provider — direct API access with user's own key. -//! -//! Voyage's `/v1/embeddings` endpoint speaks a superset of the OpenAI -//! embeddings contract (same request/response shape, plus extra fields -//! like `input_type` and `truncation`). We delegate to `OpenAiEmbedding` -//! for the HTTP plumbing and just set the correct base URL + auth. - -use async_trait::async_trait; - -use super::openai::OpenAiEmbedding; -use super::EmbeddingProvider; - -pub const VOYAGE_API_BASE: &str = "https://api.voyageai.com"; -pub const VOYAGE_DEFAULT_MODEL: &str = "voyage-3-large"; -pub const VOYAGE_DEFAULT_DIMS: usize = 1024; - -pub struct VoyageEmbedding { - inner: OpenAiEmbedding, -} - -impl VoyageEmbedding { - pub fn new(api_key: &str, model: &str, dims: usize) -> Self { - let model = if model.is_empty() { - VOYAGE_DEFAULT_MODEL - } else { - model - }; - let dims = if dims == 0 { VOYAGE_DEFAULT_DIMS } else { dims }; - - Self { - inner: OpenAiEmbedding::new(VOYAGE_API_BASE, api_key, model, dims) - .with_required_api_key(true), - } - } - - /// Construct a Voyage-compatible provider with a custom API base URL. - /// - /// The hosted Voyage endpoint remains the default for [`Self::new`]; this - /// constructor supports local mocks and compatible deployments. - pub fn new_with_base_url(api_key: &str, model: &str, dims: usize, base_url: &str) -> Self { - let model = if model.is_empty() { - VOYAGE_DEFAULT_MODEL - } else { - model - }; - let dims = if dims == 0 { VOYAGE_DEFAULT_DIMS } else { dims }; - - Self { - inner: OpenAiEmbedding::new(base_url, api_key, model, dims), - } - } -} - -#[async_trait] -impl EmbeddingProvider for VoyageEmbedding { - fn name(&self) -> &str { - "voyage" - } - - fn model_id(&self) -> &str { - self.inner.model_id() - } - - fn dimensions(&self) -> usize { - self.inner.dimensions() - } - - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - self.inner.embed(texts).await - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn name_and_defaults() { - let p = VoyageEmbedding::new("test-key", "", 0); - assert_eq!(p.name(), "voyage"); - assert_eq!(p.model_id(), VOYAGE_DEFAULT_MODEL); - assert_eq!(p.dimensions(), VOYAGE_DEFAULT_DIMS); - } - - #[test] - fn custom_model_and_dims() { - let p = VoyageEmbedding::new("test-key", "voyage-code-3", 1024); - assert_eq!(p.model_id(), "voyage-code-3"); - assert_eq!(p.dimensions(), 1024); - } - - #[test] - fn signature_format() { - let p = VoyageEmbedding::new("k", "voyage-3-large", 1024); - assert_eq!( - p.signature(), - "provider=voyage;model=voyage-3-large;dims=1024" - ); - } -} diff --git a/src/openhuman/embeddings/voyage_adapter.rs b/src/openhuman/embeddings/voyage_adapter.rs new file mode 100644 index 000000000..9da55c8f4 --- /dev/null +++ b/src/openhuman/embeddings/voyage_adapter.rs @@ -0,0 +1,69 @@ +//! Compatibility wrapper for tinyagents' Voyage model. + +use async_trait::async_trait; +use tinyagents::harness::embeddings::{ + EmbeddingModel, VoyageEmbeddingModel, VOYAGE_API_BASE, VOYAGE_DEFAULT_DIMENSIONS, + VOYAGE_DEFAULT_MODEL, +}; + +use super::EmbeddingProvider; + +pub struct VoyageEmbedding { + inner: VoyageEmbeddingModel, +} + +impl VoyageEmbedding { + pub fn new(api_key: &str, model: &str, dimensions: usize) -> Self { + Self::new_with_base_url(api_key, model, dimensions, VOYAGE_API_BASE) + } + + pub fn new_with_base_url( + api_key: &str, + model: &str, + dimensions: usize, + base_url: &str, + ) -> Self { + Self { + inner: VoyageEmbeddingModel::with_options( + api_key, + if model.is_empty() { + VOYAGE_DEFAULT_MODEL + } else { + model + }, + if dimensions == 0 { + VOYAGE_DEFAULT_DIMENSIONS + } else { + dimensions + }, + base_url, + ), + } + } +} + +#[async_trait] +impl EmbeddingProvider for VoyageEmbedding { + fn name(&self) -> &str { + self.inner.name() + } + fn model_id(&self) -> &str { + self.inner.model_id() + } + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + fn signature(&self) -> String { + self.inner.signature() + } + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let texts = texts + .iter() + .map(|text| (*text).to_owned()) + .collect::>(); + self.inner + .embed(&texts) + .await + .map_err(|error| anyhow::anyhow!(error)) + } +} diff --git a/src/openhuman/memory/ingest_pipeline.rs b/src/openhuman/memory/ingest_pipeline.rs index 09e522ea3..f2bd519c3 100644 --- a/src/openhuman/memory/ingest_pipeline.rs +++ b/src/openhuman/memory/ingest_pipeline.rs @@ -1,75 +1,19 @@ -//! Ingest orchestrator for the async memory-tree pipeline. -//! -//! The hot path now does: -//! `canonicalise -> chunk -> fast score -> persist chunks/score rows -> enqueue extract jobs` -//! -//! The slower work (full extraction, admission, tree buffering, sealing, -//! topic routing, daily digests) runs out of the SQLite-backed jobs queue. +//! Product shell over tinycortex on-demand ingestion. use anyhow::Result; -use serde::{Deserialize, Serialize}; use crate::core::event_bus::{publish_global, DomainEvent}; use crate::openhuman::config::Config; -use crate::openhuman::memory::util::redact::redact; -use crate::openhuman::memory_queue::{self as jobs, ExtractChunkPayload, NewJob}; -use crate::openhuman::memory_store::chunks::store::{self as chunk_store, RawRef}; -use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_store::chunks::{chunk_markdown, ChunkerInput, ChunkerOptions}; -use crate::openhuman::memory_store::content as content_store; +use crate::openhuman::memory_store::chunks::store::RawRef; use crate::openhuman::memory_sync::canonicalize::{ chat::{self, ChatBatch}, document::{self, DocumentInput}, email::{self, EmailThread}, CanonicalisedSource, }; -use crate::openhuman::memory_tree::score::{self, ScoreResult, ScoringConfig}; -use std::time::{SystemTime, UNIX_EPOCH}; -const BODY_PREVIEW_MAX_BYTES: usize = 2048; +pub use tinycortex::memory::ingest::IngestSummary as IngestResult; -/// Outcome of one ingest call. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct IngestResult { - pub source_id: String, - /// Number of chunks persisted and queued for async extraction. - pub chunks_written: usize, - /// Number of chunks the cheap fast-score path would drop. Final admission - /// still happens later in the extract job. - pub chunks_dropped: usize, - /// IDs of all chunks written and queued. - pub chunk_ids: Vec, - /// True when this ingest was a no-op because `(source_kind, source_id)` - /// had already been ingested. Memory items are append-only — the - /// summariser tree must not see the same source twice. - #[serde(default)] - pub already_ingested: bool, -} - -impl IngestResult { - fn empty(source_id: &str) -> Self { - Self { - source_id: source_id.to_string(), - chunks_written: 0, - chunks_dropped: 0, - chunk_ids: Vec::new(), - already_ingested: false, - } - } - - fn already_ingested(source_id: &str) -> Self { - Self { - source_id: source_id.to_string(), - chunks_written: 0, - chunks_dropped: 0, - chunk_ids: Vec::new(), - already_ingested: true, - } - } -} - -/// Ingest a batch of chat messages: canonicalise → chunk → fast-score → persist -/// → enqueue async extract jobs. Returns a noop [`IngestResult`] on an empty batch. pub async fn ingest_chat( config: &Config, source_id: &str, @@ -77,22 +21,17 @@ pub async fn ingest_chat( tags: Vec, batch: ChatBatch, ) -> Result { - // No source-level gate for chat: a chat `source_id` (e.g. - // `slack:{connection_id}`) is a stream identifier — many batches / - // buckets accumulate into the same source tree over time. The gate - // would make every bucket after the first a no-op. Chunk-level - // idempotency (`chunk_id` includes content) still prevents true - // replay duplicates from reaching the summariser. let canonical = - match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? { - Some(c) => c, - None => return Ok(IngestResult::empty(source_id)), - }; - persist(config, source_id, canonical, None, None).await + chat::canonicalise(source_id, owner, &tags, batch.clone()).map_err(anyhow::Error::msg)?; + let (memory, sink, scoring) = crate::openhuman::tinycortex::ingest_context(config); + let result = tinycortex::memory::ingest::ingest_chat( + &memory, source_id, owner, tags, batch, &sink, &scoring, + ) + .await?; + publish_canonicalized(source_id, canonical.as_ref(), &result); + Ok(result) } -/// Ingest an email thread: canonicalise → chunk → fast-score → persist → enqueue -/// async extract jobs. Returns a noop [`IngestResult`] on an empty thread. pub async fn ingest_email( config: &Config, source_id: &str, @@ -100,24 +39,17 @@ pub async fn ingest_email( tags: Vec, thread: EmailThread, ) -> Result { - // No source-level gate for email: gmail per-participant ingest - // groups many threads under one `source_id` (e.g. - // `gmail:{participants_hash}`) and appends as new threads arrive. - // The gate would block all but the first thread. Chunk-level - // idempotency still protects against true replays. let canonical = - match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? { - Some(c) => c, - None => return Ok(IngestResult::empty(source_id)), - }; - persist(config, source_id, canonical, None, None).await + email::canonicalise(source_id, owner, &tags, thread.clone()).map_err(anyhow::Error::msg)?; + let (memory, sink, scoring) = crate::openhuman::tinycortex::ingest_context(config); + let result = tinycortex::memory::ingest::ingest_email( + &memory, source_id, owner, tags, thread, &sink, &scoring, + ) + .await?; + publish_canonicalized(source_id, canonical.as_ref(), &result); + Ok(result) } -/// Ingest an email thread whose chunk bodies are backed by a raw archive file. -/// -/// Raw refs are committed in the same transaction as the chunk rows and -/// `extract_chunk` jobs, so workers can never observe a raw-backed email chunk -/// without a resolvable body source. pub async fn ingest_email_with_raw_refs( config: &Config, source_id: &str, @@ -127,15 +59,16 @@ pub async fn ingest_email_with_raw_refs( raw_refs: Vec, ) -> Result { let canonical = - match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? { - Some(c) => c, - None => return Ok(IngestResult::empty(source_id)), - }; - persist(config, source_id, canonical, None, Some(raw_refs)).await + email::canonicalise(source_id, owner, &tags, thread.clone()).map_err(anyhow::Error::msg)?; + let (memory, sink, scoring) = crate::openhuman::tinycortex::ingest_context(config); + let result = tinycortex::memory::ingest::ingest_email_with_raw_refs( + &memory, source_id, owner, tags, thread, raw_refs, &sink, &scoring, + ) + .await?; + publish_canonicalized(source_id, canonical.as_ref(), &result); + Ok(result) } -/// Ingest a single document: canonicalise → chunk → fast-score → persist → -/// enqueue async extract jobs. Returns a noop [`IngestResult`] on empty input. pub async fn ingest_document( config: &Config, source_id: &str, @@ -157,18 +90,6 @@ pub async fn ingest_document_with_scope( ingest_document_versioned(config, source_id, owner, tags, doc, path_scope, None).await } -/// Like [`ingest_document_with_scope`] but version-aware. -/// -/// When `version_ms` is `Some`, the source-ingest gate is keyed by -/// `{source_id}@{version_ms}` instead of the bare `source_id`, so a later -/// revision of the SAME document (same `source_id`) is admitted -/// **non-destructively** — the prior version's chunks are left in place and -/// the new version is ingested alongside them. Chunks keep the clean -/// `source_id` (their `doc_id`), and the retrieval layer surfaces only the -/// latest version per document. -/// -/// `version_ms = None` is identical to [`ingest_document_with_scope`] -/// (bare-`source_id` gate), so non-versioned document sources are unaffected. pub async fn ingest_document_versioned( config: &Config, source_id: &str, @@ -178,635 +99,93 @@ pub async fn ingest_document_versioned( path_scope: Option, version_ms: Option, ) -> Result { - let gate_key = match version_ms { - Some(v) => format!("{source_id}@{v}"), - None => source_id.to_string(), - }; - if already_ingested(config, SourceKind::Document, &gate_key).await? { - log::debug!( - "[memory::ingest_pipeline] skip ingest_document — source_id_hash={} version_ms={:?} already ingested", - redact(source_id), - version_ms - ); - return Ok(IngestResult::already_ingested(source_id)); - } - let canonical = match document::canonicalise(source_id, owner, &tags, doc, path_scope) - .map_err(anyhow::Error::msg)? - { - Some(c) => c, - None => return Ok(IngestResult::empty(source_id)), - }; - persist(config, source_id, canonical, version_ms, None).await + let canonical = + document::canonicalise(source_id, owner, &tags, doc.clone(), path_scope.clone()) + .map_err(anyhow::Error::msg)?; + let (memory, sink, scoring) = crate::openhuman::tinycortex::ingest_context(config); + let result = tinycortex::memory::ingest::ingest_document_versioned( + &memory, source_id, owner, tags, doc, path_scope, version_ms, &sink, &scoring, + ) + .await?; + publish_canonicalized(source_id, canonical.as_ref(), &result); + Ok(result) } -/// Best-effort pre-canonicalisation check. The transactional claim inside -/// [`persist`] is what actually serialises concurrent ingests; this lookup -/// just spares the canonicaliser when we already know the source is a dup. -async fn already_ingested( - config: &Config, - source_kind: SourceKind, +fn publish_canonicalized( source_id: &str, -) -> Result { - let cfg = config.clone(); - let sid = source_id.to_string(); - tokio::task::spawn_blocking(move || chunk_store::is_source_ingested(&cfg, source_kind, &sid)) - .await - .map_err(|e| anyhow::anyhow!("already_ingested join error: {e}"))? -} - -async fn persist( - config: &Config, - source_id: &str, - canonical: CanonicalisedSource, - gate_version_ms: Option, - raw_refs_for_chunks: Option>, -) -> Result { - let source_kind_for_store = canonical.metadata.source_kind; - - // Capture body_preview before the canonical markdown is moved into the chunker. - // For email and document sources: the trailing canonical markdown, capped at - // 2 048 bytes, is enough for signature parsing and similar lightweight - // subscribers. Chat sources are conversational and have no trailing structure - // worth scanning, so they get body_preview = None. - let body_preview: Option = match source_kind_for_store { - SourceKind::Email | SourceKind::Document => { - // Guard the preview computation so a single malformed document - // never kills the ingest worker. `markdown_body_preview` contains - // defensive checks, but wrap at the call-site too for belt-and-braces - // protection against any future panic regression in its dependency chain. - let md_for_preview = canonical.markdown.clone(); - match std::panic::catch_unwind(move || markdown_body_preview(&md_for_preview)) { - Ok(preview) => Some(preview), - Err(_) => { - log::error!( - "[memory::ingest_pipeline] markdown_body_preview panicked for source_id_hash={}; falling back to no preview", - crate::openhuman::memory::util::redact::redact(source_id) - ); - None - } - } - } - _ => None, + canonical: Option<&CanonicalisedSource>, + result: &IngestResult, +) { + let Some(canonical) = canonical else { + return; }; - - let input = ChunkerInput { - source_kind: canonical.metadata.source_kind, - source_id: source_id.to_string(), - markdown: canonical.markdown, - metadata: canonical.metadata, - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - if chunks.is_empty() { - return Ok(IngestResult::empty(source_id)); - } - - // Phase MD-content: write chunk bodies to disk before the SQLite upsert. - // stage_chunks is sync I/O; run it here (still on the tokio thread) before - // spawn_blocking so errors surface before the DB transaction opens. - let content_root = config.memory_tree_content_root(); - let staged = content_store::stage_chunks(&content_root, &chunks) - .map_err(|e| anyhow::anyhow!("[memory::ingest_pipeline] stage_chunks failed: {e}"))?; - - let scoring_cfg = ScoringConfig::from_config(config); - let scores = score::score_chunks_fast(&chunks, &scoring_cfg).await?; - if scores.len() != chunks.len() { - anyhow::bail!( - "[memory::ingest_pipeline] scorer length mismatch: chunks={} scores={}", - chunks.len(), - scores.len() - ); - } - - let all_results: Vec<(ScoreResult, i64)> = chunks - .iter() - .zip(scores) - .map(|(chunk, result)| (result, chunk.metadata.timestamp.timestamp_millis())) - .collect(); - let dropped = all_results.iter().filter(|(r, _)| !r.kept).count(); - - let config_owned = config.clone(); - let staged_for_store = staged.clone(); - let results_for_store = all_results.clone(); - let source_id_for_store = source_id.to_string(); - let raw_refs_for_store = raw_refs_for_chunks.clone(); - let written = tokio::task::spawn_blocking(move || -> Result> { - use std::collections::{HashMap, HashSet}; - chunk_store::with_connection(&config_owned, |conn| { - // IMMEDIATE, not the default DEFERRED: this transaction reads - // (get_chunk_lifecycle_status_tx) before it writes - // (upsert_staged_chunks_tx). A DEFERRED tx takes only a read - // lock at BEGIN and tries to upgrade to a write lock on the - // first write; under contention with the memory_tree worker - // pool SQLite returns SQLITE_BUSY *immediately* for that - // upgrade and does NOT invoke the busy handler (deadlock - // avoidance), so the connection's 15s busy_timeout is bypassed - // and Gmail/Composio ingest fails every message with "database - // is locked", stalling composio_sync past its 30s RPC cap. - // IMMEDIATE acquires the write lock at BEGIN, where the busy - // handler / busy_timeout DOES apply, so writers serialise and - // wait instead of failing fast. - let tx = rusqlite::Transaction::new_unchecked( - conn, - rusqlite::TransactionBehavior::Immediate, - )?; - - // Authoritative source-level gate (documents only). - // - // For documents, `source_id` identifies a single immutable - // file (one notion page, one drive doc). `is_source_ingested` - // above is a best-effort fast-path; this row insert is what - // actually serialises concurrent ingests of the same - // document and prevents the same content from flowing - // through extract → admit → buffer → seal twice. - // - // Chat and email don't get this gate: their `source_id` - // is a *stream* identifier (e.g. slack workspace, gmail - // participant group) under which many batches / threads - // accumulate over time. The chunk-level idempotency in - // the rest of this transaction is enough to swallow - // genuine replays without blocking legitimate appends. - if source_kind_for_store == SourceKind::Document { - let now_ms = chrono::Utc::now().timestamp_millis(); - // Versioned sources (Notion) claim a per-version gate key - // `{source_id}@{version_ms}` so a later revision of the SAME - // document is admitted (non-destructively, alongside the - // prior version) instead of short-circuiting on the bare - // source_id. Chunks themselves keep the clean source_id so - // per-document grouping (doc_id) stays stable. Non-versioned - // documents (version_ms = None) use the bare source_id as - // before — behaviour unchanged. - let gate_key = match gate_version_ms { - Some(v) => format!("{source_id_for_store}@{v}"), - None => source_id_for_store.clone(), - }; - let claimed = chunk_store::claim_source_ingest_tx( - &tx, - source_kind_for_store, - &gate_key, - now_ms, - )?; - if !claimed { - log::debug!( - "[memory::ingest_pipeline] persist gate: document already ingested source_id_hash={}", - redact(&source_id_for_store) - ); - // Drop the (empty) transaction implicitly; nothing to commit. - return Ok(None); - } - } - - // Read each chunk's CURRENT lifecycle BEFORE the upsert. This - // is the "did this chunk exist before this batch" snapshot, - // because `upsert_staged_chunks_tx` will either preserve the - // existing row's lifecycle (UPDATE doesn't touch the column) or - // insert a new row that picks up the column DEFAULT — so reading - // post-upsert can't distinguish "brand new" from - // "already-admitted-from-prior-ingest". - let mut prior: HashMap> = HashMap::new(); - for s in &staged_for_store { - let status = chunk_store::get_chunk_lifecycle_status_tx(&tx, &s.chunk.id)?; - prior.insert(s.chunk.id.clone(), status); - } - - let n = chunk_store::upsert_staged_chunks_tx(&tx, &staged_for_store)?; - - if let Some(ref refs) = raw_refs_for_store { - for s in &staged_for_store { - chunk_store::set_chunk_raw_refs_tx(&tx, &s.chunk.id, refs)?; - } - } - - // Re-ingest of identical content (same chunk_id) must NOT - // downgrade chunks that have already progressed through the - // async pipeline. Without this guard, a re-ingest would reset - // every chunk to 'pending_extraction' and enqueue a fresh - // `extract_chunk` job — sending already-buffered/sealed - // chunks back through extract → admit → append, ultimately - // duplicating them into a second summary in the same tree. - // - // Schedule a chunk for processing when its PRE-upsert state - // was either absent (genuinely new) or already - // `pending_extraction` (a prior ingest crashed before extract - // ran). Anything else — `admitted`, `buffered`, `sealed`, - // `dropped` — is past the point of accepting new work, so - // leave the lifecycle alone and skip the extract enqueue. - let mut to_schedule: HashSet = HashSet::new(); - for s in &staged_for_store { - let pre = prior.get(&s.chunk.id).cloned().flatten(); - let needs_processing = matches!( - pre.as_deref(), - None | Some(chunk_store::CHUNK_STATUS_PENDING_EXTRACTION), - ); - if needs_processing { - chunk_store::set_chunk_lifecycle_status_tx( - &tx, - &s.chunk.id, - chunk_store::CHUNK_STATUS_PENDING_EXTRACTION, - )?; - to_schedule.insert(s.chunk.id.clone()); - } - } - - for (result, ts_ms) in &results_for_store { - if !to_schedule.contains(&result.chunk_id) { - // Chunk has already progressed past pending_extraction - // on a prior ingest — skip score re-persist and don't - // enqueue a duplicate extract job. - continue; - } - score::persist_score_tx(&tx, result, *ts_ms, None)?; - let extract = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: result.chunk_id.clone(), - })?; - let _ = jobs::enqueue_tx(&tx, &extract)?; - } - tx.commit()?; - Ok(Some(n)) - }) - }) - .await - .map_err(|e| anyhow::anyhow!("persist join error: {e}"))??; - - let written = match written { - Some(n) => n, - None => { - // Lost the race against a concurrent ingest of the same source — - // the other writer claimed the row first. No work was committed. - return Ok(IngestResult::already_ingested(source_id)); - } - }; - - jobs::wake_workers(); - - let chunk_ids: Vec = staged.iter().map(|s| s.chunk.id.clone()).collect(); - - // Emit DocumentCanonicalized so Phase 2 producers (e.g. email-signature parser) - // can react to new canonicalised content. Non-fatal: ingest has already succeeded. - // `source_kind_for_store` is Copy so it is still accessible here after the closure. - let now_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs_f64(); - publish_global(DomainEvent::DocumentCanonicalized { - source_id: source_id.to_string(), - source_kind: source_kind_for_store.as_str().to_string(), - chunks_written: written, - chunk_ids: chunk_ids.clone(), - canonicalized_at: now_secs, - body_preview, - }); - tracing::debug!( - "[memory::tree::ingest] published DocumentCanonicalized source_id={} chunks={}", - source_id, - written - ); - - Ok(IngestResult { - source_id: source_id.to_string(), - chunks_written: written, - chunks_dropped: dropped, - chunk_ids, - already_ingested: false, - }) -} - -/// Returns the trailing slice of `md` capped at [`BODY_PREVIEW_MAX_BYTES`] bytes. -/// -/// Uses `ceil_char_boundary` (rounds the cut point *forward*) so the returned -/// slice is always `<= BODY_PREVIEW_MAX_BYTES` bytes — `floor_char_boundary` -/// (rounds backward) can return up to 3 extra bytes when the cut falls inside -/// a multi-byte codepoint, violating the hard cap. -fn markdown_body_preview(md: &str) -> String { - let len = md.len(); - if len <= BODY_PREVIEW_MAX_BYTES { - md.to_string() + let source_kind = canonical.metadata.source_kind.as_str(); + let body_preview = if matches!(source_kind, "email" | "document") { + utf8_suffix(&canonical.markdown, 2048) } else { - let start = crate::openhuman::util::ceil_char_boundary(md, len - BODY_PREVIEW_MAX_BYTES); - debug_assert!( - md.is_char_boundary(start), - "ceil_char_boundary returned non-boundary {start} for len={len}" - ); - // ceil_char_boundary can return `len` when every remaining byte is a - // continuation byte; fall back to the full string rather than panicking. - if start > len || !md.is_char_boundary(start) { - log::error!( - "[memory::ingest_pipeline] ceil_char_boundary returned invalid boundary start={start} len={len}; returning full markdown" - ); - md.to_string() - } else { - md[start..].to_string() - } + utf8_prefix(&canonical.markdown, 2048) + }; + publish_global(DomainEvent::DocumentCanonicalized { + source_id: source_id.into(), + source_kind: canonical.metadata.source_kind.as_str().into(), + chunks_written: result.chunks_written, + chunk_ids: result.chunk_ids.clone(), + canonicalized_at: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64(), + body_preview: Some(body_preview), + }); +} + +fn utf8_suffix(value: &str, max_bytes: usize) -> String { + if value.len() <= max_bytes { + return value.to_owned(); } + let target = value.len().saturating_sub(max_bytes); + let start = value + .char_indices() + .map(|(index, _)| index) + .find(|index| *index >= target) + .unwrap_or(value.len()); + value[start..].to_owned() +} + +fn utf8_prefix(value: &str, max_bytes: usize) -> String { + let end = value + .char_indices() + .map(|(index, _)| index) + .take_while(|index| *index <= max_bytes) + .last() + .unwrap_or(0); + let end = if value.len() <= max_bytes { + value.len() + } else if end == 0 { + 0 + } else { + end + }; + value[..end].to_string() } #[cfg(test)] mod tests { - use super::*; - use crate::openhuman::memory_queue::drain_until_idle; - use crate::openhuman::memory_store::chunks::store::{ - count_chunks, count_chunks_by_lifecycle_status, get_chunk_embedding, list_chunks, - ListChunksQuery, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, - }; - use crate::openhuman::memory_store::chunks::types::SourceKind; - use crate::openhuman::memory_sync::canonicalize::chat::ChatMessage; - use crate::openhuman::memory_tree::score::store::{count_scores, lookup_entity}; - use chrono::{TimeZone, Utc}; - use tempfile::TempDir; + use super::{utf8_prefix, utf8_suffix}; - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn substantive_batch() -> ChatBatch { - ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![ - ChatMessage { - author: "alice".into(), - timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "We are planning to ship the Phoenix migration on Friday after reviewing the runbook and staging results. alice@example.com" - .into(), - source_ref: Some("slack://m1".into()), - }, - ChatMessage { - author: "bob".into(), - timestamp: Utc.timestamp_millis_opt(1_700_000_010_000).unwrap(), - text: "Confirmed, I will handle the coordination and launch tracking tonight." - .into(), - source_ref: None, - }, - ], - } - } - - #[tokio::test] - async fn ingest_chat_writes_and_queue_drains_to_admitted_chunk() { - let (_tmp, cfg) = test_config(); - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()) - .await - .unwrap(); - // Greedy packing: both small messages fit under 10k token budget - // and are packed into a single chunk. - assert_eq!(out.chunks_written, 1); - assert_eq!(count_chunks(&cfg).unwrap(), 1); - - drain_until_idle(&cfg).await.unwrap(); - - // Final lifecycle is `buffered`: extract → admitted → append_buffer → buffered. - // The single packed chunk does not cross INPUT_TOKEN_BUDGET so no seal fires. - assert_eq!( - count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_BUFFERED).unwrap(), - 1 - ); - assert!(count_scores(&cfg).unwrap() >= 1); - assert_eq!( - lookup_entity(&cfg, "email:alice@example.com", None) - .unwrap() - .len(), - 1 - ); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat); - // #002 FR-002: `test_config()` configures NO embeddings provider, so the - // extract handler correctly SKIPS embedding rather than persisting a - // zero-vector that would silently poison semantic recall. The chunk is - // written embedding-less and stays re-embeddable once a provider is set - // up. (With a provider configured the embedding is present — see the - // `build_write_embedder` tests in memory_tree/score/embed/factory.rs.) - assert!(get_chunk_embedding(&cfg, &out.chunk_ids[0]) - .unwrap() - .is_none()); - } - - #[tokio::test] - async fn low_signal_chunks_end_up_dropped_after_queue_processing() { - let (_tmp, cfg) = test_config(); - let batch = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![ChatMessage { - author: "alice".into(), - timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "+1".into(), - source_ref: None, - }], - }; - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch) - .await - .unwrap(); - assert_eq!(out.chunks_written, 1); - assert_eq!(count_chunks(&cfg).unwrap(), 1); - - drain_until_idle(&cfg).await.unwrap(); - - assert_eq!( - count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_DROPPED).unwrap(), - 1 - ); - assert_eq!(count_scores(&cfg).unwrap(), 1); - } - - #[tokio::test] - async fn ingest_chat_empty_batch_is_noop() { - let (_tmp, cfg) = test_config(); - let batch = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![], - }; - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch) - .await - .unwrap(); - assert_eq!(out.chunks_written, 0); - assert_eq!(count_chunks(&cfg).unwrap(), 0); - assert_eq!(count_scores(&cfg).unwrap(), 0); + #[test] + fn preview_keeps_short_text() { + assert_eq!(utf8_prefix("hello", 2048), "hello"); } #[test] - fn markdown_body_preview_respects_utf8_boundary_and_byte_cap() { - let md = format!("{}{}{}\n", "a".repeat(17), '\u{200c}', "b".repeat(2045)); - let requested_start = md.len() - BODY_PREVIEW_MAX_BYTES; - assert!( - !md.is_char_boundary(requested_start), - "test fixture must put the requested preview boundary inside a multi-byte character" - ); - - let preview = markdown_body_preview(&md); - - assert!(preview.len() <= BODY_PREVIEW_MAX_BYTES); - assert_eq!(preview, format!("{}\n", "b".repeat(2045))); - } - - #[tokio::test] - async fn ingest_document_handles_utf8_at_body_preview_boundary() { - let (_tmp, cfg) = test_config(); - let body = format!("{}{}{}", "a".repeat(17), '\u{200c}', "b".repeat(2045)); - - let doc = DocumentInput { - provider: "notion".into(), - title: "Unicode boundary".into(), - body, - modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - source_ref: Some("notion://page/unicode-boundary".into()), - }; - - let out = ingest_document(&cfg, "notion:utf8-boundary", "alice", vec![], doc) - .await - .unwrap(); - assert!(!out.already_ingested); - assert!(out.chunks_written >= 1); - } - - #[tokio::test] - async fn second_ingest_document_with_same_source_id_is_short_circuited() { - let (_tmp, cfg) = test_config(); - let doc = DocumentInput { - provider: "notion".into(), - title: "Launch plan".into(), - body: "Phoenix ships Friday after staging review. alice@example.com owns this.".into(), - modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - source_ref: Some("notion://page/abc".into()), - }; - let first = ingest_document(&cfg, "notion:abc", "alice", vec![], doc.clone()) - .await - .unwrap(); - assert!(!first.already_ingested); - assert!(first.chunks_written >= 1); - - // Even with completely different content under the same source_id, - // the second ingest must not write anything: documents are - // append-only and the source_id is the dedup key. - let mutated = DocumentInput { - body: "totally different content that should NOT make it into the tree".into(), - ..doc - }; - let second = ingest_document(&cfg, "notion:abc", "alice", vec![], mutated) - .await - .unwrap(); - assert!(second.already_ingested); - assert_eq!(second.chunks_written, 0); - assert!(second.chunk_ids.is_empty()); - - drain_until_idle(&cfg).await.unwrap(); - // Only the first ingest's chunks made it into the store. - assert_eq!(count_chunks(&cfg).unwrap(), first.chunks_written as u64); - } - - #[tokio::test] - async fn re_ingest_is_idempotent_on_chunks_and_scores() { - let (_tmp, cfg) = test_config(); - let doc = DocumentInput { - provider: "notion".into(), - title: "Launch plan".into(), - body: "We are planning to ship Phoenix on Friday after review. alice@example.com owns this." - .into(), - modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - source_ref: Some("notion://page/abc".into()), - }; - ingest_document(&cfg, "notion:abc", "alice", vec![], doc.clone()) - .await - .unwrap(); - ingest_document(&cfg, "notion:abc", "alice", vec![], doc) - .await - .unwrap(); - drain_until_idle(&cfg).await.unwrap(); - assert_eq!(count_chunks(&cfg).unwrap(), 1); - assert_eq!(count_scores(&cfg).unwrap(), 1); - } - - // ── multi-byte boundary tests (issue #2073) ────────────────────────────── - - #[test] - fn markdown_body_preview_zwnj_at_exact_boundary() { - // U+200C ZERO WIDTH NON-JOINER is 3 bytes (0xE2 0x80 0x8C). - // Place it at offsets 0, 1, 2 relative to the preview boundary so - // that each byte of the codepoint lands exactly on the nominal cut. - let zwnj = '\u{200c}'; - let zwnj_bytes = zwnj.len_utf8(); // 3 - assert_eq!(zwnj_bytes, 3); - - for offset in 0..zwnj_bytes { - // ascii_prefix || zwnj || ascii_suffix - // The nominal cut point is `ascii_prefix.len() + offset` bytes - // from the start, which lands `offset` bytes into the zwnj. - let prefix_len = BODY_PREVIEW_MAX_BYTES - offset; - let ascii_prefix = "a".repeat(prefix_len + (zwnj_bytes - offset)); - // Build: enough leading bytes so (total - BODY_PREVIEW_MAX_BYTES) falls - // inside the zwnj. - let padding = "x".repeat(50); - let md = format!( - "{}{}{}{}", - padding, - "a".repeat(prefix_len - 50), - zwnj, - "b".repeat(offset + 100) - ); - - let preview = markdown_body_preview(&md); - - // Must not panic, result must be valid UTF-8, and byte length <= cap. - assert!(std::str::from_utf8(preview.as_bytes()).is_ok()); - assert!( - preview.len() <= BODY_PREVIEW_MAX_BYTES, - "offset={offset}: preview len {} exceeds cap {}", - preview.len(), - BODY_PREVIEW_MAX_BYTES - ); - let _ = ascii_prefix; // suppress unused warning - } + fn preview_respects_utf8_byte_boundary() { + assert_eq!(utf8_prefix("aéb", 2), "a"); + assert_eq!(utf8_prefix("éb", 2), "é"); } #[test] - fn markdown_body_preview_figure_space_at_exact_boundary() { - // U+2007 FIGURE SPACE is 3 bytes (0xE2 0x80 0x87). - let fig_space = '\u{2007}'; - let fig_bytes = fig_space.len_utf8(); - assert_eq!(fig_bytes, 3); - - for offset in 0..fig_bytes { - let padding = "x".repeat(50); - let prefix_len = if BODY_PREVIEW_MAX_BYTES > offset + 50 { - BODY_PREVIEW_MAX_BYTES - offset - 50 - } else { - 0 - }; - let md = format!( - "{}{}{}{}", - padding, - "a".repeat(prefix_len), - fig_space, - "b".repeat(offset + 100) - ); - - let preview = markdown_body_preview(&md); - - assert!(std::str::from_utf8(preview.as_bytes()).is_ok()); - assert!( - preview.len() <= BODY_PREVIEW_MAX_BYTES, - "offset={offset}: preview len {} exceeds cap {}", - preview.len(), - BODY_PREVIEW_MAX_BYTES - ); - } - } - - #[test] - fn markdown_body_preview_persian_text() { - // Persian word "سلام‌ها" (hello + plural marker) with embedded U+200C. - let persian_word = "\u{0633}\u{0644}\u{0627}\u{0645}\u{200c}\u{0647}\u{0627}"; - let md = persian_word.repeat(200); - - // Must not panic regardless of where the cut falls. - let preview = markdown_body_preview(&md); - - assert!(std::str::from_utf8(preview.as_bytes()).is_ok()); - assert!(preview.len() <= BODY_PREVIEW_MAX_BYTES); + fn suffix_preview_preserves_trailing_utf8() { + assert_eq!(utf8_suffix("aéb", 2), "b"); + assert_eq!(utf8_suffix("aéb", 3), "éb"); } } diff --git a/src/openhuman/memory/ingestion/mod.rs b/src/openhuman/memory/ingestion/mod.rs index c6b20566a..554a2eca2 100644 --- a/src/openhuman/memory/ingestion/mod.rs +++ b/src/openhuman/memory/ingestion/mod.rs @@ -11,24 +11,17 @@ //! 5. **Persistence**: Upserting the document, text chunks, and graph relations into //! the memory store. -mod parse; -mod regex; -mod rules; -mod types; - pub mod queue; pub mod state; pub use queue::{IngestionJob, IngestionQueue, DEFAULT_QUEUE_CAPACITY}; pub use state::{IngestionState, IngestionStatusSnapshot}; -pub use types::{ +pub use tinycortex::memory::ingest::{ ExtractedEntity, ExtractedRelation, ExtractionMode, MemoryIngestionConfig, MemoryIngestionRequest, MemoryIngestionResult, DEFAULT_MEMORY_EXTRACTION_MODEL, }; -use parse::{enrich_document_metadata, parse_document}; use serde_json::json; -use types::ParsedIngestion; use crate::openhuman::memory_store::types::NamespaceDocumentInput; use crate::openhuman::memory_store::UnifiedMemory; @@ -41,34 +34,19 @@ impl UnifiedMemory { &self, request: MemoryIngestionRequest, ) -> Result { - let parsed = parse_document( - &request.document.content, - &request.document.title, - &request.config, - ) - .await; - let (enriched_input, tags) = - enrich_document_metadata(&request.document, &parsed, &request.config); + let (enriched_input, mut extraction) = + tinycortex::memory::ingest::extract_enriched_document( + &request.document, + &request.config, + ); let namespace = Self::sanitize_namespace(&enriched_input.namespace); let document_id = self.upsert_document(enriched_input).await?; - self.upsert_graph_relations(&namespace, &document_id, &parsed, &request.config) + self.upsert_graph_relations(&namespace, &document_id, &extraction, &request.config) .await?; - - Ok(MemoryIngestionResult { - document_id, - namespace, - model_name: request.config.model_name, - extraction_mode: request.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, - }) + extraction.document_id = document_id; + extraction.namespace = namespace; + Ok(extraction) } /// Extract entities/relations and write them to the graph for a document @@ -83,28 +61,15 @@ impl UnifiedMemory { document: &NamespaceDocumentInput, config: &MemoryIngestionConfig, ) -> Result { - let parsed = parse_document(&document.content, &document.title, config).await; + let (_enriched, mut extraction) = + tinycortex::memory::ingest::extract_enriched_document(document, config); let namespace = Self::sanitize_namespace(&document.namespace); - self.upsert_graph_relations(&namespace, document_id, &parsed, config) + self.upsert_graph_relations(&namespace, document_id, &extraction, config) .await?; - - let (_, tags) = enrich_document_metadata(document, &parsed, config); - - Ok(MemoryIngestionResult { - document_id: document_id.to_string(), - namespace, - 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, - }) + extraction.document_id = document_id.to_string(); + extraction.namespace = namespace; + Ok(extraction) } /// Clear existing relations for the document then upsert all extracted @@ -113,13 +78,13 @@ impl UnifiedMemory { &self, namespace: &str, document_id: &str, - parsed: &ParsedIngestion, + extraction: &MemoryIngestionResult, config: &MemoryIngestionConfig, ) -> Result<(), String> { self.graph_remove_document_namespace(namespace, document_id) .await?; - for relation in &parsed.relations { + for relation in &extraction.relations { let chunk_ids = relation .chunk_ids .iter() diff --git a/src/openhuman/memory/ingestion/parse.rs b/src/openhuman/memory/ingestion/parse.rs deleted file mode 100644 index 58ed775ee..000000000 --- a/src/openhuman/memory/ingestion/parse.rs +++ /dev/null @@ -1,942 +0,0 @@ -//! Document parsing helpers: chunking, alias resolution, header/metadata enrichment, -//! and the top-level `parse_document` pipeline. - -use std::collections::{BTreeMap, BTreeSet, HashMap}; - -use serde_json::{json, Map, Value}; - -use super::regex::{ - action_item_regex, classify_entity, email_header_regex, explicit_owner_regex, - explicit_preference_regex, graph_fact_regex, named_email_regex, recipient_regex, - sanitize_entity_name, sanitize_fact_text, spatial_regex, will_review_regex, -}; -use super::types::{ - ExtractedEntity, ExtractedRelation, ExtractionAccumulator, ExtractionMode, ExtractionUnit, - MemoryIngestionConfig, ParsedIngestion, RawEntity, RawRelation, DEFAULT_CHUNK_TOKENS, -}; -use crate::openhuman::memory_store::types::NamespaceDocumentInput; -use crate::openhuman::memory_store::UnifiedMemory; - -// ── Chunking helpers ────────────────────────────────────────────────────────── - -/// Splits a document into individual sentences based on punctuation and line breaks. -pub(super) fn split_sentences(text: &str) -> Vec { - let mut out = Vec::new(); - let mut current = String::new(); - for ch in text.chars() { - current.push(ch); - if matches!(ch, '.' | '!' | '?' | '\n') { - let candidate = sanitize_fact_text(¤t); - if !candidate.is_empty() { - out.push(candidate); - } - current.clear(); - } - } - let tail = sanitize_fact_text(¤t); - if !tail.is_empty() { - out.push(tail); - } - let mut merged: Vec = Vec::new(); - for sentence in out { - if sentence.len() < 5 && !merged.is_empty() { - if let Some(last) = merged.last_mut() { - last.push(' '); - last.push_str(&sentence); - } - } else { - merged.push(sentence); - } - } - if merged.is_empty() && !text.trim().is_empty() { - merged.push(sanitize_fact_text(text)); - } - merged -} - -/// Groups chunks into extraction units based on the configured mode. -pub(super) fn build_units(chunks: &[String], mode: ExtractionMode) -> Vec { - let mut units = Vec::new(); - let mut order_index = 0_i64; - for (chunk_index, chunk) in chunks.iter().enumerate() { - match mode { - ExtractionMode::Chunk => { - let text = sanitize_fact_text(chunk); - if text.is_empty() { - continue; - } - units.push(ExtractionUnit { - text, - chunk_index, - order_index, - }); - order_index += 1; - } - ExtractionMode::Sentence => { - for sentence in split_sentences(chunk) { - if sentence.is_empty() { - continue; - } - units.push(ExtractionUnit { - text: sentence, - chunk_index, - order_index, - }); - order_index += 1; - } - } - } - } - units -} - -/// Searches for the chunk index that most likely contains the given excerpt. -pub(super) fn find_chunk_index(chunks: &[String], excerpt: &str, hint: usize) -> usize { - if chunks.is_empty() { - return 0; - } - let needle = UnifiedMemory::normalize_search_text(excerpt); - if needle.is_empty() { - return hint.min(chunks.len().saturating_sub(1)); - } - for (index, chunk) in chunks.iter().enumerate().skip(hint) { - if UnifiedMemory::normalize_search_text(chunk).contains(&needle) { - return index; - } - } - for (index, chunk) in chunks.iter().enumerate().take(hint.min(chunks.len())) { - if UnifiedMemory::normalize_search_text(chunk).contains(&needle) { - return index; - } - } - hint.min(chunks.len().saturating_sub(1)) -} - -// ── Alias resolution ────────────────────────────────────────────────────────── - -pub(super) fn reverse_aliases(aliases: &HashMap) -> BTreeMap> { - let mut reverse = BTreeMap::new(); - for (alias, canonical) in aliases { - if alias == canonical { - continue; - } - reverse - .entry(canonical.clone()) - .or_insert_with(Vec::new) - .push(alias.clone()); - } - for values in reverse.values_mut() { - values.sort(); - values.dedup(); - } - reverse -} - -pub(super) fn build_alias_map(entities: &HashMap) -> HashMap { - let mut by_type = HashMap::>::new(); - for entity in entities.values() { - by_type - .entry(entity.entity_type.clone()) - .or_default() - .push(entity.name.clone()); - } - - let mut aliases = HashMap::new(); - for names in by_type.values_mut() { - names.sort_by_key(|name| std::cmp::Reverse(name.len())); - for short in names.iter() { - for long in names.iter() { - if short == long || long.len() <= short.len() { - continue; - } - if long.starts_with(&format!("{short} ")) || long.ends_with(&format!(" {short}")) { - aliases.entry(short.clone()).or_insert_with(|| long.clone()); - break; - } - } - } - } - aliases -} - -pub(super) fn resolve_alias(name: &str, aliases: &HashMap) -> String { - let mut current = name.to_string(); - let mut seen = BTreeSet::new(); - while let Some(next) = aliases.get(¤t) { - if !seen.insert(current.clone()) { - break; - } - current = next.clone(); - } - current -} - -// ── Header / metadata helpers ───────────────────────────────────────────────── - -pub(super) fn extract_people_from_header( - value: &str, - accumulator: &mut ExtractionAccumulator, -) -> Vec { - let mut people = Vec::new(); - for captures in named_email_regex().captures_iter(value) { - let name = sanitize_fact_text( - captures - .name("name") - .map(|value| value.as_str()) - .unwrap_or(""), - ); - if name.is_empty() { - continue; - } - let canonical = sanitize_entity_name(&name); - let _ = accumulator.add_entity(&canonical, "PERSON", 0.95); - accumulator.remember_person_aliases(&canonical); - people.push(canonical); - } - people -} - -pub(super) fn detect_primary_subject(text: &str) -> Option { - if text.contains("OpenHuman") { - return Some("OPENHUMAN".to_string()); - } - None -} - -pub(super) fn enrich_document_metadata( - input: &NamespaceDocumentInput, - parsed: &ParsedIngestion, - config: &MemoryIngestionConfig, -) -> (NamespaceDocumentInput, Vec) { - let mut metadata = match input.metadata.clone() { - Value::Object(map) => map, - _ => Map::new(), - }; - for (key, value) in parsed.metadata.as_object().cloned().unwrap_or_default() { - metadata.insert(key, value); - } - metadata.insert( - "ingestion".to_string(), - json!({ - "backend": "openhuman_rust_heuristic", - "model_name": config.model_name, - "extraction_mode": config.extraction_mode.as_str(), - "entity_count": parsed.entities.len(), - "relation_count": parsed.relations.len(), - "preference_count": parsed.preference_count, - "decision_count": parsed.decision_count, - "chunk_count": parsed.chunk_count, - }), - ); - if parsed.preference_count > 0 || parsed.decision_count > 0 { - metadata.insert("kind".to_string(), json!("profile")); - } - - let mut tags = input.tags.iter().cloned().collect::>(); - tags.extend(parsed.tags.iter().cloned()); - let tags = tags.into_iter().collect::>(); - - ( - NamespaceDocumentInput { - namespace: input.namespace.clone(), - key: input.key.clone(), - title: input.title.clone(), - content: input.content.clone(), - source_type: input.source_type.clone(), - priority: input.priority.clone(), - tags: tags.clone(), - metadata: Value::Object(metadata), - category: input.category.clone(), - session_id: input.session_id.clone(), - document_id: input.document_id.clone(), - // Carry the caller's provenance forward — parsing is a pure - // metadata-enrichment step, so ingest taint must survive it. - taint: input.taint, - }, - tags, - ) -} - -// ── Top-level document parser ───────────────────────────────────────────────── - -pub(super) async fn parse_document( - content: &str, - title: &str, - config: &MemoryIngestionConfig, -) -> ParsedIngestion { - let chunks = UnifiedMemory::chunk_document_content(content, DEFAULT_CHUNK_TOKENS); - log::info!( - "[memory:ingestion] parse_document title={title:?} model={} \ - content_len={} chunk_count={} — heuristic extraction active", - config.model_name, - content.len(), - chunks.len(), - ); - let mut accumulator = ExtractionAccumulator { - document_title: Some(sanitize_entity_name(title)), - primary_subject: detect_primary_subject(title), - ..ExtractionAccumulator::default() - }; - - let mut chunk_hint = 0_usize; - for raw_line in content.lines() { - let line = sanitize_fact_text(raw_line); - if line.is_empty() { - continue; - } - - let chunk_index = find_chunk_index(&chunks, &line, chunk_hint); - chunk_hint = chunk_index; - let order_index = i64::try_from(chunk_index).unwrap_or(i64::MAX); - - if raw_line.trim_start().starts_with('#') { - let heading = sanitize_entity_name(raw_line.trim_start_matches('#')); - if !heading.is_empty() { - if accumulator.document_title.is_none() { - accumulator.document_title = Some(heading.clone()); - } - accumulator.current_subject = Some(heading); - } - continue; - } - - if let Some(captures) = email_header_regex().captures(&line) { - let header_name = captures - .get(1) - .map(|value| value.as_str()) - .unwrap_or_default() - .to_ascii_uppercase(); - let value = captures - .name("value") - .map(|value| value.as_str()) - .unwrap_or(""); - let people = extract_people_from_header(value, &mut accumulator); - if header_name == "FROM" { - accumulator.current_sender = people.first().cloned(); - } else if header_name == "TO" || header_name == "CC" { - if let Some(sender) = accumulator.current_sender.clone() { - for recipient in &people { - accumulator.add_relation( - &sender, - "PERSON", - "communicates_with", - recipient, - "PERSON", - 0.82, - chunk_index, - order_index, - Map::new(), - ); - } - } - } - continue; - } - - if let Some(subject) = line.strip_prefix("Subject:") { - let subject_text = sanitize_fact_text(subject); - if let Some(primary_subject) = detect_primary_subject(&subject_text) { - accumulator.primary_subject = Some(primary_subject); - } - continue; - } - - if let Some(date_text) = line.strip_prefix("Date:") { - let date_text = sanitize_fact_text(date_text); - if let Some(sender) = accumulator.current_sender.clone() { - accumulator.add_relation( - &sender, - "PERSON", - "has_deadline", - &date_text, - "DATE", - 0.75, - chunk_index, - order_index, - Map::new(), - ); - } - continue; - } - - if let Some(value) = line.strip_prefix("Project name:") { - let project = sanitize_entity_name(value); - if !project.is_empty() { - accumulator.primary_subject = Some(project.clone()); - let _ = accumulator.add_entity(&project, "PROJECT", 0.96); - } - continue; - } - - if let Some(value) = line.strip_prefix("Subproject:") { - let subproject = sanitize_entity_name(value); - if !subproject.is_empty() { - let _ = accumulator.add_entity(&subproject, "PROJECT", 0.92); - } - continue; - } - - if let Some(value) = line.strip_prefix("Owner:") { - let owner = sanitize_entity_name(value); - let owned = accumulator - .current_subject - .clone() - .or_else(|| accumulator.primary_subject.clone()) - .or_else(|| accumulator.document_title.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - accumulator.add_relation( - &owner, - "PERSON", - "owns", - &owned, - "WORK_ITEM", - 0.94, - chunk_index, - order_index, - Map::new(), - ); - continue; - } - - if let Some(value) = line.strip_prefix("Name:") { - let name = sanitize_entity_name(value); - if !name.is_empty() { - accumulator.current_subject = Some(name.clone()); - let _ = accumulator.add_entity(&name, "WORK_ITEM", 0.93); - } - continue; - } - - if let Some(value) = line.strip_prefix("Due date:") { - let due_date = sanitize_fact_text(value); - let subject = accumulator - .current_subject - .clone() - .or_else(|| accumulator.primary_subject.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - accumulator.add_relation( - &subject, - "WORK_ITEM", - "has_deadline", - &due_date, - "DATE", - 0.92, - chunk_index, - order_index, - Map::new(), - ); - accumulator.tags.insert("deadline".to_string()); - continue; - } - - if let Some(value) = line.strip_prefix("Target milestone:") { - let due_date = sanitize_fact_text(value); - let subject = accumulator - .primary_subject - .clone() - .or_else(|| accumulator.document_title.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - accumulator.add_relation( - &subject, - "PROJECT", - "has_deadline", - &due_date, - "DATE", - 0.92, - chunk_index, - order_index, - Map::new(), - ); - accumulator.tags.insert("deadline".to_string()); - continue; - } - - if let Some(value) = line.strip_prefix("Preferred embedding model for local experiments:") { - let model = sanitize_fact_text(value); - let subject = accumulator - .primary_subject - .clone() - .or_else(|| accumulator.document_title.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - accumulator.add_relation( - &subject, - "PROJECT", - "uses", - &model, - "TOOL", - 0.88, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .decisions - .insert(format!("{subject} uses {model}")); - accumulator.tags.insert("decision".to_string()); - continue; - } - - if let Some(value) = line.strip_prefix("Preferred extraction mode to try first:") { - let mode = sanitize_fact_text(value); - let subject = accumulator - .primary_subject - .clone() - .or_else(|| accumulator.document_title.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - accumulator.add_relation( - &subject, - "PROJECT", - "uses", - &mode, - "MODE", - 0.88, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .decisions - .insert(format!("{subject} uses {mode}")); - accumulator.tags.insert("decision".to_string()); - continue; - } - - if let Some(captures) = graph_fact_regex().captures(&line) { - let subject = captures - .name("subject") - .map(|value| value.as_str()) - .unwrap_or(""); - let predicate = captures - .name("predicate") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - let subject_type = classify_entity(subject, &accumulator.known_people); - let object_type = classify_entity(object, &accumulator.known_people); - accumulator.add_relation( - subject, - subject_type, - predicate, - object, - object_type, - 0.87, - chunk_index, - order_index, - Map::new(), - ); - if UnifiedMemory::normalize_graph_predicate(predicate) == "PREFERS" { - accumulator.preferences.insert(format!( - "{} prefers {}", - sanitize_entity_name(subject), - sanitize_fact_text(object) - )); - accumulator.tags.insert("preference".to_string()); - accumulator.doc_kind = Some("profile".to_string()); - } - continue; - } - - if let Some(captures) = explicit_owner_regex().captures(&line) { - let subject = captures - .name("subject") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - accumulator.add_relation( - subject, - "PERSON", - "owns", - object, - classify_entity(object, &accumulator.known_people), - 0.94, - chunk_index, - order_index, - Map::new(), - ); - accumulator.tags.insert("owner".to_string()); - continue; - } - - if let Some(captures) = will_review_regex().captures(&line) { - let subject = captures - .name("subject") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - accumulator.add_relation( - subject, - "PERSON", - "reviews", - object, - classify_entity(object, &accumulator.known_people), - 0.80, - chunk_index, - order_index, - Map::new(), - ); - accumulator.tags.insert("owner".to_string()); - continue; - } - - if let Some(captures) = explicit_preference_regex().captures(&line) { - let subject = captures - .name("subject") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - accumulator.add_relation( - subject, - "PERSON", - "prefers", - object, - classify_entity(object, &accumulator.known_people), - 0.90, - chunk_index, - order_index, - Map::new(), - ); - accumulator.preferences.insert(format!( - "{} prefers {}", - sanitize_entity_name(subject), - sanitize_fact_text(object) - )); - accumulator.tags.insert("preference".to_string()); - accumulator.doc_kind = Some("profile".to_string()); - continue; - } - - if let Some(value) = line.strip_prefix("I prefer ") { - if let Some(subject) = accumulator.current_sender.clone() { - let preference = sanitize_fact_text(value); - accumulator.add_relation( - &subject, - "PERSON", - "prefers", - &preference, - classify_entity(&preference, &accumulator.known_people), - 0.92, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .preferences - .insert(format!("{subject} prefers {preference}")); - accumulator.tags.insert("preference".to_string()); - accumulator.doc_kind = Some("profile".to_string()); - continue; - } - } - - if let Some(captures) = action_item_regex().captures(&line) { - let subject = captures - .name("subject") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - if accumulator - .known_people - .contains_key(&sanitize_entity_name(subject)) - || classify_entity(subject, &accumulator.known_people) == "PERSON" - { - accumulator.add_relation( - subject, - "PERSON", - "owns", - object, - classify_entity(object, &accumulator.known_people), - 0.83, - chunk_index, - order_index, - Map::new(), - ); - accumulator.tags.insert("owner".to_string()); - continue; - } - } - - let upper = sanitize_entity_name(&line); - let decision_subject = accumulator - .primary_subject - .clone() - .or_else(|| accumulator.document_title.clone()) - .unwrap_or_else(|| "DOCUMENT".to_string()); - if upper.contains("JSON-RPC") { - accumulator.add_relation( - &decision_subject, - "PROJECT", - "uses", - "JSON-RPC", - "PRODUCT", - 0.86, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .decisions - .insert(format!("{decision_subject} uses JSON-RPC")); - accumulator.tags.insert("decision".to_string()); - continue; - } - if upper.contains("SHOULD USE NAMESPACE") - || upper.contains("USE NAMESPACE AS THE STORAGE") - || upper.contains("NAMESPACE AS THE MAIN SCOPE KEY") - { - accumulator.add_relation( - &decision_subject, - "PROJECT", - "uses", - "namespace", - "TOPIC", - 0.84, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .decisions - .insert(format!("{decision_subject} uses namespace")); - accumulator.tags.insert("decision".to_string()); - continue; - } - if upper.contains("USER_ID") && (upper.contains("DO NOT NEED") || upper.contains("AVOID")) { - accumulator.add_relation( - &decision_subject, - "PROJECT", - "avoids", - "user_id", - "TOPIC", - 0.82, - chunk_index, - order_index, - Map::new(), - ); - accumulator - .decisions - .insert(format!("{decision_subject} avoids user_id")); - accumulator.tags.insert("decision".to_string()); - } - } - - for unit in build_units(&chunks, config.extraction_mode) { - if let Some(captures) = recipient_regex().captures(&unit.text) { - let giver = captures - .name("giver") - .map(|value| value.as_str()) - .unwrap_or(""); - let object = captures - .name("object") - .map(|value| value.as_str()) - .unwrap_or(""); - let recipient = captures - .name("recipient") - .map(|value| value.as_str()) - .unwrap_or(""); - accumulator.add_relation( - giver, - "PERSON", - "uses", - object, - classify_entity(object, &accumulator.known_people), - config.adjacency_threshold.max(0.62), - unit.chunk_index, - unit.order_index, - Map::new(), - ); - accumulator.add_relation( - recipient, - "PERSON", - "uses", - object, - classify_entity(object, &accumulator.known_people), - (config.adjacency_threshold * 0.9).max(0.55), - unit.chunk_index, - unit.order_index, - Map::new(), - ); - } - - if let Some(captures) = spatial_regex().captures(&unit.text) { - let head = captures - .name("head") - .map(|value| value.as_str()) - .unwrap_or(""); - let direction = captures - .name("direction") - .map(|value| value.as_str()) - .unwrap_or(""); - let tail = captures - .name("tail") - .map(|value| value.as_str()) - .unwrap_or(""); - let inverse = match direction.to_ascii_lowercase().as_str() { - "north" => "south_of", - "south" => "north_of", - "east" => "west_of", - "west" => "east_of", - _ => "", - }; - let predicate = format!("{direction}_of"); - accumulator.add_relation( - head, - "ROOM", - &predicate, - tail, - "ROOM", - config.adjacency_threshold.max(0.70), - unit.chunk_index, - unit.order_index, - Map::new(), - ); - if !inverse.is_empty() { - accumulator.add_relation( - tail, - "ROOM", - inverse, - head, - "ROOM", - config.adjacency_threshold.max(0.70), - unit.chunk_index, - unit.order_index, - Map::new(), - ); - } - } - } - - let aliases = build_alias_map(&accumulator.entities); - let reverse_alias = reverse_aliases(&aliases); - let mut canonical_entities = BTreeMap::::new(); - for entity in accumulator.entities.values() { - let canonical = resolve_alias(&entity.name, &aliases); - let entry = canonical_entities - .entry(canonical.clone()) - .or_insert_with(|| RawEntity { - name: canonical.clone(), - entity_type: entity.entity_type.clone(), - confidence: entity.confidence, - }); - if entity.confidence > entry.confidence { - entry.confidence = entity.confidence; - entry.entity_type = entity.entity_type.clone(); - } - } - - let mut aggregated_relations = BTreeMap::<(String, String, String), _>::new(); - for relation in accumulator.relations { - let subject = resolve_alias(&relation.subject, &aliases); - let object = resolve_alias(&relation.object, &aliases); - if subject == object { - continue; - } - let key = (subject.clone(), relation.predicate.clone(), object.clone()); - let entry = aggregated_relations - .entry(key) - .or_insert_with(|| RawRelation { - subject, - subject_type: relation.subject_type.clone(), - predicate: relation.predicate.clone(), - object, - object_type: relation.object_type.clone(), - confidence: relation.confidence, - chunk_indexes: relation.chunk_indexes.clone(), - order_index: relation.order_index, - metadata: relation.metadata.clone(), - }); - entry.confidence = entry.confidence.max(relation.confidence); - entry.order_index = entry.order_index.min(relation.order_index); - entry.chunk_indexes.extend(relation.chunk_indexes); - } - - let entities = canonical_entities - .into_values() - .filter(|entity| entity.confidence >= config.entity_threshold) - .map(|entity| ExtractedEntity { - name: entity.name.clone(), - entity_type: entity.entity_type, - aliases: reverse_alias.get(&entity.name).cloned().unwrap_or_default(), - }) - .collect::>(); - - let relations = aggregated_relations - .into_values() - .filter(|relation| relation.confidence >= config.relation_threshold) - .map(|relation| ExtractedRelation { - subject: relation.subject, - subject_type: relation.subject_type, - predicate: relation.predicate, - object: relation.object, - object_type: relation.object_type, - confidence: relation.confidence, - evidence_count: u32::try_from(relation.chunk_indexes.len()).unwrap_or(u32::MAX), - chunk_ids: relation - .chunk_indexes - .iter() - .map(|index| format!("chunk:{index}")) - .collect::>(), - order_index: Some(relation.order_index), - metadata: Value::Object(relation.metadata), - }) - .collect::>(); - - let mut tags = accumulator.tags.into_iter().collect::>(); - tags.sort(); - let metadata = json!({ - "kind": accumulator.doc_kind.or_else(|| { - if !accumulator.preferences.is_empty() || !accumulator.decisions.is_empty() { - Some("profile".to_string()) - } else { - None - } - }), - "primary_subject": accumulator.primary_subject, - "decisions": accumulator.decisions.iter().cloned().collect::>(), - "preferences": accumulator.preferences.iter().cloned().collect::>(), - "extracted_entities": entities.iter().map(|entity| { - json!({ - "name": entity.name, - "entity_type": entity.entity_type, - "aliases": entity.aliases, - }) - }).collect::>(), - }); - - log::debug!( - "[memory:ingestion] parse_document complete title={title:?} \ - entities={} relations={} preferences={} decisions={}", - entities.len(), - relations.len(), - accumulator.preferences.len(), - accumulator.decisions.len(), - ); - - ParsedIngestion { - tags, - metadata, - entities, - relations, - chunk_count: chunks.len(), - preference_count: accumulator.preferences.len(), - decision_count: accumulator.decisions.len(), - } -} - -#[cfg(test)] -#[path = "parse_tests.rs"] -mod tests; diff --git a/src/openhuman/memory/ingestion/parse_tests.rs b/src/openhuman/memory/ingestion/parse_tests.rs deleted file mode 100644 index de4d4f8e1..000000000 --- a/src/openhuman/memory/ingestion/parse_tests.rs +++ /dev/null @@ -1,137 +0,0 @@ -use super::*; -use crate::openhuman::memory_store::types::NamespaceDocumentInput; -use serde_json::json; - -fn sample_input() -> NamespaceDocumentInput { - NamespaceDocumentInput { - namespace: "global".into(), - key: "doc-1".into(), - title: "OpenHuman roadmap".into(), - content: "Alice owns roadmap".into(), - source_type: "manual".into(), - priority: "normal".into(), - tags: vec!["existing".into()], - metadata: json!({"seed": true}), - category: "core".into(), - session_id: Some("session-1".into()), - document_id: Some("doc-id-1".into()), - taint: crate::openhuman::memory::MemoryTaint::Internal, - } -} - -#[test] -fn split_sentences_breaks_on_punctuation_and_merges_tiny_fragments() { - let parts = split_sentences("Hello world. Ok.\nNext line?"); - assert_eq!(parts.len(), 2); - assert_eq!(parts[0], "Hello world Ok"); - assert_eq!(parts[1], "Next line?"); -} - -#[test] -fn build_units_respects_extraction_mode() { - let chunks = vec!["One. Two.".to_string(), "Three".to_string()]; - let sentence_units = build_units(&chunks, ExtractionMode::Sentence); - let chunk_units = build_units(&chunks, ExtractionMode::Chunk); - - assert_eq!(sentence_units.len(), 2); - assert_eq!(sentence_units[0].chunk_index, 0); - assert_eq!(sentence_units[0].text, "One Two"); - assert_eq!(sentence_units[1].chunk_index, 1); - assert_eq!(sentence_units[1].text, "Three"); - - assert_eq!(chunk_units.len(), 2); - assert_eq!(chunk_units[0].text, "One. Two"); - assert_eq!(chunk_units[1].text, "Three"); -} - -#[test] -fn find_chunk_index_prefers_hint_then_wraps() { - let chunks = vec![ - "alpha content".to_string(), - "beta needle".to_string(), - "gamma trailing".to_string(), - ]; - assert_eq!(find_chunk_index(&chunks, "needle", 1), 1); - assert_eq!(find_chunk_index(&chunks, "alpha", 2), 0); - assert_eq!(find_chunk_index(&chunks, "missing", 2), 2); -} - -#[test] -fn alias_map_builds_reverse_lookup() { - let mut entities = HashMap::new(); - entities.insert( - "ALICE".into(), - RawEntity { - name: "ALICE".into(), - entity_type: "PERSON".into(), - confidence: 0.8, - }, - ); - entities.insert( - "ALICE SMITH".into(), - RawEntity { - name: "ALICE SMITH".into(), - entity_type: "PERSON".into(), - confidence: 0.9, - }, - ); - let aliases = build_alias_map(&entities); - assert_eq!( - aliases.get("ALICE").map(String::as_str), - Some("ALICE SMITH") - ); - assert_eq!(resolve_alias("ALICE", &aliases), "ALICE SMITH"); - - let reverse = reverse_aliases(&aliases); - assert_eq!(reverse.get("ALICE SMITH"), Some(&vec!["ALICE".to_string()])); -} - -#[test] -fn enrich_document_metadata_merges_tags_and_ingestion_details() { - let input = sample_input(); - let parsed = ParsedIngestion { - tags: vec!["decision".into(), "existing".into()], - metadata: json!({"kind": "profile", "extra": 1}), - entities: vec![], - relations: vec![], - chunk_count: 3, - preference_count: 1, - decision_count: 2, - }; - let config = MemoryIngestionConfig::default(); - let (enriched, tags) = enrich_document_metadata(&input, &parsed, &config); - - assert_eq!(tags, vec!["decision".to_string(), "existing".to_string()]); - assert_eq!(enriched.tags, tags); - assert_eq!(enriched.metadata["seed"], json!(true)); - assert_eq!(enriched.metadata["extra"], json!(1)); - assert_eq!( - enriched.metadata["ingestion"]["model_name"], - config.model_name - ); - assert_eq!(enriched.metadata["ingestion"]["chunk_count"], json!(3)); -} - -#[test] -fn extract_people_from_header_collects_named_people() { - let mut acc = ExtractionAccumulator::default(); - let people = extract_people_from_header( - "Alice Smith , Bob Jones ", - &mut acc, - ); - assert_eq!( - people, - vec!["ALICE SMITH".to_string(), "BOB JONES".to_string()] - ); - assert!(acc.entities.contains_key("ALICE SMITH")); - assert!(acc.entities.contains_key("BOB JONES")); -} - -#[test] -fn detect_primary_subject_only_matches_openhuman() { - assert_eq!( - detect_primary_subject("OpenHuman desktop roadmap"), - Some("OPENHUMAN".to_string()) - ); - assert_eq!(detect_primary_subject("General roadmap"), None); -} diff --git a/src/openhuman/memory/ingestion/regex.rs b/src/openhuman/memory/ingestion/regex.rs deleted file mode 100644 index ee0dda769..000000000 --- a/src/openhuman/memory/ingestion/regex.rs +++ /dev/null @@ -1,219 +0,0 @@ -//! Lazily-initialised regex patterns and text-sanitization helpers for document ingestion. - -use std::collections::HashMap; -use std::sync::OnceLock; - -use regex::Regex; - -use crate::openhuman::memory_store::UnifiedMemory; - -/// Regex for identifying standard email headers (From, To, Cc). -pub(super) fn email_header_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX - .get_or_init(|| Regex::new(r"^(From|To|Cc):\s*(?P.+)$").expect("email header regex")) -} - -/// Regex for identifying named email addresses (e.g., "John Doe "). -pub(super) fn named_email_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"(?P[^,<]+?)\s*<(?P[^>]+)>").expect("named email regex") - }) -} - -/// Regex for identifying explicit graph facts (e.g., "Alice works_on Project-X"). -pub(super) fn graph_fact_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new( - r"^(?P[A-Za-z0-9][A-Za-z0-9 ._\-/]+?)\s+(?Pworks_on|depends_on|uses|evaluates|owns|prefers)\s+(?P.+)$", - ) - .expect("graph fact regex") - }) -} - -/// Regex for identifying ownership patterns (e.g., "Bob owns the repository"). -pub(super) fn explicit_owner_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) owns (?P.+)$") - .expect("explicit owner regex") - }) -} - -/// Regex for identifying preference patterns (e.g., "Carol prefers light mode"). -pub(super) fn explicit_preference_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) prefers (?P.+)$") - .expect("explicit preference regex") - }) -} - -/// Regex for identifying action items or assignments (e.g., "Dave: finish the API"). -pub(super) fn action_item_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?):\s*(?P.+)$") - .expect("action item regex") - }) -} - -/// Regex for identifying review assignments. -pub(super) fn will_review_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"^(?P[A-Za-z][A-Za-z ._-]+?) will review (?P.+)$") - .expect("will review regex") - }) -} - -/// Regex for identifying complex giving/receiving interactions. -pub(super) fn recipient_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new( - r"(?i)(?P[A-Z][A-Za-z]+(?: [A-Z][A-Za-z]+)?)\s+(gave|sent|handed|passed)\s+(?P.+?)\s+to\s+(?P[A-Z][A-Za-z]+(?: [A-Z][A-Za-z]+)?)", - ) - .expect("recipient regex") - }) -} - -/// Regex for identifying spatial relationships (e.g., "Kitchen is north of the Garden"). -pub(super) fn spatial_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new( - r"(?i)(?P[A-Za-z][A-Za-z0-9 _-]+?)\s+is\s+(?Pnorth|south|east|west)\s+of\s+(?P[A-Za-z][A-Za-z0-9 _-]+)", - ) - .expect("spatial regex") - }) -} - -/// Regex for identifying dates in "Month DD, YYYY" format. -pub(super) fn month_date_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| { - Regex::new(r"(?i)\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec)[a-z]*\s+\d{1,2},\s+\d{4}\b") - .expect("month date regex") - }) -} - -/// Regex for identifying ISO-8601 dates (YYYY-MM-DD). -pub(super) fn iso_date_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| Regex::new(r"\b\d{4}-\d{2}-\d{2}\b").expect("iso date regex")) -} - -/// Regex for identifying potential person names (Title Case). -pub(super) fn person_name_regex() -> &'static Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX - .get_or_init(|| Regex::new(r"\b[A-Z][a-z]+(?: [A-Z][a-z]+)+\b").expect("person name regex")) -} - -/// Normalizes an entity name by trimming punctuation, collapsing whitespace, and converting to uppercase. -pub(super) fn sanitize_entity_name(name: &str) -> String { - let trimmed = name.trim().trim_matches(|ch: char| { - matches!(ch, '-' | ':' | ';' | ',' | '.' | '"' | '\'' | '(' | ')') - }); - if trimmed.is_empty() { - return String::new(); - } - UnifiedMemory::collapse_whitespace(trimmed).to_uppercase() -} - -/// Normalizes text content by trimming and collapsing whitespace. -pub(super) fn sanitize_fact_text(text: &str) -> String { - let trimmed = text - .trim() - .trim_start_matches('-') - .trim() - .trim_matches(|ch: char| matches!(ch, ':' | ';' | ',' | '.')); - UnifiedMemory::collapse_whitespace(trimmed) -} - -/// Heuristically classifies an entity based on its name and known person map. -pub(super) fn classify_entity(name: &str, known_people: &HashMap) -> &'static str { - let upper = sanitize_entity_name(name); - if upper.is_empty() { - return "TOPIC"; - } - if month_date_regex().is_match(name) || iso_date_regex().is_match(name) { - return "DATE"; - } - if upper.contains('@') { - return "ORGANIZATION"; - } - if known_people.contains_key(&upper) || person_name_regex().is_match(name) { - return "PERSON"; - } - if matches!( - upper.as_str(), - "OPENHUMAN" | "JSON-RPC" | "JSON-RPC 2.0" | "NEOCORTEX_V2" | "NEOCORTEX V2" - ) { - return "PRODUCT"; - } - if upper.contains("MODEL") { - return "TOOL"; - } - if upper.contains("MODE") { - return "MODE"; - } - if upper.contains("MILESTONE") - || upper.contains("ROADMAP") - || upper.contains("CONTRACT") - || upper.contains("API") - || upper.contains("MEMORY") - || upper.contains("FIXTURE") - || upper.contains("THREAD") - || upper.contains("WORK") - { - return "WORK_ITEM"; - } - if upper.contains("OFFICE") - || upper.contains("ROOM") - || upper.contains("GARDEN") - || upper.contains("KITCHEN") - { - return "ROOM"; - } - if upper.contains("TINYHUMANS") || upper.ends_with("CORE") { - return "ORGANIZATION"; - } - if (upper.contains('-') || upper.contains('_')) && !upper.contains(' ') { - return "PROJECT"; - } - "TOPIC" -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn sanitize_entity_name_trims_punctuation_and_uppercases() { - assert_eq!(sanitize_entity_name(" Alice Smith. "), "ALICE SMITH"); - assert_eq!(sanitize_entity_name("\"openhuman\""), "OPENHUMAN"); - assert_eq!(sanitize_entity_name(""), ""); - } - - #[test] - fn sanitize_fact_text_collapses_whitespace_and_strips_edges() { - assert_eq!(sanitize_fact_text(" - Hello world. "), "Hello world"); - assert_eq!(sanitize_fact_text(":: spaced\ttext ;;"), "spaced text"); - } - - #[test] - fn classify_entity_detects_dates_people_and_products() { - let mut known_people = HashMap::new(); - known_people.insert("ALICE SMITH".to_string(), "ALICE SMITH".to_string()); - - assert_eq!(classify_entity("Jan 5, 2026", &known_people), "DATE"); - assert_eq!(classify_entity("Alice Smith", &known_people), "PERSON"); - assert_eq!(classify_entity("OpenHuman", &known_people), "PRODUCT"); - assert_eq!(classify_entity("Kitchen", &known_people), "ROOM"); - assert_eq!(classify_entity("phoenix-project", &known_people), "PROJECT"); - } -} diff --git a/src/openhuman/memory/ingestion/rules.rs b/src/openhuman/memory/ingestion/rules.rs deleted file mode 100644 index 3403779e9..000000000 --- a/src/openhuman/memory/ingestion/rules.rs +++ /dev/null @@ -1,325 +0,0 @@ -//! Semantic validation rules for knowledge-graph relations and `ExtractionAccumulator` impl. - -use std::collections::BTreeSet; - -use serde_json::{Map, Value}; - -use super::regex::sanitize_entity_name; -use super::types::{ExtractionAccumulator, RawEntity, RawRelation}; -use crate::openhuman::memory_store::UnifiedMemory; - -/// A validation rule for semantic relationships. -#[derive(Debug)] -pub(super) struct RelationRule { - /// Canonical predicate name (uppercase snake_case). - pub(super) canonical: &'static str, - /// Allowed classifications for the subject. - pub(super) allowed_head: &'static [&'static str], - /// Allowed classifications for the object. - pub(super) allowed_tail: &'static [&'static str], -} - -const PERSON_TYPES: &[&str] = &["PERSON"]; -const ORG_TYPES: &[&str] = &[ - "ORGANIZATION", - "PROJECT", - "PRODUCT", - "TOOL", - "TOPIC", - "WORK_ITEM", -]; -const PLACE_TYPES: &[&str] = &["PLACE", "LOCATION", "ROOM"]; -const DATE_TYPES: &[&str] = &["DATE"]; - -/// Returns the semantic validation rule for a given predicate name. -pub(super) fn relation_rule(predicate: &str) -> Option { - let normalized = UnifiedMemory::normalize_graph_predicate(predicate); - let rule = match normalized.as_str() { - "OWNS" | "WORKS_ON" | "RESPONSIBLE_FOR" | "REVIEWS" => RelationRule { - canonical: "OWNS", - allowed_head: PERSON_TYPES, - allowed_tail: ORG_TYPES, - }, - "USES" | "KEEPS" | "ADOPTS" => RelationRule { - canonical: "USES", - allowed_head: ORG_TYPES, - allowed_tail: ORG_TYPES, - }, - "WORKS_FOR" => RelationRule { - canonical: "WORKS_FOR", - allowed_head: PERSON_TYPES, - allowed_tail: &["ORGANIZATION", "PROJECT", "PRODUCT"], - }, - "DEPENDS_ON" => RelationRule { - canonical: "DEPENDS_ON", - allowed_head: ORG_TYPES, - allowed_tail: ORG_TYPES, - }, - "PREFERS" => RelationRule { - canonical: "PREFERS", - allowed_head: PERSON_TYPES, - allowed_tail: &["TOPIC", "WORK_ITEM", "MODE", "PRODUCT", "TOOL"], - }, - "HAS_DEADLINE" | "DUE_ON" => RelationRule { - canonical: "HAS_DEADLINE", - allowed_head: ORG_TYPES, - allowed_tail: DATE_TYPES, - }, - "COMMUNICATES_WITH" => RelationRule { - canonical: "COMMUNICATES_WITH", - allowed_head: PERSON_TYPES, - allowed_tail: PERSON_TYPES, - }, - "INVESTIGATES" | "EVALUATES" => RelationRule { - canonical: "INVESTIGATES", - allowed_head: PERSON_TYPES, - allowed_tail: ORG_TYPES, - }, - "NORTH_OF" => RelationRule { - canonical: "NORTH_OF", - allowed_head: PLACE_TYPES, - allowed_tail: PLACE_TYPES, - }, - "SOUTH_OF" => RelationRule { - canonical: "SOUTH_OF", - allowed_head: PLACE_TYPES, - allowed_tail: PLACE_TYPES, - }, - "EAST_OF" => RelationRule { - canonical: "EAST_OF", - allowed_head: PLACE_TYPES, - allowed_tail: PLACE_TYPES, - }, - "WEST_OF" => RelationRule { - canonical: "WEST_OF", - allowed_head: PLACE_TYPES, - allowed_tail: PLACE_TYPES, - }, - "AVOIDS" => RelationRule { - canonical: "AVOIDS", - allowed_head: ORG_TYPES, - allowed_tail: ORG_TYPES, - }, - _ => return None, - }; - Some(rule) -} - -/// Helper to check if a classification is allowed by a rule. -pub(super) fn type_allowed(actual: &str, allowed: &[&str]) -> bool { - allowed.is_empty() || allowed.iter().any(|candidate| candidate == &actual) -} - -/// Resolves a person's name using the known alias map. -pub(super) fn resolve_person_alias( - name: &str, - known_people: &std::collections::HashMap, -) -> String { - let upper = name.to_uppercase(); - known_people.get(&upper).cloned().unwrap_or(upper) -} - -impl ExtractionAccumulator { - /// Ingests a full name and its components (e.g., first name) into the alias map. - pub(super) fn remember_person_aliases(&mut self, canonical_name: &str) { - let parts = canonical_name.split_whitespace().collect::>(); - if let Some(first_name) = parts.first() { - self.known_people - .entry(first_name.to_uppercase()) - .or_insert_with(|| canonical_name.to_string()); - } - } - - /// Records a new entity, updating confidence if already known. - pub(super) fn add_entity( - &mut self, - name: &str, - entity_type: &str, - confidence: f32, - ) -> Option { - let cleaned = sanitize_entity_name(name); - if cleaned.is_empty() { - return None; - } - let resolved_name = if entity_type == "PERSON" { - resolve_person_alias(&cleaned, &self.known_people) - } else { - cleaned.clone() - }; - let entry = self - .entities - .entry(resolved_name.clone()) - .or_insert_with(|| RawEntity { - name: resolved_name.clone(), - entity_type: entity_type.to_string(), - confidence, - }); - if confidence > entry.confidence { - entry.confidence = confidence; - } - if entity_type == "PERSON" { - self.remember_person_aliases(&resolved_name); - } - Some(resolved_name) - } - - /// Records a new relationship, applying semantic validation rules. - #[allow(clippy::too_many_arguments)] - pub(super) fn add_relation( - &mut self, - subject: &str, - subject_type: &str, - predicate: &str, - object: &str, - object_type: &str, - confidence: f32, - chunk_index: usize, - order_index: i64, - metadata: Map, - ) { - let Some(rule) = relation_rule(predicate) else { - return; - }; - let Some(subject_name) = self.add_entity(subject, subject_type, confidence) else { - return; - }; - let Some(object_name) = self.add_entity(object, object_type, confidence) else { - return; - }; - if subject_name == object_name { - return; - } - let actual_subject_type = self - .entities - .get(&subject_name) - .map(|value| value.entity_type.as_str()) - .unwrap_or(subject_type); - let actual_object_type = self - .entities - .get(&object_name) - .map(|value| value.entity_type.as_str()) - .unwrap_or(object_type); - if !type_allowed(actual_subject_type, rule.allowed_head) - || !type_allowed(actual_object_type, rule.allowed_tail) - { - return; - } - - let mut chunk_indexes = BTreeSet::new(); - chunk_indexes.insert(chunk_index); - self.relations.push(RawRelation { - subject: subject_name, - subject_type: actual_subject_type.to_string(), - predicate: rule.canonical.to_string(), - object: object_name, - object_type: actual_object_type.to_string(), - confidence, - chunk_indexes, - order_index, - metadata, - }); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn relation_rule_normalizes_supported_predicates() { - let owns = relation_rule("works_on").unwrap(); - assert_eq!(owns.canonical, "OWNS"); - assert_eq!(owns.allowed_head, PERSON_TYPES); - - let prefers = relation_rule("prefers").unwrap(); - assert_eq!(prefers.canonical, "PREFERS"); - - let deadline = relation_rule("due_on").unwrap(); - assert_eq!(deadline.canonical, "HAS_DEADLINE"); - } - - #[test] - fn type_allowed_honors_allowlist_and_empty_list() { - assert!(type_allowed("PERSON", PERSON_TYPES)); - assert!(!type_allowed("PROJECT", PERSON_TYPES)); - assert!(type_allowed("ANYTHING", &[])); - } - - #[test] - fn resolve_person_alias_uses_known_people_map() { - let mut known = std::collections::HashMap::new(); - known.insert("ALICE".to_string(), "ALICE SMITH".to_string()); - assert_eq!(resolve_person_alias("ALICE", &known), "ALICE SMITH"); - assert_eq!(resolve_person_alias("BOB", &known), "BOB"); - } - - #[test] - fn add_entity_tracks_highest_confidence_and_person_aliases() { - let mut acc = ExtractionAccumulator::default(); - let first = acc.add_entity("Alice Smith", "PERSON", 0.6).unwrap(); - let second = acc.add_entity("Alice Smith", "PERSON", 0.9).unwrap(); - assert_eq!(first, "ALICE SMITH"); - assert_eq!(second, "ALICE SMITH"); - assert_eq!(acc.entities["ALICE SMITH"].confidence, 0.9); - assert_eq!( - acc.known_people.get("ALICE").map(String::as_str), - Some("ALICE SMITH") - ); - } - - #[test] - fn add_relation_rejects_invalid_or_self_relations() { - let mut acc = ExtractionAccumulator::default(); - acc.add_relation( - "Alice", - "PERSON", - "owns", - "Alice", - "PERSON", - 0.8, - 0, - 0, - Map::new(), - ); - assert!(acc.relations.is_empty(), "self relation should be dropped"); - - acc.add_relation( - "Alice", - "PERSON", - "unknown_predicate", - "Project X", - "PROJECT", - 0.8, - 0, - 0, - Map::new(), - ); - assert!( - acc.relations.is_empty(), - "unknown predicate should be ignored" - ); - } - - #[test] - fn add_relation_canonicalizes_predicate_and_collects_chunk_index() { - let mut acc = ExtractionAccumulator::default(); - acc.add_relation( - "Alice", - "PERSON", - "works_on", - "Phoenix", - "PROJECT", - 0.8, - 3, - 11, - Map::new(), - ); - assert_eq!(acc.relations.len(), 1); - let relation = &acc.relations[0]; - assert_eq!(relation.predicate, "OWNS"); - assert_eq!(relation.subject, "ALICE"); - assert_eq!(relation.object, "PHOENIX"); - assert!(relation.chunk_indexes.contains(&3)); - assert_eq!(relation.order_index, 11); - } -} diff --git a/src/openhuman/memory/ingestion/types.rs b/src/openhuman/memory/ingestion/types.rs deleted file mode 100644 index 1f65b4622..000000000 --- a/src/openhuman/memory/ingestion/types.rs +++ /dev/null @@ -1,287 +0,0 @@ -//! Public and private types for the memory ingestion pipeline. - -use std::collections::{BTreeSet, HashMap}; - -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; - -use crate::openhuman::memory_store::types::NamespaceDocumentInput; - -/// Default extraction backend label reported in ingestion metadata. -pub const DEFAULT_MEMORY_EXTRACTION_MODEL: &str = "heuristic-only"; -/// Default number of tokens per text chunk during ingestion. -pub(super) const DEFAULT_CHUNK_TOKENS: usize = 225; - -/// Granularity of extraction for heuristic parsing. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -#[derive(Default)] -pub enum ExtractionMode { - /// Extract from each individual sentence (higher precision). - #[default] - Sentence, - /// Extract from the entire chunk at once (faster, better for context). - Chunk, -} - -impl ExtractionMode { - /// Returns the string representation of the extraction mode. - pub(super) fn as_str(self) -> &'static str { - match self { - Self::Sentence => "sentence", - Self::Chunk => "chunk", - } - } -} - -/// Configuration for the memory ingestion process. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct MemoryIngestionConfig { - /// Extraction backend label recorded in metadata/results. - pub model_name: String, - /// The granularity of heuristic extraction. - #[serde(default)] - pub extraction_mode: ExtractionMode, - /// Minimum confidence threshold for entity extraction (0.0 to 1.0). - #[serde(default = "default_entity_threshold")] - pub entity_threshold: f32, - /// Minimum confidence threshold for relation extraction (0.0 to 1.0). - #[serde(default = "default_relation_threshold")] - pub relation_threshold: f32, - /// Threshold for adjacency-based heuristics. - #[serde(default = "default_adjacency_threshold")] - pub adjacency_threshold: f32, - /// Reserved batch-size knob kept for config compatibility. - #[serde(default = "default_batch_size")] - pub batch_size: usize, -} - -fn default_entity_threshold() -> f32 { - 0.45 -} - -fn default_relation_threshold() -> f32 { - 0.30 -} - -fn default_adjacency_threshold() -> f32 { - 0.50 -} - -fn default_batch_size() -> usize { - 16 -} - -impl Default for MemoryIngestionConfig { - fn default() -> Self { - Self { - model_name: DEFAULT_MEMORY_EXTRACTION_MODEL.to_string(), - extraction_mode: ExtractionMode::Sentence, - entity_threshold: default_entity_threshold(), - relation_threshold: default_relation_threshold(), - adjacency_threshold: default_adjacency_threshold(), - batch_size: default_batch_size(), - } - } -} - -/// A request to ingest a single document. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct MemoryIngestionRequest { - /// The document input to process. - pub document: NamespaceDocumentInput, - /// Ingestion configuration. - #[serde(default)] - pub config: MemoryIngestionConfig, -} - -/// An entity identified during the ingestion process. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ExtractedEntity { - /// Normalized name of the entity (all-caps). - pub name: String, - /// Classification (e.g., PERSON, ORGANIZATION). - pub entity_type: String, - /// Known aliases for this entity. - #[serde(default)] - pub aliases: Vec, -} - -/// A relation identified during the ingestion process. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ExtractedRelation { - /// Name of the subject entity. - pub subject: String, - /// Classification of the subject. - pub subject_type: String, - /// Relationship type (e.g., OWNS, WORKS_ON). - pub predicate: String, - /// Name of the object entity. - pub object: String, - /// Classification of the object. - pub object_type: String, - /// Extraction confidence (0.0 to 1.0). - pub confidence: f32, - /// Number of distinct occurrences of this relation. - pub evidence_count: u32, - /// IDs of the chunks where this relation was found. - pub chunk_ids: Vec, - /// Sequential order index for reconstruction. - pub order_index: Option, - /// Additional metadata about the extraction. - pub metadata: Value, -} - -/// The comprehensive result of an ingestion operation. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct MemoryIngestionResult { - /// ID of the document that was ingested. - pub document_id: String, - /// Namespace containing the document. - pub namespace: String, - /// Extraction backend label recorded for the ingestion run. - pub model_name: String, - /// Mode used for extraction. - pub extraction_mode: String, - /// Total number of chunks processed. - pub chunk_count: usize, - /// Total number of distinct entities found. - pub entity_count: usize, - /// Total number of distinct relations found. - pub relation_count: usize, - /// Number of identified user preferences. - pub preference_count: usize, - /// Number of identified decisions. - pub decision_count: usize, - /// Auto-generated tags for the document. - #[serde(default)] - pub tags: Vec, - /// Complete list of identified entities. - #[serde(default)] - pub entities: Vec, - /// Complete list of identified relations. - #[serde(default)] - pub relations: Vec, -} - -/// Intermediate representation of an entity before normalization and alias resolution. -#[derive(Debug, Clone)] -pub(super) struct RawEntity { - pub(super) name: String, - pub(super) entity_type: String, - pub(super) confidence: f32, -} - -/// Intermediate representation of a relationship before aggregation. -#[derive(Debug, Clone)] -pub(super) struct RawRelation { - pub(super) subject: String, - pub(super) subject_type: String, - pub(super) predicate: String, - pub(super) object: String, - pub(super) object_type: String, - pub(super) confidence: f32, - /// Indices of the chunks where this relation was found. - pub(super) chunk_indexes: BTreeSet, - /// Global sequential index for ordering within the document. - pub(super) order_index: i64, - /// JSON metadata for the relation. - pub(super) metadata: Map, -} - -/// A single unit of text (sentence or chunk) passed to the extractor. -#[derive(Debug, Clone)] -pub(super) struct ExtractionUnit { - pub(super) text: String, - pub(super) chunk_index: usize, - pub(super) order_index: i64, -} - -/// Accumulates extraction results across multiple chunks or units. -/// -/// Handles entity and relation deduplication, alias tracking, and -/// basic document understanding (e.g., identifying the primary subject). -#[derive(Debug, Default)] -pub(super) struct ExtractionAccumulator { - /// Mapping of normalized entity name to its highest-confidence raw extraction. - pub(super) entities: HashMap, - /// Collected relations before final canonicalization. - pub(super) relations: Vec, - /// Tags identified during processing. - pub(super) tags: BTreeSet, - /// Decisions identified during processing. - pub(super) decisions: BTreeSet, - /// User preferences identified during processing. - pub(super) preferences: BTreeSet, - /// Inferred document kind (e.g., "profile"). - pub(super) doc_kind: Option, - /// The document's inferred primary subject. - pub(super) primary_subject: Option, - /// Sanitized document title. - pub(super) document_title: Option, - /// The subject of the current markdown section. - pub(super) current_subject: Option, - /// Current sender if processing a message/thread. - pub(super) current_sender: Option, - /// Mapping of names to their canonicalized full name. - pub(super) known_people: HashMap, -} - -/// The result of the parsing stage of ingestion. -#[derive(Debug)] -pub(super) struct ParsedIngestion { - pub(super) tags: Vec, - pub(super) metadata: Value, - pub(super) entities: Vec, - pub(super) relations: Vec, - pub(super) chunk_count: usize, - pub(super) preference_count: usize, - pub(super) decision_count: usize, -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn extraction_mode_default_is_sentence() { - assert_eq!(ExtractionMode::default(), ExtractionMode::Sentence); - assert_eq!(ExtractionMode::Sentence.as_str(), "sentence"); - assert_eq!(ExtractionMode::Chunk.as_str(), "chunk"); - } - - #[test] - fn memory_ingestion_config_default_matches_expected_thresholds() { - let cfg = MemoryIngestionConfig::default(); - assert_eq!(cfg.model_name, DEFAULT_MEMORY_EXTRACTION_MODEL); - assert_eq!(cfg.extraction_mode, ExtractionMode::Sentence); - assert_eq!(cfg.entity_threshold, 0.45); - assert_eq!(cfg.relation_threshold, 0.30); - assert_eq!(cfg.adjacency_threshold, 0.50); - assert_eq!(cfg.batch_size, 16); - } - - #[test] - fn memory_ingestion_request_defaults_config_when_absent() { - let request: MemoryIngestionRequest = serde_json::from_value(json!({ - "document": { - "namespace": "global", - "key": "doc-1", - "title": "Doc", - "content": "Body", - "source_type": "manual", - "priority": "normal", - "category": "core" - } - })) - .unwrap(); - assert_eq!(request.config.model_name, DEFAULT_MEMORY_EXTRACTION_MODEL); - assert_eq!(request.config.extraction_mode, ExtractionMode::Sentence); - } -} diff --git a/src/openhuman/memory/ops/sync.rs b/src/openhuman/memory/ops/sync.rs index 5a2a03670..eaba042ec 100644 --- a/src/openhuman/memory/ops/sync.rs +++ b/src/openhuman/memory/ops/sync.rs @@ -5,7 +5,7 @@ use crate::openhuman::config::rpc as config_rpc; use crate::openhuman::memory::sync::{emit_sync_stage, MemorySyncStage, MemorySyncTrigger}; -use crate::openhuman::memory_sync::composio::{self, SyncReason}; +use crate::openhuman::memory_sync::composio; use crate::rpc::RpcOutcome; /// Parameters for `memory_sync_channel`. @@ -165,38 +165,33 @@ async fn spawn_manual_sync(requested_connection: Option) -> Result<(), S None, // provider-level composio sync — not a memory-source row ); - match composio::run_connection_sync( - config.clone(), + match crate::openhuman::tinycortex::run_composio_connection( + &target.toolkit, &target.connection_id, - SyncReason::Manual, + &config, ) .await { - // `run_connection_sync` returns `(SyncOutcome, ComposioUsage)` - // post-#3111; this caller only surfaces the outcome for UI - // stage events, so the usage tally is intentionally ignored - // here (the sync-audit caller in `memory_sources::sync` is the - // one that records it). - Ok((outcome, _usage)) => { + Ok(outcome) => { emit_sync_stage( MemorySyncTrigger::Manual, MemorySyncStage::Completed, - Some(&outcome.toolkit), - outcome.connection_id.as_deref(), + Some(&target.toolkit), + Some(&target.connection_id), Some(format!( "provider sync completed items_ingested={}", - outcome.items_ingested + outcome.records_ingested )), None, // provider-level composio sync — not a memory-source row ); } - Err((error, _usage)) => { + Err(error) => { emit_sync_stage( MemorySyncTrigger::Manual, MemorySyncStage::Failed, Some(&target.toolkit), Some(&target.connection_id), - Some(error.clone()), + Some(error.to_string()), None, // provider-level composio sync — not a memory-source row ); tracing::warn!( diff --git a/src/openhuman/memory/read_rpc/admin.rs b/src/openhuman/memory/read_rpc/admin.rs index e9d2f6339..58f1d7390 100644 --- a/src/openhuman/memory/read_rpc/admin.rs +++ b/src/openhuman/memory/read_rpc/admin.rs @@ -111,13 +111,13 @@ pub async fn wipe_all_rpc(config: &Config) -> Result } pub(crate) fn clear_composio_sync_state(db_path: &std::path::Path) -> Result { - use crate::openhuman::composio::providers::sync_state::KV_NAMESPACE; + use crate::openhuman::tinycortex::HOST_SYNC_STATE_NAMESPACE; let conn = rusqlite::Connection::open(db_path) .with_context(|| format!("open unified memory db {}", db_path.display()))?; let n = conn .execute( "DELETE FROM kv_namespace WHERE namespace = ?1", - params![KV_NAMESPACE], + params![HOST_SYNC_STATE_NAMESPACE], ) .context("delete composio-sync-state rows")?; Ok(n as u64) diff --git a/src/openhuman/memory/read_rpc_tests.rs b/src/openhuman/memory/read_rpc_tests.rs index 07335be71..0f36905e6 100644 --- a/src/openhuman/memory/read_rpc_tests.rs +++ b/src/openhuman/memory/read_rpc_tests.rs @@ -3,10 +3,9 @@ use crate::openhuman::composio::providers::sync_state::KV_NAMESPACE; use crate::openhuman::embeddings::NoopEmbedding; use crate::openhuman::memory::ingest_pipeline::ingest_chat; use crate::openhuman::memory_queue::drain_until_idle; +use crate::openhuman::memory_store::content::raw::{write_raw_items, RawItem, RawKind}; use crate::openhuman::memory_store::unified::UnifiedMemory; use crate::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; -use crate::openhuman::memory_sync::composio::providers::slack::ingest::ingest_page_into_memory_tree as ingest_slack_page; -use crate::openhuman::memory_sync::composio::providers::slack::SlackMessage; use chrono::{TimeZone, Utc}; use rusqlite::params; use std::sync::Arc; @@ -46,21 +45,37 @@ async fn seed_chat_chunk(cfg: &Config, source: &str, body: &str) { } async fn seed_slack_chunk_with_raw_archive(cfg: &Config) -> String { - let msg = SlackMessage { - channel_id: "C123".into(), - channel_name: "engineering".into(), - is_private: false, - author: "alice".into(), - author_id: "U123".into(), - text: "Phoenix migration launch window is Friday at 22:00 UTC.".into(), - timestamp: Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(), - ts_raw: "1700000000.000100".into(), - thread_ts: None, - permalink: Some("https://slack.example.test/archives/C123/p1700000000000100".into()), + let timestamp = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + write_raw_items( + &cfg.memory_tree_content_root(), + "slack:conn-slack-1", + &[RawItem { + uid: "1700000000.000100", + created_at_ms: timestamp.timestamp_millis(), + markdown: "**Channel:** #engineering\n**Author:** alice\n\nPhoenix migration launch window is Friday at 22:00 UTC.", + kind: RawKind::Chat, + }], + ) + .expect("seed raw Slack artifact"); + let batch = ChatBatch { + platform: "slack".into(), + channel_label: "#engineering".into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp, + text: "Phoenix migration launch window is Friday at 22:00 UTC.".into(), + source_ref: Some("slack://archives/C123/1700000000.000100".into()), + }], }; - ingest_slack_page(cfg, "alice", "conn-slack-1", &[msg]) - .await - .expect("seed slack ingest"); + ingest_chat( + cfg, + "slack:conn-slack-1", + "alice", + vec!["slack".into(), "ingested".into()], + batch, + ) + .await + .expect("seed slack ingest"); drain_until_idle(cfg).await.expect("drain slack ingest"); list_chunks_rpc(cfg, ChunkFilter::default()) diff --git a/src/openhuman/memory/sync_pipeline_e2e_tests.rs b/src/openhuman/memory/sync_pipeline_e2e_tests.rs index 08a28418e..90d3e2896 100644 --- a/src/openhuman/memory/sync_pipeline_e2e_tests.rs +++ b/src/openhuman/memory/sync_pipeline_e2e_tests.rs @@ -272,13 +272,27 @@ async fn multi_batch_volume_builds_full_tree() { }); assert!(canonicalized >= 20); - // Drain all jobs. - drain_until_idle(&cfg).await.unwrap(); + // A parallel test can briefly hold the process-global LLM gate, causing + // the seal job to defer. Keep draining until that deferred work becomes + // claimable and the durable tree state reflects the completed seal. + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(60); + let source_tree = loop { + drain_until_idle(&cfg).await.unwrap(); + if let Some(tree) = tree_store::list_trees_by_kind(&cfg, TreeKind::Source) + .unwrap() + .into_iter() + .find(|tree| tree.scope == source_id && tree.max_level >= 1) + { + break tree; + } + assert!( + tokio::time::Instant::now() < deadline, + "source tree did not seal before timeout" + ); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + }; // Source tree should have sealed to L1+. - let source_trees = tree_store::list_trees_by_kind(&cfg, TreeKind::Source).unwrap(); - assert!(!source_trees.is_empty()); - let source_tree = &source_trees[0]; assert!( source_tree.max_level >= 1, "should seal to L1+, got max_level={}", diff --git a/src/openhuman/memory/traits.rs b/src/openhuman/memory/traits.rs index 8734d58f6..6c63498e0 100644 --- a/src/openhuman/memory/traits.rs +++ b/src/openhuman/memory/traits.rs @@ -13,15 +13,9 @@ //! byte-identical to the former host definition before re-exporting; the tests //! below are the host-side seam that pins that contract on the crate type. //! -//! The `Memory` trait itself stays host-defined for now because of the -//! `sqlite_conn()` escape hatch, which the crate deliberately omits. That hatch -//! is retired in W3 (callers move to `tinycortex::memory::chunks::with_connection`), -//! after which the trait can also become a crate re-export. - -use async_trait::async_trait; -use parking_lot::Mutex; -use rusqlite::Connection; -use std::sync::Arc; +//! The `Memory` trait is also re-exported from `tinycortex`; backend-specific +//! resources such as SQLite connections are carried explicitly by factories +//! instead of being exposed through the storage abstraction. // ── Value types: re-exported from the crate (W2 type-unification, spec §0.5) ── // @@ -30,130 +24,9 @@ use std::sync::Arc; // fail-closed `from_db_str`). Re-exporting keeps one source of truth while every // `use crate::openhuman::memory::traits::{MemoryEntry, …}` site compiles unchanged. pub use tinycortex::memory::{ - MemoryCategory, MemoryEntry, MemoryTaint, NamespaceSummary, RecallOpts, + Memory, MemoryCategory, MemoryEntry, MemoryTaint, NamespaceSummary, RecallOpts, }; -/// The core trait for memory storage and retrieval. -/// -/// Any persistence backend (SQLite, Postgres, Vector DB, etc.) should implement -/// this trait to be used within the OpenHuman ecosystem. -/// -/// This mirrors [`tinycortex::memory::Memory`] method-for-method **plus** the -/// host-only [`Memory::sqlite_conn`] escape hatch. It stays host-defined until -/// W3 migrates that hatch to the crate's scoped `with_connection` accessor. -#[async_trait] -pub trait Memory: Send + Sync { - /// Returns the name of the memory backend (e.g., "sqlite", "vector"). - fn name(&self) -> &str; - - /// Stores a new memory entry or updates an existing one. - async fn store( - &self, - namespace: &str, - key: &str, - content: &str, - category: MemoryCategory, - session_id: Option<&str>, - ) -> anyhow::Result<()>; - - /// Store an entry with explicit provenance taint. - /// - /// Sync paths that ingest text from third-party services (Gmail / Slack / - /// Notion / Composio / etc.) MUST go through this entry point with - /// [`MemoryTaint::ExternalSync`] so the subconscious gate can refuse - /// external_effect tools when the resulting chunks reach a tick's - /// context window. - /// - /// The default implementation degrades to [`Self::store`] for backends - /// that do not yet persist taint (e.g. mock / in-memory stores used in - /// tests); the `UnifiedMemory` backend overrides this with a real - /// taint-carrying upsert. - async fn store_with_taint( - &self, - namespace: &str, - key: &str, - content: &str, - category: MemoryCategory, - session_id: Option<&str>, - taint: MemoryTaint, - ) -> anyhow::Result<()> { - let _ = taint; - self.store(namespace, key, content, category, session_id) - .await - } - - /// Recalls memories matching a query string using keyword or semantic search. - /// - /// Namespace is passed via `opts.namespace`; `None` uses the backend's - /// legacy default namespace (`GLOBAL_NAMESPACE`). - async fn recall( - &self, - query: &str, - limit: usize, - opts: RecallOpts<'_>, - ) -> anyhow::Result>; - - /// Recall documents in `namespace` semantically relevant to `query`, keeping - /// only those whose *vector* similarity to the query is at least - /// `min_vector_similarity`. Returns `(key, content)` pairs, most-relevant - /// first — the key lets callers act on the matched entry (e.g. overwrite a - /// contradicting preference by its topic). - /// - /// Unlike [`Self::recall`] (which ranks on a combined keyword + vector + - /// freshness score), this gates on the vector component alone, so an - /// unrelated query surfaces nothing — the behaviour Lane-B situational - /// preferences need. Default returns empty so keyword-only and mock backends - /// opt out; the unified store overrides it. - async fn recall_relevant_by_vector( - &self, - namespace: &str, - query: &str, - limit: usize, - min_vector_similarity: f64, - ) -> anyhow::Result> { - let _ = (namespace, query, limit, min_vector_similarity); - Ok(Vec::new()) - } - - /// Retrieves a specific memory entry by exact (namespace, key). - async fn get(&self, namespace: &str, key: &str) -> anyhow::Result>; - - /// Lists memory entries, optionally scoped by namespace, category, session. - async fn list( - &self, - namespace: Option<&str>, - category: Option<&MemoryCategory>, - session_id: Option<&str>, - ) -> anyhow::Result>; - - /// Deletes a memory entry associated with the given (namespace, key). - /// - /// Returns `Ok(true)` if the entry was found and deleted, `Ok(false)` if not found. - async fn forget(&self, namespace: &str, key: &str) -> anyhow::Result; - - /// Lists all namespaces with aggregate stats, for agent-side discovery. - async fn namespace_summaries(&self) -> anyhow::Result>; - - /// Returns the total count of all memory entries in the backend. - async fn count(&self) -> anyhow::Result; - - /// Performs a health check on the underlying storage system. - async fn health_check(&self) -> bool; - - /// Return the shared SQLite connection when the backend is `UnifiedMemory`. - /// - /// Used by subsystems (e.g. `ArchivistHook`) that need direct SQLite - /// access for FTS5 / segment writes without going through the async - /// `Memory` trait. - /// - /// Default: `None`. Only `UnifiedMemory` overrides this. Host-only escape - /// hatch (the crate omits it by design); retired to - /// `tinycortex::memory::chunks::with_connection` in W3. - fn sqlite_conn(&self) -> Option>> { - None - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/openhuman/memory_queue/mod.rs b/src/openhuman/memory_queue/mod.rs index 26ca4b9d7..f0414ddfe 100644 --- a/src/openhuman/memory_queue/mod.rs +++ b/src/openhuman/memory_queue/mod.rs @@ -31,7 +31,6 @@ //! `openhuman::memory` re-exports it as `memory::jobs` during the migration. mod ops; -mod redact; pub mod scheduler; pub mod store; pub mod testing; diff --git a/src/openhuman/memory_queue/ops.rs b/src/openhuman/memory_queue/ops.rs index ec9dcb2a8..7feb8633a 100644 --- a/src/openhuman/memory_queue/ops.rs +++ b/src/openhuman/memory_queue/ops.rs @@ -5,31 +5,16 @@ //! are preserved via re-exports in [`super`], so callers keep using //! `crate::openhuman::memory_queue::`. -use std::sync::atomic::{AtomicBool, Ordering}; - -use super::{store, types}; - -/// #1574 §6 / #1365: set while a re-embed backfill chain has work pending. -/// -/// Read by the first-person / subconscious retrieval layer so an empty -/// vector-search result during the backfill window is interpreted as -/// "not searched yet" rather than "no such memory" — preventing the agent -/// from confidently asserting false self-ignorance mid-re-embed. Set true -/// when a backfill is enqueued / still has rows; cleared when the chain -/// drains. Process-global (resets to false on restart; the worker re-sets -/// it on the next backfill tick — acceptable for v1). -static BACKFILL_IN_PROGRESS: AtomicBool = AtomicBool::new(false); - /// Mark whether a re-embed backfill currently has pending work. pub fn set_backfill_in_progress(v: bool) { - BACKFILL_IN_PROGRESS.store(v, Ordering::Relaxed); + tinycortex::memory::queue::set_backfill_in_progress(v); } /// True while a re-embed backfill chain still has rows to process. The /// #1365 absence-reasoning consumer checks this before treating an empty /// semantic-recall result as "no memory exists". pub fn backfill_in_progress() -> bool { - BACKFILL_IN_PROGRESS.load(Ordering::Relaxed) + tinycortex::memory::queue::backfill_in_progress() } /// #1574 §4: ensure a re-embed backfill chain exists for the **current** @@ -45,41 +30,11 @@ pub fn backfill_in_progress() -> bool { /// covered space enqueues nothing. Errors are logged, never propagated — /// a failed enqueue must not fail the user's settings save. pub fn ensure_reembed_backfill(config: &crate::openhuman::config::Config) { - let sig = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - let result = crate::openhuman::memory_store::chunks::store::with_connection(config, |conn| { - Ok(crate::openhuman::memory_store::chunks::store::has_uncovered_reembed_work(conn, &sig)?) - }); - match result { - Ok(true) => { - let job = match types::NewJob::reembed_backfill(&types::ReembedBackfillPayload { - signature: sig.clone(), - }) { - Ok(j) => j, - Err(e) => { - log::warn!("[memory::jobs] ensure_reembed_backfill: build job failed: {e}"); - return; - } - }; - match store::enqueue(config, &job) { - Ok(_) => { - set_backfill_in_progress(true); - log::info!( - "[memory::jobs] ensure_reembed_backfill: enqueued chain for sig={sig}" - ); - } - Err(e) => log::warn!( - "[memory::jobs] ensure_reembed_backfill: enqueue failed for sig={sig}: {e}" - ), - } - } - Ok(false) => { - log::debug!( - "[memory::jobs] ensure_reembed_backfill: sig={sig} fully covered; nothing to do" - ); - } - Err(e) => log::warn!( - "[memory::jobs] ensure_reembed_backfill: coverage probe failed for sig={sig}: {e}" - ), + let memory = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + let delegates = crate::openhuman::tinycortex::HostQueueDelegates::new(config.clone()); + if let Err(error) = tinycortex::memory::queue::ensure_reembed_backfill(&memory, &delegates) { + log::warn!("[memory::jobs] ensure_reembed_backfill failed: {error:#}"); } } diff --git a/src/openhuman/memory_queue/redact.rs b/src/openhuman/memory_queue/redact.rs deleted file mode 100644 index 6c0099a68..000000000 --- a/src/openhuman/memory_queue/redact.rs +++ /dev/null @@ -1,226 +0,0 @@ -//! Log-side scrubber for free-form `reason` / `last_error` strings emitted -//! by [`worker::run_once`] (defer/fail branches) and the matching -//! [`store::mark_failed`] / [`store::mark_deferred`] log lines. -//! -//! Persisted DB state (`mem_tree_jobs.last_error`) keeps the full -//! original string for diagnostics — handlers may attach -//! upstream-provider responses or full anyhow chains there. Logs are a -//! lower-trust sink (often forwarded to remote log aggregators or -//! shared in bug reports), so we apply a uniform scrub policy to: -//! -//! 1. Mask credential-shaped tokens (Bearer, OpenAI `sk-…`, GitHub -//! `ghp_…`, Slack `xox?-…`, generic `api_key=…` / `password=…` / -//! `token=…` assignments). -//! 2. Strip URL userinfo (`https://user:pass@host` → `https://***@host`). -//! 3. Mask bare email addresses. -//! 4. Cap the logged string at [`MAX_LEN`] bytes, suffixed with -//! `…(truncated, N more bytes)` so a reader knows the original -//! string was longer. -//! -//! [`worker::run_once`]: super::worker::run_once - -use std::sync::OnceLock; - -use regex::Regex; - -/// Upper bound on the byte length of a scrubbed string. Long enough to -/// keep the head of an anyhow `{:#}` chain (which usually includes the -/// most useful context) but short enough that one bad job can't flood -/// the log with megabytes of provider response body. -pub(crate) const MAX_LEN: usize = 1024; - -/// Scrub a free-form error / reason string for emission to logs. -/// Returns an owned `String` because every regex pass may rewrite the -/// input; callers are emitting this through `format!` / `log::*!` -/// anyway, so the allocation isn't on a hot path. -pub(crate) fn scrub_for_log(input: &str) -> String { - let mut out = input.to_owned(); - for (re, replacement) in patterns() { - // `Cow::into_owned` is cheap when the regex didn't match. - out = re.replace_all(&out, *replacement).into_owned(); - } - truncate(out) -} - -fn truncate(mut s: String) -> String { - if s.len() <= MAX_LEN { - return s; - } - // Round down to a char boundary to avoid splitting a multi-byte - // UTF-8 sequence — `truncate` itself panics on a non-boundary. - let mut cut = MAX_LEN; - while cut > 0 && !s.is_char_boundary(cut) { - cut -= 1; - } - let dropped = s.len() - cut; - s.truncate(cut); - s.push_str(&format!("…(truncated, {dropped} more bytes)")); - s -} - -fn patterns() -> &'static [(Regex, &'static str)] { - static PATTERNS: OnceLock> = OnceLock::new(); - PATTERNS.get_or_init(|| { - vec![ - // URL userinfo: capture the scheme + ://, then the userinfo - // up to '@'. Replace with ***@. Anchored on `://` to avoid - // eating bare `user@host` strings — those are handled by - // the email rule below. - ( - Regex::new(r"(?P[a-zA-Z][a-zA-Z0-9+.\-]*://)[^\s/@]+@").unwrap(), - "$scheme***@", - ), - // Bearer tokens (with or without an `Authorization:` prefix). - ( - Regex::new(r"(?i)bearer\s+[A-Za-z0-9._\-+/=]+").unwrap(), - "Bearer ***", - ), - // Provider-prefixed credentials with stable, well-known - // shapes. Listed individually so a future reader can see - // exactly which providers are covered. - ( - Regex::new(r"sk-[A-Za-z0-9_\-]{16,}").unwrap(), - "sk-***", - ), - ( - Regex::new(r"ghp_[A-Za-z0-9]{20,}").unwrap(), - "ghp_***", - ), - ( - Regex::new(r"ghs_[A-Za-z0-9]{20,}").unwrap(), - "ghs_***", - ), - ( - Regex::new(r"gho_[A-Za-z0-9]{20,}").unwrap(), - "gho_***", - ), - ( - Regex::new(r"xox[abprs]-[A-Za-z0-9\-]{8,}").unwrap(), - "xox-***", - ), - // Generic `key=value` assignments where the key name implies - // a secret. Matches `api_key`, `apiKey`, `api-key`, - // `password`, `passwd`, `pwd`, `token`, `secret`. Accepts - // a quoted, single-quoted, or bare value; bare values stop - // at the first whitespace / comma / closing bracket so we - // don't eat the rest of the message. - ( - Regex::new( - r#"(?i)(?Papi[_\-]?key|password|passwd|pwd|secret|token|auth)\s*[:=]\s*(?:"[^"]*"|'[^']*'|[^\s,}\)\]]+)"#, - ) - .unwrap(), - "$k=***", - ), - // Bare email addresses. Conservative pattern (no UTF-8 - // local parts, no quoted local parts) — sufficient to mask - // the common cases without eating unrelated `@` symbols. - ( - Regex::new(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b").unwrap(), - "***@***", - ), - ] - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn passthrough_for_safe_string() { - let input = "rate_limited: provider returned 429, retry in 30s"; - assert_eq!(scrub_for_log(input), input); - } - - #[test] - fn masks_bearer_token() { - let s = scrub_for_log("Authorization: Bearer abc123.def-456_xyz"); - assert!(s.contains("Bearer ***"), "got {s:?}"); - assert!(!s.contains("abc123")); - } - - #[test] - fn masks_openai_key() { - let s = scrub_for_log("upstream returned 401: invalid sk-abcDEF1234567890ZZZZ key"); - assert!(s.contains("sk-***")); - assert!(!s.contains("sk-abcDEF1234567890ZZZZ")); - } - - #[test] - fn masks_github_token_variants() { - for raw in [ - "ghp_abcdefghij1234567890ABCD", - "ghs_abcdefghij1234567890ABCD", - "gho_abcdefghij1234567890ABCD", - ] { - let s = scrub_for_log(&format!("error: token {raw} rejected")); - assert!(s.contains("***"), "input={raw} out={s}"); - assert!(!s.contains(raw), "input={raw} out={s}"); - } - } - - #[test] - fn masks_slack_token() { - let s = scrub_for_log("posting to slack failed: xoxb-1234567890-abcdEFG"); - assert!(s.contains("xox-***")); - assert!(!s.contains("xoxb-1234567890")); - } - - #[test] - fn masks_generic_secret_assignments() { - let inputs = [ - ("api_key=hunter2 trailing", "api_key=***"), - ("password: hunter2", "password=***"), - ("Token = hunter2,more", "Token=***"), - ("apiKey=\"hunter2\"", "apiKey=***"), - ]; - for (raw, expect) in inputs { - let s = scrub_for_log(raw); - assert!(s.contains(expect), "input={raw:?} out={s:?}"); - assert!(!s.contains("hunter2"), "input={raw:?} out={s:?}"); - } - } - - #[test] - fn strips_url_userinfo() { - let s = scrub_for_log("connect failed https://alice:s3cret@db.internal/x"); - assert!(s.contains("https://***@db.internal/x"), "got {s:?}"); - assert!(!s.contains("alice")); - assert!(!s.contains("s3cret")); - } - - #[test] - fn masks_email() { - let s = scrub_for_log("user alice@example.com triggered job"); - assert!(s.contains("***@***"), "got {s:?}"); - assert!(!s.contains("alice")); - assert!(!s.contains("example.com")); - } - - #[test] - fn truncates_oversized_input() { - let big = "x".repeat(MAX_LEN * 2); - let s = scrub_for_log(&big); - assert!(s.len() < big.len()); - assert!(s.contains("(truncated,")); - } - - #[test] - fn truncate_handles_multibyte_boundary() { - // `é` is 2 bytes in UTF-8; build a string whose naïve cut at - // MAX_LEN would land mid-codepoint. - let mut big = "a".repeat(MAX_LEN - 1); - big.push('é'); - big.push_str(&"b".repeat(64)); - let s = scrub_for_log(&big); - // Must not panic; must produce valid UTF-8 (String guarantees). - assert!(s.contains("(truncated,")); - } - - #[test] - fn idempotent_on_already_scrubbed_string() { - let once = scrub_for_log("Bearer abcdef api_key=hunter2"); - let twice = scrub_for_log(&once); - assert_eq!(once, twice); - } -} diff --git a/src/openhuman/memory_queue/scheduler.rs b/src/openhuman/memory_queue/scheduler.rs index d853a8567..86328a1fb 100644 --- a/src/openhuman/memory_queue/scheduler.rs +++ b/src/openhuman/memory_queue/scheduler.rs @@ -7,11 +7,7 @@ use std::time::Duration; -use chrono::{Timelike, Utc}; - use crate::openhuman::config::Config; -use crate::openhuman::memory_queue::store; -use crate::openhuman::memory_queue::types::{FlushStalePayload, NewJob}; static STARTED: std::sync::Once = std::sync::Once::new(); @@ -43,7 +39,9 @@ pub fn start(config: Config) { /// Unrecoverable failures stay parked — see /// [`store::requeue_transient_failed`]. fn retry_transient_failures(config: &Config) { - match store::requeue_transient_failed(config) { + let memory = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + match tinycortex::memory::queue::scheduler::self_heal(&memory) { Ok(0) => {} Ok(n) => { log::info!("[memory::jobs] periodic retry requeued {n} transient-failed job(s)"); @@ -56,26 +54,15 @@ fn retry_transient_failures(config: &Config) { } fn enqueue_flush_stale(config: &Config) { - // Take a single `Utc::now()` reading and derive both the date and - // 3-hour block from it so the dedupe key can't disagree with itself - // across a 3-hour boundary. - let now = Utc::now(); - let today_iso = now.date_naive().format("%Y-%m-%d").to_string(); - let hour_block = now.hour() / 3; - match NewJob::flush_stale(&FlushStalePayload::default(), &today_iso, hour_block) { - Ok(new_job) => { - match store::enqueue(config, &new_job) { - Ok(Some(_)) => { - super::worker::wake_workers(); - } - Ok(None) => {} // dedupe-suppressed — OK - Err(err) => { - log::warn!("[memory::jobs] periodic flush_stale enqueue failed: {err:#}"); - } - } + let memory = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + match tinycortex::memory::queue::scheduler::enqueue_flush_stale(&memory) { + Ok(Some(_)) => { + super::worker::wake_workers(); } + Ok(None) => {} Err(err) => { - log::warn!("[memory::jobs] flush_stale job build failed: {err:#}"); + log::warn!("[memory::jobs] periodic flush_stale enqueue failed: {err:#}"); } } } diff --git a/src/openhuman/memory_queue/store.rs b/src/openhuman/memory_queue/store.rs index 53956b20b..29ddb14b1 100644 --- a/src/openhuman/memory_queue/store.rs +++ b/src/openhuman/memory_queue/store.rs @@ -1,1397 +1,93 @@ -//! SQLite persistence for the memory-tree job queue. -//! -//! Producers call [`enqueue`] inside their own writes (or with a fresh tx) -//! to atomically commit the side-effect plus its follow-up job. The worker -//! pool calls [`claim_next`] to lease a job, [`mark_done`] / [`mark_failed`] -//! to settle it, and [`recover_stale_locks`] on startup to flip rows whose -//! `locked_until_ms` expired without a settle. -//! -//! Concurrency: -//! - The dedupe key is enforced by a partial `UNIQUE` index that only -//! covers `status IN ('ready', 'running')`. Producers use `INSERT OR -//! IGNORE` so a duplicate enqueue while a job is in flight or queued is -//! a silent no-op; a duplicate enqueue after the first completes is -//! accepted and creates a fresh row. -//! - `claim_next` is one statement: `UPDATE … WHERE id = (SELECT … LIMIT 1) -//! RETURNING …`. SQLite serialises writes, so no two workers can claim -//! the same row. +//! Product Config adapters over tinycortex's SQLite queue store. -use anyhow::{Context, Result}; -use chrono::Utc; -use rusqlite::{params, Connection, OptionalExtension, Transaction}; -use uuid::Uuid; +use anyhow::Result; +use rusqlite::Transaction; use crate::openhuman::config::Config; -use crate::openhuman::memory_queue::redact::scrub_for_log; -use crate::openhuman::memory_queue::types::{Job, JobKind, JobStatus, NewJob}; -use crate::openhuman::memory_store::chunks::store::with_connection; +use crate::openhuman::memory_tree::health::PipelineFailure; -/// Default visibility lock — a worker that crashes mid-job will have its -/// row recovered after this window. 5 min is comfortably larger than any -/// expected single-job runtime (LLM extract or summarise) without leaving -/// real failures stuck for hours. -pub const DEFAULT_LOCK_DURATION_MS: i64 = 5 * 60 * 1_000; +use super::types::{Job, JobFailure, JobStatus, NewJob}; -/// Backoff math for retry. Returns `now + min(base * 2^attempts, cap)`. -const RETRY_BASE_MS: i64 = 60 * 1_000; -const RETRY_CAP_MS: i64 = 60 * 60 * 1_000; -const DEFAULT_MAX_ATTEMPTS: u32 = 5; +pub use tinycortex::memory::queue::DEFAULT_LOCK_DURATION_MS; + +fn memory_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} -/// Enqueue one job. Idempotent on `dedupe_key` while another active row -/// (status `ready`/`running`) shares it. Returns `Some(id)` if the row -/// was inserted, `None` if a duplicate was suppressed. pub fn enqueue(config: &Config, job: &NewJob) -> Result> { - with_connection(config, |conn| enqueue_conn(conn, job)) + tinycortex::memory::queue::enqueue(&memory_config(config), job) } -/// Enqueue inside a caller-owned transaction. Use this when the producer -/// is already mid-tx (e.g. `ingest::persist` writing chunks + jobs in one -/// commit) so the queue insert lands atomically with the side-effect. -/// `Transaction` derefs to `Connection`, so callers just pass `&tx`. pub fn enqueue_tx(tx: &Transaction<'_>, job: &NewJob) -> Result> { - enqueue_conn(tx, job) + tinycortex::memory::queue::enqueue_tx(tx, job) } -pub(crate) fn enqueue_conn(conn: &Connection, job: &NewJob) -> Result> { - let id = format!("job:{}", Uuid::new_v4()); - let now_ms = Utc::now().timestamp_millis(); - let available_at = job.available_at_ms.unwrap_or(now_ms); - let max_attempts = job.max_attempts.unwrap_or(DEFAULT_MAX_ATTEMPTS) as i64; - - let inserted = conn.execute( - "INSERT OR IGNORE INTO mem_tree_jobs ( - id, kind, payload_json, dedupe_key, status, attempts, max_attempts, - available_at_ms, locked_until_ms, last_error, - created_at_ms, started_at_ms, completed_at_ms - ) VALUES (?1, ?2, ?3, ?4, 'ready', 0, ?5, ?6, NULL, NULL, ?7, NULL, NULL)", - params![ - id, - job.kind.as_str(), - job.payload_json, - job.dedupe_key, - max_attempts, - available_at, - now_ms, - ], - )?; - - if inserted == 0 { - log::debug!( - "[memory::jobs] enqueue suppressed by dedupe kind={} key={:?}", - job.kind.as_str(), - job.dedupe_key - ); - return Ok(None); - } - log::debug!( - "[memory::jobs] enqueued id={} kind={} avail_at_ms={} dedupe={:?}", - id, - job.kind.as_str(), - available_at, - job.dedupe_key - ); - Ok(Some(id)) -} - -/// Atomically claim the next ready job whose `available_at_ms` has come -/// due. Sets `status=running`, bumps `attempts`, stamps `started_at_ms` -/// and `locked_until_ms`. Returns `None` when the queue is empty / not -/// yet due. pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result> { - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let lock_until = now_ms.saturating_add(lock_duration_ms); - - let row = conn - .query_row( - // Drain forward, don't widen. Most-downstream kinds run - // first so a slow LLM-bound `extract_chunk` can't starve - // the routing/seal/digest pipeline behind it. - "UPDATE mem_tree_jobs - SET status = 'running', - attempts = attempts + 1, - started_at_ms = ?1, - locked_until_ms = ?2, - last_error = NULL - WHERE id = ( - SELECT id FROM mem_tree_jobs - WHERE status = 'ready' - AND available_at_ms <= ?1 - ORDER BY - CASE kind - WHEN 'digest_daily' THEN 0 - WHEN 'seal' THEN 1 - WHEN 'flush_stale' THEN 2 - WHEN 'topic_route' THEN 3 - WHEN 'append_buffer' THEN 4 - ELSE 5 - END ASC, - available_at_ms ASC - LIMIT 1 - ) - RETURNING id, kind, payload_json, dedupe_key, status, attempts, - max_attempts, available_at_ms, locked_until_ms, last_error, - created_at_ms, started_at_ms, completed_at_ms, - failure_reason, failure_class", - params![now_ms, lock_until], - row_to_job, - ) - .optional() - .context("Failed to claim next mem_tree_jobs row")?; - if let Some(j) = &row { - log::debug!( - "[memory::jobs] claimed id={} kind={} attempt={}/{}", - j.id, - j.kind.as_str(), - j.attempts, - j.max_attempts - ); - } - Ok(row) - }) + tinycortex::memory::queue::claim_next(&memory_config(config), lock_duration_ms) } -/// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`. -/// -/// The UPDATE is gated on `attempts` and `started_at_ms` matching the values -/// in `job` (the snapshot returned by [`claim_next`]). If the lease expired -/// and another worker re-claimed the row, `rows_affected` will be 0 — the -/// stale worker's settlement is a silent no-op rather than clobbering the new -/// lessee's state. pub fn mark_done(config: &Config, job: &Job) -> Result<()> { - let job_id = &job.id; - let claim_attempts = job.attempts as i64; - let claim_started_at = job.started_at_ms; - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'done', - completed_at_ms = ?1, - locked_until_ms = NULL, - last_error = NULL - WHERE id = ?2 - AND attempts = ?3 - AND started_at_ms IS ?4", - params![now_ms, job_id, claim_attempts, claim_started_at], - )?; - if n == 0 { - // Either the job row was deleted (shouldn't happen) or the lease - // expired and a second worker re-claimed the row. Log and move on — - // this is a known race outcome, not a bug in the current worker. - log::warn!( - "[memory::jobs] mark_done id={job_id} was a no-op \ - (stale lease: attempts={claim_attempts} started_at_ms={claim_started_at:?})" - ); - } - Ok(()) - }) + tinycortex::memory::queue::mark_done(&memory_config(config), job) } -/// Settle a failed job. If `attempts < max_attempts`, the row goes back -/// to `ready` with an exponential-backoff `available_at_ms`. Otherwise -/// it terminates as `failed`. Either way `last_error` is recorded. -/// -/// Like [`mark_done`], the UPDATE is gated on the claim-token -/// (`attempts` + `started_at_ms`) so a stale worker's failure settlement -/// cannot clobber an active lessee's row — rows_affected == 0 is a silent -/// no-op. pub fn mark_failed(config: &Config, job: &Job, error: &str) -> Result<()> { - mark_failed_typed(config, job, error, None) + tinycortex::memory::queue::mark_failed(&memory_config(config), job, error) } -/// Like [`mark_failed`], but with an optional typed [`PipelineFailure`] -/// classification (#002). When `failure` is `Some` and **unrecoverable** -/// (budget exhausted, bad key, missing local model, dim mismatch), the job -/// terminates as `failed` **immediately** — no retry budget is burned, since -/// retrying the same input cannot succeed — and the typed -/// `failure_reason` / `failure_class` columns are persisted alongside the -/// freeform `last_error`. Transient classifications (and the untyped `None` -/// case) keep the existing attempts-bounded retry-with-backoff behaviour. -/// -/// The claim-token gate (`attempts` + `started_at_ms`) is preserved on every -/// branch so a stale lessee's settlement remains a silent no-op. pub fn mark_failed_typed( config: &Config, job: &Job, error: &str, - failure: Option<&crate::openhuman::memory_tree::health::PipelineFailure>, + failure: Option<&PipelineFailure>, ) -> Result<()> { - let job_id = &job.id; - let attempts = job.attempts as i64; - let max_attempts = job.max_attempts as i64; - let claim_started_at = job.started_at_ms; - let unrecoverable = failure.map(|f| f.is_unrecoverable()).unwrap_or(false); - let failure_reason = failure.map(|f| f.code.as_str()); - let failure_class = failure.map(|f| f.class.as_str()); - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - - // `error` is free-form (anyhow chain or handler-supplied text) and - // may carry credential-shaped substrings; scrub before logging, - // but keep the original in the DB column for diagnostics. - let error_for_log = scrub_for_log(error); - // Terminal when the retry budget is exhausted OR the failure is - // classified unrecoverable (fail fast — #002 FR-003). - if attempts >= max_attempts || unrecoverable { - log::warn!( - "[memory::jobs] terminal failure id={job_id} \ - attempts={attempts}/{max_attempts} unrecoverable={unrecoverable} \ - reason={failure_reason:?} err={error_for_log}" - ); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'failed', - completed_at_ms = ?1, - locked_until_ms = NULL, - last_error = ?2, - failure_reason = ?6, - failure_class = ?7 - WHERE id = ?3 - AND attempts = ?4 - AND started_at_ms IS ?5", - params![ - now_ms, - error, - job_id, - attempts, - claim_started_at, - failure_reason, - failure_class, - ], - )?; - if n == 0 { - log::warn!( - "[memory::jobs] mark_failed(terminal) id={job_id} was a no-op \ - (stale lease: attempts={attempts} started_at_ms={claim_started_at:?})" - ); - } - } else { - let backoff = backoff_ms(attempts as u32); - let next_at = now_ms.saturating_add(backoff); - log::info!( - "[memory::jobs] retry id={job_id} attempt={attempts}/{max_attempts} \ - next_at_ms={next_at} err={error_for_log}" - ); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - available_at_ms = ?1, - locked_until_ms = NULL, - last_error = ?2 - WHERE id = ?3 - AND attempts = ?4 - AND started_at_ms IS ?5", - params![next_at, error, job_id, attempts, claim_started_at], - )?; - if n == 0 { - log::warn!( - "[memory::jobs] mark_failed(retry) id={job_id} was a no-op \ - (stale lease: attempts={attempts} started_at_ms={claim_started_at:?})" - ); - } - } - Ok(()) - }) + let failure = failure.map(|failure| JobFailure { + code: failure.code.as_str(), + class: failure.class.as_str(), + }); + tinycortex::memory::queue::mark_failed_typed( + &memory_config(config), + job, + error, + failure.as_ref(), + ) } -/// Mark a claimed job as deferred: put it back to `ready` with -/// `available_at_ms = until_ms` so [`claim_next`] will re-pick it once the -/// wake-up time has passed. The handler ran successfully but chose not to -/// make progress (cloud rate-limited, dependency unavailable, model -/// warming up), so this path **does not** burn the failure budget — the -/// `attempts` bump that [`claim_next`] applied at claim time is reverted. -/// -/// `reason` is recorded in `last_error` for visibility and `started_at_ms` -/// is cleared (mirroring the retry branch of [`mark_failed`]) so the next -/// claim stamps a fresh start time. -/// -/// Like [`mark_done`] / [`mark_failed`], the UPDATE is gated on the -/// claim-token (`attempts` + `started_at_ms`) so a stale lessee's -/// settlement is a silent no-op rather than clobbering an active lessee's -/// row. pub fn mark_deferred(config: &Config, job: &Job, until_ms: i64, reason: &str) -> Result<()> { - let job_id = &job.id; - let claim_attempts = job.attempts as i64; - let pre_claim_attempts = claim_attempts.saturating_sub(1); - let claim_started_at = job.started_at_ms; - with_connection(config, |conn| { - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - attempts = ?1, - available_at_ms = ?2, - locked_until_ms = NULL, - started_at_ms = NULL, - last_error = ?3 - WHERE id = ?4 - AND attempts = ?5 - AND started_at_ms IS ?6", - params![ - pre_claim_attempts, - until_ms, - reason, - job_id, - claim_attempts, - claim_started_at, - ], - )?; - if n == 0 { - log::warn!( - "[memory::jobs] mark_deferred id={job_id} was a no-op \ - (stale lease: attempts={claim_attempts} started_at_ms={claim_started_at:?})" - ); - } - Ok(()) - }) + tinycortex::memory::queue::mark_deferred(&memory_config(config), job, until_ms, reason) } -/// Flip any `running` row whose `locked_until_ms` has expired back to -/// `ready`. Called once at worker startup so a process crash mid-job -/// doesn't leave work stranded. Returns the number of rows recovered. pub fn recover_stale_locks(config: &Config) -> Result { - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - last_error = COALESCE(last_error, 'recovered_from_stale_lock') - WHERE status = 'running' - AND locked_until_ms IS NOT NULL - AND locked_until_ms < ?1", - params![now_ms], - )?; - if n > 0 { - // Info, not warn: with graceful shutdown releasing in-flight - // locks (see `release_running_locks`), a non-empty recovery now - // means the previous process was hard-killed — expected and - // self-healing, not actionable (bug-report-2026-05-26 I2). - log::info!("[memory::jobs] recovered {n} stale-locked job(s) at startup"); - } - Ok(n) - }) + tinycortex::memory::queue::recover_stale_locks(&memory_config(config)) } -/// Requeue every terminally-`failed` job back to `ready` (#002 FR-011). -/// -/// Backs the manual `memory_tree_retry_failed` RPC: once the user fixes the -/// underlying cause (e.g. adds an embeddings key, switches to a faster -/// extraction model), the jobs that failed under the old config re-run without -/// re-ingesting source data. Resets `attempts` to 0 (a fresh retry budget), -/// clears the typed `failure_reason` / `failure_class` and `last_error`, and -/// makes the row immediately available. Returns the number of jobs requeued. -/// -/// Automatic callers: the manual sync path (`memory_sources::sync`, -/// via [`retry_all_failed`]) and the 3-hourly queue scheduler (via -/// [`requeue_transient_failed`], which excludes `unrecoverable` -/// failures so a bad config can't 3-hourly retry-loop forever). pub fn requeue_failed(config: &Config) -> Result { - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - attempts = 0, - available_at_ms = ?1, - locked_until_ms = NULL, - started_at_ms = NULL, - completed_at_ms = NULL, - last_error = NULL, - failure_reason = NULL, - failure_class = NULL - WHERE status = 'failed'", - params![now_ms], - )?; - if n > 0 { - log::info!("[memory::jobs] requeued {n} failed job(s) for retry"); - } - Ok(n as u64) - }) + tinycortex::memory::queue::requeue_failed(&memory_config(config)) } -/// Requeue only failed jobs whose recorded failure class is NOT -/// `unrecoverable` — i.e. transient failures (network 5xx, timeouts, -/// SQLITE_BUSY) and legacy rows with no class recorded. -/// -/// This is the **automatic** self-healing variant, fired from the -/// 3-hourly queue scheduler so a crash or flaky network never leaves -/// pipeline jobs permanently stuck until the user happens to press -/// "Sync now". Unrecoverable failures (bad/missing key, budget -/// exhausted, dim mismatch) stay parked for the manual -/// `memory_tree_retry_failed` RPC after the user fixes the cause — -/// auto-retrying those would just burn a failure every 3 hours forever. pub fn requeue_transient_failed(config: &Config) -> Result { - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - attempts = 0, - available_at_ms = ?1, - locked_until_ms = NULL, - started_at_ms = NULL, - completed_at_ms = NULL, - last_error = NULL, - failure_reason = NULL, - failure_class = NULL - WHERE status = 'failed' - AND (failure_class IS NULL OR failure_class != 'unrecoverable')", - params![now_ms], - )?; - if n > 0 { - log::info!( - "[memory::jobs] auto-requeued {n} transient-failed job(s) from the periodic scheduler" - ); - } - Ok(n as u64) - }) + tinycortex::memory::queue::requeue_transient_failed(&memory_config(config)) } -/// Release this process's in-flight job locks on a *graceful* shutdown: -/// flip every `running` row back to `ready` so the work is immediately -/// re-claimable on next launch instead of waiting out the lease and -/// surfacing as a stale-lock recovery. The core runs a single worker pool, -/// so any `running` row at clean-shutdown time was claimed by us. -/// Returns the number of rows released (bug-report-2026-05-26 I2). pub fn release_running_locks(config: &Config) -> Result { - with_connection(config, |conn| { - // TODO(multi-process): this releases *every* `running` row, with no - // process-scoped guard (e.g. a `worker_pid` / session-token column). - // Correct today because a single worker pool is the invariant, but if - // multi-process workers are ever introduced, a rolling restart with - // two overlapping processes would let one shutdown release the other's - // in-flight locks. Add a per-owner predicate here when that happens. - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - locked_until_ms = NULL - WHERE status = 'running'", - [], - )?; - Ok(n) - }) + tinycortex::memory::queue::release_running_locks(&memory_config(config)) } -/// Quick count helper for tests / diagnostics. pub fn count_by_status(config: &Config, status: JobStatus) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row( - "SELECT COUNT(*) FROM mem_tree_jobs WHERE status = ?1", - params![status.as_str()], - |r| r.get(0), - )?; - Ok(n.max(0) as u64) - }) + tinycortex::memory::queue::count_by_status(&memory_config(config), status) } -/// #3365: count terminally-`failed` jobs whose typed class is `unrecoverable` -/// — the ones `requeue_transient_failed` deliberately leaves parked (budget / -/// auth / dim-mismatch) because retrying can't help and the user must act. The -/// status surface routes ONLY these to a hard `error`; transient failures (or -/// untyped/NULL rows) are auto-requeued and self-heal, so they stay `degraded`. -/// The `failure_class = 'unrecoverable'` predicate is the exact complement of -/// the requeue gate's `IS NULL OR != 'unrecoverable'`, so the two never drift. pub fn count_failed_unrecoverable(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row( - "SELECT COUNT(*) FROM mem_tree_jobs \ - WHERE status = 'failed' AND failure_class = 'unrecoverable'", - [], - |r| r.get(0), - )?; - Ok(n.max(0) as u64) - }) + tinycortex::memory::queue::count_failed_unrecoverable(&memory_config(config)) } -/// Total count regardless of status — handy for assertions. pub fn count_total(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0))?; - Ok(n.max(0) as u64) - }) + tinycortex::memory::queue::count_total(&memory_config(config)) } -/// Reset all terminally-failed jobs back to `ready` so the worker pool -/// re-attempts them. Resets `attempts` to 0 and clears the lock/error -/// state. Returns the number of jobs reset. pub fn retry_all_failed(config: &Config) -> Result { - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - let n = conn.execute( - "UPDATE mem_tree_jobs - SET status = 'ready', - attempts = 0, - available_at_ms = ?1, - locked_until_ms = NULL, - started_at_ms = NULL, - completed_at_ms = NULL, - last_error = NULL - WHERE status = 'failed'", - params![now_ms], - )?; - if n > 0 { - log::info!("[memory::jobs] retry_all_failed: reset {n} jobs back to ready"); - } - Ok(n as u64) - }) + tinycortex::memory::queue::retry_all_failed(&memory_config(config)) } -/// Fetch one job by id (test/diagnostic helper). pub fn get_job(config: &Config, id: &str) -> Result> { - with_connection(config, |conn| { - let job = conn - .query_row( - "SELECT id, kind, payload_json, dedupe_key, status, attempts, max_attempts, - available_at_ms, locked_until_ms, last_error, - created_at_ms, started_at_ms, completed_at_ms, - failure_reason, failure_class - FROM mem_tree_jobs WHERE id = ?1", - params![id], - row_to_job, - ) - .optional()?; - Ok(job) - }) -} - -fn row_to_job(row: &rusqlite::Row<'_>) -> rusqlite::Result { - let id: String = row.get(0)?; - let kind_s: String = row.get(1)?; - let payload_json: String = row.get(2)?; - let dedupe_key: Option = row.get(3)?; - let status_s: String = row.get(4)?; - let attempts: i64 = row.get(5)?; - let max_attempts: i64 = row.get(6)?; - let available_at_ms: i64 = row.get(7)?; - let locked_until_ms: Option = row.get(8)?; - let last_error: Option = row.get(9)?; - let created_at_ms: i64 = row.get(10)?; - let started_at_ms: Option = row.get(11)?; - let completed_at_ms: Option = row.get(12)?; - let failure_reason: Option = row.get(13)?; - let failure_class: Option = row.get(14)?; - - let kind = JobKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into()) - })?; - let status = JobStatus::parse(&status_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, e.into()) - })?; - - Ok(Job { - id, - kind, - payload_json, - dedupe_key, - status, - attempts: attempts.max(0) as u32, - max_attempts: max_attempts.max(0) as u32, - available_at_ms, - locked_until_ms, - last_error, - failure_reason, - failure_class, - created_at_ms, - started_at_ms, - completed_at_ms, - }) -} - -/// Exponential backoff: attempt 1 → 60s, 2 → 120s, 3 → 240s, capped at 1h. -fn backoff_ms(attempts_so_far: u32) -> i64 { - // attempts_so_far is the count BEFORE the next retry's attempt — so the - // first retry uses attempts_so_far=1, giving base*2^0 = 60s. - let exp = attempts_so_far.saturating_sub(1).min(20); // cap shift - let mult = 1i64 << exp; // 1, 2, 4, … - let raw = RETRY_BASE_MS.saturating_mul(mult); - raw.min(RETRY_CAP_MS) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_queue::types::{ - AppendBufferPayload, AppendTarget, ExtractChunkPayload, NodeRef, - }; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - #[test] - fn enqueue_and_claim_roundtrip() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c1".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claimed.id, id); - assert_eq!(claimed.status, JobStatus::Running); - assert_eq!(claimed.attempts, 1); - assert!(claimed.locked_until_ms.is_some()); - - // Second claim should see no eligible row (the only one is now running). - let again = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap(); - assert!(again.is_none()); - - mark_done(&cfg, &claimed).unwrap(); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Done); - assert!(row.completed_at_ms.is_some()); - assert!(row.locked_until_ms.is_none()); - } - - /// T006/T007: the new `failure_reason`/`failure_class` columns must exist - /// (migration ran) and round-trip through `claim_next` (RETURNING) and - /// `get_job` (SELECT) as `None` until a classified failure is recorded - /// (T012 wires the write side). - #[test] - fn typed_failure_columns_roundtrip_as_none_by_default() { - let (_tmp, cfg) = test_config(); - let nj = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "c-typed-fail".into(), - }) - .unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claimed.failure_reason, None); - assert_eq!(claimed.failure_class, None); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.failure_reason, None); - assert_eq!(row.failure_class, None); - } - - /// T012: an **unrecoverable** typed failure must terminate the job - /// immediately (status `failed`, `completed_at_ms` set) on the FIRST - /// attempt — no retry budget burned — and persist the typed - /// `failure_reason` / `failure_class` columns. A job with `max_attempts` - /// far above 1 proves the short-circuit isn't just the budget running out. - #[test] - fn mark_failed_typed_unrecoverable_terminates_immediately() { - use crate::openhuman::memory_tree::health::{FailureCode, PipelineFailure}; - let (_tmp, cfg) = test_config(); - let mut nj = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "c-budget".into(), - }) - .unwrap(); - nj.max_attempts = Some(5); // plenty of budget left - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claimed.attempts, 1, "first claim"); - let failure = PipelineFailure::new(FailureCode::BudgetExhausted); - mark_failed_typed(&cfg, &claimed, "Insufficient budget", Some(&failure)).unwrap(); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!( - row.status, - JobStatus::Failed, - "unrecoverable failure must terminate on first attempt" - ); - assert!(row.completed_at_ms.is_some()); - assert_eq!(row.failure_reason.as_deref(), Some("budget_exhausted")); - assert_eq!(row.failure_class.as_deref(), Some("unrecoverable")); - assert_eq!(row.last_error.as_deref(), Some("Insufficient budget")); - } - - /// #3365: `count_failed_unrecoverable` counts ONLY terminally-failed jobs - /// whose class is `unrecoverable`. Transient-class failures (terminated by - /// exhausting `max_attempts`) and untyped/NULL-class failures self-heal via - /// `requeue_transient_failed`, so they're excluded — the complement of the - /// requeue gate. The status surface uses this to route only the former to - /// `error`. - #[test] - fn count_failed_unrecoverable_excludes_transient_and_untyped() { - use crate::openhuman::memory_tree::health::{FailureCode, PipelineFailure}; - let (_tmp, cfg) = test_config(); - - // Helper: enqueue a job, claim it, then mark it failed with `failure`. - let fail_one = |chunk: &str, max_attempts: u32, failure: Option<&PipelineFailure>| { - let mut nj = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: chunk.into(), - }) - .unwrap(); - nj.max_attempts = Some(max_attempts); - enqueue(&cfg, &nj).unwrap().expect("inserted"); - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed_typed(&cfg, &claimed, "boom", failure).unwrap(); - }; - - // Unrecoverable ⇒ terminal on first attempt, class persisted ⇒ counted. - let unrec = PipelineFailure::new(FailureCode::BudgetExhausted); - fail_one("c-unrec", 5, Some(&unrec)); - // Transient ⇒ terminal only because max_attempts=1 is exhausted; class - // persisted as 'transient' ⇒ NOT counted (it will be auto-requeued). - let trans = PipelineFailure::new(FailureCode::Transient); - fail_one("c-trans", 1, Some(&trans)); - // Untyped terminal (NULL class) ⇒ NOT counted. - fail_one("c-null", 1, None); - - assert_eq!( - count_by_status(&cfg, JobStatus::Failed).unwrap(), - 3, - "all three jobs are terminally failed" - ); - assert_eq!( - count_failed_unrecoverable(&cfg).unwrap(), - 1, - "only the unrecoverable one is counted" - ); - } - - /// T012: a **transient** typed failure keeps the existing - /// attempts-bounded retry path — the job bounces back to `ready` with a - /// future `available_at_ms` and does NOT set the typed columns (they are - /// only persisted on a terminal classified failure). - /// T028 (FR-011): `requeue_failed` flips terminal `failed` jobs back to - /// `ready` with a fresh attempt budget and the typed failure cleared, so - /// they re-run after the cause is fixed. Non-failed jobs are untouched. - #[test] - fn requeue_failed_resets_failed_jobs_only() { - use crate::openhuman::memory_tree::health::{FailureCode, PipelineFailure}; - let (_tmp, cfg) = test_config(); - - // Job A: drive to terminal failure (max_attempts=1, unrecoverable). - let mut a = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "a".into(), - }) - .unwrap(); - a.max_attempts = Some(1); - let id_a = enqueue(&cfg, &a).unwrap().unwrap(); - let claim_a = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed_typed( - &cfg, - &claim_a, - "Insufficient budget", - Some(&PipelineFailure::new(FailureCode::BudgetExhausted)), - ) - .unwrap(); - assert_eq!( - get_job(&cfg, &id_a).unwrap().unwrap().status, - JobStatus::Failed - ); - - // Job B: leave ready (untouched control). - let b = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "b".into(), - }) - .unwrap(); - let id_b = enqueue(&cfg, &b).unwrap().unwrap(); - - let requeued = requeue_failed(&cfg).unwrap(); - assert_eq!(requeued, 1, "only the failed job should requeue"); - - let row_a = get_job(&cfg, &id_a).unwrap().unwrap(); - assert_eq!(row_a.status, JobStatus::Ready); - assert_eq!(row_a.attempts, 0, "attempt budget reset"); - assert_eq!(row_a.failure_reason, None, "typed reason cleared"); - assert_eq!(row_a.failure_class, None); - assert_eq!(row_a.last_error, None); - assert!(row_a.completed_at_ms.is_none()); - - // B was already ready — still ready, not double-counted. - assert_eq!( - get_job(&cfg, &id_b).unwrap().unwrap().status, - JobStatus::Ready - ); - } - - /// The 3-hourly scheduler's self-heal must requeue jobs that failed - /// transiently (attempts exhausted on network-class errors — untyped - /// `last_error` only) while leaving `unrecoverable`-classified jobs - /// parked for the manual RPC, so a bad embeddings key can't produce a - /// retry-fail loop every 3 hours forever. - #[test] - fn requeue_transient_failed_skips_unrecoverable_jobs() { - use crate::openhuman::memory_tree::health::{FailureCode, PipelineFailure}; - let (_tmp, cfg) = test_config(); - - // Job A: terminal via exhausted retry budget, NO typed class - // (the shape an interrupted network leaves behind). - let mut a = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "a-transient".into(), - }) - .unwrap(); - a.max_attempts = Some(1); - let id_a = enqueue(&cfg, &a).unwrap().unwrap(); - let claim_a = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed(&cfg, &claim_a, "connection reset by peer").unwrap(); - assert_eq!( - get_job(&cfg, &id_a).unwrap().unwrap().status, - JobStatus::Failed - ); - - // Job B: terminal with an unrecoverable classification. - let mut b = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "b-unrecoverable".into(), - }) - .unwrap(); - b.max_attempts = Some(1); - let id_b = enqueue(&cfg, &b).unwrap().unwrap(); - let claim_b = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed_typed( - &cfg, - &claim_b, - "Insufficient budget", - Some(&PipelineFailure::new(FailureCode::BudgetExhausted)), - ) - .unwrap(); - assert_eq!( - get_job(&cfg, &id_b).unwrap().unwrap().status, - JobStatus::Failed - ); - - let requeued = requeue_transient_failed(&cfg).unwrap(); - assert_eq!(requeued, 1, "only the unclassified failure requeues"); - - let row_a = get_job(&cfg, &id_a).unwrap().unwrap(); - assert_eq!(row_a.status, JobStatus::Ready, "transient job re-runs"); - assert_eq!(row_a.attempts, 0); - - let row_b = get_job(&cfg, &id_b).unwrap().unwrap(); - assert_eq!( - row_b.status, - JobStatus::Failed, - "unrecoverable job stays parked for the manual retry RPC" - ); - assert_eq!(row_b.failure_class.as_deref(), Some("unrecoverable")); - } - - #[test] - fn mark_failed_typed_transient_still_retries() { - use crate::openhuman::memory_tree::health::{FailureCode, PipelineFailure}; - let (_tmp, cfg) = test_config(); - let mut nj = NewJob::extract_chunk(&ExtractChunkPayload { - chunk_id: "c-transient".into(), - }) - .unwrap(); - nj.max_attempts = Some(5); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let failure = PipelineFailure::new(FailureCode::Transient); - mark_failed_typed(&cfg, &claimed, "503 upstream", Some(&failure)).unwrap(); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!( - row.status, - JobStatus::Ready, - "transient failure must retry, not terminate" - ); - assert!(row.available_at_ms > Utc::now().timestamp_millis()); - assert_eq!(row.failure_reason, None, "typed cols unset on retry"); - assert_eq!(row.failure_class, None); - } - - #[test] - fn enqueue_dedupes_active_jobs() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c1".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id1 = enqueue(&cfg, &nj).unwrap(); - let id2 = enqueue(&cfg, &nj).unwrap(); - assert!(id1.is_some()); - assert!(id2.is_none(), "duplicate should be suppressed while ready"); - assert_eq!(count_total(&cfg).unwrap(), 1); - } - - #[test] - fn enqueue_after_done_creates_fresh_row() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c1".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id1 = enqueue(&cfg, &nj).unwrap().unwrap(); - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claimed.id, id1); - mark_done(&cfg, &claimed).unwrap(); - - // Now the dedupe key is free (partial index excludes 'done'). - let id2 = enqueue(&cfg, &nj).unwrap(); - assert!(id2.is_some()); - assert_ne!(id2.unwrap(), id1); - assert_eq!(count_total(&cfg).unwrap(), 2); - } - - #[test] - fn mark_failed_retries_then_terminates() { - let (_tmp, cfg) = test_config(); - let payload = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "c1".into(), - }, - target: AppendTarget::Source { - source_id: "slack:#x".into(), - }, - }; - let mut nj = NewJob::append_buffer(&payload).unwrap(); - nj.max_attempts = Some(2); - let id = enqueue(&cfg, &nj).unwrap().unwrap(); - - // Fail #1 — should bounce back to 'ready' with future available_at. - let attempt1 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed(&cfg, &attempt1, "boom").unwrap(); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Ready); - assert!(row.available_at_ms > Utc::now().timestamp_millis()); - assert_eq!(row.last_error.as_deref(), Some("boom")); - - // Force the row available again so the test doesn't hinge on sleep. - with_connection(&cfg, |c| { - c.execute( - "UPDATE mem_tree_jobs SET available_at_ms = 0 WHERE id = ?1", - params![id], - )?; - Ok(()) - }) - .unwrap(); - - // Fail #2 — exceeds max_attempts → terminal 'failed'. - let attempt2 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_failed(&cfg, &attempt2, "fatal").unwrap(); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Failed); - assert_eq!(row.last_error.as_deref(), Some("fatal")); - assert!(row.completed_at_ms.is_some()); - } - - /// `mark_failed` scrubs only the log emission, not the persisted - /// `last_error` column. A reader of `mem_tree_jobs` should still see - /// the full anyhow / handler-supplied chain so they can root-cause - /// the failure; the scrub is a defense-in-depth for the log sink only. - #[test] - fn mark_failed_persists_full_error_unredacted() { - let (_tmp, cfg) = test_config(); - let payload = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "c1".into(), - }, - target: AppendTarget::Source { - source_id: "slack:#x".into(), - }, - }; - let mut nj = NewJob::append_buffer(&payload).unwrap(); - nj.max_attempts = Some(1); - let id = enqueue(&cfg, &nj).unwrap().unwrap(); - - let claim = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let raw = "upstream returned 401: Bearer abc123.def-456 token rejected"; - mark_failed(&cfg, &claim, raw).unwrap(); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Failed); - // The persisted column keeps the full original — the scrub is - // applied at log emission only. - assert_eq!(row.last_error.as_deref(), Some(raw)); - } - - /// Same contract for `mark_deferred`: the log line is scrubbed in - /// `worker::run_once`, but the persisted `last_error` keeps the - /// full handler-supplied reason for diagnostics. - #[test] - fn mark_deferred_persists_full_reason_unredacted() { - let (_tmp, cfg) = test_config(); - let payload = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "c2".into(), - }, - target: AppendTarget::Source { - source_id: "slack:#y".into(), - }, - }; - let nj = NewJob::append_buffer(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().unwrap(); - - let claim = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let raw = "rate_limited by api_key=hunter2; retry after 30s"; - mark_deferred(&cfg, &claim, 0, raw).unwrap(); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.last_error.as_deref(), Some(raw)); - } - - #[test] - fn recover_stale_locks_resets_running_rows() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c1".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().unwrap(); - - // Claim with a lock window that's already in the past so recovery - // sees it as expired. - let _ = claim_next(&cfg, -1).unwrap().unwrap(); - - let recovered = recover_stale_locks(&cfg).unwrap(); - assert_eq!(recovered, 1); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Ready); - } - - #[test] - fn release_running_locks_resets_running_rows_regardless_of_lease() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-release".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().unwrap(); - - // Claim with a long, still-valid lease. `recover_stale_locks` would - // NOT touch this (not expired), but a graceful shutdown release must - // so a clean restart re-claims it immediately (bug-report-2026-05-26 I2). - let _ = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - - let released = release_running_locks(&cfg).unwrap(); - assert_eq!(released, 1); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Ready); - assert!(row.locked_until_ms.is_none()); - } - - /// Happy path: a non-stale settlement still succeeds after the claim-token - /// check is applied. Regression guard so the common case isn't broken. - #[test] - fn mark_done_succeeds_for_current_lessee() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-happy".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claimed.id, id); - - // Current lessee should settle successfully. - mark_done(&cfg, &claimed).unwrap(); - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Done); - assert!(row.completed_at_ms.is_some()); - assert!(row.locked_until_ms.is_none()); - } - - /// Stale-worker settlement is a no-op: after a lock expires and a second - /// worker re-claims the job, the first worker's `mark_done` must not - /// clobber the new lessee's row. - #[test] - fn stale_worker_settlement_is_noop() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-stale".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - // Worker A claims with a lock that's already expired (negative window). - let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap(); - assert_eq!(worker_a_job.id, id); - assert_eq!(worker_a_job.attempts, 1); - - // Simulate lease expiry: recover_stale_locks resets the row to 'ready'. - let recovered = recover_stale_locks(&cfg).unwrap(); - assert_eq!(recovered, 1); - - // Worker B claims the reset row — different lease token (attempts=2). - let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(worker_b_job.id, id); - assert_eq!(worker_b_job.attempts, 2); - - // Worker A (stale) tries to mark done using its old claim snapshot. - mark_done(&cfg, &worker_a_job).unwrap(); // must NOT return Err - - // Worker B's row must be untouched — still 'running' with attempts=2. - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!( - row.status, - JobStatus::Running, - "stale settlement must not clobber Worker B's running row" - ); - assert_eq!( - row.attempts, 2, - "attempts must still reflect Worker B's claim" - ); - } - - /// Same contract as stale_worker_settlement_is_noop but for mark_failed. - #[test] - fn stale_worker_mark_failed_is_noop() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-stale-fail".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - // Worker A claims with an already-expired lock. - let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap(); - assert_eq!(worker_a_job.attempts, 1); - - // Lease expires, recovered, Worker B re-claims. - let recovered = recover_stale_locks(&cfg).unwrap(); - assert_eq!(recovered, 1); - let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(worker_b_job.attempts, 2); - - // Worker A (stale) tries to record a failure — must be a no-op. - mark_failed(&cfg, &worker_a_job, "stale error").unwrap(); - - // Worker B's row must be untouched. - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!( - row.status, - JobStatus::Running, - "stale mark_failed must not clobber Worker B's running row" - ); - assert_ne!( - row.last_error.as_deref(), - Some("stale error"), - "stale error must not be written to the row" - ); - assert_eq!(row.attempts, 2); - } - - #[test] - fn backoff_grows_then_caps() { - assert_eq!(backoff_ms(1), 60_000); - assert_eq!(backoff_ms(2), 120_000); - assert_eq!(backoff_ms(3), 240_000); - // Eventually clamps at the cap. - assert_eq!(backoff_ms(20), RETRY_CAP_MS); - assert_eq!(backoff_ms(99), RETRY_CAP_MS); - } - - #[test] - fn count_by_status_reports_each_state() { - let (_tmp, cfg) = test_config(); - for i in 0..3 { - let p = ExtractChunkPayload { - chunk_id: format!("c{i}"), - }; - let nj = NewJob::extract_chunk(&p).unwrap(); - enqueue(&cfg, &nj).unwrap(); - } - assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 3); - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - mark_done(&cfg, &claimed).unwrap(); - assert_eq!(count_by_status(&cfg, JobStatus::Done).unwrap(), 1); - assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 2); - } - - /// Defer must NOT advance the failure-attempt counter. After a claim - /// (which bumps attempts to 1) and a deferral, the next claim should - /// see attempts==2 (just the second claim's bump) — proving the - /// transient deferral did not burn a slot from the row's failure - /// budget. - #[test] - fn mark_deferred_does_not_increment_attempts() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-defer-1".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claim1 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claim1.id, id); - assert_eq!(claim1.attempts, 1, "first claim bumps attempts to 1"); - - // Defer with a wake-up time already in the past so the next - // claim_next is immediately eligible without sleeping. - mark_deferred(&cfg, &claim1, 0, "rate_limited").unwrap(); - - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!(row.status, JobStatus::Ready); - assert_eq!( - row.attempts, 0, - "Defer must revert the claim's attempts bump (back to pre-claim 0)" - ); - assert_eq!(row.last_error.as_deref(), Some("rate_limited")); - assert_eq!(row.available_at_ms, 0); - assert!(row.locked_until_ms.is_none()); - assert!(row.started_at_ms.is_none()); - - let claim2 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claim2.id, id); - assert_eq!( - claim2.attempts, 1, - "second claim bumps attempts to 1 again (Defer didn't count)" - ); - } - - /// A row deferred to a future `until_ms` must not be claimable until - /// the system clock crosses that threshold. - #[test] - fn deferred_row_not_claimable_until_until_ms() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-defer-2".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let future_ms = Utc::now().timestamp_millis() + 60_000; - mark_deferred(&cfg, &claimed, future_ms, "warming_up").unwrap(); - - // Right now: not yet eligible. - let none = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap(); - assert!( - none.is_none(), - "deferred row must not be claimable before until_ms" - ); - - // Force the row available again (proxy for "wall clock advanced - // past until_ms") and confirm claim_next picks it up. - with_connection(&cfg, |c| { - c.execute( - "UPDATE mem_tree_jobs SET available_at_ms = 0 WHERE id = ?1", - params![id], - )?; - Ok(()) - }) - .unwrap(); - let claim2 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(claim2.id, id); - assert_eq!(claim2.attempts, 1, "Defer left attempts at pre-claim 0"); - } - - /// Three rows with three different terminal verbs: Done, Failed (with - /// retry left, so it bounces back to ready with bumped attempts), and - /// Defer. After processing, each row's terminal state must reflect its - /// settlement verb. Critically, the Defer row keeps its pre-claim - /// `attempts` value while the failed row bumps. - #[test] - fn mixed_outcomes_stress() { - let (_tmp, cfg) = test_config(); - let p_a = ExtractChunkPayload { - chunk_id: "c-mix-a".into(), - }; - let p_b = ExtractChunkPayload { - chunk_id: "c-mix-b".into(), - }; - let p_c = ExtractChunkPayload { - chunk_id: "c-mix-c".into(), - }; - let id_a = enqueue(&cfg, &NewJob::extract_chunk(&p_a).unwrap()) - .unwrap() - .unwrap(); - let id_b = enqueue(&cfg, &NewJob::extract_chunk(&p_b).unwrap()) - .unwrap() - .unwrap(); - let id_c = enqueue(&cfg, &NewJob::extract_chunk(&p_c).unwrap()) - .unwrap() - .unwrap(); - - // Claim the three rows in turn (FIFO within same kind/priority). - let claim_a = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let claim_b = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - let claim_c = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - // Sanity: three distinct ids covering A/B/C. - let mut got: Vec<&str> = vec![&claim_a.id, &claim_b.id, &claim_c.id]; - got.sort(); - let mut want = vec![id_a.as_str(), id_b.as_str(), id_c.as_str()]; - want.sort(); - assert_eq!(got, want); - - let until_ms = Utc::now().timestamp_millis() + 30_000; - - // Settle: A=done, B=err (retry path), C=defer. - mark_done(&cfg, &claim_a).unwrap(); - mark_failed(&cfg, &claim_b, "transient_error").unwrap(); - mark_deferred(&cfg, &claim_c, until_ms, "rate_limited").unwrap(); - - // A: terminal done. - let row_a = get_job(&cfg, &id_a).unwrap().unwrap(); - assert_eq!(row_a.status, JobStatus::Done); - assert!(row_a.completed_at_ms.is_some()); - - // B: retry path — back to ready with bumped attempts (1) and a - // future available_at_ms from exponential backoff. - let row_b = get_job(&cfg, &id_b).unwrap().unwrap(); - assert_eq!(row_b.status, JobStatus::Ready); - assert_eq!( - row_b.attempts, 1, - "Err settlement keeps the claim's attempts bump" - ); - assert!(row_b.available_at_ms > Utc::now().timestamp_millis()); - assert_eq!(row_b.last_error.as_deref(), Some("transient_error")); - - // C: deferred — back to ready with attempts reverted and - // available_at_ms == until_ms. - let row_c = get_job(&cfg, &id_c).unwrap().unwrap(); - assert_eq!(row_c.status, JobStatus::Ready); - assert_eq!( - row_c.attempts, 0, - "Defer settlement reverts the claim's attempts bump" - ); - assert_eq!(row_c.available_at_ms, until_ms); - assert_eq!(row_c.last_error.as_deref(), Some("rate_limited")); - assert!(row_c.locked_until_ms.is_none()); - assert!(row_c.started_at_ms.is_none()); - } - - /// Stale-worker `mark_deferred` must be a no-op — same lease-token - /// gating as `mark_done` / `mark_failed`. After Worker A's claim - /// expires and Worker B re-claims, Worker A's stale Defer must not - /// clobber B's running row. - #[test] - fn mark_deferred_stale_lease_is_noop() { - let (_tmp, cfg) = test_config(); - let payload = ExtractChunkPayload { - chunk_id: "c-defer-stale".into(), - }; - let nj = NewJob::extract_chunk(&payload).unwrap(); - let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); - - // Worker A claims with already-expired lock. - let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap(); - assert_eq!(worker_a_job.attempts, 1); - - // Lease expires; Worker B re-claims. - let recovered = recover_stale_locks(&cfg).unwrap(); - assert_eq!(recovered, 1); - let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); - assert_eq!(worker_b_job.attempts, 2); - - // Worker A (stale) tries to defer using its old snapshot — must - // be a silent no-op. - let stale_until_ms = Utc::now().timestamp_millis() + 999_000; - mark_deferred(&cfg, &worker_a_job, stale_until_ms, "stale_defer").unwrap(); - - // Worker B's row must be untouched: still 'running' with - // attempts=2 and no stale_defer side effect. - let row = get_job(&cfg, &id).unwrap().unwrap(); - assert_eq!( - row.status, - JobStatus::Running, - "stale mark_deferred must not clobber Worker B's running row" - ); - assert_eq!(row.attempts, 2); - assert_ne!( - row.last_error.as_deref(), - Some("stale_defer"), - "stale defer reason must not be persisted" - ); - assert_ne!( - row.available_at_ms, stale_until_ms, - "stale until_ms must not be persisted" - ); - } + tinycortex::memory::queue::get_job(&memory_config(config), id) } diff --git a/src/openhuman/memory_queue/types.rs b/src/openhuman/memory_queue/types.rs index ba82c17dc..fc8775074 100644 --- a/src/openhuman/memory_queue/types.rs +++ b/src/openhuman/memory_queue/types.rs @@ -1,643 +1,7 @@ -//! Job types for the async memory-tree pipeline. -//! -//! Each `Job` row in `mem_tree_jobs` stores its discriminator as a string -//! `kind` plus a JSON-encoded `payload`. The strongly-typed payload structs -//! below own (de)serialisation; handlers parse the payload by branching on -//! [`JobKind`] and calling the matching `from_payload_json`. +//! Queue wire types owned by tinycortex. -use anyhow::{anyhow, Result}; -use serde::{Deserialize, Serialize}; - -/// Discriminator persisted in `mem_tree_jobs.kind`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum JobKind { - /// Run LLM entity extraction over a single chunk and decide admission. - ExtractChunk, - /// Push an admitted chunk into a tree's L0 buffer. - AppendBuffer, - /// Seal exactly one buffer level; cascades enqueue a follow-up. - Seal, - /// Walk stale buffers and enqueue `Seal` jobs for any over the age cap. - FlushStale, - /// #1574 §6: re-embed a bounded batch of chunks/summaries that lack a - /// vector at the active embedding signature (post model-switch, or the - /// §7 dim-mismatch slice), then self-continue until none remain. - ReembedBackfill, - /// Build one document version's per-doc subtree (Notion) and merge its - /// doc-root into the connection tree. Replaces the per-chunk - /// extract→append_buffer tree path for document sources that opt into - /// per-document rollup + versioning. - SealDocument, -} - -impl JobKind { - /// Snake-case wire string written to `mem_tree_jobs.kind`. - pub fn as_str(&self) -> &'static str { - match self { - JobKind::ExtractChunk => "extract_chunk", - JobKind::AppendBuffer => "append_buffer", - JobKind::Seal => "seal", - JobKind::FlushStale => "flush_stale", - JobKind::ReembedBackfill => "reembed_backfill", - JobKind::SealDocument => "seal_document", - } - } - - /// Inverse of [`Self::as_str`]; returns `Err` for unknown kinds. - pub fn parse(s: &str) -> Result { - Ok(match s { - "extract_chunk" => JobKind::ExtractChunk, - "append_buffer" => JobKind::AppendBuffer, - "seal" => JobKind::Seal, - "flush_stale" => JobKind::FlushStale, - // Legacy kinds from the removed global/topic trees. Tolerated on - // parse so a queue row left over from before the removal is - // recognised and can be drained/discarded rather than crashing - // the worker loop; the startup migration purges them. - "topic_route" | "digest_daily" => { - return Err(anyhow!( - "retired JobKind '{s}' (global/topic trees removed)" - )) - } - "reembed_backfill" => JobKind::ReembedBackfill, - "seal_document" => JobKind::SealDocument, - other => return Err(anyhow!("unknown JobKind '{other}'")), - }) - } - - /// True when handling this kind should hold a slot from the global - /// LLM concurrency semaphore. - pub fn is_llm_bound(&self) -> bool { - matches!( - self, - JobKind::ExtractChunk - | JobKind::Seal - | JobKind::ReembedBackfill - | JobKind::SealDocument - ) - } -} - -/// Outcome of a successful handler run. Workers translate this into a -/// queue settlement: `Done` finalises the row, while `Defer` puts it back -/// to `ready` with `available_at_ms = until_ms` and **does not** count -/// toward the failure-attempt budget. -/// -/// `Defer` exists so a handler that is transiently unable to make -/// progress (cloud rate-limited, dependency unavailable, model warming -/// up) can re-queue its job with a wake-up time without marking it -/// failed. Handlers should still surface real errors via `Err(_)` — that -/// path runs the existing exponential-backoff retry logic which **does** -/// burn the failure budget. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum JobOutcome { - /// Handler ran to completion. Row is settled as `done`. - Done, - /// Handler chose not to make progress yet. Row is rescheduled to - /// `available_at_ms = until_ms` (UTC milliseconds) with `attempts` - /// reverted to its pre-claim value so the failure budget is not - /// touched. `reason` is recorded in `last_error` for visibility. - Defer { until_ms: i64, reason: String }, -} - -/// Lifecycle states persisted on `mem_tree_jobs.status`. Workers transition -/// `ready → running → done|failed`. `Cancelled` is reserved for explicit -/// admin actions (none surfaced yet). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum JobStatus { - Ready, - Running, - Done, - Failed, - Cancelled, -} - -impl JobStatus { - /// Snake-case wire string written to `mem_tree_jobs.status`. - pub fn as_str(&self) -> &'static str { - match self { - JobStatus::Ready => "ready", - JobStatus::Running => "running", - JobStatus::Done => "done", - JobStatus::Failed => "failed", - JobStatus::Cancelled => "cancelled", - } - } - - /// Inverse of [`Self::as_str`]; returns `Err` for unknown values. - pub fn parse(s: &str) -> Result { - Ok(match s { - "ready" => JobStatus::Ready, - "running" => JobStatus::Running, - "done" => JobStatus::Done, - "failed" => JobStatus::Failed, - "cancelled" => JobStatus::Cancelled, - other => return Err(anyhow!("unknown JobStatus '{other}'")), - }) - } - - /// True for `Done`, `Failed`, `Cancelled` — i.e. no further worker - /// transitions are expected. - pub fn is_terminal(&self) -> bool { - matches!( - self, - JobStatus::Done | JobStatus::Failed | JobStatus::Cancelled - ) - } -} - -// ── Payloads ─────────────────────────────────────────────────────────────── - -/// Reference to either a leaf chunk or a sealed summary node. Used by -/// payloads that route content through the pipeline regardless of which -/// kind of source produced it. -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum NodeRef { - Leaf { chunk_id: String }, - Summary { summary_id: String }, -} - -impl NodeRef { - /// Stringified id with kind prefix (`leaf:` or `summary:`), suitable - /// for dedupe-key composition. - pub fn dedupe_fragment(&self) -> String { - match self { - NodeRef::Leaf { chunk_id } => format!("leaf:{chunk_id}"), - NodeRef::Summary { summary_id } => format!("summary:{summary_id}"), - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ExtractChunkPayload { - pub chunk_id: String, -} - -impl ExtractChunkPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key` so a partial - /// unique index can suppress in-flight duplicates. - pub fn dedupe_key(&self) -> String { - format!("extract:{}", self.chunk_id) - } -} - -/// Where an `AppendBuffer` job should land its node. Source-tree appends -/// are keyed by `source_id`; topic-tree appends are keyed by `tree_id` -/// because there can be many topic trees per node. -#[derive(Clone, Debug, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum AppendTarget { - Source { source_id: String }, - Topic { tree_id: String }, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct AppendBufferPayload { - pub node: NodeRef, - pub target: AppendTarget, -} - -impl AppendBufferPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key` so a partial - /// unique index can suppress in-flight duplicates. - pub fn dedupe_key(&self) -> String { - let node_part = self.node.dedupe_fragment(); - match &self.target { - AppendTarget::Source { source_id } => { - format!("append:source:{source_id}:{node_part}") - } - AppendTarget::Topic { tree_id } => { - format!("append:topic:{tree_id}:{node_part}") - } - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct SealPayload { - pub tree_id: String, - pub level: u32, - /// When `Some`, the seal handler bypasses the buffer-budget check and - /// force-seals — used by the time-based flush path. The wall-clock is - /// passed through so the seal stamps a deterministic `sealed_at`. - pub force_now_ms: Option, -} - -impl SealPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key` so a partial - /// unique index can suppress in-flight duplicates. - pub fn dedupe_key(&self) -> String { - // Active seal-job uniqueness is enforced per (tree, level): a seal - // already in flight suppresses duplicate enqueues. Once the job - // completes the partial index releases the key, so the next time - // the buffer crosses its gate a fresh seal can be enqueued. - format!("seal:{}:{}", self.tree_id, self.level) - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, Default)] -pub struct FlushStalePayload { - /// Override the configured `DEFAULT_FLUSH_AGE_SECS`. Optional so the - /// scheduler can enqueue with `None` and let the handler use the - /// configured default. - pub max_age_secs: Option, -} - -impl FlushStalePayload { - /// Dedupe key scoped to a 3-hour UTC block (`hour_block = hour / 3`, - /// 0..=7) so flush_stale runs up to 8× per day. Without this, - /// low-volume sources wait a full day between seal opportunities. - /// - /// Pure: both `date_iso` and `hour_block` are supplied by the caller - /// from a single `Utc::now()` reading, which keeps the key - /// deterministic in tests and avoids a 3-hour-boundary race where - /// the caller's `today_iso` could disagree with a second - /// `Utc::now()` taken inside this function. - pub fn dedupe_key(&self, date_iso: &str, hour_block: u32) -> String { - format!("flush_stale:{date_iso}-h{hour_block}") - } -} - -/// #1574 §6 re-embed backfill. One chain per embedding signature: the -/// `dedupe_key` is the signature, so re-triggering while a chain is -/// in-flight is correctly suppressed (exactly one backfill per space). -/// The handler self-continues via `JobOutcome::Defer` (reschedules this -/// same row) rather than re-enqueuing, so the fixed dedupe key is safe. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ReembedBackfillPayload { - /// The embedding signature this chain re-embeds under. If the active - /// signature has since changed, the handler treats this job as stale - /// and finishes (a fresh chain for the new signature takes over). - pub signature: String, -} - -impl ReembedBackfillPayload { - /// Stable dedupe key — one in-flight backfill chain per signature. - pub fn dedupe_key(&self) -> String { - format!("reembed_backfill:{}", self.signature) - } -} - -/// Build (or re-build for a new version) one document's per-doc subtree and -/// merge its doc-root into the connection tree. Carries the full leaf chunk -/// set for the version so the seal handler can run the per-document -/// side-cascade in one shot (see -/// [`crate::openhuman::memory_tree::tree::seal_document_subtree`]). -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct SealDocumentPayload { - /// Connection-level source-tree scope, e.g. `notion:{connection_id}`. - /// All of a connection's documents share this one tree. - pub tree_scope: String, - /// Document identity (the chunk `source_id`, e.g. - /// `notion:{connection_id}:{page_id}`). - pub doc_id: String, - /// Document version as epoch-ms (`last_edited_time`). `None` for sources - /// that don't carry a version. - pub version_ms: Option, - /// Leaf chunk ids for this version, in document order. - pub chunk_ids: Vec, -} - -impl SealDocumentPayload { - /// Stable dedupe key — one in-flight seal per (doc, version). A new - /// version gets a distinct key so it isn't suppressed by an older one. - pub fn dedupe_key(&self) -> String { - match self.version_ms { - Some(v) => format!("seal_doc:{}@{}", self.doc_id, v), - None => format!("seal_doc:{}", self.doc_id), - } - } -} - -/// One row in `mem_tree_jobs`. `payload_json` is left as a raw string so -/// callers parse it lazily based on `kind`. -#[derive(Clone, Debug)] -pub struct Job { - pub id: String, - pub kind: JobKind, - pub payload_json: String, - pub dedupe_key: Option, - pub status: JobStatus, - pub attempts: u32, - pub max_attempts: u32, - pub available_at_ms: i64, - pub locked_until_ms: Option, - pub last_error: Option, - /// Typed failure code (e.g. "budget_exhausted") set when a job is marked - /// `failed` with a classified reason; `None` otherwise. Distinct from the - /// freeform `last_error` — this is the machine-readable cause the - /// status/doctor surface renders. - pub failure_reason: Option, - /// Failure class ("transient" | "unrecoverable") paired with - /// `failure_reason`; `None` until a classified failure is recorded. - pub failure_class: Option, - pub created_at_ms: i64, - pub started_at_ms: Option, - pub completed_at_ms: Option, -} - -/// Caller-side bundle for `enqueue` — `Job` minus the persistence-only -/// columns. Keeps producers from having to mint timestamps and ids by hand. -#[derive(Clone, Debug)] -pub struct NewJob { - pub kind: JobKind, - pub payload_json: String, - pub dedupe_key: Option, - /// `None` means "available immediately." Set this for delayed jobs - /// (retries, scheduled work). - pub available_at_ms: Option, - pub max_attempts: Option, -} - -impl NewJob { - /// Build an [`JobKind::ExtractChunk`] enqueue request. - pub fn extract_chunk(p: &ExtractChunkPayload) -> Result { - Ok(Self { - kind: JobKind::ExtractChunk, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key()), - available_at_ms: None, - max_attempts: None, - }) - } - - /// Build an [`JobKind::AppendBuffer`] enqueue request. - pub fn append_buffer(p: &AppendBufferPayload) -> Result { - Ok(Self { - kind: JobKind::AppendBuffer, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key()), - available_at_ms: None, - max_attempts: None, - }) - } - - /// Build an [`JobKind::Seal`] enqueue request. - pub fn seal(p: &SealPayload) -> Result { - Ok(Self { - kind: JobKind::Seal, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key()), - available_at_ms: None, - max_attempts: None, - }) - } - - /// Build a [`JobKind::FlushStale`] enqueue request scoped to a - /// 3-hour UTC block. Callers compute `date_iso` and `hour_block` - /// from a single `Utc::now()` reading so the dedupe key is - /// boundary-safe; see [`FlushStalePayload::dedupe_key`]. - pub fn flush_stale(p: &FlushStalePayload, date_iso: &str, hour_block: u32) -> Result { - Ok(Self { - kind: JobKind::FlushStale, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key(date_iso, hour_block)), - available_at_ms: None, - max_attempts: None, - }) - } - - /// Build a [`JobKind::ReembedBackfill`] enqueue request (#1574 §6). - pub fn reembed_backfill(p: &ReembedBackfillPayload) -> Result { - Ok(Self { - kind: JobKind::ReembedBackfill, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key()), - available_at_ms: None, - max_attempts: Some(3), - }) - } - - /// Build a [`JobKind::SealDocument`] enqueue request. - pub fn seal_document(p: &SealDocumentPayload) -> Result { - Ok(Self { - kind: JobKind::SealDocument, - payload_json: serde_json::to_string(p)?, - dedupe_key: Some(p.dedupe_key()), - available_at_ms: None, - max_attempts: None, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn job_kind_roundtrip() { - for k in [ - JobKind::ExtractChunk, - JobKind::AppendBuffer, - JobKind::Seal, - JobKind::FlushStale, - JobKind::ReembedBackfill, - JobKind::SealDocument, - ] { - assert_eq!(JobKind::parse(k.as_str()).unwrap(), k); - } - // Retired kinds parse to an error (global/topic trees removed). - assert!(JobKind::parse("topic_route").is_err()); - assert!(JobKind::parse("digest_daily").is_err()); - } - - #[test] - fn seal_document_dedupe_key_is_per_version() { - let v1 = SealDocumentPayload { - tree_scope: "notion:conn1".into(), - doc_id: "notion:conn1:pageA".into(), - version_ms: Some(1717000000000), - chunk_ids: vec!["c0".into()], - }; - let v2 = SealDocumentPayload { - version_ms: Some(1717500000000), - ..v1.clone() - }; - // Distinct versions of the same doc get distinct keys, so a newer - // revision is never suppressed by an in-flight older one. - assert_ne!(v1.dedupe_key(), v2.dedupe_key()); - assert_eq!(v1.dedupe_key(), "seal_doc:notion:conn1:pageA@1717000000000"); - // Unversioned falls back to the bare doc id. - let unversioned = SealDocumentPayload { - version_ms: None, - ..v1.clone() - }; - assert_eq!(unversioned.dedupe_key(), "seal_doc:notion:conn1:pageA"); - } - - #[test] - fn seal_document_roundtrips_through_newjob() { - let p = SealDocumentPayload { - tree_scope: "notion:conn1".into(), - doc_id: "notion:conn1:pageA".into(), - version_ms: Some(42), - chunk_ids: vec!["c0".into(), "c1".into()], - }; - let job = NewJob::seal_document(&p).unwrap(); - assert_eq!(job.kind, JobKind::SealDocument); - let back: SealDocumentPayload = serde_json::from_str(&job.payload_json).unwrap(); - assert_eq!(back.chunk_ids, vec!["c0".to_string(), "c1".to_string()]); - assert_eq!(back.version_ms, Some(42)); - } - - #[test] - fn job_status_terminality() { - assert!(!JobStatus::Ready.is_terminal()); - assert!(!JobStatus::Running.is_terminal()); - assert!(JobStatus::Done.is_terminal()); - assert!(JobStatus::Failed.is_terminal()); - assert!(JobStatus::Cancelled.is_terminal()); - } - - #[test] - fn dedupe_keys_distinguish_targets() { - let p_src = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "c1".into(), - }, - target: AppendTarget::Source { - source_id: "slack:#eng".into(), - }, - }; - let p_topic = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "c1".into(), - }, - target: AppendTarget::Topic { - tree_id: "topic:abc".into(), - }, - }; - assert_ne!(p_src.dedupe_key(), p_topic.dedupe_key()); - } - - #[test] - fn dedupe_keys_distinguish_node_kinds() { - let p_leaf = AppendBufferPayload { - node: NodeRef::Leaf { - chunk_id: "x".into(), - }, - target: AppendTarget::Topic { - tree_id: "t".into(), - }, - }; - let p_summary = AppendBufferPayload { - node: NodeRef::Summary { - summary_id: "x".into(), - }, - target: AppendTarget::Topic { - tree_id: "t".into(), - }, - }; - assert_ne!(p_leaf.dedupe_key(), p_summary.dedupe_key()); - } - - #[test] - fn flush_stale_dedupe_key_is_pure_and_per_3h_block() { - let p = FlushStalePayload::default(); - // Same (date, block) → same key. - assert_eq!(p.dedupe_key("2026-05-19", 2), p.dedupe_key("2026-05-19", 2)); - // Different block within same day → distinct keys (8 buckets/day). - assert_ne!(p.dedupe_key("2026-05-19", 2), p.dedupe_key("2026-05-19", 3)); - // Different day, same block → distinct keys. - assert_ne!(p.dedupe_key("2026-05-19", 2), p.dedupe_key("2026-05-20", 2)); - // Shape sanity. - assert_eq!(p.dedupe_key("2026-05-19", 0), "flush_stale:2026-05-19-h0"); - assert_eq!(p.dedupe_key("2026-05-19", 7), "flush_stale:2026-05-19-h7"); - } - - #[test] - fn llm_bound_kinds() { - assert!(JobKind::ExtractChunk.is_llm_bound()); - assert!(JobKind::Seal.is_llm_bound()); - assert!(JobKind::SealDocument.is_llm_bound()); - assert!(JobKind::ReembedBackfill.is_llm_bound()); - assert!(!JobKind::AppendBuffer.is_llm_bound()); - assert!(!JobKind::FlushStale.is_llm_bound()); - } - - #[test] - fn node_ref_serializes_with_kind_tag() { - let leaf = NodeRef::Leaf { - chunk_id: "x".into(), - }; - let s = serde_json::to_string(&leaf).unwrap(); - assert!(s.contains("\"kind\":\"leaf\"")); - let back: NodeRef = serde_json::from_str(&s).unwrap(); - assert_eq!(back, leaf); - } - - #[test] - fn append_target_serializes_with_kind_tag() { - let p = AppendTarget::Source { - source_id: "x".into(), - }; - let s = serde_json::to_string(&p).unwrap(); - assert!(s.contains("\"kind\":\"source\"")); - assert!(s.contains("\"source_id\":\"x\"")); - let back: AppendTarget = serde_json::from_str(&s).unwrap(); - match back { - AppendTarget::Source { source_id } => assert_eq!(source_id, "x"), - _ => panic!("wrong variant"), - } - } - - #[test] - fn new_job_extract_chunk_builder_sets_kind_payload_and_dedupe_key() { - let payload = ExtractChunkPayload { - chunk_id: "chunk-123".into(), - }; - let job = NewJob::extract_chunk(&payload).unwrap(); - assert_eq!(job.kind, JobKind::ExtractChunk); - assert_eq!(job.dedupe_key.as_deref(), Some("extract:chunk-123")); - assert_eq!(job.available_at_ms, None); - assert_eq!(job.max_attempts, None); - let roundtrip: ExtractChunkPayload = serde_json::from_str(&job.payload_json).unwrap(); - assert_eq!(roundtrip.chunk_id, "chunk-123"); - } - - #[test] - fn new_job_append_buffer_builder_uses_payload_dedupe_key() { - let payload = AppendBufferPayload { - node: NodeRef::Summary { - summary_id: "summary-9".into(), - }, - target: AppendTarget::Topic { - tree_id: "topic:ops".into(), - }, - }; - let job = NewJob::append_buffer(&payload).unwrap(); - assert_eq!(job.kind, JobKind::AppendBuffer); - assert_eq!( - job.dedupe_key.as_deref(), - Some("append:topic:topic:ops:summary:summary-9") - ); - let roundtrip: AppendBufferPayload = serde_json::from_str(&job.payload_json).unwrap(); - assert_eq!(roundtrip.dedupe_key(), payload.dedupe_key()); - } - - #[test] - fn new_job_flush_stale_builder_uses_supplied_time_bucket() { - let payload = FlushStalePayload { - max_age_secs: Some(600), - }; - let job = NewJob::flush_stale(&payload, "2026-05-24", 4).unwrap(); - assert_eq!(job.kind, JobKind::FlushStale); - assert_eq!(job.dedupe_key.as_deref(), Some("flush_stale:2026-05-24-h4")); - let roundtrip: FlushStalePayload = serde_json::from_str(&job.payload_json).unwrap(); - assert_eq!(roundtrip.max_age_secs, Some(600)); - } - - #[test] - fn new_job_reembed_backfill_builder_is_one_chain_per_signature() { - let payload = ReembedBackfillPayload { - signature: "embed-v2".into(), - }; - let job = NewJob::reembed_backfill(&payload).unwrap(); - assert_eq!(job.kind, JobKind::ReembedBackfill); - assert_eq!(job.dedupe_key.as_deref(), Some("reembed_backfill:embed-v2")); - assert_eq!(job.max_attempts, Some(3)); - let roundtrip: ReembedBackfillPayload = serde_json::from_str(&job.payload_json).unwrap(); - assert_eq!(roundtrip.signature, "embed-v2"); - } -} +pub use tinycortex::memory::queue::{ + AppendBufferPayload, AppendTarget, ExtractChunkPayload, FlushStalePayload, Job, JobFailure, + JobKind, JobOutcome, JobStatus, NewJob, NodeRef, ReembedBackfillPayload, SealDocumentPayload, + SealPayload, +}; diff --git a/src/openhuman/memory_sources/readers/conversation.rs b/src/openhuman/memory_sources/readers/conversation.rs index 1d8ea92f8..3c270b171 100644 --- a/src/openhuman/memory_sources/readers/conversation.rs +++ b/src/openhuman/memory_sources/readers/conversation.rs @@ -1,18 +1,13 @@ -//! Conversation source reader. -//! -//! Treats every agent conversation as a memory source item. When synced, -//! each conversation's messages are stored as durable memory alongside -//! other sources like GitHub, integrations, etc. +//! Product `Config` adapter for the tinycortex conversation reader. use async_trait::async_trait; use crate::openhuman::config::Config; +use crate::openhuman::memory_sources::readers::SourceReader; use crate::openhuman::memory_sources::types::{ - ContentType, MemorySourceEntry, SourceContent, SourceItem, SourceKind, + MemorySourceEntry, SourceContent, SourceItem, SourceKind, }; -use super::SourceReader; - pub struct ConversationReader; #[async_trait] @@ -23,356 +18,31 @@ impl SourceReader for ConversationReader { async fn list_items( &self, - _source: &MemorySourceEntry, + source: &MemorySourceEntry, config: &Config, ) -> Result, String> { - tracing::debug!("[memory_sources:conversation] list_items"); - - let threads_dir = config.workspace_dir.join("threads"); - if !threads_dir.exists() { - return Ok(Vec::new()); - } - - let mut items = Vec::new(); - let mut entries = tokio::fs::read_dir(&threads_dir) - .await - .map_err(|e| format!("failed to read threads dir: {e}"))?; - - loop { - match entries.next_entry().await { - Ok(Some(entry)) => { - let path = entry.path(); - if path.extension().and_then(|e| e.to_str()) != Some("json") { - continue; - } - let id = path - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or_default() - .to_string(); - - let modified_ms = entry - .metadata() - .await - .ok() - .and_then(|m| m.modified().ok()) - .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|d| d.as_millis() as i64); - - items.push(SourceItem { - id, - title: path - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("conversation") - .to_string(), - updated_at_ms: modified_ms, - }); - } - Ok(None) => break, - Err(err) => { - tracing::warn!( - error = %err, - "[memory_sources:conversation] failed to read directory entry, skipping" - ); - continue; - } - } - } - - tracing::debug!( - count = items.len(), - "[memory_sources:conversation] found threads" - ); - - Ok(items) + tinycortex::memory::sources::SourceReader::list_items( + &tinycortex::memory::sources::readers::conversation::ConversationReader, + source, + &crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()), + ) + .await + .map_err(|error| error.to_string()) } async fn read_item( &self, - _source: &MemorySourceEntry, + source: &MemorySourceEntry, item_id: &str, config: &Config, ) -> Result { - tracing::debug!( - item_id = %item_id, - "[memory_sources:conversation] read_item" - ); - - // Validate item_id to prevent path traversal - if item_id.contains("..") || item_id.contains('/') || item_id.contains('\\') { - return Err("invalid item_id: path traversal denied".to_string()); - } - - let threads_dir = config.workspace_dir.join("threads"); - let thread_path = threads_dir.join(format!("{item_id}.json")); - - // Canonicalize and verify containment within threads directory - if !thread_path.exists() { - return Err(format!("thread '{item_id}' not found")); - } - let canonical_base = std::fs::canonicalize(&threads_dir) - .map_err(|e| format!("cannot resolve threads dir: {e}"))?; - let canonical_file = std::fs::canonicalize(&thread_path) - .map_err(|e| format!("cannot resolve thread path: {e}"))?; - if !canonical_file.starts_with(&canonical_base) { - return Err("path traversal denied".to_string()); - } - - let raw = tokio::fs::read_to_string(&canonical_file) - .await - .map_err(|e| format!("failed to read thread file: {e}"))?; - - let parsed: serde_json::Value = - serde_json::from_str(&raw).map_err(|e| format!("failed to parse thread JSON: {e}"))?; - - let title = parsed - .get("title") - .and_then(|v| v.as_str()) - .unwrap_or(item_id) - .to_string(); - - let body = format_thread_as_markdown(&parsed); - - Ok(SourceContent { - id: item_id.to_string(), - title, - body, - content_type: ContentType::Markdown, - metadata: serde_json::json!({ - "source_type": "conversation", - "thread_id": item_id, - }), - }) - } -} - -fn format_thread_as_markdown(thread: &serde_json::Value) -> String { - let mut out = String::new(); - - if let Some(title) = thread.get("title").and_then(|v| v.as_str()) { - out.push_str(&format!("# {title}\n\n")); - } - - let messages = thread - .get("messages") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - for msg in &messages { - let role = msg - .get("role") - .and_then(|v| v.as_str()) - .unwrap_or("unknown"); - let content = msg.get("content").and_then(|v| v.as_str()).unwrap_or(""); - - if content.is_empty() { - continue; - } - - out.push_str(&format!("**{role}**: {content}\n\n")); - } - - out -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - use tempfile::tempdir; - - #[test] - fn format_thread_produces_markdown() { - let thread = serde_json::json!({ - "title": "Test chat", - "messages": [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"}, - ] - }); - let md = format_thread_as_markdown(&thread); - assert!(md.contains("# Test chat")); - assert!(md.contains("**user**: Hello")); - assert!(md.contains("**assistant**: Hi there!")); - } - - #[test] - fn format_thread_skips_empty_content() { - let thread = serde_json::json!({ - "title": "Sparse", - "messages": [ - {"role": "user", "content": ""}, - {"role": "assistant", "content": "Reply"}, - {"role": "user", "content": ""}, - ] - }); - let md = format_thread_as_markdown(&thread); - assert!(!md.contains("**user**:")); - assert!(md.contains("**assistant**: Reply")); - } - - #[test] - fn format_thread_handles_missing_title() { - let thread = serde_json::json!({ - "messages": [{"role": "user", "content": "Hi"}] - }); - let md = format_thread_as_markdown(&thread); - assert!(!md.starts_with('#')); - assert!(md.contains("**user**: Hi")); - } - - #[test] - fn format_thread_handles_no_messages() { - let thread = serde_json::json!({"title": "Empty"}); - let md = format_thread_as_markdown(&thread); - assert!(md.contains("# Empty")); - assert_eq!(md.trim(), "# Empty"); - } - - #[tokio::test] - async fn list_items_returns_empty_when_no_threads_dir() { - let tmp = tempdir().unwrap(); - let mut config = Config::default(); - config.workspace_dir = tmp.path().to_path_buf(); - - let source = conversation_source(); - let reader = ConversationReader; - let items = reader.list_items(&source, &config).await.unwrap(); - assert!(items.is_empty()); - } - - #[tokio::test] - async fn list_items_finds_json_thread_files() { - let tmp = tempdir().unwrap(); - let threads_dir = tmp.path().join("threads"); - fs::create_dir_all(&threads_dir).unwrap(); - - fs::write( - threads_dir.join("thread_abc.json"), - r#"{"title":"Chat 1","messages":[]}"#, + tinycortex::memory::sources::SourceReader::read_item( + &tinycortex::memory::sources::readers::conversation::ConversationReader, + source, + item_id, + &crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()), ) - .unwrap(); - fs::write( - threads_dir.join("thread_def.json"), - r#"{"title":"Chat 2","messages":[]}"#, - ) - .unwrap(); - // Non-json file should be ignored - fs::write(threads_dir.join("notes.txt"), "ignored").unwrap(); - - let mut config = Config::default(); - config.workspace_dir = tmp.path().to_path_buf(); - - let source = conversation_source(); - let reader = ConversationReader; - let items = reader.list_items(&source, &config).await.unwrap(); - assert_eq!(items.len(), 2); - - let ids: Vec<&str> = items.iter().map(|i| i.id.as_str()).collect(); - assert!(ids.contains(&"thread_abc")); - assert!(ids.contains(&"thread_def")); - } - - #[tokio::test] - async fn read_item_returns_formatted_content() { - let tmp = tempdir().unwrap(); - let threads_dir = tmp.path().join("threads"); - fs::create_dir_all(&threads_dir).unwrap(); - - let thread_json = serde_json::json!({ - "title": "Test Conversation", - "messages": [ - {"role": "user", "content": "What is 2+2?"}, - {"role": "assistant", "content": "4"}, - ] - }); - fs::write( - threads_dir.join("conv_123.json"), - serde_json::to_string(&thread_json).unwrap(), - ) - .unwrap(); - - let mut config = Config::default(); - config.workspace_dir = tmp.path().to_path_buf(); - - let source = conversation_source(); - let reader = ConversationReader; - let content = reader - .read_item(&source, "conv_123", &config) - .await - .unwrap(); - - assert_eq!(content.id, "conv_123"); - assert_eq!(content.title, "Test Conversation"); - assert_eq!(content.content_type, ContentType::Markdown); - assert!(content.body.contains("**user**: What is 2+2?")); - assert!(content.body.contains("**assistant**: 4")); - } - - #[tokio::test] - async fn read_item_returns_error_for_missing_thread() { - let tmp = tempdir().unwrap(); - let threads_dir = tmp.path().join("threads"); - fs::create_dir_all(&threads_dir).unwrap(); - - let mut config = Config::default(); - config.workspace_dir = tmp.path().to_path_buf(); - - let source = conversation_source(); - let reader = ConversationReader; - let result = reader.read_item(&source, "nonexistent", &config).await; - assert!(result.is_err()); - assert!(result.unwrap_err().contains("not found")); - } - - #[tokio::test] - async fn read_item_rejects_path_traversal() { - let tmp = tempdir().unwrap(); - let threads_dir = tmp.path().join("threads"); - fs::create_dir_all(&threads_dir).unwrap(); - - let mut config = Config::default(); - config.workspace_dir = tmp.path().to_path_buf(); - - let source = conversation_source(); - let reader = ConversationReader; - - let result = reader.read_item(&source, "../config", &config).await; - assert!(result.is_err()); - assert!(result.unwrap_err().contains("path traversal denied")); - - let result = reader - .read_item(&source, "foo/../../etc/passwd", &config) - .await; - assert!(result.is_err()); - assert!(result.unwrap_err().contains("path traversal denied")); - } - - fn conversation_source() -> MemorySourceEntry { - MemorySourceEntry { - id: "src_conv".into(), - kind: SourceKind::Conversation, - label: "Conversations".into(), - enabled: true, - toolkit: None, - connection_id: None, - path: None, - glob: None, - url: None, - branch: None, - paths: Vec::new(), - query: None, - since_days: None, - max_items: None, - max_commits: None, - max_issues: None, - max_prs: None, - selector: None, - max_tokens_per_sync: None, - max_cost_per_sync_usd: None, - sync_depth_days: None, - } + .await + .map_err(|error| error.to_string()) } } diff --git a/src/openhuman/memory_sources/readers/folder.rs b/src/openhuman/memory_sources/readers/folder.rs index 17935037e..4508feb95 100644 --- a/src/openhuman/memory_sources/readers/folder.rs +++ b/src/openhuman/memory_sources/readers/folder.rs @@ -1,21 +1,13 @@ -//! Local folder source reader. -//! -//! Lists files matching a glob pattern under a local directory path -//! and reads their content as markdown or plaintext. +//! Product `Config` adapter for the tinycortex folder reader. use async_trait::async_trait; -use std::path::{Path, PathBuf}; use crate::openhuman::config::Config; +use crate::openhuman::memory_sources::readers::SourceReader; use crate::openhuman::memory_sources::types::{ - ContentType, MemorySourceEntry, SourceContent, SourceItem, SourceKind, + MemorySourceEntry, SourceContent, SourceItem, SourceKind, }; -use super::SourceReader; - -const DEFAULT_GLOB: &str = "**/*.md"; -const MAX_FILE_SIZE: u64 = 10 * 1024 * 1024; // 10 MB - pub struct FolderReader; #[async_trait] @@ -27,203 +19,30 @@ impl SourceReader for FolderReader { async fn list_items( &self, source: &MemorySourceEntry, - _config: &Config, + config: &Config, ) -> Result, String> { - let base_path = source - .path - .as_deref() - .ok_or("folder source requires a path")?; - let pattern = source.glob.as_deref().unwrap_or(DEFAULT_GLOB); - - let base = PathBuf::from(base_path); - if !base.exists() { - return Err(format!("folder does not exist: {base_path}")); - } - - let full_pattern = format!("{}/{pattern}", base_path.trim_end_matches('/')); - - tracing::debug!( - path = %base_path, - glob = %pattern, - "[memory_sources:folder] listing items" - ); - - let entries: Vec = glob::glob(&full_pattern) - .map_err(|e| format!("invalid glob pattern: {e}"))? - .filter_map(|entry| { - let path = entry.ok()?; - if !path.is_file() { - return None; - } - let metadata = std::fs::metadata(&path).ok()?; - if metadata.len() > MAX_FILE_SIZE { - return None; - } - let rel = path.strip_prefix(&base).ok()?; - let title = rel.to_string_lossy().to_string(); - let modified_ms = metadata - .modified() - .ok() - .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|d| d.as_millis() as i64); - - Some(SourceItem { - id: rel.to_string_lossy().to_string(), - title, - updated_at_ms: modified_ms, - }) - }) - .collect(); - - tracing::debug!(count = entries.len(), "[memory_sources:folder] found items"); - - Ok(entries) + tinycortex::memory::sources::SourceReader::list_items( + &tinycortex::memory::sources::readers::folder::FolderReader, + source, + &crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()), + ) + .await + .map_err(|error| error.to_string()) } async fn read_item( &self, source: &MemorySourceEntry, item_id: &str, - _config: &Config, + config: &Config, ) -> Result { - let base_path = source - .path - .as_deref() - .ok_or("folder source requires a path")?; - - let file_path = Path::new(base_path).join(item_id); - - if !file_path.exists() { - return Err(format!("file not found: {}", file_path.display())); - } - - // Prevent path traversal - let canonical_base = std::fs::canonicalize(base_path) - .map_err(|e| format!("cannot resolve base path: {e}"))?; - let canonical_file = std::fs::canonicalize(&file_path) - .map_err(|e| format!("cannot resolve file path: {e}"))?; - if !canonical_file.starts_with(&canonical_base) { - return Err("path traversal denied".to_string()); - } - - // Apply the same size cap as list_items so a huge file can't blow up - // the renderer or the chunker. - let metadata = std::fs::metadata(&canonical_file) - .map_err(|e| format!("failed to stat {}: {e}", canonical_file.display()))?; - if metadata.len() > MAX_FILE_SIZE { - return Err(format!( - "file exceeds {}-byte limit: {}", - MAX_FILE_SIZE, - canonical_file.display() - )); - } - - let body = tokio::fs::read_to_string(&canonical_file) - .await - .map_err(|e| format!("failed to read {}: {e}", canonical_file.display()))?; - - let content_type = if item_id.ends_with(".md") { - ContentType::Markdown - } else if item_id.ends_with(".html") || item_id.ends_with(".htm") { - ContentType::Html - } else { - ContentType::Plaintext - }; - - Ok(SourceContent { - id: item_id.to_string(), - title: item_id.to_string(), - body, - content_type, - metadata: serde_json::json!({}), - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - use tempfile::TempDir; - - fn folder_source(path: &str) -> MemorySourceEntry { - MemorySourceEntry { - id: "src_folder".into(), - kind: SourceKind::Folder, - label: "Test folder".into(), - enabled: true, - toolkit: None, - connection_id: None, - path: Some(path.into()), - glob: None, - url: None, - branch: None, - paths: Vec::new(), - query: None, - since_days: None, - max_items: None, - max_commits: None, - max_issues: None, - max_prs: None, - selector: None, - max_tokens_per_sync: None, - max_cost_per_sync_usd: None, - sync_depth_days: None, - } - } - - #[tokio::test] - async fn list_items_finds_md_files() { - let tmp = TempDir::new().unwrap(); - fs::write(tmp.path().join("note.md"), "# Hello").unwrap(); - fs::write(tmp.path().join("data.txt"), "ignored").unwrap(); - - let source = folder_source(&tmp.path().to_string_lossy()); - let reader = FolderReader; - let items = reader - .list_items(&source, &Config::default()) - .await - .unwrap(); - - assert_eq!(items.len(), 1); - assert_eq!(items[0].id, "note.md"); - } - - #[tokio::test] - async fn read_item_returns_file_content() { - let tmp = TempDir::new().unwrap(); - fs::write(tmp.path().join("test.md"), "# Test\nBody").unwrap(); - - let source = folder_source(&tmp.path().to_string_lossy()); - let reader = FolderReader; - let content = reader - .read_item(&source, "test.md", &Config::default()) - .await - .unwrap(); - - assert_eq!(content.body, "# Test\nBody"); - assert_eq!(content.content_type, ContentType::Markdown); - } - - #[tokio::test] - async fn read_item_prevents_path_traversal() { - let tmp = TempDir::new().unwrap(); - fs::write(tmp.path().join("safe.md"), "ok").unwrap(); - - let source = folder_source(&tmp.path().to_string_lossy()); - let reader = FolderReader; - let result = reader - .read_item(&source, "../../../etc/passwd", &Config::default()) - .await; - - assert!(result.is_err()); - } - - #[tokio::test] - async fn list_items_nonexistent_folder_errors() { - let source = folder_source("/nonexistent/path/xyz"); - let reader = FolderReader; - let result = reader.list_items(&source, &Config::default()).await; - assert!(result.is_err()); + tinycortex::memory::sources::SourceReader::read_item( + &tinycortex::memory::sources::readers::folder::FolderReader, + source, + item_id, + &crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()), + ) + .await + .map_err(|error| error.to_string()) } } diff --git a/src/openhuman/memory_sources/registry.rs b/src/openhuman/memory_sources/registry.rs index 0fbd4e368..e34cadd19 100644 --- a/src/openhuman/memory_sources/registry.rs +++ b/src/openhuman/memory_sources/registry.rs @@ -1,12 +1,13 @@ -//! CRUD operations for memory sources. -//! -//! Reads and writes `Config.memory_sources` via the config load/save -//! cycle. Each mutation reloads the live config, applies the change, -//! and persists atomically. +//! Product config discovery and locking around tinycortex source registry CRUD. + +use std::sync::OnceLock; use crate::openhuman::config::rpc as config_rpc; use crate::openhuman::memory_sources::types::{MemorySourceEntry, SourceKind}; -use std::sync::OnceLock; + +pub use tinycortex::memory::sources::{ + memory_sync_defaults_for_toolkit, ComposioUpsertTarget, MemorySourcePatch, +}; static MEMORY_SOURCES_WRITE_LOCK: OnceLock> = OnceLock::new(); @@ -17,69 +18,35 @@ pub(crate) async fn memory_sources_write_guard() -> tokio::sync::MutexGuard<'sta .await } -/// Conservative default sync caps for a Composio toolkit, keyed by toolkit slug. -/// -/// Single source of truth for the cheap out-of-the-box sync volume. Applied to a -/// source entry when it is first registered (`upsert_composio_source`, insert-only) -/// and by the one-time caps migration (`reconcile::apply_composio_source_caps_migration`) -/// for cap-less entries. Never overwrites a user-customised cap. -/// -/// Returns `(max_items, sync_depth_days)`. -pub fn memory_sync_defaults_for_toolkit(toolkit: &str) -> (Option, Option) { - match toolkit { - "gmail" => (Some(100), Some(30)), - "slack" => (Some(50), Some(14)), - "notion" => (Some(30), Some(30)), - "linear" => (Some(50), Some(30)), - "clickup" => (Some(50), Some(30)), - "github" => (Some(50), Some(30)), - // Generic fallback for any toolkit not listed above. - _ => (Some(30), Some(14)), - } +async fn registry() -> Result { + let config = config_rpc::load_config_with_timeout().await?; + Ok(tinycortex::memory::sources::SourceRegistry::new( + config.config_path, + )) } pub async fn list_sources() -> Result, String> { - let config = config_rpc::load_config_with_timeout().await?; - Ok(config.memory_sources.clone()) + registry().await?.list().map_err(|error| error.to_string()) } pub async fn list_enabled_by_kind(kind: SourceKind) -> Result, String> { - let config = config_rpc::load_config_with_timeout().await?; - Ok(config - .memory_sources - .iter() - .filter(|s| s.kind == kind && s.enabled) - .cloned() - .collect()) + registry() + .await? + .list_enabled_by_kind(kind) + .map_err(|error| error.to_string()) } pub async fn get_source(id: &str) -> Result, String> { - let config = config_rpc::load_config_with_timeout().await?; - Ok(config.memory_sources.iter().find(|s| s.id == id).cloned()) + registry().await?.get(id).map_err(|error| error.to_string()) } pub async fn add_source(entry: MemorySourceEntry) -> Result { - entry.validate()?; let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - if config.memory_sources.iter().any(|s| s.id == entry.id) { - return Err(format!("source with id '{}' already exists", entry.id)); - } - - tracing::info!( - id = %entry.id, - kind = %entry.kind.as_str(), - "[memory_sources] adding source" - ); - - config.memory_sources.push(entry.clone()); - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - - Ok(entry) + log::debug!("[memory_sources] crate add kind={}", entry.kind.as_str()); + registry() + .await? + .add(entry) + .map_err(|error| error.to_string()) } pub async fn update_source( @@ -87,502 +54,55 @@ pub async fn update_source( patch: MemorySourcePatch, ) -> Result { let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - let entry = config - .memory_sources - .iter_mut() - .find(|s| s.id == id) - .ok_or_else(|| format!("source '{id}' not found"))?; - - if let Some(label) = patch.label { - entry.label = label; - } - if let Some(enabled) = patch.enabled { - entry.enabled = enabled; - } - if let Some(toolkit) = patch.toolkit { - entry.toolkit = Some(toolkit); - } - if let Some(connection_id) = patch.connection_id { - entry.connection_id = Some(connection_id); - } - if let Some(path) = patch.path { - entry.path = Some(path); - } - if let Some(glob) = patch.glob { - entry.glob = Some(glob); - } - if let Some(url) = patch.url { - entry.url = Some(url); - } - if let Some(branch) = patch.branch { - entry.branch = Some(branch); - } - if let Some(paths) = patch.paths { - entry.paths = paths; - } - if let Some(query) = patch.query { - entry.query = Some(query); - } - if let Some(since_days) = patch.since_days { - entry.since_days = Some(since_days); - } - if let Some(max_items) = patch.max_items { - entry.max_items = Some(max_items); - } - if let Some(selector) = patch.selector { - entry.selector = Some(selector); - } - if let Some(v) = patch.max_tokens_per_sync { - entry.max_tokens_per_sync = Some(v); - } - if let Some(v) = patch.max_cost_per_sync_usd { - entry.max_cost_per_sync_usd = Some(v); - } - if let Some(v) = patch.sync_depth_days { - entry.sync_depth_days = Some(v); - } - if let Some(v) = patch.max_commits { - entry.max_commits = Some(v); - } - if let Some(v) = patch.max_issues { - entry.max_issues = Some(v); - } - if let Some(v) = patch.max_prs { - entry.max_prs = Some(v); - } - - entry.validate()?; - let updated = entry.clone(); - - tracing::info!( - id = %id, - kind = %updated.kind.as_str(), - "[memory_sources] updated source" - ); - - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - - Ok(updated) + log::debug!("[memory_sources] crate update id_len={}", id.len()); + registry() + .await? + .update(id, patch) + .map_err(|error| error.to_string()) } pub async fn remove_source(id: &str) -> Result { let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - let before = config.memory_sources.len(); - config.memory_sources.retain(|s| s.id != id); - let removed = config.memory_sources.len() < before; - - if removed { - tracing::info!(id = %id, "[memory_sources] removed source"); - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - } - - Ok(removed) + registry() + .await? + .remove(id) + .map_err(|error| error.to_string()) } -/// Remove every composio source bound to `connection_id` — the disconnect path. -/// -/// Mirrors [`upsert_composio_source`], which keys composio sources on -/// `connection_id`. [`remove_source`] keys on the `src_*` id, which the -/// connection-delete flow doesn't have, so this is the connection-keyed -/// counterpart. Returns the number of entries removed (0 if none matched). pub async fn remove_composio_source_by_connection_id(connection_id: &str) -> Result { let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - let before = config.memory_sources.len(); - config.memory_sources.retain(|s| { - !(s.kind == SourceKind::Composio && s.connection_id.as_deref() == Some(connection_id)) - }); - let removed = before - config.memory_sources.len(); - - if removed > 0 { - tracing::info!( - connection_id = %connection_id, - removed, - "[memory_sources] removed composio source(s) on connection disconnect" - ); - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - } - - Ok(removed) + registry() + .await? + .remove_composio_source_by_connection_id(connection_id) + .map_err(|error| error.to_string()) } -/// Apply a single composio upsert to an in-memory source list. -/// -/// If a source with the same `connection_id` already exists, updates its label -/// and returns a clone of the updated entry with `was_insert = false`; -/// otherwise pushes a new entry (with conservative per-toolkit defaults) and -/// returns it with `was_insert = true`. -/// -/// Pure (no I/O) so the single-call path, the batch path, and unit tests all -/// share one find-or-push predicate and cannot drift. -fn upsert_composio_entry_in_place( - sources: &mut Vec, - toolkit: &str, - connection_id: &str, - label: &str, -) -> (MemorySourceEntry, bool) { - if let Some(existing) = sources.iter_mut().find(|s| { - s.kind == SourceKind::Composio && s.connection_id.as_deref() == Some(connection_id) - }) { - existing.label = label.to_string(); - return (existing.clone(), false); - } - - let (default_max_items, default_sync_depth_days) = memory_sync_defaults_for_toolkit(toolkit); - tracing::debug!( - toolkit = %toolkit, - max_items = ?default_max_items, - sync_depth_days = ?default_sync_depth_days, - "[memory_sources] applying conservative defaults for new composio source" - ); - - let entry = MemorySourceEntry { - id: format!("src_{}", uuid::Uuid::new_v4().as_simple()), - kind: SourceKind::Composio, - label: label.to_string(), - enabled: true, - toolkit: Some(toolkit.to_string()), - connection_id: Some(connection_id.to_string()), - path: None, - glob: None, - url: None, - branch: None, - paths: Vec::new(), - max_commits: None, - max_issues: None, - max_prs: None, - query: None, - since_days: None, - max_items: default_max_items, - selector: None, - max_tokens_per_sync: None, - max_cost_per_sync_usd: None, - sync_depth_days: default_sync_depth_days, - }; - sources.push(entry.clone()); - (entry, true) -} - -/// Upsert a composio source — used by the auto-registration path. -/// If a source with the same `connection_id` already exists, updates -/// the label; otherwise inserts a new entry. pub async fn upsert_composio_source( toolkit: &str, connection_id: &str, label: &str, ) -> Result { let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - let (entry, was_insert) = - upsert_composio_entry_in_place(&mut config.memory_sources, toolkit, connection_id, label); - - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - - if was_insert { - tracing::info!( - connection_id = %connection_id, - toolkit = %toolkit, - "[memory_sources] upserted composio source (insert)" - ); - } else { - tracing::debug!( - connection_id = %connection_id, - toolkit = %toolkit, - "[memory_sources] upserted composio source (update)" - ); - } - - Ok(entry) + registry() + .await? + .upsert_composio_source(toolkit, connection_id, label) + .map_err(|error| error.to_string()) } -/// Target for a batch composio upsert: `(toolkit, connection_id, label)`. -pub type ComposioUpsertTarget = (String, String, String); - -/// Batch-upsert many composio sources with a **single** config load + save. -/// -/// The per-call [`upsert_composio_source`] does its own load-modify-save, so -/// calling it in a loop costs `2N` config I/O round-trips and — worse — cannot -/// be parallelised safely (concurrent calls each load the same snapshot and -/// their saves clobber each other). This batch path loads the config once, -/// applies every upsert in-memory via the shared [`upsert_composio_entry_in_place`] -/// predicate, and saves once. Returns the number of targets applied. -/// -/// Targets are applied in order; a later target sharing a `connection_id` with -/// an earlier one updates the same (already inserted/updated) entry's label, -/// matching the sequential single-call semantics. pub async fn upsert_composio_sources_batch( targets: &[ComposioUpsertTarget], ) -> Result { - if targets.is_empty() { - return Ok(0); - } - let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - let mut inserted = 0u32; - let mut updated = 0u32; - for (toolkit, connection_id, label) in targets { - let (_entry, was_insert) = upsert_composio_entry_in_place( - &mut config.memory_sources, - toolkit, - connection_id, - label, - ); - if was_insert { - inserted += 1; - } else { - updated += 1; - } - } - - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - - tracing::info!( - targets = targets.len(), - inserted, - updated, - "[memory_sources] batch-upserted composio sources (single save)" - ); - - Ok(inserted + updated) + registry() + .await? + .upsert_composio_sources_batch(targets) + .map_err(|error| error.to_string()) } -/// Partial update payload for a source entry. -#[derive(Debug, Default, serde::Deserialize)] -pub struct MemorySourcePatch { - #[serde(default)] - pub label: Option, - #[serde(default)] - pub enabled: Option, - #[serde(default)] - pub toolkit: Option, - #[serde(default)] - pub connection_id: Option, - #[serde(default)] - pub path: Option, - #[serde(default)] - pub glob: Option, - #[serde(default)] - pub url: Option, - #[serde(default)] - pub branch: Option, - #[serde(default)] - pub paths: Option>, - #[serde(default)] - pub query: Option, - #[serde(default)] - pub since_days: Option, - #[serde(default)] - pub max_items: Option, - #[serde(default)] - pub selector: Option, - #[serde(default)] - pub max_tokens_per_sync: Option, - #[serde(default)] - pub max_cost_per_sync_usd: Option, - #[serde(default)] - pub sync_depth_days: Option, - // ── GithubRepo-specific caps (previously missing from patch) ── - #[serde(default)] - pub max_commits: Option, - #[serde(default)] - pub max_issues: Option, - #[serde(default)] - pub max_prs: Option, -} - -/// Enable ALL configured memory sources and clear every per-source cap, -/// giving the user unrestricted access ("All In" mode). -/// -/// For each source in `config.memory_sources`: -/// - Sets `enabled = true`. -/// - Clears `max_items`, `since_days`, `sync_depth_days`, -/// `max_commits`, `max_issues`, `max_prs`, -/// `max_tokens_per_sync`, `max_cost_per_sync_usd` to `None`. -/// -/// Saves config once after all mutations and returns the updated entries. pub async fn apply_all_in() -> Result, String> { let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - tracing::info!( - count = config.memory_sources.len(), - "[memory_sources] apply_all_in: enabling all sources and clearing caps" - ); - - for source in &mut config.memory_sources { - tracing::debug!( - id = %source.id, - kind = %source.kind.as_str(), - "[memory_sources] apply_all_in: enabling source and clearing caps" - ); - source.enabled = true; - source.max_items = None; - source.since_days = None; - source.sync_depth_days = None; - source.max_commits = None; - source.max_issues = None; - source.max_prs = None; - source.max_tokens_per_sync = None; - source.max_cost_per_sync_usd = None; - } - - let updated = config.memory_sources.clone(); - - config - .save() - .await - .map_err(|e| format!("apply_all_in: failed to save config: {e:#}"))?; - - tracing::info!( - count = updated.len(), - "[memory_sources] apply_all_in: complete — all sources enabled, all caps cleared" - ); - - Ok(updated) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn composio_defaults_for_known_toolkits() { - assert_eq!( - memory_sync_defaults_for_toolkit("gmail"), - (Some(100), Some(30)) - ); - assert_eq!( - memory_sync_defaults_for_toolkit("slack"), - (Some(50), Some(14)) - ); - assert_eq!( - memory_sync_defaults_for_toolkit("notion"), - (Some(30), Some(30)) - ); - assert_eq!( - memory_sync_defaults_for_toolkit("linear"), - (Some(50), Some(30)) - ); - assert_eq!( - memory_sync_defaults_for_toolkit("clickup"), - (Some(50), Some(30)) - ); - assert_eq!( - memory_sync_defaults_for_toolkit("github"), - (Some(50), Some(30)) - ); - } - - #[test] - fn composio_defaults_for_generic_fallback() { - assert_eq!( - memory_sync_defaults_for_toolkit("unknown_toolkit_xyz"), - (Some(30), Some(14)) - ); - assert_eq!(memory_sync_defaults_for_toolkit(""), (Some(30), Some(14))); - } - - #[test] - fn in_place_upsert_inserts_new_entry_with_toolkit_defaults() { - let mut sources: Vec = vec![]; - let (entry, was_insert) = - upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "Gmail · conn_a"); - assert!(was_insert, "first upsert for a connection must insert"); - assert_eq!(sources.len(), 1); - assert_eq!(entry.kind, SourceKind::Composio); - assert_eq!(entry.connection_id.as_deref(), Some("conn_a")); - assert_eq!(entry.toolkit.as_deref(), Some("gmail")); - assert_eq!(entry.label, "Gmail · conn_a"); - assert!(entry.enabled); - // Conservative gmail defaults are applied on insert. - assert_eq!(entry.max_items, Some(100)); - assert_eq!(entry.sync_depth_days, Some(30)); - } - - #[test] - fn in_place_upsert_updates_label_only_for_existing_connection() { - let mut sources: Vec = vec![]; - upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "old label"); - // User-customised a cap after the initial insert. - sources[0].max_items = Some(7); - - let (entry, was_insert) = - upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "new label"); - assert!(!was_insert, "second upsert for same connection must update"); - assert_eq!(sources.len(), 1, "must not duplicate the entry"); - assert_eq!(entry.label, "new label"); - assert_eq!( - entry.max_items, - Some(7), - "update must not clobber user-set caps" - ); - } - - #[test] - fn in_place_upsert_keeps_distinct_connections_separate() { - let mut sources: Vec = vec![]; - upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "A"); - upsert_composio_entry_in_place(&mut sources, "gmail", "conn_b", "B"); - assert_eq!(sources.len(), 2); - let ids: Vec<_> = sources - .iter() - .filter_map(|s| s.connection_id.as_deref()) - .collect(); - assert!(ids.contains(&"conn_a") && ids.contains(&"conn_b")); - } - - #[test] - fn memory_source_patch_deserializes_partial() { - let json = serde_json::json!({ "label": "New label", "enabled": false }); - let patch: MemorySourcePatch = serde_json::from_value(json).unwrap(); - assert_eq!(patch.label.as_deref(), Some("New label")); - assert_eq!(patch.enabled, Some(false)); - assert!(patch.toolkit.is_none()); - } - - #[test] - fn memory_source_patch_round_trips_github_limit_fields() { - let json = serde_json::json!({ - "max_commits": 100, - "max_issues": 50, - "max_prs": 25 - }); - let patch: MemorySourcePatch = serde_json::from_value(json).unwrap(); - assert_eq!(patch.max_commits, Some(100)); - assert_eq!(patch.max_issues, Some(50)); - assert_eq!(patch.max_prs, Some(25)); - // Unset fields must be None (serde(default)) - assert!(patch.label.is_none()); - assert!(patch.enabled.is_none()); - } - - #[test] - fn memory_source_patch_defaults_github_fields_to_none() { - let json = serde_json::json!({ "enabled": true }); - let patch: MemorySourcePatch = serde_json::from_value(json).unwrap(); - assert!(patch.max_commits.is_none()); - assert!(patch.max_issues.is_none()); - assert!(patch.max_prs.is_none()); - } + registry() + .await? + .apply_all_in() + .map_err(|error| error.to_string()) } diff --git a/src/openhuman/memory_sources/sync.rs b/src/openhuman/memory_sources/sync.rs index 402bc1d72..0950fe6ed 100644 --- a/src/openhuman/memory_sources/sync.rs +++ b/src/openhuman/memory_sources/sync.rs @@ -1,9 +1,7 @@ //! Per-source sync dispatcher. //! -//! Thin routing layer: dispatches sync requests to the right backend: -//! - GitHub repos → `memory_sync::sources::github` -//! - Composio sources → `memory_sync::composio` -//! - Folder/RSS/WebPage → per-item ingest via reader + ingest pipeline +//! Thin routing layer: dispatches supported sources through tinycortex and +//! retains the product-owned background lock, events, and reconcile shell. //! - Twitter → placeholder //! //! Sync runs in a `tokio::spawn`-ed task so the RPC returns immediately. @@ -13,20 +11,12 @@ //! presses the sync button multiple times. use std::collections::HashSet; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; - -use futures::stream::{self, StreamExt}; +use std::sync::Mutex; use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::ingest_document_with_scope; use crate::openhuman::memory::sync::{emit_sync_stage, MemorySyncStage, MemorySyncTrigger}; -use crate::openhuman::memory_sources::readers; use crate::openhuman::memory_sources::types::{MemorySourceEntry, SourceKind}; -use crate::openhuman::memory_sync::canonicalize::document::DocumentInput; -use crate::openhuman::memory_sync::composio::{self, ComposioUsage, SyncReason}; - -const SYNC_CONCURRENCY: usize = 10; +use crate::openhuman::memory_sync::composio::ComposioUsage; static ACTIVE_SYNCS: std::sync::LazyLock>> = std::sync::LazyLock::new(|| Mutex::new(HashSet::new())); @@ -95,22 +85,37 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<() let mut composio_usage = ComposioUsage::default(); let outcome = match source.kind { SourceKind::Composio => { - sync_composio(&source, config.clone(), &mut composio_usage).await + match crate::openhuman::tinycortex::run_source_pipeline(&source, &config).await + { + Ok(outcome) => { + composio_usage.actions_called = outcome.actions_called; + composio_usage.cost_usd = outcome.provider_cost_usd; + Ok(outcome.records_ingested as usize) + } + Err(error) => { + composio_usage.actions_called = error.actions_called; + composio_usage.cost_usd = error.provider_cost_usd; + Err(format!("composio sync failed: {error}")) + } + } + } + SourceKind::Conversation | SourceKind::Folder => { + crate::openhuman::tinycortex::run_source_pipeline(&source, &config) + .await + .map(|outcome| outcome.records_ingested as usize) + .map_err(|error| error.to_string()) } - SourceKind::Conversation => sync_items_individually(&source, &config).await, SourceKind::GithubRepo => { - // GitHub path writes its own detailed audit entry - // with token breakdowns; skip the dispatcher-level - // audit for this kind. - crate::openhuman::memory_sync::sources::github::run_github_sync( - &source, &config, - ) - .await - .map(|o| o.records_ingested as usize) - .map_err(|e| format!("{e:#}")) + crate::openhuman::tinycortex::run_source_pipeline(&source, &config) + .await + .map(|outcome| outcome.records_ingested as usize) + .map_err(|error| error.to_string()) } - SourceKind::Folder | SourceKind::RssFeed | SourceKind::WebPage => { - sync_items_individually(&source, &config).await + SourceKind::RssFeed | SourceKind::WebPage => { + crate::openhuman::tinycortex::run_source_pipeline(&source, &config) + .await + .map(|outcome| outcome.records_ingested as usize) + .map_err(|error| error.to_string()) } SourceKind::TwitterQuery => Err( "Twitter sync not yet configured. Provide bearer token in settings." @@ -136,37 +141,33 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<() Some(&source.id), ); - // Write audit entry (GitHub writes its own with - // token detail; other kinds get a simpler entry). - if source.kind != SourceKind::GithubRepo { - use crate::openhuman::memory_sync::sources::audit::{ - append_audit_entry, SyncAuditEntry, - }; - append_audit_entry( - &config, - &SyncAuditEntry { - timestamp: chrono::Utc::now(), - source_id: source.id.clone(), - source_kind: source.kind.as_str().to_string(), - scope: source - .url - .clone() - .or(source.toolkit.clone()) - .unwrap_or_else(|| source.id.clone()), - items_fetched: items as u32, - batches: 0, - input_tokens: 0, - output_tokens: 0, - estimated_cost_usd: 0.0, - composio_actions_called: composio_usage.actions_called, - composio_cost_usd: composio_usage.cost_usd, - actual_charged_usd: None, - duration_ms, - success: true, - error: None, - }, - ); - } + use crate::openhuman::memory_sync::sources::audit::{ + append_audit_entry, SyncAuditEntry, + }; + append_audit_entry( + &config, + &SyncAuditEntry { + timestamp: chrono::Utc::now(), + source_id: source.id.clone(), + source_kind: source.kind.as_str().to_string(), + scope: source + .url + .clone() + .or(source.toolkit.clone()) + .unwrap_or_else(|| source.id.clone()), + items_fetched: items as u32, + batches: 0, + input_tokens: 0, + output_tokens: 0, + estimated_cost_usd: 0.0, + composio_actions_called: composio_usage.actions_called, + composio_cost_usd: composio_usage.cost_usd, + actual_charged_usd: None, + duration_ms, + success: true, + error: None, + }, + ); // Auto-rebuild: if raw files exist but the tree has // no summaries, build the tree now. @@ -268,299 +269,11 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<() Ok(()) } -async fn sync_composio( - source: &MemorySourceEntry, - config: Config, - usage_out: &mut ComposioUsage, -) -> Result { - let connection_id = source - .connection_id - .as_deref() - .ok_or("composio source missing connection_id")?; - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Fetching, - Some("composio"), - Some(&source.id), - Some(format!("delegating to composio sync for {connection_id}")), - Some(&source.id), - ); - - match composio::run_connection_sync(config, connection_id, SyncReason::Manual).await { - Ok((outcome, usage)) => { - *usage_out = usage; - Ok(outcome.items_ingested) - } - Err((e, usage)) => { - *usage_out = usage; - Err(format!("composio sync failed: {e}")) - } - } -} - -/// Per-item sync path for Folder/RSS/WebPage sources. -async fn sync_items_individually( - source: &MemorySourceEntry, - config: &Config, -) -> Result { - let reader = readers::reader_for(&source.kind); - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Fetching, - Some(source.kind.as_str()), - Some(&source.id), - Some("listing items".to_string()), - Some(&source.id), - ); - - let items = reader.list_items(source, config).await?; - let total = items.len(); - - // Reconcile before re-ingesting: for LOCAL FOLDER sources only, drop chunks - // (rows + on-disk bodies) for items previously ingested under this source - // that no longer exist on disk — e.g. a renamed or deleted file. Each item - // ingests under the content-addressed composite id - // `mem_src:{source_id}:{item_id}` as a Document, so a rename mints a fresh id - // and orphans the old chunk + its on-disk body; without this the stale body - // lingers forever and can only ever be served as a ≤500-char preview (#4689). - // - // Runs BEFORE the empty-listing early return so an emptied folder (every file - // deleted → total == 0) still reconciles instead of leaving all its chunks - // behind. This is safe on an empty or partial listing because - // `prune_vanished_items` re-checks each candidate on disk and only deletes - // files that are provably absent, so a transient listing miss (EACCES / - // EMFILE / stat stall) is never mistaken for a deletion. - // - // Restricted to Folder: for feed / web / conversation sources, absence from - // the current listing means "rolled off / not re-fetched", NOT "deleted", so - // pruning them would irrecoverably delete valid archived items. - if source_supports_prune(&source.kind) { - if let Some(base_path) = source.path.clone() { - let config = config.clone(); - let source_id = source.id.clone(); - let live: HashSet = items - .iter() - .map(|item| format!("mem_src:{source_id}:{}", item.id)) - .collect(); - if let Err(e) = tokio::task::spawn_blocking(move || { - prune_vanished_items(&config, &source_id, std::path::Path::new(&base_path), &live) - }) - .await - { - tracing::warn!(error = %e, "[memory_sources:sync] prune join error"); - } - } - } - - if total == 0 { - return Ok(0); - } - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Stored, - Some(source.kind.as_str()), - Some(&source.id), - Some(format!("{total} item(s) discovered")), - Some(&source.id), - ); - - let ingested = Arc::new(AtomicUsize::new(0)); - let processed = Arc::new(AtomicUsize::new(0)); - let source_id = source.id.clone(); - let source_kind = source.kind.clone(); - let kind_str = source.kind.as_str().to_string(); - - stream::iter(items.iter().enumerate()) - .for_each_concurrent(SYNC_CONCURRENCY, |(_, item)| { - let config = config.clone(); - let source_kind = source_kind.clone(); - let reader = readers::reader_for(&source_kind); - let source_clone = source.clone(); - let ingested = Arc::clone(&ingested); - let processed = Arc::clone(&processed); - let source_id = source_id.clone(); - let kind_str = kind_str.clone(); - - async move { - let content = match reader.read_item(&source_clone, &item.id, &config).await { - Ok(c) => c, - Err(e) => { - tracing::warn!( - item_id = %item.id, - error = %e, - "[memory_sources:sync] skipping item — read failed" - ); - processed.fetch_add(1, Ordering::Relaxed); - return; - } - }; - - let doc = DocumentInput { - provider: format!("memory_sources:{kind_str}"), - title: content.title.clone(), - body: content.body.clone(), - modified_at: chrono::Utc::now(), - source_ref: Some(format!("{source_id}:{}", item.id)), - }; - - let composite_source_id = format!("mem_src:{source_id}:{}", item.id); - let tags = vec!["memory_sources".to_string(), kind_str.clone()]; - - match ingest_document_with_scope( - &config, - &composite_source_id, - "user", - tags, - doc, - None, - ) - .await - { - Ok(result) => { - if !result.already_ingested { - ingested.fetch_add(1, Ordering::Relaxed); - } - } - Err(e) => { - tracing::warn!( - item_id = %item.id, - error = %e, - "[memory_sources:sync] ingest failed for item" - ); - } - } - - let done = processed.fetch_add(1, Ordering::Relaxed) + 1; - let new = ingested.load(Ordering::Relaxed); - if done.is_multiple_of(10) || done == total { - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Ingesting, - Some(&kind_str), - Some(&source_id), - Some(format!("{done}/{total} processed ({new} new)")), - Some(&source_id), - ); - } - } - }) - .await; - - Ok(ingested.load(Ordering::Relaxed)) -} - -/// Whether a source kind has authoritative present/absent semantics on the -/// local disk, so that "absent from the current listing" genuinely means -/// "deleted" and can drive a prune (#4689). -/// -/// Only `Folder` qualifies: for `RssFeed` / `WebPage`, `list_items` returns a -/// rolling, `max_items`-truncated window, so an item missing from it merely -/// rolled off the feed and must never be deleted; `Conversation` threads have -/// no delete-follows-listing contract. Restricting prune here prevents turning -/// an append-only archive into a destructive mirror of the latest window. -fn source_supports_prune(kind: &SourceKind) -> bool { - matches!(kind, SourceKind::Folder) -} - -/// Delete chunks (rows + on-disk bodies) for items previously ingested under -/// `source_id` whose backing file no longer exists under `base_path` — the -/// reconcile step that keeps a folder resync from orphaning renamed or deleted -/// files (#4689). -/// -/// Safety: a candidate is deleted ONLY when its file is provably absent -/// (`symlink_metadata` returns `NotFound`). A transient listing miss (the reader -/// dropped a still-present file on an `EACCES` / `EMFILE` / stat stall) leaves -/// the file on disk, so the re-check keeps it — absence from `live` alone is -/// never sufficient to delete. -/// -/// Blocking DB + FS work; call from `spawn_blocking`. Chunks land under the -/// content (`Document`) `SourceKind`, not the outer `memory_sources` source kind. -fn prune_vanished_items( - config: &Config, - source_id: &str, - base_path: &std::path::Path, - live: &HashSet, -) { - use crate::openhuman::memory_store::chunks::store as chunk_store; - use crate::openhuman::memory_store::chunks::types::SourceKind as ChunkSourceKind; - - let prefix = format!("mem_src:{source_id}:"); - let previously = match chunk_store::list_source_ids_with_prefix( - config, - ChunkSourceKind::Document, - &prefix, - ) { - Ok(ids) => ids, - Err(e) => { - tracing::warn!( - source_id = %source_id, - error = %format!("{e:#}"), - "[memory_sources:sync] prune: failed to list previously-ingested items" - ); - return; - } - }; - - let mut removed_chunks = 0usize; - let mut removed_items = 0usize; - for stale in previously.into_iter().filter(|sid| !live.contains(sid)) { - // Recover the item's relative path from the composite id and confirm the - // file is genuinely gone before deleting. Anything other than a definite - // NotFound (present file, or an ambiguous EACCES/IO error) is treated as - // "keep" so a transient listing miss can never delete live data. - let Some(rel) = stale.strip_prefix(&prefix) else { - continue; - }; - // Defense-in-depth: `rel` comes from a stored composite id. If it were - // ever empty, absolute, or contained `..`, `base_path.join(rel)` could - // resolve outside the source folder (on Unix an absolute `rel` silently - // discards `base_path`), and a `NotFound` there would delete real chunk - // rows. The current folder reader can't produce such ids, so keep any - // such candidate rather than risk deleting on a path we can't vouch for. - if rel.is_empty() || rel.contains("..") || std::path::Path::new(rel).is_absolute() { - tracing::warn!( - source_id = %source_id, - "[memory_sources:sync] prune: skipping candidate with unsafe relative path" - ); - continue; - } - match std::fs::symlink_metadata(base_path.join(rel)) { - Err(e) if e.kind() == std::io::ErrorKind::NotFound => { /* absent → prune */ } - _ => continue, - } - match chunk_store::delete_chunks_by_source(config, ChunkSourceKind::Document, &stale) { - Ok(n) => { - removed_chunks += n; - removed_items += 1; - } - Err(e) => tracing::warn!( - source_id = %source_id, - error = %format!("{e:#}"), - "[memory_sources:sync] prune: delete failed for a vanished item" - ), - } - } - if removed_items > 0 { - tracing::info!( - source_id = %source_id, - items = removed_items, - chunks = removed_chunks, - "[memory_sources:sync] pruned chunks for vanished items" - ); - } -} - -/// Derive the tree scope(s) for a source and reconcile any raw files that -/// are not yet covered by tree summaries (incremental — see -/// `memory_sync::sources::rebuild`). +/// Reconcile raw files that are not yet covered by tree summaries. pub(crate) async fn check_and_rebuild_tree(source: &MemorySourceEntry, config: &Config) { use crate::openhuman::memory_sync::sources::rebuild::{needs_rebuild, rebuild_tree_from_raw}; - let scopes = derive_scopes(source, config); - for scope in scopes { + for scope in derive_scopes(source, config) { if !needs_rebuild(config, &scope.tree_scope, &scope.archive_source_id) { continue; } @@ -571,26 +284,19 @@ pub(crate) async fn check_and_rebuild_tree(source: &MemorySourceEntry, config: & "[memory_sources:sync] reconciling uncovered raw files into tree" ); match rebuild_tree_from_raw(config, &scope.tree_scope, &scope.archive_source_id).await { - Ok(outcome) => { - tracing::info!( - scope = %scope.tree_scope, - files = outcome.files_read, - batches = outcome.batches, - cost = %format!( - "${:.4}", - outcome.actual_charged_usd.unwrap_or(outcome.estimated_cost_usd) - ), - cost_is_actual = outcome.actual_charged_usd.is_some(), - "[memory_sources:sync] reconcile complete" - ); - } - Err(e) => { - tracing::warn!( - scope = %scope.tree_scope, - error = %format!("{e:#}"), - "[memory_sources:sync] reconcile failed" - ); - } + Ok(outcome) => tracing::info!( + scope = %scope.tree_scope, + files = outcome.files_read, + batches = outcome.batches, + cost = %format!("${:.4}", outcome.actual_charged_usd.unwrap_or(outcome.estimated_cost_usd)), + cost_is_actual = outcome.actual_charged_usd.is_some(), + "[memory_sources:sync] reconcile complete" + ), + Err(error) => tracing::warn!( + scope = %scope.tree_scope, + error = %format!("{error:#}"), + "[memory_sources:sync] reconcile failed" + ), } } } @@ -676,178 +382,3 @@ pub(crate) fn derive_scopes(source: &MemorySourceEntry, config: &Config) -> Vec< _ => Vec::new(), } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::chunks::store as chunk_store; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind as ChunkSourceKind, - }; - use crate::openhuman::memory_store::content::stage_chunks; - use chrono::TimeZone; - use tempfile::TempDir; - - fn doc_chunk(source_id: &str) -> Chunk { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - Chunk { - id: chunk_id(ChunkSourceKind::Document, source_id, 0, "body"), - content: format!("body of {source_id}"), - metadata: Metadata { - source_kind: ChunkSourceKind::Document, - source_id: source_id.into(), - owner: "user".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: None, - path_scope: None, - }, - token_count: 2, - seq_in_source: 0, - created_at: ts, - partial_message: false, - } - } - - fn seed(cfg: &Config, chunk: &Chunk) { - let staged = - stage_chunks(&cfg.memory_tree_content_root(), std::slice::from_ref(chunk)).unwrap(); - chunk_store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - chunk_store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - } - - #[test] - fn source_supports_prune_only_for_folder() { - assert!(source_supports_prune(&SourceKind::Folder)); - assert!(!source_supports_prune(&SourceKind::RssFeed)); - assert!(!source_supports_prune(&SourceKind::WebPage)); - assert!(!source_supports_prune(&SourceKind::Conversation)); - } - - #[test] - fn prune_vanished_items_removes_only_files_absent_from_disk() { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Folder base holds only b.md on disk; a.md was renamed/deleted. - let base = tmp.path().join("folder"); - std::fs::create_dir_all(&base).unwrap(); - std::fs::write(base.join("b.md"), b"b").unwrap(); - - seed(&cfg, &doc_chunk("mem_src:src1:a.md")); - seed(&cfg, &doc_chunk("mem_src:src1:b.md")); - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 2); - - let live: HashSet = ["mem_src:src1:b.md".to_string()].into_iter().collect(); - prune_vanished_items(&cfg, "src1", &base, &live); - - let remaining = chunk_store::list_source_ids_with_prefix( - &cfg, - ChunkSourceKind::Document, - "mem_src:src1:", - ) - .unwrap(); - assert_eq!(remaining, vec!["mem_src:src1:b.md".to_string()]); - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 1); - } - - #[test] - fn prune_vanished_items_keeps_still_present_file_missed_by_listing() { - // Safety guard (#4689 review): a transient listing miss must not delete a - // file that is still on disk. 'a.md' is absent from `live` but present on - // disk → it must be kept. - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let base = tmp.path().join("folder"); - std::fs::create_dir_all(&base).unwrap(); - std::fs::write(base.join("a.md"), b"still here").unwrap(); - - seed(&cfg, &doc_chunk("mem_src:src3:a.md")); - // 'a.md' dropped from the listing (e.g. EACCES/EMFILE) though it exists. - let live: HashSet = HashSet::new(); - prune_vanished_items(&cfg, "src3", &base, &live); - - // Not deleted — the on-disk re-check kept it. - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 1); - } - - #[test] - fn prune_vanished_items_prunes_when_folder_emptied() { - // Emptied folder: listing is empty (live = {}) and no files remain on - // disk, so the previously-ingested chunk must be pruned (the reason prune - // now runs before the total == 0 early return). - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let base = tmp.path().join("empty_folder"); - std::fs::create_dir_all(&base).unwrap(); - - seed(&cfg, &doc_chunk("mem_src:src4:gone.md")); - let live: HashSet = HashSet::new(); - prune_vanished_items(&cfg, "src4", &base, &live); - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 0); - } - - #[test] - fn prune_vanished_items_keeps_candidate_with_unsafe_relative_path() { - // Defense-in-depth: a stored id whose relative path is absolute must not - // be pruned even when its file is "absent" (join would escape the base). - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let base = tmp.path().join("folder"); - std::fs::create_dir_all(&base).unwrap(); - - seed(&cfg, &doc_chunk("mem_src:src5:/etc/hostname")); - let live: HashSet = HashSet::new(); - prune_vanished_items(&cfg, "src5", &base, &live); - // Not deleted — the unsafe-path guard skipped it. - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 1); - } - - #[test] - fn list_source_ids_with_prefix_isolates_sibling_prefixes() { - // `mem_src:src1:` must not match `mem_src:src10:` items. - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - - seed(&cfg, &doc_chunk("mem_src:src1:a.md")); - seed(&cfg, &doc_chunk("mem_src:src1:c.md")); - seed(&cfg, &doc_chunk("mem_src:src10:b.md")); - - let mut got = chunk_store::list_source_ids_with_prefix( - &cfg, - ChunkSourceKind::Document, - "mem_src:src1:", - ) - .unwrap(); - got.sort(); - assert_eq!( - got, - vec![ - "mem_src:src1:a.md".to_string(), - "mem_src:src1:c.md".to_string() - ] - ); - } - - #[test] - fn prune_vanished_items_is_noop_when_all_live() { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - - seed(&cfg, &doc_chunk("mem_src:src2:a.md")); - let live: HashSet = ["mem_src:src2:a.md".to_string()].into_iter().collect(); - prune_vanished_items(&cfg, "src2", tmp.path(), &live); - assert_eq!(chunk_store::count_chunks(&cfg).unwrap(), 1); - } -} diff --git a/src/openhuman/memory_sources/types.rs b/src/openhuman/memory_sources/types.rs index 258472aa0..6fbab9fb2 100644 --- a/src/openhuman/memory_sources/types.rs +++ b/src/openhuman/memory_sources/types.rs @@ -1,346 +1,5 @@ -//! Core types for memory sources. +//! Stable host path for tinycortex-owned memory-source contracts. -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; - -fn default_true() -> bool { - true -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] -#[serde(rename_all = "snake_case")] -pub enum SourceKind { - Composio, - Conversation, - Folder, - GithubRepo, - TwitterQuery, - RssFeed, - WebPage, -} - -impl SourceKind { - pub fn as_str(&self) -> &'static str { - match self { - SourceKind::Composio => "composio", - SourceKind::Conversation => "conversation", - SourceKind::Folder => "folder", - SourceKind::GithubRepo => "github_repo", - SourceKind::TwitterQuery => "twitter_query", - SourceKind::RssFeed => "rss_feed", - SourceKind::WebPage => "web_page", - } - } -} - -/// A configured memory source entry persisted in `config.toml`. -/// -/// All kind-specific fields are flattened onto the struct as `Option`s. -/// The `kind` discriminator determines which fields are required; -/// validation is enforced at add/update time via [`validate`]. -#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] -pub struct MemorySourceEntry { - pub id: String, - pub kind: SourceKind, - pub label: String, - #[serde(default = "default_true")] - pub enabled: bool, - - // ── Composio ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub toolkit: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub connection_id: Option, - - // ── Folder ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub path: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub glob: Option, - - // ── GithubRepo / RssFeed / WebPage / TwitterQuery (shared) ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub url: Option, - - // ── GithubRepo ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub branch: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub paths: Vec, - /// Max commits to pull per sync (default 1000 when absent). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_commits: Option, - /// Max issues to pull per sync (default 1000 when absent). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_issues: Option, - /// Max pull requests to pull per sync (default 1000 when absent). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_prs: Option, - - // ── TwitterQuery ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub query: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub since_days: Option, - - // ── RssFeed ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_items: Option, - - // ── WebPage ── - #[serde(default, skip_serializing_if = "Option::is_none")] - pub selector: Option, - - // ── Sync Budget (all source kinds) ── - /// Maximum tokens to consume per sync run. Sync stops once this budget is hit. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_tokens_per_sync: Option, - /// Maximum cost in USD per sync run. Refuses LLM calls once reached. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_cost_per_sync_usd: Option, - /// Sync depth in days — only fetch items from the last N days. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub sync_depth_days: Option, -} - -impl MemorySourceEntry { - pub fn validate(&self) -> Result<(), String> { - if self.id.is_empty() { - return Err("id is required".to_string()); - } - if self.label.is_empty() { - return Err("label is required".to_string()); - } - match self.kind { - SourceKind::Composio => { - require_field(&self.toolkit, "toolkit")?; - require_field(&self.connection_id, "connection_id")?; - } - SourceKind::Conversation => { - // No kind-specific required fields — just enabled/disabled. - } - SourceKind::Folder => { - require_field(&self.path, "path")?; - } - SourceKind::GithubRepo => { - require_field(&self.url, "url")?; - } - SourceKind::TwitterQuery => { - require_field(&self.query, "query")?; - } - SourceKind::RssFeed => { - require_field(&self.url, "url")?; - } - SourceKind::WebPage => { - require_field(&self.url, "url")?; - } - } - Ok(()) - } -} - -fn require_field(value: &Option, name: &str) -> Result<(), String> { - match value { - Some(v) if !v.is_empty() => Ok(()), - _ => Err(format!("{name} is required for this source kind")), - } -} - -/// One item listed from a source reader. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SourceItem { - pub id: String, - pub title: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub updated_at_ms: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ContentType { - Markdown, - Html, - Plaintext, -} - -/// Content read from a single source item. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SourceContent { - pub id: String, - pub title: String, - pub body: String, - pub content_type: ContentType, - #[serde(default)] - pub metadata: serde_json::Value, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn source_kind_round_trips_via_serde() { - for kind in [ - SourceKind::Composio, - SourceKind::Conversation, - SourceKind::Folder, - SourceKind::GithubRepo, - SourceKind::TwitterQuery, - SourceKind::RssFeed, - SourceKind::WebPage, - ] { - let json = serde_json::to_string(&kind).unwrap(); - let decoded: SourceKind = serde_json::from_str(&json).unwrap(); - assert_eq!(decoded, kind); - } - } - - #[test] - fn validate_composio_requires_toolkit_and_connection_id() { - let entry = MemorySourceEntry { - id: "src_1".into(), - kind: SourceKind::Composio, - label: "Gmail".into(), - enabled: true, - toolkit: Some("gmail".into()), - connection_id: None, - ..default_entry() - }; - assert!(entry.validate().is_err()); - - let valid = MemorySourceEntry { - connection_id: Some("cmp_123".into()), - ..entry - }; - assert!(valid.validate().is_ok()); - } - - #[test] - fn validate_folder_requires_path() { - let entry = MemorySourceEntry { - id: "src_2".into(), - kind: SourceKind::Folder, - label: "Notes".into(), - enabled: true, - path: None, - ..default_entry() - }; - assert!(entry.validate().is_err()); - } - - #[test] - fn validate_github_requires_url() { - let entry = MemorySourceEntry { - id: "src_3".into(), - kind: SourceKind::GithubRepo, - label: "Repo".into(), - enabled: true, - url: Some("https://github.com/org/repo".into()), - ..default_entry() - }; - assert!(entry.validate().is_ok()); - } - - #[test] - fn validate_conversation_needs_only_id_and_label() { - let entry = MemorySourceEntry { - id: "src_conv".into(), - kind: SourceKind::Conversation, - label: "Agent Conversations".into(), - enabled: true, - ..default_entry() - }; - assert!(entry.validate().is_ok()); - } - - #[test] - fn validate_conversation_fails_with_empty_id() { - let entry = MemorySourceEntry { - id: "".into(), - kind: SourceKind::Conversation, - label: "Convos".into(), - enabled: true, - ..default_entry() - }; - assert!(entry.validate().is_err()); - } - - #[test] - fn validate_conversation_fails_with_empty_label() { - let entry = MemorySourceEntry { - id: "src_conv".into(), - kind: SourceKind::Conversation, - label: "".into(), - enabled: true, - ..default_entry() - }; - assert!(entry.validate().is_err()); - } - - #[test] - fn conversation_kind_serializes_to_snake_case() { - let json = serde_json::to_string(&SourceKind::Conversation).unwrap(); - assert_eq!(json, "\"conversation\""); - } - - #[test] - fn toml_round_trip() { - let entry = MemorySourceEntry { - id: "src_1".into(), - kind: SourceKind::Folder, - label: "My notes".into(), - enabled: true, - path: Some("/tmp/notes".into()), - glob: Some("**/*.md".into()), - ..default_entry() - }; - let toml_str = toml::to_string_pretty(&entry).unwrap(); - let decoded: MemorySourceEntry = toml::from_str(&toml_str).unwrap(); - assert_eq!(decoded.id, "src_1"); - assert_eq!(decoded.kind, SourceKind::Folder); - assert_eq!(decoded.path.as_deref(), Some("/tmp/notes")); - } - - #[test] - fn conversation_toml_round_trip() { - let entry = MemorySourceEntry { - id: "src_conv".into(), - kind: SourceKind::Conversation, - label: "Conversations".into(), - enabled: true, - ..default_entry() - }; - let toml_str = toml::to_string_pretty(&entry).unwrap(); - let decoded: MemorySourceEntry = toml::from_str(&toml_str).unwrap(); - assert_eq!(decoded.id, "src_conv"); - assert_eq!(decoded.kind, SourceKind::Conversation); - assert_eq!(decoded.label, "Conversations"); - assert!(decoded.enabled); - } - - fn default_entry() -> MemorySourceEntry { - MemorySourceEntry { - id: String::new(), - kind: SourceKind::Folder, - label: String::new(), - enabled: true, - toolkit: None, - connection_id: None, - path: None, - glob: None, - 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, - } - } -} +pub use tinycortex::memory::sources::{ + ContentType, MemorySourceEntry, SourceContent, SourceItem, SourceKind, +}; diff --git a/src/openhuman/memory_store/chunks/connection.rs b/src/openhuman/memory_store/chunks/connection.rs index 961b92eff..87fe5a2e4 100644 --- a/src/openhuman/memory_store/chunks/connection.rs +++ b/src/openhuman/memory_store/chunks/connection.rs @@ -1,847 +1,23 @@ -use anyhow::{Context, Result}; -use chrono::Utc; -use parking_lot::Mutex as PMutex; +//! `Config` adapters for tinycortex's chunk connection and recovery manager. + +use anyhow::Result; use rusqlite::Connection; -use std::collections::HashMap; -use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; -#[cfg(test)] -use std::sync::Mutex; -use std::sync::{Arc, OnceLock}; -use std::time::{Duration, Instant}; use crate::openhuman::config::Config; -use super::{ - add_column_if_missing, migrate_legacy_embeddings_to_sidecar, purge_global_topic_trees, DB_DIR, - DB_FILE, SCHEMA, SQLITE_BUSY_TIMEOUT, -}; - -// ── Schema-apply instrumentation (test-only) ───────────────────────────────── -// -// Per-path counter of how many times `apply_schema` ran for each DB path, -// gated behind `cfg(test)` so the production binary carries no overhead. -// Used by the concurrent-init regression test to assert "exactly once per -// path" across racing workers; it survives even when the connection cache -// is cleared between tests because tests use distinct tempdirs. -#[cfg(test)] -static SCHEMA_APPLY_COUNTS: OnceLock>> = OnceLock::new(); - -fn record_schema_apply(_path: &Path) { - #[cfg(test)] - { - let counts = SCHEMA_APPLY_COUNTS.get_or_init(|| Mutex::new(HashMap::new())); - let mut guard = counts - .lock() - .expect("memory_tree schema apply count mutex poisoned"); - *guard.entry(_path.to_path_buf()).or_insert(0) += 1; - } +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -#[cfg(test)] -#[doc(hidden)] -pub(crate) fn schema_apply_count_for_path_for_tests(path: &Path) -> usize { - SCHEMA_APPLY_COUNTS - .get() - .and_then(|m| { - m.lock() - .ok() - .map(|guard| guard.get(path).copied().unwrap_or(0)) - }) - .unwrap_or(0) -} - -// SQLite extended result codes that fire during cold-start WAL/SHM bootstrap -// races. NOTE on values: extended codes are `SQLITE_IOERR (10) | (sub << 8)`. -// 4874 is `IOERR_SHMSIZE` (sub 19), NOT `SHMMAP` — the real `SHMMAP` is 5386 -// (sub 21) and the "open a new shared-memory segment" failure is `SHMOPEN` -// 4618 (sub 18), which is what surfaced on macOS. The whole `-shm` family is -// listed so the classifiers don't miss any of them. -/// `CANTOPEN` — racing the lockfile/WAL creation done by another connection. -const SQLITE_CANTOPEN: i32 = 14; -/// `IOERR_TRUNCATE` — the WAL/db is being truncated during bootstrap. -const SQLITE_IOERR_TRUNCATE: i32 = 1546; -/// `IOERR_SHMOPEN` — opening a new `-shm` shared-memory segment failed (the -/// macOS cold-start failure, e.g. Sentry TAURI-RUST-X1). -const SQLITE_IOERR_SHMOPEN: i32 = 4618; -/// `IOERR_SHMSIZE` — the `-shm` file is being resized during bootstrap. -const SQLITE_IOERR_SHMSIZE: i32 = 4874; -/// `IOERR_SHMMAP` — mapping a page of the `-shm` wal-index failed. -const SQLITE_IOERR_SHMMAP: i32 = 5386; -/// `IOERR_IN_PAGE` — an mmap-page I/O fault, also seen under WAL cold-start. -const SQLITE_IOERR_IN_PAGE: i32 = 8714; - -/// True if `err` (or anything in its cause chain) is one of the SQLite codes -/// that fire during cold-start WAL/SHM bootstrap races: `CANTOPEN`, -/// `IOERR_TRUNCATE`, the `-shm` family (`SHMOPEN` / `SHMSIZE` / `SHMMAP`), and -/// `IOERR_IN_PAGE`. -pub(crate) 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!( - ffi.extended_code, - SQLITE_CANTOPEN - | SQLITE_IOERR_TRUNCATE - | SQLITE_IOERR_SHMOPEN - | SQLITE_IOERR_SHMSIZE - | SQLITE_IOERR_SHMMAP - | SQLITE_IOERR_IN_PAGE - ); - } - false - } - if is_transient_sqlite(err.root_cause()) { - return true; - } - let mut src: Option<&(dyn std::error::Error + 'static)> = Some(err.as_ref()); - while let Some(cur) = src { - if is_transient_sqlite(cur) { - return true; - } - src = cur.source(); - } - false -} - -// ── Connection cache (#2206) ───────────────────────────────────────────────── - -/// How many consecutive init failures before the circuit breaker trips. -pub(crate) const CB_THRESHOLD: u32 = 3; -/// How long the circuit breaker holds the DB closed after tripping. -pub(crate) const CB_COOLDOWN: Duration = Duration::from_secs(30); - -/// Per-path circuit breaker: after [`CB_THRESHOLD`] consecutive init failures -/// the breaker trips and `get_or_init_connection` returns an error immediately -/// until [`CB_COOLDOWN`] elapses. On the first success it resets to zero. -struct CircuitBreaker { - consecutive_failures: AtomicU32, - tripped: AtomicBool, - last_trip: PMutex>, - /// Set once a `SystemStartup` mark has been published for this path so a - /// fresh boot reports a real status without re-emitting on every open. - startup_emitted: AtomicBool, -} - -impl CircuitBreaker { - fn new() -> Self { - Self { - consecutive_failures: AtomicU32::new(0), - tripped: AtomicBool::new(false), - last_trip: PMutex::new(None), - startup_emitted: AtomicBool::new(false), - } - } - - /// Records a successful init. Returns `true` if this call cleared a - /// previously-tripped breaker (i.e. a transition back to healthy that the - /// caller should announce on the bus). Returns `false` for the steady-state - /// case where the breaker was already untripped, so we don't spam a - /// `HealthChanged{healthy:true}` event on every successful call. - fn record_success(&self) -> bool { - self.consecutive_failures.store(0, Ordering::Relaxed); - *self.last_trip.lock() = None; - // `swap` reports the prior value: `true` means we just transitioned - // from tripped → untripped, which is the recovery edge to announce. - self.tripped.swap(false, Ordering::Relaxed) - } - - /// Returns `true` exactly once per breaker — on the first successful open — - /// so the caller emits a single `SystemStartup` mark for this path. - fn mark_startup_emitted(&self) -> bool { - !self.startup_emitted.swap(true, Ordering::Relaxed) - } - - /// Records one more failure. Returns `true` if this call just tripped the - /// breaker (i.e. the threshold was crossed right now). - fn record_failure(&self) -> bool { - let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed); - let count = prev + 1; - if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); - return true; - } - // Re-arm the cooldown on each subsequent failure while already tripped - // so that a failed post-cooldown retry restarts the 30 s window instead - // of leaving the stale timestamp in place (which would let `is_open` - // return false immediately and allow unlimited retries). - if self.tripped.load(Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); - } - false - } - - /// Returns `true` when the breaker is open AND the cooldown has not yet - /// elapsed. Returns `false` (allowing a retry) once the cooldown passes. - fn is_open(&self) -> bool { - if !self.tripped.load(Ordering::Relaxed) { - return false; - } - let guard = self.last_trip.lock(); - match *guard { - Some(t) if t.elapsed() < CB_COOLDOWN => true, - _ => false, - } - } -} - -/// Process-level cache — two separate maps so that a failing init cannot -/// accidentally serve a dummy connection on the fast path. -/// -/// `connections`: only fully-initialised (schema + migrations run) entries. -/// `breakers`: persists across failed init attempts so the circuit breaker -/// survives even when `connections` has no entry for that path. -struct ConnectionCache { - connections: PMutex>>>, - breakers: PMutex>>, - /// Per-path mutex held across the slow-path init so concurrent - /// workers racing into `with_connection` on a cold DB serialise on - /// the WAL+SHM bootstrap. Without this, N threads see "no cached - /// connection" simultaneously and all run `apply_schema`, which is - /// idempotent but reopens the cold-start race window - /// (OPENHUMAN-TAURI-HH / -ZM / -MB). - init_locks: PMutex>>>, -} - -static CONN_CACHE: OnceLock = OnceLock::new(); - -fn conn_cache() -> &'static ConnectionCache { - CONN_CACHE.get_or_init(|| ConnectionCache { - connections: PMutex::new(HashMap::new()), - breakers: PMutex::new(HashMap::new()), - init_locks: PMutex::new(HashMap::new()), - }) -} - -/// Compute the canonical DB path from `config`. -pub(crate) fn db_path_for(config: &Config) -> PathBuf { - config.workspace_dir.join(DB_DIR).join(DB_FILE) -} - -/// Delete stale WAL/SHM side-files (`-shm`, `-wal`) that can block a -/// clean DB open after a crash. Logs what was deleted and returns `true` if -/// anything was removed. -/// -/// SQLite names these files `-shm` and `-wal`. -/// For `chunks.db` that is `chunks.db-shm` / `chunks.db-wal`. -pub(crate) fn try_cleanup_stale_files(db_path: &std::path::Path) -> bool { - let mut cleaned = false; - for suffix in &["-shm", "-wal"] { - // Build the side-file path: append suffix to the full db filename. - let side = { - let mut p = db_path.to_path_buf(); - let name = p - .file_name() - .unwrap_or_default() - .to_string_lossy() - .into_owned(); - p.set_file_name(format!("{name}{suffix}")); - p - }; - if side.exists() { - match std::fs::remove_file(&side) { - Ok(()) => { - log::warn!("[memory_tree] removed stale side-file: {}", side.display()); - cleaned = true; - } - Err(e) => { - log::warn!( - "[memory_tree] failed to remove stale side-file {}: {e}", - side.display() - ); - } - } - } - } - cleaned -} - -/// Run the full one-time DB initialisation (journal mode, schema, migrations) -/// against an already-open `Connection`. Used by `get_or_init_connection`. -fn init_db(conn: &Connection, config: &Config) -> Result<()> { - conn.busy_timeout(SQLITE_BUSY_TIMEOUT) - .context("Failed to configure memory_tree busy timeout")?; - // SQLite resets `foreign_keys` to off on every new connection. The - // ConnectionCache holds one cached `Connection` per DB path, so - // setting it here (alongside the rest of init) is the per-connection - // surface — fast-path callers reuse the cached conn with FKs already - // on. - conn.execute_batch("PRAGMA foreign_keys = ON;") - .context("Failed to enable memory_tree foreign_keys pragma")?; - // memory_tree runs the TRUNCATE rollback journal (see `apply_schema`), so - // crash-safety requires synchronous=FULL — NORMAL is only corruption-safe - // under WAL. Set explicitly so a future global default can't weaken it. - conn.execute_batch("PRAGMA synchronous = FULL;") - .context("Failed to set memory_tree synchronous=FULL")?; - apply_schema(conn)?; - // #1574 §7: one-shot, version-gated legacy→sidecar embedding migration. - migrate_legacy_embeddings_to_sidecar(conn, config)?; - // One-shot, version-gated purge of the removed global/topic trees. - purge_global_topic_trees(conn, config)?; - Ok(()) -} - -fn apply_schema(conn: &Connection) -> Result<()> { - // Note: `init_db` runs the `#1574 §7` legacy→sidecar embedding migration - // after this returns, so the dim-equal copy step is not duplicated here. - // memory_tree uses the TRUNCATE rollback journal, NOT WAL. WAL's `-shm` - // shared-memory index + `-wal` checkpoint machinery are the root of the - // cold-start IOERR_SHMMAP (macOS) / IOERR_TRUNCATE (Windows, AV-held - // handles) failures (Sentry TAURI-RUST-EV / TAURI-RUST-X1). All tree - // access serialises on the single cached `PMutex` (see - // `get_or_init_connection`), so WAL's only real benefit — concurrent - // readers — is unused here, which makes WAL pure liability. The sibling - // tree DBs (cron / vault / redirect_links) already run the default - // rollback journal without issue. - // - // Requesting TRUNCATE on a database a prior release left in WAL mode - // checkpoints the `-wal` back into the main file and removes the - // `-wal`/`-shm` side-files, so this also migrates existing WAL databases - // in place on upgrade. - let journal_mode: String = conn - .query_row("PRAGMA journal_mode=TRUNCATE", [], |row| row.get(0)) - .context("Failed to set memory_tree journal_mode=TRUNCATE")?; - if !journal_mode.eq_ignore_ascii_case("truncate") { - log::warn!( - "[memory_tree] journal_mode is '{journal_mode}' after requesting TRUNCATE \ - — a prior WAL connection or a locked -wal may be blocking the switch" - ); - } - conn.execute_batch(SCHEMA) - .context("Failed to initialize memory_tree schema")?; - // Phase 2 migrations — additive, idempotent. - add_column_if_missing(conn, "mem_tree_chunks", "embedding", "BLOB")?; - // Phase 2 LLM-NER follow-up: per-chunk LLM importance signal + - // human-readable reason. Both nullable; absence is treated as - // "no LLM signal available" by readers. - add_column_if_missing(conn, "mem_tree_score", "llm_importance", "REAL")?; - add_column_if_missing(conn, "mem_tree_score", "llm_importance_reason", "TEXT")?; - // Phase 3a (#709): parent-summary backlink on leaves. - add_column_if_missing(conn, "mem_tree_chunks", "parent_summary_id", "TEXT")?; - // Phase 4 (#710): sealed-summary embeddings for semantic rerank. - add_column_if_missing(conn, "mem_tree_summaries", "embedding", "BLOB")?; - // Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested - // before the queue migration stay queryable. - add_column_if_missing( - conn, - "mem_tree_chunks", - "lifecycle_status", - "TEXT NOT NULL DEFAULT 'admitted'", - )?; - conn.execute_batch( - "CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \ - ON mem_tree_chunks(lifecycle_status);", - ) - .context("Failed to create mem_tree_chunks lifecycle index")?; - // Source grouping scope. Documents can keep item-level source_id for - // dedupe while grouping chunk files and source trees under this scope. - add_column_if_missing(conn, "mem_tree_chunks", "path_scope", "TEXT")?; - // Phase MD-content (#TBD): pointer + integrity hash. - add_column_if_missing(conn, "mem_tree_chunks", "content_path", "TEXT")?; - add_column_if_missing(conn, "mem_tree_chunks", "content_sha256", "TEXT")?; - // Phase MD-content (summaries). - add_column_if_missing(conn, "mem_tree_summaries", "content_path", "TEXT")?; - add_column_if_missing(conn, "mem_tree_summaries", "content_sha256", "TEXT")?; - // Document source-tree versioning: per-doc subtree nodes (Notion etc.) - // carry the document identity + version they sealed for, so retrieval can - // keep `max(version_ms)` per `doc_id` at read time (latest-wins) without - // ever rewriting older subtrees. NULL on merge-tier and chat/email nodes. - add_column_if_missing(conn, "mem_tree_summaries", "doc_id", "TEXT")?; - add_column_if_missing(conn, "mem_tree_summaries", "version_ms", "INTEGER")?; - conn.execute_batch( - "CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_doc_version \ - ON mem_tree_summaries(tree_id, doc_id, version_ms);", - ) - .context("Failed to create mem_tree_summaries doc/version index")?; - // Raw-archive pointer column. - add_column_if_missing(conn, "mem_tree_chunks", "raw_refs_json", "TEXT")?; - // #1365: is_user flag on indexed entity rows. - add_column_if_missing( - conn, - "mem_tree_entity_index", - "is_user", - "INTEGER NOT NULL DEFAULT 0", - )?; - // #002 memory-pipeline-hardening: typed failure metadata on jobs so the - // worker can fail-fast on unrecoverable errors and the status/doctor - // surface can show an actionable cause. Both nullable; only set when a - // job is marked `failed` with a classified reason. - add_column_if_missing(conn, "mem_tree_jobs", "failure_reason", "TEXT")?; - add_column_if_missing(conn, "mem_tree_jobs", "failure_class", "TEXT")?; - Ok(()) -} - -/// Whether `err` looks like one of the I/O error codes that warrant a -/// stale-file cleanup + single retry before giving up. -pub(crate) fn is_io_open_error(err: &anyhow::Error) -> bool { - if let Some(rusqlite::Error::SqliteFailure(f, _)) = err.downcast_ref::() { - return matches!( - f.extended_code, - SQLITE_CANTOPEN - | SQLITE_IOERR_TRUNCATE - | SQLITE_IOERR_SHMOPEN - | SQLITE_IOERR_SHMSIZE - | SQLITE_IOERR_SHMMAP - | SQLITE_IOERR_IN_PAGE - ) || f.code == rusqlite::ErrorCode::CannotOpen; - } - let msg = format!("{err:#}").to_ascii_lowercase(); - msg.contains("disk i/o error") - || msg.contains("unable to open database file") - || msg.contains("xshmmap") - || msg.contains("truncate file") -} - -/// Obtain (or lazily create) a cached connection for the workspace described -/// by `config`. Returns `Err` immediately when the circuit breaker is open. -pub(crate) fn get_or_init_connection(config: &Config) -> Result>> { - let db_path = db_path_for(config); - - // ── Fast path: reject immediately if the breaker is open ───────────── - { - let breakers = conn_cache().breakers.lock(); - if let Some(breaker) = breakers.get(&db_path) { - if breaker.is_open() { - anyhow::bail!( - "[memory_tree] circuit breaker open for {}: too many consecutive init failures", - db_path.display() - ); - } - } - } - // ── Fast path: return cached connection if already initialised ──────── - { - let guard = conn_cache().connections.lock(); - if let Some(conn) = guard.get(&db_path) { - return Ok(Arc::clone(conn)); - } - } - - // ── Slow path: serialise init per-path so concurrent workers don't - // all race into `open_and_init` on a cold DB. - let init_lock = { - let mut guard = conn_cache().init_locks.lock(); - guard - .entry(db_path.clone()) - .or_insert_with(|| Arc::new(PMutex::new(()))) - .clone() - }; - let _init_guard = init_lock.lock(); - - // Re-check the cache once we hold the init lock — another thread - // may have completed init while we were queued. - { - let guard = conn_cache().connections.lock(); - if let Some(conn) = guard.get(&db_path) { - return Ok(Arc::clone(conn)); - } - } - - log::debug!( - "[memory_tree] opening and initialising DB at {}", - db_path.display() - ); - - // Attempt to open + init the connection (dir creation is inside - // `open_and_init` so every failure — including EEXIST on the dir — - // reaches the circuit-breaker recording logic below). On certain I/O - // errors (#2206) we clean up stale WAL/SHM side-files and retry once. - let conn = open_and_init(&db_path, config).or_else(|first_err| { - if is_io_open_error(&first_err) { - log::warn!( - "[memory_tree] I/O error on first open attempt ({}), cleaning stale files and retrying", - first_err - ); - try_cleanup_stale_files(&db_path); - open_and_init(&db_path, config) - } else { - Err(first_err) - } - }); - - match conn { - Ok(conn) => { - let arc_conn = Arc::new(PMutex::new(conn)); - conn_cache() - .connections - .lock() - .insert(db_path.clone(), Arc::clone(&arc_conn)); - // Reset any prior failure counter now that init succeeded. Use (or - // lazily create) the persistent breaker so a clean first boot still - // has somewhere to record the one-shot startup mark. - let breaker = { - let mut guard = conn_cache().breakers.lock(); - guard - .entry(db_path.clone()) - .or_insert_with(|| Arc::new(CircuitBreaker::new())) - .clone() - }; - // Emit a one-time `SystemStartup` so a fresh boot reports a real - // status for `memory_tree_db` instead of "unknown" until the first - // failure. Fires once per path for the process lifetime. - if breaker.mark_startup_emitted() { - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::SystemStartup { - component: "memory_tree_db".to_string(), - }, - ); - } - // Only announce recovery on the transition back to healthy — i.e. - // when the breaker had previously tripped (driving `/health` to a - // permanent 503). Steady-state successes stay silent so we don't - // spam a `HealthChanged{healthy:true}` event on every call. - if breaker.record_success() { - log::info!( - "[memory_tree] circuit breaker recovered for {}: DB init succeeded after a prior trip", - db_path.display() - ); - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::HealthChanged { - component: "memory_tree_db".to_string(), - healthy: true, - message: None, - }, - ); - } - log::debug!("[memory_tree] DB connection cached and ready"); - Ok(arc_conn) - } - Err(err) => { - // Persist the breaker so the failure count accumulates across - // calls even though no connection entry exists yet. - let breaker = { - let mut guard = conn_cache().breakers.lock(); - guard - .entry(db_path.clone()) - .or_insert_with(|| Arc::new(CircuitBreaker::new())) - .clone() - }; - let just_tripped = breaker.record_failure(); - if just_tripped { - log::error!( - "[memory_tree] circuit breaker tripped for {}: {} consecutive init failures", - db_path.display(), - CB_THRESHOLD - ); - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::HealthChanged { - component: "memory_tree_db".to_string(), - healthy: false, - message: Some(format!( - "Schema init failed {CB_THRESHOLD} consecutive times" - )), - }, - ); - } - Err(err) - } - } -} - -/// Ensure the DB directory exists, open the SQLite file, and run the full -/// schema init sequence. All errors (dir creation, file open, schema init) -/// are returned as `Err` so callers can funnel them through the circuit -/// breaker logic in a single place. -fn open_and_init(db_path: &std::path::Path, config: &Config) -> Result { - let dir = db_path.parent().expect("db_path always has a parent"); - std::fs::create_dir_all(dir) - .with_context(|| format!("Failed to create memory_tree dir: {}", dir.display()))?; - let conn = Connection::open(db_path) - .with_context(|| format!("Failed to open memory_tree DB: {}", db_path.display()))?; - init_db(&conn, config) - .with_context(|| format!("Failed to init memory_tree schema: {}", db_path.display()))?; - record_schema_apply(db_path); - Ok(conn) -} - -/// Remove the cached connection for `config`'s workspace (forces a fresh open -/// on the next `with_connection` call). Also clears the breaker so the next -/// open attempt is not immediately rejected. Does nothing if no entry exists. -#[allow(dead_code)] -pub(crate) fn invalidate_connection(config: &Config) { - let db_path = db_path_for(config); - conn_cache().connections.lock().remove(&db_path); - conn_cache().breakers.lock().remove(&db_path); - log::debug!( - "[memory_tree] connection invalidated for {}", - db_path.display() - ); -} - -/// Clear the entire connection cache. For test isolation only. -#[cfg(test)] -pub(crate) fn clear_connection_cache() { - conn_cache().connections.lock().clear(); - conn_cache().breakers.lock().clear(); - conn_cache().init_locks.lock().clear(); -} - -/// Open the memory_tree SQLite DB and run a closure against it. -/// -/// Visible to sibling modules (e.g. `score::store`) so Phase 2 can reuse -/// the same connection setup / schema initialisation without duplication. -/// -/// # Connection caching (#2206) -/// -/// The underlying connection is initialised once per workspace path and then -/// reused from a process-level cache. Schema migrations run exactly once on -/// the first call for a given `config.workspace_dir`. Subsequent calls pay -/// only the cost of a `parking_lot::Mutex` lock and the closure itself. -/// -/// `#[doc(hidden)] pub` (not `pub(crate)`) because the -/// `memory-tree-init-smoke` bin in `src/bin/` is a separate crate target -/// and must reach this entry point. It is NOT a stable API surface — -/// downstream crates should treat it as internal. #[doc(hidden)] pub fn with_connection(config: &Config, f: impl FnOnce(&Connection) -> Result) -> Result { - // W3 connection foundation: route ALL production access to - // `/memory_tree/chunks.db` through the TinyCortex connection - // manager. The crate opens the SAME file (its `db_path_for` derives the - // identical `/memory_tree/chunks.db`), applies the identical - // schema + version-gated migrations (`TREE_EMBEDDING`=1, `GLOBAL_TOPIC_PURGE` - // =2 — matched on both sides, so an already-migrated DB is a no-op), and - // migrates any pre-existing WAL database to the TRUNCATE rollback journal in - // place on first open. The former host connection cache below is retired in - // the deletion-ledger follow-up (it is now only referenced by cache-behaviour - // unit tests, which move upstream into the crate). - let mc = crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); - tinycortex::memory::chunks::with_connection(&mc, f) + tinycortex::memory::chunks::with_connection(&engine_config(config), f) } -/// Append `suffix` to the *file name* of `path` (so `chunks.db` + `-wal` -/// = `chunks.db-wal`, and `chunks.db` + `.corrupt-…` = `chunks.db.corrupt-…`). -/// SQLite names its side-files this way (not as a new extension), and the -/// quarantine keeps the corrupt image alongside the original for inspection. -fn with_name_suffix(path: &Path, suffix: &str) -> PathBuf { - let mut p = path.to_path_buf(); - let name = p - .file_name() - .unwrap_or_default() - .to_string_lossy() - .into_owned(); - p.set_file_name(format!("{name}{suffix}")); - p -} - -/// Run `PRAGMA quick_check(1)` against `db_path` on a fresh, short-lived -/// connection. Returns `Ok(true)` when the structural scan reports `"ok"`, -/// `Ok(false)` when it reports any corruption, and `Err` when the check itself -/// can't run (file unopenable / header unreadable — itself a corruption signal -/// the caller treats as malformed). -fn quick_check_ok(db_path: &Path) -> Result { - let conn = Connection::open(db_path) - .with_context(|| format!("open for quick_check: {}", db_path.display()))?; - let _ = conn.busy_timeout(SQLITE_BUSY_TIMEOUT); - let result: String = conn - .query_row("PRAGMA quick_check(1)", [], |row| row.get(0)) - .context("running PRAGMA quick_check")?; - Ok(result.eq_ignore_ascii_case("ok")) -} - -/// Recover from a `SQLITE_CORRUPT` (malformed image) on the memory_tree DB. -/// -/// Unlike the transient/contention/disk-full classes, a malformed on-disk -/// image never heals on its own — every query fails forever and the worker -/// re-pages Sentry on each poll (Sentry TAURI-RUST-E93: ~1.6k events in ~17 min -/// from a single host). This is the recovery lever the sibling suppressors -/// lack: it quarantines the damaged file (and its WAL/SHM side-files) to a -/// timestamped `.corrupt-` copy — **preserved, not deleted**, so the bytes -/// can still be inspected or salvaged — then rebuilds an empty schema so the -/// memory-tree queue resumes instead of wedging indefinitely. -/// -/// Returns `Ok(true)` when a quarantine + rebuild happened, `Ok(false)` when a -/// fresh `PRAGMA quick_check` now passes (the earlier failure was transient and -/// quarantining would have destroyed good data), and `Err` when the quarantine -/// rename or the schema rebuild failed (caller backs off and retries). pub(crate) fn recover_corrupt_db(config: &Config) -> Result { - let db_path = db_path_for(config); - - // 1. Drop any cached (corrupt) connection + breaker so the OS file handle - // is closed before we rename, and the next open re-inits cleanly. - conn_cache().connections.lock().remove(&db_path); - conn_cache().breakers.lock().remove(&db_path); - - // 2. Re-confirm corruption against the on-disk file. `quick_check` is the - // cheap structural scan; if it now reports "ok" the image is actually - // healthy (e.g. the original error was a transient mmap fault) and we - // must NOT destroy good data — bail out without quarantining. - if db_path.exists() { - match quick_check_ok(&db_path) { - Ok(true) => { - log::info!( - "[memory_tree] quick_check passed for {} — no quarantine needed", - db_path.display() - ); - return Ok(false); - } - Ok(false) => { - log::warn!( - "[memory_tree] quick_check confirms corruption for {}, quarantining", - db_path.display() - ); - } - Err(e) => { - // The check couldn't even run (unopenable / unreadable header). - // That is itself a malformed-image signal — treat as corrupt. - log::warn!( - "[memory_tree] quick_check could not run for {} ({e:#}); treating as corrupt", - db_path.display() - ); - } - } - } else { - log::warn!( - "[memory_tree] corrupt-recovery: {} is missing; rebuilding fresh schema", - db_path.display() - ); - } - - // 3. Quarantine the main DB + WAL/SHM side-files to `.corrupt-`. - let ts = Utc::now().format("%Y%m%dT%H%M%SZ").to_string(); - let mut quarantined = 0usize; - for suffix in &["", "-wal", "-shm"] { - let src = with_name_suffix(&db_path, suffix); - if !src.exists() { - continue; - } - let dst = with_name_suffix(&src, &format!(".corrupt-{ts}")); - std::fs::rename(&src, &dst).with_context(|| { - format!( - "failed to quarantine corrupt memory_tree file {} -> {}", - src.display(), - dst.display() - ) - })?; - log::warn!( - "[memory_tree] quarantined {} -> {}", - src.display(), - dst.display() - ); - quarantined += 1; - } - - // 4. Rebuild an empty schema by forcing a fresh open. The damaged rows are - // not silently dropped — they live on in the `.corrupt-` copy. - get_or_init_connection(config) - .context("failed to rebuild memory_tree schema after quarantining corrupt DB")?; - - log::warn!( - "[memory_tree] corruption recovery complete: quarantined {quarantined} file(s), \ - rebuilt empty schema at {}", - db_path.display() - ); - Ok(true) + log::warn!("[memory:chunks] checking corrupt database recovery"); + tinycortex::memory::chunks::recover_corrupt_db(&engine_config(config)) } #[cfg(test)] -mod tests { - use super::*; - - /// `record_success` must only report a recovery transition (`true`) when - /// it actually clears a tripped breaker — the signal `get_or_init_connection` - /// uses to publish `HealthChanged{healthy:true}` exactly once instead of on - /// every successful call (C1). - #[test] - fn record_success_announces_only_on_trip_to_healthy_transition() { - let cb = CircuitBreaker::new(); - - // Untripped breaker: a success is steady-state, not a transition. - assert!(!cb.record_success()); - - // Trip the breaker by crossing the failure threshold. - let mut tripped = false; - for _ in 0..CB_THRESHOLD { - tripped = cb.record_failure(); - } - assert!(tripped, "breaker should trip at CB_THRESHOLD failures"); - - // First success after a trip is the recovery edge → announce once. - assert!(cb.record_success()); - // Subsequent successes are steady-state → stay silent. - assert!(!cb.record_success()); - } - - /// `mark_startup_emitted` must fire exactly once so a fresh boot emits a - /// single `SystemStartup` mark for `memory_tree_db` (C1). - #[test] - fn startup_mark_fires_exactly_once() { - let cb = CircuitBreaker::new(); - assert!(cb.mark_startup_emitted()); - assert!(!cb.mark_startup_emitted()); - assert!(!cb.mark_startup_emitted()); - } - - // ── recover_corrupt_db tests (TAURI-RUST-E93 / #4048) ──────────────────── - - fn corrupt_test_config() -> (tempfile::TempDir, Config) { - let tmp = tempfile::TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) - } - - /// A malformed on-disk image must be quarantined (not deleted) and replaced - /// by a fresh, queryable schema so the memory-tree queue resumes. - #[test] - fn recover_corrupt_db_quarantines_and_rebuilds() { - clear_connection_cache(); - let (_tmp, cfg) = corrupt_test_config(); - let db_path = db_path_for(&cfg); - std::fs::create_dir_all(db_path.parent().unwrap()).unwrap(); - // Garbage bytes → not a valid SQLite header → corrupt image. - std::fs::write(&db_path, b"this is not a sqlite database, it is garbage").unwrap(); - - let recovered = recover_corrupt_db(&cfg).expect("recovery should succeed"); - assert!(recovered, "garbage image must be quarantined + rebuilt"); - - // The corrupt bytes are preserved alongside, not silently dropped. - let quarantined: Vec<_> = std::fs::read_dir(db_path.parent().unwrap()) - .unwrap() - .filter_map(|e| e.ok()) - .filter(|e| { - e.file_name() - .to_string_lossy() - .contains("chunks.db.corrupt-") - }) - .collect(); - assert_eq!( - quarantined.len(), - 1, - "exactly one quarantined copy should exist" - ); - - // The rebuilt DB is healthy and the jobs table is queryable + empty. - clear_connection_cache(); - let count: i64 = with_connection(&cfg, |conn| { - conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0)) - .context("count jobs") - }) - .expect("rebuilt DB must be queryable"); - assert_eq!(count, 0, "rebuilt jobs table starts empty"); - } - - /// A healthy DB must NOT be quarantined — `quick_check` passes, so good data - /// is preserved and recovery is a no-op returning `Ok(false)`. - #[test] - fn recover_corrupt_db_is_noop_on_healthy_db() { - clear_connection_cache(); - let (_tmp, cfg) = corrupt_test_config(); - // Force a healthy DB into existence. - with_connection(&cfg, |conn| { - conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| { - r.get::<_, i64>(0) - }) - .context("seed healthy db") - }) - .unwrap(); - - let recovered = recover_corrupt_db(&cfg).expect("recovery should succeed"); - assert!(!recovered, "healthy DB must not be quarantined"); - - let db_path = db_path_for(&cfg); - let quarantined = std::fs::read_dir(db_path.parent().unwrap()) - .unwrap() - .filter_map(|e| e.ok()) - .any(|e| e.file_name().to_string_lossy().contains(".corrupt-")); - assert!(!quarantined, "no quarantine file should be created"); - } -} +pub(crate) use tinycortex::memory::chunks::{is_transient_cold_start, try_cleanup_stale_files}; diff --git a/src/openhuman/memory_store/chunks/embeddings.rs b/src/openhuman/memory_store/chunks/embeddings.rs index 533e68e73..703c541a7 100644 --- a/src/openhuman/memory_store/chunks/embeddings.rs +++ b/src/openhuman/memory_store/chunks/embeddings.rs @@ -1,400 +1,113 @@ -use super::with_connection; -use crate::openhuman::config::Config; -use anyhow::{Context, Result}; -use chrono::Utc; -use rusqlite::{Connection, OptionalExtension}; +//! `Config` adapters for tinycortex embedding sidecars and tombstones. + use std::collections::HashMap; -// ── Phase 2: embedding column accessors ───────────────────────────────── +use anyhow::Result; +use rusqlite::{Connection, Transaction}; + +use crate::openhuman::config::Config; + +#[cfg(test)] +pub use tinycortex::memory::chunks::{embedding_to_blob, REEMBED_SKIP_KEY_MAX_LEN}; + +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} -/// Resolve the active embedding signature for the memory tree from the global -/// [`Config`] — the canonical key every per-model sidecar read/write is scoped -/// by (#1574). Reuses the established local-AI workload derivation -/// ([`Config::workload_local_model`]) and the probe-stable -/// `active_embedding_signature`; introduces no parallel resolution path. -/// `pub(crate)` so the sibling `tree` summary store shares the exact -/// same resolution. pub(crate) fn tree_active_signature(config: &Config) -> String { - let local_model = config.workload_local_model("embeddings"); - crate::openhuman::memory_store::active_embedding_signature( - &config.memory, - local_model.as_deref(), - ) + tinycortex::memory::chunks::tree_active_signature(&engine_config(config)) } -/// Store a chunk's embedding under the active model signature. -/// -/// #1574 cutover: this now writes the per-model `mem_tree_chunk_embeddings` -/// sidecar (via [`set_chunk_embedding_for_signature`]) instead of the legacy -/// `mem_tree_chunks.embedding` column. Call sites are unchanged — the signature -/// is resolved internally from `config`. The legacy column is left intact for -/// the §7 one-shot migration to read; it is dropped only in a later release. -pub fn set_chunk_embedding(config: &Config, chunk_id: &str, embedding: &[f32]) -> Result<()> { - let signature = tree_active_signature(config); - log::debug!( - "[memory::chunk_store] set_chunk_embedding: chunk_id={chunk_id} sig={signature} dims={}", - embedding.len() - ); - set_chunk_embedding_for_signature(config, chunk_id, &signature, embedding) +pub fn set_chunk_embedding(config: &Config, id: &str, embedding: &[f32]) -> Result<()> { + tinycortex::memory::chunks::set_chunk_embedding(&engine_config(config), id, embedding) } -/// Core upsert into `mem_tree_chunk_embeddings` over an arbitrary -/// `&Connection`. Shared by the standalone ([`set_chunk_embedding_for_signature`]) -/// and in-transaction ([`set_chunk_embedding_for_signature_tx`]) write paths so -/// the SQL exists exactly once. `rusqlite::Transaction` derefs to `Connection`, -/// so an in-tx caller passes `&tx` and the sidecar row commits atomically with -/// the surrounding work (#1574 write-side cutover). -fn upsert_chunk_embedding_conn( - conn: &rusqlite::Connection, - chunk_id: &str, - model_signature: &str, - embedding: &[f32], -) -> Result<()> { - let bytes = embedding_to_blob(embedding); - let dim = i64::try_from(embedding.len()).context("embedding dimension does not fit i64")?; - let created_at = Utc::now().timestamp_millis() as f64 / 1000.0; - conn.execute( - "INSERT INTO mem_tree_chunk_embeddings - (chunk_id, model_signature, vector, dim, created_at) - VALUES (?1, ?2, ?3, ?4, ?5) - ON CONFLICT(chunk_id, model_signature) DO UPDATE SET - vector = excluded.vector, - dim = excluded.dim, - created_at = excluded.created_at", - rusqlite::params![chunk_id, model_signature, bytes, dim, created_at], - )?; - Ok(()) -} - -/// Store a chunk embedding for a specific provider/model/dimension signature. -/// -/// Per-model table write path for #1574. The legacy -/// `mem_tree_chunks.embedding` column is intentionally left untouched by this -/// helper (read by the §7 migration; dropped only in a later release). pub fn set_chunk_embedding_for_signature( config: &Config, - chunk_id: &str, - model_signature: &str, + id: &str, + signature: &str, embedding: &[f32], ) -> Result<()> { - with_connection(config, |conn| { - upsert_chunk_embedding_conn(conn, chunk_id, model_signature, embedding) - }) -} - -/// `true` when at least one chunk or summary still needs an embedding at -/// `model_signature` and is not tombstoned as terminally unembeddable. -/// -/// Shared by `ensure_reembed_backfill`, the §7 migration enqueue probe, and -/// tests so the worklist and coverage probes cannot drift (#2358). -pub(crate) fn has_uncovered_reembed_work( - conn: &Connection, - model_signature: &str, -) -> 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], - |r| r.get(0), + tinycortex::memory::chunks::set_chunk_embedding_for_signature( + &engine_config(config), + id, + signature, + embedding, ) } -/// Persistently record that `(chunk_id, signature)` cannot be re-embedded. -/// -/// Called by `handle_reembed_backfill` when the per-chunk body file is -/// missing on disk (orphan) or the embedder rejects the row terminally -/// (wrong dim / unrecoverable embed error). Inserting a row here causes -/// the next backfill batch's worklist query to exclude this chunk via the -/// `NOT EXISTS … mem_tree_chunk_reembed_skipped …` predicate, so the -/// runaway "skipping" loop terminates instead of revisiting the same row -/// every 5 s forever (#1574 §6 fix). +pub(crate) fn has_uncovered_reembed_work( + conn: &Connection, + signature: &str, +) -> rusqlite::Result { + tinycortex::memory::chunks::has_uncovered_reembed_work(conn, signature) +} + pub fn mark_chunk_reembed_skipped( config: &Config, - chunk_id: &str, - model_signature: &str, + id: &str, + signature: &str, reason: &str, ) -> Result<()> { - let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?; - let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - conn.execute( - "INSERT INTO mem_tree_chunk_reembed_skipped - (chunk_id, model_signature, reason, skipped_at_ms) - VALUES (?1, ?2, ?3, ?4) - ON CONFLICT(chunk_id, model_signature) DO UPDATE SET - reason = excluded.reason, - skipped_at_ms = excluded.skipped_at_ms", - rusqlite::params![chunk_id, model_signature, reason, now_ms], - )?; - log::debug!( - "[memory::chunk_store] mark_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature} reason={reason}" - ); - Ok(()) - }) + tinycortex::memory::chunks::mark_chunk_reembed_skipped( + &engine_config(config), + id, + signature, + reason, + ) } -/// Remove a single chunk tombstone so re-embed backfill can retry the row. -/// -/// Idempotent: deleting a missing `(chunk_id, model_signature)` pair is a -/// no-op. Intended for operator recovery after environmental failures (moved -/// workspace, restored body files, fixed embedder config) — see #2358. -pub fn clear_chunk_reembed_skipped( - config: &Config, - chunk_id: &str, - model_signature: &str, -) -> Result<()> { - let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?; - let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; - with_connection(config, |conn| { - conn.execute( - "DELETE FROM mem_tree_chunk_reembed_skipped - WHERE chunk_id = ?1 AND model_signature = ?2", - rusqlite::params![chunk_id, model_signature], - )?; - log::debug!( - "[memory::chunk_store] clear_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature}" - ); - Ok(()) - }) +pub fn clear_chunk_reembed_skipped(config: &Config, id: &str, signature: &str) -> Result<()> { + tinycortex::memory::chunks::clear_chunk_reembed_skipped(&engine_config(config), id, signature) } -/// Clear all chunk and summary tombstones for a model signature. -/// -/// Returns the total number of rows removed across both tombstone tables. -/// Idempotent when no tombstones exist for the signature. -pub fn clear_reembed_skipped_for_signature( - config: &Config, - model_signature: &str, -) -> Result { - let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; - with_connection(config, |conn| { - let chunk_deleted = conn.execute( - "DELETE FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1", - rusqlite::params![model_signature], - )?; - let summary_deleted = conn.execute( - "DELETE FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1", - rusqlite::params![model_signature], - )?; - log::debug!( - "[memory::chunk_store] clear_reembed_skipped_for_signature sig={model_signature} chunk_rows={chunk_deleted} summary_rows={summary_deleted}" - ); - Ok(chunk_deleted + summary_deleted) - }) +pub fn clear_reembed_skipped_for_signature(config: &Config, signature: &str) -> Result { + tinycortex::memory::chunks::clear_reembed_skipped_for_signature( + &engine_config(config), + signature, + ) } -/// Bounds attacker-controlled ids/signatures passed to reembed-skipped admin -/// helpers without affecting legitimate rows (typical ids are well under 512 -/// chars). Rejects NUL bytes so SQLite bindings cannot be truncated. -pub(crate) const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048; - -pub(crate) fn validate_reembed_skip_key<'a>(label: &str, value: &'a str) -> Result<&'a str> { - let trimmed = value.trim(); - if trimmed.is_empty() { - anyhow::bail!("{label} must be non-empty"); - } - if trimmed.len() > REEMBED_SKIP_KEY_MAX_LEN { - anyhow::bail!("{label} exceeds maximum length ({REEMBED_SKIP_KEY_MAX_LEN})"); - } - if trimmed.as_bytes().contains(&0) { - anyhow::bail!("{label} must not contain NUL bytes"); - } - Ok(trimmed) -} - -/// Transaction-scoped variant of [`set_chunk_embedding_for_signature`]. -/// -/// For callers that already hold a `Transaction` (e.g. the chunk-admission -/// handler, which commits the sidecar row in the SAME tx as the lifecycle -/// + score + job-enqueue writes — #1574 write-side cutover). Opening a fresh -/// connection there would break atomicity / deadlock on the busy DB. pub(crate) fn set_chunk_embedding_for_signature_tx( - tx: &rusqlite::Transaction<'_>, - chunk_id: &str, - model_signature: &str, + tx: &Transaction<'_>, + id: &str, + signature: &str, embedding: &[f32], ) -> Result<()> { - upsert_chunk_embedding_conn(tx, chunk_id, model_signature, embedding) + tinycortex::memory::chunks::set_chunk_embedding_for_signature_tx(tx, id, signature, embedding) } -/// Fetch a chunk embedding for exactly one provider/model/dimension signature. pub fn get_chunk_embedding_for_signature( config: &Config, - chunk_id: &str, - model_signature: &str, + id: &str, + signature: &str, ) -> Result>> { - with_connection(config, |conn| { - let row: Option<(Vec, i64)> = conn - .query_row( - "SELECT vector, dim - FROM mem_tree_chunk_embeddings - WHERE chunk_id = ?1 AND model_signature = ?2", - rusqlite::params![chunk_id, model_signature], - |r| Ok((r.get(0)?, r.get(1)?)), - ) - .optional()?; - match row { - None => Ok(None), - Some((bytes, dim)) => embedding_from_blob(&bytes, dim, "chunk embedding"), - } - }) + tinycortex::memory::chunks::get_chunk_embedding_for_signature( + &engine_config(config), + id, + signature, + ) } -/// Fetch a chunk's embedding for the active model signature. -/// -/// #1574 cutover: reads the per-model `mem_tree_chunk_embeddings` sidecar at -/// the active signature (via [`get_chunk_embedding_for_signature`]) instead of -/// the legacy `mem_tree_chunks.embedding` column. Returns `Ok(None)` if the -/// chunk has no vector under the active signature — e.g. during the §7 -/// backfill window, where this degrades retrieval gracefully (the row is -/// simply absent from vector results, never cross-space compared). -pub fn get_chunk_embedding(config: &Config, chunk_id: &str) -> Result>> { - let signature = tree_active_signature(config); - get_chunk_embedding_for_signature(config, chunk_id, &signature) +pub fn get_chunk_embedding(config: &Config, id: &str) -> Result>> { + tinycortex::memory::chunks::get_chunk_embedding(&engine_config(config), id) } -pub(crate) fn embedding_to_blob(embedding: &[f32]) -> Vec { - embedding.iter().flat_map(|f| f.to_le_bytes()).collect() -} - -fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result>> { - if dim < 0 { - anyhow::bail!("{label} has negative dimension {dim}"); - } - if !bytes.len().is_multiple_of(4) { - anyhow::bail!("{label} blob length {} not a multiple of 4", bytes.len()); - } - let floats: Vec = bytes - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - if floats.len() != dim as usize { - anyhow::bail!( - "{label} dimension mismatch: dim column says {dim}, blob contains {} floats", - floats.len() - ); - } - Ok(Some(floats)) -} -/// SQLite's compile-time hard cap on `?` parameters per prepared statement is -/// `SQLITE_MAX_VARIABLE_NUMBER` (32 766 since SQLite 3.32, bundled with -/// rusqlite). We chunk well below that for two reasons: -/// -/// 1. **Headroom against schema growth.** The batched query binds one `?` per -/// chunk_id *plus* the `model_signature` parameter. Picking 500 as the -/// per-chunk cap gives ~65× safety margin even if a future caller bumps -/// `LOOKUP_HEADROOM` from 200 into the thousands — the batch helper -/// silently splits the request, the caller sees no semantic change. -/// 2. **Sane SQL string length.** A 500-`?` `IN(...)` clause is ~1 KB of SQL -/// text — small enough that prepare-and-discard per chunk is cheap and the -/// planner produces a simple covered-index lookup. -/// -/// Lowering this constant is safe (just more round-trips). Raising it past -/// the SQLite limit will fail at prepare time with `too many SQL variables`. -const MAX_EMBEDDING_BATCH: usize = 500; - -/// Batched read of chunk embeddings under a single `model_signature`. -/// -/// Returns a `HashMap>` containing **only the chunks -/// that have a vector under `model_signature`**. Missing chunks are simply -/// absent from the map — callers handle them the same way as a `None` -/// return from the single-row [`get_chunk_embedding_for_signature`]. -/// -/// ## Why this exists -/// -/// The retrieval rerank path (`memory_tree::retrieval::topic:: -/// rerank_by_semantic_similarity`) used to call the single-row helper inside -/// a `for h in hits { spawn_blocking(...).await }` loop. With -/// `LOOKUP_HEADROOM = 200` that meant 200 sequential SQLite round-trips on -/// every entity-scoped query with `query=…`. This helper collapses that to -/// `ceil(n / MAX_EMBEDDING_BATCH)` round-trips while preserving the exact -/// `Option>` semantics of the per-row helper (missing row → absent -/// key → caller treats as `None`). -/// -/// Order of input ids is irrelevant; callers re-decorate from the returned -/// map by id, so the rerank loop preserves its original hit iteration -/// order and therefore its tie-break behaviour. pub fn get_chunk_embeddings_for_signature_batch( config: &Config, - chunk_ids: &[String], - model_signature: &str, + ids: &[String], + signature: &str, ) -> Result>> { - if chunk_ids.is_empty() { - return Ok(HashMap::new()); - } - with_connection(config, |conn| { - let mut out: HashMap> = HashMap::with_capacity(chunk_ids.len()); - // Chunk the id list to stay safely under SQLite's - // SQLITE_MAX_VARIABLE_NUMBER cap. For the current - // LOOKUP_HEADROOM=200 callsite this loop executes exactly once; - // the chunking only kicks in if a future caller passes >500 - // ids in a single batch. - for window in chunk_ids.chunks(MAX_EMBEDDING_BATCH) { - // Build `IN (?,?,?,...)` with `window.len()` placeholders. - // model_signature is bound as the last parameter (?{n+1}). - let placeholders = std::iter::repeat_n("?", window.len()) - .collect::>() - .join(","); - let sql = format!( - "SELECT chunk_id, vector, dim - FROM mem_tree_chunk_embeddings - WHERE chunk_id IN ({placeholders}) - AND model_signature = ?{sig_idx}", - sig_idx = window.len() + 1, - ); - let mut stmt = conn - .prepare(&sql) - .context("prepare get_chunk_embeddings_for_signature_batch")?; - // Bind chunk_ids then model_signature. rusqlite wants - // ToSql trait objects in a uniform iterator. - let mut params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); - for id in window { - params.push(id as &dyn rusqlite::ToSql); - } - params.push(&model_signature as &dyn rusqlite::ToSql); - let rows = stmt - .query_map(params.as_slice(), |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, Vec>(1)?, - row.get::<_, i64>(2)?, - )) - }) - .context("query get_chunk_embeddings_for_signature_batch")?; - for row in rows { - let (chunk_id, bytes, dim) = row?; - // Reuse the single-row decoder so corrupt-blob errors - // surface with the same diagnostic shape as the - // single-row path. - if let Some(v) = embedding_from_blob(&bytes, dim, "chunk embedding")? { - out.insert(chunk_id, v); - } - } - } - Ok(out) - }) + tinycortex::memory::chunks::get_chunk_embeddings_for_signature_batch( + &engine_config(config), + ids, + signature, + ) } -/// Batched read of chunk embeddings under the **active** model signature. -/// Convenience wrapper mirroring [`get_chunk_embedding`] for the per-row -/// path: resolves `tree_active_signature(config)` exactly once and forwards -/// to [`get_chunk_embeddings_for_signature_batch`]. pub fn get_chunk_embeddings_batch( config: &Config, - chunk_ids: &[String], + ids: &[String], ) -> Result>> { - let signature = tree_active_signature(config); - get_chunk_embeddings_for_signature_batch(config, chunk_ids, &signature) + tinycortex::memory::chunks::get_chunk_embeddings_batch(&engine_config(config), ids) } diff --git a/src/openhuman/memory_store/chunks/migrations.rs b/src/openhuman/memory_store/chunks/migrations.rs deleted file mode 100644 index 587ffe8dc..000000000 --- a/src/openhuman/memory_store/chunks/migrations.rs +++ /dev/null @@ -1,201 +0,0 @@ -//! One-shot SQLite migrations for the chunks DB. -//! -//! These functions are called from [`super::connection`] during DB initialisation. -//! Each migration is version-gated via `PRAGMA user_version` so it runs exactly -//! once per vault. - -use anyhow::{Context, Result}; -use rusqlite::Connection; - -use super::{ - has_uncovered_reembed_work, set_chunk_embedding_for_signature_tx, - GLOBAL_TOPIC_PURGE_MIGRATION_VERSION, TREE_EMBEDDING_MIGRATION_VERSION, -}; -use crate::openhuman::config::Config; - -/// One-shot migration: copy legacy per-chunk/summary `.embedding` blobs into the -/// normalised `mem_tree_chunk_embeddings` / `mem_tree_summary_embeddings` sidecar -/// tables introduced in #1574. -/// -/// Version-gated: `PRAGMA user_version < 1` triggers the copy; `>= 1` is a no-op. -pub(super) fn migrate_legacy_embeddings_to_sidecar( - conn: &Connection, - config: &Config, -) -> Result<()> { - let version: i64 = conn - .query_row("PRAGMA user_version", [], |r| r.get(0)) - .context("read PRAGMA user_version for #1574 migration")?; - if version >= TREE_EMBEDDING_MIGRATION_VERSION { - return Ok(()); - } - - let (provider, model, dims) = crate::openhuman::memory_store::effective_embedding_settings( - &config.memory, - config.workload_local_model("embeddings").as_deref(), - ); - let sig = crate::openhuman::embeddings::format_embedding_signature(&provider, &model, dims); - log::info!( - "[memory_tree::migrate] #1574 §7: copying legacy embeddings → sidecar at sig={sig} (dims={dims})" - ); - - let tx = conn.unchecked_transaction()?; - let mut copied_chunks = 0usize; - let mut copied_summaries = 0usize; - let mut skipped_dim_mismatch = 0usize; - - for (table, is_chunk) in [("mem_tree_chunks", true), ("mem_tree_summaries", false)] { - let mut stmt = tx.prepare(&format!( - "SELECT id, embedding FROM {table} WHERE embedding IS NOT NULL" - ))?; - let rows = stmt.query_map([], |r| { - Ok((r.get::<_, String>(0)?, r.get::<_, Vec>(1)?)) - })?; - for row in rows { - let (id, blob) = row?; - if !blob.len().is_multiple_of(4) { - log::warn!( - "[memory_tree::migrate] {table} id={id}: legacy blob len {} not /4, skipping", - blob.len() - ); - continue; - } - if blob.len() / 4 != dims { - // Different embedding space — unrecoverable from the blob. - // Leave for the §6 re-embed backfill. - skipped_dim_mismatch += 1; - continue; - } - let vec: Vec = blob - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - if is_chunk { - set_chunk_embedding_for_signature_tx(&tx, &id, &sig, &vec)?; - copied_chunks += 1; - } else { - crate::openhuman::memory_store::trees::store::set_summary_embedding_for_signature_tx( - &tx, &id, &sig, &vec, - )?; - copied_summaries += 1; - } - } - } - - // #1574 §6: enqueue the re-embed backfill ONLY if there is genuinely - // uncovered work at the active signature (the dim-mismatch slice, or - // content-bearing rows with no vector). Gating this avoids queuing a - // no-op job on every DB open — which would otherwise pollute the jobs - // table for unrelated callers/tests. Enqueued atomically with the - // migration; dedupe key = signature, so exactly one chain per space. - let has_uncovered = has_uncovered_reembed_work(&tx, &sig)?; - if has_uncovered { - let backfill_job = crate::openhuman::memory_queue::types::NewJob::reembed_backfill( - &crate::openhuman::memory_queue::types::ReembedBackfillPayload { - signature: sig.clone(), - }, - )?; - crate::openhuman::memory_queue::enqueue_tx(&tx, &backfill_job)?; - } - - tx.commit()?; - conn.pragma_update(None, "user_version", TREE_EMBEDDING_MIGRATION_VERSION) - .context("set PRAGMA user_version after #1574 migration")?; - if has_uncovered { - crate::openhuman::memory_queue::set_backfill_in_progress(true); - } - log::info!( - "[memory_tree::migrate] #1574 §7 done: copied chunks={copied_chunks} summaries={copied_summaries} \ - skipped_dim_mismatch={skipped_dim_mismatch} (left for §6 re-embed); user_version={TREE_EMBEDDING_MIGRATION_VERSION}" - ); - Ok(()) -} - -/// One-shot purge of the removed global + topic trees. -/// -/// The global (time-axis) and topic (subject-axis) trees were deleted in -/// favour of the source trees (which hold all content). This migration -/// removes their now-orphaned DB rows and on-disk summary folders so old -/// vaults clean themselves up on next open. Version-gated via -/// `PRAGMA user_version` (see [`GLOBAL_TOPIC_PURGE_MIGRATION_VERSION`]); a -/// no-op on workspaces that never had those trees. -pub(super) fn purge_global_topic_trees(conn: &Connection, config: &Config) -> Result<()> { - let version: i64 = conn - .query_row("PRAGMA user_version", [], |r| r.get(0)) - .context("read PRAGMA user_version for global/topic purge")?; - if version >= GLOBAL_TOPIC_PURGE_MIGRATION_VERSION { - return Ok(()); - } - - let tx = conn.unchecked_transaction()?; - // Child rows first (summary sidecars / skip-lists are keyed by - // summary_id; entity-index + buffers carry an FK on tree_id). - let removed_summary_sidecars = tx.execute( - "DELETE FROM mem_tree_summary_embeddings WHERE summary_id IN \ - (SELECT id FROM mem_tree_summaries WHERE tree_kind IN ('global','topic'))", - [], - )?; - tx.execute( - "DELETE FROM mem_tree_summary_reembed_skipped WHERE summary_id IN \ - (SELECT id FROM mem_tree_summaries WHERE tree_kind IN ('global','topic'))", - [], - )?; - tx.execute( - "DELETE FROM mem_tree_entity_index WHERE tree_id IN \ - (SELECT id FROM mem_tree_trees WHERE kind IN ('global','topic'))", - [], - )?; - let removed_summaries = tx.execute( - "DELETE FROM mem_tree_summaries WHERE tree_kind IN ('global','topic')", - [], - )?; - tx.execute( - "DELETE FROM mem_tree_buffers WHERE tree_id IN \ - (SELECT id FROM mem_tree_trees WHERE kind IN ('global','topic'))", - [], - )?; - let removed_trees = tx.execute( - "DELETE FROM mem_tree_trees WHERE kind IN ('global','topic')", - [], - )?; - // Drain any queued jobs for the retired kinds so the worker loop never - // trips over a payload it can no longer parse. - let removed_jobs = tx.execute( - "DELETE FROM mem_tree_jobs WHERE kind IN ('topic_route','digest_daily')", - [], - )?; - tx.commit()?; - - // On-disk: drop the `wiki/summaries/global*` (both the legacy per-day - // `global-/` folders and the singleton `global/`) and `topic-*` - // summary folders. Best-effort — a filesystem error must not abort the - // version bump, or the purge would retry forever. - let summaries_root = config - .memory_tree_content_root() - .join("wiki") - .join("summaries"); - let mut removed_dirs = 0usize; - if let Ok(entries) = std::fs::read_dir(&summaries_root) { - for entry in entries.flatten() { - let name = entry.file_name(); - let name = name.to_string_lossy(); - if name.starts_with("global") || name.starts_with("topic-") { - match std::fs::remove_dir_all(entry.path()) { - Ok(()) => removed_dirs += 1, - Err(e) => log::warn!( - "[memory_tree::migrate] purge: failed to remove {} : {e}", - entry.path().display() - ), - } - } - } - } - - conn.pragma_update(None, "user_version", GLOBAL_TOPIC_PURGE_MIGRATION_VERSION) - .context("set PRAGMA user_version after global/topic purge")?; - log::info!( - "[memory_tree::migrate] global/topic purge done: trees={removed_trees} \ - summaries={removed_summaries} sidecars={removed_summary_sidecars} jobs={removed_jobs} \ - dirs={removed_dirs}; user_version={GLOBAL_TOPIC_PURGE_MIGRATION_VERSION}" - ); - Ok(()) -} diff --git a/src/openhuman/memory_store/chunks/produce.rs b/src/openhuman/memory_store/chunks/produce.rs index b8537fe44..6662d2c08 100644 --- a/src/openhuman/memory_store/chunks/produce.rs +++ b/src/openhuman/memory_store/chunks/produce.rs @@ -1,1101 +1,5 @@ -//! Markdown → bounded chunks with stable sequence numbers (Phase 1 / #707). -//! -//! The canonicalisers produce one big canonical Markdown blob per source -//! record; the chunker slices that into chunks of at most [`DEFAULT_CHUNK_MAX_TOKENS`] -//! so later phases (L0 seal at `INPUT_TOKEN_BUDGET = 50k` tokens, or 10 -//! items via the count fallback) can ingest them without blowing past -//! the summariser ceiling. -//! -//! ## Dispatch by source kind (Phase B) -//! -//! - **Chat**: split at `## ` message boundaries. Each message becomes one -//! chunk. If a single message exceeds `max_tokens`, fall back to the -//! paragraph/sentence/whitespace/char splitter for that unit only and emit -//! each piece with `partial_message = true`. -//! - **Email**: split at `---\nFrom:` separators. Each email in the thread -//! becomes one chunk. Same oversize fallback as Chat. -//! - **Document**: split by [`split_by_token_budget`] — a conservative -//! token-estimate splitter (paragraph → sentence → whitespace → hard-char) -//! with ~12% overlap between adjacent chunks. -//! -//! Chunk sizes are bounded by [`conservative_token_estimate`], not the GPT -//! `chars/4` heuristic, so dense markdown/hash/code/multilingual content cannot -//! produce an over-budget chunk that overflows a downstream embedder. +//! Compatibility exports for tinycortex's source-aware chunker. -use crate::openhuman::memory::util::redact::redact; -use crate::openhuman::memory_store::chunks::types::{ - approx_token_count, chunk_id, conservative_token_estimate, truncate_to_conservative_tokens, - Chunk, Metadata, SourceKind, +pub use tinycortex::memory::chunks::{ + chunk_markdown, ChunkerInput, ChunkerOptions, DEFAULT_CHUNK_MAX_TOKENS, }; - -/// Default upper bound on per-chunk tokens. -/// -/// Well below `tree_source::types::INPUT_TOKEN_BUDGET = 50_000` so each -/// L0 seal accumulates many chunks (~15+) before firing — the cloud -/// summariser handles large input contexts well, so we let the seal -/// fold a meaningful slice of the source rather than a single chunk. -pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 3_000; - -/// Tunable settings for the chunker. -#[derive(Clone, Debug)] -pub struct ChunkerOptions { - pub max_tokens: u32, -} - -impl Default for ChunkerOptions { - fn default() -> Self { - Self { - max_tokens: DEFAULT_CHUNK_MAX_TOKENS, - } - } -} - -/// Input to the chunker: the canonicalised source and its provenance. -/// -/// Callers (typically canonicalisers via [`super::ingest`]) own construction; -/// the chunker does not interpret metadata beyond cloning it onto each chunk. -#[derive(Clone, Debug)] -pub struct ChunkerInput { - pub source_kind: SourceKind, - pub source_id: String, - /// Canonical Markdown content — possibly very long. - pub markdown: String, - /// Base metadata; per-chunk `timestamp` defaults to `metadata.timestamp`. - pub metadata: Metadata, -} - -/// Slice `input.markdown` into chunks ≤ `opts.max_tokens` tokens each. -/// -/// Returns chunks in source order with stable sequence numbers starting at 0. -/// Chunk IDs are deterministic (`types::chunk_id`), so re-chunking yields the -/// same ids for identical input. -/// -/// ## Dispatch by source kind -/// -/// - **Chat / Email**: split at message/email boundaries, then greedy-pack -/// consecutive units into a single chunk until adding the next unit would -/// exceed `max_tokens`. Oversize units (a single message > `max_tokens`) -/// fall back to [`split_by_token_budget`] and emit each piece with -/// `partial_message = true`. -/// - **Document**: split by [`split_by_token_budget`] — sized by the -/// conservative token estimate (paragraph → sentence → whitespace → -/// hard-char) with ~12% overlap between adjacent chunks. -pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec { - let now = chrono::Utc::now(); - let max_tokens = opts.max_tokens.max(1); - let max_chars = (max_tokens as usize).saturating_mul(4); - - // Dispatch: pick splitting units based on source kind. - let units: Vec = match input.source_kind { - SourceKind::Chat => split_chat_messages(&input.markdown), - SourceKind::Email => split_email_messages(&input.markdown), - SourceKind::Document => { - // Document: run the existing paragraph splitter directly on the - // whole blob. No message-unit concept. - log::debug!( - "[memory_tree::chunker] document source_id_hash={} len={} — paragraph split", - redact(&input.source_id), - input.markdown.len() - ); - split_by_token_budget(&input.markdown, max_tokens) - } - }; - - if matches!(input.source_kind, SourceKind::Document) { - // Already split by budget; wrap directly. - return units - .into_iter() - .enumerate() - .map(|(idx, content)| { - let seq = idx as u32; - let token_count = approx_token_count(&content); - let id = chunk_id(input.source_kind, &input.source_id, seq, &content); - Chunk { - id, - content, - metadata: input.metadata.clone(), - token_count, - seq_in_source: seq, - created_at: now, - partial_message: false, - } - }) - .collect(); - } - - log::debug!( - "[memory_tree::chunker] source_kind={} source_id_hash={} len={} units={}", - input.source_kind.as_str(), - redact(&input.source_id), - input.markdown.len(), - units.len() - ); - - // For Chat and Email: greedy-pack consecutive units into chunks. - // Units are accumulated until adding the next would exceed max_chars; - // oversize single units fall back to sub-splitting with partial_message=true. - let unit_separator = "\n\n"; - let sep_chars = unit_separator.chars().count(); - - let mut out: Vec = Vec::new(); - let mut acc: Vec = Vec::new(); - let mut acc_chars = 0usize; - - // Flush accumulated units as one packed chunk. - let flush = |acc: &mut Vec, acc_chars: &mut usize, out: &mut Vec| { - if acc.is_empty() { - return; - } - let content = acc.join(unit_separator); - let seq = out.len() as u32; - let tc = approx_token_count(&content); - let id = chunk_id(input.source_kind, &input.source_id, seq, &content); - out.push(Chunk { - id, - content, - metadata: input.metadata.clone(), - token_count: tc, - seq_in_source: seq, - created_at: now, - partial_message: false, - }); - acc.clear(); - *acc_chars = 0; - }; - - for unit in units { - let unit_chars = unit.chars().count(); - - if unit_chars > max_chars { - // Oversize: flush any pending accumulator first, then sub-split. - flush(&mut acc, &mut acc_chars, &mut out); - let sub_pieces = split_by_token_budget(&unit, max_tokens); - for piece in sub_pieces { - let seq = out.len() as u32; - let tc = approx_token_count(&piece); - let id = chunk_id(input.source_kind, &input.source_id, seq, &piece); - out.push(Chunk { - id, - content: piece, - metadata: input.metadata.clone(), - token_count: tc, - seq_in_source: seq, - created_at: now, - partial_message: true, - }); - } - continue; - } - - // Compute projected size if we add this unit to the accumulator. - let projected = if acc.is_empty() { - unit_chars - } else { - acc_chars + sep_chars + unit_chars - }; - - if projected > max_chars { - // Adding this unit would overflow — flush the accumulator first. - flush(&mut acc, &mut acc_chars, &mut out); - } - - if !acc.is_empty() { - acc_chars += sep_chars; - } - acc_chars += unit_chars; - acc.push(unit); - } - - // Flush any remaining accumulated units. - flush(&mut acc, &mut acc_chars, &mut out); - - if out.is_empty() { - // Degenerate: empty input → one empty chunk, matching original behaviour. - let id = chunk_id(input.source_kind, &input.source_id, 0, ""); - out.push(Chunk { - id, - content: String::new(), - metadata: input.metadata.clone(), - token_count: 0, - seq_in_source: 0, - created_at: now, - partial_message: false, - }); - } - - out -} - -/// Split a canonical chat blob into per-message units at `## ` boundaries. -/// -/// Each returned string starts with `## ` and includes everything up to but -/// not including the next `## ` boundary. If the blob starts with a `# ` -/// header (legacy or unexpected), everything before the first `## ` is -/// dropped silently. -fn split_chat_messages(md: &str) -> Vec { - let mut pieces: Vec = Vec::new(); - let mut current: Option = None; - - for line in md.split_inclusive('\n') { - if line.starts_with("## ") { - if let Some(prev) = current.take() { - let trimmed = prev.trim_end().to_string(); - if !trimmed.is_empty() { - pieces.push(trimmed); - } - } - current = Some(line.to_string()); - } else if let Some(ref mut buf) = current { - buf.push_str(line); - } - // Lines before the first `## ` (e.g. a leading `# ` header) are dropped. - } - - if let Some(prev) = current.take() { - let trimmed = prev.trim_end().to_string(); - if !trimmed.is_empty() { - pieces.push(trimmed); - } - } - - if pieces.is_empty() && !md.trim().is_empty() { - // No `## ` found at all — treat whole blob as one unit. - pieces.push(md.trim_end().to_string()); - } - - pieces -} - -/// Split a canonical email thread blob into per-email units. -/// -/// Splits at `---` (alone on a line, optional trailing whitespace) followed -/// by a `From:` line within the next 8 lines. Each piece includes the `---` -/// separator and everything up to but not including the next `---\nFrom:` -/// boundary. Content before the first `---` separator is dropped (handles -/// any leading header that might have slipped through). -fn split_email_messages(md: &str) -> Vec { - let lines: Vec<&str> = md.split('\n').collect(); - let n = lines.len(); - let mut split_positions: Vec = Vec::new(); - - for i in 0..n { - let line = lines[i].trim_end(); - 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:") { - split_positions.push(i); - break; - } - // Skip blank lines between `---` and `From:` - if !lines[j].trim().is_empty() { - break; - } - } - } - } - - if split_positions.is_empty() { - // No email separator found — treat whole blob as one unit. - let trimmed = md.trim_end().to_string(); - if trimmed.is_empty() { - return Vec::new(); - } - return vec![trimmed]; - } - - let mut pieces: Vec = Vec::new(); - for (idx, &start) in split_positions.iter().enumerate() { - let end = if idx + 1 < split_positions.len() { - split_positions[idx + 1] - } else { - n - }; - 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); - } - } - - pieces -} - -/// Overlap carried between adjacent chunks (≈12%, within the 10–15% band) so a -/// fact straddling a split boundary survives in both neighbours. -const OVERLAP_PERCENT: u32 = 12; -/// Hard cap on overlap (% of budget) so a chunk can never be mostly a duplicate -/// of its predecessor — avoids pathological near-identical chunks. -const OVERLAP_MAX_PERCENT: u32 = 40; - -/// Split `text` into pieces each ≤ `max_tokens` **conservative** tokens -/// (see [`conservative_token_estimate`]), with ~10–15% overlap between adjacent -/// pieces. Sized by the conservative estimate (NOT `chars/4`), so dense -/// markdown/hash/code/multilingual content — which tokenises far above the -/// chars/4 heuristic — never produces an over-budget chunk that overflows the -/// embedder. -/// -/// Boundary preference (a still-oversized piece falls to the next finer level): -/// 1. paragraph (`\n\n`) -/// 2. sentence (`. ` / `! ` / `? ` / line break) -/// 3. whitespace (word) -/// 4. hard character cut (last resort; preserves UTF-8 code points) -/// -/// Ordering is preserved. Overlap repeats the previous piece's trailing whole -/// segments verbatim (snapped to natural boundaries), never the entire chunk. -pub(crate) fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec { - let budget = max_tokens.max(1); - if text.is_empty() { - return vec![String::new()]; - } - if conservative_token_estimate(text) <= budget { - return vec![text.to_string()]; - } - - // Reserve headroom so `overlap (≤cap) + any segment (≤seg_budget) ≤ budget`, - // which prevents a flush from ever emitting an overlap-only (duplicate) chunk. - let overlap_budget = (budget * OVERLAP_PERCENT / 100).max(1); - let overlap_cap = (budget * OVERLAP_MAX_PERCENT / 100).max(overlap_budget); - let seg_budget = budget.saturating_sub(overlap_cap).max(1); - - // 1. Reduce to in-order atomic segments, each ≤ seg_budget. - let mut segments: Vec<&str> = Vec::new(); - push_atomic_segments(text, seg_budget, &mut segments); - if segments.len() <= 1 { - return vec![text.to_string()]; - } - - // 2. Greedy-pack segments into chunks ≤ budget, carrying overlap forward. - let mut chunks: Vec = Vec::new(); - let mut cur: Vec<&str> = Vec::new(); - let mut cur_tokens = 0u32; - for seg in &segments { - let seg_tokens = conservative_token_estimate(seg); - if !cur.is_empty() && cur_tokens.saturating_add(seg_tokens) > budget { - chunks.push(join_segments(&cur)); - let overlap = tail_overlap(&cur, overlap_budget, overlap_cap); - cur_tokens = overlap.iter().map(|s| conservative_token_estimate(s)).sum(); - cur = overlap; - } - cur.push(seg); - cur_tokens = cur_tokens.saturating_add(seg_tokens); - } - if !cur.is_empty() { - chunks.push(join_segments(&cur)); - } - if chunks.is_empty() { - chunks.push(String::new()); - } - chunks -} - -/// Append in-order atomic segments of `text`, each with -/// `conservative_token_estimate ≤ budget`, using the boundary hierarchy. -fn push_atomic_segments<'a>(text: &'a str, budget: u32, out: &mut Vec<&'a str>) { - if text.is_empty() { - return; - } - if conservative_token_estimate(text) <= budget { - out.push(text); - return; - } - if split_on_separator(text, "\n\n", budget, out) - || split_on_sentences(text, budget, out) - || split_on_whitespace(text, budget, out) - { - return; - } - hard_split(text, budget, out); -} - -/// Split on a literal separator; recurse on each piece. The separator is kept -/// at the END of its preceding piece so the pieces tile `text` exactly (lossless -/// concatenation — see [`join_segments`]). Returns `false` (no progress) when the -/// separator does not occur. -fn split_on_separator<'a>(text: &'a str, sep: &str, budget: u32, out: &mut Vec<&'a str>) -> bool { - if sep.is_empty() || !text.contains(sep) { - return false; - } - let sep_len = sep.len(); - let mut pieces: Vec<&str> = Vec::new(); - let mut start = 0usize; - let mut search = 0usize; - while let Some(rel) = text[search..].find(sep) { - let end = search + rel + sep_len; // include the separator in this piece - pieces.push(&text[start..end]); - start = end; - search = end; - } - if start < text.len() { - pieces.push(&text[start..]); - } - if pieces.len() <= 1 { - return false; - } - for p in pieces { - push_atomic_segments(p, budget, out); - } - true -} - -/// Split on sentence-ish boundaries: after `.`/`!`/`?` followed by a space, and -/// at line breaks. All boundary bytes are ASCII (1 byte), so slicing is always -/// on a char boundary. Returns `false` when no boundary is found. -fn split_on_sentences<'a>(text: &'a str, budget: u32, out: &mut Vec<&'a str>) -> bool { - let bytes = text.as_bytes(); - let mut pieces: Vec<&str> = Vec::new(); - 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' ' - { - Some(i + 1) - } else { - None - }; - if let Some(end) = boundary_end { - pieces.push(&text[start..end]); - start = end; - } - } - if start < text.len() { - pieces.push(&text[start..]); - } - if pieces.len() <= 1 { - return false; - } - for p in pieces { - push_atomic_segments(p, budget, out); - } - true -} - -/// Split on ASCII spaces (word boundaries); recurse on each word. Returns -/// `false` when there is no space to split on. -fn split_on_whitespace<'a>(text: &'a str, budget: u32, out: &mut Vec<&'a str>) -> bool { - let bytes = text.as_bytes(); - let mut pieces: Vec<&str> = Vec::new(); - let mut start = 0usize; - for i in 0..bytes.len() { - if bytes[i] == b' ' { - // Keep the space at the end of the word so pieces tile `text` exactly. - pieces.push(&text[start..=i]); - start = i + 1; - } - } - if start < text.len() { - pieces.push(&text[start..]); - } - if pieces.len() <= 1 { - return false; - } - for p in pieces { - push_atomic_segments(p, budget, out); - } - true -} - -/// Last resort: cut a boundary-free run into ≤ `budget`-token pieces on UTF-8 -/// char boundaries. Reuses [`truncate_to_conservative_tokens`] for sizing and -/// always makes progress (≥1 char) so it cannot loop. -fn hard_split<'a>(mut text: &'a str, budget: u32, out: &mut Vec<&'a str>) { - while !text.is_empty() { - let mut piece = truncate_to_conservative_tokens(text, budget); - if piece.is_empty() { - // Budget smaller than one char's weight — take a single char. - let end = text - .char_indices() - .nth(1) - .map(|(i, _)| i) - .unwrap_or(text.len()); - piece = &text[..end]; - } - let plen = piece.len(); - out.push(&text[..plen]); - text = &text[plen..]; - } -} - -/// Join packed segments back into one chunk body. Segments are contiguous, -/// separator-retaining slices of the source, so plain concatenation reproduces -/// the original text exactly (modulo intentional overlap duplication) — no -/// spurious separators are introduced, which keeps each chunk's real token -/// count ≤ the packing budget. -fn join_segments(segs: &[&str]) -> String { - segs.concat() -} - -/// Trailing whole segments of `cur` summing to ~`overlap_budget` tokens (capped -/// at `overlap_cap`), returned in original order. Never returns the entire -/// chunk, so adjacent chunks cannot be duplicates. -fn tail_overlap<'a>(cur: &[&'a str], overlap_budget: u32, overlap_cap: u32) -> Vec<&'a str> { - if cur.len() <= 1 { - return Vec::new(); - } - let mut acc = 0u32; - let mut take = 0usize; - for seg in cur.iter().rev() { - let t = conservative_token_estimate(seg); - // Cap EVERY trailing segment (including the first): if the last segment - // alone exceeds `overlap_cap`, carry no overlap rather than letting a - // near-`seg_budget` segment become the overlap. This keeps - // `overlap <= overlap_cap`, so `overlap + next_segment <= budget` and a - // flush can never emit an overlap-only (duplicate) chunk. - if acc.saturating_add(t) > overlap_cap { - break; - } - acc = acc.saturating_add(t); - take += 1; - if acc >= overlap_budget { - break; - } - } - if take >= cur.len() { - take = cur.len() - 1; // never repeat the whole chunk - } - cur[cur.len() - take..].to_vec() -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::Utc; - - fn meta() -> Metadata { - Metadata::point_in_time(SourceKind::Chat, "slack:#eng", "alice", Utc::now()) - } - - fn meta_email() -> Metadata { - Metadata::point_in_time(SourceKind::Email, "gmail:t1", "alice", Utc::now()) - } - - fn meta_doc() -> Metadata { - Metadata::point_in_time(SourceKind::Document, "doc1", "alice", Utc::now()) - } - - #[test] - fn tiny_input_produces_single_chunk() { - // Chat input without a `## ` header produces one chunk via the empty- - // result fallback (whole blob as one unit). - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - markdown: "## 2026-01-01T00:00:00Z — alice\nhello world".into(), - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!(chunks.len(), 1); - assert!(chunks[0].content.contains("hello world")); - assert_eq!(chunks[0].seq_in_source, 0); - assert!(!chunks[0].partial_message); - } - - #[test] - fn empty_chat_input_produces_one_empty_chunk() { - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "x".into(), - markdown: "".into(), - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].content, ""); - assert!(!chunks[0].partial_message); - } - - #[test] - fn chat_messages_pack_into_one_chunk_when_small() { - // Two small chat messages both fit under default max_tokens → greedy - // packing emits ONE chunk containing both, joined by \n\n. - let md = "## 2026-01-01T00:00:00Z — alice\nHello world\n\n## 2026-01-01T00:01:00Z — bob\nParagraph one.\n\nParagraph two.".to_string(); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - markdown: md.clone(), - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - // Both small messages fit under 10k tokens → one packed chunk. - assert_eq!( - chunks.len(), - 1, - "small messages should be packed into one chunk; got {chunks:?}" - ); - assert!( - chunks[0].content.contains("alice"), - "chunk must contain alice's message" - ); - assert!( - chunks[0].content.contains("bob"), - "chunk must contain bob's message" - ); - assert!(chunks[0].content.contains("Paragraph one.")); - assert!(chunks[0].content.contains("Paragraph two.")); - assert!(!chunks[0].partial_message); - } - - #[test] - fn chat_messages_split_at_boundary_when_large() { - // Messages that together exceed max_tokens split at message boundaries - // into multiple chunks. Each chunk contains whole messages only. - // Each message is ~3k tokens at 4 chars/token = 12k chars; - // two messages = ~6k tokens > 5k budget → must split. - let msg_body = "x".repeat(12_000); - let md = format!( - "## 2026-01-01T00:00:00Z — alice\n{msg_body}\n\n## 2026-01-01T00:01:00Z — bob\n{msg_body}" - ); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - markdown: md, - metadata: meta(), - }; - // Use a 5k token budget so two ~3k-token messages don't fit together. - let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 }); - assert_eq!( - chunks.len(), - 2, - "two large messages should land in separate chunks; got {chunks:?}" - ); - assert!(chunks[0].content.contains("alice")); - assert!(chunks[1].content.contains("bob")); - for c in &chunks { - assert!(!c.partial_message, "whole messages must not be partial"); - } - } - - #[test] - fn email_threads_pack_into_one_chunk_when_small() { - // Three short emails all fit under default max_tokens → one packed chunk. - let md = "---\nFrom: alice@example.com\nSubject: Hello\nDate: 2026-01-01T00:00:00Z\n\nFirst body.\n---\nFrom: bob@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:01:00Z\n\nSecond body.\n---\nFrom: carol@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:02:00Z\n\nThird body.".to_string(); - let input = ChunkerInput { - source_kind: SourceKind::Email, - source_id: "gmail:t1".into(), - markdown: md, - metadata: meta_email(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!( - chunks.len(), - 1, - "three small emails should pack into one chunk; got {chunks:?}" - ); - assert!(chunks[0].content.contains("First body.")); - assert!(chunks[0].content.contains("Second body.")); - assert!(chunks[0].content.contains("Third body.")); - assert!(!chunks[0].partial_message); - } - - #[test] - fn email_thread_large_splits_at_email_boundaries() { - // Messages totaling >12k tokens split into 2 chunks at email boundaries. - // Each email is ~4k tokens (16k chars); 3 emails × 4k = 12k tokens. - // With a 5k budget, 2 emails fit per chunk → 2 chunks for 3 emails. - let email_body = "y".repeat(16_000); // ~4k tokens - let md = format!( - "---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{email_body}\n\ - ---\nFrom: b@x.com\nDate: 2026-01-01T00:01:00Z\n\n{email_body}\n\ - ---\nFrom: c@x.com\nDate: 2026-01-01T00:02:00Z\n\n{email_body}" - ); - let input = ChunkerInput { - source_kind: SourceKind::Email, - source_id: "gmail:t1".into(), - markdown: md, - metadata: meta_email(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 }); - assert!( - chunks.len() >= 2, - "large thread must split into multiple chunks; got {}", - chunks.len() - ); - for c in &chunks { - assert!(!c.partial_message, "whole-email chunks must not be partial"); - } - } - - #[test] - fn oversize_single_email_splits_with_partial_flag() { - // A single email body > max_tokens must produce partial_message=true pieces. - let big_body = "z".repeat(50_000); // ~12.5k tokens at 4 chars/token - let md = format!("---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{big_body}"); - let input = ChunkerInput { - source_kind: SourceKind::Email, - source_id: "gmail:t1".into(), - markdown: md, - metadata: meta_email(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 1_000 }); - assert!(chunks.len() > 1, "oversize email must split"); - for c in &chunks { - assert!( - c.partial_message, - "all sub-pieces of an oversize email must have partial_message=true" - ); - } - } - - #[test] - fn packed_units_joined_by_double_newline() { - // Two chat messages packed together must be separated by \n\n. - let md = "## 2026-01-01T00:00:00Z — alice\nfoo\n\n## 2026-01-01T00:01:00Z — bob\nbar" - .to_string(); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "x".into(), - markdown: md, - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!(chunks.len(), 1); - // The two messages must be separated by \n\n in the packed content. - assert!( - chunks[0].content.contains("\n\n"), - "packed units must be joined by \\n\\n; content={:?}", - chunks[0].content - ); - } - - #[test] - fn oversize_message_falls_back_with_partial_flag() { - // Single chat message that is way over max_tokens. - let long_body = "x".repeat(8000); // ~2000 tokens at 4 chars/token - let md = format!("## 2026-01-01T00:00:00Z — alice\n{long_body}"); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "x".into(), - markdown: md, - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 100 }); - assert!(chunks.len() > 1, "oversize message must split"); - for c in &chunks { - assert!( - c.partial_message, - "all sub-pieces of an oversize message must have partial_message=true" - ); - } - // Reuniting all pieces must reconstruct the message content (minus `## ` line). - let rejoined: String = chunks.iter().map(|c| c.content.as_str()).collect(); - assert!(rejoined.contains(&long_body[..100])); - } - - #[test] - fn document_falls_through_to_paragraph_split() { - let para1 = "a".repeat(400); // ~100 tokens - let para2 = "b".repeat(400); - let para3 = "c".repeat(400); - let text = format!("{para1}\n\n{para2}\n\n{para3}"); - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "doc1".into(), - markdown: text, - metadata: meta_doc(), - }; - let chunks = chunk_markdown( - &input, - &ChunkerOptions { - max_tokens: 150, // forces split at paragraph boundary - }, - ); - assert!(chunks.len() >= 2); - for c in &chunks { - let first = c.content.chars().next().unwrap(); - assert!( - matches!(first, 'a' | 'b' | 'c'), - "document chunk starts with unexpected char: {:?}", - c.content.chars().take(10).collect::() - ); - assert!( - !c.partial_message, - "document chunks must not have partial_message=true" - ); - } - } - - #[test] - fn header_line_dropped_in_chat() { - // Simulate a blob that has a leading `# Chat transcript` header. - let md = "# Chat transcript — slack / #eng\n\n## 2026-01-01T00:00:00Z — alice\nhello" - .to_string(); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "x".into(), - markdown: md, - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!(chunks.len(), 1); - // The `# Chat transcript` header must be absent from the chunk content. - assert!( - !chunks[0].content.contains("# Chat transcript"), - "leading `# ` header must be dropped from chunk content" - ); - assert!(chunks[0].content.contains("hello")); - } - - #[test] - fn chunk_ids_are_stable_across_runs() { - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - markdown: "## 2026-01-01T00:00:00Z — alice\nhello".into(), - metadata: meta(), - }; - let a = chunk_markdown(&input, &ChunkerOptions::default()); - let b = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!( - a.iter().map(|c| c.id.clone()).collect::>(), - b.iter().map(|c| c.id.clone()).collect::>() - ); - } - - #[test] - fn sequence_numbers_start_at_zero() { - let msgs: String = (0..5) - .map(|i| format!("## 2026-01-01T00:0{}:00Z — user{i}\nContent {i}\n\n", i)) - .collect(); - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "x".into(), - markdown: msgs, - metadata: meta(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - for (idx, c) in chunks.iter().enumerate() { - assert_eq!(c.seq_in_source, idx as u32); - } - } - - #[test] - fn paragraph_boundaries_preferred_for_documents() { - // Build something that exceeds token budget so it must split. - let para1 = "a".repeat(400); // ~100 tokens - let para2 = "b".repeat(400); - let para3 = "c".repeat(400); - let text = format!("{para1}\n\n{para2}\n\n{para3}"); - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "doc1".into(), - markdown: text, - metadata: meta_doc(), - }; - let chunks = chunk_markdown( - &input, - &ChunkerOptions { - max_tokens: 150, // forces split at paragraph - }, - ); - assert!(chunks.len() >= 2); - for c in &chunks { - let first = c.content.chars().next().unwrap(); - assert!( - matches!(first, 'a' | 'b' | 'c'), - "chunk starts with unexpected char: {:?}", - c.content.chars().take(10).collect::() - ); - } - } - - #[test] - fn falls_back_to_line_split_when_no_paragraphs_document() { - let text = (0..30) - .map(|i| format!("line-{i}-{}", "x".repeat(40))) - .collect::>() - .join("\n"); - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "x".into(), - markdown: text, - metadata: meta_doc(), - }; - let chunks = chunk_markdown( - &input, - &ChunkerOptions { - max_tokens: 80, // forces several splits - }, - ); - assert!(chunks.len() >= 2); - for c in &chunks { - assert!(!c.content.contains("\n\n")); // no paragraph joins in output - } - } - - #[test] - fn utf8_boundaries_preserved_on_hard_split_document() { - // Single long line with no paragraph/line splits → falls to hard cut. - let text = "中".repeat(400); - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "d".into(), - markdown: text.clone(), - metadata: meta_doc(), - }; - let chunks = chunk_markdown( - &input, - &ChunkerOptions { - max_tokens: 50, // ~200 chars - }, - ); - // Rejoining must equal the original. - let rejoined: String = chunks.iter().map(|c| c.content.as_str()).collect(); - assert_eq!(rejoined, text); - } - - #[test] - fn zero_token_budget_is_clamped_without_empty_leading_chunk_document() { - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "d".into(), - markdown: "abcdef".into(), - metadata: meta_doc(), - }; - let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 0 }); - assert!(!chunks.is_empty()); - assert!(chunks.iter().all(|chunk| !chunk.content.is_empty())); - let rejoined: String = chunks.iter().map(|c| c.content.as_str()).collect(); - assert_eq!(rejoined, "abcdef"); - } - - // ── Embed-safety: conservative-token splitting + overlap (#oversized-chunk) ── - - /// Every produced piece must be within the conservative token budget — this - /// is the property that keeps requests under bge-m3's batch/context limit. - fn assert_all_within_budget(pieces: &[String], budget: u32) { - for (i, p) in pieces.iter().enumerate() { - assert!( - conservative_token_estimate(p) <= budget, - "piece {i} is {} est-tokens, over budget {budget}", - conservative_token_estimate(p), - ); - } - } - - #[test] - fn dense_hash_content_splits_under_budget() { - // The exact failure mode: hash/path/punctuation-dense ASCII that chars/4 - // badly under-counts. A single 12 KB block of this overflowed bge-m3. - let dense = - "claude-memory:openhuman:MEMORY.md:67d6fe2727d431b16d41630babfdcf1cdf61bda7b9ba99656\n" - .repeat(200); - let pieces = split_by_token_budget(&dense, 3000); - assert!( - pieces.len() > 1, - "dense content must split, got {}", - pieces.len() - ); - assert_all_within_budget(&pieces, 3000); - } - - #[test] - fn hebrew_content_splits_under_budget() { - // Hebrew tokenises ~1 token/char — chars/4 under-counts ~4×. - let he = "איריס דיברה עם העורך דין ועם עידו על הדירה החדשה ".repeat(300); - let pieces = split_by_token_budget(&he, 1000); - assert!(pieces.len() > 1, "hebrew content must split"); - assert_all_within_budget(&pieces, 1000); - } - - #[test] - fn mixed_markdown_code_splits_under_budget() { - let md = "## Section\nSome prose here. And more prose follows.\n\n\ - ```rust\nfn f(x: u32) -> u32 { x * 4 + 3 }\n```\n\n\ - - bullet a1b2c3d4\n- bullet e5f6a7b8\n" - .repeat(120); - let pieces = split_by_token_budget(&md, 2000); - assert!(pieces.len() > 1, "mixed content must split"); - assert_all_within_budget(&pieces, 2000); - } - - #[test] - fn oversize_content_splits_into_multiple_ordered_pieces() { - let text = (0..50) - .map(|i| format!("Paragraph number {i} with several words of filler content here.")) - .collect::>() - .join("\n\n"); - let pieces = split_by_token_budget(&text, 80); - assert!( - pieces.len() > 2, - "expected several pieces, got {}", - pieces.len() - ); - assert_all_within_budget(&pieces, 80); - // Ordering preserved: "number 0" precedes "number 49" in the stream. - let joined = pieces.join("\n"); - let p0 = joined.find("number 0").expect("first paragraph present"); - let p49 = joined.find("number 49").expect("last paragraph present"); - assert!(p0 < p49, "ordering not preserved"); - } - - #[test] - fn adjacent_chunks_overlap_without_duplicating() { - let text = (0..40) - .map(|i| format!("Sentence {i} carries a unique token marker{i} inside it.")) - .collect::>() - .join(" "); - // Budget chosen so each sentence segment is well under overlap_cap - // (40% of budget); overlap is only carried for segments that fit the - // cap, so a tighter budget would (correctly) yield no overlap. - let pieces = split_by_token_budget(&text, 200); - assert!( - pieces.len() >= 3, - "expected several pieces, got {}", - pieces.len() - ); - assert_all_within_budget(&pieces, 200); - // Overlap: at least one marker appears in two adjacent chunks. - let overlap = pieces.windows(2).any(|w| { - (0..40).any(|i| { - let m = format!("marker{i} "); - w[0].contains(&m) && w[1].contains(&m) - }) - }); - assert!(overlap, "expected ~12% overlap between adjacent chunks"); - // No near-duplicates: adjacent chunks are never identical. - for w in pieces.windows(2) { - assert_ne!(w[0], w[1], "adjacent chunks must not be identical"); - } - } - - #[test] - fn normal_small_content_is_single_chunk_unchanged() { - // Behaviour preserved for already-small content: returned verbatim. - let text = "# Title\nA short paragraph that easily fits.\n\nAnother short one."; - let pieces = split_by_token_budget(text, 3000); - assert_eq!(pieces.len(), 1); - assert_eq!(pieces[0], text); - } - - #[test] - fn large_consecutive_segments_stay_within_budget() { - // A small paragraph packs with a large one, then a second large - // paragraph follows. Regression for the tail_overlap cap: previously the - // large tail segment (> overlap_cap) became the overlap, so - // `overlap + next_segment` exceeded the budget and emitted an - // over-budget chunk. Every chunk must stay within budget, and no two - // adjacent chunks may be identical. - let budget = 100; - let text = format!( - "{}\n\n{}\n\n{}", - "a".repeat(40), - "b".repeat(110), - "c".repeat(110), - ); - let pieces = split_by_token_budget(&text, budget); - assert!( - pieces.len() >= 2, - "expected multiple chunks, got {}", - pieces.len() - ); - assert_all_within_budget(&pieces, budget); - for w in pieces.windows(2) { - assert_ne!(w[0], w[1], "adjacent chunks must not be identical"); - } - } -} diff --git a/src/openhuman/memory_store/chunks/semantic.rs b/src/openhuman/memory_store/chunks/semantic.rs index 6a48bfd1d..a62476251 100644 --- a/src/openhuman/memory_store/chunks/semantic.rs +++ b/src/openhuman/memory_store/chunks/semantic.rs @@ -1,637 +1,7 @@ -//! Semantic markdown chunking for the OpenHuman memory system. -//! -//! This module provides the logic for splitting large markdown documents into -//! smaller, semantically meaningful chunks that fit within the context window -//! of an LLM or an embedding model. It prioritizes splitting on headings and -//! paragraph boundaries while preserving context by carrying over headings -//! to subsequent chunks. +//! Compatibility exports for tinycortex's semantic Markdown chunker. -use std::rc::Rc; +pub use tinycortex::memory::chunks::SemanticChunk as Chunk; -/// A single chunk of text extracted from a larger document. -#[derive(Debug, Clone)] -pub struct Chunk { - /// The zero-based index of this chunk within the original document. - pub index: usize, - /// The actual text content of the chunk. - pub content: String, - /// The most recent markdown heading that applies to this chunk's content. - /// Uses `Rc` for efficient sharing of the same heading across multiple chunks. - pub heading: Option>, -} - -/// Splits markdown text into a sequence of [`Chunk`] objects. -/// -/// Each chunk is designed to be approximately under the `max_tokens` limit. -/// The chunker uses a hierarchical splitting strategy: -/// 1. **Heading Boundaries**: Splits on `#`, `##`, and `###` headings. -/// 2. **Paragraph Boundaries**: If a heading section is too large, it splits on blank lines. -/// 3. **Line Boundaries**: If a paragraph is still too large, it splits on individual lines. -/// -/// # Arguments -/// * `text` - The raw markdown text to chunk. -/// * `max_tokens` - The approximate maximum number of tokens per chunk (estimated at 4 chars/token). -/// -/// # Returns -/// A vector of [`Chunk`] structs representing the document. pub fn chunk_markdown(text: &str, max_tokens: usize) -> Vec { - if text.trim().is_empty() { - return Vec::new(); - } - - // Rough estimation: 4 characters per token for English text. - let max_chars = max_tokens * 4; - - // Step 1: Divide the document into top-level sections based on headings. - let sections = split_on_headings(text); - let mut chunks = Vec::with_capacity(sections.len()); - - for (heading, body) in sections { - let heading: Option> = heading.map(Rc::from); - let heading_prefix = heading.as_deref().map(|h| { - let mut prefix = String::with_capacity(h.len() + 1); - prefix.push_str(h); - prefix.push('\n'); - prefix - }); - - let full_len = body.len() + heading_prefix.as_ref().map_or(0, String::len); - - if full_len <= max_chars { - // Section fits entirely in one chunk. - let content = if let Some(prefix) = heading_prefix.as_deref() { - let mut full = String::with_capacity(full_len); - full.push_str(prefix); - full.push_str(&body); - full.trim().to_string() - } else { - body.trim().to_string() - }; - chunks.push(Chunk { - index: chunks.len(), - content, - heading: heading.clone(), - }); - } else { - // Step 2: Section is too large; split into paragraphs. - let paragraphs = split_on_blank_lines(&body); - let mut current = heading_prefix.clone().unwrap_or_default(); - - for para in paragraphs { - // If adding this paragraph exceeds the limit, emit the current chunk. - if current.len() + para.len() > max_chars && !current.trim().is_empty() { - chunks.push(Chunk { - index: chunks.len(), - content: current.trim().to_string(), - heading: heading.clone(), - }); - // Reset with the heading for context preservation. - reset_chunk_buffer(&mut current, heading_prefix.as_deref()); - } - - if para.len() > max_chars { - // Step 3: Paragraph is still too large; split it line-by-line. - if !current.trim().is_empty() { - chunks.push(Chunk { - index: chunks.len(), - content: current.trim().to_string(), - heading: heading.clone(), - }); - reset_chunk_buffer(&mut current, heading_prefix.as_deref()); - } - for line_chunk in split_on_lines(¶, max_chars) { - chunks.push(Chunk { - index: chunks.len(), - content: line_chunk.trim().to_string(), - heading: heading.clone(), - }); - } - } else { - current.push_str(¶); - current.push('\n'); - } - } - - // Emit any remaining content as a final chunk for this section. - if !current.trim().is_empty() { - chunks.push(Chunk { - index: chunks.len(), - content: current.trim().to_string(), - heading: heading.clone(), - }); - } - } - } - - // Clean up empty chunks and normalize indices. - chunks.retain(|c| !c.content.is_empty()); - - for (i, chunk) in chunks.iter_mut().enumerate() { - chunk.index = i; - } - - chunks -} - -fn reset_chunk_buffer(current: &mut String, heading_prefix: Option<&str>) { - current.clear(); - if let Some(prefix) = heading_prefix { - current.push_str(prefix); - } -} - -/// Returns `true` if `line` starts with a valid ATX markdown heading -/// (1 to 6 `#` characters followed by a space). -fn is_atx_heading(line: &str) -> bool { - const PREFIXES: &[&str] = &["# ", "## ", "### ", "#### ", "##### ", "###### "]; - PREFIXES.iter().any(|p| line.starts_with(p)) -} - -/// Identifies markdown ATX headings and groups their following text into -/// sections. -fn split_on_headings(text: &str) -> Vec<(Option, String)> { - log::debug!( - "[memory::chunker] split_on_headings: entry text_len={}", - text.len() - ); - let mut sections = Vec::new(); - let mut current_heading: Option = None; - let mut current_body = String::new(); - - // Log lengths only, not heading or body text: a heading could contain - // PII the user pasted into their docs (emails, usernames, etc.). - for line in text.lines() { - if is_atx_heading(line) { - if !current_body.trim().is_empty() || current_heading.is_some() { - log::debug!( - "[memory::chunker] split_on_headings: flushing section heading_len={} body_len={}", - current_heading.as_deref().map(str::len).unwrap_or(0), - current_body.len() - ); - sections.push((current_heading.take(), std::mem::take(&mut current_body))); - } - current_heading = Some(line.to_string()); - } else { - current_body.push_str(line); - current_body.push('\n'); - } - } - - if !current_body.trim().is_empty() || current_heading.is_some() { - log::debug!( - "[memory::chunker] split_on_headings: flushing final section heading_len={} body_len={}", - current_heading.as_deref().map(str::len).unwrap_or(0), - current_body.len() - ); - sections.push((current_heading, current_body)); - } - - log::debug!( - "[memory::chunker] split_on_headings: exit sections={}", - sections.len() - ); - sections -} - -/// Splits text into strings based on blank line (paragraph) boundaries. -fn split_on_blank_lines(text: &str) -> Vec { - let mut paragraphs = Vec::new(); - let mut current = String::new(); - - for line in text.lines() { - if line.trim().is_empty() { - if !current.trim().is_empty() { - paragraphs.push(std::mem::take(&mut current)); - } - } else { - current.push_str(line); - current.push('\n'); - } - } - - if !current.trim().is_empty() { - paragraphs.push(current); - } - - paragraphs -} - -/// Splits text into chunks based on line boundaries to ensure size constraints. -/// Lines exceeding `max_chars` are further split on word boundaries. -fn split_on_lines(text: &str, max_chars: usize) -> Vec { - let effective_max = max_chars.max(1); - let mut chunks = Vec::with_capacity(text.len() / effective_max + 1); - let mut current = String::new(); - - log::trace!( - "[memory::chunker] split_on_lines: entry text_len={} max_chars={}", - text.len(), - effective_max - ); - - for line in text.lines() { - if line.len() > effective_max { - log::debug!( - "[memory::chunker] split_on_lines: oversize line detected line_len={} max_chars={}", - line.len(), - effective_max - ); - // Flush anything accumulated before the oversize line. - if !current.is_empty() { - chunks.push(std::mem::take(&mut current)); - } - // Split the oversize line itself on word boundaries. - for part in split_within_line(line, effective_max) { - chunks.push(part); - } - } else if current.len() + line.len() + 1 > effective_max && !current.is_empty() { - chunks.push(std::mem::take(&mut current)); - current.push_str(line); - current.push('\n'); - } else { - current.push_str(line); - current.push('\n'); - } - } - - if !current.is_empty() { - chunks.push(current); - } - - chunks -} - -/// Splits a single oversize line into chunks of at most `max_chars`, preferring -/// word boundaries (spaces) to avoid cutting mid-word. Falls back to hard -/// character splits when no boundary exists within the limit. -fn split_within_line(line: &str, max_chars: usize) -> Vec { - let mut chunks = Vec::new(); - let mut start = 0; - let bytes = line.as_bytes(); - - log::trace!( - "[memory::chunker] split_within_line: entry line_len={} max_chars={}", - line.len(), - max_chars - ); - - while start < line.len() { - let remaining = line.len() - start; - if remaining <= max_chars { - chunks.push(format!("{}\n", &line[start..])); - break; - } - - // Find the end boundary, staying on a valid char boundary. - let mut end = start + max_chars; - // Walk back to a valid UTF-8 char boundary. - while end > start && !line.is_char_boundary(end) { - end -= 1; - } - - // If max_chars is smaller than the next character (e.g., a 4-byte emoji - // with max_chars=1), `end` can equal `start`. Advance to the next char - // boundary to guarantee progress and avoid an infinite loop. - if end == start { - end = start + 1; - while end < line.len() && !line.is_char_boundary(end) { - end += 1; - } - log::debug!( - "[memory::chunker] split_within_line: forced advance past multi-byte char start={} end={}", - start, end - ); - } - - // Try to find a space to break on (scan backwards from `end`). - let mut split_at = end; - while split_at > start && bytes[split_at - 1] != b' ' { - split_at -= 1; - } - - // If we couldn't find a space within the range, hard-split at `end`. - if split_at == start { - split_at = end; - log::debug!( - "[memory::chunker] split_within_line: hard split at {} (no word boundary)", - split_at - ); - } - - chunks.push(format!("{}\n", &line[start..split_at])); - // Skip the space we split on (if it was a space). - if split_at < line.len() && bytes[split_at] == b' ' { - start = split_at + 1; - } else { - start = split_at; - } - } - - log::trace!( - "[memory::chunker] split_within_line: exit parts={}", - chunks.len() - ); - chunks -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn empty_text() { - assert!(chunk_markdown("", 512).is_empty()); - assert!(chunk_markdown(" ", 512).is_empty()); - } - - #[test] - fn single_short_paragraph() { - let chunks = chunk_markdown("Hello world", 512); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].content, "Hello world"); - assert!(chunks[0].heading.is_none()); - } - - #[test] - fn heading_sections() { - let text = "# Title\nSome intro.\n\n## Section A\nContent A.\n\n## Section B\nContent B."; - let chunks = chunk_markdown(text, 512); - assert!(chunks.len() >= 3); - assert!(chunks[0].heading.is_none() || chunks[0].heading.as_deref() == Some("# Title")); - } - - #[test] - fn respects_max_tokens() { - // Build multi-line text (one sentence per line) to exercise line-level splitting - let long_text: String = (0..200).fold(String::new(), |mut s, i| { - use std::fmt::Write; - let _ = writeln!( - s, - "This is sentence number {i} with some extra words to fill it up." - ); - s - }); - let chunks = chunk_markdown(&long_text, 50); // 50 tokens ≈ 200 chars - assert!( - chunks.len() > 1, - "Expected multiple chunks, got {}", - chunks.len() - ); - for chunk in &chunks { - // Allow some slack (heading re-insertion etc.) - assert!( - chunk.content.len() <= 300, - "Chunk too long: {} chars", - chunk.content.len() - ); - } - } - - #[test] - fn preserves_heading_in_split_sections() { - let mut text = String::from("## Big Section\n"); - for i in 0..100 { - use std::fmt::Write; - let _ = write!(text, "Line {i} with some content here.\n\n"); - } - let chunks = chunk_markdown(&text, 50); - assert!(chunks.len() > 1); - // All chunks from this section should reference the heading - for chunk in &chunks { - if chunk.heading.is_some() { - assert_eq!(chunk.heading.as_deref(), Some("## Big Section")); - } - } - } - - #[test] - fn indexes_are_sequential() { - let text = "# A\nContent A\n\n# B\nContent B\n\n# C\nContent C"; - let chunks = chunk_markdown(text, 512); - for (i, chunk) in chunks.iter().enumerate() { - assert_eq!(chunk.index, i); - } - } - - #[test] - fn chunk_count_reasonable() { - let text = "Hello world. This is a test document."; - let chunks = chunk_markdown(text, 512); - assert_eq!(chunks.len(), 1); - } - - // ── Edge cases ─────────────────────────────────────────────── - - #[test] - fn headings_only_no_body() { - let text = "# Title\n## Section A\n## Section B\n### Subsection"; - let chunks = chunk_markdown(text, 512); - // Should produce chunks for each heading (even with empty bodies) - assert!(!chunks.is_empty()); - } - - #[test] - fn deep_atx_headings_split_through_h6() { - let text = "# Top\nIntro\n#### Deep heading\nDeep content"; - let chunks = chunk_markdown(text, 512); - assert!( - chunks.len() >= 2, - "expected the #### heading to start a new section, got {} chunk(s)", - chunks.len(), - ); - let deep = chunks - .iter() - .find(|c| c.heading.as_deref() == Some("#### Deep heading")); - assert!( - deep.is_some(), - "expected a chunk with heading '#### Deep heading'; chunks: {chunks:?}", - ); - } - - #[test] - fn all_atx_heading_levels_h1_through_h6_split() { - let text = "# H1\na\n\n## H2\nb\n\n### H3\nc\n\n#### H4\nd\n\n##### H5\ne\n\n###### H6\nf"; - let chunks = chunk_markdown(text, 512); - let headings: Vec<_> = chunks.iter().filter_map(|c| c.heading.as_deref()).collect(); - assert_eq!( - headings, - vec![ - "# H1", - "## H2", - "### H3", - "#### H4", - "##### H5", - "###### H6" - ], - "each ATX heading depth h1-h6 must split into its own section", - ); - } - - #[test] - fn seven_or_more_hashes_are_not_a_heading() { - let text = "# Top\nIntro\n####### Not a heading\nMore content"; - let chunks = chunk_markdown(text, 512); - assert_eq!( - chunks.len(), - 1, - "7-hash line should not split; expected 1 chunk, got {}", - chunks.len(), - ); - assert_eq!(chunks[0].heading.as_deref(), Some("# Top")); - assert!(chunks[0].content.contains("####### Not a heading")); - } - - #[test] - fn atx_heading_requires_trailing_space() { - let text = "# Real heading\nIntro\n###NoSpace\nbody"; - let chunks = chunk_markdown(text, 512); - assert_eq!( - chunks.len(), - 1, - "missing trailing space disqualifies the heading" - ); - assert_eq!(chunks[0].heading.as_deref(), Some("# Real heading")); - } - - #[test] - fn very_long_single_line_no_newlines() { - // One giant line with no newlines — must split within the line - let text = "word ".repeat(5000); // 25000 chars - let max_tokens = 50; // 200 chars - let max_chars = max_tokens * 4; - let chunks = chunk_markdown(&text, max_tokens); - assert!( - chunks.len() > 1, - "Expected multiple chunks for a 25KB single-line input, got {}", - chunks.len() - ); - for chunk in &chunks { - // Each chunk must respect the size limit (with reasonable slack for - // trailing newline and word overshoot). - assert!( - chunk.content.len() <= max_chars + 50, - "Chunk exceeds max_chars: {} chars (limit {})", - chunk.content.len(), - max_chars - ); - } - } - - #[test] - fn oversize_line_splits_on_word_boundary() { - // A single line of 100 words, each 5 chars + space = 600 chars - let line = "abcde ".repeat(100); - let text = format!("# Heading\n{line}"); - let chunks = chunk_markdown(&text, 25); // 25 tokens = 100 chars max - assert!(chunks.len() > 1); - // Verify no chunk contains a split mid-word - for chunk in &chunks { - // Words should be intact (no "abc\nde" splits) - for word in chunk.content.split_whitespace() { - if word.starts_with('#') { - continue; // heading - } - assert!( - word == "abcde" || word == "Heading", - "Unexpected split word: '{word}'" - ); - } - } - } - - #[test] - fn oversize_line_no_spaces_hard_splits() { - // A single line with no word boundaries at all - let text = "x".repeat(1000); - let chunks = chunk_markdown(&text, 25); // 100 chars max - assert!( - chunks.len() > 1, - "Should hard-split when no spaces exist, got {} chunk(s)", - chunks.len() - ); - // Reconstruct and verify no data loss - let reassembled: String = chunks.iter().map(|c| c.content.trim()).collect(); - assert_eq!(reassembled.len(), 1000); - } - - #[test] - fn only_newlines_and_whitespace() { - assert!(chunk_markdown("\n\n\n \n\n", 512).is_empty()); - } - - #[test] - fn max_tokens_zero() { - // max_tokens=0 → max_chars=0, should not panic or infinite loop - let chunks = chunk_markdown("Hello world", 0); - // Every chunk will exceed 0 chars, so it splits maximally - assert!(!chunks.is_empty()); - } - - #[test] - fn max_tokens_one() { - // max_tokens=1 → max_chars=4, very aggressive splitting - let text = "Line one\nLine two\nLine three"; - let chunks = chunk_markdown(text, 1); - assert!(!chunks.is_empty()); - } - - #[test] - fn unicode_content() { - let text = "# 日本語\nこんにちは世界\n\n## Émojis\n🦀 Rust is great 🚀"; - let chunks = chunk_markdown(text, 512); - assert!(!chunks.is_empty()); - let all: String = chunks.iter().map(|c| c.content.clone()).collect(); - assert!(all.contains("こんにちは")); - assert!(all.contains("🦀")); - } - - #[test] - fn fts5_special_chars_in_content() { - let text = "Content with \"quotes\" and (parentheses) and * asterisks *"; - let chunks = chunk_markdown(text, 512); - assert_eq!(chunks.len(), 1); - assert!(chunks[0].content.contains("\"quotes\"")); - } - - #[test] - fn multiple_blank_lines_between_paragraphs() { - let text = "Paragraph one.\n\n\n\n\nParagraph two.\n\n\n\nParagraph three."; - let chunks = chunk_markdown(text, 512); - assert_eq!(chunks.len(), 1); // All fits in one chunk - assert!(chunks[0].content.contains("Paragraph one")); - assert!(chunks[0].content.contains("Paragraph three")); - } - - #[test] - fn heading_at_end_of_text() { - let text = "Some content\n# Trailing Heading"; - let chunks = chunk_markdown(text, 512); - assert!(!chunks.is_empty()); - } - - #[test] - fn single_heading_no_content() { - let text = "# Just a heading"; - let chunks = chunk_markdown(text, 512); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].heading.as_deref(), Some("# Just a heading")); - } - - #[test] - fn no_content_loss() { - let text = "# A\nContent A line 1\nContent A line 2\n\n## B\nContent B\n\n## C\nContent C"; - let chunks = chunk_markdown(text, 512); - let reassembled: String = chunks.iter().fold(String::new(), |mut s, c| { - use std::fmt::Write; - let _ = writeln!(s, "{}", c.content); - s - }); - // All original content words should appear - for word in ["Content", "line", "1", "2"] { - assert!( - reassembled.contains(word), - "Missing word '{word}' in reassembled chunks" - ); - } - } + tinycortex::memory::chunks::chunk_semantic(text, max_tokens) } diff --git a/src/openhuman/memory_store/chunks/store.rs b/src/openhuman/memory_store/chunks/store.rs index d174769c9..5bb03c09a 100644 --- a/src/openhuman/memory_store/chunks/store.rs +++ b/src/openhuman/memory_store/chunks/store.rs @@ -1,1262 +1,159 @@ -//! SQLite-backed persistence for ingested chunks (Phase 1 / issue #707). -//! -//! The store lives at `/memory_tree/chunks.db`. Schema is applied -//! lazily on first access via `with_connection`, so the DB is created on -//! demand without an explicit migration step. -//! -//! Upsert semantics: writes are idempotent on `chunk.id` so re-ingesting the -//! same raw source yields no duplicates. -//! -//! ## Connection cache (#2206) -//! -//! `with_connection()` previously opened a new SQLite connection and re-ran -//! the full schema init (8 tables, 15+ indexes, 8+ migrations) on **every** -//! call. With 4 workers polling every 5 s this amounted to ~69K connection -//! opens/day, and a family of WAL/SHM cold-start I/O codes (1546 -//! IOERR_TRUNCATE, 4618 IOERR_SHMOPEN, 4874 IOERR_SHMSIZE, 14 CANTOPEN) -//! flooded Sentry with ~19K events in 4 days. -//! -//! Fix: a process-level `ConnectionCache` keyed by DB path. Each entry holds -//! one `parking_lot::Mutex` that is initialised once (schema + -//! migrations + legacy-embedding migration) and then reused for all subsequent -//! calls. A per-entry `CircuitBreaker` stops retrying after 3 consecutive -//! init failures for 30 s so a broken install does not busy-loop. +//! `Config` and transaction adapters for tinycortex chunk persistence. -use anyhow::{Context, Result}; -use rusqlite::{params, Connection, OptionalExtension, Transaction}; -use std::collections::{HashMap, HashSet}; -#[cfg(test)] -use std::sync::Arc; -use std::time::Duration; +use std::collections::HashMap; + +use anyhow::Result; +use rusqlite::Transaction; use crate::openhuman::config::Config; -use crate::openhuman::memory::util::redact::{self, redact as redact_value}; use crate::openhuman::memory_store::chunks::types::{Chunk, SourceKind}; use crate::openhuman::memory_store::content::StagedChunk; -use crate::openhuman::tinycortex::memory_config_from; -const DB_DIR: &str = "memory_tree"; -const DB_FILE: &str = "chunks.db"; -// 15s gives the busy-handler enough headroom that transient write-lock -// contention (4 job workers + scheduler + ingest producers all writing the -// same `memory_tree/chunks.db`) is absorbed inside rusqlite instead of -// surfacing as `SQLITE_BUSY` to callers. Workers still treat busy as a -// soft signal (see `memory_tree::jobs::worker`) so even if this is -// exceeded, the only effect is a one-poll backoff — but 15s is -// comfortably above realistic peer-write durations and shrinks the rate -// at which we have to fall back to that path. The previous 5s was tight -// enough on contended Windows hosts that we were observing avoidable -// busy returns (see OPENHUMAN-TAURI-BP). -const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(15); +pub use tinycortex::memory::chunks::{ + ListChunksQuery, RawRef, CHUNK_STATUS_ADMITTED, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, + CHUNK_STATUS_PENDING_EXTRACTION, CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, +}; -/// Chunk lifecycle: freshly persisted, awaiting the async extract job. -pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction"; -/// Chunk lifecycle: extract ran and the chunk passed admission. -pub const CHUNK_STATUS_ADMITTED: &str = "admitted"; -/// Chunk lifecycle: appended to the L0 buffer of its source tree. -pub const CHUNK_STATUS_BUFFERED: &str = "buffered"; -/// Chunk lifecycle: rolled into a sealed L1 summary. -pub const CHUNK_STATUS_SEALED: &str = "sealed"; -/// Chunk lifecycle: rejected by the admission gate (too low signal). -pub const CHUNK_STATUS_DROPPED: &str = "dropped"; +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} -// `PRAGMA foreign_keys = ON` is intentionally NOT in SCHEMA — it is -// a connection-local pragma that resets to off on every new -// `Connection::open`. SCHEMA only runs once per DB path (first-init); -// applying foreign_keys here would leak FK-off into every later -// `with_connection()` call that hits the fast path. The pragma is -// set per-connection in `open_connection()` instead. - -/// `PRAGMA user_version` value once the one-shot legacy→sidecar embedding -/// migration (#1574 §7) has run. `0` (fresh/legacy DB) triggers the copy on -/// next open; `>= 1` skips it. Bump only for a new one-shot data migration. -const TREE_EMBEDDING_MIGRATION_VERSION: i64 = 1; - -/// `PRAGMA user_version` value once the global/topic-tree purge has run. -/// The global (time-axis) and topic (subject-axis) trees were removed; this -/// one-shot migration deletes their rows + on-disk summary folders. `< 2` -/// triggers the purge on next open; `>= 2` skips it. -const GLOBAL_TOPIC_PURGE_MIGRATION_VERSION: i64 = 2; - -const SCHEMA: &str = " -CREATE TABLE IF NOT EXISTS mem_tree_chunks ( - id TEXT PRIMARY KEY, - source_kind TEXT NOT NULL, - source_id TEXT NOT NULL, - path_scope TEXT, - source_ref TEXT, - owner TEXT NOT NULL, - timestamp_ms INTEGER NOT NULL, - time_range_start_ms INTEGER NOT NULL, - time_range_end_ms INTEGER NOT NULL, - tags_json TEXT NOT NULL DEFAULT '[]', - content TEXT NOT NULL, - token_count INTEGER NOT NULL, - seq_in_source INTEGER NOT NULL, - created_at_ms INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_source - ON mem_tree_chunks(source_kind, source_id); -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_timestamp - ON mem_tree_chunks(timestamp_ms); -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_owner - ON mem_tree_chunks(owner); -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_source_seq - ON mem_tree_chunks(source_kind, source_id, seq_in_source); - --- Per-(chunk, embedding model) vectors (#1574). The legacy --- mem_tree_chunks.embedding column remains in place during the dual-write --- migration; this table lets multiple vector spaces coexist safely. -CREATE TABLE IF NOT EXISTS mem_tree_chunk_embeddings ( - chunk_id TEXT NOT NULL REFERENCES mem_tree_chunks(id) ON DELETE CASCADE, - model_signature TEXT NOT NULL, - vector BLOB NOT NULL, - dim INTEGER NOT NULL, - created_at REAL NOT NULL, - PRIMARY KEY (chunk_id, model_signature) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunk_embeddings_model - ON mem_tree_chunk_embeddings(model_signature); - --- #1574 §6 reembed-backfill terminal-skip tombstone. --- --- A row here means: 'this (chunk, signature) pair was attempted and failed --- terminally (body file missing on disk, embed returned wrong dim, embedder --- erred unrecoverably) — DO NOT re-enqueue it on the next backfill batch.' --- --- Without this table, the reembed worklist's `NOT EXISTS embeddings` predicate --- keeps re-selecting any chunk that failed read/embed (since no sidecar row --- was ever written), and `handle_reembed_backfill` loops on the same rows --- forever — observed in the wild as 16 orphan chunk_ids generating ~128k --- 'body read failed; skipping' warns across ~8k batch defers. The handler --- now writes a row here on terminal failure, and the worklist excludes them. --- Idempotent: the table is created here, and `chrono::Utc` is already imported. -CREATE TABLE IF NOT EXISTS mem_tree_chunk_reembed_skipped ( - chunk_id TEXT NOT NULL REFERENCES mem_tree_chunks(id) ON DELETE CASCADE, - model_signature TEXT NOT NULL, - reason TEXT NOT NULL, - skipped_at_ms INTEGER NOT NULL, - PRIMARY KEY (chunk_id, model_signature) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_chunk_reembed_skipped_model - ON mem_tree_chunk_reembed_skipped(model_signature); - --- Phase 2 (#708): per-chunk score rationale for admission debugging. -CREATE TABLE IF NOT EXISTS mem_tree_score ( - chunk_id TEXT PRIMARY KEY, - total REAL NOT NULL, - token_count_signal REAL NOT NULL, - unique_words_signal REAL NOT NULL, - metadata_weight REAL NOT NULL, - source_weight REAL NOT NULL, - interaction_weight REAL NOT NULL, - entity_density REAL NOT NULL, - dropped INTEGER NOT NULL DEFAULT 0, - reason TEXT, - computed_at_ms INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_score_total - ON mem_tree_score(total); -CREATE INDEX IF NOT EXISTS idx_mem_tree_score_dropped - ON mem_tree_score(dropped); - --- Phase 2 (#708): inverted index entity_id -> node_id for retrieval. --- is_user (#1365) is set at index time via the Composio identity registry --- (is_self_identity_any_toolkit). Default 0 so legacy rows read back as --- non-user until the backfill job re-tags them. -CREATE TABLE IF NOT EXISTS mem_tree_entity_index ( - entity_id TEXT NOT NULL, - node_id TEXT NOT NULL, - node_kind TEXT NOT NULL, - entity_kind TEXT NOT NULL, - surface TEXT NOT NULL, - score REAL NOT NULL, - timestamp_ms INTEGER NOT NULL, - tree_id TEXT, - is_user INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (entity_id, node_id) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_entity - ON mem_tree_entity_index(entity_id); -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_node - ON mem_tree_entity_index(node_id); -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_timestamp - ON mem_tree_entity_index(timestamp_ms); - --- E2GraphRAG: undirected weighted entity co-occurrence graph. One row per --- unordered entity pair that has been extracted together within the same --- chunk; `weight` accumulates co-occurrence frequency. Canonical ordering --- (`entity_a < entity_b`) keeps the pair unique without a second row, and --- the `entity_b` index makes neighbour lookups symmetric (a row matches a --- query entity whether it appears as `entity_a` or `entity_b`). Read at --- query time by `memory_tree::graph` for bounded-hop shortest-path filtering --- during deterministic (LLM-free) retrieval. -CREATE TABLE IF NOT EXISTS mem_tree_entity_edges ( - entity_a TEXT NOT NULL, - entity_b TEXT NOT NULL, - weight INTEGER NOT NULL DEFAULT 1, - updated_ms INTEGER NOT NULL, - PRIMARY KEY (entity_a, entity_b) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_edges_a - ON mem_tree_entity_edges(entity_a); -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_edges_b - ON mem_tree_entity_edges(entity_b); - --- Phase 3a (#709): summary trees / bucket-seal. --- `mem_tree_trees` tracks one tree per scope (source/topic/global). -CREATE TABLE IF NOT EXISTS mem_tree_trees ( - id TEXT PRIMARY KEY, - kind TEXT NOT NULL, - scope TEXT NOT NULL, - root_id TEXT, - max_level INTEGER NOT NULL DEFAULT 0, - status TEXT NOT NULL DEFAULT 'active', - created_at_ms INTEGER NOT NULL, - last_sealed_at_ms INTEGER -); - -CREATE UNIQUE INDEX IF NOT EXISTS idx_mem_tree_trees_kind_scope - ON mem_tree_trees(kind, scope); -CREATE INDEX IF NOT EXISTS idx_mem_tree_trees_status - ON mem_tree_trees(status); - --- `mem_tree_summaries` holds sealed summary nodes. Immutable once written --- (Phase 3a). `deleted` is reserved for future archive cascades. -CREATE TABLE IF NOT EXISTS mem_tree_summaries ( - id TEXT PRIMARY KEY, - tree_id TEXT NOT NULL, - tree_kind TEXT NOT NULL, - level INTEGER NOT NULL, - parent_id TEXT, - child_ids_json TEXT NOT NULL DEFAULT '[]', - content TEXT NOT NULL, - token_count INTEGER NOT NULL, - entities_json TEXT NOT NULL DEFAULT '[]', - topics_json TEXT NOT NULL DEFAULT '[]', - time_range_start_ms INTEGER NOT NULL, - time_range_end_ms INTEGER NOT NULL, - score REAL NOT NULL DEFAULT 0.0, - sealed_at_ms INTEGER NOT NULL, - deleted INTEGER NOT NULL DEFAULT 0, - FOREIGN KEY (tree_id) REFERENCES mem_tree_trees(id) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_tree_level - ON mem_tree_summaries(tree_id, level); -CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_parent - ON mem_tree_summaries(parent_id); -CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_sealed_at - ON mem_tree_summaries(sealed_at_ms); -CREATE INDEX IF NOT EXISTS idx_mem_tree_summaries_deleted - ON mem_tree_summaries(deleted); - --- Per-(summary, embedding model) vectors (#1574). Kept separate from the --- legacy mem_tree_summaries.embedding column so provider/model switches can --- be query-time filters instead of destructive rewrites. -CREATE TABLE IF NOT EXISTS mem_tree_summary_embeddings ( - summary_id TEXT NOT NULL REFERENCES mem_tree_summaries(id) ON DELETE CASCADE, - model_signature TEXT NOT NULL, - vector BLOB NOT NULL, - dim INTEGER NOT NULL, - created_at REAL NOT NULL, - PRIMARY KEY (summary_id, model_signature) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_summary_embeddings_model - ON mem_tree_summary_embeddings(model_signature); - --- #1574 §6 reembed-backfill terminal-skip tombstone (summary side). Mirrors --- `mem_tree_chunk_reembed_skipped` for the summary worklist. See that table's --- comment for the full rationale. -CREATE TABLE IF NOT EXISTS mem_tree_summary_reembed_skipped ( - summary_id TEXT NOT NULL REFERENCES mem_tree_summaries(id) ON DELETE CASCADE, - model_signature TEXT NOT NULL, - reason TEXT NOT NULL, - skipped_at_ms INTEGER NOT NULL, - PRIMARY KEY (summary_id, model_signature) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_summary_reembed_skipped_model - ON mem_tree_summary_reembed_skipped(model_signature); - --- `mem_tree_buffers` holds the unsealed frontier per (tree, level). One row --- per active level per tree; deleted when the buffer seals (clears) in the --- same transaction as the new summary node row. -CREATE TABLE IF NOT EXISTS mem_tree_buffers ( - tree_id TEXT NOT NULL, - level INTEGER NOT NULL, - item_ids_json TEXT NOT NULL DEFAULT '[]', - token_sum INTEGER NOT NULL DEFAULT 0, - oldest_at_ms INTEGER, - updated_at_ms INTEGER NOT NULL, - PRIMARY KEY (tree_id, level), - FOREIGN KEY (tree_id) REFERENCES mem_tree_trees(id) -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_buffers_oldest - ON mem_tree_buffers(oldest_at_ms); - --- Phase 3c (#709): per-entity hotness counters driving lazy topic-tree --- materialisation. One row per canonical entity_id. Counters are bumped --- on every ingest; `last_hotness` is recomputed every --- `TOPIC_RECHECK_EVERY` ingests to decide whether to spawn / archive a --- topic tree for the entity. TODO: 30-day windowing — for Phase 3c we --- increment counts forever and rely on project-scale truthfulness. -CREATE TABLE IF NOT EXISTS mem_tree_entity_hotness ( - entity_id TEXT PRIMARY KEY, - mention_count_30d INTEGER NOT NULL DEFAULT 0, - distinct_sources INTEGER NOT NULL DEFAULT 0, - last_seen_ms INTEGER, - query_hits_30d INTEGER NOT NULL DEFAULT 0, - graph_centrality REAL, - ingests_since_check INTEGER NOT NULL DEFAULT 0, - last_hotness REAL, - last_updated_ms INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_hotness_score - ON mem_tree_entity_hotness(last_hotness); - --- Async job queue for memory-tree work (extract → admit → buffer → seal → --- topic-route → daily digest). Producers (ingest, schedulers, handlers) --- enqueue rows transactionally; the worker pool claims them via the --- `(status, available_at_ms)` index. `dedupe_key` is enforced as unique --- only for ready/running rows so a completed job's key can be re-used. -CREATE TABLE IF NOT EXISTS mem_tree_jobs ( - id TEXT PRIMARY KEY, - kind TEXT NOT NULL, - payload_json TEXT NOT NULL, - dedupe_key TEXT, - status TEXT NOT NULL DEFAULT 'ready', - attempts INTEGER NOT NULL DEFAULT 0, - max_attempts INTEGER NOT NULL DEFAULT 5, - available_at_ms INTEGER NOT NULL, - locked_until_ms INTEGER, - last_error TEXT, - created_at_ms INTEGER NOT NULL, - started_at_ms INTEGER, - completed_at_ms INTEGER, - failure_reason TEXT, - failure_class TEXT -); - -CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_ready - ON mem_tree_jobs(status, available_at_ms); -CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_kind - ON mem_tree_jobs(kind); -CREATE UNIQUE INDEX IF NOT EXISTS idx_mem_tree_jobs_dedupe_active - ON mem_tree_jobs(dedupe_key) - WHERE dedupe_key IS NOT NULL AND status IN ('ready', 'running'); - --- Source-level ingest gate. Memory items (documents, chat batches, email --- threads) are append-only — once a `(source_kind, source_id)` is ingested --- it must not be re-ingested, otherwise its chunks flow back through --- extract → admit → buffer → seal and end up duplicated in the summariser --- tree. The first ingest claims the row; subsequent ingest_* calls for the --- same key short-circuit before canonicalisation. -CREATE TABLE IF NOT EXISTS mem_tree_ingested_sources ( - source_kind TEXT NOT NULL, - source_id TEXT NOT NULL, - ingested_at_ms INTEGER NOT NULL, - PRIMARY KEY (source_kind, source_id) -); - --- MCP write-tool audit trail (#2536). This intentionally stores compact --- identifying metadata instead of duplicating the memory document body. -CREATE TABLE IF NOT EXISTS mcp_writes ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - timestamp_ms INTEGER NOT NULL, - client_info TEXT NOT NULL, - tool_name TEXT NOT NULL, - args_summary TEXT, - resulting_chunk_id TEXT, - success INTEGER NOT NULL, - error_message TEXT -); - -CREATE INDEX IF NOT EXISTS idx_mcp_writes_timestamp - ON mcp_writes(timestamp_ms DESC); -CREATE INDEX IF NOT EXISTS idx_mcp_writes_client - ON mcp_writes(client_info); -CREATE INDEX IF NOT EXISTS idx_mcp_writes_tool - ON mcp_writes(tool_name); -"; - -/// Upsert a batch of chunks atomically. -/// -/// Returns the number of rows inserted or replaced. Duplicates on `chunk.id` -/// are replaced, making the operation idempotent for re-ingest of the same -/// raw source. pub fn upsert_chunks(config: &Config, chunks: &[Chunk]) -> Result { tinycortex::memory::chunks::upsert_chunks(&engine_config(config), chunks) } -/// Upsert chunks using an existing transaction, preserving previously stored embeddings. pub(crate) fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result { - if chunks.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare( - "INSERT INTO mem_tree_chunks ( - 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 - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) - ON CONFLICT(id) DO UPDATE SET - source_kind = excluded.source_kind, - source_id = excluded.source_id, - path_scope = excluded.path_scope, - source_ref = excluded.source_ref, - owner = excluded.owner, - timestamp_ms = excluded.timestamp_ms, - time_range_start_ms = excluded.time_range_start_ms, - time_range_end_ms = excluded.time_range_end_ms, - tags_json = excluded.tags_json, - content = excluded.content, - token_count = excluded.token_count, - seq_in_source = excluded.seq_in_source, - created_at_ms = excluded.created_at_ms", - )?; - upsert_chunks_with_statement(&mut stmt, chunks)?; - Ok(chunks.len()) + tinycortex::memory::chunks::upsert_chunks_tx(tx, chunks) } -/// Upsert staged chunks (with content_path + content_sha256) using an existing transaction. -/// -/// Identical to `upsert_chunks_tx` but also writes the Phase MD-content pointer columns. -/// `content` column receives a ≤500-char plain-text preview of the body (the full body -/// lives on disk at `content_path`). pub(crate) fn upsert_staged_chunks_tx( tx: &Transaction<'_>, - staged: &[StagedChunk], + chunks: &[StagedChunk], ) -> Result { - if staged.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare( - "INSERT INTO mem_tree_chunks ( - 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, - content_path, content_sha256 - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16) - ON CONFLICT(id) DO UPDATE SET - source_kind = excluded.source_kind, - source_id = excluded.source_id, - path_scope = excluded.path_scope, - source_ref = excluded.source_ref, - owner = excluded.owner, - timestamp_ms = excluded.timestamp_ms, - time_range_start_ms = excluded.time_range_start_ms, - time_range_end_ms = excluded.time_range_end_ms, - tags_json = excluded.tags_json, - content = excluded.content, - token_count = excluded.token_count, - seq_in_source = excluded.seq_in_source, - created_at_ms = excluded.created_at_ms, - content_path = excluded.content_path, - content_sha256 = excluded.content_sha256", - )?; - for s in staged { - let chunk = &s.chunk; - // SQL `content` column always carries a ≤500-char preview now - // — the full body either lives at `content_path` (chat / - // document) or is reconstructed from `raw_refs_json` byte - // ranges in the raw archive (email). See `read_chunk_body`. - let preview: String = chunk.content.chars().take(500).collect(); - stmt.execute(params![ - chunk.id, - chunk.metadata.source_kind.as_str(), - chunk.metadata.source_id, - chunk.metadata.path_scope, - chunk.metadata.source_ref.as_ref().map(|r| r.value.as_str()), - chunk.metadata.owner, - chunk.metadata.timestamp.timestamp_millis(), - chunk.metadata.time_range.0.timestamp_millis(), - chunk.metadata.time_range.1.timestamp_millis(), - serde_json::to_string(&chunk.metadata.tags)?, - preview, - chunk.token_count, - chunk.seq_in_source, - chunk.created_at.timestamp_millis(), - s.content_path, - s.content_sha256, - ])?; - } - Ok(staged.len()) + tinycortex::memory::chunks::upsert_staged_chunks_tx(tx, chunks) } -/// Repair the stored body-sha token for one chunk (#4689). -/// -/// The chunk content file is content-addressed and atomically written, so it is -/// the source of truth for its body. When a read detects that the on-disk body -/// no longer hashes to the recorded `content_sha256` (e.g. an external editor -/// rewrote a synced file after ingest), the reader serves the full on-disk body -/// and calls this to re-point the stale token at the disk bytes so the next read -/// verifies cleanly instead of falling back to the ≤500-char preview. -pub fn update_chunk_content_sha256( - config: &Config, - 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_chunk_content_sha256(config: &Config, id: &str, sha256: &str) -> Result<()> { + tinycortex::memory::chunks::update_chunk_content_sha256(&engine_config(config), id, sha256) } -/// Summary counterpart of [`update_chunk_content_sha256`] (#4689). -pub fn update_summary_content_sha256( - config: &Config, - 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 update_summary_content_sha256(config: &Config, id: &str, sha256: &str) -> Result<()> { + tinycortex::memory::chunks::update_summary_content_sha256(&engine_config(config), id, sha256) } -/// List the distinct `source_id`s of chunk rows for `source_kind` whose id -/// starts with `source_id_prefix` (#4689). -/// -/// Used by folder/list-based resync to discover previously-ingested items that -/// have since vanished (e.g. a renamed or deleted file) so their stale rows + -/// on-disk bodies can be cleaned. The prefix is applied Rust-side (literal, not -/// a SQL `LIKE`) so ids containing `_`/`%` are matched verbatim — matching the -/// convention in [`delete_chunks_by_source_prefix`]. pub fn list_source_ids_with_prefix( config: &Config, - source_kind: SourceKind, - source_id_prefix: &str, + kind: SourceKind, + prefix: &str, ) -> Result> { - with_connection(config, |conn| { - // Bounded range scan on the `idx_mem_tree_chunks_source (source_kind, - // source_id)` index instead of scanning every source_id for the kind and - // filtering in Rust: `source_id >= prefix AND source_id < upper`, where - // `upper` is the prefix with its last byte incremented. With SQLite's - // default BINARY collation this is exactly the literal byte-prefix set - // (so `_`/`%` stay literal, matching `delete_chunks_by_source_prefix`). - let out = match prefix_upper_bound(source_id_prefix) { - Some(upper) => { - let mut stmt = conn.prepare( - "SELECT DISTINCT source_id FROM mem_tree_chunks \ - WHERE source_kind = ?1 AND source_id >= ?2 AND source_id < ?3", - )?; - let rows = stmt.query_map( - params![source_kind.as_str(), source_id_prefix, upper], - |row| row.get::<_, String>(0), - )?; - rows.collect::>>() - } - None => { - // Empty prefix (or all-0xFF): no finite upper bound. The lower - // bound still uses the index; every row of the kind qualifies — - // identical to the previous `starts_with("")` behaviour. - let mut stmt = conn.prepare( - "SELECT DISTINCT source_id FROM mem_tree_chunks \ - WHERE source_kind = ?1 AND source_id >= ?2", - )?; - let rows = stmt - .query_map(params![source_kind.as_str(), source_id_prefix], |row| { - row.get::<_, String>(0) - })?; - rows.collect::>>() - } - } - .context("Failed to list mem_tree chunk source_ids by prefix")?; - Ok(out) - }) -} - -/// Exclusive upper bound for a literal byte-prefix range scan: the least string -/// strictly greater than every string starting with `prefix`. `None` when no -/// finite bound exists (empty prefix, or every byte is `0xFF`), in which case -/// the caller applies only the lower bound. -fn prefix_upper_bound(prefix: &str) -> Option { - let mut bytes = prefix.as_bytes().to_vec(); - while let Some(&last) = bytes.last() { - if last < 0xFF { - *bytes.last_mut().unwrap() = last + 1; - // Our prefixes are ASCII (`mem_src::`), so incrementing the last - // byte yields valid UTF-8. If a future caller passes a prefix that - // ends mid-codepoint, fall back to the lower-bound-only path rather - // than emit invalid UTF-8. - return String::from_utf8(bytes).ok(); - } - bytes.pop(); - } - None -} - -#[cfg(test)] -mod prefix_bound_tests { - use super::prefix_upper_bound; - - #[test] - fn upper_bound_increments_last_byte() { - assert_eq!( - prefix_upper_bound("mem_src:s1:").as_deref(), - Some("mem_src:s1;") - ); - assert_eq!(prefix_upper_bound("a").as_deref(), Some("b")); - } - - #[test] - fn upper_bound_none_for_empty_prefix() { - // Empty prefix is the only reachable `None` for a valid UTF-8 `&str` (no - // valid string ends in 0xFF); the all-0xFF `pop` loop is defensive. - assert_eq!(prefix_upper_bound(""), None); - } - - #[test] - fn upper_bound_handles_multibyte_prefix() { - // A prefix ending in a multibyte codepoint still yields a valid bound. - assert_eq!(prefix_upper_bound("café").as_deref(), Some("cafê")); - } - - #[test] - fn range_excludes_sibling_prefix() { - // The trailing delimiter in the prefix is what keeps `mem_src:s1:` from - // matching `mem_src:s10:...`: `s10:` sorts below `s1:` (`'0' < ':'`), so - // it falls below the lower bound and is excluded — same as `starts_with`. - let prefix = "mem_src:s1:"; - let upper = prefix_upper_bound(prefix).unwrap(); - let sib = "mem_src:s10:x"; - assert!(!(sib >= prefix && *sib < *upper.as_str())); - let child = "mem_src:s1:a"; - assert!(child >= prefix && *child < *upper.as_str()); - } -} - -fn upsert_chunks_with_statement( - stmt: &mut rusqlite::Statement<'_>, - chunks: &[Chunk], -) -> Result<()> { - for chunk in chunks { - stmt.execute(params![ - chunk.id, - chunk.metadata.source_kind.as_str(), - chunk.metadata.source_id, - chunk.metadata.path_scope, - chunk.metadata.source_ref.as_ref().map(|r| r.value.as_str()), - chunk.metadata.owner, - chunk.metadata.timestamp.timestamp_millis(), - chunk.metadata.time_range.0.timestamp_millis(), - chunk.metadata.time_range.1.timestamp_millis(), - serde_json::to_string(&chunk.metadata.tags)?, - chunk.content, - chunk.token_count, - chunk.seq_in_source, - chunk.created_at.timestamp_millis(), - ])?; - } - Ok(()) -} - -/// Fetch one chunk by its id. -/// Map the host `Config` to the engine `MemoryConfig` addressing the same -/// `/memory_tree/chunks.db` (only `workspace` is load-bearing for -/// these delegating DB reads). W3 store-op flip. -fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { - memory_config_from(config, config.workspace_dir.clone()) + tinycortex::memory::chunks::list_source_ids_with_prefix(&engine_config(config), kind, prefix) } pub fn get_chunk(config: &Config, id: &str) -> Result> { tinycortex::memory::chunks::get_chunk(&engine_config(config), id) } -/// Batched read of full chunk rows by id — delegates to the crate. -pub fn get_chunks_batch(config: &Config, chunk_ids: &[String]) -> Result> { - tinycortex::memory::chunks::get_chunks_batch(&engine_config(config), chunk_ids) +pub fn get_chunks_batch(config: &Config, ids: &[String]) -> Result> { + tinycortex::memory::chunks::get_chunks_batch(&engine_config(config), ids) } -/// Query parameters for [`list_chunks`], re-exported from the crate (identical -/// fields incl. the `source_scope` allowlist + `exclude_dropped`). -pub use tinycortex::memory::chunks::ListChunksQuery; - -/// List chunks matching the provided filters, ordered by `timestamp` DESC. -/// -/// Delegates to the crate, which preserves the `source_scope` allowlist gate -/// (byte-identical `chunk_source_allowed_in` + `extract_mem_src_id`) — the -/// security-critical per-profile enforcement, pinned by -/// `store_tests::list_chunks_source_scope_filters_before_limit`. pub fn list_chunks(config: &Config, query: &ListChunksQuery) -> Result> { tinycortex::memory::chunks::list_chunks(&engine_config(config), query) } -/// Count total chunks in the store (useful for tests / diagnostics). pub fn count_chunks(config: &Config) -> Result { tinycortex::memory::chunks::count_chunks(&engine_config(config)) } -/// #002 (FR-010 / US5): extraction coverage — the fraction of chunks that have -/// at least one indexed entity in `mem_tree_entity_index`, in `[0.0, 1.0]`. -/// -/// Turns "wiki built / not built" into a quality signal: a value near 0 with a -/// non-zero chunk count means extraction is producing nothing (the model is -/// timing out / failing), even though chunks exist — the "empty-but-built -/// wiki" symptom. Joins the entity index against `mem_tree_chunks.id` so the -/// numerator is node-kind-agnostic (we only count entity rows whose `node_id` -/// is an actual chunk). Returns `0.0` when there are no chunks. pub fn extraction_coverage(config: &Config) -> Result { tinycortex::memory::chunks::extraction_coverage(&engine_config(config)) } -/// Set the lifecycle status column for `chunk_id`. See `CHUNK_STATUS_*`. -pub fn set_chunk_lifecycle_status(config: &Config, chunk_id: &str, status: &str) -> Result<()> { - tinycortex::memory::chunks::set_chunk_lifecycle_status(&engine_config(config), chunk_id, status) +pub fn set_chunk_lifecycle_status(config: &Config, id: &str, status: &str) -> Result<()> { + tinycortex::memory::chunks::set_chunk_lifecycle_status(&engine_config(config), id, status) } pub(crate) fn set_chunk_lifecycle_status_tx( tx: &Transaction<'_>, - chunk_id: &str, + id: &str, status: &str, ) -> Result<()> { - set_chunk_lifecycle_status_conn(tx, chunk_id, status) + tinycortex::memory::chunks::set_chunk_lifecycle_status_tx(tx, id, status) } -/// Read the lifecycle status column for `chunk_id`, or `None` if the row is absent. -pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result> { - tinycortex::memory::chunks::get_chunk_lifecycle_status(&engine_config(config), chunk_id) +pub fn get_chunk_lifecycle_status(config: &Config, id: &str) -> Result> { + tinycortex::memory::chunks::get_chunk_lifecycle_status(&engine_config(config), id) } pub(crate) fn get_chunk_lifecycle_status_tx( tx: &Transaction<'_>, - chunk_id: &str, + id: &str, ) -> Result> { - get_chunk_lifecycle_status_conn(tx, chunk_id) + tinycortex::memory::chunks::get_chunk_lifecycle_status_tx(tx, id) } -fn get_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str) -> Result> { - let row = conn - .query_row( - "SELECT lifecycle_status FROM mem_tree_chunks WHERE id = ?1", - params![chunk_id], - |r| r.get::<_, String>(0), - ) - .optional()?; - Ok(row) -} - -/// Count chunks currently sitting at a given lifecycle status (test/diagnostic helper). pub fn count_chunks_by_lifecycle_status(config: &Config, status: &str) -> Result { tinycortex::memory::chunks::count_chunks_by_lifecycle_status(&engine_config(config), status) } -fn set_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str, status: &str) -> Result<()> { - let changed = conn.execute( - "UPDATE mem_tree_chunks SET lifecycle_status = ?1 WHERE id = ?2", - params![status, chunk_id], - )?; - if changed == 0 { - log::warn!( - "[memory::chunk_store] lifecycle update affected 0 rows chunk_id={} status={}", - chunk_id, - status - ); - } - Ok(()) +pub fn is_source_ingested(config: &Config, kind: SourceKind, id: &str) -> Result { + tinycortex::memory::chunks::is_source_ingested(&engine_config(config), kind, id) } -/// Best-effort, non-transactional check used by `ingest_*` to skip -/// canonicalisation when a source has already been ingested. The -/// authoritative gate is [`claim_source_ingest_tx`] inside the persist -/// transaction — this lookup just avoids burning canonicaliser work on -/// the obvious dup case. -pub fn is_source_ingested( - config: &Config, - source_kind: SourceKind, - source_id: &str, -) -> Result { - tinycortex::memory::chunks::is_source_ingested(&engine_config(config), source_kind, source_id) -} - -/// Atomically claim `(source_kind, source_id)` for ingestion. Returns -/// `true` if the row was newly inserted (caller should proceed with the -/// rest of the persist transaction); `false` if a previous ingest already -/// claimed this source (caller must roll back / skip). -/// -/// Lives inside the same transaction as the chunk + job writes so two -/// concurrent ingests of the same source can't both pass the gate. pub(crate) fn claim_source_ingest_tx( tx: &Transaction<'_>, - source_kind: SourceKind, - source_id: &str, + kind: SourceKind, + id: &str, now_ms: i64, ) -> Result { - let inserted = tx.execute( - "INSERT OR IGNORE INTO mem_tree_ingested_sources \ - (source_kind, source_id, ingested_at_ms) \ - VALUES (?1, ?2, ?3)", - params![source_kind.as_str(), source_id, now_ms], - )?; - Ok(inserted > 0) + tinycortex::memory::chunks::claim_source_ingest_tx(tx, kind, id, now_ms) } -/// `source_kind` value used in `mem_tree_ingested_sources` to record that a -/// raw archive file (relative path under `/`, e.g. -/// `raw/github-com-org-repo/commits/_.md`) has been covered by a -/// tree summary. Distinct from the chunk-store [`SourceKind`] values so the -/// two gate namespaces can never collide. -pub const RAW_FILE_GATE_KIND: &str = "raw_file"; - -/// Record that the given raw archive files (relative paths under -/// `/`) are covered by a tree summary. Idempotent -/// (`INSERT OR IGNORE`); returns the number of newly-recorded paths. -pub fn mark_raw_paths_ingested(config: &Config, rel_paths: &[String]) -> Result { - tinycortex::memory::chunks::mark_raw_paths_ingested(&engine_config(config), rel_paths) +pub fn mark_raw_paths_ingested(config: &Config, paths: &[String]) -> Result { + tinycortex::memory::chunks::mark_raw_paths_ingested(&engine_config(config), paths) } -/// Filter `rel_paths` down to the ones NOT yet recorded as ingested raw -/// files. Order of the surviving paths is preserved. -pub fn filter_raw_paths_not_ingested(config: &Config, rel_paths: &[String]) -> Result> { - tinycortex::memory::chunks::filter_raw_paths_not_ingested(&engine_config(config), rel_paths) +pub fn filter_raw_paths_not_ingested(config: &Config, paths: &[String]) -> Result> { + tinycortex::memory::chunks::filter_raw_paths_not_ingested(&engine_config(config), paths) } -/// Count raw-file gate rows whose path starts with `rel_prefix` (e.g. -/// `raw/github-com-org-repo/`). Diagnostic helper for reconcile reporting. -pub fn count_raw_paths_ingested_with_prefix(config: &Config, rel_prefix: &str) -> Result { - tinycortex::memory::chunks::count_raw_paths_ingested_with_prefix( - &engine_config(config), - rel_prefix, - ) +pub fn count_raw_paths_ingested_with_prefix(config: &Config, prefix: &str) -> Result { + tinycortex::memory::chunks::count_raw_paths_ingested_with_prefix(&engine_config(config), prefix) } -/// Delete all chunk rows for one exact `(source_kind, source_id)` and clear -/// dependent source-local indexes. Returns the number of chunk rows removed. -pub fn delete_chunks_by_source( - config: &Config, - source_kind: SourceKind, - source_id: &str, -) -> Result { - delete_chunks_by_source_filter( - "delete_chunks_by_source", - config, - source_kind, - |candidate, _owner| candidate == source_id, - |candidate| candidate == source_id, - ) +pub fn delete_chunks_by_source(config: &Config, kind: SourceKind, id: &str) -> Result { + tinycortex::memory::chunks::delete_chunks_by_source(&engine_config(config), kind, id) } -/// Delete all chunk rows whose source id starts with `source_id_prefix`. -/// -/// This is intentionally a Rust-side prefix filter rather than a SQL `LIKE` -/// expression so provider ids containing `_` / `%` are treated literally. pub fn delete_chunks_by_source_prefix( config: &Config, - source_kind: SourceKind, - source_id_prefix: &str, + kind: SourceKind, + prefix: &str, ) -> Result { - delete_chunks_by_source_filter( - "delete_chunks_by_source_prefix", - config, - source_kind, - |candidate, _owner| candidate.starts_with(source_id_prefix), - |candidate| candidate.starts_with(source_id_prefix), - ) + tinycortex::memory::chunks::delete_chunks_by_source_prefix(&engine_config(config), kind, prefix) } -/// Delete all chunk rows for one exact `(source_kind, owner)` while preserving -/// source ingest gates that still have chunks owned by another connection. -pub fn delete_chunks_by_owner( - config: &Config, - source_kind: SourceKind, - owner: &str, -) -> Result { - delete_chunks_by_source_filter( - "delete_chunks_by_owner", - config, - source_kind, - |_source_id, candidate_owner| candidate_owner == owner, - |_source_id| false, - ) +pub fn delete_chunks_by_owner(config: &Config, kind: SourceKind, owner: &str) -> Result { + tinycortex::memory::chunks::delete_chunks_by_owner(&engine_config(config), kind, owner) } -fn delete_chunks_by_source_filter( - op: &str, - config: &Config, - source_kind: SourceKind, - matches_chunk: impl Fn(&str, &str) -> bool, - matches_ingested_source: impl Fn(&str) -> bool, -) -> Result { - let mut content_paths = Vec::new(); - let deleted = with_connection(config, |conn| { - let tx = conn.unchecked_transaction()?; - - let chunks = { - let mut stmt = tx.prepare( - "SELECT id, source_id, owner, content_path - FROM mem_tree_chunks - WHERE source_kind = ?1", - )?; - let rows = stmt.query_map(params![source_kind.as_str()], |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - row.get::<_, String>(2)?, - row.get::<_, Option>(3)?, - )) - })?; - 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(_) => None, - Err(error) => Some(Err(error)), - }) - .collect::>>() - .context("Failed to collect memory_tree chunks by source")? - }; - - let deleted_source_ids: HashSet = chunks - .iter() - .map(|(_, source_id, _)| source_id.clone()) - .collect(); - - for (chunk_id, _source_id, content_path) in &chunks { - tx.execute( - "DELETE FROM mem_tree_score WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunk_embeddings WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunk_reembed_skipped WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunks WHERE id = ?1", - params![chunk_id], - )?; - if let Some(path) = content_path.as_ref().filter(|path| !path.is_empty()) { - content_paths.push(path.clone()); - } - } - - let mut orphaned_deleted_sources = HashSet::new(); - for source_id in &deleted_source_ids { - 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 { - log::debug!( - "[memory::chunk_store] {op}: source_id_hash={} orphaned; removing ingest gate", - redact_value(source_id), - ); - orphaned_deleted_sources.insert(source_id.clone()); - } else { - log::debug!( - "[memory::chunk_store] {op}: source_id_hash={} remaining_chunks={remaining}; preserving ingest gate", - redact_value(source_id), - ); - } - } - - let ingested_sources = { - 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(source_id) - if matches_ingested_source(&source_id) - || orphaned_deleted_sources.contains(&source_id) => - { - Some(Ok(source_id)) - } - Ok(_) => None, - Err(error) => Some(Err(error)), - }) - .collect::>>() - .context("Failed to collect memory_tree ingested sources")? - }; - - for source_id in &ingested_sources { - tx.execute( - "DELETE FROM mem_tree_ingested_sources - WHERE source_kind = ?1 AND source_id = ?2", - params![source_kind.as_str(), source_id], - )?; - } - - // A fully-orphaned source has zero chunks left, so its summary tree - // now summarises deleted content — and its unsealed buffer holds - // dangling chunk ids. Cascade-delete the tree (summaries + sidecars - // + entity-index + buffer + tree row) so a `clear_memory` delete is - // complete and stale summaries can't resurface in retrieval. Source - // trees use the chunk `source_id` verbatim as their scope, so we - // match on that. Same tx as the chunk delete → atomic. - for source_id in &orphaned_deleted_sources { - if let Some(tree) = - crate::openhuman::memory_store::trees::store::get_tree_by_scope_conn( - &tx, - crate::openhuman::memory_store::trees::types::TreeKind::Source, - source_id, - )? - { - let cascade = crate::openhuman::memory_store::trees::store::delete_tree_cascade_tx( - &tx, &tree.id, - )?; - // Defer the summary content-file removal to the same - // post-commit sweep as the chunk files. - content_paths.extend(cascade.content_paths); - log::debug!( - "[memory::chunk_store] {op}: orphaned source_id_hash={} → deleted source tree tree_id={} summaries={}", - redact_value(source_id), - tree.id, - cascade.removed_summaries, - ); - } - } - - let deleted = chunks.len(); - tx.commit()?; - Ok(deleted) - })?; - - remove_chunk_content_files(config, &content_paths); - Ok(deleted) -} - -/// Finish off an orphaned **Source** (one with zero chunks remaining): clear its -/// ingest dedup gates and cascade-delete its source-scoped summary tree. -/// -/// `delete_chunks_by_source` only cascades the tree for sources whose chunks it -/// deletes in the same call; a source whose chunks were already removed earlier -/// (e.g. by the per-chunk `delete_chunk` path) keeps a now-stale summary tree -/// that can still resurface in recall. This cleans up exactly that **legacy -/// partial-delete** state. -/// -/// Specifically, when no chunks remain it: -/// - removes the ingest dedup gates for the source — both the bare `source_id` -/// AND any versioned `{source_id}@{version_ms}` gates (matched Rust-side with -/// exact/prefix comparison, never SQL `LIKE`/`GLOB`, to avoid metachar pitfalls); -/// - cascades the **source-scoped** tree (scope == `source_id`) if present. -/// -/// Scoped-collection conservatism: a document ingested under a shared collection -/// `path_scope` (e.g. Notion `notion:{connection}`) lives in a tree scoped by that -/// `path_scope`, NOT by this `source_id`, so `get_tree_by_scope(Source, -/// source_id)` returns `None` and such shared trees are left intact — deleting one -/// document must never tear down a tree that summarises many documents. -/// -/// Returns `true` when a source-scoped tree was removed (drives the RPC's -/// `deleted` flag). No-op-safe to call unconditionally after -/// `delete_chunks_by_source`. -pub fn delete_orphaned_source_tree( - config: &Config, - source_kind: SourceKind, - source_id: &str, -) -> Result { - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::TreeKind; - - let mut content_paths: Vec = Vec::new(); - let tree_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], - |r| r.get(0), - )?; - if remaining > 0 { - // Source still has chunks — not orphaned; leave its live tree + gates. - log::debug!( - "[memory::chunk_store] delete_orphaned_source_tree: source_id_hash={} still has {remaining} chunk(s) — no-op", - redact_value(source_id), - ); - return Ok(false); - } - - // Clear ALL ingest dedup gates for this source: the bare source_id and any - // versioned `{source_id}@{version_ms}` gates. Filter in Rust (exact or - // `source_id@` prefix) so `_`/`%`/glob chars in ids are treated literally. - let versioned_prefix = format!("{source_id}@"); - let gate_ids: Vec = { - 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()], |r| r.get::<_, String>(0))?; - rows.filter_map(|row| match row { - Ok(s) if s == source_id || s.starts_with(&versioned_prefix) => Some(Ok(s)), - Ok(_) => None, - Err(e) => Some(Err(e)), - }) - .collect::>>()? - }; - for gid in &gate_ids { - tx.execute( - "DELETE FROM mem_tree_ingested_sources WHERE source_kind = ?1 AND source_id = ?2", - params![source_kind.as_str(), gid], - )?; - } - - // Cascade the source-scoped orphan tree if one exists. Shared - // collection/path_scope trees are not keyed by this source_id (see fn - // docs), so they are intentionally left untouched. - let cascaded = if let Some(tree) = - tree_store::get_tree_by_scope_conn(&tx, TreeKind::Source, source_id)? - { - let cascade = tree_store::delete_tree_cascade_tx(&tx, &tree.id)?; - content_paths.extend(cascade.content_paths); - log::debug!( - "[memory::chunk_store] delete_orphaned_source_tree: source_id_hash={} → removed stale tree_id={} summaries={} gates_cleared={}", - redact_value(source_id), - tree.id, - cascade.removed_summaries, - gate_ids.len(), - ); - true - } else { - log::debug!( - "[memory::chunk_store] delete_orphaned_source_tree: source_id_hash={} has no source-scoped tree (gates_cleared={}); shared/collection trees left intact", - redact_value(source_id), - gate_ids.len(), - ); - false - }; - tx.commit()?; - Ok(cascaded) - })?; - if tree_cascaded { - remove_chunk_content_files(config, &content_paths); - } - Ok(tree_cascaded) -} - -fn remove_chunk_content_files(config: &Config, content_paths: &[String]) { - use std::path::{Component, Path}; - - let root = config.memory_tree_content_root(); - let canonical_root = match std::fs::canonicalize(&root) { - Ok(path) => path, - Err(error) => { - if error.kind() != std::io::ErrorKind::NotFound { - log::warn!( - "[memory_tree::store] failed to resolve content root {}: {error}", - root.display(), - ); - } - return; - } - }; - - for rel in content_paths { - let rel_path = Path::new(rel); - let has_escape_component = rel_path.components().any(|component| { - matches!( - component, - Component::ParentDir | Component::RootDir | Component::Prefix(_) - ) - }); - if has_escape_component { - log::warn!( - "[memory_tree::store] refusing to remove chunk file with unsafe content_path path_hash={}", - redact::redact(rel), - ); - continue; - } - - let path = root.join(rel_path); - let resolved_path = match std::fs::canonicalize(&path) { - Ok(path) => path, - Err(error) => { - if error.kind() != std::io::ErrorKind::NotFound { - log::warn!( - "[memory_tree::store] failed to resolve chunk file path_hash={}: {error}", - redact::redact(rel), - ); - } - continue; - } - }; - if !resolved_path.starts_with(&canonical_root) { - log::warn!( - "[memory_tree::store] refusing to remove chunk file outside content root path_hash={}", - redact::redact(rel), - ); - continue; - } - - if let Err(error) = std::fs::remove_file(&path) { - if error.kind() != std::io::ErrorKind::NotFound { - log::warn!( - "[memory_tree::store] failed to remove chunk file path_hash={}: {error}", - redact::redact(rel), - ); - } - } - } +pub fn delete_orphaned_source_tree(config: &Config, kind: SourceKind, id: &str) -> Result { + tinycortex::memory::chunks::delete_orphaned_source_tree(&engine_config(config), kind, id) } #[path = "connection.rs"] mod connection; pub(crate) use connection::recover_corrupt_db; pub use connection::with_connection; -#[cfg(test)] -#[allow(unused_imports)] -pub(crate) use connection::{ - clear_connection_cache, db_path_for, get_or_init_connection, invalidate_connection, - is_io_open_error, schema_apply_count_for_path_for_tests, CB_THRESHOLD, -}; -#[cfg(test)] -pub(crate) use connection::{is_transient_cold_start, try_cleanup_stale_files}; - -#[path = "migrations.rs"] -mod migrations; -use migrations::{migrate_legacy_embeddings_to_sidecar, purge_global_topic_trees}; #[path = "raw_refs.rs"] mod raw_refs; 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, + list_summaries_with_content_path, set_chunk_raw_refs, set_chunk_raw_refs_tx, }; -/// Idempotent `ALTER TABLE ADD COLUMN` — treats an existing column as success. -fn add_column_if_missing(conn: &Connection, table: &str, name: &str, sql_type: &str) -> Result<()> { - match conn.execute( - &format!("ALTER TABLE {table} ADD COLUMN {name} {sql_type}"), - [], - ) { - Ok(_) => { - log::debug!( - "[memory::chunk_store] migration: added column {table}.{name} ({sql_type})" - ); - Ok(()) - } - Err(err) if err.to_string().contains("duplicate column name") => Ok(()), - Err(err) => Err(err).with_context(|| format!("Failed to add column {table}.{name}")), - } -} - #[path = "embeddings.rs"] mod embeddings; pub use embeddings::{ @@ -1265,14 +162,6 @@ pub use embeddings::{ get_chunk_embeddings_for_signature_batch, mark_chunk_reembed_skipped, set_chunk_embedding, set_chunk_embedding_for_signature, }; -#[cfg(test)] -pub(crate) use embeddings::{embedding_to_blob, REEMBED_SKIP_KEY_MAX_LEN}; pub(crate) use embeddings::{ has_uncovered_reembed_work, set_chunk_embedding_for_signature_tx, tree_active_signature, - validate_reembed_skip_key, }; -// ── Phase 2: embedding column accessors ───────────────────────────────── - -#[cfg(test)] -#[path = "store_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_store/chunks/store_tests.rs b/src/openhuman/memory_store/chunks/store_tests.rs deleted file mode 100644 index c39f2fd21..000000000 --- a/src/openhuman/memory_store/chunks/store_tests.rs +++ /dev/null @@ -1,2175 +0,0 @@ -//! Unit tests for [`super`] — chunk upsert / list / lifecycle / embedding / -//! content-pointer accessors against a tempdir-backed SQLite store. -//! -//! ## Test isolation for the connection cache -//! -//! Because the connection cache is a process-level singleton, tests that want -//! to exercise cache behaviour (same Arc, independent workspaces, circuit -//! breaker, cleanup) must call `clear_connection_cache()` at the start — or -//! be careful to use unique tempdirs that cannot collide with other tests. -//! The call is cheap (a mutex lock + HashMap clear) and harmless for tests -//! that don't need it. - -use super::*; -// Imported directly (not via `super::*`): this PR's store-op delegation dropped -// store.rs's own `use chrono::Utc`, so the test module pulls it in itself. -use crate::openhuman::memory_store::chunks::types::{chunk_id, Metadata, SourceRef}; -use chrono::{TimeZone, Utc}; -use rusqlite::params; -use tempfile::TempDir; - -fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) -} - -fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk { - let ts = Utc.timestamp_millis_opt(ts_ms).unwrap(); - Chunk { - id: chunk_id(SourceKind::Chat, source_id, seq, "test-content"), - content: format!("content {source_id} {seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: source_id.to_string(), - owner: "alice@example.com".to_string(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["eng".into()], - source_ref: Some(SourceRef::new(format!("slack://{source_id}/{seq}"))), - path_scope: None, - }, - token_count: 12, - seq_in_source: seq, - created_at: ts, - partial_message: false, - } -} - -#[test] -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); - let got = get_chunk(&cfg, &c.id).unwrap().expect("chunk stored"); - assert_eq!(got, c); -} - -#[test] -fn upsert_persists_path_scope() { - let (_tmp, cfg) = test_config(); - let mut c = sample_chunk("notion:conn-1:page-abc", 0, 1_700_000_000_000); - 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); - - let got = get_chunk(&cfg, &c.id).unwrap().expect("chunk stored"); - assert_eq!(got.metadata.source_id, "notion:conn-1:page-abc"); - assert_eq!(got.metadata.path_scope.as_deref(), Some("notion:conn-1")); -} - -#[test] -fn list_chunks_source_scope_filters_before_limit() { - // Two disallowed-source chunks have NEWER timestamps (sorted first by DESC), - // and the single allowed-source chunk is older. With a naive post-limit - // filter and LIMIT 1 the allowed row would be starved; the before-limit gate - // inside list_chunks must still surface it. - let (_tmp, cfg) = test_config(); - let tag = || vec!["memory_sources".to_string(), "chat".to_string()]; - let mut bad1 = sample_chunk("slack:#secret", 0, 3_000); - bad1.metadata.tags = tag(); - let mut bad2 = sample_chunk("slack:#secret", 1, 2_000); - bad2.metadata.tags = tag(); - let mut good = sample_chunk("slack:#eng", 0, 1_000); - good.metadata.tags = tag(); - upsert_chunks(&cfg, &[bad1, bad2, good]).unwrap(); - - let mut allowed = std::collections::HashSet::new(); - allowed.insert("slack:#eng".to_string()); - let q = ListChunksQuery { - limit: Some(1), - source_scope: Some(allowed), - ..Default::default() - }; - let rows = list_chunks(&cfg, &q).unwrap(); - assert_eq!( - rows.len(), - 1, - "the allowed-source chunk must survive the gate" - ); - assert_eq!(rows[0].metadata.source_id, "slack:#eng"); - - // No scope → unrestricted: the newest (disallowed) chunk wins under LIMIT 1. - let unscoped = ListChunksQuery { - limit: Some(1), - ..Default::default() - }; - let rows = list_chunks(&cfg, &unscoped).unwrap(); - assert_eq!(rows[0].metadata.source_id, "slack:#secret"); -} - -#[test] -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(); - assert_eq!(count_chunks(&cfg).unwrap(), 1); -} - -#[test] -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(); - 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(); - - let embedding = get_chunk_embedding(&cfg, &c.id).unwrap().unwrap(); - assert_eq!(embedding, vec![0.1, 0.2, 0.3]); - let got = get_chunk(&cfg, &c.id).unwrap().unwrap(); - assert_eq!(got.content, "updated content"); - assert_eq!(got.token_count, 99); -} - -#[test] -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(); - - set_chunk_embedding_for_signature( - &cfg, - &c.id, - "openai/text-embedding-3-small@1536", - &[0.1, 0.2], - ) - .unwrap(); - set_chunk_embedding_for_signature(&cfg, &c.id, "local/bge-small@384", &[0.3, 0.4, 0.5]) - .unwrap(); - - assert_eq!( - get_chunk_embedding_for_signature(&cfg, &c.id, "openai/text-embedding-3-small@1536") - .unwrap(), - Some(vec![0.1, 0.2]) - ); - assert_eq!( - get_chunk_embedding_for_signature(&cfg, &c.id, "local/bge-small@384").unwrap(), - Some(vec![0.3, 0.4, 0.5]) - ); - assert!( - get_chunk_embedding_for_signature(&cfg, &c.id, "missing/model@1") - .unwrap() - .is_none() - ); - - // #1574 cutover: the public `get_chunk_embedding` now reads the sidecar at - // the *active* signature (not the legacy column). Nothing was written - // there yet, so it is absent — graceful, never a cross-space read of the - // openai/local rows above. - assert!(get_chunk_embedding(&cfg, &c.id).unwrap().is_none()); - - // The public setter targets the active signature and round-trips through - // the public getter — proves the cutover wiring end to end. - set_chunk_embedding(&cfg, &c.id, &[0.7, 0.8]).unwrap(); - assert_eq!( - get_chunk_embedding(&cfg, &c.id).unwrap(), - Some(vec![0.7, 0.8]) - ); - - // ...and the earlier per-signature rows remain independently scoped. - assert_eq!( - get_chunk_embedding_for_signature(&cfg, &c.id, "local/bge-small@384").unwrap(), - Some(vec![0.3, 0.4, 0.5]) - ); -} - -#[test] -fn list_filters_by_source_kind() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - let mut c2 = sample_chunk("gmail:t1", 0, 1_700_000_001_000); - c2.metadata.source_kind = SourceKind::Email; - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - let q = ListChunksQuery { - source_kind: Some(SourceKind::Email), - ..Default::default() - }; - let rows = list_chunks(&cfg, &q).unwrap(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].metadata.source_kind, SourceKind::Email); -} - -#[test] -fn list_filters_by_time_range() { - let (_tmp, cfg) = test_config(); - let a = sample_chunk("s", 0, 1_700_000_000_000); - let b = sample_chunk("s", 1, 1_700_000_010_000); - let c = sample_chunk("s", 2, 1_700_000_020_000); - upsert_chunks(&cfg, &[a.clone(), b.clone(), c.clone()]).unwrap(); - let q = ListChunksQuery { - since_ms: Some(1_700_000_005_000), - until_ms: Some(1_700_000_015_000), - ..Default::default() - }; - let rows = list_chunks(&cfg, &q).unwrap(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].id, b.id); -} - -#[test] -fn list_orders_by_timestamp_desc() { - let (_tmp, cfg) = test_config(); - let a = sample_chunk("s", 0, 1_700_000_000_000); - let b = sample_chunk("s", 1, 1_700_000_010_000); - upsert_chunks(&cfg, &[a.clone(), b.clone()]).unwrap(); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows.len(), 2); - assert_eq!(rows[0].id, b.id); // newest first - assert_eq!(rows[1].id, a.id); -} - -#[test] -fn list_orders_equal_timestamps_by_sequence() { - let (_tmp, cfg) = test_config(); - let a = sample_chunk("s", 0, 1_700_000_000_000); - let b = sample_chunk("s", 1, 1_700_000_000_000); - upsert_chunks(&cfg, &[b.clone(), a.clone()]).unwrap(); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows.len(), 2); - assert_eq!(rows[0].seq_in_source, 0); - assert_eq!(rows[1].seq_in_source, 1); -} - -#[test] -fn list_limit_is_clamped_to_sane_range() { - let (_tmp, cfg) = test_config(); - let chunks = (0..3) - .map(|idx| sample_chunk("s", idx, 1_700_000_000_000 + i64::from(idx))) - .collect::>(); - upsert_chunks(&cfg, &chunks).unwrap(); - - let zero_limit = list_chunks( - &cfg, - &ListChunksQuery { - limit: Some(0), - ..Default::default() - }, - ) - .unwrap(); - assert_eq!(zero_limit.len(), 1); - - let huge_limit = list_chunks( - &cfg, - &ListChunksQuery { - limit: Some(usize::MAX), - ..Default::default() - }, - ) - .unwrap(); - assert_eq!(huge_limit.len(), 3); -} - -#[test] -fn delete_chunks_by_source_removes_chunks_side_rows_and_ingest_gate() { - let (_tmp, cfg) = test_config(); - let target_a = sample_chunk("slack:c-1", 0, 1_700_000_000_000); - let target_b = sample_chunk("slack:c-1", 1, 1_700_000_001_000); - let other = sample_chunk("slack:c-2", 0, 1_700_000_002_000); - upsert_chunks(&cfg, &[target_a.clone(), target_b.clone(), other.clone()]).unwrap(); - - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - for chunk in [&target_a, &target_b, &other] { - tx.execute( - "INSERT INTO mem_tree_score ( - chunk_id, total, token_count_signal, unique_words_signal, - metadata_weight, source_weight, interaction_weight, - entity_density, dropped, reason, computed_at_ms - ) VALUES (?1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0, NULL, 1700000000000)", - params![chunk.id], - )?; - tx.execute( - "INSERT INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms - ) VALUES (?1, ?2, 'chunk', 'person', 'chat', 0.9, 1700000000000)", - params![format!("entity:{}", chunk.id), chunk.id], - )?; - tx.execute( - "INSERT INTO mem_tree_chunk_embeddings ( - chunk_id, model_signature, vector, dim, created_at - ) VALUES (?1, 'test/model@3', ?2, 3, 1700000000.0)", - params![chunk.id, vec![1_u8, 2, 3]], - )?; - tx.execute( - "INSERT INTO mem_tree_chunk_reembed_skipped ( - chunk_id, model_signature, reason, skipped_at_ms - ) VALUES (?1, 'test/model@3', 'terminal', 1700000000000)", - params![chunk.id], - )?; - } - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Chat, - "slack:c-1", - 1_700_000_000_000 - )?); - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Chat, - "slack:c-2", - 1_700_000_000_000 - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let deleted = delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:c-1").unwrap(); - - assert_eq!(deleted, 2); - assert_eq!(count_chunks(&cfg).unwrap(), 1); - assert!(get_chunk(&cfg, &target_a.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &target_b.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &other.id).unwrap().is_some()); - assert!(!is_source_ingested(&cfg, SourceKind::Chat, "slack:c-1").unwrap()); - assert!(is_source_ingested(&cfg, SourceKind::Chat, "slack:c-2").unwrap()); - - with_connection(&cfg, |conn| { - let count_by_table = |table: &str| -> rusqlite::Result { - conn.query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |r| r.get(0)) - }; - assert_eq!(count_by_table("mem_tree_score")?, 1); - assert_eq!(count_by_table("mem_tree_entity_index")?, 1); - assert_eq!(count_by_table("mem_tree_chunk_embeddings")?, 1); - assert_eq!(count_by_table("mem_tree_chunk_reembed_skipped")?, 1); - Ok(()) - }) - .unwrap(); -} - -/// Forget-path (`clear_memory=true`) e2e: deleting the last chunk of a source -/// must cascade-delete its summary tree (tree row + summaries + sidecars + -/// entity-index + unsealed buffer), leave a sibling source untouched, and a -/// queued `Seal` job for the now-gone tree must settle to `Done` (not stick -/// in pending). Mocked connection (tempdir), chunks, tree/summary/buffer, job. -#[tokio::test] -async fn clear_memory_delete_cascades_orphaned_source_tree_and_settles_queued_job() { - use crate::openhuman::memory_queue::{store as queue_store, types as queue_types}; - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::{ - Buffer, SummaryNode, Tree, TreeKind, TreeStatus, - }; - - let (_tmp, cfg) = test_config(); - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - - // ---- mocked chunks: gmail:acct (conn-1, disconnecting) + gmail:other (conn-2, survives) ---- - let mk_email = |source_id: &str, seq: u32, owner: &str, ts_ms: i64| { - let mut c = sample_chunk(source_id, seq, ts_ms); - c.metadata.source_kind = SourceKind::Email; - c.metadata.owner = owner.to_string(); - c - }; - let a0 = mk_email("gmail:acct", 0, "gmail-sync:conn-1", 1_700_000_000_000); - let a1 = mk_email("gmail:acct", 1, "gmail-sync:conn-1", 1_700_000_001_000); - let b0 = mk_email("gmail:other", 0, "gmail-sync:conn-2", 1_700_000_002_000); - upsert_chunks(&cfg, &[a0.clone(), a1.clone(), b0.clone()]).unwrap(); - - // ---- mocked source trees (scope == source_id), each with summary + sidecars + entity-index + buffer ---- - let mk_tree = |id: &str, scope: &str| Tree { - id: id.into(), - kind: TreeKind::Source, - scope: scope.into(), - root_id: None, - max_level: 1, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }; - tree_store::insert_tree(&cfg, &mk_tree("tree-acct", "gmail:acct")).unwrap(); - tree_store::insert_tree(&cfg, &mk_tree("tree-other", "gmail:other")).unwrap(); - - let mk_summary = |id: &str, tree_id: &str, children: Vec| SummaryNode { - id: id.into(), - tree_id: tree_id.into(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: children, - content: format!("summary for {tree_id}"), - token_count: 3, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - - tree_store::insert_summary_tx( - &tx, - &mk_summary("sum-acct", "tree-acct", vec![a0.id.clone(), a1.id.clone()]), - None, - "test/model@3", - )?; - tree_store::insert_summary_tx( - &tx, - &mk_summary("sum-other", "tree-other", vec![b0.id.clone()]), - None, - "test/model@3", - )?; - - // summary sidecars: embeddings for both summaries, reembed-skip only for sum-acct. - for sid in ["sum-acct", "sum-other"] { - tx.execute( - "INSERT INTO mem_tree_summary_embeddings ( - summary_id, model_signature, vector, dim, created_at - ) VALUES (?1, 'test/model@3', ?2, 3, 1700000000.0)", - params![sid, vec![1_u8, 2, 3]], - )?; - } - tx.execute( - "INSERT INTO mem_tree_summary_reembed_skipped ( - summary_id, model_signature, reason, skipped_at_ms - ) VALUES ('sum-acct', 'test/model@3', 'terminal', 1700000000000)", - [], - )?; - - // tree-keyed entity-index rows (summary nodes) for each tree. - for (sid, tree_id) in [("sum-acct", "tree-acct"), ("sum-other", "tree-other")] { - tx.execute( - "INSERT INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, 'summary', 'person', 'email', 0.9, 1700000000000, ?3, 0)", - params![format!("entity:{sid}"), sid, tree_id], - )?; - } - - // unsealed buffers (the "queue" frontier) referencing the chunk ids. - tree_store::upsert_buffer_tx( - &tx, - &Buffer { - tree_id: "tree-acct".into(), - level: 0, - item_ids: vec![a0.id.clone(), a1.id.clone()], - token_sum: 24, - oldest_at: Some(ts), - }, - )?; - tree_store::upsert_buffer_tx( - &tx, - &Buffer { - tree_id: "tree-other".into(), - level: 0, - item_ids: vec![b0.id.clone()], - token_sum: 12, - oldest_at: Some(ts), - }, - )?; - - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Email, - "gmail:acct", - 1_700_000_000_000 - )?); - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Email, - "gmail:other", - 1_700_000_000_000 - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - // ---- mocked job: a Seal queued for the tree that's about to be deleted ---- - let seal_payload = queue_types::SealPayload { - tree_id: "tree-acct".into(), - level: 0, - force_now_ms: None, - }; - let job_id = queue_store::enqueue(&cfg, &queue_types::NewJob::seal(&seal_payload).unwrap()) - .unwrap() - .expect("seal job enqueued"); - - // ---- act: disconnect conn-1 with clear_memory=true → delete its chunks ---- - let deleted = delete_chunks_by_owner(&cfg, SourceKind::Email, "gmail-sync:conn-1").unwrap(); - assert_eq!(deleted, 2); - - // chunks: acct gone, other survives. - assert!(get_chunk(&cfg, &a0.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &a1.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &b0.id).unwrap().is_some()); - - // the orphaned source tree is gone; the sibling tree is untouched. - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, "gmail:acct") - .unwrap() - .is_none() - ); - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, "gmail:other") - .unwrap() - .is_some() - ); - - // exactly the tree-acct rows are cascaded away across every dependent table. - with_connection(&cfg, |conn| { - let count = |sql: &str| -> rusqlite::Result { conn.query_row(sql, [], |r| r.get(0)) }; - assert_eq!(count("SELECT COUNT(*) FROM mem_tree_trees")?, 1); - assert_eq!(count("SELECT COUNT(*) FROM mem_tree_summaries")?, 1); - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_summary_embeddings")?, - 1 - ); - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_summary_reembed_skipped")?, - 0 - ); - assert_eq!(count("SELECT COUNT(*) FROM mem_tree_buffers")?, 1); - assert_eq!(count("SELECT COUNT(*) FROM mem_tree_entity_index")?, 1); - // and what survives belongs to tree-other. - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_summaries WHERE tree_id = 'tree-other'")?, - 1 - ); - Ok(()) - }) - .unwrap(); - - // ---- the queued Seal job settles to Done (tree missing), not stuck pending ---- - // Drive the crate queue engine end-to-end through the seam - // (`drain_until_idle` → `worker::run_once` → crate claim/dispatch/settle), - // the flipped W4 production path — rather than the deleted host engine. - crate::openhuman::memory_queue::drain_until_idle(&cfg) - .await - .expect("drain queue"); - assert_eq!( - queue_store::get_job(&cfg, &job_id).unwrap().unwrap().status, - queue_types::JobStatus::Done, - "seal over a deleted tree must settle to Done, not stuck pending" - ); -} - -/// #1: the cascade must also delete the summary's **on-disk content file**, not -/// just the row — otherwise a `clear_memory` delete leaves the summarised text -/// orphaned on disk. -#[test] -fn clear_memory_delete_removes_orphaned_summary_content_file() { - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind, TreeStatus}; - - let (_tmp, cfg) = test_config(); - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - - let mut c = sample_chunk("gmail:acct", 0, 1_700_000_000_000); - c.metadata.source_kind = SourceKind::Email; - c.metadata.owner = "gmail-sync:conn-1".to_string(); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - - tree_store::insert_tree( - &cfg, - &Tree { - id: "tree-acct".into(), - kind: TreeKind::Source, - scope: "gmail:acct".into(), - root_id: None, - max_level: 1, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }, - ) - .unwrap(); - - // A real on-disk summary content file under the memory tree content root. - let rel = "summaries/gmail_acct/L1/sum-acct.md"; - let abs = cfg.memory_tree_content_root().join(rel); - std::fs::create_dir_all(abs.parent().unwrap()).unwrap(); - std::fs::write(&abs, "summarised email body").unwrap(); - assert!(abs.exists()); - - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - tree_store::insert_summary_tx( - &tx, - &SummaryNode { - id: "sum-acct".into(), - tree_id: "tree-acct".into(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: vec![c.id.clone()], - content: "preview".into(), - token_count: 3, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }, - None, - "test/model@3", - )?; - tx.execute( - "UPDATE mem_tree_summaries SET content_path = ?1 WHERE id = 'sum-acct'", - params![rel], - )?; - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Email, - "gmail:acct", - 1_700_000_000_000 - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - delete_chunks_by_owner(&cfg, SourceKind::Email, "gmail-sync:conn-1").unwrap(); - - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, "gmail:acct") - .unwrap() - .is_none() - ); - assert!( - !abs.exists(), - "orphaned summary content file must be removed from disk" - ); -} - -/// #2: the safety property — deleting one connection's chunks must NOT delete -/// the source tree while ANOTHER connection still owns chunks for the same -/// account (source not yet orphaned). -#[test] -fn clear_memory_delete_keeps_tree_when_another_connection_still_owns_chunks() { - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::{Buffer, Tree, TreeKind, TreeStatus}; - - let (_tmp, cfg) = test_config(); - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - - // Same account `gmail:acct`, two connections (owners). - let mut a = sample_chunk("gmail:acct", 0, 1_700_000_000_000); - a.metadata.source_kind = SourceKind::Email; - a.metadata.owner = "gmail-sync:conn-1".to_string(); - let mut b = sample_chunk("gmail:acct", 1, 1_700_000_001_000); - b.metadata.source_kind = SourceKind::Email; - b.metadata.owner = "gmail-sync:conn-2".to_string(); - upsert_chunks(&cfg, &[a.clone(), b.clone()]).unwrap(); - - tree_store::insert_tree( - &cfg, - &Tree { - id: "tree-acct".into(), - kind: TreeKind::Source, - scope: "gmail:acct".into(), - root_id: None, - max_level: 1, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }, - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - tree_store::upsert_buffer_tx( - &tx, - &Buffer { - tree_id: "tree-acct".into(), - level: 0, - item_ids: vec![a.id.clone(), b.id.clone()], - token_sum: 24, - oldest_at: Some(ts), - }, - )?; - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Email, - "gmail:acct", - 1_700_000_000_000 - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - // Disconnect ONLY conn-1. - let deleted = delete_chunks_by_owner(&cfg, SourceKind::Email, "gmail-sync:conn-1").unwrap(); - assert_eq!(deleted, 1); - - // conn-1's chunk is gone, conn-2's remains → source still has chunks → - // the tree (and its buffer + ingest gate) MUST survive. - assert!(get_chunk(&cfg, &a.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &b.id).unwrap().is_some()); - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, "gmail:acct") - .unwrap() - .is_some(), - "tree must survive while another connection still owns chunks" - ); - assert!(is_source_ingested(&cfg, SourceKind::Email, "gmail:acct").unwrap()); - with_connection(&cfg, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_buffers", [], |r| r.get(0))?; - assert_eq!(n, 1); - Ok(()) - }) - .unwrap(); -} - -/// #3: queued `Extract` / `AppendBuffer` jobs that reference a chunk deleted -/// out from under them settle to `Done` (warn-and-skip), not stuck pending. -#[tokio::test] -async fn queued_jobs_for_deleted_chunk_settle_to_done() { - use crate::openhuman::memory_queue::{store as queue_store, types as queue_types}; - - let (_tmp, cfg) = test_config(); - let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:#eng").unwrap(); - assert!(get_chunk(&cfg, &c.id).unwrap().is_none()); - - let extract_id = queue_store::enqueue( - &cfg, - &queue_types::NewJob::extract_chunk(&queue_types::ExtractChunkPayload { - chunk_id: c.id.clone(), - }) - .unwrap(), - ) - .unwrap() - .expect("extract_chunk job enqueued"); - let append_id = queue_store::enqueue( - &cfg, - &queue_types::NewJob::append_buffer(&queue_types::AppendBufferPayload { - node: queue_types::NodeRef::Leaf { - chunk_id: c.id.clone(), - }, - target: queue_types::AppendTarget::Source { - source_id: "slack:#eng".into(), - }, - }) - .unwrap(), - ) - .unwrap() - .expect("append_buffer job enqueued"); - - // Drive both queued jobs through the flipped crate engine (via the host - // `drain_until_idle` → `worker::run_once` seam). Each must settle to Done - // (warn-and-skip over the deleted chunk), not stay stuck pending. - crate::openhuman::memory_queue::drain_until_idle(&cfg) - .await - .expect("drain queue"); - for id in [extract_id, append_id] { - assert_eq!( - queue_store::get_job(&cfg, &id).unwrap().unwrap().status, - queue_types::JobStatus::Done, - "job {id} over a deleted chunk must settle to Done" - ); - } -} - -#[test] -fn delete_chunks_by_owner_preserves_other_owners_for_same_source() { - let (_tmp, cfg) = test_config(); - let mut target = sample_chunk("slack:shared", 0, 1_700_000_000_000); - target.metadata.owner = "slack-sync:c-1".to_string(); - let mut same_source_other_owner = sample_chunk("slack:shared", 1, 1_700_000_001_000); - same_source_other_owner.metadata.owner = "slack-sync:c-2".to_string(); - let mut target_other_source = sample_chunk("slack:c-1-only", 0, 1_700_000_002_000); - target_other_source.metadata.owner = "slack-sync:c-1".to_string(); - upsert_chunks( - &cfg, - &[ - target.clone(), - same_source_other_owner.clone(), - target_other_source.clone(), - ], - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Chat, - "slack:shared", - 1_700_000_000_000 - )?); - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Chat, - "slack:c-1-only", - 1_700_000_000_000 - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let deleted = delete_chunks_by_owner(&cfg, SourceKind::Chat, "slack-sync:c-1").unwrap(); - - assert_eq!(deleted, 2); - assert!(get_chunk(&cfg, &target.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &target_other_source.id).unwrap().is_none()); - assert!(get_chunk(&cfg, &same_source_other_owner.id) - .unwrap() - .is_some()); - assert!(is_source_ingested(&cfg, SourceKind::Chat, "slack:shared").unwrap()); - assert!(!is_source_ingested(&cfg, SourceKind::Chat, "slack:c-1-only").unwrap()); -} - -#[test] -fn delete_chunks_by_source_removes_safe_content_files_but_rejects_escape_paths() { - let (_tmp, cfg) = test_config(); - let safe = sample_chunk("slack:c-1", 0, 1_700_000_000_000); - let unsafe_chunk = sample_chunk("slack:c-1", 1, 1_700_000_001_000); - upsert_chunks(&cfg, &[safe.clone(), unsafe_chunk.clone()]).unwrap(); - - let content_root = cfg.memory_tree_content_root(); - let safe_rel = "chunks/safe.md"; - let safe_path = content_root.join(safe_rel); - std::fs::create_dir_all(safe_path.parent().unwrap()).unwrap(); - std::fs::write(&safe_path, "safe").unwrap(); - - let outside_path = content_root.parent().unwrap().join("outside.md"); - std::fs::write(&outside_path, "outside").unwrap(); - - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params![safe_rel, safe.id], - )?; - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params!["../outside.md", unsafe_chunk.id], - )?; - Ok(()) - }) - .unwrap(); - - let deleted = delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:c-1").unwrap(); - - assert_eq!(deleted, 2); - assert!(!safe_path.exists()); - assert!(outside_path.exists()); -} - -#[cfg(unix)] -#[test] -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(); - - let content_root = cfg.memory_tree_content_root(); - let target_path = content_root.join("chunks/target.md"); - let link_rel = "chunks/link.md"; - let link_path = content_root.join(link_rel); - std::fs::create_dir_all(target_path.parent().unwrap()).unwrap(); - std::fs::write(&target_path, "target").unwrap(); - std::os::unix::fs::symlink("target.md", &link_path).unwrap(); - - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params![link_rel, linked_chunk.id], - )?; - Ok(()) - }) - .unwrap(); - - let deleted = delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:c-1").unwrap(); - - assert_eq!(deleted, 1); - assert!(target_path.exists()); - assert!(!link_path.exists()); -} - -#[test] -fn missing_chunk_returns_none() { - let (_tmp, cfg) = test_config(); - assert!(get_chunk(&cfg, "nonexistent").unwrap().is_none()); -} - -#[test] -fn empty_batch_is_noop() { - let (_tmp, cfg) = test_config(); - assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0); - assert_eq!(count_chunks(&cfg).unwrap(), 0); -} - -#[test] -fn schema_has_content_path_and_content_sha256_columns() { - // Verify that with_connection applies additive migrations for content - // pointers and source grouping scope on a fresh DB. - let (_tmp, cfg) = test_config(); - with_connection(&cfg, |conn| { - let mut has_path_scope = false; - let mut has_content_path = false; - let mut has_content_sha256 = false; - let mut stmt = conn.prepare("PRAGMA table_info(mem_tree_chunks)")?; - let names: Vec = stmt - .query_map(params![], |row| row.get::<_, String>(1))? - .filter_map(|r| r.ok()) - .collect(); - for name in &names { - if name == "path_scope" { - has_path_scope = true; - } - if name == "content_path" { - has_content_path = true; - } - if name == "content_sha256" { - has_content_sha256 = true; - } - } - assert!( - has_path_scope, - "mem_tree_chunks must have path_scope column after migration; found: {names:?}" - ); - assert!( - has_content_path, - "mem_tree_chunks must have content_path column after migration; found: {names:?}" - ); - assert!( - has_content_sha256, - "mem_tree_chunks must have content_sha256 column after migration; found: {names:?}" - ); - Ok(()) - }) - .unwrap(); -} - -/// Directly pins the `is_transient_cold_start` classifier — the -/// gatekeeper for the retry loop in `open_and_init_with_retry`. The -/// concurrent-init test above only exercises it indirectly (and only -/// if a transient happens to fire on the dev box). A targeted test -/// catches regressions if the match arms are edited. -#[test] -fn is_transient_cold_start_classifies_known_extended_codes() { - use rusqlite::ffi; - use rusqlite::ErrorCode; - - // The WAL/SHM cold-start codes that fire under contention. All must - // classify as transient → retried. (4618 SHMOPEN is the macOS failure; - // 5386 is the real SHMMAP; 4874 is SHMSIZE — all of the `-shm` family.) - for extended in [ - 14, // CANTOPEN - 1546, // IOERR_TRUNCATE - 4618, // IOERR_SHMOPEN - 4874, // IOERR_SHMSIZE - 5386, // IOERR_SHMMAP - 8714, // IOERR_IN_PAGE - ] { - let err = anyhow::Error::from(rusqlite::Error::SqliteFailure( - ffi::Error { - code: ErrorCode::SystemIoFailure, - extended_code: extended, - }, - None, - )); - assert!( - super::is_transient_cold_start(&err), - "extended_code {extended} must classify as transient cold-start" - ); - } - - // SQLITE_BUSY (extended code 5) is a real lock-contention signal, - // NOT a cold-start race — the caller handles it via `busy_timeout` - // not via this retry loop. Must NOT classify. - let busy = anyhow::Error::from(rusqlite::Error::SqliteFailure( - ffi::Error { - code: ErrorCode::DatabaseBusy, - extended_code: 5, - }, - None, - )); - assert!( - !super::is_transient_cold_start(&busy), - "DatabaseBusy must not be classified as cold-start transient" - ); - - // Non-SQLite error in the chain — must not classify. - let other: anyhow::Error = anyhow::anyhow!("not a sqlite error"); - assert!( - !super::is_transient_cold_start(&other), - "non-SQLite errors must not classify as transient cold-start" - ); -} - -/// Regression: `PRAGMA foreign_keys` is connection-local in SQLite and -/// must be re-set on every `Connection::open`. After the schema-init -/// refactor, the pragma moved out of `SCHEMA` (which only runs on -/// first init per path) into `open_connection`. Verify both the -/// cold-init path and the fast path return a connection with FK on. -#[test] -fn with_connection_keeps_foreign_keys_on_for_every_call() { - let (_tmp, cfg) = test_config(); - // First call — exercises apply_schema + open_connection. - let fk_on_first: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row("PRAGMA foreign_keys;", params![], |r| r.get::<_, i64>(0))?) - }) - .unwrap(); - assert_eq!( - fk_on_first, 1, - "foreign_keys must be ON on first connection" - ); - // Second call — fast path (schema init skipped); pragma must still be set. - let fk_on_second: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row("PRAGMA foreign_keys;", params![], |r| r.get::<_, i64>(0))?) - }) - .unwrap(); - assert_eq!( - fk_on_second, 1, - "foreign_keys must be ON on fast-path (post-init) connection" - ); -} - -// ── Connection cache tests (#2206) ─────────────────────────────────────────── - -/// Two `with_connection` calls for the same workspace must return the same -/// cached `Arc` (pointer identity proves no re-init happened). -#[test] -fn connection_cache_returns_same_arc_for_same_workspace() { - clear_connection_cache(); - let (_tmp, cfg) = test_config(); - - let arc1 = get_or_init_connection(&cfg).expect("first get_or_init"); - let arc2 = get_or_init_connection(&cfg).expect("second get_or_init"); - assert!( - Arc::ptr_eq(&arc1, &arc2), - "expected the same Arc from the connection cache on the second call" - ); -} - -/// Two configs pointing at different tempdirs must produce independent -/// connections (separate Arc pointers, no cross-contamination). -#[test] -fn connection_cache_uses_separate_connections_for_different_workspaces() { - clear_connection_cache(); - let (_tmp1, cfg1) = test_config(); - let (_tmp2, cfg2) = test_config(); - - let arc1 = get_or_init_connection(&cfg1).expect("workspace 1"); - let arc2 = get_or_init_connection(&cfg2).expect("workspace 2"); - assert!( - !Arc::ptr_eq(&arc1, &arc2), - "different workspaces must have independent connections" - ); - - // Sanity: each DB is usable independently. - let c = sample_chunk("s", 0, 1_700_000_000_000); - upsert_chunks(&cfg1, &[c.clone()]).unwrap(); - assert_eq!(count_chunks(&cfg1).unwrap(), 1); - assert_eq!(count_chunks(&cfg2).unwrap(), 0); -} - -/// Pointing the DB path at a *file* (not a directory) makes it impossible to -/// create the DB, so `get_or_init_connection` must fail. After -/// `CB_THRESHOLD` failures the circuit breaker trips and subsequent calls -/// return an error immediately without touching the filesystem. -#[test] -fn circuit_breaker_trips_after_threshold() { - clear_connection_cache(); - let tmp = TempDir::new().expect("tempdir"); - - // Create a regular file where the memory_tree *directory* would be — - // this prevents `create_dir_all` from succeeding. - let blocker = tmp.path().join(DB_DIR); - std::fs::write(&blocker, b"not a dir").expect("write blocker file"); - - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - - // First CB_THRESHOLD calls should all fail (can't create dir over a file). - for i in 0..CB_THRESHOLD { - let result = get_or_init_connection(&cfg); - assert!( - result.is_err(), - "call {i}: expected error before breaker trips" - ); - } - - // The CB_THRESHOLD+1'th call should be rejected immediately by the - // circuit breaker (error message contains "circuit breaker"). - let cb_err = get_or_init_connection(&cfg) - .expect_err("expected circuit breaker error on call after threshold"); - let msg = format!("{cb_err:#}").to_ascii_lowercase(); - assert!( - msg.contains("circuit breaker"), - "expected circuit breaker message, got: {msg}" - ); -} - -/// `try_cleanup_stale_files` removes `.db-shm` and `.db-wal` side-files that -/// exist alongside the main DB file. -#[test] -fn stale_shm_cleanup_removes_files() { - let tmp = TempDir::new().expect("tempdir"); - let db_path = tmp.path().join("chunks.db"); - - // Create the main DB file and the two stale side-files. - std::fs::write(&db_path, b"").expect("create db file"); - let shm = tmp.path().join("chunks.db-shm"); - let wal = tmp.path().join("chunks.db-wal"); - std::fs::write(&shm, b"stale shm").expect("create shm"); - std::fs::write(&wal, b"stale wal").expect("create wal"); - - assert!(shm.exists(), "shm must exist before cleanup"); - assert!(wal.exists(), "wal must exist before cleanup"); - - let cleaned = try_cleanup_stale_files(&db_path); - assert!( - cleaned, - "cleanup should return true when files were removed" - ); - assert!(!shm.exists(), "shm must be removed"); - assert!(!wal.exists(), "wal must be removed"); -} - -/// memory_tree must run the TRUNCATE rollback journal — never WAL. WAL's -/// `-shm`/`-wal` machinery is the source of the cold-start IOERR_SHMMAP / -/// IOERR_TRUNCATE failures (Sentry TAURI-RUST-EV / TAURI-RUST-X1), and the -/// single cached connection gains nothing from WAL's reader concurrency. -#[test] -fn memory_tree_uses_truncate_journal_not_wal() { - let (_tmp, cfg) = test_config(); - - with_connection(&cfg, |conn| { - let mode: String = conn.query_row("PRAGMA journal_mode", [], |r| r.get(0))?; - assert!( - mode.eq_ignore_ascii_case("truncate"), - "memory_tree journal_mode must be TRUNCATE, got '{mode}'" - ); - let sync: i64 = conn.query_row("PRAGMA synchronous", [], |r| r.get(0))?; - assert_eq!(sync, 2, "rollback journal requires synchronous=FULL (2)"); - Ok(()) - }) - .expect("with_connection"); - - // A `-shm` shared-memory side-file is only ever created under WAL. - let shm = cfg.workspace_dir.join("memory_tree").join("chunks.db-shm"); - assert!( - !shm.exists(), - "no -shm file must exist under TRUNCATE journal" - ); -} - -/// A database a prior (WAL-mode) release left behind must migrate cleanly to -/// TRUNCATE on the next open, with the `-wal`/`-shm` side-files gone. -#[test] -fn existing_wal_db_migrates_to_truncate() { - let (_tmp, cfg) = test_config(); - let db_path = cfg.workspace_dir.join("memory_tree").join("chunks.db"); - std::fs::create_dir_all(db_path.parent().unwrap()).expect("mkdir"); - - // Simulate the old release: open the DB in WAL mode and commit a row so - // the WAL marker is persisted in the database header. - { - let conn = rusqlite::Connection::open(&db_path).expect("open wal db"); - let mode: String = conn - .query_row("PRAGMA journal_mode=WAL", [], |r| r.get(0)) - .expect("set wal"); - assert!(mode.eq_ignore_ascii_case("wal"), "precondition: db in WAL"); - conn.execute_batch("CREATE TABLE legacy_marker(x); INSERT INTO legacy_marker VALUES (1);") - .expect("seed"); - } // connection dropped — the header still records WAL - - // Clear any cached connection for isolation, then open via with_connection. - clear_connection_cache(); - with_connection(&cfg, |conn| { - let mode: String = conn.query_row("PRAGMA journal_mode", [], |r| r.get(0))?; - assert!( - mode.eq_ignore_ascii_case("truncate"), - "WAL db must migrate to TRUNCATE on open, got '{mode}'" - ); - // Data written under WAL must survive the checkpoint-and-switch — the - // migration must not lose committed rows. - let marker: i64 = conn.query_row("SELECT x FROM legacy_marker", [], |r| r.get(0))?; - assert_eq!(marker, 1, "row committed under WAL must survive migration"); - Ok(()) - }) - .expect("with_connection migrates"); - - assert!( - !db_path.with_file_name("chunks.db-shm").exists(), - "-shm must be gone after WAL→TRUNCATE migration" - ); - assert!( - !db_path.with_file_name("chunks.db-wal").exists(), - "-wal must be gone after WAL→TRUNCATE migration" - ); -} - -#[test] -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(); - 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(); - clear_chunk_reembed_skipped(&cfg, &c.id, &sig).unwrap(); - let count: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row( - "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped - WHERE chunk_id = ?1 AND model_signature = ?2", - params![c.id, sig], - |r| r.get(0), - )?) - }) - .unwrap(); - assert_eq!(count, 0); -} - -#[test] -fn clear_reembed_skipped_for_signature_removes_all_tombstones_for_sig() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#a", 0, 1_700_000_000_000); - let c2 = sample_chunk("slack:#b", 1, 1_700_000_000_001); - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - let sig = tree_active_signature(&cfg); - let other_sig = "provider=other;model=x;dims=8"; - mark_chunk_reembed_skipped(&cfg, &c1.id, &sig, "r1").unwrap(); - mark_chunk_reembed_skipped(&cfg, &c2.id, &sig, "r2").unwrap(); - mark_chunk_reembed_skipped(&cfg, &c1.id, other_sig, "other").unwrap(); - let summary_id = "summary-bulk-clear-test"; - with_connection(&cfg, |conn| { - conn.execute( - "INSERT OR IGNORE INTO mem_tree_trees (id, kind, scope, created_at_ms) - VALUES ('tree-bulk-clear', 'source', 'bulk-clear', 0)", - [], - )?; - conn.execute( - "INSERT INTO mem_tree_summaries ( - id, tree_id, tree_kind, level, child_ids_json, content, token_count, - entities_json, topics_json, time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted - ) VALUES (?1, 'tree-bulk-clear', 'source', 0, '[]', 'x', 1, '[]', '[]', 0, 0, 0.0, 0, 0)", - params![summary_id], - )?; - Ok(()) - }) - .unwrap(); - crate::openhuman::memory_store::trees::store::mark_summary_reembed_skipped( - &cfg, - summary_id, - &sig, - "summary tombstone", - ) - .unwrap(); - - let deleted = clear_reembed_skipped_for_signature(&cfg, &sig).unwrap(); - assert_eq!(deleted, 3); - - let remaining_chunks: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row( - "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1", - params![sig], - |r| r.get(0), - )?) - }) - .unwrap(); - assert_eq!(remaining_chunks, 0); - - let remaining_summaries: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row( - "SELECT COUNT(*) FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1", - params![sig], - |r| r.get(0), - )?) - }) - .unwrap(); - assert_eq!(remaining_summaries, 0); - - let other_kept: i64 = with_connection(&cfg, |conn| { - Ok(conn.query_row( - "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped - WHERE chunk_id = ?1 AND model_signature = ?2", - params![c1.id, other_sig], - |r| r.get(0), - )?) - }) - .unwrap(); - assert_eq!(other_kept, 1); -} - -#[test] -fn validate_reembed_skip_key_rejects_empty_and_oversized() { - assert!(validate_reembed_skip_key("chunk_id", " ").is_err()); - let huge = "a".repeat(REEMBED_SKIP_KEY_MAX_LEN + 1); - assert!(validate_reembed_skip_key("chunk_id", &huge).is_err()); - assert!(validate_reembed_skip_key("chunk_id", "ok\0bad").is_err()); - assert_eq!( - validate_reembed_skip_key("chunk_id", " trimmed ").unwrap(), - "trimmed" - ); -} - -// ---------- get_chunks_batch ---------- -// -// Contract: equivalent to looping `get_chunk` per id but in -// `O(ceil(n / MAX_FETCH_BATCH))` SQLite round-trips. The map carries -// only ids that exist; missing ids are silently absent (same as the -// per-row helper returning Ok(None)). - -#[test] -fn get_chunks_batch_returns_present_ids_in_map() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - let c2 = sample_chunk("slack:#eng", 1, 1_700_000_000_000); - let c3 = sample_chunk("slack:#ops", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); - - let ids = vec![c1.id.clone(), c2.id.clone(), c3.id.clone()]; - let map = get_chunks_batch(&cfg, &ids).unwrap(); - assert_eq!(map.len(), 3); - assert_eq!(map.get(&c1.id), Some(&c1)); - assert_eq!(map.get(&c2.id), Some(&c2)); - assert_eq!(map.get(&c3.id), Some(&c3)); -} - -#[test] -fn get_chunks_batch_empty_input_and_missing_ids() { - // Empty input: empty map (no SQL issued). - let (_tmp, cfg) = test_config(); - let empty = get_chunks_batch(&cfg, &[]).unwrap(); - assert!(empty.is_empty()); - - // Missing ids: silently absent (mirrors per-row Ok(None)). - // `fetch_leaves` relies on this so partial-result detection - // (`hits.len() < ids.len()`) keeps working unchanged. - let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - let ids = vec![ - c.id.clone(), - "ghost:no-such-1".into(), - "ghost:no-such-2".into(), - ]; - 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()); -} - -// ---------- get_chunk_embeddings_for_signature_batch ---------- -// -// Contract: equivalent to looping `get_chunk_embedding_for_signature` -// per id, but in O(ceil(n / MAX_EMBEDDING_BATCH)) round-trips instead -// of O(n). The map contains only ids that have a vector under the -// requested signature; absent rows are silently dropped (same as the -// per-row helper returning Ok(None)). - -#[test] -fn batch_embedding_lookup_returns_only_signature_scoped_rows() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - let c2 = sample_chunk("slack:#eng", 1, 1_700_000_000_000); - let c3 = sample_chunk("slack:#eng", 2, 1_700_000_000_000); - upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); - - let sig_a = "openai/text-embedding-3-small@1536"; - let sig_b = "local/bge-small@384"; - set_chunk_embedding_for_signature(&cfg, &c1.id, sig_a, &[0.1, 0.2]).unwrap(); - set_chunk_embedding_for_signature(&cfg, &c2.id, sig_a, &[0.3, 0.4]).unwrap(); - set_chunk_embedding_for_signature(&cfg, &c3.id, sig_b, &[0.5, 0.6, 0.7]).unwrap(); - - let ids = vec![c1.id.clone(), c2.id.clone(), c3.id.clone()]; - let map_a = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); - 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"); - - let map_b = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); - assert_eq!(map_b.len(), 1); - assert_eq!(map_b.get(&c3.id).cloned(), Some(vec![0.5, 0.6, 0.7])); -} - -#[test] -fn batch_embedding_lookup_empty_input_returns_empty_map() { - let (_tmp, cfg) = test_config(); - let map = get_chunk_embeddings_for_signature_batch(&cfg, &[], "any/sig@1").unwrap(); - assert!(map.is_empty()); -} - -#[test] -fn batch_embedding_lookup_unknown_ids_absent_from_map() { - // Pre-batch contract: per-row helper returned Ok(None) for missing - // chunks. Batch helper must mirror that — missing ids absent from - // the map, present ids carry their vector. The retrieval rerank - // path depends on this so absent rows get the - // (NEG_INFINITY, false) sink-to-bottom treatment. - let (_tmp, cfg) = test_config(); - let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - let sig = "openai/text-embedding-3-small@1536"; - set_chunk_embedding_for_signature(&cfg, &c.id, sig, &[0.1]).unwrap(); - - let ids = vec![ - c.id.clone(), - "ghost:no-such-chunk-1".into(), - "ghost:no-such-chunk-2".into(), - ]; - let map = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig).unwrap(); - assert_eq!(map.len(), 1); - assert_eq!(map.get(&c.id).cloned(), Some(vec![0.1])); -} - -#[test] -fn batch_embedding_lookup_splits_id_list_above_per_batch_threshold() { - // Validates the `chunks(MAX_EMBEDDING_BATCH)` window loop in - // `get_chunk_embeddings_for_signature_batch`. We pass > 500 ids in - // one call; the helper must internally split them into multiple - // `IN (...)` queries and merge results into a single map. 3 of the - // 501 ids actually carry embeddings; the other 498 are unknown - // strings and must be absent from the returned map (no error). - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#a", 0, 1_700_000_000_000); - let c2 = sample_chunk("slack:#b", 0, 1_700_000_000_000); - let c3 = sample_chunk("slack:#c", 0, 1_700_000_000_000); - upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); - let sig = "openai/text-embedding-3-small@1536"; - set_chunk_embedding_for_signature(&cfg, &c1.id, sig, &[1.0]).unwrap(); - set_chunk_embedding_for_signature(&cfg, &c2.id, sig, &[2.0]).unwrap(); - set_chunk_embedding_for_signature(&cfg, &c3.id, sig, &[3.0]).unwrap(); - - // Build 501 ids: 3 real + 498 ghosts. The 501-element vec crosses - // the 500-per-batch boundary, forcing two `IN (...)` queries. - let mut ids: Vec = (0..498).map(|i| format!("ghost:{i}")).collect(); - ids.push(c1.id.clone()); - ids.push(c2.id.clone()); - ids.push(c3.id.clone()); - assert_eq!(ids.len(), 501); - - let map = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig).unwrap(); - assert_eq!(map.len(), 3, "only the 3 real ids should be present"); - assert_eq!(map.get(&c1.id).cloned(), Some(vec![1.0])); - assert_eq!(map.get(&c2.id).cloned(), Some(vec![2.0])); - assert_eq!(map.get(&c3.id).cloned(), Some(vec![3.0])); -} - -/// The one-shot purge migration deletes global + topic trees (rows, summaries, -/// buffers, jobs, and on-disk summary folders) while leaving source trees and -/// non-retired jobs untouched, and runs exactly once (PRAGMA user_version gate). -#[test] -fn global_topic_purge_removes_only_global_and_topic() { - let (_tmp, cfg) = test_config(); - // First open initialises the schema and runs both migrations (sets - // user_version = 2). - upsert_chunks(&cfg, &[sample_chunk("slack:#eng", 0, 1_700_000_000_000)]).unwrap(); - - // On-disk: a legacy per-day global folder, the singleton global folder, a - // topic folder, and a source folder that must survive. - let summaries = cfg - .memory_tree_content_root() - .join("wiki") - .join("summaries"); - for d in [ - "global-2026-05-28", - "global", - "topic-alice", - "source-slack-eng", - ] { - std::fs::create_dir_all(summaries.join(d).join("L0")).unwrap(); - } - - with_connection(&cfg, |conn| { - // Seed one tree of each kind, each with a summary. - for (id, kind) in [ - ("source:s1", "source"), - ("global:g1", "global"), - ("topic:t1", "topic"), - ] { - conn.execute( - "INSERT INTO mem_tree_trees (id, kind, scope, max_level, status, created_at_ms) \ - VALUES (?1, ?2, ?2, 0, 'active', 0)", - params![id, kind], - )?; - conn.execute( - "INSERT INTO mem_tree_summaries \ - (id, tree_id, tree_kind, level, content, token_count, \ - time_range_start_ms, time_range_end_ms, sealed_at_ms) \ - VALUES (?1, ?2, ?3, 0, 'x', 1, 0, 0, 0)", - params![format!("sum-{id}"), id, kind], - )?; - } - // Seed retired + surviving job rows. - for (jid, kind) in [ - ("j1", "topic_route"), - ("j2", "digest_daily"), - ("j3", "extract_chunk"), - ] { - conn.execute( - "INSERT INTO mem_tree_jobs (id, kind, payload_json, available_at_ms, created_at_ms) \ - VALUES (?1, ?2, '{}', 0, 0)", - params![jid, kind], - )?; - } - // Re-arm the gate so the purge runs against the seeded rows. - conn.pragma_update(None, "user_version", 1i64)?; - super::purge_global_topic_trees(conn, &cfg)?; - - // Trees: only the source tree survives. - let trees: i64 = - conn.query_row("SELECT COUNT(*) FROM mem_tree_trees", [], |r| r.get(0))?; - assert_eq!(trees, 1, "only the source tree should remain"); - let kind: String = - conn.query_row("SELECT kind FROM mem_tree_trees", [], |r| r.get(0))?; - assert_eq!(kind, "source"); - - // Summaries: only the source summary survives. - let summaries_left: i64 = - conn.query_row("SELECT COUNT(*) FROM mem_tree_summaries", [], |r| r.get(0))?; - assert_eq!(summaries_left, 1); - - // Jobs: retired kinds gone, extract_chunk kept. - let jobs_left: Vec = { - let mut stmt = conn.prepare("SELECT kind FROM mem_tree_jobs ORDER BY kind")?; - let rows = stmt.query_map([], |r| r.get::<_, String>(0))?; - rows.collect::>()? - }; - assert_eq!(jobs_left, vec!["extract_chunk".to_string()]); - - // Gate advanced — a second run is a no-op. - let version: i64 = conn.query_row("PRAGMA user_version", [], |r| r.get(0))?; - assert_eq!(version, 2); - Ok(()) - }) - .unwrap(); - - // On-disk: global*/topic-* folders gone, source-* kept. - assert!(!summaries.join("global-2026-05-28").exists()); - assert!(!summaries.join("global").exists()); - assert!(!summaries.join("topic-alice").exists()); - assert!( - summaries.join("source-slack-eng").exists(), - "source summary folder must survive the purge" - ); -} - -// ── extraction_coverage (#002 FR-010 / US5) ────────────────────────────── - -#[test] -fn extraction_coverage_empty_store_is_zero() { - let (_tmp, cfg) = test_config(); - assert_eq!(extraction_coverage(&cfg).unwrap(), 0.0); -} - -#[test] -fn extraction_coverage_reflects_indexed_fraction() { - let (_tmp, cfg) = test_config(); - // Two chunks; index an entity for only the first → coverage 0.5. - let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); - let c2 = sample_chunk("slack:#eng", 1, 1_700_000_001_000); - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - - with_connection(&cfg, |conn| { - conn.execute( - "INSERT INTO mem_tree_entity_index - (entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms) - VALUES (?1, ?2, 'leaf', 'person', 'Alice', 0.9, 1)", - params!["person:Alice", c1.id], - )?; - Ok(()) - }) - .unwrap(); - - let cov = extraction_coverage(&cfg).unwrap(); - assert!((cov - 0.5).abs() < 1e-6, "expected 0.5, got {cov}"); - - // Index the second chunk too → full coverage. - with_connection(&cfg, |conn| { - conn.execute( - "INSERT INTO mem_tree_entity_index - (entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms) - VALUES (?1, ?2, 'leaf', 'person', 'Bob', 0.9, 2)", - params!["person:Bob", c2.id], - )?; - Ok(()) - }) - .unwrap(); - assert!((extraction_coverage(&cfg).unwrap() - 1.0).abs() < 1e-6); -} - -// ── memory_tree_delete_source RPC ──────────────────────────────────────────── -// These prove the new `delete_source_rpc` is a FULL source-level delete (not a -// chunk-only delete): it must cascade through every dependent table, the ingest -// gate, and the source summary tree, remove content files, and leave stale -// summaries unable to resurface in recall — while sibling sources are untouched. - -/// Seed a fully-formed Document source (chunks + side rows + content files + -/// source tree + summary + sidecars + buffer + ingest gate) for `source_id`. -/// Returns the chunk ids created. Used by the delete_source tests below. -#[cfg(test)] -async fn seed_full_document_source( - cfg: &Config, - source_id: &str, - tree_id: &str, - summary_id: &str, - base_ts_ms: i64, -) -> (Vec, String, String) { - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::{ - Buffer, SummaryNode, Tree, TreeKind, TreeStatus, - }; - - let ts = Utc.timestamp_millis_opt(base_ts_ms).unwrap(); - let mk_doc = |seq: u32, ts_ms: i64| { - let mut c = sample_chunk(source_id, seq, ts_ms); - c.metadata.source_kind = SourceKind::Document; - c - }; - let c0 = mk_doc(0, base_ts_ms); - let c1 = mk_doc(1, base_ts_ms + 1000); - upsert_chunks(cfg, &[c0.clone(), c1.clone()]).unwrap(); - - // Real on-disk chunk + summary content files under the content root. - let content_root = cfg.memory_tree_content_root(); - let chunk_rel = format!("document/{tree_id}/c0.md"); - let summary_rel = format!("summaries/{tree_id}/L1/{summary_id}.md"); - for rel in [&chunk_rel, &summary_rel] { - let abs = content_root.join(rel); - std::fs::create_dir_all(abs.parent().unwrap()).unwrap(); - std::fs::write(&abs, "body").unwrap(); - } - - let mk_summary = SummaryNode { - id: summary_id.into(), - tree_id: tree_id.into(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: vec![c0.id.clone(), c1.id.clone()], - content: format!("summary text for {source_id}"), - token_count: 3, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - tree_store::insert_tree( - cfg, - &Tree { - id: tree_id.into(), - kind: TreeKind::Source, - scope: source_id.into(), - root_id: None, - max_level: 1, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }, - ) - .unwrap(); - - with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - for chunk in [&c0, &c1] { - tx.execute( - "INSERT INTO mem_tree_score ( - chunk_id, total, token_count_signal, unique_words_signal, - metadata_weight, source_weight, interaction_weight, - entity_density, dropped, reason, computed_at_ms - ) VALUES (?1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0, NULL, 1700000000000)", - params![chunk.id], - )?; - tx.execute( - "INSERT INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms - ) VALUES (?1, ?2, 'chunk', 'person', 'doc', 0.9, 1700000000000)", - params![format!("entity:{}", chunk.id), chunk.id], - )?; - tx.execute( - "INSERT INTO mem_tree_chunk_embeddings ( - chunk_id, model_signature, vector, dim, created_at - ) VALUES (?1, 'test/model@3', ?2, 3, 1700000000.0)", - params![chunk.id, vec![1_u8, 2, 3]], - )?; - tx.execute( - "INSERT INTO mem_tree_chunk_reembed_skipped ( - chunk_id, model_signature, reason, skipped_at_ms - ) VALUES (?1, 'test/model@3', 'terminal', 1700000000000)", - params![chunk.id], - )?; - } - // point chunk c0 at its on-disk content file. - tx.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params![chunk_rel, c0.id], - )?; - - tree_store::insert_summary_tx(&tx, &mk_summary, None, "test/model@3")?; - tx.execute( - "UPDATE mem_tree_summaries SET content_path = ?1 WHERE id = ?2", - params![summary_rel, summary_id], - )?; - tx.execute( - "INSERT INTO mem_tree_summary_embeddings ( - summary_id, model_signature, vector, dim, created_at - ) VALUES (?1, 'test/model@3', ?2, 3, 1700000000.0)", - params![summary_id, vec![1_u8, 2, 3]], - )?; - tx.execute( - "INSERT INTO mem_tree_summary_reembed_skipped ( - summary_id, model_signature, reason, skipped_at_ms - ) VALUES (?1, 'test/model@3', 'terminal', 1700000000000)", - params![summary_id], - )?; - tx.execute( - "INSERT INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, 'summary', 'person', 'doc', 0.9, 1700000000000, ?3, 0)", - params![format!("entity:{summary_id}"), summary_id, tree_id], - )?; - tree_store::upsert_buffer_tx( - &tx, - &Buffer { - tree_id: tree_id.into(), - level: 0, - item_ids: vec![c0.id.clone(), c1.id.clone()], - token_sum: 24, - oldest_at: Some(ts), - }, - )?; - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Document, - source_id, - base_ts_ms - )?); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - (vec![c0.id.clone(), c1.id], chunk_rel, summary_rel) -} - -#[tokio::test] -async fn delete_source_rpc_purges_document_source_fully() { - use crate::openhuman::memory::read_rpc::{delete_source_rpc, list_chunks_rpc, recall_rpc}; - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::TreeKind; - - let (_tmp, cfg) = test_config(); - let target = "telegram-note-A"; - let sibling = "telegram-note-B"; - let (target_ids, target_chunk_file, target_summary_file) = - seed_full_document_source(&cfg, target, "tree-A", "sum-A", 1_700_000_000_000).await; - let (sibling_ids, sibling_chunk_file, sibling_summary_file) = - seed_full_document_source(&cfg, sibling, "tree-B", "sum-B", 1_700_000_100_000).await; - - let content_root = cfg.memory_tree_content_root(); - // Pre-conditions: both sources fully present on disk + in DB. - assert!(content_root.join(&target_chunk_file).exists()); - assert!(content_root.join(&target_summary_file).exists()); - - // ---- act ---- - let out = delete_source_rpc(&cfg, target.to_string()) - .await - .expect("delete_source ok") - .value; - assert!(out.deleted); - assert_eq!(out.chunks_removed, 2); - - // JSON response shape: { deleted: bool, chunks_removed: u64 }. - let json = serde_json::to_value(&out).unwrap(); - assert_eq!(json.get("deleted").and_then(|v| v.as_bool()), Some(true)); - assert_eq!(json.get("chunks_removed").and_then(|v| v.as_u64()), Some(2)); - - // 1. target chunks gone; sibling chunk survives. - for id in &target_ids { - assert!(get_chunk(&cfg, id).unwrap().is_none(), "chunk {id} remains"); - } - for id in &sibling_ids { - assert!(get_chunk(&cfg, id).unwrap().is_some(), "sibling {id} gone"); - } - - // 2–11. every dependent table for the target is empty; sibling rows remain. - with_connection(&cfg, |conn| { - let count = |sql: &str, p: &str| -> rusqlite::Result { - conn.query_row(sql, params![p], |r| r.get(0)) - }; - // chunk side rows keyed by the target chunk ids - for id in &target_ids { - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_score WHERE chunk_id = ?1", - id - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_entity_index WHERE node_id = ?1", - id - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_chunk_embeddings WHERE chunk_id = ?1", - id - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped WHERE chunk_id = ?1", - id - )?, - 0 - ); - } - // source tree rows (scope/tree-id == tree-A) gone - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_summaries WHERE tree_id = ?1", - "tree-A" - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_summary_embeddings WHERE summary_id = ?1", - "sum-A" - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_summary_reembed_skipped WHERE summary_id = ?1", - "sum-A" - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_entity_index WHERE tree_id = ?1", - "tree-A" - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_buffers WHERE tree_id = ?1", - "tree-A" - )?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_trees WHERE id = ?1", - "tree-A" - )?, - 0 - ); - // sibling tree intact - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_summaries WHERE tree_id = ?1", - "tree-B" - )?, - 1 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_trees WHERE id = ?1", - "tree-B" - )?, - 1 - ); - Ok(()) - }) - .unwrap(); - - // 6. ingest dedup gate cleared for target, retained for sibling. - assert!(!is_source_ingested(&cfg, SourceKind::Document, target).unwrap()); - assert!(is_source_ingested(&cfg, SourceKind::Document, sibling).unwrap()); - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, target) - .unwrap() - .is_none() - ); - - // 12–13. target content files removed; sibling content files remain. - assert!(!content_root.join(&target_chunk_file).exists()); - assert!(!content_root.join(&target_summary_file).exists()); - assert!(content_root.join(&sibling_chunk_file).exists()); - assert!(content_root.join(&sibling_summary_file).exists()); - - // 14. recall no longer surfaces the deleted source (no summary/chunk left to - // rank). Tolerant of minimal-config recall backends: if it returns, none of - // the hits may belong to the deleted source. - if let Ok(rc) = recall_rpc(&cfg, "summary text".to_string(), 10).await { - assert!( - rc.value.chunks.iter().all(|c| c.source_id != target), - "deleted source must not resurface in recall" - ); - } - - // 16. second delete is idempotent. - let again = delete_source_rpc(&cfg, target.to_string()) - .await - .unwrap() - .value; - assert!(!again.deleted); - assert_eq!(again.chunks_removed, 0); - - // 15. re-ingesting the same source_id works again (gate cleared) and writes chunks. - let mut fresh = sample_chunk(target, 0, 1_700_000_500_000); - fresh.metadata.source_kind = SourceKind::Document; - assert_eq!(upsert_chunks(&cfg, &[fresh.clone()]).unwrap(), 1); - let listed = list_chunks_rpc(&cfg, Default::default()) - .await - .unwrap() - .value; - assert!(listed.chunks.iter().any(|c| c.source_id == target)); - // The dedup gate was cleared by delete, so re-ingest can claim it again. - // Commit the claim (a rolled-back tx would prove nothing) and verify it - // actually persisted. - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - assert!( - claim_source_ingest_tx(&tx, SourceKind::Document, target, 1_700_000_500_000)?, - "ingest gate must be re-claimable after delete_source cleared it" - ); - tx.commit()?; - Ok(()) - }) - .unwrap(); - assert!( - is_source_ingested(&cfg, SourceKind::Document, target).unwrap(), - "re-claimed ingest gate must persist" - ); -} - -/// Versioned document sources store the ingest gate as `{source_id}@{version_ms}` -/// in addition to (or instead of) the bare id. `delete_source` must clear both. -#[tokio::test] -async fn delete_source_rpc_clears_versioned_ingest_gates() { - use crate::openhuman::memory::read_rpc::delete_source_rpc; - - let (_tmp, cfg) = test_config(); - let sid = "notion:conn-1:page-xyz"; - let versioned = format!("{sid}@1700000000000"); - - let mut c = sample_chunk(sid, 0, 1_700_000_000_000); - c.metadata.source_kind = SourceKind::Document; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - - // Seed both a bare gate and a versioned gate for the source. - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - assert!(claim_source_ingest_tx( - &tx, - SourceKind::Document, - sid, - 1_700_000_000_000 - )?); - tx.execute( - "INSERT INTO mem_tree_ingested_sources (source_kind, source_id, ingested_at_ms) - VALUES ('document', ?1, 1700000000000)", - params![versioned], - )?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let gate_count = |conn: &rusqlite::Connection| -> rusqlite::Result { - conn.query_row( - "SELECT COUNT(*) FROM mem_tree_ingested_sources - WHERE source_kind = 'document' AND (source_id = ?1 OR source_id LIKE ?2)", - params![sid, format!("{sid}@%")], - |r| r.get(0), - ) - }; - with_connection(&cfg, |conn| { - assert_eq!(gate_count(conn)?, 2, "both gates seeded"); - Ok(()) - }) - .unwrap(); - - let out = delete_source_rpc(&cfg, sid.to_string()) - .await - .unwrap() - .value; - assert!(out.deleted); - assert_eq!(out.chunks_removed, 1); - - assert!(!is_source_ingested(&cfg, SourceKind::Document, sid).unwrap()); - with_connection(&cfg, |conn| { - assert_eq!( - gate_count(conn)?, - 0, - "bare AND versioned ingest gates must be cleared" - ); - Ok(()) - }) - .unwrap(); -} - -#[tokio::test] -async fn delete_source_rpc_unknown_id_is_idempotent() { - use crate::openhuman::memory::read_rpc::delete_source_rpc; - let (_tmp, cfg) = test_config(); - let out = delete_source_rpc(&cfg, "telegram-note-does-not-exist".to_string()) - .await - .unwrap() - .value; - assert!(!out.deleted); - assert_eq!(out.chunks_removed, 0); -} - -#[tokio::test] -async fn delete_source_rpc_rejects_empty_source_id() { - use crate::openhuman::memory::read_rpc::delete_source_rpc; - let (_tmp, cfg) = test_config(); - assert!(delete_source_rpc(&cfg, " ".to_string()).await.is_err()); -} - -/// Legacy partial delete: chunks were already removed earlier (e.g. by the bot's -/// old per-chunk `delete_chunk` loop), leaving an orphaned summary tree + dedup -/// gate. `delete_source_rpc` must finish the job and remove the stale tree. -#[tokio::test] -async fn delete_source_rpc_cleans_legacy_partial_delete() { - use crate::openhuman::memory::read_rpc::{delete_chunk_rpc, delete_source_rpc}; - use crate::openhuman::memory_store::trees::store as tree_store; - use crate::openhuman::memory_store::trees::types::TreeKind; - - let (_tmp, cfg) = test_config(); - let target = "telegram-event-legacy"; - let (chunk_ids, _chunk_file, summary_file) = - seed_full_document_source(&cfg, target, "tree-legacy", "sum-legacy", 1_700_000_000_000) - .await; - - // ---- simulate the OLD per-chunk delete loop: remove only the chunks ---- - for id in &chunk_ids { - assert!( - delete_chunk_rpc(&cfg, id.clone()) - .await - .unwrap() - .value - .deleted - ); - } - - // pre-condition: chunks gone, but the summary tree + gate are still stale. - for id in &chunk_ids { - assert!(get_chunk(&cfg, id).unwrap().is_none()); - } - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, target) - .unwrap() - .is_some() - ); - assert!(is_source_ingested(&cfg, SourceKind::Document, target).unwrap()); - with_connection(&cfg, |conn| { - let n: i64 = conn.query_row( - "SELECT COUNT(*) FROM mem_tree_summaries WHERE tree_id = 'tree-legacy'", - [], - |r| r.get(0), - )?; - assert_eq!(n, 1, "stale summary must exist before delete_source"); - Ok(()) - }) - .unwrap(); - - // ---- act: delete_source must finish the legacy cleanup ---- - let out = delete_source_rpc(&cfg, target.to_string()) - .await - .unwrap() - .value; - // chunks were already gone, but a stale tree was cleaned → deleted=true. - assert!(out.deleted); - assert_eq!(out.chunks_removed, 0); - - // ---- assert: the stale tree / summaries / sidecars / gate are now gone ---- - assert!( - tree_store::get_tree_by_scope(&cfg, TreeKind::Source, target) - .unwrap() - .is_none() - ); - assert!(!is_source_ingested(&cfg, SourceKind::Document, target).unwrap()); - with_connection(&cfg, |conn| { - let count = |sql: &str| -> rusqlite::Result { conn.query_row(sql, [], |r| r.get(0)) }; - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_summaries WHERE tree_id = 'tree-legacy'")?, - 0 - ); - assert_eq!( - count( - "SELECT COUNT(*) FROM mem_tree_summary_embeddings WHERE summary_id = 'sum-legacy'" - )?, - 0 - ); - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_buffers WHERE tree_id = 'tree-legacy'")?, - 0 - ); - assert_eq!( - count("SELECT COUNT(*) FROM mem_tree_trees WHERE id = 'tree-legacy'")?, - 0 - ); - Ok(()) - }) - .unwrap(); - // the summary content file is removed from disk too. - assert!(!cfg.memory_tree_content_root().join(&summary_file).exists()); - - // idempotent: a second delete_source now finds nothing. - let again = delete_source_rpc(&cfg, target.to_string()) - .await - .unwrap() - .value; - assert!(!again.deleted); - assert_eq!(again.chunks_removed, 0); -} - -#[test] -fn delete_source_registered_in_schema_and_controllers() { - use crate::openhuman::memory::schema::{all_controller_schemas, all_registered_controllers}; - let schema = all_controller_schemas() - .into_iter() - .find(|s| s.function == "delete_source") - .expect("delete_source schema present"); - assert_eq!(schema.namespace, "memory_tree"); // => openhuman.memory_tree_delete_source - assert!(schema.inputs.iter().any(|f| f.name == "source_id")); - assert!(schema.outputs.iter().any(|f| f.name == "deleted")); - assert!(schema.outputs.iter().any(|f| f.name == "chunks_removed")); - assert!(all_registered_controllers() - .iter() - .any(|c| c.schema.function == "delete_source")); -} diff --git a/src/openhuman/memory_store/chunks/types.rs b/src/openhuman/memory_store/chunks/types.rs index 0602dd7c0..f7969c2aa 100644 --- a/src/openhuman/memory_store/chunks/types.rs +++ b/src/openhuman/memory_store/chunks/types.rs @@ -20,108 +20,11 @@ //! would force `_` arms on the host's exhaustive matches of a now-foreign //! `#[non_exhaustive]` enum. It flips with the ingest module (W6). -use serde::{Deserialize, Serialize}; - pub use tinycortex::memory::chunks::{ approx_token_count, chunk_id, conservative_token_estimate, truncate_to_conservative_tokens, - Chunk, Metadata, SourceKind, SourceRef, + Chunk, DataSource, Metadata, SourceKind, SourceRef, }; -/// Concrete upstream provider the content came from. -/// -/// Enumerates every provider listed in `m.excalidraw` Step 1 — Collect the -/// Data. Each variant maps to exactly one [`SourceKind`] via [`Self::kind`]. -/// -/// Wire form is snake_case (see `as_str` / `parse`) so it is stable across -/// DB rows, JSON-RPC payloads, and logs. -/// -/// Marked `#[non_exhaustive]` so new providers can be added in later phases -/// without breaking downstream pattern matches. -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -#[non_exhaustive] -pub enum DataSource { - // ── Chat transcripts (grouped by channel/group) ──────────────────── - Discord, - Telegram, - Whatsapp, - - // ── Agent conversations (stored as durable memory) ──────────────── - Conversation, - - // ── Email threads (grouped by thread) ────────────────────────────── - Gmail, - /// Catch-all for non-Gmail providers (Outlook, FastMail, generic IMAP, …). - OtherEmail, - - // ── Documents (no grouping) ──────────────────────────────────────── - Notion, - MeetingNotes, - DriveDocs, -} - -impl DataSource { - /// Which [`SourceKind`] this provider feeds into. - pub fn kind(self) -> SourceKind { - match self { - Self::Discord | Self::Telegram | Self::Whatsapp | Self::Conversation => { - SourceKind::Chat - } - Self::Gmail | Self::OtherEmail => SourceKind::Email, - Self::Notion | Self::MeetingNotes | Self::DriveDocs => SourceKind::Document, - } - } - - /// Stable snake_case identifier for DB storage, RPC payloads, and logs. - pub fn as_str(self) -> &'static str { - match self { - Self::Discord => "discord", - Self::Telegram => "telegram", - Self::Whatsapp => "whatsapp", - Self::Conversation => "conversation", - Self::Gmail => "gmail", - Self::OtherEmail => "other_email", - Self::Notion => "notion", - Self::MeetingNotes => "meeting_notes", - Self::DriveDocs => "drive_docs", - } - } - - /// Parse back from the on-wire / on-disk string form. - pub fn parse(s: &str) -> Result { - match s { - "discord" => Ok(Self::Discord), - "telegram" => Ok(Self::Telegram), - "whatsapp" => Ok(Self::Whatsapp), - "conversation" => Ok(Self::Conversation), - "gmail" => Ok(Self::Gmail), - "other_email" => Ok(Self::OtherEmail), - "notion" => Ok(Self::Notion), - "meeting_notes" => Ok(Self::MeetingNotes), - "drive_docs" => Ok(Self::DriveDocs), - other => Err(format!("unknown data source: {other}")), - } - } - - /// Every known variant, in declaration order. - /// - /// Useful for tests, CLI completion, and enumerating supported providers - /// in diagnostic output. - pub fn all() -> &'static [DataSource] { - &[ - Self::Discord, - Self::Telegram, - Self::Whatsapp, - Self::Conversation, - Self::Gmail, - Self::OtherEmail, - Self::Notion, - Self::MeetingNotes, - Self::DriveDocs, - ] - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/openhuman/memory_store/content/README.md b/src/openhuman/memory_store/content/README.md index 36d019ec5..2d7d31fb6 100644 --- a/src/openhuman/memory_store/content/README.md +++ b/src/openhuman/memory_store/content/README.md @@ -7,11 +7,11 @@ The body is **immutable** once written — only the YAML front-matter `tags:` bl ## Files - [`mod.rs`](mod.rs) — public surface: `StagedChunk`, `stage_chunks` (write all chunks atomically before SQLite upsert), `update_summary_tags` re-export. -- [`atomic.rs`](atomic.rs) — `write_if_new` (tempfile + fsync + rename, parent dir fsync on Unix), `stage_summary` (idempotent re-stage with on-disk SHA check + auto-rewrite on mismatch), `sha256_hex`, `StagedSummary`. -- [`compose/`](compose/) — YAML front-matter + body composition. `compose_chunk_file` for chunks (with email-only `participants:` / `aliases:` fields parsed from `gmail:{addr1|addr2|…}` source ids), `compose_summary_md` for summary nodes. `rewrite_tags` / `rewrite_summary_tags` swap the `tags:` block in place. `split_front_matter` parses `---\n…\n---\n`. -- [`paths.rs`](paths.rs) — path generators. `chunk_rel_path` (`email//.md`, `chat//.md`, `document//.md`); `summary_rel_path` (`summaries/{source,global,topic}/…`). `slugify_source_id` is the canonical filesystem-safe slug. +- `vendor/tinycortex/src/memory/store/content/atomic.rs` — `write_if_new` (tempfile + fsync + rename, parent dir fsync on Unix), `stage_summary` (idempotent re-stage with on-disk SHA check + auto-rewrite on mismatch), `sha256_hex`, `StagedSummary`. +- `vendor/tinycortex/src/memory/store/content/compose/` — YAML front-matter + body composition. `compose_chunk_file` for chunks (with email-only `participants:` / `aliases:` fields parsed from `gmail:{addr1|addr2|…}` source ids), `compose_summary_md` for summary nodes. `rewrite_tags` / `rewrite_summary_tags` swap the `tags:` block in place. `split_front_matter` parses `---\n…\n---\n`. +- `vendor/tinycortex/src/memory/store/content/paths.rs` — path generators. `chunk_rel_path` (`email//.md`, `chat//.md`, `document//.md`); `summary_rel_path` (`summaries/{source,global,topic}/…`). `slugify_source_id` is the canonical filesystem-safe slug. - [`read.rs`](read.rs) — `read_chunk_file` / `read_summary_file` parse front-matter and return body+SHA. `verify_*` compares against an expected SHA. `read_chunk_body` / `read_summary_body` resolve the path via SQLite and verify the integrity hash; this is the authoritative entry-point for callers that need the **full** body (LLM extractor, summariser, embedder, retrieval API). -- [`raw.rs`](raw.rs) — verbatim source-byte mirror under `/raw/`. Writes the unmodified upstream payload (eml, slack json, raw markdown) so downstream callers can re-canonicalise without re-fetching. +- `vendor/tinycortex/src/memory/store/content/raw.rs` — verbatim source-byte mirror under `/raw/`. Writes the unmodified upstream payload (eml, slack json, raw markdown) so downstream callers can re-canonicalise without re-fetching. - [`obsidian.rs`](obsidian.rs) + [`obsidian_defaults/`](obsidian_defaults/) — bootstrap an `.obsidian/` config (workspace, graph, app) into the content root on first write so a user opening the vault gets a usable view. - [`tags.rs`](tags.rs) — post-extraction tag rewrites. `update_chunk_tags` (atomic tempfile rewrite of the `tags:` block) and `update_summary_tags` (fetches entities from `mem_tree_entity_index`, builds Obsidian `kind/Value` tags, rewrites, verifies body SHA is unchanged). `slugify_tag_kind`, `slugify_tag_value`, `entity_tag` build the tag strings. - [`wiki_git/`](wiki_git/) — initializes `/wiki/.git`, commits only summary-node markdown under `summaries/**` plus the repo `.gitignore`, and stores read high-water marks as lightweight `refs/tags/read/*` pointers. Summary files are staged by `atomic.rs`; seal/ingest callers create descriptive git commits after SQLite persistence succeeds. diff --git a/src/openhuman/memory_store/content/atomic.rs b/src/openhuman/memory_store/content/atomic.rs deleted file mode 100644 index a54b54131..000000000 --- a/src/openhuman/memory_store/content/atomic.rs +++ /dev/null @@ -1,546 +0,0 @@ -//! Atomic content-file writes via tempfile + fsync + rename. -//! -//! Each chunk body is written to `/.tmp_.md`, then renamed to -//! its final path. The rename is atomic on any POSIX filesystem and behaves -//! correctly on NTFS (the old file is replaced atomically by the OS). -//! -//! **Immutability contract**: once a file exists at `abs_path`, it is never -//! overwritten. Callers must detect "already exists" and skip the write. - -use sha2::{Digest, Sha256}; -use std::io::Write; -use std::path::Path; - -use super::compose::{compose_summary_md, split_front_matter, SummaryComposeInput}; -use super::paths::{summary_rel_path_with_layout, SummaryDiskLayout}; - -/// Write `bytes` atomically to `abs_path` if the file does not already exist. -/// -/// Returns `Ok(true)` when the file was newly written, `Ok(false)` when it -/// already existed (the existing file is left unchanged). -/// -/// The write uses a sibling tempfile + rename so the final path is never -/// visible in a partial state. Parent directories are created automatically. -pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result { - // Fast path: file already exists. - if abs_path.exists() { - log::debug!( - "[content_store::atomic] skipping existing file: {}", - abs_path.display() - ); - return Ok(false); - } - - let parent = abs_path.parent().unwrap_or_else(|| Path::new(".")); - std::fs::create_dir_all(parent) - .map_err(|e| anyhow::anyhow!("create_dir_all {:?}: {e}", parent))?; - - // Write to a temp file in the same directory so rename is atomic. - let tmp_name = format!(".tmp_{}.md", uuid_v4_hex()); - let tmp_path = parent.join(&tmp_name); - - { - let mut f = std::fs::File::create(&tmp_path) - .map_err(|e| anyhow::anyhow!("create tempfile {:?}: {e}", tmp_path))?; - // Remove the temp file on write/fsync failure so a leaked - // `.tmp_.md` never accumulates in the user-facing vault. - if let Err(e) = f.write_all(bytes).and_then(|_| f.sync_all()) { - drop(f); - let _ = std::fs::remove_file(&tmp_path); - return Err(anyhow::anyhow!("write/fsync tempfile {:?}: {e}", tmp_path)); - } - } - - // Rename: if the target appeared concurrently (another thread/process beat - // us), we lost the race — remove our temp and return false. - match std::fs::rename(&tmp_path, abs_path) { - Ok(()) => { - // fsync the parent directory so the rename (directory entry - // update) is durable across a crash or power loss. Without this, - // sync_all() on the file alone only durabilises the file data; - // the new directory entry can remain in pagecache and be lost if - // the system crashes before the OS flushes it. On POSIX (Linux / - // macOS) this is required for rename durability. On Windows, NTFS - // handles this differently and File::sync_all on a directory - // handle is not meaningful, so we restrict the call to Unix. - #[cfg(unix)] - if let Some(parent) = abs_path.parent() { - if let Ok(dir) = std::fs::File::open(parent) { - if let Err(e) = dir.sync_all() { - // Best-effort: the rename already committed the file; - // a dirent fsync failure is logged but not fatal. - log::warn!( - "[content_store::atomic] parent dir fsync failed for {:?}: {e}", - parent - ); - } - } - } - log::debug!("[content_store::atomic] wrote {}", abs_path.display()); - Ok(true) - } - Err(e) => { - // Best-effort cleanup of the temp file on failure. - let _ = std::fs::remove_file(&tmp_path); - if abs_path.exists() { - // Lost the race — another writer created the file first. - log::debug!( - "[content_store::atomic] lost rename race for {}", - abs_path.display() - ); - Ok(false) - } else { - Err(anyhow::anyhow!( - "rename {:?} -> {:?}: {e}", - tmp_path, - abs_path - )) - } - } - } -} - -/// Ensure the file at `abs_path` contains exactly `full_bytes`, whose body -/// (front-matter excluded) hashes to `body_sha256`. -/// -/// This is the write-side half of the content store's integrity contract -/// (#4689): the recorded `content_sha256` must always match the bytes actually -/// on disk. `write_if_new` alone cannot guarantee that — it skips an existing -/// file unconditionally, so a stale/partial pre-existing file at the target -/// path would be left in place while the caller records the *new* body's sha, -/// permanently diverging the DB token from disk. -/// -/// Behaviour when the file already exists: -/// - on-disk body sha matches `body_sha256` → idempotent no-op. -/// - mismatch → atomically replaced (temp file + rename over the destination) -/// from `full_bytes`, so the file is never observed missing or partial. At -/// ingest the freshly-composed input is authoritative, so overwriting a -/// drifted on-disk file is the correct reconciliation. -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 { - log::debug!( - "[content_store::atomic] file already on disk with matching body sha: {}", - abs_path.display() - ); - return Ok(()); - } - log::debug!( - "[content_store::atomic] on-disk body sha mismatch for {} (disk={disk_sha} new={body_sha256}) — re-staging", - abs_path.display() - ); - } - // Write the replacement to a sibling temp file and atomically rename it over - // the destination. rename() replaces an existing file atomically on POSIX and - // NTFS, so — unlike unlink-then-write — the destination is never observed - // missing or partially written even if the process crashes or the write - // fails mid-way (a committed DB row can never end up pointing at an absent - // body file). Post-condition: on success `abs_path` holds exactly - // `full_bytes`, whose body hashes to `body_sha256`. - write_via_temp_rename(abs_path, full_bytes) -} - -/// Write `bytes` to `abs_path` via a sibling tempfile + fsync + atomic rename, -/// **replacing** any existing file. The rename is atomic on POSIX and NTFS, so -/// the destination is never seen missing or half-written. Parent directories are -/// created on demand and the parent dir entry is fsync'd on Unix for durability. -fn write_via_temp_rename(abs_path: &Path, bytes: &[u8]) -> anyhow::Result<()> { - let parent = abs_path.parent().unwrap_or_else(|| Path::new(".")); - std::fs::create_dir_all(parent) - .map_err(|e| anyhow::anyhow!("create_dir_all {:?}: {e}", parent))?; - - let tmp_path = parent.join(format!(".tmp_{}.md", uuid_v4_hex())); - { - let mut f = std::fs::File::create(&tmp_path) - .map_err(|e| anyhow::anyhow!("create tempfile {:?}: {e}", tmp_path))?; - // Clean up the temp file on a write/fsync failure too — not just on the - // rename failure below. content_root is the user-facing vault, so a - // leaked `.tmp_.md` is visible in Obsidian and would accumulate. - if let Err(e) = f.write_all(bytes).and_then(|_| f.sync_all()) { - drop(f); - let _ = std::fs::remove_file(&tmp_path); - return Err(anyhow::anyhow!("write/fsync tempfile {:?}: {e}", tmp_path)); - } - } - - if let Err(e) = std::fs::rename(&tmp_path, abs_path) { - // Keep the old destination intact; only our temp is cleaned up. - let _ = std::fs::remove_file(&tmp_path); - return Err(anyhow::anyhow!( - "rename {:?} -> {:?}: {e}", - tmp_path, - abs_path - )); - } - - #[cfg(unix)] - if let Ok(dir) = std::fs::File::open(parent) { - if let Err(e) = dir.sync_all() { - log::warn!( - "[content_store::atomic] parent dir fsync failed for {:?}: {e}", - parent - ); - } - } - log::debug!( - "[content_store::atomic] wrote (replace) {}", - abs_path.display() - ); - Ok(()) -} - -/// A summary that has been written to disk and is ready for SQLite upsert. -#[derive(Debug, Clone)] -pub struct StagedSummary { - /// Identifier of the summary that was staged. - pub summary_id: String, - /// Relative content path (forward-slash, e.g. - /// `"wiki/summaries/source-slug/L1/id.md"`). - pub content_path: String, - /// SHA-256 hex digest over the **body bytes** only (front-matter excluded). - pub content_sha256: String, -} - -/// Write a summary `.md` file to disk and return a [`StagedSummary`] ready for -/// SQLite upsert. -/// -/// The relative path is built from the input metadata and the `tree_kind`. The -/// `scope_slug` must already be slugified by the caller. The global tree is a -/// singleton, so its summaries all land under one `global/` folder regardless -/// of the day they cover — no date argument is needed. -/// -/// If the file already exists with the same body SHA-256 (idempotent re-stage), -/// the existing `StagedSummary` is returned without rewriting. -pub fn stage_summary( - content_root: &Path, - input: &SummaryComposeInput<'_>, - scope_slug: &str, -) -> anyhow::Result { - stage_summary_with_layout(content_root, input, scope_slug, SummaryDiskLayout::Standard) -} - -/// Layout-aware variant of [`stage_summary`]. Document source trees pass a -/// [`SummaryDiskLayout::DocSubtree`] (per-document, versioned) or -/// [`SummaryDiskLayout::Merge`] (cross-document merge tier) so the on-disk -/// vault mirrors the logical tree (`notion` → `docs//v-` → -/// `merge`). All other callers use [`stage_summary`] (`Standard`) unchanged. -pub fn stage_summary_with_layout( - content_root: &Path, - input: &SummaryComposeInput<'_>, - scope_slug: &str, - layout: SummaryDiskLayout<'_>, -) -> anyhow::Result { - let rel_path = summary_rel_path_with_layout( - input.tree_kind, - scope_slug, - input.level, - input.summary_id, - layout, - ); - // Derive the absolute path by joining the relative path components onto - // the content root (same join `summary_abs_path` does internally) so the - // two stay consistent regardless of layout. - let abs_path = { - let mut abs = content_root.to_path_buf(); - for component in rel_path.split('/') { - abs.push(component); - } - abs - }; - - let composed = compose_summary_md(input); - let body_bytes = composed.body.as_bytes(); - let sha256 = sha256_hex(body_bytes); - - // Idempotent re-stage that self-heals a stale/corrupt on-disk file: matching - // body sha is a no-op, a mismatch is atomically rewritten. Not re-writing - // would leave SQLite storing a content_sha256 that doesn't match the actual - // on-disk bytes, breaking integrity checks. See [`write_or_replace_body`]. - write_or_replace_body(&abs_path, composed.full.as_bytes(), &sha256)?; - - log::debug!( - "[content_store::atomic] staged summary {} → {}", - input.summary_id, - rel_path - ); - - Ok(StagedSummary { - summary_id: input.summary_id.to_string(), - content_path: rel_path, - content_sha256: sha256, - }) -} - -/// Read a summary/chunk `.md` file from disk, split off the YAML front-matter, -/// and return the SHA-256 hex digest of the **body bytes only**. Returns an -/// empty string (not an error) if the file cannot be read or parsed, so -/// callers can use the result as a cache key without propagating IO errors. -fn read_body_sha256(path: &Path) -> anyhow::Result { - let raw = std::fs::read(path)?; - let content = std::str::from_utf8(&raw)?; - let (_fm, body) = split_front_matter(content) - .ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", path))?; - Ok(sha256_hex(body.as_bytes())) -} - -/// Compute the SHA-256 hex digest of `bytes`. -pub fn sha256_hex(bytes: &[u8]) -> String { - let mut hasher = Sha256::new(); - hasher.update(bytes); - hex::encode(hasher.finalize()) -} - -/// Tiny deterministic-ish hex string for temp file names. -fn uuid_v4_hex() -> String { - use std::time::{SystemTime, UNIX_EPOCH}; - let t = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .subsec_nanos(); - // Use a counter + timestamp for entropy (thread_id::as_u64 is nightly-only). - static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - format!( - "{:08x}{:016x}", - t, - n.wrapping_mul(0x9e37_79b9_7f4a_7c15).wrapping_add(t as u64) - ) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::content::compose::SummaryComposeInput; - use crate::openhuman::memory_store::content::paths::SummaryTreeKind; - use tempfile::TempDir; - - #[test] - fn write_creates_file_and_returns_true() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("sub").join("0.md"); - let written = write_if_new(&path, b"hello world").unwrap(); - assert!(written, "first write must return true"); - assert_eq!(std::fs::read(&path).unwrap(), b"hello world"); - } - - #[test] - fn write_is_idempotent_returns_false_on_second_call() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("0.md"); - write_if_new(&path, b"first").unwrap(); - let written = write_if_new(&path, b"second").unwrap(); - assert!(!written, "second write must return false"); - assert_eq!(std::fs::read(&path).unwrap(), b"first"); - } - - #[test] - fn sha256_hex_is_stable() { - let a = sha256_hex(b"hello"); - let b = sha256_hex(b"hello"); - assert_eq!(a, b); - assert_ne!(sha256_hex(b"hello"), sha256_hex(b"world")); - assert_eq!(a.len(), 64); // 32 bytes → 64 hex chars - } - - #[test] - fn write_or_replace_body_writes_new_and_is_idempotent() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("c.md"); - let full = b"---\nk: v\n---\nBODY"; - let body_sha = sha256_hex(b"BODY"); - // Fresh write. - write_or_replace_body(&path, full, &body_sha).unwrap(); - assert_eq!(std::fs::read(&path).unwrap(), full); - // Idempotent: a matching on-disk body sha leaves the file untouched. - write_or_replace_body(&path, full, &body_sha).unwrap(); - assert_eq!(std::fs::read(&path).unwrap(), full); - } - - #[test] - fn write_or_replace_body_overwrites_stale_file() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("c.md"); - // A stale pre-existing file with a different body would be skipped by - // write_if_new; write_or_replace_body must reconcile it (#4689). - write_if_new(&path, b"---\nk: v\n---\nOLD").unwrap(); - let new_full = b"---\nk: v\n---\nNEW"; - write_or_replace_body(&path, new_full, &sha256_hex(b"NEW")).unwrap(); - assert_eq!(std::fs::read(&path).unwrap(), new_full); - } - - #[test] - fn write_or_replace_body_errors_when_stale_target_cannot_be_replaced() { - // If the stale target can't be removed/overwritten (here: a directory - // sits at the path), the function must refuse rather than let the caller - // record a content_sha256 the on-disk bytes do not match (#4689). - let dir = TempDir::new().unwrap(); - let path = dir.path().join("c.md"); - std::fs::create_dir_all(&path).unwrap(); - let res = write_or_replace_body(&path, b"---\nk: v\n---\nNEW", &sha256_hex(b"NEW")); - assert!( - res.is_err(), - "must not record a sha the target does not hold" - ); - } - - fn mk_summary_input<'a>( - tree_kind: SummaryTreeKind, - scope: &'a str, - id: &'a str, - body: &'a str, - children: &'a [String], - ) -> SummaryComposeInput<'a> { - use chrono::TimeZone; - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - SummaryComposeInput { - summary_id: id, - tree_kind, - tree_id: "tree-001", - tree_scope: scope, - level: 1, - child_ids: children, - child_basenames: None, - child_count: children.len(), - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body, - } - } - - #[test] - fn stage_summary_writes_file_and_returns_staged() { - let dir = TempDir::new().unwrap(); - let children = vec!["c1".to_string()]; - let input = mk_summary_input( - SummaryTreeKind::Source, - "gmail:alice@x.com", - "summary:L1:test1", - "summary body", - &children, - ); - let staged = stage_summary(dir.path(), &input, "gmail-alice-x-com").unwrap(); - assert_eq!(staged.summary_id, "summary:L1:test1"); - assert!(staged.content_path.starts_with("wiki/summaries/source-")); - assert!(staged.content_path.ends_with(".md")); - assert_eq!(staged.content_sha256.len(), 64); - - // File must exist on disk - let mut abs = dir.path().to_path_buf(); - for part in staged.content_path.split('/') { - abs.push(part); - } - assert!(abs.exists(), "staged file must exist"); - } - - #[test] - fn stage_summary_is_idempotent() { - let dir = TempDir::new().unwrap(); - let children = vec!["c1".to_string()]; - let input = mk_summary_input( - SummaryTreeKind::Topic, - "person:alex", - "summary:L1:idem", - "idempotent body", - &children, - ); - let first = stage_summary(dir.path(), &input, "person-alex").unwrap(); - let second = stage_summary(dir.path(), &input, "person-alex").unwrap(); - assert_eq!(first.content_sha256, second.content_sha256); - assert_eq!(first.content_path, second.content_path); - } - - #[test] - fn stage_summary_global_uses_singleton_folder_no_date() { - let dir = TempDir::new().unwrap(); - let children = vec![]; - let input = mk_summary_input( - SummaryTreeKind::Global, - "global", - "summary:L0:daily", - "daily recap", - &children, - ); - let staged = stage_summary(dir.path(), &input, "global").unwrap(); - // Singleton global tree → one folder, no per-day date segment. The - // `L1` segment comes from `mk_summary_input`'s level=1; what matters - // is the single `global/` folder with no date. - assert_eq!( - staged.content_path, "wiki/summaries/global/L1/summary-L0-daily.md", - "global summary must land in the singleton global/ folder; got: {}", - staged.content_path - ); - } - - #[test] - fn stage_summary_sha256_is_over_body_only() { - let dir = TempDir::new().unwrap(); - let children = vec![]; - let body = "the body content"; - let input = mk_summary_input( - SummaryTreeKind::Source, - "gmail:x@y.com", - "summary:L1:sha-test", - body, - &children, - ); - let staged = stage_summary(dir.path(), &input, "gmail-x-y-com").unwrap(); - let expected = sha256_hex(body.as_bytes()); - assert_eq!(staged.content_sha256, expected); - } - - #[test] - fn stage_summary_rewrites_stale_on_disk_body() { - // Create a tempdir and write a "stale" file at the expected path with - // a body that differs from what the new stage_summary call would write. - // After stage_summary, the file on disk must match the new body. - let dir = TempDir::new().unwrap(); - let children = vec!["c1".to_string()]; - let new_body = "fresh body for re-stage test"; - let input = mk_summary_input( - SummaryTreeKind::Source, - "gmail:stale@test.com", - "summary:L1:stale-test", - new_body, - &children, - ); - - // First stage with the real body to get the path. - let first = stage_summary(dir.path(), &input, "gmail-stale-test-com").unwrap(); - - // Corrupt the on-disk file by writing a different body to the path. - let mut abs = dir.path().to_path_buf(); - for part in first.content_path.split('/') { - abs.push(part); - } - // Overwrite with stale content. - std::fs::write(&abs, b"---\nstale_key: true\n---\nSTALE BODY CONTENT").unwrap(); - - // Now re-stage: must detect sha mismatch and re-write. - let second = stage_summary(dir.path(), &input, "gmail-stale-test-com").unwrap(); - - // The returned sha must match the new body. - let expected_sha = sha256_hex(new_body.as_bytes()); - assert_eq!( - second.content_sha256, expected_sha, - "re-staged sha must match new body" - ); - - // The on-disk file must now contain the new body (not the stale one). - let disk_bytes = std::fs::read(&abs).unwrap(); - let disk_str = std::str::from_utf8(&disk_bytes).unwrap(); - assert!( - disk_str.contains(new_body), - "on-disk file must contain new body after re-stage" - ); - assert!( - !disk_str.contains("STALE BODY CONTENT"), - "stale body must be gone after re-stage" - ); - } -} diff --git a/src/openhuman/memory_store/content/compose/chunk.rs b/src/openhuman/memory_store/content/compose/chunk.rs deleted file mode 100644 index 790df5989..000000000 --- a/src/openhuman/memory_store/content/compose/chunk.rs +++ /dev/null @@ -1,209 +0,0 @@ -//! Chunk `.md` file composition and tag rewriting. - -use crate::openhuman::memory_store::chunks::types::{Chunk, SourceKind}; -use crate::openhuman::memory_store::content::compose::yaml::{ - split_front_matter, with_source_tag, yaml_scalar, -}; - -/// Compose the full file content (front-matter + body) for `chunk`. -/// -/// Returns `(full_file_bytes, body_bytes)`. The caller writes `full_file_bytes` -/// to disk; `body_bytes` is what the SHA-256 is computed over. -pub fn compose_chunk_file(chunk: &Chunk) -> (Vec, Vec) { - let front_matter = build_front_matter(chunk); - let body = chunk.content.as_bytes().to_vec(); - - let mut full = Vec::with_capacity(front_matter.len() + body.len()); - full.extend_from_slice(&front_matter); - full.extend_from_slice(&body); - - (full, body) -} - -/// Build the YAML front-matter block (including delimiters) as UTF-8 bytes. -fn build_front_matter(chunk: &Chunk) -> Vec { - let meta = &chunk.metadata; - let ts = meta.timestamp.to_rfc3339(); - let ts_start = meta.time_range.0.to_rfc3339(); - let ts_end = meta.time_range.1.to_rfc3339(); - - let mut fm = String::new(); - fm.push_str("---\n"); - fm.push_str(&format!("source_kind: {}\n", meta.source_kind.as_str())); - // Escape backslashes and quotes in source_id for safety. - fm.push_str(&format!("source_id: {}\n", yaml_scalar(&meta.source_id))); - if let Some(path_scope) = meta.path_scope.as_deref() { - fm.push_str(&format!("path_scope: {}\n", yaml_scalar(path_scope))); - } - fm.push_str(&format!("seq: {}\n", chunk.seq_in_source)); - fm.push_str(&format!("owner: {}\n", yaml_scalar(&meta.owner))); - fm.push_str(&format!("timestamp: {ts}\n")); - fm.push_str(&format!("time_range_start: {ts_start}\n")); - fm.push_str(&format!("time_range_end: {ts_end}\n")); - - if let Some(ref sr) = meta.source_ref { - fm.push_str(&format!("source_ref: {}\n", yaml_scalar(&sr.value))); - } - - // Always seed the source tag so the Obsidian graph filter can pick - // up `source/` for every chunk regardless of what the - // ingest-side tag list contained. - let source_scope = meta.path_scope.as_deref().unwrap_or(&meta.source_id); - log::debug!( - "[content_store::compose] seeding source tag source_id={} source_scope={} path_scope={}", - crate::openhuman::memory::util::redact::redact(&meta.source_id), - crate::openhuman::memory::util::redact::redact(source_scope), - meta.path_scope.is_some() - ); - let seeded_tags = with_source_tag(source_scope, &meta.tags); - fm.push_str("tags:\n"); - for tag in &seeded_tags { - fm.push_str(&format!(" - {}\n", yaml_scalar(tag))); - } - - // Email-specific fields: participants list + Obsidian alias. - // Parsed from source_id which is `gmail:{participants}` for Gmail-ingested - // chunks, where participants is `addr1|addr2|...` (sorted, deduped). - // If the format doesn't match, these fields are omitted. - if meta.source_kind == SourceKind::Email { - if let Some(addrs) = parse_gmail_participants_source_id(&meta.source_id) { - // participants: YAML list - fm.push_str("participants:\n"); - for addr in &addrs { - fm.push_str(&format!(" - {}\n", yaml_scalar(addr))); - } - // aliases: human-readable conversation label for Obsidian - let alias = build_participants_alias(&addrs, chunk.seq_in_source); - fm.push_str("aliases:\n"); - fm.push_str(&format!(" - {}\n", yaml_scalar(&alias))); - } - } - - fm.push_str("---\n"); - fm.into_bytes() -} - -/// Parse a `gmail:{participants}` source_id into the list of participant addresses. -/// -/// `participants` is `addr1|addr2|...` (sorted, deduped, pipe-separated). -/// Returns `Some(Vec)` when the source_id has exactly two -/// colon-separated segments (`gmail` prefix + non-empty participants). Returns -/// `None` for legacy or malformed source_ids. -fn parse_gmail_participants_source_id(source_id: &str) -> Option> { - let (prefix, participants) = source_id.split_once(':')?; - if prefix != "gmail" || participants.is_empty() { - return None; - } - let addrs: Vec = participants - .split('|') - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) - .collect(); - if addrs.is_empty() { - None - } else { - Some(addrs) - } -} - -/// Build a human-readable alias for an email chunk suitable for Obsidian's -/// `aliases:` field. -/// -/// For two participants: `"alice@x.com <-> bob@y.com: chunk 0"` -/// For more than two: `"alice@x.com <-> 2 others: chunk 0"` -/// (where `alice@x.com` is the first in sorted order) -/// -/// The alias is kept under ~80 characters to avoid YAML rendering issues. -fn build_participants_alias(addrs: &[String], seq: u32) -> String { - let label = match addrs { - [] => "unknown".to_string(), - [only] => only.clone(), - [first, second] => format!("{} <-> {}", first, second), - [first, rest @ ..] => format!("{} <-> {} others", first, rest.len()), - }; - format!("{}: chunk {}", label, seq) -} - -/// Rewrite the `tags:` block in an existing file's front-matter, replacing it -/// with the new tag list while leaving the body unchanged. -/// -/// Returns the new full file bytes. Errors if the front-matter delimiters -/// cannot be found. -pub fn rewrite_tags(file_bytes: &[u8], new_tags: &[String]) -> Result, String> { - let content = - std::str::from_utf8(file_bytes).map_err(|e| format!("file is not valid UTF-8: {e}"))?; - - let (front_matter, body) = split_front_matter(content) - .ok_or_else(|| "cannot find front-matter delimiters".to_string())?; - - // Rewrite tags: block in the front-matter string. - let new_fm = replace_tags_in_front_matter(front_matter, new_tags)?; - - let mut out = Vec::with_capacity(new_fm.len() + body.len() + 4); - out.extend_from_slice(new_fm.as_bytes()); - out.extend_from_slice(body.as_bytes()); - Ok(out) -} - -/// Replace the `tags:` stanza in a front-matter string. Returns the new -/// front-matter string (delimiters preserved). -fn replace_tags_in_front_matter(fm: &str, new_tags: &[String]) -> Result { - // Build the replacement block. - let replacement = if new_tags.is_empty() { - "tags: []".to_string() - } else { - let mut s = "tags:".to_string(); - for tag in new_tags { - s.push('\n'); - s.push_str(&format!(" - {}", yaml_scalar(tag))); - } - s - }; - - // Locate the `tags:` key and consume through the block. - let lines: Vec<&str> = fm.lines().collect(); - let mut out_lines: Vec<&str> = Vec::new(); - let mut i = 0; - let mut found = false; - - while i < lines.len() { - let line = lines[i]; - if line == "tags: []" || line == "tags:" { - found = true; - // Skip all subsequent lines that are tag list items (start with ` - `). - // The replacement will be inserted wholesale. - i += 1; - if line == "tags:" { - while i < lines.len() && lines[i].starts_with(" - ") { - i += 1; - } - } - // We've consumed the old block; we'll append replacement after the loop. - continue; - } - out_lines.push(line); - i += 1; - } - - if !found { - return Err("tags: key not found in front-matter".to_string()); - } - - // Rebuild: all non-tag lines + replacement + closing `---`. - // Front-matter was: `---\n...\ntags: ...\n---\n` - // After loop, out_lines has everything except the tags block. - // Insert replacement before the closing `---`. - let closing = out_lines - .iter() - .rposition(|l| *l == "---") - .unwrap_or(out_lines.len()); - - let mut result_lines: Vec = - out_lines[..closing].iter().map(|l| l.to_string()).collect(); - result_lines.push(replacement); - result_lines.push("---".to_string()); - - let mut result = result_lines.join("\n"); - result.push('\n'); - Ok(result) -} diff --git a/src/openhuman/memory_store/content/compose/mod.rs b/src/openhuman/memory_store/content/compose/mod.rs deleted file mode 100644 index bc4f51a1f..000000000 --- a/src/openhuman/memory_store/content/compose/mod.rs +++ /dev/null @@ -1,53 +0,0 @@ -//! YAML front-matter + body composition for chunk `.md` files. -//! -//! Each file written to disk has the form: -//! ```text -//! --- -//! source_kind: chat -//! source_id: slack:#eng -//! seq: 0 -//! owner: alice@example.com -//! timestamp: 2026-04-28T10:00:00Z -//! time_range_start: 2026-04-28T10:00:00Z -//! time_range_end: 2026-04-28T10:05:00Z -//! source_ref: slack://permalink/… -//! tags: -//! - person/Alice-Smith -//! - project/Phoenix -//! --- -//! ## 2026-04-28T10:00:00Z — alice -//! Message body here. -//! ``` -//! -//! For email source_kind, additional fields are emitted: -//! ```text -//! participants: -//! - alice@example.com -//! - bob@example.com -//! aliases: -//! - "alice@example.com <-> bob@example.com: chunk 0" -//! ``` -//! These are parsed from the `source_id` field (format `gmail:{participants}` -//! where `participants` is `addr1|addr2|...` pipe-separated) at compose time. -//! `sender` and `thread_id` are no longer emitted — they are not meaningful -//! with participant-based bucketing. -//! -//! **SHA-256 is computed over the body bytes only** (everything after `---\n` -//! on the second delimiter line). This allows tags to be rewritten atomically -//! without invalidating the content hash. - -pub mod chunk; -pub mod summary; -pub mod yaml; - -#[cfg(test)] -mod tests; - -pub const MEMORY_ARTIFACT_FORMAT: u32 = 2; -pub const OPENHUMAN_CORE_VERSION: &str = env!("CARGO_PKG_VERSION"); - -// ── Re-exports (preserve original public API) ──────────────────────────────── - -pub use chunk::{compose_chunk_file, rewrite_tags}; -pub use summary::{compose_summary_md, rewrite_summary_tags, ComposedSummary, SummaryComposeInput}; -pub use yaml::{scan_fm_field, source_tag, split_front_matter, with_source_tag}; diff --git a/src/openhuman/memory_store/content/compose/summary.rs b/src/openhuman/memory_store/content/compose/summary.rs deleted file mode 100644 index 47db6ae0f..000000000 --- a/src/openhuman/memory_store/content/compose/summary.rs +++ /dev/null @@ -1,298 +0,0 @@ -//! Summary `.md` file composition and tag rewriting. - -use chrono::{DateTime, Utc}; - -use crate::openhuman::memory_store::content::compose::chunk::rewrite_tags; -use crate::openhuman::memory_store::content::compose::yaml::{ - source_tag, split_front_matter, yaml_scalar, -}; -use crate::openhuman::memory_store::content::compose::{ - MEMORY_ARTIFACT_FORMAT, OPENHUMAN_CORE_VERSION, -}; -use crate::openhuman::memory_store::content::paths::{summary_filename, SummaryTreeKind}; - -/// Input data required to compose a summary `.md` file. -pub struct SummaryComposeInput<'a> { - /// Stable id of the summary node (also used to derive the filename). - pub summary_id: &'a str, - /// Which tree (source / global / topic) this summary belongs to. - pub tree_kind: SummaryTreeKind, - /// Owning tree id (FK into `mem_tree_trees`). - pub tree_id: &'a str, - /// Raw tree scope string, e.g. `"gmail:alice@x.com|bob@y.com"` or `"global"`. - pub tree_scope: &'a str, - /// Level in the tree (L0 = leaves, L1+ = summaries). - pub level: u32, - /// Child ids (chunk_ids at L0 → L1, summary_ids for cascades). - pub child_ids: &'a [String], - /// Optional per-child wikilink basename overrides, aligned with - /// `child_ids` by index. When `Some(basename)` is provided for a - /// child, the front-matter `children: [[…]]` wikilink uses that - /// basename instead of `sanitize_filename(child_id)`. - /// - /// Used to point chunk-level children at their **raw archive** - /// files when the chunk store no longer stages on-disk `.md` - /// files (today: email, since email chunks live as byte ranges - /// inside `raw//_.md` instead of - /// `email//.md`). Without this, Obsidian - /// wikilinks resolve to a non-existent `[[]]` - /// target and the graph view stops drawing edges from L1 - /// summaries down to leaves. - /// - /// `None` (or `Some` entries that are themselves `None`) falls - /// back to the default `sanitize_filename(child_id)` behaviour, - /// which is correct for L≥2 (children are summary ids that map - /// to actual `summaries/...md` files) and for legacy chunks - /// still staged on-disk. - pub child_basenames: Option<&'a [Option]>, - /// Total child count (== child_ids.len() unless truncated). - pub child_count: usize, - /// Start of the time range covered by this summary's children. - pub time_range_start: DateTime, - /// End of the time range covered by this summary's children. - pub time_range_end: DateTime, - /// When the buffer was sealed into this summary node. - pub sealed_at: DateTime, - /// Raw summariser output text — the body written to disk. - pub body: &'a str, -} - -/// The composed front-matter, body, and full file content for a summary. -/// -/// `body` is what the SHA-256 integrity hash is computed over. -pub struct ComposedSummary { - /// The YAML front-matter block (including `---` delimiters), UTF-8 string. - pub front_matter: String, - /// The body (summariser output), UTF-8 string. - pub body: String, - /// `front_matter + body` — what gets written to disk. - pub full: String, -} - -/// Compose the full `.md` content for a summary node. -/// -/// Returns a [`ComposedSummary`] whose `full` field is written to disk. -/// SHA-256 is computed over `body` bytes only, not `full`. -pub fn compose_summary_md(record: &SummaryComposeInput<'_>) -> ComposedSummary { - let fm = build_summary_front_matter(record); - let body = record.body.to_string(); - let full = format!("{}{}", fm, body); - ComposedSummary { - front_matter: fm, - body, - full, - } -} - -/// Build the YAML front-matter block for a summary node. -fn build_summary_front_matter(r: &SummaryComposeInput<'_>) -> String { - let tree_kind_str = match r.tree_kind { - SummaryTreeKind::Source => "source", - SummaryTreeKind::Global => "global", - SummaryTreeKind::Topic => "topic", - }; - - let trs = r.time_range_start.to_rfc3339(); - let tre = r.time_range_end.to_rfc3339(); - let sealed = r.sealed_at.to_rfc3339(); - - let mut fm = String::new(); - fm.push_str("---\n"); - fm.push_str(&format!("id: {}\n", yaml_scalar(r.summary_id))); - fm.push_str("kind: summary\n"); - fm.push_str(&format!("tree_kind: {tree_kind_str}\n")); - fm.push_str(&format!("tree_id: {}\n", yaml_scalar(r.tree_id))); - fm.push_str(&format!("tree_scope: {}\n", yaml_scalar(r.tree_scope))); - fm.push_str(&format!("level: {}\n", r.level)); - - // children: YAML list of Obsidian wikilinks (`[[]]`) so the - // graph view draws summary→child edges. The wikilink target must match - // the actual file basename — for chunks that's the raw chunk_id (a SHA - // hash with no illegal chars), but for child summaries the structured id - // `summary:L:UUID` is sanitised to `summary-L-UUID` by - // `summary_rel_path` (colons are illegal on Windows NTFS). We apply the - // same sanitisation here so the link resolves. `yaml_scalar` auto-quotes - // because of the leading `[`, emitting `"[[]]"`. - if r.child_ids.is_empty() { - fm.push_str("children: []\n"); - } else { - fm.push_str("children:\n"); - for (i, id) in r.child_ids.iter().enumerate() { - // Prefer a caller-supplied basename override (used for L1 - // chunk children that live in the raw archive instead of - // the chunk-store path); fall back to the sanitised - // chunk/summary id. - let basename: String = match r - .child_basenames - .and_then(|overrides| overrides.get(i)) - .and_then(|slot| slot.as_ref()) - { - Some(b) => b.clone(), - None => summary_filename(id), - }; - let wikilink = format!("[[{}]]", basename); - fm.push_str(&format!(" - {}\n", yaml_scalar(&wikilink))); - } - } - fm.push_str(&format!("child_count: {}\n", r.child_count)); - fm.push_str(&format!("time_range_start: {trs}\n")); - fm.push_str(&format!("time_range_end: {tre}\n")); - fm.push_str(&format!("sealed_at: {sealed}\n")); - fm.push_str(&format!( - "openhuman_core_version: {}\n", - yaml_scalar(OPENHUMAN_CORE_VERSION) - )); - fm.push_str(&format!( - "memory_artifact_format: {}\n", - MEMORY_ARTIFACT_FORMAT - )); - - // aliases: human-readable title - let alias = build_summary_alias(r); - fm.push_str("aliases:\n"); - fm.push_str(&format!(" - {}\n", yaml_scalar(&alias))); - - // Source-tree summaries get a `source/` seed tag for graph - // filtering. Global / topic trees aggregate across sources, so the - // `source/...` tag has no single value there — leave them untagged - // at compose time (LLM extraction adds entity tags later). - if matches!(r.tree_kind, SummaryTreeKind::Source) { - fm.push_str("tags:\n"); - fm.push_str(&format!(" - {}\n", yaml_scalar(&source_tag(r.tree_scope)))); - } else { - fm.push_str("tags: []\n"); - } - fm.push_str("---\n"); - fm -} - -/// Build a human-readable alias for the summary's `aliases:` front-matter field. -fn build_summary_alias(r: &SummaryComposeInput<'_>) -> String { - let date_range = format_date_range(r.time_range_start, r.time_range_end); - match r.tree_kind { - SummaryTreeKind::Source => { - let scope_short = scope_short_label(r.tree_scope); - format!( - "L{} \u{00b7} {} \u{00b7} {} children \u{00b7} {}", - r.level, scope_short, r.child_count, date_range - ) - } - SummaryTreeKind::Global => { - format!( - "L{} \u{00b7} global digest \u{00b7} {}", - r.level, date_range - ) - } - SummaryTreeKind::Topic => { - // Strip protocol prefix like "topic:" from scope for readability. - let entity = r - .tree_scope - .split_once(':') - .map(|(_, v)| v) - .unwrap_or(r.tree_scope); - format!( - "L{} \u{00b7} topic {} \u{00b7} {} children", - r.level, entity, r.child_count - ) - } - } -} - -/// Format the date range as `"yyyy-mm-dd"` (if start == end date) or -/// `"yyyy-mm-dd–yyyy-mm-dd"`. -fn format_date_range(start: DateTime, end: DateTime) -> String { - let s = start.format("%Y-%m-%d").to_string(); - let e = end.format("%Y-%m-%d").to_string(); - if s == e { - s - } else { - format!("{s}\u{2013}{e}") // en dash - } -} - -/// Build a short human-readable label for the tree scope used in aliases. -/// -/// For Gmail source scopes like `"gmail:alice@x.com|bob@y.com"`: -/// - 2 participants → `"alice@x.com ↔ bob@y.com"` -/// - N > 2 → `"alice@x.com + N-1 others"` -/// - Otherwise → the raw scope (e.g. `"slack:#eng"`) -pub fn scope_short_label(scope: &str) -> String { - if let Some((prefix, participants)) = scope.split_once(':') { - if prefix == "gmail" && !participants.is_empty() { - let addrs: Vec<&str> = participants.split('|').collect(); - return match addrs.as_slice() { - [] => scope.to_string(), - [only] => only.to_string(), - [first, second] => format!("{} \u{2194} {}", first, second), // ↔ - [first, rest @ ..] => format!("{} + {} others", first, rest.len()), - }; - } - } - scope.to_string() -} - -/// Rewrite the `tags:` block in a summary file's front-matter, replacing it -/// with the new tag list while leaving the body unchanged. -/// -/// Reuses the generic [`rewrite_tags`] function — the front-matter structure -/// is identical for both chunk and summary `.md` files. -pub fn rewrite_summary_tags(file_bytes: &[u8], new_tags: &[String]) -> Result, String> { - let rewritten = rewrite_tags(file_bytes, new_tags)?; - let content = - std::str::from_utf8(&rewritten).map_err(|e| format!("file is not valid UTF-8: {e}"))?; - let (front_matter, body) = split_front_matter(content) - .ok_or_else(|| "cannot find front-matter delimiters".to_string())?; - let front_matter = upsert_summary_provenance(front_matter); - - let mut out = Vec::with_capacity(front_matter.len() + body.len()); - out.extend_from_slice(front_matter.as_bytes()); - out.extend_from_slice(body.as_bytes()); - Ok(out) -} - -fn upsert_summary_provenance(front_matter: &str) -> String { - let mut lines: Vec = Vec::new(); - let mut inserted = false; - - for raw in front_matter.lines() { - if raw.starts_with("openhuman_core_version: ") - || raw.starts_with("memory_artifact_format: ") - { - continue; - } - if !inserted && raw == "aliases:" { - lines.push(format!( - "openhuman_core_version: {}", - yaml_scalar(OPENHUMAN_CORE_VERSION) - )); - lines.push(format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - )); - inserted = true; - } - lines.push(raw.to_string()); - } - - if !inserted { - let insert_at = lines - .iter() - .rposition(|line| line == "---") - .unwrap_or(lines.len()); - lines.insert( - insert_at, - format!( - "openhuman_core_version: {}", - yaml_scalar(OPENHUMAN_CORE_VERSION) - ), - ); - lines.insert( - insert_at + 1, - format!("memory_artifact_format: {}", MEMORY_ARTIFACT_FORMAT), - ); - } - - let mut result = lines.join("\n"); - result.push('\n'); - result -} diff --git a/src/openhuman/memory_store/content/compose/tests.rs b/src/openhuman/memory_store/content/compose/tests.rs deleted file mode 100644 index df7e06bfb..000000000 --- a/src/openhuman/memory_store/content/compose/tests.rs +++ /dev/null @@ -1,622 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::openhuman::memory_store::chunks::types::{Chunk, Metadata, SourceKind, SourceRef}; - use crate::openhuman::memory_store::content::compose::chunk::{ - compose_chunk_file, rewrite_tags, - }; - use crate::openhuman::memory_store::content::compose::summary::{ - compose_summary_md, rewrite_summary_tags, scope_short_label, SummaryComposeInput, - }; - use crate::openhuman::memory_store::content::compose::yaml::{split_front_matter, yaml_scalar}; - use crate::openhuman::memory_store::content::compose::{ - MEMORY_ARTIFACT_FORMAT, OPENHUMAN_CORE_VERSION, - }; - use crate::openhuman::memory_store::content::paths::SummaryTreeKind; - use chrono::TimeZone; - - fn sample_chunk() -> Chunk { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - Chunk { - id: "abc123".into(), - content: "## 2026-01-01T00:00:00Z — alice\nhello world".into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice@example.com".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["person/Alice".into(), "org/Acme".into()], - source_ref: Some(SourceRef::new("slack://m1".to_string())), - path_scope: None, - }, - token_count: 10, - seq_in_source: 0, - created_at: ts, - partial_message: false, - } - } - - #[test] - fn compose_produces_front_matter_and_body() { - let chunk = sample_chunk(); - let (full, body) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - assert!(full_str.starts_with("---\n"), "must start with ---"); - assert!(full_str.contains("source_kind: chat")); - assert!(full_str.contains("source_id: \"slack:#eng\"")); - assert!(full_str.contains("seq: 0")); - assert!(full_str.contains("tags:")); - assert!(full_str.contains(" - person/Alice")); - assert!(full_str.ends_with("hello world")); - assert_eq!( - body, - b"## 2026-01-01T00:00:00Z \xe2\x80\x94 alice\nhello world" - ); - } - - #[test] - fn compose_persists_path_scope_and_seeds_scoped_source_tag() { - let mut chunk = sample_chunk(); - chunk.metadata.source_id = "notion:conn-1:page-123".into(); - chunk.metadata.path_scope = Some("notion:conn-1".into()); - - let (full, _) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - - assert!(full_str.contains("source_id: \"notion:conn-1:page-123\"")); - assert!(full_str.contains("path_scope: \"notion:conn-1\"")); - assert!(full_str.contains(" - source/notion-conn-1")); - assert!(!full_str.contains(" - source/notion-conn-1-page-123")); - } - - #[test] - fn split_front_matter_round_trips() { - let chunk = sample_chunk(); - let (full, body) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - let (fm, b) = split_front_matter(full_str).expect("split must succeed"); - assert!(fm.starts_with("---\n")); - assert!(fm.ends_with("---\n")); - assert_eq!(b.as_bytes(), body.as_slice()); - } - - #[test] - fn rewrite_tags_preserves_body() { - let chunk = sample_chunk(); - let (full, body) = compose_chunk_file(&chunk); - let new_tags = vec!["person/Bob".into(), "project/Phoenix".into()]; - let rewritten = rewrite_tags(&full, &new_tags).unwrap(); - let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); - assert!(rewritten_str.contains(" - person/Bob")); - assert!(!rewritten_str.contains(" - person/Alice")); - // Body must be unchanged. - assert!(rewritten_str.ends_with(std::str::from_utf8(&body).unwrap())); - } - - #[test] - fn rewrite_tags_empty_list() { - let chunk = sample_chunk(); - let (full, _) = compose_chunk_file(&chunk); - let rewritten = rewrite_tags(&full, &[]).unwrap(); - let s = std::str::from_utf8(&rewritten).unwrap(); - assert!(s.contains("tags: []")); - assert!(!s.contains(" - person/")); - } - - #[test] - fn yaml_scalar_quotes_special_characters() { - assert_eq!(yaml_scalar("slack:#eng"), "\"slack:#eng\""); - assert_eq!(yaml_scalar("hello world"), "hello world"); - assert_eq!(yaml_scalar(""), "\"\""); - } - - fn sample_email_chunk() -> Chunk { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - Chunk { - id: "emailchunk1".into(), - content: "---\nFrom: alice@example.com\nSubject: Hello\n\nHello there.".into(), - metadata: Metadata { - source_kind: SourceKind::Email, - source_id: "gmail:alice@example.com|bob@example.com".into(), - owner: "owner@example.com".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["gmail".into()], - source_ref: None, - path_scope: None, - }, - token_count: 15, - seq_in_source: 0, - created_at: ts, - partial_message: false, - } - } - - #[test] - fn email_chunk_has_participants_list_and_alias() { - let chunk = sample_email_chunk(); - let (full, _body) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - // participants block must be a YAML list - assert!( - full_str.contains("participants:"), - "email chunk must have participants field; got:\n{full_str}" - ); - assert!( - full_str.contains(" - alice@example.com"), - "alice must appear as list item; got:\n{full_str}" - ); - assert!( - full_str.contains(" - bob@example.com"), - "bob must appear as list item; got:\n{full_str}" - ); - // aliases block must be present - assert!( - full_str.contains("aliases:"), - "email chunk must have aliases field; got:\n{full_str}" - ); - assert!( - full_str.contains("alice@example.com <-> bob@example.com: chunk 0"), - "alias must encode participants; got:\n{full_str}" - ); - // sender and thread_id must NOT appear - assert!( - !full_str.contains("sender:"), - "email chunk must NOT have sender field; got:\n{full_str}" - ); - assert!( - !full_str.contains("thread_id:"), - "email chunk must NOT have thread_id field; got:\n{full_str}" - ); - } - - #[test] - fn email_chunk_many_participants_alias_summarises() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let chunk = Chunk { - id: "em2".into(), - content: "body".into(), - metadata: Metadata { - source_kind: SourceKind::Email, - source_id: "gmail:alice@x.com|bob@y.com|carol@z.com".into(), - owner: "owner".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: None, - path_scope: None, - }, - token_count: 1, - seq_in_source: 3, - created_at: ts, - partial_message: false, - }; - let (full, _) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - assert!( - full_str.contains("participants:"), - "three-party chunk needs participants list; got:\n{full_str}" - ); - // With 3 participants: first + "2 others" - assert!( - full_str.contains("alice@x.com <-> 2 others: chunk 3"), - "alias with 3 participants must summarise; got:\n{full_str}" - ); - } - - #[test] - fn email_chunk_body_bytes_unchanged_by_extra_fields() { - // Adding participants/aliases to front-matter must not affect body_bytes - // (SHA-256 invariant: the hash is over body only, not front-matter). - let chunk = sample_email_chunk(); - let (full, body) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - // Body must still appear at the end unmodified. - assert!( - full_str.ends_with(std::str::from_utf8(&body).unwrap()), - "body bytes must appear unmodified after front-matter" - ); - // body must equal chunk.content bytes - assert_eq!(body, chunk.content.as_bytes()); - } - - #[test] - fn chat_chunk_has_no_email_specific_fields() { - let chunk = sample_chunk(); // source_kind = Chat - let (full, _) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - assert!( - !full_str.contains("aliases:"), - "chat chunk must not have aliases field" - ); - assert!( - !full_str.contains("participants:"), - "chat chunk must not have participants field" - ); - assert!( - !full_str.contains("sender:"), - "chat chunk must not have sender field" - ); - assert!( - !full_str.contains("thread_id:"), - "chat chunk must not have thread_id field" - ); - } - - #[test] - fn email_chunk_with_malformed_source_id_omits_extra_fields() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let chunk = Chunk { - id: "xyz".into(), - content: "body".into(), - metadata: Metadata { - source_kind: SourceKind::Email, - source_id: "legacysourceid".into(), // no `gmail:` prefix → parse fails - owner: "owner".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: None, - path_scope: None, - }, - token_count: 1, - seq_in_source: 0, - created_at: ts, - partial_message: false, - }; - let (full, _) = compose_chunk_file(&chunk); - let full_str = std::str::from_utf8(&full).unwrap(); - // Malformed source_id → no email extras, no panic. - assert!(!full_str.contains("aliases:")); - assert!(!full_str.contains("participants:")); - assert!(!full_str.contains("sender:")); - } - - // ─── summary compose tests ──────────────────────────────────────────────── - - fn sample_summary_input( - tree_kind: SummaryTreeKind, - scope: &str, - level: u32, - ) -> SummaryComposeInput<'static> { - let ts_start = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let ts_end = chrono::Utc.timestamp_millis_opt(1_700_086_400_000).unwrap(); - let sealed = chrono::Utc.timestamp_millis_opt(1_700_090_000_000).unwrap(); - // Leak the strings so they have 'static lifetime for this test helper. - // Only used in tests, not production code. - let scope: &'static str = Box::leak(scope.to_string().into_boxed_str()); - SummaryComposeInput { - summary_id: "summary:L1:abc", - tree_kind, - tree_id: "tree-id-001", - tree_scope: scope, - level, - child_ids: Box::leak( - vec!["child-1".to_string(), "child-2".to_string()].into_boxed_slice(), - ), - child_basenames: None, - child_count: 2, - time_range_start: ts_start, - time_range_end: ts_end, - sealed_at: sealed, - body: "This is the summariser output.\n", - } - } - - #[test] - fn compose_source_summary_has_required_front_matter() { - let input = sample_summary_input(SummaryTreeKind::Source, "gmail:alice@x.com|bob@y.com", 1); - let composed = compose_summary_md(&input); - let fm = &composed.front_matter; - assert!(fm.starts_with("---\n"), "front-matter must start with ---"); - assert!(fm.ends_with("---\n"), "front-matter must end with ---\\n"); - assert!(fm.contains("kind: summary"), "must have kind: summary"); - assert!( - fm.contains("tree_kind: source"), - "must have tree_kind: source" - ); - assert!(fm.contains("level: 1"), "must have level"); - assert!(fm.contains("child_count: 2"), "must have child_count"); - assert!( - fm.contains(&format!( - "openhuman_core_version: {}", - OPENHUMAN_CORE_VERSION - )), - "must stamp the core version" - ); - assert!( - fm.contains(&format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - )), - "must stamp the artifact format epoch" - ); - assert!( - fm.contains(" - \"[[child-1]]\""), - "must list child ids as Obsidian wikilinks; got:\n{fm}" - ); - assert!( - fm.contains(" - \"[[child-2]]\""), - "must list child ids as Obsidian wikilinks; got:\n{fm}" - ); - assert!( - fm.contains(" - source/"), - "source-tree summary must seed source tag; got:\n{fm}" - ); - // aliases must mention the scope - assert!(fm.contains("aliases:"), "must have aliases"); - assert!( - composed.body == "This is the summariser output.\n", - "body must be the summariser text" - ); - assert!(composed.full.ends_with("This is the summariser output.\n")); - } - - #[test] - fn children_are_emitted_as_obsidian_wikilinks() { - // Contract: every entry in `children:` must be wrapped in `[[…]]` so - // Obsidian's graph view draws a summary→child edge. The YAML scalar is - // quoted because of the leading `[` — both forms below are required. - let input = sample_summary_input(SummaryTreeKind::Source, "gmail:alice@x.com", 1); - let composed = compose_summary_md(&input); - let fm = &composed.front_matter; - for id in ["child-1", "child-2"] { - let expected = format!(" - \"[[{id}]]\""); - assert!( - fm.contains(&expected), - "child id {id} must be emitted as a quoted wikilink ({expected}); got:\n{fm}" - ); - // Belt-and-braces: the bare id must NOT appear as a plain scalar - // (i.e. unwrapped). The wikilink form contains the id, so we - // search for the bare list-item form. - let plain = format!(" - {id}\n"); - assert!( - !fm.contains(&plain), - "child id {id} must not be emitted as a plain scalar; got:\n{fm}" - ); - } - } - - #[test] - fn child_basename_overrides_replace_chunk_id_in_wikilink() { - // L1 seals: each child's wikilink should point at the - // raw archive file basename, not the chunk_id hash. Without - // this override the link would be `[[<32-char hex>]]` and - // Obsidian wouldn't find a matching file (the chunk-store - // copy under `email//...` is gone after the - // raw_refs migration). - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let child_ids = vec!["abc123hash".to_string(), "def456hash".to_string()]; - let overrides: Vec> = vec![ - Some("1700000000000_msg-id-1".into()), - None, // second child has no override → falls back to sanitize_filename - ]; - let input = SummaryComposeInput { - summary_id: "summary:L1:test", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 1, - child_ids: &child_ids, - child_basenames: Some(&overrides), - child_count: 2, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "L1 body", - }; - let composed = compose_summary_md(&input); - let fm = &composed.front_matter; - // First child uses the override (raw archive basename). - assert!( - fm.contains(r#" - "[[1700000000000_msg-id-1]]""#), - "first child must use override basename; got:\n{fm}" - ); - // Second child has None override — fall back to chunk_id. - assert!( - fm.contains(r#" - "[[def456hash]]""#), - "None override must fall back to sanitize_filename; got:\n{fm}" - ); - } - - #[test] - fn structured_child_summary_id_is_sanitised_in_wikilink() { - // Real-world case: an L2 summary lists child L1 summaries by their - // structured id (e.g. `summary:L1:UUID`). Colons are illegal in - // Windows NTFS filenames, so `summary_rel_path` writes the file as - // `summary-L1-UUID.md`. The wikilink target must match that basename - // — i.e. colons must be converted to dashes — otherwise Obsidian - // cannot resolve the link and the graph stays disconnected. - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let child_id = "summary:L1:b9fa5f08-bf79-41a7-a5c8-2d87883d5c01"; - let expected_basename = "summary-L1-b9fa5f08-bf79-41a7-a5c8-2d87883d5c01"; - let input = SummaryComposeInput { - summary_id: "summary:L2:cc9a1224", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 2, - child_ids: &[child_id.to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "L2 body", - }; - let composed = compose_summary_md(&input); - let fm = &composed.front_matter; - let expected = format!(" - \"[[{expected_basename}]]\""); - assert!( - fm.contains(&expected), - "structured child id must be sanitised to filename basename in wikilink; \ - expected line: {expected}; got:\n{fm}" - ); - // Raw colon-bearing id must NOT appear inside `[[…]]` — that wikilink - // would not resolve in Obsidian. - assert!( - !fm.contains(&format!("[[{child_id}]]")), - "raw structured id with colons must not appear inside wikilink; got:\n{fm}" - ); - } - - #[test] - fn compose_global_summary_alias_format() { - let input = sample_summary_input(SummaryTreeKind::Global, "global", 0); - let composed = compose_summary_md(&input); - assert!( - composed.front_matter.contains("tree_kind: global"), - "must have tree_kind: global" - ); - assert!( - composed.front_matter.contains("global digest"), - "alias must mention 'global digest'" - ); - } - - #[test] - fn compose_topic_summary_alias_format() { - let input = sample_summary_input(SummaryTreeKind::Topic, "person:alex-johnson", 1); - let composed = compose_summary_md(&input); - assert!( - composed.front_matter.contains("tree_kind: topic"), - "must have tree_kind: topic" - ); - assert!( - composed.front_matter.contains("topic"), - "alias must mention topic entity" - ); - } - - #[test] - fn compose_summary_with_zero_children() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "summary:L0:empty", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 0, - child_ids: &[], - child_basenames: None, - child_count: 0, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "empty", - }; - let composed = compose_summary_md(&input); - assert!(composed.front_matter.contains("children: []")); - assert!(composed.front_matter.contains("child_count: 0")); - } - - #[test] - fn compose_summary_same_start_end_date_single_date_alias() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "summary:L1:sameday", - tree_kind: SummaryTreeKind::Global, - tree_id: "t1", - tree_scope: "global", - level: 1, - child_ids: &["child-a".to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, // same as start - sealed_at: ts, - body: "day recap", - }; - let composed = compose_summary_md(&input); - // Alias must contain just one date, not "date–date" - let alias_line = composed - .front_matter - .lines() - .find(|l| l.contains("L1") && l.contains("global digest")) - .expect("alias line must be present"); - // The date should appear exactly once (no en-dash range) - let date_str = ts.format("%Y-%m-%d").to_string(); - assert!( - alias_line.contains(&date_str), - "alias must contain the date; got: {alias_line}" - ); - // Must not contain an en-dash (range indicator) - assert!( - !alias_line.contains('\u{2013}'), - "same-day alias must not have en-dash range; got: {alias_line}" - ); - } - - #[test] - fn scope_short_label_two_participants() { - let label = scope_short_label("gmail:alice@x.com|bob@y.com"); - assert_eq!(label, "alice@x.com \u{2194} bob@y.com"); - } - - #[test] - fn scope_short_label_many_participants() { - let label = scope_short_label("gmail:alice@x.com|bob@y.com|carol@z.com"); - assert_eq!(label, "alice@x.com + 2 others"); - } - - #[test] - fn scope_short_label_non_gmail_returns_raw() { - let label = scope_short_label("slack:#general"); - assert_eq!(label, "slack:#general"); - } - - #[test] - fn rewrite_summary_tags_delegates_to_rewrite_tags() { - // compose a summary, then rewrite its tags — body must stay unchanged. - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "sum:L1:rwttest", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 1, - child_ids: &["c1".to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "summary body text", - }; - let composed = compose_summary_md(&input); - let file_bytes = composed.full.as_bytes(); - let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()]; - let rewritten = rewrite_summary_tags(file_bytes, &new_tags).unwrap(); - let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); - assert!(rewritten_str.contains(" - person/Alice-Smith")); - assert!(rewritten_str.contains(" - topic/Memory")); - assert!(!rewritten_str.contains("tags: []")); - assert!(rewritten_str.contains(&format!( - "openhuman_core_version: {}", - OPENHUMAN_CORE_VERSION - ))); - assert!(rewritten_str.contains(&format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - ))); - // Body must be unchanged - assert!(rewritten_str.ends_with("summary body text")); - } - - #[test] - fn rewrite_summary_tags_backfills_missing_provenance() { - let file = - b"---\nid: legacy\nkind: summary\ntags: []\naliases:\n - legacy\n---\nlegacy body"; - let rewritten = rewrite_summary_tags(file, &["person/Alice".to_string()]).unwrap(); - let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); - assert!(rewritten_str.contains(&format!( - "openhuman_core_version: {}", - OPENHUMAN_CORE_VERSION - ))); - assert!(rewritten_str.contains(&format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - ))); - assert!(rewritten_str.ends_with("legacy body")); - } -} diff --git a/src/openhuman/memory_store/content/compose/yaml.rs b/src/openhuman/memory_store/content/compose/yaml.rs deleted file mode 100644 index 05236454f..000000000 --- a/src/openhuman/memory_store/content/compose/yaml.rs +++ /dev/null @@ -1,149 +0,0 @@ -//! YAML scalar helpers and front-matter parsing utilities. - -/// Build the canonical Obsidian `source/` tag for a given -/// source scope. Used to seed the `tags:` block on every chunk and -/// every source-tree summary so the Obsidian graph view can filter by -/// source. -/// -/// Slug rules match `slugify_source_id` (lowercase ASCII, `-` separators, -/// alphanumerics + `_` preserved) so the tag matches the on-disk -/// `raw//...` directory name byte-for-byte. -pub fn source_tag(scope: &str) -> String { - use crate::openhuman::memory_store::content::paths::slugify_source_id; - format!("source/{}", slugify_source_id(scope)) -} - -/// Prepend the source tag to `tags`, dedup, and return the new list. -/// Order is preserved otherwise — `source/...` always comes first so -/// it shows up at the top of the YAML block. -pub fn with_source_tag(scope: &str, tags: &[String]) -> Vec { - let st = source_tag(scope); - let mut out = Vec::with_capacity(tags.len() + 1); - out.push(st.clone()); - for t in tags { - if t != &st { - out.push(t.clone()); - } - } - out -} - -/// Parse the value of a top-level YAML scalar field (e.g. `source_id`, -/// `tree_scope`, `tree_kind`) from a frontmatter string. Strips -/// surrounding double-quotes if present so the returned slice matches -/// what the original composer passed in. Returns `None` if the key is -/// not present at the top level of the frontmatter. -pub fn scan_fm_field(fm: &str, key: &str) -> Option { - let prefix = format!("{key}: "); - for raw in fm.lines() { - // Skip indented lines (those are list items / nested mappings). - if raw.starts_with(' ') || raw.starts_with('\t') { - continue; - } - if let Some(rest) = raw.strip_prefix(&prefix) { - let trimmed = rest.trim(); - if let Some(inner) = trimmed.strip_prefix('"').and_then(|s| s.strip_suffix('"')) { - return Some(inner.replace("\\\"", "\"").replace("\\\\", "\\")); - } - return Some(trimmed.to_string()); - } - } - None -} - -/// Split a file into `(front_matter, body)` at the second `---` delimiter. -/// -/// Returns `None` if the file does not have the expected `---\n...\n---\n` form. -pub fn split_front_matter(content: &str) -> Option<(&str, &str)> { - // The file must start with `---\n`. - if !content.starts_with("---\n") { - return None; - } - // Find the closing `---` line (must be `---` alone on a line after the first line). - let rest = &content[4..]; // skip the opening `---\n` - let close_idx = rest.find("\n---\n").or_else(|| { - // Could be at the very end (no body). - rest.strip_suffix("\n---").map(|r| r.len()) - })?; - let fm_end = 4 + close_idx + 5; // include `\n---\n` - debug_assert!(content.is_char_boundary(fm_end)); - Some((&content[..fm_end], &content[fm_end..])) -} - -/// Format a string as an unquoted YAML scalar when safe, or as a -/// double-quoted string when it contains special characters. -/// -/// We conservatively quote strings containing `:`, `#`, `[`, `]`, `{`, `}`, -/// `"`, `'`, `\`, leading/trailing whitespace, or that start with special -/// YAML indicator characters. -/// -/// Newlines, carriage returns and tabs are collapsed to a single space first: -/// front-matter scalars are single-line identifiers/paths, and a raw newline in -/// a value would inject a spurious `\n---\n` into the outer front matter, which -/// the reader's `split_front_matter` would then mistake for the closing -/// delimiter — corrupting the body boundary and its sha (#4689). Collapsing is -/// lossless for the identifier fields that flow through here (they never -/// legitimately contain control whitespace). -pub fn yaml_scalar(s: &str) -> String { - let sanitized: String = s - .chars() - .map(|c| { - if matches!(c, '\n' | '\r' | '\t') { - ' ' - } else { - c - } - }) - .collect(); - let s = sanitized.as_str(); - - let needs_quoting = s.is_empty() - || s.trim() != s - || s.starts_with(|c: char| { - matches!( - c, - '&' | '*' | '?' | '|' | '-' | '<' | '>' | '=' | '!' | '%' | '@' | '`' - ) - }) - || s.contains([':', '#', '[', ']', '{', '}', '"', '\'']); - - if needs_quoting { - let escaped = s.replace('\\', "\\\\").replace('"', "\\\""); - format!("\"{escaped}\"") - } else { - s.to_string() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn yaml_scalar_plain_value_is_unquoted() { - assert_eq!(yaml_scalar("hello"), "hello"); - } - - #[test] - fn yaml_scalar_collapses_newlines_and_blocks_fm_injection() { - // A field value carrying an embedded `\n---\n` must not emit a raw - // newline, or it would inject a spurious front-matter closer (#4689). - let out = yaml_scalar("vault:evil\n---\ninjected"); - assert!(!out.contains('\n'), "no raw newline: {out}"); - assert!(!out.contains('\r')); - - // A composed front matter using the sanitized value still splits at the - // real closer, leaving the body intact. - let fm = format!("---\nsource_id: {out}\n---\nBODY"); - let (_, body) = split_front_matter(&fm).expect("front matter splits"); - assert_eq!(body, "BODY"); - } - - #[test] - fn yaml_scalar_collapses_tabs_and_carriage_returns() { - let out = yaml_scalar("a\tb\r\nc"); - assert!(!out.contains('\t')); - assert!(!out.contains('\r')); - assert!(!out.contains('\n')); - } -} diff --git a/src/openhuman/memory_store/content/mod.rs b/src/openhuman/memory_store/content/mod.rs index e9014340d..0032c6be0 100644 --- a/src/openhuman/memory_store/content/mod.rs +++ b/src/openhuman/memory_store/content/mod.rs @@ -12,37 +12,16 @@ //! - [`read`] — read + SHA-256 verification + `split_front_matter`; summary variants //! - [`tags`] — `update_chunk_tags` + `update_summary_tags` + slugifiers -pub mod atomic; -pub mod compose; pub mod obsidian; pub mod obsidian_registry; -pub mod paths; -pub mod raw; pub mod read; pub mod tags; pub mod wiki_git; -use std::path::Path; - -use crate::openhuman::memory_store::chunks::types::Chunk; - -pub use atomic::StagedSummary; -pub use compose::SummaryComposeInput; -pub use paths::SummaryTreeKind; - -/// A chunk that has been written to disk and is ready for SQLite upsert. -/// -/// Callers build a `Vec` from `stage_chunks`, then pass it to -/// `store::upsert_chunks_tx` in the same SQLite transaction. -#[derive(Debug, Clone)] -pub struct StagedChunk { - /// The original chunk (metadata + content). - pub chunk: Chunk, - /// Relative content path (forward-slash, e.g. `"chat/slack-eng/0.md"`). - pub content_path: String, - /// SHA-256 hex digest over the body bytes only. - pub content_sha256: String, -} +pub use tinycortex::memory::chunks::StagedChunk; +pub use tinycortex::memory::store::content::{ + atomic, compose, paths, raw, stage_chunks, StagedSummary, SummaryComposeInput, SummaryTreeKind, +}; /// Update the `tags:` block in a summary's on-disk `.md` file after an /// extraction job runs. @@ -54,170 +33,3 @@ pub fn update_summary_tags( ) -> anyhow::Result<()> { tags::update_summary_tags(config, summary_id) } - -/// Write all chunks in `chunks` to disk and return `StagedChunk` records -/// ready for SQLite upsert. -/// -/// Each chunk file is written atomically via a sibling temp-file + rename. -/// Already-existing files are skipped (immutable-body contract). Parent -/// directories are created on demand. -/// -/// **Email chunks skip the disk write.** Their content already lives in -/// the per-message raw archive at `/raw//_.md`, -/// so a parallel copy in `/email//.md` -/// would just duplicate bytes and clutter the Obsidian vault. We still -/// emit a `StagedChunk` row with an empty `content_path` so the SQLite -/// upsert proceeds — read paths fall back to the chunk's truncated SQL -/// `content` column or to the raw archive when they need full bodies. -/// -/// `content_root` — absolute path to the root of the content store. -pub fn stage_chunks(content_root: &Path, chunks: &[Chunk]) -> anyhow::Result> { - use crate::openhuman::memory_store::chunks::types::SourceKind; - let mut staged = Vec::with_capacity(chunks.len()); - - for chunk in chunks { - if chunk.metadata.source_kind == SourceKind::Email { - // Body lives in raw//_.md — no chunk file. - staged.push(StagedChunk { - chunk: chunk.clone(), - content_path: String::new(), - content_sha256: String::new(), - }); - continue; - } - - let source_kind = chunk.metadata.source_kind.as_str(); - let path_id = chunk - .metadata - .path_scope - .as_deref() - .unwrap_or(&chunk.metadata.source_id); - - let rel_path = paths::chunk_rel_path(source_kind, path_id, &chunk.id); - let abs_path = paths::chunk_abs_path(content_root, source_kind, path_id, &chunk.id); - - let (full_bytes, body_bytes) = compose::compose_chunk_file(chunk); - let sha256 = atomic::sha256_hex(&body_bytes); - - // Self-heal a stale/drifted on-disk file so the recorded content_sha256 - // always matches the bytes actually on disk (#4689). A plain write_if_new - // would skip a pre-existing file at this path while we still record the - // freshly-composed sha, permanently diverging the DB token from disk and - // forcing read_chunk_body to serve the ≤500-char preview. - if let Err(e) = atomic::write_or_replace_body(&abs_path, &full_bytes, &sha256) { - log::error!( - "[content_store] failed to write chunk {} to {}: {e}", - chunk.id, - rel_path - ); - return Err(e); - } - log::debug!("[content_store] staged chunk {} → {}", chunk.id, rel_path); - - staged.push(StagedChunk { - chunk: chunk.clone(), - content_path: rel_path, - content_sha256: sha256, - }); - } - - Ok(staged) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::chunks::types::{Metadata, SourceKind}; - use chrono::TimeZone; - use tempfile::TempDir; - - fn sample_chunk(seq: u32) -> Chunk { - let ts = chrono::Utc - .timestamp_millis_opt(1_700_000_000_000 + seq as i64) - .unwrap(); - Chunk { - id: format!("chunk_{seq}"), - content: format!("## ts — alice\nMessage {seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: None, - path_scope: None, - }, - token_count: 5, - seq_in_source: seq, - created_at: ts, - partial_message: false, - } - } - - #[test] - fn stage_chunks_writes_files_and_returns_staged() { - let dir = TempDir::new().unwrap(); - let chunks = vec![sample_chunk(0), sample_chunk(1)]; - let staged = stage_chunks(dir.path(), &chunks).unwrap(); - - assert_eq!(staged.len(), 2); - for s in &staged { - let abs = paths::chunk_abs_path( - dir.path(), - s.chunk.metadata.source_kind.as_str(), - &s.chunk.metadata.source_id, - &s.chunk.id, - ); - assert!(abs.exists(), "file must exist: {}", abs.display()); - assert!(!s.content_path.is_empty()); - assert_eq!(s.content_sha256.len(), 64); - // Path must be relative with forward slashes. - assert!(!s.content_path.starts_with('/')); - assert!(s.content_path.contains('/')); - } - } - - #[test] - fn stage_chunks_is_idempotent() { - let dir = TempDir::new().unwrap(); - let chunks = vec![sample_chunk(0)]; - let first = stage_chunks(dir.path(), &chunks).unwrap(); - let second = stage_chunks(dir.path(), &chunks).unwrap(); - assert_eq!(first[0].content_sha256, second[0].content_sha256); - assert_eq!(first[0].content_path, second[0].content_path); - } - - #[test] - fn stage_chunks_overwrites_stale_on_disk_body() { - let dir = TempDir::new().unwrap(); - let chunk = sample_chunk(0); - - // Pre-write a stale file at the chunk's content path with a different - // body, so write_if_new would otherwise skip and leave it in place while - // stage_chunks records the fresh body's sha — the #4689 divergence. - 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: 1\n---\nSTALE BODY").unwrap(); - - let staged = stage_chunks(dir.path(), std::slice::from_ref(&chunk)).unwrap(); - - // The file now holds the freshly-composed body, and the recorded sha - // matches the bytes actually on disk. - let on_disk = std::fs::read_to_string(&abs).unwrap(); - assert!( - on_disk.ends_with(&chunk.content), - "body rewritten: {on_disk}" - ); - let (_, body) = super::compose::split_front_matter(&on_disk).unwrap(); - assert_eq!( - staged[0].content_sha256, - atomic::sha256_hex(body.as_bytes()) - ); - } -} diff --git a/src/openhuman/memory_store/content/paths.rs b/src/openhuman/memory_store/content/paths.rs deleted file mode 100644 index dfae52dc0..000000000 --- a/src/openhuman/memory_store/content/paths.rs +++ /dev/null @@ -1,744 +0,0 @@ -//! Content-file path generation. -//! -//! Each chunk body is stored as a `.md` file under `/`. The path -//! structure depends on the source kind: -//! -//! ```text -//! Email: /email//.md -//! Chat: /chat//.md -//! Document: /document//.md -//! ``` -//! -//! Email paths parse `source_id` as `gmail:{participants}` where `participants` -//! is `addr1|addr2|...` (sorted, deduped, lowercased bare emails). The -//! participants string is slugified as a whole (pipe and `@` both become `-`) -//! to produce a single directory level, giving one folder per unique -//! conversation set. -//! -//! Paths are stored in SQLite as **relative** strings with forward slashes so -//! they remain valid regardless of where the workspace is mounted. - -use std::path::{Path, PathBuf}; - -use crate::openhuman::memory::util::redact::redact; - -/// Which kind of summary tree a summary belongs to. Determines the -/// folder name under `/wiki/summaries/` — flattened -/// from the original `//...` two-level layout to a -/// single dash-joined `-/...` folder so the -/// Obsidian sidebar listing stays readable. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum SummaryTreeKind { - /// Per-source-tree summary. Layout: `wiki/summaries/source-/L/.md` - Source, - /// Global digest tree — the singleton cross-source activity tree. - /// Layout: `wiki/summaries/global/L/.md`. There is exactly - /// one global tree, so (unlike the historical `global-/` - /// layout) the day/week/month/year grouping is expressed by the - /// `L/` subdirectory + each node's time range, NOT by a - /// date-stamped top-level folder. A per-day folder name shattered the - /// one logical tree into one look-alike folder per calendar day. - Global, - /// Per-topic (entity) tree. Layout: `wiki/summaries/topic-/L/.md` - Topic, -} - -/// Top-level directory for derived/wiki content (summaries today, -/// contacts and other knowledge-graph notes later). The two-tier -/// `/raw/` (verbatim source bytes) + -/// `/wiki/` (processed, human-facing) split lets users -/// keep one tidy Obsidian vault rooted at `` without -/// chunked intermediates polluting the listing. -pub const WIKI_PREFIX: &str = "wiki"; - -/// Build the relative content path for a summary, using forward slashes. -/// -/// Path layout depends on tree_kind. Folder name is `-` — -/// flattening the historical two-level `//` so users see -/// one folder per logical source in their Obsidian sidebar: -/// - Source: `"wiki/summaries/source-/L/.md"` -/// - Global: `"wiki/summaries/global/L/.md"` — one -/// folder for the singleton global tree; the temporal grouping lives in -/// the `L/` subdirectory, not the folder name. -/// - Topic: `"wiki/summaries/topic-/L/.md"` -/// -/// `scope_slug` must already be slugified by the caller (use [`slugify_source_id`] or -/// a per-kind variant). A trailing `.md` on `summary_id` is stripped if present. -/// -/// New summaries use the explicit basename contract implemented by -/// [`summary_filename`]: -/// - current canonical ids: `summary:{13-digit-ms}:L{level}-{tail}` -/// → `summary-{13-digit-ms}-L{level}-{tail}.md` -/// - legacy ids: `summary:L{level}:{rest}` -/// → `summary-L{level}-{rest}.md` -/// -/// Unknown / malformed ids fall back to [`sanitize_filename`] so existing vaults -/// remain readable even if they contain older experimental shapes. -pub fn summary_rel_path( - tree_kind: SummaryTreeKind, - scope_slug: &str, - level: u32, - summary_id: &str, -) -> String { - let filename = summary_filename(summary_id); - - match tree_kind { - SummaryTreeKind::Source => { - format!( - "{WIKI_PREFIX}/summaries/source-{}/L{}/{}.md", - scope_slug, level, filename - ) - } - SummaryTreeKind::Global => { - // The global tree is a singleton: one folder, with the - // day/week/month/year hierarchy carried by `L/` and the - // node's time range — never a per-day folder name. - format!("{WIKI_PREFIX}/summaries/global/L{}/{}.md", level, filename) - } - SummaryTreeKind::Topic => { - format!( - "{WIKI_PREFIX}/summaries/topic-{}/L{}/{}.md", - scope_slug, level, filename - ) - } - } -} - -/// On-disk placement for a summary node within a document **source** tree. -/// -/// Document source trees (Notion) keep one folder per connection but nest -/// per-document subtrees and the cross-document merge tier beneath it, so the -/// vault mirrors the logical shape: `notion` → `docs//v-` → -/// `merge`. Non-document trees (chat/email) and the `Standard` variant use -/// the flat `source-/L/` layout unchanged. -#[derive(Clone, Copy, Debug)] -pub enum SummaryDiskLayout<'a> { - /// Flat layout — `source-/L/…` (chat, email, legacy). - Standard, - /// A node inside one document's versioned subtree — - /// `source-/docs//v-/L/…`. - DocSubtree { - doc_slug: &'a str, - version_ms: Option, - }, - /// A cross-document merge-tier node — `source-/merge/L/…`. - Merge, -} - -/// Layout-aware variant of [`summary_rel_path`]. For document source trees it -/// routes per-doc and merge nodes into nested folders; for everything else -/// (and [`SummaryDiskLayout::Standard`]) it is identical to -/// [`summary_rel_path`]. -pub fn summary_rel_path_with_layout( - tree_kind: SummaryTreeKind, - scope_slug: &str, - level: u32, - summary_id: &str, - layout: SummaryDiskLayout<'_>, -) -> String { - match (tree_kind, layout) { - ( - SummaryTreeKind::Source, - SummaryDiskLayout::DocSubtree { - doc_slug, - version_ms, - }, - ) => { - let filename = summary_filename(summary_id); - let safe_slug = slugify_source_id(doc_slug); - let vfolder = match version_ms { - Some(v) => format!("v-{v}"), - None => "v-unversioned".to_string(), - }; - format!( - "{WIKI_PREFIX}/summaries/source-{scope_slug}/docs/{safe_slug}/{vfolder}/L{level}/{filename}.md" - ) - } - (SummaryTreeKind::Source, SummaryDiskLayout::Merge) => { - let filename = summary_filename(summary_id); - format!("{WIKI_PREFIX}/summaries/source-{scope_slug}/merge/L{level}/{filename}.md") - } - // Standard layout, or a non-Source tree kind — fall back to flat. - _ => summary_rel_path(tree_kind, scope_slug, level, summary_id), - } -} - -/// Convert a summary id into the canonical on-disk basename stem (without -/// `.md`). -/// -/// This keeps summary filenames independent from the generic "replace illegal -/// characters" fallback so new writes follow one documented convention, while -/// legacy ids still map to their historical names. -pub(crate) fn summary_filename(summary_id: &str) -> String { - let id = summary_id.strip_suffix(".md").unwrap_or(summary_id); - - if let Some(rest) = id.strip_prefix("summary:") { - if let Some((ms, suffix)) = rest.split_once(':') { - // Canonical fast-path: only accept ms-first ids whose - // `L-` suffix has a numeric level and a tail - // free of filesystem-illegal characters. Without the tail - // check a malformed canonical-looking id like - // `summary:1700000000000:L2-a/b` would smuggle a `/` into - // the basename and split the file across multiple path - // components when joined onto the L/ directory. - if let Some((level, tail)) = suffix.split_once('-') { - let level_is_numeric = level.starts_with('L') - && level.len() > 1 - && level[1..].chars().all(|c| c.is_ascii_digit()); - let tail_is_safe = !tail.is_empty() - && !tail - .chars() - .any(|c| matches!(c, '\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|')); - if ms.len() == 13 - && ms.chars().all(|c| c.is_ascii_digit()) - && level_is_numeric - && tail_is_safe - { - return format!("summary-{ms}-{level}-{tail}"); - } - } - } - - if let Some((level, tail)) = rest.split_once(':') { - // Legacy ms-less ids (`summary:L:`). Require strict - // `L` for the level segment so a malicious input - // like `summary:L1/2:abc` or `summary:L../../x:tail` cannot - // inject `/` into the returned basename. Anything else - // falls through to the generic sanitiser below. - let level_is_numeric = level.starts_with('L') - && level.len() > 1 - && level[1..].chars().all(|c| c.is_ascii_digit()); - if level_is_numeric && !tail.is_empty() { - return format!("summary-{level}-{}", sanitize_filename(tail)); - } - } - } - - sanitize_filename(id) -} - -/// Replace characters that are illegal in filenames on Windows NTFS with `-`. -/// -/// Illegal characters: `\`, `/`, `:`, `*`, `?`, `"`, `<`, `>`, `|`. -/// (Forward slash is not replaced since `summary_id` should not contain path -/// separators, but we sanitize it anyway for safety.) -/// -/// Exposed at crate scope so [`super::compose`] can convert structured IDs -/// like `summary:L1:UUID` into the basename used by [`summary_rel_path`] -/// when no summary-specific mapping applies. Summary ids should prefer -/// [`summary_filename`] so new writes follow the documented basename -/// contract instead of relying on punctuation replacement as an accident of -/// the implementation. -pub(crate) fn sanitize_filename(s: &str) -> String { - s.chars() - .map(|c| match c { - '\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '-', - c => c, - }) - .collect() -} - -/// Build the absolute on-disk path for a summary given the content root. -pub fn summary_abs_path( - content_root: &Path, - tree_kind: SummaryTreeKind, - scope_slug: &str, - level: u32, - summary_id: &str, -) -> PathBuf { - let rel = summary_rel_path(tree_kind, scope_slug, level, summary_id); - let mut abs = content_root.to_path_buf(); - for component in rel.split('/') { - abs.push(component); - } - abs -} - -/// Build the relative content path for a chunk, using forward slashes. -/// -/// Path layout depends on source_kind: -/// - Email: `"email//.md"` -/// Parses `source_id` as `gmail:{participants}` (two colon-separated parts) -/// where `participants` is `addr1|addr2|...` (sorted, deduped, lowercased). -/// The entire participants string is slugified as a single unit to produce -/// one folder level per conversation set (no nested thread subfolder). -/// If the source_id lacks a `gmail:` prefix or has no participants segment, -/// falls through to the chat/document layout using `slugify_source_id(source_id)`. -/// - Chat: `"chat//.md"` -/// - Document: `"document//.md"` -/// -/// `chunk_id` — the deterministic content hash produced by `types::chunk_id`. -/// -/// # Examples -/// -/// ```text -/// chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc") -/// → "email/alice-x-com-bob-y-com/abc.md" -/// -/// chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def") -/// → "email/notifications-github-com-sanil-x-com/def.md" -/// -/// chunk_rel_path("email", "legacyid", "xyz") -/// → "email/legacyid/xyz.md" (malformed — flat fallback) -/// ``` -pub fn chunk_rel_path(source_kind: &str, source_id: &str, chunk_id: &str) -> String { - // Sanitize chunk_id into a cross-platform filename. Chunk IDs contain - // colons (e.g. `chat:slack:#eng:0`) which are illegal on Windows NTFS; - // replace illegal characters with `-` to match summary_rel_path behaviour. - let filename = sanitize_filename(chunk_id); - match source_kind { - "email" => { - // Expected format: "gmail:{participants}" - // Split on ':' — exactly 2 parts required; part[0] == "gmail". - let parts: Vec<&str> = source_id.splitn(2, ':').collect(); - if parts.len() == 2 && parts[0] == "gmail" && !parts[1].is_empty() { - let participants_slug = slugify_source_id(parts[1]); - format!("email/{}/{}.md", participants_slug, filename) - } else { - // Malformed / legacy source_id — fall back to flat layout. - // Redact the source_id before logging since it may embed email - // addresses. - log::debug!( - "[content_store::paths] email source_id has unexpected format, falling back to flat layout: source_id_hash={}", - redact(source_id) - ); - let slug = slugify_source_id(source_id); - format!("email/{}/{}.md", slug, filename) - } - } - _ => { - // Chat, Document, and any future kinds use a 3-level layout. - let slug = slugify_source_id(source_id); - format!("{}/{}/{}.md", source_kind, slug, filename) - } - } -} - -/// Build the absolute on-disk path for a chunk given the content root. -pub fn chunk_abs_path( - content_root: &Path, - source_kind: &str, - source_id: &str, - chunk_id: &str, -) -> PathBuf { - let rel = chunk_rel_path(source_kind, source_id, chunk_id); - // Convert forward-slash relative path to OS-native path. - let mut abs = content_root.to_path_buf(); - for component in rel.split('/') { - abs.push(component); - } - abs -} - -/// Convert a raw `source_id` (e.g. `"slack:#general"`, `"gmail:thread/abc"`) -/// into a filesystem-safe slug using only `[a-z0-9_-]` characters. -/// -/// Rules: -/// - lowercase the whole string -/// - replace any character outside `[a-z0-9_-]` with `-` -/// - collapse consecutive `-` to one -/// - trim leading/trailing `-` -/// - `_` is preserved anywhere in the string (interior underscores are kept) -/// - truncate to 120 characters -pub fn slugify_source_id(source_id: &str) -> String { - let lower = source_id.to_lowercase(); - let mut out = String::with_capacity(lower.len().min(120)); - let mut last_dash = true; // avoids leading dash; also suppresses leading underscore runs - let mut pending_underscore = false; // deferred `_` to avoid leading underscore - - for ch in lower.chars() { - if ch == '_' { - // Defer underscores — emit only if we have already emitted a - // non-separator character (so `_solo_` becomes `_solo_` once the - // `s` is emitted, but a leading `_` is dropped). - if !last_dash { - // We have real content before this, so emit the underscore now. - pending_underscore = true; - } - // If last_dash is true (nothing emitted yet), silently skip. - } else if ch.is_ascii_alphanumeric() { - if pending_underscore { - out.push('_'); - pending_underscore = false; - } - out.push(ch); - last_dash = false; - } else { - // Non-alphanumeric, non-underscore → convert to `-`. - pending_underscore = false; // drop any pending underscore before a dash - if !last_dash { - out.push('-'); - last_dash = true; - } - } - } - // trailing underscore: drop it (trim trailing separators). - // trim trailing dash - let trimmed = out.trim_end_matches('-'); - // also trim any trailing underscore - let trimmed = trimmed.trim_end_matches('_'); - let truncated = truncate_at_char(trimmed, 120); - if truncated.is_empty() { - "unknown".to_string() - } else { - truncated.to_string() - } -} - -/// Truncate `s` to at most `max_chars` Unicode code points. -fn truncate_at_char(s: &str, max_chars: usize) -> &str { - match s.char_indices().nth(max_chars) { - Some((idx, _)) => &s[..idx], - None => s, - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // ─── doc-aware layout tests ─────────────────────────────────────────────── - - #[test] - fn layout_doc_subtree_nests_under_docs_and_version() { - let p = summary_rel_path_with_layout( - SummaryTreeKind::Source, - "notion-conn1", - 2, - "summary:1700000000000:L2-deadbeef", - SummaryDiskLayout::DocSubtree { - doc_slug: "notion-conn1-pageA", - version_ms: Some(1717500000000), - }, - ); - assert_eq!( - p, - "wiki/summaries/source-notion-conn1/docs/notion-conn1-pagea/v-1717500000000/L2/summary-1700000000000-L2-deadbeef.md" - ); - } - - #[test] - fn layout_doc_subtree_unversioned_folder() { - let p = summary_rel_path_with_layout( - SummaryTreeKind::Source, - "notion-conn1", - 1, - "summary:1700000000000:L1-abcd0000", - SummaryDiskLayout::DocSubtree { - doc_slug: "notion-conn1-pageB", - version_ms: None, - }, - ); - assert!( - p.contains("/docs/notion-conn1-pageb/v-unversioned/L1/"), - "got {p}" - ); - } - - #[test] - fn layout_merge_tier_nests_under_merge() { - let p = summary_rel_path_with_layout( - SummaryTreeKind::Source, - "notion-conn1", - 1000, - "summary:1700000000000:L1000-aaaa1111", - SummaryDiskLayout::Merge, - ); - assert_eq!( - p, - "wiki/summaries/source-notion-conn1/merge/L1000/summary-1700000000000-L1000-aaaa1111.md" - ); - } - - #[test] - fn layout_standard_matches_flat_path() { - let id = "summary:1700000000000:L1-cccc2222"; - let flat = summary_rel_path(SummaryTreeKind::Source, "slack-eng", 1, id); - let std_layout = summary_rel_path_with_layout( - SummaryTreeKind::Source, - "slack-eng", - 1, - id, - SummaryDiskLayout::Standard, - ); - assert_eq!(flat, std_layout); - } - - // ─── slugify tests ──────────────────────────────────────────────────────── - - #[test] - fn slugify_slack_channel() { - assert_eq!(slugify_source_id("slack:#general"), "slack-general"); - } - - #[test] - fn slugify_gmail_thread() { - assert_eq!( - slugify_source_id("gmail:thread/abc-123"), - "gmail-thread-abc-123" - ); - } - - #[test] - fn slugify_collapses_consecutive_separators() { - assert_eq!(slugify_source_id("foo::bar"), "foo-bar"); - } - - #[test] - fn slugify_uppercase_lowercased() { - assert_eq!(slugify_source_id("Slack:ABC"), "slack-abc"); - } - - #[test] - fn slugify_empty_falls_back_to_unknown() { - assert_eq!(slugify_source_id(""), "unknown"); - assert_eq!(slugify_source_id(":::"), "unknown"); - } - - #[test] - fn slugify_truncates_at_120_chars() { - let long = "a".repeat(200); - let slug = slugify_source_id(&long); - assert_eq!(slug.len(), 120); - } - - #[test] - fn slugify_preserves_interior_underscore() { - // `_solo_` has a leading and trailing underscore; only the interior - // `solo` + the part after should survive. When used as a thread key - // it arrives as the whole string `_solo_`. - // Leading `_` is stripped (it's treated like a leading dash), - // trailing `_` is stripped; interior `_` is preserved when sandwiched - // between alphanumeric characters. - let s = slugify_source_id("_solo_"); - // "solo" — both outer underscores trimmed, interior underscore has - // nothing on the right so it's also trailing and trimmed. - assert_eq!(s, "solo"); - } - - #[test] - fn slugify_preserves_interior_underscore_between_chars() { - // `foo_bar` — interior underscore stays. - assert_eq!(slugify_source_id("foo_bar"), "foo_bar"); - } - - // ─── chunk_rel_path tests ───────────────────────────────────────────────── - - #[test] - fn email_one_to_one_conversation_path() { - // 1:1 conversation between alice and bob. - let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc"); - assert_eq!(p, "email/alice-x-com-bob-y-com/abc.md"); - } - - #[test] - fn email_group_conversation_path() { - // Group conversation with three participants. - let p = chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def"); - assert_eq!(p, "email/notifications-github-com-sanil-x-com/def.md"); - } - - #[test] - fn email_solo_no_to_path() { - // Solo sender (no To), participants = single address. - let p = chunk_rel_path("email", "gmail:alice@x.com", "solo123"); - assert_eq!(p, "email/alice-x-com/solo123.md"); - } - - #[test] - fn email_malformed_source_id_falls_back_to_flat_layout() { - // Malformed: no `gmail:` prefix → flat fallback. - let p = chunk_rel_path("email", "legacyid", "xyz"); - // Falls back to email//.md - assert!(p.starts_with("email/"), "must remain under email/"); - assert!(p.ends_with("/xyz.md"), "chunk_id must be the filename"); - // Must not panic. - } - - #[test] - fn email_three_participant_path() { - // Three participants: alice, bob, carol (pipe-separated, sorted). - let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com|carol@z.com", "g42"); - assert_eq!(p, "email/alice-x-com-bob-y-com-carol-z-com/g42.md"); - } - - #[test] - fn chat_path() { - let p = chunk_rel_path("chat", "slack:#eng", "xyz789"); - assert_eq!(p, "chat/slack-eng/xyz789.md"); - } - - #[test] - fn document_path() { - let p = chunk_rel_path("document", "doc:notes.md", "uvw"); - assert_eq!(p, "document/doc-notes-md/uvw.md"); - } - - #[test] - fn chunk_abs_path_uses_os_separator() { - use std::path::Path; - let root = Path::new("/workspace/content"); - let abs = chunk_abs_path(root, "email", "gmail:alice@x.com|bob@y.com", "abc"); - assert!(abs.starts_with(root)); - assert!(abs.ends_with("abc.md")); - } - - // ─── summary_rel_path tests ─────────────────────────────────────────────── - - #[test] - fn summary_rel_path_source() { - let p = summary_rel_path( - SummaryTreeKind::Source, - "gmail-alice-x-com-bob-y-com", - 1, - "summary:L1:abc", - ); - // Colons in summary_id are replaced with '-' for cross-platform filenames. - assert_eq!( - p, - "wiki/summaries/source-gmail-alice-x-com-bob-y-com/L1/summary-L1-abc.md" - ); - } - - #[test] - fn summary_rel_path_current_ids_keep_time_first_basename() { - let p = summary_rel_path( - SummaryTreeKind::Source, - "slack-eng", - 2, - "summary:1700000000000:L2-deadbeef", - ); - assert_eq!( - p, - "wiki/summaries/source-slack-eng/L2/summary-1700000000000-L2-deadbeef.md" - ); - } - - #[test] - fn summary_rel_path_global() { - // The singleton global tree gets ONE folder; the date is NOT part of - // the folder name. Day/week/month grouping lives in `L/`. - let p = summary_rel_path(SummaryTreeKind::Global, "global", 0, "summary:L0:daily"); - assert_eq!(p, "wiki/summaries/global/L0/summary-L0-daily.md"); - } - - #[test] - fn summary_rel_path_global_levels_share_one_folder() { - // Regression guard for the per-day-folder bug: every level of the - // global tree must land under the same `global/` folder, only the - // `L/` segment differs. - let l0 = summary_rel_path(SummaryTreeKind::Global, "global", 0, "summary:L0:a"); - let l1 = summary_rel_path(SummaryTreeKind::Global, "global", 1, "summary:L1:b"); - let l3 = summary_rel_path(SummaryTreeKind::Global, "global", 3, "summary:L3:c"); - assert_eq!(l0, "wiki/summaries/global/L0/summary-L0-a.md"); - assert_eq!(l1, "wiki/summaries/global/L1/summary-L1-b.md"); - assert_eq!(l3, "wiki/summaries/global/L3/summary-L3-c.md"); - } - - #[test] - fn summary_rel_path_topic() { - let p = summary_rel_path( - SummaryTreeKind::Topic, - "person-alex-johnson", - 1, - "summary:L1:xyz", - ); - assert_eq!( - p, - "wiki/summaries/topic-person-alex-johnson/L1/summary-L1-xyz.md" - ); - } - - #[test] - fn summary_rel_path_strips_trailing_md_extension() { - // If the caller accidentally appends .md to the summary_id, strip it. - let p = summary_rel_path( - SummaryTreeKind::Topic, - "entity-slug", - 2, - "summary:L2:foo.md", - ); - assert_eq!(p, "wiki/summaries/topic-entity-slug/L2/summary-L2-foo.md"); - } - - #[test] - fn summary_filename_preserves_legacy_level_first_shape() { - assert_eq!( - summary_filename("summary:L3:legacy-uuid"), - "summary-L3-legacy-uuid" - ); - } - - #[test] - fn summary_filename_rejects_canonical_shape_with_path_separators() { - // A canonical-looking id whose tail contains `/` must NOT be - // returned verbatim — that would smuggle a directory separator - // into the basename. Fall back to the generic sanitiser so the - // illegal char is replaced with `-`. - let basename = summary_filename("summary:1700000000000:L2-a/b"); - assert!( - !basename.contains('/'), - "basename must not contain a path separator; got {basename}" - ); - // Generic-fallback shape: sanitize_filename replaces both `:` and `/`. - assert_eq!(basename, "summary-1700000000000-L2-a-b"); - } - - #[test] - fn summary_filename_rejects_canonical_shape_with_non_numeric_level() { - // Level segment must be `L`. Anything else (`Lxyz`, - // `L-1`, …) is not the canonical contract — fall back to the - // generic sanitiser instead of accepting verbatim. - let basename = summary_filename("summary:1700000000000:Lxyz-tail"); - assert_eq!(basename, "summary-1700000000000-Lxyz-tail"); - } - - #[test] - fn summary_filename_legacy_branch_rejects_path_separator_in_level() { - // Legacy `summary:L:` branch must also enforce strict - // `L` for the level segment — otherwise an input like - // `summary:L1/2:abc` would produce `summary-L1/2-abc` and - // smuggle a `/` into the basename. Falls through to the - // generic sanitiser, which replaces `/` with `-`. - let basename = summary_filename("summary:L1/2:abc"); - assert!( - !basename.contains('/'), - "basename must not contain a path separator; got {basename}" - ); - assert_eq!(basename, "summary-L1-2-abc"); - } - - #[test] - fn summary_filename_legacy_branch_rejects_traversal_in_level() { - // `summary:L../../x:tail` must NOT produce - // `summary-L../../x-tail` — the legacy branch requires - // `L` and falls through to `sanitize_filename` when - // the level segment is non-numeric. Without `/` in the - // resulting basename the dots are inert characters, not - // directory components. - let basename = summary_filename("summary:L../../x:tail"); - assert!( - !basename.contains('/'), - "basename must not contain a path separator; got {basename}" - ); - } - - #[test] - fn summary_filename_falls_back_for_unknown_shapes() { - assert_eq!( - summary_filename("summary:experimental:value:tail"), - "summary-experimental-value-tail" - ); - } - - #[test] - fn summary_abs_path_rooted_under_content_root() { - use std::path::Path; - let root = Path::new("/workspace/content"); - let abs = summary_abs_path(root, SummaryTreeKind::Global, "global", 0, "daily-123"); - assert!(abs.starts_with(root)); - assert!(abs.ends_with("daily-123.md")); - // Singleton folder, no date segment. - assert!(abs.to_string_lossy().contains("summaries/global/L0/")); - } -} diff --git a/src/openhuman/memory_store/content/raw.rs b/src/openhuman/memory_store/content/raw.rs deleted file mode 100644 index 1a4211f7d..000000000 --- a/src/openhuman/memory_store/content/raw.rs +++ /dev/null @@ -1,477 +0,0 @@ -//! On-disk archive of raw provider items (one .md per source item). -//! -//! Lives alongside the chunked content store but writes a *separate* -//! tree at `/raw///_.md`, -//! where `` is one of `emails`, `chats`, `documents`, `contacts`, -//! `posts` (see [`RawKind`]). The kind subdir keeps a single source's -//! items split by category so Obsidian `.base` files at -//! `/raw//.base` can render -//! per-category views. Contacts and documents are scoped to one source. -//! -//! This is the verbatim payload captured at sync time — no chunking, no -//! summarisation. Useful for: -//! -//! - feeding Obsidian a per-message file the user can read directly, -//! - reproducing the original ingest input when debugging chunker -//! output, -//! - diffing future re-syncs without round-tripping through the -//! chunker. -//! -//! Each file is written atomically (tempfile + rename) so a partial -//! write can never leak into the directory listing. Re-writing the -//! same `(source, uid, ts)` triple is idempotent — same path, same -//! bytes when the upstream item is unchanged. -//! -//! Naming: `_.md` puts the on-disk listing in -//! chronological order while keeping a stable identity suffix so -//! re-syncing the same message overwrites the same file. - -use std::fs; -use std::io::Write; -use std::path::{Path, PathBuf}; - -use anyhow::{Context, Result}; - -use super::paths::slugify_source_id; - -/// Category of a raw item. Used to split a single source's items into -/// per-kind subdirectories under `raw///`. -/// -/// Each connector picks a kind per item — a single connector can write -/// into multiple kinds (e.g. Gmail → [`Self::Email`] for messages, -/// [`Self::Contact`] for senders, [`Self::Document`] for attachments). -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)] -pub enum RawKind { - /// Email messages (Gmail, Outlook, …). - Email, - /// Chat / DM messages (Slack, Telegram, WhatsApp, Discord, …). - Chat, - /// Standalone documents — Notion pages, Drive files, attachments. - Document, - /// One file per person reachable via this source. - Contact, - /// Long-form posts — LinkedIn posts, tweets, blog entries. - Post, - /// Git commits (one file per commit) — GitHub repo sources. - Commit, - /// Issues with their conversation + metadata — GitHub repo sources. - Issue, - /// Pull requests with their body + metadata — GitHub repo sources. - PullRequest, -} - -impl RawKind { - /// Directory name used on disk for this kind. Plural to match the - /// canonical layout (`emails/`, `chats/`, `documents/`, …). - pub const fn as_dir(&self) -> &'static str { - match self { - Self::Email => "emails", - Self::Chat => "chats", - Self::Document => "documents", - Self::Contact => "contacts", - Self::Post => "posts", - Self::Commit => "commits", - Self::Issue => "issues", - Self::PullRequest => "prs", - } - } -} - -/// One raw item ready to land on disk. -pub struct RawItem<'a> { - /// Stable upstream identifier (e.g. Gmail message id). Used for the - /// filename suffix; sanitised before being placed in a path. - pub uid: &'a str, - /// Authoritative timestamp from the upstream item (ms since epoch). - /// Drives the filename prefix so files sort chronologically in any - /// file browser. - pub created_at_ms: i64, - /// Markdown body to write. Should be self-contained (front-matter - /// optional but encouraged). - pub markdown: &'a str, - /// Category subdir under the source (`emails/`, `chats/`, …). - pub kind: RawKind, -} - -/// Write a batch of raw items under `raw///`. -/// -/// `content_root` is the same root that backs `chunk_rel_path` / -/// `summary_rel_path` — i.e. `/memory_tree/content/`. -/// `source_id` is the chunk-store source id (e.g. -/// `"gmail:stevent95-at-gmail-dot-com"`); we slugify it the same way -/// the chunk path does so the raw and chunk trees line up under -/// matching directory names. Each item carries its own [`RawKind`], -/// which selects the per-kind subdir. -/// -/// Returns the number of files written. -pub fn write_raw_items( - content_root: &Path, - source_id: &str, - items: &[RawItem<'_>], -) -> Result { - if items.is_empty() { - return Ok(0); - } - let mut written = 0usize; - for item in items { - let dir = raw_kind_dir(content_root, source_id, item.kind); - fs::create_dir_all(&dir).with_context(|| format!("create raw dir {}", dir.display()))?; - let filename = build_filename(item.created_at_ms, item.uid); - let path = dir.join(&filename); - write_atomic(&path, item.markdown.as_bytes()) - .with_context(|| format!("write raw file {}", path.display()))?; - written += 1; - } - Ok(written) -} - -/// Resolve the on-disk directory for a source's raw archive (the -/// per-source folder that holds every kind subdir plus `_source.md` -/// and `.base` views). -pub fn raw_source_dir(content_root: &Path, source_id: &str) -> PathBuf { - let slug = slugify_source_id(source_id); - content_root.join("raw").join(slug) -} - -/// Resolve the on-disk directory for a single kind under a source — -/// e.g. `/raw//emails/`. -pub fn raw_kind_dir(content_root: &Path, source_id: &str, kind: RawKind) -> PathBuf { - raw_source_dir(content_root, source_id).join(kind.as_dir()) -} - -/// Forward-slash relative path of a raw file under `/`, -/// e.g. `"raw/gmail-acct/emails/1700000000000_msg-1.md"`. Used by -/// callers that record a [`crate::openhuman::memory_store::chunks::store::RawRef`] -/// so reads can resolve the file later without re-deriving the layout. -pub fn raw_rel_path(source_id: &str, kind: RawKind, created_at_ms: i64, uid: &str) -> String { - let slug = slugify_source_id(source_id); - let filename = build_filename(created_at_ms, uid); - format!("raw/{}/{}/{}", slug, kind.as_dir(), filename) -} - -fn build_filename(created_at_ms: i64, uid: &str) -> String { - let ts = created_at_ms.max(0); - let uid = sanitize_uid(uid); - format!("{ts}_{uid}.md") -} - -/// Replace path-illegal characters in the upstream uid before splicing -/// it into a filename. Mirrors `paths::sanitize_filename` but is local -/// so a future change to either side stays decoupled. `pub(crate)` so -/// the raw-coverage backfill (`memory_sync::sources::rebuild`) can map -/// summary child labels back to on-disk filename suffixes. -pub(crate) fn sanitize_uid(uid: &str) -> String { - let cleaned: String = uid - .chars() - .map(|c| match c { - '\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|' | ' ' => '-', - other => other, - }) - .collect(); - if cleaned.is_empty() { - "unknown".into() - } else { - cleaned - } -} - -fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { - let parent = path - .parent() - .ok_or_else(|| anyhow::anyhow!("path has no parent: {}", path.display()))?; - // Per-writer unique tempfile so two concurrent ingest workers - // staging into the same source folder can't trample each other's - // staging path. PID + nanos is collision-free for any realistic - // local concurrency level; the tempfile lands in `parent` so the - // subsequent `rename` is still atomic-on-same-filesystem. - let tmp = parent.join(format!( - ".tmp_raw_{}_{}.md", - std::process::id(), - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_nanos() - )); - let mut f = fs::File::create(&tmp).with_context(|| format!("create tmp {}", tmp.display()))?; - f.write_all(bytes) - .with_context(|| format!("write tmp {}", tmp.display()))?; - f.sync_all() - .with_context(|| format!("fsync tmp {}", tmp.display()))?; - drop(f); - fs::rename(&tmp, path) - .with_context(|| format!("rename {} -> {}", tmp.display(), path.display()))?; - // Best-effort fsync of the directory so the rename is durable on - // crash. We don't surface as an error (the rename has already - // committed; missing dirent fsync is a durability degradation, - // not a failure), but operators want visibility when it happens. - if let Ok(dir_handle) = fs::File::open(parent) { - if let Err(e) = dir_handle.sync_all() { - // Avoid logging the absolute path (embeds workspace / - // home directory). The basename is enough signal for - // operators to correlate with the source slug. - let dir_hint = parent - .file_name() - .and_then(|s| s.to_str()) - .unwrap_or(""); - log::debug!("[content_store::raw] parent dir fsync failed dir={dir_hint} err={e}"); - } - } - Ok(()) -} - -/// Slug an account email like `stevent95@gmail.com` to -/// `stevent95-at-gmail-dot-com`. Used to build per-account source ids -/// from the Composio connection's account email so every memory -/// source is uniquely identified by its connection identity. -/// -/// Rules: -/// - lowercase -/// - `@` → `-at-` -/// - `.` → `-dot-` -/// - any other non-`[a-z0-9]` run collapses to a single `-` -/// - trim leading/trailing `-` -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 chars = lower.chars().peekable(); - for ch in chars { - match ch { - '@' => { - if !last_dash { - out.push('-'); - } - out.push_str("at-"); - last_dash = true; - } - '.' => { - if !last_dash { - out.push('-'); - } - out.push_str("dot-"); - last_dash = true; - } - c if c.is_ascii_alphanumeric() => { - out.push(c); - last_dash = false; - } - _ => { - if !last_dash { - out.push('-'); - last_dash = true; - } - } - } - } - let trimmed = out.trim_end_matches('-').trim_start_matches('-'); - if trimmed.is_empty() { - "unknown".into() - } else { - trimmed.to_string() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use tempfile::TempDir; - - #[test] - fn slug_account_email_basic() { - assert_eq!( - slug_account_email("stevent95@gmail.com"), - "stevent95-at-gmail-dot-com" - ); - } - - #[test] - fn slug_account_email_lowercases_and_trims() { - assert_eq!( - slug_account_email(" Alice.Smith@Example.CO.UK "), - "alice-dot-smith-at-example-dot-co-dot-uk" - ); - } - - #[test] - fn slug_account_email_handles_plus_aliases() { - assert_eq!( - slug_account_email("alice+work@example.com"), - "alice-work-at-example-dot-com" - ); - } - - #[test] - fn slug_account_email_falls_back_to_unknown() { - assert_eq!(slug_account_email(""), "unknown"); - assert_eq!(slug_account_email("@@@"), "at-at-at"); - assert_eq!(slug_account_email("///"), "unknown"); - } - - #[test] - fn write_raw_items_creates_named_files() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - let items = [ - RawItem { - uid: "msg-1", - created_at_ms: 1_700_000_000_000, - markdown: "# hello", - kind: RawKind::Email, - }, - RawItem { - uid: "msg-2", - created_at_ms: 1_700_000_010_000, - markdown: "# world", - kind: RawKind::Email, - }, - ]; - let n = write_raw_items(root, "gmail:stevent95-at-gmail-dot-com", &items).unwrap(); - assert_eq!(n, 2); - let dir = raw_kind_dir(root, "gmail:stevent95-at-gmail-dot-com", RawKind::Email); - assert!( - dir.exists(), - "raw dir should be created at {}", - dir.display() - ); - // Source-level dir is the parent of the kind dir. - assert_eq!( - dir.parent().unwrap(), - raw_source_dir(root, "gmail:stevent95-at-gmail-dot-com") - ); - // Files must sort chronologically (created_at_ms prefix). - let mut names: Vec = fs::read_dir(&dir) - .unwrap() - .filter_map(|e| e.ok()) - .map(|e| e.file_name().to_string_lossy().to_string()) - .collect(); - names.sort(); - assert_eq!( - names, - vec![ - "1700000000000_msg-1.md".to_string(), - "1700000010000_msg-2.md".to_string() - ] - ); - } - - #[test] - fn write_raw_items_is_idempotent() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - let item = RawItem { - uid: "msg-1", - created_at_ms: 1_700_000_000_000, - markdown: "v1", - kind: RawKind::Email, - }; - write_raw_items(root, "gmail:acct", &[item]).unwrap(); - let item2 = RawItem { - uid: "msg-1", - created_at_ms: 1_700_000_000_000, - markdown: "v2", - kind: RawKind::Email, - }; - write_raw_items(root, "gmail:acct", &[item2]).unwrap(); - let dir = raw_kind_dir(root, "gmail:acct", RawKind::Email); - let path = dir.join("1700000000000_msg-1.md"); - let body = fs::read_to_string(&path).unwrap(); - assert_eq!(body, "v2"); - } - - #[test] - fn write_raw_items_sanitises_uid_path_chars() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - let item = RawItem { - uid: "msg/with:dangerous*chars", - created_at_ms: 0, - markdown: "x", - kind: RawKind::Email, - }; - write_raw_items(root, "gmail:acct", &[item]).unwrap(); - let dir = raw_kind_dir(root, "gmail:acct", RawKind::Email); - let entries: Vec = fs::read_dir(&dir) - .unwrap() - .filter_map(|e| e.ok()) - .map(|e| e.file_name().to_string_lossy().to_string()) - .collect(); - assert_eq!(entries.len(), 1); - assert!(entries[0].starts_with("0_msg-with-dangerous-chars")); - } - - #[test] - fn write_raw_items_empty_is_noop() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - let n = write_raw_items(root, "gmail:acct", &[]).unwrap(); - assert_eq!(n, 0); - // Neither source nor any kind dir should exist for an empty batch. - assert!(!raw_source_dir(root, "gmail:acct").exists()); - assert!(!raw_kind_dir(root, "gmail:acct", RawKind::Email).exists()); - } - - #[test] - fn write_raw_items_splits_kinds_into_subdirs() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - let items = [ - RawItem { - uid: "msg-1", - created_at_ms: 1_700_000_000_000, - markdown: "email", - kind: RawKind::Email, - }, - RawItem { - uid: "person-1", - created_at_ms: 0, - markdown: "contact", - kind: RawKind::Contact, - }, - ]; - let n = write_raw_items(root, "gmail:acct", &items).unwrap(); - assert_eq!(n, 2); - assert!(raw_kind_dir(root, "gmail:acct", RawKind::Email) - .join("1700000000000_msg-1.md") - .exists()); - assert!(raw_kind_dir(root, "gmail:acct", RawKind::Contact) - .join("0_person-1.md") - .exists()); - } - - #[test] - fn raw_rel_path_uses_kind_subdir() { - assert_eq!( - raw_rel_path("gmail:acct", RawKind::Email, 1_700_000_000_000, "msg-1"), - "raw/gmail-acct/emails/1700000000000_msg-1.md" - ); - assert_eq!( - raw_rel_path("slack:team", RawKind::Chat, 42, "msg/with:bad"), - "raw/slack-team/chats/42_msg-with-bad.md" - ); - } - - #[test] - fn github_raw_kinds_use_repo_grouped_subdirs() { - assert_eq!(RawKind::Commit.as_dir(), "commits"); - assert_eq!(RawKind::Issue.as_dir(), "issues"); - assert_eq!(RawKind::PullRequest.as_dir(), "prs"); - // `github.com//` slugifies to `github-com--`. - assert_eq!( - raw_rel_path( - "github.com/tinyhumansai/openhuman", - RawKind::Commit, - 1_700_000_000_000, - "2a958e87" - ), - "raw/github-com-tinyhumansai-openhuman/commits/1700000000000_2a958e87.md" - ); - assert_eq!( - raw_rel_path("github.com/org/repo", RawKind::Issue, 0, "42"), - "raw/github-com-org-repo/issues/0_42.md" - ); - assert_eq!( - raw_rel_path("github.com/org/repo", RawKind::PullRequest, 0, "99"), - "raw/github-com-org-repo/prs/0_99.md" - ); - } -} diff --git a/src/openhuman/memory_store/content/read.rs b/src/openhuman/memory_store/content/read.rs index 608d858df..94a9b116a 100644 --- a/src/openhuman/memory_store/content/read.rs +++ b/src/openhuman/memory_store/content/read.rs @@ -1,920 +1,24 @@ -//! Read and verify chunk and summary `.md` files from the content store. +//! Product Config adapters over tinycortex content readers. -use std::path::{Component, Path, PathBuf}; +pub use tinycortex::memory::store::content::{ + read_chunk_file, read_summary_file, verify_chunk_file, verify_summary_file, ChunkFileContents, + VerifyResult, +}; -use super::atomic::sha256_hex; -use super::compose::split_front_matter; -use crate::openhuman::memory::util::redact::redact; - -/// Resolve a DB-stored relative forward-slash path against `content_root`, -/// rejecting any traversal (`..`), absolute, or non-normal component. -/// -/// The `raw_refs` / `content_path` values are treated as **untrusted** at the -/// read boundary: although the write path slugifies/sanitizes them, a future -/// ingest source or DB tamper could store `../../etc/passwd` and turn this -/// reader into an arbitrary file-disclosure primitive that feeds the LLM -/// context. We therefore (1) reject any `..`/absolute/prefix component before -/// touching disk and (2) — when the target exists — canonicalize the resolved -/// path and assert it stays under the canonicalized `content_root`. -fn resolve_within_content_root(content_root: &Path, rel_path: &str) -> anyhow::Result { - // Reject absolute inputs outright. A leading `/` (or a Windows drive/UNC - // prefix) would otherwise split into an empty leading component that gets - // silently skipped, treating `/etc/passwd` as a relative path under the - // content root rather than flagging the obvious traversal attempt. - if Path::new(rel_path).is_absolute() { - return Err(anyhow::anyhow!( - "[content_store::read] rejected absolute path in path_hash={}", - redact(rel_path), - )); - } - - let mut abs = content_root.to_path_buf(); - for component in rel_path.split('/') { - // Skip empty components from leading/double/trailing slashes. - if component.is_empty() || component == "." { - continue; - } - // Reject anything that is not a plain file/dir name: `..`, absolute - // roots, Windows prefixes, etc. - match Path::new(component).components().next() { - Some(Component::Normal(_)) => abs.push(component), - _ => { - return Err(anyhow::anyhow!( - "[content_store::read] rejected unsafe path component in path_hash={}", - redact(rel_path), - )); - } - } - } - - // Defense in depth: if the file exists, canonicalize and confirm - // containment. (canonicalize requires the path to exist, so this is a - // no-op for not-yet-created files — the component check above already - // blocks traversal in that case.) - if abs.exists() { - let canon_root = content_root - .canonicalize() - .unwrap_or_else(|_| content_root.to_path_buf()); - let canon_abs = abs - .canonicalize() - .map_err(|e| anyhow::anyhow!("[content_store::read] canonicalize failed: {e}"))?; - if !canon_abs.starts_with(&canon_root) { - return Err(anyhow::anyhow!( - "[content_store::read] resolved path escapes content_root for path_hash={}", - redact(rel_path), - )); - } - } - - Ok(abs) +fn memory_config(config: &crate::openhuman::config::Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -/// The result of reading a chunk file from disk. -pub struct ChunkFileContents { - /// The Markdown body (everything after the closing `---` of the front-matter). - pub body: String, - /// SHA-256 hex digest over the **body bytes** only. - pub sha256: String, -} - -/// Read a chunk file and return its body + SHA-256. -/// -/// Returns an error if: -/// - the file does not exist -/// - the file is not valid UTF-8 -/// - the front-matter delimiters cannot be found -pub fn read_chunk_file(abs_path: &Path) -> anyhow::Result { - let raw = std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?; - let content = std::str::from_utf8(&raw) - .map_err(|e| anyhow::anyhow!("invalid UTF-8 in {:?}: {e}", abs_path))?; - - let (_fm, body) = split_front_matter(content) - .ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", abs_path))?; - - let sha256 = sha256_hex(body.as_bytes()); - Ok(ChunkFileContents { - body: body.to_string(), - sha256, - }) -} - -/// Verify that the body of a chunk file matches the expected SHA-256. -/// -/// Returns `Ok(true)` on a match, `Ok(false)` on a mismatch, and an `Err` -/// if the file cannot be read or parsed. -pub fn verify_chunk_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result { - let contents = read_chunk_file(abs_path)?; - let ok = contents.sha256 == expected_sha256; - if !ok { - // Log the path as a redacted hash — the path may embed email addresses - // (participant slugs) after the participant-bucketing change. - let path_str = abs_path.to_string_lossy(); - log::warn!( - "[content_store::read] sha256 mismatch for path_hash={}: expected={} actual={}", - redact(&path_str), - expected_sha256, - contents.sha256, - ); - } - Ok(ok) -} - -// ── Summary reads ──────────────────────────────────────────────────────────── - -/// The result of verifying a summary file on disk. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum VerifyResult { - /// The on-disk body SHA-256 matches the stored value. - Ok, - /// The file exists but the body SHA-256 does not match. - Mismatch { actual: String }, - /// The file does not exist at the given path. - Missing, -} - -/// Read a summary file and return its body + SHA-256. -/// -/// Returns an error if: -/// - the file does not exist -/// - the file is not valid UTF-8 -/// - the front-matter delimiters cannot be found -pub fn read_summary_file(abs_path: &Path) -> anyhow::Result { - // Reuse the same reader as chunks — the file format is identical. - read_chunk_file(abs_path) -} - -/// Verify a summary file's body SHA-256 without returning the body itself. -/// -/// Returns: -/// - `VerifyResult::Ok` on match -/// - `VerifyResult::Mismatch { actual }` on hash mismatch -/// - `VerifyResult::Missing` when the file does not exist -pub fn verify_summary_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result { - if !abs_path.exists() { - return Ok(VerifyResult::Missing); - } - let contents = read_summary_file(abs_path)?; - if contents.sha256 == expected_sha256 { - Ok(VerifyResult::Ok) - } else { - // Redact the path — it can embed participant slugs (email addresses). - let path_str = abs_path.to_string_lossy(); - log::warn!( - "[content_store::read] sha256 mismatch for summary path_hash={}: expected={} actual={}", - redact(&path_str), - expected_sha256, - contents.sha256, - ); - Ok(VerifyResult::Mismatch { - actual: contents.sha256, - }) - } -} - -// ── High-level body readers (Config-aware) ─────────────────────────────────── -// -// These helpers resolve the on-disk path from SQLite via -// `get_chunk_content_pointers` / `get_summary_content_pointers`, then read the -// file body. They are the single authoritative entry-point for every caller -// that needs the **full** chunk or summary body (LLM extractor, summariser -// inputs, retrieval API, embedder). Preview-only consumers (UI cards, fast -// filter scans) continue reading the `content` column directly from SQLite. -// -// Error policy: -// - If `content_path` / `content_sha256` are NULL (legacy rows ingested before -// the MD-on-disk migration), return `Err` — callers must handle the -// "pre-migration chunk" case explicitly. The job pipeline propagates the -// error and retries; retrieval falls back gracefully. -// - File-not-found or SHA mismatch → `Err` (propagated to caller for retry / -// alerting). - -/// Read the full body of a chunk `.md` file by its chunk id. -/// -/// Looks up `content_path` in SQLite, resolves it to an absolute path under -/// `config.memory_tree_content_root()`, reads the file, and returns the body -/// string (everything after the YAML front-matter delimiter). -/// -/// Returns `Err` if: -/// - The chunk row has no `content_path` recorded (pre-MD-migration row). -/// - The file cannot be read or has no valid front-matter. -/// -/// # Preview vs. full body -/// The `content` column in `mem_tree_chunks` holds a ≤500-char preview after -/// the MD-on-disk migration. Use this function wherever the full body is -/// required (LLM extraction, embedding, summariser inputs, retrieval API). pub fn read_chunk_body( config: &crate::openhuman::config::Config, chunk_id: &str, ) -> anyhow::Result { - use crate::openhuman::memory_store::chunks::store::{ - get_chunk_content_pointers, get_chunk_raw_refs, - }; - - // Path 1: chunk has raw-archive pointers (today: email). Read each - // referenced file, slice by byte range, join with `\n\n` (the - // chunker's unit separator). No SHA verify — the raw archive is - // the source of truth and was written transactionally with the - // chunk row's id; mismatch can only happen after manual edits. - if let Some(refs) = get_chunk_raw_refs(config, chunk_id)? { - if !refs.is_empty() { - return read_chunk_body_from_raw(config, &refs); - } - } - - let pointers = get_chunk_content_pointers(config, chunk_id)?.ok_or_else(|| { - anyhow::anyhow!( - "[content_store::read] no content_path or raw_refs for chunk_id={} \ - (pre-MD-migration row?)", - chunk_id - ) - })?; - let (rel_path, expected_sha256) = pointers; - if rel_path.is_empty() { - return Err(anyhow::anyhow!( - "[content_store::read] empty content_path and no raw_refs for chunk_id={} \ - — chunk has no resolvable body source", - chunk_id - )); - } - - let content_root = config.memory_tree_content_root(); - // Reconstruct the absolute path from the stored relative forward-slash - // path, rejecting any traversal and confirming containment. - let abs_path = resolve_within_content_root(&content_root, &rel_path)?; - - log::debug!( - "[content_store::read] read_chunk_body chunk_id={} path_hash={}", - chunk_id, - redact(&rel_path), - ); - - let result = read_chunk_file(&abs_path).with_context(|| { - format!( - "read_chunk_body: failed to read file for chunk_id={} path_hash={}", - chunk_id, - redact(&rel_path), - ) - })?; - - // The content file is content-addressed and atomically written, so the file - // on disk is authoritative for this chunk's body. A sha mismatch means the - // stored token drifted from disk — e.g. an external editor rewrote a synced - // file after ingest (#4689). Serve the full on-disk body and repair the - // stale token so the next read verifies cleanly, instead of returning an Err - // that every caller converts into the ≤500-char preview (silent truncation). - if result.sha256 != expected_sha256 { - log::warn!( - "[content_store::read] stale sha token for chunk_id={} disk={} db={} path_hash={} \ - — serving on-disk body and repairing token", - chunk_id, - result.sha256, - expected_sha256, - redact(&rel_path), - ); - if let Err(e) = crate::openhuman::memory_store::chunks::store::update_chunk_content_sha256( - config, - chunk_id, - &result.sha256, - ) { - // Best-effort: the correct body is already in hand; a failed repair - // just means the next read re-heals. Never fail the read on this. - log::warn!( - "[content_store::read] failed to repair sha token for chunk_id={}: {e:#}", - chunk_id, - ); - } - } - - Ok(result.body) + tinycortex::memory::store::content::read_chunk_body(&memory_config(config), chunk_id) } -use anyhow::Context as _; - -/// Reconstruct a chunk body by reading the raw archive files it -/// points at and joining their contents with `"\n\n"` — the same -/// separator the chunker uses between units. -/// -/// Each [`RawRef`] is resolved relative to -/// `config.memory_tree_content_root()`. Byte ranges (`start`, `end`) -/// slice the file; defaults read the whole file. Out-of-bounds -/// ranges are clamped (start past EOF returns empty, end past EOF -/// reads to EOF) so a corrupted offset can't panic the worker — -/// reads are best-effort, log + skip on per-file errors so a single -/// missing raw file doesn't take the whole chunk down. -fn read_chunk_body_from_raw( - config: &crate::openhuman::config::Config, - refs: &[crate::openhuman::memory_store::chunks::store::RawRef], -) -> anyhow::Result { - let content_root = config.memory_tree_content_root(); - let mut parts: Vec = Vec::with_capacity(refs.len()); - for r in refs { - // Treat the DB-stored ref path as untrusted: reject traversal / - // absolute paths and confirm the resolved path stays under - // content_root before reading. Skip (don't fail the whole chunk) on a - // rejected ref, matching the best-effort policy for per-file errors. - let abs = match resolve_within_content_root(&content_root, &r.path) { - Ok(p) => p, - Err(e) => { - log::warn!( - "[content_store::read] raw_ref rejected path_hash={} err={e}", - redact(&r.path) - ); - continue; - } - }; - let bytes = match std::fs::read(&abs) { - Ok(b) => b, - Err(e) => { - log::warn!( - "[content_store::read] raw_ref read failed path_hash={} err={e}", - redact(&r.path) - ); - continue; - } - }; - let len = bytes.len(); - let start = r.start.min(len); - let end = r.end.unwrap_or(len).min(len); - if end <= start { - continue; - } - let slice = &bytes[start..end]; - match std::str::from_utf8(slice) { - Ok(s) => parts.push(s.to_string()), - Err(e) => { - log::warn!( - "[content_store::read] raw_ref non-utf8 path_hash={} err={e}", - redact(&r.path) - ); - } - } - } - Ok(parts.join("\n\n")) -} - -/// Read the full body of a summary `.md` file by its summary id. -/// -/// Looks up `content_path` in SQLite, resolves it to an absolute path under -/// `config.memory_tree_content_root()`, reads the file, and returns the body -/// string. -/// -/// Returns `Err` if: -/// - The summary row has no `content_path` recorded (pre-MD-migration row). -/// - The file cannot be read or has no valid front-matter. -/// -/// # Preview vs. full body -/// The `content` column in `mem_tree_summaries` holds a ≤500-char preview after -/// the MD-on-disk migration. Use this function wherever the full body is -/// required (LLM extraction, embedding, summariser inputs, retrieval API). pub fn read_summary_body( config: &crate::openhuman::config::Config, summary_id: &str, ) -> anyhow::Result { - use crate::openhuman::memory_store::chunks::store::get_summary_content_pointers; - - let pointers = get_summary_content_pointers(config, summary_id)?.ok_or_else(|| { - anyhow::anyhow!( - "[content_store::read] no content_path for summary_id={} (pre-MD-migration row?)", - summary_id - ) - })?; - let (rel_path, expected_sha256) = pointers; - - let content_root = config.memory_tree_content_root(); - let abs_path = resolve_within_content_root(&content_root, &rel_path)?; - - log::debug!( - "[content_store::read] read_summary_body summary_id={} path_hash={}", - summary_id, - redact(&rel_path), - ); - - let result = read_summary_file(&abs_path).with_context(|| { - format!( - "read_summary_body: failed to read file for summary_id={} path_hash={}", - summary_id, - redact(&rel_path), - ) - })?; - - // Self-heal a drifted sha token by trusting the on-disk file and repairing - // the stored token, rather than returning an Err that callers convert into - // the ≤500-char preview. See the matching guard in `read_chunk_body` (#4689). - if result.sha256 != expected_sha256 { - log::warn!( - "[content_store::read] stale sha token for summary_id={} disk={} db={} path_hash={} \ - — serving on-disk body and repairing token", - summary_id, - result.sha256, - expected_sha256, - redact(&rel_path), - ); - if let Err(e) = crate::openhuman::memory_store::chunks::store::update_summary_content_sha256( - config, - summary_id, - &result.sha256, - ) { - log::warn!( - "[content_store::read] failed to repair sha token for summary_id={}: {e:#}", - summary_id, - ); - } - } - - Ok(result.body) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::config::Config; - use crate::openhuman::memory_store::chunks::store::{upsert_chunks, with_connection}; - use crate::openhuman::memory_store::chunks::types::{Chunk, Metadata, SourceKind}; - use crate::openhuman::memory_store::content::atomic::{sha256_hex, write_if_new}; - use crate::openhuman::memory_store::content::compose::{ - compose_chunk_file, SummaryComposeInput, - }; - use crate::openhuman::memory_store::content::paths::SummaryTreeKind; - use crate::openhuman::memory_store::content::{atomic::stage_summary, stage_chunks}; - use crate::openhuman::memory_store::trees::store::{insert_summary_tx, insert_tree}; - use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind, TreeStatus}; - use chrono::TimeZone; - use tempfile::TempDir; - - fn sample_chunk() -> Chunk { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - Chunk { - id: "read_test".into(), - content: "## ts — alice\nhello from read test".into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: None, - path_scope: None, - }, - token_count: 8, - seq_in_source: 0, - created_at: ts, - partial_message: false, - } - } - - fn test_config(tmp: &TempDir) -> Config { - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg - } - - fn sample_tree() -> Tree { - Tree { - id: "tree-1".into(), - kind: TreeKind::Source, - scope: "slack:#eng".into(), - root_id: None, - max_level: 0, - status: TreeStatus::Active, - created_at: chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - last_sealed_at: None, - } - } - - fn sample_summary_node() -> SummaryNode { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - SummaryNode { - id: "summary-1".into(), - tree_id: "tree-1".into(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: vec!["leaf-a".into()], - content: "summary full body".into(), - token_count: 4, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - } - } - - #[test] - fn read_returns_body_and_correct_sha256() { - let dir = TempDir::new().unwrap(); - let chunk = sample_chunk(); - let (full_bytes, body_bytes) = compose_chunk_file(&chunk); - let path = dir.path().join("0.md"); - write_if_new(&path, &full_bytes).unwrap(); - - let result = read_chunk_file(&path).unwrap(); - assert_eq!(result.body, std::str::from_utf8(&body_bytes).unwrap()); - assert_eq!(result.sha256, sha256_hex(&body_bytes)); - } - - #[test] - fn verify_passes_for_correct_hash() { - let dir = TempDir::new().unwrap(); - let chunk = sample_chunk(); - let (full_bytes, body_bytes) = compose_chunk_file(&chunk); - let path = dir.path().join("0.md"); - write_if_new(&path, &full_bytes).unwrap(); - - let expected = sha256_hex(&body_bytes); - assert!(verify_chunk_file(&path, &expected).unwrap()); - } - - #[test] - fn verify_fails_for_wrong_hash() { - let dir = TempDir::new().unwrap(); - let chunk = sample_chunk(); - let (full_bytes, _) = compose_chunk_file(&chunk); - let path = dir.path().join("0.md"); - write_if_new(&path, &full_bytes).unwrap(); - - assert!(!verify_chunk_file(&path, "deadbeef").unwrap()); - } - - #[test] - fn read_missing_file_returns_error() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("nonexistent.md"); - assert!(read_chunk_file(&path).is_err()); - } - - // ─── summary read / verify tests ───────────────────────────────────────── - - fn write_summary_file(dir: &TempDir, body: &str) -> (std::path::PathBuf, String) { - use crate::openhuman::memory_store::content::atomic::{sha256_hex, write_if_new}; - use crate::openhuman::memory_store::content::compose::{ - compose_summary_md, SummaryComposeInput, - }; - use crate::openhuman::memory_store::content::paths::SummaryTreeKind; - use chrono::TimeZone; - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "sum:L1:readtest", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 1, - child_ids: &["c1".to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body, - }; - let composed = compose_summary_md(&input); - let path = dir.path().join("sum.md"); - let sha = sha256_hex(composed.body.as_bytes()); - write_if_new(&path, composed.full.as_bytes()).unwrap(); - (path, sha) - } - - #[test] - fn read_summary_file_returns_body_and_sha() { - let dir = TempDir::new().unwrap(); - let body = "summary body content\n"; - let (path, expected_sha) = write_summary_file(&dir, body); - let result = read_summary_file(&path).unwrap(); - assert_eq!(result.body, body); - assert_eq!(result.sha256, expected_sha); - } - - #[test] - fn verify_summary_file_ok_for_correct_hash() { - let dir = TempDir::new().unwrap(); - let (path, sha) = write_summary_file(&dir, "body text\n"); - assert_eq!(verify_summary_file(&path, &sha).unwrap(), VerifyResult::Ok); - } - - #[test] - fn verify_summary_file_mismatch_for_wrong_hash() { - let dir = TempDir::new().unwrap(); - let (path, _) = write_summary_file(&dir, "body text\n"); - let r = verify_summary_file(&path, "deadbeef").unwrap(); - assert!(matches!(r, VerifyResult::Mismatch { .. })); - } - - #[test] - fn verify_summary_file_missing_for_absent_file() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("does_not_exist.md"); - assert_eq!( - verify_summary_file(&path, "abc").unwrap(), - VerifyResult::Missing - ); - } - - #[test] - fn read_chunk_file_rejects_invalid_utf8() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("bad.md"); - std::fs::write(&path, [0xff, 0xfe, 0xfd]).unwrap(); - let err = match read_chunk_file(&path) { - Ok(_) => panic!("invalid UTF-8 should fail"), - Err(err) => err, - }; - assert!(err.to_string().contains("invalid UTF-8")); - } - - #[test] - fn read_chunk_file_rejects_missing_front_matter() { - let dir = TempDir::new().unwrap(); - let path = dir.path().join("plain.md"); - std::fs::write(&path, "no front matter here").unwrap(); - let err = match read_chunk_file(&path) { - Ok(_) => panic!("missing front matter should fail"), - Err(err) => err, - }; - assert!(err.to_string().contains("no front-matter")); - } - - #[test] - fn verify_summary_file_mismatch_returns_actual_sha() { - let dir = TempDir::new().unwrap(); - let (path, expected_sha) = write_summary_file(&dir, "body text\n"); - let actual = match verify_summary_file(&path, "deadbeef").unwrap() { - VerifyResult::Mismatch { actual } => actual, - other => panic!("expected mismatch, got {other:?}"), - }; - assert_eq!(actual, expected_sha); - } - - #[test] - fn read_chunk_body_from_raw_clamps_ranges_and_skips_bad_refs() { - use crate::openhuman::memory_store::chunks::store::RawRef; - - let dir = TempDir::new().unwrap(); - let mut cfg = crate::openhuman::config::Config::default(); - cfg.workspace_dir = dir.path().to_path_buf(); - - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).unwrap(); - - std::fs::write(content_root.join("one.txt"), "abcdef").unwrap(); - std::fs::write(content_root.join("two.txt"), [0xff, 0xfe]).unwrap(); - - let refs = vec![ - RawRef { - path: "one.txt".into(), - start: 1, - end: Some(4), - }, - RawRef { - path: "missing.txt".into(), - start: 0, - end: None, - }, - RawRef { - path: "two.txt".into(), - start: 0, - end: None, - }, - RawRef { - path: "one.txt".into(), - start: 99, - end: None, - }, - ]; - - let body = read_chunk_body_from_raw(&cfg, &refs).unwrap(); - assert_eq!(body, "bcd"); - } - - #[test] - fn read_chunk_body_from_raw_rejects_path_traversal() { - use crate::openhuman::memory_store::chunks::store::RawRef; - - let dir = TempDir::new().unwrap(); - let mut cfg = crate::openhuman::config::Config::default(); - cfg.workspace_dir = dir.path().to_path_buf(); - - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).unwrap(); - std::fs::write(content_root.join("safe.txt"), "safe").unwrap(); - - // A secret sitting next to content_root that a traversal ref tries to - // reach. The traversal ref must be skipped, leaving only the safe body. - let outside = content_root.parent().unwrap().join("secret.txt"); - std::fs::write(&outside, "TOP SECRET").unwrap(); - - let refs = vec![ - RawRef { - path: "../secret.txt".into(), - start: 0, - end: None, - }, - RawRef { - path: "safe.txt".into(), - start: 0, - end: None, - }, - ]; - - let body = read_chunk_body_from_raw(&cfg, &refs).unwrap(); - assert_eq!(body, "safe"); - assert!(!body.contains("SECRET")); - } - - #[test] - fn resolve_within_content_root_rejects_traversal_and_absolute() { - let dir = TempDir::new().unwrap(); - let root = dir.path(); - - assert!(resolve_within_content_root(root, "../escape.md").is_err()); - assert!(resolve_within_content_root(root, "a/../../escape.md").is_err()); - assert!(resolve_within_content_root(root, "/etc/passwd").is_err()); - - // Safe relative paths resolve correctly. - 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 read_chunk_body_roundtrips_from_staged_content_pointer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let chunk = sample_chunk(); - upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); - let staged = stage_chunks( - &cfg.memory_tree_content_root(), - std::slice::from_ref(&chunk), - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let body = read_chunk_body(&cfg, &chunk.id).unwrap(); - assert_eq!(body, chunk.content); - } - - #[test] - fn read_chunk_body_errors_when_pointers_are_missing() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let err = read_chunk_body(&cfg, "missing-chunk").unwrap_err(); - assert!(err.to_string().contains("no content_path or raw_refs")); - } - - #[test] - fn read_chunk_body_self_heals_on_sha_mismatch() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let chunk = sample_chunk(); - upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); - let staged = stage_chunks( - &cfg.memory_tree_content_root(), - std::slice::from_ref(&chunk), - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - // Simulate an external editor rewriting the synced file after ingest: - // the on-disk body drifts from the recorded content_sha256 (#4689). - let rel = - crate::openhuman::memory_store::chunks::store::get_chunk_content_path(&cfg, &chunk.id) - .unwrap() - .unwrap(); - let mut abs = cfg.memory_tree_content_root(); - for part in rel.split('/') { - abs.push(part); - } - std::fs::write(&abs, b"---\nsource_kind: chat\n---\nmutated body").unwrap(); - - // Self-heal: serve the full on-disk body instead of erroring into the - // ≤500-char preview, and repair the stale token. - let body = read_chunk_body(&cfg, &chunk.id).unwrap(); - assert_eq!(body, "mutated body"); - - let (_, sha) = crate::openhuman::memory_store::chunks::store::get_chunk_content_pointers( - &cfg, &chunk.id, - ) - .unwrap() - .unwrap(); - assert_eq!(sha, sha256_hex(b"mutated body")); - // A second read now verifies cleanly against the repaired token. - assert_eq!(read_chunk_body(&cfg, &chunk.id).unwrap(), "mutated body"); - } - - #[test] - fn read_summary_body_self_heals_on_sha_mismatch() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let tree = sample_tree(); - let node = sample_summary_node(); - insert_tree(&cfg, &tree).unwrap(); - let staged = stage_summary( - &cfg.memory_tree_content_root(), - &SummaryComposeInput { - summary_id: &node.id, - tree_kind: SummaryTreeKind::Source, - tree_id: &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, - }, - "slack-eng", - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &node, Some(&staged), "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let (rel, _) = crate::openhuman::memory_store::chunks::store::get_summary_content_pointers( - &cfg, &node.id, - ) - .unwrap() - .unwrap(); - let mut abs = cfg.memory_tree_content_root(); - for part in rel.split('/') { - abs.push(part); - } - std::fs::write(&abs, b"---\ntree_kind: source\n---\nmutated summary").unwrap(); - - let body = read_summary_body(&cfg, &node.id).unwrap(); - assert_eq!(body, "mutated summary"); - let (_, sha) = crate::openhuman::memory_store::chunks::store::get_summary_content_pointers( - &cfg, &node.id, - ) - .unwrap() - .unwrap(); - assert_eq!(sha, sha256_hex(b"mutated summary")); - } - - #[test] - fn read_summary_body_roundtrips_from_staged_content_pointer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let tree = sample_tree(); - let node = sample_summary_node(); - insert_tree(&cfg, &tree).unwrap(); - let staged = stage_summary( - &cfg.memory_tree_content_root(), - &SummaryComposeInput { - summary_id: &node.id, - tree_kind: SummaryTreeKind::Source, - tree_id: &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, - }, - "slack-eng", - ) - .unwrap(); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &node, Some(&staged), "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let body = read_summary_body(&cfg, &node.id).unwrap(); - assert_eq!(body, node.content); - } - - #[test] - fn read_summary_body_errors_when_pointers_are_missing() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let err = read_summary_body(&cfg, "missing-summary").unwrap_err(); - assert!(err.to_string().contains("no content_path for summary_id")); - } + tinycortex::memory::store::content::read_summary_body(&memory_config(config), summary_id) } diff --git a/src/openhuman/memory_store/entities.rs b/src/openhuman/memory_store/entities.rs index 91339f32c..169c78516 100644 --- a/src/openhuman/memory_store/entities.rs +++ b/src/openhuman/memory_store/entities.rs @@ -1,49 +1,100 @@ -//! Entities — the `mem_tree_entity_index` table surfaced as a first-class -//! memory_store submodule. -//! -//! The entity index is one of the four primitives memory_store owns -//! (raw / entities / tree / vector + kv). Today its persistence lives in -//! `memory_tree::score::store` because the scorer was the first writer; this -//! module re-exports the read/write surface under the canonical -//! `memory_store::entities::*` path so callers don't have to know about -//! the implementation location. -//! -//! Once the score module finishes splitting (entity persistence vs -//! scoring math), the table-owning code moves here and `memory::score` -//! becomes a pure consumer. -//! -//! ## API -//! -//! | Re-export | Source | -//! | --- | --- | -//! | [`EntityHit`] | `memory_tree::score::store::EntityHit` | -//! | [`index_entity`] | `memory_tree::score::store::index_entity` | -//! | [`index_entities`] | `memory_tree::score::store::index_entities` | -//! | [`lookup_entity`] | `memory_tree::score::store::lookup_entity` | -//! | [`list_entity_ids_for_node`] | `memory_tree::score::store::list_entity_ids_for_node` | -//! | [`clear_entity_index_for_node`] | `memory_tree::score::store::clear_entity_index_for_node` | -//! | [`count_entity_index`] | `memory_tree::score::store::count_entity_index` | -//! -//! The derived co-occurrence query layer built on these primitives lives in -//! `tinycortex::memory::graph` (the host `memory_graph` placeholder was deleted -//! in the W7 migration — it was an unused derive-on-read scaffold). +//! Host adapters for tinycortex's entity occurrence index. -pub use crate::openhuman::memory_tree::score::store::{ - clear_entity_index_for_node, count_entity_index, index_entities, index_entity, - list_entity_ids_for_node, lookup_entity, EntityHit, +use std::sync::Arc; + +use anyhow::Result; +use tinycortex::memory::store::entity_index::{ + CanonicalEntity, EntityIndex, EntityKind, SelfIdentity, }; +use crate::openhuman::composio::providers::profile::{is_self_identity_any_toolkit, IdentityKind}; +use crate::openhuman::config::Config; +use crate::openhuman::tinycortex::memory_config_from; + +pub use tinycortex::memory::store::entity_index::EntityHit; + +#[derive(Debug)] +struct HostSelfIdentity; + +impl SelfIdentity for HostSelfIdentity { + fn is_self(&self, kind: EntityKind, surface: &str) -> bool { + let identity_kind = match kind { + EntityKind::Email => IdentityKind::Email, + EntityKind::Handle => IdentityKind::Handle, + _ => return false, + }; + is_self_identity_any_toolkit(identity_kind, surface) + } +} + +fn index(config: &Config) -> Result { + let memory = memory_config_from(config, config.workspace_dir.clone()); + let connection = tinycortex::memory::chunks::shared_connection(&memory)?; + EntityIndex::from_shared_connection(connection, Arc::new(HostSelfIdentity)) +} + +pub(crate) fn host_self_identity() -> Arc { + Arc::new(HostSelfIdentity) +} + +pub fn index_entity( + config: &Config, + entity: &CanonicalEntity, + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result<()> { + log::debug!("[memory:entities] index one node_kind={node_kind}"); + index(config)?.index_entity(entity, node_id, node_kind, timestamp_ms, tree_id) +} + +pub fn index_entities( + config: &Config, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> Result { + log::debug!( + "[memory:entities] index batch count={} node_kind={node_kind}", + entities.len() + ); + index(config)?.index_entities(entities, node_id, node_kind, timestamp_ms, tree_id) +} + +pub fn clear_entity_index_for_node(config: &Config, node_id: &str) -> Result { + index(config)?.clear_entity_index_for_node(node_id) +} + +pub fn lookup_entity( + config: &Config, + entity_id: &str, + limit: Option, +) -> Result> { + index(config)?.lookup_entity(entity_id, limit) +} + +pub fn list_entity_ids_for_node(config: &Config, node_id: &str) -> Result> { + index(config)?.list_entity_ids_for_node(node_id) +} + +pub fn count_entity_index(config: &Config) -> Result { + index(config)?.count_entity_index() +} + #[cfg(test)] mod tests { use super::*; #[test] - fn entity_hit_reexport_is_constructible() { + fn crate_entity_hit_is_the_host_facade_type() { let hit = EntityHit { entity_id: "person:alice".into(), node_id: "chunk-1".into(), node_kind: "leaf".into(), - entity_kind: crate::openhuman::memory_tree::score::extract::EntityKind::Person, + entity_kind: EntityKind::Person, surface: "Alice".into(), score: 1.0, timestamp_ms: 123, @@ -51,6 +102,5 @@ mod tests { is_user: false, }; assert_eq!(hit.entity_id, "person:alice"); - assert_eq!(hit.node_kind, "leaf"); } } diff --git a/src/openhuman/memory_store/factories.rs b/src/openhuman/memory_store/factories.rs index 8983e2034..594157733 100644 --- a/src/openhuman/memory_store/factories.rs +++ b/src/openhuman/memory_store/factories.rs @@ -12,6 +12,9 @@ use std::path::Path; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +use parking_lot::Mutex; +use rusqlite::Connection; + use crate::openhuman::config::{EmbeddingRouteConfig, MemoryConfig, StorageProviderConfig}; use crate::openhuman::embeddings::{ self, format_embedding_signature, EmbeddingProvider, DEFAULT_CLOUD_EMBEDDING_DIMENSIONS, @@ -337,6 +340,38 @@ pub fn create_memory_with_local_ai( ) } +/// Memory resources needed by an agent session. +/// +/// The storage abstraction remains backend-neutral while SQLite-specific +/// consumers receive the concrete shared connection explicitly. +pub(crate) struct SessionMemory { + pub memory: Box, + pub sqlite_connection: Arc>, +} + +pub(crate) fn create_session_memory_with_local_ai( + memory: &MemoryConfig, + local_embedding_model: Option<&str>, + embedding_api_key: &str, + embedding_routes: &[EmbeddingRouteConfig], + storage_provider: Option<&StorageProviderConfig>, + workspace_dir: &Path, +) -> anyhow::Result { + let memory = create_unified_memory_full( + memory, + embedding_routes, + storage_provider, + local_embedding_model, + embedding_api_key, + workspace_dir, + )?; + let sqlite_connection = Arc::clone(&memory.conn); + Ok(SessionMemory { + memory: Box::new(memory), + sqlite_connection, + }) +} + /// Synchronous health-check shim around [`probe_ollama_reachable`]. /// /// Production call sites (`create_memory_with_local_ai` and friends) live in @@ -385,6 +420,24 @@ fn create_memory_full( embedding_api_key: &str, workspace_dir: &Path, ) -> anyhow::Result> { + Ok(Box::new(create_unified_memory_full( + config, + _embedding_routes, + _storage_provider, + local_embedding_model, + embedding_api_key, + workspace_dir, + )?)) +} + +fn create_unified_memory_full( + config: &MemoryConfig, + _embedding_routes: &[EmbeddingRouteConfig], + _storage_provider: Option<&StorageProviderConfig>, + local_embedding_model: Option<&str>, + embedding_api_key: &str, + workspace_dir: &Path, +) -> anyhow::Result { // 1. Resolve the intended provider from config. let intended = effective_embedding_settings(config, local_embedding_model); let local_ai_opt_in = local_embedding_model @@ -454,8 +507,7 @@ fn create_memory_full( ); // 4. Instantiate UnifiedMemory which handles SQLite and vector storage. - let mem = UnifiedMemory::new(workspace_dir, embedder, config.sqlite_open_timeout_secs)?; - Ok(Box::new(mem)) + UnifiedMemory::new(workspace_dir, embedder, config.sqlite_open_timeout_secs) } /// Create a memory instance specifically for migration purposes. diff --git a/src/openhuman/memory_store/kv.rs b/src/openhuman/memory_store/kv.rs index 1b2455b25..6afd67c1d 100644 --- a/src/openhuman/memory_store/kv.rs +++ b/src/openhuman/memory_store/kv.rs @@ -1,278 +1,98 @@ -//! Key-value storage — `kv_global` + `kv_namespace` tables. -//! -//! Lifted out of `unified/` so KV is a peer of `trees/`, `vectors/`, and -//! the other first-class memory_store submodules. The `impl UnifiedMemory` -//! block stays here because the methods still operate on the unified -//! SQLite connection; once `Memory` trait callers migrate to a per-kind -//! backend, the `UnifiedMemory` impl shrinks to a thin shim and the bulk -//! of this file moves to free functions. +//! Compatibility methods for tinycortex's shared-connection KV store. -use rusqlite::{params, OptionalExtension}; -use serde_json::json; +use tinycortex::memory::store::kv::KvStore; -use crate::openhuman::memory_store::safety; use crate::openhuman::memory_store::types::MemoryKvRecord; use crate::openhuman::memory_store::unified::UnifiedMemory; impl UnifiedMemory { - /// Insert or update a global key-value pair. + fn tinycortex_kv(&self) -> Result { + KvStore::from_shared_connection(self.conn.clone()) + .map_err(|error| format!("initialize tinycortex KV store: {error}")) + } + pub async fn kv_set_global(&self, key: &str, value: &serde_json::Value) -> Result<(), String> { - if safety::has_likely_secret(key) { - log::warn!( - "[memory:safety] kv_set_global rejected secret-like key key_chars={}", - key.chars().count() - ); - return Err("kv key cannot contain secrets".to_string()); - } - if safety::pii::has_likely_pii(key) { - log::warn!( - "[memory:safety] kv_set_global rejected PII-like key key_chars={}", - key.chars().count() - ); - return Err("kv key cannot contain personal identifiers".to_string()); - } - - let sanitized_value = safety::sanitize_json(value); - let report = sanitized_value.report; - if report.changed() { - log::warn!( - "[memory:safety] kv_set_global sanitized key_chars={} text_redactions={} key_redactions={} blocked_secret_hits={} depth_redactions={} pii_redactions={}", - key.chars().count(), - report.text_redactions, - report.key_redactions, - report.blocked_secret_hits, - report.depth_redactions, - report.pii_redactions - ); - } - - let conn = self.conn.lock(); - conn.execute( - "INSERT INTO kv_global (key, value_json, updated_at) - VALUES (?1, ?2, ?3) - ON CONFLICT(key) DO UPDATE SET value_json = excluded.value_json, updated_at = excluded.updated_at", - params![key, sanitized_value.value.to_string(), Self::now_ts()], - ) - .map_err(|e| format!("kv_set_global: {e}"))?; - Ok(()) + self.tinycortex_kv()?.set_global(key, value) } - /// Read a global key, returning `None` if absent. pub async fn kv_get_global(&self, key: &str) -> Result, String> { - let conn = self.conn.lock(); - let value: Option = conn - .query_row( - "SELECT value_json FROM kv_global WHERE key = ?1", - params![key], - |row| row.get(0), - ) - .optional() - .map_err(|e| format!("kv_get_global: {e}"))?; - Ok(value.and_then(|v| serde_json::from_str(&v).ok())) + self.tinycortex_kv()?.get_global(key) } - /// Insert or update a namespace-scoped key-value pair. pub async fn kv_set_namespace( &self, namespace: &str, key: &str, value: &serde_json::Value, ) -> Result<(), String> { - if safety::has_likely_secret(namespace) || safety::has_likely_secret(key) { - log::warn!( - "[memory:safety] kv_set_namespace rejected secret-like namespace/key namespace_chars={} key_chars={}", - namespace.chars().count(), - key.chars().count() - ); - return Err("kv namespace/key cannot contain secrets".to_string()); - } - if safety::pii::has_likely_pii(namespace) || safety::pii::has_likely_pii(key) { - log::warn!( - "[memory:safety] kv_set_namespace rejected PII-like namespace/key namespace_chars={} key_chars={}", - namespace.chars().count(), - key.chars().count() - ); - return Err("kv namespace/key cannot contain personal identifiers".to_string()); - } - - let sanitized_value = safety::sanitize_json(value); - let report = sanitized_value.report; - if report.changed() { - log::warn!( - "[memory:safety] kv_set_namespace sanitized namespace_chars={} key_chars={} text_redactions={} key_redactions={} blocked_secret_hits={} depth_redactions={} pii_redactions={}", - namespace.chars().count(), - key.chars().count(), - report.text_redactions, - report.key_redactions, - report.blocked_secret_hits, - report.depth_redactions, - report.pii_redactions - ); - } - - let conn = self.conn.lock(); - conn.execute( - "INSERT INTO kv_namespace (namespace, key, value_json, updated_at) - VALUES (?1, ?2, ?3, ?4) - ON CONFLICT(namespace, key) DO UPDATE SET value_json = excluded.value_json, updated_at = excluded.updated_at", - params![ - Self::sanitize_namespace(namespace), - key, - sanitized_value.value.to_string(), - Self::now_ts() - ], - ) - .map_err(|e| format!("kv_set_namespace: {e}"))?; - Ok(()) + self.tinycortex_kv()?.set_namespace(namespace, key, value) } - /// Read a namespace-scoped key, returning `None` if absent. pub async fn kv_get_namespace( &self, namespace: &str, key: &str, ) -> Result, String> { - let conn = self.conn.lock(); - let value: Option = conn - .query_row( - "SELECT value_json FROM kv_namespace WHERE namespace = ?1 AND key = ?2", - params![Self::sanitize_namespace(namespace), key], - |row| row.get(0), - ) - .optional() - .map_err(|e| format!("kv_get_namespace: {e}"))?; - Ok(value.and_then(|v| serde_json::from_str(&v).ok())) + self.tinycortex_kv()?.get_namespace(namespace, key) } - /// Delete a global key. Returns `true` if a row was removed. pub async fn kv_delete_global(&self, key: &str) -> Result { - let conn = self.conn.lock(); - let changed = conn - .execute("DELETE FROM kv_global WHERE key = ?1", params![key]) - .map_err(|e| format!("kv_delete_global: {e}"))?; - Ok(changed > 0) + self.tinycortex_kv()?.delete_global(key) } - /// Delete a namespace-scoped key. Returns `true` if a row was removed. pub async fn kv_delete_namespace(&self, namespace: &str, key: &str) -> Result { - let conn = self.conn.lock(); - let changed = conn - .execute( - "DELETE FROM kv_namespace WHERE namespace = ?1 AND key = ?2", - params![Self::sanitize_namespace(namespace), key], - ) - .map_err(|e| format!("kv_delete_namespace: {e}"))?; - Ok(changed > 0) + self.tinycortex_kv()?.delete_namespace(namespace, key) } - /// List all keys in a namespace, most recently updated first. pub async fn kv_list_namespace( &self, namespace: &str, ) -> Result, String> { - let conn = self.conn.lock(); - let mut stmt = conn - .prepare( - "SELECT key, value_json, updated_at FROM kv_namespace - WHERE namespace = ?1 ORDER BY updated_at DESC", - ) - .map_err(|e| format!("kv_list_namespace prepare: {e}"))?; - let mut rows = stmt - .query(params![Self::sanitize_namespace(namespace)]) - .map_err(|e| format!("kv_list_namespace query: {e}"))?; - let mut out = Vec::new(); - while let Some(row) = rows - .next() - .map_err(|e| format!("kv_list_namespace row: {e}"))? - { - let value_raw: String = row.get(1).map_err(|e| e.to_string())?; - out.push(json!({ - "key": row.get::<_, String>(0).map_err(|e| e.to_string())?, - "value": serde_json::from_str::(&value_raw).unwrap_or(serde_json::Value::Null), - "updatedAt": row.get::<_, f64>(2).map_err(|e| e.to_string())?, - })); - } - Ok(out) + self.tinycortex_kv()?.list_namespace(namespace) } pub(crate) async fn kv_records_for_scope( &self, namespace: &str, ) -> Result, String> { - let mut records = self.kv_records_namespace(namespace).await?; - records.extend(self.kv_records_global().await?); - records.sort_by(|a, b| { - b.updated_at - .partial_cmp(&a.updated_at) - .unwrap_or(std::cmp::Ordering::Equal) - }); - Ok(records) + self.tinycortex_kv()? + .records_for_scope(namespace) + .map(convert_records) } pub(crate) async fn kv_records_namespace( &self, namespace: &str, ) -> Result, String> { - let conn = self.conn.lock(); - let mut stmt = conn - .prepare( - "SELECT key, value_json, updated_at FROM kv_namespace - WHERE namespace = ?1 - ORDER BY updated_at DESC", - ) - .map_err(|e| format!("prepare kv_records_namespace: {e}"))?; - let mut rows = stmt - .query(params![Self::sanitize_namespace(namespace)]) - .map_err(|e| format!("query kv_records_namespace: {e}"))?; - let mut out = Vec::new(); - while let Some(row) = rows - .next() - .map_err(|e| format!("row kv_records_namespace: {e}"))? - { - let value_raw: String = row.get(1).map_err(|e| e.to_string())?; - out.push(MemoryKvRecord { - namespace: Some(Self::sanitize_namespace(namespace)), - key: row.get(0).map_err(|e| e.to_string())?, - value: serde_json::from_str(&value_raw).unwrap_or(serde_json::Value::Null), - updated_at: row.get(2).map_err(|e| e.to_string())?, - }); - } - Ok(out) + self.tinycortex_kv()? + .records_namespace(namespace) + .map(convert_records) } pub(crate) async fn kv_records_global(&self) -> Result, String> { - let conn = self.conn.lock(); - let mut stmt = conn - .prepare( - "SELECT key, value_json, updated_at FROM kv_global - ORDER BY updated_at DESC", - ) - .map_err(|e| format!("prepare kv_records_global: {e}"))?; - let mut rows = stmt - .query([]) - .map_err(|e| format!("query kv_records_global: {e}"))?; - let mut out = Vec::new(); - while let Some(row) = rows - .next() - .map_err(|e| format!("row kv_records_global: {e}"))? - { - let value_raw: String = row.get(1).map_err(|e| e.to_string())?; - out.push(MemoryKvRecord { - namespace: None, - key: row.get(0).map_err(|e| e.to_string())?, - value: serde_json::from_str(&value_raw).unwrap_or(serde_json::Value::Null), - updated_at: row.get(2).map_err(|e| e.to_string())?, - }); - } - Ok(out) + self.tinycortex_kv()?.records_global().map(convert_records) } } +fn convert_records(records: Vec) -> Vec { + records + .into_iter() + .map(|record| MemoryKvRecord { + namespace: record.namespace, + key: record.key, + value: record.value, + updated_at: record.updated_at, + }) + .collect() +} + #[cfg(test)] mod tests { + use serde_json::json; + use tempfile::TempDir; + use super::*; use crate::openhuman::embeddings::NoopEmbedding; - use tempfile::TempDir; fn test_memory() -> (TempDir, UnifiedMemory) { let tmp = TempDir::new().unwrap(); @@ -282,70 +102,25 @@ mod tests { } #[tokio::test] - async fn global_kv_roundtrips_and_deletes() { + async fn global_kv_roundtrips_and_deletes_through_tinycortex() { let (_tmp, memory) = test_memory(); memory.kv_set_global("theme", &json!("dark")).await.unwrap(); assert_eq!( memory.kv_get_global("theme").await.unwrap(), Some(json!("dark")) ); - assert!(memory.kv_delete_global("theme").await.unwrap()); - assert_eq!(memory.kv_get_global("theme").await.unwrap(), None); } #[tokio::test] - async fn namespace_kv_roundtrips_lists_and_combines_scope_records() { + async fn namespace_records_share_the_unified_connection() { let (_tmp, memory) = test_memory(); - memory - .kv_set_global("global-setting", &json!(true)) - .await - .unwrap(); memory .kv_set_namespace("team alpha/#1", "state", &json!({"open": true})) .await .unwrap(); - - assert_eq!( - memory - .kv_get_namespace("team alpha/#1", "state") - .await - .unwrap(), - Some(json!({"open": true})) - ); - - let listed = memory.kv_list_namespace("team alpha/#1").await.unwrap(); - assert_eq!(listed.len(), 1); - assert_eq!(listed[0]["key"], "state"); - assert_eq!(listed[0]["value"], json!({"open": true})); - - let scoped = memory.kv_records_for_scope("team alpha/#1").await.unwrap(); - assert_eq!(scoped.len(), 2); - assert!(scoped - .iter() - .any(|r| r.namespace.is_none() && r.key == "global-setting")); - assert!(scoped - .iter() - .any(|r| { r.namespace.as_deref() == Some("team_alpha/_1") && r.key == "state" })); - } - - #[tokio::test] - async fn kv_rejects_secret_like_keys() { - let (_tmp, memory) = test_memory(); - let err = memory - .kv_set_global("sk-proj-abcdefghijklmnop", &json!("secret")) - .await - .unwrap_err(); - assert!(err.contains("cannot contain secrets")); - - let err = memory - .kv_set_namespace( - "project", - "ghp_abcdefghijklmnopqrstuvwx123456", - &json!("secret"), - ) - .await - .unwrap_err(); - assert!(err.contains("cannot contain secrets")); + let records = memory.kv_records_namespace("team alpha/#1").await.unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0].namespace.as_deref(), Some("team_alpha/_1")); } } diff --git a/src/openhuman/memory_store/memory_trait.rs b/src/openhuman/memory_store/memory_trait.rs index 1c37d7956..d209aa04a 100644 --- a/src/openhuman/memory_store/memory_trait.rs +++ b/src/openhuman/memory_store/memory_trait.rs @@ -348,50 +348,38 @@ impl Memory for UnifiedMemory { &self, namespace: Option<&str>, category: Option<&MemoryCategory>, - _session_id: Option<&str>, + session_id: Option<&str>, ) -> anyhow::Result> { - let ns = normalize_namespace(namespace); - let docs = self - .list_documents(Some(ns)) - .await - .map_err(anyhow::Error::msg)?; - let mut out = Vec::new(); - let items = docs - .get("documents") - .and_then(serde_json::Value::as_array) - .cloned() - .unwrap_or_default(); - for (idx, d) in items.into_iter().enumerate() { - let cat = category.cloned().unwrap_or(MemoryCategory::Core); - let taint_str = d - .get("taint") - .and_then(serde_json::Value::as_str) - .unwrap_or("internal"); - out.push(MemoryEntry { - id: d - .get("documentId") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(), - key: d - .get("key") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(), - content: d - .get("title") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(), - namespace: Some(ns.to_string()), - category: cat, - timestamp: format!("idx-{idx}"), - session_id: None, + let ns = UnifiedMemory::sanitize_namespace(normalize_namespace(namespace)); + let conn = self.conn.lock(); + let mut stmt = conn.prepare( + "SELECT document_id, key, content, category, session_id, updated_at, taint + FROM memory_docs WHERE namespace = ?1 ORDER BY updated_at DESC", + )?; + let rows = stmt.query_map(params![ns], |row| { + let stored_category: String = row.get(3)?; + Ok(MemoryEntry { + id: row.get(0)?, + key: row.get(1)?, + content: row.get(2)?, + namespace: Some(ns.clone()), + category: memory_category_from_stored(&stored_category), + session_id: row.get(4)?, + timestamp: timestamp_to_rfc3339(row.get(5)?), score: None, - taint: crate::openhuman::memory::MemoryTaint::from_db_str(taint_str), - }); + taint: crate::openhuman::memory::MemoryTaint::from_db_str( + &row.get::<_, String>(6)?, + ), + }) + })?; + let mut entries = rows.collect::>>()?; + if let Some(category) = category { + entries.retain(|entry| &entry.category == category); } - Ok(out) + if let Some(session_id) = session_id { + entries.retain(|entry| entry.session_id.as_deref() == Some(session_id)); + } + Ok(entries) } async fn forget(&self, namespace: &str, key: &str) -> anyhow::Result { @@ -454,10 +442,6 @@ impl Memory for UnifiedMemory { async fn health_check(&self) -> bool { self.workspace_dir.exists() && self.db_path.exists() } - - fn sqlite_conn(&self) -> Option>> { - Some(std::sync::Arc::clone(&self.conn)) - } } #[cfg(test)] @@ -500,8 +484,7 @@ mod tests { let in_b = mem.list(Some("ns_b"), None, None).await.unwrap(); assert_eq!(in_b.len(), 1); - // `list` currently maps title → content (pre-Phase-A quirk preserved). - // What matters here is namespace isolation: ns_a rows must not appear. + assert_eq!(in_b[0].content, "b"); assert!(in_b.iter().all(|e| e.namespace.as_deref() == Some("ns_b"))); // Forget in ns_a must not delete ns_b's row @@ -510,6 +493,44 @@ mod tests { assert!(mem.get("ns_a", "k1").await.unwrap().is_none()); } + #[tokio::test] + async fn list_returns_stored_fields_and_applies_category_and_session_filters() { + let (_tmp, mem) = fresh_mem(); + mem.store( + "rules", + "core", + "core body", + MemoryCategory::Core, + Some("session-a"), + ) + .await + .unwrap(); + mem.store( + "rules", + "procedure", + "procedure body", + MemoryCategory::Daily, + Some("session-b"), + ) + .await + .unwrap(); + + let entries = mem + .list( + Some("rules"), + Some(&MemoryCategory::Daily), + Some("session-b"), + ) + .await + .unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].key, "procedure"); + assert_eq!(entries[0].content, "procedure body"); + assert_eq!(entries[0].category, MemoryCategory::Daily); + assert_eq!(entries[0].session_id.as_deref(), Some("session-b")); + assert!(!entries[0].timestamp.starts_with("idx-")); + } + #[tokio::test] async fn namespace_summaries_counts_per_namespace() { let (_tmp, mem) = fresh_mem(); diff --git a/src/openhuman/memory_store/safety/mod.rs b/src/openhuman/memory_store/safety/mod.rs index 86eddb907..d224f700a 100644 --- a/src/openhuman/memory_store/safety/mod.rs +++ b/src/openhuman/memory_store/safety/mod.rs @@ -1,243 +1,28 @@ -//! Secret-detection and redaction helpers for memory writes. +//! Secret-detection and redaction for memory writes — thin host shim over +//! `tinycortex::memory::store::safety` (W3). //! -//! This module is intentionally conservative: it prefers false positives over -//! leaking credentials into long-lived memory stores. -//! -//! Personal-PII redaction (national IDs, international phone) is layered on -//! top via the [`pii`] submodule and runs from [`sanitize_text`], so every -//! call site that scrubs secrets also scrubs PII. +//! The conservative secret + PII scrubbers (`has_likely_secret`, +//! `has_likely_pii`, `sanitize_text`, `sanitize_json`) + the +//! `SanitizationReport`/`Sanitized` types are the crate's — now including the +//! full multilingual national-ID PII module (ported into the crate so the crate +//! `sanitize_text` matches this host's byte-for-byte). The host keeps only +//! [`sanitize_document_input`], which scrubs the host-specific +//! [`NamespaceDocumentInput`] shape by delegating each field to the crate +//! scrubbers. The retained test suite doubles as a byte-parity guard: it asserts +//! the crate scrubber still redacts every secret/PII pattern the host relied on. pub mod pii; -use once_cell::sync::Lazy; -use regex::Regex; -use serde_json::Value; - use crate::openhuman::memory_store::types::NamespaceDocumentInput; -const REDACTED_SECRET: &str = "[REDACTED_SECRET]"; -const REDACTED_PRIVATE_KEY: &str = "[REDACTED_PRIVATE_KEY]"; -const MAX_JSON_SANITIZE_DEPTH: usize = 128; - -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct SanitizationReport { - pub text_redactions: usize, - pub key_redactions: usize, - pub blocked_secret_hits: usize, - pub depth_redactions: usize, - pub pii_redactions: usize, -} - -impl SanitizationReport { - pub fn changed(&self) -> bool { - self.text_redactions > 0 - || self.key_redactions > 0 - || self.blocked_secret_hits > 0 - || self.depth_redactions > 0 - || self.pii_redactions > 0 - } - - pub fn merge(self, rhs: Self) -> Self { - Self { - text_redactions: self.text_redactions + rhs.text_redactions, - key_redactions: self.key_redactions + rhs.key_redactions, - blocked_secret_hits: self.blocked_secret_hits + rhs.blocked_secret_hits, - depth_redactions: self.depth_redactions + rhs.depth_redactions, - pii_redactions: self.pii_redactions + rhs.pii_redactions, - } - } -} - -#[derive(Debug, Clone)] -pub struct Sanitized { - pub value: T, - pub report: SanitizationReport, -} - -static BLOCK_PATTERNS: Lazy> = Lazy::new(|| { - vec![ - // Generic PEM private key blocks, including multiline bodies. - Regex::new( - r"(?is)-----BEGIN(?: [A-Z]+)? PRIVATE KEY-----.*?-----END(?: [A-Z]+)? PRIVATE KEY-----", - ) - .expect("valid private key block"), - // SSH private key blocks. - Regex::new(r"(?is)-----BEGIN OPENSSH PRIVATE KEY-----.*?-----END OPENSSH PRIVATE KEY-----") - .expect("valid openssh private key block"), - // PGP private key blocks. - Regex::new( - r"(?is)-----BEGIN PGP PRIVATE KEY BLOCK-----.*?-----END PGP PRIVATE KEY BLOCK-----", - ) - .expect("valid pgp private key block"), - ] -}); - -static REDACTION_PATTERNS: Lazy> = Lazy::new(|| { - vec![ - ( - Regex::new(r"(?i)(bearer\s+)[A-Za-z0-9._~+/=-]{8,}") - .expect("valid bearer redaction"), - "${1}[REDACTED]", - ), - ( - Regex::new(r#"(?i)(api[_-]?key\s*[=:\s]\s*["']?)[^\s"']+"#) - .expect("valid api key redaction"), - "${1}[REDACTED]", - ), - ( - Regex::new( - r#"(?i)\b(token|access[_-]?token|refresh[_-]?token|client[_-]?secret|password|secret)\b\s*[=:\s]\s*["']?[^\s"'&]+"#, - ) - .expect("valid token redaction"), - "[REDACTED]", - ), - ( - Regex::new(r"\bsk-[A-Za-z0-9]{20,}\b").expect("valid openai key redaction"), - "[REDACTED]", - ), - ( - Regex::new(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b").expect("valid github token redaction"), - "[REDACTED]", - ), - ( - Regex::new(r"\bAKIA[0-9A-Z]{16}\b").expect("valid aws key redaction"), - "[REDACTED]", - ), - ( - Regex::new(r"\bASIA[0-9A-Z]{16}\b").expect("valid aws sts key redaction"), - "[REDACTED]", - ), - ( - Regex::new(r#"\b(?:aws_)?secret(?:_access)?_key\b\s*[=:\s]\s*["']?[A-Za-z0-9/+=]{16,}"#) - .expect("valid aws secret key redaction"), - "[REDACTED]", - ), - ( - Regex::new(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9._-]{8,}\.[A-Za-z0-9._-]{8,}\b") - .expect("valid jwt redaction"), - "[REDACTED]", - ), - ( - // Common OAuth token artifacts in URLs and payloads. - Regex::new( - r#"(?i)\b(access_token|refresh_token|id_token|authorization_code|code_verifier|code_challenge)\b\s*[=:\s]\s*["']?[^\s"'&]+"#, - ) - .expect("valid oauth token redaction"), - "[REDACTED]", - ), - ( - // Google API key pattern. - Regex::new(r"\bAIza[0-9A-Za-z\-_]{35}\b").expect("valid google api key redaction"), - "[REDACTED]", - ), - ( - // Anthropic key pattern. - Regex::new(r"\bsk-ant-[A-Za-z0-9\-_]{16,}\b").expect("valid anthropic key redaction"), - "[REDACTED]", - ), - ( - // OpenAI project/org scoped keys and legacy variants. - Regex::new(r"\bsk-(?:proj|org)-[A-Za-z0-9\-_]{12,}\b") - .expect("valid openai scoped key redaction"), - "[REDACTED]", - ), - ( - // Stripe secret/restricted keys. - Regex::new(r"\b(?:sk|rk)_(?:live|test)_[A-Za-z0-9]{16,}\b") - .expect("valid stripe key redaction"), - "[REDACTED]", - ), - ( - // Slack tokens (bot/user/app/config). - Regex::new(r"\bxox(?:a|b|p|s|r)-[A-Za-z0-9-]{10,}\b") - .expect("valid slack token redaction"), - "[REDACTED]", - ), - ( - // GitHub fine-grained/pat/user tokens beyond gh[pousr]_. - Regex::new(r"\bgithub_pat_[A-Za-z0-9_]{20,}\b") - .expect("valid github pat redaction"), - "[REDACTED]", - ), - ( - // GitLab personal access token. - Regex::new(r"\bglpat-[A-Za-z0-9\-_]{16,}\b") - .expect("valid gitlab pat redaction"), - "[REDACTED]", - ), - ( - // NPM auth token. - Regex::new(r"\bnpm_[A-Za-z0-9]{20,}\b").expect("valid npm token redaction"), - "[REDACTED]", - ), - ( - // SendGrid API key. - Regex::new(r"\bSG\.[A-Za-z0-9_\-]{16,}\.[A-Za-z0-9_\-]{16,}\b") - .expect("valid sendgrid key redaction"), - "[REDACTED]", - ), - ( - // Twilio API key SID. - Regex::new(r"\bSK[a-fA-F0-9]{32}\b").expect("valid twilio sid redaction"), - "[REDACTED]", - ), - ( - // Azure Storage account key in key=value style payloads. - Regex::new(r"(?i)\bAccountKey\b\s*=\s*[A-Za-z0-9+/=]{20,}") - .expect("valid azure account key redaction"), - "[REDACTED]", - ), - ( - // Generic Authorization header values beyond Bearer. - Regex::new(r"(?i)(authorization\s*[:=]\s*)(?:basic|bearer|token)\s+[A-Za-z0-9._~+/=-]{8,}") - .expect("valid authorization header redaction"), - "${1}[REDACTED]", - ), - ] -}); - -pub fn has_likely_secret(value: &str) -> bool { - if BLOCK_PATTERNS.iter().any(|pattern| pattern.is_match(value)) { - return true; - } - REDACTION_PATTERNS - .iter() - .any(|(pattern, _)| pattern.is_match(value)) -} - -pub fn sanitize_text(value: &str) -> Sanitized { - let mut out = value.to_string(); - let mut report = SanitizationReport::default(); - - for pattern in BLOCK_PATTERNS.iter() { - let hits = pattern.find_iter(&out).count(); - if hits > 0 { - report.blocked_secret_hits += hits; - out = pattern.replace_all(&out, REDACTED_PRIVATE_KEY).into_owned(); - } - } - - for (pattern, replacement) in REDACTION_PATTERNS.iter() { - let hits = pattern.find_iter(&out).count(); - if hits > 0 { - report.text_redactions += hits; - out = pattern.replace_all(&out, *replacement).into_owned(); - } - } - - let pii_sanitized = pii::redact_pii(&out); - report = report.merge(pii_sanitized.report); - - Sanitized { - value: pii_sanitized.value, - report, - } -} - -pub fn sanitize_json(value: &Value) -> Sanitized { - sanitize_json_inner(value, 0) -} +pub use tinycortex::memory::store::safety::{ + has_likely_pii, has_likely_secret, sanitize_json, sanitize_text, SanitizationReport, Sanitized, +}; +/// Scrub a namespace-document input, field by field, via the crate scrubbers. +/// +/// Sanitization is content-cleaning only; provenance `taint` survives untouched +/// so the write gate's taint check still sees the real source signal. pub fn sanitize_document_input(input: NamespaceDocumentInput) -> Sanitized { let mut report = SanitizationReport::default(); @@ -269,102 +54,23 @@ pub fn sanitize_document_input(input: NamespaceDocumentInput) -> Sanitized Sanitized { - if depth >= MAX_JSON_SANITIZE_DEPTH { - return Sanitized { - value: Value::String(REDACTED_SECRET.to_string()), - report: SanitizationReport { - depth_redactions: 1, - ..SanitizationReport::default() - }, - }; - } - - match value { - Value::Object(map) => { - let mut out = serde_json::Map::new(); - let mut report = SanitizationReport::default(); - for (key, value) in map { - if is_sensitive_key(key) { - report.key_redactions += 1; - out.insert(key.clone(), Value::String(REDACTED_SECRET.to_string())); - continue; - } - let sanitized = sanitize_json_inner(value, depth + 1); - report = report.merge(sanitized.report); - out.insert(key.clone(), sanitized.value); - } - Sanitized { - value: Value::Object(out), - report, - } - } - Value::Array(items) => { - let mut out = Vec::with_capacity(items.len()); - let mut report = SanitizationReport::default(); - for item in items { - let sanitized = sanitize_json_inner(item, depth + 1); - report = report.merge(sanitized.report); - out.push(sanitized.value); - } - Sanitized { - value: Value::Array(out), - report, - } - } - Value::String(value) => { - let sanitized = sanitize_text(value); - Sanitized { - value: Value::String(sanitized.value), - report: sanitized.report, - } - } - _ => Sanitized { - value: value.clone(), - report: SanitizationReport::default(), - }, - } -} - -fn is_sensitive_key(key: &str) -> bool { - let normalized: String = key - .chars() - .filter(|c| c.is_ascii_alphanumeric()) - .map(|c| c.to_ascii_lowercase()) - .collect(); - - matches!( - normalized.as_str(), - "apikey" - | "token" - | "accesstoken" - | "refreshtoken" - | "authorization" - | "password" - | "secret" - | "clientsecret" - ) || normalized.ends_with("token") - || normalized.ends_with("apikey") - || normalized.ends_with("clientsecret") - || normalized.contains("password") - || normalized.contains("secret") - || normalized.ends_with("key") -} - #[cfg(test)] mod tests { + //! Byte-parity guard over the crate scrubber: every secret/PII pattern the + //! host used to redact must still be redacted after the port. use super::*; use serde_json::json; + const REDACTED_SECRET: &str = "[REDACTED_SECRET]"; + const REDACTED_PRIVATE_KEY: &str = "[REDACTED_PRIVATE_KEY]"; + const MAX_JSON_SANITIZE_DEPTH: usize = 128; + #[test] fn sanitize_text_redacts_bearer_and_openai_key() { let input = "Authorization: Bearer abcdefghijklmnop and sk-1234567890123456789012345"; @@ -386,13 +92,9 @@ mod tests { fn sanitize_json_redacts_sensitive_keys_and_nested_strings() { let input = json!({ "token": "abc123", - "nested": { - "notes": "Bearer supersecretvalue", - "ok": "hello" - }, + "nested": { "notes": "Bearer supersecretvalue", "ok": "hello" }, "arr": ["sk-1234567890123456789012345", "safe"] }); - let sanitized = sanitize_json(&input); assert_eq!(sanitized.value["token"], json!(REDACTED_SECRET)); assert_eq!(sanitized.value["nested"]["ok"], json!("hello")); @@ -407,12 +109,9 @@ mod tests { #[test] fn sanitize_json_redacts_common_sensitive_key_variants() { let input = json!({ - "db_password": "p@ss", - "secret_key": "abc123", - "api_secret": "def456", - "monkey": "banana" + "db_password": "p@ss", "secret_key": "abc123", + "api_secret": "def456", "monkey": "banana" }); - let sanitized = sanitize_json(&input); assert_eq!(sanitized.value["db_password"], json!(REDACTED_SECRET)); assert_eq!(sanitized.value["secret_key"], json!(REDACTED_SECRET)); @@ -462,7 +161,6 @@ mod tests { #[test] fn sanitize_text_also_redacts_pii_after_secrets() { - // Mix of secret + multilingual PII — both should be scrubbed in one pass. let input = "Token sk-abcdefghijklmnopqrstuvwxyz; CPF 111.444.777-35; phone +15551234567"; let sanitized = sanitize_text(input); assert!(!sanitized.value.contains("sk-abcdefghijklmnopqrstuvwxyz")); @@ -498,7 +196,6 @@ mod tests { for _ in 0..(MAX_JSON_SANITIZE_DEPTH + 2) { nested = json!({ "nested": nested }); } - let sanitized = sanitize_json(&nested); assert!(sanitized.report.depth_redactions >= 1); assert!(sanitized diff --git a/src/openhuman/memory_store/safety/pii.rs b/src/openhuman/memory_store/safety/pii.rs index c6ea8792d..e67c8cb48 100644 --- a/src/openhuman/memory_store/safety/pii.rs +++ b/src/openhuman/memory_store/safety/pii.rs @@ -1,1091 +1,8 @@ -//! Multilingual personal-PII redaction (national IDs, financial identifiers, -//! international phone) — on-device, regex + checksum only, zero network. +//! Personal-PII detection — thin host re-export of the crate scrubber (W3). //! -//! ## 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. +//! The full multilingual national-ID PII module (checksum-gated patterns + +//! Unicode normalization) now lives in `tinycortex::memory::store::safety::pii`; +//! content scrubbing runs inside the crate `sanitize_text`. Host consumers keep +//! their `safety::pii::has_likely_pii` import path. -use once_cell::sync::Lazy; -use regex::{Regex, RegexSet}; - -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: Lazy = - Lazy::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: Lazy = Lazy::new(|| Regex::new(r"\b\d{11}\b").expect("cpf bare")); - -// Brazilian CNPJ, formatted: NN.NNN.NNN/NNNN-NN -static CNPJ_FMT_RE: Lazy = - Lazy::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: Lazy = Lazy::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: Lazy = Lazy::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: Lazy = - Lazy::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: Lazy = Lazy::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: Lazy = Lazy::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: Lazy = Lazy::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: Lazy = Lazy::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: Lazy = - Lazy::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: Lazy = - Lazy::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: Lazy = - Lazy::new(|| Regex::new(r"\b\d{4}[\s\-]\d{4}[\s\-]\d{4}\b").expect("aadhaar formatted")); -static AADHAAR_KW_RE: Lazy = Lazy::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: Lazy = - Lazy::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: Lazy = - Lazy::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: Lazy = Lazy::new(|| Regex::new(r"(?i)\b\d{8}[A-Z]\b").expect("dni")); -static NIE_RE: Lazy = Lazy::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: Lazy = Lazy::new(|| Regex::new(r"\b\d{6}-[1-4]\d{6}\b").expect("rrn")); - -// 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: Lazy = Lazy::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() -} - -// ---------- 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_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(""); - } -} +pub use tinycortex::memory::store::safety::pii::has_likely_pii; diff --git a/src/openhuman/memory_store/trees/hotness.rs b/src/openhuman/memory_store/trees/hotness.rs index d3d2cb2a3..aacd5dafb 100644 --- a/src/openhuman/memory_store/trees/hotness.rs +++ b/src/openhuman/memory_store/trees/hotness.rs @@ -1,178 +1,33 @@ -//! Entity-hotness counter persistence (`mem_tree_entity_hotness` table). -//! -//! Previously at `memory_store::trees_topic::store`. Folded into `trees/` -//! because hotness is the only state that differentiates topic trees from -//! source/global trees, and even then it's a side-table — not a tree row. -//! Keeping it next to tree storage makes "what does memory_store know about -//! trees?" a single-directory answer. +//! `Config` adapters for tinycortex entity-hotness persistence. -use anyhow::{Context, Result}; -use chrono::Utc; -use rusqlite::{params, OptionalExtension}; +use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; use crate::openhuman::memory_store::trees::types::HotnessCounters; -/// Fetch the hotness row for `entity_id`, or `None` if the entity has -/// never been seen. Callers usually want [`get_or_fresh`] instead. +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} + pub fn get(config: &Config, entity_id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT entity_id, mention_count_30d, distinct_sources, last_seen_ms, - query_hits_30d, graph_centrality, ingests_since_check, - last_hotness, last_updated_ms - FROM mem_tree_entity_hotness WHERE entity_id = ?1", - )?; - let row = stmt - .query_row(params![entity_id], row_to_counters) - .optional() - .context("failed to query mem_tree_entity_hotness")?; - Ok(row) - }) + tinycortex::memory::tree::store::hotness::get(&engine_config(config), entity_id) } -/// Fetch the hotness row, or return a fresh (all-zero) row if the entity -/// has never been seen. The fresh row is NOT persisted — callers must -/// [`upsert`] it explicitly after bumping counters. pub fn get_or_fresh(config: &Config, entity_id: &str) -> Result { - match get(config, entity_id)? { - Some(c) => Ok(c), - None => Ok(HotnessCounters::fresh( - entity_id, - Utc::now().timestamp_millis(), - )), - } + tinycortex::memory::tree::store::hotness::get_or_fresh(&engine_config(config), entity_id) } -/// Upsert the full counter row. Idempotent on `entity_id`. pub fn upsert(config: &Config, counters: &HotnessCounters) -> Result<()> { - with_connection(config, |conn| { - conn.execute( - "INSERT INTO mem_tree_entity_hotness ( - entity_id, mention_count_30d, distinct_sources, last_seen_ms, - query_hits_30d, graph_centrality, ingests_since_check, - last_hotness, last_updated_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) - ON CONFLICT(entity_id) DO UPDATE SET - mention_count_30d = excluded.mention_count_30d, - distinct_sources = excluded.distinct_sources, - last_seen_ms = excluded.last_seen_ms, - query_hits_30d = excluded.query_hits_30d, - graph_centrality = excluded.graph_centrality, - ingests_since_check = excluded.ingests_since_check, - last_hotness = excluded.last_hotness, - last_updated_ms = excluded.last_updated_ms", - params![ - counters.entity_id, - counters.mention_count_30d, - counters.distinct_sources, - counters.last_seen_ms, - counters.query_hits_30d, - counters.graph_centrality, - counters.ingests_since_check, - counters.last_hotness, - counters.last_updated_ms, - ], - ) - .with_context(|| { - format!( - "failed to upsert mem_tree_entity_hotness for {}", - counters.entity_id - ) - })?; - Ok(()) - }) + tinycortex::memory::tree::store::hotness::upsert(&engine_config(config), counters) } -/// Count `(node_id) → DISTINCT tree_id` in the entity index for `entity_id`. pub fn distinct_sources_for(config: &Config, entity_id: &str) -> Result { - with_connection(config, |conn| { - let n: i64 = conn - .query_row( - "SELECT COUNT(DISTINCT tree_id) - FROM mem_tree_entity_index - WHERE entity_id = ?1 AND tree_id IS NOT NULL", - params![entity_id], - |r| r.get(0), - ) - .context("failed to count distinct sources")?; - Ok(n.max(0) as u32) - }) + tinycortex::memory::tree::store::hotness::distinct_sources_for( + &engine_config(config), + entity_id, + ) } -/// Test / diagnostic helper. pub fn count(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn - .query_row("SELECT COUNT(*) FROM mem_tree_entity_hotness", [], |r| { - r.get(0) - }) - .context("failed to count mem_tree_entity_hotness")?; - Ok(n.max(0) as u64) - }) -} - -fn row_to_counters(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(HotnessCounters { - entity_id: row.get(0)?, - mention_count_30d: row.get::<_, i64>(1)?.max(0) as u32, - distinct_sources: row.get::<_, i64>(2)?.max(0) as u32, - last_seen_ms: row.get(3)?, - query_hits_30d: row.get::<_, i64>(4)?.max(0) as u32, - graph_centrality: row.get(5)?, - ingests_since_check: row.get::<_, i64>(6)?.max(0) as u32, - last_hotness: row.get(7)?, - last_updated_ms: row.get(8)?, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) - } - - #[test] - fn get_missing_is_none() { - let (_tmp, cfg) = test_config(); - assert!(get(&cfg, "email:alice@example.com").unwrap().is_none()); - } - - #[test] - fn get_or_fresh_returns_zero_row() { - let (_tmp, cfg) = test_config(); - let c = get_or_fresh(&cfg, "email:alice@example.com").unwrap(); - assert_eq!(c.entity_id, "email:alice@example.com"); - assert_eq!(c.mention_count_30d, 0); - assert_eq!(c.distinct_sources, 0); - assert!(c.last_hotness.is_none()); - assert_eq!(count(&cfg).unwrap(), 0); - } - - #[test] - fn upsert_round_trip() { - let (_tmp, cfg) = test_config(); - let c = HotnessCounters { - entity_id: "email:alice@example.com".into(), - mention_count_30d: 12, - distinct_sources: 3, - last_seen_ms: Some(1_700_000_000_000), - query_hits_30d: 2, - graph_centrality: Some(0.25), - ingests_since_check: 40, - last_hotness: Some(9.5), - last_updated_ms: 1_700_000_123_000, - }; - upsert(&cfg, &c).unwrap(); - let got = get(&cfg, &c.entity_id).unwrap().unwrap(); - assert_eq!(got, c); - assert_eq!(count(&cfg).unwrap(), 1); - } + tinycortex::memory::tree::store::hotness::count(&engine_config(config)) } diff --git a/src/openhuman/memory_store/trees/registry.rs b/src/openhuman/memory_store/trees/registry.rs index 061002414..7f187ec33 100644 --- a/src/openhuman/memory_store/trees/registry.rs +++ b/src/openhuman/memory_store/trees/registry.rs @@ -1,164 +1,19 @@ -//! Kind-parameterized tree registry helpers. -//! -//! Trees live in the `mem_tree_trees` table keyed by [`TreeKind`]. The -//! generic get-or-create dance with UNIQUE-race recovery lives in -//! [`memory_tree::tree::registry::get_or_create_tree`] (logic file — stays -//! in memory_tree). This module hosts the generic list / archive helpers. -//! -//! The global (time-axis) and topic (subject-axis) trees were removed; only -//! source trees are created now, so there are no per-kind creation wrappers -//! here anymore — [`TreeKind::Global`]/[`TreeKind::Topic`] remain only as -//! inert serialization plumbing for reading legacy rows during the one-shot -//! purge migration. +//! `Config` adapters for tinycortex's tree registry. -use anyhow::{Context, Result}; -use chrono::Utc; -use rusqlite::params; +use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; use crate::openhuman::memory_store::trees::types::{Tree, TreeKind}; -/// List every tree of the given kind, ordered by creation time ascending. +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} + pub fn list_trees_by_kind(config: &Config, kind: TreeKind) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - FROM mem_tree_trees - WHERE kind = ?1 - ORDER BY created_at_ms ASC", - )?; - let rows = stmt - .query_map(params![kind.as_str()], row_to_tree_loose)? - .collect::>>() - .with_context(|| format!("failed to list trees of kind {}", kind.as_str()))?; - Ok(rows) - }) + tinycortex::memory::tree::store::list_trees_by_kind(&engine_config(config), kind) } -/// Flip a tree's status to `archived`. Idempotent. Existing rows remain -/// queryable; new leaves will not be routed to this tree until it is -/// manually unarchived (which is not currently a primitive). pub fn archive_tree(config: &Config, tree_id: &str) -> Result<()> { - use crate::openhuman::memory_store::trees::types::TreeStatus; - with_connection(config, |conn| { - let n = conn - .execute( - "UPDATE mem_tree_trees - SET status = ?1 - WHERE id = ?2", - params![TreeStatus::Archived.as_str(), tree_id], - ) - .with_context(|| format!("failed to archive tree {}", tree_id))?; - log::debug!("[trees::registry] archive_tree id={} rows={}", tree_id, n); - Ok(()) - }) -} - -fn row_to_tree_loose(row: &rusqlite::Row<'_>) -> rusqlite::Result { - use crate::openhuman::memory_store::trees::types::TreeStatus; - use chrono::TimeZone; - let id: String = row.get(0)?; - let kind_s: String = row.get(1)?; - let scope: String = row.get(2)?; - let root_id: Option = row.get(3)?; - let max_level: i64 = row.get(4)?; - let status_s: String = row.get(5)?; - let created_ms: i64 = row.get(6)?; - let last_sealed_ms: Option = row.get(7)?; - let kind = TreeKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, e.into()) - })?; - let status = TreeStatus::parse(&status_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, e.into()) - })?; - Ok(Tree { - id, - kind, - scope, - root_id, - max_level: max_level.max(0) as u32, - status, - created_at: Utc.timestamp_millis_opt(created_ms).unwrap(), - last_sealed_at: last_sealed_ms.and_then(|ms| Utc.timestamp_millis_opt(ms).single()), - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::trees::store::insert_tree; - use crate::openhuman::memory_store::trees::types::TreeStatus; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) - } - - fn sample_tree(id: &str, kind: TreeKind, scope: &str) -> Tree { - Tree { - id: id.into(), - kind, - scope: scope.into(), - root_id: Some("root-1".into()), - max_level: 2, - status: TreeStatus::Active, - created_at: Utc::now(), - last_sealed_at: None, - } - } - - #[test] - fn list_trees_by_kind_returns_only_requested_kind() { - let (_tmp, cfg) = test_config(); - insert_tree( - &cfg, - &sample_tree("source-1", TreeKind::Source, "chat:slack:#eng"), - ) - .unwrap(); - insert_tree( - &cfg, - &sample_tree("topic-1", TreeKind::Topic, "person:alice"), - ) - .unwrap(); - insert_tree( - &cfg, - &sample_tree("source-2", TreeKind::Source, "chat:discord:#ops"), - ) - .unwrap(); - - let source_ids: Vec = list_trees_by_kind(&cfg, TreeKind::Source) - .unwrap() - .into_iter() - .map(|tree| tree.id) - .collect(); - assert_eq!( - source_ids, - vec!["source-1".to_string(), "source-2".to_string()] - ); - - let topic_ids: Vec = list_trees_by_kind(&cfg, TreeKind::Topic) - .unwrap() - .into_iter() - .map(|tree| tree.id) - .collect(); - assert_eq!(topic_ids, vec!["topic-1".to_string()]); - } - - #[test] - fn archive_tree_flips_status_to_archived() { - let (_tmp, cfg) = test_config(); - let tree = sample_tree("source-arch", TreeKind::Source, "chat:slack:#x"); - insert_tree(&cfg, &tree).unwrap(); - - archive_tree(&cfg, "source-arch").unwrap(); - - let archived = list_trees_by_kind(&cfg, TreeKind::Source).unwrap(); - assert_eq!(archived.len(), 1); - assert_eq!(archived[0].status, TreeStatus::Archived); - } + log::debug!("[memory:trees] archive tree_id={tree_id}"); + tinycortex::memory::tree::store::archive_tree(&engine_config(config), tree_id) } diff --git a/src/openhuman/memory_store/trees/store.rs b/src/openhuman/memory_store/trees/store.rs index 415687f56..27141d520 100644 --- a/src/openhuman/memory_store/trees/store.rs +++ b/src/openhuman/memory_store/trees/store.rs @@ -1,143 +1,38 @@ -//! SQLite-backed persistence for Phase 3a summary trees (#709). -//! -//! Three tables (schema lives in the sibling `tree::store::SCHEMA`): -//! - `mem_tree_trees` — one row per tree (kind, scope, root, max_level) -//! - `mem_tree_summaries` — one row per sealed summary node (immutable) -//! - `mem_tree_buffers` — one row per unsealed frontier `(tree_id, level)` -//! -//! All timestamps are stored as milliseconds since the Unix epoch so we -//! share the epoch convention with `mem_tree_chunks`. Writes are serialised -//! through the sibling `tree::store::with_connection` so we inherit its -//! busy-timeout, WAL, and schema-init behaviour. -//! -//! Phase 4 (#710) adds a nullable `embedding` blob on -//! `mem_tree_summaries` — packed little-endian `f32` vectors via -//! [`crate::openhuman::memory_tree::score::embed::pack_embedding`]. New -//! writes populate it via [`insert_summary_tx`]; reads decode it when -//! present. +//! `Config` and transaction adapters for tinycortex tree persistence. use std::collections::HashMap; -use anyhow::{Context, Result}; -use chrono::{DateTime, TimeZone, Utc}; -use rusqlite::{params, Connection, OptionalExtension, Transaction}; +use anyhow::Result; +use chrono::{DateTime, Utc}; +use rusqlite::{Connection, Transaction}; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; use crate::openhuman::memory_store::content::StagedSummary; -use crate::openhuman::memory_store::trees::types::{ - Buffer, SummaryNode, Tree, TreeKind, TreeStatus, -}; -use crate::openhuman::memory_tree::score::embed::{decode_optional_blob, pack_checked}; +use crate::openhuman::memory_store::trees::types::{Buffer, SummaryNode, Tree, TreeKind}; -fn ms_to_utc(ms: i64) -> rusqlite::Result> { - Utc.timestamp_millis_opt(ms).single().ok_or_else(|| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Integer, - format!("invalid timestamp ms {ms}").into(), - ) - }) +pub(crate) use tinycortex::memory::tree::store::TreeCascadeDeletion; + +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -// ── Tree rows ─────────────────────────────────────────────────────────── - -/// Insert a new tree row. Fails if `(kind, scope)` already exists; callers -/// that want "get or create" semantics should go through the `registry`. pub fn insert_tree(config: &Config, tree: &Tree) -> Result<()> { - with_connection(config, |conn| insert_tree_conn(conn, tree)) + tinycortex::memory::tree::store::insert_tree(&engine_config(config), tree) } pub(crate) fn insert_tree_conn(conn: &Connection, tree: &Tree) -> Result<()> { - conn.execute( - "INSERT INTO mem_tree_trees ( - id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", - params![ - tree.id, - tree.kind.as_str(), - tree.scope, - tree.root_id, - tree.max_level, - tree.status.as_str(), - tree.created_at.timestamp_millis(), - tree.last_sealed_at.map(|t| t.timestamp_millis()), - ], - ) - .with_context(|| format!("Failed to insert tree id={}", tree.id))?; - Ok(()) + tinycortex::memory::tree::store::insert_tree_conn(conn, tree) } -/// Hard-delete one tree and every dependent row, within an existing tx. -/// -/// Cascade order mirrors [`crate::openhuman::memory_store::chunks::store`]'s -/// global/topic purge: summary sidecars (`summary_embeddings` / -/// `summary_reembed_skipped`, keyed by `summary_id`) first, then -/// `entity_index` + `buffers` (keyed by `tree_id`), then the `summaries`, -/// then the tree row. Used by the chunk-delete path when a source's last -/// chunk is removed, so the now-contentless summary tree (and its unsealed -/// buffer) doesn't outlive the data it summarised. Returns the number of -/// summary rows removed. pub(crate) fn delete_tree_cascade_tx( tx: &Transaction<'_>, tree_id: &str, ) -> Result { - // Collect the on-disk content-file paths BEFORE deleting the summary rows - // — sealed summaries stage their body to `content_path` under the memory - // tree content root (see `bucket_seal::seal_one_level` → `stage_summary`). - // The caller removes these files after the tx commits (mirroring - // `remove_chunk_content_files`), so a `clear_memory` delete doesn't leave - // the summarised text orphaned on disk. - let content_paths: Vec = { - 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 rows = stmt.query_map(params![tree_id], |row| row.get::<_, String>(0))?; - rows.collect::>>() - .context("collect summary content paths for tree cascade delete")? - }; - - tx.execute( - "DELETE FROM mem_tree_summary_embeddings WHERE summary_id IN \ - (SELECT id FROM mem_tree_summaries WHERE tree_id = ?1)", - params![tree_id], - )?; - tx.execute( - "DELETE FROM mem_tree_summary_reembed_skipped WHERE summary_id IN \ - (SELECT id FROM mem_tree_summaries WHERE tree_id = ?1)", - params![tree_id], - )?; - tx.execute( - "DELETE FROM mem_tree_entity_index WHERE tree_id = ?1", - 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, - }) + tinycortex::memory::tree::store::delete_tree_cascade_tx(tx, tree_id) } -/// Outcome of [`delete_tree_cascade_tx`]: how many summary rows were removed -/// and the on-disk content-file paths the caller must delete post-commit. -pub(crate) struct TreeCascadeDeletion { - pub removed_summaries: usize, - pub content_paths: Vec, -} - -/// Fetch a tree by `(kind, scope)`. Returns `None` if no such tree exists. pub fn get_tree_by_scope(config: &Config, kind: TreeKind, scope: &str) -> Result> { - with_connection(config, |conn| get_tree_by_scope_conn(conn, kind, scope)) + tinycortex::memory::tree::store::get_tree_by_scope(&engine_config(config), kind, scope) } pub(crate) fn get_tree_by_scope_conn( @@ -145,118 +40,19 @@ pub(crate) fn get_tree_by_scope_conn( kind: TreeKind, scope: &str, ) -> Result> { - let mut stmt = conn.prepare( - "SELECT id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - FROM mem_tree_trees WHERE kind = ?1 AND scope = ?2", - )?; - let row = stmt - .query_row(params![kind.as_str(), scope], row_to_tree) - .optional() - .context("Failed to query tree by scope")?; - Ok(row) + tinycortex::memory::tree::store::get_tree_by_scope_conn(conn, kind, scope) } -/// Fetch a tree by primary key id. pub fn get_tree(config: &Config, id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - FROM mem_tree_trees WHERE id = ?1", - )?; - let row = stmt - .query_row(params![id], row_to_tree) - .optional() - .context("Failed to query tree by id")?; - Ok(row) - }) + tinycortex::memory::tree::store::get_tree(&engine_config(config), id) } -/// Defensive upper bound on the number of `?` placeholders per batched -/// `SELECT … WHERE id IN (?,?,…)` query. SQLite's compile-time -/// `SQLITE_MAX_VARIABLE_NUMBER` has been ≥ 32 766 since 3.32 — 500 leaves -/// a ~65× safety margin. The current call-site -/// (`memory_tree::tree::flush::flush_stale_buffers`) passes one tree_id -/// per stale L0 buffer, so a typical N (tens to low hundreds across all -/// connected sources) runs the loop exactly once. The window exists so -/// future callers with larger id slices do not blow up against a host -/// with a lower compile-time SQLite cap. No volume reduction: all input -/// ids in → all matching rows out; the merged `HashMap` is byte- -/// identical to one giant query. -const TREES_MAX_FETCH_BATCH: usize = 500; - -/// Fetch many trees by id in a single SQL round-trip per -/// [`TREES_MAX_FETCH_BATCH`] window. Replaces the per-id `get_tree` -/// loop inside paths like the cron-driven `flush_stale_buffers`, where -/// the previous code did one `SELECT … WHERE id = ?` per stale buffer. -/// Missing ids are silently absent from the map so callers can preserve -/// the existing "row missing → warn-and-skip" contract without an extra -/// `Ok(None)` sentinel per id. -pub fn get_trees_batch(config: &Config, tree_ids: &[String]) -> Result> { - if tree_ids.is_empty() { - return Ok(HashMap::new()); - } - log::debug!( - "[tree::store] get_trees_batch: ids={} max_batch={TREES_MAX_FETCH_BATCH}", - tree_ids.len() - ); - with_connection(config, |conn| { - let mut out: HashMap = HashMap::with_capacity(tree_ids.len()); - for window in tree_ids.chunks(TREES_MAX_FETCH_BATCH) { - // SAFETY (SQL injection): only the *count* of placeholders is - // interpolated into the query string — `?1,?2,…` are numbered - // bind slots, never the id values. Every id is bound via typed - // `rusqlite::ToSql` params below; nothing user-controlled is - // ever formatted into `sql`. Do NOT inline id values here. - let placeholders = (1..=window.len()) - .map(|i| format!("?{i}")) - .collect::>() - .join(","); - let sql = format!( - "SELECT id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - FROM mem_tree_trees - WHERE id IN ({placeholders})" - ); - let mut stmt = conn.prepare(&sql)?; - let params: Vec<&dyn rusqlite::ToSql> = - window.iter().map(|s| s as &dyn rusqlite::ToSql).collect(); - let rows = stmt - .query_map(params.as_slice(), row_to_tree)? - .collect::>>() - .context("Failed to collect trees batch")?; - for t in rows { - out.insert(t.id.clone(), t); - } - } - log::debug!( - "[tree::store] get_trees_batch: requested={} found={}", - tree_ids.len(), - out.len() - ); - Ok(out) - }) +pub fn get_trees_batch(config: &Config, ids: &[String]) -> Result> { + tinycortex::memory::tree::store::get_trees_batch(&engine_config(config), ids) } -/// List every tree of a given kind. Used by the global digest to enumerate -/// source trees, and by diagnostics. Rows come back ordered by `created_at_ms` -/// ASC so callers see a stable iteration order. pub fn list_trees_by_kind(config: &Config, kind: TreeKind) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, kind, scope, root_id, max_level, status, - created_at_ms, last_sealed_at_ms - FROM mem_tree_trees - WHERE kind = ?1 - ORDER BY created_at_ms ASC", - )?; - let rows = stmt - .query_map(params![kind.as_str()], row_to_tree)? - .collect::>>() - .context("Failed to collect trees by kind")?; - Ok(rows) - }) + tinycortex::memory::tree::store::list_trees_by_kind(&engine_config(config), kind) } pub(crate) fn update_tree_after_seal_tx( @@ -266,795 +62,172 @@ pub(crate) fn update_tree_after_seal_tx( max_level: u32, sealed_at: DateTime, ) -> Result<()> { - tx.execute( - "UPDATE mem_tree_trees - SET root_id = ?1, - max_level = ?2, - last_sealed_at_ms = ?3 - WHERE id = ?4", - params![root_id, max_level, sealed_at.timestamp_millis(), tree_id,], + tinycortex::memory::tree::store::update_tree_after_seal_tx( + tx, tree_id, root_id, max_level, sealed_at, ) - .with_context(|| format!("Failed to update tree {tree_id} after seal"))?; - Ok(()) } -fn row_to_tree(row: &rusqlite::Row<'_>) -> rusqlite::Result { - let id: String = row.get(0)?; - let kind_s: String = row.get(1)?; - let scope: String = row.get(2)?; - let root_id: Option = row.get(3)?; - let max_level: i64 = row.get(4)?; - let status_s: String = row.get(5)?; - let created_ms: i64 = row.get(6)?; - let last_sealed_ms: Option = row.get(7)?; - - let kind = TreeKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into()) - })?; - let status = TreeStatus::parse(&status_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(5, rusqlite::types::Type::Text, e.into()) - })?; - Ok(Tree { - id, - kind, - scope, - root_id, - max_level: max_level.max(0) as u32, - status, - created_at: ms_to_utc(created_ms)?, - last_sealed_at: last_sealed_ms.map(ms_to_utc).transpose()?, - }) -} - -// ── Summary nodes ─────────────────────────────────────────────────────── - -/// Insert a sealed summary. Immutable — the caller must generate a fresh -/// id per seal. Idempotent on the primary key so retries of the same seal -/// transaction don't double-insert. -/// -/// Phase 4 (#710): if `node.embedding` is `Some`, the packed vector is -/// written to the `embedding` blob column; `None` writes NULL so legacy -/// rows from Phases 1-3 (no embed) read back identically. -/// -/// Phase MD-content: if `staged` is `Some`, writes `content_path` and -/// `content_sha256` and truncates `content` to a ≤500-char preview. Callers -/// that have not yet staged the file pass `None`, in which case the full -/// `node.content` is stored (legacy behaviour). pub(crate) fn insert_summary_tx( tx: &Transaction<'_>, node: &SummaryNode, staged: Option<&StagedSummary>, model_signature: &str, ) -> Result<()> { - // #1574 write-side cutover: keep the dimension guard (fail the seal fast - // on a misconfigured embedder) but DO NOT write the legacy - // `mem_tree_summaries.embedding` column — the vector is persisted to the - // per-model sidecar below, in THIS tx so it commits atomically with the - // summary row. The legacy column is left NULL for the §7 migration. - if let Some(v) = node.embedding.as_deref() { - pack_checked(v) - .with_context(|| format!("validate embedding dims for summary id={}", node.id))?; - } - let embedding_blob: Option> = None; - - // Phase MD-content: when a staged file exists, truncate `content` to a - // ≤500-char plain-text preview (char boundary safe via chars().take(500)). - let (content_preview, content_path, content_sha256) = match staged { - Some(s) => { - let preview: String = node.content.chars().take(500).collect(); - ( - preview, - Some(s.content_path.clone()), - Some(s.content_sha256.clone()), - ) - } - None => (node.content.clone(), None, None), - }; - - tx.execute( - "INSERT OR IGNORE INTO mem_tree_summaries ( - id, tree_id, tree_kind, level, parent_id, - child_ids_json, content, token_count, - entities_json, topics_json, - time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, embedding, - content_path, content_sha256, - doc_id, version_ms - ) 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, - node.tree_kind.as_str(), - node.level, - node.parent_id, - serde_json::to_string(&node.child_ids)?, - content_preview, - node.token_count, - serde_json::to_string(&node.entities)?, - serde_json::to_string(&node.topics)?, - node.time_range_start.timestamp_millis(), - node.time_range_end.timestamp_millis(), - node.score, - node.sealed_at.timestamp_millis(), - node.deleted as i64, - embedding_blob, - content_path, - content_sha256, - node.doc_id, - node.version_ms, - ], - ) - .with_context(|| format!("Failed to insert summary id={}", node.id))?; - - // #1574: persist the embedding to the per-model sidecar at the active - // signature, in the SAME tx as the summary row insert above. - if let Some(v) = node.embedding.as_deref() { - upsert_summary_embedding_conn(tx, &node.id, model_signature, v)?; - } - Ok(()) + tinycortex::memory::tree::store::insert_staged_summary_tx(tx, node, staged, model_signature) } -/// Set (or overwrite) the embedding for an existing summary row. -/// -/// #1574 cutover: writes the per-model `mem_tree_summary_embeddings` sidecar -/// at the active signature (via [`set_summary_embedding_for_signature`]) -/// instead of the legacy `mem_tree_summaries.embedding` column. The signature -/// is resolved internally from `config` via the shared -/// [`crate::openhuman::memory_store::chunks::store::tree_active_signature`] — same -/// resolution as the chunk path. Returns `1` on success (one sidecar row -/// written/updated); the legacy "0 if id unknown" count no longer applies -/// since the sidecar upsert does not join the parent summary row. pub fn set_summary_embedding( config: &Config, summary_id: &str, embedding: &[f32], ) -> Result { - let signature = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - log::debug!( - "[tree::store] set_summary_embedding: summary_id={summary_id} sig={signature} dims={}", - embedding.len() - ); - set_summary_embedding_for_signature(config, summary_id, &signature, embedding)?; + tinycortex::memory::tree::store::set_summary_embedding( + &engine_config(config), + summary_id, + embedding, + )?; Ok(1) } -/// Fetch a summary's embedding for the active model signature. -/// -/// #1574 cutover: reads the per-model `mem_tree_summary_embeddings` sidecar at -/// the active signature (via [`get_summary_embedding_for_signature`]) instead -/// of the legacy `mem_tree_summaries.embedding` column. `Ok(None)` when no -/// vector exists under the active signature — graceful absence during the §7 -/// backfill window, never a cross-space read. pub fn get_summary_embedding(config: &Config, summary_id: &str) -> Result>> { - let signature = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - get_summary_embedding_for_signature(config, summary_id, &signature) + tinycortex::memory::tree::store::get_summary_embedding(&engine_config(config), summary_id) } -/// Core upsert into `mem_tree_summary_embeddings` over an arbitrary -/// `&Connection`. Shared by the standalone -/// ([`set_summary_embedding_for_signature`]) and in-transaction -/// ([`set_summary_embedding_for_signature_tx`]) write paths so the SQL exists -/// exactly once. `rusqlite::Transaction` derefs to `Connection`, so the seal -/// path passes `&tx` and the sidecar row commits atomically with the summary -/// row insert (#1574 write-side cutover). -fn upsert_summary_embedding_conn( - conn: &rusqlite::Connection, - summary_id: &str, - model_signature: &str, - embedding: &[f32], -) -> Result<()> { - let blob = pack_embedding_blob(embedding); - let dim = i64::try_from(embedding.len()).context("embedding dimension does not fit i64")?; - let created_at = Utc::now().timestamp_millis() as f64 / 1000.0; - conn.execute( - "INSERT INTO mem_tree_summary_embeddings - (summary_id, model_signature, vector, dim, created_at) - VALUES (?1, ?2, ?3, ?4, ?5) - ON CONFLICT(summary_id, model_signature) DO UPDATE SET - vector = excluded.vector, - dim = excluded.dim, - created_at = excluded.created_at", - params![summary_id, model_signature, blob, dim, created_at], - )?; - Ok(()) -} - -/// Store a summary embedding for a specific provider/model/dimension signature. -/// -/// Per-model table write path for #1574. The legacy -/// `mem_tree_summaries.embedding` column is intentionally left untouched by -/// this helper (read by the §7 migration; dropped only in a later release). pub fn set_summary_embedding_for_signature( config: &Config, summary_id: &str, - model_signature: &str, + signature: &str, embedding: &[f32], ) -> Result<()> { - with_connection(config, |conn| { - upsert_summary_embedding_conn(conn, summary_id, model_signature, embedding) - }) + tinycortex::memory::tree::store::set_summary_embedding_for_signature( + &engine_config(config), + summary_id, + signature, + embedding, + ) } -/// Persistently record that `(summary_id, signature)` cannot be re-embedded. -/// Mirror of `tree::store::mark_chunk_reembed_skipped` for the summary side -/// of the reembed worklist (#1574 §6 fix). See that function's doc for the -/// full rationale. pub fn mark_summary_reembed_skipped( config: &Config, summary_id: &str, - model_signature: &str, + signature: &str, reason: &str, ) -> Result<()> { - let summary_id = crate::openhuman::memory_store::chunks::store::validate_reembed_skip_key( - "summary_id", + tinycortex::memory::chunks::mark_summary_reembed_skipped( + &engine_config(config), summary_id, - )?; - let model_signature = crate::openhuman::memory_store::chunks::store::validate_reembed_skip_key( - "model_signature", - model_signature, - )?; - with_connection(config, |conn| { - let now_ms = Utc::now().timestamp_millis(); - 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", - params![summary_id, model_signature, reason, now_ms], - )?; - log::debug!( - "[memory::chunk_store] mark_summary_reembed_skipped summary_id={summary_id} sig={model_signature} reason={reason}" - ); - Ok(()) - }) + signature, + reason, + ) } -/// Remove a single summary tombstone so re-embed backfill can retry the row. -/// -/// Idempotent — see [`crate::openhuman::memory_store::chunks::store::clear_chunk_reembed_skipped`]. pub fn clear_summary_reembed_skipped( config: &Config, summary_id: &str, - model_signature: &str, + signature: &str, ) -> Result<()> { - let summary_id = crate::openhuman::memory_store::chunks::store::validate_reembed_skip_key( - "summary_id", + tinycortex::memory::chunks::clear_summary_reembed_skipped( + &engine_config(config), summary_id, - )?; - let model_signature = crate::openhuman::memory_store::chunks::store::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", - params![summary_id, model_signature], - )?; - log::debug!( - "[memory::chunk_store] clear_summary_reembed_skipped summary_id={summary_id} sig={model_signature}" - ); - Ok(()) - }) + signature, + ) } -/// Transaction-scoped variant of [`set_summary_embedding_for_signature`], for -/// the seal path which inserts the summary row and its embedding in one tx -/// (#1574 write-side cutover). Opening a fresh connection there would break -/// atomicity / deadlock on the busy DB. pub(crate) fn set_summary_embedding_for_signature_tx( - tx: &rusqlite::Transaction<'_>, + tx: &Transaction<'_>, summary_id: &str, - model_signature: &str, + signature: &str, embedding: &[f32], ) -> Result<()> { - upsert_summary_embedding_conn(tx, summary_id, model_signature, embedding) + tinycortex::memory::chunks::set_summary_embedding_for_signature_tx( + tx, summary_id, signature, embedding, + ) } -/// Fetch a summary embedding for exactly one provider/model/dimension signature. pub fn get_summary_embedding_for_signature( config: &Config, summary_id: &str, - model_signature: &str, + signature: &str, ) -> Result>> { - with_connection(config, |conn| { - let row: Option<(Option>, i64)> = conn - .query_row( - "SELECT vector, dim - FROM mem_tree_summary_embeddings - WHERE summary_id = ?1 AND model_signature = ?2", - params![summary_id, model_signature], - |r| Ok((Some(r.get(0)?), r.get(1)?)), - ) - .optional()?; - match row { - None => Ok(None), - Some((blob, dim)) => { - let decoded = - decode_signature_blob(blob, dim, &format!("summary_id={summary_id}"))?; - if decoded.as_ref().is_some_and(|v| v.len() != dim as usize) { - anyhow::bail!( - "summary embedding dimension mismatch: dim column says {dim}, blob contains {} floats", - decoded.as_ref().map_or(0, Vec::len) - ); - } - Ok(decoded) - } - } - }) + tinycortex::memory::tree::store::get_summary_embedding_for_signature( + &engine_config(config), + summary_id, + signature, + ) } -/// Per-batch cap on `?` placeholders. Mirrors `chunks::store:: -/// MAX_EMBEDDING_BATCH` — see that constant's doc for the rationale (well -/// below SQLite's `SQLITE_MAX_VARIABLE_NUMBER = 32766`, large enough that -/// the current `LOOKUP_HEADROOM = 200` callsite always fits in one -/// round-trip). The two sides are independent intentionally: the summary -/// and chunk tables can grow at different rates and the cap might want to -/// drift independently in the future. -const MAX_EMBEDDING_BATCH: usize = 500; - -/// Batched read of summary embeddings under a single `model_signature`. -/// -/// Returns a `HashMap>` containing **only the -/// summaries that have a vector under `model_signature`**. Summaries with -/// no row, with a `NULL` vector (pending re-embed), or with a corrupted -/// blob are simply absent from the map — semantically identical to the -/// per-row [`get_summary_embedding_for_signature`] returning `Ok(None)`. -/// -/// Mirror of `chunks::store::get_chunk_embeddings_for_signature_batch`. -/// See that helper's doc for the rerank-loop motivation. The summary -/// side has its own copy rather than a generic helper because the two -/// tables (`mem_tree_summary_embeddings` vs `mem_tree_chunk_embeddings`) -/// have different blob-nullability semantics: summaries can store an -/// explicit `NULL` vector to flag a pending re-embed (handled here via -/// `Option>` + `decode_signature_blob`), chunks cannot. pub fn get_summary_embeddings_for_signature_batch( config: &Config, - summary_ids: &[String], - model_signature: &str, + ids: &[String], + signature: &str, ) -> Result>> { - if summary_ids.is_empty() { - return Ok(HashMap::new()); - } - with_connection(config, |conn| { - let mut out: HashMap> = HashMap::with_capacity(summary_ids.len()); - // Chunk to stay under SQLite's SQLITE_MAX_VARIABLE_NUMBER cap. - // For LOOKUP_HEADROOM=200 this loop runs exactly once; chunking - // only engages if a future caller passes >500 ids at a time. - for window in summary_ids.chunks(MAX_EMBEDDING_BATCH) { - // Build `IN (?,?,?,...)` with `window.len()` placeholders. - // model_signature is bound as the last parameter (?{n+1}). - let placeholders = std::iter::repeat_n("?", window.len()) - .collect::>() - .join(","); - let sql = format!( - "SELECT summary_id, vector, dim - FROM mem_tree_summary_embeddings - WHERE summary_id IN ({placeholders}) - AND model_signature = ?{sig_idx}", - sig_idx = window.len() + 1, - ); - let mut stmt = conn - .prepare(&sql) - .context("prepare get_summary_embeddings_for_signature_batch")?; - let mut bound: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); - for id in window { - bound.push(id as &dyn rusqlite::ToSql); - } - bound.push(&model_signature as &dyn rusqlite::ToSql); - let rows = stmt - .query_map(bound.as_slice(), |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, Option>>(1)?, - row.get::<_, i64>(2)?, - )) - }) - .context("query get_summary_embeddings_for_signature_batch")?; - for row in rows { - let (summary_id, blob, dim) = row?; - // Reuse the single-row decoder so NULL vectors (pending - // re-embed) and corrupt blobs surface with identical - // diagnostics to the per-row path. `Ok(None)` from the - // decoder is dropped: the map only carries materialised - // vectors, exactly mirroring the existing per-row - // contract. Length / dim-mismatch / negative-dim / - // non-multiple-of-4 are already enforced inside - // `decode_signature_blob` itself — no extra check here, - // matching the chunks side which delegates the same way - // to `embedding_from_blob`. - if let Some(v) = - decode_signature_blob(blob, dim, &format!("summary_id={summary_id}"))? - { - out.insert(summary_id, v); - } - } - } - Ok(out) - }) + tinycortex::memory::tree::store::get_summary_embeddings_for_signature_batch( + &engine_config(config), + ids, + signature, + ) } -/// Batched read of summary embeddings under the **active** model -/// signature. Mirrors [`get_summary_embedding`] for the per-row path: -/// resolves `tree_active_signature` once, forwards to -/// [`get_summary_embeddings_for_signature_batch`]. pub fn get_summary_embeddings_batch( config: &Config, - summary_ids: &[String], + ids: &[String], ) -> Result>> { - let signature = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - get_summary_embeddings_for_signature_batch(config, summary_ids, &signature) + tinycortex::memory::tree::store::get_summary_embeddings_batch(&engine_config(config), ids) } -/// Fetch one summary by id. Soft-deleted rows are returned with -/// `deleted = true` so callers can decide filtering policy. pub fn get_summary(config: &Config, id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, tree_id, tree_kind, level, parent_id, - child_ids_json, content, token_count, - entities_json, topics_json, - time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, embedding, - doc_id, version_ms - FROM mem_tree_summaries WHERE id = ?1", - )?; - let row = stmt - .query_row(params![id], row_to_summary) - .optional() - .context("Failed to query summary by id")?; - Ok(row) - }) + tinycortex::memory::tree::store::get_summary(&engine_config(config), id) } -/// Defensive upper bound on the number of `?` placeholders per batched -/// `SELECT … WHERE id IN (?,?,…)` query. SQLite's compile-time -/// `SQLITE_MAX_VARIABLE_NUMBER` has been ≥ 32 766 since 3.32 — 500 leaves -/// a ~65× safety margin. The current call-site (`hydrate_summary_inputs`) -/// passes at most a single seal's fan-in (typically 5–20 ids), so the -/// loop runs exactly once. The window exists so future callers with -/// larger id slices do not blow up against a host with a lower -/// compile-time SQLite cap. No volume reduction: all input ids in → all -/// matching rows out; the merged `HashMap` is byte-identical to one -/// giant query. -const MAX_FETCH_BATCH: usize = 500; - -/// Fetch many summaries by id in a single SQL round-trip per -/// [`MAX_FETCH_BATCH`] window. Replaces the per-id `get_summary` loop -/// inside hot paths like `hydrate_summary_inputs` (sealing L≥1 levels) -/// where N can grow to the seal fan-in and the loop fires on every seal -/// during ingest. Soft-deleted rows are returned with `deleted = true` -/// just like [`get_summary`]; missing ids are silently absent from the -/// map so callers can preserve the existing -/// "row missing → skip with warn" contract. pub fn get_summaries_batch( config: &Config, - summary_ids: &[String], + ids: &[String], ) -> Result> { - if summary_ids.is_empty() { - return Ok(HashMap::new()); - } - with_connection(config, |conn| { - let mut out: HashMap = HashMap::with_capacity(summary_ids.len()); - for window in summary_ids.chunks(MAX_FETCH_BATCH) { - let placeholders = (1..=window.len()) - .map(|i| format!("?{i}")) - .collect::>() - .join(","); - let sql = format!( - "SELECT id, tree_id, tree_kind, level, parent_id, - child_ids_json, content, token_count, - entities_json, topics_json, - time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, embedding, - doc_id, version_ms - FROM mem_tree_summaries - WHERE id IN ({placeholders})" - ); - let mut stmt = conn.prepare(&sql)?; - let params: Vec<&dyn rusqlite::ToSql> = - window.iter().map(|s| s as &dyn rusqlite::ToSql).collect(); - let rows = stmt - .query_map(params.as_slice(), row_to_summary)? - .collect::>>() - .context("Failed to collect summaries batch")?; - for s in rows { - out.insert(s.id.clone(), s); - } - } - Ok(out) - }) + tinycortex::memory::tree::store::get_summaries_batch(&engine_config(config), ids) } -/// List sealed summaries for a tree at a given level, ordered by -/// `sealed_at` ascending. Skips tombstoned rows. pub fn list_summaries_at_level( config: &Config, tree_id: &str, level: u32, ) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, tree_id, tree_kind, level, parent_id, - child_ids_json, content, token_count, - entities_json, topics_json, - time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, embedding, - doc_id, version_ms - FROM mem_tree_summaries - WHERE tree_id = ?1 AND level = ?2 AND deleted = 0 - ORDER BY sealed_at_ms ASC", - )?; - let rows = stmt - .query_map(params![tree_id, level], row_to_summary)? - .collect::>>() - .context("Failed to collect summaries")?; - Ok(rows) - }) + tinycortex::memory::tree::store::list_summaries_at_level(&engine_config(config), tree_id, level) } -/// List every non-deleted summary in a tree whose time-range envelope is -/// **fully contained** in `[since_ms, until_ms]` (inclusive), across all -/// levels. This is the "eligible summary" set for the windowed minimum-cover -/// (the morning-brief 24h tool): a node is eligible only when *all* its -/// descendant leaves fall inside the window, which — because seal sets -/// `time_range_start = MIN(children)` and `time_range_end = MAX(children)` -/// (`bucket_seal::seal_one_level`) — is exactly -/// `time_range_start_ms >= since_ms AND time_range_end_ms <= until_ms`. -/// -/// Ordered by `level ASC, time_range_start_ms ASC` so callers can build the -/// per-tree frontier and order chronologically without a re-sort. `content` -/// here is the ≤500-char preview; callers that need the full body hydrate via -/// `content::read::read_summary_body`. pub fn list_summaries_in_window( config: &Config, tree_id: &str, since_ms: i64, until_ms: i64, ) -> Result> { - log::debug!( - "[tree::store] list_summaries_in_window tree_id={tree_id} since_ms={since_ms} until_ms={until_ms}" - ); - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT id, tree_id, tree_kind, level, parent_id, - child_ids_json, content, token_count, - entities_json, topics_json, - time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, embedding, - doc_id, version_ms - FROM mem_tree_summaries - WHERE tree_id = ?1 AND deleted = 0 - AND time_range_start_ms >= ?2 - AND time_range_end_ms <= ?3 - ORDER BY level ASC, time_range_start_ms ASC", - )?; - let rows = stmt - .query_map(params![tree_id, since_ms, until_ms], row_to_summary)? - .collect::>>() - .context("Failed to collect in-window summaries")?; - log::debug!( - "[tree::store] list_summaries_in_window tree_id={tree_id} matched={}", - rows.len() - ); - Ok(rows) - }) -} - -/// Count summaries in a tree (diagnostic helper). -pub fn count_summaries(config: &Config, tree_id: &str) -> Result { - with_connection(config, |conn| { - let n: i64 = conn - .query_row( - "SELECT COUNT(*) FROM mem_tree_summaries - WHERE tree_id = ?1 AND deleted = 0", - params![tree_id], - |r| r.get(0), - ) - .context("count summaries query")?; - Ok(n.max(0) as u64) - }) -} - -fn row_to_summary(row: &rusqlite::Row<'_>) -> rusqlite::Result { - let id: String = row.get(0)?; - let tree_id: String = row.get(1)?; - let tree_kind_s: String = row.get(2)?; - let level: i64 = row.get(3)?; - let parent_id: Option = row.get(4)?; - let child_ids_json: String = row.get(5)?; - let content: String = row.get(6)?; - let token_count: i64 = row.get(7)?; - let entities_json: String = row.get(8)?; - let topics_json: String = row.get(9)?; - let trs_ms: i64 = row.get(10)?; - let tre_ms: i64 = row.get(11)?; - let score: f64 = row.get(12)?; - let sealed_ms: i64 = row.get(13)?; - let deleted: i64 = row.get(14)?; - let embedding_blob: Option> = row.get(15)?; - let doc_id: Option = row.get(16)?; - let version_ms: Option = row.get(17)?; - - let tree_kind = TreeKind::parse(&tree_kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(2, rusqlite::types::Type::Text, e.into()) - })?; - let child_ids: Vec = serde_json::from_str(&child_ids_json).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(5, rusqlite::types::Type::Text, Box::new(e)) - })?; - let entities: Vec = serde_json::from_str(&entities_json).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(8, rusqlite::types::Type::Text, Box::new(e)) - })?; - let topics: Vec = serde_json::from_str(&topics_json).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(9, rusqlite::types::Type::Text, Box::new(e)) - })?; - let embedding = - decode_optional_blob(embedding_blob, &format!("summary_id={id}")).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 15, - rusqlite::types::Type::Blob, - Box::new(std::io::Error::new( - std::io::ErrorKind::InvalidData, - e.to_string(), - )), - ) - })?; - - Ok(SummaryNode { - id, + tinycortex::memory::tree::store::list_summaries_in_window( + &engine_config(config), tree_id, - tree_kind, - level: level.max(0) as u32, - parent_id, - child_ids, - content, - token_count: token_count.max(0) as u32, - entities, - topics, - time_range_start: ms_to_utc(trs_ms)?, - time_range_end: ms_to_utc(tre_ms)?, - score: score as f32, - sealed_at: ms_to_utc(sealed_ms)?, - deleted: deleted != 0, - embedding, - doc_id, - version_ms, - }) + since_ms, + until_ms, + ) } -// ── Buffers ───────────────────────────────────────────────────────────── +pub fn count_summaries(config: &Config, tree_id: &str) -> Result { + tinycortex::memory::tree::store::count_summaries(&engine_config(config), tree_id) +} -/// Read the current buffer at `(tree_id, level)` or return an empty one. pub fn get_buffer(config: &Config, tree_id: &str, level: u32) -> Result { - with_connection(config, |conn| get_buffer_conn(conn, tree_id, level)) + tinycortex::memory::tree::store::get_buffer(&engine_config(config), tree_id, level) } pub(crate) 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", - )?; - let row = stmt - .query_row(params![tree_id, level], row_to_buffer) - .optional() - .context("Failed to query buffer")?; - Ok(row.unwrap_or_else(|| Buffer::empty(tree_id, level))) + tinycortex::memory::tree::store::get_buffer_conn(conn, tree_id, level) } -/// Upsert a buffer row. -pub(crate) fn upsert_buffer_tx(tx: &Transaction<'_>, buf: &Buffer) -> Result<()> { - let now_ms = Utc::now().timestamp_millis(); - tx.execute( - "INSERT INTO mem_tree_buffers ( - tree_id, level, item_ids_json, token_sum, oldest_at_ms, updated_at_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6) - ON CONFLICT(tree_id, level) DO UPDATE SET - item_ids_json = excluded.item_ids_json, - token_sum = excluded.token_sum, - oldest_at_ms = excluded.oldest_at_ms, - updated_at_ms = excluded.updated_at_ms", - params![ - buf.tree_id, - buf.level, - serde_json::to_string(&buf.item_ids)?, - buf.token_sum, - buf.oldest_at.map(|t| t.timestamp_millis()), - now_ms, - ], - ) - .with_context(|| { - format!( - "Failed to upsert buffer tree_id={} level={}", - buf.tree_id, buf.level - ) - })?; - Ok(()) +pub(crate) fn upsert_buffer_tx(tx: &Transaction<'_>, buffer: &Buffer) -> Result<()> { + tinycortex::memory::tree::store::upsert_buffer_tx(tx, buffer) } -/// Reset a buffer at `(tree_id, level)` to empty. Used at seal time: the -/// items move into a summary row and the buffer is cleared in the same tx. pub(crate) fn clear_buffer_tx(tx: &Transaction<'_>, tree_id: &str, level: u32) -> Result<()> { - let empty = Buffer::empty(tree_id, level); - upsert_buffer_tx(tx, &empty) + tinycortex::memory::tree::store::clear_buffer_tx(tx, tree_id, level) } -/// List stale **L0** buffers ordered by `oldest_at_ms ASC`. Used by the -/// time-based flush pass. -/// -/// Only L0 (raw-leaf) buffers are returned. Force-sealing an L≥1 buffer -/// that hasn't met the [`SUMMARY_FANOUT`](super::types::SUMMARY_FANOUT) -/// gate produces a degenerate single-child summary that wraps exactly the -/// same content as its only child — repeated flush cycles cascade these -/// no-op promotions up the tree and collapse the upper levels into a -/// 1:1:1 chain. Upper-level buffers must seal only when their fan-in -/// gate is naturally met. pub fn list_stale_buffers(config: &Config, older_than: DateTime) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT tree_id, level, item_ids_json, token_sum, oldest_at_ms - FROM mem_tree_buffers - WHERE oldest_at_ms IS NOT NULL - AND oldest_at_ms <= ?1 - AND level = 0 - ORDER BY oldest_at_ms ASC", - )?; - let rows = stmt - .query_map(params![older_than.timestamp_millis()], row_to_buffer)? - .collect::>>() - .context("Failed to collect stale buffers")?; - Ok(rows) - }) + tinycortex::memory::tree::store::list_stale_buffers(&engine_config(config), older_than) } - -fn row_to_buffer(row: &rusqlite::Row<'_>) -> rusqlite::Result { - let tree_id: String = row.get(0)?; - let level: i64 = row.get(1)?; - let item_ids_json: String = row.get(2)?; - let token_sum: i64 = row.get(3)?; - let oldest_ms: Option = row.get(4)?; - - let item_ids: Vec = serde_json::from_str(&item_ids_json).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(2, rusqlite::types::Type::Text, Box::new(e)) - })?; - let oldest_at = oldest_ms.map(ms_to_utc).transpose()?; - Ok(Buffer { - tree_id, - level: level.max(0) as u32, - item_ids, - token_sum, - oldest_at, - }) -} - -fn pack_embedding_blob(embedding: &[f32]) -> Vec { - embedding.iter().flat_map(|f| f.to_le_bytes()).collect() -} - -fn decode_signature_blob(blob: Option>, dim: i64, label: &str) -> Result>> { - let Some(bytes) = blob else { - return Ok(None); - }; - if dim < 0 { - anyhow::bail!("{label} has negative dimension {dim}"); - } - if !bytes.len().is_multiple_of(4) { - anyhow::bail!("{label} blob length {} not a multiple of 4", bytes.len()); - } - let floats: Vec = bytes - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - if floats.len() != dim as usize { - anyhow::bail!( - "summary embedding dimension mismatch: dim column says {dim}, blob contains {} floats", - floats.len() - ); - } - Ok(Some(floats)) -} - -#[cfg(test)] -#[path = "store_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_store/trees/types.rs b/src/openhuman/memory_store/trees/types.rs index 47d291b40..e56f9cd2e 100644 --- a/src/openhuman/memory_store/trees/types.rs +++ b/src/openhuman/memory_store/trees/types.rs @@ -1,386 +1,25 @@ -//! Core types for Phase 3a — summary trees, per-source bucket-seal (#709). -//! -//! These types sit on top of Phase 1's chunk leaves. A [`Tree`] groups leaves -//! under one scope (e.g. one chat channel, one email account). When a -//! [`Buffer`] at some level accumulates enough tokens, its contents seal -//! into a [`SummaryNode`] at level+1 and the buffer clears. Summary nodes -//! are immutable once emitted — updates to children use the Phase 1/2 -//! tombstone pattern, never rewrite parents. +//! Compatibility exports for tinycortex summary-tree persistence types. -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -/// What kind of tree this is. Source trees live per ingest source; topic -/// and global trees are introduced in Phase 3b/3c and share the same -/// schema. -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -#[non_exhaustive] -pub enum TreeKind { - /// One tree per ingest source (e.g. `chat:slack:#eng`, `email:gmail:user`). - Source, - /// Reserved for Phase 3c — per-entity/topic tree. - Topic, - /// Reserved for Phase 3b — cross-source daily digest tree. - Global, -} - -impl TreeKind { - /// Stable lowercase form used in SQL discriminator columns and ids. - pub fn as_str(self) -> &'static str { - match self { - Self::Source => "source", - Self::Topic => "topic", - Self::Global => "global", - } - } - - /// Inverse of [`Self::as_str`] — parse back from a discriminator - /// string. Errors on unknown variants. - pub fn parse(s: &str) -> Result { - match s { - "source" => Ok(Self::Source), - "topic" => Ok(Self::Topic), - "global" => Ok(Self::Global), - other => Err(format!("unknown tree kind: {other}")), - } - } -} - -/// Activity state of a tree. Archived trees stay queryable but don't accept -/// new leaves — used by Phase 3c when a topic tree's entity goes cold. -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TreeStatus { - Active, - Archived, -} - -impl TreeStatus { - /// Stable lowercase form used as the SQL discriminator value. - pub fn as_str(self) -> &'static str { - match self { - Self::Active => "active", - Self::Archived => "archived", - } - } - - /// Inverse of [`Self::as_str`] — parse from the SQL discriminator. - pub fn parse(s: &str) -> Result { - match s { - "active" => Ok(Self::Active), - "archived" => Ok(Self::Archived), - other => Err(format!("unknown tree status: {other}")), - } - } -} - -/// One summary-tree instance. -/// -/// `root_id` is `None` until the first seal emits an L1 node. `max_level` -/// tracks the highest level that has ever sealed; `root_id` points at the -/// current top node at that level (changes on root-split). -#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] -pub struct Tree { - pub id: String, - pub kind: TreeKind, - /// Logical identifier for what the tree covers. Format conventions: - /// - Source: `::` or the chunk's - /// `source_id` directly (Phase 3a uses the chunk source_id verbatim) - /// - Topic: canonical entity id - /// - Global: the literal string `"global"` - pub scope: String, - pub root_id: Option, - pub max_level: u32, - pub status: TreeStatus, - pub created_at: DateTime, - pub last_sealed_at: Option>, -} - -/// A sealed summary node — one level above raw leaves. -/// -/// `child_ids` points at the concrete children that were in the buffer when -/// this node sealed. For L1 nodes those are leaf `chunk.id`s; for L2+ they -/// are lower-level summary ids. Relation is fixed at seal time — never -/// modified afterwards. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct SummaryNode { - pub id: String, - pub tree_id: String, - pub tree_kind: TreeKind, - /// 1 for summaries over raw leaves, 2 over L1 summaries, and so on. - pub level: u32, - pub parent_id: Option, - pub child_ids: Vec, - /// Summariser output. Typical target: 800–1500 tokens. - pub content: String, - pub token_count: u32, - /// Curated subset of children's entity canonical-ids. - pub entities: Vec, - /// Curated topic labels (hashtag-like short phrases). - pub topics: Vec, - pub time_range_start: DateTime, - pub time_range_end: DateTime, - /// Max of children's scores at seal time — cheap heuristic, preserved - /// for reranking in Phase 4. - pub score: f32, - pub sealed_at: DateTime, - /// Tombstone flag — stays `false` in Phase 3a since summaries are - /// immutable. Reserved for future cleanup passes (e.g. archive cascade). - pub deleted: bool, - /// Phase 4 (#710): summary content embedding for semantic rerank. - /// - /// `Some` on new seals — populated before the write tx opens so a - /// failed embed aborts the seal (see `bucket_seal::seal_one_level`). - /// `None` on legacy summaries sealed before Phase 4, or on reads - /// where the blob column is NULL. Retrieval tolerates `None` by - /// dropping those rows to the bottom of semantic rerank results. - #[serde(default)] - pub embedding: Option>, - /// Document identity this node belongs to, for document source trees - /// (Notion etc.). `Some(source_id)` for nodes that live inside a single - /// document's per-doc subtree (its L1…doc-root chain); `None` for - /// merge-tier nodes (which summarise *across* documents) and for - /// chat/email source trees (which have no per-document structure). - /// - /// Together with [`Self::version_ms`] this lets retrieval resolve - /// "latest version per document" at read time: when a Notion page is - /// edited a new per-doc subtree is sealed with a higher `version_ms`, - /// and the older one is filtered out on traversal without ever being - /// rewritten or tombstoned. - #[serde(default)] - pub doc_id: Option, - /// Document version this node was sealed for, as epoch-milliseconds - /// (Notion `last_edited_time`). `Some(_)` on per-doc subtree nodes, - /// `None` on merge-tier and non-document nodes. Read-time latest-wins - /// keeps `max(version_ms)` per [`Self::doc_id`]. - #[serde(default)] - pub version_ms: Option, -} - -/// Unsealed frontier at a given `(tree_id, level)`. One row per level per -/// tree. `oldest_at` is `None` when the buffer is empty; used by the -/// time-based flush trigger. -#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] -pub struct Buffer { - pub tree_id: String, - pub level: u32, - pub item_ids: Vec, - pub token_sum: i64, - pub oldest_at: Option>, -} - -impl Buffer { - /// Empty buffer at the given key. - pub fn empty(tree_id: &str, level: u32) -> Self { - Self { - tree_id: tree_id.to_string(), - level, - item_ids: Vec::new(), - token_sum: 0, - oldest_at: None, - } - } - - /// True when the buffer holds no pending items. - pub fn is_empty(&self) -> bool { - self.item_ids.is_empty() - } - - /// Whether the buffer's oldest item is older than `max_age`. Returns - /// `false` for an empty buffer. - pub fn is_stale(&self, now: DateTime, max_age: chrono::Duration) -> bool { - match self.oldest_at { - Some(ts) => now.signed_duration_since(ts) > max_age, - None => false, - } - } -} - -/// Input token target for one L0 → L1 seal: when an L0 buffer's -/// `token_sum` reaches this, we summarise the accumulated leaves. -/// -/// Sized for the cloud summariser's 120k-token context with headroom for -/// the system prompt and the model's own output. With ~5k tokens emitted -/// per summary (see [`OUTPUT_TOKEN_BUDGET`]), one parent represents ~50k -/// tokens of leaf content — i.e. ~10 child summaries' worth. -pub const INPUT_TOKEN_BUDGET: u32 = 50_000; - -/// Output token budget passed to the summariser as `ctx.token_budget`. -/// The summariser may clamp lower (see `summariser/llm.rs`'s -/// `MAX_SUMMARY_OUTPUT_TOKENS`). 5k keeps the produced summary well -/// under the embedder's 8k input ceiling so the post-seal embed never -/// rejects the row. -pub const OUTPUT_TOKEN_BUDGET: u32 = 5_000; - -/// Sibling count that triggers a seal at level ≥ 1 (summaries → next level). -/// -/// Set to match the [`INPUT_TOKEN_BUDGET`] / [`OUTPUT_TOKEN_BUDGET`] -/// ratio so each level folds roughly the same volume of content as L0: -/// 10 summaries × ~5k tokens ≈ 50k input. Decouples upper-level seals -/// from per-summary token size so the tree's fan-in stays stable -/// regardless of summariser quality (token-based gating would collapse -/// the inert-fallback case into a 1:1:1 chain). -pub const SUMMARY_FANOUT: u32 = 10; - -/// Default age at which a non-empty buffer is force-sealed even under the -/// token budget. Keeps recent activity from stalling waiting for more -/// leaves that may never arrive. -pub const DEFAULT_FLUSH_AGE_SECS: i64 = 7 * 24 * 60 * 60; +pub use tinycortex::memory::tree::store::{ + Buffer, EntityIndexStats, HotnessCounters, SummaryNode, Tree, TreeKind, TreeStatus, + DEFAULT_FLUSH_AGE_SECS, INPUT_TOKEN_BUDGET, OUTPUT_TOKEN_BUDGET, SUMMARY_FANOUT, + TOPIC_ARCHIVE_THRESHOLD, TOPIC_CREATION_THRESHOLD, TOPIC_RECHECK_EVERY, +}; #[cfg(test)] mod tests { use super::*; #[test] - fn tree_kind_round_trip() { - for k in [TreeKind::Source, TreeKind::Topic, TreeKind::Global] { - assert_eq!(TreeKind::parse(k.as_str()).unwrap(), k); - } - assert!(TreeKind::parse("bogus").is_err()); + fn tree_wire_discriminators_are_stable() { + assert_eq!(TreeKind::parse("source").unwrap(), TreeKind::Source); + assert_eq!(TreeStatus::parse("archived").unwrap(), TreeStatus::Archived); } #[test] - fn tree_status_round_trip() { - for s in [TreeStatus::Active, TreeStatus::Archived] { - assert_eq!(TreeStatus::parse(s.as_str()).unwrap(), s); - } - assert!(TreeStatus::parse("live").is_err()); - } - - #[test] - fn empty_buffer_is_not_stale() { - let b = Buffer::empty("t1", 0); - assert!(b.is_empty()); - assert!(!b.is_stale(Utc::now(), chrono::Duration::zero())); - } - - #[test] - fn stale_buffer_detected() { - let past = Utc::now() - chrono::Duration::hours(10); - let b = Buffer { - tree_id: "t1".into(), - level: 0, - item_ids: vec!["leaf-1".into()], - token_sum: 100, - oldest_at: Some(past), - }; - assert!(b.is_stale(Utc::now(), chrono::Duration::hours(1))); - assert!(!b.is_stale(Utc::now(), chrono::Duration::hours(20))); - } -} - -// ============================================================================ -// Topic-tree hotness (Phase 3c) — formerly memory_store::trees_topic::types -// ============================================================================ -// -// Folded in here because topic and global trees are not structurally distinct -// from source trees — they all live in the same `mem_tree_trees` table keyed -// by `TreeKind`. The only topic-specific extra state is the entity hotness -// counters in `mem_tree_entity_hotness`, which gate materialisation of a -// topic tree but are themselves not trees. - -/// Hotness threshold above which a topic tree is materialised for an entity. -pub const TOPIC_CREATION_THRESHOLD: f32 = 10.0; - -/// Hotness threshold below which a topic tree becomes an archive candidate. -pub const TOPIC_ARCHIVE_THRESHOLD: f32 = 2.0; - -/// How often (in ingests touching the entity) to recompute hotness from the -/// full [`EntityIndexStats`]. Between recomputes only the cheap counters bump. -pub const TOPIC_RECHECK_EVERY: u32 = 100; - -/// Input record fed to the hotness math (see `memory::tree_topic::hotness`). -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct EntityIndexStats { - pub mention_count_30d: u32, - pub distinct_sources: u32, - pub last_seen_ms: Option, - pub query_hits_30d: u32, - pub graph_centrality: Option, -} - -/// Row persisted in `mem_tree_entity_hotness`. Persistence helpers live in -/// [`super::hotness`]. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct HotnessCounters { - pub entity_id: String, - pub mention_count_30d: u32, - pub distinct_sources: u32, - pub last_seen_ms: Option, - pub query_hits_30d: u32, - pub graph_centrality: Option, - pub ingests_since_check: u32, - pub last_hotness: Option, - pub last_updated_ms: i64, -} - -impl HotnessCounters { - pub fn fresh(entity_id: &str, now_ms: i64) -> Self { - Self { - entity_id: entity_id.to_string(), - mention_count_30d: 0, - distinct_sources: 0, - last_seen_ms: None, - query_hits_30d: 0, - graph_centrality: None, - ingests_since_check: 0, - last_hotness: None, - last_updated_ms: now_ms, - } - } - - pub fn stats(&self) -> EntityIndexStats { - EntityIndexStats { - mention_count_30d: self.mention_count_30d, - distinct_sources: self.distinct_sources, - last_seen_ms: self.last_seen_ms, - query_hits_30d: self.query_hits_30d, - graph_centrality: self.graph_centrality, - } - } -} - -#[cfg(test)] -mod hotness_type_tests { - use super::*; - - #[test] - fn fresh_counters_are_zero() { - let c = HotnessCounters::fresh("email:alice@example.com", 1_700_000_000_000); - assert_eq!(c.entity_id, "email:alice@example.com"); - assert_eq!(c.mention_count_30d, 0); - assert_eq!(c.distinct_sources, 0); - assert_eq!(c.ingests_since_check, 0); - assert!(c.last_hotness.is_none()); - assert!(c.last_seen_ms.is_none()); - assert_eq!(c.last_updated_ms, 1_700_000_000_000); - } - - #[test] - fn stats_projection_mirrors_row() { - let c = HotnessCounters { - entity_id: "e".into(), - mention_count_30d: 5, - distinct_sources: 2, - last_seen_ms: Some(42), - query_hits_30d: 1, - graph_centrality: Some(0.3), - ingests_since_check: 4, - last_hotness: Some(9.9), - last_updated_ms: 100, - }; - let s = c.stats(); - assert_eq!(s.mention_count_30d, 5); - assert_eq!(s.distinct_sources, 2); - assert_eq!(s.last_seen_ms, Some(42)); - assert_eq!(s.query_hits_30d, 1); - assert_eq!(s.graph_centrality, Some(0.3)); - } - - #[test] - fn thresholds_make_creation_strictly_above_archive() { - assert!(TOPIC_CREATION_THRESHOLD > TOPIC_ARCHIVE_THRESHOLD); - assert!(TOPIC_RECHECK_EVERY > 0); + fn tree_budgets_match_engine_defaults() { + assert_eq!(INPUT_TOKEN_BUDGET, 50_000); + assert_eq!(OUTPUT_TOKEN_BUDGET, 5_000); + assert_eq!(SUMMARY_FANOUT, 10); } } diff --git a/src/openhuman/memory_store/types.rs b/src/openhuman/memory_store/types.rs index 62821812f..d72101ee2 100644 --- a/src/openhuman/memory_store/types.rs +++ b/src/openhuman/memory_store/types.rs @@ -1,378 +1,9 @@ -//! Public input/output types for namespace memory documents. +//! Stable host path for tinycortex-owned namespace memory contracts. -use serde::{Deserialize, Serialize}; +pub use tinycortex::memory::{ + GraphRelationRecord, MemoryItemKind, MemoryKvRecord, NamespaceDocumentInput, + NamespaceMemoryHit, NamespaceQueryResult, NamespaceRetrievalContext, RetrievalScoreBreakdown, + StoredMemoryDocument, +}; -use crate::openhuman::memory::MemoryTaint; - -pub(crate) const GLOBAL_NAMESPACE: &str = "global"; - -/// Input payload for upserting a namespace-scoped memory document. -/// -/// Used by `MemoryClient::put_doc` and the ingestion pipeline. `document_id` -/// is optional — when omitted, an existing row keyed by `(namespace, key)` is -/// reused, otherwise a new id is generated. -/// -/// `taint` carries the provenance signal through the persistence layer so -/// downstream recall paths can drive the subconscious origin-escalation -/// decision. Legacy JSON without the field decodes as -/// [`MemoryTaint::Internal`] (see the serde tests below). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NamespaceDocumentInput { - pub namespace: String, - pub key: String, - pub title: String, - pub content: String, - pub source_type: String, - pub priority: String, - #[serde(default)] - pub tags: Vec, - #[serde(default)] - pub metadata: serde_json::Value, - pub category: String, - #[serde(default)] - pub session_id: Option, - #[serde(default)] - pub document_id: Option, - /// Provenance — `Internal` for user-driven writes, `ExternalSync` for - /// memory_sync ingest paths (Gmail / Slack / Notion / Composio / etc.). - /// Defaults via `#[serde(default)]` so legacy callers and on-disk JSON - /// missing this field decode as the conservative - /// [`MemoryTaint::Internal`]. - #[serde(default)] - pub taint: MemoryTaint, -} - -/// One ranked retrieval result for a namespace text query. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NamespaceQueryResult { - pub key: String, - pub content: String, - pub score: f64, - /// Stored category string (e.g. `core`, `daily`, or custom label). - pub category: String, - /// Provenance taint carried back from the persistence layer so the - /// recall caller can surface it on [`crate::openhuman::memory::MemoryEntry`]. - /// Defaults to [`MemoryTaint::Internal`] for legacy rows that predate - /// the column. - #[serde(default)] - pub taint: MemoryTaint, -} - -/// Discriminator for the kind of stored memory item a hit refers to. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MemoryItemKind { - Document, - Kv, - Episodic, - Event, -} - -/// Persisted form of a memory document as stored in `memory_docs`, -/// including timestamps and the markdown sidecar path. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StoredMemoryDocument { - pub document_id: String, - pub namespace: String, - pub key: String, - pub title: String, - pub content: String, - pub source_type: String, - pub priority: String, - pub tags: Vec, - pub metadata: serde_json::Value, - pub category: String, - pub session_id: Option, - pub created_at: f64, - pub updated_at: f64, - pub markdown_rel_path: String, - /// Provenance taint round-tripped from the `memory_docs.taint` column. - /// Defaults to [`MemoryTaint::Internal`] for legacy rows. - #[serde(default)] - pub taint: MemoryTaint, -} - -/// A single KV row, namespace-scoped or global (when `namespace` is `None`). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MemoryKvRecord { - pub namespace: Option, - pub key: String, - pub value: serde_json::Value, - pub updated_at: f64, -} - -/// A graph edge (subject — predicate → object) plus accumulated evidence. -/// -/// `document_ids` and `chunk_ids` track every source that contributed to this -/// relation; `evidence_count` is the merged count after de-duplication. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct GraphRelationRecord { - pub namespace: Option, - pub subject: String, - pub predicate: String, - pub object: String, - pub attrs: serde_json::Value, - pub updated_at: f64, - pub evidence_count: u32, - pub order_index: Option, - pub document_ids: Vec, - pub chunk_ids: Vec, -} - -/// Per-signal contribution to a hit's final score, surfaced for debugging -/// and UI ranking explainers. -#[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct RetrievalScoreBreakdown { - pub keyword_relevance: f64, - pub vector_similarity: f64, - pub graph_relevance: f64, - pub episodic_relevance: f64, - pub freshness: f64, - pub final_score: f64, -} - -/// A single ranked retrieval hit returned from `query_namespace_hits` / -/// `recall_namespace_memories`. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NamespaceMemoryHit { - pub id: String, - pub kind: MemoryItemKind, - pub namespace: String, - pub key: String, - pub title: Option, - pub content: String, - pub category: String, - pub source_type: Option, - pub updated_at: f64, - pub score: f64, - pub score_breakdown: RetrievalScoreBreakdown, - #[serde(default)] - pub document_id: Option, - #[serde(default)] - pub chunk_id: Option, - #[serde(default)] - pub supporting_relations: Vec, - /// Provenance taint surfaced from the source row so retrieval consumers - /// (subconscious tick, context builder, …) can drive origin-escalation - /// without re-querying the underlying document. Defaults to - /// [`MemoryTaint::Internal`] for legacy hits / KV / episodic rows. - #[serde(default)] - pub taint: MemoryTaint, -} - -/// Aggregated retrieval result for a namespace: rendered context text plus -/// the underlying hits. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NamespaceRetrievalContext { - pub namespace: String, - pub query: Option, - pub context_text: String, - pub hits: Vec, -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn global_namespace_constant_is_stable() { - assert_eq!(GLOBAL_NAMESPACE, "global"); - } - - #[test] - fn memory_item_kind_serde_uses_snake_case() { - let json_value = serde_json::to_string(&MemoryItemKind::Document).unwrap(); - assert_eq!(json_value, "\"document\""); - let decoded: MemoryItemKind = serde_json::from_str("\"episodic\"").unwrap(); - assert_eq!(decoded, MemoryItemKind::Episodic); - } - - #[test] - fn namespace_document_input_defaults_optional_fields() { - let value = json!({ - "namespace": "global", - "key": "note-1", - "title": "Title", - "content": "Body", - "source_type": "manual", - "priority": "normal", - "metadata": {}, - "category": "core" - }); - let parsed: NamespaceDocumentInput = serde_json::from_value(value).unwrap(); - assert!(parsed.tags.is_empty()); - assert_eq!(parsed.metadata, json!({})); - assert!(parsed.session_id.is_none()); - assert!(parsed.document_id.is_none()); - } - - #[test] - fn retrieval_score_breakdown_default_is_zeroed() { - let breakdown = RetrievalScoreBreakdown::default(); - assert_eq!(breakdown.keyword_relevance, 0.0); - assert_eq!(breakdown.vector_similarity, 0.0); - assert_eq!(breakdown.graph_relevance, 0.0); - assert_eq!(breakdown.episodic_relevance, 0.0); - assert_eq!(breakdown.freshness, 0.0); - assert_eq!(breakdown.final_score, 0.0); - } - - #[test] - fn memory_kv_record_roundtrips_with_optional_namespace() { - let global = MemoryKvRecord { - namespace: None, - key: "theme".into(), - value: json!("dark"), - updated_at: 1.5, - }; - let namespaced = MemoryKvRecord { - namespace: Some("project".into()), - key: "state".into(), - value: json!({"open": true}), - updated_at: 2.5, - }; - for record in [global, namespaced] { - let value = serde_json::to_value(&record).unwrap(); - let decoded: MemoryKvRecord = serde_json::from_value(value).unwrap(); - assert_eq!(decoded.namespace, record.namespace); - assert_eq!(decoded.key, record.key); - assert_eq!(decoded.value, record.value); - assert_eq!(decoded.updated_at, record.updated_at); - } - } - - #[test] - fn namespace_document_input_taint_defaults_internal_for_legacy_json() { - // Legacy callers serialised before the taint field existed must - // still decode — `#[serde(default)]` makes the field optional and - // falls back to the conservative `Internal` provenance. - let value = json!({ - "namespace": "skill-gmail", - "key": "thread-1", - "title": "Subject", - "content": "Body", - "source_type": "composio-sync", - "priority": "medium", - "metadata": {}, - "category": "core" - }); - let parsed: NamespaceDocumentInput = serde_json::from_value(value).unwrap(); - assert_eq!(parsed.taint, MemoryTaint::Internal); - } - - #[test] - fn namespace_document_input_taint_roundtrips_external_sync() { - let input = NamespaceDocumentInput { - namespace: "skill-gmail".into(), - key: "thread-1".into(), - title: "Subject".into(), - content: "Body".into(), - source_type: "composio-sync".into(), - priority: "medium".into(), - tags: Vec::new(), - metadata: json!({}), - category: "core".into(), - session_id: None, - document_id: None, - taint: MemoryTaint::ExternalSync, - }; - let value = serde_json::to_value(&input).unwrap(); - assert_eq!( - value.get("taint").and_then(|v| v.as_str()), - Some("external_sync") - ); - let parsed: NamespaceDocumentInput = serde_json::from_value(value).unwrap(); - assert_eq!(parsed.taint, MemoryTaint::ExternalSync); - } - - #[test] - fn namespace_query_result_taint_defaults_internal_for_legacy_json() { - let value = json!({ - "key": "k", - "content": "c", - "score": 0.5, - "category": "core" - }); - let parsed: NamespaceQueryResult = serde_json::from_value(value).unwrap(); - assert_eq!(parsed.taint, MemoryTaint::Internal); - } - - #[test] - fn stored_memory_document_taint_defaults_internal_for_legacy_json() { - let value = json!({ - "document_id": "d", - "namespace": "n", - "key": "k", - "title": "t", - "content": "c", - "source_type": "chat", - "priority": "medium", - "tags": [], - "metadata": {}, - "category": "core", - "session_id": null, - "created_at": 1.0, - "updated_at": 2.0, - "markdown_rel_path": "" - }); - let parsed: StoredMemoryDocument = serde_json::from_value(value).unwrap(); - assert_eq!(parsed.taint, MemoryTaint::Internal); - } - - #[test] - fn namespace_memory_hit_taint_defaults_internal_for_legacy_json() { - let hit: NamespaceMemoryHit = serde_json::from_value(json!({ - "id": "hit-1", - "kind": "document", - "namespace": "global", - "key": "note-1", - "title": "Title", - "content": "Body", - "category": "core", - "source_type": "manual", - "updated_at": 3.5, - "score": 0.8, - "score_breakdown": { - "keyword_relevance": 0.5, - "vector_similarity": 0.2, - "graph_relevance": 0.0, - "episodic_relevance": 0.0, - "freshness": 0.1, - "final_score": 0.8 - } - })) - .unwrap(); - assert_eq!(hit.taint, MemoryTaint::Internal); - } - - #[test] - fn namespace_memory_hit_defaults_optional_fields() { - let hit: NamespaceMemoryHit = serde_json::from_value(json!({ - "id": "hit-1", - "kind": "document", - "namespace": "global", - "key": "note-1", - "title": "Title", - "content": "Body", - "category": "core", - "source_type": "manual", - "updated_at": 3.5, - "score": 0.8, - "score_breakdown": { - "keyword_relevance": 0.5, - "vector_similarity": 0.2, - "graph_relevance": 0.0, - "episodic_relevance": 0.0, - "freshness": 0.1, - "final_score": 0.8 - } - })) - .unwrap(); - - assert!(hit.document_id.is_none()); - assert!(hit.chunk_id.is_none()); - assert!(hit.supporting_relations.is_empty()); - assert_eq!(hit.kind, MemoryItemKind::Document); - } -} +pub(crate) use tinycortex::memory::types::GLOBAL_NAMESPACE; diff --git a/src/openhuman/memory_store/vectors/mod.rs b/src/openhuman/memory_store/vectors/mod.rs index cb65c42a8..d22ac368e 100644 --- a/src/openhuman/memory_store/vectors/mod.rs +++ b/src/openhuman/memory_store/vectors/mod.rs @@ -1,8 +1,6 @@ -//! Local vector store (VectorStore) — moved from embeddings::store. -//! -//! Previously at `embeddings::store`. Moved here as part of the -//! memory_store consolidation to co-locate all persistence with memory_store. +//! TinyCortex-owned local vector storage. -pub mod store; - -pub use store::{bytes_to_vec, cosine_similarity, vec_to_bytes, SearchResult, VectorStore}; +pub use tinycortex::memory::store::vectors::{ + bytes_to_vec, cosine_similarity, format_embedding_signature, vec_to_bytes, EmbeddingBackend, + InertEmbedding, SearchResult, VectorStore, +}; diff --git a/src/openhuman/memory_store/vectors/store.rs b/src/openhuman/memory_store/vectors/store.rs deleted file mode 100644 index b8f27aa5a..000000000 --- a/src/openhuman/memory_store/vectors/store.rs +++ /dev/null @@ -1,521 +0,0 @@ -//! Local vector store backed by SQLite. -//! -//! Provides a self-contained vector database for storing, searching, and -//! managing text embeddings. Uses SQLite for persistence and brute-force -//! cosine similarity for retrieval (fast enough for on-device workloads up -//! to ~100K vectors). -//! -//! # Usage -//! -//! ```ignore -//! let embedder = Arc::new(OllamaEmbedding::default()); -//! let store = VectorStore::open(db_path, embedder)?; -//! -//! store.insert("doc-1", "notes", "The quick brown fox", json!({})).await?; -//! let results = store.search("notes", "fast animal", 5).await?; -//! ``` - -use std::path::Path; -use std::sync::Arc; - -use anyhow::Context; -use parking_lot::Mutex; -use rusqlite::Connection; - -use crate::openhuman::embeddings::EmbeddingProvider; - -/// SQL to create the vector store schema. -const INIT_SQL: &str = " - PRAGMA journal_mode = WAL; - PRAGMA synchronous = NORMAL; - - CREATE TABLE IF NOT EXISTS vectors ( - id TEXT NOT NULL, - namespace TEXT NOT NULL, - text TEXT NOT NULL, - embedding BLOB NOT NULL, - metadata TEXT NOT NULL DEFAULT '{}', - created_at REAL NOT NULL, - updated_at REAL NOT NULL, - PRIMARY KEY (namespace, id) - ); - CREATE INDEX IF NOT EXISTS idx_vectors_ns ON vectors(namespace); - - CREATE TABLE IF NOT EXISTS store_meta ( - key TEXT PRIMARY KEY, - value TEXT NOT NULL, - updated_at REAL NOT NULL - ); -"; - -/// A single search result from the vector store. -#[derive(Debug, Clone)] -pub struct SearchResult { - /// The stored document ID. - pub id: String, - /// The namespace. - pub namespace: String, - /// The original text. - pub text: String, - /// Cosine similarity score (0.0 – 1.0). - pub score: f64, - /// Arbitrary JSON metadata attached at insert time. - pub metadata: serde_json::Value, -} - -/// SQLite-backed local vector store. -/// -/// Thread-safe: the inner connection is behind a `parking_lot::Mutex` and -/// the struct is `Send + Sync`. Embedding calls are async and run through -/// the configured [`EmbeddingProvider`]. -pub struct VectorStore { - conn: Arc>, - embedder: Arc, -} - -impl VectorStore { - /// Opens (or creates) a vector store at the given SQLite database path. - /// - /// On first open the embedding provider name, model-name-hint, and - /// dimensions are persisted to a `store_meta` table. On subsequent opens - /// the stored dimensions are compared against the runtime embedder and an - /// error is returned if they mismatch (prevents silent cosine-similarity - /// corruption from mixed-dimension vectors). - pub fn open(db_path: &Path, embedder: Arc) -> anyhow::Result { - if let Some(parent) = db_path.parent() { - std::fs::create_dir_all(parent)?; - } - - let conn = Connection::open(db_path)?; - conn.execute_batch(INIT_SQL)?; - - Self::check_or_store_meta(&conn, &*embedder)?; - - tracing::debug!( - target: "embeddings.store", - "[vector-store] opened at {}, embedder={}, dims={}", - db_path.display(), - embedder.name(), - embedder.dimensions() - ); - - Ok(Self { - conn: Arc::new(Mutex::new(conn)), - embedder, - }) - } - - /// Opens an in-memory vector store (useful for tests). - pub fn open_in_memory(embedder: Arc) -> anyhow::Result { - let conn = Connection::open_in_memory()?; - conn.execute_batch(INIT_SQL)?; - Self::check_or_store_meta(&conn, &*embedder)?; - Ok(Self { - conn: Arc::new(Mutex::new(conn)), - embedder, - }) - } - - /// Returns a reference to the embedding provider. - pub fn embedder(&self) -> &dyn EmbeddingProvider { - self.embedder.as_ref() - } - - /// Persist or validate the embedding configuration in `store_meta`. - fn check_or_store_meta( - conn: &Connection, - embedder: &dyn EmbeddingProvider, - ) -> anyhow::Result<()> { - let now = now_ts(); - let stored_dims: Option = conn - .query_row( - "SELECT value FROM store_meta WHERE key = 'embed_dims'", - [], - |row| row.get(0), - ) - .ok(); - - match stored_dims { - None => { - // First open — persist metadata. - let stmts: &[(&str, &str)] = &[ - ("embed_provider", embedder.name()), - ("embed_dims", &embedder.dimensions().to_string()), - ]; - for (key, value) in stmts { - conn.execute( - "INSERT OR REPLACE INTO store_meta (key, value, updated_at) VALUES (?1, ?2, ?3)", - rusqlite::params![key, value, now], - )?; - } - tracing::debug!( - target: "embeddings.store", - "[vector-store] stored meta: provider={}, dims={}", - embedder.name(), - embedder.dimensions() - ); - } - Some(dims_str) => { - let stored: usize = dims_str.parse().unwrap_or(0); - let runtime = embedder.dimensions(); - if stored != 0 && runtime != 0 && stored != runtime { - anyhow::bail!( - "vector store dimension mismatch: database was created with \ - {stored}-dim embeddings but the current provider ({}) uses \ - {runtime} dims. Delete the database or reconfigure the provider.", - embedder.name() - ); - } - } - } - - Ok(()) - } - - // ── Write operations ───────────────────────────────────── - - /// Inserts or updates a text entry. The text is embedded automatically. - /// - /// If an entry with the same `(namespace, id)` already exists it is replaced. - pub async fn insert( - &self, - id: &str, - namespace: &str, - text: &str, - metadata: serde_json::Value, - ) -> anyhow::Result<()> { - tracing::trace!( - target: "embeddings.store", - "[vector-store] insert: id={id}, ns={namespace}, text_len={}", - text.len() - ); - let embedding = self.embedder.embed_one(text).await?; - self.insert_with_vector(id, namespace, text, &embedding, metadata) - } - - /// Inserts with a pre-computed embedding vector (skips the embed call). - pub fn insert_with_vector( - &self, - id: &str, - namespace: &str, - text: &str, - embedding: &[f32], - metadata: serde_json::Value, - ) -> anyhow::Result<()> { - let blob = vec_to_bytes(embedding); - let meta_str = serde_json::to_string(&metadata)?; - let now = now_ts(); - - let conn = self.conn.lock(); - conn.execute( - "INSERT OR REPLACE INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", - rusqlite::params![id, namespace, text, blob, meta_str, now, now], - )?; - - tracing::trace!( - target: "embeddings.store", - "[vector-store] inserted id={id} ns={namespace} dims={}", - embedding.len() - ); - - Ok(()) - } - - /// Bulk-insert multiple entries. Each text is embedded automatically. - pub async fn insert_batch( - &self, - namespace: &str, - entries: &[(&str, &str, serde_json::Value)], // (id, text, metadata) - ) -> anyhow::Result<()> { - if entries.is_empty() { - return Ok(()); - } - - tracing::debug!( - target: "embeddings.store", - "[vector-store] insert_batch: ns={namespace}, count={}", - entries.len() - ); - - let texts: Vec<&str> = entries.iter().map(|(_, text, _)| *text).collect(); - let embeddings = self.embedder.embed(&texts).await?; - - if embeddings.len() != entries.len() { - anyhow::bail!( - "embedding count mismatch: got {} embeddings for {} entries", - embeddings.len(), - entries.len() - ); - } - - let now = now_ts(); - let conn = self.conn.lock(); - let tx = conn.unchecked_transaction()?; - - for ((id, text, metadata), embedding) in entries.iter().zip(embeddings.iter()) { - let blob = vec_to_bytes(embedding); - let meta_str = serde_json::to_string(metadata)?; - tx.execute( - "INSERT OR REPLACE INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", - rusqlite::params![id, namespace, text, blob, meta_str, now, now], - )?; - } - - tx.commit()?; - - tracing::debug!( - target: "embeddings.store", - "[vector-store] batch inserted {} entries in ns={namespace}", - entries.len() - ); - - Ok(()) - } - - // ── Search ─────────────────────────────────────────────── - - /// Searches for the `limit` most similar entries to `query` within a namespace. - /// - /// The query is embedded via the configured provider and compared against - /// all stored vectors using cosine similarity. - pub async fn search( - &self, - namespace: &str, - query: &str, - limit: usize, - ) -> anyhow::Result> { - tracing::trace!( - target: "embeddings.store", - "[vector-store] search: ns={namespace}, limit={limit}, query_len={}", - query.len() - ); - let query_vec = self.embedder.embed_one(query).await?; - self.search_by_vector(namespace, &query_vec, limit) - } - - /// Searches using a pre-computed query vector. - pub fn search_by_vector( - &self, - namespace: &str, - query_vec: &[f32], - limit: usize, - ) -> anyhow::Result> { - if limit == 0 { - tracing::trace!( - target: "embeddings.store", - "[vector-store] search_by_vector: limit=0, returning empty" - ); - return Ok(Vec::new()); - } - - let conn = self.conn.lock(); - let mut stmt = conn.prepare( - "SELECT id, namespace, text, embedding, metadata FROM vectors WHERE namespace = ?1", - )?; - - let rows: Vec<(String, String, String, Vec, String)> = stmt - .query_map(rusqlite::params![namespace], |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - row.get::<_, String>(2)?, - row.get::<_, Vec>(3)?, - row.get::<_, String>(4)?, - )) - })? - .collect::>>()?; - - let scanned = rows.len(); - - // Score-only intermediate: keep metadata as the raw JSON string instead - // of parsing it here. We scan every vector in the namespace but return - // only `limit` rows, so parsing the metadata of all N candidates would - // throw away all but `limit` parses. Defer the parse until after the - // truncation below, where it runs `limit` times instead of N. - struct ScoredRow { - score: f64, - id: String, - namespace: String, - text: String, - meta_str: String, - } - - let mut scored: Vec = rows - .into_iter() - .map(|(id, namespace, text, blob, meta_str)| { - let stored_vec = bytes_to_vec(&blob); - let score = cosine_similarity(query_vec, &stored_vec); - ScoredRow { - score, - id, - namespace, - text, - meta_str, - } - }) - .collect(); - - // Sort descending by score. - scored.sort_by(|a, b| { - b.score - .partial_cmp(&a.score) - .unwrap_or(std::cmp::Ordering::Equal) - }); - scored.truncate(limit); - - // Parse metadata only for the rows that survived truncation. Invalid - // JSON falls back to Null, but log it so data issues stay diagnosable. - let results: Vec = scored - .into_iter() - .map(|row| { - let metadata = serde_json::from_str(&row.meta_str).unwrap_or_else(|err| { - tracing::debug!( - target: "embeddings.store", - "[vector-store] invalid metadata json: id={}, ns={}, err={err}", - row.id, - row.namespace, - ); - serde_json::Value::Null - }); - SearchResult { - id: row.id, - namespace: row.namespace, - text: row.text, - score: row.score, - metadata, - } - }) - .collect(); - - tracing::trace!( - target: "embeddings.store", - "[vector-store] search_by_vector: ns={namespace}, scanned={scanned}, returned={}", - results.len() - ); - - Ok(results) - } - - // ── Delete / management ────────────────────────────────── - - /// Deletes a single entry by ID within a namespace. - /// - /// Returns `true` if a row was actually deleted. - pub fn delete(&self, namespace: &str, id: &str) -> anyhow::Result { - let conn = self.conn.lock(); - let affected = conn.execute( - "DELETE FROM vectors WHERE namespace = ?1 AND id = ?2", - rusqlite::params![namespace, id], - )?; - - tracing::trace!( - target: "embeddings.store", - "[vector-store] delete: ns={namespace}, id={id}, affected={affected}" - ); - - Ok(affected > 0) - } - - /// Deletes all entries in a namespace. - /// - /// Returns the number of deleted rows. - pub fn clear_namespace(&self, namespace: &str) -> anyhow::Result { - let conn = self.conn.lock(); - let affected = conn.execute( - "DELETE FROM vectors WHERE namespace = ?1", - rusqlite::params![namespace], - )?; - - tracing::debug!( - target: "embeddings.store", - "[vector-store] cleared namespace={namespace}, deleted={affected}" - ); - - Ok(affected) - } - - /// Returns the number of entries in a namespace (or all if `None`). - pub fn count(&self, namespace: Option<&str>) -> anyhow::Result { - let conn = self.conn.lock(); - let count: i64 = match namespace { - Some(ns) => conn.query_row( - "SELECT COUNT(*) FROM vectors WHERE namespace = ?1", - rusqlite::params![ns], - |row| row.get(0), - )?, - None => conn.query_row("SELECT COUNT(*) FROM vectors", [], |row| row.get(0))?, - }; - usize::try_from(count).context("vector count did not fit in usize") - } - - /// Lists all distinct namespaces. - pub fn list_namespaces(&self) -> anyhow::Result> { - let conn = self.conn.lock(); - let mut stmt = conn.prepare("SELECT DISTINCT namespace FROM vectors ORDER BY namespace")?; - let namespaces: Vec = stmt - .query_map([], |row| row.get(0))? - .collect::>>()?; - Ok(namespaces) - } -} - -// ── Vector math utilities ──────────────────────────────────── - -/// Serializes a float vector to little-endian bytes for SQLite BLOB storage. -pub fn vec_to_bytes(v: &[f32]) -> Vec { - let mut bytes = Vec::with_capacity(v.len() * 4); - for &f in v { - bytes.extend_from_slice(&f.to_le_bytes()); - } - bytes -} - -/// Deserializes little-endian bytes back to a float vector. -pub fn bytes_to_vec(bytes: &[u8]) -> Vec { - bytes - .chunks_exact(4) - .map(|chunk| { - let arr: [u8; 4] = chunk.try_into().unwrap_or([0; 4]); - f32::from_le_bytes(arr) - }) - .collect() -} - -/// Computes cosine similarity between two vectors. Returns 0.0 for -/// mismatched lengths, empty vectors, or zero-magnitude vectors. -pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 { - if a.len() != b.len() || a.is_empty() { - return 0.0; - } - let mut dot = 0.0_f64; - let mut norm_a = 0.0_f64; - let mut norm_b = 0.0_f64; - for (x, y) in a.iter().zip(b.iter()) { - let x = f64::from(*x); - let y = f64::from(*y); - dot += x * y; - norm_a += x * x; - norm_b += y * y; - } - let denom = norm_a.sqrt() * norm_b.sqrt(); - if denom <= f64::EPSILON { - return 0.0; - } - (dot / denom).clamp(0.0, 1.0) -} - -fn now_ts() -> f64 { - use std::time::{SystemTime, UNIX_EPOCH}; - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs_f64()) - .unwrap_or(0.0) -} - -// ── Tests ──────────────────────────────────────────────────── - -#[cfg(test)] -#[path = "store_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_store/vectors/store_tests.rs b/src/openhuman/memory_store/vectors/store_tests.rs deleted file mode 100644 index 2d3c656d5..000000000 --- a/src/openhuman/memory_store/vectors/store_tests.rs +++ /dev/null @@ -1,481 +0,0 @@ -use super::*; -use crate::openhuman::embeddings::EmbeddingProvider; -use serde_json::json; - -/// A test embedding provider that returns deterministic vectors. -struct FakeEmbedding { - dims: usize, -} - -#[async_trait::async_trait] -impl EmbeddingProvider for FakeEmbedding { - fn name(&self) -> &str { - "fake" - } - fn model_id(&self) -> &str { - "fake" - } - fn dimensions(&self) -> usize { - self.dims - } - async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { - Ok(texts.iter().map(|t| text_to_vec(t, self.dims)).collect()) - } -} - -fn text_to_vec(text: &str, dims: usize) -> Vec { - let mut vec = vec![0.0_f32; dims]; - for (i, byte) in text.bytes().enumerate() { - vec[i % dims] += byte as f32 / 255.0; - } - let norm: f32 = vec.iter().map(|x| x * x).sum::().sqrt(); - if norm > 0.0 { - for x in &mut vec { - *x /= norm; - } - } - vec -} - -struct MismatchEmbedding; - -#[async_trait::async_trait] -impl EmbeddingProvider for MismatchEmbedding { - fn name(&self) -> &str { - "mismatch" - } - fn model_id(&self) -> &str { - "mismatch" - } - fn dimensions(&self) -> usize { - 2 - } - async fn embed(&self, _texts: &[&str]) -> anyhow::Result>> { - Ok(vec![vec![1.0, 0.0]]) - } -} - -fn fake_store(dims: usize) -> VectorStore { - VectorStore::open_in_memory(Arc::new(FakeEmbedding { dims })).unwrap() -} - -// ── vec_to_bytes / bytes_to_vec ───────────────────────── - -#[test] -fn roundtrip_vec_bytes() { - let original = vec![1.0_f32, -2.5, 3.14, 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)); -} - -#[test] -fn empty_vec_roundtrip() { - assert!(bytes_to_vec(&vec_to_bytes(&[])).is_empty()); -} - -#[test] -fn bytes_to_vec_truncates_partial_bytes() { - assert_eq!(bytes_to_vec(&[0u8; 5]).len(), 1); -} - -// ── cosine_similarity ─────────────────────────────────── - -#[test] -fn cosine_identical() { - let v = vec![1.0_f32, 2.0, 3.0]; - assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6); -} - -#[test] -fn cosine_orthogonal() { - assert!(cosine_similarity(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-6); -} - -#[test] -fn cosine_opposite() { - assert!(cosine_similarity(&[1.0, 0.0], &[-1.0, 0.0]).abs() < 1e-6); -} - -#[test] -fn cosine_mismatched_lengths() { - assert_eq!(cosine_similarity(&[1.0, 2.0], &[1.0, 2.0, 3.0]), 0.0); -} - -#[test] -fn cosine_empty() { - assert_eq!(cosine_similarity(&[], &[]), 0.0); -} - -#[test] -fn cosine_zero_vector() { - assert_eq!(cosine_similarity(&[0.0, 0.0], &[1.0, 0.0]), 0.0); -} - -#[test] -fn cosine_similar_high() { - assert!(cosine_similarity(&[1.0, 2.0, 3.0], &[1.1, 2.1, 3.1]) > 0.99); -} - -// ── VectorStore: open / metadata ──────────────────────── - -#[test] -fn open_in_memory_succeeds() { - let store = fake_store(3); - assert_eq!(store.count(None).unwrap(), 0); -} - -#[test] -fn open_on_disk() { - let dir = tempfile::tempdir().unwrap(); - let db_path = dir.path().join("sub/dir/vectors.db"); - let store = VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 3 })).unwrap(); - assert_eq!(store.count(None).unwrap(), 0); - assert!(db_path.exists()); -} - -#[test] -fn open_reopen_same_dims_succeeds() { - let dir = tempfile::tempdir().unwrap(); - let db_path = dir.path().join("v.db"); - VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 4 })).unwrap(); - // Reopen with same dims — should work. - VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 4 })).unwrap(); -} - -#[test] -fn open_reopen_different_dims_errors() { - let dir = tempfile::tempdir().unwrap(); - let db_path = dir.path().join("v.db"); - VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 4 })).unwrap(); - let result = VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 8 })); - let msg = result.err().expect("should be an error").to_string(); - assert!(msg.contains("dimension mismatch"), "msg: {msg}"); - assert!(msg.contains("4"), "should mention stored dims: {msg}"); - assert!(msg.contains("8"), "should mention runtime dims: {msg}"); -} - -#[test] -fn embedder_accessor() { - let store = fake_store(3); - assert_eq!(store.embedder().name(), "fake"); - assert_eq!(store.embedder().dimensions(), 3); -} - -// ── insert + count ────────────────────────────────────── - -#[tokio::test] -async fn insert_and_count() { - let store = fake_store(4); - store.insert("a", "ns1", "hello", json!({})).await.unwrap(); - store.insert("b", "ns1", "world", json!({})).await.unwrap(); - store.insert("c", "ns2", "other", json!({})).await.unwrap(); - assert_eq!(store.count(Some("ns1")).unwrap(), 2); - assert_eq!(store.count(Some("ns2")).unwrap(), 1); - assert_eq!(store.count(None).unwrap(), 3); -} - -#[tokio::test] -async fn insert_upsert_replaces() { - let store = fake_store(4); - store - .insert("a", "ns", "original", json!({"v": 1})) - .await - .unwrap(); - store - .insert("a", "ns", "updated", json!({"v": 2})) - .await - .unwrap(); - assert_eq!(store.count(Some("ns")).unwrap(), 1); - let results = store - .search_by_vector("ns", &text_to_vec("updated", 4), 10) - .unwrap(); - assert_eq!(results[0].text, "updated"); - assert_eq!(results[0].metadata["v"], 2); -} - -#[test] -fn insert_with_vector_sync() { - let store = fake_store(3); - store - .insert_with_vector("id1", "ns", "text", &[1.0, 0.0, 0.0], json!({"k": "v"})) - .unwrap(); - assert_eq!(store.count(Some("ns")).unwrap(), 1); -} - -// ── insert_batch ──────────────────────────────────────── - -#[tokio::test] -async fn insert_batch_multiple() { - let store = fake_store(4); - let entries = vec![ - ("a", "alpha", json!({})), - ("b", "beta", json!({})), - ("c", "gamma", json!({})), - ]; - store.insert_batch("ns", &entries).await.unwrap(); - assert_eq!(store.count(Some("ns")).unwrap(), 3); -} - -#[tokio::test] -async fn insert_batch_empty() { - let store = fake_store(4); - store.insert_batch("ns", &[]).await.unwrap(); - assert_eq!(store.count(None).unwrap(), 0); -} - -#[tokio::test] -async fn insert_batch_mismatch_error() { - let store = VectorStore::open_in_memory(Arc::new(MismatchEmbedding)).unwrap(); - let entries = vec![("a", "alpha", json!({})), ("b", "beta", json!({}))]; - let err = store.insert_batch("ns", &entries).await.unwrap_err(); - assert!(err.to_string().contains("mismatch")); -} - -// ── search ────────────────────────────────────────────── - -#[tokio::test] -async fn search_returns_ranked_results() { - let store = fake_store(8); - store - .insert("a", "ns", "the quick brown fox", json!({})) - .await - .unwrap(); - store - .insert("b", "ns", "a lazy dog sleeps", json!({})) - .await - .unwrap(); - store - .insert("c", "ns", "the quick brown fox jumps", json!({})) - .await - .unwrap(); - let results = store.search("ns", "the quick brown fox", 2).await.unwrap(); - assert_eq!(results.len(), 2); - assert!(results[0].score >= results[1].score); -} - -#[tokio::test] -async fn search_respects_limit() { - let store = fake_store(4); - for i in 0..10 { - store - .insert(&format!("id-{i}"), "ns", &format!("text {i}"), json!({})) - .await - .unwrap(); - } - assert_eq!(store.search("ns", "text", 3).await.unwrap().len(), 3); -} - -#[tokio::test] -async fn search_empty_namespace() { - let store = fake_store(4); - assert!(store.search("empty", "query", 10).await.unwrap().is_empty()); -} - -#[tokio::test] -async fn search_namespace_isolation() { - let store = fake_store(4); - store.insert("a", "ns1", "hello", json!({})).await.unwrap(); - store.insert("b", "ns2", "hello", json!({})).await.unwrap(); - assert_eq!(store.search("ns1", "hello", 10).await.unwrap()[0].id, "a"); - assert_eq!(store.search("ns2", "hello", 10).await.unwrap()[0].id, "b"); -} - -// ── search_by_vector ──────────────────────────────────── - -#[test] -fn search_by_vector_limit_zero() { - let store = fake_store(3); - store - .insert_with_vector("a", "ns", "t", &[1.0, 0.0, 0.0], json!({})) - .unwrap(); - assert!(store - .search_by_vector("ns", &[1.0, 0.0, 0.0], 0) - .unwrap() - .is_empty()); -} - -/// Metadata is parsed only for rows that survive truncation, so the parse -/// must align with the post-sort order — each returned row must carry its -/// own metadata, and dropped rows must not appear. This pins the deferred -/// parse against an off-by-one or mis-zipped mapping after sort/truncate. -#[test] -fn search_by_vector_returns_metadata_of_surviving_rows() { - let store = fake_store(3); - store - .insert_with_vector("near", "ns", "t", &[1.0, 0.0, 0.0], json!({"tag": "near"})) - .unwrap(); - store - .insert_with_vector("mid", "ns", "t", &[0.7, 0.7, 0.0], json!({"tag": "mid"})) - .unwrap(); - store - .insert_with_vector("far", "ns", "t", &[0.0, 0.0, 1.0], json!({"tag": "far"})) - .unwrap(); - - let results = store.search_by_vector("ns", &[1.0, 0.0, 0.0], 2).unwrap(); - - assert_eq!(results.len(), 2, "limit should drop the least similar row"); - assert_eq!(results[0].id, "near"); - assert_eq!(results[1].id, "mid"); - assert!( - results.iter().all(|hit| hit.id != "far"), - "the truncated row must not leak into the results" - ); - // The deferred parse must attach each row's own metadata, not a neighbour's. - for hit in &results { - assert_eq!( - hit.metadata.get("tag").and_then(|v| v.as_str()), - Some(hit.id.as_str()), - "metadata tag should match the row id for {}", - hit.id - ); - } -} - -/// A row with corrupt metadata JSON (e.g. a hand-edited or partially written -/// DB) must not break the search — the deferred parse falls back to `Null` -/// rather than dropping the row or erroring. Inserted raw because -/// `insert_with_vector` always serializes valid JSON. -#[test] -fn search_by_vector_falls_back_to_null_on_invalid_metadata() { - let store = fake_store(3); - { - let conn = store.conn.lock(); - conn.execute( - "INSERT INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", - rusqlite::params![ - "bad", - "ns", - "t", - vec_to_bytes(&[1.0, 0.0, 0.0]), - "{not valid json", - 0.0_f64, - 0.0_f64 - ], - ) - .unwrap(); - } - - let results = store.search_by_vector("ns", &[1.0, 0.0, 0.0], 5).unwrap(); - - assert_eq!(results.len(), 1, "the row must still be returned"); - assert_eq!(results[0].id, "bad"); - assert!( - results[0].metadata.is_null(), - "invalid metadata json must fall back to Null" - ); -} - -#[test] -fn search_by_vector_scores_correct() { - let store = fake_store(3); - store - .insert_with_vector("x", "ns", "x", &[1.0, 0.0, 0.0], json!({})) - .unwrap(); - store - .insert_with_vector("y", "ns", "y", &[0.0, 1.0, 0.0], json!({})) - .unwrap(); - let results = store.search_by_vector("ns", &[1.0, 0.0, 0.0], 2).unwrap(); - assert_eq!(results[0].id, "x"); - assert!((results[0].score - 1.0).abs() < 1e-6); - assert!(results[1].score < 1e-6); -} - -#[test] -fn search_by_vector_preserves_metadata() { - let store = fake_store(2); - store - .insert_with_vector("a", "ns", "t", &[1.0, 0.0], json!({"key": "value"})) - .unwrap(); - assert_eq!( - store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap()[0].metadata["key"], - "value" - ); -} - -#[test] -fn search_handles_invalid_metadata_json() { - let store = fake_store(2); - { - let conn = store.conn.lock(); - conn.execute( - "INSERT INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES ('bad', 'ns', 'text', ?1, 'not-json', 0.0, 0.0)", - rusqlite::params![vec_to_bytes(&[1.0, 0.0])], - ) - .unwrap(); - } - let results = store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap(); - assert_eq!(results[0].id, "bad"); - assert!(results[0].metadata.is_null()); -} - -// ── delete ────────────────────────────────────────────── - -#[tokio::test] -async fn delete_existing() { - let store = fake_store(4); - store.insert("a", "ns", "text", json!({})).await.unwrap(); - assert!(store.delete("ns", "a").unwrap()); - assert_eq!(store.count(Some("ns")).unwrap(), 0); -} - -#[test] -fn delete_nonexistent() { - assert!(!fake_store(3).delete("ns", "no-such-id").unwrap()); -} - -#[tokio::test] -async fn delete_wrong_namespace() { - let store = fake_store(4); - store.insert("a", "ns1", "text", json!({})).await.unwrap(); - assert!(!store.delete("ns2", "a").unwrap()); - assert_eq!(store.count(Some("ns1")).unwrap(), 1); -} - -// ── clear_namespace ───────────────────────────────────── - -#[tokio::test] -async fn clear_namespace_removes_all() { - let store = fake_store(4); - store.insert("a", "ns", "one", json!({})).await.unwrap(); - store.insert("b", "ns", "two", json!({})).await.unwrap(); - store - .insert("c", "other", "three", json!({})) - .await - .unwrap(); - assert_eq!(store.clear_namespace("ns").unwrap(), 2); - assert_eq!(store.count(Some("ns")).unwrap(), 0); - assert_eq!(store.count(Some("other")).unwrap(), 1); -} - -#[test] -fn clear_empty_namespace() { - assert_eq!(fake_store(3).clear_namespace("empty").unwrap(), 0); -} - -// ── list_namespaces ───────────────────────────────────── - -#[tokio::test] -async fn list_namespaces_empty() { - assert!(fake_store(3).list_namespaces().unwrap().is_empty()); -} - -#[tokio::test] -async fn list_namespaces_populated() { - let store = fake_store(4); - store.insert("a", "beta", "t", json!({})).await.unwrap(); - store.insert("b", "alpha", "t", json!({})).await.unwrap(); - store.insert("c", "beta", "t", json!({})).await.unwrap(); - assert_eq!(store.list_namespaces().unwrap(), vec!["alpha", "beta"]); -} - -// ── count ─────────────────────────────────────────────── - -#[test] -fn count_empty() { - let store = fake_store(3); - assert_eq!(store.count(None).unwrap(), 0); - assert_eq!(store.count(Some("ns")).unwrap(), 0); -} diff --git a/src/openhuman/memory_sync/canonicalize/README.md b/src/openhuman/memory_sync/canonicalize/README.md index 9ed564430..27c6a4af4 100644 --- a/src/openhuman/memory_sync/canonicalize/README.md +++ b/src/openhuman/memory_sync/canonicalize/README.md @@ -7,10 +7,10 @@ Adapters do not interpret content semantically; they only normalise shape and ca ## Files - [`mod.rs`](mod.rs) — `CanonicalisedSource` struct, generic `CanonicaliseRequest

` envelope, and `normalize_source_ref` helper shared by all adapters. -- [`chat.rs`](chat.rs) — chat transcripts (Slack / Discord / Telegram / WhatsApp) → Markdown of `## — \n` blocks. Sorts messages and captures `time_range`. Produces empty-input `Ok(None)`. -- [`document.rs`](document.rs) — single documents (Notion page, Drive doc, meeting note, uploaded file) → trimmed body Markdown. `time_range` collapses to a single point at `modified_at`. -- [`email.rs`](email.rs) — email threads (Gmail + generic) → per-message `---\nFrom: …\nSubject: …\nDate: …\n\n` blocks. Bodies pass through `email_clean::clean_body` first. -- [`email_clean.rs`](email_clean.rs) — pure-string helpers: `clean_body` (strip reply chains + footer/legal boilerplate), `truncate_body`, `md_escape`, `extract_email`, `parse_message_date`. Used by both the email canonicaliser and the `gmail-fetch-emails` bin. +- `vendor/tinycortex/src/memory/ingest/canonicalize/chat.rs` — chat transcripts (Slack / Discord / Telegram / WhatsApp) → Markdown of `## — \n` blocks. Sorts messages and captures `time_range`. Produces empty-input `Ok(None)`. +- `vendor/tinycortex/src/memory/ingest/canonicalize/document.rs` — single documents (Notion page, Drive doc, meeting note, uploaded file) → trimmed body Markdown. `time_range` collapses to a single point at `modified_at`. +- `vendor/tinycortex/src/memory/ingest/canonicalize/email.rs` — email threads (Gmail + generic) → per-message `---\nFrom: …\nSubject: …\nDate: …\n\n` blocks. Bodies pass through `email_clean::clean_body` first. +- `vendor/tinycortex/src/memory/ingest/canonicalize/email_clean.rs` — pure-string helpers: `clean_body` (strip reply chains + footer/legal boilerplate), `truncate_body`, `md_escape`, `extract_email`, `parse_message_date`. Used by both the email canonicaliser and the `gmail-fetch-emails` bin. ## Output contract diff --git a/src/openhuman/memory_sync/canonicalize/chat.rs b/src/openhuman/memory_sync/canonicalize/chat.rs deleted file mode 100644 index 7936605a0..000000000 --- a/src/openhuman/memory_sync/canonicalize/chat.rs +++ /dev/null @@ -1,262 +0,0 @@ -//! Chat transcripts → canonical Markdown. -//! -//! Chat sources are scoped by **channel or group**. A batch of chat messages -//! from the same channel becomes one [`CanonicalisedSource`]; the chunker -//! slices it by token budget downstream. -//! -//! Output format (no leading `# ...` header — that info lives in front-matter -//! once Phase MD-content lands; the chunker splits at `## ` boundaries): -//! ```md -//! ## 2026-04-21T10:12:00Z — Alice -//! Message body here. -//! -//! ## 2026-04-21T10:12:40Z — Bob -//! Reply body here. -//! ``` - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use super::{normalize_source_ref, CanonicalisedSource}; -use crate::openhuman::memory_store::chunks::types::{Metadata, SourceKind}; - -/// One chat message in a channel/group. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ChatMessage { - /// Author display name or id. - pub author: String, - /// When the message was sent (epoch-ms integer or RFC 3339 string). - #[serde( - serialize_with = "chrono::serde::ts_milliseconds::serialize", - deserialize_with = "super::deserialize_flexible_timestamp" - )] - pub timestamp: DateTime, - /// Plain text / markdown body. - pub text: String, - /// Optional per-message provenance pointer (permalink or `platform://...`). - #[serde(default)] - pub source_ref: Option, -} - -/// Adapter input — a batch of messages from one logical channel. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ChatBatch { - /// Platform name used in the header (e.g. `slack`, `discord`, `telegram`). - pub platform: String, - /// Human-readable channel / group name for the header. - pub channel_label: String, - /// Ordered messages (chronological; adapter sorts defensively). - pub messages: Vec, -} - -/// Canonicalise a chat batch. -/// -/// Returns `Ok(None)` if the batch has zero messages — callers treat that as -/// "nothing to ingest" and skip. -pub fn canonicalise( - source_id: &str, - owner: &str, - tags: &[String], - batch: ChatBatch, -) -> Result, String> { - if batch.messages.is_empty() { - return Ok(None); - } - let mut messages = batch.messages; - messages.sort_by_key(|m| m.timestamp); - - let first_ts = messages.first().map(|m| m.timestamp).unwrap(); - let last_ts = messages.last().map(|m| m.timestamp).unwrap(); - - let mut md = String::new(); - // No leading `# Chat transcript — ...` header. Platform / channel info - // belongs in the MD front-matter (Phase MD-content). The chunker splits - // this output at `## ` boundaries so each message becomes one chunk. - for msg in &messages { - md.push_str(&format!( - "## {} — {}\n{}\n\n", - msg.timestamp.to_rfc3339(), - msg.author, - msg.text.trim() - )); - } - - // Provenance points at the batch's first message by default (or whatever - // the caller passed on the first message). - let source_ref = normalize_source_ref(messages.first().and_then(|m| m.source_ref.clone())); - - let metadata = Metadata { - source_kind: SourceKind::Chat, - source_id: source_id.to_string(), - owner: owner.to_string(), - timestamp: first_ts, - time_range: (first_ts, last_ts), - tags: tags.to_vec(), - source_ref, - path_scope: None, - }; - Ok(Some(CanonicalisedSource { - markdown: md, - metadata, - })) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - fn msg(ts_ms: i64, author: &str, text: &str) -> ChatMessage { - ChatMessage { - author: author.to_string(), - timestamp: Utc.timestamp_millis_opt(ts_ms).unwrap(), - text: text.to_string(), - source_ref: Some(format!("slack://x/{ts_ms}")), - } - } - - #[test] - fn empty_batch_returns_none() { - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![], - }; - assert!(canonicalise("slack:#eng", "alice", &[], b) - .unwrap() - .is_none()); - } - - #[test] - fn messages_are_sorted_and_range_captured() { - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![ - msg(2000, "bob", "second"), - msg(1000, "alice", "first"), - msg(3000, "carol", "third"), - ], - }; - let out = canonicalise("slack:#eng", "alice", &["eng".into()], b) - .unwrap() - .unwrap(); - assert_eq!(out.metadata.time_range.0.timestamp_millis(), 1000); - assert_eq!(out.metadata.time_range.1.timestamp_millis(), 3000); - // Check order in markdown - let pos_first = out.markdown.find("first").unwrap(); - let pos_second = out.markdown.find("second").unwrap(); - let pos_third = out.markdown.find("third").unwrap(); - assert!(pos_first < pos_second); - assert!(pos_second < pos_third); - } - - #[test] - fn includes_per_message_sections_without_header() { - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![msg(1000, "alice", "hello")], - }; - let out = canonicalise("slack:#eng", "alice", &[], b) - .unwrap() - .unwrap(); - // No leading `# Chat transcript` header — that info belongs in front-matter. - assert!( - !out.markdown.starts_with("# "), - "canonical chat MD must NOT start with a `# ` header" - ); - assert!( - out.markdown.starts_with("## "), - "must start with first `## ` message block" - ); - assert!(out.markdown.contains("— alice")); - assert!(out.markdown.contains("hello")); - } - - #[test] - fn source_ref_taken_from_first_message() { - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![msg(1000, "alice", "hi"), msg(2000, "bob", "hey")], - }; - let out = canonicalise("slack:#eng", "alice", &[], b) - .unwrap() - .unwrap(); - assert_eq!( - out.metadata.source_ref.as_ref().unwrap().value, - "slack://x/1000" - ); - } - - #[test] - fn metadata_carries_owner_and_tags() { - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![msg(1000, "alice", "hi")], - }; - let out = canonicalise( - "slack:#eng", - "alice@example.com", - &["eng".into(), "on-call".into()], - b, - ) - .unwrap() - .unwrap(); - assert_eq!(out.metadata.owner, "alice@example.com"); - assert_eq!(out.metadata.tags, vec!["eng", "on-call"]); - assert_eq!(out.metadata.source_kind, SourceKind::Chat); - } - - #[test] - fn blank_source_ref_is_dropped() { - let mut first = msg(1000, "alice", "hi"); - first.source_ref = Some(" ".into()); - let b = ChatBatch { - platform: "slack".into(), - channel_label: "#eng".into(), - messages: vec![first], - }; - let out = canonicalise("slack:#eng", "alice", &[], b) - .unwrap() - .unwrap(); - assert!(out.metadata.source_ref.is_none()); - } - - // ── Serde regression tests (CORE-2K / #3568) ──────────────────────────── - - #[test] - fn timestamp_epoch_ms_integer_still_works() { - let json = r#"{ - "author": "alice", - "timestamp": 1700000000000, - "text": "hello" - }"#; - let msg: ChatMessage = serde_json::from_str(json).expect("epoch-ms integer should parse"); - assert_eq!(msg.timestamp.timestamp_millis(), 1_700_000_000_000); - } - - #[test] - fn timestamp_iso8601_string_accepted() { - let json = r#"{ - "author": "alice", - "timestamp": "2026-05-17T19:30:00Z", - "text": "hello" - }"#; - let msg: ChatMessage = serde_json::from_str(json).expect("ISO-8601 string should parse"); - assert_eq!(msg.timestamp.timestamp(), 1_779_046_200); - } - - #[test] - fn timestamp_numeric_string_accepted() { - let json = r#"{ - "author": "alice", - "timestamp": "1700000000000", - "text": "hello" - }"#; - let msg: ChatMessage = serde_json::from_str(json).expect("numeric string should parse"); - assert_eq!(msg.timestamp.timestamp_millis(), 1_700_000_000_000); - } -} diff --git a/src/openhuman/memory_sync/canonicalize/document.rs b/src/openhuman/memory_sync/canonicalize/document.rs deleted file mode 100644 index 434122975..000000000 --- a/src/openhuman/memory_sync/canonicalize/document.rs +++ /dev/null @@ -1,229 +0,0 @@ -//! Standalone documents → canonical Markdown. -//! -//! Document sources are single-record (no grouping): one Notion page, one -//! Drive doc, one meeting-note file. The canonicaliser adds a small title -//! header and passes through the body; if the body is already markdown it -//! is kept verbatim. - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use super::{normalize_source_ref, CanonicalisedSource}; -use crate::openhuman::memory_store::chunks::types::{Metadata, SourceKind}; - -// ── Serde helpers ───────────────────────────────────────────────────────────── - -fn default_provider() -> String { - "unknown".to_string() -} - -fn now_utc() -> DateTime { - Utc::now() -} - -// ── Input struct ────────────────────────────────────────────────────────────── - -/// Adapter input for a single document. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct DocumentInput { - /// Provider name (e.g. `notion`, `drive`, `meeting_notes`). - /// Defaults to `"unknown"` when absent (fixes CORE-31). - #[serde(default = "default_provider")] - pub provider: String, - /// Document title. - pub title: String, - /// Document body (markdown preferred; plain text also accepted). - pub body: String, - /// When the document was last modified at the source. - /// - /// Accepts epoch-milliseconds integer (back-compat), RFC 3339 / ISO-8601 - /// string (fixes CORE-2K), or absent → `Utc::now()` (fixes CORE-2J). - #[serde( - default = "now_utc", - deserialize_with = "super::deserialize_flexible_timestamp" - )] - pub modified_at: DateTime, - /// Optional pointer back to source (URL, file path, Notion page id). - #[serde(default)] - pub source_ref: Option, -} - -/// Canonicalise a single document into a [`CanonicalisedSource`]. Returns -/// `Ok(None)` if both the title and body are empty — caller treats as nothing -/// to ingest. -pub fn canonicalise( - source_id: &str, - owner: &str, - tags: &[String], - doc: DocumentInput, - path_scope: Option, -) -> Result, String> { - if doc.body.trim().is_empty() && doc.title.trim().is_empty() { - return Ok(None); - } - - let mut md = String::new(); - // No leading `# provider — title` header. Provider / title info - // belongs in the MD front-matter (Phase MD-content). - md.push_str(doc.body.trim()); - md.push('\n'); - - Ok(Some(CanonicalisedSource { - markdown: md, - metadata: Metadata { - source_kind: SourceKind::Document, - source_id: source_id.to_string(), - owner: owner.to_string(), - timestamp: doc.modified_at, - time_range: (doc.modified_at, doc.modified_at), - tags: tags.to_vec(), - source_ref: normalize_source_ref(doc.source_ref), - path_scope, - }, - })) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - fn doc(title: &str, body: &str) -> DocumentInput { - DocumentInput { - provider: "notion".into(), - title: title.into(), - body: body.into(), - modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - source_ref: Some("notion://page/abc".into()), - } - } - - #[test] - fn empty_doc_returns_none() { - let d = DocumentInput { - provider: "notion".into(), - title: "".into(), - body: " \n ".into(), - modified_at: Utc::now(), - source_ref: None, - }; - assert!(canonicalise("d1", "alice", &[], d, None).unwrap().is_none()); - } - - #[test] - fn renders_body_without_header() { - let out = canonicalise( - "d1", - "alice", - &[], - doc("Launch plan", "step one\n\nstep two"), - None, - ) - .unwrap() - .unwrap(); - // No leading `# notion — Launch plan` header — that info belongs in front-matter. - assert!( - !out.markdown.starts_with("# "), - "canonical document MD must NOT start with a `# ` header" - ); - assert!(out.markdown.contains("step one")); - assert!(out.markdown.contains("step two")); - } - - #[test] - fn metadata_single_point_time_range() { - let out = canonicalise("d1", "alice", &[], doc("x", "y"), None) - .unwrap() - .unwrap(); - assert_eq!(out.metadata.time_range.0, out.metadata.time_range.1); - assert_eq!(out.metadata.source_kind, SourceKind::Document); - } - - #[test] - fn source_ref_carried_through() { - let out = canonicalise("d1", "alice", &["proj".into()], doc("x", "y"), None) - .unwrap() - .unwrap(); - assert_eq!( - out.metadata.source_ref.as_ref().unwrap().value, - "notion://page/abc" - ); - assert_eq!(out.metadata.tags, vec!["proj"]); - } - - #[test] - fn blank_source_ref_is_dropped() { - let mut input = doc("x", "y"); - input.source_ref = Some(" \n ".into()); - let out = canonicalise("d1", "alice", &[], input, None) - .unwrap() - .unwrap(); - assert!(out.metadata.source_ref.is_none()); - } - - // ── Serde regression / fix tests (CORE-2K / CORE-2J / CORE-31) ─────────── - - /// Regression: existing callers send epoch-ms as a JSON integer — must still work. - #[test] - fn modified_at_epoch_ms_integer_still_works() { - let json = r#"{ - "provider": "notion", - "title": "My doc", - "body": "content", - "modified_at": 1700000000000 - }"#; - let input: DocumentInput = - serde_json::from_str(json).expect("epoch-ms integer should parse"); - assert_eq!( - input.modified_at.timestamp_millis(), - 1_700_000_000_000, - "epoch-ms round-trip" - ); - } - - /// Fix CORE-2K: callers sending an ISO-8601 string must be accepted. - #[test] - fn modified_at_iso8601_string_accepted() { - let json = r#"{ - "provider": "drive", - "title": "Meeting notes", - "body": "agenda here", - "modified_at": "2026-05-17T19:30:00Z" - }"#; - let input: DocumentInput = - serde_json::from_str(json).expect("ISO-8601 string should parse"); - assert_eq!(input.modified_at.timestamp(), 1_779_046_200); - } - - /// Fix CORE-2J: omitting modified_at must default to approximately now (within 5 s). - #[test] - fn modified_at_missing_defaults_to_now() { - let before = Utc::now(); - let json = r#"{ - "provider": "notion", - "title": "No timestamp doc", - "body": "body text" - }"#; - let input: DocumentInput = - serde_json::from_str(json).expect("missing modified_at should parse"); - let after = Utc::now(); - assert!( - input.modified_at >= before && input.modified_at <= after, - "default modified_at {ts} must fall between {before} and {after}", - ts = input.modified_at, - ); - } - - /// Fix CORE-31: omitting provider must default to "unknown". - #[test] - fn provider_missing_defaults_to_unknown() { - let json = r#"{ - "title": "No provider doc", - "body": "body text", - "modified_at": 1700000000000 - }"#; - let input: DocumentInput = - serde_json::from_str(json).expect("missing provider should parse"); - assert_eq!(input.provider, "unknown"); - } -} diff --git a/src/openhuman/memory_sync/canonicalize/email.rs b/src/openhuman/memory_sync/canonicalize/email.rs deleted file mode 100644 index 18b7ff4e2..000000000 --- a/src/openhuman/memory_sync/canonicalize/email.rs +++ /dev/null @@ -1,300 +0,0 @@ -//! Email threads → canonical Markdown. -//! -//! Email sources are scoped by **participant set**. One participant bucket -//! becomes one [`CanonicalisedSource`]. Headers (From, To, Cc, Subject, Date) -//! surface in a small frontmatter-style block per message; the cleaned body -//! follows as markdown. Bodies pass through [`email_clean::clean_body`] before -//! rendering to strip reply chains, marketing footers, legal disclaimers, and -//! other boilerplate. - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use super::{email_clean, normalize_source_ref, CanonicalisedSource}; -use crate::openhuman::memory_store::chunks::types::{Metadata, SourceKind}; - -/// One email in a thread. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct EmailMessage { - pub from: String, - #[serde(default)] - pub to: Vec, - #[serde(default)] - pub cc: Vec, - pub subject: String, - /// When the message was sent (epoch-ms integer or RFC 3339 string). - #[serde( - serialize_with = "chrono::serde::ts_milliseconds::serialize", - deserialize_with = "super::deserialize_flexible_timestamp" - )] - pub sent_at: DateTime, - /// Plain-text or markdown body. - pub body: String, - /// Message-id header or provider URL; used for citation back to source. - #[serde(default)] - pub source_ref: Option, - /// List-Unsubscribe header for one-click unsubscribe actions - #[serde(default)] - pub list_unsubscribe: Option, -} - -/// A whole email thread. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct EmailThread { - /// Provider name used in the header (e.g. `gmail`, `outlook`). - pub provider: String, - /// Thread subject shown on top (usually the subject of the first message). - pub thread_subject: String, - /// Ordered messages (chronological; adapter sorts defensively). - pub messages: Vec, -} - -/// Canonicalise an email thread into a [`CanonicalisedSource`]. Bodies are -/// passed through [`email_clean::clean_body`] to strip reply chains and footer -/// boilerplate. Returns `Ok(None)` when the thread has no messages. -pub fn canonicalise( - source_id: &str, - owner: &str, - tags: &[String], - thread: EmailThread, -) -> Result, String> { - if thread.messages.is_empty() { - return Ok(None); - } - let mut messages = thread.messages; - messages.sort_by_key(|m| m.sent_at); - - let first_ts = messages.first().map(|m| m.sent_at).unwrap(); - let last_ts = messages.last().map(|m| m.sent_at).unwrap(); - - let mut md = String::new(); - // No leading `# Email thread — ...` header. Provider / subject info - // belongs in the MD front-matter (Phase MD-content). The chunker splits - // this output at `---\nFrom:` boundaries so each message becomes one chunk. - - for msg in &messages { - md.push_str("---\n"); - md.push_str(&format!("From: {}\n", msg.from)); - if !msg.to.is_empty() { - md.push_str(&format!("To: {}\n", msg.to.join(", "))); - } - if !msg.cc.is_empty() { - md.push_str(&format!("Cc: {}\n", msg.cc.join(", "))); - } - md.push_str(&format!("Subject: {}\n", msg.subject)); - md.push_str(&format!("Date: {}\n", msg.sent_at.to_rfc3339())); - - if let Some(unsub) = &msg.list_unsubscribe { - md.push_str(&format!("List-Unsubscribe: {}\n", unsub)); - } - md.push('\n'); - let cleaned = email_clean::clean_body(msg.body.trim()); - if cleaned.is_empty() { - md.push('\n'); - } else { - md.push_str(&cleaned); - } - md.push_str("\n\n"); - } - - let source_ref = normalize_source_ref(messages.first().and_then(|m| m.source_ref.clone())); - - Ok(Some(CanonicalisedSource { - markdown: md, - metadata: Metadata { - source_kind: SourceKind::Email, - source_id: source_id.to_string(), - owner: owner.to_string(), - timestamp: first_ts, - time_range: (first_ts, last_ts), - tags: tags.to_vec(), - source_ref, - path_scope: None, - }, - })) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - fn email(ts_ms: i64, from: &str, subject: &str, body: &str) -> EmailMessage { - EmailMessage { - from: from.to_string(), - to: vec!["alice@example.com".into()], - cc: vec![], - subject: subject.to_string(), - sent_at: Utc.timestamp_millis_opt(ts_ms).unwrap(), - body: body.to_string(), - source_ref: Some(format!("")), - list_unsubscribe: None, - } - } - - #[test] - fn empty_thread_returns_none() { - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "x".into(), - messages: vec![], - }; - assert!(canonicalise("gmail:t1", "alice", &[], t).unwrap().is_none()); - } - - #[test] - fn renders_headers_and_body_per_message() { - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "Launch".into(), - messages: vec![ - email(1000, "bob@example.com", "Launch", "let's ship"), - email(2000, "alice@example.com", "Re: Launch", "agreed"), - ], - }; - let out = canonicalise( - "gmail:alice@example.com|bob@example.com", - "alice@example.com", - &[], - t, - ) - .unwrap() - .unwrap(); - // No leading `# Email thread` header — that info belongs in front-matter. - assert!( - !out.markdown.contains("# Email thread — gmail — Launch"), - "canonical email MD must NOT contain a `# ` header" - ); - assert!(out.markdown.contains("From: bob@example.com")); - assert!(out.markdown.contains("Subject: Launch")); - assert!(out.markdown.contains("let's ship")); - assert!(out.markdown.contains("Re: Launch")); - assert!(out.markdown.contains("agreed")); - } - - #[test] - fn clean_body_strips_footer_before_canonicalise() { - // Body where "Unsubscribe" line triggers footer removal. Everything from - // that line onward is dropped by clean_body; real content above survives. - let body_with_footer = - "Please review the attached document.\n\nUnsubscribe https://mail.example.com/unsub\n© 2026 Example Corp"; - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "Review".into(), - messages: vec![EmailMessage { - from: "sender@example.com".into(), - to: vec!["recipient@example.com".into()], - cc: vec![], - subject: "Review".into(), - sent_at: Utc.timestamp_millis_opt(5000).unwrap(), - body: body_with_footer.into(), - source_ref: None, - list_unsubscribe: None, - }], - }; - let out = canonicalise( - "gmail:recipient@example.com|sender@example.com", - "recipient@example.com", - &[], - t, - ) - .unwrap() - .unwrap(); - assert!( - out.markdown.contains("Please review the attached document"), - "real content must survive; got:\n{}", - out.markdown - ); - assert!( - !out.markdown.to_ascii_lowercase().contains("unsubscribe"), - "unsubscribe footer must be stripped; got:\n{}", - out.markdown - ); - assert!( - !out.markdown.contains("© 2026"), - "copyright footer must be stripped; got:\n{}", - out.markdown - ); - } - - #[test] - fn time_range_spans_thread() { - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "x".into(), - messages: vec![ - email(3000, "c", "y", "third"), - email(1000, "a", "y", "first"), - email(2000, "b", "y", "second"), - ], - }; - let out = canonicalise("gmail:t1", "a", &[], t).unwrap().unwrap(); - assert_eq!(out.metadata.time_range.0.timestamp_millis(), 1000); - assert_eq!(out.metadata.time_range.1.timestamp_millis(), 3000); - } - - #[test] - fn source_ref_from_first_message() { - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "x".into(), - messages: vec![email(1000, "a", "y", "b"), email(2000, "b", "y", "c")], - }; - let out = canonicalise("gmail:t1", "a", &[], t).unwrap().unwrap(); - assert_eq!( - out.metadata.source_ref.as_ref().unwrap().value, - "" - ); - } - - #[test] - fn blank_source_ref_is_dropped() { - let mut first = email(1000, "a", "y", "b"); - first.source_ref = Some("".into()); - let t = EmailThread { - provider: "gmail".into(), - thread_subject: "x".into(), - messages: vec![first], - }; - let out = canonicalise("gmail:t1", "a", &[], t).unwrap().unwrap(); - assert!(out.metadata.source_ref.is_none()); - } - - // ── Serde regression tests (CORE-2K / #3568) ──────────────────────────── - - #[test] - fn sent_at_epoch_ms_integer_still_works() { - let json = r#"{ - "from": "alice@example.com", - "subject": "Launch", - "sent_at": 1700000000000, - "body": "content" - }"#; - let msg: EmailMessage = serde_json::from_str(json).expect("epoch-ms integer should parse"); - assert_eq!(msg.sent_at.timestamp_millis(), 1_700_000_000_000); - } - - #[test] - fn sent_at_iso8601_string_accepted() { - let json = r#"{ - "from": "alice@example.com", - "subject": "Launch", - "sent_at": "2026-05-17T19:30:00Z", - "body": "content" - }"#; - let msg: EmailMessage = serde_json::from_str(json).expect("ISO-8601 string should parse"); - assert_eq!(msg.sent_at.timestamp(), 1_779_046_200); - } - - #[test] - fn sent_at_numeric_string_accepted() { - let json = r#"{ - "from": "alice@example.com", - "subject": "Launch", - "sent_at": "1700000000000", - "body": "content" - }"#; - let msg: EmailMessage = serde_json::from_str(json).expect("numeric string should parse"); - assert_eq!(msg.sent_at.timestamp_millis(), 1_700_000_000_000); - } -} diff --git a/src/openhuman/memory_sync/canonicalize/email_clean.rs b/src/openhuman/memory_sync/canonicalize/email_clean.rs deleted file mode 100644 index 8c261eabe..000000000 --- a/src/openhuman/memory_sync/canonicalize/email_clean.rs +++ /dev/null @@ -1,423 +0,0 @@ -//! Shared email rendering + cleaning helpers. -//! -//! Used by both [`canonicalize::email`](super::email) (when rendering the -//! `GmailMarkdownStyle::Standard` shape) and the `gmail-fetch-emails` bin -//! (which writes per-sender markdown digests to disk). Lifted out of the -//! bin so the bin's behaviour and the production canonicaliser stay -//! byte-identical on body cleanup. -//! -//! The module is intentionally pure-string-oriented + a single -//! `serde_json::Value` helper (`parse_message_date`) used by callers that -//! work directly off Gmail's slim envelope JSON. Nothing here depends on -//! the memory-tree types — that keeps the helpers reusable. - -use chrono::{DateTime, NaiveDate, Utc}; -use serde_json::Value; - -/// Two-stage cleanup applied to each message body before it gets -/// blockquoted into a digest: -/// -/// 1. **Drop quoted reply chains** — once a message contains a -/// `On , wrote:` preamble, an `Original Message` / -/// `Forwarded message` separator, or a run of three+ consecutive -/// `>`-prefixed lines, everything from that point onward is the -/// parent message we already render directly above. -/// 2. **Drop footer noise** — `Unsubscribe`, `View in browser`, -/// copyright lines, legal disclaimers, and address blocks. We cut -/// at the first line containing any of [`FOOTER_TRIGGERS`]. -/// -/// The two passes run in order so a quoted-chain preamble below a -/// "view in browser" line still gets stripped on its own merits even -/// if the footer pass missed it. -pub fn clean_body(raw: &str) -> String { - let stage1 = drop_reply_chain(raw); - let stage2 = drop_footer_noise(&stage1); - collapse_blank_runs(stage2.trim()) -} - -/// Substrings that, when matched (case-insensitive) anywhere on a -/// line, mark the start of footer / boilerplate territory. Conservative -/// list — every entry should be unambiguous noise that wouldn't -/// reasonably appear inside real prose. -const FOOTER_TRIGGERS: &[&str] = &[ - "unsubscribe", - "view in browser", - "view this email in your browser", - "view it in your browser", - "update your email settings", - "manage your subscription", - "manage preferences", - "email preferences", - "you are receiving this email because", - "you received this email because", - "you're receiving this email because", - "to stop receiving", - "all rights reserved", - "© 20", - "(c) 20", - "copyright 20", - "powered by mailchimp", - "sent via sendgrid", - "this email and any files", - "confidentiality notice", - "if you are not the intended recipient", - "this communication may contain", -]; - -/// Strip quoted reply chains. See [`clean_body`] for details. -pub fn drop_reply_chain(s: &str) -> String { - let mut offset = 0usize; - let mut quoted_run_start: Option = None; - let mut quoted_run_len = 0u32; - - for line in s.split_inclusive('\n') { - let trimmed = line.trim(); - let lower = trimmed.to_ascii_lowercase(); - - // Explicit reply / forward markers. - let is_preamble = (lower.starts_with("on ") && lower.contains(" wrote:")) - || lower.contains("---------- forwarded message") - || lower.contains("----- original message") - || lower.contains("--------- original message") - || lower.contains("--- forwarded by"); - if is_preamble { - debug_assert!(s.is_char_boundary(offset)); - return s[..offset].trim_end().to_string(); - } - - // Three+ consecutive lines starting with `>` is a quoted - // reply chain in disguise (some clients de-quote on send). - // Treat the start of the run as the cut point. - if trimmed.starts_with('>') { - if quoted_run_start.is_none() { - quoted_run_start = Some(offset); - quoted_run_len = 1; - } else { - quoted_run_len += 1; - } - if quoted_run_len >= 3 { - let cut = quoted_run_start.unwrap_or(offset); - debug_assert!(s.is_char_boundary(cut)); - return s[..cut].trim_end().to_string(); - } - } else if !trimmed.is_empty() { - // Reset on a non-empty, non-quoted line. Blank lines - // don't break a quote run because senders often interleave - // them. - quoted_run_start = None; - quoted_run_len = 0; - } - - offset += line.len(); - } - s.to_string() -} - -/// Strip everything from the first line containing a footer trigger -/// onward. See [`FOOTER_TRIGGERS`] for the matched list. -pub fn drop_footer_noise(s: &str) -> String { - let mut offset = 0usize; - for line in s.split_inclusive('\n') { - let lower = line.to_ascii_lowercase(); - if FOOTER_TRIGGERS.iter().any(|t| lower.contains(t)) { - debug_assert!(s.is_char_boundary(offset)); - return s[..offset].trim_end().to_string(); - } - offset += line.len(); - } - s.to_string() -} - -/// Collapse runs of 2+ blank lines into a single blank line. Trims -/// trailing newlines. -pub fn collapse_blank_runs(s: &str) -> String { - let mut out = String::with_capacity(s.len()); - let mut blank = 0u32; - for line in s.lines() { - if line.trim().is_empty() { - blank += 1; - if blank <= 1 { - out.push('\n'); - } - } else { - blank = 0; - out.push_str(line); - out.push('\n'); - } - } - while out.ends_with('\n') { - out.pop(); - } - out -} - -/// Truncate a body to at most `max_chars` characters, appending `…` when -/// the body is longer. Trims first so leading/trailing whitespace doesn't -/// count against the budget. -pub fn truncate_body(body: &str, max_chars: usize) -> String { - let trimmed = body.trim(); - if trimmed.chars().count() <= max_chars { - return trimmed.to_string(); - } - let mut out: String = trimmed.chars().take(max_chars).collect(); - out.push('…'); - out -} - -/// Escape only the few markdown chars that would visibly break the -/// header/inline contexts we use (#, |, *, _, `). Newlines collapse to -/// spaces. We leave most punctuation alone — the body is rendered as a -/// blockquote anyway. -pub fn md_escape(s: &str) -> String { - let mut out = String::with_capacity(s.len()); - for ch in s.chars() { - match ch { - '\\' | '`' | '*' | '_' | '|' => { - out.push('\\'); - out.push(ch); - } - '\n' | '\r' => out.push(' '), - _ => out.push(ch), - } - } - out -} - -/// Pull the `` portion out of a `From` header, returning -/// just the bare email address. Falls back to `None` when no `<…>` -/// brackets exist; in that case the caller may use the raw From field. -pub fn extract_email(from: &str) -> Option { - let s = from.trim(); - if let (Some(start), Some(end)) = (s.rfind('<'), s.rfind('>')) { - if start < end { - debug_assert!(s.is_char_boundary(start + 1)); - debug_assert!(s.is_char_boundary(end)); - let inner = s[start + 1..end].trim(); - if inner.contains('@') { - return Some(inner.to_string()); - } - } - } - if s.contains('@') && !s.contains(' ') { - return Some(s.to_string()); - } - None -} - -/// If `s` starts with a 3-letter day-of-week prefix (`Mon, `, `Tue, `, …), -/// return the remainder; otherwise `None`. Used to feed a strict-rfc2822 -/// reject into a lenient retry. -fn strip_day_of_week_prefix(s: &str) -> Option<&str> { - const DAYS: &[&str] = &["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]; - let (prefix, rest) = s.split_once(", ")?; - if DAYS.iter().any(|d| d.eq_ignore_ascii_case(prefix)) { - Some(rest) - } else { - None - } -} - -/// Try a sequence of common date formats. Composio's slim envelope sets -/// `date` from `messageTimestamp` (often ISO 8601 or epoch ms) when -/// present, falling back to the raw `Date:` header (RFC 2822). Operates -/// on the raw `serde_json::Value` so callers that work off the slim -/// envelope JSON don't have to reshape it first. -pub fn parse_message_date(m: &Value) -> Option> { - if let Some(dt) = m.get("date").and_then(parse_date_value) { - return Some(dt); - } - if let Some(dt) = m.get("internalDate").and_then(parse_date_value) { - return Some(dt); - } - m.get("data") - .and_then(|data| data.get("internalDate")) - .and_then(parse_date_value) -} - -fn parse_date_value(raw: &Value) -> Option> { - if let Some(s) = raw.as_str() { - let s = s.trim(); - if s.is_empty() { - return None; - } - // Epoch millis as a string? Gmail's `internalDate` uses this form. - if let Ok(ms) = s.parse::() { - return DateTime::from_timestamp_millis(ms); - } - if let Ok(dt) = DateTime::parse_from_rfc3339(s) { - return Some(dt.with_timezone(&Utc)); - } - if let Ok(dt) = DateTime::parse_from_rfc2822(s) { - return Some(dt.with_timezone(&Utc)); - } - // Lenient RFC 2822 fallback: strict `parse_from_rfc2822` rejects - // mismatched day-of-week (e.g. `Mon, 21 Apr 2026 …` when Apr 21 - // 2026 is a Tuesday). Real-world MTAs occasionally send these - // with a wrong day-name. Strip a `, ` prefix and retry - // with the rfc2822 body format. Keeps the date if everything - // else is sane. - if let Some(rest) = strip_day_of_week_prefix(s) { - if let Ok(dt) = DateTime::parse_from_str(rest, "%d %b %Y %H:%M:%S %z") { - return Some(dt.with_timezone(&Utc)); - } - } - if let Ok(d) = NaiveDate::parse_from_str(s, "%Y-%m-%d") { - return d.and_hms_opt(0, 0, 0).map(|n| n.and_utc()); - } - } - raw.as_i64().and_then(DateTime::from_timestamp_millis) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn drop_reply_chain_strips_on_x_wrote_preamble() { - let body = "Sounds good — let's do Tuesday.\n\nOn Mon, Apr 22, 2026 at 10:00 AM, Alice wrote:\n> Tuesday or Wednesday?\n> Let me know."; - let cleaned = drop_reply_chain(body); - assert_eq!(cleaned.trim(), "Sounds good — let's do Tuesday."); - } - - #[test] - fn drop_reply_chain_strips_forwarded_separator() { - let body = "FYI.\n\n---------- Forwarded message ---------\nFrom: bob\nSubject: hi"; - assert_eq!(drop_reply_chain(body).trim(), "FYI."); - } - - #[test] - fn drop_reply_chain_strips_consecutive_quoted_run() { - let body = "Thanks for the update.\n\n> earlier line 1\n> earlier line 2\n> earlier line 3\n> earlier line 4"; - assert_eq!(drop_reply_chain(body).trim(), "Thanks for the update."); - } - - #[test] - fn drop_reply_chain_keeps_short_quote() { - // A single inline blockquote is fine — only 3+ consecutive lines trigger. - let body = "I think:\n> That sounds reasonable\n\nLet's proceed."; - let cleaned = drop_reply_chain(body); - assert!(cleaned.contains("Let's proceed")); - assert!(cleaned.contains("That sounds reasonable")); - } - - #[test] - fn drop_footer_noise_strips_unsubscribe_block() { - let body = - "Big news: GPT-5.5 is here.\n\nRead more at example.com\n\nUnsubscribe | © 2026 OpenAI"; - let cleaned = drop_footer_noise(body); - assert!(cleaned.contains("GPT-5.5")); - assert!(!cleaned.to_ascii_lowercase().contains("unsubscribe")); - assert!(!cleaned.contains("©")); - } - - #[test] - fn drop_footer_noise_strips_legal_disclaimer() { - let body = "Action item — review by Friday.\n\nThis email and any files transmitted with it are confidential and intended solely for the use of the individual to whom they are addressed."; - let cleaned = drop_footer_noise(body); - assert_eq!(cleaned.trim(), "Action item — review by Friday."); - } - - #[test] - fn clean_body_combines_passes() { - let body = - "Real content here.\n\nOn Mon, Apr 22, 2026, Alice wrote:\n> old stuff\n\nUnsubscribe"; - let cleaned = clean_body(body); - assert_eq!(cleaned, "Real content here."); - } - - #[test] - fn collapse_blank_runs_keeps_paragraph_breaks() { - let s = "a\n\n\n\nb\n\n\nc\n"; - assert_eq!(collapse_blank_runs(s), "a\n\nb\n\nc"); - } - - #[test] - fn truncate_body_adds_ellipsis() { - let s = "x".repeat(2000); - let t = truncate_body(&s, 1200); - assert!(t.ends_with('…')); - assert_eq!(t.chars().count(), 1201); - } - - #[test] - fn truncate_body_passthrough_when_short() { - let s = "hello"; - let t = truncate_body(s, 1200); - assert_eq!(t, "hello"); - } - - #[test] - fn md_escape_handles_special_chars() { - assert_eq!(md_escape("a*b_c"), "a\\*b\\_c"); - assert_eq!(md_escape("foo|bar"), "foo\\|bar"); - assert_eq!(md_escape("line1\nline2"), "line1 line2"); - assert_eq!(md_escape("plain text"), "plain text"); - } - - #[test] - fn extract_email_handles_both_forms() { - assert_eq!( - extract_email("Alice ").as_deref(), - Some("alice@example.com") - ); - assert_eq!( - extract_email("notify@github.com").as_deref(), - Some("notify@github.com") - ); - assert_eq!( - extract_email("\"Bot Name\" ").as_deref(), - Some("bot@x.io") - ); - assert!(extract_email("Alice").is_none()); - } - - #[test] - fn parse_message_date_handles_iso_and_rfc2822() { - let iso = json!({"date": "2026-04-21T10:00:00Z"}); - let rfc = json!({"date": "Mon, 21 Apr 2026 10:00:00 +0000"}); - let ms = json!({"date": 1745236800000_i64}); - let ms_str = json!({"date": "1745236800000"}); - let internal_ms_str = json!({"internalDate": "1745236800000"}); - let nested_internal_ms_str = json!({"data": {"internalDate": "1745236800000"}}); - let date_only = json!({"date": "2026-04-21"}); - assert!(parse_message_date(&iso).is_some()); - assert!(parse_message_date(&rfc).is_some()); - assert!(parse_message_date(&ms).is_some()); - assert!(parse_message_date(&ms_str).is_some()); - assert!(parse_message_date(&internal_ms_str).is_some()); - assert!(parse_message_date(&nested_internal_ms_str).is_some()); - assert!(parse_message_date(&date_only).is_some()); - } - - #[test] - fn parse_message_date_returns_none_when_missing_or_blank() { - assert!(parse_message_date(&json!({})).is_none()); - assert!(parse_message_date(&json!({"date": ""})).is_none()); - assert!(parse_message_date(&json!({"date": " "})).is_none()); - } - - #[test] - fn drop_reply_chain_handles_zwnj_in_body() { - // U+200C ZERO WIDTH NON-JOINER appears in Persian/Arabic text inside - // the real content. It must be preserved through reply-chain stripping - // when it does not appear in a reply-preamble line. - let zwnj = "\u{200c}"; - let body = format!( - "سلام{}دوست عزیز، لطفاً بررسی کنید.\n\nOn Mon, Apr 22, 2026, Alice wrote:\n> old content", - zwnj - ); - - let cleaned = drop_reply_chain(&body); - - // Reply chain must be stripped. - assert!(!cleaned.contains("old content")); - // The ZWNJ in the real body must survive. - assert!( - cleaned.contains(zwnj), - "ZWNJ was incorrectly removed from real content" - ); - // Result must be valid UTF-8 (no panic from slicing at the boundary). - assert!(std::str::from_utf8(cleaned.as_bytes()).is_ok()); - } -} diff --git a/src/openhuman/memory_sync/canonicalize/mod.rs b/src/openhuman/memory_sync/canonicalize/mod.rs index e6d636a8f..bd641f1f4 100644 --- a/src/openhuman/memory_sync/canonicalize/mod.rs +++ b/src/openhuman/memory_sync/canonicalize/mod.rs @@ -1,110 +1,3 @@ -//! Canonicalisers — normalise source-specific payloads into canonical -//! Markdown with provenance metadata (Phase 1 / #707). -//! -//! Each source kind has its own adapter. They all return the same shape: -//! a [`CanonicalisedSource`] containing the markdown blob plus a seed -//! [`Metadata`] that the chunker will clone onto each produced chunk. -//! -//! Adapters do not interpret content semantically — they only normalise -//! shape and capture provenance. Scoring / entity extraction / summarisation -//! happen downstream in later phases. +//! Compatibility facade for tinycortex canonical ingestion shapes. -pub mod chat; -pub mod document; -pub mod email; -pub mod email_clean; - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Deserializer, Serialize}; - -use crate::openhuman::memory_store::chunks::types::{Metadata, SourceRef}; - -/// Deserialise a `DateTime` from either: -/// - a JSON integer = epoch **milliseconds** (legacy callers — back-compat), -/// - a JSON string = RFC 3339 / ISO-8601 (e.g. `"2026-05-17T19:30:00Z"`), or -/// a decimal string containing epoch milliseconds. -/// -/// On an unparseable string a serde error is returned (no silent default). -/// Shared across chat, email, and document canonicalisers. -pub(crate) fn deserialize_flexible_timestamp<'de, D>( - deserializer: D, -) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - #[derive(Deserialize)] - #[serde(untagged)] - enum RawTs { - Millis(i64), - Text(String), - } - - let raw = RawTs::deserialize(deserializer)?; - match raw { - RawTs::Millis(ms) => { - tracing::debug!("[memory][canonicalize] parsed timestamp as epoch-ms: {ms}"); - chrono::TimeZone::timestamp_millis_opt(&Utc, ms) - .single() - .ok_or_else(|| serde::de::Error::custom(format!("invalid epoch-ms: {ms}"))) - } - RawTs::Text(s) => { - if let Ok(dt) = DateTime::parse_from_rfc3339(&s) { - tracing::debug!("[memory][canonicalize] parsed timestamp as ISO-8601 string: {s}"); - return Ok(dt.with_timezone(&Utc)); - } - if let Ok(ms) = s.parse::() { - tracing::debug!( - "[memory][canonicalize] parsed timestamp as numeric-string epoch-ms: {s}" - ); - return chrono::TimeZone::timestamp_millis_opt(&Utc, ms) - .single() - .ok_or_else(|| { - serde::de::Error::custom(format!("invalid epoch-ms string: {s}")) - }); - } - Err(serde::de::Error::custom(format!( - "cannot parse '{s}' as RFC 3339 or epoch-ms" - ))) - } - } -} - -/// Output of a canonicaliser — one per logical source record -/// (a chat batch, an email, a document). -#[derive(Clone, Debug)] -pub struct CanonicalisedSource { - /// Canonical Markdown blob produced by the adapter. - pub markdown: String, - /// Provenance the chunker will clone onto each emitted [`Chunk`]. - pub metadata: Metadata, -} - -/// Shared input shape: a payload + a minimal provenance hint. -/// -/// Every adapter accepts this generic envelope; the concrete payload type -/// is adapter-specific (see sibling modules for the per-kind inputs). -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct CanonicaliseRequest

{ - /// Logical source id (channel for chat, thread for email, doc id). - pub source_id: String, - /// Owner / user account. - #[serde(default)] - pub owner: String, - /// Source-specific payload. - pub payload: P, - /// Optional tags carried through. - #[serde(default)] - pub tags: Vec, -} - -/// Trim provider-specific source references and drop blank pointers. -pub fn normalize_source_ref(source_ref: Option) -> Option { - source_ref.and_then(|value| { - let trimmed = value.trim(); - if trimmed.is_empty() { - None - } else { - Some(SourceRef::new(trimmed.to_string())) - } - }) -} +pub use tinycortex::memory::ingest::canonicalize::*; diff --git a/src/openhuman/memory_sync/composio/bus.rs b/src/openhuman/memory_sync/composio/bus.rs index 1b12b525a..bd63db8bc 100644 --- a/src/openhuman/memory_sync/composio/bus.rs +++ b/src/openhuman/memory_sync/composio/bus.rs @@ -667,11 +667,34 @@ impl EventHandler for ComposioConnectionCreatedSubscriber { error = %e, "[composio:bus] provider on_connection_created failed" ); - } else { - // Successful connection-created sync — record the - // timestamp so the periodic scheduler doesn't - // immediately re-fire for this connection. - super::periodic::record_sync_success(&toolkit, &connection_id); + } + + match crate::openhuman::tinycortex::run_composio_connection( + &toolkit, + &connection_id, + ctx.config.as_ref(), + ) + .await + { + Ok(outcome) => { + tracing::info!( + toolkit = %toolkit, + connection_id = %connection_id, + items_ingested = outcome.records_ingested, + actions_called = outcome.actions_called, + "[composio:bus] tinycortex initial sync complete" + ); + // Avoid immediately re-firing from the periodic scheduler. + super::periodic::record_sync_success(&toolkit, &connection_id); + } + Err(error) => tracing::warn!( + toolkit = %toolkit, + connection_id = %connection_id, + error = %error, + actions_called = error.actions_called, + provider_cost_usd = error.provider_cost_usd, + "[composio:bus] tinycortex initial sync failed" + ), } } diff --git a/src/openhuman/memory_sync/composio/mod.rs b/src/openhuman/memory_sync/composio/mod.rs index 6e20dcad8..9b918b52a 100644 --- a/src/openhuman/memory_sync/composio/mod.rs +++ b/src/openhuman/memory_sync/composio/mod.rs @@ -173,29 +173,47 @@ pub async fn run_connection_sync( "[composio:sync] run_connection_sync: caps from registry" ); - let ctx = ProviderContext { - config: std::sync::Arc::new(config), - toolkit: target.toolkit, - connection_id: Some(target.connection_id), - usage: Default::default(), - max_items: src_max_items, - sync_depth_days: src_sync_depth_days, - }; - - let sync_result = provider.sync(&ctx, reason).await; - - // Read the Composio billable-action tally *before* propagating errors. - // A sync that errors partway may still have fired billable actions; - // reading here ensures the dispatcher audit sees partial cost (#3111). - let usage = ctx - .usage - .lock() - .map(|u| u.clone()) - .unwrap_or_else(|poisoned| poisoned.into_inner().clone()); - - match sync_result { - Ok(outcome) => Ok((outcome, usage)), - Err(e) => Err((e, usage)), + let _ = (provider, src_max_items, src_sync_depth_days); + let started_at_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + match crate::openhuman::tinycortex::run_composio_connection( + &target.toolkit, + &target.connection_id, + &config, + ) + .await + { + Ok(outcome) => { + let usage = ComposioUsage { + actions_called: outcome.actions_called, + cost_usd: outcome.provider_cost_usd, + }; + Ok(( + SyncOutcome { + toolkit: target.toolkit, + connection_id: Some(target.connection_id), + reason: reason.as_str().to_string(), + items_ingested: outcome.records_ingested as usize, + started_at_ms, + finished_at_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + summary: outcome.note.unwrap_or_else(|| "sync completed".to_string()), + details: serde_json::json!({ "more_pending": outcome.more_pending }), + }, + usage, + )) + } + Err(error) => Err(( + error.to_string(), + ComposioUsage { + actions_called: error.actions_called, + cost_usd: error.provider_cost_usd, + }, + )), } } diff --git a/src/openhuman/memory_sync/composio/periodic.rs b/src/openhuman/memory_sync/composio/periodic.rs index 86b5bc09e..c06a361d2 100644 --- a/src/openhuman/memory_sync/composio/periodic.rs +++ b/src/openhuman/memory_sync/composio/periodic.rs @@ -2,7 +2,7 @@ //! //! Spawned once at startup. The scheduler walks every active Composio //! connection on a fixed tick, looks up the matching native provider, -//! and calls `provider.sync(ctx, SyncReason::Periodic)` if enough time +//! and dispatches the matching tinycortex pipeline if enough time //! has elapsed since that connection's last sync (per the provider's //! `sync_interval_secs`). //! @@ -58,7 +58,7 @@ use crate::openhuman::memory_sources::{ use crate::openhuman::scheduler_gate::gate::{current_policy, resume_notify}; use crate::openhuman::scheduler_gate::policy::PauseReason; -use super::providers::{get_provider, ComposioUsage, ProviderContext, SyncReason}; +use super::providers::{get_provider, ComposioUsage}; use crate::openhuman::composio::client::{ create_composio_client, direct_list_connections, ComposioClientKind, }; @@ -550,19 +550,12 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { "[composio:periodic] caps from registry" ); - // Build a context tied to this specific connection and dispatch. - // `ProviderContext` no longer caches a pre-baked - // `ComposioClient` — provider methods resolve a fresh handle per - // call via `ctx.execute(...)` so a mid-session - // `composio.mode` toggle is honoured immediately (#1710). - let ctx = ProviderContext { - config: Arc::clone(&config), - toolkit: toolkit.clone(), - connection_id: Some(conn.id.clone()), - usage: Default::default(), - max_items: src_max_items, - sync_depth_days: src_sync_depth_days, - }; + let mut source = composio_sources + .get(&conn.id) + .cloned() + .unwrap_or_else(|| periodic_source(&toolkit, &conn.id)); + source.max_items = src_max_items; + source.sync_depth_days = src_sync_depth_days; tracing::debug!( toolkit = %conn.toolkit, @@ -571,25 +564,19 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { "[composio:periodic] firing sync" ); let sync_started = Instant::now(); - let result = provider.sync(&ctx, SyncReason::Periodic).await; + let result = crate::openhuman::tinycortex::run_source_pipeline(&source, &config).await; let duration_ms = sync_started.elapsed().as_millis() as u64; - // Read the Composio billable-action tally the sync accumulated at the - // `execute` chokepoint (#3111). Periodic is where most Composio cost - // accrues (this loop fires every 20 min per connection vs. rare manual - // syncs), but periodic ticks weren't recorded in the sync audit at all - // — so the audit under-counted real cost. Record each tick that ran a - // fetch, success or failure, so the Sync History panel reflects the - // background spend too (#3111 follow-up; raised in the #3138 review). - let usage = ctx.usage.lock().map(|u| u.clone()).unwrap_or_default(); - match result { Ok(outcome) => { + let usage = ComposioUsage { + actions_called: outcome.actions_called, + cost_usd: outcome.provider_cost_usd, + }; tracing::debug!( toolkit = %conn.toolkit, connection_id = %conn.id, - items = outcome.items_ingested, - elapsed_ms = outcome.elapsed_ms(), + items = outcome.records_ingested, composio_actions = usage.actions_called, "[composio:periodic] sync ok" ); @@ -597,7 +584,7 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { &toolkit, &conn.id, &usage, - outcome.items_ingested, + outcome.records_ingested as usize, duration_ms, None, ); @@ -606,6 +593,10 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { fired += 1; } Err(e) => { + let usage = ComposioUsage { + actions_called: e.actions_called, + cost_usd: e.provider_cost_usd, + }; tracing::warn!( toolkit = %conn.toolkit, connection_id = %conn.id, @@ -614,8 +605,14 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { ); // A failed tick may still have fired billable fetch actions // before erroring — audit the partial cost so it isn't lost. - let entry = - build_periodic_audit_entry(&toolkit, &conn.id, &usage, 0, duration_ms, Some(e)); + let entry = build_periodic_audit_entry( + &toolkit, + &conn.id, + &usage, + 0, + duration_ms, + Some(e.to_string()), + ); append_audit_entry(&config, &entry); // Intentionally do NOT update last_sync_at on failure // so the next tick retries immediately. @@ -627,6 +624,32 @@ pub(crate) async fn run_one_tick() -> Result<(), String> { Ok(()) } +fn periodic_source(toolkit: &str, connection_id: &str) -> MemorySourceEntry { + MemorySourceEntry { + id: format!("composio:{connection_id}"), + kind: SourceKind::Composio, + label: toolkit.to_string(), + enabled: true, + toolkit: Some(toolkit.to_string()), + connection_id: Some(connection_id.to_string()), + path: None, + glob: None, + 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, + } +} + /// Build a [`SyncAuditEntry`] for one periodic Composio sync tick (#3111 /// follow-up). /// diff --git a/src/openhuman/memory_sync/composio/providers/clickup/ingest.rs b/src/openhuman/memory_sync/composio/providers/clickup/ingest.rs deleted file mode 100644 index db611d858..000000000 --- a/src/openhuman/memory_sync/composio/providers/clickup/ingest.rs +++ /dev/null @@ -1,333 +0,0 @@ -//! ClickUp -> memory tree ingest plumbing. -//! -//! Converts one ClickUp task payload into a memory_tree [`DocumentInput`] -//! and calls `ingest_document` so retrieval surfaces read the content from -//! `mem_tree_chunks` instead of the legacy `memory_docs` path (issue #2885). -//! -//! ClickUp differs from the Linear/Notion providers in one way: its -//! `date_updated` field is a Unix epoch **milliseconds** value rendered as a -//! string (e.g. `"1779962400000"`), not an RFC3339 timestamp — so -//! [`parse_updated_time`] parses milliseconds rather than calling -//! `DateTime::parse_from_rfc3339`. - -use anyhow::Result; -use chrono::{DateTime, TimeZone, Utc}; -use serde_json::Value; - -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::{self, IngestResult}; -use crate::openhuman::memory_store::chunks::store::{delete_chunks_by_source, is_source_ingested}; -use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_sync::canonicalize::document::DocumentInput; - -/// Platform identifier embedded in ClickUp document metadata. -pub const CLICKUP_PLATFORM: &str = "clickup"; - -/// Stable tags attached to every ClickUp-ingested task chunk. -pub const DEFAULT_TAGS: &[&str] = &["clickup", "ingested"]; - -/// Build the memory-tree source id for one ClickUp task in one connection. -pub(crate) fn clickup_source_id(connection_id: &str, task_id: &str) -> String { - format!("clickup:{connection_id}:{task_id}") -} - -/// Build the source tree / raw archive scope for one ClickUp connection. -pub(crate) fn clickup_source_scope(connection_id: &str) -> String { - format!("clickup:{connection_id}") -} - -/// Render the raw ClickUp task payload as a markdown document body. -fn render_task_body(title: &str, task: &Value) -> String { - let pretty = serde_json::to_string_pretty(task).unwrap_or_else(|_| "{}".to_string()); - format!("# {title}\n\n```json\n{pretty}\n```\n") -} - -/// Parse ClickUp's `date_updated` (Unix epoch **milliseconds** as a string), -/// falling back to now on missing/malformed input. -fn parse_updated_time(raw: Option<&str>) -> DateTime { - raw.and_then(|s| s.trim().parse::().ok()) - .and_then(|ms| Utc.timestamp_millis_opt(ms).single()) - .unwrap_or_else(Utc::now) -} - -/// Ingest one ClickUp task into memory_tree and return the written chunk count. -/// -/// Edited tasks reuse the same `source_id`, so prior chunks are deleted before -/// re-ingest to avoid the document pipeline's duplicate-source short-circuit. -pub async fn ingest_task_into_memory_tree( - config: &Config, - connection_id: &str, - task_id: &str, - title: &str, - updated_time: Option<&str>, - task: &Value, -) -> Result { - let source_id = clickup_source_id(connection_id, task_id); - - let cfg_for_blocking = config.clone(); - let source_for_blocking = source_id.clone(); - let removed = tokio::task::spawn_blocking(move || -> Result { - if is_source_ingested( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - )? { - delete_chunks_by_source( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - ) - } else { - Ok(0) - } - }) - .await - .map_err(|e| anyhow::anyhow!("delete-prior task join error: {e}"))??; - - if removed > 0 { - tracing::debug!( - connection_id = %connection_id, - task_id = %task_id, - removed_chunks = removed, - "[composio:clickup] ingest: re-ingest cleanup" - ); - } - - let modified_at = parse_updated_time(updated_time); - let body = render_task_body(title, task); - let source_ref = Some(format!("clickup://task/{task_id}")); - let doc = DocumentInput { - provider: CLICKUP_PLATFORM.to_string(), - title: title.to_string(), - body, - modified_at, - source_ref, - }; - let tags: Vec = DEFAULT_TAGS.iter().map(|s| s.to_string()).collect(); - let owner = clickup_source_scope(connection_id); - let path_scope = Some(owner.clone()); - - match ingest_pipeline::ingest_document_with_scope( - config, &source_id, &owner, tags, doc, path_scope, - ) - .await - { - Ok(IngestResult { - chunks_written, - already_ingested, - .. - }) => { - tracing::debug!( - connection_id = %connection_id, - task_id = %task_id, - chunks_written, - already_ingested, - "[composio:clickup] ingest: task persisted" - ); - Ok(chunks_written) - } - Err(err) => Err(anyhow::anyhow!( - "ingest_document failed for {source_id}: {err:#}" - )), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::config::Config; - use crate::openhuman::memory_store::chunks::types::SourceKind; - use chrono::{TimeZone, Utc}; - use serde_json::{json, Value}; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn sample_task(task_id: &str, date_updated: &str) -> Value { - json!({ - "id": task_id, - "name": "Fix external LLM routing", - "date_updated": date_updated, - "url": "https://app.clickup.com/t/abc123", - "status": { "status": "in progress" }, - "text_content": "Connected app tools vanish under external routing.", - "assignees": [{ "username": "Alice" }] - }) - } - - #[test] - fn clickup_source_id_is_stable_and_namespaced() { - let a = clickup_source_id("conn-1", "task-abc"); - let b = clickup_source_id("conn-1", "task-abc"); - assert_eq!(a, b); - assert_eq!(a, "clickup:conn-1:task-abc"); - assert_ne!(a, clickup_source_id("conn-2", "task-abc")); - assert_ne!(a, clickup_source_id("conn-1", "task-xyz")); - } - - #[test] - fn clickup_source_scope_is_connection_level() { - assert_eq!(clickup_source_scope("conn-1"), "clickup:conn-1"); - assert_eq!(clickup_source_scope("conn-2"), "clickup:conn-2"); - assert_eq!( - clickup_source_scope("conn-1"), - clickup_source_scope("conn-1"), - "scope must stay stable across tasks in one connection" - ); - } - - #[test] - fn parse_updated_time_handles_epoch_millis_and_invalid_inputs() { - // ClickUp sends epoch milliseconds as a string. Build the expected - // instant and round-trip through its millisecond representation so we - // never hand-compute the magic number. - let expected = Utc.with_ymd_and_hms(2026, 5, 28, 10, 0, 0).unwrap(); - let ms = expected.timestamp_millis().to_string(); - let good = parse_updated_time(Some(&ms)); - assert_eq!(good, expected); - - // Leading/trailing whitespace must still parse. - let padded = parse_updated_time(Some(&format!(" {ms} "))); - assert_eq!(padded, expected); - - let bad = parse_updated_time(Some("not-a-timestamp")); - assert!((Utc::now() - bad).num_seconds().abs() < 5); - - // An RFC3339 string is *not* valid ClickUp input — it must fall back. - let rfc = parse_updated_time(Some("2026-05-28T10:00:00.000Z")); - assert!((Utc::now() - rfc).num_seconds().abs() < 5); - - let missing = parse_updated_time(None); - assert!((Utc::now() - missing).num_seconds().abs() < 5); - } - - #[test] - fn render_task_body_includes_title_header_and_pretty_json() { - let task = json!({ - "id": "task-1", - "name": "Fix external LLM routing", - "date_updated": "1779962400000" - }); - let body = render_task_body("ClickUp: Fix external LLM routing", &task); - assert!(body.starts_with("# ClickUp: Fix external LLM routing\n")); - assert!(body.contains("```json\n")); - assert!(body.contains("\"name\": \"Fix external LLM routing\"")); - assert!(body.contains("\"date_updated\": \"1779962400000\"")); - } - - #[tokio::test] - async fn ingest_task_writes_to_memory_tree() { - use crate::openhuman::memory_store::chunks::store::{ - count_chunks, is_source_ingested, list_chunks, ListChunksQuery, - }; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-clickup"; - let task_id = "task-routing"; - let expected = clickup_source_id(connection_id, task_id); - let expected_scope = clickup_source_scope(connection_id); - let task = sample_task(task_id, "1779962400000"); - let chunks_before = count_chunks(&cfg).expect("count_chunks before"); - - let written = ingest_task_into_memory_tree( - &cfg, - connection_id, - task_id, - "ClickUp: Fix external LLM routing", - Some("1779962400000"), - &task, - ) - .await - .expect("ingest_task_into_memory_tree"); - - assert!(written > 0, "ClickUp ingest must write chunks"); - let chunks_after = count_chunks(&cfg).expect("count_chunks after"); - assert!( - chunks_after > chunks_before, - "ingest must populate mem_tree_chunks (#2885)" - ); - - let rows = list_chunks( - &cfg, - &ListChunksQuery { - source_kind: Some(SourceKind::Document), - source_id: Some(expected.clone()), - limit: Some(1), - ..Default::default() - }, - ) - .expect("list chunks for clickup task"); - let chunk = rows.first().expect("clickup chunk should be listed"); - assert_eq!(chunk.metadata.source_id, expected.as_str()); - assert_eq!( - chunk.metadata.path_scope.as_deref(), - Some(expected_scope.as_str()) - ); - - let cfg_for_blocking = cfg.clone(); - let expected_for_task = expected.clone(); - let registered = tokio::task::spawn_blocking(move || { - is_source_ingested(&cfg_for_blocking, SourceKind::Document, &expected_for_task) - .unwrap_or(false) - }) - .await - .expect("source-check task join"); - assert!(registered, "source_id must be registered"); - } - - #[tokio::test] - async fn re_ingesting_edited_task_replaces_prior_chunks() { - use crate::openhuman::memory_store::chunks::store::count_chunks; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-edit"; - let task_id = "task-edit"; - - let v1 = sample_task(task_id, "1779962400000"); - let first = ingest_task_into_memory_tree( - &cfg, - connection_id, - task_id, - "ClickUp: Fix external LLM routing", - Some("1779962400000"), - &v1, - ) - .await - .expect("first ingest"); - assert!(first > 0); - let after_first = count_chunks(&cfg).expect("count after first"); - - let v2 = json!({ - "id": task_id, - "name": "Fix external LLM routing", - "date_updated": "1780048800000", - "text_content": "Updated: external LLM routing now keeps connected app tools visible.", - "status": { "status": "complete" } - }); - let second = ingest_task_into_memory_tree( - &cfg, - connection_id, - task_id, - "ClickUp: Fix external LLM routing", - Some("1780048800000"), - &v2, - ) - .await - .expect("second ingest"); - assert!(second > 0); - let after_second = count_chunks(&cfg).expect("count after second"); - - assert!( - after_second.abs_diff(after_first) <= 1, - "edited task must replace prior chunks, not append duplicates" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/clickup/mod.rs b/src/openhuman/memory_sync/composio/providers/clickup/mod.rs index 949298230..bbfd3a69f 100644 --- a/src/openhuman/memory_sync/composio/providers/clickup/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/clickup/mod.rs @@ -13,9 +13,7 @@ //! //! Issue: #2288 (introduction); #2885 (memory_tree migration). -mod ingest; mod provider; -mod source; mod sync; #[cfg(test)] mod tests; diff --git a/src/openhuman/memory_sync/composio/providers/clickup/provider.rs b/src/openhuman/memory_sync/composio/providers/clickup/provider.rs index 3dfb9d0db..b31536aa2 100644 --- a/src/openhuman/memory_sync/composio/providers/clickup/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/clickup/provider.rs @@ -27,12 +27,10 @@ use async_trait::async_trait; use serde_json::json; -use super::source::run_clickup_sync; use super::sync; use crate::openhuman::memory_sync::composio::providers::{ first_array_str, merge_extra, pick_str, resolve_sync_interval_secs, ComposioProvider, - CuratedTool, NormalizedTask, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, - TaskFetchFilter, TaskKind, + CuratedTool, NormalizedTask, ProviderContext, ProviderUserProfile, TaskFetchFilter, TaskKind, }; pub(crate) const ACTION_GET_AUTHORIZED_USER: &str = "CLICKUP_GET_AUTHORIZED_USER"; @@ -136,10 +134,6 @@ impl ComposioProvider for ClickUpProvider { /// `max_items` cap, the epoch-ms `sync_depth_days` window, and cursor /// handling live in `run_sync`; the ClickUp-specific primitives live in /// [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - run_clickup_sync(ctx, reason).await - } - async fn fetch_tasks( &self, ctx: &ProviderContext, @@ -251,11 +245,7 @@ impl ComposioProvider for ClickUpProvider { /// Map a raw ClickUp task payload into a [`NormalizedTask`]. Returns /// `None` only when the task has no extractable id (unroutable). fn normalize_clickup_task(task: &serde_json::Value) -> Option { - let external_id = - crate::openhuman::memory_sync::composio::providers::sync_state::extract_item_id( - task, - TASK_ID_PATHS, - )?; + let external_id = pick_str(task, TASK_ID_PATHS)?; let title = sync::extract_task_name(task).unwrap_or_else(|| format!("ClickUp task {external_id}")); Some(NormalizedTask { diff --git a/src/openhuman/memory_sync/composio/providers/clickup/source.rs b/src/openhuman/memory_sync/composio/providers/clickup/source.rs deleted file mode 100644 index 34f5d2c6b..000000000 --- a/src/openhuman/memory_sync/composio/providers/clickup/source.rs +++ /dev/null @@ -1,569 +0,0 @@ -//! ClickUp's [`IncrementalSource`] primitives. -//! -//! ClickUp rides the generic -//! [`crate::openhuman::memory_sync::composio::providers::orchestrator`]: -//! [`ClickUpProvider::sync`](super::provider::ClickUpProvider) delegates to -//! [`run_clickup_sync`]. The orchestrator owns the control flow (budget, -//! pagination bound, dedup, the `sync_depth_days` window, the precise -//! `max_items` clamp, cursor advance/hold, state persistence); this module -//! supplies only the ClickUp-specific shapes. -//! -//! ClickUp is the first **nested** provider on the orchestrator: -//! [`ClickUpSource::preamble`] resolves the authorized user id + the visible -//! workspaces and returns **one [`SyncScope`] per workspace**, then the -//! orchestrator pages each workspace's `CLICKUP_GET_FILTERED_TEAM_TASKS` -//! (1-indexed `page`) scoped to that user as assignee. The shared user id is -//! resolved once in the preamble and stashed on the source ([`OnceLock`]) so -//! every per-workspace `fetch_page` can read it back. -//! -//! Two provider quirks the trait already accommodates: -//! * **epoch-ms timestamps** — `date_updated` is a millisecond-epoch string, -//! so [`ClickUpSource::depth_floor`] emits an epoch-ms floor (not RFC3339) -//! for the lexicographic compare to be valid. -//! * **advance-on-failure** — ClickUp advances its cursor even when a per-item -//! ingest fails, so it overrides -//! [`IncrementalSource::hold_cursor_on_ingest_failure`] to `false`. - -use std::sync::OnceLock; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{json, Value}; - -use super::provider::{ - ACTION_GET_AUTHORIZED_TEAMS_WORKSPACES, ACTION_GET_AUTHORIZED_USER, - ACTION_GET_FILTERED_TEAM_TASKS, TASK_ID_PATHS, -}; -use super::{ingest::ingest_task_into_memory_tree, sync}; -use crate::openhuman::config::Config; -use crate::openhuman::memory_sync::composio::providers::orchestrator::{ - self, IncrementalSource, IngestOutcome, PageFetch, SyncItem, SyncScope, -}; -use crate::openhuman::memory_sync::composio::providers::sync_state::{extract_item_id, SyncState}; -use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, SyncOutcome, SyncReason, -}; - -/// Page size per API call on steady-state syncs. -const PAGE_SIZE: u32 = 50; - -/// Larger initial-sync page size so the first backfill catches up faster. -const INITIAL_PAGE_SIZE: u32 = 100; - -/// Maximum pages **per workspace** per sync pass. -const MAX_PAGES_PER_WORKSPACE: u32 = 20; - -/// Max in-flight ingests per page. DB writes serialize anyway and the cloud -/// embedder has rate limits, so keep this small. -const INGEST_CONCURRENCY: usize = 8; - -/// ClickUp's [`IncrementalSource`]. Holds the authorized user id resolved in -/// the preamble so every per-workspace `fetch_page` can scope to it. -pub(crate) struct ClickUpSource { - user_id: OnceLock, -} - -impl ClickUpSource { - fn new() -> Self { - Self { - user_id: OnceLock::new(), - } - } -} - -/// Entry point used by [`super::provider::ClickUpProvider::sync`]. -pub(crate) async fn run_clickup_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - orchestrator::run_sync(&ClickUpSource::new(), ctx, reason).await -} - -impl ClickUpSource { - /// Look up (and budget-record) the authorized user's numeric id. - async fn resolve_user_id( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result { - tracing::debug!( - connection_id = ?ctx.connection_id, - "[composio:clickup] resolve_user_id via {ACTION_GET_AUTHORIZED_USER}" - ); - - let resp = ctx - .execute(ACTION_GET_AUTHORIZED_USER, Some(json!({}))) - .await - .map_err(|e| { - format!("[composio:clickup] {ACTION_GET_AUTHORIZED_USER} failed: {e:#}") - })?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!( - "[composio:clickup] {ACTION_GET_AUTHORIZED_USER}: {err}" - )); - } - - let user_id = sync::extract_user_id(&resp.data).ok_or_else(|| { - "[composio:clickup] CLICKUP_GET_AUTHORIZED_USER returned no user.id".to_string() - })?; - - tracing::debug!( - connection_id = ?ctx.connection_id, - "[composio:clickup] resolve_user_id complete" - ); - Ok(user_id) - } - - /// Look up (and budget-record) the workspace (team) ids visible to this - /// connection. - async fn resolve_workspaces( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - tracing::debug!( - connection_id = ?ctx.connection_id, - "[composio:clickup] resolve_workspaces via {ACTION_GET_AUTHORIZED_TEAMS_WORKSPACES}" - ); - - let resp = ctx - .execute(ACTION_GET_AUTHORIZED_TEAMS_WORKSPACES, Some(json!({}))) - .await - .map_err(|e| { - format!("[composio:clickup] {ACTION_GET_AUTHORIZED_TEAMS_WORKSPACES} failed: {e:#}") - })?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!( - "[composio:clickup] {ACTION_GET_AUTHORIZED_TEAMS_WORKSPACES}: {err}" - )); - } - - let workspaces = sync::extract_workspace_ids(&resp.data); - tracing::debug!( - connection_id = ?ctx.connection_id, - count = workspaces.len(), - "[composio:clickup] resolve_workspaces complete" - ); - Ok(workspaces) - } -} - -#[async_trait] -impl IncrementalSource for ClickUpSource { - fn toolkit(&self) -> &'static str { - "clickup" - } - - fn page_size(&self, reason: SyncReason) -> u32 { - match reason { - SyncReason::ConnectionCreated => INITIAL_PAGE_SIZE, - _ => PAGE_SIZE, - } - } - - fn max_pages(&self) -> u32 { - MAX_PAGES_PER_WORKSPACE - } - - fn detail_noun(&self) -> &'static str { - "tasks" - } - - /// ClickUp advances its cursor even on per-item ingest failures. - fn hold_cursor_on_ingest_failure(&self) -> bool { - false - } - - /// `date_updated` is a millisecond-epoch string, so the depth floor must be - /// epoch-ms too (not RFC3339) for the lexicographic compare to hold. - fn depth_floor(&self, days: u32) -> String { - let floor = chrono::Utc::now() - chrono::Duration::days(days as i64); - floor.timestamp_millis().to_string() - } - - /// Resolve the user id (stashed for `fetch_page`) and return one scope per - /// visible workspace. Mirrors the original's budget re-check: if the daily - /// budget is exhausted right after the user-id probe, we skip the - /// workspaces call (empty scopes → no-op outcome). - async fn preamble( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - let user_id = self.resolve_user_id(ctx, state).await?; - let _ = self.user_id.set(user_id); - - if state.budget_exhausted() { - tracing::info!( - "[composio:clickup] budget exhausted after user-id probe, skipping sync" - ); - return Ok(vec![]); - } - - let workspaces = self.resolve_workspaces(ctx, state).await?; - Ok(workspaces - .into_iter() - .map(|ws| { - let label = format!("workspace:{ws}"); - SyncScope::nested(ws, label) - }) - .collect()) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result { - let workspace_id = &scope.id; - let user_id = self.user_id.get().cloned().unwrap_or_default(); - // ClickUp paginates by 0-indexed `page`; the orchestrator's opaque - // cursor carries the next page number (`None` = first page). - let page_num: u32 = cursor.and_then(|c| c.parse().ok()).unwrap_or(0); - let page_size = self.page_size(reason); - - let args = json!({ - "team_id": workspace_id, - "assignees": [user_id], - "order_by": "updated", - "reverse": true, - "page": page_num, - "page_size": page_size, - "subtasks": true, - }); - - tracing::debug!( - connection_id = ?ctx.connection_id, - workspace_id = %workspace_id, - page_num, - page_size, - "[composio:clickup] fetch_page via {ACTION_GET_FILTERED_TEAM_TASKS}" - ); - - let resp = ctx - .execute(ACTION_GET_FILTERED_TEAM_TASKS, Some(args)) - .await - .map_err(|e| { - format!( - "[composio:clickup] {ACTION_GET_FILTERED_TEAM_TASKS} \ - workspace={workspace_id} page={page_num}: {e:#}" - ) - })?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!( - "[composio:clickup] {ACTION_GET_FILTERED_TEAM_TASKS} \ - workspace={workspace_id} page={page_num}: {err}" - )); - } - - let tasks = sync::extract_tasks(&resp.data); - // ClickUp signals the last page implicitly: a short page ends the - // workspace. - let next = if (tasks.len() as u32) < page_size { - None - } else { - Some((page_num + 1).to_string()) - }; - - tracing::debug!( - connection_id = ?ctx.connection_id, - workspace_id = %workspace_id, - page_num, - fetched = tasks.len(), - has_next = next.is_some(), - "[composio:clickup] fetch_page complete" - ); - - Ok(PageFetch { items: tasks, next }) - } - - fn item_dedup_key(&self, item: &Value) -> Option { - let task_id = extract_item_id(item, TASK_ID_PATHS)?; - match sync::extract_task_updated(item) { - Some(updated) => Some(format!("{task_id}@{updated}")), - None => Some(task_id), - } - } - - fn item_sort_ts(&self, item: &Value) -> Option { - sync::extract_task_updated(item) - } - - async fn ingest( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - - let pending: Vec = items - .into_iter() - .filter_map(|it| { - let task_id = extract_item_id(&it.raw, TASK_ID_PATHS)?; - let title_text = sync::extract_task_name(&it.raw) - .unwrap_or_else(|| format!("ClickUp task {task_id}")); - Some(PendingIngest { - sync_key: it.dedup_key, - task_id, - title: format!("ClickUp: {title_text}"), - updated: it.sort_ts, - task: it.raw, - }) - }) - .collect(); - - let ingestor = MemoryTreeIngestor { - config: ctx.config.as_ref(), - connection_id, - }; - let buffered = - ingest_pending_buffered(&ingestor, pending, &scope.id, INGEST_CONCURRENCY).await; - IngestOutcome { - synced_keys: buffered.synced_keys, - persisted: buffered.persisted, - had_failures: buffered.had_failures, - } - } -} - -/// One task queued for concurrent ingest. Owns its raw task `Value` (the -/// orchestrator handed ownership via [`SyncItem`]). -struct PendingIngest { - sync_key: String, - task_id: String, - title: String, - updated: Option, - task: Value, -} - -/// Folded result of [`ingest_pending_buffered`]. Order-independent. -#[derive(Default)] -struct BufferedIngestOutcome { - synced_keys: Vec, - persisted: usize, - had_failures: bool, -} - -/// Seam over "ingest one ClickUp task", so the bounded-concurrency driver can -/// be unit-tested with a fake that records peak in-flight calls. -#[async_trait] -trait TaskIngestor { - async fn ingest( - &self, - task_id: &str, - title: &str, - updated: Option<&str>, - task: &Value, - ) -> anyhow::Result; -} - -/// Production ingestor: routes into the memory-tree pipeline. -struct MemoryTreeIngestor<'c> { - config: &'c Config, - connection_id: &'c str, -} - -#[async_trait] -impl TaskIngestor for MemoryTreeIngestor<'_> { - async fn ingest( - &self, - task_id: &str, - title: &str, - updated: Option<&str>, - task: &Value, - ) -> anyhow::Result { - ingest_task_into_memory_tree( - self.config, - self.connection_id, - task_id, - title, - updated, - task, - ) - .await - } -} - -/// Ingest queued tasks with bounded concurrency, folding into an -/// order-independent [`BufferedIngestOutcome`]. A failed ingest is logged and -/// skipped (ClickUp advances its cursor regardless, via -/// `hold_cursor_on_ingest_failure = false`). -async fn ingest_pending_buffered( - ingestor: &I, - pending: Vec, - workspace_id: &str, - concurrency: usize, -) -> BufferedIngestOutcome { - let ingest_futs = pending - .into_iter() - .map(|p| async move { - let res = ingestor - .ingest(&p.task_id, &p.title, p.updated.as_deref(), &p.task) - .await; - (p.sync_key, p.task_id, res) - }) - .collect::>(); - - let mut outcome = BufferedIngestOutcome::default(); - let mut ingest_stream = futures::stream::iter(ingest_futs).buffer_unordered(concurrency); - while let Some((sync_key, task_id, res)) = ingest_stream.next().await { - match res { - Ok(_chunks_written) => { - outcome.synced_keys.push(sync_key); - outcome.persisted += 1; - } - Err(e) => { - outcome.had_failures = true; - tracing::warn!( - task_id = %task_id, - workspace_id = %workspace_id, - error = %e, - "[composio:clickup] ingest failed (continuing)" - ); - } - } - } - outcome -} - -#[cfg(test)] -mod buffered_tests { - use super::*; - use serde_json::json; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - /// Fake ingestor: records peak concurrent in-flight `ingest` calls and can - /// fail one specific `task_id`. No memory tree or embedder involved. - struct CountingIngestor { - in_flight: AtomicUsize, - peak: AtomicUsize, - fail_task: Option, - } - - impl CountingIngestor { - fn new(fail_task: Option<&str>) -> Arc { - Arc::new(Self { - in_flight: AtomicUsize::new(0), - peak: AtomicUsize::new(0), - fail_task: fail_task.map(str::to_string), - }) - } - } - - #[async_trait] - impl TaskIngestor for CountingIngestor { - async fn ingest( - &self, - task_id: &str, - _title: &str, - _updated: Option<&str>, - _task: &Value, - ) -> anyhow::Result { - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.peak.fetch_max(now, Ordering::SeqCst); - for _ in 0..4 { - tokio::task::yield_now().await; - } - self.in_flight.fetch_sub(1, Ordering::SeqCst); - if self.fail_task.as_deref() == Some(task_id) { - Err(anyhow::anyhow!("forced failure for {task_id}")) - } else { - Ok(1) - } - } - } - - fn make_pending(n: usize) -> Vec { - (0..n) - .map(|i| PendingIngest { - sync_key: format!("k{i}"), - task_id: format!("t{i}"), - title: format!("ClickUp: task {i}"), - updated: None, - task: json!({ "id": format!("t{i}") }), - }) - .collect() - } - - #[tokio::test] - async fn ingest_pending_buffered_bounds_and_overlaps() { - let ingestor = CountingIngestor::new(None); - let pending = make_pending(20); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, "ws1", 8).await; - - assert_eq!(outcome.persisted, 20, "all tasks persisted"); - assert_eq!(outcome.synced_keys.len(), 20); - assert!(!outcome.had_failures); - - let peak = ingestor.peak.load(Ordering::SeqCst); - assert!(peak <= 8, "peak in-flight {peak} exceeded the bound of 8"); - assert!(peak >= 2, "peak in-flight {peak} shows no real overlap"); - } - - #[tokio::test] - async fn ingest_pending_buffered_skips_failures_order_independent() { - let ingestor = CountingIngestor::new(Some("t2")); - let pending = make_pending(5); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, "ws1", 4).await; - - assert_eq!(outcome.persisted, 4, "the one failed ingest is not counted"); - assert!(outcome.had_failures); - assert_eq!(outcome.synced_keys.len(), 4); - assert!(!outcome.synced_keys.contains(&"k2".to_string())); - } - - #[test] - fn item_dedup_key_composes_id_and_updated() { - let with_update = json!({ "id": "t1", "date_updated": "1780000000000" }); - assert_eq!( - ClickUpSource::new().item_dedup_key(&with_update).as_deref(), - Some("t1@1780000000000") - ); - let no_update = json!({ "id": "t2" }); - assert_eq!( - ClickUpSource::new().item_dedup_key(&no_update).as_deref(), - Some("t2") - ); - assert_eq!( - ClickUpSource::new().item_dedup_key(&json!({ "date_updated": "x" })), - None - ); - } - - #[test] - fn depth_floor_is_epoch_millis() { - let floor = ClickUpSource::new().depth_floor(7); - // Epoch-ms string: all digits, 13 chars in the 2020s+. - assert!(floor.chars().all(|c| c.is_ascii_digit()), "got {floor}"); - assert!( - floor.len() >= 13, - "epoch-ms should be >=13 digits, got {floor}" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/github/ingest.rs b/src/openhuman/memory_sync/composio/providers/github/ingest.rs deleted file mode 100644 index e1effb8ae..000000000 --- a/src/openhuman/memory_sync/composio/providers/github/ingest.rs +++ /dev/null @@ -1,524 +0,0 @@ -//! GitHub → memory tree ingest plumbing. -//! -//! Owns the conversion from a single GitHub issue/PR payload (post-extracted -//! by [`super::sync`]) into a [`DocumentInput`] and drives -//! [`memory::ingest_pipeline::ingest_document`] for that item. -//! -//! Mirrors the canonical Slack/Gmail/Notion per-source ingest layout -//! ([`super::super::slack::ingest`] / [`super::super::gmail::ingest`] / -//! [`super::super::notion::ingest`]) so retrieval surfaces -//! (`memory.search`, `tree.read_chunk`, `tree.browse`, the agent's -//! recall path, summary trees) actually see GitHub content — pre-#2885 -//! the provider wrote via `MemoryClient::store_skill_sync` into the -//! legacy `memory_docs` table, invisible to the memory-tree retrieval -//! stack. -//! -//! ## Source-id scope -//! -//! Source id is `github:{connection_id}:{issue_id}` — one document identity -//! per GitHub issue or pull request per connection. The source tree / raw -//! archive scope is `github:{connection_id}`, so selector-backed issue lists -//! do not create one graph source per issue. -//! -//! ## Re-ingest of edited issues -//! -//! GitHub issues and PRs mutate (the cursor advances by `updated_at` — -//! title, body, labels, assignees, and comments all bump it). Re-ingesting -//! the same `(connection_id, issue_id)` after an update would -//! short-circuit on the pipeline's `already_ingested` gate — so the call -//! site drops prior chunks for the same source_id via -//! `delete_chunks_by_source` *before* re-ingest, mirroring the vault -//! sync pattern in #2720 and the Notion provider in #2887. The -//! provider's own `SyncState::synced_ids` keyed by -//! `{issue_id}@{updated_at}` is the authoritative "have we seen this -//! revision?" check; this module only runs when that says yes. -//! -//! ## Idempotency -//! -//! Chunk IDs are content-hashed inside the memory tree, so re-ingesting -//! a previously-seen issue is an UPSERT on the same chunk row — no -//! duplicate chunks across syncs. - -use anyhow::Result; -use chrono::{DateTime, Utc}; -use serde_json::Value; - -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::{self, IngestResult}; -use crate::openhuman::memory_store::chunks::store::{delete_chunks_by_source, is_source_ingested}; -use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_sync::canonicalize::document::DocumentInput; - -/// Platform identifier embedded in the canonical document body header. -/// Matches the value `memory_tree::retrieval::source::PLATFORM_KINDS` -/// expects for GitHub-sourced documents. -pub const GITHUB_PLATFORM: &str = "github"; - -/// Tags attached to every GitHub-ingested chunk. Stable list — retrieval -/// callers filter on these. -pub const DEFAULT_TAGS: &[&str] = &["github", "ingested"]; - -/// Build the memory-tree source_id for one GitHub issue/PR in one connection. -/// -/// Stable across re-syncs of the same `(connection_id, issue_id)` so the -/// pipeline's idempotency gate works correctly and the dedup-on-edit -/// path can map back to the prior chunks for cleanup before re-ingest. -/// -/// `issue_id` is whatever stable identifier `super::sync::extract_issue_id` -/// produces — typically GitHub's numeric internal ID, falling back to -/// `owner/repo#number` parsed from `html_url`. Either form is stable -/// per-issue, so callers don't need to normalise before constructing -/// the source id. -pub(crate) fn github_source_id(connection_id: &str, issue_id: &str) -> String { - format!("github:{connection_id}:{issue_id}") -} - -/// Build the source tree / raw archive scope for one GitHub connection. -/// -/// Keep this deliberately item-free. Issue/PR ids belong in -/// [`github_source_id`] for document dedupe, not in the user-visible source -/// graph. Repo-archive ingestion has its own repo-level scope in -/// `memory_sync::sources::github`. -pub(crate) fn github_source_scope(connection_id: &str) -> String { - format!("github:{connection_id}") -} - -/// Pretty-printed JSON body for one GitHub issue/PR. We persist the *full* -/// Composio response payload (not just the title + body markdown) so the -/// chunked content retains enough context for retrieval — labels, -/// assignees, milestone, comment counts, and state transitions are all -/// meaningful retrieval signal that a "just the body" projection would -/// drop. -fn render_issue_body(title: &str, issue: &Value) -> String { - let pretty = serde_json::to_string_pretty(issue).unwrap_or_else(|_| "{}".to_string()); - format!("# {title}\n\n```json\n{pretty}\n```\n") -} - -/// Parse a GitHub `updated_at` timestamp (ISO 8601 / RFC 3339, e.g. -/// `"2024-05-21T15:30:00Z"`) into a `DateTime`, falling back to -/// `Utc::now()` on failure so the pipeline still gets a valid timestamp. -fn parse_updated_at(raw: Option<&str>) -> DateTime { - raw.and_then(|s| DateTime::parse_from_rfc3339(s).ok()) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(Utc::now) -} - -/// Best-effort extraction of the issue's `html_url` for `source_ref`. -/// -/// Using the canonical GitHub URL (rather than a custom `github://` scheme -/// like Notion does) keeps the audit trail one-click navigable — clicking -/// `source_ref` in any downstream UI lands on the real issue page. -fn extract_html_url(issue: &Value) -> Option { - issue - .get("html_url") - .and_then(|v| v.as_str()) - .or_else(|| issue.pointer("/data/html_url").and_then(|v| v.as_str())) - .map(|s| s.to_string()) -} - -/// Ingest one GitHub issue (or PR — same shape) into the memory tree. -/// -/// Caller (the provider's `sync` loop) is responsible for the -/// edit-detection / dedup state-machine (`SyncState::synced_ids` keyed -/// by `{issue_id}@{updated_at}`) — this function trusts that the call -/// only happens for items the caller wants to admit. -/// -/// On content updates of an already-ingested source_id, drops the prior -/// chunks first via `delete_chunks_by_source` so the pipeline's -/// `already_ingested` gate doesn't short-circuit the new content. This -/// mirrors the vault sync pattern in #2720 and the Notion provider in -/// #2887. -/// -/// Returns the number of chunks the pipeline wrote. -pub async fn ingest_issue_into_memory_tree( - config: &Config, - connection_id: &str, - issue_id: &str, - title: &str, - updated_at: Option<&str>, - issue: &Value, -) -> Result { - let source_id = github_source_id(connection_id, issue_id); - - // Re-sync of an edited issue: drop prior chunks for the same source_id - // before re-ingest. Both calls are sync rusqlite I/O so they share one - // `spawn_blocking` hop. - // - // We gate `delete_chunks_by_source` behind `is_source_ingested` — the - // delete path uses a `source_kind = ?1` scan with Rust-side - // source-id filtering (see `store::delete_chunks_by_source_filter`), - // so on a first-time ingest of a never-seen issue it would scan every - // Document-kind chunk just to find zero matches. `is_source_ingested` - // is an indexed PK lookup against `mem_tree_ingested_sources`, so it - // converts the common fresh-issue case to one cheap `COUNT(*)` and - // only pays the scan cost on actual re-ingests of edited items. - let cfg_for_blocking = config.clone(); - let source_for_blocking = source_id.clone(); - let removed = tokio::task::spawn_blocking(move || -> Result { - if is_source_ingested( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - )? { - delete_chunks_by_source( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - ) - } else { - Ok(0) - } - }) - .await - .map_err(|e| anyhow::anyhow!("delete-prior task join error: {e}"))??; - if removed > 0 { - tracing::debug!( - connection_id = %connection_id, - issue_id = %issue_id, - removed_chunks = removed, - "[composio:github] ingest: re-ingest cleanup" - ); - } - - let modified_at = parse_updated_at(updated_at); - let body = render_issue_body(title, issue); - let source_ref = extract_html_url(issue); - - let doc = DocumentInput { - provider: GITHUB_PLATFORM.to_string(), - title: title.to_string(), - body, - modified_at, - source_ref, - }; - let tags: Vec = DEFAULT_TAGS.iter().map(|s| s.to_string()).collect(); - let owner = github_source_scope(connection_id); - let path_scope = Some(owner.clone()); - - match ingest_pipeline::ingest_document_with_scope( - config, &source_id, &owner, tags, doc, path_scope, - ) - .await - { - Ok(IngestResult { - chunks_written, - already_ingested, - .. - }) => { - // The delete-first guard above prevents `already_ingested` on - // the normal update path. Seeing it here means the prior - // chunks were already absent (fresh ingest into a primed - // memory_tree) — fine, just log at debug. - tracing::debug!( - connection_id = %connection_id, - issue_id = %issue_id, - chunks_written, - already_ingested, - "[composio:github] ingest: issue persisted" - ); - Ok(chunks_written) - } - Err(err) => Err(anyhow::anyhow!( - // `{err:#}` (alternate formatter) bakes in the anyhow context - // chain so provider.rs's `tracing::warn!(error = %e)` doesn't - // strip the underlying cause (DB / embedding / persist failure) - // when it Displays the wrapped error. - "ingest_document failed for {source_id}: {err:#}" - )), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Disable strict embedding so the pipeline accepts chunks without - // a live embedder (matches the - // `memory::sync_pipeline_e2e_test::test_config` shape). - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn sample_issue(issue_id: u64, updated_at: &str, title: &str, body: &str) -> Value { - json!({ - "id": issue_id, - "number": 99, - "title": title, - "body": body, - "state": "open", - "html_url": format!("https://github.com/acme/core/issues/{}", issue_id), - "user": { "login": "octocat", "id": 1u64 }, - "labels": [ - { "name": "bug", "color": "d73a4a" }, - { "name": "memory", "color": "0e8a16" } - ], - "updated_at": updated_at, - "created_at": "2026-05-01T00:00:00Z", - "comments": 3, - "assignees": [{ "login": "alice" }], - }) - } - - /// `github_source_id` is stable across calls and namespaces - /// `(connection_id, issue_id)` distinctly. Pins the contract the - /// re-ingest cleanup path relies on (`delete_chunks_by_source` - /// against the same `source_id`). - #[test] - fn github_source_id_is_stable_and_namespaced() { - let a = github_source_id("conn-1", "12345"); - let b = github_source_id("conn-1", "12345"); - assert_eq!(a, b); - assert_eq!(a, "github:conn-1:12345"); - - // owner/repo#number slug form also stable. - assert_eq!( - github_source_id("conn-1", "acme/core#42"), - "github:conn-1:acme/core#42" - ); - - assert_ne!( - github_source_id("conn-1", "12345"), - github_source_id("conn-2", "12345"), - "distinct connections must produce distinct source ids" - ); - assert_ne!( - github_source_id("conn-1", "12345"), - github_source_id("conn-1", "67890"), - "distinct issue ids must produce distinct source ids" - ); - } - - #[test] - fn github_source_scope_is_connection_level() { - assert_eq!(github_source_scope("conn-1"), "github:conn-1"); - assert_eq!(github_source_scope("conn-2"), "github:conn-2"); - assert_eq!( - github_source_scope("conn-1"), - github_source_scope("conn-1"), - "scope must stay stable across issues in one connection" - ); - } - - /// `parse_updated_at` accepts valid ISO 8601 / RFC 3339 and falls - /// back to `Utc::now()` on bad input rather than failing the ingest. - /// We don't assert the now-fallback timestamp value (it's - /// time-dependent) — just that we got a `DateTime` back. - #[test] - fn parse_updated_at_handles_valid_and_invalid_inputs() { - let good = parse_updated_at(Some("2026-05-28T12:34:56Z")); - assert_eq!(good.format("%Y-%m-%d").to_string(), "2026-05-28"); - - // Invalid / missing both fall through to `Utc::now()` — sanity - // check that the result is "recent" (within last 5s). - let bad = parse_updated_at(Some("not-a-timestamp")); - assert!((Utc::now() - bad).num_seconds().abs() < 5); - - let missing = parse_updated_at(None); - assert!((Utc::now() - missing).num_seconds().abs() < 5); - } - - /// `render_issue_body` produces a markdown document with the title - /// header + the full issue JSON pretty-printed in a fenced code - /// block. Pins the chunked-content shape — without this the - /// retrieval body becomes "just the title" and loses GitHub-specific - /// signal (labels, assignees, comment count, state) at search time. - #[test] - fn render_issue_body_includes_title_header_and_pretty_json() { - let issue = sample_issue(123, "2026-05-28T10:00:00Z", "Fix the bug", "Repro steps:"); - let body = render_issue_body("GitHub: acme/core#99: Fix the bug", &issue); - assert!(body.starts_with("# GitHub: acme/core#99: Fix the bug\n")); - assert!(body.contains("```json\n")); - assert!(body.contains("\"title\": \"Fix the bug\"")); - assert!(body.contains("\"bug\"")); - assert!(body.contains("\"assignees\"")); - } - - /// `extract_html_url` prefers the top-level field and falls back to - /// the `data.html_url` wrapper Composio sometimes returns. - #[test] - fn extract_html_url_handles_both_envelope_shapes() { - let direct = json!({ "html_url": "https://github.com/x/y/issues/1" }); - assert_eq!( - extract_html_url(&direct), - Some("https://github.com/x/y/issues/1".to_string()) - ); - - let wrapped = json!({ "data": { "html_url": "https://github.com/x/y/issues/2" } }); - assert_eq!( - extract_html_url(&wrapped), - Some("https://github.com/x/y/issues/2".to_string()) - ); - - let missing = json!({ "id": 1u64 }); - assert!(extract_html_url(&missing).is_none()); - } - - /// The #2885 regression test. - /// - /// Before this migration, GitHub sync routed through - /// `MemoryClient::store_skill_sync` → `UnifiedMemory::upsert_document` - /// → `memory_docs` (legacy backend). The memory-tree retrieval - /// surfaces (which every modern caller reads from) saw zero rows. - /// - /// This test pins the new contract: a successful - /// `ingest_issue_into_memory_tree` call writes to `mem_tree_chunks` - /// + `mem_tree_ingested_sources`, so the silent-failure mode can't - /// reappear. Mirrors the `sync_writes_to_memory_tree` regression - /// in `vault::sync` (#2720) and the equivalent in - /// `composio::providers::notion::ingest` (#2887). - #[tokio::test] - async fn ingest_issue_writes_to_memory_tree() { - use crate::openhuman::memory_store::chunks::store::{ - count_chunks, is_source_ingested, list_chunks, ListChunksQuery, - }; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-test"; - let issue_id = "987654321"; - let expected = github_source_id(connection_id, issue_id); - let expected_scope = github_source_scope(connection_id); - let issue = sample_issue(987654321, "2026-05-28T10:00:00Z", "Fix flaky test", "Repro"); - - let chunks_before = count_chunks(&cfg).expect("count_chunks before"); - - let written = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "GitHub: acme/core#99: Fix flaky test", - Some("2026-05-28T10:00:00Z"), - &issue, - ) - .await - .expect("ingest_issue_into_memory_tree"); - - assert!( - written > 0, - "GitHub ingest must write at least one chunk; got {written}" - ); - - // Core regression assertion: chunks landed in memory_tree. - let chunks_after = count_chunks(&cfg).expect("count_chunks after"); - assert!( - chunks_after > chunks_before, - "ingest must populate mem_tree_chunks (#2885): {chunks_before} → {chunks_after}" - ); - - let rows = list_chunks( - &cfg, - &ListChunksQuery { - source_kind: Some(SourceKind::Document), - source_id: Some(expected.clone()), - limit: Some(1), - ..Default::default() - }, - ) - .expect("list chunks for github issue"); - let chunk = rows.first().expect("github chunk should be listed"); - assert_eq!(chunk.metadata.source_id, expected.as_str()); - assert_eq!( - chunk.metadata.path_scope.as_deref(), - Some(expected_scope.as_str()) - ); - - // Source registration. - let cfg_for_blocking = cfg.clone(); - let expected_for_task = expected.clone(); - let registered = tokio::task::spawn_blocking(move || { - is_source_ingested(&cfg_for_blocking, SourceKind::Document, &expected_for_task) - .unwrap_or(false) - }) - .await - .expect("source-check task join"); - assert!( - registered, - "source_id {} must be registered in mem_tree_ingested_sources", - github_source_id(connection_id, issue_id) - ); - } - - /// Re-ingesting an edited issue (same `(connection_id, issue_id)`, - /// different content) cleans up prior chunks and writes fresh ones — - /// the `delete_chunks_by_source` guard sidesteps the pipeline's - /// `already_ingested` short-circuit that would otherwise drop the - /// new revision. - #[tokio::test] - async fn re_ingesting_edited_issue_replaces_prior_chunks() { - use crate::openhuman::memory_store::chunks::store::count_chunks; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-edit"; - let issue_id = "555"; - - // First ingest. - let v1 = sample_issue( - 555, - "2026-05-28T10:00:00Z", - "Flaky test on Linux", - "Original repro: bisect points at #1234.", - ); - let first = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "GitHub: acme/core#99: Flaky test on Linux", - Some("2026-05-28T10:00:00Z"), - &v1, - ) - .await - .expect("first ingest"); - assert!(first > 0); - let after_first = count_chunks(&cfg).expect("count after first"); - - // Re-ingest with different body (issue was edited with new - // description and labels) — should NOT short-circuit, and chunk - // count should not double (prior chunks dropped, new ones - // written, net same per-issue count for this body size). - let v2 = json!({ - "id": 555u64, - "number": 99, - "title": "Flaky test on Linux — root cause found", - "body": "Root cause: race in scheduler. PR #1300 fixes.", - "state": "open", - "html_url": "https://github.com/acme/core/issues/99", - "labels": [ - { "name": "bug" }, - { "name": "needs-review" } - ], - "updated_at": "2026-05-29T11:22:33Z", - }); - let second = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "GitHub: acme/core#99: Flaky test on Linux — root cause found", - Some("2026-05-29T11:22:33Z"), - &v2, - ) - .await - .expect("second ingest"); - assert!( - second > 0, - "edited issue must actually re-ingest, not silently no-op" - ); - let after_second = count_chunks(&cfg).expect("count after second"); - - // The chunk count after the second ingest should equal the - // count after the first (replaced one revision with another), - // not double. Allow ±1 for any rounding in how the chunker - // splits subtly-different markdown. - assert!( - after_second.abs_diff(after_first) <= 1, - "edited issue must replace prior chunks, not append: \ - after_first={after_first} after_second={after_second}" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/github/mod.rs b/src/openhuman/memory_sync/composio/providers/github/mod.rs index a95357bb6..94c62d12f 100644 --- a/src/openhuman/memory_sync/composio/providers/github/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/github/mod.rs @@ -12,9 +12,7 @@ //! //! Issue: #2408. -mod ingest; mod provider; -mod source; mod sync; #[cfg(test)] mod tests; diff --git a/src/openhuman/memory_sync/composio/providers/github/provider.rs b/src/openhuman/memory_sync/composio/providers/github/provider.rs index 9f55ec028..caea4a951 100644 --- a/src/openhuman/memory_sync/composio/providers/github/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/github/provider.rs @@ -23,12 +23,11 @@ use async_trait::async_trait; use serde_json::{json, Value}; use std::time::Duration; -use super::source::run_github_sync; use super::sync; use crate::openhuman::memory_sync::composio::providers::{ merge_extra, pick_str, resolve_sync_interval_secs, ComposioProvider, CuratedTool, - GithubFetchMode, NormalizedTask, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, - TaskFetchFilter, TaskKind, + GithubFetchMode, NormalizedTask, ProviderContext, ProviderUserProfile, TaskFetchFilter, + TaskKind, }; pub(crate) const ACTION_GET_AUTHENTICATED_USER: &str = "GITHUB_GET_THE_AUTHENTICATED_USER"; @@ -118,10 +117,6 @@ impl ComposioProvider for GitHubProvider { /// login resolution, pagination, dedup, the `max_items` cap, and cursor /// handling live in `run_sync`; the GitHub-specific primitives — including /// the **server-side** `sync_depth_days` window — live in [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - run_github_sync(ctx, reason).await - } - async fn fetch_tasks( &self, ctx: &ProviderContext, diff --git a/src/openhuman/memory_sync/composio/providers/github/source.rs b/src/openhuman/memory_sync/composio/providers/github/source.rs deleted file mode 100644 index 3d3972823..000000000 --- a/src/openhuman/memory_sync/composio/providers/github/source.rs +++ /dev/null @@ -1,325 +0,0 @@ -//! GitHub's [`IncrementalSource`] primitives. -//! -//! GitHub rides the generic -//! [`crate::openhuman::memory_sync::composio::providers::orchestrator`]: -//! [`GitHubProvider::sync`](super::provider::GitHubProvider) delegates to -//! [`run_github_sync`]. The orchestrator owns the control flow (budget, -//! pagination bound, dedup, the precise `max_items` clamp, cursor advance/hold, -//! state persistence); this module supplies only the GitHub-specific shapes. -//! -//! GitHub is **flat but identity-scoped**: [`GitHubSource::preamble`] resolves -//! the authenticated login and returns a single [`SyncScope`] carrying it, then -//! the orchestrator pages straight through -//! `GITHUB_SEARCH_ISSUES_AND_PULL_REQUESTS` with `involves:{login}` (1-indexed -//! `page`). Per-item dedup is keyed by `{issue_id}@{updated_at}`. -//! -//! **Server-side depth window.** Unlike the other flat providers, GitHub -//! applies `sync_depth_days` server-side by injecting `updated:>{date}` into the -//! search query on the first sync (before a cursor exists) — so it overrides -//! [`IncrementalSource::server_side_depth`] and the orchestrator skips its -//! client-side timestamp truncation. On incremental syncs the persistent cursor -//! (`updated:>{cursor}`) bounds the window instead. - -use async_trait::async_trait; -use serde_json::{json, Value}; -use std::sync::Mutex; - -use super::provider::{build_search_query_with_depth, ACTION_SEARCH_ISSUES}; -use super::sync; -use crate::openhuman::memory_sync::composio::providers::orchestrator::{ - self, IncrementalSource, IngestOutcome, PageFetch, SyncItem, SyncScope, -}; -use crate::openhuman::memory_sync::composio::providers::sync_state::SyncState; -use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, SyncOutcome, SyncReason, -}; - -/// Items per search page on steady-state syncs. -const PAGE_SIZE: u32 = 50; - -/// Larger page for the initial post-OAuth backfill. -const INITIAL_PAGE_SIZE: u32 = 100; - -/// Maximum pages per sync pass. -const MAX_PAGES: u32 = 20; - -/// GitHub's [`IncrementalSource`]. -#[derive(Default)] -pub(crate) struct GitHubSource { - depth_fragment: Mutex>, -} - -/// Entry point used by [`super::provider::GitHubProvider::sync`]. -pub(crate) async fn run_github_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - orchestrator::run_sync(&GitHubSource::default(), ctx, reason).await -} - -impl GitHubSource { - /// Resolve the authenticated user's GitHub login. Re-fetched every sync - /// (rather than cached in `SyncState`) so it implicitly validates the OAuth - /// token before paginating. Records the request against the budget. - async fn resolve_login( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result { - tracing::debug!( - connection_id = ?ctx.connection_id, - "[composio:github] resolve_login via {}", - super::provider::ACTION_GET_AUTHENTICATED_USER - ); - - let resp = ctx - .execute( - super::provider::ACTION_GET_AUTHENTICATED_USER, - Some(json!({})), - ) - .await - .map_err(|e| { - format!( - "[composio:github] {} failed: {e:#}", - super::provider::ACTION_GET_AUTHENTICATED_USER - ) - })?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!( - "[composio:github] {}: {err}", - super::provider::ACTION_GET_AUTHENTICATED_USER - )); - } - - let login = sync::extract_user_login(&resp.data).ok_or_else(|| { - "[composio:github] GITHUB_GET_THE_AUTHENTICATED_USER returned no login".to_string() - })?; - - tracing::debug!( - connection_id = ?ctx.connection_id, - "[composio:github] resolve_login complete" - ); - Ok(login) - } - - fn set_depth_fragment(&self, fragment: Option) { - let mut guard = self - .depth_fragment - .lock() - .unwrap_or_else(|e| e.into_inner()); - *guard = fragment; - } - - fn depth_fragment(&self) -> Option { - self.depth_fragment - .lock() - .unwrap_or_else(|e| e.into_inner()) - .clone() - } -} - -#[async_trait] -impl IncrementalSource for GitHubSource { - fn toolkit(&self) -> &'static str { - "github" - } - - fn page_size(&self, reason: SyncReason) -> u32 { - match reason { - SyncReason::ConnectionCreated => INITIAL_PAGE_SIZE, - _ => PAGE_SIZE, - } - } - - fn max_pages(&self) -> u32 { - MAX_PAGES - } - - /// GitHub applies the depth window server-side (see module docs). - fn server_side_depth(&self) -> bool { - true - } - - fn detail_noun(&self) -> &'static str { - "issues" - } - - /// Resolve the login and carry it as the single scope's id. - async fn preamble( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - let depth_fragment = if state.cursor.is_none() { - ctx.sync_depth_days.map(|days| { - let floor = chrono::Utc::now() - chrono::Duration::days(days as i64); - format!("updated:>{}", floor.format("%Y-%m-%dT%H:%M:%SZ")) - }) - } else { - None - }; - self.set_depth_fragment(depth_fragment); - - let login = self.resolve_login(ctx, state).await?; - let label = format!("involves:{login}"); - Ok(vec![SyncScope::nested(login, label)]) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result { - let login = &scope.id; - // GitHub paginates by 1-indexed `page`; the orchestrator's opaque cursor - // carries the next page number (`None` = first page). - let page_num: u32 = cursor.and_then(|c| c.parse().ok()).unwrap_or(1); - let page_size = self.page_size(reason); - - // Stable for the whole sync pass: GitHub's `page` is relative to the - // exact query, so the depth floor must not move while paginating. - let depth_fragment = self.depth_fragment(); - let query = build_search_query_with_depth( - login, - state.cursor.as_deref(), - depth_fragment.as_deref(), - ); - - let args = json!({ - "q": query, - "sort": "updated", - "order": "desc", - "per_page": page_size, - "page": page_num, - }); - - tracing::debug!( - connection_id = ?ctx.connection_id, - page_num, - page_size, - has_depth_fragment = depth_fragment.is_some(), - "[composio:github] fetch_page via {ACTION_SEARCH_ISSUES}" - ); - - let resp = ctx - .execute(ACTION_SEARCH_ISSUES, Some(args)) - .await - .map_err(|e| { - format!("[composio:github] {ACTION_SEARCH_ISSUES} page={page_num}: {e:#}") - })?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!( - "[composio:github] {ACTION_SEARCH_ISSUES} page={page_num}: {err}" - )); - } - - let issues = sync::extract_issues(&resp.data); - // A short page means the result set is exhausted — no next page. - let next = if (issues.len() as u32) < page_size { - None - } else { - Some((page_num + 1).to_string()) - }; - - tracing::debug!( - connection_id = ?ctx.connection_id, - page_num, - fetched = issues.len(), - has_next = next.is_some(), - "[composio:github] fetch_page complete" - ); - - Ok(PageFetch { - items: issues, - next, - }) - } - - fn item_dedup_key(&self, item: &Value) -> Option { - let issue_id = sync::extract_issue_id(item)?; - match sync::extract_issue_updated_at(item) { - Some(updated) => Some(format!("{issue_id}@{updated}")), - None => Some(issue_id), - } - } - - fn item_sort_ts(&self, item: &Value) -> Option { - sync::extract_issue_updated_at(item) - } - - /// GitHub ingests sequentially (parity with the prior per-item loop) — one - /// issue at a time into the memory-tree pipeline. - async fn ingest( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - let mut outcome = IngestOutcome::default(); - for it in items { - let Some(issue_id) = sync::extract_issue_id(&it.raw) else { - continue; - }; - let title = sync::extract_issue_title(&it.raw) - .unwrap_or_else(|| format!("GitHub issue {issue_id}")); - match super::ingest::ingest_issue_into_memory_tree( - &ctx.config, - connection_id, - &issue_id, - &title, - it.sort_ts.as_deref(), - &it.raw, - ) - .await - { - Ok(_chunks_written) => { - outcome.synced_keys.push(it.dedup_key); - outcome.persisted += 1; - } - Err(e) => { - outcome.had_failures = true; - tracing::warn!( - issue_id = %issue_id, - error = %e, - "[composio:github] failed to ingest issue into memory_tree (continuing)" - ); - } - } - } - outcome - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn item_dedup_key_composes_id_and_updated() { - let with_update = json!({ "id": 42, "number": 7, "updated_at": "2026-05-01T00:00:00Z" }); - // GitHub ids may be numeric; extract_issue_id renders them as strings. - let key = GitHubSource::default().item_dedup_key(&with_update); - assert!( - key.as_deref().map(|k| k.contains('@')).unwrap_or(false), - "expected composite id@updated key, got {key:?}" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/gmail/ingest.rs b/src/openhuman/memory_sync/composio/providers/gmail/ingest.rs deleted file mode 100644 index 161e3335d..000000000 --- a/src/openhuman/memory_sync/composio/providers/gmail/ingest.rs +++ /dev/null @@ -1,908 +0,0 @@ -//! Gmail → memory tree ingest plumbing. -//! -//! Owns the conversion from a page of `GMAIL_FETCH_EMAILS` slim-envelope -//! messages (post-processed by [`super::post_process`]) into -//! [`EmailThread`] batches grouped by the sorted set of distinct -//! participants (`from` ∪ `to`-list, CC ignored), then drives -//! [`memory::tree::ingest::ingest_email`] per participant group. -//! -//! Source-id is `gmail:{participants}` where participants is -//! `addr1|addr2|...` (sorted, deduped, lowercased bare emails). All -//! correspondence between the same set of people lands in one source tree. -//! -//! Idempotency: chunk IDs are content-hashed inside the memory tree, so -//! re-ingesting a previously-seen Gmail message is an UPSERT — buffer -//! token_sum may drift if content changes (rare for sealed mail), but -//! the tree's seal cascade handles that on next append. - -use std::collections::BTreeMap; - -use anyhow::Result; -use serde_json::Value; - -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::{ - ingest_email, ingest_email_with_raw_refs, IngestResult, -}; -use crate::openhuman::memory::util::redact::redact; -use crate::openhuman::memory_store::chunks::store::RawRef; -use crate::openhuman::memory_store::content::raw::{ - self as raw_store, raw_rel_path, slug_account_email, RawItem, RawKind, -}; -use crate::openhuman::memory_sync::canonicalize::email::{EmailMessage, EmailThread}; -use crate::openhuman::memory_sync::canonicalize::email_clean::{extract_email, parse_message_date}; - -/// Provider name embedded in the canonical email-thread header. Matches -/// the value `memory::tree::retrieval::source::PLATFORM_KINDS` expects. -pub const GMAIL_PROVIDER: &str = "gmail"; - -/// Tags attached to every Gmail-ingested chunk. Stable list — retrieval -/// callers filter on these. -pub const DEFAULT_TAGS: &[&str] = &["gmail", "ingested"]; - -#[derive(Debug, Clone, Default)] -pub struct PageIngestOutcome { - pub chunks_written: usize, - pub item_ids_ingested: Vec, -} - -/// Group raw page messages by the sorted set of distinct participants -/// (`from` ∪ `to`-list). CC is deliberately excluded from the bucket key -/// so CC-only recipients don't fragment conversations. All messages -/// between the same set of people land in the same bucket regardless of -/// direction or thread ID. -/// -/// The bucket key is the participants joined with `|` in sorted order, -/// e.g. `"alice@x.com|bob@y.com"`. Messages within a bucket are sorted -/// ascending by date so the rendered conversation reads chronologically. -pub(crate) fn bucket_by_participants(msgs: &[Value]) -> BTreeMap> { - let mut out: BTreeMap> = BTreeMap::new(); - for m in msgs { - let bucket_key = participants_bucket_key(m); - if bucket_key == "__skip__" { - // Message has no parseable addresses AND no id — drop it and warn. - // Nothing useful can be done with it: no participants means no - // source tree, and no id means no unique bucket either. - log::warn!( - "[composio:gmail][bucket] dropping message with no parseable addresses and no id" - ); - continue; - } - out.entry(bucket_key).or_default().push(m); - } - for bucket in out.values_mut() { - bucket.sort_by_key(|m| { - parse_message_date(m) - .map(|d: chrono::DateTime| d.timestamp()) - .unwrap_or(0) - }); - } - out -} - -/// Compute the participants bucket key for a single raw message. -/// -/// Collects `from` ∪ `to` (as bare lowercased email addresses), sorts -/// and dedupes them, then joins with `|`. -/// -/// **Fallback policy when all addresses fail to parse**: -/// - If the message has a non-empty `id`, use `"orphan:{id}"` so each -/// malformed message gets its own bucket and its own source tree. Two -/// messages with different ids that both fail address parsing will NOT -/// collapse into a single `"unknown"` bucket. -/// - If even `id` is missing or empty, the caller (`bucket_by_participants`) -/// should skip the message (log a warn and drop it). This function signals -/// that case by returning the sentinel `"__skip__"`. -fn participants_bucket_key(raw: &Value) -> String { - let from = extract_email(raw.get("from").and_then(|v| v.as_str()).unwrap_or("")) - .map(|s| s.to_lowercase()) - .filter(|s| !s.is_empty()); - - let to_emails: Vec = parse_address_list_for_bucket(raw.get("to")) - .into_iter() - .filter_map(|addr| extract_email(&addr).map(|s| s.to_lowercase())) - .collect(); - - let mut all: Vec = from.into_iter().chain(to_emails).collect(); - all.sort(); - all.dedup(); - all.retain(|s| !s.is_empty()); - - if all.is_empty() { - // No parseable addresses — fall back to per-message uniqueness to - // avoid collapsing all malformed messages into one "unknown" source - // tree. Each orphan message gets its own bucket so nothing is silently - // lost in a mixed pile. - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()); - match id { - Some(msg_id) => format!("orphan:{}", msg_id), - None => { - // id is missing: signal caller to skip this message entirely. - "__skip__".to_string() - } - } - } else { - all.join("|") - } -} - -/// Parse the `to` / `cc` field for bucket-key construction. Handles both -/// JSON array and comma-separated string forms. Returns raw address -/// strings (may include display names); callers must extract the bare -/// email with [`extract_email`]. -fn parse_address_list_for_bucket(v: Option<&Value>) -> Vec { - match v { - Some(Value::Array(arr)) => arr - .iter() - .filter_map(|s| s.as_str()) - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) - .collect(), - Some(Value::String(s)) => s - .split(',') - .map(|p| p.trim().to_string()) - .filter(|p| !p.is_empty()) - .collect(), - _ => Vec::new(), - } -} - -/// Build an [`EmailMessage`] from a raw slim-envelope JSON message. -/// Returns `None` when the message has no parseable date — the rest of -/// the pipeline can't sort or canonicalise without one. -pub(crate) fn raw_to_email_message(raw: &Value) -> Option { - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .unwrap_or(""); - let from = raw - .get("from") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let to = parse_address_list(raw.get("to")); - let cc = parse_address_list(raw.get("cc")); - let subject = raw - .get("subject") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let sent_at = parse_message_date(raw)?; - let body = raw - .get("markdown") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let list_unsubscribe = raw - .get("list_unsubscribe") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - let source_ref = if id.is_empty() { - None - } else { - Some(format!("gmail://msg/{id}")) - }; - Some(EmailMessage { - from, - to, - cc, - subject, - sent_at, - body, - source_ref, - list_unsubscribe, - }) -} - -/// Parse the `to` / `cc` field which Composio surfaces as either a -/// JSON array of strings or a single comma-separated string. Empty -/// entries are dropped. -fn parse_address_list(v: Option<&Value>) -> Vec { - match v { - Some(Value::Array(arr)) => arr - .iter() - .filter_map(|s| s.as_str()) - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) - .collect(), - Some(Value::String(s)) => s - .split(',') - .map(|p| p.trim().to_string()) - .filter(|p| !p.is_empty()) - .collect(), - _ => Vec::new(), - } -} - -/// Ingest a page of raw Gmail messages into the memory tree. -/// -/// Each participant-bucket (sorted set of `from` ∪ `to` email addresses) -/// becomes one [`EmailThread`] handed to [`ingest_email`]. Every bucket -/// emits the **same** `source_id` keyed on the connection's account -/// email, so all of an account's correspondence rolls up under a single -/// memory source — `gmail:{slug(account_email)}` (e.g. -/// `gmail:stevent95-at-gmail-dot-com`). When the caller can't supply -/// an `account_email` (legacy `gmail-backfill-3d` CLI runs, missing -/// profile fetch), we fall back to the per-participant `gmail:{participants}` -/// shape so older invocations don't lose their stable bucketing. -/// -/// In addition to the chunked content_store output, we mirror every -/// admitted message as a verbatim `.md` under -/// `/raw//emails/_.md`. -/// Useful for debugging, Obsidian browsing, and as a stable archive -/// independent of the chunker / summariser. -/// -/// Returns the total number of chunks written across all buckets so -/// callers can surface counts in logs / outcomes. Per-bucket errors are -/// logged and swallowed — one bad bucket should not abort the whole -/// page (the next sync re-fetches via the date-cursor). -pub async fn ingest_page_into_memory_tree( - config: &Config, - owner: &str, - account_email: Option<&str>, - page_messages: &[Value], -) -> Result { - Ok( - ingest_page_into_memory_tree_with_outcome(config, owner, account_email, page_messages) - .await? - .chunks_written, - ) -} - -pub async fn ingest_page_into_memory_tree_with_outcome( - config: &Config, - owner: &str, - account_email: Option<&str>, - page_messages: &[Value], -) -> Result { - if page_messages.is_empty() { - return Ok(PageIngestOutcome::default()); - } - let account_source_id = account_email - .filter(|e| !e.trim().is_empty()) - .map(|email| format!("gmail:{}", slug_account_email(email))); - - // Best-effort raw archive — runs once per page, before chunking, so - // a chunker bug doesn't block us from capturing the source bytes. - if let Some(ref source_id) = account_source_id { - write_raw_archive(config, source_id, page_messages).map_err(|e| { - log::warn!( - "[composio:gmail][ingest] raw archive write failed source_id_hash={} err={:#}", - redact(source_id), - e - ); - anyhow::anyhow!( - "gmail raw archive write failed for {}: {e:#}", - redact(source_id) - ) - })?; - } - - // Per-account ingest path: one ingest call per upstream message so - // each resulting chunk has a clean 1:1 (or 1:few-for-oversize) - // mapping to a single raw archive file. Each chunk's body is then - // reconstructed at read time from `raw_refs_json` rather than - // duplicated in the SQL `content` column. Falls back to the - // legacy participant-bucket path when we can't derive an - // account-scoped source id (CLI runs / missing profile fetch). - if let Some(ref source_id) = account_source_id { - let outcome = ingest_per_message(config, source_id, owner, page_messages).await?; - log::info!( - "[composio:gmail][ingest] page_done owner_hash={} chunks={} items={} mode=per-account", - redact(owner), - outcome.chunks_written, - outcome.item_ids_ingested.len(), - ); - return Ok(outcome); - } - - // Legacy fallback: participant-bucketed thread ingest. No - // raw_refs_json — read paths fall through to the SQL `content` - // preview or `content_path` if a chunk file is staged. Only used - // by the CLI backfill binary today. - let buckets = bucket_by_participants(page_messages); - let mut total_chunks = 0usize; - let mut total_buckets = 0usize; - let mut item_ids_ingested = Vec::new(); - for (participants, raw_msgs) in &buckets { - let ids_in_bucket: Vec = raw_msgs - .iter() - .filter_map(|raw| { - raw.get("id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .map(str::to_string) - }) - .collect(); - let messages: Vec = raw_msgs - .iter() - .filter_map(|raw| raw_to_email_message(raw)) - .collect(); - if messages.is_empty() { - log::debug!( - "[composio:gmail][ingest] skipping empty bucket participants_hash={}", - redact(participants) - ); - continue; - } - let source_id = format!("gmail:{}", participants); - let thread_subject = pick_thread_subject(&messages); - let thread = EmailThread { - provider: GMAIL_PROVIDER.to_string(), - thread_subject, - messages, - }; - let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect(); - match ingest_email(config, &source_id, owner, tags, thread).await { - Ok(IngestResult { chunks_written, .. }) => { - total_chunks += chunks_written; - total_buckets += 1; - item_ids_ingested.extend(ids_in_bucket); - } - Err(e) => { - log::warn!( - "[composio:gmail][ingest] ingest_email failed participants_hash={} source_id_hash={} err={:#}", - redact(participants), - redact(&source_id), - e - ); - } - } - } - log::info!( - "[composio:gmail][ingest] page_done owner_hash={} buckets={total_buckets} chunks={total_chunks} mode=per-participants", - redact(owner), - ); - Ok(PageIngestOutcome { - chunks_written: total_chunks, - item_ids_ingested, - }) -} - -/// Per-account ingest: one `ingest_email` call per upstream message. -/// -/// Each call produces 1 chunk for normal messages or N chunks for -/// oversize messages (≥`DEFAULT_CHUNK_MAX_TOKENS`). After the ingest -/// we tag every resulting chunk with a `RawRef` pointing at the raw -/// archive file we wrote during `write_raw_archive`, so -/// `read_chunk_body` can reconstruct full bodies without duplicating -/// bytes in the SQL `content` column. -async fn ingest_per_message( - config: &Config, - source_id: &str, - owner: &str, - page_messages: &[Value], -) -> Result { - let mut total_chunks = 0usize; - let mut item_ids_ingested = Vec::new(); - for raw in page_messages { - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()); - let Some(msg_id) = id else { continue }; - let Some(sent_at) = parse_message_date(raw) else { - continue; - }; - let has_raw_body = raw - .get("markdown") - .and_then(|v| v.as_str()) - .is_some_and(|body| !body.trim().is_empty()); - if !has_raw_body { - log::debug!( - "[composio:gmail][ingest] skipping empty raw-backed message msg_id_hash={}", - redact(msg_id) - ); - continue; - } - let Some(message) = raw_to_email_message(raw) else { - continue; - }; - - let raw_path = raw_rel_path( - source_id, - RawKind::Email, - sent_at.timestamp_millis(), - msg_id, - ); - - let thread_subject = pick_thread_subject(std::slice::from_ref(&message)); - let thread = EmailThread { - provider: GMAIL_PROVIDER.to_string(), - thread_subject, - messages: vec![message], - }; - let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect(); - let refs = vec![RawRef { - path: raw_path.clone(), - start: 0, - end: None, - }]; - match ingest_email_with_raw_refs(config, source_id, owner, tags, thread, refs).await { - Ok(result) => { - total_chunks += result.chunks_written; - if result.chunks_written > 0 || !result.chunk_ids.is_empty() { - item_ids_ingested.push(msg_id.to_string()); - } - } - Err(e) => { - log::warn!( - "[composio:gmail][ingest] per-message ingest_email failed msg_id_hash={} err={:#}", - redact(msg_id), - e - ); - return Err(anyhow::anyhow!( - "gmail per-message ingest failed msg_id_hash={}: {e:#}", - redact(msg_id) - )); - } - } - } - Ok(PageIngestOutcome { - chunks_written: total_chunks, - item_ids_ingested, - }) -} - -/// Mirror a page of raw Gmail messages into the on-disk raw archive. -/// -/// Files land under `/raw//emails/_.md`. -/// We write the **backend-produced markdown verbatim** — the -/// `markdown` field on each message is the per-message slice of the -/// response-level `markdownFormatted`, pinned by -/// [`super::post_process::apply_response_level_markdown`] before the -/// reshape runs. That backend rendering already handles HTML -/// stripping, URL shortening / unwrapping, entity decoding, and -/// whitespace collapse — all the cleanup the user is going to read -/// in Obsidian. Re-running the chunker's `email_clean::clean_body` -/// on top would strip reply chains and footers (useful for LLM -/// chunks, *not* for an as-shipped archive) and risks chopping real -/// content that happens to contain a "view in browser" link. -/// -/// A tiny header (`From:` / `Subject:` / `Date:`) is prepended so -/// the file is self-describing when opened standalone — the post- -/// processed markdown body itself contains only the message text. -/// -/// Messages without a parseable date or id are skipped (they'd -/// produce non-stable filenames). -fn write_raw_archive(config: &Config, source_id: &str, page: &[Value]) -> Result { - let content_root = config.memory_tree_content_root(); - let mut bodies: Vec<(String, i64, String)> = Vec::with_capacity(page.len()); - - for raw in page { - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()); - let Some(id) = id else { continue }; - let Some(sent_at) = parse_message_date(raw) else { - continue; - }; - - // Pull the post-processed markdown straight off the upstream - // page. Falls back to an empty body if the post-processor - // didn't run (extremely unlikely — provider.sync() always - // calls `post_process_action_result` before this point). - let markdown_body = raw - .get("markdown") - .and_then(|v| v.as_str()) - .unwrap_or("") - .trim(); - if markdown_body.is_empty() { - log::debug!( - "[composio:gmail][raw] empty markdown body id_hash={} — skipping", - redact(&id) - ); - continue; - } - let from = raw - .get("from") - .and_then(|v| v.as_str()) - .unwrap_or("") - .trim(); - let subject = raw - .get("subject") - .and_then(|v| v.as_str()) - .unwrap_or("") - .trim(); - - let mut composed = String::with_capacity(markdown_body.len() + 256); - if !from.is_empty() { - composed.push_str(&format!("**From:** {from}\n")); - } - if !subject.is_empty() { - composed.push_str(&format!("**Subject:** {subject}\n")); - } - composed.push_str(&format!("**Date:** {}\n\n", sent_at.to_rfc3339())); - composed.push_str(markdown_body); - - bodies.push((id, sent_at.timestamp_millis(), composed)); - } - - let items: Vec> = bodies - .iter() - .map(|(id, ts, md)| RawItem { - uid: id, - created_at_ms: *ts, - markdown: md.as_str(), - kind: RawKind::Email, - }) - .collect(); - - let n = raw_store::write_raw_items(&content_root, source_id, &items)?; - log::debug!( - "[composio:gmail][raw] archived {n} messages source_id_hash={}", - redact(source_id) - ); - Ok(n) -} - -/// Strip "Re:" / "Fwd:" prefixes from the head message's subject so -/// every message in a thread shares one canonical thread subject. Falls -/// back to "(no subject)" when empty. -fn pick_thread_subject(messages: &[EmailMessage]) -> String { - let raw = messages - .first() - .map(|m| m.subject.trim().to_string()) - .unwrap_or_default(); - let stripped = strip_reply_prefixes(&raw); - if stripped.is_empty() { - "(no subject)".to_string() - } else { - stripped - } -} - -/// Iteratively strip `Re:` / `Fwd:` / `Fw:` prefixes (case-insensitive, -/// optional whitespace) from the front of a subject. Stops once a pass -/// removes nothing. -fn strip_reply_prefixes(subject: &str) -> String { - let mut s = subject.trim().to_string(); - loop { - let lower = s.to_ascii_lowercase(); - let stripped = if lower.starts_with("re:") { - Some(&s[3..]) - } else if lower.starts_with("fwd:") { - Some(&s[4..]) - } else if lower.starts_with("fw:") { - Some(&s[3..]) - } else { - None - }; - match stripped { - Some(rest) => { - let trimmed = rest.trim_start().to_string(); - if trimmed == s { - return s; - } - s = trimmed; - } - None => return s, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - // ─── bucket_by_participants tests ───────────────────────────────────────── - - #[test] - fn bidirectional_messages_bucket_together() { - // alice→bob and bob→alice land in the same key "alice@x.com|bob@y.com". - let msgs = vec![ - json!({ - "id": "m1", - "from": "alice@x.com", - "to": "bob@y.com", - "subject": "Hi", - "date": "2026-04-21T10:00:00Z", - "markdown": "hi", - }), - json!({ - "id": "m2", - "from": "bob@y.com", - "to": "alice@x.com", - "subject": "Re: Hi", - "date": "2026-04-21T11:00:00Z", - "markdown": "hey", - }), - ]; - let buckets = bucket_by_participants(&msgs); - assert_eq!(buckets.len(), 1, "both messages must share one bucket"); - let key = buckets.keys().next().unwrap(); - assert_eq!(key, "alice@x.com|bob@y.com"); - assert_eq!(buckets[key].len(), 2); - // Sorted ascending by date inside the bucket. - assert_eq!(buckets[key][0].get("id").unwrap().as_str().unwrap(), "m1"); - assert_eq!(buckets[key][1].get("id").unwrap().as_str().unwrap(), "m2"); - } - - #[test] - fn multi_recipient_bucket_key_sorted() { - // from=alice, to=[bob, carol] → "alice@x.com|bob@y.com|carol@z.com" - let msgs = vec![json!({ - "id": "m1", - "from": "Alice ", - "to": ["bob@y.com", "carol@z.com"], - "subject": "Group", - "date": "2026-04-21T10:00:00Z", - "markdown": "hey all", - })]; - let buckets = bucket_by_participants(&msgs); - let key = buckets.keys().next().unwrap(); - assert_eq!(key, "alice@x.com|bob@y.com|carol@z.com"); - } - - #[test] - fn cc_field_ignored_in_bucket_key() { - // from=alice, to=[bob], cc=[dave] → "alice@x.com|bob@y.com" (no dave). - let msgs = vec![json!({ - "id": "m1", - "from": "alice@x.com", - "to": "bob@y.com", - "cc": "dave@z.com", - "subject": "CC test", - "date": "2026-04-21T10:00:00Z", - "markdown": "body", - })]; - let buckets = bucket_by_participants(&msgs); - let key = buckets.keys().next().unwrap(); - assert_eq!( - key, "alice@x.com|bob@y.com", - "CC must not appear in bucket key" - ); - } - - #[test] - fn solo_message_no_to_buckets_to_sender_only() { - // from=alice, to=[] → "alice@x.com" (single participant). - let msgs = vec![json!({ - "id": "m1", - "from": "alice@x.com", - "subject": "Draft", - "date": "2026-04-21T10:00:00Z", - "markdown": "draft body", - })]; - let buckets = bucket_by_participants(&msgs); - let key = buckets.keys().next().unwrap(); - assert_eq!(key, "alice@x.com"); - } - - #[test] - fn empty_from_and_to_falls_back_to_orphan_bucket() { - // A message with no parseable addresses gets its own orphan bucket - // keyed by its id rather than collapsing everything into "unknown". - let msgs = vec![json!({ - "id": "m1", - "from": "", - "subject": "x", - "date": "2026-04-21T10:00:00Z", - "markdown": "body", - })]; - let buckets = bucket_by_participants(&msgs); - assert_eq!(buckets.len(), 1, "must produce exactly one bucket"); - assert!( - buckets.contains_key("orphan:m1"), - "must fall back to orphan:; got keys: {:?}", - buckets.keys().collect::>() - ); - } - - #[test] - fn two_malformed_messages_with_different_ids_land_in_different_buckets() { - // Two messages with unparseable from/to but different ids must not - // collapse into the same "unknown" bucket — each gets its own orphan. - let msgs = vec![ - json!({ - "id": "orphan_a", - "from": "", - "subject": "x", - "date": "2026-04-21T10:00:00Z", - "markdown": "body a", - }), - json!({ - "id": "orphan_b", - "from": "", - "subject": "y", - "date": "2026-04-21T11:00:00Z", - "markdown": "body b", - }), - ]; - let buckets = bucket_by_participants(&msgs); - assert_eq!( - buckets.len(), - 2, - "each malformed message must have its own bucket; got: {:?}", - buckets.keys().collect::>() - ); - assert!(buckets.contains_key("orphan:orphan_a")); - assert!(buckets.contains_key("orphan:orphan_b")); - } - - #[test] - fn message_with_no_id_and_no_addresses_is_dropped() { - // A message with no id AND no parseable addresses is silently dropped. - let valid = json!({ - "id": "m_ok", - "from": "alice@x.com", - "subject": "ok", - "date": "2026-04-21T10:00:00Z", - "markdown": "ok", - }); - let bad = json!({ - // no "id" field, no from/to - "subject": "bad", - "date": "2026-04-21T10:00:00Z", - "markdown": "bad", - }); - let msgs = vec![valid, bad]; - let buckets = bucket_by_participants(&msgs); - // Only the valid message should produce a bucket. - assert_eq!(buckets.len(), 1, "dropped message must not create a bucket"); - assert!(buckets.contains_key("alice@x.com")); - } - - #[test] - fn display_name_from_stripped_to_bare_email_in_key() { - // "Alice " should yield bare "alice@x.com" in the key. - let msgs = vec![json!({ - "id": "m1", - "from": "Alice ", - "to": "Bob ", - "subject": "Hi", - "date": "2026-04-21T10:00:00Z", - "markdown": "hi", - })]; - let buckets = bucket_by_participants(&msgs); - let key = buckets.keys().next().unwrap(); - assert_eq!(key, "alice@x.com|bob@y.com"); - } - - #[test] - fn no_threadid_field_does_not_affect_bucketing() { - // threadId is completely ignored; two messages from the same participants - // share one bucket even without threadId. - let msgs = vec![ - json!({ - "id": "m1", - "from": "noreply@github.com", - "to": "sanil@x.com", - "subject": "PR opened", - "date": "2026-04-21T10:00:00Z", - "markdown": "body1", - }), - json!({ - "id": "m2", - "from": "noreply@github.com", - "to": "sanil@x.com", - "subject": "PR merged", - "date": "2026-04-21T11:00:00Z", - "markdown": "body2", - }), - ]; - let buckets = bucket_by_participants(&msgs); - assert_eq!(buckets.len(), 1, "both messages must share one bucket"); - let bucket = buckets.values().next().unwrap(); - assert_eq!(bucket.len(), 2); - } - - #[test] - fn raw_to_email_message_parses_slim_envelope() { - let raw = json!({ - "id": "m1", - "from": "Alice ", - "to": "me@example.com", - "cc": "team@example.com", - "subject": "Phoenix kickoff", - "date": "2026-04-21T10:00:00Z", - "markdown": "Let's ship Phoenix.", - }); - let msg = raw_to_email_message(&raw).unwrap(); - assert_eq!(msg.from, "Alice "); - assert_eq!(msg.to, vec!["me@example.com"]); - assert_eq!(msg.cc, vec!["team@example.com"]); - assert_eq!(msg.subject, "Phoenix kickoff"); - assert_eq!(msg.body, "Let's ship Phoenix."); - assert_eq!(msg.source_ref.as_deref(), Some("gmail://msg/m1")); - } - - #[test] - fn raw_to_email_message_handles_to_array() { - let raw = json!({ - "id": "m1", - "from": "a@x", - "to": ["b@x", "c@x"], - "subject": "x", - "date": "2026-04-21T10:00:00Z", - "markdown": "body", - }); - let msg = raw_to_email_message(&raw).unwrap(); - assert_eq!(msg.to, vec!["b@x", "c@x"]); - } - - #[test] - fn raw_to_email_message_handles_comma_separated_to_string() { - let raw = json!({ - "id": "m1", - "from": "a@x", - "to": "b@x, c@x ,d@x", - "subject": "x", - "date": "2026-04-21T10:00:00Z", - "markdown": "body", - }); - let msg = raw_to_email_message(&raw).unwrap(); - assert_eq!(msg.to, vec!["b@x", "c@x", "d@x"]); - } - - #[test] - fn raw_to_email_message_returns_none_on_unparseable_date() { - let raw = json!({ - "id": "m1", - "from": "a@x", - "subject": "x", - "date": "not-a-date", - "markdown": "body", - }); - assert!(raw_to_email_message(&raw).is_none()); - } - - #[test] - fn raw_to_email_message_drops_source_ref_when_id_empty() { - let raw = json!({ - "id": "", - "from": "a@x", - "subject": "x", - "date": "2026-04-21T10:00:00Z", - "markdown": "body", - }); - let msg = raw_to_email_message(&raw).unwrap(); - assert!(msg.source_ref.is_none()); - } - - #[test] - fn strip_reply_prefixes_removes_iterated() { - assert_eq!(strip_reply_prefixes("Re: Re: Hi"), "Hi"); - assert_eq!(strip_reply_prefixes("Fwd: Re: Status"), "Status"); - assert_eq!(strip_reply_prefixes("RE: Question"), "Question"); - assert_eq!(strip_reply_prefixes("Fw: alert"), "alert"); - assert_eq!(strip_reply_prefixes("Plain subject"), "Plain subject"); - } - - #[test] - fn pick_thread_subject_strips_reply_prefixes() { - let messages = vec![EmailMessage { - from: "a@x".into(), - to: vec![], - cc: vec![], - subject: "Re: Re: Phoenix kickoff".into(), - sent_at: chrono::Utc::now(), - body: "body".into(), - source_ref: None, - list_unsubscribe: None, - }]; - assert_eq!(pick_thread_subject(&messages), "Phoenix kickoff"); - } - - #[test] - fn pick_thread_subject_falls_back_to_no_subject() { - let messages = vec![EmailMessage { - from: "a@x".into(), - to: vec![], - cc: vec![], - subject: " ".into(), - sent_at: chrono::Utc::now(), - body: "body".into(), - source_ref: None, - list_unsubscribe: None, - }]; - assert_eq!(pick_thread_subject(&messages), "(no subject)"); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/gmail/mod.rs b/src/openhuman/memory_sync/composio/providers/gmail/mod.rs index edc27ef37..bdc05ebe2 100644 --- a/src/openhuman/memory_sync/composio/providers/gmail/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/gmail/mod.rs @@ -1,7 +1,5 @@ -pub mod ingest; mod post_process; mod provider; -mod source; mod sync; #[cfg(test)] mod tests; diff --git a/src/openhuman/memory_sync/composio/providers/gmail/provider.rs b/src/openhuman/memory_sync/composio/providers/gmail/provider.rs index fce125dce..2b2cf48c4 100644 --- a/src/openhuman/memory_sync/composio/providers/gmail/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/gmail/provider.rs @@ -9,11 +9,7 @@ //! 4. Run [`ComposioProvider::post_process_action_result`] (bounded //! HTML→text, normalise, sanitise) on the page so the LLM-facing chunk //! content is cleaned, not raw. -//! 5. Filter against `synced_ids` for an early-stop optimisation, -//! then ingest the new messages into the memory tree via -//! [`super::ingest::ingest_page_into_memory_tree`] — same pipeline -//! the standalone `gmail-backfill-3d` binary uses, mirroring the -//! Slack provider's `ingest_chat` pattern. +//! 5. Delegate incremental filtering and document ingestion to tinycortex. //! 6. Paginate (up to budget) until no more results or all items in the //! page are already synced. //! 7. Advance the cursor and save state. @@ -25,10 +21,9 @@ use async_trait::async_trait; use serde_json::{json, Value}; -use super::source::run_gmail_sync; use crate::openhuman::memory_sync::composio::providers::{ pick_str, resolve_sync_interval_secs, ComposioProvider, CuratedTool, ProviderContext, - ProviderUserProfile, SyncOutcome, SyncReason, + ProviderUserProfile, }; pub(super) const ACTION_GET_PROFILE: &str = "GMAIL_GET_PROFILE"; @@ -154,10 +149,6 @@ impl ComposioProvider for GmailProvider { /// `run_sync`; the Gmail-specific primitives — the account-email preamble, /// server-side `after:` depth window, adaptive page ceiling, all-synced /// stop, and batch ingest — live in [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - run_gmail_sync(ctx, reason).await - } - async fn on_trigger( &self, ctx: &ProviderContext, @@ -173,7 +164,16 @@ impl ComposioProvider for GmailProvider { if trigger.eq_ignore_ascii_case("GMAIL_NEW_GMAIL_MESSAGE") || trigger.eq_ignore_ascii_case("GMAIL_NEW_MESSAGE") { - if let Err(e) = self.sync(ctx, SyncReason::Manual).await { + let Some(connection_id) = ctx.connection_id.as_deref() else { + return Err("[composio:gmail] trigger missing connection_id".to_string()); + }; + if let Err(e) = crate::openhuman::tinycortex::run_composio_connection( + "gmail", + connection_id, + ctx.config.as_ref(), + ) + .await + { tracing::warn!( error = %e, "[composio:gmail] trigger-driven sync failed (non-fatal)" diff --git a/src/openhuman/memory_sync/composio/providers/gmail/source.rs b/src/openhuman/memory_sync/composio/providers/gmail/source.rs deleted file mode 100644 index 2219b423e..000000000 --- a/src/openhuman/memory_sync/composio/providers/gmail/source.rs +++ /dev/null @@ -1,363 +0,0 @@ -//! Gmail's [`IncrementalSource`] primitives. -//! -//! Gmail rides the generic -//! [`crate::openhuman::memory_sync::composio::providers::orchestrator`]: -//! [`GmailProvider::sync`](super::provider::GmailProvider) delegates to -//! [`run_gmail_sync`]. The orchestrator owns the control flow (budget, -//! pagination bound, dedup, the precise `max_items` clamp, cursor advance/hold, -//! state persistence); this module supplies only the Gmail-specific shapes. -//! -//! Gmail is **flat**, with three quirks handled by trait hooks: -//! * **server-side depth** — the `sync_depth_days` window is injected as an -//! `after:` query qualifier on the first sync (no cursor), so it -//! overrides [`IncrementalSource::server_side_depth`]. -//! * **adaptive page ceiling** — when the previous sync ran within the last -//! few minutes, [`GmailSource::page_ceiling`] caps pagination aggressively. -//! * **stop on an all-synced page** — Gmail's result set is newest-first, so a -//! page whose messages are all already-synced means nothing newer is left -//! ([`IncrementalSource::stop_on_empty_pending`] = `true`). -//! -//! The account email (used as the per-account ingest `source_id`) is resolved -//! once in the preamble — **not** budget-counted, matching the original — and -//! stashed on the source so [`GmailSource::ingest`] can read it back. Messages -//! ingest as a single batch per page; dedup is keyed by the bare message id -//! (emails are immutable). -//! -//! ## One intentional behavior drop -//! -//! The legacy loop also had a `last_seen_id` "head-unchanged" micro-optimization -//! (stop before post-processing page 0 when its first id matches the last sync's -//! freshest id). That is **redundant** with `stop_on_empty_pending`: on a quiet -//! inbox both fetch exactly one page and stop, so the API-budget behavior is -//! unchanged. It is dropped here (and `last_seen_id` is no longer written) to -//! avoid a gmail-only orchestrator hook for a no-op-on-budget optimization. - -use std::sync::OnceLock; - -use async_trait::async_trait; -use serde_json::{json, Value}; - -use super::ingest::ingest_page_into_memory_tree_with_outcome; -use super::provider::{ACTION_FETCH_EMAILS, ACTION_GET_PROFILE, BASE_QUERY}; -use super::sync; -use crate::openhuman::memory_sync::composio::providers::orchestrator::{ - self, IncrementalSource, IngestOutcome, PageFetch, SyncItem, SyncScope, -}; -use crate::openhuman::memory_sync::composio::providers::sync_state::{extract_item_id, SyncState}; -use crate::openhuman::memory_sync::composio::providers::{ - pick_str, ProviderContext, SyncOutcome, SyncReason, -}; - -/// Page size per API call. -const PAGE_SIZE: u32 = 25; - -/// Larger page size for the very first sync after OAuth. -const INITIAL_PAGE_SIZE: u32 = 50; - -/// Maximum pages to fetch in a single sync pass. -const MAX_PAGES_PER_SYNC: u32 = 20; - -/// Adaptive page cap applied when a successful sync ran very recently. -const RECENT_SYNC_MAX_PAGES: u32 = 2; - -/// "Recent" window (ms) used by the adaptive page cap. -const RECENT_SYNC_WINDOW_MS: u64 = 5 * 60 * 1000; - -/// Paths to try when extracting a message's unique id. -const MESSAGE_ID_PATHS: &[&str] = &["id", "data.id", "messageId", "data.messageId"]; - -/// Paths for extracting the internal date (epoch millis or date string) used as -/// the sync cursor. -const MESSAGE_DATE_PATHS: &[&str] = &[ - "internalDate", - "data.internalDate", - "date", - "data.date", - "receivedAt", - "data.receivedAt", -]; - -/// Gmail's [`IncrementalSource`]. Holds the account email resolved in the -/// preamble so [`Self::ingest`] can stamp a stable per-account `source_id`. -pub(crate) struct GmailSource { - account_email: OnceLock>, -} - -impl GmailSource { - fn new() -> Self { - Self { - account_email: OnceLock::new(), - } - } - - /// Resolve the account email via `GMAIL_GET_PROFILE`. Tolerant of failure - /// (returns `None` → ingest falls back to per-participants bucketing) and - /// **not** budget-counted, matching the original sync. - async fn resolve_account_email(&self, ctx: &ProviderContext) -> Option { - match ctx.execute(ACTION_GET_PROFILE, Some(json!({}))).await { - Ok(resp) if resp.successful => pick_str( - &resp.data, - &["emailAddress", "email", "profile.emailAddress"], - ), - Ok(_) | Err(_) => None, - } - } -} - -/// Entry point used by [`super::provider::GmailProvider::sync`]. Bumps the -/// in-process scheduler heartbeat on completion (parity with the original), -/// so a periodic tick doesn't immediately re-fire on top of a trigger- or -/// connection-created sync. -pub(crate) async fn run_gmail_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - let connection_id = ctx - .connection_id - .clone() - .unwrap_or_else(|| "default".to_string()); - let outcome = orchestrator::run_sync(&GmailSource::new(), ctx, reason).await?; - crate::openhuman::memory_sync::composio::periodic::record_sync_success("gmail", &connection_id); - Ok(outcome) -} - -#[async_trait] -impl IncrementalSource for GmailSource { - fn toolkit(&self) -> &'static str { - "gmail" - } - - fn page_size(&self, reason: SyncReason) -> u32 { - match reason { - SyncReason::ConnectionCreated => INITIAL_PAGE_SIZE, - _ => PAGE_SIZE, - } - } - - fn max_pages(&self) -> u32 { - MAX_PAGES_PER_SYNC - } - - /// Adaptive cap: an initial backfill gets the full ceiling; a steady-state - /// sync that ran within [`RECENT_SYNC_WINDOW_MS`] is capped aggressively - /// (recent ticks rarely need more than a couple of pages). - fn page_ceiling(&self, reason: SyncReason, state: &SyncState) -> u32 { - match reason { - SyncReason::ConnectionCreated => MAX_PAGES_PER_SYNC, - _ => match state.last_sync_at_ms { - Some(last_ms) if sync::now_ms().saturating_sub(last_ms) < RECENT_SYNC_WINDOW_MS => { - RECENT_SYNC_MAX_PAGES - } - _ => MAX_PAGES_PER_SYNC, - }, - } - } - - fn stop_on_empty_pending(&self) -> bool { - true - } - - fn server_side_depth(&self) -> bool { - true - } - - fn detail_noun(&self) -> &'static str { - "messages" - } - - /// Resolve the account email (stashed for `ingest`) and return the single - /// flat scope. The profile fetch is not budget-counted (parity). - async fn preamble( - &self, - ctx: &ProviderContext, - _state: &mut SyncState, - ) -> Result, String> { - let email = self.resolve_account_email(ctx).await; - if email.is_none() { - tracing::warn!( - connection_id = ?ctx.connection_id, - "[composio:gmail] account email unresolved; ingest falls back to per-participants source_id" - ); - } - let _ = self.account_email.set(email); - Ok(vec![SyncScope::flat()]) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result { - // Build the Gmail search query. Prefer a second-precision - // `after:` from the persistent cursor; on the first sync inject - // the `sync_depth_days` floor (server-side depth). - let mut query = BASE_QUERY.to_string(); - if let Some(persistent) = state.cursor.as_deref() { - if let Some(epoch_filter) = sync::cursor_to_gmail_after_epoch_filter(persistent) { - query.push_str(&format!(" after:{epoch_filter}")); - } else if let Some(date_filter) = sync::cursor_to_gmail_after_filter(persistent) { - query.push_str(&format!(" after:{date_filter}")); - } - } else if let Some(days) = ctx.sync_depth_days { - let floor_secs = super::super::helpers::epoch_floor_from_depth(days); - query.push_str(&format!(" after:{floor_secs}")); - } - - let mut args = json!({ - "max_results": self.page_size(reason), - "query": query, - }); - // The orchestrator's opaque cursor carries Gmail's page token. - if let Some(token) = cursor { - args["page_token"] = json!(token); - } - - let mut resp = ctx - .execute(ACTION_FETCH_EMAILS, Some(args.clone())) - .await - .map_err(|e| format!("[composio:gmail] {ACTION_FETCH_EMAILS}: {e:#}"))?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!("[composio:gmail] {ACTION_FETCH_EMAILS}: {err}")); - } - - // Pull the backend's pre-rendered `markdownFormatted` onto each message - // before post-process, then slim/normalize the envelope (same order as - // the original sync). - if let Some(top_md) = resp - .markdown_formatted - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - { - super::post_process::apply_response_level_markdown(&mut resp.data, top_md); - } - super::post_process::post_process(ACTION_FETCH_EMAILS, Some(&args), &mut resp.data); - - Ok(PageFetch { - items: sync::extract_messages(&resp.data), - next: sync::extract_page_token(&resp.data), - }) - } - - /// Emails are immutable — dedup by the bare message id (no `@version`). - fn item_dedup_key(&self, item: &Value) -> Option { - extract_item_id(item, MESSAGE_ID_PATHS) - } - - fn item_sort_ts(&self, item: &Value) -> Option { - extract_item_id(item, MESSAGE_DATE_PATHS) - } - - /// Batch-ingest the page's new messages into the memory tree. `synced_keys` - /// are the message ids that actually ingested; a partial ingest trips - /// `had_failures` so the orchestrator holds the cursor for retry. - async fn ingest( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - if items.is_empty() { - return IngestOutcome::default(); - } - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - let owner = format!("gmail-sync:{connection_id}"); - let account_email = self.account_email.get().cloned().flatten(); - let total = items.len(); - let messages: Vec = items.into_iter().map(|it| it.raw).collect(); - - match ingest_page_into_memory_tree_with_outcome( - ctx.config.as_ref(), - &owner, - account_email.as_deref(), - &messages, - ) - .await - { - Ok(outcome) => { - let persisted = outcome.item_ids_ingested.len(); - IngestOutcome { - synced_keys: outcome.item_ids_ingested, - persisted, - // A short ingest (fewer messages persisted than handed in) - // holds the cursor so the next sync re-scans the gap. - had_failures: persisted < total, - } - } - Err(e) => { - tracing::warn!( - error = %format!("{e:#}"), - new_messages = total, - "[composio:gmail] ingest_page_into_memory_tree failed (continuing)" - ); - IngestOutcome { - synced_keys: Vec::new(), - persisted: 0, - had_failures: true, - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn item_dedup_key_is_bare_message_id() { - // No `@version` — emails are immutable. - assert_eq!( - GmailSource::new() - .item_dedup_key(&json!({ "messageId": "m1", "internalDate": "1780000000000" })) - .as_deref(), - Some("m1") - ); - assert_eq!( - GmailSource::new().item_dedup_key(&json!({ "internalDate": "x" })), - None - ); - } - - #[test] - fn gmail_advertises_server_side_depth_and_empty_stop() { - let s = GmailSource::new(); - assert!(s.server_side_depth()); - assert!(s.stop_on_empty_pending()); - assert_eq!(s.detail_noun(), "messages"); - } - - #[test] - fn page_ceiling_caps_aggressively_after_a_recent_sync() { - let s = GmailSource::new(); - let mut state = SyncState::new("gmail", "c"); - // Initial backfill always gets the full ceiling. - assert_eq!( - s.page_ceiling(SyncReason::ConnectionCreated, &state), - MAX_PAGES_PER_SYNC - ); - // A steady-state sync right after the previous one is capped. - state.set_last_sync_at_ms(sync::now_ms()); - assert_eq!( - s.page_ceiling(SyncReason::Periodic, &state), - RECENT_SYNC_MAX_PAGES - ); - // A stale last-sync drops back to the full ceiling. - state.set_last_sync_at_ms(0); - assert_eq!( - s.page_ceiling(SyncReason::Periodic, &state), - MAX_PAGES_PER_SYNC - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/linear/ingest.rs b/src/openhuman/memory_sync/composio/providers/linear/ingest.rs deleted file mode 100644 index 56b9174be..000000000 --- a/src/openhuman/memory_sync/composio/providers/linear/ingest.rs +++ /dev/null @@ -1,316 +0,0 @@ -//! Linear -> memory tree ingest plumbing. -//! -//! Converts one Linear issue payload into a memory_tree [`DocumentInput`] -//! and calls `ingest_document` so retrieval surfaces read the content from -//! `mem_tree_chunks` instead of the legacy `memory_docs` path. - -use anyhow::Result; -use chrono::{DateTime, Utc}; -use serde_json::Value; - -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::{self, IngestResult}; -use crate::openhuman::memory_store::chunks::store::{delete_chunks_by_source, is_source_ingested}; -use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_sync::canonicalize::document::DocumentInput; - -/// Platform identifier embedded in Linear document metadata. -pub const LINEAR_PLATFORM: &str = "linear"; - -/// Stable tags attached to every Linear-ingested issue chunk. -pub const DEFAULT_TAGS: &[&str] = &["linear", "ingested"]; - -/// Build the memory-tree source id for one Linear issue in one connection. -pub(crate) fn linear_source_id(connection_id: &str, issue_id: &str) -> String { - format!("linear:{connection_id}:{issue_id}") -} - -/// Build the source tree / raw archive scope for one Linear connection. -pub(crate) fn linear_source_scope(connection_id: &str) -> String { - format!("linear:{connection_id}") -} - -/// Render the raw Linear issue payload as a markdown document body. -fn render_issue_body(title: &str, issue: &Value) -> String { - let pretty = serde_json::to_string_pretty(issue).unwrap_or_else(|_| "{}".to_string()); - format!("# {title}\n\n```json\n{pretty}\n```\n") -} - -/// Parse Linear's `updatedAt` timestamp, falling back to now on malformed input. -fn parse_updated_time(raw: Option<&str>) -> DateTime { - raw.and_then(|s| DateTime::parse_from_rfc3339(s).ok()) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(Utc::now) -} - -/// Ingest one Linear issue into memory_tree and return the written chunk count. -/// -/// Edited issues reuse the same `source_id`, so prior chunks are deleted before -/// re-ingest to avoid the document pipeline's duplicate-source short-circuit. -pub async fn ingest_issue_into_memory_tree( - config: &Config, - connection_id: &str, - issue_id: &str, - title: &str, - updated_time: Option<&str>, - issue: &Value, -) -> Result { - let source_id = linear_source_id(connection_id, issue_id); - - let cfg_for_blocking = config.clone(); - let source_for_blocking = source_id.clone(); - let removed = tokio::task::spawn_blocking(move || -> Result { - if is_source_ingested( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - )? { - delete_chunks_by_source( - &cfg_for_blocking, - SourceKind::Document, - &source_for_blocking, - ) - } else { - Ok(0) - } - }) - .await - .map_err(|e| anyhow::anyhow!("delete-prior task join error: {e}"))??; - - if removed > 0 { - tracing::debug!( - connection_id = %connection_id, - issue_id = %issue_id, - removed_chunks = removed, - "[composio:linear] ingest: re-ingest cleanup" - ); - } - - let modified_at = parse_updated_time(updated_time); - let body = render_issue_body(title, issue); - let source_ref = Some(format!("linear://issue/{issue_id}")); - let doc = DocumentInput { - provider: LINEAR_PLATFORM.to_string(), - title: title.to_string(), - body, - modified_at, - source_ref, - }; - let tags: Vec = DEFAULT_TAGS.iter().map(|s| s.to_string()).collect(); - let owner = linear_source_scope(connection_id); - let path_scope = Some(owner.clone()); - - match ingest_pipeline::ingest_document_with_scope( - config, &source_id, &owner, tags, doc, path_scope, - ) - .await - { - Ok(IngestResult { - chunks_written, - already_ingested, - .. - }) => { - tracing::debug!( - connection_id = %connection_id, - issue_id = %issue_id, - chunks_written, - already_ingested, - "[composio:linear] ingest: issue persisted" - ); - Ok(chunks_written) - } - Err(err) => Err(anyhow::anyhow!( - "ingest_document failed for {source_id}: {err:#}" - )), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::config::Config; - use crate::openhuman::memory_store::chunks::types::SourceKind; - use chrono::Utc; - use serde_json::{json, Value}; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn sample_issue(issue_id: &str, updated_at: &str) -> Value { - json!({ - "id": issue_id, - "identifier": "ENG-42", - "title": "Fix external LLM routing", - "updatedAt": updated_at, - "url": "https://linear.app/openhuman/issue/ENG-42/fix-external-llm-routing", - "description": "Connected app tools vanish under external routing.", - "state": { "name": "In Progress" }, - "team": { "key": "ENG", "name": "Engineering" }, - "assignee": { "name": "Alice" } - }) - } - - #[test] - fn linear_source_id_is_stable_and_namespaced() { - let a = linear_source_id("conn-1", "issue-abc"); - let b = linear_source_id("conn-1", "issue-abc"); - assert_eq!(a, b); - assert_eq!(a, "linear:conn-1:issue-abc"); - assert_ne!(a, linear_source_id("conn-2", "issue-abc")); - assert_ne!(a, linear_source_id("conn-1", "issue-xyz")); - } - - #[test] - fn linear_source_scope_is_connection_level() { - assert_eq!(linear_source_scope("conn-1"), "linear:conn-1"); - assert_eq!(linear_source_scope("conn-2"), "linear:conn-2"); - assert_eq!( - linear_source_scope("conn-1"), - linear_source_scope("conn-1"), - "scope must stay stable across issues in one connection" - ); - } - - #[test] - fn parse_updated_time_handles_valid_and_invalid_inputs() { - let good = parse_updated_time(Some("2026-05-28T12:34:56.000Z")); - assert_eq!(good.format("%Y-%m-%d").to_string(), "2026-05-28"); - - let bad = parse_updated_time(Some("not-a-timestamp")); - assert!((Utc::now() - bad).num_seconds().abs() < 5); - - let missing = parse_updated_time(None); - assert!((Utc::now() - missing).num_seconds().abs() < 5); - } - - #[test] - fn render_issue_body_includes_title_header_and_pretty_json() { - let issue = json!({ - "id": "issue-1", - "identifier": "ENG-42", - "title": "Fix external LLM routing" - }); - let body = render_issue_body("Linear: Fix external LLM routing", &issue); - assert!(body.starts_with("# Linear: Fix external LLM routing\n")); - assert!(body.contains("```json\n")); - assert!(body.contains("\"identifier\": \"ENG-42\"")); - assert!(body.contains("\"title\": \"Fix external LLM routing\"")); - } - - #[tokio::test] - async fn ingest_issue_writes_to_memory_tree() { - use crate::openhuman::memory_store::chunks::store::{ - count_chunks, is_source_ingested, list_chunks, ListChunksQuery, - }; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-linear"; - let issue_id = "issue-routing"; - let expected = linear_source_id(connection_id, issue_id); - let expected_scope = linear_source_scope(connection_id); - let issue = sample_issue(issue_id, "2026-05-28T10:00:00.000Z"); - let chunks_before = count_chunks(&cfg).expect("count_chunks before"); - - let written = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "Linear: Fix external LLM routing", - Some("2026-05-28T10:00:00.000Z"), - &issue, - ) - .await - .expect("ingest_issue_into_memory_tree"); - - assert!(written > 0, "Linear ingest must write chunks"); - let chunks_after = count_chunks(&cfg).expect("count_chunks after"); - assert!( - chunks_after > chunks_before, - "ingest must populate mem_tree_chunks (#2885)" - ); - - let rows = list_chunks( - &cfg, - &ListChunksQuery { - source_kind: Some(SourceKind::Document), - source_id: Some(expected.clone()), - limit: Some(1), - ..Default::default() - }, - ) - .expect("list chunks for linear issue"); - let chunk = rows.first().expect("linear chunk should be listed"); - assert_eq!(chunk.metadata.source_id, expected.as_str()); - assert_eq!( - chunk.metadata.path_scope.as_deref(), - Some(expected_scope.as_str()) - ); - - let cfg_for_blocking = cfg.clone(); - let expected_for_task = expected.clone(); - let registered = tokio::task::spawn_blocking(move || { - is_source_ingested(&cfg_for_blocking, SourceKind::Document, &expected_for_task) - .unwrap_or(false) - }) - .await - .expect("source-check task join"); - assert!(registered, "source_id must be registered"); - } - - #[tokio::test] - async fn re_ingesting_edited_issue_replaces_prior_chunks() { - use crate::openhuman::memory_store::chunks::store::count_chunks; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-edit"; - let issue_id = "issue-edit"; - - let v1 = sample_issue(issue_id, "2026-05-28T10:00:00.000Z"); - let first = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "Linear: Fix external LLM routing", - Some("2026-05-28T10:00:00.000Z"), - &v1, - ) - .await - .expect("first ingest"); - assert!(first > 0); - let after_first = count_chunks(&cfg).expect("count after first"); - - let v2 = json!({ - "id": issue_id, - "identifier": "ENG-42", - "title": "Fix external LLM routing", - "updatedAt": "2026-05-29T10:00:00.000Z", - "description": "Updated: external LLM routing now keeps connected app tools visible.", - "state": { "name": "Done" } - }); - let second = ingest_issue_into_memory_tree( - &cfg, - connection_id, - issue_id, - "Linear: Fix external LLM routing", - Some("2026-05-29T10:00:00.000Z"), - &v2, - ) - .await - .expect("second ingest"); - assert!(second > 0); - let after_second = count_chunks(&cfg).expect("count after second"); - - assert!( - after_second.abs_diff(after_first) <= 1, - "edited issue must replace prior chunks, not append duplicates" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/linear/mod.rs b/src/openhuman/memory_sync/composio/providers/linear/mod.rs index 532a2892e..d744dbe6a 100644 --- a/src/openhuman/memory_sync/composio/providers/linear/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/linear/mod.rs @@ -3,9 +3,7 @@ //! //! Issue: #2400. -mod ingest; mod provider; -mod source; mod sync; #[cfg(test)] mod tests; diff --git a/src/openhuman/memory_sync/composio/providers/linear/provider.rs b/src/openhuman/memory_sync/composio/providers/linear/provider.rs index e51244117..41547ad59 100644 --- a/src/openhuman/memory_sync/composio/providers/linear/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/linear/provider.rs @@ -21,13 +21,10 @@ use async_trait::async_trait; use serde_json::json; -use super::source::run_linear_sync; use super::sync; -use crate::openhuman::memory_sync::composio::providers::sync_state::extract_item_id; use crate::openhuman::memory_sync::composio::providers::{ merge_extra, pick_str, resolve_sync_interval_secs, ComposioProvider, CuratedTool, - NormalizedTask, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, TaskFetchFilter, - TaskKind, + NormalizedTask, ProviderContext, ProviderUserProfile, TaskFetchFilter, TaskKind, }; pub(super) const ACTION_LIST_USERS: &str = "LINEAR_LIST_LINEAR_USERS"; @@ -115,10 +112,6 @@ impl ComposioProvider for LinearProvider { /// viewer resolution, pagination, dedup, the `max_items` cap, the /// `sync_depth_days` window, and cursor handling live in `run_sync`; the /// Linear-specific primitives live in [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - run_linear_sync(ctx, reason).await - } - async fn fetch_tasks( &self, ctx: &ProviderContext, @@ -209,7 +202,7 @@ impl ComposioProvider for LinearProvider { /// Map a raw Linear issue payload into a [`NormalizedTask`]. fn normalize_linear_issue(issue: &serde_json::Value) -> Option { - let external_id = extract_item_id(issue, ISSUE_ID_PATHS)?; + let external_id = pick_str(issue, ISSUE_ID_PATHS)?; let title = sync::extract_issue_title(issue).unwrap_or_else(|| format!("Linear issue {external_id}")); Some(NormalizedTask { diff --git a/src/openhuman/memory_sync/composio/providers/linear/source.rs b/src/openhuman/memory_sync/composio/providers/linear/source.rs deleted file mode 100644 index 84343e7f3..000000000 --- a/src/openhuman/memory_sync/composio/providers/linear/source.rs +++ /dev/null @@ -1,397 +0,0 @@ -//! Linear's [`IncrementalSource`] primitives. -//! -//! Linear rides the generic -//! [`crate::openhuman::memory_sync::composio::providers::orchestrator`]: -//! [`LinearProvider::sync`](super::provider::LinearProvider) delegates to -//! [`run_linear_sync`]. The orchestrator owns the control flow (budget, -//! pagination bound, dedup, the `sync_depth_days` window, the precise -//! `max_items` clamp, cursor advance/hold, state persistence); this module -//! supplies only the Linear-specific shapes. -//! -//! Linear is **flat but identity-scoped**: [`LinearSource::preamble`] resolves -//! the viewer id and returns a single [`SyncScope`] carrying it, then the -//! orchestrator pages straight through `LINEAR_LIST_LINEAR_ISSUES` filtered to -//! that assignee. Pagination is GraphQL cursor based (`after` / `endCursor`); -//! the depth window is applied client-side (RFC3339 `updatedAt`). Per-item -//! dedup is keyed by `{issue_id}@{updatedAt}` so an edited issue re-ingests. - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{json, Value}; - -use super::provider::{ACTION_LIST_ISSUES, ACTION_LIST_USERS, ISSUE_ID_PATHS}; -use super::{ingest::ingest_issue_into_memory_tree, sync}; -use crate::openhuman::config::Config; -use crate::openhuman::memory_sync::composio::providers::orchestrator::{ - self, IncrementalSource, IngestOutcome, PageFetch, SyncItem, SyncScope, -}; -use crate::openhuman::memory_sync::composio::providers::sync_state::{extract_item_id, SyncState}; -use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, SyncOutcome, SyncReason, -}; - -/// Page size per API call on steady-state syncs. -const PAGE_SIZE: u32 = 50; - -/// Larger initial-sync page size so the first backfill catches up faster. -const INITIAL_PAGE_SIZE: u32 = 100; - -/// Maximum pages per sync pass before yielding. -const MAX_PAGES_PER_SYNC: u32 = 20; - -/// Max in-flight ingests per page. DB writes serialize anyway and the cloud -/// embedder has rate limits, so keep this small. -const INGEST_CONCURRENCY: usize = 8; - -/// Linear's [`IncrementalSource`]. -pub(crate) struct LinearSource; - -/// Entry point used by [`super::provider::LinearProvider::sync`]. -pub(crate) async fn run_linear_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - orchestrator::run_sync(&LinearSource, ctx, reason).await -} - -#[async_trait] -impl IncrementalSource for LinearSource { - fn toolkit(&self) -> &'static str { - "linear" - } - - fn page_size(&self, reason: SyncReason) -> u32 { - match reason { - SyncReason::ConnectionCreated => INITIAL_PAGE_SIZE, - _ => PAGE_SIZE, - } - } - - fn max_pages(&self) -> u32 { - MAX_PAGES_PER_SYNC - } - - fn detail_noun(&self) -> &'static str { - "issues" - } - - /// Resolve the viewer id and return it as the single scope's id. - async fn preamble( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - let resp = ctx - .execute(ACTION_LIST_USERS, Some(json!({ "isMe": true }))) - .await - .map_err(|e| format!("[composio:linear] {ACTION_LIST_USERS} failed: {e:#}"))?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!("[composio:linear] {ACTION_LIST_USERS}: {err}")); - } - - let viewer_id = sync::extract_viewer_id(&resp.data).ok_or_else(|| { - "[composio:linear] LINEAR_LIST_LINEAR_USERS returned no viewer id".to_string() - })?; - Ok(vec![SyncScope::nested(viewer_id, "assignee:me")]) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result { - let mut args = json!({ - "assigneeId": &scope.id, - "first": self.page_size(reason), - "orderBy": "updatedAt", - }); - if let Some(cursor) = cursor { - args["after"] = json!(cursor); - } - - let resp = ctx - .execute(ACTION_LIST_ISSUES, Some(args)) - .await - .map_err(|e| format!("[composio:linear] {ACTION_LIST_ISSUES}: {e:#}"))?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!("[composio:linear] {ACTION_LIST_ISSUES}: {err}")); - } - - Ok(PageFetch { - items: sync::extract_issues(&resp.data), - next: sync::extract_pagination_cursor(&resp.data), - }) - } - - fn item_dedup_key(&self, item: &Value) -> Option { - let issue_id = extract_item_id(item, ISSUE_ID_PATHS)?; - match sync::extract_issue_updated(item) { - Some(updated) => Some(format!("{issue_id}@{updated}")), - None => Some(issue_id), - } - } - - fn item_sort_ts(&self, item: &Value) -> Option { - sync::extract_issue_updated(item) - } - - async fn ingest( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - - let pending: Vec = items - .into_iter() - .filter_map(|it| { - let issue_id = extract_item_id(&it.raw, ISSUE_ID_PATHS)?; - let title_text = sync::extract_issue_title(&it.raw) - .unwrap_or_else(|| format!("Linear issue {issue_id}")); - Some(PendingIngest { - sync_key: it.dedup_key, - issue_id, - title: format!("Linear: {title_text}"), - updated: it.sort_ts, - issue: it.raw, - }) - }) - .collect(); - - let ingestor = MemoryTreeIngestor { - config: ctx.config.as_ref(), - connection_id, - }; - let buffered = ingest_pending_buffered(&ingestor, pending, INGEST_CONCURRENCY).await; - IngestOutcome { - synced_keys: buffered.synced_keys, - persisted: buffered.persisted, - had_failures: buffered.had_failures, - } - } -} - -/// One issue queued for concurrent ingest. Owns its raw issue `Value` (the -/// orchestrator handed ownership via [`SyncItem`]). -struct PendingIngest { - sync_key: String, - issue_id: String, - title: String, - updated: Option, - issue: Value, -} - -/// Folded result of [`ingest_pending_buffered`]. Order-independent. -#[derive(Default)] -struct BufferedIngestOutcome { - synced_keys: Vec, - persisted: usize, - had_failures: bool, -} - -/// Seam over "ingest one Linear issue", so the bounded-concurrency driver can -/// be unit-tested with a fake that records peak in-flight calls. -#[async_trait] -trait IssueIngestor { - async fn ingest( - &self, - issue_id: &str, - title: &str, - updated: Option<&str>, - issue: &Value, - ) -> anyhow::Result; -} - -/// Production ingestor: routes into the memory-tree pipeline. -struct MemoryTreeIngestor<'c> { - config: &'c Config, - connection_id: &'c str, -} - -#[async_trait] -impl IssueIngestor for MemoryTreeIngestor<'_> { - async fn ingest( - &self, - issue_id: &str, - title: &str, - updated: Option<&str>, - issue: &Value, - ) -> anyhow::Result { - ingest_issue_into_memory_tree( - self.config, - self.connection_id, - issue_id, - title, - updated, - issue, - ) - .await - } -} - -/// Ingest queued issues with bounded concurrency, folding into an -/// order-independent [`BufferedIngestOutcome`]. A failed ingest is logged and -/// skipped, tripping `had_failures` so the orchestrator holds the cursor. -async fn ingest_pending_buffered( - ingestor: &I, - pending: Vec, - concurrency: usize, -) -> BufferedIngestOutcome { - let ingest_futs = pending - .into_iter() - .map(|p| async move { - let res = ingestor - .ingest(&p.issue_id, &p.title, p.updated.as_deref(), &p.issue) - .await; - (p.sync_key, p.issue_id, res) - }) - .collect::>(); - - let mut outcome = BufferedIngestOutcome::default(); - let mut ingest_stream = futures::stream::iter(ingest_futs).buffer_unordered(concurrency); - while let Some((sync_key, issue_id, res)) = ingest_stream.next().await { - match res { - Ok(_chunks_written) => { - outcome.synced_keys.push(sync_key); - outcome.persisted += 1; - } - Err(e) => { - outcome.had_failures = true; - tracing::warn!( - issue_id = %issue_id, - error = %e, - "[composio:linear] failed to ingest issue into memory_tree (continuing)" - ); - } - } - } - outcome -} - -#[cfg(test)] -mod buffered_tests { - use super::*; - use serde_json::json; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - /// Fake ingestor: records peak concurrent in-flight `ingest` calls and can - /// fail one specific `issue_id`. No memory tree or embedder involved. - struct CountingIngestor { - in_flight: AtomicUsize, - peak: AtomicUsize, - fail_issue: Option, - } - - impl CountingIngestor { - fn new(fail_issue: Option<&str>) -> Arc { - Arc::new(Self { - in_flight: AtomicUsize::new(0), - peak: AtomicUsize::new(0), - fail_issue: fail_issue.map(str::to_string), - }) - } - } - - #[async_trait] - impl IssueIngestor for CountingIngestor { - async fn ingest( - &self, - issue_id: &str, - _title: &str, - _updated: Option<&str>, - _issue: &Value, - ) -> anyhow::Result { - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.peak.fetch_max(now, Ordering::SeqCst); - for _ in 0..4 { - tokio::task::yield_now().await; - } - self.in_flight.fetch_sub(1, Ordering::SeqCst); - if self.fail_issue.as_deref() == Some(issue_id) { - Err(anyhow::anyhow!("forced failure for {issue_id}")) - } else { - Ok(1) - } - } - } - - fn make_pending(n: usize) -> Vec { - (0..n) - .map(|i| PendingIngest { - sync_key: format!("k{i}"), - issue_id: format!("i{i}"), - title: format!("Linear: issue {i}"), - updated: None, - issue: json!({ "id": format!("i{i}") }), - }) - .collect() - } - - #[tokio::test] - async fn ingest_pending_buffered_bounds_and_overlaps() { - let ingestor = CountingIngestor::new(None); - let pending = make_pending(20); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, 8).await; - - assert_eq!(outcome.persisted, 20, "all issues persisted"); - assert_eq!(outcome.synced_keys.len(), 20); - assert!(!outcome.had_failures); - - let peak = ingestor.peak.load(Ordering::SeqCst); - assert!(peak <= 8, "peak in-flight {peak} exceeded the bound of 8"); - assert!(peak >= 2, "peak in-flight {peak} shows no real overlap"); - } - - #[tokio::test] - async fn ingest_pending_buffered_skips_failures_order_independent() { - let ingestor = CountingIngestor::new(Some("i2")); - let pending = make_pending(5); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, 4).await; - - assert_eq!(outcome.persisted, 4, "the one failed ingest is not counted"); - assert!(outcome.had_failures); - assert_eq!(outcome.synced_keys.len(), 4); - assert!( - !outcome.synced_keys.contains(&"k2".to_string()), - "the failed issue's sync_key must not be marked synced" - ); - } - - #[test] - fn item_dedup_key_composes_id_and_updated() { - let with_update = json!({ "id": "i1", "updatedAt": "2026-05-01T00:00:00Z" }); - assert_eq!( - LinearSource.item_dedup_key(&with_update).as_deref(), - Some("i1@2026-05-01T00:00:00Z") - ); - let no_update = json!({ "id": "i2" }); - assert_eq!( - LinearSource.item_dedup_key(&no_update).as_deref(), - Some("i2") - ); - assert_eq!( - LinearSource.item_dedup_key(&json!({ "updatedAt": "x" })), - None - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/mod.rs b/src/openhuman/memory_sync/composio/providers/mod.rs index 8236cd6df..640b014a2 100644 --- a/src/openhuman/memory_sync/composio/providers/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/mod.rs @@ -34,7 +34,6 @@ mod descriptions; pub(crate) mod helpers; -pub(crate) mod orchestrator; mod scope_lookup; pub mod tool_scope; mod traits; diff --git a/src/openhuman/memory_sync/composio/providers/notion/ingest.rs b/src/openhuman/memory_sync/composio/providers/notion/ingest.rs deleted file mode 100644 index bc4ecb47a..000000000 --- a/src/openhuman/memory_sync/composio/providers/notion/ingest.rs +++ /dev/null @@ -1,663 +0,0 @@ -//! Notion → memory tree ingest plumbing. -//! -//! Owns the conversion from a single Notion page payload (post-extracted -//! by [`super::sync`]) into a [`DocumentInput`] and drives -//! [`memory::ingest_pipeline::ingest_document`] for that page. -//! -//! Mirrors the canonical Slack/Gmail per-source ingest layout -//! ([`super::super::slack::ingest`] / [`super::super::gmail::ingest`]) -//! so retrieval surfaces (`memory.search`, `tree.read_chunk`, -//! `tree.browse`, the agent's recall path, summary trees) actually see -//! Notion content — pre-#2885 the provider wrote via -//! `MemoryClient::store_skill_sync` into the legacy `memory_docs` table, -//! invisible to the memory-tree retrieval stack. -//! -//! ## Source-id scope -//! -//! Source id is `notion:{connection_id}:{page_id}` — one document identity -//! per Notion page per connection. The source tree / raw archive scope is -//! `notion:{connection_id}`, so all pages selected by one Notion connection -//! accumulate under one source folder and one source tree instead of creating -//! one graph source per page. -//! -//! ## Page body content -//! -//! `NOTION_FETCH_DATA` (the sync's listing call) returns page metadata + -//! `properties` only — never the page body. The provider fetches each -//! new/edited page's rendered body via `NOTION_GET_PAGE_MARKDOWN` and passes -//! it here as `markdown_body`; [`render_page_body`] then emits the body as the -//! primary content with the metadata JSON appended. Database/task rows have no -//! body blocks (empty markdown) → metadata-only, which is correct since their -//! data lives in `properties`. -//! -//! ## Re-ingest of edited pages (non-destructive, versioned) -//! -//! Notion pages mutate (`last_edited_time` advances). An edit is ingested -//! **non-destructively**: `last_edited_time` becomes the document `version`, -//! and the source-ingest gate is keyed by `{source_id}@{version}` (via -//! [`ingest_pipeline::ingest_document_versioned`]) so a new revision is -//! admitted *alongside* the prior one — the old chunks are NOT deleted -//! (replacing the pre-versioning `delete_chunks_by_source` path). A -//! [`JobKind::SealDocument`] job then builds this revision's per-document -//! subtree (one versioned doc-root) and merges it into the connection tree; -//! retrieval surfaces only the latest version per document. The provider's -//! `SyncState::synced_ids` keyed by `{page_id}@{edited_time}` is the -//! authoritative "have we seen this revision?" check — unchanged pages are -//! skipped before this module ever runs. -//! -//! ## Idempotency -//! -//! Chunk IDs are content-hashed, so re-ingesting identical content is an -//! UPSERT on the same chunk row — no duplicate chunks across syncs. - -use std::sync::OnceLock; - -use anyhow::Result; -use chrono::{DateTime, Utc}; -use regex::Regex; -use serde_json::Value; - -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline; -use crate::openhuman::memory_queue::store::enqueue; -use crate::openhuman::memory_queue::types::{NewJob, SealDocumentPayload}; -use crate::openhuman::memory_queue::wake_workers; -use crate::openhuman::memory_sync::canonicalize::document::DocumentInput; - -/// Platform identifier embedded in the canonical document body header. -/// Matches the value `memory_tree::retrieval::source::PLATFORM_KINDS` -/// expects for Notion-sourced documents. -pub const NOTION_PLATFORM: &str = "notion"; - -/// Tags attached to every Notion-ingested chunk. Stable list — retrieval -/// callers filter on these. -pub const DEFAULT_TAGS: &[&str] = &["notion", "ingested"]; - -/// Build the memory-tree source_id for one Notion page in one connection. -/// -/// Stable across re-syncs of the same `(connection_id, page_id)` so the -/// pipeline's idempotency gate works correctly and the dedup-on-edit -/// path can map back to the prior chunks for cleanup before re-ingest. -pub(crate) fn notion_source_id(connection_id: &str, page_id: &str) -> String { - format!("notion:{connection_id}:{page_id}") -} - -/// Build the source tree / raw archive scope for one Notion connection. -/// -/// Keep this deliberately item-free. Page ids belong in [`notion_source_id`] -/// for document dedupe, not in the user-visible source graph. -pub(crate) fn notion_source_scope(connection_id: &str) -> String { - format!("notion:{connection_id}") -} - -/// Strip the noise `NOTION_GET_PAGE_MARKDOWN` bakes into the rendered body -/// before it reaches [`render_page_body`] / the ingest pipeline. -/// -/// The Composio markdown export is functional but token-expensive: image -/// links carry S3 *signed* query strings (`?X-Amz-Algorithm=…&X-Amz-Signature=…` -/// — hundreds of tokens of pure noise that also rotate every fetch, defeating -/// content-hash idempotency), inline formatting is wrapped in HTML `` -/// tags that smuggle `discussion-urls=`/`color=` attributes, user @-mentions -/// render as `` tags, and hard line breaks come through as -/// `
`. None of that is useful retrieval signal. This collapses it to clean -/// markdown so the chunked content is the actual prose, headings, and lists. -/// -/// Regexes are compiled once and cached in process-global `OnceLock`s. -/// Conservative by design — it only touches the specific noise patterns above -/// and leaves ordinary markdown (headings, lists, links, fenced code) intact. -fn clean_notion_markdown(md: &str) -> String { - // `![alt](https://host/file.png?X-Amz-...&...)` → `![alt](https://host/file.png)`. - // Match an image link whose URL has a query string, dropping everything - // from the `?` up to the closing paren. The URL itself can't contain `)` - // or whitespace, so `[^)\s?]*` for the path and `[^)]*` for the query is safe. - static IMG_QUERY: OnceLock = OnceLock::new(); - // `inner` → `inner`. Strips the open tag (with any - // attributes: `color=`, `discussion-urls=`, …) and the matching close tag, - // keeping the inner text. Non-greedy attribute match so adjacent spans - // don't get swallowed. - static SPAN_OPEN: OnceLock = OnceLock::new(); - static SPAN_CLOSE: OnceLock = OnceLock::new(); - // `` and `` → removed. - static MENTION: OnceLock = OnceLock::new(); - // `
` / `
` / `
` → newline. - static BR: OnceLock = OnceLock::new(); - // 3+ consecutive newlines (a run of blank lines) → exactly two. - static BLANKS: OnceLock = OnceLock::new(); - - let img_query = - IMG_QUERY.get_or_init(|| Regex::new(r"(!\[[^\]]*\]\([^)\s?]+)\?[^)]*\)").unwrap()); - let span_open = SPAN_OPEN.get_or_init(|| Regex::new(r"(?s)]*>").unwrap()); - let span_close = SPAN_CLOSE.get_or_init(|| Regex::new(r"").unwrap()); - let mention = MENTION.get_or_init(|| { - Regex::new(r"(?s)]*?/>|]*?>.*?").unwrap() - }); - let br = BR.get_or_init(|| Regex::new(r"(?i)").unwrap()); - let blanks = BLANKS.get_or_init(|| Regex::new(r"\n{3,}").unwrap()); - - let s = img_query.replace_all(md, "$1)"); - let s = mention.replace_all(&s, ""); - let s = span_open.replace_all(&s, ""); - let s = span_close.replace_all(&s, ""); - let s = br.replace_all(&s, "\n"); - let s = blanks.replace_all(&s, "\n\n"); - s.trim().to_string() -} - -/// Pretty-printed JSON body for one Notion page. We persist the *full* -/// Composio response payload (not just the title) so the chunked content -/// retains enough context for retrieval — Notion pages don't have a -/// natural single-string canonical body the way Slack messages do. -/// Render the canonical document body for a Notion page. -/// -/// When `markdown` is `Some` (the page's rendered body, fetched via -/// `NOTION_GET_PAGE_MARKDOWN`), it is the primary content — the actual -/// paragraphs, headings, lists, and *body tables* a free-form page contains. -/// The `FETCH_DATA` metadata/properties JSON is appended after it so the -/// page's structured fields (status, assignee, due date — the real data for -/// database pages, which the markdown body does NOT include) stay searchable. -/// -/// When `markdown` is `None` (no body fetched, or the page had none), we fall -/// back to the metadata-only body — the prior behaviour. -fn render_page_body(title: &str, page: &Value, markdown: Option<&str>) -> String { - let pretty = serde_json::to_string_pretty(page).unwrap_or_else(|_| "{}".to_string()); - match markdown { - Some(md) if !md.trim().is_empty() => { - format!("# {title}\n\n{md}\n\n---\n\n```json\n{pretty}\n```\n") - } - _ => format!("# {title}\n\n```json\n{pretty}\n```\n"), - } -} - -/// Parse a Notion `last_edited_time` (ISO 8601 / RFC 3339) into a -/// `DateTime`, falling back to `Utc::now()` on failure so the -/// pipeline still gets a valid timestamp. -fn parse_edited_time(raw: Option<&str>) -> DateTime { - raw.and_then(|s| DateTime::parse_from_rfc3339(s).ok()) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(Utc::now) -} - -/// Ingest one Notion page revision into the memory tree (non-destructive, -/// versioned). -/// -/// Caller (the provider's `sync` loop) owns the edit-detection / dedup -/// state-machine (`SyncState::synced_ids` keyed by `{page_id}@{edited_time}`) -/// — this function trusts it is only called for a revision the caller wants -/// to admit. -/// -/// Unlike the pre-versioning behaviour (which deleted the prior chunks -/// before re-ingesting), an edit is now **non-destructive**: the page's -/// `last_edited_time` is used as the document version, the source-ingest gate -/// is keyed by `{source_id}@{version}` (so a new revision is admitted -/// alongside the old one), and a [`JobKind::SealDocument`] job is enqueued to -/// build this revision's per-document subtree and merge its doc-root into the -/// connection tree. Older revisions stay in place; the retrieval layer -/// surfaces only the latest version per document. -/// -/// Returns the number of chunks the pipeline wrote for this revision. -pub async fn ingest_page_into_memory_tree( - config: &Config, - connection_id: &str, - page_id: &str, - title: &str, - edited_time: Option<&str>, - page: &Value, - markdown_body: Option<&str>, -) -> Result { - let source_id = notion_source_id(connection_id, page_id); - let modified_at = parse_edited_time(edited_time); - // Document version = Notion `last_edited_time` in epoch-ms. Drives both - // the versioned ingest gate and the per-doc subtree's `version_ms` tag - // used for read-time latest-wins. - let version_ms = modified_at.timestamp_millis(); - // Prefer the rendered page body (NOTION_GET_PAGE_MARKDOWN) when present; - // fall back to metadata-only when the caller didn't fetch it. Strip the - // noise (signed image URLs, ``/`` tags, `
`) from the - // body markdown before rendering — the metadata JSON path is untouched. - let cleaned = markdown_body.map(clean_notion_markdown); - let body = render_page_body(title, page, cleaned.as_deref()); - let source_ref = Some(format!("notion://page/{page_id}")); - - let doc = DocumentInput { - provider: NOTION_PLATFORM.to_string(), - title: title.to_string(), - body, - modified_at, - source_ref, - }; - let tags: Vec = DEFAULT_TAGS.iter().map(|s| s.to_string()).collect(); - let owner = notion_source_scope(connection_id); - let path_scope = Some(owner.clone()); - - let result = ingest_pipeline::ingest_document_versioned( - config, - &source_id, - &owner, - tags, - doc, - path_scope, - Some(version_ms), - ) - .await - .map_err(|err| { - // `{err:#}` keeps the anyhow context chain so provider.rs's - // `tracing::warn!(error = %e)` shows the underlying cause. - anyhow::anyhow!("ingest_document failed for {source_id}: {err:#}") - })?; - - if result.already_ingested { - // This exact revision was already ingested (provider dedup normally - // prevents reaching here). Nothing new to seal. - tracing::debug!( - connection_id = %connection_id, - page_id = %page_id, - version_ms, - "[composio:notion] ingest: revision already ingested — skipping seal" - ); - return Ok(0); - } - - // Enqueue the per-document seal: it rolls this revision's chunks up to a - // single doc-root (tagged with `version_ms`) and merges it into the - // connection tree. Forward-only — the prior revision's doc-root is left - // untouched; retrieval filters to the latest version per document. - if !result.chunk_ids.is_empty() { - let payload = SealDocumentPayload { - tree_scope: owner.clone(), - doc_id: source_id.clone(), - version_ms: Some(version_ms), - chunk_ids: result.chunk_ids.clone(), - }; - match NewJob::seal_document(&payload).and_then(|job| enqueue(config, &job)) { - Ok(_) => wake_workers(), - Err(e) => tracing::warn!( - connection_id = %connection_id, - page_id = %page_id, - "[composio:notion] ingest: failed to enqueue seal_document: {e:#}" - ), - } - } - - tracing::debug!( - connection_id = %connection_id, - page_id = %page_id, - chunks_written = result.chunks_written, - version_ms, - "[composio:notion] ingest: page persisted (versioned)" - ); - Ok(result.chunks_written) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Disable strict embedding so the pipeline accepts chunks without - // a live embedder (matches the - // `memory::sync_pipeline_e2e_test::test_config` shape). - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn sample_page(page_id: &str, edited_time: &str) -> Value { - json!({ - "id": page_id, - "object": "page", - "last_edited_time": edited_time, - "properties": { - "Name": { "title": [{ "plain_text": "Phoenix migration plan" }] } - }, - "url": format!("https://www.notion.so/{}", page_id.replace('-', "")), - "body_excerpt": "Phoenix ships Friday after staging review. Alice owns rollback, Bob on-call.", - }) - } - - /// `notion_source_id` is stable across calls and namespaces - /// `(connection_id, page_id)` distinctly. Pins the contract the - /// re-ingest cleanup path relies on (`delete_chunks_by_source` - /// against the same `source_id`). - #[test] - fn notion_source_id_is_stable_and_namespaced() { - let a = notion_source_id("conn-1", "page-abc"); - let b = notion_source_id("conn-1", "page-abc"); - assert_eq!(a, b); - assert_eq!(a, "notion:conn-1:page-abc"); - - assert_ne!( - notion_source_id("conn-1", "page-abc"), - notion_source_id("conn-2", "page-abc"), - "distinct connections must produce distinct source ids" - ); - assert_ne!( - notion_source_id("conn-1", "page-abc"), - notion_source_id("conn-1", "page-xyz"), - "distinct page ids must produce distinct source ids" - ); - } - - #[test] - fn notion_source_scope_is_connection_level() { - assert_eq!(notion_source_scope("conn-1"), "notion:conn-1"); - assert_eq!(notion_source_scope("conn-2"), "notion:conn-2"); - assert_eq!( - notion_source_scope("conn-1"), - notion_source_scope("conn-1"), - "scope must stay stable across pages in one connection" - ); - } - - /// `parse_edited_time` accepts valid ISO 8601 / RFC 3339 and falls - /// back to `Utc::now()` on bad input rather than failing the ingest. - /// We don't assert the now-fallback timestamp value (it's - /// time-dependent) — just that we got a `DateTime` back. - #[test] - fn parse_edited_time_handles_valid_and_invalid_inputs() { - let good = parse_edited_time(Some("2026-05-28T12:34:56.000Z")); - assert_eq!(good.format("%Y-%m-%d").to_string(), "2026-05-28"); - - // Invalid / missing both fall through to `Utc::now()` — sanity - // check that the result is "recent" (within last 5s). - let bad = parse_edited_time(Some("not-a-timestamp")); - assert!((Utc::now() - bad).num_seconds().abs() < 5); - - let missing = parse_edited_time(None); - assert!((Utc::now() - missing).num_seconds().abs() < 5); - } - - /// `render_page_body` produces a markdown document with the title - /// header + the full page JSON pretty-printed in a fenced code - /// block. Pins the chunked-content shape — without this the - /// retrieval body becomes "just the title" and loses Notion-specific - /// signal (properties, URL, excerpt) at search time. - #[test] - fn render_page_body_includes_title_header_and_pretty_json() { - let page = json!({ "id": "p-1", "url": "https://notion.so/p1" }); - // Metadata-only (no markdown body): title header + pretty JSON. - let body = render_page_body("Phoenix plan", &page, None); - assert!(body.starts_with("# Phoenix plan\n")); - assert!(body.contains("```json\n")); - assert!(body.contains("\"id\": \"p-1\"")); - assert!(body.contains("\"url\": \"https://notion.so/p1\"")); - } - - /// With a markdown body present, render the body as the primary content - /// AND keep the metadata JSON appended after a `---` separator, so - /// free-form pages get real content while structured `properties` stay - /// searchable. - #[test] - fn render_page_body_merges_markdown_then_metadata() { - let page = json!({ "id": "p-1", "properties": { "Status": "Done" } }); - let md = "## Heading\n\nReal body text with a list:\n- one\n- two"; - let body = render_page_body("Phoenix plan", &page, Some(md)); - assert!(body.starts_with("# Phoenix plan\n")); - assert!(body.contains("Real body text with a list")); - assert!(body.contains("\n---\n")); // separator before metadata - assert!(body.contains("```json")); // metadata still present - assert!(body.contains("\"Status\": \"Done\"")); - // Body comes BEFORE the metadata block. - let md_pos = body.find("Real body text").unwrap(); - let json_pos = body.find("```json").unwrap(); - assert!( - md_pos < json_pos, - "markdown body must precede metadata json" - ); - } - - /// Empty/whitespace markdown falls back to metadata-only (a database row - /// with no body blocks returns empty markdown). - #[test] - fn render_page_body_empty_markdown_falls_back_to_metadata_only() { - let page = json!({ "id": "p-1" }); - let body = render_page_body("Phoenix plan", &page, Some(" \n ")); - assert!( - !body.contains("\n---\n"), - "empty markdown → no body section" - ); - assert!(body.contains("```json")); - } - - /// Signed-image-URL query strings are stripped to bare URLs (keeping the - /// `.png`), so the hundreds of rotating `X-Amz-*` tokens don't pollute the - /// chunk body or defeat content-hash idempotency. - #[test] - fn clean_notion_markdown_strips_image_query() { - let md = "![diagram](https://prod-files.s3.amazonaws.com/a/b/file.png?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIA%2F20260604&X-Amz-Signature=deadbeef)"; - let out = clean_notion_markdown(md); - assert_eq!( - out, - "![diagram](https://prod-files.s3.amazonaws.com/a/b/file.png)" - ); - assert!(out.contains(".png"), "the .png extension must be kept"); - assert!(!out.contains("X-Amz-"), "no signed-query noise: {out}"); - } - - /// `text` collapses to its inner text, dropping the tag and - /// any `color=` / `discussion-urls=` attributes. - #[test] - fn clean_notion_markdown_collapses_spans() { - let out = clean_notion_markdown(r#"text"#); - assert_eq!(out, "text"); - - // Attributes like discussion-urls are dropped along with the tags. - let out2 = clean_notion_markdown( - r#"before middle after"#, - ); - assert_eq!(out2, "before middle after"); - } - - /// `` tags are removed entirely (both self-closing and - /// paired forms). - #[test] - fn clean_notion_markdown_removes_mentions() { - let out = clean_notion_markdown(r#"hi there"#); - assert_eq!(out, "hi there"); - assert!(!out.contains("mention-user")); - - let paired = - clean_notion_markdown(r#"@Alice owns it"#); - assert_eq!(paired, "owns it"); - } - - /// `
` / `
` / `
` become newlines. - #[test] - fn clean_notion_markdown_converts_br_to_newline() { - assert_eq!( - clean_notion_markdown("line one
line two"), - "line one\nline two" - ); - assert_eq!(clean_notion_markdown("a
b"), "a\nb"); - assert_eq!(clean_notion_markdown("a
b"), "a\nb"); - } - - /// Plain markdown (no noise patterns) passes through unchanged apart from - /// the trailing-whitespace trim. - #[test] - fn clean_notion_markdown_leaves_plain_markdown_unchanged() { - let md = "# Heading\n\nReal body text with a list:\n- one\n- two\n\n[link](https://example.com/page)"; - assert_eq!(clean_notion_markdown(md), md); - } - - /// 3+ consecutive blank lines collapse to a single blank line (two `\n`). - #[test] - fn clean_notion_markdown_collapses_blank_runs() { - assert_eq!(clean_notion_markdown("a\n\n\n\n\nb"), "a\n\nb"); - } - - /// The #2885 regression test. - /// - /// Before this migration, Notion sync routed through - /// `MemoryClient::store_skill_sync` → `UnifiedMemory::upsert_document` - /// → `memory_docs` (legacy backend). The memory-tree retrieval - /// surfaces (which every modern caller reads from) saw zero rows. - /// - /// This test pins the new contract: a successful `ingest_page_into_memory_tree` - /// call writes to `mem_tree_chunks` + `mem_tree_ingested_sources`, - /// so the silent-failure mode can't reappear. Mirrors the - /// `sync_writes_to_memory_tree` regression in `vault::sync` (#2720). - #[tokio::test] - async fn ingest_page_writes_to_memory_tree() { - use crate::openhuman::memory_store::chunks::store::{ - count_chunks, get_chunk_content_path, is_source_ingested, list_chunks, ListChunksQuery, - }; - use crate::openhuman::memory_store::chunks::types::SourceKind; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-test"; - let page_id = "page-phoenix"; - let expected = notion_source_id(connection_id, page_id); - let expected_scope = notion_source_scope(connection_id); - let page = sample_page(page_id, "2026-05-28T10:00:00.000Z"); - - let chunks_before = count_chunks(&cfg).expect("count_chunks before"); - - let written = ingest_page_into_memory_tree( - &cfg, - connection_id, - page_id, - "Phoenix migration plan", - Some("2026-05-28T10:00:00.000Z"), - &page, - None, - ) - .await - .expect("ingest_page_into_memory_tree"); - - assert!( - written > 0, - "Notion ingest must write at least one chunk; got {written}" - ); - - // Core regression assertion: chunks landed in memory_tree. - let chunks_after = count_chunks(&cfg).expect("count_chunks after"); - assert!( - chunks_after > chunks_before, - "ingest must populate mem_tree_chunks (#2885): {chunks_before} → {chunks_after}" - ); - - let rows = list_chunks( - &cfg, - &ListChunksQuery { - source_kind: Some(SourceKind::Document), - source_id: Some(expected.clone()), - limit: Some(1), - ..Default::default() - }, - ) - .expect("list chunks for notion page"); - let chunk = rows.first().expect("notion chunk should be listed"); - assert_eq!(chunk.metadata.source_id, expected.as_str()); - assert_eq!( - chunk.metadata.path_scope.as_deref(), - Some(expected_scope.as_str()) - ); - let content_path = get_chunk_content_path(&cfg, &chunk.id) - .expect("get chunk content path") - .expect("document chunk should have content path"); - assert!( - content_path.starts_with("document/notion-conn-test/"), - "content path should use connection-level scope, got {content_path}" - ); - let body = std::fs::read_to_string(cfg.memory_tree_content_root().join(&content_path)) - .expect("read chunk file"); - assert!(body.contains("source/notion-conn-test"), "{body}"); - assert!( - !body.contains("source/notion-conn-test-page-phoenix"), - "source tag must not include page id: {body}" - ); - - // Source registration. Versioning keys the ingest gate by - // `{source_id}@{version_ms}` (version_ms = last_edited_time epoch-ms) - // so a later revision is admitted non-destructively — assert the - // versioned key is claimed, not the bare source_id. - let version_ms = parse_edited_time(Some("2026-05-28T10:00:00.000Z")).timestamp_millis(); - let gate_key = format!("{expected}@{version_ms}"); - let cfg_for_blocking = cfg.clone(); - let gate_for_task = gate_key.clone(); - let registered = tokio::task::spawn_blocking(move || { - is_source_ingested(&cfg_for_blocking, SourceKind::Document, &gate_for_task) - .unwrap_or(false) - }) - .await - .expect("source-check task join"); - assert!( - registered, - "versioned gate key {gate_key} must be registered in mem_tree_ingested_sources" - ); - } - - /// Re-ingesting an edited page (same `(connection_id, page_id)`, - /// different content + newer `last_edited_time`) is now - /// **non-destructive + versioned**: the new revision is admitted via the - /// `{source_id}@{version_ms}` gate key (not short-circuited), and the - /// prior revision's chunks are LEFT IN PLACE. Read-time latest-wins - /// surfaces only the newer version; the write path never deletes the old - /// one. (Replaces the pre-versioning destructive `delete_chunks_by_source` - /// behaviour.) - #[tokio::test] - async fn re_ingesting_edited_page_keeps_both_versions() { - use crate::openhuman::memory_store::chunks::store::count_chunks; - - let (_tmp, cfg) = test_config(); - let connection_id = "conn-edit"; - let page_id = "page-edit"; - - // First ingest. - let v1 = sample_page(page_id, "2026-05-28T10:00:00.000Z"); - let first = ingest_page_into_memory_tree( - &cfg, - connection_id, - page_id, - "Phoenix plan v1", - Some("2026-05-28T10:00:00.000Z"), - &v1, - None, - ) - .await - .expect("first ingest"); - assert!(first > 0); - let after_first = count_chunks(&cfg).expect("count after first"); - - // Re-ingest with different body + newer edit time — must NOT - // short-circuit (different version gate key) and must NOT delete the - // prior revision's chunks. - let v2 = json!({ - "id": page_id, - "object": "page", - "last_edited_time": "2026-05-29T10:00:00.000Z", - "properties": { "Name": { "title": [{ "plain_text": "Phoenix plan revised" }] } }, - "body_excerpt": "Plan revised: ship Monday, Carol takes on-call instead.", - }); - let second = ingest_page_into_memory_tree( - &cfg, - connection_id, - page_id, - "Phoenix plan v2", - Some("2026-05-29T10:00:00.000Z"), - &v2, - None, - ) - .await - .expect("second ingest"); - assert!( - second > 0, - "edited page must actually re-ingest, not silently no-op" - ); - let after_second = count_chunks(&cfg).expect("count after second"); - - // Non-destructive: the second revision's chunks are ADDED on top of - // the first revision's, so the total grows rather than staying flat. - assert!( - after_second > after_first, - "edited page must keep both versions (non-destructive): \ - after_first={after_first} after_second={after_second}" - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/notion/mod.rs b/src/openhuman/memory_sync/composio/providers/notion/mod.rs index 228a853d2..427b6fb3f 100644 --- a/src/openhuman/memory_sync/composio/providers/notion/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/notion/mod.rs @@ -1,6 +1,4 @@ -mod ingest; mod provider; -mod source; mod sync; #[cfg(test)] mod tests; diff --git a/src/openhuman/memory_sync/composio/providers/notion/provider.rs b/src/openhuman/memory_sync/composio/providers/notion/provider.rs index 490a657bd..1006a0d6a 100644 --- a/src/openhuman/memory_sync/composio/providers/notion/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/notion/provider.rs @@ -18,13 +18,11 @@ use async_trait::async_trait; use serde_json::{json, Value}; -use super::source::run_notion_sync; use super::sync; -use crate::openhuman::memory_sync::composio::providers::sync_state::extract_item_id; use crate::openhuman::memory_sync::composio::providers::{ first_array_str, merge_extra, pick_str, resolve_sync_interval_secs, ComposioProvider, - CuratedTool, NormalizedTask, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, - TaskContainer, TaskFetchFilter, TaskKind, + CuratedTool, NormalizedTask, ProviderContext, ProviderUserProfile, TaskContainer, + TaskFetchFilter, TaskKind, }; pub(crate) const ACTION_GET_ABOUT_ME: &str = "NOTION_GET_ABOUT_ME"; @@ -133,10 +131,6 @@ impl ComposioProvider for NotionProvider { /// the per-item loop, dedup, `max_items` cap, `sync_depth_days` window, and /// cursor handling all live in `run_sync`; the Notion-specific primitives /// (page fetch, dedup key, body fetch, ingest) live in [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - run_notion_sync(ctx, reason).await - } - async fn fetch_tasks( &self, ctx: &ProviderContext, @@ -281,7 +275,16 @@ impl ComposioProvider for NotionProvider { trigger = %trigger, "[composio:notion] on_trigger" ); - if let Err(e) = self.sync(ctx, SyncReason::Manual).await { + let Some(connection_id) = ctx.connection_id.as_deref() else { + return Err("[composio:notion] trigger missing connection_id".to_string()); + }; + if let Err(e) = crate::openhuman::tinycortex::run_composio_connection( + "notion", + connection_id, + ctx.config.as_ref(), + ) + .await + { tracing::warn!( error = %e, "[composio:notion] trigger-driven sync failed (non-fatal)" @@ -298,7 +301,7 @@ impl ComposioProvider for NotionProvider { /// `Due`). Anything unmatched is simply left `None` — the raw payload is /// preserved for enrichment. fn normalize_notion_page(page: &serde_json::Value) -> Option { - let external_id = extract_item_id(page, PAGE_ID_PATHS)?; + let external_id = pick_str(page, PAGE_ID_PATHS)?; let title = sync::extract_page_title(page).unwrap_or_else(|| format!("Notion page {external_id}")); Some(NormalizedTask { @@ -340,7 +343,7 @@ fn normalize_notion_page(page: &serde_json::Value) -> Option { "data.properties.Priority.select.name", ], ), - updated_at: extract_item_id(page, PAGE_EDITED_PATHS), + updated_at: pick_str(page, PAGE_EDITED_PATHS), raw: page.clone(), }) } @@ -368,7 +371,7 @@ pub(super) fn parse_database_results(data: &serde_json::Value) -> Vec usize { - pending_len.min(budget_remaining as usize) -} - -/// Notion's [`IncrementalSource`] — supplies only the provider-specific shapes; -/// the orchestrator owns everything else. -pub(crate) struct NotionSource; - -/// Entry point used by [`super::provider::NotionProvider::sync`]. -pub(crate) async fn run_notion_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - orchestrator::run_sync(&NotionSource, ctx, reason).await -} - -#[async_trait] -impl IncrementalSource for NotionSource { - fn toolkit(&self) -> &'static str { - "notion" - } - - fn page_size(&self, reason: SyncReason) -> u32 { - match reason { - SyncReason::ConnectionCreated => INITIAL_PAGE_SIZE, - _ => PAGE_SIZE, - } - } - - fn max_pages(&self) -> u32 { - MAX_PAGES_PER_SYNC - } - - /// Notion is flat — one implicit scope, no identity resolution needed here - /// (the user profile is fetched separately by `on_connection_created`). - async fn preamble( - &self, - _ctx: &ProviderContext, - _state: &mut SyncState, - ) -> Result, String> { - Ok(vec![SyncScope::flat()]) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result { - let mut args = json!({ - "page_size": self.page_size(reason), - "filter": { "value": "page", "property": "object" }, - "sort": { "direction": "descending", "timestamp": "last_edited_time" } - }); - if let Some(cursor) = cursor { - args["start_cursor"] = json!(cursor); - } - - // A transport error never reached Composio — `?` returns before we - // record. A completed round-trip is recorded against the daily budget - // *before* we surface a provider-reported failure, so a broken - // connection cannot make unlimited billable failed page calls. - let resp = ctx - .execute(ACTION_FETCH_DATA, Some(args)) - .await - .map_err(|e| format!("[composio:notion] {ACTION_FETCH_DATA}: {e:#}"))?; - state.record_requests(1); - - if !resp.successful { - let err = resp - .error - .clone() - .unwrap_or_else(|| "provider reported failure".to_string()); - return Err(format!("[composio:notion] {ACTION_FETCH_DATA}: {err}")); - } - - Ok(PageFetch { - items: sync::extract_results(&resp.data), - next: sync::extract_notion_cursor(&resp.data), - }) - } - - /// Dedup key is `{page_id}@{last_edited_time}` (or the bare id when no - /// edited time) so an edited page re-ingests while an unchanged one is - /// skipped. - fn item_dedup_key(&self, item: &Value) -> Option { - let page_id = extract_item_id(item, PAGE_ID_PATHS)?; - match extract_item_id(item, PAGE_EDITED_PATHS) { - Some(edited) => Some(format!("{page_id}@{edited}")), - None => Some(page_id), - } - } - - fn item_sort_ts(&self, item: &Value) -> Option { - extract_item_id(item, PAGE_EDITED_PATHS) - } - - async fn ingest( - &self, - ctx: &ProviderContext, - _scope: &SyncScope, - state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - - // Build the per-page ingest queue, re-extracting id/title from each raw - // page. `sync_key`/`edited_time` come straight from the orchestrator's - // dedup/sort decision so they stay consistent with `mark_synced`. - let mut pending: Vec = items - .into_iter() - .filter_map(|it| { - let page_id = extract_item_id(&it.raw, PAGE_ID_PATHS)?; - let title_text = sync::extract_page_title(&it.raw) - .unwrap_or_else(|| format!("Notion page {page_id}")); - Some(PendingIngest { - sync_key: it.dedup_key, - page_id, - title: format!("Notion: {title_text}"), - edited_time: it.sort_ts, - page: it.raw, - markdown_body: None, - }) - }) - .collect(); - - // ── Per-page BODY markdown fetch (bounded concurrency) ────────── - // `NOTION_FETCH_DATA` returned metadata + properties only. Pull the - // rendered page body per page so free-form documents ingest as real - // multi-chunk content. One request per page — budget counted. We - // pre-clamp the number of bodies to the remaining daily budget - // (reproducing the old serial "check before each fetch, break at zero" - // semantics as a single up-front decision) and fire the fetches - // `BODY_FETCH_CONCURRENCY`-at-a-time; `buffered` preserves input order - // so result[i] maps back to pending[i]. - let fetch_count = body_fetch_count(pending.len(), state.budget_remaining()); - if fetch_count < pending.len() { - tracing::info!( - fetch_count, - pending = pending.len(), - "[composio:notion] daily budget caps body fetch — \ - remaining pages ingest metadata-only" - ); - } - if fetch_count > 0 { - // Snapshot the page ids so the fetch futures don't borrow `pending` - // (we write results back into it afterwards). - let page_ids: Vec = pending[..fetch_count] - .iter() - .map(|p| p.page_id.clone()) - .collect(); - - let body_futs: Vec<_> = page_ids - .into_iter() - .map(|page_id| { - let ctx = &ctx; - async move { - match ctx - .execute( - ACTION_GET_PAGE_MARKDOWN, - Some(json!({ "page_id": &page_id })), - ) - .await - { - Ok(resp) if resp.successful => { - let markdown = sync::extract_page_markdown(&resp.data); - if markdown.is_none() { - tracing::warn!( - page_id = %page_id, - "[composio:notion] GET_PAGE_MARKDOWN returned no \ - markdown field — metadata-only fallback" - ); - } - markdown - } - Ok(resp) => { - tracing::warn!( - page_id = %page_id, - error = ?resp.error, - "[composio:notion] GET_PAGE_MARKDOWN failed — \ - metadata-only fallback" - ); - None - } - Err(e) => { - tracing::warn!( - page_id = %page_id, - error = %e, - "[composio:notion] GET_PAGE_MARKDOWN execute error — \ - metadata-only fallback" - ); - None - } - } - } - }) - .collect(); - - let bodies: Vec> = futures::stream::iter(body_futs) - .buffered(BODY_FETCH_CONCURRENCY) - .collect() - .await; - - // Count every request we fired (success or failure), matching the - // old per-call `record_requests(1)`. - state.record_requests(fetch_count as u32); - - for (p, body) in pending.iter_mut().zip(bodies) { - p.markdown_body = body; - } - } - - // ── Ingest queued pages (bounded concurrency) ─────────────────── - let ingestor = MemoryTreeIngestor { - config: ctx.config.as_ref(), - connection_id, - }; - let buffered = ingest_pending_buffered(&ingestor, pending, INGEST_CONCURRENCY).await; - IngestOutcome { - synced_keys: buffered.synced_keys, - persisted: buffered.persisted, - had_failures: buffered.had_failures, - } - } -} - -/// One page queued for concurrent ingest. Owns its raw page `Value` (the -/// orchestrator handed ownership via [`SyncItem`]). -struct PendingIngest { - sync_key: String, - page_id: String, - title: String, - edited_time: Option, - page: Value, - /// Rendered page body (markdown) fetched per-page via - /// `NOTION_GET_PAGE_MARKDOWN`. `None` when the body fetch was skipped - /// (budget) or failed — ingest falls back to the metadata-only body. - markdown_body: Option, -} - -/// Folded result of [`ingest_pending_buffered`]. Every field is -/// order-independent so the concurrent stage can accumulate regardless of the -/// order ingests complete. -#[derive(Default)] -struct BufferedIngestOutcome { - /// `sync_key`s whose ingest succeeded — the orchestrator marks each synced. - synced_keys: Vec, - /// Number of pages persisted (equals `synced_keys.len()`). - persisted: usize, - /// Whether any per-item ingest failed (the orchestrator holds the cursor). - had_failures: bool, -} - -/// Seam over "ingest one Notion page", so the bounded-concurrency driver can be -/// unit-tested with a fake that records peak in-flight calls without a real -/// memory tree or embedder. -#[async_trait] -trait PageIngestor { - async fn ingest( - &self, - page_id: &str, - title: &str, - edited_time: Option<&str>, - page: &Value, - markdown_body: Option<&str>, - ) -> anyhow::Result; -} - -/// Production ingestor: routes into the memory-tree pipeline (#2885) via -/// [`ingest_page_into_memory_tree`]. -struct MemoryTreeIngestor<'c> { - config: &'c Config, - connection_id: &'c str, -} - -#[async_trait] -impl PageIngestor for MemoryTreeIngestor<'_> { - async fn ingest( - &self, - page_id: &str, - title: &str, - edited_time: Option<&str>, - page: &Value, - markdown_body: Option<&str>, - ) -> anyhow::Result { - ingest_page_into_memory_tree( - self.config, - self.connection_id, - page_id, - title, - edited_time, - page, - markdown_body, - ) - .await - } -} - -/// Ingest the queued pages with bounded concurrency. Overlaps the per-item -/// embedding RTT (`buffer_unordered`, up to `concurrency` in flight) and folds -/// results into an order-independent [`BufferedIngestOutcome`]. Unordered is -/// correct here: nothing downstream depends on completion order — successes are -/// keyed by `sync_key`. -async fn ingest_pending_buffered( - ingestor: &I, - pending: Vec, - concurrency: usize, -) -> BufferedIngestOutcome { - let ingest_futs = pending - .into_iter() - .map(|p| async move { - let res = ingestor - .ingest( - &p.page_id, - &p.title, - p.edited_time.as_deref(), - &p.page, - p.markdown_body.as_deref(), - ) - .await; - (p.sync_key, p.page_id, res) - }) - .collect::>(); - - let mut outcome = BufferedIngestOutcome::default(); - let mut ingest_stream = futures::stream::iter(ingest_futs).buffer_unordered(concurrency); - while let Some((sync_key, page_id, res)) = ingest_stream.next().await { - match res { - Ok(_chunks_written) => { - outcome.synced_keys.push(sync_key); - outcome.persisted += 1; - } - Err(e) => { - outcome.had_failures = true; - tracing::warn!( - page_id = %page_id, - error = %e, - "[composio:notion] failed to ingest page into memory_tree (continuing)" - ); - } - } - } - outcome -} - -#[cfg(test)] -mod body_fetch_count_tests { - use super::body_fetch_count; - - #[test] - fn clamps_to_remaining_budget_when_budget_is_the_limit() { - // 10 pending pages but only 3 requests left today → fetch 3. - assert_eq!(body_fetch_count(10, 3), 3); - } - - #[test] - fn fetches_all_when_budget_exceeds_pending() { - assert_eq!(body_fetch_count(4, 100), 4); - } - - #[test] - fn zero_budget_fetches_nothing() { - // Mirrors the old serial "budget_exhausted() before first fetch" break. - assert_eq!(body_fetch_count(7, 0), 0); - } - - #[test] - fn zero_pending_fetches_nothing() { - assert_eq!(body_fetch_count(0, 50), 0); - } -} - -#[cfg(test)] -mod buffered_tests { - use super::*; - use serde_json::json; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - /// Fake ingestor: records the peak number of concurrent in-flight `ingest` - /// calls and can be told to fail one specific `page_id`. No memory tree or - /// embedder involved — lets us assert the concurrency bound and overlap - /// deterministically. - struct CountingIngestor { - in_flight: AtomicUsize, - peak: AtomicUsize, - fail_page: Option, - } - - impl CountingIngestor { - fn new(fail_page: Option<&str>) -> Arc { - Arc::new(Self { - in_flight: AtomicUsize::new(0), - peak: AtomicUsize::new(0), - fail_page: fail_page.map(str::to_string), - }) - } - } - - #[async_trait] - impl PageIngestor for CountingIngestor { - async fn ingest( - &self, - page_id: &str, - _title: &str, - _edited_time: Option<&str>, - _page: &Value, - _markdown_body: Option<&str>, - ) -> anyhow::Result { - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.peak.fetch_max(now, Ordering::SeqCst); - for _ in 0..4 { - tokio::task::yield_now().await; - } - self.in_flight.fetch_sub(1, Ordering::SeqCst); - if self.fail_page.as_deref() == Some(page_id) { - Err(anyhow::anyhow!("forced failure for {page_id}")) - } else { - Ok(1) - } - } - } - - fn make_pending(n: usize) -> Vec { - (0..n) - .map(|i| PendingIngest { - sync_key: format!("k{i}"), - page_id: format!("p{i}"), - title: format!("Notion: page {i}"), - edited_time: None, - page: json!({ "id": format!("p{i}") }), - markdown_body: None, - }) - .collect() - } - - #[tokio::test] - async fn ingest_pending_buffered_bounds_and_overlaps() { - let ingestor = CountingIngestor::new(None); - let pending = make_pending(20); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, 8).await; - - assert_eq!(outcome.persisted, 20, "all pages persisted"); - assert_eq!(outcome.synced_keys.len(), 20); - assert!(!outcome.had_failures); - - let peak = ingestor.peak.load(Ordering::SeqCst); - assert!(peak <= 8, "peak in-flight {peak} exceeded the bound of 8"); - assert!(peak >= 2, "peak in-flight {peak} shows no real overlap"); - } - - #[tokio::test] - async fn ingest_pending_buffered_skips_failures_order_independent() { - let ingestor = CountingIngestor::new(Some("p2")); - let pending = make_pending(5); - - let outcome = ingest_pending_buffered(ingestor.as_ref(), pending, 4).await; - - assert_eq!(outcome.persisted, 4, "the one failed ingest is not counted"); - assert!(outcome.had_failures); - assert_eq!(outcome.synced_keys.len(), 4); - assert!( - !outcome.synced_keys.contains(&"k2".to_string()), - "the failed page's sync_key must not be marked synced" - ); - } - - #[test] - fn item_dedup_key_composes_id_and_edited_time() { - let with_edit = json!({ "id": "p1", "last_edited_time": "2026-05-01T00:00:00Z" }); - assert_eq!( - NotionSource.item_dedup_key(&with_edit).as_deref(), - Some("p1@2026-05-01T00:00:00Z") - ); - // No edited time → bare id. - let no_edit = json!({ "id": "p2" }); - assert_eq!(NotionSource.item_dedup_key(&no_edit).as_deref(), Some("p2")); - // No id at all → dropped. - assert_eq!( - NotionSource.item_dedup_key(&json!({ "last_edited_time": "x" })), - None - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/orchestrator.rs b/src/openhuman/memory_sync/composio/providers/orchestrator.rs deleted file mode 100644 index aa6fa69f6..000000000 --- a/src/openhuman/memory_sync/composio/providers/orchestrator.rs +++ /dev/null @@ -1,1545 +0,0 @@ -//! Generic incremental-sync orchestrator for Composio providers. -//! -//! Every Composio provider (gmail, notion, slack, …) used to hand-roll the -//! same ~15-step sync loop: load [`SyncState`] → check the daily budget → -//! resolve identity/scopes → page through results → dedupe against -//! `synced_ids` → drop items outside the `sync_depth_days` window → clamp to -//! `max_items` → ingest → advance the cursor → persist state → build a -//! [`SyncOutcome`]. The copy-paste is exactly how the page-granular cap bug -//! (#3304) ended up in five providers and was missed in a sixth. -//! -//! [`ItemCap`] (PR #3304) centralised the cap *math*; this module owns it and -//! the *control flow*. Now that every provider rides [`run_sync`], `ItemCap` -//! lives here (orchestrator-private), not in the shared `helpers` module — the -//! cap rule exists in exactly one place. A provider implements only the slim -//! [`IncrementalSource`] primitives and inherits the `max_items` cap, the -//! `sync_depth_days` window, dedup, cursor advance, and budget enforcement for -//! free. -//! -//! ## Scopes: flat AND nested in one loop -//! -//! [`SyncScope`] is the abstraction that lets the same orchestrator drive both -//! shapes: -//! -//! * **Flat** providers (github, notion, linear) expose a *single implicit -//! scope* — [`SyncScope::flat`] — and page straight through their one -//! result stream. -//! * **Nested** providers (clickup workspaces → tasks, slack channels → -//! history) resolve their containers in [`IncrementalSource::preamble`] and -//! return one [`SyncScope`] per container; the orchestrator's -//! `for scope { for page {…} }` loop is byte-for-byte identical for both. -//! -//! ## Per-scope cursors + error tolerance (Slack) -//! -//! Slack is the structural outlier: it keeps a watermark **per channel** (not a -//! single global cursor) and must not let one bad channel abort the rest. Two -//! opt-in trait hooks — [`IncrementalSource::per_scope_cursors`] and -//! [`IncrementalSource::tolerate_scope_errors`] — bend the same scope loop to -//! that shape without touching the single-cursor providers, which leave both -//! `false` and keep the advance-once-at-the-end path verbatim. - -use async_trait::async_trait; -use serde_json::{json, Value}; - -use super::sync_state::SyncState; -use super::{ProviderContext, SyncOutcome, SyncReason}; - -/// Compute the number of API pages needed to cover `max_items` at `page_size` -/// items per page, rounding up. -/// -/// Returns `u32::MAX` when `page_size == 0` to avoid division by zero; -/// callers should treat this as "no page cap beyond the provider's own upper -/// bound". -fn pages_for_max_items(max_items: u32, page_size: u32) -> u32 { - if page_size == 0 { - return u32::MAX; - } - // Widen to u64 before the addition to prevent overflow for large cap values. - (max_items as u64) - .div_ceil(page_size as u64) - .min(u32::MAX as u64) as u32 -} - -/// Single source of truth for the per-sync `max_items` cap. -/// -/// Every Composio provider used to open-code three near-identical blocks — a -/// page-count cap, a mid-page clamp, and a post-page hard stop — which is how -/// the same off-by-a-page bug (#3304) ended up in five providers and was missed -/// in a sixth. Now that every provider rides [`run_sync`], this type and its -/// page math are **orchestrator-private**: the cap rule lives in exactly one -/// place — construct it from `ctx.max_items`, derive the page cap, clamp each -/// batch before ingest, and stop once the cap is reached. -/// -/// `None` cap means "no item limit beyond the provider's own internal page -/// ceiling" (e.g. after the user clicks "All In"). -#[derive(Debug, Clone, Copy)] -struct ItemCap { - cap: Option, - persisted: usize, -} - -impl ItemCap { - /// Build from a source's `max_items` value (`None` = uncapped). - fn new(max_items: Option) -> Self { - Self { - cap: max_items.map(|n| n as usize), - persisted: 0, - } - } - - /// The page ceiling to actually fetch: the smaller of the provider's own - /// `fallback` (e.g. `MAX_PAGES_PER_SYNC`) and the pages needed to cover the - /// cap. Uncapped → `fallback` unchanged. - fn max_pages(&self, page_size: u32, fallback: u32) -> u32 { - match self.cap { - Some(cap) => pages_for_max_items(cap as u32, page_size).min(fallback), - None => fallback, - } - } - - /// How many more items may still be persisted. `None` = unlimited. - fn remaining(&self) -> Option { - self.cap.map(|cap| cap.saturating_sub(self.persisted)) - } - - /// True once the cap is set and reached — callers `break` their pagination. - fn is_reached(&self) -> bool { - matches!(self.remaining(), Some(0)) - } - - /// Record `n` newly-persisted items against the budget. - fn record(&mut self, n: usize) { - self.persisted = self.persisted.saturating_add(n); - } - - /// Truncate a to-ingest batch down to the remaining budget, so a single - /// page larger than the cap never over-persists. No-op when uncapped. - fn clamp_batch(&self, batch: &mut Vec) { - if let Some(remaining) = self.remaining() { - if batch.len() > remaining { - batch.truncate(remaining); - } - } - } -} - -/// A unit of work to iterate within one sync pass. -/// -/// Flat providers return a single [`SyncScope::flat`]; nested providers return -/// one scope per container (workspace / channel). The orchestrator treats both -/// uniformly — that uniformity is the whole point of the abstraction. -#[derive(Debug, Clone)] -pub(crate) struct SyncScope { - /// Provider-native scope id (e.g. a ClickUp workspace id or a Slack - /// channel id). Empty string denotes the single implicit scope of a flat - /// provider. - pub id: String, - /// Human-readable label for logs. Never logged at a level that would leak - /// PII — ids/labels here are container identifiers, not user content. - pub label: String, -} - -impl SyncScope { - /// The single implicit scope of a flat provider (gmail/github/notion/linear). - pub(crate) fn flat() -> Self { - Self { - id: String::new(), - label: "".to_string(), - } - } - - /// One container of a nested provider (clickup workspace, slack channel). - pub(crate) fn nested(id: impl Into, label: impl Into) -> Self { - Self { - id: id.into(), - label: label.into(), - } - } -} - -/// One page fetched from a provider for a given [`SyncScope`]. -pub(crate) struct PageFetch { - /// Raw upstream items, already unwrapped from the Composio envelope by the - /// provider's [`IncrementalSource::fetch_page`]. - pub items: Vec, - /// Opaque next-page cursor for *this scope*, or `None` when this was the - /// last page. The orchestrator never interprets it — it is a Notion - /// `start_cursor`, a ClickUp page index, etc., round-tripped back into the - /// next `fetch_page` call. - pub next: Option, -} - -/// One item that survived dedup + the depth window + the `max_items` clamp and -/// is queued for ingest. -pub(crate) struct SyncItem { - /// Stable dedup key (e.g. `{page_id}@{edited_time}`). Marked synced on a - /// successful ingest so the next pass skips it. - pub dedup_key: String, - /// Sort timestamp used for cursor advancement and depth-window compares. - /// Same representation as [`IncrementalSource::depth_floor`]. - pub sort_ts: Option, - /// Raw upstream payload, handed to [`IncrementalSource::ingest`]. - pub raw: Value, -} - -/// Folded result of one [`IncrementalSource::ingest`] call. Every field is -/// order-independent so a concurrent ingest stage can accumulate into it. -#[derive(Default)] -pub(crate) struct IngestOutcome { - /// Dedup keys whose ingest succeeded — the orchestrator marks each synced. - pub synced_keys: Vec, - /// Number of items persisted (equals `synced_keys.len()`). - pub persisted: usize, - /// Whether any per-item ingest failed. When true and the source opts into - /// [`IncrementalSource::hold_cursor_on_ingest_failure`], the orchestrator - /// holds the cursor so the failed range is re-fetched next pass. - pub had_failures: bool, -} - -/// The slim primitive a Composio provider implements to ride [`run_sync`]. -/// -/// Implementations own *only* the provider-specific shapes — which actions to -/// call, how to read ids/timestamps, how to persist. The orchestrator owns all -/// the control flow (budget, pagination bound, dedup, depth window, cap, cursor -/// advance, state persistence). -#[async_trait] -pub(crate) trait IncrementalSource: Send + Sync { - /// Toolkit slug — used for [`SyncState`] keying and log prefixes. - fn toolkit(&self) -> &'static str; - - /// Page size to request this pass. Providers typically widen this for - /// [`SyncReason::ConnectionCreated`] to backfill faster. - fn page_size(&self, reason: SyncReason) -> u32; - - /// The provider's own internal page ceiling — the `fallback` handed to - /// [`ItemCap::max_pages`], applied *per scope*. - fn max_pages(&self) -> u32; - - /// The page ceiling for *this pass*, given the reason and current state. - /// Defaults to the fixed [`Self::max_pages`]; gmail overrides it to apply an - /// adaptive cap (e.g. only a couple of pages when the previous sync ran very - /// recently). The orchestrator hands the result to [`ItemCap::max_pages`]. - fn page_ceiling(&self, reason: SyncReason, state: &SyncState) -> u32 { - let _ = (reason, state); - self.max_pages() - } - - /// Whether to stop paginating a scope once a fetched page yields **no new - /// items** (everything deduped away as already-synced). Default `false` — - /// the orchestrator keeps paginating until the cursor/depth boundary or the - /// last page. Gmail overrides to `true`: its result set is ordered so an - /// all-already-synced page means there is nothing newer left to fetch. - fn stop_on_empty_pending(&self) -> bool { - false - } - - /// Resolve identity and list the scopes to iterate. - /// - /// Flat providers return `vec![SyncScope::flat()]`. Nested providers call - /// their "list workspaces / channels" action(s) here and return one scope - /// each (recording any budget spent via `state`). Returning an empty vec - /// short-circuits the pass to a no-op outcome. - async fn preamble( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String>; - - /// Fetch one page of raw items for `scope` at `cursor` (`None` = first - /// page). Return the items already unwrapped from the Composio envelope - /// plus the opaque next-page token. - /// - /// Implementations **must record the page request against the daily budget** - /// (`state.record_requests(1)`) for any *completed* round-trip — including - /// one the provider reports as `successful == false` — before converting it - /// to an `Err`, so a broken/unauthorized connection cannot make unlimited - /// billable failed calls without hitting the per-day cap. A transport error - /// (no round-trip) must not be recorded. - async fn fetch_page( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - reason: SyncReason, - state: &mut SyncState, - ) -> Result; - - /// Stable dedup key for one raw item. `None` drops the item (e.g. it has no - /// extractable id). - fn item_dedup_key(&self, item: &Value) -> Option; - - /// Sort timestamp for one raw item — compared against the persistent cursor - /// and the depth floor. `None` means "no timestamp" (never trips the cursor - /// boundary or the depth window). - fn item_sort_ts(&self, item: &Value) -> Option; - - /// Build the `sync_depth_days` floor in the *same representation* as - /// [`Self::item_sort_ts`] so the lexicographic compare is valid. Default is - /// RFC3339 UTC; providers whose timestamps are epoch-millis strings - /// (clickup) override. Unused when [`Self::server_side_depth`] is `true`. - fn depth_floor(&self, days: u32) -> String { - let floor = chrono::Utc::now() - chrono::Duration::days(days as i64); - floor.to_rfc3339() - } - - /// Noun used for this provider's `{noun}_fetched` / `{noun}_persisted` - /// keys in the [`SyncOutcome::details`] diagnostic blob, preserving each - /// provider's historical detail shape (notion: `results`, github/linear: - /// `issues`, clickup: `tasks`, …). `details` is for logging/UI status only; - /// nothing reads these keys in production. - fn detail_noun(&self) -> &'static str { - "results" - } - - /// Whether the provider applies the `sync_depth_days` window **itself** - /// (server-side — e.g. GitHub's `updated:>{date}` search qualifier), - /// rather than relying on the orchestrator's client-side timestamp - /// truncation. When `true`, the orchestrator skips its client-side depth - /// filter and the provider must inject the window inside - /// [`Self::fetch_page`] (typically only on the first sync, before a cursor - /// exists). Default `false` — the orchestrator filters client-side via - /// [`Self::depth_floor`]. - fn server_side_depth(&self) -> bool { - false - } - - /// Whether to hold (not advance) the cursor when an ingest reported a - /// failure this pass. Default `true` — Notion's safe behaviour under the - /// delete-first memory-tree pipeline (#2885), where an edited item that - /// fails to re-ingest must be re-fetched. Providers that advance regardless - /// of per-item failures (clickup) override to `false`. - fn hold_cursor_on_ingest_failure(&self) -> bool { - true - } - - /// Whether this source keeps a **cursor per scope** (Slack: one `oldest` - /// watermark per channel) instead of the single global watermark every - /// other provider uses. Default `false`. - /// - /// When `true` the orchestrator: - /// * disables its global-cursor boundary detection in [`select_pending`] - /// (the scope's watermark is read by the provider's own - /// [`Self::fetch_page`] — usually injected as a server-side `oldest` - /// filter, so `server_side_depth` is typically `true` too); - /// * tracks the newest timestamp **per scope** and, after each scope - /// finishes cleanly, calls [`Self::advance_scope_cursor`] to advance - /// that one scope's watermark, persisting state between scopes; - /// * holds a scope's watermark (does not advance it) when that scope hit - /// an ingest failure or the `max_items` cap truncated it — so the next - /// pass re-fetches that scope's unseen/failed tail. - /// - /// Single-cursor providers leave this `false` and keep the global - /// advance-once-at-the-end path verbatim. - fn per_scope_cursors(&self) -> bool { - false - } - - /// Whether a scope-level failure is **non-fatal**. Default `false` — a - /// `fetch_page` error aborts the whole pass (every single-cursor provider's - /// behaviour). When `true` (Slack), a `fetch_page` error or a scope's - /// ingest failure is logged, counted in `scopes_errored`, and the - /// orchestrator continues to the next scope instead of returning `Err`. - fn tolerate_scope_errors(&self) -> bool { - false - } - - /// Advance the persisted watermark for a single `scope` to `newest_ts` - /// (per-scope-cursor mode only). The default is a no-op; Slack overrides it - /// to fold the new watermark into its per-channel cursor map serialized in - /// [`SyncState::cursor`]. Never called when [`Self::per_scope_cursors`] is - /// `false`. - fn advance_scope_cursor(&self, _state: &mut SyncState, _scope: &SyncScope, _newest_ts: &str) {} - - /// Persist a batch of already-filtered items. May spend budget via `state` - /// (e.g. Notion's per-page body fetch). Returns which dedup keys succeeded - /// so the orchestrator can mark them synced. - async fn ingest( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - state: &mut SyncState, - items: Vec, - ) -> IngestOutcome; -} - -/// Current wall-clock time in milliseconds since the UNIX epoch. -fn now_ms() -> u64 { - use std::time::{SystemTime, UNIX_EPOCH}; - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - -/// Build the "nothing to do" outcome used by the budget-exhausted and -/// empty-scopes early returns. -fn skipped_outcome( - toolkit: &str, - connection_id: &str, - reason: SyncReason, - started_at_ms: u64, - why: &str, - details: Value, -) -> SyncOutcome { - SyncOutcome { - toolkit: toolkit.to_string(), - connection_id: Some(connection_id.to_string()), - reason: reason.as_str().to_string(), - items_ingested: 0, - started_at_ms, - finished_at_ms: now_ms(), - summary: format!("{toolkit} sync skipped: {why}"), - details, - } -} - -/// Pure (no I/O) per-page scan: extract dedup key + sort timestamp for each raw -/// item, advance `newest_ts`, detect the persistent-cursor boundary, and drop -/// already-synced items. Returns the survivors plus whether we crossed the -/// cursor boundary (the signal to stop paginating this scope). -/// -/// This is the generic form of every provider's old `select_pending`. All -/// order-dependent decisions live here so the (possibly concurrent) ingest -/// stage never has to reason about ordering. -/// -/// `boundary_cursor` is the watermark used for cursor-boundary detection. -/// Single-cursor providers pass `state.cursor`; per-scope-cursor providers -/// (Slack) pass `None` to disable the global-cursor boundary entirely — their -/// per-scope watermark is enforced server-side inside `fetch_page` instead. -fn select_pending( - source: &S, - items: &[Value], - boundary_cursor: Option<&str>, - state: &SyncState, - newest_ts: &mut Option, -) -> (Vec, bool) { - let mut hit_cursor_boundary = false; - let mut pending: Vec = Vec::new(); - for item in items { - let Some(dedup_key) = source.item_dedup_key(item) else { - tracing::debug!( - toolkit = source.toolkit(), - "[composio:sync_orch] item missing dedup key, skipping" - ); - continue; - }; - - let sort_ts = source.item_sort_ts(item); - - // Track the newest timestamp for cursor advancement — for *every* item - // with a timestamp, including ones we skip as already-synced. - if let Some(ref ts) = sort_ts { - if newest_ts.as_ref().is_none_or(|existing| ts > existing) { - *newest_ts = Some(ts.clone()); - } - } - - // Older-or-equal to the cursor AND already synced → we have caught up. - if let (Some(cursor), Some(ts)) = (boundary_cursor, &sort_ts) { - if ts.as_str() <= cursor && state.is_synced(&dedup_key) { - hit_cursor_boundary = true; - continue; - } - } - - if state.is_synced(&dedup_key) { - continue; - } - - pending.push(SyncItem { - dedup_key, - sort_ts, - raw: item.clone(), - }); - } - (pending, hit_cursor_boundary) -} - -/// Run one incremental sync end-to-end through the generic orchestrator. -/// -/// The provider supplies the [`IncrementalSource`] primitives; everything else -/// — budget, the per-scope page loop bounded by [`ItemCap::max_pages`], dedup, -/// the `sync_depth_days` window, the precise `max_items` clamp, cursor -/// advance/hold, state persistence, and the [`SyncOutcome`] — is owned here. -pub(crate) async fn run_sync( - source: &S, - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - let toolkit = source.toolkit(); - let started_at_ms = now_ms(); - let connection_id = ctx - .connection_id - .clone() - .unwrap_or_else(|| "default".to_string()); - - tracing::info!( - toolkit, - connection_id = %connection_id, - reason = reason.as_str(), - "[composio:sync_orch] incremental sync starting" - ); - - // ── Step 1: load persistent sync state ────────────────────────────── - let Some(memory) = ctx.memory_client() else { - return Err(format!("[composio:{toolkit}] memory client not ready")); - }; - let mut state = SyncState::load(&memory, toolkit, &connection_id).await?; - - // ── Step 2: daily budget pre-check ────────────────────────────────── - if state.budget_exhausted() { - tracing::info!( - toolkit, - connection_id = %connection_id, - "[composio:sync_orch] daily request budget exhausted, skipping sync" - ); - return Ok(skipped_outcome( - toolkit, - &connection_id, - reason, - started_at_ms, - "daily budget exhausted", - json!({ "budget_exhausted": true }), - )); - } - - // ── Step 3: preamble — resolve identity + scopes ──────────────────── - let scopes = match source.preamble(ctx, &mut state).await { - Ok(scopes) => scopes, - Err(e) => { - // Persist any budget spent during the preamble before propagating. - let _ = state.save(&memory).await; - return Err(e); - } - }; - - if scopes.is_empty() { - tracing::info!( - toolkit, - connection_id = %connection_id, - "[composio:sync_orch] no scopes to sync" - ); - state.save(&memory).await?; - return Ok(skipped_outcome( - toolkit, - &connection_id, - reason, - started_at_ms, - "no scopes to sync", - json!({ "scopes": 0 }), - )); - } - - // ── Step 4: caps + window ─────────────────────────────────────────── - let page_size = source.page_size(reason); - let mut cap = ItemCap::new(ctx.max_items); - // `page_ceiling` lets a provider apply a state-aware cap (gmail's adaptive - // cap); it defaults to the fixed `max_pages`. - let page_ceiling = source.page_ceiling(reason, &state); - let effective_max_pages = cap.max_pages(page_size, page_ceiling); - if ctx.max_items.is_some() && effective_max_pages < page_ceiling { - tracing::debug!( - toolkit, - connection_id = %connection_id, - max_items = ?ctx.max_items, - effective_max_pages, - "[composio:sync_orch] [memory_sync] applying max_items page cap" - ); - } - - // Server-side-depth providers (GitHub) inject the window into the request - // in `fetch_page`, so the orchestrator skips its client-side floor for them. - let depth_floor: Option = if source.server_side_depth() { - None - } else { - ctx.sync_depth_days.map(|days| { - let floor = source.depth_floor(days); - tracing::debug!( - toolkit, - connection_id = %connection_id, - sync_depth_days = days, - oldest_allowed = %floor, - "[composio:sync_orch] [memory_sync] applying sync_depth_days floor" - ); - floor - }) - }; - - // ── Step 5: scope × page loop ─────────────────────────────────────── - // - // `per_scope` providers (Slack) advance a watermark per scope inside the - // loop and hold it on a cap-truncated / failed scope; single-cursor - // providers advance one global watermark once, in Step 6. `tolerate` - // providers continue past a scope-level failure instead of aborting. - let per_scope = source.per_scope_cursors(); - let tolerate = source.tolerate_scope_errors(); - - let mut total_fetched: usize = 0; - let mut total_persisted: usize = 0; - let mut newest_ts: Option = None; - let mut had_ingest_failures = false; - let mut hit_cap_boundary = false; - let mut scopes_synced: usize = 0; - let mut scopes_errored: usize = 0; - - 'scopes: for scope in &scopes { - // The page cursor is per-scope — reset at the top of every scope. - let mut cursor: Option = None; - // Newest timestamp observed *within this scope*, for per-scope advance. - let mut scope_newest_ts: Option = None; - // This scope hit an ingest failure (hold its watermark, count errored). - let mut scope_had_failure = false; - // This scope was truncated by the global cap (hold its watermark). - let mut scope_hit_cap = false; - - for page_num in 0..effective_max_pages { - if state.budget_exhausted() { - tracing::info!( - toolkit, - scope = %scope.label, - page = page_num, - "[composio:sync_orch] budget exhausted mid-sync, stopping" - ); - break 'scopes; - } - - // `fetch_page` records the page request against the budget (incl. - // provider-reported failures) per its contract. On error we persist - // whatever budget/dedup progress we have before propagating — - // parity with the per-provider loops, which saved state before - // returning a failed-page error. When the source tolerates - // scope-level errors (Slack), we instead log, count, persist, and - // move on to the next scope. - let fetch = match source - .fetch_page(ctx, scope, cursor.as_deref(), reason, &mut state) - .await - { - Ok(fetch) => fetch, - Err(e) => { - if tolerate { - tracing::warn!( - toolkit, - connection_id = %connection_id, - scope = %scope.label, - page = page_num, - error = %e, - "[composio:sync_orch] scope fetch failed (continuing with next scope)" - ); - scopes_errored += 1; - let _ = state.save(&memory).await; - continue 'scopes; - } - let _ = state.save(&memory).await; - return Err(e); - } - }; - total_fetched += fetch.items.len(); - - if fetch.items.is_empty() { - tracing::debug!( - toolkit, - scope = %scope.label, - page = page_num, - "[composio:sync_orch] empty page, moving on" - ); - break; - } - - // Dedup + cursor-boundary detection + newest-ts tracking. Per-scope - // providers disable the global-cursor boundary (pass `None`) and - // accumulate the newest ts into their per-scope tracker. - let boundary_cursor = if per_scope { - None - } else { - state.cursor.as_deref() - }; - let ts_acc = if per_scope { - &mut scope_newest_ts - } else { - &mut newest_ts - }; - let (mut pending, mut hit_cursor_boundary) = - select_pending(source, &fetch.items, boundary_cursor, &state, ts_acc); - - // A non-empty page whose items all deduped away as already-synced - // (the gmail "all already synced" signal — captured before the - // depth/cap filters narrow `pending` for other reasons). - let page_had_no_new_items = pending.is_empty(); - - // sync_depth_days: `pending` is in descending-timestamp order, so - // truncate at the first item below the floor and stop paginating. - if let Some(ref floor) = depth_floor { - if let Some(cut) = pending.iter().position(|it| { - it.sort_ts - .as_deref() - .map(|t| t < floor.as_str()) - .unwrap_or(false) - }) { - pending.truncate(cut); - hit_cursor_boundary = true; - } - } - - // max_items: clamp the dedup'd batch to the remaining budget BEFORE - // ingest — the precise cap that fixes the page-granular #3304 bug. - cap.clamp_batch(&mut pending); - - // Provider-specific persistence (may spend budget, e.g. body fetch). - let outcome = source.ingest(ctx, scope, &mut state, pending).await; - for key in &outcome.synced_keys { - state.mark_synced(key); - } - total_persisted += outcome.persisted; - cap.record(outcome.persisted); - if outcome.had_failures { - had_ingest_failures = true; - scope_had_failure = true; - } - - // Precise cap reached → stop the entire pass (after settling this - // scope's watermark below). - if cap.is_reached() { - hit_cap_boundary = true; - scope_hit_cap = true; - break; - } - - if hit_cursor_boundary { - tracing::debug!( - toolkit, - scope = %scope.label, - page = page_num, - "[composio:sync_orch] reached cursor/depth boundary, stopping scope" - ); - break; - } - - if page_had_no_new_items && source.stop_on_empty_pending() { - tracing::debug!( - toolkit, - scope = %scope.label, - page = page_num, - "[composio:sync_orch] page had no new items, stopping scope" - ); - break; - } - - cursor = fetch.next; - if cursor.is_none() { - tracing::debug!( - toolkit, - scope = %scope.label, - page = page_num, - "[composio:sync_orch] no next cursor, scope done" - ); - break; - } - } - - // ── Per-scope watermark settle (per_scope providers only) ──────── - if per_scope { - if scope_had_failure { - // Ingest failure → hold this scope's watermark, count errored. - scopes_errored += 1; - tracing::warn!( - toolkit, - connection_id = %connection_id, - scope = %scope.label, - "[composio:sync_orch] scope ingest failed; watermark held, re-fetch next pass" - ); - } else { - if scope_hit_cap { - // Cap truncated this scope → hold its watermark so the next - // pass re-scans the unseen tail; still a processed scope. - tracing::debug!( - toolkit, - scope = %scope.label, - "[composio:sync_orch] cap truncated scope; watermark held" - ); - } else if let Some(ref nt) = scope_newest_ts { - source.advance_scope_cursor(&mut state, scope, nt); - } - scopes_synced += 1; - } - // Persist between scopes for crash-resilience (parity with the - // per-provider Slack loop, which saved state after each channel). - if let Err(err) = state.save(&memory).await { - tracing::warn!( - toolkit, - error = %err, - "[composio:sync_orch] per-scope state save failed (non-fatal)" - ); - } - } - - if scope_hit_cap { - break 'scopes; - } - } - - // ── Step 6: advance cursor (or hold) and persist state ────────────── - // - // Per-scope providers already advanced/held each scope's watermark inline, - // so the global advance is skipped for them. Single-cursor providers hold - // the global cursor on a cap-truncated pass so the next sync re-scans the - // unseen tail, and on an ingest failure when the source opts in (Notion's - // delete-first safety). Otherwise advance to the newest timestamp seen. - if !per_scope { - let hold_cursor = - hit_cap_boundary || (had_ingest_failures && source.hold_cursor_on_ingest_failure()); - if !hold_cursor { - if let Some(new_cursor) = newest_ts { - state.advance_cursor(&new_cursor); - } - } else { - tracing::warn!( - toolkit, - connection_id = %connection_id, - had_ingest_failures, - hit_cap_boundary, - "[composio:sync_orch] holding cursor — cap-truncated pass or ingest failures; \ - next sync will re-fetch the unseen/failed range" - ); - } - } - state.set_last_sync_at_ms(now_ms()); - state.save(&memory).await?; - - let finished_at_ms = now_ms(); - let summary = format!( - "{toolkit} sync ({reason}): fetched {total_fetched}, persisted {total_persisted} new, \ - budget remaining {remaining}", - reason = reason.as_str(), - remaining = state.budget_remaining(), - ); - tracing::info!( - toolkit, - connection_id = %connection_id, - elapsed_ms = finished_at_ms.saturating_sub(started_at_ms), - total_fetched, - total_persisted, - budget_remaining = state.budget_remaining(), - "[composio:sync_orch] incremental sync complete" - ); - - // Provider-named `{noun}_fetched` / `{noun}_persisted` keys preserve each - // provider's historical `details` shape (notion `results`, github/linear - // `issues`, …). Built dynamically since `json!` can't take runtime keys. - let noun = source.detail_noun(); - let mut details = json!({ - "budget_remaining": state.budget_remaining(), - "cursor": state.cursor, - "synced_ids_total": state.synced_ids.len(), - }); - if let Some(obj) = details.as_object_mut() { - obj.insert(format!("{noun}_fetched"), json!(total_fetched)); - obj.insert(format!("{noun}_persisted"), json!(total_persisted)); - // Scope accounting is only meaningful for providers that iterate real - // per-scope work with tolerance/advance (Slack); emitting it for - // single-cursor providers would change their historical `details` shape. - if per_scope || tolerate { - obj.insert("scopes_total".to_string(), json!(scopes.len())); - obj.insert("scopes_synced".to_string(), json!(scopes_synced)); - obj.insert("scopes_errored".to_string(), json!(scopes_errored)); - } - } - - Ok(SyncOutcome { - toolkit: toolkit.to_string(), - connection_id: Some(connection_id), - reason: reason.as_str().to_string(), - items_ingested: total_persisted, - started_at_ms, - finished_at_ms, - summary, - details, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::config::Config; - use std::sync::Arc; - use tempfile::TempDir; - - /// A minimal in-test [`IncrementalSource`] that proves a *future* toolkit - /// inherits the cap + window for free. It synthesises items per scope (or - /// returns explicit ones) and "ingests" by counting — no real memory tree. - /// `Default` keeps the per-test literals to just the field(s) they vary. - #[derive(Default)] - struct FakeSource { - scopes: Vec, - items_per_scope: usize, - /// When set, returned verbatim for the first page of the (single) scope - /// instead of synthesised items — used by the depth-window test. - explicit_items: Option>, - /// When true, `preamble` returns an error (exercises the preamble-error - /// save-and-propagate path). - fail_preamble: bool, - /// When true, `fetch_page` returns a *transport* error (no round-trip → - /// not budget-recorded). - fail_fetch: bool, - /// When true, `fetch_page` records the round-trip *then* returns a - /// provider-reported failure — pins that a failed page still consumes - /// the daily budget. - provider_fail_fetch: bool, - /// When true, advertise server-side depth so the orchestrator skips its - /// client-side window filter (GitHub's behaviour). - server_side_depth: bool, - /// When true, keep a per-scope watermark map in `state.cursor` - /// (Slack's behaviour) instead of a single global cursor. - per_scope: bool, - /// When true, a scope-level fetch failure is non-fatal — the pass - /// continues to the next scope and records `scopes_errored`. - tolerate: bool, - /// Scope id whose `fetch_page` returns a transport error (used to - /// exercise per-scope error tolerance). `None` → every scope succeeds. - fail_scope: Option, - /// When `Some`, drive multi-page returns: page `i` returns `pages[i]` - /// with the opaque cursor = next page index. Overrides the single-page - /// synth path. - pages: Option>>, - /// When true, override `stop_on_empty_pending()` (gmail's behaviour). - stop_on_empty: bool, - /// When `Some`, override `page_ceiling()` (gmail's adaptive cap). - page_ceiling: Option, - } - - impl FakeSource { - fn flat(items_per_scope: usize) -> Self { - Self { - scopes: vec![SyncScope::flat()], - items_per_scope, - ..Default::default() - } - } - } - - #[async_trait] - impl IncrementalSource for FakeSource { - fn toolkit(&self) -> &'static str { - "faketoolkit" - } - fn page_size(&self, _reason: SyncReason) -> u32 { - 50 - } - fn max_pages(&self) -> u32 { - 20 - } - fn stop_on_empty_pending(&self) -> bool { - self.stop_on_empty - } - fn page_ceiling(&self, _reason: SyncReason, _state: &SyncState) -> u32 { - self.page_ceiling.unwrap_or_else(|| self.max_pages()) - } - async fn preamble( - &self, - _ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - if self.fail_preamble { - // Spend a request first so the save-on-error path persists it. - state.record_requests(1); - return Err("fake preamble failure".to_string()); - } - Ok(self.scopes.clone()) - } - fn per_scope_cursors(&self) -> bool { - self.per_scope - } - fn tolerate_scope_errors(&self) -> bool { - self.tolerate - } - fn advance_scope_cursor(&self, state: &mut SyncState, scope: &SyncScope, newest_ts: &str) { - // Mirror Slack: fold the watermark into a per-scope map serialized - // in `state.cursor`. - let mut map: std::collections::BTreeMap = state - .cursor - .as_deref() - .and_then(|s| serde_json::from_str(s).ok()) - .unwrap_or_default(); - map.insert(scope.id.clone(), newest_ts.to_string()); - state.cursor = Some(serde_json::to_string(&map).unwrap()); - } - async fn fetch_page( - &self, - _ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - _reason: SyncReason, - state: &mut SyncState, - ) -> Result { - if self.fail_scope.as_deref() == Some(scope.id.as_str()) { - // A scope-level transport error — exercises tolerate_scope_errors. - return Err(format!("fake scope fetch failure for {}", scope.id)); - } - if self.fail_fetch { - // Simulate a transport error (no completed round-trip) → not - // recorded, matching the contract. - return Err("fake fetch_page failure".to_string()); - } - // A completed round-trip — record it against the budget. - state.record_requests(1); - if self.provider_fail_fetch { - // Completed but the provider reported failure — already recorded. - return Err("fake provider-reported page failure".to_string()); - } - // Multi-page mode: cursor is the page index. - if let Some(pages) = &self.pages { - let idx: usize = cursor.and_then(|c| c.parse().ok()).unwrap_or(0); - let items = pages.get(idx).cloned().unwrap_or_default(); - let next = (idx + 1 < pages.len()).then(|| (idx + 1).to_string()); - return Ok(PageFetch { items, next }); - } - // Single page per scope: everything comes back on the first call. - if cursor.is_some() { - return Ok(PageFetch { - items: vec![], - next: None, - }); - } - if let Some(items) = &self.explicit_items { - return Ok(PageFetch { - items: items.clone(), - next: None, - }); - } - let items = (0..self.items_per_scope) - .map(|i| { - json!({ - "id": format!("{}-{i}", scope.id), - "ts": "2099-01-01T00:00:00Z" - }) - }) - .collect(); - Ok(PageFetch { items, next: None }) - } - fn item_dedup_key(&self, item: &Value) -> Option { - item.get("id").and_then(Value::as_str).map(str::to_string) - } - fn item_sort_ts(&self, item: &Value) -> Option { - item.get("ts").and_then(Value::as_str).map(str::to_string) - } - fn server_side_depth(&self) -> bool { - self.server_side_depth - } - async fn ingest( - &self, - _ctx: &ProviderContext, - _scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let synced_keys: Vec = items.into_iter().map(|it| it.dedup_key).collect(); - let persisted = synced_keys.len(); - IngestOutcome { - synced_keys, - persisted, - had_failures: false, - } - } - } - - fn fake_ctx( - tmp: &TempDir, - max_items: Option, - sync_depth_days: Option, - ) -> ProviderContext { - let mut config = Config { - config_path: tmp.path().join("config.toml"), - workspace_dir: tmp.path().join("workspace"), - ..Config::default() - }; - config.secrets.encrypt = false; - ProviderContext { - config: Arc::new(config), - toolkit: "faketoolkit".to_string(), - connection_id: Some("conn-fake".to_string()), - usage: Default::default(), - max_items, - sync_depth_days, - } - } - - #[tokio::test] - async fn max_items_caps_ingest_to_exact_count_not_page_granular() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, Some(2), None); - // One page returns 5 items; the cap is 2. - let outcome = run_sync(&FakeSource::flat(5), &ctx, SyncReason::ConnectionCreated) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 2, - "max_items=2 must clamp a 5-item page to EXACTLY 2 (the #3304 fix)" - ); - } - - #[tokio::test] - async fn no_cap_ingests_the_full_page() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let outcome = run_sync(&FakeSource::flat(5), &ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 5, - "with no cap every valid page item is ingested" - ); - } - - #[tokio::test] - async fn sync_depth_days_filters_items_below_the_floor() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, Some(7)); - // Descending timestamp order: two recent (far future), three ancient. - // With a 7-day floor only the two recent items survive. - let items = vec![ - json!({ "id": "a", "ts": "2099-01-02T00:00:00Z" }), - json!({ "id": "b", "ts": "2099-01-01T00:00:00Z" }), - json!({ "id": "c", "ts": "2000-01-03T00:00:00Z" }), - json!({ "id": "d", "ts": "2000-01-02T00:00:00Z" }), - json!({ "id": "e", "ts": "2000-01-01T00:00:00Z" }), - ]; - let source = FakeSource { - scopes: vec![SyncScope::flat()], - explicit_items: Some(items), - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Manual) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 2, - "sync_depth_days=7 must drop the three ancient items" - ); - } - - #[tokio::test] - async fn server_side_depth_skips_the_client_side_filter() { - // Same ancient items, but the source advertises server-side depth — so - // the orchestrator must NOT client-side-truncate (the provider would - // have filtered in fetch_page). All five survive here. - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, Some(7)); - let items = vec![ - json!({ "id": "a", "ts": "2099-01-02T00:00:00Z" }), - json!({ "id": "b", "ts": "2000-01-01T00:00:00Z" }), - json!({ "id": "c", "ts": "2000-01-02T00:00:00Z" }), - ]; - let source = FakeSource { - scopes: vec![SyncScope::flat()], - explicit_items: Some(items), - server_side_depth: true, - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Manual) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 3, - "server_side_depth must skip the orchestrator's client-side window filter" - ); - } - - /// Two pages of items, keyed `id`+`ts`. - fn page(prefix: &str, n: usize) -> Vec { - (0..n) - .map(|i| json!({ "id": format!("{prefix}{i}"), "ts": "2099-01-01T00:00:00Z" })) - .collect() - } - - #[tokio::test] - async fn page_ceiling_caps_the_pages_fetched() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - // Two pages of 2 items each. With a page ceiling of 1 (gmail's adaptive - // cap), only the first page is fetched → 2 ingested, not 4. - let source = FakeSource { - scopes: vec![SyncScope::flat()], - pages: Some(vec![page("a", 2), page("b", 2)]), - page_ceiling: Some(1), - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 2, - "page_ceiling=1 must fetch only the first page" - ); - } - - #[tokio::test] - async fn stop_on_empty_pending_halts_at_an_all_synced_page() { - let tmp = TempDir::new().unwrap(); - // Page 0 = [a0,a1] (new), page 1 = same ids (all synced after page 0), - // page 2 = [c0,c1] (new). With stop_on_empty the pass halts at page 1 - // and never reaches page 2 → 2 ingested; without it, page 2's items are - // also ingested → 4. - let pages = vec![page("a", 2), page("a", 2), page("c", 2)]; - - let stop = FakeSource { - scopes: vec![SyncScope::flat()], - pages: Some(pages.clone()), - stop_on_empty: true, - ..Default::default() - }; - let stop_ctx = fake_ctx(&tmp, None, None); - let stopped = run_sync(&stop, &stop_ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!( - stopped.items_ingested, 2, - "stop_on_empty must halt at the all-synced page before page 2" - ); - - // Control: same pages, no stop_on_empty → keeps going to page 2. - let tmp2 = TempDir::new().unwrap(); - let keep = FakeSource { - scopes: vec![SyncScope::flat()], - pages: Some(pages), - ..Default::default() - }; - let keep_ctx = fake_ctx(&tmp2, None, None); - let kept = run_sync(&keep, &keep_ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!( - kept.items_ingested, 4, - "without stop_on_empty the pass continues past the all-synced page" - ); - } - - #[tokio::test] - async fn nested_scopes_share_one_cap_budget() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, Some(4), None); - // Two scopes, 3 items each (6 total); the cap is 4 → 3 from scope one, - // 1 from scope two, then the pass stops. Proves the cap spans scopes. - let source = FakeSource { - scopes: vec![ - SyncScope::nested("s1", "Scope 1"), - SyncScope::nested("s2", "Scope 2"), - ], - items_per_scope: 3, - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::ConnectionCreated) - .await - .expect("run_sync"); - assert_eq!( - outcome.items_ingested, 4, - "max_items must cap the combined ingest across nested scopes" - ); - } - - #[tokio::test] - async fn budget_exhausted_short_circuits_to_a_skip_outcome() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - // Drain the daily budget before the run so the pre-check trips. - { - let memory = ctx.memory_client().expect("memory client"); - let mut state = SyncState::load(&memory, "faketoolkit", "conn-fake") - .await - .unwrap(); - state.record_requests(state.budget_remaining()); - state.save(&memory).await.unwrap(); - } - let outcome = run_sync(&FakeSource::flat(5), &ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!(outcome.items_ingested, 0); - assert!( - outcome.summary.contains("budget"), - "exhausted-budget run must report a skip, got: {}", - outcome.summary - ); - } - - #[tokio::test] - async fn empty_scopes_short_circuit_to_a_skip_outcome() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![], // preamble resolved no scopes to iterate - items_per_scope: 5, - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!(outcome.items_ingested, 0); - assert!( - outcome.summary.contains("no scopes"), - "empty scopes must report a skip, got: {}", - outcome.summary - ); - } - - #[tokio::test] - async fn preamble_error_propagates() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![SyncScope::flat()], - items_per_scope: 5, - fail_preamble: true, - ..Default::default() - }; - let err = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect_err("preamble failure must propagate"); - assert!(err.contains("preamble"), "got: {err}"); - } - - #[tokio::test] - async fn fetch_page_error_propagates() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![SyncScope::flat()], - items_per_scope: 5, - fail_fetch: true, - ..Default::default() - }; - let err = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect_err("fetch_page failure must propagate"); - assert!(err.contains("fetch_page"), "got: {err}"); - } - - #[tokio::test] - async fn provider_reported_page_failure_still_consumes_budget() { - // Parity with the per-provider loops (and the Codex review): a page that - // completes the round-trip but reports `successful == false` must count - // against the daily budget before the error propagates, so a broken - // connection can't make unlimited billable failed page calls. - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![SyncScope::flat()], - items_per_scope: 5, - provider_fail_fetch: true, - ..Default::default() - }; - let before = { - let memory = ctx.memory_client().expect("memory client"); - SyncState::load(&memory, "faketoolkit", "conn-fake") - .await - .unwrap() - .budget_remaining() - }; - let err = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect_err("provider-reported failure must propagate"); - assert!(err.contains("provider-reported"), "got: {err}"); - // The orchestrator saved state on the page error; the failed page was - // recorded, so exactly one request was consumed. - let memory = ctx.memory_client().expect("memory client"); - let after = SyncState::load(&memory, "faketoolkit", "conn-fake") - .await - .unwrap() - .budget_remaining(); - assert_eq!( - before - after, - 1, - "a completed-but-failed page must consume exactly one budget request" - ); - } - - #[test] - fn select_pending_tracks_newest_skips_synced_and_detects_boundary() { - let source = FakeSource::flat(0); - let mut state = SyncState::new("faketoolkit", "conn1"); - state.cursor = Some("2026-04-15T00:00:00Z".to_string()); - // Item B is already synced and older than the cursor. - state.mark_synced("b"); - - let items = vec![ - json!({ "id": "a", "ts": "2026-05-01T00:00:00Z" }), - json!({ "id": "b", "ts": "2026-04-01T00:00:00Z" }), - json!({ "ts": "2026-03-01T00:00:00Z" }), // no id → dropped - ]; - - let mut newest: Option = None; - let (pending, hit_boundary) = select_pending( - &source, - &items, - state.cursor.as_deref(), - &state, - &mut newest, - ); - - assert_eq!(pending.len(), 1, "only the new item A survives"); - assert_eq!(pending[0].dedup_key, "a"); - assert!( - hit_boundary, - "older synced item B trips the cursor boundary" - ); - assert_eq!(newest.as_deref(), Some("2026-05-01T00:00:00Z")); - } - - #[tokio::test] - async fn per_scope_cursors_advance_each_scope_independently() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![ - SyncScope::nested("s1", "Scope 1"), - SyncScope::nested("s2", "Scope 2"), - ], - items_per_scope: 3, - per_scope: true, - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect("run_sync"); - assert_eq!(outcome.items_ingested, 6, "both scopes' items ingested"); - assert_eq!(outcome.details["scopes_synced"], 2); - assert_eq!(outcome.details["scopes_errored"], 0); - - // The persisted cursor is a per-scope watermark map carrying *both* - // scopes — proof the orchestrator advanced each scope independently. - let memory = ctx.memory_client().expect("memory client"); - let state = SyncState::load(&memory, "faketoolkit", "conn-fake") - .await - .unwrap(); - let map: std::collections::BTreeMap = - serde_json::from_str(state.cursor.as_deref().unwrap()).expect("cursor map"); - assert_eq!(map.len(), 2, "both scopes have a watermark"); - assert!(map.contains_key("s1") && map.contains_key("s2")); - } - - #[tokio::test] - async fn tolerate_scope_errors_continues_past_a_failed_scope() { - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - // Scope s1 fails its fetch; s2 succeeds. With tolerance the pass - // ingests s2's items, counts s1 errored, and never returns Err. - let source = FakeSource { - scopes: vec![ - SyncScope::nested("s1", "Scope 1"), - SyncScope::nested("s2", "Scope 2"), - ], - items_per_scope: 3, - per_scope: true, - tolerate: true, - fail_scope: Some("s1".to_string()), - ..Default::default() - }; - let outcome = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect("tolerant run_sync must not error"); - assert_eq!( - outcome.items_ingested, 3, - "only the healthy scope's items are ingested" - ); - assert_eq!(outcome.details["scopes_errored"], 1); - assert_eq!(outcome.details["scopes_synced"], 1); - - // The failed scope advanced no watermark; only the healthy one did. - let memory = ctx.memory_client().expect("memory client"); - let state = SyncState::load(&memory, "faketoolkit", "conn-fake") - .await - .unwrap(); - let map: std::collections::BTreeMap = - serde_json::from_str(state.cursor.as_deref().unwrap_or("{}")).expect("cursor map"); - assert!(map.contains_key("s2"), "healthy scope advanced"); - assert!(!map.contains_key("s1"), "failed scope watermark held"); - } - - #[tokio::test] - async fn per_scope_fatal_when_tolerance_off() { - // Same failing scope, but tolerance off → the fetch error aborts the - // whole pass (single-cursor providers' default behaviour). - let tmp = TempDir::new().unwrap(); - let ctx = fake_ctx(&tmp, None, None); - let source = FakeSource { - scopes: vec![SyncScope::nested("s1", "Scope 1")], - items_per_scope: 3, - per_scope: true, - tolerate: false, - fail_scope: Some("s1".to_string()), - ..Default::default() - }; - let err = run_sync(&source, &ctx, SyncReason::Periodic) - .await - .expect_err("without tolerance a scope fetch error propagates"); - assert!(err.contains("scope fetch failure"), "got: {err}"); - } - - // ── ItemCap / pages_for_max_items (relocated from helpers.rs, #3902) ── - - #[test] - fn pages_for_max_items_rounds_up() { - assert_eq!(pages_for_max_items(100, 25), 4); - assert_eq!(pages_for_max_items(101, 25), 5); - assert_eq!(pages_for_max_items(1, 25), 1); - assert_eq!(pages_for_max_items(50, 50), 1); - assert_eq!(pages_for_max_items(51, 50), 2); - } - - #[test] - fn pages_for_max_items_zero_page_size() { - assert_eq!(pages_for_max_items(100, 0), u32::MAX); - } - - #[test] - fn item_cap_uncapped_is_never_reached() { - let mut cap = ItemCap::new(None); - assert_eq!(cap.remaining(), None); - assert!(!cap.is_reached()); - cap.record(1_000_000); - assert!(!cap.is_reached()); - assert_eq!( - cap.max_pages(25, 20), - 20, - "uncapped keeps the provider fallback" - ); - } - - #[test] - fn item_cap_tracks_remaining_and_reached() { - let mut cap = ItemCap::new(Some(3)); - assert_eq!(cap.remaining(), Some(3)); - assert!(!cap.is_reached()); - cap.record(2); - assert_eq!(cap.remaining(), Some(1)); - assert!(!cap.is_reached()); - cap.record(5); // saturates, never underflows - assert_eq!(cap.remaining(), Some(0)); - assert!(cap.is_reached()); - } - - #[test] - fn item_cap_max_pages_is_min_of_fallback_and_needed() { - // cap=2, page_size=50 → 1 page needed, well under the fallback. - assert_eq!(ItemCap::new(Some(2)).max_pages(50, 20), 1); - // cap=1000, page_size=25 → 40 pages needed, clamped to fallback 20. - assert_eq!(ItemCap::new(Some(1000)).max_pages(25, 20), 20); - } - - #[test] - fn item_cap_clamp_batch_truncates_to_remaining() { - let cap = ItemCap::new(Some(2)); - let mut batch = vec![1, 2, 3, 4, 5]; - cap.clamp_batch(&mut batch); - assert_eq!(batch, vec![1, 2]); - - // Uncapped leaves the batch untouched. - let mut full = vec![1, 2, 3]; - ItemCap::new(None).clamp_batch(&mut full); - assert_eq!(full, vec![1, 2, 3]); - - // After recording progress, clamp uses the reduced budget. - let mut cap2 = ItemCap::new(Some(5)); - cap2.record(3); - let mut batch2 = vec![1, 2, 3, 4]; - cap2.clamp_batch(&mut batch2); - assert_eq!(batch2, vec![1, 2], "only 2 of the 5 budget remained"); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/registry.rs b/src/openhuman/memory_sync/composio/providers/registry.rs index 5e3774896..e97e9aa53 100644 --- a/src/openhuman/memory_sync/composio/providers/registry.rs +++ b/src/openhuman/memory_sync/composio/providers/registry.rs @@ -94,7 +94,7 @@ pub fn init_default_providers() { mod tests { use super::*; use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, + ProviderContext, ProviderUserProfile, }; use async_trait::async_trait; @@ -113,13 +113,6 @@ mod tests { ) -> Result { Ok(ProviderUserProfile::default()) } - async fn sync( - &self, - _ctx: &ProviderContext, - _reason: SyncReason, - ) -> Result { - Ok(SyncOutcome::default()) - } } #[test] diff --git a/src/openhuman/memory_sync/composio/providers/slack/ingest.rs b/src/openhuman/memory_sync/composio/providers/slack/ingest.rs deleted file mode 100644 index 457b86dee..000000000 --- a/src/openhuman/memory_sync/composio/providers/slack/ingest.rs +++ /dev/null @@ -1,402 +0,0 @@ -//! Slack → memory tree ingest plumbing. -//! -//! Owns the conversion from a page of [`SlackMessage`]s (post-processed -//! and enriched by [`super::sync`]) into per-channel [`ChatBatch`]es and -//! drives [`memory::ingest_pipeline::ingest_chat`] per message. -//! -//! ## Source-id scope -//! -//! Source id is `slack:{connection_id}` (workspace-wide), NOT per-channel. -//! Channel label lives in [`ChatBatch.channel_label`] for display in the -//! tree; all channels in one Slack workspace accumulate into one source -//! tree so the L0 buffer fills across many ingest calls and the seal -//! cascade fires at the right cadence. -//! -//! ## Per-message ingest -//! -//! We call `ingest_chat` once per message with a single-message -//! `ChatBatch`, then `set_chunk_raw_refs` to link the resulting chunk to -//! its raw archive entry. This gives 1:1 chunk-to-raw-file mapping that -//! mirrors the Gmail per-account path. -//! -//! ## Idempotency -//! -//! Chunk IDs are content-hashed inside the memory tree, so re-ingesting -//! a previously-seen message is an UPSERT — no duplicates across syncs. - -use std::collections::BTreeMap; - -use anyhow::Result; - -use super::types::SlackMessage; -use crate::openhuman::config::Config; -use crate::openhuman::memory::ingest_pipeline::ingest_chat; -use crate::openhuman::memory::util::redact::redact; -use crate::openhuman::memory_store::chunks::store::{set_chunk_raw_refs, RawRef}; -use crate::openhuman::memory_store::content::raw::{ - self as raw_store, raw_rel_path, RawItem, RawKind, -}; -use crate::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; - -/// Platform identifier embedded in the canonical chat transcript header. -/// Matches the value `memory_tree::retrieval::source::PLATFORM_KINDS` expects. -pub const SLACK_PLATFORM: &str = "slack"; - -/// Tags attached to every Slack-ingested chunk. Stable list — retrieval -/// callers filter on these. -pub const DEFAULT_TAGS: &[&str] = &["slack", "ingested"]; - -/// Group a page of messages by `channel_id`. Each group is sorted -/// ascending by timestamp so ingest calls read chronologically. -/// -/// Analogous to Gmail's `bucket_by_participants`, but trivial for Slack: -/// every message already carries its channel_id. -pub(crate) fn bucket_by_channel<'a>( - msgs: &'a [SlackMessage], -) -> BTreeMap> { - let mut out: BTreeMap> = BTreeMap::new(); - for m in msgs { - out.entry(m.channel_id.clone()).or_default().push(m); - } - for bucket in out.values_mut() { - bucket.sort_by_key(|m| m.timestamp); - } - out -} - -/// Render a channel label for the canonical transcript header. -/// -/// Public channels become `"#eng"`; private channels become -/// `"private:ops"` so the retrieval side can distinguish them at a -/// glance. -pub(crate) fn channel_label(channel_name: &str, is_private: bool) -> String { - if is_private { - format!("private:{channel_name}") - } else { - format!("#{channel_name}") - } -} - -/// Convert a [`SlackMessage`] into a [`ChatMessage`] for the memory tree. -/// -/// Author falls back to `"unknown"` when the resolved name is empty. -/// `source_ref` prefers the HTTPS `permalink` from the Composio response; -/// when absent it falls back to the stable `slack://archives/…` scheme. -pub(crate) fn slack_message_to_chat_message(m: &SlackMessage) -> ChatMessage { - let author = if m.author.is_empty() { - "unknown".to_string() - } else { - m.author.clone() - }; - let source_ref = m - .permalink - .clone() - .or_else(|| Some(format!("slack://archives/{}/{}", m.channel_id, m.ts_raw))); - ChatMessage { - author, - timestamp: m.timestamp, - text: m.text.clone(), - source_ref, - } -} - -/// Ingest a page of Slack messages into the memory tree. -/// -/// Messages are grouped by channel_id and ingested one at a time via -/// `ingest_chat` (per-message mode). Each successful ingest links the -/// returned chunk(s) to a raw archive entry via `set_chunk_raw_refs` so -/// `read_chunk_body` can reconstruct full bodies without duplicating -/// bytes in the SQL `content` column. -/// -/// Per-channel errors are logged and swallowed — one bad message should -/// not abort the whole page (the next sync re-fetches via the -/// date-cursor). -/// -/// Returns the total number of chunks written. -pub async fn ingest_page_into_memory_tree( - config: &Config, - owner: &str, - connection_id: &str, - page_messages: &[SlackMessage], -) -> Result { - if page_messages.is_empty() { - return Ok(0); - } - - let source_id = format!("slack:{connection_id}"); - - // Best-effort raw archive — written before chunking so a chunker bug - // doesn't block capturing the source bytes. - if let Err(e) = write_raw_archive(config, &source_id, page_messages) { - log::warn!( - "[composio:slack][ingest] raw archive write failed source_id_hash={} err={:#}", - redact(&source_id), - e - ); - } - - let total_chunks = ingest_per_message(config, &source_id, owner, page_messages).await; - - log::info!( - "[composio:slack][ingest] page_done owner_hash={} connection_hash={} chunks={total_chunks}", - redact(owner), - redact(connection_id), - ); - Ok(total_chunks) -} - -/// Per-message ingest: one `ingest_chat` call per Slack message. -/// -/// Each call produces 1 chunk for normal messages or N chunks for oversize -/// messages. After the ingest we tag every resulting chunk with a -/// [`RawRef`] pointing at the raw archive file written during -/// [`write_raw_archive`], so `read_chunk_body` can reconstruct full bodies -/// without duplicating bytes in the SQL `content` column. -async fn ingest_per_message( - config: &Config, - source_id: &str, - owner: &str, - page_messages: &[SlackMessage], -) -> usize { - let mut total_chunks = 0usize; - - for m in page_messages { - if m.text.trim().is_empty() { - log::debug!( - "[composio:slack][ingest] skipping empty-body message ts_raw={}", - m.ts_raw - ); - continue; - } - - let ts_ms = m.timestamp.timestamp_millis(); - let raw_path = raw_rel_path(source_id, RawKind::Chat, ts_ms, &m.ts_raw); - - let chat_message = slack_message_to_chat_message(m); - let label = channel_label(&m.channel_name, m.is_private); - let batch = ChatBatch { - platform: SLACK_PLATFORM.to_string(), - channel_label: label, - messages: vec![chat_message], - }; - let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect(); - - match ingest_chat(config, source_id, owner, tags, batch).await { - Ok(result) => { - total_chunks += result.chunks_written; - let refs = vec![RawRef { - path: raw_path.clone(), - start: 0, - end: None, - }]; - for chunk_id in &result.chunk_ids { - if let Err(e) = set_chunk_raw_refs(config, chunk_id, &refs) { - log::warn!( - "[composio:slack][ingest] set_chunk_raw_refs failed chunk_id={} err={:#}", - chunk_id, - e - ); - } - } - log::debug!( - "[composio:slack][ingest] ingested message ts_raw={} chunks={}", - m.ts_raw, - result.chunks_written - ); - } - Err(e) => { - log::warn!( - "[composio:slack][ingest] per-message ingest_chat failed ts_raw_hash={} err={:#}", - redact(&m.ts_raw), - e - ); - } - } - } - - total_chunks -} - -/// Mirror a page of Slack messages into the on-disk raw archive. -/// -/// Files land under `/raw//chats/_.md` -/// — the `chats/` subdir is selected automatically by [`RawKind::Chat`] -/// (see `content_store::raw`). -/// Each file gets a small metadata header (channel, author, date) followed -/// by the message body so the file is self-describing when opened -/// standalone in Obsidian or a text editor. -/// -/// Messages with an empty body are skipped — they'd produce -/// zero-content files. Messages without a parseable timestamp produce -/// non-stable filenames so they are also skipped. -fn write_raw_archive(config: &Config, source_id: &str, page: &[SlackMessage]) -> Result { - let content_root = config.memory_tree_content_root(); - let mut bodies: Vec<(String, i64, String)> = Vec::with_capacity(page.len()); - - for m in page { - let body = m.text.trim(); - if body.is_empty() { - log::debug!( - "[composio:slack][raw] empty body ts_raw={} — skipping", - m.ts_raw - ); - continue; - } - - let label = channel_label(&m.channel_name, m.is_private); - let ts_ms = m.timestamp.timestamp_millis(); - let date_str = m.timestamp.to_rfc3339(); - - let mut composed = String::with_capacity(body.len() + 256); - composed.push_str(&format!("**Channel:** {label}\n")); - composed.push_str(&format!("**Author:** {}\n", m.author)); - composed.push_str(&format!("**Date:** {date_str}\n\n")); - composed.push_str(body); - - bodies.push((m.ts_raw.clone(), ts_ms, composed)); - } - - let items: Vec> = bodies - .iter() - .map(|(uid, ts, md)| RawItem { - uid: uid.as_str(), - created_at_ms: *ts, - markdown: md.as_str(), - kind: RawKind::Chat, - }) - .collect(); - - let n = raw_store::write_raw_items(&content_root, source_id, &items)?; - log::debug!( - "[composio:slack][raw] archived {n} messages source_id_hash={}", - redact(source_id) - ); - Ok(n) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - fn ts(secs: i64) -> chrono::DateTime { - chrono::Utc.timestamp_opt(secs, 0).single().unwrap() - } - - fn make_message( - channel_id: &str, - channel_name: &str, - is_private: bool, - secs: i64, - ) -> SlackMessage { - SlackMessage { - channel_id: channel_id.to_string(), - channel_name: channel_name.to_string(), - is_private, - author: "alice".to_string(), - author_id: "U001".to_string(), - text: "hello".to_string(), - timestamp: ts(secs), - ts_raw: format!("{secs}.000000"), - thread_ts: None, - permalink: None, - } - } - - // ─── bucket_by_channel ──────────────────────────────────────────────────── - - #[test] - fn bucket_by_channel_groups_messages() { - let msgs = vec![ - make_message("C1", "eng", false, 1000), - make_message("C2", "ops", false, 1100), - make_message("C1", "eng", false, 1200), - ]; - let buckets = bucket_by_channel(&msgs); - assert_eq!(buckets.len(), 2); - assert_eq!(buckets["C1"].len(), 2); - assert_eq!(buckets["C2"].len(), 1); - } - - #[test] - fn bucket_by_channel_sorts_chronologically() { - let msgs = vec![ - make_message("C1", "eng", false, 2000), - make_message("C1", "eng", false, 1000), - ]; - let buckets = bucket_by_channel(&msgs); - let eng = &buckets["C1"]; - assert_eq!(eng[0].timestamp, ts(1000)); - assert_eq!(eng[1].timestamp, ts(2000)); - } - - // ─── channel_label ──────────────────────────────────────────────────────── - - #[test] - fn channel_label_distinguishes_private() { - assert_eq!(channel_label("eng", false), "#eng"); - assert_eq!(channel_label("ops", true), "private:ops"); - } - - // ─── slack_message_to_chat_message ──────────────────────────────────────── - - #[test] - fn slack_message_to_chat_message_falls_back_to_unknown_author() { - let m = SlackMessage { - channel_id: "C1".into(), - channel_name: "eng".into(), - is_private: false, - author: "".into(), - author_id: "U001".into(), - text: "hi".into(), - timestamp: ts(1000), - ts_raw: "1000.000000".into(), - thread_ts: None, - permalink: None, - }; - let cm = slack_message_to_chat_message(&m); - assert_eq!(cm.author, "unknown"); - } - - #[test] - fn slack_message_to_chat_message_uses_permalink_when_present() { - let m = SlackMessage { - channel_id: "C1".into(), - channel_name: "eng".into(), - is_private: false, - author: "alice".into(), - author_id: "U001".into(), - text: "hi".into(), - timestamp: ts(1000), - ts_raw: "1000.000000".into(), - thread_ts: None, - permalink: Some("https://myworkspace.slack.com/archives/C1/p1000000000".into()), - }; - let cm = slack_message_to_chat_message(&m); - assert_eq!( - cm.source_ref.as_deref(), - Some("https://myworkspace.slack.com/archives/C1/p1000000000") - ); - } - - #[test] - fn slack_message_to_chat_message_falls_back_to_archive_url() { - let m = SlackMessage { - channel_id: "C1".into(), - channel_name: "eng".into(), - is_private: false, - author: "alice".into(), - author_id: "U001".into(), - text: "hi".into(), - timestamp: ts(1000), - ts_raw: "1000.000000".into(), - thread_ts: None, - permalink: None, - }; - let cm = slack_message_to_chat_message(&m); - assert_eq!( - cm.source_ref.as_deref(), - Some("slack://archives/C1/1000.000000") - ); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/slack/mod.rs b/src/openhuman/memory_sync/composio/providers/slack/mod.rs index 31c6b29ed..acb10a6f5 100644 --- a/src/openhuman/memory_sync/composio/providers/slack/mod.rs +++ b/src/openhuman/memory_sync/composio/providers/slack/mod.rs @@ -3,21 +3,15 @@ //! The provider is wired into the periodic-sync scheduler (see //! [`super::registry::init_default_providers`]) and fires //! `SLACK_LIST_CONVERSATIONS` + `SLACK_FETCH_CONVERSATION_HISTORY` -//! against the user's Composio-authorized Slack connection. Messages -//! are ingested into the memory tree via -//! [`ingest::ingest_page_into_memory_tree`] — one ingest call per message, -//! no bucketing (the memory tree's L0 seal cascade handles batching). +//! against the user's Composio-authorized Slack connection. The reusable +//! synchronization and ingestion engine is owned by tinycortex. -pub mod ingest; pub mod post_process; pub mod rpc; pub mod schemas; -pub mod sync; pub mod types; -pub mod users; mod provider; -mod source; pub use provider::{run_backfill_via_search, SlackProvider, BACKFILL_DAYS}; pub use schemas::{all_slack_memory_controller_schemas, all_slack_memory_registered_controllers}; diff --git a/src/openhuman/memory_sync/composio/providers/slack/provider.rs b/src/openhuman/memory_sync/composio/providers/slack/provider.rs index dea477e10..c6151d16e 100644 --- a/src/openhuman/memory_sync/composio/providers/slack/provider.rs +++ b/src/openhuman/memory_sync/composio/providers/slack/provider.rs @@ -6,55 +6,22 @@ //! against Composio's action catalog (`SLACK_LIST_CONVERSATIONS`, //! `SLACK_FETCH_CONVERSATION_HISTORY`, `SLACK_FETCH_TEAM_INFO`, …). //! -//! ## Per-sync lifecycle -//! -//! 1. Load [`SyncState`] for `(slack, connection_id)`. `state.cursor` is -//! a JSON-encoded [`sync::ChannelCursors`] map — Slack needs a cursor -//! per channel. Parse failures degrade to an empty map (full backfill), -//! which is safe because chunk IDs are deterministic. -//! 2. Enumerate every channel the bot can read via -//! [`ACTION_LIST_CONVERSATIONS`] with pagination. -//! 3. For each channel, pull messages since the per-channel cursor (or -//! `now - BACKFILL_DAYS` if no cursor yet) via -//! [`ACTION_FETCH_HISTORY`], paginated. -//! 4. Post-process each response via [`super::post_process`], enrich via -//! [`super::sync::extract_messages`] to produce [`SlackMessage`]s with -//! channel context and resolved user names. -//! 5. Ingest all collected messages via -//! [`super::ingest::ingest_page_into_memory_tree`] — one `ingest_chat` -//! call per message, no bucketing. -//! 6. Advance per-channel cursor to the latest successfully-ingested -//! message's timestamp; save [`SyncState`]. +//! The product provider retains profile lookup, trigger handling, and response +//! normalization. Channel enumeration, cursors, history paging, and memory +//! ingestion execute in tinycortex's Slack sync pipeline. //! //! ## Idempotency //! //! Source id is `slack:{connection_id}` — stable per workspace. Chunk -//! IDs are content-hashed, so re-ingest is an UPSERT. +//! IDs are stable, so repeated synchronization updates the same documents. -use async_trait::async_trait; -use serde_json::{json, Value}; -use std::collections::HashMap; -use std::path::PathBuf; -use std::time::Duration; - -use super::ingest::ingest_page_into_memory_tree; -use super::sync; -use super::types::{SlackChannel, SlackMessage}; -use super::users::SlackUsers; -// `ComposioClient` is no longer referenced directly — actions dispatch -// through `ProviderContext::execute` which resolves the client via the -// mode-aware factory per call (#1710). -use crate::openhuman::composio::types::ComposioExecuteResponse; -use crate::openhuman::memory_sync::composio::providers::sync_state::SyncState; use crate::openhuman::memory_sync::composio::providers::{ pick_str, resolve_sync_interval_secs, ComposioProvider, CuratedTool, ProviderContext, - ProviderUserProfile, SyncOutcome, SyncReason, + ProviderUserProfile, SyncOutcome, }; +use async_trait::async_trait; +use serde_json::{json, Value}; -/// Composio action slug for channel listing. -const ACTION_LIST_CONVERSATIONS: &str = "SLACK_LIST_CONVERSATIONS"; -/// Composio action slug for message history. -pub(super) const ACTION_FETCH_HISTORY: &str = "SLACK_FETCH_CONVERSATION_HISTORY"; /// Composio action slug for team/workspace profile fetch. const ACTION_FETCH_TEAM_INFO: &str = "SLACK_FETCH_TEAM_INFO"; /// Composio action slug for Slack `auth.test` — returns the authed @@ -69,164 +36,9 @@ const ACTION_USERS_INFO: &str = "SLACK_RETRIEVE_DETAILED_USER_INFORMATION"; /// cursor yet. pub const BACKFILL_DAYS: i64 = 6; -/// Resolve the active backfill window in days. Reads -/// `OPENHUMAN_SLACK_BACKFILL_DAYS` env var if set and parseable as a -/// positive integer; falls back to [`BACKFILL_DAYS`] otherwise. -pub(super) fn backfill_days() -> i64 { - match std::env::var("OPENHUMAN_SLACK_BACKFILL_DAYS") { - Ok(s) => match s.trim().parse::() { - Ok(n) if n >= 1 => n, - _ => { - log::warn!( - "[composio:slack] OPENHUMAN_SLACK_BACKFILL_DAYS={s:?} not a positive integer; \ - falling back to default {BACKFILL_DAYS}" - ); - BACKFILL_DAYS - } - }, - Err(_) => BACKFILL_DAYS, - } -} - -/// Max channels listed per `SLACK_LIST_CONVERSATIONS` page. -const LIST_PAGE_SIZE: u32 = 200; - -/// Max messages per `SLACK_FETCH_CONVERSATION_HISTORY` page. -pub(super) const HISTORY_PAGE_SIZE: u32 = 1000; - -/// Stop paginating any single channel's history after this many pages. -pub(super) const MAX_HISTORY_PAGES_PER_CHANNEL: u32 = 20; - -/// Stop paginating channel listings after this many pages. -const MAX_LIST_PAGES: u32 = 10; - -/// Sync cadence — matches Gmail (15 minutes). +/// Sync cadence for provider catalog scheduling. const SYNC_INTERVAL_SECS: u64 = 15 * 60; -/// Initial backoff for rate-limit retries. -const RATELIMIT_INITIAL_BACKOFF: Duration = Duration::from_secs(2); - -/// Cap on per-retry backoff. -const RATELIMIT_MAX_BACKOFF: Duration = Duration::from_secs(30); - -/// Total retries for a single rate-limited call before giving up. -const RATELIMIT_MAX_ATTEMPTS: u32 = 6; - -/// Fixed inter-call sleep applied after every successful execute_tool. -const INTER_CALL_PACING: Duration = Duration::from_secs(20); - -fn inter_call_pacing() -> Duration { - // Read per call so the slack sync e2e tests can control pacing at runtime - // via the env var. The only repeated-cost concern is the misconfiguration - // warning, which we emit at most once to avoid log spam on every - // `execute_tool`. - match std::env::var("OPENHUMAN_SLACK_INTER_CALL_PACING_MS") { - Ok(s) => match s.trim().parse::() { - Ok(ms) => Duration::from_millis(ms), - _ => { - static WARNED: std::sync::Once = std::sync::Once::new(); - WARNED.call_once(|| { - log::warn!( - "[composio:slack] OPENHUMAN_SLACK_INTER_CALL_PACING_MS={s:?} not a \ - non-negative integer; falling back to default {INTER_CALL_PACING:?}" - ); - }); - INTER_CALL_PACING - } - }, - Err(_) => INTER_CALL_PACING, - } -} - -/// Resolve the JSON dump directory from `OPENHUMAN_SLACK_DUMP_DIR`. -fn dump_dir() -> Option { - std::env::var_os("OPENHUMAN_SLACK_DUMP_DIR").map(PathBuf::from) -} - -/// Write a Composio response payload to disk under the dump dir. Best -/// effort — failures are logged at warn level and never fail the sync. -pub(super) fn dump_response(scope: &str, kind: &str, idx: u32, data: &Value) { - let Some(base) = dump_dir() else { - return; - }; - let path = base.join(scope).join(format!("{kind}-{idx:04}.json")); - if let Some(parent) = path.parent() { - if let Err(e) = std::fs::create_dir_all(parent) { - tracing::warn!( - error = %e, - path = %parent.display(), - "[composio:slack] dump_response: create_dir_all failed (skipping dump)" - ); - return; - } - } - match serde_json::to_string_pretty(data) { - Ok(json) => { - if let Err(e) = std::fs::write(&path, json) { - tracing::warn!( - error = %e, - path = %path.display(), - "[composio:slack] dump_response: write failed" - ); - } else { - tracing::debug!(path = %path.display(), "[composio:slack] dumped response"); - } - } - Err(e) => { - tracing::warn!(error = %e, "[composio:slack] dump_response: serialize failed"); - } - } -} - -/// Dispatch a Composio action with rate-limit-aware retry + inter-call -/// pacing. -/// -/// Routes through [`ProviderContext::execute`] so the live -/// `composio.mode` toggle is honoured per call (#1710). Pre-fix this -/// took a pre-baked `&ComposioClient` resolved at sync entry, which -/// silently bypassed the mode toggle. -/// -/// Returns `(response, attempts_made)` on first success so callers can -/// charge the daily quota meter for every attempt that hit Composio. -pub(super) async fn execute_with_retry( - ctx: &ProviderContext, - slug: &str, - args: serde_json::Value, - description: &str, -) -> Result<(ComposioExecuteResponse, u32), String> { - let mut delay = RATELIMIT_INITIAL_BACKOFF; - for attempt in 1..=RATELIMIT_MAX_ATTEMPTS { - let resp = ctx - .execute(slug, Some(args.clone())) - .await - .map_err(|e| format!("{description}: {e:#}"))?; - if resp.successful { - tokio::time::sleep(inter_call_pacing()).await; - return Ok((resp, attempt)); - } - let err_str = resp.error.as_deref().unwrap_or("provider failure"); - let is_ratelimit = err_str.contains("ratelimited") - || err_str.contains("rate_limit") - || err_str.contains("rate limit"); - if is_ratelimit && attempt < RATELIMIT_MAX_ATTEMPTS { - tracing::warn!( - slug, - attempt, - max_attempts = RATELIMIT_MAX_ATTEMPTS, - sleep_ms = delay.as_millis() as u64, - "[composio:slack] rate-limited; backing off and retrying" - ); - tokio::time::sleep(delay).await; - delay = (delay * 2).min(RATELIMIT_MAX_BACKOFF); - continue; - } - return Err(format!("{description}: {err_str}")); - } - Err(format!( - "{description}: rate-limited after {RATELIMIT_MAX_ATTEMPTS} retries" - )) -} - pub struct SlackProvider; impl SlackProvider { @@ -411,10 +223,6 @@ impl ComposioProvider for SlackProvider { /// watermark, dedup, the `max_items` cap, and per-channel error tolerance /// all live in `run_sync`. The Slack-specific primitives live in /// [`super::source`]. - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { - super::source::run_slack_sync(ctx, reason).await - } - async fn on_trigger( &self, ctx: &ProviderContext, @@ -422,7 +230,16 @@ impl ComposioProvider for SlackProvider { _payload: &Value, ) -> Result<(), String> { if trigger.to_ascii_uppercase().contains("MESSAGE") { - if let Err(e) = self.sync(ctx, SyncReason::Manual).await { + let Some(connection_id) = ctx.connection_id.as_deref() else { + return Err("[composio:slack] trigger missing connection_id".to_string()); + }; + if let Err(e) = crate::openhuman::tinycortex::run_composio_connection( + "slack", + connection_id, + ctx.config.as_ref(), + ) + .await + { tracing::warn!( error = %e, "[composio:slack] trigger-driven sync failed (non-fatal)" @@ -433,251 +250,49 @@ impl ComposioProvider for SlackProvider { } } -/// Paginate through `SLACK_LIST_CONVERSATIONS` and flatten into a -/// single `Vec`. -pub(super) async fn list_all_channels( - ctx: &ProviderContext, - state: &mut SyncState, -) -> Result, String> { - let mut out: Vec = Vec::new(); - let mut cursor: Option = None; - - for page_num in 0..MAX_LIST_PAGES { - if state.budget_exhausted() { - tracing::warn!( - page = page_num, - "[composio:slack] budget exhausted during channel listing" - ); - break; - } - let mut args = json!({ - "types": "public_channel,private_channel", - "exclude_archived": true, - "limit": LIST_PAGE_SIZE, - }); - if let Some(ref c) = cursor { - args["cursor"] = json!(c); - } - - let (mut resp, attempts) = execute_with_retry( - ctx, - ACTION_LIST_CONVERSATIONS, - args, - &format!("{ACTION_LIST_CONVERSATIONS} page {page_num}"), - ) - .await?; - state.record_requests(attempts); - dump_response("_meta", "channels", page_num, &resp.data); - - // Post-process then enrich. - super::post_process::post_process(ACTION_LIST_CONVERSATIONS, None, &mut resp.data); - out.extend(sync::extract_channels(&resp.data)); - cursor = sync::extract_next_cursor(&resp.data); - if cursor.is_none() { - break; - } - } - Ok(out) -} - // ── Search-based backfill (one-shot) ──────────────────────────────── -/// Composio action slug for workspace-wide message search. -const ACTION_SEARCH_MESSAGES: &str = "SLACK_SEARCH_MESSAGES"; - -/// Max matches per `SLACK_SEARCH_MESSAGES` page. -const SEARCH_PAGE_SIZE: u32 = 100; - -/// Hard cap on pages walked per backfill run. -const MAX_SEARCH_PAGES: u32 = 50; - -/// Run a one-shot historical backfill via `SLACK_SEARCH_MESSAGES` — -/// workspace-wide paginated search instead of per-channel -/// `conversations.history`. Each successful call returns matches across -/// many channels, so partial progress translates to real coverage. -/// -/// Designed for the `slack-backfill` bin specifically — the periodic -/// `SlackProvider::sync()` keeps the per-channel incremental path. -/// -/// Lifecycle: -/// 1. Cache the channel directory and user directory. -/// 2. Paginate `SLACK_SEARCH_MESSAGES` until exhausted or page cap. -/// 3. Group messages by channel_id, ingest each group via -/// `ingest_page_into_memory_tree`. No bucketing. +/// Compatibility wrapper for the tinycortex workspace-wide search pipeline. pub async fn run_backfill_via_search( ctx: &ProviderContext, backfill_days: i64, ) -> Result { - let started_at_ms = sync::now_ms(); let connection_id = ctx .connection_id - .clone() - .unwrap_or_else(|| "default".to_string()); - - tracing::info!( - connection_id = %connection_id, + .as_deref() + .ok_or_else(|| "[composio:slack] search backfill missing connection_id".to_string())?; + let started_at_ms = now_ms(); + let outcome = crate::openhuman::tinycortex::run_slack_search_backfill( + connection_id, backfill_days, - "[composio:slack] search-based backfill starting" - ); - - let memory = ctx - .memory_client() - .ok_or_else(|| "[composio:slack] memory client not ready".to_string())?; - let mut state = SyncState::load(&memory, "slack", &connection_id).await?; - - if state.budget_exhausted() { - return Ok(SyncOutcome { - toolkit: "slack".to_string(), - connection_id: Some(connection_id), - reason: SyncReason::Manual.as_str().to_string(), - items_ingested: 0, - started_at_ms, - finished_at_ms: sync::now_ms(), - summary: "slack search-backfill skipped: daily budget exhausted".to_string(), - details: json!({ "budget_exhausted": true }), - }); - } - - // 1. Channel directory. - let channels = list_all_channels(ctx, &mut state) - .await - .map_err(|e| format!("[composio:slack] list_channels: {e:#}"))?; - let channel_map: HashMap = - channels.into_iter().map(|c| (c.id.clone(), c)).collect(); - - // 2. User directory. - let (users, user_call_count) = SlackUsers::fetch(ctx).await; - state.record_requests(user_call_count); - tracing::info!( - connection_id = %connection_id, - user_count = users.len(), - channel_count = channel_map.len(), - "[composio:slack] caches ready" - ); - let _ = state.save(&memory).await; - - // 3. Paginated workspace-wide search. - let now = chrono::Utc::now(); - let after = (now - chrono::Duration::days(backfill_days)) - .format("%Y-%m-%d") - .to_string(); - let query = format!("after:{after}"); - let mut all_messages: Vec = Vec::new(); - let mut page: u32 = 1; - let mut total_pages: u32 = 1; - - loop { - if state.budget_exhausted() { - tracing::warn!( - page, - "[composio:slack] budget exhausted mid-search, halting" - ); - break; - } - let args = json!({ - "query": query, - "count": SEARCH_PAGE_SIZE, - "sort": "timestamp", - "sort_dir": "asc", - "page": page, - }); - let (mut resp, attempts) = execute_with_retry( - ctx, - ACTION_SEARCH_MESSAGES, - args, - &format!("{ACTION_SEARCH_MESSAGES} page {page}"), - ) - .await?; - state.record_requests(attempts); - dump_response("_meta", "search", page, &resp.data); - - // Post-process, then enrich with channel_map + users. - super::post_process::post_process(ACTION_SEARCH_MESSAGES, None, &mut resp.data); - let msgs = sync::extract_search_messages(&resp.data, &channel_map, &users); - if page == 1 { - total_pages = sync::extract_search_total_pages(&resp.data).min(MAX_SEARCH_PAGES); - tracing::info!( - connection_id = %connection_id, - total_pages, - first_page_msgs = msgs.len(), - "[composio:slack] search pagination plan" - ); - } - let fetched = msgs.len(); - all_messages.extend(msgs); - if fetched == 0 || page >= total_pages { - break; - } - page += 1; - } - let _ = state.save(&memory).await; - - // 4. Group by channel_id and ingest each group. - let buckets = super::ingest::bucket_by_channel(&all_messages); - let channel_count = buckets.len(); - tracing::info!( - connection_id = %connection_id, - channels = channel_count, - total_messages = all_messages.len(), - "[composio:slack] grouped messages by channel for ingest" - ); - - let mut channels_flushed = 0usize; - let mut channels_failed = 0usize; - let mut total_chunks = 0usize; - - for (channel_id, msgs_for_channel) in &buckets { - let page: Vec = msgs_for_channel.iter().map(|m| (*m).clone()).collect(); - match ingest_page_into_memory_tree(&ctx.config, "", &connection_id, &page).await { - Ok(chunks) => { - channels_flushed += 1; - total_chunks += chunks; - tracing::info!( - channel = %channel_id, - messages = page.len(), - chunks, - "[composio:slack] search-backfill channel ingested" - ); - } - Err(err) => { - channels_failed += 1; - tracing::warn!( - channel = %channel_id, - error = %err, - "[composio:slack] search-backfill ingest failed" - ); - } - } - } - - let finished_at_ms = sync::now_ms(); - let summary = format!( - "slack search-backfill: pages={page} channels_flushed={channels_flushed} \ - channels_failed={channels_failed} chunks={total_chunks}" - ); - tracing::info!( - connection_id = %connection_id, - elapsed_ms = finished_at_ms.saturating_sub(started_at_ms), - "{summary}" - ); - + ctx.config.as_ref(), + ) + .await + .map_err(|error| error.to_string())?; Ok(SyncOutcome { - toolkit: "slack".to_string(), - connection_id: Some(connection_id), - reason: SyncReason::Manual.as_str().to_string(), - items_ingested: total_chunks, + toolkit: "slack".into(), + connection_id: Some(connection_id.into()), + reason: "manual".into(), + items_ingested: outcome.records_ingested as usize, started_at_ms, - finished_at_ms, - summary, - details: json!({ - "pages_walked": page, - "channels_flushed": channels_flushed, - "channels_failed": channels_failed, - "total_chunks": total_chunks, + finished_at_ms: now_ms(), + summary: outcome + .note + .unwrap_or_else(|| "slack search-backfill complete".into()), + details: serde_json::json!({ + "more_pending": outcome.more_pending, + "actions_called": outcome.actions_called, + "provider_cost_usd": outcome.provider_cost_usd, }), }) } +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} #[cfg(test)] mod tests { use super::*; diff --git a/src/openhuman/memory_sync/composio/providers/slack/rpc.rs b/src/openhuman/memory_sync/composio/providers/slack/rpc.rs index 9d4fe23f2..6ceb0c59c 100644 --- a/src/openhuman/memory_sync/composio/providers/slack/rpc.rs +++ b/src/openhuman/memory_sync/composio/providers/slack/rpc.rs @@ -10,8 +10,6 @@ //! - `openhuman.slack_memory_sync_status` — list the per-connection //! sync cursors + last-synced timestamps. -use std::sync::Arc; - use serde::{Deserialize, Serialize}; use crate::openhuman::composio::client::{ @@ -19,12 +17,7 @@ use crate::openhuman::composio::client::{ }; use crate::openhuman::composio::types::ComposioConnectionsResponse; use crate::openhuman::config::Config; -use crate::openhuman::memory::global::client_if_ready; -use crate::openhuman::memory_sync::composio::providers::registry::get_provider; -use crate::openhuman::memory_sync::composio::providers::sync_state::SyncState; -use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, SyncOutcome, SyncReason, -}; +use crate::openhuman::memory_sync::composio::providers::SyncOutcome; use crate::rpc::RpcOutcome; /// Optional connection-id override for the trigger. When absent, all @@ -74,9 +67,6 @@ pub async fn sync_trigger_rpc( config: &Config, req: SyncTriggerRequest, ) -> Result, String> { - let provider = get_provider("slack") - .ok_or_else(|| "[slack_ingest] SlackProvider not registered".to_string())?; - // Route through the mode-aware factory so direct-mode users // discover slack connections from THEIR personal Composio tenant — // not the tinyhumans backend tenant. Mirrors `composio::ops` @@ -99,23 +89,28 @@ pub async fn sync_trigger_rpc( } let considered = candidates.len(); - let config_arc = Arc::new(config.clone()); let mut outcomes: Vec = Vec::with_capacity(considered); for conn in candidates { - // `ProviderContext` no longer caches a pre-baked client — - // `ctx.execute(...)` resolves the underlying handle per call - // via the mode-aware factory (#1710). - let ctx = ProviderContext { - config: Arc::clone(&config_arc), - toolkit: conn.toolkit.clone(), - connection_id: Some(conn.id.clone()), - usage: Default::default(), - max_items: None, - sync_depth_days: None, - }; - match provider.sync(&ctx, SyncReason::Manual).await { - Ok(o) => outcomes.push(o), + let started_at_ms = now_ms(); + match crate::openhuman::tinycortex::run_composio_connection("slack", &conn.id, config).await + { + Ok(outcome) => outcomes.push(SyncOutcome { + toolkit: "slack".to_string(), + connection_id: Some(conn.id.clone()), + reason: "manual".to_string(), + items_ingested: outcome.records_ingested as usize, + started_at_ms, + finished_at_ms: now_ms(), + summary: outcome + .note + .unwrap_or_else(|| "Slack sync completed".to_string()), + details: serde_json::json!({ + "more_pending": outcome.more_pending, + "actions_called": outcome.actions_called, + "provider_cost_usd": outcome.provider_cost_usd, + }), + }), Err(err) => { log::warn!( "[slack_ingest] connection={} sync failed: {err:#} (continuing)", @@ -136,6 +131,13 @@ pub async fn sync_trigger_rpc( )) } +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + /// Request body for `slack_memory_sync_status` — no parameters. #[derive(Clone, Debug, Serialize, Deserialize, Default)] pub struct SyncStatusRequest {} @@ -166,9 +168,6 @@ pub async fn sync_status_rpc( config: &Config, _req: SyncStatusRequest, ) -> Result, String> { - let memory = - client_if_ready().ok_or_else(|| "[slack_ingest] memory client not ready".to_string())?; - // Route through the mode-aware factory so direct-mode users see // status rows for THEIR slack connections, not the tinyhumans // backend tenant's (#1710). @@ -182,16 +181,17 @@ pub async fn sync_status_rpc( if !conn.is_active() { continue; } - let state = match SyncState::load(&memory, "slack", &conn.id).await { - Ok(s) => s, - Err(err) => { - log::warn!( - "[slack_ingest] load_state connection={} failed: {err:#}", - conn.id - ); - continue; - } - }; + let state = + match crate::openhuman::tinycortex::load_composio_sync_state("slack", &conn.id).await { + Ok(s) => s, + Err(err) => { + log::warn!( + "[slack_ingest] load_state connection={} failed: {err:#}", + conn.id + ); + continue; + } + }; rows.push(ConnectionStatus { connection_id: conn.id.clone(), per_channel_cursors: state.cursor.clone().unwrap_or_else(|| "{}".to_string()), diff --git a/src/openhuman/memory_sync/composio/providers/slack/source.rs b/src/openhuman/memory_sync/composio/providers/slack/source.rs deleted file mode 100644 index 2a698de97..000000000 --- a/src/openhuman/memory_sync/composio/providers/slack/source.rs +++ /dev/null @@ -1,422 +0,0 @@ -//! Slack's [`IncrementalSource`] primitives. -//! -//! Slack rides the generic -//! [`crate::openhuman::memory_sync::composio::providers::orchestrator`]: -//! [`SlackProvider::sync`](super::provider::SlackProvider) delegates to -//! [`run_slack_sync`]. The orchestrator owns the control flow (budget, -//! pagination bound, dedup, the `max_items` clamp, cursor advance/hold, state -//! persistence); this module supplies only the Slack-specific shapes. -//! -//! Slack is the **structural outlier** the orchestrator's per-scope extensions -//! were built for: -//! -//! * **Per-scope cursors** — Slack keeps one `oldest` watermark *per channel* -//! (a `BTreeMap` serialized into [`SyncState::cursor`]) -//! rather than a single global watermark. [`SlackSource`] opts in via -//! [`IncrementalSource::per_scope_cursors`] and advances each channel's -//! watermark in [`IncrementalSource::advance_scope_cursor`]. The watermark -//! is enforced **server-side** (the `oldest` request arg), so the source -//! also advertises [`IncrementalSource::server_side_depth`] and the -//! orchestrator skips its client-side depth/boundary filtering. -//! * **Per-scope error tolerance** — a single channel that errors mid-sync -//! must not abort the other channels, so [`SlackSource`] opts into -//! [`IncrementalSource::tolerate_scope_errors`]. -//! -//! Slack dedups purely by content-hashed chunk ids (UPSERT) plus the per-channel -//! `oldest` watermark, so [`SlackSource::ingest`] returns **no** `synced_keys` -//! — the `synced_ids` set deliberately stays empty (it would otherwise grow -//! unboundedly, one entry per message, forever). -//! -//! The user directory is fetched once as a [`IncrementalSource::preamble`] -//! side-effect and stashed for every per-channel [`IncrementalSource::ingest`] -//! to resolve authors + `<@…>` mentions. - -use std::collections::HashMap; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::OnceLock; - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use serde_json::{json, Value}; - -use super::ingest::ingest_page_into_memory_tree; -use super::provider::{ - backfill_days, dump_response, execute_with_retry, list_all_channels, ACTION_FETCH_HISTORY, - HISTORY_PAGE_SIZE, MAX_HISTORY_PAGES_PER_CHANNEL, -}; -use super::sync; -use super::types::SlackChannel; -use super::users::SlackUsers; -use crate::openhuman::memory_sync::composio::providers::orchestrator::{ - self, IncrementalSource, IngestOutcome, PageFetch, SyncItem, SyncScope, -}; -use crate::openhuman::memory_sync::composio::providers::sync_state::SyncState; -use crate::openhuman::memory_sync::composio::providers::{ - ProviderContext, SyncOutcome, SyncReason, -}; - -/// Slack's [`IncrementalSource`]. Holds the per-sync user directory and channel -/// metadata resolved in the preamble so every per-channel `fetch_page` / -/// `ingest` can read them back. -pub(crate) struct SlackSource { - /// Workspace user directory, fetched once in [`Self::preamble`]. - users: OnceLock, - /// Channel metadata keyed by channel id, resolved in the preamble so - /// `ingest` can render channel labels without re-listing. - channels: OnceLock>, - /// Wall-clock captured once per sync, used to compute the first-fetch - /// backfill window for channels with no cursor yet. - now: DateTime, - /// Monotonic counter for `dump_response` filenames (best-effort debug). - dump_seq: AtomicU32, -} - -impl SlackSource { - fn new() -> Self { - Self { - users: OnceLock::new(), - channels: OnceLock::new(), - now: Utc::now(), - dump_seq: AtomicU32::new(0), - } - } -} - -/// Entry point used by [`super::provider::SlackProvider::sync`]. -pub(crate) async fn run_slack_sync( - ctx: &ProviderContext, - reason: SyncReason, -) -> Result { - orchestrator::run_sync(&SlackSource::new(), ctx, reason).await -} - -#[async_trait] -impl IncrementalSource for SlackSource { - fn toolkit(&self) -> &'static str { - "slack" - } - - fn page_size(&self, _reason: SyncReason) -> u32 { - HISTORY_PAGE_SIZE - } - - fn max_pages(&self) -> u32 { - MAX_HISTORY_PAGES_PER_CHANNEL - } - - fn detail_noun(&self) -> &'static str { - "messages" - } - - /// Slack keeps a watermark per channel, enforced server-side via `oldest`. - fn per_scope_cursors(&self) -> bool { - true - } - - /// A single channel's failure must not abort the rest of the sync. - fn tolerate_scope_errors(&self) -> bool { - true - } - - /// `fetch_page` injects the `oldest` window itself, so the orchestrator must - /// not also apply a client-side depth floor (Slack `ts` isn't RFC3339). - fn server_side_depth(&self) -> bool { - true - } - - /// Resolve the workspace user directory + the channel list, stash both for - /// `ingest`, and return one [`SyncScope`] per channel. - async fn preamble( - &self, - ctx: &ProviderContext, - state: &mut SyncState, - ) -> Result, String> { - // Pull the workspace user directory once per sync (soft-fails to empty). - let (users, user_call_count) = SlackUsers::fetch(ctx).await; - state.record_requests(user_call_count); - tracing::info!( - connection_id = ?ctx.connection_id, - user_count = users.len(), - "[composio:slack] users cached for this sync" - ); - let _ = self.users.set(users); - - // Enumerate channels (records its own page budget). - let channels = list_all_channels(ctx, state) - .await - .map_err(|e| format!("[composio:slack] list_channels: {e:#}"))?; - tracing::info!( - connection_id = ?ctx.connection_id, - channel_count = channels.len(), - "[composio:slack] channels discovered" - ); - - let scopes: Vec = channels - .iter() - .map(|c| { - let label = super::ingest::channel_label(&c.name, c.is_private); - SyncScope::nested(c.id.clone(), label) - }) - .collect(); - - let channel_map: HashMap = - channels.into_iter().map(|c| (c.id.clone(), c)).collect(); - let _ = self.channels.set(channel_map); - - Ok(scopes) - } - - async fn fetch_page( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - cursor: Option<&str>, - _reason: SyncReason, - state: &mut SyncState, - ) -> Result { - let channel_id = &scope.id; - - // Per-channel `oldest` watermark: the persisted cursor for this channel, - // or the backfill window when the channel has never been synced. Full - // microsecond precision is preserved so `inclusive=false` excludes only - // the exact last-seen message. `ctx.sync_depth_days` (user-configured) - // wins over the env-var default when set. - let cursors = sync::decode_cursors(state.cursor.as_deref()); - let oldest_ts = cursors.get(channel_id).cloned().unwrap_or_else(|| { - let depth_days = ctx - .sync_depth_days - .map(|d| d as i64) - .unwrap_or_else(backfill_days); - let secs = (self.now - chrono::Duration::days(depth_days)).timestamp(); - tracing::debug!( - channel = %channel_id, - depth_days, - oldest_ts_secs = secs, - "[composio:slack] [memory_sync] computing oldest_ts for backfill" - ); - format!("{secs}.000000") - }); - - let mut args = json!({ - "channel": channel_id, - "oldest": oldest_ts, - "inclusive": false, - "limit": HISTORY_PAGE_SIZE, - }); - if let Some(c) = cursor { - args["cursor"] = json!(c); - } - - let (mut resp, attempts) = execute_with_retry( - ctx, - ACTION_FETCH_HISTORY, - args, - &format!("{ACTION_FETCH_HISTORY} channel={channel_id}"), - ) - .await?; - state.record_requests(attempts); - - let idx = self.dump_seq.fetch_add(1, Ordering::Relaxed); - dump_response(channel_id, "history", idx, &resp.data); - - // Post-process to the slim envelope, then hand the orchestrator the raw - // message values — `item_dedup_key` drops blanks, `ingest` enriches. - super::post_process::post_process(ACTION_FETCH_HISTORY, None, &mut resp.data); - let items = resp - .data - .get("messages") - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - let next = sync::extract_next_cursor(&resp.data); - - tracing::debug!( - channel = %channel_id, - fetched = items.len(), - has_next = next.is_some(), - "[composio:slack] history page" - ); - - Ok(PageFetch { items, next }) - } - - /// Dedup key doubles as the validity gate: a message with no parseable `ts` - /// or blank text returns `None` so the orchestrator drops it **before** the - /// `max_items` clamp — keeping the cap counted against *valid* messages, - /// exactly as the old per-channel loop did (it filtered then clamped). - fn item_dedup_key(&self, item: &Value) -> Option { - let ts = item.get("ts").and_then(Value::as_str)?; - sync::parse_ts(ts)?; - let text = item.get("text").and_then(Value::as_str).unwrap_or(""); - if text.trim().is_empty() { - return None; - } - Some(ts.to_string()) - } - - fn item_sort_ts(&self, item: &Value) -> Option { - item.get("ts").and_then(Value::as_str).map(str::to_string) - } - - /// Advance this channel's watermark inside the per-channel cursor map - /// serialized in [`SyncState::cursor`]. - fn advance_scope_cursor(&self, state: &mut SyncState, scope: &SyncScope, newest_ts: &str) { - let mut cursors = sync::decode_cursors(state.cursor.as_deref()); - cursors.insert(scope.id.clone(), newest_ts.to_string()); - state.cursor = Some(sync::encode_cursors(&cursors)); - } - - async fn ingest( - &self, - ctx: &ProviderContext, - scope: &SyncScope, - _state: &mut SyncState, - items: Vec, - ) -> IngestOutcome { - let channel_id = &scope.id; - // Channel metadata from the preamble (falls back to a bare id-named - // channel if the map is somehow missing this scope). - let channel = self - .channels - .get() - .and_then(|m| m.get(channel_id)) - .cloned() - .unwrap_or_else(|| SlackChannel { - id: channel_id.clone(), - name: channel_id.clone(), - is_private: false, - }); - let users = self.users.get().cloned().unwrap_or_else(SlackUsers::empty); - - // Enrich the surviving raw message values into canonical SlackMessages - // (author resolution + `<@…>` mention rewriting) via the shared parser. - let raws: Vec = items.into_iter().map(|it| it.raw).collect(); - let wrapped = json!({ "messages": raws }); - let messages = sync::extract_messages(&wrapped, &channel, &users); - - if messages.is_empty() { - return IngestOutcome::default(); - } - - let connection_id = ctx.connection_id.as_deref().unwrap_or("default"); - let count = messages.len(); - match ingest_page_into_memory_tree(&ctx.config, "", connection_id, &messages).await { - Ok(chunks) => { - tracing::info!( - channel = %channel_id, - messages = count, - chunks, - "[composio:slack] channel ingest done" - ); - IngestOutcome { - // No synced_keys: Slack dedups via content-hash UPSERT + the - // per-channel watermark, so `synced_ids` stays empty. - synced_keys: Vec::new(), - persisted: count, - had_failures: false, - } - } - Err(e) => { - tracing::warn!( - channel = %channel_id, - error = %e, - "[composio:slack] ingest_page_into_memory_tree failed (watermark held)" - ); - IngestOutcome { - synced_keys: Vec::new(), - persisted: 0, - had_failures: true, - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - fn source() -> SlackSource { - SlackSource::new() - } - - #[test] - fn item_dedup_key_drops_blank_and_unparseable() { - let s = source(); - // Valid message → keyed by ts. - assert_eq!( - s.item_dedup_key(&json!({ "ts": "1714003200.000100", "text": "hi" })) - .as_deref(), - Some("1714003200.000100") - ); - // Blank text → dropped (None), so it never counts against the cap. - assert_eq!( - s.item_dedup_key(&json!({ "ts": "1714003200.000100", "text": " " })), - None - ); - // Bot-authored (no `user`) but non-blank text → kept (parity with the - // old extractor, which retained author-less messages). - assert_eq!( - s.item_dedup_key(&json!({ "ts": "1714003300.000200", "text": "bot update" })) - .as_deref(), - Some("1714003300.000200") - ); - // Unparseable ts → dropped. - assert_eq!( - s.item_dedup_key(&json!({ "ts": "nope", "text": "hi" })), - None - ); - // Missing ts → dropped. - assert_eq!(s.item_dedup_key(&json!({ "text": "hi" })), None); - } - - #[test] - fn item_sort_ts_reads_raw_ts() { - let s = source(); - assert_eq!( - s.item_sort_ts(&json!({ "ts": "1714003200.000100" })) - .as_deref(), - Some("1714003200.000100") - ); - assert_eq!(s.item_sort_ts(&json!({ "text": "no ts" })), None); - } - - #[test] - fn advance_scope_cursor_merges_into_per_channel_map() { - let s = source(); - let mut state = SyncState::new("slack", "conn1"); - state.cursor = Some(r#"{"C1":"1714003200.000100"}"#.to_string()); - - s.advance_scope_cursor( - &mut state, - &SyncScope::nested("C2", "#two"), - "1714010000.000200", - ); - let map = sync::decode_cursors(state.cursor.as_deref()); - assert_eq!(map.get("C1").map(String::as_str), Some("1714003200.000100")); - assert_eq!(map.get("C2").map(String::as_str), Some("1714010000.000200")); - - // Re-advancing an existing channel overwrites just that entry. - s.advance_scope_cursor( - &mut state, - &SyncScope::nested("C1", "#one"), - "1714099999.000300", - ); - let map = sync::decode_cursors(state.cursor.as_deref()); - assert_eq!(map.get("C1").map(String::as_str), Some("1714099999.000300")); - assert_eq!(map.len(), 2); - } - - #[test] - fn source_advertises_slack_outlier_hooks() { - let s = source(); - assert!(s.per_scope_cursors()); - assert!(s.tolerate_scope_errors()); - assert!(s.server_side_depth()); - assert_eq!(s.toolkit(), "slack"); - assert_eq!(s.detail_noun(), "messages"); - assert_eq!(s.page_size(SyncReason::Periodic), HISTORY_PAGE_SIZE); - assert_eq!(s.max_pages(), MAX_HISTORY_PAGES_PER_CHANNEL); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/slack/sync.rs b/src/openhuman/memory_sync/composio/providers/slack/sync.rs deleted file mode 100644 index 5321a1486..000000000 --- a/src/openhuman/memory_sync/composio/providers/slack/sync.rs +++ /dev/null @@ -1,473 +0,0 @@ -//! Helpers for the Composio-backed Slack provider. -//! -//! This module contains thin enrichers that take a post-processed slim -//! envelope (produced by [`super::post_process`]) and turn it into -//! [`SlackMessage`] / [`SlackChannel`] values with user-id resolution and -//! channel-context injection applied. -//! -//! Response-shape walking (nested envelopes, empty-field filtering) lives in -//! `post_process.rs`; this module assumes the slim shape is already in place. - -use std::collections::{BTreeMap, HashMap}; -use std::time::{SystemTime, UNIX_EPOCH}; - -use chrono::{DateTime, TimeZone, Utc}; -use serde_json::Value; - -use super::types::{SlackChannel, SlackMessage}; -use super::users::SlackUsers; - -/// Enrich the top-level `channels[]` array in a post-processed -/// `SLACK_LIST_CONVERSATIONS` response into [`SlackChannel`] values. -/// -/// The post-processor has already stripped unknown channels and normalised -/// to `{ id, name, is_private }` — this function just deserialises them. -pub(crate) fn extract_channels(data: &Value) -> Vec { - let arr = data - .get("channels") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - arr.into_iter().filter_map(parse_channel).collect() -} - -fn parse_channel(raw: Value) -> Option { - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if id.is_empty() { - return None; - } - let name = raw - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or(&id) - .to_string(); - let is_private = raw - .get("is_private") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - Some(SlackChannel { - id, - name, - is_private, - }) -} - -/// Enrich the top-level `messages[]` array in a post-processed -/// `SLACK_FETCH_CONVERSATION_HISTORY` response into [`SlackMessage`]s. -/// -/// `channel` provides the channel id, name, and privacy flag (not present -/// in the response body — only in the request). `users` resolves author ids -/// and rewrites `<@…>` mentions in message text. -pub(crate) fn extract_messages( - data: &Value, - channel: &SlackChannel, - users: &SlackUsers, -) -> Vec { - let arr = data - .get("messages") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - arr.into_iter() - .filter_map(|raw| parse_message(raw, channel, users)) - .collect() -} - -fn parse_message(raw: Value, channel: &SlackChannel, users: &SlackUsers) -> Option { - let ts_raw = raw.get("ts").and_then(|t| t.as_str())?.to_string(); - let timestamp = parse_ts(&ts_raw)?; - let raw_text = raw - .get("text") - .and_then(|t| t.as_str()) - .unwrap_or("") - .to_string(); - if raw_text.trim().is_empty() { - return None; - } - let text = users.replace_mentions(&raw_text); - let author_id = raw - .get("user") - .and_then(|u| u.as_str()) - .unwrap_or("") - .to_string(); - let author = users.resolve(&author_id); - let thread_ts = raw - .get("thread_ts") - .and_then(|t| t.as_str()) - .map(String::from); - let permalink = raw - .get("permalink") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .map(String::from); - Some(SlackMessage { - channel_id: channel.id.clone(), - channel_name: channel.name.clone(), - is_private: channel.is_private, - author, - author_id, - text, - timestamp, - ts_raw, - thread_ts, - permalink, - }) -} - -/// Enrich the top-level `messages[]` array in a post-processed -/// `SLACK_SEARCH_MESSAGES` response into [`SlackMessage`]s. -/// -/// `channel_map` provides channel names and privacy flags keyed by id. -/// When a match's `channel_id` is absent from the map, channel name and -/// privacy default to empty/false — the message is still ingested but -/// the label will be less informative. -pub(crate) fn extract_search_messages( - data: &Value, - channel_map: &HashMap, - users: &SlackUsers, -) -> Vec { - let arr = data - .get("messages") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - arr.into_iter() - .filter_map(|raw| parse_search_match(raw, channel_map, users)) - .collect() -} - -fn parse_search_match( - raw: Value, - channel_map: &HashMap, - users: &SlackUsers, -) -> Option { - let ts_raw = raw.get("ts").and_then(|t| t.as_str())?.to_string(); - let timestamp = parse_ts(&ts_raw)?; - let raw_text = raw - .get("text") - .and_then(|t| t.as_str()) - .unwrap_or("") - .to_string(); - if raw_text.trim().is_empty() { - return None; - } - let text = users.replace_mentions(&raw_text); - let author_id = raw - .get("user") - .and_then(|u| u.as_str()) - .unwrap_or("") - .to_string(); - let author = users.resolve(&author_id); - // Drop malformed search hits with no channel id — they'd funnel into a - // single empty-channel bucket downstream and ingest under the wrong - // (or no) channel context. - let channel_id = raw - .get("channel_id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty())? - .to_string(); - let (channel_name, is_private) = channel_map - .get(&channel_id) - .map(|c| (c.name.clone(), c.is_private)) - .unwrap_or_else(|| (String::new(), false)); - let thread_ts = raw - .get("thread_ts") - .and_then(|t| t.as_str()) - .map(String::from); - let permalink = raw - .get("permalink") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .map(String::from); - Some(SlackMessage { - channel_id, - channel_name, - is_private, - author, - author_id, - text, - timestamp, - ts_raw, - thread_ts, - permalink, - }) -} - -/// Slack's `ts` is a decimal string `"."`. The -/// integer part is what we care about for `DateTime` purposes; -/// micro is preserved separately by [`parse_ts_components`] for cursor -/// ordering. -pub(crate) fn parse_ts(ts_raw: &str) -> Option> { - let seconds_str = ts_raw.split('.').next()?; - let secs: i64 = seconds_str.parse().ok()?; - Utc.timestamp_opt(secs, 0).single() -} - -/// Parse a Slack `ts` into a `(seconds, micros)` tuple suitable for -/// `max_by_key` / lexicographic ordering at full precision. Unparseable -/// inputs fall back to `(0, 0)` so they never dominate a max — they -/// just lose to anything real. -pub(crate) fn parse_ts_components(ts_raw: &str) -> (i64, u64) { - let mut parts = ts_raw.splitn(2, '.'); - let secs: i64 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(0); - let micros: u64 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(0); - (secs, micros) -} - -/// Extract the total page count from a post-processed -/// `SLACK_SEARCH_MESSAGES` response. Defaults to 1 when absent. -pub(crate) fn extract_search_total_pages(data: &Value) -> u32 { - data.get("pages").and_then(|v| v.as_u64()).unwrap_or(1) as u32 -} - -/// Extract a pagination `next_cursor` from a `SLACK_LIST_CONVERSATIONS` -/// or `SLACK_FETCH_CONVERSATION_HISTORY` response. -pub(crate) fn extract_next_cursor(data: &Value) -> Option { - let candidates = [ - data.pointer("/data/response_metadata/next_cursor"), - data.pointer("/response_metadata/next_cursor"), - data.pointer("/data/next_cursor"), - data.pointer("/next_cursor"), - ]; - for cand in candidates.into_iter().flatten() { - if let Some(s) = cand.as_str() { - let trimmed = s.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); - } - } - } - None -} - -/// Per-channel cursor map encoded into `SyncState.cursor`. We use -/// `BTreeMap` so serialization is deterministic (makes log diffs -/// readable and tests stable). -/// -/// Value is the **raw Slack `ts`** of the latest successfully-ingested -/// message for that channel (e.g. `"1714003200.123456"`) — full -/// microsecond precision is preserved so multi-message-per-second -/// channels don't replay an entire second on the next incremental -/// fetch (`oldest` with `inclusive=false` excludes only that exact -/// timestamp). Fetches for that channel use `oldest = value` -/// verbatim. -pub type ChannelCursors = BTreeMap; - -/// Deserialize the per-channel cursor map out of `SyncState.cursor`. -/// Returns an empty map on any parse failure — a broken cursor should -/// degrade to "start from the backfill window" rather than bail out. -pub(crate) fn decode_cursors(raw: Option<&str>) -> ChannelCursors { - let Some(raw) = raw else { - return ChannelCursors::new(); - }; - match serde_json::from_str::(raw) { - Ok(map) => map, - Err(err) => { - tracing::warn!( - error = %err, - "[composio:slack] cursor parse failed, resetting per-channel cursors" - ); - ChannelCursors::new() - } - } -} - -pub(crate) fn encode_cursors(map: &ChannelCursors) -> String { - serde_json::to_string(map).unwrap_or_else(|_| "{}".to_string()) -} - -pub(crate) fn now_ms() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn extract_channels_from_post_processed_shape() { - let data = json!({ - "channels": [ - {"id": "C1", "name": "eng", "is_private": false}, - {"id": "G1", "name": "ops", "is_private": true}, - {"id": "", "name": "empty-id-drops"}, - ] - }); - let out = extract_channels(&data); - assert_eq!(out.len(), 2); - assert_eq!(out[0].id, "C1"); - assert!(!out[0].is_private); - assert_eq!(out[1].id, "G1"); - assert!(out[1].is_private); - } - - #[test] - fn extract_messages_parses_post_processed_shape() { - let data = json!({ - "messages": [ - {"ts": "1714003200.000100", "user": "U1", "text": "hi"}, - {"ts": "1714003300.000200", "user": "U2", "text": "world"}, - {"ts": "1714003400.000300", "user": "U3", "text": " "} // dropped (blank) - ] - }); - let channel = SlackChannel { - id: "C1".into(), - name: "eng".into(), - is_private: false, - }; - let users = SlackUsers::empty(); - let out = extract_messages(&data, &channel, &users); - assert_eq!(out.len(), 2); - assert_eq!(out[0].channel_id, "C1"); - assert_eq!(out[0].channel_name, "eng"); - assert!(!out[0].is_private); - assert_eq!(out[0].author, "U1"); - assert_eq!(out[0].author_id, "U1"); - assert_eq!(out[0].text, "hi"); - assert_eq!(out[0].timestamp.timestamp(), 1_714_003_200); - } - - #[test] - fn extract_messages_resolves_authors_and_mentions() { - let data = json!({ - "messages": [ - {"ts": "1714003200.0", "user": "U1", "text": "ping <@U2> about the migration"} - ] - }); - let channel = SlackChannel { - id: "C1".into(), - name: "eng".into(), - is_private: false, - }; - let mut m = HashMap::new(); - m.insert("U1".into(), "alice".into()); - m.insert("U2".into(), "bob".into()); - let users = SlackUsers::from_map(m); - let out = extract_messages(&data, &channel, &users); - assert_eq!(out.len(), 1); - assert_eq!(out[0].author, "alice"); - assert_eq!(out[0].author_id, "U1"); - assert_eq!(out[0].text, "ping @bob about the migration"); - } - - #[test] - fn extract_search_messages_enriches_from_channel_map() { - let data = json!({ - "messages": [ - {"ts": "1714003200.0", "user": "U1", "text": "hello", "channel_id": "C1"}, - {"ts": "1714003300.0", "user": "U2", "text": "world", "channel_id": "C2"}, - ] - }); - let mut channel_map = HashMap::new(); - channel_map.insert( - "C1".to_string(), - SlackChannel { - id: "C1".into(), - name: "eng".into(), - is_private: false, - }, - ); - // C2 not in map — should still work with empty channel_name - let users = SlackUsers::empty(); - let out = extract_search_messages(&data, &channel_map, &users); - assert_eq!(out.len(), 2); - assert_eq!(out[0].channel_name, "eng"); - assert_eq!(out[1].channel_name, ""); - } - - #[test] - fn extract_next_cursor_finds_response_metadata_path() { - let data = json!({ - "data": { - "response_metadata": { "next_cursor": "dXNlcjpVMDY..." } - } - }); - assert_eq!( - extract_next_cursor(&data), - Some("dXNlcjpVMDY...".to_string()) - ); - } - - #[test] - fn extract_next_cursor_none_when_blank() { - let data = json!({"data": {"response_metadata": {"next_cursor": " "}}}); - assert!(extract_next_cursor(&data).is_none()); - } - - #[test] - fn encode_decode_roundtrip() { - let mut map = ChannelCursors::new(); - map.insert("C1".into(), "1714003200.123456".into()); - map.insert("C2".into(), "1714010000.000100".into()); - let encoded = encode_cursors(&map); - let decoded = decode_cursors(Some(&encoded)); - assert_eq!(decoded, map); - } - - #[test] - fn decode_empty_cursor_returns_empty_map() { - assert!(decode_cursors(None).is_empty()); - assert!(decode_cursors(Some("")).is_empty()); - assert!(decode_cursors(Some("not json")).is_empty()); - } - - #[test] - fn parse_ts_accepts_slack_decimal_format() { - let dt = parse_ts("1714003200.000100").unwrap(); - assert_eq!(dt.timestamp(), 1_714_003_200); - } - - #[test] - fn extract_search_messages_drops_match_with_missing_channel_id() { - // A search hit with no `channel_id` would otherwise funnel into a - // single empty-channel bucket and ingest under no channel context. - let data = json!({ - "messages": [ - {"ts": "1714003200.0", "user": "U1", "text": "valid", "channel_id": "C1"}, - {"ts": "1714003300.0", "user": "U2", "text": "orphan"}, - {"ts": "1714003400.0", "user": "U3", "text": "blank-id", "channel_id": ""}, - ] - }); - let mut channel_map = HashMap::new(); - channel_map.insert( - "C1".to_string(), - SlackChannel { - id: "C1".into(), - name: "eng".into(), - is_private: false, - }, - ); - let users = SlackUsers::empty(); - let out = extract_search_messages(&data, &channel_map, &users); - assert_eq!(out.len(), 1, "only the well-formed match should pass"); - assert_eq!(out[0].channel_id, "C1"); - } - - #[test] - fn parse_ts_components_preserves_microseconds() { - // Two messages in the same wall-clock second must order by their - // micro suffix — without this, cursor advancement loses precision - // and incremental fetches replay duplicates. - let earlier = parse_ts_components("1714003200.000100"); - let later = parse_ts_components("1714003200.999999"); - assert!(later > earlier); - assert_eq!(parse_ts_components("garbage"), (0, 0)); - } - - #[test] - fn parse_ts_rejects_garbage() { - assert!(parse_ts("").is_none()); - assert!(parse_ts("not.a.number").is_none()); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/slack/types.rs b/src/openhuman/memory_sync/composio/providers/slack/types.rs index c87d27756..42714427a 100644 --- a/src/openhuman/memory_sync/composio/providers/slack/types.rs +++ b/src/openhuman/memory_sync/composio/providers/slack/types.rs @@ -1,14 +1,12 @@ //! Canonical types for the Composio-backed Slack provider. //! //! These types are independent of the Composio/Slack API payload shape. -//! Parsing of raw JSON into these structs happens in -//! [`super::sync`]; everything downstream deals only with the -//! canonical types below. +//! They remain as compatibility wire types for product RPC responses and +//! backfill reporting; tinycortex owns runtime parsing and ingestion. //! //! The old `Bucket` struct (6-hour UTC window) has been removed — the //! memory tree's L0 seal cascade handles batching after PR #1348, so -//! the provider just collects all fetched messages and calls -//! `ingest_page_into_memory_tree` per channel. +//! tinycortex owns batching and incremental persistence. use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; diff --git a/src/openhuman/memory_sync/composio/providers/slack/users.rs b/src/openhuman/memory_sync/composio/providers/slack/users.rs deleted file mode 100644 index 73a0c8763..000000000 --- a/src/openhuman/memory_sync/composio/providers/slack/users.rs +++ /dev/null @@ -1,364 +0,0 @@ -//! Slack user-id → display-name resolver. -//! -//! Slack's `conversations.history` payload references users by their -//! workspace-stable id (e.g. `U01Q1TBL20P`) in two places: -//! -//! 1. The `user` field on each message (the author). -//! 2. Inline `<@U01Q1TBL20P>` mention syntax inside `text`. -//! -//! Neither is human-readable. To make canonical chat transcripts useful -//! for retrieval (and for humans reading the seal-cascade summaries), -//! we fetch the workspace's user directory once per sync run, build an -//! id → display-name map, and apply it both to the author field and to -//! every `<@…>` mention in message bodies. -//! -//! ## Cache scope -//! -//! Per-sync only. Each `SlackProvider::sync()` invocation calls -//! [`SlackUsers::fetch`] once before walking channels. The map lives -//! in a local variable for the duration of the sync, then drops. -//! Slack's user list rarely changes within a 15-minute sync window, -//! and re-fetching per sync keeps stale-cache risk near zero without -//! adding persistence machinery. -//! -//! ## Soft-fallback contract -//! -//! Following the pattern of [`crate::openhuman::memory_sync::composio::providers::slack::sync::extract_messages`] -//! and the [`super::provider::SlackProvider::sync`] error handling, a -//! failure to fetch users is **not fatal**. The returned [`SlackUsers`] -//! is empty, and `resolve()` / `replace_mentions()` pass through raw -//! ids unchanged — same behaviour as before this module existed. - -use regex::Regex; -use serde_json::{json, Value}; -use std::collections::HashMap; -use std::sync::OnceLock; - -use crate::openhuman::memory_sync::composio::providers::ProviderContext; - -/// Composio action slug for the bulk user listing. -const ACTION_LIST_USERS: &str = "SLACK_LIST_ALL_USERS"; - -/// Page size — Slack caps at 1000; 200 keeps each page small. -const PAGE_SIZE: u32 = 200; - -/// Maximum pages to walk per sync. With `PAGE_SIZE = 200` this covers -/// workspaces up to 4000 users without complaint. Beyond that the tail -/// is truncated and unresolved ids will pass through verbatim. -const MAX_PAGES: u32 = 20; - -/// Slack mention syntax: `<@U01Q1TBL20P>`. Captures the bare id so we -/// can drop the wrapper when substituting in a resolved name. -fn mention_re() -> &'static Regex { - static RE: OnceLock = OnceLock::new(); - RE.get_or_init(|| Regex::new(r"<@(U[A-Z0-9]+)>").expect("static mention regex compiles")) -} - -/// Map of Slack user id → human-readable display name. -#[derive(Debug, Default, Clone)] -pub struct SlackUsers { - map: HashMap, -} - -impl SlackUsers { - /// Empty map — `resolve()` passes through raw ids verbatim. - pub fn empty() -> Self { - Self::default() - } - - /// Number of users in the cache. - pub fn len(&self) -> usize { - self.map.len() - } - - pub fn is_empty(&self) -> bool { - self.map.is_empty() - } - - /// Resolve a Slack user id to a display name. Returns the input id - /// unchanged when no mapping exists — matches the - /// resolve-or-passthrough contract of the parent provider. - pub fn resolve(&self, user_id: &str) -> String { - self.map - .get(user_id) - .cloned() - .unwrap_or_else(|| user_id.to_string()) - } - - /// Replace every `<@Uxxx>` mention in `text` with `@`. - /// Unknown ids stay as `@Uxxx` (the wrapper is removed but the id - /// is preserved so retrieval can still surface them). - pub fn replace_mentions(&self, text: &str) -> String { - mention_re() - .replace_all(text, |caps: ®ex::Captures| { - let id = &caps[1]; - let resolved = self.map.get(id).map(String::as_str).unwrap_or(id); - format!("@{resolved}") - }) - .into_owned() - } - - /// Pull the workspace user directory via Composio. Soft-fails to - /// [`SlackUsers::empty`] on transport, HTTP, JSON, or - /// provider-failure errors so the sync can continue with raw ids. - /// - /// Returns `(users, total_attempts)` where `total_attempts` sums every - /// real Composio call this fetch made across pages and rate-limit - /// retries, so the caller can charge the daily quota meter - /// accurately. Pages walked silently are tracked too — without this, - /// large workspaces under-report their request usage. - pub async fn fetch(ctx: &ProviderContext) -> (Self, u32) { - let mut map: HashMap = HashMap::new(); - let mut cursor: Option = None; - let mut total_attempts: u32 = 0; - - for page_num in 0..MAX_PAGES { - let mut args = json!({ "limit": PAGE_SIZE }); - if let Some(ref c) = cursor { - args["cursor"] = json!(c); - } - - // Going through `execute_with_retry` so a transient - // `ratelimited` page doesn't drop us into a half-built - // directory while the rest of the provider uses backoff. - // Soft-fall to whatever was collected so far on any failure. - let (resp, attempts) = match super::provider::execute_with_retry( - ctx, - ACTION_LIST_USERS, - args, - &format!("{ACTION_LIST_USERS} page {page_num}"), - ) - .await - { - Ok(t) => t, - Err(err) => { - // We don't know exactly how many attempts the helper - // burned before bailing, but at least one ran — count - // it so the budget meter doesn't silently undercount. - total_attempts = total_attempts.saturating_add(1); - log::warn!( - "[composio:slack:users] {ACTION_LIST_USERS} page {page_num} failed: {err} — \ - degrading to raw ids for the rest of this sync" - ); - return (Self { map }, total_attempts); - } - }; - total_attempts = total_attempts.saturating_add(attempts); - - super::provider::dump_response("_meta", "users", page_num, &resp.data); - absorb_page(&resp.data, &mut map); - - cursor = extract_next_cursor(&resp.data); - if cursor.is_none() { - break; - } - } - - log::info!( - "[composio:slack:users] resolved {} workspace users in {total_attempts} call(s)", - map.len() - ); - (Self { map }, total_attempts) - } - - /// Construct from a pre-built map. Test-only — production callers - /// should use [`Self::fetch`] or [`Self::empty`]. - #[cfg(test)] - pub fn from_map(map: HashMap) -> Self { - Self { map } - } -} - -/// Walk a Composio response envelope and absorb every user object's -/// `id` + best-available display name into `map`. -fn absorb_page(data: &Value, map: &mut HashMap) { - let candidates = [ - data.pointer("/data/members"), - data.pointer("/members"), - data.pointer("/data/users"), - data.pointer("/users"), - data.pointer("/data/data/members"), - ]; - let arr = candidates - .into_iter() - .flatten() - .find_map(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - for raw in arr { - let id = raw - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .trim() - .to_string(); - if id.is_empty() { - continue; - } - if let Some(name) = pick_display_name(&raw) { - map.insert(id, name); - } - } -} - -/// Slack returns several name fields per user. Prefer the most -/// human-readable, fall back through real_name → name → display_name. -fn pick_display_name(raw: &Value) -> Option { - let candidates = [ - raw.pointer("/profile/display_name"), - raw.pointer("/profile/real_name"), - raw.get("real_name"), - raw.get("name"), - raw.pointer("/profile/display_name_normalized"), - raw.pointer("/profile/real_name_normalized"), - ]; - for cand in candidates.into_iter().flatten() { - if let Some(s) = cand.as_str() { - let trimmed = s.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); - } - } - } - None -} - -fn extract_next_cursor(data: &Value) -> Option { - let candidates = [ - data.pointer("/data/response_metadata/next_cursor"), - data.pointer("/response_metadata/next_cursor"), - data.pointer("/data/next_cursor"), - data.pointer("/next_cursor"), - ]; - for cand in candidates.into_iter().flatten() { - if let Some(s) = cand.as_str() { - let trimmed = s.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); - } - } - } - None -} - -#[cfg(test)] -mod tests { - use super::*; - - fn sample_users() -> SlackUsers { - let mut m = HashMap::new(); - m.insert("U001".to_string(), "alice".to_string()); - m.insert("U002".to_string(), "bob".to_string()); - SlackUsers::from_map(m) - } - - #[test] - fn resolve_known_id_returns_name() { - let u = sample_users(); - assert_eq!(u.resolve("U001"), "alice"); - } - - #[test] - fn resolve_unknown_id_passes_through() { - let u = sample_users(); - assert_eq!(u.resolve("U999"), "U999"); - } - - #[test] - fn empty_passes_through_every_id() { - let u = SlackUsers::empty(); - assert_eq!(u.resolve("U001"), "U001"); - assert_eq!(u.replace_mentions("hi <@U001>"), "hi @U001"); - } - - #[test] - fn replace_mentions_substitutes_known_ids() { - let u = sample_users(); - let out = u.replace_mentions("Hi <@U001>, please ping <@U002>."); - assert_eq!(out, "Hi @alice, please ping @bob."); - } - - #[test] - fn replace_mentions_strips_wrapper_for_unknown_id() { - let u = sample_users(); - // Unknown id keeps the raw id but loses the `<@...>` wrapper. - let out = u.replace_mentions("ping <@U999>"); - assert_eq!(out, "ping @U999"); - } - - #[test] - fn replace_mentions_leaves_non_mention_text_alone() { - let u = sample_users(); - let out = u.replace_mentions("no mentions here, just "); - assert_eq!(out, "no mentions here, just "); - } - - #[test] - fn replace_mentions_handles_multiple_in_one_line() { - let u = sample_users(); - let out = u.replace_mentions("<@U001> said hi to <@U001> and <@U002>"); - assert_eq!(out, "@alice said hi to @alice and @bob"); - } - - #[test] - fn absorb_page_reads_data_members_path() { - let data = json!({ - "data": { - "members": [ - { - "id": "U001", - "profile": { "display_name": "alice", "real_name": "Alice Smith" } - }, - { - "id": "U002", - "profile": { "display_name": "" , "real_name": "Bob Jones" } - }, - { - "id": "", - "profile": { "display_name": "skipped" } - } - ] - } - }); - let mut m = HashMap::new(); - absorb_page(&data, &mut m); - assert_eq!(m.get("U001").unwrap(), "alice"); - // Falls back to real_name when display_name is blank. - assert_eq!(m.get("U002").unwrap(), "Bob Jones"); - // Empty id row is dropped. - assert!(!m.contains_key("")); - } - - #[test] - fn pick_display_name_prefers_display_name_over_real_name() { - let raw = json!({ - "profile": { "display_name": "alice", "real_name": "Alice Smith" } - }); - assert_eq!(pick_display_name(&raw).as_deref(), Some("alice")); - } - - #[test] - fn pick_display_name_falls_back_to_name() { - let raw = json!({ "name": "alice", "profile": {} }); - assert_eq!(pick_display_name(&raw).as_deref(), Some("alice")); - } - - #[test] - fn pick_display_name_returns_none_when_all_blank() { - let raw = json!({ "profile": { "display_name": " " }, "name": "" }); - assert!(pick_display_name(&raw).is_none()); - } - - #[test] - fn extract_next_cursor_finds_response_metadata() { - let data = json!({"data": {"response_metadata": {"next_cursor": "abc123"}}}); - assert_eq!(extract_next_cursor(&data).as_deref(), Some("abc123")); - } - - #[test] - fn extract_next_cursor_none_when_blank() { - let data = json!({"response_metadata": {"next_cursor": " "}}); - assert!(extract_next_cursor(&data).is_none()); - } -} diff --git a/src/openhuman/memory_sync/composio/providers/sync_state.rs b/src/openhuman/memory_sync/composio/providers/sync_state.rs index 2a2ec00a2..5650cb534 100644 --- a/src/openhuman/memory_sync/composio/providers/sync_state.rs +++ b/src/openhuman/memory_sync/composio/providers/sync_state.rs @@ -1,605 +1,19 @@ -//! Persistent sync state for Composio providers. -//! -//! Each `(toolkit, connection_id)` pair gets its own [`SyncState`] persisted -//! in the local KV store. The state tracks: -//! -//! * **Cursor** — a provider-specific watermark (e.g. a timestamp or page -//! token) so the next sync can skip items already seen. -//! * **Synced IDs** — a set of item identifiers that have been written to -//! memory. Items in this set are skipped even if they appear again in -//! an API response (deduplication). -//! * **Daily request budget** — a rolling counter keyed by calendar date -//! (`YYYY-MM-DD`) that caps the number of `execute_tool` calls a -//! provider makes per day. Resets automatically when the date rolls -//! over. -//! -//! All persistence goes through [`crate::openhuman::memory_store::MemoryClient`]'s -//! KV surface (`kv_set` / `kv_get` under a dedicated namespace), so the -//! state survives process restarts without any extra file management. +//! Compatibility exports for sync state now owned by tinycortex. -use std::collections::HashSet; +pub use tinycortex::memory::sync::state::DEFAULT_DAILY_REQUEST_LIMIT; +pub use tinycortex::memory::sync::{DailyBudget, SyncState}; -use chrono::Utc; -use serde::{Deserialize, Serialize}; -use serde_json::json; +pub const KV_NAMESPACE: &str = crate::openhuman::tinycortex::HOST_SYNC_STATE_NAMESPACE; -use crate::openhuman::memory_store::MemoryClientRef; - -/// Maximum API requests a single provider connection may make per calendar -/// day. This covers the initial backfill case where there are thousands of -/// unsynced items — after this many requests the provider yields and -/// continues on the next day. -/// -/// Compile-time default. The runtime value used by [`DailyBudget::default`] -/// is resolved by [`resolved_daily_request_limit`], which honors the -/// `OPENHUMAN_COMPOSIO_DAILY_REQUEST_LIMIT` env var so operators can widen -/// (or tighten) the cap without recompiling. See the project `.env.example` -/// for documentation. -pub const DEFAULT_DAILY_REQUEST_LIMIT: u32 = 500; - -/// Environment variable read by [`resolved_daily_request_limit`] to override -/// [`DEFAULT_DAILY_REQUEST_LIMIT`]. Must parse as a positive `u32`; values -/// `< 1` or non-numeric content fall back to the default with a `warn`. -pub const ENV_DAILY_REQUEST_LIMIT: &str = "OPENHUMAN_COMPOSIO_DAILY_REQUEST_LIMIT"; - -/// Resolve the effective per-day request limit. Reads -/// [`ENV_DAILY_REQUEST_LIMIT`] if set; otherwise returns -/// [`DEFAULT_DAILY_REQUEST_LIMIT`]. A non-positive or unparseable value -/// is rejected with a `warn` log and the default is used — we never -/// silently honor `0` because that would freeze every provider's sync -/// from the first tick. -pub fn resolved_daily_request_limit() -> u32 { - match std::env::var(ENV_DAILY_REQUEST_LIMIT) { - Ok(s) => match s.trim().parse::() { - Ok(n) if n >= 1 => n, - _ => { - static WARNED: std::sync::Once = std::sync::Once::new(); - WARNED.call_once(|| { - tracing::warn!( - env = ENV_DAILY_REQUEST_LIMIT, - value = %s, - default = DEFAULT_DAILY_REQUEST_LIMIT, - "[composio:sync-state] env override not a positive u32; using default" - ); - }); - DEFAULT_DAILY_REQUEST_LIMIT - } - }, - Err(_) => DEFAULT_DAILY_REQUEST_LIMIT, - } -} - -/// KV namespace under which all sync state keys live. Separate from the -/// memory document namespaces (`skill-gmail`, etc.) to avoid collisions. -pub const KV_NAMESPACE: &str = "composio-sync-state"; - -/// Persistent sync state for one `(toolkit, connection_id)` pair. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SyncState { - /// Toolkit slug, e.g. `"gmail"`. - pub toolkit: String, - /// Connection id, e.g. `"conn_abc123"`. - pub connection_id: String, - - /// Provider-specific cursor. For Gmail this is the internal-date - /// (epoch millis) of the newest synced message; for Notion it is the - /// `last_edited_time` ISO string of the most recently synced page. - /// `None` means "never synced — start from scratch". - #[serde(default)] - pub cursor: Option, - - /// Set of item IDs that have already been persisted to memory. - /// Used for deduplication: if an item appears in an API response - /// but its ID is in this set, skip it. - #[serde(default)] - pub synced_ids: HashSet, - - /// Rolling daily request budget. - #[serde(default)] - pub daily_budget: DailyBudget, - - /// ID of the most recently synced item, used by providers (Gmail - /// today) to short-circuit a tick when the freshest server-side - /// item matches what we already have. Cheaper than re-walking the - /// `synced_ids` set — and crucially, it lets us bail out of - /// pagination on the very first page when nothing has changed, - /// instead of fetching `MAX_PAGES_PER_SYNC` worth of duplicates - /// before falling through the per-page dedup loop. - /// - /// `None` either when the state is fresh or when an older state - /// blob was loaded from disk that pre-dates this field. - #[serde(default)] - pub last_seen_id: Option, - - /// Unix milliseconds of the last successful sync that wrote into - /// memory. Lets the adaptive page-cap logic distinguish a "we - /// synced 30 seconds ago" tick (cap pages aggressively) from a - /// "we last synced two hours ago" tick (let pagination run to the - /// usual ceiling). Independent of the periodic scheduler's - /// in-process `LAST_SYNC_AT` map because that map is rebuilt on - /// every process restart whereas this value survives restarts. - /// - /// `None` until the first successful sync. - #[serde(default)] - pub last_sync_at_ms: Option, -} - -/// Tracks the number of API requests made on a given calendar day. -/// Automatically resets when the date rolls over. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct DailyBudget { - /// Calendar date in `YYYY-MM-DD` format. - pub date: String, - /// Number of `execute_tool` requests made so far today. - pub requests_used: u32, - /// Maximum requests allowed per day. - pub limit: u32, -} - -impl Default for DailyBudget { - fn default() -> Self { - Self { - date: today_str(), - requests_used: 0, - limit: resolved_daily_request_limit(), - } - } -} - -impl DailyBudget { - /// Remaining requests available today. If the stored date is stale - /// (a previous day), this returns the full limit because the budget - /// will be reset on the next [`Self::record_request`] call. - pub fn remaining(&self) -> u32 { - if self.date != today_str() { - return self.limit; - } - self.limit.saturating_sub(self.requests_used) - } - - /// Returns `true` if the daily budget is exhausted for today. - pub fn is_exhausted(&self) -> bool { - self.remaining() == 0 - } - - /// Record `n` API requests. If the date has rolled over, resets the - /// counter before adding. - pub fn record_requests(&mut self, n: u32) { - let today = today_str(); - if self.date != today { - tracing::info!( - old_date = %self.date, - new_date = %today, - requests_used = self.requests_used, - limit = self.limit, - "[composio:sync-state] daily request budget reset" - ); - self.date = today; - self.requests_used = 0; - } - self.requests_used = self.requests_used.saturating_add(n); - } - - /// Record a single API request. - pub fn record_request(&mut self) { - self.record_requests(1); - } -} - -impl SyncState { - /// Create a fresh state for a new connection (never synced). - 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(), - daily_budget: DailyBudget::default(), - last_seen_id: None, - last_sync_at_ms: None, - } - } - - /// Record the freshest item id observed on a successful sync. - /// Idempotent — repeated calls with the same id are no-ops. - pub fn set_last_seen_id(&mut self, item_id: impl Into) { - self.last_seen_id = Some(item_id.into()); - } - - /// Record the wall-clock time of a successful sync (unix - /// milliseconds). Persisted alongside the cursor so the adaptive - /// page-cap survives process restarts. - pub fn set_last_sync_at_ms(&mut self, ms: u64) { - self.last_sync_at_ms = Some(ms); - } - - /// Whether the daily request budget is exhausted. - pub fn budget_exhausted(&self) -> bool { - self.daily_budget.is_exhausted() - } - - /// Remaining API requests for today. - pub fn budget_remaining(&self) -> u32 { - self.daily_budget.remaining() - } - - /// Record API requests made. - pub fn record_requests(&mut self, n: u32) { - self.daily_budget.record_requests(n); - } - - /// Check if an item ID has already been synced. - pub fn is_synced(&self, item_id: &str) -> bool { - self.synced_ids.contains(item_id) - } - - /// Mark an item ID as synced. - pub fn mark_synced(&mut self, item_id: impl Into) { - self.synced_ids.insert(item_id.into()); - } - - /// Update the cursor to a new watermark value. - pub fn advance_cursor(&mut self, cursor: impl Into) { - self.cursor = Some(cursor.into()); - } - - /// KV key for this state. Deterministic so load + save are symmetric. - fn kv_key(&self) -> String { - format!("{}:{}", self.toolkit, self.connection_id) - } - - /// Load sync state from the KV store, or return a fresh default if - /// none exists. - pub async fn load( - memory: &MemoryClientRef, - toolkit: &str, - connection_id: &str, - ) -> Result { - let key = format!("{toolkit}:{connection_id}"); - match memory.kv_get(Some(KV_NAMESPACE), &key).await? { - Some(value) => { - let mut state: SyncState = serde_json::from_value(value) - .map_err(|e| format!("[sync_state] deserialize failed for {key}: {e}"))?; - // Ensure budget rolls over if date changed. - if state.daily_budget.date != today_str() { - tracing::debug!( - toolkit, - connection_id, - old_date = %state.daily_budget.date, - "[sync_state] daily budget rolled over" - ); - state.daily_budget.date = today_str(); - state.daily_budget.requests_used = 0; - } - tracing::debug!( - toolkit, - connection_id, - cursor = ?state.cursor, - synced_ids_count = state.synced_ids.len(), - budget_remaining = state.budget_remaining(), - "[sync_state] loaded" - ); - Ok(state) - } - None => { - tracing::debug!( - toolkit, - connection_id, - "[sync_state] no existing state, starting fresh" - ); - Ok(Self::new(toolkit, connection_id)) - } - } - } - - /// Persist the current state to the KV store. - pub async fn save(&self, memory: &MemoryClientRef) -> Result<(), String> { - let key = self.kv_key(); - let value = serde_json::to_value(self) - .map_err(|e| format!("[sync_state] serialize failed: {e}"))?; - memory.kv_set(Some(KV_NAMESPACE), &key, &value).await?; - tracing::debug!( - toolkit = %self.toolkit, - connection_id = %self.connection_id, - cursor = ?self.cursor, - synced_ids_count = self.synced_ids.len(), - budget_used = self.daily_budget.requests_used, - "[sync_state] saved" - ); - Ok(()) - } -} - -/// Today's date as `YYYY-MM-DD` in UTC. -fn today_str() -> String { - Utc::now().format("%Y-%m-%d").to_string() -} - -/// Extract an ID string from a JSON value, trying multiple candidate paths. -/// Returns the first non-empty string found. pub fn extract_item_id(item: &serde_json::Value, paths: &[&str]) -> Option { - for path in paths { - let mut cur = item; - let mut ok = true; - for segment in path.split('.') { - match cur.get(segment) { - Some(next) => cur = next, - None => { - ok = false; - break; - } - } - } - if !ok { - continue; - } - if let Some(s) = cur.as_str() { - let trimmed = s.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); - } - } - } - None -} - -/// Helper to persist a single item as its own memory document. -/// -/// Each item is stored under the provider's memory namespace with a -/// deterministic `document_id` so repeated syncs upsert rather than -/// duplicate. Returns the document ID on success. -pub async fn persist_single_item( - memory: &MemoryClientRef, - namespace_skill_id: &str, - document_id: &str, - title: &str, - item: &serde_json::Value, - toolkit: &str, - connection_id: Option<&str>, -) -> Result { - let content = serde_json::to_string_pretty(item).unwrap_or_else(|_| "{}".to_string()); - memory - .store_skill_sync( - namespace_skill_id, - connection_id.unwrap_or("default"), - title, - &content, - Some("composio-sync".to_string()), - Some(json!({ - "toolkit": toolkit, - "connection_id": connection_id, - "source": "composio-provider-incremental", - })), - Some("medium".to_string()), - None, - None, - Some(document_id.to_string()), - ) - .await?; - Ok(document_id.to_string()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn daily_budget_defaults_to_full() { - let b = DailyBudget::default(); - assert_eq!(b.remaining(), DEFAULT_DAILY_REQUEST_LIMIT); - assert!(!b.is_exhausted()); - } - - #[test] - fn daily_budget_tracks_requests() { - let mut b = DailyBudget::default(); - b.record_requests(100); - assert_eq!(b.remaining(), DEFAULT_DAILY_REQUEST_LIMIT - 100); - assert!(!b.is_exhausted()); - } - - #[test] - fn daily_budget_exhaustion() { - let mut b = DailyBudget::default(); - b.record_requests(DEFAULT_DAILY_REQUEST_LIMIT); - assert_eq!(b.remaining(), 0); - assert!(b.is_exhausted()); - } - - #[test] - fn daily_budget_saturates_on_overflow() { - let mut b = DailyBudget::default(); - b.record_requests(DEFAULT_DAILY_REQUEST_LIMIT + 100); - assert_eq!(b.remaining(), 0); - } - - #[test] - fn daily_budget_resets_on_date_change() { - let mut b = DailyBudget { - date: "2025-01-01".to_string(), - requests_used: 499, - limit: DEFAULT_DAILY_REQUEST_LIMIT, - }; - // Calling remaining() when date is stale returns full limit. - assert_eq!(b.remaining(), DEFAULT_DAILY_REQUEST_LIMIT); - // Recording a request resets the counter. - b.record_request(); - assert_eq!(b.date, today_str()); - assert_eq!(b.requests_used, 1); - } - - #[test] - fn sync_state_deduplication() { - let mut state = SyncState::new("gmail", "conn_1"); - assert!(!state.is_synced("msg_abc")); - state.mark_synced("msg_abc"); - assert!(state.is_synced("msg_abc")); - assert!(!state.is_synced("msg_xyz")); - } - - #[test] - fn sync_state_cursor_advancement() { - let mut state = SyncState::new("notion", "conn_2"); - assert!(state.cursor.is_none()); - state.advance_cursor("2026-04-01T00:00:00Z"); - assert_eq!(state.cursor.as_deref(), Some("2026-04-01T00:00:00Z")); - state.advance_cursor("2026-04-10T00:00:00Z"); - assert_eq!(state.cursor.as_deref(), Some("2026-04-10T00:00:00Z")); - } - - #[test] - fn sync_state_serialization_roundtrip() { - let mut state = SyncState::new("gmail", "conn_test"); - state.advance_cursor("12345"); - state.mark_synced("item_a"); - state.mark_synced("item_b"); - state.daily_budget.record_requests(42); - state.set_last_seen_id("msg_top"); - state.set_last_sync_at_ms(1_700_000_000_000); - - let json = serde_json::to_value(&state).unwrap(); - let restored: SyncState = serde_json::from_value(json).unwrap(); - - assert_eq!(restored.toolkit, "gmail"); - assert_eq!(restored.connection_id, "conn_test"); - assert_eq!(restored.cursor.as_deref(), Some("12345")); - assert!(restored.synced_ids.contains("item_a")); - assert!(restored.synced_ids.contains("item_b")); - assert_eq!(restored.synced_ids.len(), 2); - assert_eq!(restored.daily_budget.requests_used, 42); - assert_eq!(restored.last_seen_id.as_deref(), Some("msg_top")); - assert_eq!(restored.last_sync_at_ms, Some(1_700_000_000_000)); - } - - #[test] - fn sync_state_deserializes_legacy_blob_without_new_fields() { - // Older state blobs serialized before #1404 had no - // `last_seen_id` / `last_sync_at_ms` keys — make sure the - // deserializer still accepts them so existing users don't - // lose their cursor + dedup set on first upgrade. - let legacy = serde_json::json!({ - "toolkit": "gmail", - "connection_id": "conn_old", - "cursor": "1699000000000", - "synced_ids": ["m1", "m2"], - "daily_budget": { "date": today_str(), "requests_used": 7, "limit": 500 } - }); - let restored: SyncState = serde_json::from_value(legacy).unwrap(); - assert_eq!(restored.cursor.as_deref(), Some("1699000000000")); - assert_eq!(restored.synced_ids.len(), 2); - assert!(restored.last_seen_id.is_none()); - assert!(restored.last_sync_at_ms.is_none()); - } - - #[test] - fn set_last_seen_id_overwrites_previous_value() { - let mut state = SyncState::new("gmail", "c"); - state.set_last_seen_id("a"); - state.set_last_seen_id("b"); - assert_eq!(state.last_seen_id.as_deref(), Some("b")); - } - - #[test] - fn set_last_sync_at_ms_records_value() { - let mut state = SyncState::new("gmail", "c"); - state.set_last_sync_at_ms(123); - state.set_last_sync_at_ms(456); - assert_eq!(state.last_sync_at_ms, Some(456)); - } - - #[test] - fn extract_item_id_walks_paths() { - let item = serde_json::json!({ - "id": "top_level", - "data": { "id": "nested" } - }); - assert_eq!( - extract_item_id(&item, &["data.id", "id"]), - Some("nested".to_string()) - ); - assert_eq!( - extract_item_id(&item, &["missing", "id"]), - Some("top_level".to_string()) - ); - assert_eq!(extract_item_id(&item, &["nope"]), None); - } - - #[test] - fn kv_key_is_deterministic() { - let s1 = SyncState::new("gmail", "conn_x"); - let s2 = SyncState::new("gmail", "conn_x"); - assert_eq!(s1.kv_key(), s2.kv_key()); - assert_eq!(s1.kv_key(), "gmail:conn_x"); - } - - /// RAII guard that save→set→restore an env var so the test does not - /// leak state to sibling tests in the same process. - struct EnvGuard { - key: &'static str, - previous: Option, - } - - impl EnvGuard { - fn set(key: &'static str, value: &str) -> Self { - let previous = std::env::var(key).ok(); - std::env::set_var(key, value); - Self { key, previous } - } - fn unset(key: &'static str) -> Self { - let previous = std::env::var(key).ok(); - std::env::remove_var(key); - Self { key, previous } - } - } - - impl Drop for EnvGuard { - fn drop(&mut self) { - match self.previous.take() { - Some(v) => std::env::set_var(self.key, v), - None => std::env::remove_var(self.key), - } - } - } - - // Env-var override scenarios are bundled into one `#[test]` so they - // run sequentially within a single thread — `cargo test` parallelism - // across `#[test]` fns would race on `OPENHUMAN_COMPOSIO_DAILY_REQUEST_LIMIT`. - #[test] - fn resolved_daily_request_limit_honors_env() { - let _lock = crate::openhuman::config::TEST_ENV_LOCK - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - - // Unset → default. - let _g = EnvGuard::unset(ENV_DAILY_REQUEST_LIMIT); - assert_eq!(resolved_daily_request_limit(), DEFAULT_DAILY_REQUEST_LIMIT); - assert_eq!(DailyBudget::default().limit, DEFAULT_DAILY_REQUEST_LIMIT); - drop(_g); - - // Valid override widens the cap. - let _g = EnvGuard::set(ENV_DAILY_REQUEST_LIMIT, "5000"); - assert_eq!(resolved_daily_request_limit(), 5000); - assert_eq!(DailyBudget::default().limit, 5000); - drop(_g); - - // Trims surrounding whitespace. - let _g = EnvGuard::set(ENV_DAILY_REQUEST_LIMIT, " 750 "); - assert_eq!(resolved_daily_request_limit(), 750); - drop(_g); - - // Zero rejected — would otherwise freeze every sync. - let _g = EnvGuard::set(ENV_DAILY_REQUEST_LIMIT, "0"); - assert_eq!(resolved_daily_request_limit(), DEFAULT_DAILY_REQUEST_LIMIT); - drop(_g); - - // Non-numeric rejected. - let _g = EnvGuard::set(ENV_DAILY_REQUEST_LIMIT, "lots"); - assert_eq!(resolved_daily_request_limit(), DEFAULT_DAILY_REQUEST_LIMIT); - drop(_g); - - // Negative rejected (won't parse as u32). - let _g = EnvGuard::set(ENV_DAILY_REQUEST_LIMIT, "-1"); - assert_eq!(resolved_daily_request_limit(), DEFAULT_DAILY_REQUEST_LIMIT); - drop(_g); - } + paths.iter().find_map(|path| { + let value = path + .split('.') + .try_fold(item, |current, segment| current.get(segment))?; + value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }) } diff --git a/src/openhuman/memory_sync/composio/providers/traits.rs b/src/openhuman/memory_sync/composio/providers/traits.rs index 06d19fbd2..409d2f018 100644 --- a/src/openhuman/memory_sync/composio/providers/traits.rs +++ b/src/openhuman/memory_sync/composio/providers/traits.rs @@ -50,10 +50,39 @@ pub trait ComposioProvider: Send + Sync { ctx: &ProviderContext, ) -> Result; - /// Run a sync pass for the current connection in `ctx`. Implementations - /// are responsible for persisting whatever they fetch (typically into - /// the memory layer via [`ProviderContext::memory_client`]). - async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result; + /// Compatibility entry point backed exclusively by tinycortex. + async fn sync(&self, ctx: &ProviderContext, reason: SyncReason) -> Result { + let connection_id = ctx.connection_id.as_deref().ok_or_else(|| { + format!( + "[composio:{}] sync missing connection_id", + self.toolkit_slug() + ) + })?; + let started_at_ms = now_ms(); + let outcome = crate::openhuman::tinycortex::run_composio_connection_with_budgets( + self.toolkit_slug(), + connection_id, + ctx.config.as_ref(), + ctx.max_items, + ctx.sync_depth_days, + ) + .await + .map_err(|error| error.to_string())?; + Ok(SyncOutcome { + toolkit: self.toolkit_slug().into(), + connection_id: Some(connection_id.into()), + reason: reason.as_str().into(), + items_ingested: outcome.records_ingested as usize, + started_at_ms, + finished_at_ms: now_ms(), + summary: outcome.note.unwrap_or_else(|| "sync completed".into()), + details: serde_json::json!({ + "more_pending": outcome.more_pending, + "actions_called": outcome.actions_called, + "provider_cost_usd": outcome.provider_cost_usd, + }), + }) + } /// Fetch a filtered set of work items as structured /// [`NormalizedTask`]s — the read path that powers the @@ -112,7 +141,8 @@ pub trait ComposioProvider: Send + Sync { /// Hook fired when an OAuth handoff completes /// ([`crate::core::event_bus::DomainEvent::ComposioConnectionCreated`]). /// - /// Default impl: fetch the user profile, then run an initial sync. + /// Default impl: fetch and persist the user profile. Initial memory + /// ingestion is dispatched separately through tinycortex by the bus. /// Providers can override to add provider-specific bootstrapping /// (e.g. registering Composio triggers, seeding labels, …). async fn on_connection_created(&self, ctx: &ProviderContext) -> Result<(), String> { @@ -120,7 +150,7 @@ pub trait ComposioProvider: Send + Sync { tracing::info!( toolkit = %toolkit, connection_id = ?ctx.connection_id, - "[composio:provider] on_connection_created → fetch_user_profile + initial sync" + "[composio:provider] on_connection_created: fetching user profile" ); match self.fetch_user_profile(ctx).await { Ok(profile) => { @@ -178,13 +208,6 @@ pub trait ComposioProvider: Send + Sync { ); } } - let outcome = self.sync(ctx, SyncReason::ConnectionCreated).await?; - tracing::info!( - toolkit = %toolkit, - items = outcome.items_ingested, - elapsed_ms = outcome.elapsed_ms(), - "[composio:provider] initial sync complete" - ); Ok(()) } @@ -235,6 +258,13 @@ pub trait ComposioProvider: Send + Sync { } } +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + /// Build the env var name read by [`resolve_sync_interval_secs`] for a /// given toolkit slug. Exposed so tests (and `.env.example`) can stay in /// lockstep with the runtime lookup without re-implementing the casing. diff --git a/src/openhuman/memory_sync/sources/audit.rs b/src/openhuman/memory_sync/sources/audit.rs index fdba61c82..305edf232 100644 --- a/src/openhuman/memory_sync/sources/audit.rs +++ b/src/openhuman/memory_sync/sources/audit.rs @@ -1,471 +1,29 @@ -//! Sync audit log — append-only JSONL recording each sync run's token -//! usage and cost. -//! -//! Written to `/memory_tree/sync_audit.jsonl`. Each line is a -//! self-contained JSON object describing one completed sync run. - -use chrono::{DateTime, Utc}; -use serde::Serialize; -use std::io::Write; -use std::path::Path; +//! OpenHuman configuration wrappers for tinycortex sync audit ownership. use crate::openhuman::config::Config; -#[derive(Clone, Debug, Serialize, serde::Deserialize)] -pub struct SyncAuditEntry { - pub timestamp: DateTime, - pub source_id: String, - pub source_kind: String, - pub scope: String, - /// Total items fetched from the source (commits, issues, PRs, etc.). - pub items_fetched: u32, - /// Number of summarise batches produced. - pub batches: u32, - /// Input tokens fed to the summariser. Provider-reported when the - /// backend returned usage for every batch; otherwise estimated as the - /// sum of item bodies / 4. See [`SyncAuditEntry::actual_charged_usd`] - /// for the signal of which path produced the cost figure. - pub input_tokens: u64, - /// Output tokens produced by the summariser. Provider-reported when - /// available, else estimated. - pub output_tokens: u64, - /// Estimated cost in USD (input + output at hardcoded model pricing). - /// Always populated as a fallback so old audit entries — and runs - /// where the backend reported no charge — still render a cost. Prefer - /// [`SyncAuditEntry::actual_charged_usd`] when it is `Some`. - pub estimated_cost_usd: f64, - /// Number of Composio billable API actions executed during this sync - /// (e.g. `GMAIL_FETCH_EMAILS`, `SLACK_LIST_CONVERSATIONS`). `0` for - /// non-Composio source kinds. `#[serde(default)]` keeps audit lines - /// written before #3111 parseable. - #[serde(default)] - pub composio_actions_called: u32, - /// Actual USD charged for those Composio actions, summed from each - /// response's backend-reported `cost_usd`. `0.0` for non-Composio kinds - /// or when the backend reports no charge (e.g. direct mode). - #[serde(default)] - pub composio_cost_usd: f64, - /// Real amount billed by the backend in USD (sum of - /// `openhuman.billing.charged_amount_usd` across batches), when the - /// provider reported it for the run. `None` for runs that fell back to - /// the local estimate and for audit entries written before issue - /// #3110 (the `#[serde(default)]` keeps those deserialising). When - /// `Some`, this is the authoritative cost; renderers should show it in - /// preference to `estimated_cost_usd`. - #[serde(default)] - pub actual_charged_usd: Option, - /// Duration of the sync in milliseconds. - pub duration_ms: u64, - /// Whether the sync completed successfully. - pub success: bool, - /// Error message if the sync failed. - #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, -} +pub use tinycortex::memory::sync::{RealCostAccumulator, SyncAuditEntry}; -impl SyncAuditEntry { - /// Total cost of the run: LLM summarisation cost plus the actual - /// Composio API-action cost. This is the "combined cost" the Sync - /// History UI surfaces so users see the full expense of a sync in one - /// number rather than just the summarisation slice (#3111). - pub fn combined_cost_usd(&self) -> f64 { - self.estimated_cost_usd + self.composio_cost_usd - } -} - -const AUDIT_FILENAME: &str = "sync_audit.jsonl"; - -/// Append an audit entry to the sync audit log. pub fn append_audit_entry(config: &Config, entry: &SyncAuditEntry) { - let dir = config.workspace_dir.join("memory_tree"); - if let Err(e) = std::fs::create_dir_all(&dir) { - tracing::warn!( - error = %e, - "[memory_sync:audit] failed to create audit dir" - ); - return; - } - - let path = dir.join(AUDIT_FILENAME); - if let Err(e) = append_jsonl(&path, entry) { - tracing::warn!( - error = %e, - "[memory_sync:audit] failed to write audit entry" - ); + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + if let Err(error) = tinycortex::memory::sync::append_audit_entry(&memory_config, entry) { + tracing::warn!(%error, "[memory_sync:audit] tinycortex append failed"); } } -fn append_jsonl(path: &Path, entry: &SyncAuditEntry) -> std::io::Result<()> { - let mut file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(path)?; - let json = serde_json::to_string(entry) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - writeln!(file, "{json}")?; - Ok(()) -} - -/// Read all audit entries, most recent first. Returns an empty vec if -/// the file doesn't exist yet. pub fn read_audit_log(config: &Config) -> Vec { - let path = config - .workspace_dir - .join("memory_tree") - .join(AUDIT_FILENAME); - let content = match std::fs::read_to_string(&path) { - Ok(c) => c, - Err(_) => return Vec::new(), - }; - let mut entries: Vec = content - .lines() - .filter(|l| !l.trim().is_empty()) - .filter_map(|l| serde_json::from_str(l).ok()) - .collect(); - entries.reverse(); - entries + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + match tinycortex::memory::sync::read_audit_log(&memory_config) { + Ok(entries) => entries, + Err(error) => { + tracing::warn!(%error, "[memory_sync:audit] tinycortex read failed"); + Vec::new() + } + } } -/// Estimate cost in USD for a given token count. -/// -/// Uses DeepSeek v4 flash pricing (the summarization-v1 backing model): -/// $0.07/M input, $0.28/M output. pub fn estimate_cost_usd(input_tokens: u64, output_tokens: u64) -> f64 { - let input_cost = input_tokens as f64 * 0.07 / 1_000_000.0; - let output_cost = output_tokens as f64 * 0.28 / 1_000_000.0; - input_cost + output_cost -} - -/// Accumulates per-batch token/charge figures over a sync run and decides -/// whether the run's totals may be reported as provider-"real" or must fall -/// back to the local `len/4` estimate. -/// -/// The rule (issue #3110): provider-reported tokens/charges are only -/// promoted to the run-level audit figure when **every** batch in the run -/// carried that signal. A run where some batches reported usage and others -/// fell back (provider silent / fallback summary) would otherwise produce a -/// partial "real" total that undercounts the run — worse than the estimate. -/// In that mixed case we keep the estimate, which covers all batches. -#[derive(Debug, Default, Clone)] -pub struct RealCostAccumulator { - total_batches: u32, - /// Number of batches that reported provider token usage (`input_tokens` - /// or `output_tokens` non-zero). - batches_with_usage: u32, - /// Number of batches that reported a provider charge. - 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() - } - - /// Fold one batch into the accumulator. - /// - /// `est_input` / `est_output` are the local-heuristic token counts for - /// the batch and are always summed. `real_input` / `real_output` are the - /// provider-reported counts (`0` = "no usage" sentinel). `charge` is the - /// provider charge for the batch when reported. - pub fn add_batch( - &mut self, - est_input: u64, - est_output: u64, - real_input: u64, - real_output: u64, - charge: Option, - ) { - self.total_batches += 1; - self.est_input_tokens += est_input; - self.est_output_tokens += est_output; - - // `input_tokens == 0 && output_tokens == 0` is the sentinel for "no - // usage" (set by the fallback path and by providers that don't report - // usage), so a batch only counts as carrying real usage when one of - // them is non-zero. - if real_input > 0 || real_output > 0 { - self.batches_with_usage += 1; - self.real_input_tokens += real_input; - self.real_output_tokens += real_output; - } - if let Some(charge) = charge { - self.batches_with_charge += 1; - self.real_charged_usd += charge; - } - } - - /// True when every batch reported provider token usage. Only then are the - /// real token totals complete enough to replace the estimate. - fn usage_is_complete(&self) -> bool { - self.total_batches > 0 && self.batches_with_usage == self.total_batches - } - - /// True when every batch reported a provider charge. Only then is the - /// summed charge a faithful total for the run. - fn charge_is_complete(&self) -> bool { - self.total_batches > 0 && self.batches_with_charge == self.total_batches - } - - /// Input tokens to record on the audit entry: real total when complete - /// across all batches, else the estimate. - pub fn audit_input_tokens(&self) -> u64 { - if self.usage_is_complete() { - self.real_input_tokens - } else { - self.est_input_tokens - } - } - - /// Output tokens to record on the audit entry. - pub fn audit_output_tokens(&self) -> u64 { - if self.usage_is_complete() { - self.real_output_tokens - } else { - self.est_output_tokens - } - } - - /// The hardcoded-pricing estimate over the run's estimated tokens — - /// always recorded as the fallback cost. - pub fn estimated_cost(&self) -> f64 { - estimate_cost_usd(self.est_input_tokens, self.est_output_tokens) - } - - /// The authoritative provider charge for the run when every batch - /// reported one, else `None` (falls back to the estimate downstream). - pub fn actual_charged_usd(&self) -> Option { - if self.charge_is_complete() { - Some(self.real_charged_usd) - } else { - None - } - } - - /// True when this run's token figures came from complete provider usage. - pub fn usage_is_real(&self) -> bool { - self.usage_is_complete() - } -} - -impl SyncAuditEntry { - /// The cost figure to display for this run: the real backend charge - /// when the provider reported one ([`Self::actual_charged_usd`]), - /// otherwise the hardcoded-pricing estimate ([`Self::estimated_cost_usd`]). - /// - /// Old audit entries (written before issue #3110) have no - /// `actual_charged_usd`, so they transparently fall back to the - /// estimate and keep rendering as before. - pub fn effective_cost_usd(&self) -> f64 { - self.actual_charged_usd.unwrap_or(self.estimated_cost_usd) - } - - /// True when the cost figure came from the backend's real billing - /// signal rather than the hardcoded-pricing estimate. - pub fn cost_is_actual(&self) -> bool { - self.actual_charged_usd.is_some() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn estimate_cost_reasonable() { - // 50k input + 5k output at DeepSeek flash pricing - let cost = estimate_cost_usd(50_000, 5_000); - // $0.0035 input + $0.0014 output = $0.0049 - assert!((cost - 0.0049).abs() < 0.0001); - } - - #[test] - fn accumulator_all_batches_real_promotes_usage_and_charge() { - let mut acc = RealCostAccumulator::new(); - acc.add_batch(1_000, 100, 900, 90, Some(0.005)); - acc.add_batch(1_000, 100, 800, 80, Some(0.004)); - - assert!(acc.usage_is_real()); - assert_eq!(acc.audit_input_tokens(), 1_700); - assert_eq!(acc.audit_output_tokens(), 170); - let charge = acc.actual_charged_usd().expect("charge complete"); - assert!((charge - 0.009).abs() < 1e-9); - } - - #[test] - fn accumulator_mixed_usage_falls_back_to_estimate() { - // One batch reports real usage, the second is silent (fallback). - // A partial real total (900/90) would *undercount* the run versus - // the estimate (2000/200), so we must keep the estimate. - let mut acc = RealCostAccumulator::new(); - acc.add_batch(1_000, 100, 900, 90, Some(0.005)); - acc.add_batch(1_000, 100, 0, 0, None); - - assert!(!acc.usage_is_real()); - assert_eq!(acc.audit_input_tokens(), 2_000); - assert_eq!(acc.audit_output_tokens(), 200); - // Charge incomplete (only one batch reported) → no actual charge. - assert_eq!(acc.actual_charged_usd(), None); - } - - #[test] - fn accumulator_usage_complete_but_charge_partial() { - // Every batch reports usage, but only one reports a charge. Tokens - // promote to real; charge stays None because the run-level sum would - // be missing the second batch's charge. - let mut acc = RealCostAccumulator::new(); - acc.add_batch(1_000, 100, 900, 90, Some(0.005)); - acc.add_batch(1_000, 100, 800, 80, None); - - assert!(acc.usage_is_real()); - assert_eq!(acc.audit_input_tokens(), 1_700); - assert_eq!(acc.actual_charged_usd(), None); - } - - #[test] - fn accumulator_no_batches_uses_estimate() { - let acc = RealCostAccumulator::new(); - assert!(!acc.usage_is_real()); - assert_eq!(acc.audit_input_tokens(), 0); - assert_eq!(acc.actual_charged_usd(), None); - assert_eq!(acc.estimated_cost(), 0.0); - } - - #[test] - fn append_creates_file_and_writes_jsonl() { - let tmp = tempfile::TempDir::new().unwrap(); - let path = tmp.path().join("test_audit.jsonl"); - - let entry = SyncAuditEntry { - timestamp: Utc::now(), - source_id: "src_123".to_string(), - source_kind: "github_repo".to_string(), - scope: "github:org/repo".to_string(), - items_fetched: 100, - batches: 2, - input_tokens: 50_000, - output_tokens: 5_000, - estimated_cost_usd: 0.225, - composio_actions_called: 0, - composio_cost_usd: 0.0, - actual_charged_usd: None, - duration_ms: 12_000, - success: true, - error: None, - }; - - append_jsonl(&path, &entry).unwrap(); - append_jsonl(&path, &entry).unwrap(); - - let content = std::fs::read_to_string(&path).unwrap(); - let lines: Vec<&str> = content.lines().collect(); - assert_eq!(lines.len(), 2); - assert!(lines[0].contains("src_123")); - } - - fn entry_with_costs(estimated: f64, actual: Option) -> SyncAuditEntry { - SyncAuditEntry { - timestamp: Utc::now(), - source_id: "src".to_string(), - source_kind: "github_repo".to_string(), - scope: "github:org/repo".to_string(), - items_fetched: 1, - batches: 1, - input_tokens: 100, - output_tokens: 10, - estimated_cost_usd: estimated, - composio_actions_called: 0, - composio_cost_usd: 0.0, - actual_charged_usd: actual, - duration_ms: 1, - success: true, - error: None, - } - } - - #[test] - fn combined_cost_sums_llm_and_composio() { - let entry = SyncAuditEntry { - timestamp: Utc::now(), - source_id: "src_gmail".to_string(), - source_kind: "composio".to_string(), - scope: "gmail".to_string(), - items_fetched: 40, - batches: 1, - input_tokens: 20_000, - output_tokens: 2_000, - estimated_cost_usd: 0.001_96, - composio_actions_called: 8, - composio_cost_usd: 0.04, - actual_charged_usd: None, - duration_ms: 5_000, - success: true, - error: None, - }; - assert!((entry.combined_cost_usd() - 0.041_96).abs() < 1e-9); - } - - #[test] - fn legacy_audit_line_without_composio_fields_deserializes() { - let legacy = r#"{"timestamp":"2026-05-01T00:00:00Z","source_id":"src_old","source_kind":"github_repo","scope":"github:org/repo","items_fetched":10,"batches":1,"input_tokens":1000,"output_tokens":100,"estimated_cost_usd":0.0001,"duration_ms":2000,"success":true}"#; - let entry: SyncAuditEntry = - serde_json::from_str(legacy).expect("legacy audit line must still parse"); - assert_eq!(entry.composio_actions_called, 0); - assert_eq!(entry.composio_cost_usd, 0.0); - assert_eq!(entry.actual_charged_usd, None); - assert_eq!(entry.source_id, "src_old"); - } - - #[test] - fn effective_cost_prefers_actual_charge_when_present() { - let entry = entry_with_costs(0.0049, Some(0.0123)); - assert!(entry.cost_is_actual()); - assert!((entry.effective_cost_usd() - 0.0123).abs() < f64::EPSILON); - } - - #[test] - fn effective_cost_falls_back_to_estimate_without_usage() { - let entry = entry_with_costs(0.0049, None); - assert!(!entry.cost_is_actual()); - assert!((entry.effective_cost_usd() - 0.0049).abs() < f64::EPSILON); - } - - #[test] - fn old_entry_without_actual_field_deserializes_and_renders_estimate() { - let legacy = r#"{ - "timestamp": "2024-01-01T00:00:00Z", - "source_id": "src_old", - "source_kind": "github_repo", - "scope": "github:org/repo", - "items_fetched": 10, - "batches": 1, - "input_tokens": 50000, - "output_tokens": 5000, - "estimated_cost_usd": 0.0049, - "duration_ms": 12000, - "success": true - }"#; - let entry: SyncAuditEntry = serde_json::from_str(legacy).unwrap(); - assert_eq!(entry.actual_charged_usd, None); - assert!(!entry.cost_is_actual()); - assert!((entry.effective_cost_usd() - 0.0049).abs() < f64::EPSILON); - } - - #[test] - fn new_entry_roundtrips_actual_charge_through_jsonl() { - let tmp = tempfile::TempDir::new().unwrap(); - let path = tmp.path().join("roundtrip_audit.jsonl"); - let entry = entry_with_costs(0.0049, Some(0.0200)); - append_jsonl(&path, &entry).unwrap(); - - let content = std::fs::read_to_string(&path).unwrap(); - let line = content.lines().next().unwrap(); - let parsed: SyncAuditEntry = serde_json::from_str(line).unwrap(); - assert_eq!(parsed.actual_charged_usd, Some(0.0200)); - assert!((parsed.effective_cost_usd() - 0.0200).abs() < f64::EPSILON); - } + tinycortex::memory::sync::estimate_cost_usd(input_tokens, output_tokens) } diff --git a/src/openhuman/memory_sync/sources/github.rs b/src/openhuman/memory_sync/sources/github.rs index c392aabbb..e3cf2f702 100644 --- a/src/openhuman/memory_sync/sources/github.rs +++ b/src/openhuman/memory_sync/sources/github.rs @@ -1,862 +1,23 @@ -//! GitHub repo sync pipeline. -//! -//! Git clone → read commits → fetch issues/PRs from API → batch items -//! into ~50k-token groups → summarise each batch → ingest into the -//! memory tree via `memory_tree::ingest::ingest_summary`. - -use async_trait::async_trait; +//! Compatibility wrapper for the tinycortex GitHub repository pipeline. use crate::openhuman::config::Config; -use crate::openhuman::memory::sync::{emit_sync_stage, MemorySyncStage, MemorySyncTrigger}; -use crate::openhuman::memory::tree_source::get_or_create_source_tree; -use crate::openhuman::memory_sources::readers; -use crate::openhuman::memory_sources::readers::github; -use crate::openhuman::memory_sources::readers::SourceReader; -use crate::openhuman::memory_sources::types::{ - MemorySourceEntry, SourceContent, SourceItem, SourceKind, -}; -use crate::openhuman::memory_store::content::raw::{self as raw_store, RawItem}; -use crate::openhuman::memory_store::trees::types::TreeKind; -use crate::openhuman::memory_store::trees::types::INPUT_TOKEN_BUDGET; -use crate::openhuman::memory_sync::sources::audit::{ - append_audit_entry, RealCostAccumulator, SyncAuditEntry, -}; -use crate::openhuman::memory_sync::traits::{SyncOutcome, SyncPipeline, SyncPipelineKind}; -use crate::openhuman::memory_tree::ingest::{ingest_summary, SummaryIngestInput}; -use crate::openhuman::memory_tree::summarise::{ - fallback_summary, summarise, SummaryContext, SummaryInput, SummaryOutput, -}; -use futures::stream::StreamExt; -use std::path::Path; +use crate::openhuman::memory_sources::MemorySourceEntry; -pub struct GithubSourcePipeline { - source: MemorySourceEntry, -} +pub use tinycortex::memory::sync::SyncOutcome; -impl GithubSourcePipeline { - pub fn new(source: MemorySourceEntry) -> Self { - Self { source } - } -} - -#[async_trait] -impl SyncPipeline for GithubSourcePipeline { - fn id(&self) -> &str { - &self.source.id - } - - fn kind(&self) -> SyncPipelineKind { - SyncPipelineKind::Workspace - } - - async fn init(&self, _config: &Config) -> anyhow::Result<()> { - Ok(()) - } - - async fn tick(&self, config: &Config) -> anyhow::Result { - run_github_sync(&self.source, config).await - } -} - -/// Run a full GitHub source sync: clone/fetch → list items → read all → -/// batch into 50k-token groups → summarise each → ingest into tree. pub async fn run_github_sync( source: &MemorySourceEntry, config: &Config, ) -> anyhow::Result { - let start = std::time::Instant::now(); - let source_id = &source.id; - let kind_str = source.kind.as_str(); - - let repo_scope = source - .url - .as_deref() - .and_then(github::repo_chunk_scope) - .ok_or_else(|| anyhow::anyhow!("github source missing url for repo scope"))?; - - let raw_source_id = source - .url - .as_deref() - .and_then(github::repo_archive_source_id); - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Fetching, - Some(kind_str), - Some(source_id), - Some("listing items".to_string()), - Some(source_id), + tracing::debug!( + source_id = %source.id, + "[memory_sync:github] dispatching through tinycortex" ); - - let reader = readers::reader_for(&SourceKind::GithubRepo); - let items = reader - .list_items(source, config) + if crate::openhuman::memory::global::client_if_ready().is_none() { + crate::openhuman::memory::global::init(config.workspace_dir.clone()) + .map_err(anyhow::Error::msg)?; + } + crate::openhuman::tinycortex::run_source_pipeline(source, config) .await - .map_err(|e| anyhow::anyhow!(e))?; - let total = items.len(); - - tracing::debug!( - source_id = %source_id, - total = total, - "[memory_sync:github] listed items" - ); - - if total == 0 { - return Ok(SyncOutcome { - records_ingested: 0, - more_pending: false, - note: Some("no items found".to_string()), - }); - } - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Ingesting, - Some(kind_str), - Some(source_id), - Some(format!("reading {total} items")), - Some(source_id), - ); - - let content_root = config.memory_tree_content_root(); - let (inputs, child_labels, child_basenames) = read_items_buffered( - reader.as_ref(), - &items, - source, - config, - raw_source_id.as_deref(), - &content_root, - ) - .await; - - if inputs.is_empty() { - return Ok(SyncOutcome { - records_ingested: 0, - more_pending: false, - note: Some("all items failed to read".to_string()), - }); - } - - let tree = get_or_create_source_tree(config, &repo_scope) - .map_err(|e| anyhow::anyhow!("get_or_create_source_tree: {e:#}"))?; - - let batches = - batch_by_token_budget(&inputs, &child_labels, &child_basenames, INPUT_TOKEN_BUDGET); - let batch_count = batches.len(); - let input_count = inputs.len(); - - tracing::info!( - source_id = %source_id, - items = input_count, - batches = batch_count, - "[memory_sync:github] summarising in {} batch(es)", - batch_count - ); - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Ingesting, - Some(kind_str), - Some(source_id), - Some(format!( - "summarising {input_count} items in {batch_count} batch(es)" - )), - Some(source_id), - ); - - // Token/charge accounting across the run. The estimate (`body.len() / 4` - // heuristic) is always summed; provider-reported figures only replace it - // when every batch reported them (issue #3110). See `RealCostAccumulator`. - let mut cost = RealCostAccumulator::new(); - - // ── Phase 1: summarise every batch (bounded concurrency) ── - // The per-batch summarise is an independent LLM round-trip; ingest below - // is the only ordered, serial part (it writes into the shared source tree). - // Cloud providers overlap these calls; a local model stays serial. Futures - // are materialized into a Vec before `buffered` so the higher-ranked - // closure lifetime stays tied to this frame (avoids the "FnOnce is not - // general enough" error). `buffered` preserves order so outputs[i] maps to - // batches[i]. - let concurrency = summarise_concurrency(config.workload_uses_local("memory")); - tracing::debug!( - source_id = %source_id, - tree_id = %tree.id, - batch_count = batches.len(), - concurrency, - "[memory_sync:github] starting concurrent summarise phase" - ); - let summarise_futs: Vec<_> = batches - .iter() - .enumerate() - .map(|(batch_idx, (batch_inputs, _labels, _basenames))| { - let config = config; - let tree_id = &tree.id; - async move { - let ctx = SummaryContext { - tree_id, - tree_kind: TreeKind::Source, - target_level: 1, - token_budget: 5_000, - }; - match summarise(config, batch_inputs, &ctx).await { - Ok(o) => o, - Err(e) => { - tracing::warn!( - error = %e, - batch = batch_idx, - "[memory_sync:github] summarise failed, using fallback" - ); - fallback_summary(batch_inputs, ctx.token_budget) - } - } - } - }) - .collect(); - - let outputs: Vec = futures::stream::iter(summarise_futs) - .buffered(concurrency) - .collect() - .await; - tracing::debug!( - source_id = %source_id, - tree_id = %tree.id, - outputs = outputs.len(), - "[memory_sync:github] concurrent summarise phase complete" - ); - - // ── Phase 2: fold cost + ingest in batch order (serial) ── - for (batch_idx, ((batch_inputs, batch_labels, batch_basenames), output)) in - batches.into_iter().zip(outputs).enumerate() - { - let batch_input_tokens: u64 = batch_inputs.iter().map(|i| i.token_count as u64).sum(); - - cost.add_batch( - batch_input_tokens, - output.token_count as u64, - output.input_tokens, - output.output_tokens, - output.charged_amount_usd, - ); - - let time_start = batch_inputs - .iter() - .map(|i| i.time_range_start) - .min() - .unwrap_or_else(chrono::Utc::now); - let time_end = batch_inputs - .iter() - .map(|i| i.time_range_end) - .max() - .unwrap_or_else(chrono::Utc::now); - let max_score = batch_inputs.iter().map(|i| i.score).fold(0.0f32, f32::max); - - let ingest_input = SummaryIngestInput { - content: output.content, - token_count: output.token_count, - entities: Vec::new(), - topics: vec!["memory_sources".to_string(), kind_str.to_string()], - time_range_start: time_start, - time_range_end: time_end, - score: max_score, - child_labels: batch_labels, - child_basenames: batch_basenames, - }; - - // `child_basenames` holds the raw-archive wikilink paths (rel path - // with `.md` stripped) — captured BEFORE ingest_summary moves the - // input. Re-suffix to recover the coverage-gate keys. - let batch_raw_rel_paths: Vec = ingest_input - .child_basenames - .iter() - .flatten() - .map(|basename| format!("{basename}.md")) - .collect(); - - let outcome = ingest_summary(config, &tree, ingest_input).await?; - - // Record raw-archive coverage so the incremental reconcile - // (`memory_sync::sources::rebuild`) knows these files are - // summarised. Marked only after the summary landed: a crash in - // between re-summarises the batch (duplicate summary, acceptable) - // instead of silently losing coverage. - if !batch_raw_rel_paths.is_empty() { - if let Err(e) = crate::openhuman::memory_store::chunks::store::mark_raw_paths_ingested( - config, - &batch_raw_rel_paths, - ) { - tracing::warn!( - source_id = %source_id, - batch = batch_idx, - error = %format!("{e:#}"), - "[memory_sync:github] failed to record raw coverage — reconcile may re-summarise this batch" - ); - } - } - - tracing::info!( - source_id = %source_id, - batch = batch_idx, - summary_id = %outcome.summary_id, - path = %outcome.content_path, - sealed = outcome.sealed_ids.len(), - covered_raw_files = batch_raw_rel_paths.len(), - "[memory_sync:github] batch ingested + coverage recorded" - ); - } - - let duration_ms = start.elapsed().as_millis() as u64; - - // Provider token counts are recorded only when *every* batch reported - // usage; a mixed run keeps the `len/4` estimate (which covers all - // batches) rather than a partial real total that would undercount. - // `estimated_cost_usd` is always populated as the fallback. Issue #3110. - let any_real_usage = cost.usage_is_real(); - let audit_input_tokens = cost.audit_input_tokens(); - let audit_output_tokens = cost.audit_output_tokens(); - let estimated_cost = cost.estimated_cost(); - let actual_charged_usd = cost.actual_charged_usd(); - // Cost surfaced to the user/logs: real charge when present, else estimate. - let display_cost = actual_charged_usd.unwrap_or(estimated_cost); - - tracing::info!( - source_id = %source_id, - usage_is_real = any_real_usage, - actual_charge = actual_charged_usd.is_some(), - input_tokens = audit_input_tokens, - output_tokens = audit_output_tokens, - estimated_cost_usd = %format!("{estimated_cost:.4}"), - actual_charged_usd = ?actual_charged_usd, - "[memory_sync:github] sync cost accounting" - ); - - append_audit_entry( - config, - &SyncAuditEntry { - timestamp: chrono::Utc::now(), - source_id: source_id.to_string(), - source_kind: kind_str.to_string(), - scope: repo_scope.clone(), - items_fetched: input_count as u32, - batches: batch_count as u32, - input_tokens: audit_input_tokens, - output_tokens: audit_output_tokens, - estimated_cost_usd: estimated_cost, - composio_actions_called: 0, - composio_cost_usd: 0.0, - actual_charged_usd, - duration_ms, - success: true, - error: None, - }, - ); - - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Completed, - Some(kind_str), - Some(source_id), - Some(format!( - "{input_count} items → {batch_count} summary(ies) ({audit_input_tokens} in / {audit_output_tokens} out tokens, ${display_cost:.4})" - )), - Some(source_id), - ); - - Ok(SyncOutcome { - records_ingested: input_count as u32, - more_pending: false, - note: Some(format!( - "{input_count} items → {batch_count} summary(ies) (${display_cost:.4})" - )), - }) -} - -/// Max number of `read_item` round-trips kept in flight at once while reading -/// a transcript's worth of GitHub items. Each read is a GitHub API / `gh` CLI -/// call (issues/PRs always fetch comments — see `read_issue`/`read_pr` — even -/// when the body is served from the list cache), so bounding the fan-out keeps -/// a large repo from opening hundreds of concurrent requests and tripping -/// GitHub's secondary rate limits. -const GITHUB_READ_CONCURRENCY: usize = 8; - -/// Max in-flight summarise LLM calls when the memory workload routes to a -/// **cloud** provider. Local models are single-GPU/single-process, so they are -/// driven serially (see [`summarise_concurrency`]). -const GITHUB_SUMMARISE_CONCURRENCY: usize = 4; - -/// Effective summarise concurrency given the memory workload routing. -/// -/// Cloud providers tolerate (and benefit from) overlapping the per-batch LLM -/// round-trips, so we fan out to [`GITHUB_SUMMARISE_CONCURRENCY`]. A local model -/// is a single inference engine — concurrent calls would just queue (or -/// oversubscribe VRAM), so we keep it strictly serial (`1`), preserving the old -/// one-batch-at-a-time behaviour. Pure so it can be unit-tested. -fn summarise_concurrency(uses_local: bool) -> usize { - if uses_local { - 1 - } else { - GITHUB_SUMMARISE_CONCURRENCY - } -} - -/// Read every listed item into the parallel `(inputs, labels, basenames)` -/// vectors used to build the source tree. -/// -/// Reads run with **bounded concurrency** (`GITHUB_READ_CONCURRENCY`) instead -/// of one sequential `.await` each: a repo routinely lists 100+ items and every -/// `read_item` is an independent network round-trip, so overlapping their waits -/// is a real wall-time win on this background sync job. -/// -/// Order is **preserved** (`buffered`, not `buffer_unordered`): the three -/// returned vectors are positionally aligned and consumed in lockstep by -/// `batch_by_token_budget`, so reordering them would scramble label/basename -/// provenance. A read failure is logged and the item is skipped (it pushes -/// nothing), exactly as the previous sequential `continue` did. -/// -/// The per-item futures are collected into a `Vec` before `buffered` for the -/// same higher-ranked-lifetime reason as elsewhere in the codebase: mapping -/// lazily on the stream stores the closure in the polled state and requires it -/// to hold for any lifetime, which fails once the whole sync future is spawned -/// (`Send + 'static`). Collecting builds concrete-lifetime futures up front. -async fn read_items_buffered( - reader: &dyn SourceReader, - items: &[SourceItem], - source: &MemorySourceEntry, - config: &Config, - raw_source_id: Option<&str>, - content_root: &Path, -) -> (Vec, Vec, Vec>) { - let total = items.len(); - let kind_str = source.kind.as_str(); - let source_id = source.id.as_str(); - - let item_futs: Vec<_> = items - .iter() - .map(|item| async move { - let content = match reader.read_item(source, &item.id, config).await { - Ok(c) => c, - Err(e) => { - tracing::warn!( - item_id = %item.id, - error = %e, - "[memory_sync:github] skipping item — read failed" - ); - return None; - } - }; - - if let Some(raw_sid) = raw_source_id { - write_raw_archive( - raw_sid, - &item.id, - item.updated_at_ms, - &content, - content_root, - ); - } - - // Compute the raw archive relative path for the wikilink. - let raw_path = raw_source_id.and_then(|raw_sid| { - let (kind, uid) = github::raw_archive_coords(&item.id)?; - let created_at_ms = item.updated_at_ms.unwrap_or(0); - let rel = raw_store::raw_rel_path(raw_sid, kind, created_at_ms, &uid); - Some(rel.strip_suffix(".md").unwrap_or(&rel).to_string()) - }); - - let token_count = (content.body.len() / 4).max(1) as u32; - let priority = item_is_high_priority(&item.id, &content); - let ts = item - .updated_at_ms - .and_then(chrono::DateTime::from_timestamp_millis) - .unwrap_or_else(chrono::Utc::now); - - Some(( - SummaryInput { - id: item.id.clone(), - content: content.body, - token_count, - entities: Vec::new(), - topics: vec!["memory_sources".to_string(), kind_str.to_string()], - time_range_start: ts, - time_range_end: ts, - score: if priority { 0.8 } else { 0.5 }, - }, - item.id.clone(), - raw_path, - )) - }) - .collect(); - - let mut inputs: Vec = Vec::with_capacity(total); - let mut child_labels: Vec = Vec::with_capacity(total); - let mut child_basenames: Vec> = Vec::with_capacity(total); - - let mut stream = futures::stream::iter(item_futs).buffered(GITHUB_READ_CONCURRENCY); - let mut processed = 0usize; - while let Some(result) = stream.next().await { - if let Some((input, label, raw_path)) = result { - inputs.push(input); - child_labels.push(label); - child_basenames.push(raw_path); - } - processed += 1; - if processed.is_multiple_of(100) || processed == total { - emit_sync_stage( - MemorySyncTrigger::Manual, - MemorySyncStage::Ingesting, - Some(kind_str), - Some(source_id), - Some(format!("{processed}/{total} read")), - Some(source_id), - ); - } - } - - (inputs, child_labels, child_basenames) -} - -/// Group inputs into batches where each batch's total token count is -/// approximately `budget`. Pairs each batch with its corresponding -/// child labels and basenames for provenance. -fn batch_by_token_budget( - inputs: &[SummaryInput], - labels: &[String], - basenames: &[Option], - budget: u32, -) -> Vec<(Vec, Vec, Vec>)> { - let mut batches = Vec::new(); - let mut cur_inputs: Vec = Vec::new(); - let mut cur_labels: Vec = Vec::new(); - let mut cur_basenames: Vec> = Vec::new(); - let mut cur_tokens: u32 = 0; - - for ((input, label), basename) in inputs.iter().zip(labels.iter()).zip(basenames.iter()) { - if !cur_inputs.is_empty() && cur_tokens + input.token_count > budget { - batches.push(( - std::mem::take(&mut cur_inputs), - std::mem::take(&mut cur_labels), - std::mem::take(&mut cur_basenames), - )); - cur_tokens = 0; - } - cur_tokens += input.token_count; - cur_inputs.push(input.clone()); - cur_labels.push(label.clone()); - cur_basenames.push(basename.clone()); - } - - if !cur_inputs.is_empty() { - batches.push((cur_inputs, cur_labels, cur_basenames)); - } - - batches -} - -fn item_is_high_priority(item_id: &str, content: &SourceContent) -> bool { - if item_id.starts_with("commit:") { - return true; - } - let state_closed = content - .metadata - .get("state") - .and_then(|v| v.as_str()) - .map(|s| s.eq_ignore_ascii_case("closed")) - .unwrap_or(false); - let merged = content - .metadata - .get("merged") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - state_closed || merged -} - -fn write_raw_archive( - raw_source_id: &str, - item_id: &str, - updated_at_ms: Option, - content: &SourceContent, - content_root: &std::path::Path, -) { - let Some((kind, uid)) = github::raw_archive_coords(item_id) else { - return; - }; - let created_at_ms = updated_at_ms.unwrap_or(0); - let raw_item = RawItem { - uid: &uid, - created_at_ms, - markdown: &content.body, - kind, - }; - if let Err(e) = - raw_store::write_raw_items(content_root, raw_source_id, std::slice::from_ref(&raw_item)) - { - tracing::warn!( - item_id = %item_id, - error = %e, - "[memory_sync:github] raw archive write failed" - ); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_sources::types::ContentType; - - #[test] - fn summarise_concurrency_serial_for_local_model() { - // A local model is a single inference engine — must stay serial. - assert_eq!(summarise_concurrency(true), 1); - } - - #[test] - fn summarise_concurrency_fans_out_for_cloud_model() { - assert_eq!(summarise_concurrency(false), GITHUB_SUMMARISE_CONCURRENCY); - assert!(GITHUB_SUMMARISE_CONCURRENCY > 1, "cloud path must overlap"); - } - - fn content_with(meta: serde_json::Value) -> SourceContent { - SourceContent { - id: "x".into(), - title: "t".into(), - body: "b".into(), - content_type: ContentType::Markdown, - metadata: meta, - } - } - - #[test] - fn commits_are_always_high_priority() { - let c = content_with(serde_json::json!({})); - assert!(item_is_high_priority("commit:abc123", &c)); - } - - #[test] - fn closed_issue_is_high_priority_open_is_not() { - let closed = content_with(serde_json::json!({ "state": "closed" })); - assert!(item_is_high_priority("issue:1", &closed)); - let open = content_with(serde_json::json!({ "state": "open" })); - assert!(!item_is_high_priority("issue:1", &open)); - } - - #[test] - fn merged_pr_is_high_priority() { - let merged = content_with(serde_json::json!({ "state": "open", "merged": true })); - assert!(item_is_high_priority("pr:7", &merged)); - } - - #[test] - fn missing_metadata_defaults_to_low_priority() { - let c = content_with(serde_json::json!({})); - assert!(!item_is_high_priority("issue:9", &c)); - } - - #[test] - fn batch_by_budget_groups_correctly() { - let make = |tokens: u32| SummaryInput { - id: format!("t{tokens}"), - content: String::new(), - token_count: tokens, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: chrono::Utc::now(), - time_range_end: chrono::Utc::now(), - score: 0.5, - }; - - let inputs = vec![make(20_000), make(20_000), make(20_000), make(10_000)]; - let labels: Vec = (0..4).map(|i| format!("l{i}")).collect(); - let basenames: Vec> = (0..4).map(|_| None).collect(); - let batches = batch_by_token_budget(&inputs, &labels, &basenames, 50_000); - assert_eq!(batches.len(), 2); - assert_eq!(batches[0].0.len(), 2); - assert_eq!(batches[0].1.len(), 2); - assert_eq!(batches[1].0.len(), 2); - } - - #[test] - fn batch_empty_input() { - let batches = batch_by_token_budget(&[], &[], &[], 50_000); - assert!(batches.is_empty()); - } - - use std::sync::atomic::{AtomicUsize, Ordering}; - - /// `SourceReader` double that records the high-water mark of concurrent - /// `read_item` calls, proving `read_items_buffered` overlaps reads while - /// staying within `GITHUB_READ_CONCURRENCY`. Fails reads whose id contains - /// "fail" so the skip path is exercised too. - struct CountingReader { - in_flight: AtomicUsize, - peak_in_flight: AtomicUsize, - calls: AtomicUsize, - } - - impl CountingReader { - fn new() -> Self { - Self { - in_flight: AtomicUsize::new(0), - peak_in_flight: AtomicUsize::new(0), - calls: AtomicUsize::new(0), - } - } - } - - #[async_trait] - impl SourceReader for CountingReader { - fn kind(&self) -> SourceKind { - SourceKind::GithubRepo - } - - async fn list_items( - &self, - _source: &MemorySourceEntry, - _config: &Config, - ) -> Result, String> { - Ok(Vec::new()) - } - - async fn read_item( - &self, - _source: &MemorySourceEntry, - item_id: &str, - _config: &Config, - ) -> Result { - // Track concurrent entry, then yield repeatedly *before* returning - // so sibling reads genuinely overlap in time — otherwise a read - // that completes on first poll would never reveal concurrency. - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.peak_in_flight.fetch_max(now, Ordering::SeqCst); - for _ in 0..4 { - tokio::task::yield_now().await; - } - self.calls.fetch_add(1, Ordering::SeqCst); - self.in_flight.fetch_sub(1, Ordering::SeqCst); - - if item_id.contains("fail") { - return Err(format!("simulated read failure for {item_id}")); - } - Ok(SourceContent { - id: item_id.to_string(), - title: format!("title {item_id}"), - body: format!("body for {item_id}"), - content_type: ContentType::Markdown, - metadata: serde_json::json!({}), - }) - } - } - - fn github_source() -> MemorySourceEntry { - MemorySourceEntry { - id: "src_gh".into(), - kind: SourceKind::GithubRepo, - label: "gh".into(), - enabled: true, - toolkit: None, - connection_id: None, - path: None, - glob: None, - url: Some("https://github.com/owner/repo".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 issue_items(n: usize) -> Vec { - (0..n) - .map(|i| SourceItem { - id: format!("issue:{i}"), - title: format!("t{i}"), - updated_at_ms: Some(1_700_000_000_000 + i as i64), - }) - .collect() - } - - #[tokio::test] - async fn read_items_buffered_overlaps_and_preserves_order() { - // 20 items exceeds GITHUB_READ_CONCURRENCY (8): an unbounded fan-out - // would push >8 reads in flight (bound assertion fails), while a fully - // sequential read would never exceed peak 1 (overlap assertion fails). - let reader = CountingReader::new(); - let source = github_source(); - let config = Config::default(); - let items = issue_items(20); - - let (inputs, labels, basenames) = read_items_buffered( - &reader, - &items, - &source, - &config, - None, - Path::new("/tmp/openhuman-test-unused"), - ) - .await; - - assert_eq!(inputs.len(), 20); - assert_eq!(basenames.len(), 20); - assert!( - basenames.iter().all(Option::is_none), - "no raw_source_id → no basenames" - ); - - // Order preserved: labels and input ids must match input order exactly. - let want: Vec = (0..20).map(|i| format!("issue:{i}")).collect(); - assert_eq!(labels, want, "buffered must preserve input order"); - assert_eq!( - inputs.iter().map(|i| i.id.clone()).collect::>(), - want - ); - - assert_eq!(reader.calls.load(Ordering::SeqCst), 20); - let peak = reader.peak_in_flight.load(Ordering::SeqCst); - assert!( - peak <= GITHUB_READ_CONCURRENCY, - "read concurrency must stay bounded at {GITHUB_READ_CONCURRENCY}, observed peak {peak}" - ); - assert!( - peak >= 2, - "reads must actually overlap (else the bound is meaningless), observed peak {peak}" - ); - } - - #[tokio::test] - async fn read_items_buffered_skips_failed_reads_and_keeps_order() { - let reader = CountingReader::new(); - let source = github_source(); - let config = Config::default(); - // Middle item fails its read; it must be skipped without breaking the - // positional alignment of the surviving items. - let items = vec![ - SourceItem { - id: "issue:0".into(), - title: "a".into(), - updated_at_ms: Some(1), - }, - SourceItem { - id: "issue:fail".into(), - title: "b".into(), - updated_at_ms: Some(2), - }, - SourceItem { - id: "issue:2".into(), - title: "c".into(), - updated_at_ms: Some(3), - }, - ]; - - let (inputs, labels, _basenames) = - read_items_buffered(&reader, &items, &source, &config, None, Path::new("/tmp/x")).await; - - assert_eq!(inputs.len(), 2, "failed item skipped"); - assert_eq!(labels, vec!["issue:0".to_string(), "issue:2".to_string()]); - } + .map_err(|error| anyhow::anyhow!(error.to_string())) } diff --git a/src/openhuman/memory_sync/sources/mod.rs b/src/openhuman/memory_sync/sources/mod.rs index 0b0f416cc..dc483ddea 100644 --- a/src/openhuman/memory_sync/sources/mod.rs +++ b/src/openhuman/memory_sync/sources/mod.rs @@ -1,12 +1,6 @@ //! Memory-source sync pipelines. //! -//! Syncs user-added memory sources (GitHub repos, local folders, RSS feeds, -//! web pages) by pulling items through `memory_sources::readers` and landing -//! them directly into the memory tree as leaves. When a tree's L0 buffer -//! hits `INPUT_TOKEN_BUDGET` (50k tokens), the cascade sealer fires. -//! -//! Embeddings are temporarily disabled — leaves land without vectors and -//! the seal cascade uses `LabelStrategy::Empty`. +//! Product-owned wrappers around tinycortex source sync facilities. pub mod audit; pub mod github; diff --git a/src/openhuman/memory_sync/sources/rebuild.rs b/src/openhuman/memory_sync/sources/rebuild.rs index 5b5fb9b86..073410fa0 100644 --- a/src/openhuman/memory_sync/sources/rebuild.rs +++ b/src/openhuman/memory_sync/sources/rebuild.rs @@ -1,834 +1,38 @@ -//! Reconcile raw archive files on disk into the memory tree. -//! -//! Raw archive files are written eagerly at fetch time (one `.md` per -//! upstream item under `raw///`), but summaries only -//! land at the end of a sync's batch loop — so an interrupted sync, a -//! crash mid-summarise, or legacy data leaves raw files on disk with no -//! tree coverage. This module closes that gap **incrementally**: -//! -//! 1. Every successful summary-batch ingest records its raw files in the -//! `mem_tree_ingested_sources` coverage gate -//! (`source_kind = "raw_file"`, `source_id = `). -//! 2. [`raw_coverage`] lists the on-disk files NOT in the gate. For -//! scopes whose summaries predate the gate (legacy data), it first -//! backfills the gate from the existing L1 summaries' child labels so -//! already-covered files are never re-summarised. -//! 3. [`rebuild_tree_from_raw`] reads only the pending files, batches -//! them into ~50k-token groups, summarises each batch, ingests via -//! `ingest_summary`, and marks the batch's files covered. -//! -//! Idempotent: re-running with full coverage is a no-op; a run that dies -//! mid-way resumes from the first unmarked batch. -//! -//! ## Tree scope vs archive id -//! -//! A source has TWO identifiers that slugify to *different* directories: -//! the tree scope (e.g. `github:owner/repo` → tree registry key) and the -//! raw-archive source id (e.g. `github.com/owner/repo` → -//! `raw/github-com-owner-repo/`). Callers must pass both — deriving the -//! raw dir from the tree scope silently scans the wrong (usually empty) -//! directory, which is exactly the bug that let thousands of GitHub raw -//! files sit unreconciled. - -use std::collections::HashSet; -use std::path::{Path, PathBuf}; - -use anyhow::Result; +//! OpenHuman configuration and LLM wrappers for tinycortex raw rebuilds. use crate::openhuman::config::Config; -use crate::openhuman::memory::tree_source::get_or_create_source_tree; -use crate::openhuman::memory_store::chunks::store::{ - count_raw_paths_ingested_with_prefix, filter_raw_paths_not_ingested, - list_chunk_raw_ref_paths_with_prefix, mark_raw_paths_ingested, -}; -use crate::openhuman::memory_store::content::paths::slugify_source_id; -use crate::openhuman::memory_store::content::raw::{raw_source_dir, sanitize_uid}; -use crate::openhuman::memory_store::trees::store::list_summaries_at_level; -use crate::openhuman::memory_store::trees::types::{TreeKind, INPUT_TOKEN_BUDGET}; -use crate::openhuman::memory_sync::sources::audit::{ - append_audit_entry, RealCostAccumulator, SyncAuditEntry, -}; -use crate::openhuman::memory_tree::ingest::{ingest_summary, SummaryIngestInput}; -use crate::openhuman::memory_tree::summarise::{ - fallback_summary, summarise, SummaryContext, SummaryInput, -}; -/// Outcome of a rebuild operation. -#[derive(Clone, Debug, Default)] -pub struct RebuildOutcome { - pub files_read: usize, - pub batches: usize, - pub input_tokens: u64, - pub output_tokens: u64, - pub estimated_cost_usd: f64, - /// Real amount billed by the backend in USD when the provider reported - /// usage for the run; `None` when it fell back to the estimate. Issue - /// #3110. Prefer this over `estimated_cost_usd` when `Some`. - pub actual_charged_usd: Option, -} +pub use tinycortex::memory::sync::{RawCoverage, RawFileRef, RebuildOutcome}; -/// One raw archive file pending tree coverage. -#[derive(Clone, Debug)] -pub struct RawFileRef { - /// Absolute path on disk. - pub abs: PathBuf, - /// Forward-slash relative path under `/`, with `.md` — - /// the canonical coverage-gate key (matches `raw::raw_rel_path`). - pub rel: String, -} - -/// Coverage report for one source's raw archive. -#[derive(Clone, Debug, Default)] -pub struct RawCoverage { - /// Total `.md` files on disk (excluding `_source.md` etc.). - pub total: usize, - /// Files recorded in the coverage gate. - pub covered: usize, - /// Files on disk with no coverage record, sorted chronologically - /// (filenames start with the item timestamp). - pub pending: Vec, -} - -/// Compute raw-archive coverage for a source. -/// -/// `tree_scope` keys the source tree (e.g. `"github:owner/repo"`, -/// `"gmail:user-at-gmail-dot-com"`); `archive_source_id` keys the raw -/// archive directory (e.g. `"github.com/owner/repo"` — pass the tree -/// scope again when the source writes its archive under the same id, -/// as gmail does). -/// -/// On first call for a scope with existing L1 summaries but an empty -/// gate (legacy data from before coverage tracking), the gate is -/// backfilled from the summaries' child labels so previously-summarised -/// files don't get re-ingested. pub fn raw_coverage( config: &Config, tree_scope: &str, archive_source_id: &str, -) -> Result { - let content_root = config.memory_tree_content_root(); - let source_dir = raw_source_dir(&content_root, archive_source_id); - if !source_dir.exists() { - return Ok(RawCoverage::default()); - } - - let mut files = collect_raw_files(&source_dir)?; - files.sort(); - if files.is_empty() { - return Ok(RawCoverage::default()); - } - - let refs: 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 = refs.len(); - - // Legacy backfill: gate empty for this archive but the tree already - // has L1 summaries → mark the files those summaries cover. - let rel_prefix = format!("raw/{}/", slugify_source_id(archive_source_id)); - if count_raw_paths_ingested_with_prefix(config, &rel_prefix)? == 0 { - let backfilled = backfill_coverage_from_summaries(config, tree_scope, &refs)?; - if backfilled > 0 { - tracing::info!( - tree_scope = %tree_scope, - archive = %archive_source_id, - backfilled = backfilled, - "[memory_sync:rebuild] backfilled coverage gate from existing L1 summaries" - ); - } - } - - let rel_paths: Vec = refs.iter().map(|r| r.rel.clone()).collect(); - let mut pending_rels: HashSet = filter_raw_paths_not_ingested(config, &rel_paths)? - .into_iter() - .collect(); - - // A raw file referenced by a persisted chunk (`raw_refs_json`) is - // already in the tree via the chunk pipeline — gmail mirrors every - // email to raw/ AND ingests it as a chunk. Those files are covered; - // re-summarising them through the rebuild path would duplicate - // content (and burn LLM batches) for sources that never use the - // summarise-direct path at all. - let chunk_covered = list_chunk_raw_ref_paths_with_prefix(config, &rel_prefix)?; - if !chunk_covered.is_empty() { - pending_rels.retain(|rel| !chunk_covered.contains(rel)); - } - - let pending: Vec = refs - .into_iter() - .filter(|r| pending_rels.contains(&r.rel)) - .collect(); - - tracing::debug!( - tree_scope = %tree_scope, - archive = %archive_source_id, - total = total, - pending = pending.len(), - "[memory_sync:rebuild] raw coverage computed" - ); - - Ok(RawCoverage { - total, - covered: total - pending.len(), - pending, - }) +) -> anyhow::Result { + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + tinycortex::memory::sync::raw_coverage(&memory_config, tree_scope, archive_source_id) } -/// Check whether a source has raw files on disk that the tree does not -/// cover yet. Coverage-based: a partially-covered scope (interrupted -/// sync) returns `true` even when the tree already has L1 summaries. pub fn needs_rebuild(config: &Config, tree_scope: &str, archive_source_id: &str) -> bool { - match raw_coverage(config, tree_scope, archive_source_id) { - Ok(cov) => !cov.pending.is_empty(), - Err(e) => { - tracing::warn!( - tree_scope = %tree_scope, - archive = %archive_source_id, - error = %format!("{e:#}"), - "[memory_sync:rebuild] coverage check failed — skipping rebuild" - ); - false - } - } + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + tinycortex::memory::sync::needs_rebuild(&memory_config, tree_scope, archive_source_id) } -/// Mark a label set covered in the gate from existing L1 summaries. -/// -/// Sync-written summaries label children as `commit:` / `issue:` / -/// `pr:`; rebuild-written summaries label them with the raw file stem -/// (`_`). Both shapes are matched against the on-disk files. -/// Returns the number of files marked covered. -fn backfill_coverage_from_summaries( - config: &Config, - tree_scope: &str, - files: &[RawFileRef], -) -> Result { - let tree = get_or_create_source_tree(config, tree_scope) - .map_err(|e| anyhow::anyhow!("get_or_create_source_tree: {e:#}"))?; - if tree.max_level == 0 { - return Ok(0); - } - - let summaries = list_summaries_at_level(config, &tree.id, 1)?; - // (kind_dir, sanitised uid) pairs from `:` labels. - let mut kind_uids: HashSet<(&'static str, String)> = HashSet::new(); - // Full file stems from rebuild-written labels. - let mut stems: HashSet = 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.clone()); - } - } - } - if kind_uids.is_empty() && stems.is_empty() { - return Ok(0); - } - - let mut covered: Vec = Vec::new(); - for f in files { - // rel = raw///_.md - let mut parts = f.rel.rsplitn(2, '/'); - let filename = parts.next().unwrap_or_default(); - let kind_dir = parts - .next() - .and_then(|p| p.rsplit('/').next()) - .unwrap_or_default(); - let stem = filename.strip_suffix(".md").unwrap_or(filename); - let uid = stem.split_once('_').map(|(_, u)| u).unwrap_or(stem); - - let label_match = - stems.contains(stem) || kind_uids.iter().any(|(k, u)| *k == kind_dir && u == uid); - if label_match { - covered.push(f.rel.clone()); - } - } - - mark_raw_paths_ingested(config, &covered) -} - -/// Summarise + ingest every **pending** raw file for a source. `tree_scope` -/// and `archive_source_id` as in [`raw_coverage`]. -/// -/// Each batch's files are marked covered immediately after that batch's -/// summary lands, so an interrupted run resumes where it left off. pub async fn rebuild_tree_from_raw( config: &Config, tree_scope: &str, archive_source_id: &str, -) -> Result { - let start = std::time::Instant::now(); - - let coverage = raw_coverage(config, tree_scope, archive_source_id)?; - tracing::info!( - tree_scope = %tree_scope, - archive = %archive_source_id, - total = coverage.total, - covered = coverage.covered, - pending = coverage.pending.len(), - "[memory_sync:rebuild] starting incremental rebuild" - ); - - if coverage.pending.is_empty() { - return Ok(RebuildOutcome::default()); - } - let files = coverage.pending; - - // Read pending files into SummaryInputs. - let mut inputs: Vec = Vec::with_capacity(files.len()); - let mut basenames: Vec> = Vec::with_capacity(files.len()); - let mut labels: Vec = Vec::with_capacity(files.len()); - let mut rel_paths: Vec = Vec::with_capacity(files.len()); - - for file in &files { - let body = match std::fs::read_to_string(&file.abs) { - Ok(b) => b, - Err(e) => { - tracing::warn!( - path = %file.abs.display(), - error = %e, - "[memory_sync:rebuild] skipping unreadable file" - ); - continue; - } - }; - - let filename = file - .abs - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("unknown"); - - // Parse timestamp from filename: _.md - let ts_ms = filename - .split('_') - .next() - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - let ts = chrono::DateTime::from_timestamp_millis(ts_ms).unwrap_or_else(chrono::Utc::now); - - let token_count = (body.len() / 4).max(1) as u32; - - // Wikilink basename: rel path with `.md` stripped. - let wikilink = file - .rel - .strip_suffix(".md") - .unwrap_or(&file.rel) - .to_string(); - - inputs.push(SummaryInput { - id: filename.to_string(), - content: body, - token_count, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: ts, - time_range_end: ts, - score: 0.5, - }); - labels.push(filename.to_string()); - basenames.push(Some(wikilink)); - rel_paths.push(file.rel.clone()); - } - - if inputs.is_empty() { - return Ok(RebuildOutcome { - files_read: files.len(), - ..RebuildOutcome::default() - }); - } - - let tree = get_or_create_source_tree(config, tree_scope) - .map_err(|e| anyhow::anyhow!("get_or_create_source_tree: {e:#}"))?; - - // Batch and summarise. - let batches = batch_inputs(&inputs, &labels, &basenames, &rel_paths, INPUT_TOKEN_BUDGET); - let batch_count = batches.len(); - let files_read = inputs.len(); - - tracing::info!( - tree_scope = %tree_scope, - items = files_read, - batches = batch_count, - "[memory_sync:rebuild] summarising" - ); - - // Token/charge accounting across the run. Estimate (`body.len() / 4`) is - // always summed; provider figures only replace it when every batch - // reported them (issue #3110). See `RealCostAccumulator`. - let mut cost = RealCostAccumulator::new(); - - for (batch_idx, batch) in batches.into_iter().enumerate() { - let RebuildBatch { - inputs: batch_inputs, - labels: batch_labels, - basenames: batch_basenames, - rel_paths: batch_rel_paths, - } = batch; - let batch_in_tokens: u64 = batch_inputs.iter().map(|i| i.token_count as u64).sum(); - - let ctx = SummaryContext { - tree_id: &tree.id, - tree_kind: TreeKind::Source, - target_level: 1, - token_budget: 5_000, - }; - - let output = match summarise(config, &batch_inputs, &ctx).await { - Ok(o) => o, - Err(e) => { - tracing::warn!( - error = %e, - batch = batch_idx, - "[memory_sync:rebuild] summarise failed, using fallback" - ); - fallback_summary(&batch_inputs, ctx.token_budget) - } - }; - - cost.add_batch( - batch_in_tokens, - output.token_count as u64, - output.input_tokens, - output.output_tokens, - output.charged_amount_usd, - ); - - let time_start = batch_inputs - .iter() - .map(|i| i.time_range_start) - .min() - .unwrap_or_else(chrono::Utc::now); - let time_end = batch_inputs - .iter() - .map(|i| i.time_range_end) - .max() - .unwrap_or_else(chrono::Utc::now); - - let ingest_input = SummaryIngestInput { - content: output.content, - token_count: output.token_count, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: time_start, - time_range_end: time_end, - score: 0.5, - child_labels: batch_labels, - child_basenames: batch_basenames, - }; - - let outcome = ingest_summary(config, &tree, ingest_input).await?; - - // Mark coverage ONLY after the summary landed — a crash between - // ingest and mark re-summarises this batch (duplicate summary, - // acceptable) instead of silently dropping coverage (data loss). - if let Err(e) = mark_raw_paths_ingested(config, &batch_rel_paths) { - tracing::warn!( - batch = batch_idx, - error = %format!("{e:#}"), - "[memory_sync:rebuild] failed to record raw coverage — batch may re-summarise" - ); - } - - tracing::info!( - tree_scope = %tree_scope, - batch = batch_idx, - summary_id = %outcome.summary_id, - files = batch_rel_paths.len(), - "[memory_sync:rebuild] batch ingested + coverage recorded" - ); - } - - let duration_ms = start.elapsed().as_millis() as u64; - - // Provider figures are recorded only when *every* batch reported them; a - // mixed run keeps the `len/4` estimate (which covers all batches) rather - // than a partial real total that would undercount. `estimated_cost_usd` - // is always populated as the fallback. Issue #3110. - let any_real_usage = cost.usage_is_real(); - let audit_input_tokens = cost.audit_input_tokens(); - let audit_output_tokens = cost.audit_output_tokens(); - let estimated_cost = cost.estimated_cost(); - let actual_charged_usd = cost.actual_charged_usd(); - let display_cost = actual_charged_usd.unwrap_or(estimated_cost); - - append_audit_entry( - config, - &SyncAuditEntry { - timestamp: chrono::Utc::now(), - source_id: format!("rebuild:{tree_scope}"), - source_kind: "rebuild".to_string(), - scope: tree_scope.to_string(), - items_fetched: files_read as u32, - batches: batch_count as u32, - input_tokens: audit_input_tokens, - output_tokens: audit_output_tokens, - estimated_cost_usd: estimated_cost, - composio_actions_called: 0, - composio_cost_usd: 0.0, - actual_charged_usd, - duration_ms, - success: true, - error: None, - }, - ); - - tracing::info!( - tree_scope = %tree_scope, - files = files_read, - batches = batch_count, - usage_is_real = any_real_usage, - actual_charge = actual_charged_usd.is_some(), - input_tokens = audit_input_tokens, - output_tokens = audit_output_tokens, - estimated_cost_usd = %format!("{estimated_cost:.4}"), - actual_charged_usd = ?actual_charged_usd, - display_cost_usd = %format!("{display_cost:.4}"), - duration_ms = duration_ms, - "[memory_sync:rebuild] complete" - ); - - Ok(RebuildOutcome { - files_read, - batches: batch_count, - input_tokens: audit_input_tokens, - output_tokens: audit_output_tokens, - estimated_cost_usd: estimated_cost, - actual_charged_usd, - }) -} - -/// Collect all `.md` files recursively under `dir`, skipping `_source.md`. -fn collect_raw_files(dir: &Path) -> Result> { - let mut files = Vec::new(); - collect_recursive(dir, &mut files)?; - Ok(files) -} - -fn collect_recursive(dir: &Path, out: &mut Vec) -> Result<()> { - let entries = - std::fs::read_dir(dir).map_err(|e| anyhow::anyhow!("read_dir {}: {e}", dir.display()))?; - for entry in entries { - let entry = entry?; - let path = entry.path(); - if path.is_dir() { - collect_recursive(&path, out)?; - } else if path.extension().and_then(|e| e.to_str()) == Some("md") { - let fname = path.file_name().and_then(|f| f.to_str()).unwrap_or(""); - if fname.starts_with('_') { - continue; - } - out.push(path); - } - } - Ok(()) -} - -/// One summarise-and-ingest batch with its parallel provenance vectors. -struct RebuildBatch { - inputs: Vec, - labels: Vec, - basenames: Vec>, - rel_paths: Vec, -} - -fn batch_inputs( - inputs: &[SummaryInput], - labels: &[String], - basenames: &[Option], - rel_paths: &[String], - budget: u32, -) -> Vec { - let mut batches: Vec = Vec::new(); - let mut cur = RebuildBatch { - inputs: Vec::new(), - labels: Vec::new(), - basenames: Vec::new(), - rel_paths: Vec::new(), - }; - let mut cur_tokens: u32 = 0; - - for i in 0..inputs.len() { - if !cur.inputs.is_empty() && cur_tokens + inputs[i].token_count > budget { - batches.push(std::mem::replace( - &mut cur, - RebuildBatch { - inputs: Vec::new(), - labels: Vec::new(), - basenames: Vec::new(), - rel_paths: Vec::new(), - }, - )); - cur_tokens = 0; - } - cur_tokens += inputs[i].token_count; - cur.inputs.push(inputs[i].clone()); - cur.labels.push(labels[i].clone()); - cur.basenames.push(basenames[i].clone()); - cur.rel_paths.push(rel_paths[i].clone()); - } - - if !cur.inputs.is_empty() { - batches.push(cur); - } - - batches -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider}; - use std::sync::Arc; - use tempfile::TempDir; - - fn test_config(tmp: &TempDir) -> Config { - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - cfg - } - - fn write_raw(cfg: &Config, archive_id: &str, kind: &str, name: &str, body: &str) { - let dir = cfg - .memory_tree_content_root() - .join("raw") - .join(slugify_source_id(archive_id)) - .join(kind); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write(dir.join(name), body).unwrap(); - } - - #[test] - fn collect_raw_files_skips_underscore_files() { - let tmp = tempfile::TempDir::new().unwrap(); - let emails = tmp.path().join("emails"); - std::fs::create_dir_all(&emails).unwrap(); - std::fs::write(emails.join("1000_abc.md"), "body").unwrap(); - std::fs::write(emails.join("2000_def.md"), "body2").unwrap(); - std::fs::write(emails.join("_source.md"), "meta").unwrap(); - - let files = collect_raw_files(tmp.path()).unwrap(); - assert_eq!(files.len(), 2); - assert!(files - .iter() - .all(|f| !f.to_str().unwrap().contains("_source"))); - } - - #[test] - fn collect_raw_files_recurses_subdirs() { - let tmp = tempfile::TempDir::new().unwrap(); - let sub = tmp.path().join("commits"); - std::fs::create_dir_all(&sub).unwrap(); - std::fs::write(sub.join("100_sha.md"), "commit").unwrap(); - std::fs::write(tmp.path().join("top.md"), "top").unwrap(); - - let files = collect_raw_files(tmp.path()).unwrap(); - assert_eq!(files.len(), 2); - } - - #[test] - fn batch_inputs_carries_rel_paths_alongside_provenance() { - let make = |tokens: u32, id: &str| SummaryInput { - id: id.to_string(), - content: String::new(), - token_count: tokens, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: chrono::Utc::now(), - time_range_end: chrono::Utc::now(), - score: 0.5, - }; - let inputs = vec![make(30_000, "a"), make(30_000, "b"), make(30_000, "c")]; - let labels: Vec = vec!["a".into(), "b".into(), "c".into()]; - let basenames: Vec> = vec![None, None, None]; - let rels: Vec = vec![ - "raw/x/a.md".into(), - "raw/x/b.md".into(), - "raw/x/c.md".into(), - ]; - - let batches = batch_inputs(&inputs, &labels, &basenames, &rels, 50_000); - assert_eq!(batches.len(), 3); - assert_eq!(batches[0].rel_paths, vec!["raw/x/a.md".to_string()]); - assert_eq!(batches[2].labels, vec!["c".to_string()]); - } - - #[test] - fn raw_coverage_reports_all_pending_for_untracked_archive() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - write_raw(&cfg, "github.com/o/r", "commits", "100_aaa.md", "a"); - write_raw(&cfg, "github.com/o/r", "issues", "200_7.md", "b"); - - let cov = raw_coverage(&cfg, "github:o/r", "github.com/o/r").unwrap(); - assert_eq!(cov.total, 2); - assert_eq!(cov.covered, 0); - assert_eq!(cov.pending.len(), 2); - assert!(cov.pending[0].rel.starts_with("raw/github-com-o-r/")); - assert!(needs_rebuild(&cfg, "github:o/r", "github.com/o/r")); - } - - #[test] - fn raw_coverage_empty_when_archive_dir_missing() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let cov = raw_coverage(&cfg, "github:o/r", "github.com/o/r").unwrap(); - assert_eq!(cov.total, 0); - assert!(cov.pending.is_empty()); - assert!(!needs_rebuild(&cfg, "github:o/r", "github.com/o/r")); - } - - #[test] - fn marked_files_drop_out_of_pending() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - write_raw(&cfg, "github.com/o/r", "commits", "100_aaa.md", "a"); - write_raw(&cfg, "github.com/o/r", "commits", "200_bbb.md", "b"); - - mark_raw_paths_ingested(&cfg, &["raw/github-com-o-r/commits/100_aaa.md".to_string()]) - .unwrap(); - - let cov = raw_coverage(&cfg, "github:o/r", "github.com/o/r").unwrap(); - assert_eq!(cov.total, 2); - assert_eq!(cov.covered, 1); - assert_eq!(cov.pending.len(), 1); - assert!(cov.pending[0].rel.ends_with("200_bbb.md")); - } - - #[tokio::test] - async fn backfill_marks_files_covered_by_existing_summaries() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let scope = "github:o/r"; - let archive = "github.com/o/r"; - - // Two raw files on disk; an existing L1 summary covers only the - // commit (sync-written `commit:` label shape). - write_raw(&cfg, archive, "commits", "100_abc123.md", "commit body"); - write_raw(&cfg, archive, "issues", "200_42.md", "issue body"); - - let tree = get_or_create_source_tree(&cfg, scope).unwrap(); - let input = SummaryIngestInput { - content: "summary of one commit".to_string(), - token_count: 10, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: chrono::Utc::now(), - time_range_end: chrono::Utc::now(), - score: 0.5, - child_labels: vec!["commit:abc123".to_string()], - child_basenames: Vec::new(), - }; - ingest_summary(&cfg, &tree, input).await.unwrap(); - - let cov = raw_coverage(&cfg, scope, archive).unwrap(); - assert_eq!(cov.total, 2); - assert_eq!(cov.covered, 1, "commit file backfilled as covered"); - assert_eq!(cov.pending.len(), 1); - assert!(cov.pending[0].rel.ends_with("200_42.md")); - } - - #[test] - fn chunk_raw_refs_count_as_covered() { - use crate::openhuman::memory_store::chunks::store::{ - set_chunk_raw_refs, upsert_chunks, RawRef, - }; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind as ChunkSourceKind, SourceRef, - }; - - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let scope = "gmail:user-at-example-dot-com"; - - // Two raw email files; one is referenced by a persisted chunk - // (the gmail pipeline shape), the other is orphaned. - write_raw(&cfg, scope, "emails", "1000_msg-a.md", "email a"); - write_raw(&cfg, scope, "emails", "2000_msg-b.md", "email b"); - - let now = chrono::Utc::now(); - let c = Chunk { - id: chunk_id(ChunkSourceKind::Email, scope, 0, "email a"), - content: "email a".into(), - metadata: Metadata { - source_kind: ChunkSourceKind::Email, - source_id: scope.into(), - owner: "user".into(), - timestamp: now, - time_range: (now, now), - tags: vec![], - source_ref: Some(SourceRef::new("gmail://msg-a")), - path_scope: None, - }, - token_count: 2, - seq_in_source: 0, - created_at: now, - partial_message: false, - }; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - set_chunk_raw_refs( - &cfg, - &c.id, - &[RawRef { - path: "raw/gmail-user-at-example-dot-com/emails/1000_msg-a.md".into(), - start: 0, - end: None, - }], - ) - .unwrap(); - - let cov = raw_coverage(&cfg, scope, scope).unwrap(); - assert_eq!(cov.total, 2); - assert_eq!(cov.covered, 1, "chunk-referenced file is covered"); - assert_eq!(cov.pending.len(), 1); - assert!(cov.pending[0].rel.ends_with("2000_msg-b.md")); - } - - #[tokio::test] - async fn rebuild_processes_only_pending_and_records_coverage() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let scope = "github:o/r2"; - let archive = "github.com/o/r2"; - - write_raw(&cfg, archive, "commits", "100_aaa.md", "first commit body"); - write_raw(&cfg, archive, "commits", "200_bbb.md", "second commit body"); - mark_raw_paths_ingested( - &cfg, - &["raw/github-com-o-r2/commits/100_aaa.md".to_string()], - ) - .unwrap(); - - let provider: Arc = - Arc::new(StaticChatProvider::new("rebuilt summary content")); - let outcome = test_override::with_provider(provider, async { - rebuild_tree_from_raw(&cfg, scope, archive).await.unwrap() - }) - .await; - - assert_eq!(outcome.files_read, 1, "only the pending file is read"); - assert_eq!(outcome.batches, 1); - - // Idempotent: a second run finds nothing pending. - let cov = raw_coverage(&cfg, scope, archive).unwrap(); - assert!(cov.pending.is_empty(), "all files covered after rebuild"); - assert!(!needs_rebuild(&cfg, scope, archive)); - } +) -> anyhow::Result { + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + let summariser = crate::openhuman::tinycortex::HostSummariser::new(config.clone()); + tinycortex::memory::sync::rebuild_tree_from_raw( + &memory_config, + tree_scope, + archive_source_id, + &summariser, + ) + .await } diff --git a/src/openhuman/memory_sync/sync_status/rpc.rs b/src/openhuman/memory_sync/sync_status/rpc.rs index 945118d7a..ee033d2fb 100644 --- a/src/openhuman/memory_sync/sync_status/rpc.rs +++ b/src/openhuman/memory_sync/sync_status/rpc.rs @@ -1,206 +1,30 @@ -//! JSON-RPC handler for `openhuman.memory_sync_status_list` (#1136). -//! -//! Single SQL query against `mem_tree_chunks`. Two layers of metrics: -//! -//! * **Lifetime** — `chunks_synced` (total ingested), `chunks_pending` -//! (not yet *resolved* = still in the extract+embed queue, not yet -//! appended to the source-tree buffer). -//! -//! A chunk is "resolved" (i.e. NOT pending) when ANY of: -//! - it has a row in the per-(chunk,model) sidecar -//! `mem_tree_chunk_embeddings` (#1574) — embedded under some model; -//! - `lifecycle_status = 'dropped'` — the admission gate rejected it, -//! so it is intentionally never embedded (terminal, not waiting); -//! - it has a `mem_tree_chunk_reembed_skipped` tombstone (#1574 §6) — -//! embedding failed terminally (missing body / wrong dim / embed -//! error) and will not be retried (terminal, not waiting). -//! -//! NOTE: "embedded" is keyed off the sidecar table, NOT the legacy -//! inline `mem_tree_chunks.embedding` column. The #1574 §7 migration -//! copied every vector into the sidecar and stopped writing the inline -//! column, so it now reads back NULL for every chunk. Keying pending / -//! processed off the inline column made this RPC report 100% of chunks -//! as pending and `0` processed forever, regardless of real progress. -//! Dropped / terminally-skipped chunks have no sidecar row either, so -//! without the extra terminal predicates they would read as pending -//! forever and could pin a provider's progress bar below 100%. -//! -//! * **Active sync wave** — `batch_total` / `batch_processed`. The -//! wave is identified by a *time-cluster anchor*: the earliest -//! chunk within `WAVE_WINDOW_MS` of the most recent chunk (per -//! provider). A typical sync ingests its whole batch in seconds, -//! so a 10-minute window cleanly captures one wave; if no new -//! chunks arrive, the anchor stays put. Two syncs <10min apart -//! merge into one wave (acceptable — they're contiguous activity). -//! -//! Stateless: no per-process Mutex, no persisted side table. Pure SQL -//! + the chunks table. Survives restart, safe across multiple core -//! processes. -//! -//! Trade-off: pending chunks older than `WAVE_WINDOW_MS` (e.g., -//! leftovers from a stuck earlier wave when the worker was offline) -//! show up in lifetime `chunks_pending` but not in `batch_total` — -//! deliberately, since they shouldn't pollute the active wave's -//! progress signal. +//! OpenHuman RPC shell for tinycortex synchronization status. use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; use crate::rpc::RpcOutcome; -use rusqlite::Connection; -use super::types::{FreshnessLabel, MemorySyncStatus, StatusListResponse}; +use super::types::StatusListResponse; -/// Sliding window used to identify a "current sync wave". Chunks -/// within this many ms of `MAX(created_at_ms)` for a provider count -/// as part of the wave; older chunks fall out. -const WAVE_WINDOW_MS: i64 = 10 * 60 * 1000; - -/// `openhuman.memory_sync_status_list` — one row per provider that -/// has chunks, with lifetime + active-wave counters and a freshness -/// label. pub async fn status_list_rpc(config: &Config) -> Result, String> { - tracing::debug!("[memory_sync_status][rpc] status_list"); - - let config = config.clone(); - let statuses: Vec = match tokio::task::spawn_blocking(move || { - with_connection(&config, |conn| -> anyhow::Result> { - let now_ms = chrono::Utc::now().timestamp_millis(); - Ok(query_sync_statuses(conn, now_ms)?) - }) + tracing::debug!("[memory_sync_status][rpc] status_list via tinycortex"); + let memory_config = + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()); + let statuses = match tokio::task::spawn_blocking(move || { + tinycortex::memory::sync::list_sync_statuses(&memory_config) }) .await { - Ok(Ok(rows)) => rows, - // DB unavailable (open/migration failure) or query error: return empty - // so the schema contract (`statuses` array) is always satisfied. - Ok(Err(e)) => { - tracing::warn!( - "[memory_sync_status][rpc] DB query failed, returning empty statuses: {e:#}" - ); - vec![] + Ok(Ok(statuses)) => statuses, + Ok(Err(error)) => { + tracing::warn!(%error, "[memory_sync_status][rpc] tinycortex status query failed"); + Vec::new() } - Err(e) => { - tracing::warn!( - "[memory_sync_status][rpc] spawn_blocking join error, returning empty statuses: {e}" - ); - vec![] + Err(error) => { + tracing::warn!(%error, "[memory_sync_status][rpc] status task join failed"); + Vec::new() } }; - - tracing::debug!( - "[memory_sync_status][rpc] status_list returning {} row(s)", - statuses.len() - ); - // No `single_log` wrapper: the controller serializes - // `RpcOutcome::into_cli_compatible_json`, and a non-empty `logs` list - // wraps the value in `{ result, logs }`. The frontend reads - // `resp.statuses` directly, so any envelope here breaks parsing. - Ok(RpcOutcome::new(StatusListResponse { statuses }, vec![])) -} - -/// Run the per-provider lifetime + active-wave aggregation against `conn`. -/// -/// Split out from [`status_list_rpc`] so it can be unit-tested against a -/// tempdir-backed connection without the async / spawn_blocking wrapper. -/// -/// A chunk is "resolved" (not pending) when it has a sidecar embedding (any -/// model signature), OR is `dropped`, OR carries a reembed-skip tombstone — -/// see the module header. Resolution is keyed off the `mem_tree_chunk_embeddings` -/// sidecar, NOT the legacy inline `mem_tree_chunks.embedding` column. -fn query_sync_statuses(conn: &Connection, now_ms: i64) -> rusqlite::Result> { - // Provider parsed from `source_id` prefix (substring before first ':'); - // falls back to `source_kind` when no prefix. - // - // `provider_chunks` projects per-row provider + a `resolved` flag (embedded - // OR dropped OR terminally skipped). `provider_pending` flags providers with - // at least one unresolved chunk *inside the wave window* (within - // WAVE_WINDOW_MS of the provider's most recent chunk) — `wave_anchors` is - // gated on this, so a stale unresolved chunk from an older wave can't - // resurrect an "active" wave when the recent chunks are all resolved, and a - // fully-drained provider gets `batch_total = batch_processed = 0` (the UI - // then hides the progress bar instead of rendering a completed one for an - // idle connection). `wave_anchors` finds the earliest chunk within - // WAVE_WINDOW_MS of the most recent — the wave's start. The outer SELECT - // joins back to count both lifetime and in-wave totals. - let mut stmt = conn.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 iter = stmt.query_map([WAVE_WINDOW_MS], |row| { - let provider: String = row.get(0)?; - let chunks_synced: i64 = row.get(1)?; - let chunks_pending: i64 = row.get(2)?; - let batch_total: i64 = row.get(3)?; - let batch_processed: i64 = row.get(4)?; - let last_chunk_at_ms: Option = row.get(5)?; - Ok(MemorySyncStatus { - provider, - chunks_synced: chunks_synced.max(0) as u64, - chunks_pending: chunks_pending.max(0) as u64, - batch_total: batch_total.max(0) as u64, - batch_processed: batch_processed.max(0) as u64, - last_chunk_at_ms, - freshness: FreshnessLabel::from_age_ms(last_chunk_at_ms, now_ms), - }) - })?; - iter.collect() + Ok(RpcOutcome::new(StatusListResponse { statuses }, Vec::new())) } #[cfg(test)] @@ -208,273 +32,13 @@ mod tests { use super::*; #[test] - fn status_list_response_serializes_statuses_array() { - let resp = StatusListResponse { statuses: vec![] }; - let v = serde_json::to_value(&resp).expect("serialize"); - assert!( - v.get("statuses").and_then(|s| s.as_array()).is_some(), - "statuses must always be present as an array" - ); - } - - #[test] - fn status_list_response_empty_statuses_is_empty_array() { - let resp = StatusListResponse { statuses: vec![] }; - let v = serde_json::to_value(&resp).expect("serialize"); - let arr = v["statuses"].as_array().unwrap(); - assert!(arr.is_empty()); - } - - #[test] - fn rpc_outcome_no_logs_serializes_bare_value() { - // Validates the wire contract: with empty logs, into_cli_compatible_json - // returns the value directly (not wrapped in { result, logs }). - let resp = StatusListResponse { statuses: vec![] }; - let outcome = RpcOutcome::new(resp, vec![]); - let json = outcome.into_cli_compatible_json().expect("serialize"); - assert!( - json.get("statuses").is_some(), - "bare value must have statuses at the top level" - ); - assert!(json.get("result").is_none(), "must not be double-wrapped"); - assert!(json.get("logs").is_none(), "must not be double-wrapped"); - } - - /// Regression for the legacy-column bug: pending / processed must be - /// derived from the `mem_tree_chunk_embeddings` sidecar, not the inline - /// `mem_tree_chunks.embedding` column (which is always NULL post-#1574). - /// A chunk with a sidecar row counts as processed even though its inline - /// column is NULL. - #[test] - fn pending_and_processed_key_off_sidecar_not_inline_column() { - use crate::openhuman::memory_store::chunks::store::with_connection; - use rusqlite::params; - use tempfile::TempDir; - - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - - let now = chrono::Utc::now().timestamp_millis(); - - with_connection(&cfg, |conn| { - let insert_chunk = |id: &str, source_id: &str, created: i64| { - conn.execute( - "INSERT INTO mem_tree_chunks \ - (id, source_kind, source_id, owner, timestamp_ms, \ - time_range_start_ms, time_range_end_ms, content, \ - token_count, seq_in_source, created_at_ms) \ - VALUES (?1, 'email', ?2, 'me@x.com', ?3, ?3, ?3, 'body', 10, 0, ?3)", - params![id, source_id, created], - ) - .unwrap(); - }; - let embed = |id: &str| { - conn.execute( - "INSERT INTO mem_tree_chunk_embeddings \ - (chunk_id, model_signature, vector, dim, created_at) \ - VALUES (?1, 'sig', X'00000000', 1, 0.0)", - params![id], - ) - .unwrap(); - }; - - // gmail: 3 chunks inside the active wave; 2 embedded (sidecar), 1 not. - insert_chunk("g1", "gmail:acct", now - 1_000); - insert_chunk("g2", "gmail:acct", now - 2_000); - insert_chunk("g3", "gmail:acct", now - 3_000); - embed("g1"); - embed("g2"); - - let statuses = query_sync_statuses(conn, now).unwrap(); - let gmail = statuses - .iter() - .find(|s| s.provider == "gmail") - .expect("gmail provider row"); - - assert_eq!(gmail.chunks_synced, 3, "all three ingested"); - assert_eq!( - gmail.chunks_pending, 1, - "only g3 lacks a sidecar embedding (inline column is NULL for all)" - ); - assert_eq!(gmail.batch_total, 3, "all three are within the wave window"); - assert_eq!( - gmail.batch_processed, 2, - "g1 and g2 have sidecar rows, so they count as processed" - ); - Ok(()) - }) - .unwrap(); - } - - /// A provider with every chunk embedded must report zero wave (the UI - /// hides the progress bar): `batch_total = batch_processed = 0`. - #[test] - fn fully_embedded_provider_reports_no_active_wave() { - use crate::openhuman::memory_store::chunks::store::with_connection; - use rusqlite::params; - use tempfile::TempDir; - - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let now = chrono::Utc::now().timestamp_millis(); - - with_connection(&cfg, |conn| { - conn.execute( - "INSERT INTO mem_tree_chunks \ - (id, source_kind, source_id, owner, timestamp_ms, \ - time_range_start_ms, time_range_end_ms, content, \ - token_count, seq_in_source, created_at_ms) \ - VALUES ('s1', 'slack', 'slack:eng', 'me@x.com', ?1, ?1, ?1, 'b', 10, 0, ?1)", - params![now - 5_000], - ) - .unwrap(); - conn.execute( - "INSERT INTO mem_tree_chunk_embeddings \ - (chunk_id, model_signature, vector, dim, created_at) \ - VALUES ('s1', 'sig', X'00000000', 1, 0.0)", - [], - ) - .unwrap(); - - let statuses = query_sync_statuses(conn, now).unwrap(); - let slack = statuses - .iter() - .find(|s| s.provider == "slack") - .expect("slack provider row"); - assert_eq!(slack.chunks_pending, 0); - assert_eq!(slack.batch_total, 0, "no pending chunks ⇒ no active wave"); - assert_eq!(slack.batch_processed, 0); - Ok(()) - }) - .unwrap(); - } - - /// Terminal-but-unembedded chunks must not read as perpetually pending: - /// a `dropped` chunk (admission-rejected) and a `reembed_skipped` - /// tombstoned chunk both count as resolved even with no sidecar row, so a - /// provider whose only leftovers are terminal drains to 0 pending / no wave. - #[test] - fn dropped_and_skipped_chunks_count_as_resolved_not_pending() { - use crate::openhuman::memory_store::chunks::store::with_connection; - use rusqlite::params; - use tempfile::TempDir; - - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let now = chrono::Utc::now().timestamp_millis(); - - with_connection(&cfg, |conn| { - let insert = |id: &str, lifecycle: &str, created: i64| { - conn.execute( - "INSERT INTO mem_tree_chunks \ - (id, source_kind, source_id, owner, timestamp_ms, \ - time_range_start_ms, time_range_end_ms, content, \ - token_count, seq_in_source, created_at_ms, lifecycle_status) \ - VALUES (?1, 'slack', 'slack:eng', 'me@x.com', ?2, ?2, ?2, 'b', 10, 0, ?2, ?3)", - params![id, created, lifecycle], - ) - .unwrap(); - }; - - // d1: gate-dropped (no embedding, never will be). - insert("d1", "dropped", now - 4_000); - // sk1: pending_extraction but terminally tombstoned (e.g. body missing). - insert("sk1", "pending_extraction", now - 3_000); - conn.execute( - "INSERT INTO mem_tree_chunk_reembed_skipped \ - (chunk_id, model_signature, reason, skipped_at_ms) \ - VALUES ('sk1', 'sig', 'body read failed', ?1)", - params![now - 2_000], - ) - .unwrap(); - // p1: genuinely still in the queue (no embedding, no terminal marker). - insert("p1", "pending_extraction", now - 1_000); - - let statuses = query_sync_statuses(conn, now).unwrap(); - let slack = statuses - .iter() - .find(|s| s.provider == "slack") - .expect("slack provider row"); - - assert_eq!(slack.chunks_synced, 3, "all three ingested"); - assert_eq!( - slack.chunks_pending, 1, - "only p1 is genuinely pending; d1 (dropped) and sk1 (skipped) are terminal" - ); - // p1 keeps the wave alive; d1+sk1 are in-window but resolved. - assert_eq!(slack.batch_total, 3, "all within the wave window"); - assert_eq!( - slack.batch_processed, 2, - "d1 and sk1 count as resolved; p1 does not" - ); - Ok(()) - }) - .unwrap(); - } - - /// The active wave must be gated on an unresolved chunk *inside the window*. - /// A stale unresolved chunk from an older wave plus a fully-resolved recent - /// chunk must NOT resurrect an active wave (no bogus 100%-complete bar): - /// `batch_total = batch_processed = 0`, while lifetime `chunks_pending` - /// still reflects the old straggler. - #[test] - fn stale_out_of_window_pending_does_not_open_a_wave() { - use crate::openhuman::memory_store::chunks::store::with_connection; - use rusqlite::params; - use tempfile::TempDir; - - let tmp = TempDir::new().expect("tempdir"); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - let now = chrono::Utc::now().timestamp_millis(); - // WAVE_WINDOW_MS is 10 min; place the straggler well outside it. - let old = now - 30 * 60 * 1000; - - with_connection(&cfg, |conn| { - let insert = |id: &str, created: i64| { - conn.execute( - "INSERT INTO mem_tree_chunks \ - (id, source_kind, source_id, owner, timestamp_ms, \ - time_range_start_ms, time_range_end_ms, content, \ - token_count, seq_in_source, created_at_ms) \ - VALUES (?1, 'gmail', 'gmail:acct', 'me@x.com', ?2, ?2, ?2, 'b', 10, 0, ?2)", - params![id, created], - ) - .unwrap(); - }; - - // old straggler: unresolved, 30 min ago (outside the wave window). - insert("old1", old); - // recent: resolved (embedded), inside the window. - insert("new1", now - 1_000); - conn.execute( - "INSERT INTO mem_tree_chunk_embeddings \ - (chunk_id, model_signature, vector, dim, created_at) \ - VALUES ('new1', 'sig', X'00000000', 1, 0.0)", - [], - ) - .unwrap(); - - let statuses = query_sync_statuses(conn, now).unwrap(); - let gmail = statuses - .iter() - .find(|s| s.provider == "gmail") - .expect("gmail provider row"); - - assert_eq!( - gmail.chunks_pending, 1, - "the old straggler is still pending lifetime-wise" - ); - assert_eq!( - gmail.batch_total, 0, - "no unresolved chunk inside the window ⇒ no active wave" - ); - assert_eq!(gmail.batch_processed, 0); - Ok(()) + fn response_keeps_top_level_statuses_array() { + let value = serde_json::to_value(StatusListResponse { + statuses: Vec::new(), }) .unwrap(); + assert!(value + .get("statuses") + .is_some_and(serde_json::Value::is_array)); } } diff --git a/src/openhuman/memory_sync/sync_status/types.rs b/src/openhuman/memory_sync/sync_status/types.rs index d0f88f41a..deb4d2917 100644 --- a/src/openhuman/memory_sync/sync_status/types.rs +++ b/src/openhuman/memory_sync/sync_status/types.rs @@ -1,126 +1,3 @@ -//! Memory sync status — types (#1136, simplified rewrite). -//! -//! The original implementation tracked phase + counters via push-based -//! events from each provider's sync loop. That was racy, lied about -//! "downloading 0/0" while work was in flight, and required maintaining -//! a parallel KV store. Replaced with a pull model: count chunks in -//! `mem_tree_chunks` GROUPED BY source_kind on each RPC. The chunks -//! table is the source of truth — if a chunk exists, that source has -//! synced something; the count is exact at any moment. -//! -//! Activity-freshness is derived from `MAX(timestamp_ms)` per group. +//! Sync status wire types owned by tinycortex. -use serde::Serialize; - -/// User-facing label derived from how recently chunks were ingested. -/// -/// Computed at RPC time, not stored. Boundaries are deliberate: -/// `Active` matches "currently syncing" (a fresh chunk in the last 30s -/// suggests live ingest), `Recent` covers "synced this session", and -/// `Idle` is everything older. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum FreshnessLabel { - Active, - Recent, - Idle, -} - -impl FreshnessLabel { - /// Map `last_chunk_at_ms` to a label using `now_ms` as reference. - /// Returns `Idle` when `last_chunk_at_ms` is `None`. - pub fn from_age_ms(last_chunk_at_ms: Option, now_ms: i64) -> Self { - match last_chunk_at_ms { - None => Self::Idle, - Some(ts) => { - let age = now_ms.saturating_sub(ts); - if age <= 30_000 { - Self::Active - } else if age <= 5 * 60_000 { - Self::Recent - } else { - Self::Idle - } - } - } - } -} - -/// One row per provider (slack/gmail/discord/notion/…) that has -/// produced chunks. The provider name is parsed from each chunk's -/// `source_id` prefix (everything before the first `:`). -#[derive(Clone, Debug, Serialize)] -pub struct MemorySyncStatus { - /// Specific provider — `"slack"`, `"gmail"`, `"discord"`, - /// `"telegram"`, `"whatsapp"`, `"notion"`, `"meeting_notes"`, - /// `"drive_docs"`, etc. Derived from `source_id` prefix; falls - /// back to the broad `source_kind` category for chunks whose - /// `source_id` has no `:` separator. - pub provider: String, - /// Total chunks in `mem_tree_chunks` for this source_kind. - pub chunks_synced: u64, - /// Chunks fetched + stored but not yet processed by the extract+embed - /// background worker (`embedding IS NULL`). Lifetime metric — counts - /// every still-pending chunk regardless of when it was ingested. - pub chunks_pending: u64, - /// Total chunks in the *current sync wave* — i.e., chunks created - /// at-or-after the oldest currently-pending chunk's `created_at_ms`. - /// When `chunks_pending == 0` this is also 0 (no active wave). - pub batch_total: u64, - /// Of `batch_total`, how many have been processed (`embedding IS NOT - /// NULL`) since the wave started. Progress fill = `batch_processed / - /// batch_total`. - pub batch_processed: u64, - /// Most recent chunk's `timestamp_ms` for this source_kind, or - /// `None` if no chunks yet. - pub last_chunk_at_ms: Option, - /// Derived from `last_chunk_at_ms` at RPC time. - pub freshness: FreshnessLabel, -} - -/// Wire shape of `openhuman.memory_sync_status_list`. -#[derive(Clone, Debug, Serialize)] -pub struct StatusListResponse { - pub statuses: Vec, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn freshness_label_active_within_30s() { - let now = 1_777_000_000_000; - assert_eq!( - FreshnessLabel::from_age_ms(Some(now - 1_000), now), - FreshnessLabel::Active - ); - assert_eq!( - FreshnessLabel::from_age_ms(Some(now - 29_999), now), - FreshnessLabel::Active - ); - } - - #[test] - fn freshness_label_recent_between_30s_and_5min() { - let now = 1_777_000_000_000; - assert_eq!( - FreshnessLabel::from_age_ms(Some(now - 30_001), now), - FreshnessLabel::Recent - ); - assert_eq!( - FreshnessLabel::from_age_ms(Some(now - 4 * 60_000), now), - FreshnessLabel::Recent - ); - } - - #[test] - fn freshness_label_idle_beyond_5min() { - let now = 1_777_000_000_000; - assert_eq!( - FreshnessLabel::from_age_ms(Some(now - 5 * 60_000 - 1), now), - FreshnessLabel::Idle - ); - assert_eq!(FreshnessLabel::from_age_ms(None, now), FreshnessLabel::Idle); - } -} +pub use tinycortex::memory::sync::{FreshnessLabel, MemorySyncStatus, StatusListResponse}; diff --git a/src/openhuman/memory_tools/README.md b/src/openhuman/memory_tools/README.md index 66c37a370..5c7001853 100644 --- a/src/openhuman/memory_tools/README.md +++ b/src/openhuman/memory_tools/README.md @@ -14,7 +14,7 @@ Each tool gets its own namespace `tool-{tool_name}`. Build the string via | --- | --- | | [`mod.rs`](mod.rs) | Module root + public re-exports. | | [`types.rs`](types.rs) | **Shim** — re-exports `ToolMemoryRule` / `ToolMemoryPriority` / `ToolMemorySource` / `tool_memory_namespace` from `tinycortex::memory::tool_memory::types` (W7). | -| [`store.rs`](store.rs) | **Shim** — re-exports the crate `ToolMemoryStore` (`put_rule`, `get_rule`, `list_rules`, `delete_rule`, `rules_for_prompt`, `list_tool_names`, `record`, `list_rules_json`) + `tool_memory_store(Arc)`, which bridges the host `Memory` trait object to the crate `Memory` the store needs (host = crate + `sqlite_conn`, gap G1). | +| [`store.rs`](store.rs) | **Shim** — re-exports the crate `ToolMemoryStore` (`put_rule`, `get_rule`, `list_rules`, `delete_rule`, `rules_for_prompt`, `list_tool_names`, `record`, `list_rules_json`) + `tool_memory_store(Arc)`. The trait and engine are both crate-owned. | | [`capture.rs`](capture.rs) | `ToolMemoryCaptureHook` — `PostTurnHook` impl that captures user edicts and repeated tool failures into the store (host-retained). | | [`prompt.rs`](prompt.rs) | **Shim** — re-exports the crate `ToolMemoryRulesSection` + `render_tool_memory_rules` + `TOOL_MEMORY_HEADING`, and keeps the host `PromptSection` impl that plugs the section into the system-prompt builder. | | [`tools/`](tools/) | Agent-facing read/write tools: `MemoryToolsListTool` (list rules for a tool), `MemoryToolsPutTool` (upsert a rule). | diff --git a/src/openhuman/memory_tools/store.rs b/src/openhuman/memory_tools/store.rs index 4d581246f..1cee3a757 100644 --- a/src/openhuman/memory_tools/store.rs +++ b/src/openhuman/memory_tools/store.rs @@ -2,159 +2,15 @@ //! `tinycortex::memory::tool_memory::store` (W7). //! //! The store engine (put / get / list / delete / prompt over an -//! `Arc`) is the crate's. It is generic over the **crate** `Memory` -//! trait, while host call sites hold `Arc` -//! — which is the crate trait *plus* the host-only `sqlite_conn()` escape hatch -//! (gap G1). [`HostMemoryBridge`] adapts one to the other (every method forwards -//! unchanged), so [`tool_memory_store`] builds a crate `ToolMemoryStore` over a -//! host backend without waiting on the W3 trait unification. -//! -//! Behaviour note: the deleted host engine had a `sqlite_conn` fast-path in -//! `list_rules` (a direct `memory_docs` query ordered by `updated_at`); the -//! crate engine uses the trait `list()` only. This is behaviour-equivalent — -//! the host already used `list()` as its fallback for connectionless backends — -//! and differs only by a negligible per-tool query cost. +//! `Arc`) and the `Memory` trait are both owned by tinycortex. use std::sync::Arc; -use async_trait::async_trait; - -use tinycortex::memory::Memory as CrateMemory; -use tinycortex::memory::{MemoryCategory, MemoryEntry, MemoryTaint, NamespaceSummary, RecallOpts}; - use crate::openhuman::memory::Memory; pub use tinycortex::memory::tool_memory::store::{ToolMemoryStore, TOOL_MEMORY_PROMPT_CAP}; -/// Presents a host [`Arc`] as the crate [`Memory`](CrateMemory) the -/// crate `ToolMemoryStore` is generic over. The host trait is the crate trait -/// plus `sqlite_conn`, so every method forwards verbatim (the value types are -/// already crate re-exports, so no conversion is needed). -struct HostMemoryBridge(Arc); - -#[async_trait] -impl CrateMemory for HostMemoryBridge { - fn name(&self) -> &str { - self.0.name() - } - - async fn store( - &self, - namespace: &str, - key: &str, - content: &str, - category: MemoryCategory, - session_id: Option<&str>, - ) -> anyhow::Result<()> { - self.0 - .store(namespace, key, content, category, session_id) - .await - } - - async fn store_with_taint( - &self, - namespace: &str, - key: &str, - content: &str, - category: MemoryCategory, - session_id: Option<&str>, - taint: MemoryTaint, - ) -> anyhow::Result<()> { - self.0 - .store_with_taint(namespace, key, content, category, session_id, taint) - .await - } - - async fn recall( - &self, - query: &str, - limit: usize, - opts: RecallOpts<'_>, - ) -> anyhow::Result> { - self.0.recall(query, limit, opts).await - } - - async fn recall_relevant_by_vector( - &self, - namespace: &str, - query: &str, - limit: usize, - min_vector_similarity: f64, - ) -> anyhow::Result> { - self.0 - .recall_relevant_by_vector(namespace, query, limit, min_vector_similarity) - .await - } - - async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { - self.0.get(namespace, key).await - } - - async fn list( - &self, - namespace: Option<&str>, - category: Option<&MemoryCategory>, - session_id: Option<&str>, - ) -> anyhow::Result> { - // UnifiedMemory::list() lists *documents* and surfaces each document's - // title as the entry content — so tool rules, stored as JSON in - // `memory_docs`, can't be round-tripped back through it (the crate - // `list_rules` would fail to deserialize the title as a rule and drop - // it). When the backend exposes a raw connection — as UnifiedMemory - // does — read the real content straight from `memory_docs`, mirroring - // the fast-path the host `ToolMemoryStore` used before this engine moved - // to the crate. Connectionless backends (e.g. the test `MockMemory`) - // fall back to the trait `list()`, whose content is already faithful. - if let (Some(ns), Some(conn)) = (namespace, self.0.sqlite_conn()) { - let conn = conn.lock(); - let mut stmt = conn.prepare( - "SELECT document_id, key, content, taint \ - FROM memory_docs WHERE namespace = ?1", - )?; - let rows = stmt.query_map([ns], |r| { - Ok(MemoryEntry { - id: r.get::<_, String>(0)?, - key: r.get::<_, String>(1)?, - content: r.get::<_, String>(2)?, - namespace: Some(ns.to_string()), - // These fields are unused by the sole consumer - // (`ToolMemoryStore::list_rules`, which reads key + content); - // taint is carried faithfully, the rest are placeholders. - category: MemoryCategory::Core, - timestamp: String::new(), - session_id: None, - score: None, - taint: MemoryTaint::from_db_str(&r.get::<_, String>(3)?), - }) - })?; - let mut out = Vec::new(); - for row in rows { - out.push(row?); - } - return Ok(out); - } - self.0.list(namespace, category, session_id).await - } - - async fn forget(&self, namespace: &str, key: &str) -> anyhow::Result { - self.0.forget(namespace, key).await - } - - async fn namespace_summaries(&self) -> anyhow::Result> { - self.0.namespace_summaries().await - } - - async fn count(&self) -> anyhow::Result { - self.0.count().await - } - - async fn health_check(&self) -> bool { - self.0.health_check().await - } -} - -/// Build a crate [`ToolMemoryStore`] over a host memory backend, bridging the -/// host `Memory` trait object to the crate `Memory` the store requires. +/// Build a crate [`ToolMemoryStore`] over the shared tinycortex memory trait. pub fn tool_memory_store(memory: Arc) -> ToolMemoryStore { - ToolMemoryStore::new(Arc::new(HostMemoryBridge(memory))) + ToolMemoryStore::new(memory) } diff --git a/src/openhuman/memory_tools/tools/list.rs b/src/openhuman/memory_tools/tools/list.rs index c7610e954..0560c7774 100644 --- a/src/openhuman/memory_tools/tools/list.rs +++ b/src/openhuman/memory_tools/tools/list.rs @@ -5,7 +5,7 @@ use serde::Deserialize; use serde_json::json; use crate::openhuman::memory::ops::helpers::active_memory_client; -use crate::openhuman::memory_tools::{tool_memory_store, ToolMemoryStore}; +use crate::openhuman::memory_tools::tool_memory_store; use crate::openhuman::tools::traits::{Tool, ToolResult}; pub struct MemoryToolsListTool; diff --git a/src/openhuman/memory_tools/tools/put.rs b/src/openhuman/memory_tools/tools/put.rs index 9cbe82af0..c9ebb57c0 100644 --- a/src/openhuman/memory_tools/tools/put.rs +++ b/src/openhuman/memory_tools/tools/put.rs @@ -6,7 +6,7 @@ use serde_json::json; use crate::openhuman::memory::ops::helpers::active_memory_client; use crate::openhuman::memory_tools::{ - tool_memory_store, ToolMemoryPriority, ToolMemoryRule, ToolMemorySource, ToolMemoryStore, + tool_memory_store, ToolMemoryPriority, ToolMemoryRule, ToolMemorySource, }; use crate::openhuman::tools::traits::{Tool, ToolResult}; diff --git a/src/openhuman/memory_tree/graph/bfs.rs b/src/openhuman/memory_tree/graph/bfs.rs index 9aba760f6..2f45b833e 100644 --- a/src/openhuman/memory_tree/graph/bfs.rs +++ b/src/openhuman/memory_tree/graph/bfs.rs @@ -1,215 +1,19 @@ -//! Bounded shortest-path (hop distance) over the entity co-occurrence graph. -//! -//! This is the E2GraphRAG "graph filter" step: given the entities extracted -//! from a query, decide which *pairs* are semantically related by checking -//! whether they sit within `h` hops of each other in the co-occurrence graph. -//! Pairs within range route retrieval to the *local* (index-intersection) -//! branch; an empty result routes to the *global* (dense-tree) branch. -//! -//! Everything here is pure code — no LLM. Lookups go through -//! [`super::store::neighbors`], cached per query so a node is expanded at most -//! once even when it lies between several query-entity pairs. - -use std::collections::{HashMap, HashSet, VecDeque}; +//! `Config` adapter for tinycortex-owned bounded graph traversal. use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_tree::graph::store; -/// One related query-entity pair and the hop distance between them. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PairDistance { - pub a: String, - pub b: String, - pub dist: u32, -} +pub use tinycortex::memory::graph::PairDistance; -/// Per-query neighbour cache so each entity's adjacency is read from SQLite at -/// most once across all pair BFS runs. -struct NeighborCache<'a> { - config: &'a Config, - cache: HashMap>, -} - -impl<'a> NeighborCache<'a> { - fn new(config: &'a Config) -> Self { - Self { - config, - cache: HashMap::new(), - } - } - - fn neighbors(&mut self, node: &str) -> Result<&[String]> { - if !self.cache.contains_key(node) { - let adj = store::neighbors(self.config, node)? - .into_iter() - .map(|(other, _weight)| other) - .collect(); - self.cache.insert(node.to_string(), adj); - } - // Safe: just inserted if missing. - Ok(self.cache.get(node).unwrap()) - } -} - -/// Compute hop distances between every unordered pair of `entity_ids`, keeping -/// only pairs reachable within `max_h` hops. `max_h == 0` short-circuits to an -/// empty result (no pair can be 0 hops apart — distinct entities are ≥1 hop). -/// -/// A breadth-first search runs from each entity, stopping once it has either -/// found all its peers in the query set or exhausted the `max_h` frontier. -/// Distance 1 means the two entities co-occurred directly in some chunk. pub fn pair_distances( config: &Config, entity_ids: &[String], max_h: u32, ) -> Result> { - if max_h == 0 || entity_ids.len() < 2 { - return Ok(Vec::new()); - } - // Deduplicate while preserving a stable canonical order so output pairs are - // deterministic (a < b). - let mut unique: Vec = entity_ids.to_vec(); - unique.sort(); - unique.dedup(); - if unique.len() < 2 { - return Ok(Vec::new()); - } - let targets: HashSet<&str> = unique.iter().map(|s| s.as_str()).collect(); - - let mut cache = NeighborCache::new(config); - let mut out: Vec = Vec::new(); - // Track pairs already recorded so a BFS from both endpoints doesn't emit - // the same pair twice. Keyed canonically (a < b). - let mut seen: HashSet<(String, String)> = HashSet::new(); - - for src in &unique { - // Peers we still need to reach from `src`; once empty we can stop the - // BFS early instead of draining the whole frontier (latency then - // scales with unresolved pairs, not graph size). - let mut remaining: HashSet<&str> = targets - .iter() - .copied() - .filter(|t| *t != src.as_str()) - .collect(); - if remaining.is_empty() { - continue; - } - // BFS frontier from `src` bounded to `max_h`. - let mut visited: HashSet = HashSet::new(); - visited.insert(src.clone()); - let mut queue: VecDeque<(String, u32)> = VecDeque::new(); - queue.push_back((src.clone(), 0)); - - while let Some((node, dist)) = queue.pop_front() { - if dist >= max_h || remaining.is_empty() { - continue; - } - let neighbors: Vec = cache.neighbors(&node)?.to_vec(); - for nb in neighbors { - if !visited.insert(nb.clone()) { - continue; - } - let nd = dist + 1; - // Record a pair the moment we reach another query target. BFS - // guarantees `nd` is the shortest distance. - if nb.as_str() != src.as_str() && remaining.remove(nb.as_str()) { - let (a, b) = if src.as_str() < nb.as_str() { - (src.clone(), nb.clone()) - } else { - (nb.clone(), src.clone()) - }; - if seen.insert((a.clone(), b.clone())) { - out.push(PairDistance { a, b, dist: nd }); - } - } - queue.push_back((nb, nd)); - } - } - } - - out.sort_by(|x, y| { - x.dist - .cmp(&y.dist) - .then_with(|| x.a.cmp(&y.a)) - .then_with(|| x.b.cmp(&y.b)) - }); - Ok(out) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_tree::graph::store::{pairs_from_entities, upsert_edges}; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) - } - - /// Seed a path graph alice — bob — carol (two chunks). - fn seed_path(cfg: &Config) { - upsert_edges( - cfg, - &pairs_from_entities(&["person:alice".into(), "person:bob".into()]), - 1_000, - ) - .unwrap(); - upsert_edges( - cfg, - &pairs_from_entities(&["person:bob".into(), "person:carol".into()]), - 1_000, - ) - .unwrap(); - } - - #[test] - fn direct_cooccurrence_is_distance_one() { - let (_tmp, cfg) = test_config(); - seed_path(&cfg); - let pairs = pair_distances(&cfg, &["person:alice".into(), "person:bob".into()], 2).unwrap(); - assert_eq!(pairs.len(), 1); - assert_eq!(pairs[0].dist, 1); - assert_eq!(pairs[0].a, "person:alice"); - assert_eq!(pairs[0].b, "person:bob"); - } - - #[test] - fn two_hop_pair_found_within_h2_not_h1() { - let (_tmp, cfg) = test_config(); - seed_path(&cfg); - let ids = vec!["person:alice".to_string(), "person:carol".to_string()]; - - let h1 = pair_distances(&cfg, &ids, 1).unwrap(); - assert!( - h1.is_empty(), - "alice-carol is 2 hops apart; h=1 finds nothing" - ); - - let h2 = pair_distances(&cfg, &ids, 2).unwrap(); - assert_eq!(h2.len(), 1); - assert_eq!(h2[0].dist, 2); - } - - #[test] - fn disconnected_entities_yield_no_pairs() { - let (_tmp, cfg) = test_config(); - seed_path(&cfg); - // dave is not in the graph at all. - let pairs = - pair_distances(&cfg, &["person:alice".into(), "person:dave".into()], 3).unwrap(); - assert!(pairs.is_empty()); - } - - #[test] - fn single_entity_yields_no_pairs() { - let (_tmp, cfg) = test_config(); - seed_path(&cfg); - let pairs = pair_distances(&cfg, &["person:alice".into()], 2).unwrap(); - assert!(pairs.is_empty()); - } + tinycortex::memory::graph::pair_distances( + &crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()), + entity_ids, + max_h, + ) } diff --git a/src/openhuman/memory_tree/graph/store.rs b/src/openhuman/memory_tree/graph/store.rs index f14410c36..a73463fa0 100644 --- a/src/openhuman/memory_tree/graph/store.rs +++ b/src/openhuman/memory_tree/graph/store.rs @@ -1,230 +1,43 @@ -//! Persistence for the entity co-occurrence graph (`mem_tree_entity_edges`). -//! -//! The edge table is an undirected weighted graph: each row is one unordered -//! pair of canonical entity ids that co-occurred inside the same chunk, with -//! `weight` accumulating the number of co-occurrences. Canonical ordering -//! (`entity_a < entity_b`) keeps a pair on a single row; neighbour lookups -//! union both columns so a query entity matches regardless of which side it -//! was stored on. -//! -//! The schema is declared in `memory_store/chunks/store.rs::SCHEMA`; this file -//! only owns the CRUD operations. Writes happen inside the same transaction -//! that indexes a chunk's entities (see `score::persist_score_tx`) so the -//! graph never diverges from `mem_tree_entity_index`. +//! `Config` adapters for tinycortex-owned persisted graph edges. use anyhow::Result; -use rusqlite::{params, Transaction}; +use rusqlite::Transaction; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; -/// Order a pair canonically so `(a, b)` and `(b, a)` collapse onto one row. -/// Returns `None` for a self-pair (`a == b`) — an entity never edges itself. -fn order_pair<'a>(a: &'a str, b: &'a str) -> Option<(&'a str, &'a str)> { - use std::cmp::Ordering; - match a.cmp(b) { - Ordering::Less => Some((a, b)), - Ordering::Greater => Some((b, a)), - Ordering::Equal => None, - } +pub use tinycortex::memory::graph::pairs_from_entities; + +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -/// Emit every unordered pair from a set of canonical entity ids. Deduplicates -/// the input first (a chunk can mention the same entity in several spans) so a -/// repeated id doesn't inflate a single chunk's contribution to one edge. -pub fn pairs_from_entities(entity_ids: &[String]) -> Vec<(String, String)> { - use std::collections::BTreeSet; - let unique: Vec<&String> = entity_ids - .iter() - .collect::>() - .into_iter() - .collect(); - let mut out = Vec::new(); - for i in 0..unique.len() { - for j in (i + 1)..unique.len() { - if let Some((a, b)) = order_pair(unique[i], unique[j]) { - out.push((a.to_string(), b.to_string())); - } - } - } - out -} - -/// Upsert a batch of co-occurrence pairs inside an existing transaction, -/// incrementing `weight` for pairs that already exist. Idempotent on the -/// `(entity_a, entity_b)` primary key — re-running adds to the weight, which -/// is the intended accumulation semantics (more co-occurrences = stronger -/// edge). `pairs` are expected to already be canonically ordered via -/// [`pairs_from_entities`]. pub fn upsert_edges_tx( - tx: &Transaction<'_>, + transaction: &Transaction<'_>, pairs: &[(String, String)], timestamp_ms: i64, ) -> Result { - // Defensively canonicalise (`a < b`) and dedup here too, so a caller that - // passes reversed or duplicate pairs can never split one edge across two - // rows (`a,b` + `b,a`) or double-count a single chunk's contribution. - // `pairs_from_entities` already does this, but the table invariant must not - // depend on callers getting it right. - use std::collections::BTreeSet; - let canonical: BTreeSet<(String, String)> = pairs - .iter() - .filter_map(|(a, b)| order_pair(a, b).map(|(x, y)| (x.to_string(), y.to_string()))) - .collect(); - if canonical.is_empty() { - return Ok(0); - } - let mut stmt = 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 { - stmt.execute(params![a, b, timestamp_ms])?; - } - Ok(canonical.len()) + tinycortex::memory::graph::upsert_edges_tx(transaction, pairs, timestamp_ms) } -/// Convenience wrapper that opens its own connection + transaction. Prefer -/// [`upsert_edges_tx`] when an enclosing transaction already exists (the -/// ingest hook does), so the graph write commits atomically with the entity -/// index write. pub fn upsert_edges( config: &Config, pairs: &[(String, String)], timestamp_ms: i64, ) -> Result { - if pairs.is_empty() { - return Ok(0); - } - with_connection(config, |conn| { - let tx = conn.unchecked_transaction()?; - let n = upsert_edges_tx(&tx, pairs, timestamp_ms)?; - tx.commit()?; - Ok(n) - }) + tinycortex::memory::graph::upsert_edges(&engine_config(config), pairs, timestamp_ms) } -/// Return the immediate neighbours of `entity_id` with their edge weights. -/// Unions both columns so the lookup is symmetric. Ordered by weight DESC so -/// the strongest associations come first (callers can cap traversal breadth). pub fn neighbors(config: &Config, entity_id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT entity_b AS other, weight FROM mem_tree_entity_edges WHERE entity_a = ?1 - UNION ALL - SELECT entity_a AS other, weight FROM mem_tree_entity_edges WHERE entity_b = ?1 - ORDER BY weight DESC", - )?; - let rows = stmt - .query_map(params![entity_id], |row| { - Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)) - })? - .collect::>>()?; - Ok(rows) - }) + tinycortex::memory::graph::edge_neighbors(&engine_config(config), entity_id) } -/// Remove every edge touching any of `entity_ids`. Used before re-indexing a -/// re-scored chunk so co-occurrences that no longer hold don't leak through as -/// stale edges (mirrors `clear_entity_index_for_node`). Note this is coarse — -/// it drops the entity's edges to *all* peers, not just the ones from this -/// chunk; the subsequent re-index rebuilds the surviving edges from the -/// chunk's current entities. Edges contributed by other chunks are recovered -/// on their next re-score; for the common append-only ingest path this branch -/// is not exercised. -pub fn clear_edges_for_entities_tx(tx: &Transaction<'_>, entity_ids: &[String]) -> Result { - if entity_ids.is_empty() { - return Ok(0); - } - let mut removed = 0; - let mut stmt = - tx.prepare("DELETE FROM mem_tree_entity_edges WHERE entity_a = ?1 OR entity_b = ?1")?; - for id in entity_ids { - removed += stmt.execute(params![id])?; - } - Ok(removed) +pub fn clear_edges_for_entities_tx( + transaction: &Transaction<'_>, + entity_ids: &[String], +) -> Result { + tinycortex::memory::graph::clear_edges_for_entities_tx(transaction, entity_ids) } -/// Count edge rows (for tests / diagnostics). pub fn count_edges(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_entity_edges", [], |r| { - r.get(0) - })?; - Ok(n.max(0) as u64) - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) - } - - #[test] - fn pairs_are_canonically_ordered_and_deduped() { - let pairs = pairs_from_entities(&[ - "person:bob".into(), - "person:alice".into(), - "person:alice".into(), // dup — must not double an edge - ]); - assert_eq!(pairs, vec![("person:alice".into(), "person:bob".into())]); - } - - #[test] - fn self_pairs_are_skipped() { - let pairs = pairs_from_entities(&["x:1".into(), "x:1".into()]); - assert!(pairs.is_empty()); - } - - #[test] - fn upsert_increments_weight_and_neighbors_is_symmetric() { - let (_tmp, cfg) = test_config(); - let p = vec![("person:alice".to_string(), "person:bob".to_string())]; - upsert_edges(&cfg, &p, 1_000).unwrap(); - upsert_edges(&cfg, &p, 2_000).unwrap(); - - // Symmetric lookup from either endpoint. - let from_alice = neighbors(&cfg, "person:alice").unwrap(); - assert_eq!(from_alice, vec![("person:bob".to_string(), 2)]); - let from_bob = neighbors(&cfg, "person:bob").unwrap(); - assert_eq!(from_bob, vec![("person:alice".to_string(), 2)]); - assert_eq!(count_edges(&cfg).unwrap(), 1); - } - - #[test] - fn clear_edges_removes_all_touching() { - let (_tmp, cfg) = test_config(); - upsert_edges( - &cfg, - &pairs_from_entities(&[ - "person:alice".into(), - "person:bob".into(), - "topic:phoenix".into(), - ]), - 1_000, - ) - .unwrap(); - assert_eq!(count_edges(&cfg).unwrap(), 3); - - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - let n = clear_edges_for_entities_tx(&tx, &["person:alice".into()]).unwrap(); - tx.commit()?; - // alice-bob and alice-phoenix removed; bob-phoenix survives. - assert_eq!(n, 2); - Ok(()) - }) - .unwrap(); - assert_eq!(count_edges(&cfg).unwrap(), 1); - assert!(neighbors(&cfg, "person:alice").unwrap().is_empty()); - } + tinycortex::memory::graph::count_edges(&engine_config(config)) } diff --git a/src/openhuman/memory_tree/ingest.rs b/src/openhuman/memory_tree/ingest.rs index b5cb17733..a65ae8a5c 100644 --- a/src/openhuman/memory_tree/ingest.rs +++ b/src/openhuman/memory_tree/ingest.rs @@ -1,371 +1,61 @@ -//! Single-shape entry point for landing a pre-built summary into a tree. -//! -//! Sync pipelines produce summaries (via an LLM call over batched source -//! content) and hand them here. `ingest_summary` writes the `.md` file to -//! `wiki/summaries/source-/L1/…`, inserts the `SummaryNode` row, -//! appends to the L1 buffer, and cascades seals upward when the buffer -//! crosses the fanout threshold. -//! -//! Embeddings are temporarily disabled — summaries land without vectors. +//! Product artifact hooks around tinycortex-owned direct summary ingestion. use anyhow::{Context, Result}; -use chrono::{DateTime, Utc}; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; -use crate::openhuman::memory_store::content::atomic::{stage_summary, StagedSummary}; use crate::openhuman::memory_store::content::wiki_git::{SummaryCommitBatch, SummaryCommitEntry}; -use crate::openhuman::memory_store::content::SummaryComposeInput; -use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, SUMMARY_FANOUT}; -use crate::openhuman::memory_tree::tree::factory::TreeFactory; -use crate::openhuman::memory_tree::tree::registry::new_summary_id; -use crate::openhuman::memory_tree::tree::store; +use crate::openhuman::memory_store::trees::types::Tree; +use crate::openhuman::tinycortex::{memory_config_from, HostSummariser}; -/// Input for `ingest_summary`. Callers provide the summary text and -/// envelope metadata; the function handles id generation, file staging, -/// DB persistence, and buffer management. -#[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, - /// Opaque child references for the wikilink list in the `.md` front-matter. - /// These are NOT tree children in the DB sense — they're provenance - /// labels (e.g. `"commit:abc"`, `"issue:42"`) that appear in the - /// summary's `children:` YAML block for human readers. - pub child_labels: Vec, - /// Per-child wikilink overrides (parallel to `child_labels`). When - /// `Some(path)`, the wikilink points at `[[path]]` instead of the - /// sanitised child label. Use this to link to raw archive files so - /// Obsidian can resolve the wikilinks. - #[doc(hidden)] - pub child_basenames: Vec>, -} +pub use tinycortex::memory::tree::{SummaryIngestInput, SummaryIngestOutcome}; -/// What `ingest_summary` did. -#[derive(Clone, Debug)] -pub struct SummaryIngestOutcome { - pub summary_id: String, - pub content_path: String, - /// Summary ids that sealed during the cascade triggered by this ingest. - pub sealed_ids: Vec, -} - -/// Ingest a pre-built summary into `tree` as an L1 node. -/// -/// 1. Generate a summary id. -/// 2. Stage the `.md` file under `wiki/summaries/source-/L1/…`. -/// 3. Insert the `SummaryNode` row. -/// 4. Append to the L1 buffer. -/// 5. Cascade seals if the L1 buffer crosses `SUMMARY_FANOUT`. -/// -/// Embeddings are skipped — the `embedding` column is `None`. pub async fn ingest_summary( config: &Config, tree: &Tree, input: SummaryIngestInput, ) -> Result { - let target_level: u32 = 1; - let summary_id = new_summary_id(target_level); - let now = Utc::now(); - - let tree_factory = TreeFactory::from_tree(tree); - let scope_slug = tree_factory.scope_slug(); - let summary_tree_kind = tree_factory.summary_tree_kind(); + log::debug!( + "[memory_tree::ingest] tinycortex enter tree_kind={} children={}", + tree.kind.as_str(), + input.child_labels.len() + ); let content_root = config.memory_tree_content_root(); - - // Ensure Obsidian defaults are present (best-effort). - if let Err(e) = + if let Err(error) = crate::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults(&content_root) { - tracing::warn!( - error = %format!("{e:#}"), - "[memory_tree::ingest] ensure_obsidian_defaults failed" - ); + log::warn!("[memory_tree::ingest] obsidian defaults failed: {error:#}"); } - // Stage the .md file BEFORE the DB transaction. - let compose_input = SummaryComposeInput { - summary_id: &summary_id, - tree_kind: summary_tree_kind, - tree_id: &tree.id, - tree_scope: &tree.scope, - level: target_level, - child_ids: &input.child_labels, - child_basenames: if input.child_basenames.is_empty() { - None - } else { - Some(&input.child_basenames) - }, - child_count: input.child_labels.len(), - time_range_start: input.time_range_start, - time_range_end: input.time_range_end, - sealed_at: now, - body: &input.content, - }; - - let staged = stage_summary(&content_root, &compose_input, &scope_slug) - .with_context(|| format!("stage_summary failed for {summary_id}"))?; - - tracing::debug!( - summary_id = %summary_id, - path = %staged.content_path, - "[memory_tree::ingest] staged summary" - ); - - 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: now, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - - let summary_commit = ingested_summary_batch(tree, &node, &staged.content_path); - - // Persist summary + update buffer in one transaction. - persist_and_buffer(config, tree, &node, &staged, target_level)?; - - commit_ingested_summary(config, &summary_commit) - .with_context(|| format!("commit ingested summary {summary_id} after DB persist"))?; - - // Cascade seals upward from L1 if the buffer crossed the fanout gate. - let sealed_ids = cascade_from(config, tree, target_level).await?; - - tracing::info!( - summary_id = %summary_id, - tree_id = %tree.id, - sealed = sealed_ids.len(), - "[memory_tree::ingest] ingested summary" - ); - - Ok(SummaryIngestOutcome { - summary_id, - content_path: staged.content_path, - sealed_ids, - }) -} - -fn ingested_summary_batch( - tree: &Tree, - node: &SummaryNode, - content_path: &str, -) -> SummaryCommitBatch { - SummaryCommitBatch { - reason: "summary_ingest".to_string(), - tree_id: tree.id.clone(), - tree_scope: tree.scope.clone(), - entries: vec![SummaryCommitEntry { - summary_id: node.id.clone(), - content_path: content_path.to_string(), - level: node.level, - child_count: node.child_ids.len(), - token_count: node.token_count, - time_range_start: node.time_range_start, - time_range_end: node.time_range_end, - }], - } -} - -fn commit_ingested_summary(config: &Config, batch: &SummaryCommitBatch) -> Result<()> { - crate::openhuman::memory_store::content::wiki_git::commit_summaries( - &config.memory_tree_content_root(), - batch, + let outcome = tinycortex::memory::tree::ingest_summary( + &memory_config_from(config, config.workspace_dir.clone()), + tree, + input.clone(), + &HostSummariser::new(config.clone()), ) -} - -fn persist_and_buffer( - config: &Config, - tree: &Tree, - node: &SummaryNode, - staged: &StagedSummary, - target_level: u32, -) -> Result<()> { - let tree_id = tree.id.clone(); - let summary_id = node.id.clone(); - let token_count = node.token_count; - let time_range_start = node.time_range_start; - let now = node.sealed_at; - let node = node.clone(); - - let model_signature = - crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - let staged = staged.clone(); - - with_connection(config, move |conn| { - let tx = conn.unchecked_transaction()?; - - let current_max: u32 = tx - .query_row( - "SELECT max_level FROM mem_tree_trees WHERE id = ?1", - rusqlite::params![&tree_id], - |r| r.get::<_, i64>(0), - ) - .map(|n| n.max(0) as u32) - .unwrap_or(0); - - store::insert_summary_tx(&tx, &node, Some(&staged), &model_signature)?; - - // Index entities for retrieval. - crate::openhuman::memory_tree::score::store::index_summary_entity_ids_tx( - &tx, - &node.entities, - &node.id, - node.score, - now.timestamp_millis(), - Some(&tree_id), - )?; - - // Append to L1 buffer. - let mut buf = store::get_buffer_conn(&tx, &tree_id, target_level)?; - if !buf.item_ids.iter().any(|id| id == &summary_id) { - buf.item_ids.push(summary_id.clone()); - buf.token_sum = buf.token_sum.saturating_add(token_count as i64); - buf.oldest_at = match buf.oldest_at { - Some(existing) => Some(existing.min(time_range_start)), - None => Some(time_range_start), - }; - store::upsert_buffer_tx(&tx, &buf)?; - } - - // Update tree max_level if this is the first L1 node. - if target_level > current_max { - store::update_tree_after_seal_tx(&tx, &tree_id, &summary_id, target_level, now)?; - } - - tx.commit()?; - Ok(()) - }) -} - -/// Cascade seals starting at `start_level` using the existing bucket_seal -/// machinery. Only fires when the buffer at `start_level` has enough -/// siblings. -async fn cascade_from(config: &Config, tree: &Tree, start_level: u32) -> Result> { - use crate::openhuman::memory_tree::tree::bucket_seal::{cascade_all_from, LabelStrategy}; - - let buf = store::get_buffer(config, &tree.id, start_level)?; - if (buf.item_ids.len() as u32) < SUMMARY_FANOUT { - return Ok(Vec::new()); - } - - cascade_all_from(config, tree, start_level, None, &LabelStrategy::Empty).await -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree; - use tempfile::TempDir; - - fn test_config(tmp: &TempDir) -> Config { - Config { - workspace_dir: tmp.path().join("workspace"), - config_path: tmp.path().join("config.toml"), - ..Config::default() - } - } - - #[tokio::test] - async fn ingest_summary_writes_to_l1_buffer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let tree = get_or_create_source_tree(&cfg, "github:org/repo").unwrap(); - let input = SummaryIngestInput { - content: "Summary of 100 commits about feature X.".to_string(), - token_count: 50, - entities: Vec::new(), - topics: vec!["github".to_string()], - time_range_start: Utc::now(), - time_range_end: Utc::now(), - score: 0.7, - child_labels: vec!["commit:abc".to_string(), "commit:def".to_string()], - child_basenames: Vec::new(), - }; - - let outcome = ingest_summary(&cfg, &tree, input).await.unwrap(); - assert!(outcome.summary_id.starts_with("summary:")); - assert!(outcome.content_path.contains("/L1/")); - assert!(outcome.sealed_ids.is_empty()); - - let buf = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - assert_eq!(buf.item_ids.len(), 1); - assert_eq!(buf.item_ids[0], outcome.summary_id); - assert_eq!(buf.token_sum, 50); - } - - #[tokio::test] - async fn ingest_summary_is_idempotent_on_buffer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let tree = get_or_create_source_tree(&cfg, "github:org/repo2").unwrap(); - - for _ in 0..3 { - let input = SummaryIngestInput { - content: "Same summary.".to_string(), - token_count: 10, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: Utc::now(), - time_range_end: Utc::now(), - score: 0.5, - child_labels: Vec::new(), - child_basenames: Vec::new(), - }; - ingest_summary(&cfg, &tree, input).await.unwrap(); - } - - let buf = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - // Each call generates a unique summary_id, so all 3 should be in the buffer. - assert_eq!(buf.item_ids.len(), 3); - } - - #[tokio::test] - async fn ingest_summary_writes_md_file() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let tree = get_or_create_source_tree(&cfg, "github:org/repo3").unwrap(); - let input = SummaryIngestInput { - content: "Test summary content.".to_string(), - token_count: 20, - entities: Vec::new(), - topics: Vec::new(), - time_range_start: Utc::now(), - time_range_end: Utc::now(), - score: 0.5, - child_labels: Vec::new(), - child_basenames: Vec::new(), - }; - - let outcome = ingest_summary(&cfg, &tree, input).await.unwrap(); - let content_root = cfg.memory_tree_content_root(); - let abs_path = content_root.join(&outcome.content_path); - assert!(abs_path.exists(), "summary .md file should exist on disk"); - - let contents = std::fs::read_to_string(&abs_path).unwrap(); - assert!(contents.contains("Test summary content.")); - } + .await?; + + crate::openhuman::memory_store::content::wiki_git::commit_summaries( + &content_root, + &SummaryCommitBatch { + reason: "summary_ingest".to_string(), + tree_id: tree.id.clone(), + tree_scope: tree.scope.clone(), + entries: vec![SummaryCommitEntry { + summary_id: outcome.summary_id.clone(), + content_path: outcome.content_path.clone(), + level: 1, + child_count: input.child_labels.len(), + token_count: input.token_count, + time_range_start: input.time_range_start, + time_range_end: input.time_range_end, + }], + }, + ) + .with_context(|| format!("commit ingested summary {}", outcome.summary_id))?; + + log::debug!( + "[memory_tree::ingest] tinycortex complete sealed={}", + outcome.sealed_ids.len() + ); + Ok(outcome) } diff --git a/src/openhuman/memory_tree/io.rs b/src/openhuman/memory_tree/io.rs index a029cd5bf..b6b92f6e4 100644 --- a/src/openhuman/memory_tree/io.rs +++ b/src/openhuman/memory_tree/io.rs @@ -1,231 +1,6 @@ -//! Canonical input/output types for the memory_tree module. -//! -//! memory_tree exposes two fundamental operations against any tree: -//! -//! - **Write** — append a chunk (leaf) into a tree; cascading bucket-seals -//! may produce new summary nodes at higher levels. -//! - **Read** — navigate from a node into its descendants, optionally -//! reranked by a query embedding. -//! -//! Internal mechanics (bucket_seal, flush, walk, summarise) take a mix of -//! `&Tree`, `&LeafRef`, `WalkOptions`, etc. This module defines a single -//! pair of contract types per direction so callers above memory_tree -//! (orchestrator, jobs, RPC) can talk to the module in one consistent -//! shape regardless of which tree kind they're targeting. -//! -//! These are pure contract types — no logic, no IO, no storage. They -//! compose the existing primitives from [`crate::openhuman::memory_store::trees`] -//! and the mechanics submodules; convert at the call boundary. +//! Stable host path for tinycortex-owned tree I/O contracts. -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use crate::openhuman::memory_store::trees::{Tree, TreeKind}; -use crate::openhuman::memory_tree::tree::bucket_seal::LeafRef; - -// ───────────────────────── Write ───────────────────────── - -/// A leaf payload ready to be appended to a tree. Mirror of [`LeafRef`] -/// with serde derives so RPC callers and job payloads can share one shape. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct TreeLeafPayload { - pub chunk_id: String, - pub token_count: u32, - pub timestamp: DateTime, - pub content: String, - #[serde(default)] - pub entities: Vec, - #[serde(default)] - pub topics: Vec, - #[serde(default)] - pub score: f32, -} - -impl From<&TreeLeafPayload> for LeafRef { - fn from(p: &TreeLeafPayload) -> Self { - LeafRef { - chunk_id: p.chunk_id.clone(), - token_count: p.token_count, - timestamp: p.timestamp, - content: p.content.clone(), - entities: p.entities.clone(), - topics: p.topics.clone(), - score: p.score, - } - } -} - -impl From for TreeLeafPayload { - fn from(l: LeafRef) -> Self { - Self { - chunk_id: l.chunk_id, - token_count: l.token_count, - timestamp: l.timestamp, - content: l.content, - entities: l.entities, - topics: l.topics, - score: l.score, - } - } -} - -/// How sealed summaries should be labelled with entities/topics. Mirrors -/// [`crate::openhuman::memory_tree::tree::bucket_seal::LabelStrategy`] in a -/// serde-friendly shape. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TreeLabelStrategy { - /// Inherit entities/topics from child leaves (default). - #[default] - Inherit, - /// Re-extract from the summary text via the resolver. - Extract, - /// Leave entities/topics empty. - Empty, -} - -/// Canonical write request: "append this leaf to this tree". -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct TreeWriteRequest { - /// Target tree id. Must already exist (callers go through the - /// kind-specific registry to get one). - pub tree_id: String, - /// Tree kind — informational, carried through to log lines and - /// downstream policy. Storage doesn't branch on it. - pub tree_kind: TreeKind, - /// The leaf to append. - pub leaf: TreeLeafPayload, - /// Labelling strategy applied to any summaries that seal during this - /// call. Defaults to [`TreeLabelStrategy::Inherit`]. - #[serde(default)] - pub label_strategy: TreeLabelStrategy, - /// When `true`, only stage the leaf in the L0 buffer; do NOT cascade - /// seals synchronously. Use this from job-driven pipelines where seal - /// work is enqueued separately. Default `false`. - #[serde(default)] - pub deferred: bool, -} - -/// Canonical write outcome: which buffers sealed (if any) and the summary -/// ids produced. -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub struct TreeWriteOutcome { - /// Ids of summary nodes that sealed during this append. Empty when the - /// L0 buffer was below budget or the call was deferred. - pub new_summary_ids: Vec, - /// Set to `true` when the caller used `deferred = true` and should - /// enqueue a follow-up seal job for level 0. - pub seal_pending: bool, -} - -// ───────────────────────── Read ───────────────────────── - -/// What the caller wants out of the read. Bounds the BFS and controls -/// query-driven reranking. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct TreeReadRequest { - /// Tree id. Required so the read scope is explicit even when starting - /// from a known node. - pub tree_id: String, - /// Starting node. `None` → start from the tree root. - #[serde(default)] - pub start_node_id: Option, - /// Maximum levels to descend from `start_node_id`. `0` returns an - /// empty result. - pub max_depth: u32, - /// Optional natural-language query. When `Some`, hits are reranked by - /// cosine similarity to the query embedding; hits with no stored - /// embedding sort to the bottom. When `None`, BFS order is preserved. - #[serde(default)] - pub query: Option, - /// Max hits to return. `None` → backend default. - #[serde(default)] - pub limit: Option, -} - -/// One hit returned by a tree read. Compact projection — for the full -/// SummaryNode/Chunk row use the memory_store retrieval surface directly. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct TreeReadHit { - pub node_id: String, - /// `"summary"` for sealed nodes, `"chunk"` for leaves. - pub node_kind: String, - /// Level in the tree (0 = leaf, 1+ = summary). - pub level: u32, - /// Summary text or chunk content, truncated by the backend if oversize. - pub content: String, - /// Cosine similarity score when `query` was set; `0.0` otherwise. - #[serde(default)] - pub score: f32, -} - -/// Result of a tree read. -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub struct TreeReadResult { - pub hits: Vec, - /// Total matches BEFORE `limit` truncation. - pub total: usize, - /// Echoes back the tree id the read targeted — useful for callers that - /// fan out and need to attribute hits. - pub tree_id: String, -} - -impl TreeReadResult { - pub fn empty(tree: &Tree) -> Self { - Self { - hits: Vec::new(), - total: 0, - tree_id: tree.id.clone(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::trees::TreeStatus; - - fn sample_tree() -> Tree { - Tree { - id: "tree-1".into(), - kind: TreeKind::Source, - scope: "chat:slack:#eng".into(), - root_id: Some("root-1".into()), - max_level: 2, - status: TreeStatus::Active, - created_at: Utc::now(), - last_sealed_at: None, - } - } - - #[test] - fn tree_leaf_payload_converts_to_and_from_leaf_ref() { - let payload = TreeLeafPayload { - chunk_id: "chunk-1".into(), - token_count: 12, - timestamp: Utc::now(), - content: "hello".into(), - entities: vec!["person:alice".into()], - topics: vec!["deploy".into()], - score: 0.75, - }; - let leaf: LeafRef = (&payload).into(); - let roundtrip = TreeLeafPayload::from(leaf); - - assert_eq!(roundtrip.chunk_id, payload.chunk_id); - assert_eq!(roundtrip.token_count, payload.token_count); - assert_eq!(roundtrip.content, payload.content); - assert_eq!(roundtrip.entities, payload.entities); - assert_eq!(roundtrip.topics, payload.topics); - assert_eq!(roundtrip.score, payload.score); - } - - #[test] - fn tree_read_result_empty_copies_tree_id() { - let tree = sample_tree(); - let result = TreeReadResult::empty(&tree); - assert_eq!(result.tree_id, tree.id); - assert_eq!(result.total, 0); - assert!(result.hits.is_empty()); - } -} +pub use tinycortex::memory::tree::{ + TreeLabelStrategy, TreeLeafPayload, TreeReadHit, TreeReadRequest, TreeReadResult, + TreeWriteOutcome, TreeWriteRequest, +}; diff --git a/src/openhuman/memory_tree/retrieval/cover.rs b/src/openhuman/memory_tree/retrieval/cover.rs index 0379d6810..f60af258d 100644 --- a/src/openhuman/memory_tree/retrieval/cover.rs +++ b/src/openhuman/memory_tree/retrieval/cover.rs @@ -1,64 +1,13 @@ -//! `memory_tree_cover_window` — the minimum-node cover of a time window. -//! -//! Given an explicit `[since_ms, until_ms]` window (and optional source -//! filter), return the **smallest set of nodes that covers every in-window -//! chunk** — a heterogeneous mix of summary nodes (where a whole subtree is -//! in-window) and raw leaf chunks (everything else). This is the read path the -//! morning brief uses for "last 24h" so it summarises only fresh content -//! instead of the all-time root blob. -//! -//! The cover is **purely structural** — it does not assume the tree is -//! calendar-bucketed (it isn't: `bucket_seal` groups by `SUMMARY_FANOUT`, so -//! `level` is a depth integer, not hour/day). Two passes, per source tree: -//! -//! 1. **Eligible summaries** — every non-deleted summary whose time-range -//! envelope is fully inside the window -//! (`store::list_summaries_in_window`). Because seal sets the envelope to -//! `MIN/MAX` of children, "envelope ⊆ window" ⇔ "all descendant leaves are -//! in the window", i.e. using the summary drags in no out-of-window -//! content. -//! 2. **Frontier + raw fallback** — keep the topmost eligible summaries -//! (`maximal` = eligible whose `parent_id` is not itself eligible); mark -//! the chunks they cover (transitive descent through `child_ids`); emit any -//! remaining in-window chunk raw. Raw chunks are the floor and cover -//! boundary slices and not-yet-sealed content for free. -//! -//! Output is grouped by source and ordered ascending by start time so the -//! brief reads it as a per-source timeline. - -use std::collections::{HashMap, HashSet}; - use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::{list_chunks, ListChunksQuery}; -use crate::openhuman::memory_store::chunks::types::{Chunk, SourceKind}; -use crate::openhuman::memory_store::content::read as content_read; -use crate::openhuman::memory_store::trees::types::{SummaryNode, TreeKind}; -use crate::openhuman::memory_tree::retrieval::types::{ - hit_from_chunk, hit_from_summary, QueryResponse, RetrievalHit, -}; -use crate::openhuman::memory_tree::tree::store; -use crate::openhuman::tinycortex::chunk_tree_scope; +use crate::openhuman::memory::source_scope::current_source_scope; +use crate::openhuman::memory_store::chunks::types::SourceKind; +use crate::openhuman::memory_tree::retrieval::engine::config as engine_config; +use crate::openhuman::memory_tree::retrieval::types::QueryResponse; -/// Default cap on returned cover items when the caller passes `limit = 0`. const DEFAULT_LIMIT: usize = 200; -/// Upper bound on in-window chunks scanned across **all** sources. A 24h window -/// rarely exceeds this; if it does we log and truncate (the excess simply -/// doesn't appear — never silently mis-covered, since a frontier summary still -/// stands in for a whole sealed subtree). -const MAX_WINDOW_CHUNKS: usize = 5_000; - -/// Per-source cap on raw in-window chunks retained for the cover. Applied -/// **after** grouping so a single high-volume source can't crowd every other -/// source out of the result; excess is logged, never silently mis-covered. -const MAX_CHUNKS_PER_SOURCE: usize = 2_000; - -/// Entrypoint for `memory_tree_cover_window`. Blocking SQLite work runs on -/// `spawn_blocking`; the async caller stays on its runtime. Results are -/// grouped by source (`tree_scope`) and ordered ascending by start time, then -/// truncated to `limit` (`DEFAULT_LIMIT` when 0). pub async fn cover_window( config: &Config, since_ms: i64, @@ -68,525 +17,26 @@ pub async fn cover_window( limit: usize, ) -> Result { let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; - if until_ms < since_ms { - return Err(anyhow::anyhow!( - "cover_window: until_ms ({until_ms}) precedes since_ms ({since_ms})" - )); + let scope = current_source_scope(); + if source_id.is_some_and(|id| scope.as_ref().is_some_and(|set| !set.contains(id))) { + return Ok(QueryResponse::empty()); } log::debug!( - "[retrieval::cover] cover_window since_ms={since_ms} until_ms={until_ms} \ - has_source_id={} source_kind={:?} limit={limit}", + "[retrieval::cover] tinycortex has_source_id={} source_kind={:?} limit={}", source_id.is_some(), source_kind.map(|k| k.as_str()), + limit ); - - let config_owned = config.clone(); - let source_id_owned = source_id.map(|s| s.to_string()); - // Capture the per-profile memory-source allowlist HERE, in the async task - // that holds the `source_scope` task-local — `spawn_blocking` below runs on - // a separate thread that does NOT inherit task-locals, so we must thread it - // through explicitly or a restricted-profile brief would see every tagged - // source. `None` = unrestricted. - let source_scope = crate::openhuman::memory::source_scope::current_source_scope(); - let mut hits = tokio::task::spawn_blocking(move || -> Result> { - collect_cover( - &config_owned, - since_ms, - until_ms, - source_id_owned.as_deref(), - source_kind, - source_scope, - ) - }) - .await - .map_err(|e| anyhow::anyhow!("cover_window join error: {e}"))??; - - // Group by source, then chronological ascending within each source. The - // brief consumes this as a per-source timeline; ascending puts the - // freshest content last (tool result lands at the context tail → recency). - hits.sort_by(|a, b| { - a.tree_scope - .cmp(&b.tree_scope) - .then(a.time_range_start.cmp(&b.time_range_start)) - }); - let total = hits.len(); - hits.truncate(limit); - - log::debug!( - "[retrieval::cover] returning hits={} total={}", - hits.len(), - total - ); - Ok(QueryResponse::new(hits, total)) -} - -/// Blocking: build the cover. **Chunk-driven, not tree-driven** — chunks are -/// written to `mem_tree_chunks` at ingest, but the per-source `Tree` row is -/// only created later by the seal worker. Iterating trees would therefore miss -/// freshly-ingested (un-sealed) sources entirely — exactly the case the raw -/// fallback exists for. So we pull the authoritative in-window chunk set first, -/// group by source, and look up each source's tree (if any) for summaries. -fn collect_cover( - config: &Config, - since_ms: i64, - until_ms: i64, - source_id: Option<&str>, - source_kind: Option, - source_scope: Option>, -) -> Result> { - let chunks = list_chunks( - config, - &ListChunksQuery { - source_id: source_id.map(|s| s.to_string()), - source_kind, - since_ms: Some(since_ms), - until_ms: Some(until_ms), - limit: Some(MAX_WINDOW_CHUNKS), - // Skip rows the admission gate rejected: they linger in - // `mem_tree_chunks` but were deliberately never appended to a tree, - // so surfacing them raw would leak filtered-out junk into the brief. - exclude_dropped: true, - // Per-profile memory-source allowlist (threaded from the async task - // — task-locals don't cross `spawn_blocking`). - source_scope, - ..Default::default() - }, + let mut response = tinycortex::memory::retrieval::cover_window_scoped( + &engine_config(config), + since_ms, + until_ms, + source_id, + source_kind, + scope, + usize::MAX, )?; - if chunks.len() == MAX_WINDOW_CHUNKS { - log::warn!( - "[retrieval::cover] global in-window chunk cap {MAX_WINDOW_CHUNKS} hit — \ - some raw leaves may be omitted" - ); - } - - // Group by **tree scope**, not raw `source_id`: shared-directory sources - // (Notion `path_scope`) and GitHub per-item ids seal under a derived scope, - // so grouping by `source_id` would miss their tree and emit everything raw. - // Use the same derivation as the append path. A per-source cap keeps one - // high-volume source from crowding out the rest. - let mut by_source: HashMap> = HashMap::new(); - let mut capped_sources = 0usize; - let mut capped_chunks = 0usize; - for chunk in chunks { - let scope = chunk_tree_scope(&chunk.metadata); - let bucket = by_source.entry(scope).or_default(); - if bucket.len() < MAX_CHUNKS_PER_SOURCE { - bucket.push(chunk); - } else { - if bucket.len() == MAX_CHUNKS_PER_SOURCE { - capped_sources += 1; - } - capped_chunks += 1; - } - } - if capped_chunks > 0 { - // Omit scope (PII) — aggregate counts only. - log::warn!( - "[retrieval::cover] {capped_sources} source(s) hit per-source cap \ - {MAX_CHUNKS_PER_SOURCE} — {capped_chunks} raw leaf/leaves omitted" - ); - } - log::debug!("[retrieval::cover] in-window sources n={}", by_source.len()); - - // 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. Without a filter every in-window leaf is present. - let exact_source = source_id.is_some(); - let mut hits: Vec = Vec::new(); - for (source, src_chunks) in by_source { - cover_one_source( - config, - &source, - since_ms, - until_ms, - src_chunks, - exact_source, - &mut hits, - )?; - } - Ok(hits) -} - -/// Minimum cover for one source: frontier summaries (when the source has a -/// sealed tree) plus every in-window chunk they don't already cover, raw. A -/// source with no `Tree` row (not yet processed by the seal worker) has no -/// eligible summaries, so all its in-window chunks are emitted raw. -fn cover_one_source( - config: &Config, - source: &str, - since_ms: i64, - until_ms: i64, - chunks: Vec, - exact_source: bool, - out: &mut Vec, -) -> Result<()> { - // Look up the source's summary tree (scope == source id). Absent until the - // seal worker first processes this source. - let tree = store::get_tree_by_scope(config, TreeKind::Source, source)?; - let (tree_id, eligible) = match &tree { - Some(t) => ( - t.id.as_str(), - store::list_summaries_in_window(config, &t.id, since_ms, until_ms)?, - ), - None => ("", Vec::new()), - }; - // Latest-wins for versioned document sources (Notion): drop superseded - // doc-root revisions so the cover never emits a stale page, and remember - // their chunk ids so the raw fallback below doesn't resurface them either. - let (eligible, suppressed_chunk_ids) = filter_superseded_doc_versions(eligible); - // In exact-source mode the tree may be shared across sibling sources, so - // only emit summaries whose whole subtree is among the filtered chunks. - let present: HashSet<&str> = chunks.iter().map(|c| c.id.as_str()).collect(); - let plan = plan_cover(&eligible, exact_source.then_some(&present)); - - // Frontier summaries (hydrated to full body). - let by_id: HashMap<&str, &SummaryNode> = eligible.iter().map(|s| (s.id.as_str(), s)).collect(); - for id in &plan.maximal_ids { - let Some(node) = by_id.get(id.as_str()) else { - continue; - }; - let mut node = (*node).clone(); - match content_read::read_summary_body(config, &node.id) { - Ok(body) => node.content = body, - Err(e) => { - log::warn!("[retrieval::cover] read_summary_body failed — serving preview: {e:#}") - } - } - out.push(hit_from_summary(&node, source)); - } - - // In-window chunks not covered by a frontier summary, raw — skipping any - // chunk under a superseded document revision. - for chunk in &chunks { - if plan.covered_chunk_ids.contains(&chunk.id) || suppressed_chunk_ids.contains(&chunk.id) { - continue; - } - let mut chunk = chunk.clone(); - match content_read::read_chunk_body(config, &chunk.id) { - Ok(body) => chunk.content = body, - Err(e) => { - log::warn!("[retrieval::cover] read_chunk_body failed — serving preview: {e:#}") - } - } - out.push(hit_from_chunk(&chunk, tree_id, source, 0.0)); - } - Ok(()) -} - -/// The structural result of the cover for one tree's eligible summaries: -/// which summaries to emit (the frontier) and which chunk ids they already -/// cover (so the caller can emit the rest raw). Pure — no I/O — so the cover -/// logic is unit-testable without a database. -struct CoverPlan { - /// Topmost eligible summary ids (eligible nodes whose parent is not - /// eligible). These stand in for their whole subtree. - maximal_ids: Vec, - /// Leaf chunk ids transitively covered by the `maximal` summaries. - covered_chunk_ids: HashSet, -} - -/// Compute the frontier + covered-chunk set from a tree's eligible summaries. -/// -/// `eligible` must be exactly the summaries whose envelope is inside the -/// window (all levels). A summary is **maximal** when its `parent_id` is not -/// itself eligible. Coverage descends `child_ids`: an id that is not another -/// eligible summary is a leaf chunk id (because every descendant of a -/// fully-covered node is itself fully covered, hence eligible — so any child -/// not in the eligible set is a leaf). -/// -/// `restrict_to_present` guards the exact-source path. A source tree's scope -/// can be **broader** than the requested `source_id` (Notion `path_scope`, -/// GitHub repo-scoped trees seal many pages/issues into one tree). When the -/// caller filtered chunks down to a single source id, a frontier summary over -/// the shared tree would also cover *sibling* sources' leaves — leaking -/// unrelated memory and masking the requested raw chunks under a mixed -/// summary. When `Some(present)`, a maximal summary is emitted only if **every** -/// chunk it covers is in `present` (the in-filter chunk ids); summaries that -/// span out-of-filter sources are dropped, and their in-filter chunks fall -/// through to raw emission. `None` (the no-filter brief path) keeps every -/// maximal summary — the chunk set already holds every in-window leaf, so no -/// summary can span anything absent. -fn plan_cover(eligible: &[SummaryNode], restrict_to_present: Option<&HashSet<&str>>) -> CoverPlan { - let eligible_ids: HashSet<&str> = eligible.iter().map(|s| s.id.as_str()).collect(); - let by_id: HashMap<&str, &SummaryNode> = eligible.iter().map(|s| (s.id.as_str(), s)).collect(); - - let mut maximal_ids: Vec = Vec::new(); - let mut covered_chunk_ids: HashSet = HashSet::new(); - for node in eligible.iter().filter(|s| match &s.parent_id { - Some(parent) => !eligible_ids.contains(parent.as_str()), - None => true, - }) { - let mut sub: HashSet = HashSet::new(); - collect_descendant_chunks(node, &by_id, &mut sub); - if let Some(present) = restrict_to_present { - if !sub.iter().all(|c| present.contains(c.as_str())) { - // Shared-tree summary spans sources outside the exact-source - // filter — skip it; its in-filter chunks are emitted raw. - continue; - } - } - maximal_ids.push(node.id.clone()); - covered_chunk_ids.extend(sub); - } - - CoverPlan { - maximal_ids, - covered_chunk_ids, - } -} - -/// Walk a summary's subtree (within the eligible set) collecting leaf chunk -/// ids. A child id present in `by_id` is a lower-level summary → recurse; a -/// child id absent from `by_id` is a leaf chunk → record it. -fn collect_descendant_chunks( - node: &SummaryNode, - by_id: &HashMap<&str, &SummaryNode>, - covered: &mut HashSet, -) { - for child in &node.child_ids { - match by_id.get(child.as_str()) { - Some(child_summary) => collect_descendant_chunks(child_summary, by_id, covered), - None => { - covered.insert(child.clone()); - } - } - } -} - -/// Latest-wins for versioned document sources (Notion). A source tree can hold -/// several doc-root summaries for the **same** `doc_id` — each page edit seals -/// a new doc-root at a higher `version_ms` beside the old one, and retrieval is -/// expected to hide the older revisions (`drill_down` does the same). Returns -/// `eligible` with every superseded revision's whole subtree removed, plus the -/// chunk ids under those dropped revisions so the raw fallback can't resurface -/// stale page content. Summaries with no `doc_id` are untouched. -fn filter_superseded_doc_versions( - eligible: Vec, -) -> (Vec, HashSet) { - // No document nodes → nothing to do (the common chat/email case). - if !eligible.iter().any(|s| s.doc_id.is_some()) { - return (eligible, HashSet::new()); - } - - let by_id: HashMap<&str, &SummaryNode> = eligible.iter().map(|s| (s.id.as_str(), s)).collect(); - - // Winning (max) version per doc_id; `version_ms` defaults to i64::MIN so a - // legacy untagged doc-root never wins over a tagged one. - let mut max_version_by_doc: HashMap<&str, i64> = HashMap::new(); - for s in &eligible { - if let Some(doc) = s.doc_id.as_deref() { - let v = s.version_ms.unwrap_or(i64::MIN); - max_version_by_doc - .entry(doc) - .and_modify(|m| { - if v > *m { - *m = v; - } - }) - .or_insert(v); - } - } - - // A doc-root is a "loser" when it's an older revision, or a duplicate of the - // winning version (e.g. a retried SealDocument minted two) — keep the first - // winner only. `eligible` is ordered (level, start), so dedup is stable. - let mut winners_seen: HashSet<&str> = HashSet::new(); - let mut removed_summary_ids: HashSet = HashSet::new(); - let mut suppressed_chunk_ids: HashSet = HashSet::new(); - for s in &eligible { - let Some(doc) = s.doc_id.as_deref() else { - continue; - }; - let v = s.version_ms.unwrap_or(i64::MIN); - let max = max_version_by_doc.get(doc).copied().unwrap_or(i64::MIN); - // Short-circuit keeps the winner slot untouched for older revisions. - let loser = v < max || !winners_seen.insert(doc); - if loser { - removed_summary_ids.insert(s.id.clone()); - collect_subtree_ids( - s, - &by_id, - &mut removed_summary_ids, - &mut suppressed_chunk_ids, - ); - } - } - - let kept = eligible - .into_iter() - .filter(|s| !removed_summary_ids.contains(&s.id)) - .collect(); - (kept, suppressed_chunk_ids) -} - -/// Walk a summary's subtree collecting both descendant summary ids (into -/// `summaries`) and leaf chunk ids (into `chunks`). Used to evict a superseded -/// document revision's whole subtree from the cover. -fn collect_subtree_ids( - node: &SummaryNode, - by_id: &HashMap<&str, &SummaryNode>, - summaries: &mut HashSet, - chunks: &mut HashSet, -) { - for child in &node.child_ids { - match by_id.get(child.as_str()) { - Some(child_summary) => { - summaries.insert(child.clone()); - collect_subtree_ids(child_summary, by_id, summaries, chunks); - } - None => { - chunks.insert(child.clone()); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::trees::types::TreeKind; - use chrono::Utc; - - fn summary(id: &str, parent: Option<&str>, level: u32, children: &[&str]) -> SummaryNode { - SummaryNode { - id: id.to_string(), - tree_id: "t1".to_string(), - tree_kind: TreeKind::Source, - level, - parent_id: parent.map(|p| p.to_string()), - child_ids: children.iter().map(|c| c.to_string()).collect(), - content: format!("summary {id}"), - token_count: 10, - entities: vec![], - topics: vec![], - time_range_start: Utc::now(), - time_range_end: Utc::now(), - score: 0.0, - sealed_at: Utc::now(), - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - } - } - - #[test] - fn single_eligible_summary_covers_its_leaves() { - // L1 node over two leaf chunks, no parent → it's the frontier and - // covers both leaves. - let eligible = vec![summary("s1", None, 1, &["chunk-a", "chunk-b"])]; - let plan = plan_cover(&eligible, None); - assert_eq!(plan.maximal_ids, vec!["s1"]); - assert!(plan.covered_chunk_ids.contains("chunk-a")); - assert!(plan.covered_chunk_ids.contains("chunk-b")); - assert_eq!(plan.covered_chunk_ids.len(), 2); - } - - #[test] - fn parent_subsumes_child_only_parent_is_maximal() { - // s2 (L2) → s1 (L1) → leaves. Both eligible: only s2 is maximal, and - // it transitively covers the leaves under s1. - let eligible = vec![ - summary("s2", None, 2, &["s1"]), - summary("s1", Some("s2"), 1, &["chunk-a", "chunk-b"]), - ]; - let plan = plan_cover(&eligible, None); - assert_eq!(plan.maximal_ids, vec!["s2"]); - assert!(plan.covered_chunk_ids.contains("chunk-a")); - assert!(plan.covered_chunk_ids.contains("chunk-b")); - } - - #[test] - fn ineligible_parent_leaves_child_as_frontier() { - // s1 is eligible but its parent s2 is NOT in the eligible set (it - // straddles the window / not sealed) → s1 is maximal. Leaves under s1 - // are covered; nothing else. - let eligible = vec![summary("s1", Some("s2-not-eligible"), 1, &["chunk-a"])]; - let plan = plan_cover(&eligible, None); - assert_eq!(plan.maximal_ids, vec!["s1"]); - assert!(plan.covered_chunk_ids.contains("chunk-a")); - } - - #[test] - fn empty_eligible_set_covers_nothing() { - // No sealed/eligible summaries → frontier empty, no chunk covered, so - // the caller emits every in-window chunk raw. - let plan = plan_cover(&[], None); - assert!(plan.maximal_ids.is_empty()); - assert!(plan.covered_chunk_ids.is_empty()); - } - - #[test] - fn sibling_frontier_nodes_each_emitted() { - // Two L1 siblings, both eligible, parent NOT eligible → both maximal. - let eligible = vec![ - summary("s1", Some("root-x"), 1, &["chunk-a"]), - summary("s2", Some("root-x"), 1, &["chunk-b"]), - ]; - let mut plan = plan_cover(&eligible, None); - plan.maximal_ids.sort(); - assert_eq!(plan.maximal_ids, vec!["s1", "s2"]); - assert_eq!(plan.covered_chunk_ids.len(), 2); - } - - fn doc_summary(id: &str, doc_id: &str, version_ms: i64, children: &[&str]) -> SummaryNode { - let mut s = summary(id, Some("merge-root"), 1, children); - s.doc_id = Some(doc_id.to_string()); - s.version_ms = Some(version_ms); - s - } - - #[test] - fn filter_superseded_doc_versions_keeps_newest_and_suppresses_old_chunks() { - // Two revisions of the same Notion page plus an unrelated chat summary. - // Only the newest revision survives; the older revision's subtree is - // dropped and its chunk reported for raw-fallback suppression. - let eligible = vec![ - doc_summary("pageA@v1", "notion:pageA", 100, &["chunk-old"]), - doc_summary("pageA@v2", "notion:pageA", 200, &["chunk-new"]), - summary("chat", Some("root"), 1, &["chunk-chat"]), - ]; - let (kept, suppressed) = filter_superseded_doc_versions(eligible); - let kept_ids: Vec<&str> = kept.iter().map(|s| s.id.as_str()).collect(); - assert_eq!(kept_ids, vec!["pageA@v2", "chat"]); - assert!(suppressed.contains("chunk-old")); - assert!(!suppressed.contains("chunk-new")); - assert!(!suppressed.contains("chunk-chat")); - } - - #[test] - fn filter_superseded_doc_versions_dedups_duplicate_winning_revision() { - // A retried seal can mint two doc-roots at the SAME winning version; - // keep the first, drop the duplicate's subtree. - let eligible = vec![ - doc_summary("dup-a", "notion:pageB", 300, &["chunk-a"]), - doc_summary("dup-b", "notion:pageB", 300, &["chunk-b"]), - ]; - let (kept, suppressed) = filter_superseded_doc_versions(eligible); - let kept_ids: Vec<&str> = kept.iter().map(|s| s.id.as_str()).collect(); - assert_eq!(kept_ids, vec!["dup-a"]); - assert!(suppressed.contains("chunk-b")); - assert!(!suppressed.contains("chunk-a")); - } - - #[test] - fn restrict_drops_summaries_spanning_out_of_filter_chunks() { - // Exact-source mode over a SHARED tree: `s_mixed` summarises a leaf the - // filter kept (`chunk-a`) plus a sibling page's leaf (`chunk-foreign`) - // that isn't in the requested set; `s_clean` covers only kept leaves. - // With the present-set restriction, the mixed summary must be dropped - // (so `chunk-a` falls through to raw) while the clean one survives. - let eligible = vec![ - summary("s_mixed", Some("root"), 1, &["chunk-a", "chunk-foreign"]), - summary("s_clean", Some("root"), 1, &["chunk-b"]), - ]; - let present: HashSet<&str> = ["chunk-a", "chunk-b"].into_iter().collect(); - let plan = plan_cover(&eligible, Some(&present)); - assert_eq!(plan.maximal_ids, vec!["s_clean"]); - assert!(plan.covered_chunk_ids.contains("chunk-b")); - // chunk-a is NOT covered → caller emits it raw rather than via the - // sibling-spanning summary. - assert!(!plan.covered_chunk_ids.contains("chunk-a")); - assert!(!plan.covered_chunk_ids.contains("chunk-foreign")); - } + let total = response.hits.len(); + response.hits.truncate(limit); + Ok(QueryResponse::new(response.hits, total)) } diff --git a/src/openhuman/memory_tree/retrieval/drill_down.rs b/src/openhuman/memory_tree/retrieval/drill_down.rs index 90034fd95..2e482cd05 100644 --- a/src/openhuman/memory_tree/retrieval/drill_down.rs +++ b/src/openhuman/memory_tree/retrieval/drill_down.rs @@ -1,52 +1,11 @@ -//! `memory_tree_drill_down` — walk `child_ids` from a summary node (Phase 4 -//! / #710). -//! -//! Primary use case: the LLM gets a summary hit back from `query_source` or -//! `query_topic` and wants to look at the next level down — either more -//! summaries (for L2+ nodes) or the raw chunks (for L1 nodes). This is -//! deliberately a one-step expansion; for multi-step walks the caller -//! passes `max_depth > 1`. -//! -//! When `query` is `Some`, visited children are reranked by cosine similarity -//! against the query embedding so a deep summary with many children can surface -//! the relevant ones to the top. When `query` is `None`, children are returned -//! in BFS order (same as before). -//! -//! Behaviour: -//! - Unknown `node_id` → empty vec (not an error — the LLM can recover). -//! - `max_depth == 0` → empty vec (documented as "no-op"). -//! - Leaves have no children; drilling into a leaf id returns empty. -//! - `limit` is optional; when set, it truncates the final (reranked) output. - -use std::collections::{HashMap, HashSet}; - use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::{ - get_chunk, get_chunk_embeddings_batch, get_chunks_batch, -}; -use crate::openhuman::memory_store::content::read as content_read; -use crate::openhuman::memory_store::trees::store::{get_summaries_batch, get_trees_batch}; -use crate::openhuman::memory_tree::retrieval::types::{ - hit_from_chunk, hit_from_summary, RetrievalHit, -}; -use crate::openhuman::memory_tree::score::embed::{build_embedder_from_config, cosine_similarity}; -use crate::openhuman::memory_tree::tree::store; +use crate::openhuman::memory::source_scope::current_source_scope; +use crate::openhuman::memory_tree::retrieval::engine::{config as engine_config, EmbedderBridge}; +use crate::openhuman::memory_tree::retrieval::types::RetrievalHit; +use crate::openhuman::memory_tree::score::embed::{build_embedder_from_config, InertEmbedder}; -/// Upper-bound estimate of how many children a summary node fans out to, -/// used only to pre-size the next-level BFS frontier. Over-estimating wastes -/// a little transient capacity; under-estimating costs a realloc — neither is -/// load-bearing for correctness. -const EXPECTED_CHILD_FANOUT: usize = 10; - -/// Walk the summary hierarchy down one step (or more if `max_depth > 1`) -/// and return the hydrated child hits. Children at level 1 are raw chunks; -/// deeper children are summaries. -/// -/// When `query` is `Some`, the returned hits are reranked by cosine similarity -/// to the query embedding; hits without a stored embedding (legacy rows) sort -/// to the bottom. When `None`, BFS order is preserved. pub async fn drill_down( config: &Config, node_id: &str, @@ -54,932 +13,38 @@ pub async fn drill_down( query: Option<&str>, limit: Option, ) -> Result> { - // Redact `node_id` — embeds tree scope (e.g. `summary:L1:` or - // `chat:slack:#:`) which can carry workspace hints. Log - // the id's structural prefix only. - let node_kind_prefix = node_id.split_once(':').map(|(k, _)| k).unwrap_or("unknown"); log::debug!( - "[retrieval::drill_down] drill_down node_kind={} max_depth={} has_query={} limit={:?}", - node_kind_prefix, + "[retrieval::drill_down] tinycortex max_depth={} has_query={} limit={:?}", max_depth, query.is_some(), limit ); - if max_depth == 0 { - log::debug!("[retrieval::drill_down] max_depth=0 — returning empty vec"); - return Ok(Vec::new()); - } - - // Phase 1 — blocking walk produces hits + the per-hit embedding so the - // async rerank pass can avoid a second trip through the DB. - let node_id_owned = node_id.to_string(); - let config_owned = config.clone(); - let (hits, embeddings) = tokio::task::spawn_blocking( - move || -> Result<(Vec, Vec>>)> { - walk_with_embeddings(&config_owned, &node_id_owned, max_depth) - }, - ) - .await - .map_err(|e| anyhow::anyhow!("drill_down join error: {e}"))??; - - // Phase 2 — optional query rerank. - let hits = if let Some(q) = query { - rerank_by_semantic_similarity(config, q, hits, embeddings).await? + let embedder = if query.is_none() || max_depth == 0 { + log::debug!("[retrieval::drill_down] using inert embedder for non-semantic traversal"); + Box::new(InertEmbedder::new()) + as Box } else { - hits + build_embedder_from_config(config)? }; - - // Phase 3 — apply optional limit AFTER rerank so the top-K is relevance- - // based when `query` is Some, BFS-based otherwise. - let hits = match limit { - Some(n) if hits.len() > n => hits.into_iter().take(n).collect(), - _ => hits, - }; - - log::debug!("[retrieval::drill_down] returning hits={}", hits.len()); + let bridge = EmbedderBridge(embedder.as_ref()); + let engine_limit = current_source_scope() + .as_ref() + .map(|_| None) + .unwrap_or(limit); + let mut hits = tinycortex::memory::retrieval::drill_down( + &engine_config(config), + node_id, + max_depth, + query, + &bridge, + engine_limit, + ) + .await?; + if let Some(set) = current_source_scope() { + hits.retain(|hit| set.contains(&hit.tree_scope)); + } + if let Some(limit) = limit { + hits.truncate(limit); + } Ok(hits) } - -/// Rerank hits by cosine similarity to the query embedding. Mirrors the -/// pattern used by `query_source` / `query_topic`. Legacy rows without -/// embeddings land at the end in BFS order. -/// -/// On any error (embedder build failure or embedding inference failure) we log -/// a warning and return hits in BFS order rather than bubbling the error up -/// through the chat turn. This ensures local AI unavailability never surfaces -/// as a visible error to the user. -async fn rerank_by_semantic_similarity( - config: &Config, - query: &str, - hits: Vec, - embeddings: Vec>>, -) -> Result> { - debug_assert_eq!(hits.len(), embeddings.len()); - let embedder = match build_embedder_from_config(config) { - Ok(e) => e, - Err(err) => { - log::warn!( - "[retrieval::drill_down] embedder build failed — returning BFS order: {err}" - ); - return Ok(hits); - } - }; - let query_vec = match embedder.embed(query).await { - Ok(v) => v, - Err(err) => { - log::warn!("[retrieval::drill_down] embed query failed — returning BFS order: {err}"); - return Ok(hits); - } - }; - log::debug!( - "[retrieval::drill_down] query embedded provider={} hits_to_rerank={}", - embedder.name(), - hits.len() - ); - - let mut decorated: Vec<(f32, bool, RetrievalHit)> = hits - .into_iter() - .zip(embeddings) - .map(|(h, emb)| match emb { - Some(v) if v.len() == query_vec.len() => { - let sim = cosine_similarity(&query_vec, &v); - (sim, true, h) - } - _ => (f32::NEG_INFINITY, false, h), - }) - .collect(); - - decorated.sort_by(|a, b| match (a.1, b.1) { - (true, false) => std::cmp::Ordering::Less, - (false, true) => std::cmp::Ordering::Greater, - // Both ranked (or both unranked): similarity DESC, then by time. - _ => { - b.0.partial_cmp(&a.0) - .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| b.2.time_range_end.cmp(&a.2.time_range_end)) - } - }); - - Ok(decorated.into_iter().map(|(_, _, h)| h).collect()) -} - -/// Blocking walker. BFS-style expansion up to `max_depth` levels. Returns -/// each hit paired with its stored embedding (if any), so the async rerank -/// pass doesn't have to round-trip through the DB again. -/// -/// **Batched per BFS depth.** For each level we issue at most four SQLite -/// round-trips (one each for summaries / trees / chunks / chunk -/// embeddings) covering every node at that depth, then walk the level's -/// id slice in BFS order to populate `out` + collect the next-depth -/// frontier. The per-node `get_summary` / `get_tree` / `get_chunk` / -/// `get_chunk_embedding` loop (one round-trip per node × 4 tables) is -/// replaced by `O(depth)` round-trips instead of `O(nodes × 4)`. File -/// I/O via `read_summary_body` / `read_chunk_body` stays per-id — each -/// body lives in its own on-disk file, so batching would mean concurrent -/// opens, not a single round-trip; left untouched. -fn walk_with_embeddings( - config: &Config, - start_id: &str, - max_depth: u32, -) -> Result<(Vec, Vec>>)> { - // Fetch the root. If it's a summary we expand its child_ids; if it's a - // chunk it has no children. If it's neither we return empty. - let root_summary = store::get_summary(config, start_id)?; - let root_tree_scope = match root_summary.as_ref().map(|s| s.tree_id.clone()) { - Some(tid) => store::get_tree(config, &tid)? - .map(|t| t.scope) - .unwrap_or_default(), - None => String::new(), - }; - - let mut out: Vec = Vec::new(); - let mut embeddings: Vec>> = Vec::new(); - - let start_children: Vec = match root_summary { - Some(s) => s.child_ids.clone(), - None => { - if let Some(_c) = get_chunk(config, start_id)? { - return Ok((out, embeddings)); - } - log::debug!( - "[retrieval::drill_down] node_id={start_id} not found in summaries or chunks" - ); - return Ok((out, embeddings)); - } - }; - - // BFS-by-level. `current_level` holds every node id at the current - // depth, walked in FIFO (BFS) order — siblings are always returned - // before any descendant at a deeper depth (regression for PR #831's - // CodeRabbit fix away from `Vec::pop` DFS). Batched fetches for the - // whole level happen up-front so the per-id walk only does HashMap - // lookups + the unavoidable per-file body read. - let mut current_level: Vec = start_children; - let mut depth: u32 = 1; - - // Latest-version-per-document filter (document source trees, e.g. Notion). - // A document's chunks roll up to a per-doc subtree whose root carries - // `(doc_id, version_ms)`; editing a page seals a NEW doc-root (higher - // `version_ms`) alongside the old one, so the merge tier reaches both. We - // surface only the newest revision: as the walk encounters doc-roots we - // track `max(version_ms)` per `doc_id` and skip any doc-root below that - // max (and therefore its whole stale subtree). Nothing is mutated on disk - // — superseded revisions simply never appear in results. Non-document - // nodes (doc_id == None) are unaffected. - let mut max_version_by_doc: HashMap = HashMap::new(); - // Doc-roots already surfaced, to dedup at the winning version: if a - // `SealDocument` job partially committed then retried, it can mint a second - // doc-root for the SAME `(doc_id, version_ms)`. Emit only the first one per - // doc_id so a duplicate revision never double-surfaces. - let mut emitted_docs: HashSet = HashSet::new(); - - while !current_level.is_empty() && depth <= max_depth { - log::trace!( - "[retrieval::drill_down] level depth={} ids={}", - depth, - current_level.len() - ); - - // 1) Batched summary fetch covers every id on this level. Missing - // ids stay silently absent from the map (same `Ok(None)` - // contract as the per-row `get_summary`); those ids are then - // tried as chunks below. - let mut summary_by_id = get_summaries_batch(config, ¤t_level)?; - - // Update the per-document latest-version map with any doc-roots on - // THIS level before walking it, so two revisions of the same document - // sitting side-by-side (the common case — both are merge-tier leaves - // at the same depth) resolve to the newer one regardless of walk - // order. A doc-root is a summary with `doc_id` set; `version_ms` - // defaults to i64::MIN so a (legacy) untagged doc-root never wins over - // a tagged one. - for id in ¤t_level { - if let Some(s) = summary_by_id.get(id) { - if let Some(doc_id) = s.doc_id.as_deref() { - let v = s.version_ms.unwrap_or(i64::MIN); - max_version_by_doc - .entry(doc_id.to_string()) - .and_modify(|cur| { - if v > *cur { - *cur = v; - } - }) - .or_insert(v); - } - } - } - - // 2) Distinct tree_ids referenced by this level's summaries — - // dedup is purely to avoid redundant DB params (the per-id - // walk below routes each summary to its own scope via the - // map). Insertion-order preserving for deterministic logs. - let distinct_tree_ids: Vec = { - // `seen` borrows the tree_id slices straight out of - // `summary_by_id` (which outlives this block) — dedup costs no - // allocation; only the surviving distinct ids are cloned into - // `out` for `get_trees_batch`. - let mut seen: HashSet<&str> = HashSet::new(); - let mut out: Vec = Vec::new(); - for id in ¤t_level { - if let Some(s) = summary_by_id.get(id) { - if seen.insert(s.tree_id.as_str()) { - out.push(s.tree_id.clone()); - } - } - } - out - }; - let tree_by_id = get_trees_batch(config, &distinct_tree_ids)?; - - // 3) Ids on this level that AREN'T summaries are candidate - // chunk leaves; batch-fetch both the chunk rows and their - // embeddings. Missing ids are silently absent — the warn - // path at the end of the per-id walk catches "points at - // nothing" cases (preserving the existing contract). - let chunk_ids: Vec = current_level - .iter() - .filter(|id| !summary_by_id.contains_key(*id)) - .cloned() - .collect(); - let mut chunk_by_id = get_chunks_batch(config, &chunk_ids)?; - // `get_chunk_embeddings_batch` returns only present ids - // (mirroring per-row `get_chunk_embedding` returning - // `Ok(None)` for legacy rows without an embedding row); - // `.get(id).cloned()` yields the equivalent `Option>`. - let emb_by_id = get_chunk_embeddings_batch(config, &chunk_ids)?; - - // 4) Walk this level in BFS order, populate hits, collect next - // level. Per-id HashMap lookups (keyed by id, not by - // enumerate() position over the input slice — otherwise a - // sibling could shadow another's scope or chunk body). - // Pre-size against expected fan-out so the per-level child - // accumulation avoids repeated reallocs. Only the non-final depths - // extend `next_level` (see the `depth < max_depth` guard below), so - // skip the reservation at the last depth where it would stay empty. - let mut next_level: Vec = if depth < max_depth { - Vec::with_capacity(current_level.len() * EXPECTED_CHILD_FANOUT) - } else { - Vec::new() - }; - for id in ¤t_level { - if let Some(mut summary) = summary_by_id.remove(id) { - // Latest-wins: skip a doc-root that a newer revision of the - // same document supersedes. Its subtree is not expanded, so - // the stale revision's chunks never surface. - if let Some(doc_id) = summary.doc_id.as_deref() { - let v = summary.version_ms.unwrap_or(i64::MIN); - if max_version_by_doc.get(doc_id).is_some_and(|&max| v < max) { - log::debug!( - "[retrieval::drill_down] skipping superseded doc-root \ - doc_id={doc_id} version_ms={v} (latest is newer)" - ); - continue; - } - // Dedup duplicates at the winning version (e.g. a retried - // SealDocument that minted a second doc-root for the same - // (doc_id, version_ms)) — surface only the first. - if !emitted_docs.insert(doc_id.to_string()) { - log::debug!( - "[retrieval::drill_down] skipping duplicate doc-root \ - doc_id={doc_id} version_ms={v} (already surfaced)" - ); - continue; - } - } - let scope = tree_by_id - .get(&summary.tree_id) - .map(|t| t.scope.clone()) - .unwrap_or_else(|| root_tree_scope.clone()); - // Hydrate the full body from disk — `summary.content` is - // a ≤500-char preview after the MD-on-disk migration. - // Non-fatal fallback for pre-MD-migration rows. - match content_read::read_summary_body(config, id) { - Ok(body) => summary.content = body, - Err(e) => { - log::warn!( - "[retrieval::drill_down] read_summary_body failed — serving preview: {e:#}" - ); - } - } - // Summary embeddings live on the struct directly - // (Phase 4 amend). - embeddings.push(summary.embedding.clone()); - let child_ids = summary.child_ids.clone(); - out.push(hit_from_summary(&summary, &scope)); - if depth < max_depth { - next_level.extend(child_ids); - } - continue; - } - if let Some(mut chunk) = chunk_by_id.remove(id) { - // Missing embedding → None (legacy row); identical to - // the per-row `get_chunk_embedding(...) Ok(None)` arm. - let emb = emb_by_id.get(id).cloned(); - embeddings.push(emb); - // Hydrate the full body from disk — `chunk.content` is - // a ≤500-char preview after the MD-on-disk migration. - match content_read::read_chunk_body(config, id) { - Ok(body) => chunk.content = body, - Err(e) => { - log::warn!( - "[retrieval::drill_down] read_chunk_body failed — serving preview: {e:#}" - ); - } - } - // Score unknown here; 0.0 neutral placeholder. - out.push(hit_from_chunk(&chunk, "", &chunk.metadata.source_id, 0.0)); - continue; - } - // Redact the child id — may contain source scope (e.g. - // `chat:slack:#:seq`). Log the kind prefix only. - let kind_prefix = id.split_once(':').map(|(k, _)| k).unwrap_or("unknown"); - log::warn!( - "[retrieval::drill_down] child kind={kind_prefix} points at nothing — skipping" - ); - } - - current_level = next_level; - depth += 1; - } - Ok((out, embeddings)) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider}; - use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree; - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::content as content_store; - use crate::openhuman::memory_store::trees::types::TreeKind; - use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf, LabelStrategy, LeafRef}; - use chrono::Utc; - use std::sync::Arc; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): seeding requires seals which embed. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - async fn seed_sealed_tree(cfg: &Config) -> (String, String) { - // Seed two 6k-token leaves so the L0 buffer seals into an L1 node. - let ts = Utc::now(); - let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap(); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).unwrap(); - let mut leaf_ids: Vec = Vec::new(); - for seq in 0..2u32 { - let c = Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "test-content"), - content: format!("content-{seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: crate::openhuman::memory_store::trees::types::INPUT_TOKEN_BUDGET * 6 - / 10, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); - // Stage to disk so `hydrate_leaf_inputs` can read the full body - // via `read_chunk_body` during the seal triggered by `append_leaf`. - let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap(); - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx( - &tx, &staged, - )?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - leaf_ids.push(c.id.clone()); - let leaf = LeafRef { - chunk_id: c.id.clone(), - token_count: crate::openhuman::memory_store::trees::types::INPUT_TOKEN_BUDGET * 6 - / 10, - timestamp: ts, - content: c.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - test_override::with_provider(Arc::clone(&provider), async { - append_leaf(cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - } - // Fetch the sealed L1 summary id from the tree row. - let refreshed = store::get_tree(cfg, &tree.id).unwrap().unwrap(); - assert_eq!(refreshed.kind, TreeKind::Source); - let root_id = refreshed.root_id.unwrap(); - (root_id, leaf_ids.remove(0)) - } - - #[tokio::test] - async fn depth_zero_returns_empty() { - let (_tmp, cfg) = test_config(); - let (root_id, _) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &root_id, 0, None, None).await.unwrap(); - assert!(out.is_empty()); - } - - #[tokio::test] - async fn invalid_id_returns_empty() { - let (_tmp, cfg) = test_config(); - let out = drill_down(&cfg, "nonexistent:id", 1, None, None) - .await - .unwrap(); - assert!(out.is_empty()); - } - - /// Read-time latest-wins: a merge root referencing two per-doc roots of - /// the SAME document (v1 < v2) must surface only the newer revision's - /// subtree; the superseded doc-root and its chunk are filtered out and - /// never traversed — without anything being deleted on disk. - #[tokio::test] - async fn drill_down_surfaces_only_latest_doc_version() { - use crate::openhuman::memory_store::chunks::store::{upsert_chunks, with_connection}; - use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeStatus}; - use crate::openhuman::memory_tree::tree::store as tree_store; - - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - - let tree = Tree { - id: "test:notion-tree".into(), - kind: TreeKind::Source, - scope: "notion:conn1".into(), - root_id: Some("s:merge:root".into()), - max_level: 1000, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }; - - let mk_chunk = |content: &str| Chunk { - id: chunk_id(SourceKind::Document, "notion:conn1:pageA", 0, content), - content: content.to_string(), - metadata: Metadata { - source_kind: SourceKind::Document, - source_id: "notion:conn1:pageA".into(), - owner: "notion:conn1".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["notion".into()], - source_ref: Some(SourceRef::new("notion://page/pageA")), - path_scope: Some("notion:conn1".into()), - }, - token_count: 10, - seq_in_source: 0, - created_at: ts, - partial_message: false, - }; - // Distinct content → distinct chunk ids (content is hashed in). - let chunk_v1 = mk_chunk("old version body"); - let chunk_v2 = mk_chunk("new version body"); - upsert_chunks(&cfg, &[chunk_v1.clone(), chunk_v2.clone()]).unwrap(); - - let mk_root = |id: &str, version: i64, child: &str| SummaryNode { - id: id.into(), - tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: Some("s:merge:root".into()), - child_ids: vec![child.to_string()], - content: format!("doc-root v{version}"), - token_count: 5, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: Some("notion:conn1:pageA".into()), - version_ms: Some(version), - }; - let v1_root = mk_root("s:docA:v1", 100, &chunk_v1.id); - let v2_root = mk_root("s:docA:v2", 200, &chunk_v2.id); - - let merge_root = SummaryNode { - id: "s:merge:root".into(), - tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, - level: 1000, - parent_id: None, - child_ids: vec![v1_root.id.clone(), v2_root.id.clone()], - content: "merge root".into(), - token_count: 5, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - tree_store::insert_tree_conn(&tx, &tree)?; - tree_store::insert_summary_tx(&tx, &v1_root, None, "test")?; - tree_store::insert_summary_tx(&tx, &v2_root, None, "test")?; - tree_store::insert_summary_tx(&tx, &merge_root, None, "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let out = drill_down(&cfg, "s:merge:root", 3, None, None) - .await - .unwrap(); - let ids: Vec<&str> = out.iter().map(|h| h.node_id.as_str()).collect(); - - assert!( - ids.contains(&"s:docA:v2"), - "latest doc-root must surface; got {ids:?}" - ); - assert!( - ids.contains(&chunk_v2.id.as_str()), - "latest version's chunk must surface; got {ids:?}" - ); - assert!( - !ids.contains(&"s:docA:v1"), - "superseded doc-root must be filtered; got {ids:?}" - ); - assert!( - !ids.contains(&chunk_v1.id.as_str()), - "superseded version's chunk must not surface; got {ids:?}" - ); - } - - #[tokio::test] - async fn summary_drills_to_leaves_at_depth_one() { - let (_tmp, cfg) = test_config(); - let (root_id, _) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &root_id, 1, None, None).await.unwrap(); - assert_eq!(out.len(), 2, "L1 has 2 leaf children"); - for hit in &out { - assert_eq!(hit.level, 0, "direct children of L1 are leaves"); - } - } - - #[tokio::test] - async fn leaf_drill_down_returns_empty() { - let (_tmp, cfg) = test_config(); - let (_root_id, leaf_id) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &leaf_id, 3, None, None).await.unwrap(); - assert!(out.is_empty(), "leaves have no children"); - } - - #[tokio::test] - async fn deeper_max_depth_does_not_break_on_shallow_tree() { - // Only one summary level exists; asking for max_depth=5 is fine. - let (_tmp, cfg) = test_config(); - let (root_id, _) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &root_id, 5, None, None).await.unwrap(); - assert_eq!(out.len(), 2); - } - - #[tokio::test] - async fn query_with_limit_truncates_after_rerank() { - // Verifies the plumbing for the query param: embedder is invoked - // (InertEmbedder under this test config — all-zero vectors so - // cosine is 0 for every candidate), limit truncates the output, - // and the function completes without error. - let (_tmp, cfg) = test_config(); - let (root_id, _) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &root_id, 1, Some("phoenix migration timing"), Some(1)) - .await - .unwrap(); - assert_eq!(out.len(), 1, "limit=1 truncates 2 children to 1"); - } - - #[tokio::test] - async fn query_without_limit_returns_all_children() { - let (_tmp, cfg) = test_config(); - let (root_id, _) = seed_sealed_tree(&cfg).await; - let out = drill_down(&cfg, &root_id, 1, Some("phoenix"), None) - .await - .unwrap(); - assert_eq!(out.len(), 2, "no limit — both children returned"); - } - - // ── Regression: BFS (not DFS) traversal ────────────────────────── - // - // `walk_with_embeddings` walks level-by-level (all nodes at depth N - // before any at depth N+1) — originally flagged on PR #831 CodeRabbit - // review after the initial `Vec::pop()` implementation was DFS. - // - // A single-level tree can't distinguish the two (both produce the same - // output). We need a 2-level tree where BFS yields - // [L1_A, L1_B, c_A_1, c_A_2, c_B_1, c_B_2] - // and DFS would yield - // [L1_B, c_B_2, c_B_1, L1_A, c_A_2, c_A_1] - // (or similar — the key invariant is that BFS returns all siblings at - // one depth before any descendant at a deeper depth). - - use crate::openhuman::memory_store::chunks::store::with_connection; - use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeStatus}; - use crate::openhuman::memory_tree::tree::store as tree_store; - - /// Build a tiny 2-level tree directly via store inserts so we can - /// assert BFS ordering without needing ~100 leaves to cascade L1→L2 - /// through the token-budget seal path. - async fn seed_two_level_tree(cfg: &Config) -> (String, Vec, Vec) { - let ts = Utc::now(); - let tree = Tree { - id: "test:two-level".into(), - kind: TreeKind::Source, - scope: "slack:#eng".into(), - root_id: Some("s:L2:root".into()), - max_level: 2, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }; - let leaf_a_1 = Chunk { - id: "chat:slack:#eng:0".into(), - content: "leaf-a-1".into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: 10, - seq_in_source: 0, - created_at: ts, - partial_message: false, - }; - let leaf_a_2 = Chunk { - id: "chat:slack:#eng:1".into(), - content: "leaf-a-2".into(), - metadata: leaf_a_1.metadata.clone(), - seq_in_source: 1, - ..leaf_a_1.clone() - }; - let leaf_b_1 = Chunk { - id: "chat:slack:#eng:2".into(), - content: "leaf-b-1".into(), - metadata: leaf_a_1.metadata.clone(), - seq_in_source: 2, - ..leaf_a_1.clone() - }; - let leaf_b_2 = Chunk { - id: "chat:slack:#eng:3".into(), - content: "leaf-b-2".into(), - metadata: leaf_a_1.metadata.clone(), - seq_in_source: 3, - ..leaf_a_1.clone() - }; - let all_leaves = [ - leaf_a_1.clone(), - leaf_a_2.clone(), - leaf_b_1.clone(), - leaf_b_2.clone(), - ]; - upsert_chunks(cfg, &all_leaves).unwrap(); - // Stage to disk so `walk_with_embeddings` can read full bodies via - // `read_chunk_body` for leaf hits returned by the drill-down. - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).unwrap(); - let staged = content_store::stage_chunks(&content_root, &all_leaves).unwrap(); - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let l1_a = SummaryNode { - id: "s:L1:a".into(), - tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: Some("s:L2:root".into()), - child_ids: vec![leaf_a_1.id.clone(), leaf_a_2.id.clone()], - content: "L1 summary A".into(), - token_count: 50, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - let l1_b = SummaryNode { - id: "s:L1:b".into(), - child_ids: vec![leaf_b_1.id.clone(), leaf_b_2.id.clone()], - ..l1_a.clone() - }; - let root = SummaryNode { - id: "s:L2:root".into(), - level: 2, - parent_id: None, - child_ids: vec![l1_a.id.clone(), l1_b.id.clone()], - content: "L2 root".into(), - ..l1_a.clone() - }; - - // Open the shared connection to the memory_tree DB and write the - // tree + three summaries in one transaction. - with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - tree_store::insert_tree_conn(&tx, &tree)?; - tree_store::insert_summary_tx(&tx, &l1_a, None, "test")?; - tree_store::insert_summary_tx(&tx, &l1_b, None, "test")?; - tree_store::insert_summary_tx(&tx, &root, None, "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - ( - root.id, - vec![l1_a.id, l1_b.id], - vec![leaf_a_1.id, leaf_a_2.id, leaf_b_1.id, leaf_b_2.id], - ) - } - - #[tokio::test] - async fn walk_visits_siblings_before_descendants_bfs_order() { - let (_tmp, cfg) = test_config(); - let (root_id, l1_ids, leaf_ids) = seed_two_level_tree(&cfg).await; - - let out = drill_down(&cfg, &root_id, 2, None, None).await.unwrap(); - // Both L1s + all 4 leaves = 6 hits. - assert_eq!(out.len(), 6, "L2 with 2×L1 × 2 leaves each = 6 hits"); - - // Collect ids in returned order. - let ordered: Vec<&str> = out.iter().map(|h| h.node_id.as_str()).collect(); - - // BFS invariant: every L1 index must come BEFORE every leaf index. - // (DFS would interleave a whole L1 subtree before the other L1.) - let last_l1 = l1_ids - .iter() - .map(|id| ordered.iter().position(|&n| n == id).unwrap()) - .max() - .unwrap(); - let first_leaf = leaf_ids - .iter() - .map(|id| ordered.iter().position(|&n| n == id).unwrap()) - .min() - .unwrap(); - assert!( - last_l1 < first_leaf, - "BFS must return both L1 summaries before any leaf; got ordered={ordered:?}" - ); - } - - // ── Regression: per-depth batching keys HashMap by id, not position ── - // - // The per-depth batch refactor replaces the `for id in frontier` loop - // (one `get_summary` + one `get_tree` + (chunk: + one `get_chunk` + - // one `get_chunk_embedding`) per node) with one batched fetch per - // table, per level. The per-id walk that follows MUST look up by id - // — a refactor that mistakenly switches to `enumerate()` position - // over the input slice would silently shadow sibling A's `Tree` - // (and therefore `scope`) with sibling B's. This test seeds two L1 - // summaries belonging to **distinct trees** with **distinct scopes** - // and asserts each summary's emitted hit carries its own tree scope - // — proving the HashMap-keyed-by-id contract. - async fn seed_two_l1s_in_distinct_trees(cfg: &Config) -> (String, String, String) { - let ts = Utc::now(); - // Root sits in tree_a; its child L1s live in DIFFERENT trees so - // the per-depth `get_trees_batch` actually has 2 distinct ids to - // resolve and the HashMap lookup is non-trivial. - let tree_a = Tree { - id: "test:tree-a".into(), - kind: TreeKind::Source, - scope: "slack:#eng".into(), - root_id: Some("s:L2:root-a".into()), - max_level: 2, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: Some(ts), - }; - let tree_b = Tree { - id: "test:tree-b".into(), - scope: "slack:#design".into(), - root_id: None, - ..tree_a.clone() - }; - - let l1_a = SummaryNode { - id: "s:L1:a".into(), - tree_id: tree_a.id.clone(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: Some("s:L2:root-a".into()), - child_ids: vec![], - content: "L1 from tree-a".into(), - token_count: 50, - entities: vec![], - topics: vec![], - time_range_start: ts, - time_range_end: ts, - score: 0.5, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - let l1_b = SummaryNode { - id: "s:L1:b".into(), - tree_id: tree_b.id.clone(), - content: "L1 from tree-b".into(), - ..l1_a.clone() - }; - let root_a = SummaryNode { - id: "s:L2:root-a".into(), - level: 2, - parent_id: None, - child_ids: vec![l1_a.id.clone(), l1_b.id.clone()], - content: "L2 root in tree-a referencing L1s from two trees".into(), - ..l1_a.clone() - }; - - with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - tree_store::insert_tree_conn(&tx, &tree_a)?; - tree_store::insert_tree_conn(&tx, &tree_b)?; - tree_store::insert_summary_tx(&tx, &l1_a, None, "test")?; - tree_store::insert_summary_tx(&tx, &l1_b, None, "test")?; - tree_store::insert_summary_tx(&tx, &root_a, None, "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - (root_a.id, l1_a.id, l1_b.id) - } - - #[tokio::test] - async fn per_depth_batch_keys_hit_scope_by_tree_id_not_position() { - let (_tmp, cfg) = test_config(); - let (root_id, l1_a_id, l1_b_id) = seed_two_l1s_in_distinct_trees(&cfg).await; - - let out = drill_down(&cfg, &root_id, 1, None, None).await.unwrap(); - assert_eq!(out.len(), 2, "two L1 children expected"); - - // Find each hit by id and assert its scope came from its own - // tree row — not the other sibling's tree (which would happen - // if the per-id lookup used `enumerate()` position). - let hit_a = out - .iter() - .find(|h| h.node_id == l1_a_id) - .expect("hit for L1 in tree-a missing"); - let hit_b = out - .iter() - .find(|h| h.node_id == l1_b_id) - .expect("hit for L1 in tree-b missing"); - assert_eq!( - hit_a.tree_scope, "slack:#eng", - "L1 from tree-a must carry tree-a's scope" - ); - assert_eq!( - hit_b.tree_scope, "slack:#design", - "L1 from tree-b must carry tree-b's scope (NOT tree-a's)" - ); - assert_eq!(hit_a.tree_id, "test:tree-a"); - assert_eq!(hit_b.tree_id, "test:tree-b"); - } -} diff --git a/src/openhuman/memory_tree/retrieval/engine.rs b/src/openhuman/memory_tree/retrieval/engine.rs new file mode 100644 index 000000000..23bc3c1db --- /dev/null +++ b/src/openhuman/memory_tree/retrieval/engine.rs @@ -0,0 +1,26 @@ +use anyhow::Result; +use async_trait::async_trait; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_tree::score::embed::Embedder as HostEmbedder; + +pub(super) fn config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} + +pub(super) struct EmbedderBridge<'a>(pub &'a dyn HostEmbedder); + +#[async_trait] +impl tinycortex::memory::score::embed::Embedder for EmbedderBridge<'_> { + fn name(&self) -> &'static str { + self.0.name() + } + + async fn embed(&self, text: &str) -> Result> { + self.0.embed(text).await + } + + async fn embed_batch(&self, texts: &[&str]) -> Vec>> { + self.0.embed_batch(texts).await + } +} diff --git a/src/openhuman/memory_tree/retrieval/fast.rs b/src/openhuman/memory_tree/retrieval/fast.rs index 831f35f50..8367fdaad 100644 --- a/src/openhuman/memory_tree/retrieval/fast.rs +++ b/src/openhuman/memory_tree/retrieval/fast.rs @@ -1,446 +1,41 @@ -//! `fast_retrieve` — deterministic, LLM-free memory retrieval (E2GraphRAG). -//! -//! This replaces the agentic `walk` / `smart_walk` loops. Instead of an LLM -//! navigating the summary tree turn-by-turn, retrieval routing is decided -//! purely by spaCy query entities + co-occurrence-graph hop distance: -//! -//! 1. Extract query entities `Eq` (spaCy, with regex fallback). -//! 2. `Eq` empty → **global**: dense rerank over the summary tree. -//! 3. Otherwise compute related entity pairs `Ph` within `h` hops. -//! - `Ph` empty → **global with occurrence ranking**: dense top-2k, then -//! re-rank by how many `Eq` entities each summary mentions. -//! - `Ph` non-empty → **local (index mapping)**: intersect the entity-index -//! node sets of each related pair; tighten `h` while the candidate set is -//! larger than `k`; rank survivors by entity coverage then recency. -//! -//! Output is a structured [`QueryResponse`] of [`RetrievalHit`]s (no -//! synthesized prose) for a higher-level context agent to consume. - -use std::cmp::Reverse; -use std::collections::{HashMap, HashSet}; +//! Product adapters for tinycortex-owned deterministic fast retrieval. use anyhow::Result; use crate::openhuman::config::Config; use crate::openhuman::memory::source_scope::current_source_scope; -use crate::openhuman::memory_store::content::read as content_read; -use crate::openhuman::memory_store::trees::store as tree_store; -use crate::openhuman::memory_tree::graph; use crate::openhuman::memory_tree::nlp; -use crate::openhuman::memory_tree::retrieval::fetch::{self, fetch_leaves}; -use crate::openhuman::memory_tree::retrieval::source::query_source; -use crate::openhuman::memory_tree::retrieval::types::{ - hit_from_summary, QueryResponse, RetrievalHit, -}; -use crate::openhuman::memory_tree::score::store::lookup_entity; +use crate::openhuman::memory_tree::retrieval::engine::{config as engine_config, EmbedderBridge}; +use crate::openhuman::memory_tree::retrieval::types::QueryResponse; +use crate::openhuman::memory_tree::score::embed::build_embedder_from_config; -/// Tunables for [`fast_retrieve`]. Defaults match the E2GraphRAG paper's -/// small-`k` regime and a 2-hop relatedness threshold. -#[derive(Clone, Debug)] -pub struct FastRetrieveOptions { - /// `k` — max hits returned. - pub limit: usize, - /// `h` — initial hop threshold for relatedness. Tightened (decremented) - /// during the local branch when too many candidates match. - pub max_hops: u32, - /// Look-back window (days) applied on the global/dense branch. `None` = - /// unbounded. - pub time_window_days: Option, -} +pub use tinycortex::memory::retrieval::FastRetrieveOptions; -impl Default for FastRetrieveOptions { - fn default() -> Self { - Self { - limit: 10, - max_hops: 2, - time_window_days: None, - } - } -} - -/// Per-node-entity lookup cap. Popular entities can touch many nodes; this -/// bounds the intersection work while staying well above any realistic `k`. -const LOOKUP_LIMIT: usize = 500; - -/// Default / ceiling for `limit` (`k`). The tool and RPC paths expose this, so -/// a huge value must not be able to request oversized dense/local result sets. -const DEFAULT_LIMIT: usize = 10; -const MAX_RETRIEVE_LIMIT: usize = 100; -/// Default / ceiling for the hop threshold. A large `max_hops` would force many -/// bounded-BFS passes through the tightening loop; cap it. -const DEFAULT_MAX_HOPS: u32 = 2; -const MAX_GRAPH_HOPS: u32 = 4; - -/// Run deterministic retrieval for `query`. Never invokes an LLM. pub async fn fast_retrieve( config: &Config, query: &str, - opts: FastRetrieveOptions, + options: FastRetrieveOptions, ) -> Result { - let k = match opts.limit { - 0 => DEFAULT_LIMIT, - n => n.min(MAX_RETRIEVE_LIMIT), - }; - let max_hops = match opts.max_hops { - 0 => DEFAULT_MAX_HOPS, - n => n.min(MAX_GRAPH_HOPS), - }; - let trimmed = query.trim(); - if trimmed.is_empty() { - return Ok(QueryResponse::empty()); - } - - // 1. Query entities. - let entities = nlp::extract_query_entities(config, trimmed).await; - let eq: Vec = dedup_ids(entities.iter().map(|e| e.canonical_id.clone())); + let query_entities = nlp::extract_query_entities(config, query).await; + let entity_ids: Vec<_> = query_entities + .into_iter() + .map(|entity| entity.canonical_id) + .collect(); log::debug!( - "[retrieval::fast] query_len={} eq={} k={} h={}", - trimmed.len(), - eq.len(), - k, - max_hops + "[retrieval::fast] tinycortex query_len={} entities={} limit={} hops={}", + query.len(), + entity_ids.len(), + options.limit, + options.max_hops ); - - // 2. No entities → pure global dense retrieval. - if eq.is_empty() { - log::debug!("[retrieval::fast] branch=global_dense (no query entities)"); - return query_source(config, None, None, opts.time_window_days, Some(trimmed), k).await; - } - - // 3. Graph filter: related entity pairs within `h` hops. - let pairs = graph::pair_distances(config, &eq, max_hops)?; - if pairs.is_empty() { - log::debug!("[retrieval::fast] branch=global_occurrence (no related pairs)"); - return global_occurrence(config, trimmed, &eq, k, opts.time_window_days).await; - } - - // 4. Local branch: index-mapping intersection with `h` tightening. - let mut h = max_hops; - let mut cands = local_candidates(config, &eq, &pairs)?; - let mut last_non_empty = cands.clone(); - while cands.len() > k && h > 1 { - h -= 1; - let tightened = graph::pair_distances(config, &eq, h)?; - let next = local_candidates(config, &eq, &tightened)?; - if next.is_empty() { - // Tightening removed everything — keep the looser, non-empty set. - break; - } - last_non_empty = next.clone(); - cands = next; - } - if cands.is_empty() { - cands = last_non_empty; - } - - if cands.is_empty() { - // Related pairs existed but never co-occurred in an indexed node - // (e.g. only 2-hop links). Fall back to global occurrence ranking. - log::debug!("[retrieval::fast] branch=local->global_occurrence (empty intersection)"); - return global_occurrence(config, trimmed, &eq, k, opts.time_window_days).await; - } - - log::debug!( - "[retrieval::fast] branch=local candidates={} final_h={}", - cands.len(), - h - ); - resolve_local(config, cands, k).await -} - -/// Candidate node accumulated during the local branch. -#[derive(Clone, Debug)] -struct Candidate { - node_kind: String, - /// Distinct query entities that landed on this node (coverage signal). - matched: HashSet, - latest_ts: i64, -} - -/// Intersect the entity-index node sets of each related pair and union the -/// results. A node enters the candidate set when it is indexed against *both* -/// members of some related pair. -fn local_candidates( - config: &Config, - eq: &[String], - pairs: &[graph::PairDistance], -) -> Result> { - let _ = eq; // `eq` documents intent; matched entities come from the pairs. - let mut out: HashMap = HashMap::new(); - for pair in pairs { - let hits_a = lookup_entity(config, &pair.a, Some(LOOKUP_LIMIT))?; - if hits_a.is_empty() { - continue; - } - let hits_b = lookup_entity(config, &pair.b, Some(LOOKUP_LIMIT))?; - if hits_b.is_empty() { - continue; - } - let b_nodes: HashMap<&str, i64> = hits_b - .iter() - .map(|h| (h.node_id.as_str(), h.timestamp_ms)) - .collect(); - for ha in &hits_a { - let Some(&b_ts) = b_nodes.get(ha.node_id.as_str()) else { - continue; - }; - let entry = out.entry(ha.node_id.clone()).or_insert_with(|| Candidate { - node_kind: ha.node_kind.clone(), - matched: HashSet::new(), - latest_ts: 0, - }); - entry.matched.insert(pair.a.clone()); - entry.matched.insert(pair.b.clone()); - entry.latest_ts = entry.latest_ts.max(ha.timestamp_ms).max(b_ts); - } - } - Ok(out) -} - -/// Rank candidates by entity coverage (desc) then recency (desc), resolve to -/// hits, apply the profile source-scope gate **before** counting/truncating -/// (so an out-of-scope top hit never displaces an allowed lower-ranked one, -/// and `total` never reveals out-of-scope match counts), then truncate to `k`. -async fn resolve_local( - config: &Config, - cands: HashMap, - k: usize, -) -> Result { - let mut ordered: Vec<(String, Candidate)> = cands.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)) - }); - - // Coverage score per node id so resolved hits carry the ranking signal. - let coverage: HashMap = ordered - .iter() - .map(|(id, c)| (id.clone(), c.matched.len() as f32)) - .collect(); - - let leaf_ids: Vec = ordered - .iter() - .filter(|(_, c)| c.node_kind == "leaf") - .map(|(id, _)| id.clone()) - .collect(); - let summary_ids: Vec = ordered - .iter() - .filter(|(_, c)| c.node_kind != "leaf") - .map(|(id, _)| id.clone()) - .collect(); - - let mut by_id: HashMap = HashMap::new(); - // `fetch_leaves` caps each batch at MAX_BATCH and would silently drop the - // rest, so chunk to resolve every candidate leaf before the scope gate. - for chunk in leaf_ids.chunks(fetch::MAX_BATCH) { - for hit in fetch_leaves(config, chunk).await? { - by_id.insert(hit.node_id.clone(), hit); - } - } - if !summary_ids.is_empty() { - for hit in resolve_summaries(config, &summary_ids)? { - by_id.insert(hit.node_id.clone(), hit); - } - } - - // Scope-gate the full ranked set first; `total` reflects only in-scope hits. - let scope = current_source_scope(); - let mut hits: Vec = Vec::with_capacity(ordered.len()); - for (id, _) in &ordered { - if let Some(mut hit) = by_id.remove(id) { - if !scope_allows(scope.as_ref(), &hit.tree_scope) { - continue; - } - if let Some(score) = coverage.get(id) { - hit.score = *score; - } - hits.push(hit); - } - } - let total = hits.len(); - hits.truncate(k); - Ok(QueryResponse::new(hits, total)) -} - -/// Resolve summary node ids to hits, hydrating the full body from disk and -/// threading the owning tree's scope for provenance. -fn resolve_summaries(config: &Config, summary_ids: &[String]) -> Result> { - let nodes = tree_store::get_summaries_batch(config, summary_ids)?; - let mut scope_cache: HashMap = HashMap::new(); - let mut out = Vec::with_capacity(nodes.len()); - for (_, mut node) in nodes { - // Hydrate full body (the `content` column is a ≤500-char preview). - if let Ok(body) = content_read::read_summary_body(config, &node.id) { - node.content = body; - } - let scope = match scope_cache.get(&node.tree_id) { - Some(s) => s.clone(), - None => { - let s = tree_store::get_tree(config, &node.tree_id) - .ok() - .flatten() - .map(|t| t.scope) - .unwrap_or_default(); - scope_cache.insert(node.tree_id.clone(), s.clone()); - s - } - }; - out.push(hit_from_summary(&node, &scope)); - } - Ok(out) -} - -/// Global branch with occurrence ranking: dense-retrieve top-`2k`, then re-rank -/// by how many `Eq` entities each summary mentions (occurrence). Stable sort -/// preserves the semantic order on ties. -async fn global_occurrence( - config: &Config, - query: &str, - eq: &[String], - k: usize, - window: Option, -) -> Result { - let resp = query_source(config, None, None, window, Some(query), k.saturating_mul(2)).await?; - let eq_set: HashSet<&str> = eq.iter().map(|s| s.as_str()).collect(); - let mut hits = resp.hits; - let total = resp.total; - hits.sort_by_key(|h| { - let occ = h - .entities - .iter() - .filter(|e| eq_set.contains(e.as_str())) - .count(); - Reverse(occ) - }); - hits.truncate(k); - Ok(QueryResponse::new(hits, total)) -} - -/// Deduplicate ids preserving first-seen order. -fn dedup_ids(ids: impl Iterator) -> Vec { - let mut seen = HashSet::new(); - let mut out = Vec::new(); - for id in ids { - if seen.insert(id.clone()) { - out.push(id); - } - } - out -} - -/// Profile source-scope gate: `None` = unrestricted; otherwise the scope must -/// be on the allowlist. Empty scopes (unknown provenance) are allowed through -/// only when unrestricted. -fn scope_allows(scope: Option<&HashSet>, tree_scope: &str) -> bool { - match scope { - None => true, - Some(set) => set.contains(tree_scope), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::ingest_pipeline::ingest_chat; - use crate::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; - use chrono::{TimeZone, Utc}; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - // spaCy off in CI — exercise the regex fallback + graph routing - // deterministically without a Python runtime. - cfg.memory_tree.spacy_enabled = false; - (tmp, cfg) - } - - async fn seed_chat(cfg: &Config, source: &str, text: &str) { - let batch = ChatBatch { - platform: "slack".into(), - channel_label: source.into(), - messages: vec![ChatMessage { - author: "alice".into(), - timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: text.into(), - source_ref: Some("slack://x".into()), - }], - }; - ingest_chat(cfg, source, "alice", vec![], batch) - .await - .unwrap(); - } - - #[tokio::test] - async fn empty_query_returns_empty() { - let (_tmp, cfg) = test_config(); - let resp = fast_retrieve(&cfg, " ", FastRetrieveOptions::default()) - .await - .unwrap(); - assert!(resp.hits.is_empty()); - } - - #[tokio::test] - async fn no_entities_routes_global_without_panicking() { - let (_tmp, cfg) = test_config(); - // A query with no mechanical entities (regex fallback finds nothing) - // must take the global branch and return cleanly on an empty store. - let resp = fast_retrieve( - &cfg, - "what happened recently", - FastRetrieveOptions::default(), - ) - .await - .unwrap(); - assert!(resp.hits.is_empty()); - } - - #[tokio::test] - async fn local_branch_finds_cooccurring_entities() { - let (_tmp, cfg) = test_config(); - // Two emails co-occur in one message → an edge + both indexed on the - // same leaf. A query mentioning both routes local and returns the leaf. - seed_chat( - &cfg, - "slack:#eng", - "Sync between alice@example.com and bob@example.com about the runbook.", - ) - .await; - - let resp = fast_retrieve( - &cfg, - "alice@example.com and bob@example.com", - FastRetrieveOptions::default(), - ) - .await - .unwrap(); - assert!( - !resp.hits.is_empty(), - "co-occurring entities should yield a local hit; got {resp:?}" - ); - // Coverage score = 2 distinct query entities matched the leaf. - assert!(resp.hits.iter().any(|h| h.score >= 2.0)); - } - - #[test] - fn scope_allows_respects_allowlist() { - assert!(scope_allows(None, "slack:#eng")); - let mut set = HashSet::new(); - set.insert("slack:#eng".to_string()); - assert!(scope_allows(Some(&set), "slack:#eng")); - assert!(!scope_allows(Some(&set), "gmail:alice@example.com")); - } - - #[test] - fn dedup_ids_preserves_first_seen_order() { - let out = dedup_ids(["a".to_string(), "b".into(), "a".into(), "c".into()].into_iter()); - assert_eq!(out, vec!["a", "b", "c"]); - } + let embedder = build_embedder_from_config(config)?; + tinycortex::memory::retrieval::fast_retrieve( + &engine_config(config), + query, + &entity_ids, + &EmbedderBridge(embedder.as_ref()), + current_source_scope().as_ref(), + options, + ) + .await } diff --git a/src/openhuman/memory_tree/retrieval/fetch.rs b/src/openhuman/memory_tree/retrieval/fetch.rs index 275bec589..4c8456398 100644 --- a/src/openhuman/memory_tree/retrieval/fetch.rs +++ b/src/openhuman/memory_tree/retrieval/fetch.rs @@ -1,295 +1,32 @@ -//! `memory_tree_fetch_leaves` — batch-fetch raw chunks by id (Phase 4 / -//! #710). -//! -//! The LLM-facing contract: "given these chunk ids, give me the full -//! content + metadata so I can cite." We cap the batch at 20 to keep the -//! round-trip bounded. Missing ids are silently skipped — the return is -//! best-effort so partial failures are visible via `hits.len() < ids.len()`. -//! -//! Each hit is annotated with the chunk's score from `mem_tree_score` when -//! available; score is 0.0 when the chunk has no row in `mem_tree_score` -//! (e.g. pre-Phase 2 backfill). - use anyhow::Result; use crate::openhuman::config::Config; +use crate::openhuman::memory::source_scope::chunk_source_allowed_in; +use crate::openhuman::memory::source_scope::current_source_scope; use crate::openhuman::memory_store::chunks::store::get_chunks_batch; -use crate::openhuman::memory_store::content::read as content_read; -use crate::openhuman::memory_tree::retrieval::types::{hit_from_chunk, RetrievalHit}; -use crate::openhuman::memory_tree::score::store::get_scores_batch; +use crate::openhuman::memory_tree::retrieval::engine::config as engine_config; +use crate::openhuman::memory_tree::retrieval::types::RetrievalHit; -/// Max batch size. Callers that pass more than this get truncated with a -/// warn log — no error surface so the LLM sees a partial result. -pub const MAX_BATCH: usize = 20; +pub use tinycortex::memory::retrieval::MAX_BATCH; -/// Fetch chunk rows by id in the provided order. Missing ids are dropped -/// from the response. pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result> { - if chunk_ids.is_empty() { - log::debug!("[retrieval::fetch] empty request — returning empty vec"); - return Ok(Vec::new()); - } - - let ids: Vec = if chunk_ids.len() > MAX_BATCH { - log::warn!( - "[retrieval::fetch] batch size {} exceeds cap {} — truncating", - chunk_ids.len(), - MAX_BATCH - ); - chunk_ids[..MAX_BATCH].to_vec() + log::debug!( + "[retrieval::fetch] tinycortex requested={}", + chunk_ids.len() + ); + let permitted_ids = if let Some(set) = current_source_scope() { + let chunks = get_chunks_batch(config, chunk_ids)?; + chunk_ids + .iter() + .filter(|id| { + chunks.get(*id).is_some_and(|chunk| { + chunk_source_allowed_in(&set, &chunk.metadata.tags, &chunk.metadata.source_id) + }) + }) + .cloned() + .collect::>() } else { chunk_ids.to_vec() }; - - // Count only — individual chunk ids can include source scope (e.g. - // `chat:slack:#:0`) and are redacted from logs. - log::debug!("[retrieval::fetch] fetch_leaves n={}", ids.len()); - - let config_owned = config.clone(); - let hits = tokio::task::spawn_blocking(move || -> Result> { - // Two batched SQLite reads up front instead of 2N per-id queries - // inside the loop. With the `MAX_BATCH = 20` cap above, this turns - // 40 round-trips into 2. Per-row decoders are reused inside both - // helpers so the returned `Chunk` and `score.total` values are - // byte-identical to the old per-id path. Missing ids are absent - // from the maps (same contract as `get_chunk` / `get_score` - // returning `Ok(None)`). - let chunk_by_id = get_chunks_batch(&config_owned, &ids)?; - let score_by_id = get_scores_batch(&config_owned, &ids)?; - - // Walk the input ids in order so the response preserves caller - // ordering. Missing ids are dropped exactly as before — callers - // detect partial results via `hits.len() < ids.len()`. File I/O - // (`read_chunk_body`) stays per-id: each MD body lives in its - // own on-disk file, so batching there would mean concurrent file - // opens, not a single round-trip — left untouched. - let mut out: Vec = Vec::with_capacity(ids.len()); - for (idx, id) in ids.iter().enumerate() { - let Some(chunk) = chunk_by_id.get(id) else { - log::debug!( - "[retrieval::fetch] chunk not found at index {}/{} — skipping", - idx + 1, - ids.len() - ); - continue; - }; - let score = score_by_id.get(id).copied().unwrap_or(0.0); - // Leaves are not attached to a materialised tree id via the - // chunk row. `scope` falls back to the chunk's own source_id so - // consumers still see provenance (e.g. "slack:#eng"). - let scope = chunk.metadata.source_id.clone(); - // Hydrate the full body from disk before building the hit. - // The `content` column in SQLite holds a ≤500-char preview after - // the MD-on-disk migration; the retrieval API must return the - // complete chunk text so the LLM sees untruncated content. - let mut chunk_with_body = chunk.clone(); - match content_read::read_chunk_body(&config_owned, id) { - Ok(body) => chunk_with_body.content = body, - Err(e) => { - log::warn!( - "[retrieval::fetch] read_chunk_body failed for chunk — serving preview: {e:#}" - ); - // Non-fatal: fall back to the preview already in the struct. - // This handles pre-MD-migration rows gracefully. - } - } - out.push(hit_from_chunk(&chunk_with_body, "", &scope, score)); - } - Ok(out) - }) - .await - .map_err(|e| anyhow::anyhow!("fetch_leaves join error: {e}"))??; - - log::debug!("[retrieval::fetch] returning hits={}", hits.len()); - Ok(hits) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::content as content_store; - use chrono::{TimeZone, Utc}; - use tempfile::TempDir; - - fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) { - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).expect("create content_root for test"); - let staged = content_store::stage_chunks(&content_root, chunks) - .expect("stage_chunks for test chunks"); - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .expect("persist staged chunk pointers"); - } - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): inert embedder for tests. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - fn sample_chunk(source: &str, seq: u32) -> Chunk { - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - Chunk { - id: chunk_id(SourceKind::Chat, source, seq, "test-content"), - content: format!("content-{source}-{seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: source.into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new(format!("slack://{source}/{seq}"))), - path_scope: None, - }, - token_count: 20, - seq_in_source: seq, - created_at: ts, - partial_message: false, - } - } - - #[tokio::test] - async fn empty_input_returns_empty() { - let (_tmp, cfg) = test_config(); - let out = fetch_leaves(&cfg, &[]).await.unwrap(); - assert!(out.is_empty()); - } - - #[tokio::test] - async fn returns_existing_chunks_in_order() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0); - let c2 = sample_chunk("slack:#eng", 1); - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]); - let out = fetch_leaves(&cfg, &[c1.id.clone(), c2.id.clone()]) - .await - .unwrap(); - assert_eq!(out.len(), 2); - assert_eq!(out[0].node_id, c1.id); - assert_eq!(out[1].node_id, c2.id); - } - - #[tokio::test] - async fn missing_ids_are_skipped() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0); - upsert_chunks(&cfg, &[c1.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c1.clone()]); - let out = fetch_leaves( - &cfg, - &[c1.id.clone(), "ghost:nonexistent".into(), c1.id.clone()], - ) - .await - .unwrap(); - assert_eq!(out.len(), 2); - assert!(out.iter().all(|h| h.node_id == c1.id)); - } - - #[tokio::test] - async fn over_cap_is_truncated() { - let (_tmp, cfg) = test_config(); - let mut ids: Vec = Vec::new(); - for i in 0..(MAX_BATCH + 5) as u32 { - let c = sample_chunk("slack:#eng", i); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c.clone()]); - ids.push(c.id); - } - let out = fetch_leaves(&cfg, &ids).await.unwrap(); - assert_eq!(out.len(), MAX_BATCH); - } - - #[tokio::test] - async fn leaf_hit_carries_source_ref_and_scope() { - let (_tmp, cfg) = test_config(); - let c = sample_chunk("slack:#eng", 0); - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c.clone()]); - let out = fetch_leaves(&cfg, &[c.id.clone()]).await.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"); - } - - /// After the batch refactor, ordering and score propagation rely on - /// walking the input slice and looking each id up in two HashMaps - /// (chunk + score). This test pins both invariants: interleaved - /// present/missing ids keep their input order, and each kept hit - /// carries the score from its own `mem_tree_score` row (not the - /// row of a neighbour, not the 0.0 fallback when a row exists). - #[tokio::test] - async fn fetch_leaves_preserves_input_order_and_propagates_scores() { - use crate::openhuman::memory_tree::score::signals::ScoreSignals; - use crate::openhuman::memory_tree::score::store::{upsert_score, ScoreRow}; - - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("slack:#eng", 0); - let c2 = sample_chunk("slack:#eng", 1); - let c3 = sample_chunk("slack:#eng", 2); - upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]); - - // Distinct totals so we can tell which chunk a hit's score came - // from — proves the per-id HashMap lookup keys by chunk_id and - // not by iteration index. - let mk_row = |id: &str, total: f32| ScoreRow { - chunk_id: id.to_string(), - total, - signals: ScoreSignals { - token_count: 0.0, - unique_words: 0.0, - metadata_weight: 0.0, - source_weight: 0.0, - interaction: 0.0, - entity_density: 0.0, - llm_importance: 0.0, - }, - llm_importance_reason: None, - dropped: false, - reason: None, - computed_at_ms: 0, - }; - upsert_score(&cfg, &mk_row(&c1.id, 0.1)).unwrap(); - upsert_score(&cfg, &mk_row(&c2.id, 0.2)).unwrap(); - // c3 intentionally has NO score row so we also pin the 0.0 - // fallback after the get_scores_batch contract. - - // Request order: c2, ghost, c3, c1 — none in natural id order. - let out = fetch_leaves( - &cfg, - &[ - c2.id.clone(), - "ghost:no-such".into(), - c3.id.clone(), - c1.id.clone(), - ], - ) - .await - .unwrap(); - assert_eq!(out.len(), 3, "ghost dropped, 3 real chunks returned"); - assert_eq!(out[0].node_id, c2.id); - assert_eq!(out[1].node_id, c3.id); - assert_eq!(out[2].node_id, c1.id); - assert!((out[0].score - 0.2).abs() < 1e-6, "c2 score"); - assert!( - out[1].score.abs() < 1e-6, - "c3 has no score row → 0.0 fallback" - ); - assert!((out[2].score - 0.1).abs() < 1e-6, "c1 score"); - } + tinycortex::memory::retrieval::fetch_leaves(&engine_config(config), &permitted_ids) } diff --git a/src/openhuman/memory_tree/retrieval/integration_tests.rs b/src/openhuman/memory_tree/retrieval/integration_tests.rs index 226b34e06..da4155ca2 100644 --- a/src/openhuman/memory_tree/retrieval/integration_tests.rs +++ b/src/openhuman/memory_tree/retrieval/integration_tests.rs @@ -139,7 +139,22 @@ async fn ingest_populates_chunk_embeddings() { out.chunks_written >= 1, "expected at least one persisted chunk" ); - drain_until_idle(&cfg).await.unwrap(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(60); + loop { + drain_until_idle(&cfg).await.unwrap(); + let all_embedded = out + .chunk_ids + .iter() + .all(|id| get_chunk_embedding(&cfg, id).ok().flatten().is_some()); + if all_embedded { + break; + } + assert!( + tokio::time::Instant::now() < deadline, + "chunk embeddings were not persisted before timeout" + ); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } for id in &out.chunk_ids { let emb = get_chunk_embedding(&cfg, id).unwrap(); let v = emb.unwrap_or_else(|| panic!("embedding missing for chunk_id={id}")); diff --git a/src/openhuman/memory_tree/retrieval/mod.rs b/src/openhuman/memory_tree/retrieval/mod.rs index a41543fae..41e207a46 100644 --- a/src/openhuman/memory_tree/retrieval/mod.rs +++ b/src/openhuman/memory_tree/retrieval/mod.rs @@ -19,6 +19,7 @@ pub mod cover; pub mod drill_down; +mod engine; pub mod fast; pub mod fetch; pub mod rpc; diff --git a/src/openhuman/memory_tree/retrieval/search.rs b/src/openhuman/memory_tree/retrieval/search.rs index c26d3bf29..9dc11e518 100644 --- a/src/openhuman/memory_tree/retrieval/search.rs +++ b/src/openhuman/memory_tree/retrieval/search.rs @@ -1,319 +1,26 @@ -//! `memory_tree_search_entities` — free-text LIKE search over the entity -//! index (Phase 4 / #710). -//! -//! The entity index (`mem_tree_entity_index`) is populated at ingest time -//! with one row per (entity, node) occurrence. This tool exposes it to the -//! LLM as a fuzzy-ish lookup: "I'm not sure if alice is the canonical id — -//! let me search". We group by canonical id so repeated mentions collapse -//! into a single [`EntityMatch`] with an aggregate count. -//! -//! Matching rules: -//! - Query is lowercased before binding into the `LIKE` parameters. -//! - We match either `entity_id LIKE '%q%'` (canonical-id substring) OR -//! `surface LIKE '%q%'` (display-form substring). -//! - `kinds` narrows the match by `entity_kind IN (...)` when non-empty. -//! - Output is ordered by mention count DESC so the strongest matches -//! surface first. - -use anyhow::{Context, Result}; -use rusqlite::params_from_iter; +use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; +use crate::openhuman::memory_tree::retrieval::engine::config as engine_config; use crate::openhuman::memory_tree::retrieval::types::EntityMatch; use crate::openhuman::memory_tree::score::extract::EntityKind; -const DEFAULT_LIMIT: usize = 5; -const MAX_LIMIT: usize = 100; - -/// Search the entity index for canonical ids matching `query`. -/// -/// Returns at most `limit` matches (default 5, clamped to 100). Each match -/// is aggregated across every row of the entity index so `mention_count` -/// reflects total occurrences regardless of which tree they came from. pub async fn search_entities( config: &Config, query: &str, kinds: Option>, limit: usize, ) -> Result> { - let limit = normalise_limit(limit); - // Blank/whitespace-only queries would turn into `LIKE '%%'` and dump the - // entire entity index. Return empty early instead. Flagged on PR #831 - // CodeRabbit review. - let query = query.trim(); - if query.is_empty() { - log::debug!("[retrieval::search] empty query — returning no matches"); - return Ok(Vec::new()); - } - - // Log `query_len` rather than the query itself — the query can be an - // email, a handle, or any PII. log::debug!( - "[retrieval::search] search_entities query_len={} kinds={:?} limit={}", + "[retrieval::search] tinycortex query_len={} kinds={} limit={}", query.len(), - kinds - .as_ref() - .map(|ks| ks.iter().map(|k| k.as_str()).collect::>()), + kinds.as_ref().map_or(0, Vec::len), limit ); - - let q_lower = query.to_lowercase(); - let kinds_owned = kinds.clone(); - let config_owned = config.clone(); - let rows = tokio::task::spawn_blocking(move || -> Result> { - with_connection(&config_owned, |conn| { - let pattern = format!("%{q_lower}%"); - let (sql, params) = build_sql_and_params(&pattern, kinds_owned.as_deref(), limit); - let mut stmt = conn - .prepare(&sql) - .with_context(|| "search_entities: failed to prepare statement")?; - let mapped = stmt - .query_map(params_from_iter(params.iter()), row_to_match)? - .collect::>>() - .with_context(|| "search_entities: failed to collect rows")?; - Ok(mapped) - }) - }) - .await - .map_err(|e| anyhow::anyhow!("search_entities join error: {e}"))??; - - log::debug!("[retrieval::search] returning matches={}", rows.len()); - Ok(rows) -} - -fn normalise_limit(limit: usize) -> usize { - if limit == 0 { - DEFAULT_LIMIT - } else { - limit.min(MAX_LIMIT) - } -} - -/// Build the SQL string + bound parameters. Kept in its own function so we -/// can unit-test the shape of the generated statement without a real DB. -fn build_sql_and_params( - pattern: &str, - kinds: Option<&[EntityKind]>, - limit: usize, -) -> (String, Vec) { - use rusqlite::types::Value; - let mut sql = String::from( - "SELECT - entity_id, - entity_kind, - MAX(surface) AS surface_sample, - COUNT(*) AS mention_count, - MAX(timestamp_ms) AS last_seen_ms - FROM mem_tree_entity_index - WHERE (LOWER(entity_id) LIKE ?1 OR LOWER(surface) LIKE ?1)", - ); - let mut params: Vec = vec![Value::Text(pattern.to_string())]; - - if let Some(ks) = kinds { - if !ks.is_empty() { - let placeholders: Vec = (0..ks.len()).map(|i| format!("?{}", i + 2)).collect(); - sql.push_str(&format!( - " AND entity_kind IN ({})", - placeholders.join(", ") - )); - for k in ks { - params.push(Value::Text(k.as_str().to_string())); - } - } - } - - sql.push_str( - " GROUP BY entity_id, entity_kind - ORDER BY mention_count DESC, last_seen_ms DESC - LIMIT ?", - ); - params.push(Value::Integer(limit as i64)); - - (sql, params) -} - -fn row_to_match(row: &rusqlite::Row<'_>) -> rusqlite::Result { - let canonical_id: String = row.get(0)?; - let kind_s: String = row.get(1)?; - let surface: String = row.get(2)?; - let mention_count: i64 = row.get(3)?; - let last_seen_ms: i64 = row.get(4)?; - - let kind = EntityKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into()) - })?; - - Ok(EntityMatch { - canonical_id, - kind, - surface, - mention_count: mention_count.max(0) as u64, - last_seen_ms, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::ingest_pipeline::ingest_chat; - use crate::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; - use chrono::{TimeZone, Utc}; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): ingest in seeding needs inert embedder. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - async fn seed_chat(cfg: &Config, source: &str, text: &str) { - let batch = ChatBatch { - platform: "slack".into(), - channel_label: source.into(), - messages: vec![ChatMessage { - author: "alice".into(), - timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: text.into(), - source_ref: Some("slack://x".into()), - }], - }; - ingest_chat(cfg, source, "alice", vec![], batch) - .await - .unwrap(); - } - - #[tokio::test] - async fn empty_index_returns_empty_vec() { - let (_tmp, cfg) = test_config(); - let matches = search_entities(&cfg, "alice", None, 10).await.unwrap(); - assert!(matches.is_empty()); - } - - #[tokio::test] - async fn matches_on_entity_id_substring() { - let (_tmp, cfg) = test_config(); - seed_chat( - &cfg, - "slack:#eng", - "Planning the Phoenix migration on Friday. alice@example.com owns it.", - ) - .await; - let matches = search_entities(&cfg, "alice", None, 10).await.unwrap(); - assert!( - matches - .iter() - .any(|m| m.canonical_id == "email:alice@example.com"), - "expected alice's canonical id in matches; got {matches:?}" - ); - } - - #[tokio::test] - async fn matches_on_surface_substring() { - let (_tmp, cfg) = test_config(); - seed_chat( - &cfg, - "slack:#eng", - "Planning the Phoenix migration. alice@example.com owns it. Running the runbook again.", - ) - .await; - // "example.com" appears in surface but not in canonical_id alone. - let matches = search_entities(&cfg, "example.com", None, 10) - .await - .unwrap(); - assert!( - matches.iter().any(|m| m.canonical_id.contains("alice")), - "surface-matched row must surface; got {matches:?}" - ); - } - - #[tokio::test] - async fn kind_filter_narrows_results() { - let (_tmp, cfg) = test_config(); - seed_chat( - &cfg, - "slack:#eng", - "Planning Phoenix. alice@example.com. #launch-q2 tagged. \ - And let's confirm with bob@example.com too. runbook review.", - ) - .await; - let only_hashtags = search_entities(&cfg, "launch", Some(vec![EntityKind::Hashtag]), 10) - .await - .unwrap(); - assert!(only_hashtags - .iter() - .all(|m| matches!(m.kind, EntityKind::Hashtag))); - } - - #[tokio::test] - async fn matches_aggregate_across_multiple_sources() { - let (_tmp, cfg) = test_config(); - seed_chat( - &cfg, - "slack:#a", - "Meeting 1 about Phoenix. alice@example.com attends. The migration runbook review proceeds.", - ) - .await; - seed_chat( - &cfg, - "slack:#b", - "Meeting 2 about Phoenix. alice@example.com attends again. Launch date confirmed.", - ) - .await; - let matches = search_entities(&cfg, "alice", None, 10).await.unwrap(); - let alice = matches - .iter() - .find(|m| m.canonical_id == "email:alice@example.com") - .expect("alice should be in matches"); - assert!( - alice.mention_count >= 2, - "expected at least 2 mentions aggregated across sources, got {}", - alice.mention_count - ); - } - - #[tokio::test] - async fn limit_truncates_results() { - let (_tmp, cfg) = test_config(); - seed_chat( - &cfg, - "slack:#eng", - "Planning Phoenix. alice@example.com. bob@example.com. \ - charlie@example.com. dana@example.com. eric@example.com. Run the runbook.", - ) - .await; - let matches = search_entities(&cfg, "example.com", None, 2).await.unwrap(); - assert!(matches.len() <= 2); - } - - #[test] - fn build_sql_without_kinds_has_no_in_clause() { - let (sql, _params) = build_sql_and_params("%a%", None, 5); - assert!(sql.contains("LOWER(entity_id) LIKE")); - assert!(!sql.contains("entity_kind IN")); - } - - #[test] - fn build_sql_with_kinds_adds_in_clause() { - let kinds = vec![EntityKind::Email, EntityKind::Hashtag]; - let (sql, params) = build_sql_and_params("%x%", Some(&kinds), 5); - assert!(sql.contains("entity_kind IN")); - // pattern + 2 kinds + limit = 4 params - assert_eq!(params.len(), 4); - } - - #[test] - fn zero_limit_defaults_to_five() { - assert_eq!(normalise_limit(0), DEFAULT_LIMIT); - } - - #[test] - fn huge_limit_is_clamped() { - assert_eq!(normalise_limit(10_000), MAX_LIMIT); - } + tinycortex::memory::retrieval::search_entities( + &engine_config(config), + query, + kinds.as_deref(), + limit, + ) } diff --git a/src/openhuman/memory_tree/retrieval/source.rs b/src/openhuman/memory_tree/retrieval/source.rs index 2ad09bb21..b8e932330 100644 --- a/src/openhuman/memory_tree/retrieval/source.rs +++ b/src/openhuman/memory_tree/retrieval/source.rs @@ -1,47 +1,14 @@ -//! `memory_tree_query_source` — retrieve summary hits from per-source trees -//! (Phase 4 / #710). -//! -//! Three selection modes, in priority order: -//! 1. `source_id` Some → one tree lookup via `(kind=source, scope=source_id)` -//! 2. `source_kind` Some → every source tree whose scope prefix matches the -//! kind (chat/email/document); scope convention is the chunk's -//! `metadata.source_id` verbatim, which always embeds a platform hint. -//! 3. Neither → every source tree -//! -//! For each tree we pull the current root (if any) plus all level-1 -//! summaries. If the caller supplied `time_window_days`, we keep only -//! summaries whose `time_range_[start,end]` overlaps `[now - window, now]`. -//! Results are sorted by `time_range_end DESC` so newest-first, then -//! truncated to `limit`. -//! -//! This is deliberately a thin read-only view over `mem_tree_trees` and -//! `mem_tree_summaries`; no new indexes or tables are introduced. - use anyhow::Result; -use chrono::{Duration, Utc}; use crate::openhuman::config::Config; +use crate::openhuman::memory::source_scope::current_source_scope; use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_store::content::read as content_read; -use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind}; -use crate::openhuman::memory_tree::retrieval::types::{ - hit_from_summary, QueryResponse, RetrievalHit, -}; -use crate::openhuman::memory_tree::score::embed::{build_embedder_from_config, cosine_similarity}; -use crate::openhuman::memory_tree::tree::store; +use crate::openhuman::memory_tree::retrieval::engine::{config as engine_config, EmbedderBridge}; +use crate::openhuman::memory_tree::retrieval::types::QueryResponse; +use crate::openhuman::memory_tree::score::embed::build_embedder_from_config; const DEFAULT_LIMIT: usize = 10; -/// Public entrypoint for the tool. All parameters are optional except -/// `limit`, which defaults to 10 when 0. Blocking SQLite work is isolated -/// on `spawn_blocking` so the async caller stays on its runtime. -/// -/// When `query` is `Some`, hits are reranked by cosine similarity between -/// the query embedding and each candidate summary's stored embedding. -/// Candidates with NULL embeddings (pre-Phase-4 legacy rows) fall to the -/// bottom rather than being excluded — callers can still see them, just -/// after all semantically scored rows. When `query` is `None`, the classic -/// newest-first ordering applies. pub async fn query_source( config: &Config, source_id: Option<&str>, @@ -51,882 +18,46 @@ pub async fn query_source( limit: usize, ) -> Result { let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; - // Redact `source_id` — can be a workspace scope like `slack:#` - // that leaks organisational structure. Log only presence + kind filter. + let scope = current_source_scope(); + if source_id.is_some_and(|id| scope.as_ref().is_some_and(|set| !set.contains(id))) { + log::debug!("[retrieval::source] explicit source excluded by active scope"); + return Ok(QueryResponse::empty()); + } + log::debug!( - "[retrieval::source] query_source has_source_id={} source_kind={:?} window_days={:?} has_query={} limit={}", - source_id.is_some(), - source_kind.map(|k| k.as_str()), - time_window_days, - query.is_some(), - limit + "[retrieval::source] tinycortex query has_source_id={} source_kind={:?} window_days={:?} has_query={} limit={}", + source_id.is_some(), source_kind.map(|k| k.as_str()), time_window_days, query.is_some(), limit ); - - let source_id_owned = source_id.map(|s| s.to_string()); - let config_owned = config.clone(); - // Capture the per-profile memory-source allowlist HERE, in the async task - // that holds the `source_scope` task-local — `spawn_blocking` below runs on - // a separate thread that does NOT inherit task-locals, so we must thread it - // through explicitly. `None` = unrestricted. - let source_scope = crate::openhuman::memory::source_scope::current_source_scope(); - // We need the full SummaryNode (with embedding) when semantic rerank - // is on, so return both shapes from the blocking path. - let (hits, scored_nodes) = tokio::task::spawn_blocking( - move || -> Result<(Vec, Vec<(SummaryNode, String)>)> { - collect_hits_and_nodes( - &config_owned, - source_id_owned.as_deref(), - source_kind, - source_scope.as_ref(), - ) - }, - ) - .await - .map_err(|e| anyhow::anyhow!("query_source join error: {e}"))??; - - let filtered = if let Some(days) = time_window_days { - filter_by_window(hits, days) + let semantic_query = query.filter(|value| !value.trim().is_empty()); + let mut response = if let Some(query) = semantic_query { + let embedder = build_embedder_from_config(config)?; + let bridge = EmbedderBridge(embedder.as_ref()); + tinycortex::memory::retrieval::query_source( + &engine_config(config), + source_id, + source_kind, + time_window_days, + Some(query), + &bridge, + usize::MAX, + ) + .await? } else { - hits + tinycortex::memory::retrieval::query_source( + &engine_config(config), + source_id, + source_kind, + time_window_days, + None, + &tinycortex::memory::score::embed::InertEmbedder::new(), + usize::MAX, + ) + .await? }; - let total = filtered.len(); - - let sorted = if let Some(q) = query { - rerank_by_semantic_similarity(config, q, filtered, &scored_nodes).await? - } else { - let mut recency = filtered; - recency.sort_by(|a, b| b.time_range_end.cmp(&a.time_range_end)); - recency - }; - let mut sorted = sorted; - sorted.truncate(limit); - - log::debug!( - "[retrieval::source] returning hits={} total={}", - sorted.len(), - total - ); - Ok(QueryResponse::new(sorted, total)) -} - -/// Blocking helper: walk `mem_tree_trees` + `mem_tree_summaries` and gather -/// every summary under the selected source trees. -/// -/// Returns both the hit shape (for the final response) and the raw -/// `(SummaryNode, tree_scope)` pairs so the async path can read -/// embeddings during semantic rerank without a second DB round-trip. -fn collect_hits_and_nodes( - config: &Config, - source_id: Option<&str>, - source_kind: Option, - source_scope: Option<&std::collections::HashSet>, -) -> Result<(Vec, Vec<(SummaryNode, String)>)> { - let trees = select_trees(config, source_id, source_kind, source_scope)?; - log::debug!("[retrieval::source] selected trees n={}", trees.len()); - - let mut hits: Vec = Vec::new(); - let mut nodes: Vec<(SummaryNode, String)> = Vec::new(); - for tree in &trees { - // max_level starts at 0 before the first seal. For an un-sealed - // tree there's nothing to return. - if tree.max_level == 0 && tree.root_id.is_none() { - continue; - } - // Pull root (highest level) + all L1 summaries. L1 is always the - // finest-grained summary layer above raw leaves. - for level in 1..=tree.max_level { - let level_nodes = store::list_summaries_at_level(config, &tree.id, level)?; - for mut node in level_nodes { - // Hydrate the full body from disk — `node.content` is a - // ≤500-char preview after the MD-on-disk migration. Callers - // (including the LLM) must receive the complete summary text. - // Non-fatal fallback for pre-MD-migration rows. - match content_read::read_summary_body(config, &node.id) { - Ok(body) => node.content = body, - Err(e) => { - log::warn!( - "[retrieval::source] read_summary_body failed — serving preview: {e:#}" - ); - } - } - hits.push(hit_from_summary(&node, &tree.scope)); - nodes.push((node, tree.scope.clone())); - } - } + if let Some(set) = scope { + response.hits.retain(|hit| set.contains(&hit.tree_scope)); } - - // #1574 read-path: embeddings live in the per-model sidecar - // (`mem_tree_summary_embeddings`); the legacy in-row - // `mem_tree_summaries.embedding` column is left NULL by the write path - // after the per-model cutover. `rerank_by_semantic_similarity` reads - // `SummaryNode.embedding`, so without hydrating from the sidecar every - // summary looks un-embedded and semantic recall silently degrades to - // recency order (identical results for unrelated queries). Hydrate in a - // single batched lookup at the active signature, only for nodes the legacy - // column left empty — so pre-cutover rows that still carry an in-row vector - // are untouched. - let node_count = nodes.len(); - let unembedded: Vec = nodes - .iter() - .filter(|(node, _)| node.embedding.is_none()) - .map(|(node, _)| node.id.clone()) - .collect(); - if !unembedded.is_empty() { - log::debug!( - "[retrieval::source] sidecar hydration: unembedded_count={} node_count={} source_kind={:?}", - unembedded.len(), - node_count, - source_kind.map(|k| k.as_str()) - ); - match store::get_summary_embeddings_batch(config, &unembedded) { - Ok(by_id) => { - let mut hydrated = 0usize; - for (node, _) in nodes.iter_mut() { - if node.embedding.is_none() { - if let Some(vec) = by_id.get(&node.id) { - node.embedding = Some(vec.clone()); - hydrated += 1; - } - } - } - log::debug!( - "[retrieval::source] sidecar hydration done: hydrated_count={} missing_count={} unembedded_count={}", - hydrated, - unembedded.len() - hydrated, - unembedded.len() - ); - } - Err(e) => log::warn!( - "[retrieval::source] sidecar summary-embedding batch read failed \ - (unembedded_count={} node_count={} source_kind={:?}): {e:#} — \ - semantic rerank may degrade to recency", - unembedded.len(), - node_count, - source_kind.map(|k| k.as_str()) - ), - } - } - - Ok((hits, nodes)) -} - -/// Rerank hits by cosine similarity to the query embedding. Hits with no -/// embedding (legacy rows) sort to the bottom, preserving their relative -/// order by `time_range_end DESC` so the unranked tail still looks sane. -async fn rerank_by_semantic_similarity( - config: &Config, - query: &str, - hits: Vec, - scored_nodes: &[(SummaryNode, String)], -) -> Result> { - let embedder = build_embedder_from_config(config)?; - let query_vec = embedder.embed(query).await?; - log::debug!( - "[retrieval::source] query embedded provider={} hits_to_rerank={}", - embedder.name(), - hits.len() - ); - // Build a map node_id -> embedding option for O(n) lookup during sort. - use std::collections::HashMap; - let embedding_by_id: HashMap>> = scored_nodes - .iter() - .map(|(n, _)| (n.id.clone(), n.embedding.clone())) - .collect(); - - // Decorate each hit with (score, has_embedding). `has_embedding=false` - // rows get sorted to the bottom by returning negative infinity so - // they keep their relative recency order below the ranked rows. - let mut decorated: Vec<(f32, bool, RetrievalHit)> = hits - .into_iter() - .map(|h| { - let emb = embedding_by_id.get(&h.node_id).cloned().flatten(); - match emb { - Some(v) => { - let sim = cosine_similarity(&query_vec, &v); - (sim, true, h) - } - None => (f32::NEG_INFINITY, false, h), - } - }) - .collect(); - - decorated.sort_by(|a, b| { - // Rows with embeddings first (stable by similarity DESC, then - // recency DESC); legacy rows last (recency DESC). - match (a.1, b.1) { - (true, false) => std::cmp::Ordering::Less, - (false, true) => std::cmp::Ordering::Greater, - _ => { - b.0.partial_cmp(&a.0) - .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| b.2.time_range_end.cmp(&a.2.time_range_end)) - } - } - }); - - Ok(decorated.into_iter().map(|(_, _, h)| h).collect()) -} - -/// Resolve the set of source trees to scan. `source_id` has priority, then -/// `source_kind` (via scope prefix matching), then "all source trees". -fn select_trees( - config: &Config, - source_id: Option<&str>, - source_kind: Option, - source_scope: Option<&std::collections::HashSet>, -) -> Result> { - // Per-profile memory-source gate: `source_scope` is the active profile's - // allowlist of recallable source scopes (None = unrestricted). - let allowed = |scope: &str| source_scope.is_none_or(|set| set.contains(scope)); - - if let Some(id) = source_id { - // An explicit lookup for a source the active profile didn't allow - // surfaces nothing (fail-closed). - if !allowed(id) { - log::debug!( - "[retrieval::source] source_id={id} blocked by profile memory-source scope" - ); - return Ok(Vec::new()); - } - return match store::get_tree_by_scope(config, TreeKind::Source, id)? { - Some(t) => Ok(vec![t]), - None => { - log::debug!( - "[retrieval::source] no tree for source_id={id} — returning empty list" - ); - Ok(Vec::new()) - } - }; - } - let all = store::list_trees_by_kind(config, TreeKind::Source)?; - // Scope the candidate set to the active profile's allowlist (None = all). - let before = all.len(); - let all: Vec = all.into_iter().filter(|t| allowed(&t.scope)).collect(); - if all.len() != before { - log::debug!( - "[retrieval::source] profile memory-source scope: {before} -> {} tree(s)", - all.len() - ); - } - if let Some(kind) = source_kind { - let prefix = kind.as_str(); - let filtered: Vec = all - .into_iter() - .filter(|t| scope_matches_kind(&t.scope, prefix)) - .collect(); - return Ok(filtered); - } - Ok(all) -} - -/// Map from platform prefix → canonical `SourceKind` (as a string). Consulted -/// by [`scope_matches_kind`] so a scope like `slack:#eng` classifies as a -/// chat source. -/// -/// Centralising the mapping here means adding a new integration only touches -/// one place. Keep this list in sync with the channel/provider registry — -/// CodeRabbit on PR #831 flagged the original hardcoded 4-platform list as -/// silently excluding irc/matrix/mattermost/lark/linq/signal/imessage/ -/// dingtalk/qq chat providers. -const PLATFORM_KINDS: &[(&str, &str)] = &[ - // Chat platforms - ("slack", "chat"), - ("discord", "chat"), - ("telegram", "chat"), - ("whatsapp", "chat"), - ("irc", "chat"), - ("matrix", "chat"), - ("mattermost", "chat"), - ("lark", "chat"), - ("linq", "chat"), - ("signal", "chat"), - ("imessage", "chat"), - ("dingtalk", "chat"), - ("qq", "chat"), - ("teams", "chat"), - ("rocketchat", "chat"), - // Email platforms - ("gmail", "email"), - ("imap", "email"), - ("outlook", "email"), - ("fastmail", "email"), - ("protonmail", "email"), - // Document platforms - ("notion", "document"), - ("linear", "document"), - ("drive", "document"), - ("googledoc", "document"), - ("doc", "document"), - ("dropbox", "document"), - ("onedrive", "document"), - ("confluence", "document"), -]; - -/// Decide whether a tree's `scope` falls under `kind_prefix`. Scope is the -/// chunk's `source_id` verbatim (e.g. `slack:#eng`, `gmail:abc`). We check: -/// - Literal `:` prefix (`chat:`, `email:`, `document:`) -/// - Platform-specific prefix via [`PLATFORM_KINDS`] registry -/// -/// This is inherently heuristic — callers that need exact matching should -/// pass `source_id` directly. -fn scope_matches_kind(scope: &str, kind_prefix: &str) -> bool { - let lower = scope.to_lowercase(); - if lower.starts_with(&format!("{kind_prefix}:")) { - return true; - } - PLATFORM_KINDS - .iter() - .any(|(platform, kind)| *kind == kind_prefix && lower.starts_with(&format!("{platform}:"))) -} - -/// Keep hits whose `[time_range_start, time_range_end]` overlaps the -/// `[now - window_days, now]` window. Open-ended intervals (end == start) -/// still pass if the point falls inside. -fn filter_by_window(hits: Vec, window_days: u32) -> Vec { - let now = Utc::now(); - let window_start = now - Duration::days(window_days as i64); - hits.into_iter() - .filter(|h| h.time_range_end >= window_start && h.time_range_start <= now) - .collect() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider}; - use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree; - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::content as content_store; - use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf, LabelStrategy, LeafRef}; - use chrono::{DateTime, TimeZone}; - use std::sync::Arc; - use tempfile::TempDir; - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): seed_source / ingest triggers seals which embed. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - async fn seed_source(cfg: &Config, scope: &str, ts: DateTime) { - let tree = get_or_create_source_tree(cfg, scope).unwrap(); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).unwrap(); - for seq in 0..2u32 { - let c = Chunk { - id: chunk_id(SourceKind::Chat, scope, seq, "test-content"), - content: format!("payload-{scope}-{seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: scope.into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["eng".into()], - source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))), - path_scope: None, - }, - token_count: crate::openhuman::memory_store::trees::types::INPUT_TOKEN_BUDGET * 6 - / 10, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); - // Stage to disk so `hydrate_leaf_inputs` can read the full body - // via `read_chunk_body` during the seal triggered by `append_leaf`, - // and `collect_hits_and_nodes` can read summary bodies for the API. - let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap(); - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx( - &tx, &staged, - )?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - let leaf = LeafRef { - chunk_id: c.id.clone(), - token_count: crate::openhuman::memory_store::trees::types::INPUT_TOKEN_BUDGET * 6 - / 10, - timestamp: ts, - content: c.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - test_override::with_provider(Arc::clone(&provider), async { - append_leaf(cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - } - } - - #[tokio::test] - async fn query_by_source_id_returns_tree_summaries() { - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - - let resp = query_source(&cfg, Some("slack:#eng"), None, None, None, 10) - .await - .unwrap(); - assert_eq!( - resp.hits.len(), - 1, - "two 6k-token leaves seal into one L1 summary" - ); - assert_eq!(resp.total, 1); - assert!(!resp.truncated); - assert_eq!(resp.hits[0].tree_scope, "slack:#eng"); - assert_eq!(resp.hits[0].level, 1); - } - - #[tokio::test] - async fn query_unknown_source_id_returns_empty() { - let (_tmp, cfg) = test_config(); - let resp = query_source(&cfg, Some("slack:#does-not-exist"), None, None, None, 10) - .await - .unwrap(); - assert!(resp.hits.is_empty()); - assert_eq!(resp.total, 0); - assert!(!resp.truncated); - } - - #[tokio::test] - async fn query_by_source_kind_filters_scopes() { - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - seed_source(&cfg, "gmail:alice@example.com", ts).await; - - let chat_only = query_source(&cfg, None, Some(SourceKind::Chat), None, None, 10) - .await - .unwrap(); - assert_eq!(chat_only.hits.len(), 1); - assert_eq!(chat_only.hits[0].tree_scope, "slack:#eng"); - - let email_only = query_source(&cfg, None, Some(SourceKind::Email), None, None, 10) - .await - .unwrap(); - assert_eq!(email_only.hits.len(), 1); - assert_eq!(email_only.hits[0].tree_scope, "gmail:alice@example.com"); - } - - #[tokio::test] - async fn query_all_source_trees_when_no_filter() { - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - seed_source(&cfg, "gmail:alice@example.com", ts).await; - let resp = query_source(&cfg, None, None, None, None, 10) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 2); - } - - #[tokio::test] - async fn profile_source_scope_restricts_all_query_to_allowlisted_trees() { - use crate::openhuman::memory::source_scope::with_source_scope; - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - seed_source(&cfg, "gmail:alice@example.com", ts).await; - - // Allowlist only the Slack source: the all-trees query drops Gmail. - let resp = with_source_scope(Some(vec!["slack:#eng".into()]), async { - query_source(&cfg, None, None, None, None, 10).await - }) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 1); - assert_eq!(resp.hits[0].tree_scope, "slack:#eng"); - } - - #[tokio::test] - async fn profile_source_scope_blocks_explicit_disallowed_source_id() { - use crate::openhuman::memory::source_scope::with_source_scope; - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - seed_source(&cfg, "gmail:alice@example.com", ts).await; - - // Explicit lookup of a source outside the allowlist surfaces nothing. - let blocked = with_source_scope(Some(vec!["slack:#eng".into()]), async { - query_source(&cfg, Some("gmail:alice@example.com"), None, None, None, 10).await - }) - .await - .unwrap(); - assert!(blocked.hits.is_empty()); - - // The allowlisted source still resolves within the same scope. - let allowed = with_source_scope(Some(vec!["slack:#eng".into()]), async { - query_source(&cfg, Some("slack:#eng"), None, None, None, 10).await - }) - .await - .unwrap(); - assert_eq!(allowed.hits.len(), 1); - } - - #[tokio::test] - async fn no_source_scope_leaves_recall_unrestricted() { - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#eng", ts).await; - seed_source(&cfg, "gmail:alice@example.com", ts).await; - // Outside any with_source_scope, all trees remain visible. - let resp = query_source(&cfg, None, None, None, None, 10) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 2); - } - - #[tokio::test] - async fn query_with_time_window_filters_old_hits() { - let (_tmp, cfg) = test_config(); - let ancient = Utc.timestamp_millis_opt(1_000_000_000_000).unwrap(); - seed_source(&cfg, "slack:#ancient", ancient).await; - let recent = Utc::now(); - seed_source(&cfg, "slack:#recent", recent).await; - - let resp = query_source(&cfg, None, None, Some(7), None, 10) - .await - .unwrap(); - assert_eq!( - resp.hits.len(), - 1, - "only the recent tree's summary falls in 7d" - ); - assert_eq!(resp.hits[0].tree_scope, "slack:#recent"); - } - - #[tokio::test] - async fn query_truncates_to_limit() { - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#a", ts).await; - seed_source(&cfg, "slack:#b", ts).await; - seed_source(&cfg, "slack:#c", ts).await; - let resp = query_source(&cfg, None, None, None, None, 2).await.unwrap(); - assert_eq!(resp.hits.len(), 2); - assert_eq!(resp.total, 3); - assert!(resp.truncated); - } - - #[tokio::test] - async fn query_orders_newest_first() { - let (_tmp, cfg) = test_config(); - let older = Utc::now() - Duration::hours(1); - let newer = Utc::now(); - seed_source(&cfg, "slack:#older", older).await; - seed_source(&cfg, "slack:#newer", newer).await; - let resp = query_source(&cfg, None, None, None, None, 10) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 2); - assert_eq!(resp.hits[0].tree_scope, "slack:#newer"); - assert_eq!(resp.hits[1].tree_scope, "slack:#older"); - } - - #[test] - fn scope_prefix_matching_known_platforms() { - assert!(scope_matches_kind("slack:#eng", "chat")); - assert!(scope_matches_kind("gmail:alice", "email")); - assert!(scope_matches_kind("notion:page123", "document")); - assert!(scope_matches_kind("linear:conn-1:issue-abc", "document")); - assert!(!scope_matches_kind("slack:#eng", "email")); - assert!(scope_matches_kind("chat:custom", "chat")); - } - - #[test] - fn zero_limit_defaults_to_ten() { - // Guards against callers passing usize::MIN and quietly getting empty - // results. DEFAULT_LIMIT is the documented default surface. - assert_eq!(DEFAULT_LIMIT, 10); - } - - // ── Phase 4 (#710): semantic rerank tests ─────────────────────── - - /// Hand-craft two source trees whose L1 summaries carry specific - /// embeddings, then verify that providing a `query` string whose - /// embedding matches one tree's direction pushes that tree's hit - /// to the top. Uses a deterministic embedder that returns a - /// direction derived from the input text's first word — no Ollama, - /// no inert zeros (which would make every similarity tie). - /// - /// We override the store's summary embeddings directly after seal so - /// the test doesn't depend on the inert-embedder zero vectors that - /// the ingest path writes by default. - #[tokio::test] - async fn query_reranks_by_cosine_similarity() { - use crate::openhuman::memory_tree::score::embed::{pack_embedding, EMBEDDING_DIM}; - use crate::openhuman::memory_tree::tree::store as src_store; - - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#phoenix", ts).await; - seed_source(&cfg, "slack:#unrelated", ts).await; - - // Fetch the two summaries and give them orthogonal embeddings: - // - "phoenix" tree: [1, 0, 0, ...] padded to 768 - // - "unrelated" tree: [0, 1, 0, ...] padded to 768 - fn unit_vec(axis: usize) -> Vec { - let mut v = vec![0.0_f32; EMBEDDING_DIM]; - v[axis] = 1.0; - v - } - let phoenix_vec = unit_vec(0); - let unrelated_vec = unit_vec(1); - - // Write directly via raw UPDATE so we replace whatever the - // seal-time inert embedder wrote. - use crate::openhuman::memory_store::chunks::store::with_connection; - let phoenix_tree = src_store::get_tree_by_scope( - &cfg, - crate::openhuman::memory_store::trees::types::TreeKind::Source, - "slack:#phoenix", - ) - .unwrap() - .unwrap(); - let unrelated_tree = src_store::get_tree_by_scope( - &cfg, - crate::openhuman::memory_store::trees::types::TreeKind::Source, - "slack:#unrelated", - ) - .unwrap() - .unwrap(); - let phoenix_summaries = - src_store::list_summaries_at_level(&cfg, &phoenix_tree.id, 1).unwrap(); - let unrelated_summaries = - src_store::list_summaries_at_level(&cfg, &unrelated_tree.id, 1).unwrap(); - assert_eq!(phoenix_summaries.len(), 1); - assert_eq!(unrelated_summaries.len(), 1); - - let phoenix_blob = pack_embedding(&phoenix_vec); - let unrelated_blob = pack_embedding(&unrelated_vec); - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_summaries SET embedding = ?1 WHERE id = ?2", - rusqlite::params![phoenix_blob, &phoenix_summaries[0].id], - ) - .unwrap(); - conn.execute( - "UPDATE mem_tree_summaries SET embedding = ?1 WHERE id = ?2", - rusqlite::params![unrelated_blob, &unrelated_summaries[0].id], - ) - .unwrap(); - Ok(()) - }) - .unwrap(); - - // Override the factory: normally the test config returns an inert - // embedder. We need a non-inert embedder to get a non-zero query - // vector. Since build_embedder_from_config is called internally - // we can't easily inject — so instead we simulate via direct - // rerank using `rerank_by_semantic_similarity` indirectly by - // hand-calling `cosine_similarity` on the known vectors. - // - // The practical test here: construct a hypothetical query - // vector equal to phoenix_vec, then verify that running the - // rerank helper with that vector places phoenix first. - use crate::openhuman::memory_tree::score::embed::cosine_similarity; - let query_vec = phoenix_vec.clone(); - let phoenix_sim = cosine_similarity(&query_vec, &phoenix_vec); - let unrelated_sim = cosine_similarity(&query_vec, &unrelated_vec); - assert!( - phoenix_sim > unrelated_sim, - "query aligned to phoenix must outscore unrelated" - ); - - // And: the test-config embedder is inert so query_source's own - // call to embed(query) will yield zero vector — verify the path - // still returns both hits without panicking. - let resp = query_source( - &cfg, - None, - Some(SourceKind::Chat), - None, - Some("phoenix launch"), - 10, - ) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 2); - // With zero query vector, all cosine scores are 0 and rows with - // embeddings stay ahead of legacy rows — both have embeddings so - // they rank equally; order falls to the tiebreaker on time. - } - - /// A legacy summary (NULL embedding, pre-Phase-4) must fall below - /// summaries that do have embeddings when a `query` is supplied. - #[tokio::test] - async fn legacy_null_embedding_rows_sort_last() { - use crate::openhuman::memory_store::trees::types::TreeKind; - use crate::openhuman::memory_tree::score::embed::{pack_embedding, EMBEDDING_DIM}; - use crate::openhuman::memory_tree::tree::store as src_store; - - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#with-embedding", ts).await; - seed_source(&cfg, "slack:#legacy-null", ts).await; - - // Overwrite one tree's summary to have a real unit-vector embedding, - // and explicitly NULL out the other's to mimic a pre-Phase-4 row. - let a = src_store::get_tree_by_scope(&cfg, TreeKind::Source, "slack:#with-embedding") - .unwrap() - .unwrap(); - let b = src_store::get_tree_by_scope(&cfg, TreeKind::Source, "slack:#legacy-null") - .unwrap() - .unwrap(); - let a_sum = src_store::list_summaries_at_level(&cfg, &a.id, 1).unwrap(); - let b_sum = src_store::list_summaries_at_level(&cfg, &b.id, 1).unwrap(); - assert_eq!(a_sum.len(), 1); - assert_eq!(b_sum.len(), 1); - - let mut v = vec![0.0_f32; EMBEDDING_DIM]; - v[0] = 1.0; - let blob = pack_embedding(&v); - - use crate::openhuman::memory_store::chunks::store::with_connection; - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_summaries SET embedding = ?1 WHERE id = ?2", - rusqlite::params![blob, &a_sum[0].id], - ) - .unwrap(); - conn.execute( - "UPDATE mem_tree_summaries SET embedding = NULL WHERE id = ?1", - rusqlite::params![&b_sum[0].id], - ) - .unwrap(); - Ok(()) - }) - .unwrap(); - - let resp = query_source( - &cfg, - None, - Some(SourceKind::Chat), - None, - Some("any query here"), - 10, - ) - .await - .unwrap(); - assert_eq!(resp.hits.len(), 2); - // The embedded row must come before the NULL one. - assert_eq!(resp.hits[0].tree_scope, "slack:#with-embedding"); - assert_eq!(resp.hits[1].tree_scope, "slack:#legacy-null"); - } - - /// #1574 read-path regression: when summary embeddings live ONLY in the - /// per-model sidecar (`mem_tree_summary_embeddings`) and the legacy in-row - /// `mem_tree_summaries.embedding` column is NULL — the post-cutover state — - /// `collect_hits_and_nodes` must hydrate `SummaryNode.embedding` from the - /// sidecar so the rerank can score by query similarity. Without the fix the - /// in-row column is NULL, every node looks un-embedded, and recall collapses - /// to one recency order for all queries. - #[tokio::test] - async fn sidecar_only_embeddings_hydrate_for_rerank() { - use crate::openhuman::memory_store::chunks::store::{ - tree_active_signature, with_connection, - }; - use crate::openhuman::memory_tree::score::embed::{cosine_similarity, EMBEDDING_DIM}; - - let (_tmp, cfg) = test_config(); - let ts = Utc::now(); - seed_source(&cfg, "slack:#alpha", ts).await; - seed_source(&cfg, "slack:#beta", ts).await; - - let alpha_tree = get_or_create_source_tree(&cfg, "slack:#alpha").unwrap(); - let beta_tree = get_or_create_source_tree(&cfg, "slack:#beta").unwrap(); - let alpha = store::list_summaries_at_level(&cfg, &alpha_tree.id, 1).unwrap(); - let beta = store::list_summaries_at_level(&cfg, &beta_tree.id, 1).unwrap(); - assert_eq!(alpha.len(), 1, "alpha tree should seal one L1 summary"); - assert_eq!(beta.len(), 1, "beta tree should seal one L1 summary"); - let alpha_id = alpha[0].id.clone(); - let beta_id = beta[0].id.clone(); - - // Orthogonal unit vectors, written to the SIDECAR only. - fn unit(axis: usize) -> Vec { - let mut v = vec![0.0_f32; EMBEDDING_DIM]; - v[axis] = 1.0; - v - } - let alpha_vec = unit(0); - let beta_vec = unit(1); - let sig = tree_active_signature(&cfg); - store::set_summary_embedding_for_signature(&cfg, &alpha_id, &sig, &alpha_vec).unwrap(); - store::set_summary_embedding_for_signature(&cfg, &beta_id, &sig, &beta_vec).unwrap(); - - // Force the legacy in-row column NULL (the seal path already leaves it - // NULL, but be explicit so the precondition is pinned regardless). - with_connection(&cfg, |conn| { - conn.execute("UPDATE mem_tree_summaries SET embedding = NULL", [])?; - Ok(()) - }) - .unwrap(); - - // (1)+(2) preconditions: in-row NULL, sidecar populated. - with_connection(&cfg, |conn| { - let non_null: i64 = conn.query_row( - "SELECT COUNT(*) FROM mem_tree_summaries WHERE embedding IS NOT NULL", - [], - |r| r.get(0), - )?; - assert_eq!(non_null, 0, "in-row embedding column must be NULL"); - Ok(()) - }) - .unwrap(); - assert!( - store::get_summary_embedding_for_signature(&cfg, &alpha_id, &sig) - .unwrap() - .is_some(), - "sidecar must hold the alpha embedding" - ); - - // (3) the fix: collect_hits_and_nodes hydrates node.embedding from the sidecar. - // 4th arg `source_scope` (added upstream after this PR was opened): None = unrestricted. - let (_hits, nodes) = - collect_hits_and_nodes(&cfg, None, Some(SourceKind::Chat), None).unwrap(); - let emb_of = |id: &str| { - nodes - .iter() - .find(|(n, _)| n.id == id) - .and_then(|(n, _)| n.embedding.clone()) - }; - let alpha_emb = emb_of(&alpha_id).expect("alpha node embedding hydrated from sidecar"); - let beta_emb = emb_of(&beta_id).expect("beta node embedding hydrated from sidecar"); - assert_eq!( - alpha_emb, alpha_vec, - "hydrated vector must equal the sidecar vector" - ); - assert_eq!(beta_emb, beta_vec); - - // (4) different query vectors → different rankings over the hydrated - // embeddings (before the fix both were None and tied at the bottom). - let q_alpha = unit(0); - let q_beta = unit(1); - assert!( - cosine_similarity(&q_alpha, &alpha_emb) > cosine_similarity(&q_alpha, &beta_emb), - "alpha-aligned query must score alpha above beta" - ); - assert!( - cosine_similarity(&q_beta, &beta_emb) > cosine_similarity(&q_beta, &alpha_emb), - "beta-aligned query must score beta above alpha" - ); - } + let total = response.hits.len(); + response.hits.truncate(limit); + Ok(QueryResponse::new(response.hits, total)) } diff --git a/src/openhuman/memory_tree/retrieval/types.rs b/src/openhuman/memory_tree/retrieval/types.rs index 078e9d9ec..91b48faf0 100644 --- a/src/openhuman/memory_tree/retrieval/types.rs +++ b/src/openhuman/memory_tree/retrieval/types.rs @@ -1,263 +1,6 @@ -//! Shared types for Phase 4 retrieval tools (#710). -//! -//! These types are the wire / JSON-RPC shape returned by the six retrieval -//! primitives. They wrap the internal persistence structs (`SummaryNode`, -//! `Chunk`, `EntityHit`) into a single unified [`RetrievalHit`] shape so the -//! calling LLM sees the same schema regardless of which tool it invoked. -//! -//! Rules of the road: -//! - All types are [`serde::Serialize`] + [`serde::Deserialize`] so they -//! round-trip through JSON-RPC without bespoke conversion. -//! - Time fields use `DateTime` serialised as RFC3339 — matches the -//! global recap convention so callers comparing hits across tools don't -//! juggle epochs. -//! - `node_kind` discriminates leaf (raw chunk) vs. summary — retrieval -//! consumers frequently branch on this (e.g. "drill down only on -//! summaries"). +//! Stable host path for tinycortex-owned retrieval wire types and converters. -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -use crate::openhuman::memory_store::chunks::types::{Chunk, SourceKind}; -use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind}; -use crate::openhuman::memory_tree::score::extract::EntityKind; - -/// Whether a hit represents a leaf (raw chunk) or a summary node. -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum NodeKind { - /// Leaf = one `mem_tree_chunks` row (level 0). - Leaf, - /// Summary = one `mem_tree_summaries` row (level ≥ 1 for source/topic, - /// level ≥ 0 for global where L0 is a daily digest). - Summary, -} - -impl NodeKind { - /// Stable lowercase string form (`"leaf"` / `"summary"`) — matches the - /// serde representation and is suitable for SQL discriminator columns. - pub fn as_str(self) -> &'static str { - match self { - Self::Leaf => "leaf", - Self::Summary => "summary", - } - } -} - -/// One unit of retrieval output. Shape is identical whether the hit was -/// sourced from a source tree summary, a topic tree summary, the global -/// digest, or a raw leaf chunk. -/// -/// `tree_id` / `tree_kind` / `tree_scope` identify which tree the hit -/// belongs to so UIs can surface provenance ("from Slack #eng"). For bare -/// leaves not yet attached to any tree, `tree_id` is empty and `tree_kind` -/// is still meaningful (mirrors the leaf's source kind classification — -/// see [`leaf_tree_placeholder`] for how we synthesise it). -/// -/// `child_ids` is empty on leaves; on summaries it points at the next level -/// down (chunks for L1, summaries for L2+). -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct RetrievalHit { - pub node_id: String, - pub node_kind: NodeKind, - pub tree_id: String, - pub tree_kind: TreeKind, - pub tree_scope: String, - pub level: u32, - pub content: String, - pub entities: Vec, - pub topics: Vec, - pub time_range_start: DateTime, - pub time_range_end: DateTime, - pub score: f32, - pub child_ids: Vec, - /// Populated for leaves (chunk back-pointer); `None` for summaries. - pub source_ref: Option, -} - -/// Envelope for the four "query" tools (`query_source`, `query_global`, -/// `query_topic`, `drill_down` by wrapper). -/// -/// `total` is the pre-truncation match count so callers can tell whether a -/// high-limit follow-up would return more. `truncated` is `total > hits.len()`. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct QueryResponse { - pub hits: Vec, - pub total: usize, - pub truncated: bool, -} - -impl QueryResponse { - /// Build a response from a post-filtered, post-sorted hit list. The - /// `total_matches` is the count BEFORE applying `limit` so callers can - /// see whether truncation happened. - pub fn new(hits: Vec, total_matches: usize) -> Self { - let truncated = total_matches > hits.len(); - Self { - hits, - total: total_matches, - truncated, - } - } - - /// Empty response (no matches). `total=0`, `truncated=false`. - pub fn empty() -> Self { - Self { - hits: Vec::new(), - total: 0, - truncated: false, - } - } -} - -/// Convert a sealed [`SummaryNode`] into a [`RetrievalHit`]. `tree_scope` -/// is threaded in from the caller so we don't force a tree lookup on every -/// conversion — the caller already has the parent [`Tree`] in hand. -pub fn hit_from_summary(node: &SummaryNode, tree_scope: &str) -> RetrievalHit { - RetrievalHit { - node_id: node.id.clone(), - node_kind: NodeKind::Summary, - tree_id: node.tree_id.clone(), - tree_kind: node.tree_kind, - tree_scope: tree_scope.to_string(), - level: node.level, - content: node.content.clone(), - entities: node.entities.clone(), - topics: node.topics.clone(), - time_range_start: node.time_range_start, - time_range_end: node.time_range_end, - score: node.score, - child_ids: node.child_ids.clone(), - source_ref: None, - } -} - -/// Convert a sealed [`SummaryNode`] using a full [`Tree`] for the scope. A -/// thin convenience over [`hit_from_summary`]. -pub fn hit_from_summary_with_tree(node: &SummaryNode, tree: &Tree) -> RetrievalHit { - hit_from_summary(node, &tree.scope) -} - -/// Convert a raw [`Chunk`] (leaf) into a [`RetrievalHit`]. Because a chunk -/// may not yet be attached to a summary tree, callers can pass `tree_id` / -/// `tree_scope` as empty strings. `tree_kind` is always [`TreeKind::Source`] -/// for leaves — raw chunks belong conceptually to their originating source -/// tree even when the tree hasn't materialised yet (no seals). -pub fn hit_from_chunk(chunk: &Chunk, tree_id: &str, tree_scope: &str, score: f32) -> RetrievalHit { - let source_ref = chunk.metadata.source_ref.as_ref().map(|r| r.value.clone()); - RetrievalHit { - node_id: chunk.id.clone(), - node_kind: NodeKind::Leaf, - tree_id: tree_id.to_string(), - tree_kind: leaf_tree_placeholder(chunk.metadata.source_kind), - tree_scope: tree_scope.to_string(), - level: 0, - content: chunk.content.clone(), - entities: Vec::new(), - topics: chunk.metadata.tags.clone(), - time_range_start: chunk.metadata.time_range.0, - time_range_end: chunk.metadata.time_range.1, - score, - child_ids: Vec::new(), - source_ref, - } -} - -/// 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 call site explicit about why the classification is stable. -pub fn leaf_tree_placeholder(_source_kind: SourceKind) -> TreeKind { - TreeKind::Source -} - -/// Output shape for `search_entities`. One row per canonical id with the -/// aggregate stats across the entity index. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct EntityMatch { - /// Canonical id (e.g. `email:alice@example.com`, `topic:phoenix`). - pub canonical_id: String, - pub kind: EntityKind, - /// Example surface form that matched — useful for UI display. - pub surface: String, - /// Total rows in `mem_tree_entity_index` grouped under this canonical id. - pub mention_count: u64, - /// Epoch-millis of the newest mention across all rows. - pub last_seen_ms: i64, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn node_kind_as_str_round_trips() { - assert_eq!(NodeKind::Leaf.as_str(), "leaf"); - assert_eq!(NodeKind::Summary.as_str(), "summary"); - } - - #[test] - fn query_response_truncated_when_total_exceeds_hits() { - let hit = sample_hit(); - let resp = QueryResponse::new(vec![hit.clone()], 5); - assert_eq!(resp.hits.len(), 1); - assert_eq!(resp.total, 5); - assert!(resp.truncated); - } - - #[test] - fn query_response_not_truncated_when_all_returned() { - let hit = sample_hit(); - let resp = QueryResponse::new(vec![hit.clone()], 1); - assert!(!resp.truncated); - } - - #[test] - fn query_response_empty_is_inert() { - let resp = QueryResponse::empty(); - assert!(resp.hits.is_empty()); - assert_eq!(resp.total, 0); - assert!(!resp.truncated); - } - - #[test] - fn retrieval_hit_serde_round_trip() { - let hit = sample_hit(); - let json = serde_json::to_string(&hit).unwrap(); - let back: RetrievalHit = serde_json::from_str(&json).unwrap(); - assert_eq!(back, hit); - } - - #[test] - fn entity_match_serde_round_trip() { - let m = EntityMatch { - canonical_id: "email:alice@example.com".into(), - kind: EntityKind::Email, - surface: "alice@example.com".into(), - mention_count: 7, - last_seen_ms: 1_700_000_000_000, - }; - let json = serde_json::to_string(&m).unwrap(); - let back: EntityMatch = serde_json::from_str(&json).unwrap(); - assert_eq!(back, m); - } - - fn sample_hit() -> RetrievalHit { - RetrievalHit { - node_id: "sum-1".into(), - node_kind: NodeKind::Summary, - tree_id: "tree-1".into(), - tree_kind: TreeKind::Source, - tree_scope: "slack:#eng".into(), - level: 1, - content: "the sealed summary content".into(), - entities: vec!["email:alice@example.com".into()], - topics: vec!["#launch".into()], - time_range_start: Utc::now(), - time_range_end: Utc::now(), - score: 0.75, - child_ids: vec!["leaf-a".into(), "leaf-b".into()], - source_ref: None, - } - } -} +pub use tinycortex::memory::retrieval::{ + hit_from_chunk, hit_from_summary, hit_from_summary_with_tree, leaf_tree_placeholder, + EntityMatch, NodeKind, QueryResponse, RetrievalHit, +}; diff --git a/src/openhuman/memory_tree/score/extract/extractor.rs b/src/openhuman/memory_tree/score/extract/extractor.rs deleted file mode 100644 index 2e05e6f3b..000000000 --- a/src/openhuman/memory_tree/score/extract/extractor.rs +++ /dev/null @@ -1,111 +0,0 @@ -//! [`EntityExtractor`] trait plus the regex and composite implementations -//! used as Phase 2's default extraction stack. - -use async_trait::async_trait; - -use super::regex; -use super::types::ExtractedEntities; - -/// Interface for anything that can read a chunk's text and emit entities. -#[async_trait] -pub trait EntityExtractor: Send + Sync { - /// Human-readable name for logs and diagnostics. - fn name(&self) -> &'static str; - - /// Run extraction. Implementations should be idempotent per input. - async fn extract(&self, text: &str) -> anyhow::Result; -} - -/// Synchronous regex extractor adapted to the async [`EntityExtractor`] trait. -pub struct RegexEntityExtractor; - -#[async_trait] -impl EntityExtractor for RegexEntityExtractor { - fn name(&self) -> &'static str { - "regex" - } - - async fn extract(&self, text: &str) -> anyhow::Result { - Ok(regex::extract(text)) - } -} - -/// Runs a sequence of extractors and merges their results. -/// -/// An extractor returning an error is logged and skipped — one bad extractor -/// does not abort ingestion. -pub struct CompositeExtractor { - inner: Vec>, -} - -impl CompositeExtractor { - /// Build a composite from an explicit list of extractors. Order matters - /// only to logs — outputs are merged and deduplicated. - pub fn new(inner: Vec>) -> Self { - Self { inner } - } - - /// Convenience constructor: regex-only (the Phase 2 default). - pub fn regex_only() -> Self { - Self::new(vec![Box::new(RegexEntityExtractor)]) - } -} - -#[async_trait] -impl EntityExtractor for CompositeExtractor { - fn name(&self) -> &'static str { - "composite" - } - - async fn extract(&self, text: &str) -> anyhow::Result { - let mut out = ExtractedEntities::default(); - for ex in &self.inner { - match ex.extract(text).await { - Ok(batch) => out.merge(batch), - Err(e) => { - log::warn!( - "[memory_tree::extract] extractor `{}` failed: {e} — continuing", - ex.name() - ); - } - } - } - Ok(out) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_tree::score::extract::EntityKind; - - #[tokio::test] - async fn regex_only_extractor_works() { - let c = CompositeExtractor::regex_only(); - let out = c.extract("hi @alice a@b.com #launch").await.unwrap(); - assert!(out.entities.iter().any(|e| e.kind == EntityKind::Handle)); - assert!(out.entities.iter().any(|e| e.kind == EntityKind::Email)); - assert!(out.entities.iter().any(|e| e.kind == EntityKind::Hashtag)); - } - - struct FailingExtractor; - #[async_trait] - impl EntityExtractor for FailingExtractor { - fn name(&self) -> &'static str { - "failing" - } - async fn extract(&self, _: &str) -> anyhow::Result { - Err(anyhow::anyhow!("boom")) - } - } - - #[tokio::test] - async fn composite_survives_one_failing_extractor() { - let c = CompositeExtractor::new(vec![ - Box::new(FailingExtractor), - Box::new(RegexEntityExtractor), - ]); - let out = c.extract("@alice").await.unwrap(); - assert!(out.entities.iter().any(|e| e.kind == EntityKind::Handle)); - } -} diff --git a/src/openhuman/memory_tree/score/extract/llm.rs b/src/openhuman/memory_tree/score/extract/llm.rs deleted file mode 100644 index 842ade38b..000000000 --- a/src/openhuman/memory_tree/score/extract/llm.rs +++ /dev/null @@ -1,661 +0,0 @@ -//! LLM-based entity + importance extractor. -//! -//! Builds a (system, user) prompt asking for NER + an importance rating -//! in one structured-JSON response, hands the prompt to a -//! [`ChatProvider`], and parses the result into [`ExtractedEntities`]. -//! -//! ## Why this lives here -//! -//! Phase 2 ships a regex extractor only. Semantic NER (Person/Org/Loc/…) -//! requires a model. Originally we used a small local LLM (Ollama default: -//! `qwen2.5:0.5b`) because openhuman already ran Ollama for embeddings. -//! After the cloud-default refactor, the same prompt now routes through -//! whichever backend the workspace selected — typically the OpenHuman -//! backend's `summarization-v1`. The extractor itself is unchanged below the -//! HTTP layer; only the transport moved. -//! -//! ## Span recovery -//! -//! LLMs are unreliable about character offsets. We re-find each returned -//! entity surface in the source text via `text.find(...)` to recover spans. -//! Entities whose surface form can't be located in the source text are -//! dropped with a warn log (this catches model hallucinations). -//! -//! ## Soft fallback -//! -//! If the chat call fails (provider unavailable, malformed JSON, …), we -//! log a warn and return [`ExtractedEntities::default()`]. The -//! [`super::CompositeExtractor`] already tolerates errors from individual -//! extractors; ingestion never blocks on LLM availability. - -use std::sync::Arc; - -use async_trait::async_trait; -use serde::Deserialize; - -use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; -use super::EntityExtractor; -use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - -/// Output-token cap requested for an extraction call (`max_tokens` on the -/// wire). An extraction response is one small structured-JSON object -/// (an entities array + an importance float + a one-sentence reason), so a -/// few thousand tokens is generous. The cap exists less to bound generation -/// than to keep credit-metered providers (OpenRouter) from reserving the -/// model's *entire* output window during their balance pre-flight: with no -/// `max_tokens`, OpenRouter priced extraction against the full 64k+ window -/// and returned `402 Payment Required` to any BYO user whose balance couldn't -/// cover it, on every call (TAURI-RUST-C62). 8192 is well under any plausible -/// extraction output yet small enough to clear the pre-flight for a -/// low-balance account. -const EXTRACTION_MAX_OUTPUT_TOKENS: u32 = 8192; - -// ── Configuration ──────────────────────────────────────────────────────── - -/// Configuration for [`LlmEntityExtractor`]. -#[derive(Clone, Debug)] -pub struct LlmExtractorConfig { - /// Model identifier the chat provider should target. For cloud this - /// is e.g. `summarization-v1`; for local Ollama it's the Ollama tag - /// (`qwen2.5:0.5b`). Threaded through to [`ChatPrompt`] so the - /// provider can route to the right model. - /// - /// Stored on the extractor for diagnostic logging only — the actual - /// model selection happens inside the [`ChatProvider`]. - pub model: String, - /// Which entity kinds the LLM is allowed to emit. Anything outside this - /// set is mapped to [`EntityKind::Misc`] or dropped depending on - /// `strict_kinds`. - pub allowed_kinds: Vec, - /// If true, drop entities whose declared kind isn't in `allowed_kinds` - /// instead of falling back to [`EntityKind::Misc`]. - pub strict_kinds: bool, - /// If true, the system prompt asks the model to also emit a - /// `topics` array (free-form theme labels), and the response parser - /// populates [`ExtractedEntities::topics`]. Default `false` — the - /// extractor's primary job is named-entity extraction; topics are - /// an opt-in side-channel for callers that need a thematic - /// summary in the same call (e.g. running over a sealed summary's - /// content). Adds prompt tokens and gives the model one more - /// schema field to keep track of, so leave off unless needed. - pub emit_topics: bool, - /// Optional configured output language for natural-language values such - /// as `importance_reason` and topic labels. JSON field names and enum - /// values remain stable. - pub output_language: Option, -} - -impl Default for LlmExtractorConfig { - fn default() -> Self { - Self { - model: "qwen2.5:0.5b".to_string(), - allowed_kinds: vec![ - EntityKind::Person, - EntityKind::Organization, - EntityKind::Location, - EntityKind::Event, - EntityKind::Product, - EntityKind::Datetime, - EntityKind::Technology, - EntityKind::Artifact, - EntityKind::Quantity, - ], - strict_kinds: false, - emit_topics: false, - output_language: None, - } - } -} - -// ── Extractor ──────────────────────────────────────────────────────────── - -/// LLM-backed entity + importance extractor. -/// -/// Holds an `Arc` rather than a per-instance HTTP -/// client. The provider abstraction lets a single workspace choose -/// cloud vs local at runtime (see -/// [`crate::openhuman::memory::chat::build_chat_provider`]). Tests -/// can mock the provider to assert the prompt / parse behaviour without -/// a real Ollama or backend. -pub struct LlmEntityExtractor { - cfg: LlmExtractorConfig, - provider: Arc, -} - -impl LlmEntityExtractor { - /// Build the extractor with the supplied chat provider. Infallible — - /// the caller is responsible for provider construction. - pub fn new(cfg: LlmExtractorConfig, provider: Arc) -> Self { - Self { cfg, provider } - } - - /// Build the chat prompt sent to the provider for `text`. - fn build_prompt(&self, text: &str) -> ChatPrompt { - ChatPrompt { - system: build_system_prompt(self.cfg.emit_topics, self.cfg.output_language.as_deref()), - user: format!("Text:\n{text}\n\nReturn JSON only."), - temperature: 0.0, - kind: "memory_tree::extract", - max_tokens: Some(EXTRACTION_MAX_OUTPUT_TOKENS), - } - } -} - -#[async_trait] -impl EntityExtractor for LlmEntityExtractor { - fn name(&self) -> &'static str { - "llm-ollama" - } - - async fn extract(&self, text: &str) -> anyhow::Result { - // Soft-fallback contract: every failure path (transport, HTTP status, - // JSON parse) is logged as a warn and returns an empty - // `ExtractedEntities` rather than `Err`. This makes the extractor - // safe to call from any context, not just `score_chunk` (which - // separately catches errors from its own extractor chain). - // - // Transport failures get bounded retry-with-backoff before falling - // back to empty — see [`Self::try_extract`]. Non-transport failures - // (HTTP non-success, malformed JSON) fall back immediately because - // retrying the same input would yield the same bad response. - const MAX_ATTEMPTS: u32 = 3; - const BASE_BACKOFF_MS: u64 = 250; - - for attempt in 0..MAX_ATTEMPTS { - match self.try_extract(text).await { - AttemptOutcome::Done(extracted) => { - // #3365 + #002 (T013): the structure-degraded latch means - // "the extraction model is timing out / unreachable" — set - // ONLY after MAX_ATTEMPTS transport failures (below). A - // *completed* call disproves that: `try_extract` returns - // `Done` whenever the provider responded, including the - // valid-but-empty and malformed-JSON-as-empty cases. So a - // completed extraction clears the latch regardless of whether - // this particular text yielded entities — liveness, not - // content, is what it tracks. - // - // (Clearing only on a non-empty result left the latch stuck - // after the model recovered whenever later texts were - // entity-light, so the surface kept warning "extraction model - // is timing out" long after it had stopped — #3365.) - crate::openhuman::memory_tree::health::clear_structure_degraded(); - return Ok(extracted); - } - AttemptOutcome::Permanent(code) => { - // Non-retryable client error — e.g. a 402 (the BYO - // provider account is out of credits), a rejected API - // key, or a model that no longer exists. Retrying - // reproduces the identical error and, on a credit-metered - // provider, fires one more Sentry event per attempt - // (TAURI-RUST-C62: 12k events from a single user). Skip - // the backoff loop entirely and fall back to the same - // degraded-but-empty result the exhausted-retry path - // returns — but tag the cause precisely so doctor/status - // can steer the user ("top up credits" / "fix the key") - // instead of "extraction is timing out". - log::warn!( - "[memory_tree::extract::llm] non-retryable provider failure ({}) — \ - returning empty extraction without retry (structure degraded)", - code.as_str() - ); - crate::openhuman::memory_tree::health::mark_structure_degraded(code); - return Ok(ExtractedEntities::default()); - } - AttemptOutcome::Retryable => { - // Transport / transient failure. Retry with exponential - // backoff unless we've exhausted attempts. - if attempt + 1 < MAX_ATTEMPTS { - let delay_ms = BASE_BACKOFF_MS * 2u64.pow(attempt); - log::warn!( - "[memory_tree::extract::llm] transport failure, retrying in \ - {delay_ms}ms (attempt {}/{})", - attempt + 2, - MAX_ATTEMPTS - ); - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - } - } - } - } - - // #002 (T013): every attempt hit a transport failure (the model - // timed out / was unreachable). The soft-fallback contract still - // returns empty (ingest never blocks on a slow model), but we now - // record a structure-degraded signal with a typed cause so the - // status/doctor surface can say "extraction is timing out — switch - // the Memory extraction model" instead of presenting an empty wiki - // as success. - log::warn!( - "[memory_tree::extract::llm] transport failed after {} attempts — \ - returning empty extraction (structure degraded)", - MAX_ATTEMPTS - ); - crate::openhuman::memory_tree::health::mark_structure_degraded( - crate::openhuman::memory_tree::health::FailureCode::ExtractionTimeout, - ); - Ok(ExtractedEntities::default()) - } -} - -/// Outcome of a single [`LlmEntityExtractor::try_extract`] attempt. -/// -/// Splits the old `Option` ("`None` = retry") into three cases so the retry -/// loop can tell a transient blip apart from a permanent client error and -/// stop hammering the latter (TAURI-RUST-C62). -enum AttemptOutcome { - /// The call completed (provider returned content). Includes the - /// "malformed wrong-shape JSON → empty" case, because re-sending the - /// same input won't change the response. - Done(ExtractedEntities), - /// Transport / transient failure (unreachable backend, 5xx, truncated - /// mid-JSON). Retrying may help. - Retryable, - /// Permanent, non-retryable provider failure (4xx client error: 402 out - /// of credits, rejected key, model gone). Retrying reproduces the same - /// error. Carries the typed [`FailureCode`] to record on the degraded - /// signal so doctor/status can steer the user. - Permanent(crate::openhuman::memory_tree::health::FailureCode), -} - -impl LlmEntityExtractor { - /// Internal: one attempt at calling the chat provider. - async fn try_extract(&self, text: &str) -> AttemptOutcome { - let prompt = self.build_prompt(text); - log::debug!( - "[memory_tree::extract::llm] chat provider={} model={} text_chars={}", - self.provider.name(), - self.cfg.model, - text.chars().count() - ); - - let raw = match self.provider.chat_for_json(&prompt).await { - Ok(v) => v, - Err(e) => { - // Classify against the provider stack's single source of - // truth. A non-retryable client error (402/401/403/404/…) - // must not be retried — that is the 12k-event amplifier in - // TAURI-RUST-C62. Only genuine transport/transient failures - // earn the backoff loop. - if crate::openhuman::inference::provider::reliable::is_non_retryable(&e) { - log::warn!( - "[memory_tree::extract::llm] chat provider={} permanently failed: {e:#}", - self.provider.name() - ); - return AttemptOutcome::Permanent(permanent_failure_code(&e)); - } - log::warn!( - "[memory_tree::extract::llm] chat provider={} failed: {e:#}", - self.provider.name() - ); - return AttemptOutcome::Retryable; - } - }; - log::debug!( - "[memory_tree::extract::llm] response chars={} provider={}", - raw.len(), - self.provider.name() - ); - - let parsed: LlmExtractionOutput = match serde_json::from_str(&raw) { - Ok(v) => v, - Err(e) if e.is_eof() => { - // Truncated mid-JSON: the response stream closed before the - // closing brace (token budget or a server-side timeout cut - // it short). Unlike a wrong-shape body, this is transient — - // signal a retryable failure so `extract`'s backoff loop - // tries again rather than silently dropping the whole chunk - // (bug-report-2026-05-26 I1). - log::warn!( - "[memory_tree::extract::llm] LLM response truncated mid-JSON ({e}); \ - response_bytes={} — retrying", - raw.len() - ); - return AttemptOutcome::Retryable; - } - Err(e) => { - log::warn!( - "[memory_tree::extract::llm] LLM returned non-JSON or wrong-shape \ - response: {e}; content was: {} — returning empty extraction", - truncate_for_log(&raw, 400) - ); - return AttemptOutcome::Done(ExtractedEntities::default()); - } - }; - - AttemptOutcome::Done(parsed.into_extracted_entities(text, &self.cfg)) - } -} - -/// Map a permanent (non-retryable) extraction provider error to the most -/// honest *existing* [`FailureCode`], so the degraded-structure signal steers -/// the user without minting a new variant (and its 14-locale remediation key). -/// -/// - 402 / "payment required" / "requires more credits" / "insufficient" → -/// [`FailureCode::BudgetExhausted`] (its remediation already reads "out of -/// budget; bring your own key or top up — retrying won't help"), the -/// dominant TAURI-RUST-C62 case. -/// - 401 / 403 / auth-key rejection → [`FailureCode::AuthInvalid`]. -/// - any other permanent failure → [`FailureCode::ExtractionTimeout`], the -/// existing "extraction model isn't producing" bucket. -fn permanent_failure_code( - err: &anyhow::Error, -) -> crate::openhuman::memory_tree::health::FailureCode { - use crate::openhuman::memory_tree::health::FailureCode; - let lower = format!("{err:#}").to_lowercase(); - if lower.contains("402") - || lower.contains("payment required") - || lower.contains("requires more credits") - || lower.contains("insufficient") - // Monthly-quota exhaustion (e.g. Kiro `MONTHLY_REQUEST_COUNT`) — the - // budget-exhausted remediation ("out of budget; top up — retrying won't - // help") is the honest signal even when the proxy omits a 402 status - // (TAURI-RUST-C9A). - || lower.contains("monthly_request_count") - || lower.contains("monthly request") - || lower.contains("quota") - { - FailureCode::BudgetExhausted - } else if lower.contains("401") - || lower.contains("403") - || lower.contains("unauthorized") - || lower.contains("forbidden") - || lower.contains("invalid api key") - || lower.contains("incorrect api key") - { - FailureCode::AuthInvalid - } else { - FailureCode::ExtractionTimeout - } -} - -// ── Prompt ─────────────────────────────────────────────────────────────── - -/// Build the system prompt for the extractor. When `emit_topics` is true -/// the schema, required-fields list, and example outputs include a -/// `topics` array (free-form theme labels). When false the prompt -/// matches the pre-flag behaviour exactly — no mention of topics -/// anywhere — so the small model isn't asked to produce a field the -/// caller doesn't want. -fn build_system_prompt(emit_topics: bool, output_language: Option<&str>) -> String { - let topics_schema_line = if emit_topics { - " \"topics\": [\"\"],\n" - } else { - "" - }; - let topics_required = if emit_topics { "topics, " } else { "" }; - let fields_count = if emit_topics { "four" } else { "three" }; - let topics_guide = if emit_topics { - "Topics are short free-form theme labels for what the text is ABOUT \ - (e.g. \"rate limiting\", \"memory tree\", \"auth flow\"). They are \ - distinct from entities — entities are specific named things mentioned \ - in the text; topics are the abstract themes those things relate to.\n" - } else { - "" - }; - let example1_topics = if emit_topics { - ",\"topics\":[\"shipping\",\"auth\"]" - } else { - "" - }; - let example2_topics = if emit_topics { - ",\"topics\":[\"product launch\",\"revenue\"]" - } else { - "" - }; - let language_directive = crate::openhuman::config::output_language_directive(output_language) - .map(|directive| format!("{directive}\n\n")) - .unwrap_or_default(); - - format!( - "{language_directive}You are a named-entity extractor and importance rater. Return JSON only — \ -no prose, no markdown, no commentary. Do not summarize. Extract every named \ -entity mention you find, including duplicates, and rate the chunk's overall \ -importance as a float in [0.0, 1.0]. - -Schema: -{{ - \"entities\": [ - {{ \"kind\": \"person|organization|location|event|product|datetime|technology|artifact|quantity\", - \"text\": \"\" }} - ], -{topics_schema_line} \"importance\": 0.0, - \"importance_reason\": \"\" -}} - -Kinds guide: - person named human (\"Alice\", \"Steven Enamakel\") - organization company / team / project (\"Anthropic\", \"TinyHumans\") - location place (\"SF office\", \"London\") - event scheduled occurrence (\"Q2 launch\", \"design review\") - product commercial offering (\"Claude Code\", \"OpenHuman\") - datetime temporal expression (\"Friday\", \"Q2 2026\", \"EOD tomorrow\") - technology tool / framework / language / service (\"Rust\", \"OAuth\", \"Slack API\") - artifact code / ticket / doc reference (\"PR #934\", \"src/foo.rs\", \"OH-42\") - quantity amount / metric / money (\"$5K\", \"20/min\", \"10k tokens\") - -{topics_guide} -If a mention doesn't clearly fit a kind above, omit it rather than guessing. -Always emit ALL {fields_count} top-level fields (entities, {topics_required}importance, importance_reason), -even when entities is empty. - -Examples: - -Input: alice and bob shipped the auth migration friday. PR #42 ships OAuth refactor in src/auth/. -Output: {{\"entities\":[{{\"kind\":\"person\",\"text\":\"alice\"}},{{\"kind\":\"person\",\"text\":\"bob\"}},{{\"kind\":\"event\",\"text\":\"auth migration\"}},{{\"kind\":\"datetime\",\"text\":\"friday\"}},{{\"kind\":\"artifact\",\"text\":\"PR #42\"}},{{\"kind\":\"technology\",\"text\":\"OAuth\"}},{{\"kind\":\"artifact\",\"text\":\"src/auth/\"}}]{example1_topics},\"importance\":0.9,\"importance_reason\":\"explicit shipping commitment\"}} - -Input: Anthropic shipped Claude Code in SF — $20M ARR target by Q2. -Output: {{\"entities\":[{{\"kind\":\"organization\",\"text\":\"Anthropic\"}},{{\"kind\":\"product\",\"text\":\"Claude Code\"}},{{\"kind\":\"location\",\"text\":\"SF\"}},{{\"kind\":\"quantity\",\"text\":\"$20M ARR\"}},{{\"kind\":\"datetime\",\"text\":\"Q2\"}}]{example2_topics},\"importance\":0.85,\"importance_reason\":\"factual content with key business metric\"}} - -Importance guide: - 0.9+ actionable decisions, key information, explicit commitments - 0.6+ substantive discussion, factual content, named entities - 0.3+ ambient context, low-density prose - <0.3 reactions, acknowledgments, bots, trivial exchanges -" - ) -} - -// ── LLM JSON output ────────────────────────────────────────────────────── - -#[derive(Debug, Deserialize)] -struct LlmExtractionOutput { - #[serde(default)] - entities: Vec, - /// Free-form theme labels — populated only when the extractor is - /// configured with `emit_topics = true`. Always tolerant of absence - /// so models that ignore the field don't fail parsing. - #[serde(default)] - topics: Vec, - #[serde(default)] - importance: Option, - #[serde(default)] - importance_reason: Option, -} - -#[derive(Debug, Deserialize)] -struct LlmEntity { - kind: String, - text: String, -} - -impl LlmExtractionOutput { - fn into_extracted_entities( - self, - source_text: &str, - cfg: &LlmExtractorConfig, - ) -> ExtractedEntities { - let mut entities = Vec::with_capacity(self.entities.len()); - - // Per-surface search cursor (char offset). When the LLM returns the - // same surface text twice (deliberately — the prompt asks for - // duplicates), we resume searching AFTER the previous occurrence so - // each emitted entity points at a distinct span. Byte indices are - // tracked separately from char indices because `str::find` returns - // byte offsets while the rest of the pipeline uses char spans. - use std::collections::HashMap; - let mut cursors: HashMap = HashMap::new(); - - for raw in self.entities { - let surface = raw.text.trim(); - if surface.is_empty() { - continue; - } - - let kind = match parse_kind(&raw.kind) { - Some(k) => { - if cfg.allowed_kinds.contains(&k) { - k - } else if cfg.strict_kinds { - log::debug!( - "[memory_tree::extract::llm] dropping entity with disallowed kind: {}", - raw.kind - ); - continue; - } else { - EntityKind::Misc - } - } - None => { - if cfg.strict_kinds { - log::debug!( - "[memory_tree::extract::llm] dropping entity with unknown kind: {}", - raw.kind - ); - continue; - } - EntityKind::Misc - } - }; - - // Recover spans by string search, advancing the cursor for this - // surface so repeated mentions get distinct spans. If the model - // hallucinated a surface (or we've exhausted all of its - // occurrences), drop the entity. - let (byte_from, char_from) = cursors.get(surface).copied().unwrap_or((0, 0)); - let (span_start, span_end, byte_after) = - match find_char_span_from(source_text, surface, byte_from, char_from) { - Some(s) => s, - None => { - log::debug!( - "[memory_tree::extract::llm] dropping hallucinated or exhausted \ - entity (not found beyond cursor): {surface:?}" - ); - continue; - } - }; - cursors.insert(surface.to_string(), (byte_after, span_end)); - - entities.push(ExtractedEntity { - kind, - text: surface.to_string(), - span_start, - span_end, - score: 0.85, // LLM-derived; lower confidence than regex - }); - } - - let llm_importance = self.importance.map(|v| v.clamp(0.0, 1.0)); - - // Topics: only populated when the caller enabled `emit_topics` - // (the prompt asked for them). Otherwise this is empty by - // default — the model didn't know to emit topics, so any value - // here would be hallucination. - let topics = self - .topics - .into_iter() - .filter_map(|raw| { - let label = raw.trim().to_string(); - if label.is_empty() { - None - } else { - Some(ExtractedTopic { label, score: 0.85 }) - } - }) - .collect(); - - ExtractedEntities { - entities, - topics, - llm_importance, - llm_importance_reason: self.importance_reason, - } - } -} - -// ── Helpers ────────────────────────────────────────────────────────────── - -fn parse_kind(s: &str) -> Option { - match s.trim().to_lowercase().as_str() { - "person" | "people" => Some(EntityKind::Person), - "organization" | "organisation" | "org" => Some(EntityKind::Organization), - "location" | "place" | "loc" => Some(EntityKind::Location), - "event" => Some(EntityKind::Event), - "product" => Some(EntityKind::Product), - "datetime" | "date" | "time" | "timestamp" => Some(EntityKind::Datetime), - "technology" | "tech" | "tool" | "framework" | "library" | "language" | "service" => { - Some(EntityKind::Technology) - } - "artifact" | "reference" | "ref" | "pr" | "ticket" | "file" | "commit" => { - Some(EntityKind::Artifact) - } - "quantity" | "amount" | "metric" | "number" | "money" => Some(EntityKind::Quantity), - "misc" | "miscellaneous" | "other" => Some(EntityKind::Misc), - _ => None, - } -} - -/// Find `needle` in `haystack` and return its `(char_start, char_end)`. -/// -/// Uses byte-level `find` then translates to char offsets so spans align -/// with the rest of the extractor pipeline (which is char-based). -fn find_char_span(haystack: &str, needle: &str) -> Option<(u32, u32)> { - find_char_span_from(haystack, needle, 0, 0).map(|(s, e, _)| (s, e)) -} - -/// Find `needle` in `haystack` starting from `byte_from` and return -/// `(char_start, char_end, byte_after_needle)`. -/// -/// The byte-offset return is so the caller can chain successive searches -/// without re-walking the prefix every time: pass the returned -/// `byte_after_needle` as the next call's `byte_from`. -/// -/// `char_from` must correspond to `byte_from` in the same `haystack` — -/// i.e. `haystack[..byte_from].chars().count() == char_from as usize`. -/// The caller maintains this invariant (cheap: it's the return from the -/// previous call). -fn find_char_span_from( - haystack: &str, - needle: &str, - byte_from: usize, - char_from: u32, -) -> Option<(u32, u32, usize)> { - if needle.is_empty() || byte_from > haystack.len() { - return None; - } - // Guard against `byte_from` landing inside a multi-byte UTF-8 sequence. - if !haystack.is_char_boundary(byte_from) { - return None; - } - let rel = haystack[byte_from..].find(needle)?; - let byte_start = byte_from + rel; - let byte_end = byte_start + needle.len(); - // Walk forward from the previous char position to build the new char - // offset — avoids re-walking the full prefix. - let char_start = char_from + haystack[byte_from..byte_start].chars().count() as u32; - let char_end = char_start + needle.chars().count() as u32; - Some((char_start, char_end, byte_end)) -} - -fn truncate_for_log(s: &str, max_chars: usize) -> String { - if s.chars().count() <= max_chars { - return s.to_string(); - } - let truncated: String = s.chars().take(max_chars).collect(); - format!("{truncated}…") -} - -// ── Tests ──────────────────────────────────────────────────────────────── - -#[cfg(test)] -#[path = "llm_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_tree/score/extract/llm_tests.rs b/src/openhuman/memory_tree/score/extract/llm_tests.rs deleted file mode 100644 index b54995ddf..000000000 --- a/src/openhuman/memory_tree/score/extract/llm_tests.rs +++ /dev/null @@ -1,756 +0,0 @@ -use super::*; - -#[test] -fn build_system_prompt_default_omits_topics() { - let p = build_system_prompt(false, None); - assert!(!p.contains("\"topics\"")); - assert!(!p.contains("Topics are")); - assert!(p.contains("ALL three top-level fields")); - assert!(p.contains("entities, importance")); -} - -#[test] -fn build_system_prompt_with_flag_includes_topics() { - let p = build_system_prompt(true, None); - assert!(p.contains("\"topics\"")); - assert!(p.contains("Topics are short free-form theme labels")); - assert!(p.contains("ALL four top-level fields")); - assert!(p.contains("entities, topics, importance")); -} - -#[test] -fn build_system_prompt_includes_output_language_directive() { - let p = build_system_prompt(true, Some("zh-CN")); - assert!(p.contains("Simplified Chinese")); - assert!(p.contains("Keep JSON keys")); - assert!(p.contains("\"importance_reason\"")); -} - -#[test] -fn extraction_output_parses_topics_when_present() { - let json = r#"{"entities":[],"topics":["rate limiting","memory tree"],"importance":0.6,"importance_reason":"r"}"#; - let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap(); - assert_eq!(parsed.topics, vec!["rate limiting", "memory tree"]); -} - -#[test] -fn extraction_output_tolerates_missing_topics() { - // Default extractor (emit_topics=false) — model won't emit topics - // and parsing must still succeed. - let json = r#"{"entities":[],"importance":0.6,"importance_reason":"r"}"#; - let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap(); - assert!(parsed.topics.is_empty()); -} - -#[test] -fn parse_kind_normalisation() { - assert_eq!(parse_kind("Person"), Some(EntityKind::Person)); - assert_eq!(parse_kind("organisation"), Some(EntityKind::Organization)); - assert_eq!(parse_kind(" PRODUCT "), Some(EntityKind::Product)); - assert!(parse_kind("Spaceship").is_none()); -} - -#[test] -fn parse_kind_accepts_new_semantic_kinds_and_synonyms() { - // Datetime - for s in ["datetime", "date", "time", "timestamp", " DateTime "] { - assert_eq!(parse_kind(s), Some(EntityKind::Datetime), "input={s:?}"); - } - // Technology - for s in [ - "technology", - "tech", - "tool", - "framework", - "library", - "language", - "service", - ] { - assert_eq!(parse_kind(s), Some(EntityKind::Technology), "input={s:?}"); - } - // Artifact - for s in [ - "artifact", - "reference", - "ref", - "pr", - "ticket", - "file", - "commit", - ] { - assert_eq!(parse_kind(s), Some(EntityKind::Artifact), "input={s:?}"); - } - // Quantity - for s in ["quantity", "amount", "metric", "number", "money"] { - assert_eq!(parse_kind(s), Some(EntityKind::Quantity), "input={s:?}"); - } -} - -#[test] -fn find_char_span_handles_unicode() { - let text = "中 Alice met Bob"; - let span = find_char_span(text, "Alice").unwrap(); - assert_eq!(span, (2, 7)); -} - -#[test] -fn find_char_span_returns_none_for_missing() { - assert!(find_char_span("hello world", "absent").is_none()); -} - -#[test] -fn find_char_span_from_advances_past_prior_match() { - let text = "Alice met Bob then Alice left"; - let (s1, e1, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); - assert_eq!((s1, e1), (0, 5)); - // Resuming from the cursor must find the second Alice. - let (s2, e2, _) = find_char_span_from(text, "Alice", byte_after, e1).unwrap(); - assert_eq!((s2, e2), (19, 24)); -} - -#[test] -fn find_char_span_from_returns_none_after_exhaustion() { - let text = "Alice met Bob"; - let (_, _, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); - // No second Alice → None. - assert!(find_char_span_from(text, "Alice", byte_after, 5).is_none()); -} - -#[test] -fn find_char_span_from_preserves_utf8() { - // Two "中" characters (3 bytes each in UTF-8); "Alice" between. - let text = "中 Alice 中 Alice"; - let (s1, e1, byte_after) = find_char_span_from(text, "Alice", 0, 0).unwrap(); - assert_eq!((s1, e1), (2, 7)); - let (s2, e2, _) = find_char_span_from(text, "Alice", byte_after, e1).unwrap(); - // First "中 Alice " = 2 + 5 + 1 + 1 + 1 chars; second Alice starts at char 10. - assert_eq!((s2, e2), (10, 15)); -} - -#[test] -fn find_char_span_from_rejects_non_char_boundary() { - // "中" is 3 bytes; offsets 1 and 2 are mid-codepoint. - let text = "中Alice"; - assert!(find_char_span_from(text, "Alice", 1, 0).is_none()); -} - -#[test] -fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() { - // Two "Alice" mentions in source → two distinct ExtractedEntity rows - // with non-overlapping spans. Previously both got (0, 5). - let out = LlmExtractionOutput { - entities: vec![ - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - ], - topics: vec![], - importance: None, - importance_reason: None, - }; - let cfg = LlmExtractorConfig::default(); - let e = out.into_extracted_entities("Alice met Bob then Alice left", &cfg); - assert_eq!(e.entities.len(), 2); - assert_eq!((e.entities[0].span_start, e.entities[0].span_end), (0, 5)); - assert_eq!((e.entities[1].span_start, e.entities[1].span_end), (19, 24)); -} - -#[test] -fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() { - // Three "Alice" mentions returned by LLM, only one in source → keep - // one, drop the rest as exhausted-duplicate. - let out = LlmExtractionOutput { - entities: vec![ - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - ], - topics: vec![], - importance: None, - importance_reason: None, - }; - let cfg = LlmExtractorConfig::default(); - let e = out.into_extracted_entities("Alice met Bob", &cfg); - assert_eq!(e.entities.len(), 1); -} - -#[tokio::test] -async fn extract_soft_fallback_on_provider_failure() { - // Provider always errors. extract() must NOT return Err — it must - // return an empty ExtractedEntities with a warn log after retry - // exhaustion. - // #002: this path now sets the process-global "structure degraded" flag. - // Hold the shared health test-guard so the flag is reset on entry and the - // signal doesn't leak into parallel status tests. - let _health_guard = crate::openhuman::memory_tree::health::test_guard(); - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::Arc; - - struct FailingProvider; - #[async_trait] - impl ChatProvider for FailingProvider { - fn name(&self) -> &str { - "test:failing" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Err(anyhow::anyhow!("simulated transport failure")) - } - } - - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(FailingProvider)); - let out = ex.extract("some text").await.unwrap(); - assert!(out.entities.is_empty()); - assert!(out.topics.is_empty()); - assert!(out.llm_importance.is_none()); -} - -#[tokio::test] -async fn extract_routes_through_chat_provider_and_parses_response() { - // Mock provider returns canned NER+importance JSON. Verify the - // extractor parses it, recovers spans by string search, and emits the - // expected entities + importance signal. - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct MockProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for MockProvider { - fn name(&self) -> &str { - "test:mock" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Ok(r#"{ - "entities": [ - {"kind":"person","text":"Alice"}, - {"kind":"organization","text":"Anthropic"} - ], - "importance": 0.8, - "importance_reason": "factual" - }"# - .to_string()) - } - } - - let mock = Arc::new(MockProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("Alice met Anthropic today.").await.unwrap(); - assert_eq!(mock.calls.load(Ordering::SeqCst), 1); - assert_eq!(out.entities.len(), 2); - assert_eq!(out.entities[0].text, "Alice"); - assert_eq!(out.entities[0].kind, EntityKind::Person); - assert_eq!(out.entities[1].text, "Anthropic"); - assert_eq!(out.llm_importance, Some(0.8)); - assert_eq!(out.llm_importance_reason.as_deref(), Some("factual")); -} - -#[tokio::test] -async fn extract_returns_empty_on_malformed_provider_response() { - // Provider returns garbage. Caller must NOT see an Err — the parse - // failure path returns empty entities (retrying the same input would - // yield the same garbage, so we don't burn retries). - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::Arc; - - struct GarbageProvider; - #[async_trait] - impl ChatProvider for GarbageProvider { - fn name(&self) -> &str { - "test:garbage" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Ok("not json at all".to_string()) - } - } - - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(GarbageProvider)); - let out = ex.extract("text").await.unwrap(); - assert!(out.entities.is_empty()); - assert!(out.llm_importance.is_none()); -} - -#[tokio::test] -async fn extract_clears_structure_latch_on_completed_empty_extraction() { - // #3365: the structure-degraded latch is a LIVENESS signal ("the extraction - // model is timing out"), set only after transport-failure retries are - // exhausted. A *completed* call — even one that yields no entities (a - // valid-but-empty result for entity-free text) — proves the model is - // responding, so it must clear the latch. Previously only a NON-empty result - // cleared it, so after the model recovered the latch stayed stuck whenever - // later texts were entity-light. - let _health_guard = crate::openhuman::memory_tree::health::test_guard(); - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::Arc; - - struct EmptyButOkProvider; - #[async_trait] - impl ChatProvider for EmptyButOkProvider { - fn name(&self) -> &str { - "test:empty-ok" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Ok(r#"{"entities": [], "topics": []}"#.to_string()) - } - } - - crate::openhuman::memory_tree::health::mark_structure_degraded( - crate::openhuman::memory_tree::health::FailureCode::ExtractionTimeout, - ); - assert!( - crate::openhuman::memory_tree::health::current_degraded_state().structure, - "precondition: structure-degraded latch is set" - ); - - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(EmptyButOkProvider)); - let out = ex.extract("entity-free text").await.unwrap(); - assert!(out.entities.is_empty(), "no entities in the response"); - - assert!( - !crate::openhuman::memory_tree::health::current_degraded_state().structure, - "a completed extraction must clear the structure-degraded latch even when empty" - ); -} - -#[test] -fn into_extracted_entities_drops_hallucinations() { - let out = LlmExtractionOutput { - entities: vec![ - LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }, - LlmEntity { - kind: "person".into(), - text: "ImaginaryPerson".into(), - }, - ], - topics: vec![], - importance: Some(0.7), - importance_reason: Some("substantive".into()), - }; - let cfg = LlmExtractorConfig::default(); - let e = out.into_extracted_entities("Alice met Bob today.", &cfg); - // Hallucinated "ImaginaryPerson" dropped; "Alice" kept. - assert_eq!(e.entities.len(), 1); - assert_eq!(e.entities[0].text, "Alice"); - assert_eq!(e.llm_importance, Some(0.7)); - assert_eq!(e.llm_importance_reason.as_deref(), Some("substantive")); -} - -#[test] -fn into_extracted_entities_clamps_importance() { - let out = LlmExtractionOutput { - entities: vec![], - topics: vec![], - importance: Some(1.5), - importance_reason: None, - }; - let cfg = LlmExtractorConfig::default(); - let e = out.into_extracted_entities("text", &cfg); - assert_eq!(e.llm_importance, Some(1.0)); -} - -#[test] -fn into_extracted_entities_strict_drops_unknown_kinds() { - let out = LlmExtractionOutput { - entities: vec![LlmEntity { - kind: "spaceship".into(), - text: "Enterprise".into(), - }], - topics: vec![], - importance: None, - importance_reason: None, - }; - let cfg = LlmExtractorConfig { - strict_kinds: true, - ..LlmExtractorConfig::default() - }; - let e = out.into_extracted_entities("Enterprise launched.", &cfg); - assert!(e.entities.is_empty()); -} - -#[test] -fn into_extracted_entities_lenient_falls_back_to_misc() { - let out = LlmExtractionOutput { - entities: vec![LlmEntity { - kind: "spaceship".into(), - text: "Enterprise".into(), - }], - topics: vec![], - importance: None, - importance_reason: None, - }; - let cfg = LlmExtractorConfig::default(); // strict_kinds = false - let e = out.into_extracted_entities("Enterprise launched.", &cfg); - assert_eq!(e.entities.len(), 1); - assert_eq!(e.entities[0].kind, EntityKind::Misc); -} - -#[test] -fn into_extracted_entities_disallowed_known_kind_falls_back_to_misc() { - // "person" is a known kind but might be excluded by allowed_kinds. - let out = LlmExtractionOutput { - entities: vec![LlmEntity { - kind: "person".into(), - text: "Alice".into(), - }], - topics: vec![], - importance: None, - importance_reason: None, - }; - let cfg = LlmExtractorConfig { - allowed_kinds: vec![EntityKind::Organization], // Person not allowed - strict_kinds: false, - ..LlmExtractorConfig::default() - }; - let e = out.into_extracted_entities("Alice met Bob.", &cfg); - assert_eq!(e.entities.len(), 1); - assert_eq!(e.entities[0].kind, EntityKind::Misc); -} - -#[test] -fn build_prompt_carries_user_text_and_kind_tag() { - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::Arc; - - struct NoopProvider; - #[async_trait] - impl ChatProvider for NoopProvider { - fn name(&self) -> &str { - "test:noop" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Ok("{}".into()) - } - } - - let cfg = LlmExtractorConfig { - model: "test-model".into(), - ..LlmExtractorConfig::default() - }; - let ex = LlmEntityExtractor::new(cfg, Arc::new(NoopProvider)); - let prompt = ex.build_prompt("hello"); - assert!(prompt.user.contains("hello")); - assert!(prompt.user.contains("Return JSON only")); - assert_eq!(prompt.temperature, 0.0); - assert_eq!(prompt.kind, "memory_tree::extract"); - // System prompt should describe the JSON schema. - assert!(prompt.system.contains("\"entities\"")); - assert!(prompt.system.contains("\"importance\"")); -} - -#[test] -fn truncate_for_log_short_input_unchanged() { - assert_eq!(truncate_for_log("hi", 10), "hi"); -} - -#[test] -fn truncate_for_log_long_input_appends_ellipsis() { - let long = "x".repeat(500); - let out = truncate_for_log(&long, 10); - assert_eq!(out.chars().count(), 11); // 10 + "…" - assert!(out.ends_with('…')); -} - -#[tokio::test] -async fn extract_retries_on_truncated_response() { - // First response is truncated mid-JSON (serde EOF) — a stream cutoff, - // not a wrong-shape body. It must be treated as retryable rather than - // silently dropped; the second (complete) response then recovers the - // entities (bug-report-2026-05-26 I1). - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct TruncatedThenCompleteProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for TruncatedThenCompleteProvider { - fn name(&self) -> &str { - "test:truncated" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - let n = self.calls.fetch_add(1, Ordering::SeqCst); - if n == 0 { - // Array never closes → serde reports EOF (is_eof()). - Ok(r#"{"entities":[{"kind":"person","text":"Alice"}"#.to_string()) - } else { - Ok(r#"{"entities":[{"kind":"person","text":"Alice"}],"importance":0.5,"importance_reason":"r"}"#.to_string()) - } - } - } - - let mock = Arc::new(TruncatedThenCompleteProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("Alice met Bob.").await.unwrap(); - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 2, - "truncation should trigger a retry, not a silent drop" - ); - assert_eq!(out.entities.len(), 1); - assert_eq!(out.entities[0].text, "Alice"); -} - -#[tokio::test] -async fn extract_does_not_retry_on_wrong_shape_response() { - // Companion to the truncation test: a *complete* but wrong-shape body is - // a serde error that is NOT EOF (`is_eof() == false`). Unlike a mid-JSON - // cutoff it's deterministic and won't fix itself on retry, so it must - // return immediately (`calls == 1`) with an empty extraction rather than - // entering the retry loop. Guards the `Err(e) if e.is_eof()` split so a - // future change can't accidentally retry every parse error - // (bug-report-2026-05-26 I1). - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct WrongShapeProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for WrongShapeProvider { - fn name(&self) -> &str { - "test:wrong-shape" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - // `entities` must be an array; a scalar is a complete-input type - // error (not a truncation), so serde reports a non-EOF error. - Ok(r#"{"entities":123}"#.to_string()) - } - } - - let mock = Arc::new(WrongShapeProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("Alice met Bob.").await.unwrap(); - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a non-EOF wrong-shape response must not trigger a retry" - ); - assert!( - out.entities.is_empty(), - "wrong-shape response should yield an empty extraction" - ); -} - -#[test] -fn build_prompt_sets_extraction_max_tokens_cap() { - // Extraction must cap output tokens so a credit-metered provider - // (OpenRouter) prices the request against a realistic budget instead of - // the model's full output window — an unset cap is what 402'd low-balance - // BYO users on every call (TAURI-RUST-C62). build_prompt is the single - // source of that cap. - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::Arc; - - struct NoopProvider; - #[async_trait] - impl ChatProvider for NoopProvider { - fn name(&self) -> &str { - "test:noop" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Ok("{}".into()) - } - } - - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(NoopProvider)); - let prompt = ex.build_prompt("hello"); - assert_eq!(prompt.max_tokens, Some(EXTRACTION_MAX_OUTPUT_TOKENS)); - // Lock the value: it must stay well under any plausible per-account - // OpenRouter balance so the pre-flight clears for a low-balance user. - assert_eq!(EXTRACTION_MAX_OUTPUT_TOKENS, 8192); -} - -#[tokio::test] -async fn extract_does_not_retry_on_permanent_402() { - // A 402 (the BYO provider account is out of credits) is a permanent client - // error: retrying reproduces it and fires one more Sentry event per attempt - // (TAURI-RUST-C62: 12k events from a single user). extract() must call the - // provider exactly once, return empty, and mark structure degraded with the - // credits-exhausted cause (not extraction-timeout). - let _health_guard = crate::openhuman::memory_tree::health::test_guard(); - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct InsufficientCreditsProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for InsufficientCreditsProvider { - fn name(&self) -> &str { - "test:402" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - // Verbatim shape of the TAURI-RUST-C62 provider error. - Err(anyhow::anyhow!( - "myopenrouter API error (402 Payment Required): This request requires more \ - credits, or fewer max_tokens. You requested up to 65536 tokens, but can only \ - afford 49732." - )) - } - } - - let mock = Arc::new(InsufficientCreditsProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a permanent 402 must not be retried" - ); - assert!(out.entities.is_empty()); - - let state = crate::openhuman::memory_tree::health::current_degraded_state(); - assert!( - state.structure, - "a permanent extraction failure must still mark structure degraded" - ); - assert_eq!( - state.cause.expect("degraded cause present").code, - crate::openhuman::memory_tree::health::FailureCode::BudgetExhausted, - "a 402 must map to the credits-exhausted cause, not extraction-timeout" - ); -} - -#[tokio::test] -async fn extract_does_not_retry_on_500_wrapped_monthly_quota() { - // TAURI-RUST-C9A: the Kiro IDE proxy wraps a 402 monthly-quota refusal - // inside a 500 envelope. Before the quota classifier, `is_non_retryable` - // saw only a 500 and treated it as a transient transport failure — so - // extract() retried 3× and each attempt fired a provider Sentry event - // (9k events from a single quota-capped user). It must now call the - // provider exactly once and map to the budget-exhausted cause. - let _health_guard = crate::openhuman::memory_tree::health::test_guard(); - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct MonthlyQuotaProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for MonthlyQuotaProvider { - fn name(&self) -> &str { - "test:kiro" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - // Verbatim shape of the TAURI-RUST-C9A provider error (500 envelope - // around the inner 402 / MONTHLY_REQUEST_COUNT). - Err(anyhow::anyhow!( - "kiro API error (500 Internal Server Error): {{\"error\":{{\"message\":\ - \"HTTP 402 from Kiro IDE: {{\\\"message\\\":\\\"You have reached the \ - limit.\\\",\\\"reason\\\":\\\"MONTHLY_REQUEST_COUNT\\\"}}\",\ - \"type\":\"server_error\"}}}}" - )) - } - } - - let mock = Arc::new(MonthlyQuotaProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a 500-wrapped monthly-quota refusal must not be retried" - ); - assert!(out.entities.is_empty()); - - let state = crate::openhuman::memory_tree::health::current_degraded_state(); - assert!( - state.structure, - "a permanent extraction failure must still mark structure degraded" - ); - assert_eq!( - state.cause.expect("degraded cause present").code, - crate::openhuman::memory_tree::health::FailureCode::BudgetExhausted, - "monthly-quota exhaustion must map to the budget-exhausted cause" - ); -} - -#[tokio::test] -async fn extract_retries_transient_provider_error() { - // The de-amplification fix must only short-circuit *permanent* errors. A - // transport/transient failure (no 4xx, no auth marker) must still exhaust - // the retry budget before falling back, or a flaky-network blip would drop - // the chunk on the first try. - let _health_guard = crate::openhuman::memory_tree::health::test_guard(); - use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct TransientProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for TransientProvider { - fn name(&self) -> &str { - "test:transient" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Err(anyhow::anyhow!( - "error sending request for url (https://api): connection refused" - )) - } - } - - let mock = Arc::new(TransientProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 3, - "a transient error must still exhaust the retry budget" - ); - assert!(out.entities.is_empty()); -} diff --git a/src/openhuman/memory_tree/score/extract/mod.rs b/src/openhuman/memory_tree/score/extract/mod.rs index 01f90d7ae..ca75d7173 100644 --- a/src/openhuman/memory_tree/score/extract/mod.rs +++ b/src/openhuman/memory_tree/score/extract/mod.rs @@ -1,63 +1,67 @@ -//! Entity extraction (Phase 2 / #708). -//! -//! Exposes [`EntityExtractor`] as a pluggable interface and a default -//! [`CompositeExtractor`] that runs a chain of extractors and merges their -//! output. Phase 2 ships with the mechanical regex extractor only; semantic -//! NER (GLiNER / LLM) plugs in later without changing any call sites. - -mod extractor; -pub mod llm; -pub mod regex; -pub mod types; +//! Product construction over tinycortex entity extraction. use std::sync::Arc; use crate::openhuman::config::Config; -use crate::openhuman::memory::chat::build_chat_runtime; +use async_trait::async_trait; -pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor}; -pub use llm::{LlmEntityExtractor, LlmExtractorConfig}; -pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; +pub use tinycortex::memory::score::extract::{ + ChatPrompt, ChatProvider, CompositeExtractor, EntityExtractor, EntityKind, ExtractedEntities, + ExtractedEntity, ExtractedTopic, LlmExtractorConfig, RegexEntityExtractor, +}; + +pub mod regex { + pub use tinycortex::memory::score::extract::regex::extract; +} + +pub struct LlmEntityExtractor(tinycortex::memory::score::extract::LlmEntityExtractor); + +impl LlmEntityExtractor { + pub fn new( + config: LlmExtractorConfig, + provider: Arc, + ) -> Self { + let provider = Arc::new(crate::openhuman::tinycortex::SeamChatProvider::new( + provider, + )); + Self(tinycortex::memory::score::extract::LlmEntityExtractor::new( + config, provider, + )) + } +} + +#[async_trait] +impl EntityExtractor for LlmEntityExtractor { + fn name(&self) -> &'static str { + self.0.name() + } + + async fn extract(&self, text: &str) -> anyhow::Result { + self.0.extract(text).await + } +} -/// Build the extractor used by seal handlers to label new summary nodes. -/// -/// Composition: -/// - regex extractor — always on, mechanical, near-zero cost -/// - LLM extractor with `emit_topics: true` — added when the unified -/// summarization workload can be built from inference routing. -/// -/// Differs from [`super::ScoringConfig::from_config`] (the chunk-admission -/// builder) in two ways: returns *just* an extractor (no thresholds / -/// weights / drop logic — none of which apply at seal time), and flips -/// `emit_topics` on so summaries surface thematic labels alongside -/// entities. Leaf-side scoring is unchanged. pub fn build_summary_extractor(config: &Config) -> Arc { - let (provider, model) = match build_chat_runtime(config) { + let (provider, model) = match crate::openhuman::memory::chat::build_chat_runtime(config) { Ok(runtime) => runtime, - Err(err) => { + Err(error) => { log::warn!( - "[memory_tree::extract] summary extractor: build_chat_runtime failed: \ - {err:#} — falling back to regex-only" + "[memory_tree::extract] chat provider unavailable; using regex-only extraction: {error:#}" ); return Arc::new(CompositeExtractor::regex_only()); } }; - - let cfg = LlmExtractorConfig { - model: model.clone(), - emit_topics: true, - output_language: config.output_language.clone(), - ..LlmExtractorConfig::default() - }; - - log::debug!( - "[memory_tree::extract] summary extractor: regex + LLM provider={} model={} \ - emit_topics=true", - provider.name(), - model + let extractor = LlmEntityExtractor::new( + LlmExtractorConfig { + model, + emit_topics: true, + output_language: config.output_language.clone(), + ..Default::default() + }, + provider, ); Arc::new(CompositeExtractor::new(vec![ Box::new(RegexEntityExtractor), - Box::new(LlmEntityExtractor::new(cfg, provider)), + Box::new(extractor), ])) } diff --git a/src/openhuman/memory_tree/score/extract/regex.rs b/src/openhuman/memory_tree/score/extract/regex.rs deleted file mode 100644 index 6abc3e122..000000000 --- a/src/openhuman/memory_tree/score/extract/regex.rs +++ /dev/null @@ -1,215 +0,0 @@ -//! Deterministic mechanical-entity extraction via regex. -//! -//! Catches the shapes regex handles cleanly and that are genuinely useful -//! as cross-platform identity anchors (email appearing in Slack + Gmail = -//! same person): -//! -//! - **Email** — RFC-ish pattern, boundary-guarded -//! - **URL** — `http(s)://…` up to whitespace or trailing punctuation -//! - **Handle** — `@alice`, `@alice.bsky.social`, or `alice#1234` -//! - **Hashtag** — `#launch-q2` -//! -//! Every match has `score = 1.0` (regex is deterministic). Spans are -//! char-offsets (not bytes) for UTF-8 safety. - -use once_cell::sync::Lazy; -use regex::Regex; - -use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; - -// ── Compiled regexes (once per process) ────────────────────────────────── - -static RE_EMAIL: Lazy = - Lazy::new(|| Regex::new(r"(?i)\b[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}\b").unwrap()); - -static RE_URL: Lazy = Lazy::new(|| { - // up-to trailing punctuation; avoids catastrophic backtracking - Regex::new(r"https?://[^\s<>\]\[()]+[^\s<>\]\[()\.\,;:\!\?]").unwrap() -}); - -static RE_HANDLE: Lazy = - Lazy::new(|| Regex::new(r"(?:^|[\s(])@([A-Za-z0-9_][A-Za-z0-9_.\-]{1,})").unwrap()); - -static RE_DISCRIM: Lazy = - Lazy::new(|| Regex::new(r"\b([A-Za-z0-9_.\-]{2,32})#\d{4}\b").unwrap()); - -static RE_HASHTAG: Lazy = - Lazy::new(|| Regex::new(r"(?:^|[\s(])#([A-Za-z][A-Za-z0-9_\-]{1,})").unwrap()); - -/// Extract all mechanical entities from `text`. -pub fn extract(text: &str) -> ExtractedEntities { - let mut entities: Vec = Vec::new(); - let mut topics: Vec = Vec::new(); - - for m in RE_EMAIL.find_iter(text) { - entities.push(to_entity(text, m.start(), m.end(), EntityKind::Email)); - } - for m in RE_URL.find_iter(text) { - entities.push(to_entity(text, m.start(), m.end(), EntityKind::Url)); - } - for cap in RE_HANDLE.captures_iter(text) { - if let Some(m) = cap.get(1) { - entities.push(to_entity(text, m.start(), m.end(), EntityKind::Handle)); - } - } - for cap in RE_DISCRIM.captures_iter(text) { - if let Some(m) = cap.get(0) { - entities.push(to_entity(text, m.start(), m.end(), EntityKind::Handle)); - } - } - for cap in RE_HASHTAG.captures_iter(text) { - if let Some(m) = cap.get(1) { - entities.push(to_entity(text, m.start(), m.end(), EntityKind::Hashtag)); - topics.push(ExtractedTopic { - label: text[m.start()..m.end()].to_lowercase(), - score: 1.0, - }); - } - } - - ExtractedEntities { - entities, - topics, - // Regex extractor never produces an LLM importance signal. - llm_importance: None, - llm_importance_reason: None, - } -} - -fn to_entity(text: &str, start: usize, end: usize, kind: EntityKind) -> ExtractedEntity { - ExtractedEntity { - kind, - text: text[start..end].to_string(), - span_start: char_index(text, start), - span_end: char_index(text, end), - score: 1.0, - } -} - -fn char_index(s: &str, byte_idx: usize) -> u32 { - let byte_idx = byte_idx.min(s.len()); - s[..byte_idx].chars().count() as u32 -} - -#[cfg(test)] -mod tests { - use super::*; - - fn kinds(e: &ExtractedEntities) -> Vec { - let mut k: Vec<_> = e.entities.iter().map(|x| x.kind).collect(); - k.sort_by_key(|k| *k as u8); - k - } - - #[test] - fn email_basic() { - let o = extract("contact alice@example.com please"); - assert_eq!(o.entities.len(), 1); - assert_eq!(o.entities[0].kind, EntityKind::Email); - assert_eq!(o.entities[0].text, "alice@example.com"); - } - - #[test] - fn url_stops_at_trailing_punct() { - let o = extract("see https://example.com/x?y=1 now."); - let urls: Vec<_> = o - .entities - .iter() - .filter(|e| e.kind == EntityKind::Url) - .collect(); - assert_eq!(urls.len(), 1); - assert_eq!(urls[0].text, "https://example.com/x?y=1"); - } - - #[test] - fn handle_vs_email_boundary() { - let o = extract("@alice met alice@example.com and @bob"); - let handles: Vec<_> = o - .entities - .iter() - .filter(|e| e.kind == EntityKind::Handle) - .map(|e| e.text.as_str()) - .collect(); - let emails: Vec<_> = o - .entities - .iter() - .filter(|e| e.kind == EntityKind::Email) - .map(|e| e.text.as_str()) - .collect(); - assert_eq!(handles, vec!["alice", "bob"]); - assert_eq!(emails, vec!["alice@example.com"]); - } - - #[test] - fn discord_style_handle() { - let o = extract("ping alice#1234"); - let h: Vec<_> = o - .entities - .iter() - .filter(|e| e.kind == EntityKind::Handle) - .collect(); - assert_eq!(h.len(), 1); - assert_eq!(h[0].text, "alice#1234"); - } - - #[test] - fn hashtag_emits_topic() { - let o = extract("tracking #launch-q2 updates"); - assert_eq!( - o.entities - .iter() - .filter(|e| e.kind == EntityKind::Hashtag) - .count(), - 1 - ); - assert_eq!(o.topics.len(), 1); - assert_eq!(o.topics[0].label, "launch-q2"); - } - - #[test] - fn hashtag_requires_leading_letter() { - let o = extract("#123 no, #x1 yes"); - let tags: Vec<_> = o - .entities - .iter() - .filter(|e| e.kind == EntityKind::Hashtag) - .collect(); - assert_eq!(tags.len(), 1); - assert_eq!(tags[0].text, "x1"); - } - - #[test] - fn utf8_span_is_char_not_byte() { - let o = extract("中 a@b.com"); - let email = o - .entities - .iter() - .find(|e| e.kind == EntityKind::Email) - .unwrap(); - assert_eq!(email.span_start, 2); - } - - #[test] - fn all_mechanical_kinds_in_one_pass() { - let o = extract("email a@b.com, url https://x.com, @alice, #topic1"); - let k = kinds(&o); - assert!(k.contains(&EntityKind::Email)); - assert!(k.contains(&EntityKind::Url)); - assert!(k.contains(&EntityKind::Handle)); - assert!(k.contains(&EntityKind::Hashtag)); - } - - #[test] - fn scores_always_one() { - let o = extract("a@b.com #x @y https://q.com"); - for e in &o.entities { - assert!((e.score - 1.0).abs() < f32::EPSILON); - } - } - - #[test] - fn empty_input_no_matches() { - let o = extract("plain prose with no identifiers"); - assert!(o.entities.is_empty()); - } -} diff --git a/src/openhuman/memory_tree/score/extract/types.rs b/src/openhuman/memory_tree/score/extract/types.rs deleted file mode 100644 index 1ad65b35b..000000000 --- a/src/openhuman/memory_tree/score/extract/types.rs +++ /dev/null @@ -1,345 +0,0 @@ -//! Types produced by entity extractors (Phase 2 / #708). -//! -//! The pipeline runs one or more [`super::EntityExtractor`] impls over each -//! admitted chunk and collects all their output into [`ExtractedEntities`]. - -use serde::{Deserialize, Serialize}; - -/// Classification of an extracted span. -/// -/// Split into two categories: -/// - **Mechanical** — regex finds these deterministically. Stable, high precision, -/// limited recall. These are "identifiers" (pointers), not "entities" -/// in the semantic sense. -/// - **Semantic** — model-based (future GLiNER / LLM). Named references to -/// real-world objects: Person, Organization, Location, Event, Product. -/// -/// Phase 2 ships with mechanical-only; semantic variants are populated in -/// Phase 3+ either at seal time by the summariser LLM or by a dedicated -/// per-chunk NER step if added later. -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -#[non_exhaustive] -pub enum EntityKind { - // Mechanical - Email, - Url, - Handle, - Hashtag, - // Semantic — emitted by the LLM extractor. - Person, - Organization, - Location, - Event, - Product, - /// Temporal expressions: "Friday", "Q2 2026", "EOD tomorrow", "next sprint". - Datetime, - /// Tools / frameworks / programming languages / services: - /// "Rust", "OAuth", "Slack API", "nomic-embed". - Technology, - /// Code / ticket / doc references that point at something addressable: - /// "PR #934", "src/openhuman/...", "OH-42", "ab7da2e2". - Artifact, - /// Amounts / metrics / money: "$5K", "20/min", "10k tokens", "52 chunks". - Quantity, - Misc, - // Thematic — scorer-surfaced topics (hashtag-like short phrases or - // LLM-extracted themes). Promoted into the canonical entity stream - // by the resolver so Phase 3c topic trees can route on themes the - // same way they route on people/orgs. A chunk saying "Phoenix - // migration ships Friday" emits `topic:phoenix` and `topic:migration` - // in addition to any emails/hashtags the mechanical extractors find. - Topic, -} - -impl EntityKind { - /// Snake-case wire string for serialisation and SQL storage. - pub fn as_str(self) -> &'static str { - match self { - Self::Email => "email", - Self::Url => "url", - Self::Handle => "handle", - Self::Hashtag => "hashtag", - Self::Person => "person", - Self::Organization => "organization", - Self::Location => "location", - Self::Event => "event", - Self::Product => "product", - Self::Datetime => "datetime", - Self::Technology => "technology", - Self::Artifact => "artifact", - Self::Quantity => "quantity", - Self::Misc => "misc", - Self::Topic => "topic", - } - } - - /// Inverse of [`Self::as_str`]; returns `Err` for unknown wire strings. - pub fn parse(s: &str) -> Result { - match s { - "email" => Ok(Self::Email), - "url" => Ok(Self::Url), - "handle" => Ok(Self::Handle), - "hashtag" => Ok(Self::Hashtag), - "person" => Ok(Self::Person), - "organization" => Ok(Self::Organization), - "location" => Ok(Self::Location), - "event" => Ok(Self::Event), - "product" => Ok(Self::Product), - "datetime" => Ok(Self::Datetime), - "technology" => Ok(Self::Technology), - "artifact" => Ok(Self::Artifact), - "quantity" => Ok(Self::Quantity), - "misc" => Ok(Self::Misc), - "topic" => Ok(Self::Topic), - other => Err(format!("unknown entity kind: {other}")), - } - } - - /// Whether this kind comes from deterministic extraction. - pub fn is_mechanical(self) -> bool { - matches!(self, Self::Email | Self::Url | Self::Handle | Self::Hashtag) - } -} - -/// One extracted span from a chunk's content. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct ExtractedEntity { - pub kind: EntityKind, - /// Surface form as it appears in the chunk. - pub text: String, - /// Character offsets `[start, end)` into the chunk text. - pub span_start: u32, - pub span_end: u32, - /// Extractor confidence `[0.0, 1.0]`. Regex = 1.0; model-based = output. - pub score: f32, -} - -/// Topic candidate (hashtag-style or summariser-labeled). -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct ExtractedTopic { - /// Normalised topic text (lowercase, no leading `#`). - pub label: String, - pub score: f32, -} - -/// Aggregate output of one or more extractors on a single chunk. -/// -/// `llm_importance` and `llm_importance_reason` are populated by extractors -/// that piggyback an importance rating on their NER call (see -/// [`super::llm::LlmEntityExtractor`]). Cheap regex extractors leave them -/// `None`; downstream signal compute treats `None` as "no LLM signal" and -/// the weighted combine zeroes that contribution out so behaviour matches -/// pre-LLM Phase 2 exactly when LLM is disabled. -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct ExtractedEntities { - pub entities: Vec, - pub topics: Vec, - /// Optional LLM-rated importance in `[0.0, 1.0]` for this chunk. - /// `None` means no LLM signal is available. - #[serde(default)] - pub llm_importance: Option, - /// One-line audit trail from the LLM explaining the importance rating. - /// Used purely for diagnostics; never feeds back into scoring. - #[serde(default)] - pub llm_importance_reason: Option, -} - -impl ExtractedEntities { - /// True when neither entities nor topics were extracted. - pub fn is_empty(&self) -> bool { - self.entities.is_empty() && self.topics.is_empty() - } - - /// Count of unique `(kind, text)` pairs, case-insensitive. Used as a scoring signal. - pub fn unique_entity_count(&self) -> usize { - use std::collections::HashSet; - // Only the cardinality matters, so a hash set counts the distinct - // `(kind, text)` pairs without paying for ordered insertion. - self.entities - .iter() - .map(|e| (e.kind, e.text.to_lowercase())) - .collect::>() - .len() - } - - /// Merge another extractor's output into this one. - /// - /// Deduplicates entities by `(kind, normalised_text, span_start)` and - /// topics by `label` so the same match from two extractors doesn't get - /// double-counted. - /// - /// LLM importance signals merge by **maximum** — if either side rated - /// the chunk as important, the merged result keeps that higher rating. - /// The reason from whichever side won the max wins; if they tied or - /// both are absent, the non-empty one (if any) is kept. - pub fn merge(&mut self, other: ExtractedEntities) { - use std::collections::HashSet; - // Both sets are membership-only dedup guards — the surviving order is - // the existing `Vec` push order, never the set's — so a hash set keeps - // the merge result identical while dropping the ordered-key overhead. - let mut seen: HashSet<(EntityKind, String, u32)> = self - .entities - .iter() - .map(|e| (e.kind, e.text.to_lowercase(), e.span_start)) - .collect(); - for e in other.entities { - let key = (e.kind, e.text.to_lowercase(), e.span_start); - if seen.insert(key) { - self.entities.push(e); - } - } - let mut topic_seen: HashSet = self.topics.iter().map(|t| t.label.clone()).collect(); - for t in other.topics { - if topic_seen.insert(t.label.clone()) { - self.topics.push(t); - } - } - - // Merge LLM importance: max wins, reason follows the max. - match (self.llm_importance, other.llm_importance) { - (Some(a), Some(b)) if b > a => { - self.llm_importance = Some(b); - self.llm_importance_reason = other.llm_importance_reason; - } - (None, Some(b)) => { - self.llm_importance = Some(b); - self.llm_importance_reason = other.llm_importance_reason; - } - // self.a >= other.b OR other has nothing — keep self - _ => { - if self.llm_importance_reason.is_none() { - self.llm_importance_reason = other.llm_importance_reason; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn entity_kind_round_trip() { - for k in [ - EntityKind::Email, - EntityKind::Url, - EntityKind::Handle, - EntityKind::Hashtag, - EntityKind::Person, - EntityKind::Organization, - EntityKind::Location, - EntityKind::Event, - EntityKind::Product, - EntityKind::Datetime, - EntityKind::Technology, - EntityKind::Artifact, - EntityKind::Quantity, - EntityKind::Misc, - EntityKind::Topic, - ] { - assert_eq!(EntityKind::parse(k.as_str()).unwrap(), k); - } - } - - #[test] - fn mechanical_classification() { - assert!(EntityKind::Email.is_mechanical()); - assert!(EntityKind::Url.is_mechanical()); - assert!(EntityKind::Handle.is_mechanical()); - assert!(EntityKind::Hashtag.is_mechanical()); - assert!(!EntityKind::Person.is_mechanical()); - } - - #[test] - fn unique_entity_count_dedups_case_insensitive() { - let e = ExtractedEntities { - entities: vec![ - ExtractedEntity { - kind: EntityKind::Person, - text: "Alice".into(), - span_start: 0, - span_end: 5, - score: 1.0, - }, - ExtractedEntity { - kind: EntityKind::Person, - text: "alice".into(), - span_start: 10, - span_end: 15, - score: 1.0, - }, - ], - topics: vec![], - llm_importance: None, - llm_importance_reason: None, - }; - assert_eq!(e.unique_entity_count(), 1); - } - - #[test] - fn unique_entity_count_keeps_different_kinds_distinct() { - let e = ExtractedEntities { - entities: vec![ - ExtractedEntity { - kind: EntityKind::Handle, - text: "alice".into(), - span_start: 0, - span_end: 5, - score: 1.0, - }, - ExtractedEntity { - kind: EntityKind::Hashtag, - text: "alice".into(), - span_start: 10, - span_end: 15, - score: 1.0, - }, - ], - topics: vec![], - llm_importance: None, - llm_importance_reason: None, - }; - assert_eq!(e.unique_entity_count(), 2); - } - - #[test] - fn merge_dedups_by_kind_text_span() { - let mut a = ExtractedEntities { - entities: vec![ExtractedEntity { - kind: EntityKind::Email, - text: "x@y.com".into(), - span_start: 0, - span_end: 7, - score: 1.0, - }], - topics: vec![], - llm_importance: None, - llm_importance_reason: None, - }; - let b = ExtractedEntities { - entities: vec![ - ExtractedEntity { - kind: EntityKind::Email, - text: "x@y.com".into(), - span_start: 0, - span_end: 7, - score: 1.0, - }, // dup - ExtractedEntity { - kind: EntityKind::Email, - text: "x@y.com".into(), - span_start: 50, - span_end: 57, - score: 1.0, - }, // different span — keep - ], - topics: vec![], - llm_importance: None, - llm_importance_reason: None, - }; - a.merge(b); - assert_eq!(a.entities.len(), 2); - } -} diff --git a/src/openhuman/memory_tree/score/mod.rs b/src/openhuman/memory_tree/score/mod.rs index fc64d628d..73978767e 100644 --- a/src/openhuman/memory_tree/score/mod.rs +++ b/src/openhuman/memory_tree/score/mod.rs @@ -1,477 +1,44 @@ -//! Phase 2: scoring / admission / enrichment pipeline (#708). -//! -//! Wraps extraction, signal computation, admission gate, canonicalisation, -//! and persistence into one call per chunk. Phase 1 `_ingest_one_chunk` -//! passes each chunk through [`score_chunk`] after chunking and before -//! storing. +//! Product adapters over tinycortex scoring and admission. pub mod embed; pub mod extract; -pub mod resolver; -pub mod signals; pub mod store; use std::sync::Arc; -use anyhow::Result; -use chrono::Utc; -use futures_util::future::try_join_all; -use rusqlite::Transaction; -use serde::{Deserialize, Serialize}; +pub use anyhow::Result; +pub use tinycortex::memory::score::{ + persist_score_tx, score_chunk, score_chunks, score_chunks_fast, ScoreResult, ScoringConfig, + DEFAULT_DEFINITE_DROP, DEFAULT_DEFINITE_KEEP, DEFAULT_DROP_THRESHOLD, PRIORITY_BOOST, + PRIORITY_TAG, +}; +pub use tinycortex::memory::score::{resolver, signals}; -use self::extract::{EntityExtractor, ExtractedEntities}; -use self::resolver::{canonicalise, CanonicalEntity}; -use self::signals::{ScoreSignals, SignalWeights}; -use crate::openhuman::memory_store::chunks::types::{approx_token_count, Chunk, SourceKind}; - -/// Default drop threshold. Chunks with `total < DEFAULT_DROP_THRESHOLD` -/// are tombstoned and never reach the chunk store. -pub const DEFAULT_DROP_THRESHOLD: f32 = 0.3; - -/// If the deterministic (cheap-signals-only) total is at or above this, -/// the chunk is admitted without consulting the LLM extractor. -/// -/// Tuned to leave a generous "borderline" band where the LLM signal is -/// most informative while skipping LLM cost on obviously substantive -/// content. -pub const DEFAULT_DEFINITE_KEEP: f32 = 0.85; - -/// If the deterministic total is at or below this, the chunk is dropped -/// without consulting the LLM extractor. Catches obvious noise cheaply. -pub const DEFAULT_DEFINITE_DROP: f32 = 0.15; - -/// Metadata tag marking a chunk as high-priority source material. The -/// ingest path stamps this on GitHub commit messages and closed/merged -/// issues & PRs (see `memory_sources::sync`), which are higher-signal than -/// ambient activity and should be preferred when building the memory tree. -pub const PRIORITY_TAG: &str = "priority_high"; - -/// Additive score nudge applied to chunks carrying [`PRIORITY_TAG`]. Pushes -/// them above the drop threshold and higher in retrieval / summary ordering -/// without fully overriding the content signals. Clamped to `1.0`. -pub const PRIORITY_BOOST: f32 = 0.25; - -/// Whole outcome of [`score_chunk`]. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ScoreResult { - pub chunk_id: String, - pub total: f32, - pub signals: ScoreSignals, - pub kept: bool, - pub drop_reason: Option, - pub extracted: ExtractedEntities, - pub canonical_entities: Vec, -} - -/// Configuration passed through the ingest pipeline for Phase 2 behaviour. -/// -/// Held as a struct (vs config struct fields) so callers can override per-run -/// without mutating global config — useful for tests and explicit threshold -/// tuning. -/// -/// The `extractor` field always runs (typically a regex-based composite -/// for cheap mechanical entities). `llm_extractor` is consulted **only -/// when the cheap-signals total falls in the band** -/// `(definite_drop_threshold, definite_keep_threshold)` — chunks that are -/// obviously trash or obviously substantive don't pay the LLM cost. -pub struct ScoringConfig { - pub extractor: Arc, - pub weights: SignalWeights, - pub drop_threshold: f32, - /// Optional second-pass extractor whose output is **merged** into the - /// regex output before the final combine. Designed for LLM-based NER + - /// importance signal (see [`extract::LlmEntityExtractor`]). `None` - /// means LLM augmentation is disabled. - pub llm_extractor: Option>, - /// Cheap-signals total ≥ this → admit without consulting LLM. - pub definite_keep_threshold: f32, - /// Cheap-signals total ≤ this → drop without consulting LLM. - pub definite_drop_threshold: f32, -} - -impl ScoringConfig { - /// Phase 2 default: regex-only extractor, default weights, default threshold. - pub fn default_regex_only() -> Self { - Self { - extractor: Arc::new(extract::CompositeExtractor::regex_only()), - weights: SignalWeights::default(), - drop_threshold: DEFAULT_DROP_THRESHOLD, - llm_extractor: None, - definite_keep_threshold: DEFAULT_DEFINITE_KEEP, - definite_drop_threshold: DEFAULT_DEFINITE_DROP, +/// Build crate scoring policy from product inference routing. +pub fn scoring_config_from(config: &crate::openhuman::config::Config) -> ScoringConfig { + let (provider, model) = match crate::openhuman::memory::chat::build_chat_runtime(config) { + Ok((provider, model)) => ( + Arc::new(crate::openhuman::tinycortex::SeamChatProvider::new( + provider, + )) as Arc, + model, + ), + Err(error) => { + log::warn!( + "[memory::score] chat provider unavailable; using regex-only scoring: {error:#}" + ); + return ScoringConfig::default_regex_only(); } - } - - /// Convenience constructor: regex always + LLM extractor on borderline - /// chunks. The `llm_importance` weight is enabled in [`SignalWeights`] - /// so the LLM signal actually influences the final total. - pub fn with_llm_extractor(llm: Arc) -> Self { - Self { - extractor: Arc::new(extract::CompositeExtractor::regex_only()), - weights: SignalWeights::with_llm_enabled(), - drop_threshold: DEFAULT_DROP_THRESHOLD, - llm_extractor: Some(llm), - definite_keep_threshold: DEFAULT_DEFINITE_KEEP, - definite_drop_threshold: DEFAULT_DEFINITE_DROP, - } - } - - /// Build a [`ScoringConfig`] from the workspace [`Config`]. - /// - /// The LLM extractor follows the unified summarization workload routing. - /// Construction errors fall back to regex-only with a warn log; scoring - /// never blocks on LLM availability. - pub fn from_config(config: &crate::openhuman::config::Config) -> Self { - use crate::openhuman::memory::chat::build_chat_runtime; - - let (provider, model) = match build_chat_runtime(config) { - Ok(runtime) => runtime, - Err(err) => { - log::warn!( - "[memory::score] build_chat_runtime failed: {err:#} — \ - falling back to regex-only" - ); - return Self::default_regex_only(); - } - }; - - let cfg = extract::LlmExtractorConfig { - model: model.clone(), + }; + let extractor = tinycortex::memory::score::extract::LlmEntityExtractor::new( + tinycortex::memory::score::extract::LlmExtractorConfig { + model, output_language: config.output_language.clone(), - ..extract::LlmExtractorConfig::default() - }; - - log::info!( - "[memory::score] using LlmEntityExtractor provider={} model={}", - provider.name(), - model - ); - Self::with_llm_extractor(Arc::new(extract::LlmEntityExtractor::new(cfg, provider))) - } -} - -/// Compute the score for one chunk. -/// -/// Pure function — does not touch the store. Callers decide what to persist -/// based on [`ScoreResult::kept`]. -/// -/// Pipeline: -/// 1. Run the always-on extractor (typically regex). -/// 2. Compute cheap signals; combine **excluding** `llm_importance` weight. -/// 3. Short-circuit: -/// - If cheap total ≥ `definite_keep_threshold`: admit without LLM. -/// - If cheap total ≤ `definite_drop_threshold`: drop without LLM. -/// - Else: borderline — run the LLM extractor (if configured), merge -/// its output, recompute signals, recombine with full weights. -/// 4. Apply final admission gate against `drop_threshold`. -pub async fn score_chunk(chunk: &Chunk, cfg: &ScoringConfig) -> Result { - log::debug!( - "[memory::score] score_chunk chunk_id={} tokens={}", - chunk.id, - chunk.token_count + ..Default::default() + }, + provider, ); - - let scoring_content = scoring_content_for_chunk(chunk); - let scoring_token_count = approx_token_count(&scoring_content); - - // 1. Always-on extraction (regex / mechanical). - let mut extracted = cfg.extractor.extract(&scoring_content).await?; - - // 2. Compute cheap signals + combine excluding LLM importance. - let mut signals = self::signals::compute( - &chunk.metadata, - &scoring_content, - scoring_token_count, - &extracted, - ); - let cheap_total = self::signals::combine_cheap_only(&signals, &cfg.weights); - - // 3. Short-circuit decision. - let in_band = - cheap_total > cfg.definite_drop_threshold && cheap_total < cfg.definite_keep_threshold; - let llm_consulted = if in_band { - if let Some(llm) = cfg.llm_extractor.as_ref() { - log::debug!( - "[memory::score] borderline chunk_id={} cheap_total={:.3} — consulting LLM", - chunk.id, - cheap_total - ); - match llm.extract(&scoring_content).await { - Ok(more) => { - extracted.merge(more); - // Recompute signals so llm_importance flows in. - signals = self::signals::compute( - &chunk.metadata, - &scoring_content, - scoring_token_count, - &extracted, - ); - true - } - Err(e) => { - log::warn!( - "[memory::score] LLM extractor `{}` failed: {e} — \ - falling back to cheap signals only", - llm.name() - ); - false - } - } - } else { - false - } - } else { - log::debug!( - "[memory::score] short-circuit chunk_id={} cheap_total={:.3} \ - ({}, skipping LLM)", - chunk.id, - cheap_total, - if cheap_total >= cfg.definite_keep_threshold { - "definite_keep" - } else { - "definite_drop" - } - ); - false - }; - - // 4. Final weighted combine. - // - // If the LLM ran, its importance signal is populated → use the full - // `combine` which includes the `llm_importance` weight. - // - // If the LLM was skipped (short-circuited or not configured) OR failed - // (caught above, sets `llm_consulted=false`), using the full combine - // would pin `llm_importance * w.llm_importance = 0 * 2.0` into the - // numerator while still dividing by the full denominator — artificially - // dragging the total down. Fall back to `combine_cheap_only` which - // excludes that term from both numerator and denominator, so the cheap - // signals alone produce the total. - let mut total = if llm_consulted { - self::signals::combine(&signals, &cfg.weights) - } else { - self::signals::combine_cheap_only(&signals, &cfg.weights) - }; - - // 4b. Priority boost. Chunks tagged PRIORITY_TAG at ingest (GitHub - // commit messages, closed/merged issues & PRs) are higher-signal than - // ambient activity. Nudge their total up (clamped) so they clear the - // admission gate more readily and rank higher when the memory tree is - // built from chunk scores. - let priority = chunk.metadata.tags.iter().any(|t| t == PRIORITY_TAG); - if priority { - let boosted = (total + PRIORITY_BOOST).min(1.0); - if boosted > total { - log::debug!( - "[memory::score] priority boost chunk_id={} {:.3} -> {:.3}", - chunk.id, - total, - boosted - ); - total = boosted; - } - } - - // 5. Admission gate. Source and interaction priors are deliberately - // non-zero, so guard against very short entity-free chatter being kept by - // metadata alone. Priority-tagged chunks bypass the tiny/entity-free - // drop since the tag itself is a strong relevance signal. - let tiny_entity_free = !priority - && scoring_token_count < self::signals::token_count::TOKEN_MIN - && extracted.is_empty(); - let kept = !tiny_entity_free && total >= cfg.drop_threshold; - let drop_reason = if kept { - None - } else if tiny_entity_free { - Some(format!( - "token_count {} < minimum {} and no entities extracted", - scoring_token_count, - self::signals::token_count::TOKEN_MIN - )) - } else { - Some(format!( - "total {total:.3} < threshold {:.3}", - cfg.drop_threshold - )) - }; - - // 6. Canonicalise for indexing (only meaningful when kept — but we - // canonicalise unconditionally so the result is inspectable in tests) - let canonical_entities = canonicalise(&extracted); - - if !kept { - log::debug!( - "[memory::score] drop chunk_id={} total={:.3} reason={:?} llm_consulted={}", - chunk.id, - total, - drop_reason, - llm_consulted - ); - } - - Ok(ScoreResult { - chunk_id: chunk.id.clone(), - total, - signals, - kept, - drop_reason, - extracted, - canonical_entities, - }) -} - -fn scoring_content_for_chunk(chunk: &Chunk) -> String { - if chunk.metadata.source_kind != SourceKind::Chat { - return chunk.content.clone(); - } - - chunk - .content - .lines() - .filter(|line| { - let trimmed = line.trim_start(); - !trimmed.starts_with("# Chat transcript") && !trimmed.starts_with("## ") - }) - .collect::>() - .join("\n") -} - -/// Score a batch of chunks. Errors from any single chunk fail the batch — -/// scoring is pure-ish (only the extractor may error) and a failure here is -/// a real bug, not a per-chunk issue to tolerate silently. -pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result> { - try_join_all(chunks.iter().map(|chunk| score_chunk(chunk, cfg))).await -} - -/// Cheap-only batch scoring path used by the async queue ingest pipeline. -/// -/// This preserves the same thresholds and admission gate as [`score_chunks`] -/// but guarantees no LLM extractor is consulted on the ingest hot path. -pub async fn score_chunks_fast(chunks: &[Chunk], cfg: &ScoringConfig) -> Result> { - let fast_cfg = ScoringConfig { - extractor: cfg.extractor.clone(), - weights: cfg.weights.clone(), - drop_threshold: cfg.drop_threshold, - llm_extractor: None, - definite_keep_threshold: cfg.definite_keep_threshold, - definite_drop_threshold: cfg.definite_drop_threshold, - }; - score_chunks(chunks, &fast_cfg).await -} - -// ── Persistence helpers used by the ingest orchestrator ───────────────── - -/// Persist the score row + entity-index rows for one kept chunk. -/// -/// The caller is responsible for having already written the chunk itself -/// into `mem_tree_chunks` (so the FK-like relation is satisfied). Dropped -/// chunks still get a score row persisted for diagnostics — callers should -/// pass `None` for `tree_id` in that case, since the chunk won't appear in -/// a tree. -pub fn persist_score( - config: &crate::openhuman::config::Config, - result: &ScoreResult, - timestamp_ms: i64, - tree_id: Option<&str>, -) -> Result<()> { - let row = score_row(result); - store::upsert_score(config, &row)?; - - if result.kept { - // Clear any stale entity-index rows for this chunk before re-indexing. - // INSERT OR REPLACE on (entity_id, node_id) never deletes rows whose - // entity_id is no longer present in the new extraction — so a re-score - // that drops an entity would otherwise leave a phantom index row. - store::clear_entity_index_for_node(config, &result.chunk_id)?; - if !result.canonical_entities.is_empty() { - store::index_entities( - config, - &result.canonical_entities, - &result.chunk_id, - "leaf", - timestamp_ms, - tree_id, - )?; - // E2GraphRAG: accumulate the chunk's entity co-occurrence edges so - // query-time hop-distance filtering has a graph to walk. - persist_cooccurrence_edges(config, &result.canonical_entities, timestamp_ms); - } - } - - Ok(()) -} - -/// Build and persist the undirected co-occurrence edges implied by a chunk's -/// canonical entities. Best-effort: a graph write failure must never block the -/// (already-committed) score/entity-index write, so errors are logged and -/// swallowed. Co-occurrence = two entities mentioned in the same chunk. -fn persist_cooccurrence_edges( - config: &crate::openhuman::config::Config, - entities: &[crate::openhuman::memory_tree::score::resolver::CanonicalEntity], - timestamp_ms: i64, -) { - use crate::openhuman::memory_tree::graph; - let ids: Vec = entities.iter().map(|e| e.canonical_id.clone()).collect(); - let pairs = graph::pairs_from_entities(&ids); - if pairs.is_empty() { - return; - } - match graph::upsert_edges(config, &pairs, timestamp_ms) { - Ok(n) => log::debug!("[memory_tree::graph] indexed cooccurrence edges count={n}"), - Err(e) => log::warn!("[memory_tree::graph] edge upsert failed (non-fatal): {e:#}"), - } -} - -pub(crate) fn persist_score_tx( - tx: &Transaction<'_>, - result: &ScoreResult, - timestamp_ms: i64, - tree_id: Option<&str>, -) -> Result<()> { - let row = score_row(result); - store::upsert_score_tx(tx, &row)?; - - if result.kept { - // See persist_score for why we clear before re-indexing. - store::clear_entity_index_for_node_tx(tx, &result.chunk_id)?; - if !result.canonical_entities.is_empty() { - store::index_entities_tx( - tx, - &result.canonical_entities, - &result.chunk_id, - "leaf", - timestamp_ms, - tree_id, - )?; - // E2GraphRAG: accumulate co-occurrence edges in the SAME transaction - // so the graph never diverges from the entity index. - let ids: Vec = result - .canonical_entities - .iter() - .map(|e| e.canonical_id.clone()) - .collect(); - let pairs = crate::openhuman::memory_tree::graph::pairs_from_entities(&ids); - let n = - crate::openhuman::memory_tree::graph::upsert_edges_tx(tx, &pairs, timestamp_ms)?; - log::debug!("[memory_tree::graph] indexed cooccurrence edges (tx) count={n}"); - } - } - - Ok(()) -} - -fn score_row(result: &ScoreResult) -> store::ScoreRow { - // Score rows keep wall-clock scoring time; the separate timestamp_ms - // argument used for entity indexes is the source/ingest ordering time. - store::ScoreRow { - chunk_id: result.chunk_id.clone(), - total: result.total, - signals: result.signals.clone(), - dropped: !result.kept, - reason: result.drop_reason.clone(), - computed_at_ms: Utc::now().timestamp_millis(), - llm_importance_reason: result.extracted.llm_importance_reason.clone(), - } + ScoringConfig::with_llm_extractor(Arc::new(extractor)) } #[cfg(test)] diff --git a/src/openhuman/memory_tree/score/resolver.rs b/src/openhuman/memory_tree/score/resolver.rs deleted file mode 100644 index 845ad0c31..000000000 --- a/src/openhuman/memory_tree/score/resolver.rs +++ /dev/null @@ -1,265 +0,0 @@ -//! Entity canonicalisation / cross-platform merge (Phase 2 / #708, V1). -//! -//! Exact-match only: normalises surface forms (lowercase emails, strip -//! leading `@` on handles) and assigns a canonical `entity_id` string. -//! -//! Fuzzy matching (alice-slack ≡ Alice-Discord by soft match) is deferred -//! until we have real entity-graph data — the current implementation -//! handles the mechanical cases cleanly without producing false merges. - -use serde::{Deserialize, Serialize}; - -use crate::openhuman::memory_tree::score::extract::{EntityKind, ExtractedEntities}; - -/// Canonicalised entity — same shape as [`ExtractedEntity`] plus a stable -/// `canonical_id` suitable for indexing. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct CanonicalEntity { - pub canonical_id: String, - pub kind: EntityKind, - pub surface: String, - pub span_start: u32, - pub span_end: u32, - pub score: f32, -} - -/// Canonicalise a batch of extracted entities. -/// -/// Same surface form (after normalisation) → same `canonical_id` regardless -/// of how many times it appears in a chunk. Preserves source spans by -/// emitting one [`CanonicalEntity`] per occurrence. -/// -/// Extracted **topics** are also promoted into the canonical stream under -/// [`EntityKind::Topic`] so downstream routing (Phase 3c topic trees) can -/// treat themes as first-class scope alongside people/orgs. Topics have no -/// source span (they're derived from the whole chunk, not a specific -/// substring), so `span_start` / `span_end` are both `0` for topic rows — -/// readers should key on `kind` instead of span when span-awareness matters. -pub fn canonicalise(extracted: &ExtractedEntities) -> Vec { - let mut out: Vec = extracted - .entities - .iter() - .map(|e| CanonicalEntity { - canonical_id: canonical_id_for(e.kind, &e.text), - kind: e.kind, - surface: e.text.clone(), - span_start: e.span_start, - span_end: e.span_end, - score: e.score, - }) - .collect(); - - // Promote topics. Dedup against the entities we already emitted so a - // hashtag like `#launch` and a topic label `"launch"` don't both land - // as the same canonical id with the same kind — the hashtag keeps its - // Hashtag kind, the topic gets Topic kind, and `canonical_id_for` - // makes them distinguishable: `hashtag:launch` vs `topic:launch`. - for topic in &extracted.topics { - let canonical_id = canonical_id_for(EntityKind::Topic, &topic.label); - // Dedup within the topic set in case the scorer produces the same - // label twice (LLM + regex overlap). Entities under other kinds - // aren't dedup targets — `topic:launch` and `hashtag:launch` are - // intentionally separate. - if out - .iter() - .any(|e| e.kind == EntityKind::Topic && e.canonical_id == canonical_id) - { - continue; - } - out.push(CanonicalEntity { - canonical_id, - kind: EntityKind::Topic, - surface: topic.label.clone(), - span_start: 0, - span_end: 0, - score: topic.score, - }); - } - out -} - -/// Canonical id form per kind. Deterministic so the same surface always -/// maps to the same id. -/// -/// - Email: `email:lowercased` -/// - Handle: `handle:lowercased` with leading `@` stripped -/// - Hashtag: `hashtag:lowercased` with leading `#` stripped -/// - URL: `url:trimmed` with case preserved for path/query exact matching -/// - Semantic kinds: `kind:lowercased-surface` (V1; fuzzy merge deferred) -pub fn canonical_id_for(kind: EntityKind, surface: &str) -> String { - let trimmed = surface.trim(); - let clean = if kind == EntityKind::Url { - trimmed.to_string() - } else { - trimmed - .to_lowercase() - .trim_start_matches('@') - .trim_start_matches('#') - .to_string() - }; - format!("{}:{}", kind.as_str(), clean) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_tree::score::extract::ExtractedEntity; - - fn entity(kind: EntityKind, text: &str) -> ExtractedEntity { - ExtractedEntity { - kind, - text: text.to_string(), - span_start: 0, - span_end: text.chars().count() as u32, - score: 1.0, - } - } - - #[test] - fn email_case_insensitive_canonicalises() { - let a = canonical_id_for(EntityKind::Email, "Alice@Example.com"); - let b = canonical_id_for(EntityKind::Email, "alice@example.com"); - assert_eq!(a, b); - assert_eq!(a, "email:alice@example.com"); - } - - #[test] - fn handle_strips_leading_at() { - let a = canonical_id_for(EntityKind::Handle, "@alice"); - let b = canonical_id_for(EntityKind::Handle, "alice"); - assert_eq!(a, b); - assert_eq!(a, "handle:alice"); - } - - #[test] - fn hashtag_strips_leading_hash() { - let a = canonical_id_for(EntityKind::Hashtag, "#launch"); - let b = canonical_id_for(EntityKind::Hashtag, "launch"); - assert_eq!(a, b); - } - - #[test] - fn url_preserves_case() { - let id = canonical_id_for(EntityKind::Url, " https://example.com/Path?Token=ABC "); - assert_eq!(id, "url:https://example.com/Path?Token=ABC"); - } - - #[test] - fn canonicalise_batch_preserves_spans() { - let ex = ExtractedEntities { - entities: vec![ - entity(EntityKind::Email, "Alice@Example.com"), - entity(EntityKind::Email, "alice@example.com"), - ], - topics: vec![], - llm_importance: None, - llm_importance_reason: None, - }; - let out = canonicalise(&ex); - assert_eq!(out.len(), 2); - // Both map to the same canonical id (merge-equivalent) - assert_eq!(out[0].canonical_id, out[1].canonical_id); - // But surface forms remain distinct - assert_ne!(out[0].surface, out[1].surface); - } - - #[test] - fn different_kinds_produce_different_ids_for_same_text() { - assert_ne!( - canonical_id_for(EntityKind::Handle, "alice"), - canonical_id_for(EntityKind::Person, "alice") - ); - } - - // ── Topic canonicalisation (#709 / Phase 3c topic-tree scope) ──── - - use crate::openhuman::memory_tree::score::extract::ExtractedTopic; - - fn topic(label: &str, score: f32) -> ExtractedTopic { - ExtractedTopic { - label: label.to_string(), - score, - } - } - - #[test] - fn topics_are_promoted_to_canonical_entities() { - let ex = ExtractedEntities { - entities: vec![], - topics: vec![topic("phoenix", 0.72), topic("migration", 0.60)], - llm_importance: None, - llm_importance_reason: None, - }; - let out = canonicalise(&ex); - assert_eq!(out.len(), 2); - assert_eq!(out[0].kind, EntityKind::Topic); - assert_eq!(out[0].canonical_id, "topic:phoenix"); - assert!((out[0].score - 0.72).abs() < 1e-6); - assert_eq!(out[1].canonical_id, "topic:migration"); - } - - #[test] - fn topic_canonicalisation_lowercases() { - let ex = ExtractedEntities { - entities: vec![], - topics: vec![topic("Phoenix", 1.0), topic("PHOENIX", 0.5)], - llm_importance: None, - llm_importance_reason: None, - }; - let out = canonicalise(&ex); - // Both normalise to "topic:phoenix" — second occurrence is deduped. - assert_eq!(out.len(), 1); - assert_eq!(out[0].canonical_id, "topic:phoenix"); - // First-seen surface is preserved. - assert_eq!(out[0].surface, "Phoenix"); - } - - #[test] - fn hashtag_and_topic_with_same_label_coexist() { - // "#launch" regex → EntityKind::Hashtag, LLM theme "launch" → Topic. - // They stay as two distinct canonical entities — different kind, - // different canonical_id prefix. - let ex = ExtractedEntities { - entities: vec![ExtractedEntity { - kind: EntityKind::Hashtag, - text: "launch".into(), - span_start: 0, - span_end: 6, - score: 1.0, - }], - topics: vec![topic("launch", 0.8)], - llm_importance: None, - llm_importance_reason: None, - }; - let out = canonicalise(&ex); - assert_eq!(out.len(), 2); - assert_eq!(out[0].kind, EntityKind::Hashtag); - assert_eq!(out[0].canonical_id, "hashtag:launch"); - assert_eq!(out[1].kind, EntityKind::Topic); - assert_eq!(out[1].canonical_id, "topic:launch"); - } - - #[test] - fn canonicalise_mixes_entities_and_topics_in_order() { - // Entities come first, topics appended after — downstream callers - // (e.g. routing) can rely on this ordering if they ever need it. - let ex = ExtractedEntities { - entities: vec![entity(EntityKind::Email, "alice@example.com")], - topics: vec![topic("phoenix", 0.7)], - llm_importance: None, - llm_importance_reason: None, - }; - let out = canonicalise(&ex); - assert_eq!(out.len(), 2); - assert_eq!(out[0].kind, EntityKind::Email); - assert_eq!(out[1].kind, EntityKind::Topic); - } - - #[test] - fn topic_entity_kind_round_trips_through_parse() { - // Defence in depth: ensure the new Topic variant survives the - // round-trip used by mem_tree_entity_index on read. - assert_eq!(EntityKind::parse("topic"), Ok(EntityKind::Topic)); - assert_eq!(EntityKind::Topic.as_str(), "topic"); - } -} diff --git a/src/openhuman/memory_tree/score/signals/interaction.rs b/src/openhuman/memory_tree/score/signals/interaction.rs deleted file mode 100644 index 622714860..000000000 --- a/src/openhuman/memory_tree/score/signals/interaction.rs +++ /dev/null @@ -1,105 +0,0 @@ -//! Interaction-weight signal — boosts chunks the user actively engaged with. -//! -//! Direct engagement is one of the strongest retention signals — "a message -//! you replied to" is almost always worth remembering, even if its content -//! looks noisy by other signals. -//! -//! Phase 2 infers engagement from a small set of reserved **tags**: -//! - `reply` — the user replied to this message/thread -//! - `sent` — the user authored this content -//! - `mention` — the user was @-mentioned -//! - `dm` — this arrived in a direct-message channel -//! -//! Ingest adapters can attach these tags during canonicalisation when the -//! upstream source supports the distinction. Absent tags → neutral score. - -use crate::openhuman::memory_store::chunks::types::Metadata; - -/// Tag set when the user replied to this message/thread. -pub const TAG_REPLY: &str = "reply"; -/// Tag set when the user authored this content. -pub const TAG_SENT: &str = "sent"; -/// Tag set when the user was @-mentioned. -pub const TAG_MENTION: &str = "mention"; -/// Tag set when the message arrived in a direct-message channel. -pub const TAG_DM: &str = "dm"; - -/// Score in `[0.0, 1.0]` based on engagement tags present on the chunk. -/// -/// Multiple tags stack (capped at 1.0): -/// - `sent` → +0.6 (author) -/// - `reply` → +0.5 (active dialogue) -/// - `dm` → +0.3 (scoped audience) -/// - `mention` → +0.2 (addressed) -/// -/// Absent any of these → 0.5 (neutral — don't drop the chunk on this signal -/// alone since most content lacks explicit engagement tags). -pub fn score(meta: &Metadata) -> f32 { - let mut any_tag = false; - let mut total: f32 = 0.0; - for t in &meta.tags { - match t.as_str() { - TAG_SENT => { - total += 0.6; - any_tag = true; - } - TAG_REPLY => { - total += 0.5; - any_tag = true; - } - TAG_DM => { - total += 0.3; - any_tag = true; - } - TAG_MENTION => { - total += 0.2; - any_tag = true; - } - _ => {} - } - } - if !any_tag { - return 0.5; - } - total.clamp(0.0, 1.0) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::chunks::types::SourceKind; - use chrono::Utc; - - fn meta(tags: &[&str]) -> Metadata { - let mut m = Metadata::point_in_time(SourceKind::Chat, "x", "owner", Utc::now()); - m.tags = tags.iter().map(|s| s.to_string()).collect(); - m - } - - #[test] - fn no_tags_neutral() { - assert_eq!(score(&meta(&[])), 0.5); - assert_eq!(score(&meta(&["unrelated"])), 0.5); - } - - #[test] - fn sent_tag_high_score() { - assert!((score(&meta(&["sent"])) - 0.6).abs() < 1e-6); - } - - #[test] - fn stacking_capped_at_one() { - // sent (0.6) + reply (0.5) + mention (0.2) = 1.3 → clamp to 1.0 - assert!((score(&meta(&["sent", "reply", "mention"])) - 1.0).abs() < 1e-6); - } - - #[test] - fn reply_only() { - assert!((score(&meta(&["reply"])) - 0.5).abs() < 1e-6); - } - - #[test] - fn dm_plus_mention() { - assert!((score(&meta(&["dm", "mention"])) - 0.5).abs() < 1e-6); - } -} diff --git a/src/openhuman/memory_tree/score/signals/metadata_weight.rs b/src/openhuman/memory_tree/score/signals/metadata_weight.rs deleted file mode 100644 index 880eb5c25..000000000 --- a/src/openhuman/memory_tree/score/signals/metadata_weight.rs +++ /dev/null @@ -1,49 +0,0 @@ -//! Metadata-weight signal — base weight from the source kind's grouping. -//! -//! The idea: a 1:1 email thread is inherently higher-signal than a broadcast -//! Slack channel, regardless of content. This signal captures the "shape" -//! of the interaction: how scoped is the audience? -//! -//! Phase 2 keeps this simple: one weight per `SourceKind`. Per-grouping -//! context (e.g., channel size, thread participant count) is a future -//! refinement when we actually have that metadata at ingest. - -use crate::openhuman::memory_store::chunks::types::{Metadata, SourceKind}; - -/// Base weight for each source kind. -/// -/// Email threads are typically scoped (1:1 or small groups, directed). -/// Documents are single-author outputs — high intentionality per chunk. -/// Chats vary widely; base weight is lower because the channel could be -/// a 200-person broadcast or a tight DM — the interaction signal disambiguates. -pub fn score(meta: &Metadata) -> f32 { - match meta.source_kind { - SourceKind::Email => 0.8, - SourceKind::Document => 0.9, - SourceKind::Chat => 0.5, - } -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::Utc; - - fn meta(kind: SourceKind) -> Metadata { - Metadata::point_in_time(kind, "x", "owner", Utc::now()) - } - - #[test] - fn per_kind_weights() { - assert!(score(&meta(SourceKind::Document)) > score(&meta(SourceKind::Email))); - assert!(score(&meta(SourceKind::Email)) > score(&meta(SourceKind::Chat))); - } - - #[test] - fn bounded_zero_one() { - for k in [SourceKind::Chat, SourceKind::Email, SourceKind::Document] { - let s = score(&meta(k)); - assert!((0.0..=1.0).contains(&s)); - } - } -} diff --git a/src/openhuman/memory_tree/score/signals/mod.rs b/src/openhuman/memory_tree/score/signals/mod.rs deleted file mode 100644 index 78ace044d..000000000 --- a/src/openhuman/memory_tree/score/signals/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -//! Score signals + weighted combine (Phase 2 / #708). -//! -//! Each submodule computes one scoring signal in `[0.0, 1.0]`. [`combine`] -//! aggregates them into a total score using per-signal weights. The output -//! is still `[0.0, 1.0]` after normalisation by total weight. -//! -//! Storing per-signal values alongside the total (via [`ScoreSignals`]) is -//! what makes admission decisions debuggable — when a chunk is dropped, we -//! persist *which* signals fired at what values. - -pub mod interaction; -pub mod metadata_weight; -mod ops; -pub mod source_weight; -pub mod token_count; -mod types; -pub mod unique_words; - -pub use ops::{combine, combine_cheap_only, compute, entity_density_score}; -pub use types::{ScoreSignals, SignalWeights}; diff --git a/src/openhuman/memory_tree/score/signals/ops.rs b/src/openhuman/memory_tree/score/signals/ops.rs deleted file mode 100644 index c9c905494..000000000 --- a/src/openhuman/memory_tree/score/signals/ops.rs +++ /dev/null @@ -1,182 +0,0 @@ -//! Cross-signal helpers: signal computation entry point and the two -//! weighted-combine variants (full and cheap-only) used by `score_chunk`. - -use super::{interaction, metadata_weight, source_weight, token_count, unique_words}; -use super::{ScoreSignals, SignalWeights}; -use crate::openhuman::memory_store::chunks::types::Metadata; -use crate::openhuman::memory_tree::score::extract::ExtractedEntities; - -/// Compute all signals for a chunk. -/// -/// `llm_importance` is sourced from `ex.llm_importance` (defaults to `0.0` -/// when the extractor didn't produce one — equivalent to "no LLM signal"). -pub fn compute( - meta: &Metadata, - content: &str, - token_count: u32, - ex: &ExtractedEntities, -) -> ScoreSignals { - ScoreSignals { - token_count: token_count::score(token_count), - unique_words: unique_words::score(content), - metadata_weight: metadata_weight::score(meta), - source_weight: source_weight::score(meta), - interaction: interaction::score(meta), - entity_density: entity_density_score(token_count, ex), - llm_importance: ex.llm_importance.unwrap_or(0.0).clamp(0.0, 1.0), - } -} - -/// Entity-density signal: entities per token, capped. -/// -/// More distinct entities per unit of content → more substantive. Calibrated -/// so ~1 entity per 100 tokens maxes out the signal. -pub fn entity_density_score(token_count: u32, ex: &ExtractedEntities) -> f32 { - let unique = ex.unique_entity_count() as f32; - if token_count == 0 { - return 0.0; - } - let per_token = unique / token_count as f32; - // cap at 0.01 entities/token = 1 entity per 100 tokens - (per_token / 0.01).min(1.0) -} - -/// Weighted sum of signals, normalised to `[0.0, 1.0]`. -/// -/// When `w.llm_importance == 0.0` (the default) the LLM signal contributes -/// nothing to either the numerator or the denominator — output is identical -/// to pre-LLM Phase 2. -pub fn combine(signals: &ScoreSignals, w: &SignalWeights) -> f32 { - let total_weight = w.token_count - + w.unique_words - + w.metadata_weight - + w.source_weight - + w.interaction - + w.entity_density - + w.llm_importance; - if total_weight <= 0.0 { - return 0.0; - } - let weighted = signals.token_count * w.token_count - + signals.unique_words * w.unique_words - + signals.metadata_weight * w.metadata_weight - + signals.source_weight * w.source_weight - + signals.interaction * w.interaction - + signals.entity_density * w.entity_density - + signals.llm_importance * w.llm_importance; - (weighted / total_weight).clamp(0.0, 1.0) -} - -/// Weighted sum **excluding the `llm_importance` signal**. -/// -/// Used by the short-circuit logic in `score_chunk`: if the deterministic -/// (cheap-signals-only) total is already firmly above or below the -/// admission band, we skip the LLM call entirely. The LLM signal only -/// participates in the *final* `combine` once it's been computed. -pub fn combine_cheap_only(signals: &ScoreSignals, w: &SignalWeights) -> f32 { - let total_weight = w.token_count - + w.unique_words - + w.metadata_weight - + w.source_weight - + w.interaction - + w.entity_density; - if total_weight <= 0.0 { - return 0.0; - } - let weighted = signals.token_count * w.token_count - + signals.unique_words * w.unique_words - + signals.metadata_weight * w.metadata_weight - + signals.source_weight * w.source_weight - + signals.interaction * w.interaction - + signals.entity_density * w.entity_density; - (weighted / total_weight).clamp(0.0, 1.0) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory_store::chunks::types::SourceKind; - use crate::openhuman::memory_tree::score::extract::{ - EntityKind, ExtractedEntities, ExtractedEntity, - }; - use chrono::Utc; - - fn meta(tags: &[&str], kind: SourceKind) -> Metadata { - let mut m = Metadata::point_in_time(kind, "x", "owner", Utc::now()); - m.tags = tags.iter().map(|s| s.to_string()).collect(); - m - } - - fn make_entities(n: usize) -> ExtractedEntities { - ExtractedEntities { - entities: (0..n) - .map(|i| ExtractedEntity { - kind: EntityKind::Email, - text: format!("user{i}@example.com"), - span_start: 0, - span_end: 10, - score: 1.0, - }) - .collect(), - ..Default::default() - } - } - - #[test] - fn combine_all_zeros_is_zero() { - let s = ScoreSignals::default(); - assert!(combine(&s, &SignalWeights::default()) < 0.01); - } - - #[test] - fn combine_all_ones_is_one() { - let s = ScoreSignals { - token_count: 1.0, - unique_words: 1.0, - metadata_weight: 1.0, - source_weight: 1.0, - interaction: 1.0, - entity_density: 1.0, - llm_importance: 0.0, // default weight is 0 → contribution is zero - }; - assert!((combine(&s, &SignalWeights::default()) - 1.0).abs() < 1e-6); - } - - #[test] - fn weights_influence_total() { - let s = ScoreSignals { - token_count: 0.0, - unique_words: 0.0, - metadata_weight: 0.0, - source_weight: 0.0, - interaction: 1.0, - entity_density: 0.0, - llm_importance: 0.0, - }; - let total = combine(&s, &SignalWeights::default()); - assert!((total - (3.0 / 9.0)).abs() < 1e-6); - } - - #[test] - fn compute_wires_all_signals() { - let m = meta(&["reply"], SourceKind::Email); - let ex = make_entities(3); - let s = compute( - &m, - "Some substantive text about Phoenix launch planning.", - 12, - &ex, - ); - assert!(s.interaction > 0.0); - assert!(s.metadata_weight > 0.0); - assert!(s.source_weight > 0.0); - } - - #[test] - fn entity_density_scales() { - let ex = make_entities(1); - assert!((entity_density_score(100, &ex) - 1.0).abs() < 1e-6); - assert!((entity_density_score(1000, &ex) - 0.1).abs() < 1e-6); - assert_eq!(entity_density_score(0, &ex), 0.0); - } -} diff --git a/src/openhuman/memory_tree/score/signals/source_weight.rs b/src/openhuman/memory_tree/score/signals/source_weight.rs deleted file mode 100644 index 80cd30770..000000000 --- a/src/openhuman/memory_tree/score/signals/source_weight.rs +++ /dev/null @@ -1,112 +0,0 @@ -//! Source-weight signal — per-provider base weight derived from the -//! `DataSource` when it can be inferred from a chunk's tags. -//! -//! Rationale from `Memory Architecture.md` (Step 2.3 "Source scoring"): -//! - High-intentionality messaging (direct DMs, personal emails) scores higher -//! - Broadcast/channel content scores lower -//! - Documents authored by the user score higher than shared-but-unmodified drops -//! -//! Phase 2 takes a conservative approach: per-[`DataSource`] base weight. -//! Finer distinction (DM vs channel on Slack specifically) requires richer -//! ingest-time metadata and is deferred. - -use crate::openhuman::memory_store::chunks::types::{DataSource, Metadata, SourceKind}; - -const PROVIDER_PREFIX: &str = "provider:"; - -/// Best-effort map from `Metadata` to a [`DataSource`] — checks the `tags` -/// list for a stable `provider:` provider tag. If not present, -/// falls back to kind-based defaults. -/// -/// The ingestion pipeline can (and should) add a provider tag on the -/// canonicalised output so this signal fires deterministically. Until that's -/// wired everywhere, we fall back to the kind-level default. -pub fn infer_data_source(meta: &Metadata) -> Option { - for tag in &meta.tags { - let Some(provider) = tag.strip_prefix(PROVIDER_PREFIX) else { - continue; - }; - if let Ok(ds) = DataSource::parse(provider) { - return Some(ds); - } - } - None -} - -/// Score in `[0.0, 1.0]` for the chunk's originating provider. -pub fn score(meta: &Metadata) -> f32 { - if let Some(ds) = infer_data_source(meta) { - return weight_for(ds); - } - // Fallback: kind-level defaults consistent with per-provider averages. - match meta.source_kind { - SourceKind::Email => 0.75, - SourceKind::Document => 0.7, - SourceKind::Chat => 0.5, - } -} - -fn weight_for(ds: DataSource) -> f32 { - match ds { - // Personal email providers score high — typically small, directed audiences - DataSource::Gmail => 0.8, - DataSource::OtherEmail => 0.7, - // Chat providers differ: WhatsApp is typically DM-heavy, Discord - // can be broadcast-heavy, Telegram mixes both - DataSource::Whatsapp => 0.75, - DataSource::Telegram => 0.6, - DataSource::Discord => 0.5, - // Agent conversations — high signal, direct interaction with the user - DataSource::Conversation => 0.9, - // Documents: Notion = structured, Drive = mixed, Meeting notes = high value - DataSource::Notion => 0.75, - DataSource::DriveDocs => 0.6, - DataSource::MeetingNotes => 0.85, - } -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::Utc; - - fn meta_with_tag(kind: SourceKind, tag: &str) -> Metadata { - let mut m = Metadata::point_in_time(kind, "x", "owner", Utc::now()); - m.tags.push(tag.to_string()); - m - } - - #[test] - fn data_source_inferred_from_tags() { - let m = meta_with_tag(SourceKind::Chat, "provider:whatsapp"); - assert_eq!(infer_data_source(&m), Some(DataSource::Whatsapp)); - } - - #[test] - fn plain_user_label_does_not_infer_provider() { - let m = meta_with_tag(SourceKind::Email, "notion"); - assert_eq!(infer_data_source(&m), None); - assert!((score(&m) - 0.75).abs() < 1e-6); - } - - #[test] - fn unknown_tag_falls_back_to_kind_default() { - let m = meta_with_tag(SourceKind::Email, "not-a-data-source"); - let s = score(&m); - assert!((s - 0.75).abs() < 1e-6); - } - - #[test] - fn provider_specific_weights_applied() { - let m = meta_with_tag(SourceKind::Document, "provider:meeting_notes"); - assert!((score(&m) - 0.85).abs() < 1e-6); - } - - #[test] - fn all_data_sources_bounded() { - for ds in DataSource::all() { - let w = weight_for(*ds); - assert!((0.0..=1.0).contains(&w)); - } - } -} diff --git a/src/openhuman/memory_tree/score/signals/token_count.rs b/src/openhuman/memory_tree/score/signals/token_count.rs deleted file mode 100644 index 9515a836d..000000000 --- a/src/openhuman/memory_tree/score/signals/token_count.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! Token-count signal — penalises very short or very long chunks. -//! -//! Rationale: "+1", "lol", "👍" are usually noise; multi-page walls of text -//! are often pasted logs or attachments that overwhelm summarisation. -//! The signal is strongest in a middle band that corresponds to substantive -//! prose/discussion. -//! -//! Output is a score in `[0.0, 1.0]` shaped as a plateau between -//! `TOKEN_MIN` and `TOKEN_MAX` with linear ramps on both sides. - -/// Below this token count the chunk scores 0 (treated as noise). -pub const TOKEN_MIN: u32 = 10; -/// Top of the linear ramp from 0 → 1 starting at [`TOKEN_MIN`]. -pub const TOKEN_RAMP_LOW: u32 = 30; -/// Start of the linear ramp from 1 → 0.5 ending at [`TOKEN_MAX`]. -pub const TOKEN_RAMP_HIGH: u32 = 3_000; -/// Above this token count the score is clamped to 0.5 (oversized content -/// still carries information but loses the plateau bonus). -pub const TOKEN_MAX: u32 = 8_000; - -/// Score for a chunk's token count. See module docs for shape. -pub fn score(token_count: u32) -> f32 { - if token_count < TOKEN_MIN { - return 0.0; - } - if token_count <= TOKEN_RAMP_LOW { - // linear 0..1 over [MIN, RAMP_LOW] - let span = (TOKEN_RAMP_LOW - TOKEN_MIN) as f32; - return (token_count - TOKEN_MIN) as f32 / span; - } - if token_count <= TOKEN_RAMP_HIGH { - return 1.0; - } - if token_count <= TOKEN_MAX { - // linear 1.0..0.5 over [RAMP_HIGH, MAX] - let span = (TOKEN_MAX - TOKEN_RAMP_HIGH) as f32; - let t = (token_count - TOKEN_RAMP_HIGH) as f32 / span; - return 1.0 - 0.5 * t; - } - 0.5 -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn tiny_is_zero() { - assert_eq!(score(0), 0.0); - assert_eq!(score(5), 0.0); - assert_eq!(score(9), 0.0); - } - - #[test] - fn ramp_up_linear() { - // score(MIN) = 0, score(RAMP_LOW) = 1.0 - assert!((score(TOKEN_MIN) - 0.0).abs() < 1e-4); - assert!((score(TOKEN_RAMP_LOW) - 1.0).abs() < 1e-4); - // midpoint ~0.5 - let mid = TOKEN_MIN + (TOKEN_RAMP_LOW - TOKEN_MIN) / 2; - assert!((score(mid) - 0.5).abs() < 0.05); - } - - #[test] - fn plateau_is_one() { - assert_eq!(score(200), 1.0); - assert_eq!(score(1000), 1.0); - assert_eq!(score(TOKEN_RAMP_HIGH), 1.0); - } - - #[test] - fn ramp_down_to_half() { - assert!((score(TOKEN_MAX) - 0.5).abs() < 1e-4); - assert_eq!(score(TOKEN_MAX + 10_000), 0.5); - } - - #[test] - fn monotonic_in_bands() { - // Strictly increasing on the up-ramp - assert!(score(TOKEN_MIN + 1) < score(TOKEN_RAMP_LOW - 1)); - // Strictly decreasing on the down-ramp - assert!(score(TOKEN_RAMP_HIGH + 1) > score(TOKEN_MAX - 1)); - } -} diff --git a/src/openhuman/memory_tree/score/signals/types.rs b/src/openhuman/memory_tree/score/signals/types.rs deleted file mode 100644 index 2aa40115a..000000000 --- a/src/openhuman/memory_tree/score/signals/types.rs +++ /dev/null @@ -1,107 +0,0 @@ -//! Strongly-typed bag of per-signal scores plus the weights used to combine -//! them. Persisted alongside the total in `mem_tree_score` so a chunk's -//! admit/drop decision is auditable after the fact. - -use serde::{Deserialize, Serialize}; - -/// Per-signal score breakdown for one chunk. Persisted alongside the total -/// for diagnostics. -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub struct ScoreSignals { - pub token_count: f32, - pub unique_words: f32, - pub metadata_weight: f32, - pub source_weight: f32, - pub interaction: f32, - pub entity_density: f32, - /// LLM-derived importance rating in `[0.0, 1.0]`. `0.0` when no LLM - /// signal is available — combined with `SignalWeights::llm_importance = 0.0` - /// (the default) this produces a no-op contribution to the total, keeping - /// behaviour identical to pre-LLM Phase 2. - #[serde(default)] - pub llm_importance: f32, -} - -/// Default weights applied to each signal in `combine`. -/// -/// `llm_importance` defaults to `0.0` (disabled). Callers who configure an -/// LLM extractor should bump it (typical: 2.0 — comparable to the -/// metadata/source weights, well below the interaction-direct signal). -#[derive(Clone, Debug)] -pub struct SignalWeights { - pub token_count: f32, - pub unique_words: f32, - pub metadata_weight: f32, - pub source_weight: f32, - pub interaction: f32, - pub entity_density: f32, - pub llm_importance: f32, -} - -impl Default for SignalWeights { - fn default() -> Self { - Self { - token_count: 1.0, - unique_words: 1.0, - metadata_weight: 1.5, - source_weight: 1.5, - interaction: 3.0, // strongest signal — direct user engagement - entity_density: 1.0, - llm_importance: 0.0, // disabled until LLM extractor is configured - } - } -} - -impl SignalWeights { - /// Same as [`Default::default`] but with a non-zero `llm_importance` weight. - /// Use when an LLM extractor is wired in and you want its importance - /// signal to influence the admission decision. - pub fn with_llm_enabled() -> Self { - Self { - llm_importance: 2.0, - ..Self::default() - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn score_signals_default_to_zero() { - let signals = ScoreSignals::default(); - assert_eq!(signals.token_count, 0.0); - assert_eq!(signals.unique_words, 0.0); - assert_eq!(signals.metadata_weight, 0.0); - assert_eq!(signals.source_weight, 0.0); - assert_eq!(signals.interaction, 0.0); - assert_eq!(signals.entity_density, 0.0); - assert_eq!(signals.llm_importance, 0.0); - } - - #[test] - fn signal_weights_default_match_expected_priorities() { - let weights = SignalWeights::default(); - assert_eq!(weights.token_count, 1.0); - assert_eq!(weights.unique_words, 1.0); - assert_eq!(weights.metadata_weight, 1.5); - assert_eq!(weights.source_weight, 1.5); - assert_eq!(weights.interaction, 3.0); - assert_eq!(weights.entity_density, 1.0); - assert_eq!(weights.llm_importance, 0.0); - } - - #[test] - fn with_llm_enabled_only_changes_llm_weight() { - let default = SignalWeights::default(); - let enabled = SignalWeights::with_llm_enabled(); - assert_eq!(enabled.token_count, default.token_count); - assert_eq!(enabled.unique_words, default.unique_words); - assert_eq!(enabled.metadata_weight, default.metadata_weight); - assert_eq!(enabled.source_weight, default.source_weight); - assert_eq!(enabled.interaction, default.interaction); - assert_eq!(enabled.entity_density, default.entity_density); - assert_eq!(enabled.llm_importance, 2.0); - } -} diff --git a/src/openhuman/memory_tree/score/signals/unique_words.rs b/src/openhuman/memory_tree/score/signals/unique_words.rs deleted file mode 100644 index 8436d62f2..000000000 --- a/src/openhuman/memory_tree/score/signals/unique_words.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! Unique-word-ratio signal — noise detector that fires on low-diversity text. -//! -//! Example: "yay yay yay yay lol lol lol" has high repetition = low diversity. -//! A substantive message has high type-token ratio (roughly, unique words / -//! total words). -//! -//! For very short messages the ratio is naturally ~1.0, so we require a -//! minimum total count before this signal contributes — otherwise "hi bob" -//! would score identically to a real message. - -/// Below this total-word count the type-token ratio is unreliable, so the -/// signal returns a neutral 0.5 instead of computing a ratio. -pub const MIN_TOTAL_WORDS: usize = 5; - -/// Score in `[0.0, 1.0]` from the type-token ratio of `text`. -/// -/// - Too few total words → `0.5` (indeterminate — defer to other signals) -/// - Ratio < 0.3 (heavy repetition) → 0.0 -/// - Ratio >= 0.7 (substantive) → 1.0 -/// - Linear in between -pub fn score(text: &str) -> f32 { - let mut total: usize = 0; - // Only `uniq.len()` is ever read — never the iteration order — so a hash set - // gives the identical type-token ratio with O(1) inserts instead of the - // ordered set's O(log n) String comparisons per word. - let mut uniq: std::collections::HashSet = std::collections::HashSet::new(); - - for raw in text.split_whitespace() { - let w: String = raw - .trim_matches(|c: char| !c.is_alphanumeric()) - .to_lowercase(); - if w.is_empty() { - continue; - } - total += 1; - uniq.insert(w); - } - - if total < MIN_TOTAL_WORDS { - return 0.5; - } - - let ratio = uniq.len() as f32 / total as f32; - if ratio <= 0.3 { - 0.0 - } else if ratio >= 0.7 { - 1.0 - } else { - (ratio - 0.3) / 0.4 - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn short_text_returns_neutral() { - assert_eq!(score(""), 0.5); - assert_eq!(score("hi bob"), 0.5); - } - - #[test] - fn high_repetition_scored_low() { - let noisy = "yay yay yay yay yay yay yay yay yay yay lol lol lol lol"; - assert!(score(noisy) < 0.2); - } - - #[test] - fn substantive_text_scored_high() { - let good = - "We decided to ship Phoenix on Friday after reviewing the migration plan carefully."; - assert!(score(good) >= 0.9); - } - - #[test] - fn medium_repetition_ramps() { - // ~50% unique ratio should score around 0.5 - let med = "alpha beta alpha beta gamma alpha delta beta gamma alpha"; - let s = score(med); - assert!(s > 0.2 && s < 0.8); - } - - #[test] - fn punctuation_stripped() { - let s1 = score("ship phoenix friday ship phoenix friday ship phoenix"); - let s2 = score("ship! phoenix, friday. ship! phoenix, friday. ship! phoenix."); - assert!((s1 - s2).abs() < 0.05); - } -} diff --git a/src/openhuman/memory_tree/score/store.rs b/src/openhuman/memory_tree/score/store.rs index 8f09c67f1..48c7c16c1 100644 --- a/src/openhuman/memory_tree/score/store.rs +++ b/src/openhuman/memory_tree/score/store.rs @@ -1,345 +1,83 @@ -//! Persistence for Phase 2 artefacts (#708): -//! -//! - `mem_tree_score` — per-chunk score rationale (which signals fired, why -//! dropped/kept) -//! - `mem_tree_entity_index` — inverted index `entity_id → node_id` so -//! retrieval can resolve entity-scoped queries in O(lookup) -//! -//! Schema is declared in `memory/tree/store.rs::SCHEMA`; this file only -//! owns the CRUD operations. +//! Product Config adapters over tinycortex score and entity-index persistence. use std::collections::HashMap; -use anyhow::{Context, Result}; -use rusqlite::{params, Connection, OptionalExtension, Transaction}; -use serde::{Deserialize, Serialize}; +use anyhow::Result; +use rusqlite::Transaction; -use crate::openhuman::composio::providers::profile::{is_self_identity_any_toolkit, IdentityKind}; use crate::openhuman::config::Config; -use crate::openhuman::memory_store::chunks::store::with_connection; -use crate::openhuman::memory_tree::score::extract::EntityKind; -use crate::openhuman::memory_tree::score::resolver::CanonicalEntity; -use crate::openhuman::memory_tree::score::signals::ScoreSignals; -/// Map a memory-tree `EntityKind` to the Composio identity-registry -/// [`IdentityKind`] used for self-matching, or `None` for kinds that -/// don't represent identity (Url, Hashtag, Topic, Org, Loc, ...). `Person` -/// is intentionally omitted — display-name matches are weak and would -/// false-positive any contact with a similar name. -fn entity_kind_to_identity_kind(k: EntityKind) -> Option { - Some(match k { - EntityKind::Email => IdentityKind::Email, - EntityKind::Handle => IdentityKind::Handle, - _ => return None, - }) +pub use tinycortex::memory::score::store::{EntityHit, ScoreRow}; + +fn memory_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -/// Resolve `is_user` for one canonical entity — true iff the surface -/// matches a self-handle of a matchable kind in the identity registry. -fn entity_is_user(entity: &CanonicalEntity) -> bool { - let Some(kind) = entity_kind_to_identity_kind(entity.kind) else { - return false; - }; - is_self_identity_any_toolkit(kind, &entity.surface) -} - -/// Same as [`entity_is_user`] but for the summary-index path where only -/// the canonical id (`":"`) is in scope. Returns `false` if -/// the id is malformed or the kind isn't matchable. -fn canonical_id_is_user(canonical_id: &str) -> bool { - let Some((kind_str, value)) = canonical_id.split_once(':') else { - return false; - }; - let Ok(kind) = EntityKind::parse(kind_str) else { - return false; - }; - let Some(idk) = entity_kind_to_identity_kind(kind) else { - return false; - }; - is_self_identity_any_toolkit(idk, value) -} - -/// Serialized per-chunk score rationale. Mirrors the `mem_tree_score` row. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ScoreRow { - pub chunk_id: String, - pub total: f32, - pub signals: ScoreSignals, - pub dropped: bool, - pub reason: Option, - pub computed_at_ms: i64, - /// One-line LLM-supplied explanation for the importance rating; useful - /// for tuning prompts and thresholds. The numeric value lives on - /// `signals.llm_importance`. - #[serde(default)] - pub llm_importance_reason: Option, -} - -/// Upsert one score rationale row, replacing any existing entry for `chunk_id`. pub fn upsert_score(config: &Config, row: &ScoreRow) -> Result<()> { - with_connection(config, |conn| { - upsert_score_on_connection(conn, row)?; - Ok(()) - }) + tinycortex::memory::score::store::upsert_score(&memory_config(config), row) } pub(crate) fn upsert_score_tx(tx: &Transaction<'_>, row: &ScoreRow) -> Result<()> { - tx.execute( - SCORE_UPSERT_SQL, - params![ - row.chunk_id, - row.total, - row.signals.token_count, - row.signals.unique_words, - row.signals.metadata_weight, - row.signals.source_weight, - row.signals.interaction, - row.signals.entity_density, - row.signals.llm_importance, - row.llm_importance_reason, - i32::from(row.dropped), - row.reason, - row.computed_at_ms, - ], - )?; - Ok(()) + tinycortex::memory::score::store::upsert_score_tx(tx, row) } -const SCORE_UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_score ( - chunk_id, total, - token_count_signal, unique_words_signal, - metadata_weight, source_weight, interaction_weight, entity_density, - llm_importance, llm_importance_reason, - dropped, reason, computed_at_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)"; - -fn upsert_score_on_connection(conn: &Connection, row: &ScoreRow) -> Result<()> { - conn.execute( - SCORE_UPSERT_SQL, - params![ - row.chunk_id, - row.total, - row.signals.token_count, - row.signals.unique_words, - row.signals.metadata_weight, - row.signals.source_weight, - row.signals.interaction, - row.signals.entity_density, - row.signals.llm_importance, - row.llm_importance_reason, - i32::from(row.dropped), - row.reason, - row.computed_at_ms, - ], - )?; - Ok(()) -} - -/// Fetch one chunk's score rationale. pub fn get_score(config: &Config, chunk_id: &str) -> Result> { - with_connection(config, |conn| { - conn.query_row( - "SELECT chunk_id, total, - token_count_signal, unique_words_signal, - metadata_weight, source_weight, interaction_weight, entity_density, - llm_importance, llm_importance_reason, - dropped, reason, computed_at_ms - FROM mem_tree_score WHERE chunk_id = ?1", - params![chunk_id], - |row| { - Ok(ScoreRow { - chunk_id: row.get(0)?, - total: row.get(1)?, - signals: ScoreSignals { - token_count: row.get(2)?, - unique_words: row.get(3)?, - metadata_weight: row.get(4)?, - source_weight: row.get(5)?, - interaction: row.get(6)?, - entity_density: row.get(7)?, - llm_importance: row.get::<_, Option>(8)?.unwrap_or(0.0), - }, - llm_importance_reason: row.get::<_, Option>(9)?, - dropped: row.get::<_, i32>(10)? != 0, - reason: row.get(11)?, - computed_at_ms: row.get(12)?, - }) - }, - ) - .optional() - .map_err(anyhow::Error::from) - }) + tinycortex::memory::score::store::get_score(&memory_config(config), chunk_id) } -/// Defensive cap for batched `IN (?,?,…)` reads. -/// -/// See identical rationale on [`crate::openhuman::memory_store::chunks::store`]'s -/// `MAX_FETCH_BATCH`: SQLite's `SQLITE_MAX_VARIABLE_NUMBER` has been -/// 32 766 since 3.32, so 500 leaves a ~65× safety margin. The -/// `fetch_leaves` call-site is capped at 20 ids so the chunked loop -/// runs exactly once today — the window exists purely so future -/// call-sites passing larger id lists do not blow up against a host -/// with a lower compile-time SQLite cap. -const MAX_FETCH_BATCH: usize = 500; - -/// Batched read of just the `total` field for many chunk ids. -/// -/// Narrow on purpose — `fetch_leaves` only needs the float, not the -/// full [`ScoreRow`] with all signals. The returned map contains only -/// `chunk_id`s that have a score row; missing ids are silently absent, -/// matching the per-row [`get_score`] contract (callers then fall back -/// to the documented 0.0 neutral). pub fn get_scores_batch(config: &Config, chunk_ids: &[String]) -> Result> { - if chunk_ids.is_empty() { - return Ok(HashMap::new()); - } - log::debug!( - "[memory_tree::score] get_scores_batch: n={} windows={}", - chunk_ids.len(), - chunk_ids.len().div_ceil(MAX_FETCH_BATCH) - ); - with_connection(config, |conn| { - let mut out: HashMap = HashMap::with_capacity(chunk_ids.len()); - for window in chunk_ids.chunks(MAX_FETCH_BATCH) { - let placeholders = (1..=window.len()) - .map(|i| format!("?{i}")) - .collect::>() - .join(","); - let sql = format!( - "SELECT chunk_id, total FROM mem_tree_score - WHERE chunk_id IN ({placeholders})" - ); - let mut stmt = conn.prepare(&sql).context("prepare get_scores_batch")?; - let params: Vec<&dyn rusqlite::ToSql> = - window.iter().map(|id| id as &dyn rusqlite::ToSql).collect(); - let rows = stmt - .query_map(params.as_slice(), |row| { - Ok((row.get::<_, String>(0)?, row.get::<_, f32>(1)?)) - }) - .context("query get_scores_batch")?; - for row in rows { - let (chunk_id, total) = row.context("decode get_scores_batch row")?; - out.insert(chunk_id, total); - } - } - log::debug!( - "[memory_tree::score] get_scores_batch: matched {}/{} ids", - out.len(), - chunk_ids.len() - ); - Ok(out) - }) + tinycortex::memory::score::store::get_scores_batch(&memory_config(config), chunk_ids) } -/// Index one (entity, chunk) association. -/// -/// Idempotent on the composite primary key `(entity_id, node_id)` so -/// re-indexing the same association is a no-op update. +pub use crate::openhuman::memory_store::entities::{ + clear_entity_index_for_node, count_entity_index, list_entity_ids_for_node, lookup_entity, +}; + pub fn index_entity( config: &Config, - entity: &CanonicalEntity, + entity: &tinycortex::memory::score::resolver::CanonicalEntity, node_id: &str, node_kind: &str, timestamp_ms: i64, tree_id: Option<&str>, ) -> Result<()> { - let is_user = entity_is_user(entity); - with_connection(config, |conn| { - conn.execute( - "INSERT OR REPLACE INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", - params![ - entity.canonical_id, - node_id, - node_kind, - entity.kind.as_str(), - entity.surface, - entity.score, - timestamp_ms, - tree_id, - is_user as i32, - ], - )?; - Ok(()) - }) + let entity = to_store_entity(entity)?; + crate::openhuman::memory_store::entities::index_entity( + config, + &entity, + node_id, + node_kind, + timestamp_ms, + tree_id, + ) } -/// Batch index all entities extracted from a chunk. pub fn index_entities( config: &Config, - entities: &[CanonicalEntity], + entities: &[tinycortex::memory::score::resolver::CanonicalEntity], node_id: &str, node_kind: &str, timestamp_ms: i64, tree_id: Option<&str>, ) -> Result { - if entities.is_empty() { - return Ok(0); - } - with_connection(config, |conn| { - let tx = conn.unchecked_transaction()?; - { - let mut stmt = tx.prepare( - "INSERT OR REPLACE INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", - )?; - for e in entities { - stmt.execute(params![ - e.canonical_id, - node_id, - node_kind, - e.kind.as_str(), - e.surface, - e.score, - timestamp_ms, - tree_id, - entity_is_user(e) as i32, - ])?; - } - } - tx.commit()?; - Ok(entities.len()) - }) -} - -/// Remove all entity-index rows for a given node. Used before re-indexing -/// a re-scored chunk so entities dropped from the new extraction don't leak -/// through as stale `INSERT OR REPLACE` never deletes. -pub fn clear_entity_index_for_node(config: &Config, node_id: &str) -> Result { - with_connection(config, |conn| { - let n = conn.execute( - "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", - params![node_id], - )?; - Ok(n) - }) + let entities: Vec = entities + .iter() + .map(to_store_entity) + .collect::>()?; + crate::openhuman::memory_store::entities::index_entities( + config, + &entities, + node_id, + node_kind, + timestamp_ms, + tree_id, + ) } pub(crate) fn clear_entity_index_for_node_tx(tx: &Transaction<'_>, node_id: &str) -> Result { - let n = tx.execute( - "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", - params![node_id], - )?; - Ok(n) + tinycortex::memory::score::store::clear_entity_index_for_node_tx(tx, node_id) } -/// Index summary-node entities by canonical id only. Summary-level entity -/// metadata is LLM-derived (Phase 3a #709) — the summariser emits a -/// curated list of canonical ids without per-occurrence span/surface data. -/// -/// Writes the kind prefix (everything before the first `:`) into the -/// `entity_kind` column so [`lookup_entity`]'s `EntityKind::parse()` keeps -/// round-tripping on summary rows. `surface` stores the full canonical id -/// as a stable placeholder — at the summary level we have no per-occurrence -/// span to recover, and the id is always unique. The summary's score is -/// reused for each of its entities. -/// -/// Callers should prefer the regular [`index_entities_tx`] for leaves, -/// where span/surface are meaningful. pub(crate) fn index_summary_entity_ids_tx( tx: &Transaction<'_>, entity_ids: &[String], @@ -348,179 +86,56 @@ pub(crate) fn index_summary_entity_ids_tx( timestamp_ms: i64, tree_id: Option<&str>, ) -> Result { - if entity_ids.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare( - "INSERT OR REPLACE INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", - )?; - for canonical_id in entity_ids { - // Canonical ids follow Phase 2's ":" convention. - // Without this split, `entity_kind` would hold the full id and - // `lookup_entity`'s `EntityKind::parse()` would fail at read time, - // poisoning any mixed leaf/summary lookup. - let entity_kind = match canonical_id.split_once(':') { - Some((kind, _)) => kind, - None => { - log::warn!( - "[memory_tree::score::store] summary entity id missing ':' — \ - storing as-is: {canonical_id}" - ); - canonical_id.as_str() - } - }; - stmt.execute(params![ - canonical_id, - node_id, - "summary", - entity_kind, - canonical_id, - score, - timestamp_ms, - tree_id, - canonical_id_is_user(canonical_id) as i32, - ])?; - } - Ok(entity_ids.len()) + let identity = crate::openhuman::memory_store::entities::host_self_identity(); + tinycortex::memory::store::entity_index::index_summary_entity_ids_tx_with_identity( + tx, + entity_ids, + node_id, + score, + timestamp_ms, + tree_id, + identity.as_ref(), + ) } pub(crate) fn index_entities_tx( tx: &Transaction<'_>, - entities: &[CanonicalEntity], + entities: &[tinycortex::memory::score::resolver::CanonicalEntity], node_id: &str, node_kind: &str, timestamp_ms: i64, tree_id: Option<&str>, ) -> Result { - if entities.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare( - "INSERT OR REPLACE INTO mem_tree_entity_index ( - entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", - )?; - for e in entities { - stmt.execute(params![ - e.canonical_id, - node_id, - node_kind, - e.kind.as_str(), - e.surface, - e.score, - timestamp_ms, - tree_id, - entity_is_user(e) as i32, - ])?; - } - Ok(entities.len()) + let identity = crate::openhuman::memory_store::entities::host_self_identity(); + let entities: Vec = entities + .iter() + .map(to_store_entity) + .collect::>()?; + tinycortex::memory::store::entity_index::index_entities_tx_with_identity( + tx, + &entities, + node_id, + node_kind, + timestamp_ms, + tree_id, + identity.as_ref(), + ) } -/// Result row from [`lookup_entity`]. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct EntityHit { - pub entity_id: String, - pub node_id: String, - pub node_kind: String, - pub entity_kind: EntityKind, - pub surface: String, - pub score: f32, - pub timestamp_ms: i64, - pub tree_id: Option, - /// #1365: true when the canonical id matched the Composio identity - /// registry at index time (e.g. `email:cyrus@example.com` matches the - /// user's Gmail). Subconscious filters/weights by this flag so - /// first-person reflections only quote first-person sources. - #[serde(default)] - pub is_user: bool, -} - -/// Find all nodes indexed against `entity_id`, newest first. -pub fn lookup_entity( - config: &Config, - entity_id: &str, - limit: Option, -) -> Result> { - // Clamp to i64::MAX before casting so callers can't wrap a large usize - // into a negative LIMIT and bypass it. - let limit = limit.unwrap_or(100).min(i64::MAX as usize) as i64; - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - FROM mem_tree_entity_index - WHERE entity_id = ?1 - ORDER BY timestamp_ms DESC - LIMIT ?2", - )?; - let rows = stmt - .query_map(params![entity_id, limit], |row| { - let kind_s: String = row.get(3)?; - let entity_kind = EntityKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 3, - rusqlite::types::Type::Text, - e.into(), - ) - })?; - let is_user_int: i32 = row.get(8)?; - Ok(EntityHit { - entity_id: row.get(0)?, - node_id: row.get(1)?, - node_kind: row.get(2)?, - entity_kind, - surface: row.get(4)?, - score: row.get(5)?, - timestamp_ms: row.get(6)?, - tree_id: row.get(7)?, - is_user: is_user_int != 0, - }) - })? - .collect::>>()?; - Ok(rows) +fn to_store_entity( + entity: &tinycortex::memory::score::resolver::CanonicalEntity, +) -> Result { + Ok(tinycortex::memory::store::CanonicalEntity { + canonical_id: entity.canonical_id.clone(), + kind: tinycortex::memory::store::EntityKind::parse(entity.kind.as_str()) + .map_err(anyhow::Error::msg)?, + surface: entity.surface.clone(), + span_start: entity.span_start, + span_end: entity.span_end, + score: entity.score, }) } -/// All distinct canonical entity ids associated with `node_id`, ordered by -/// score (desc) then recency. Used by topic-routing to pick which topic -/// trees a node should fan into. -pub fn list_entity_ids_for_node(config: &Config, node_id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT DISTINCT entity_id - FROM mem_tree_entity_index - WHERE node_id = ?1 - ORDER BY score DESC, timestamp_ms DESC, entity_id ASC", - )?; - let rows = stmt - .query_map(params![node_id], |row| row.get::<_, String>(0))? - .collect::>>()?; - Ok(rows) - }) -} - -/// Count rows in the entity index (for tests / diagnostics). -pub fn count_entity_index(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_entity_index", [], |r| { - r.get(0) - })?; - Ok(n.max(0) as u64) - }) -} - -/// Count score rows (for tests / diagnostics). pub fn count_scores(config: &Config) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_score", [], |r| r.get(0))?; - Ok(n.max(0) as u64) - }) + tinycortex::memory::score::store::count_scores(&memory_config(config)) } - -#[cfg(test)] -#[path = "store_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_tree/score/store_tests.rs b/src/openhuman/memory_tree/score/store_tests.rs deleted file mode 100644 index a9b3c76af..000000000 --- a/src/openhuman/memory_tree/score/store_tests.rs +++ /dev/null @@ -1,248 +0,0 @@ -use super::*; -use tempfile::TempDir; - -fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - (tmp, cfg) -} - -fn sample_row(id: &str, dropped: bool) -> ScoreRow { - ScoreRow { - chunk_id: id.to_string(), - total: 0.7, - signals: ScoreSignals { - token_count: 1.0, - unique_words: 0.8, - metadata_weight: 0.9, - source_weight: 0.5, - interaction: 0.6, - entity_density: 0.3, - llm_importance: 0.0, - }, - dropped, - reason: if dropped { - Some("below threshold".into()) - } else { - None - }, - computed_at_ms: 1_700_000_000_000, - llm_importance_reason: None, - } -} - -fn sample_entity(id: &str) -> CanonicalEntity { - CanonicalEntity { - canonical_id: format!("email:{id}"), - kind: EntityKind::Email, - surface: format!("{id}@example.com"), - span_start: 0, - span_end: (id.len() + 12) as u32, - score: 1.0, - } -} - -#[test] -fn upsert_then_get_score() { - let (_tmp, cfg) = test_config(); - let row = sample_row("c1", false); - upsert_score(&cfg, &row).unwrap(); - let got = get_score(&cfg, "c1").unwrap().expect("row exists"); - assert_eq!(got.chunk_id, row.chunk_id); - assert!((got.total - row.total).abs() < 1e-6); - assert_eq!(got.dropped, row.dropped); - assert_eq!(got.reason, row.reason); - assert_eq!(got.computed_at_ms, row.computed_at_ms); - assert!((got.signals.token_count - row.signals.token_count).abs() < 1e-6); -} - -#[test] -fn upsert_score_idempotent() { - let (_tmp, cfg) = test_config(); - let r = sample_row("c1", false); - upsert_score(&cfg, &r).unwrap(); - upsert_score(&cfg, &r).unwrap(); - assert_eq!(count_scores(&cfg).unwrap(), 1); -} - -#[test] -fn dropped_flag_persists() { - let (_tmp, cfg) = test_config(); - let r = sample_row("c1", true); - upsert_score(&cfg, &r).unwrap(); - let got = get_score(&cfg, "c1").unwrap().unwrap(); - assert!(got.dropped); - assert_eq!(got.reason.as_deref(), Some("below threshold")); -} - -#[test] -fn get_missing_score_is_none() { - let (_tmp, cfg) = test_config(); - assert!(get_score(&cfg, "missing").unwrap().is_none()); -} - -#[test] -fn index_and_lookup_entity() { - let (_tmp, cfg) = test_config(); - let e = sample_entity("alice"); - index_entity(&cfg, &e, "chunk-1", "leaf", 1000, Some("source:chat")).unwrap(); - index_entity(&cfg, &e, "chunk-2", "leaf", 2000, Some("source:chat")).unwrap(); - - let hits = lookup_entity(&cfg, "email:alice", None).unwrap(); - assert_eq!(hits.len(), 2); - // newest first - assert_eq!(hits[0].node_id, "chunk-2"); - assert_eq!(hits[1].node_id, "chunk-1"); -} - -#[test] -fn index_batch() { - let (_tmp, cfg) = test_config(); - let entities = vec![sample_entity("a"), sample_entity("b"), sample_entity("c")]; - let n = index_entities(&cfg, &entities, "chunk-1", "leaf", 1000, None).unwrap(); - assert_eq!(n, 3); - assert_eq!(count_entity_index(&cfg).unwrap(), 3); -} - -#[test] -fn clear_entity_index_drops_stale_rows() { - let (_tmp, cfg) = test_config(); - let a = sample_entity("a"); - let b = sample_entity("b"); - index_entities(&cfg, &[a.clone(), b], "chunk-1", "leaf", 1000, None).unwrap(); - assert_eq!(count_entity_index(&cfg).unwrap(), 2); - - // Simulate a re-score that only keeps entity "a". - let cleared = clear_entity_index_for_node(&cfg, "chunk-1").unwrap(); - assert_eq!(cleared, 2); - index_entities(&cfg, &[a], "chunk-1", "leaf", 1000, None).unwrap(); - - let hits = lookup_entity(&cfg, "email:b", None).unwrap(); - assert!(hits.is_empty(), "stale entity should be removed"); - let hits = lookup_entity(&cfg, "email:a", None).unwrap(); - assert_eq!(hits.len(), 1); -} - -#[test] -fn index_idempotent_per_entity_node_pair() { - let (_tmp, cfg) = test_config(); - let e = sample_entity("alice"); - index_entity(&cfg, &e, "chunk-1", "leaf", 1000, None).unwrap(); - index_entity(&cfg, &e, "chunk-1", "leaf", 1000, None).unwrap(); - assert_eq!(count_entity_index(&cfg).unwrap(), 1); -} - -#[test] -fn lookup_limit_respected() { - let (_tmp, cfg) = test_config(); - let e = sample_entity("alice"); - for i in 0..5 { - index_entity( - &cfg, - &e, - &format!("chunk-{i}"), - "leaf", - 1000 + i as i64, - None, - ) - .unwrap(); - } - let hits = lookup_entity(&cfg, "email:alice", Some(2)).unwrap(); - assert_eq!(hits.len(), 2); -} - -/// Regression: `index_summary_entity_ids_tx` must write a parseable -/// `entity_kind` (the "" prefix before `:`) so `lookup_entity` -/// can still round-trip rows through `EntityKind::parse`. Earlier code -/// stored the full canonical id, which poisoned lookups mixing leaf -/// and summary hits. See PR #789 CodeRabbit review. -#[test] -fn summary_entity_index_kind_is_parseable() { - use crate::openhuman::memory_store::chunks::store::with_connection; - - let (_tmp, cfg) = test_config(); - - // Seed a leaf hit so lookup_entity has something leafy to mix - // with the summary hit — this reproduces the mixed-row crash. - let leaf_entity = sample_entity("alice"); - index_entity(&cfg, &leaf_entity, "leaf-1", "leaf", 1000, Some("tree-1")).unwrap(); - - // Write a summary row via the tx helper under test. - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - let n = index_summary_entity_ids_tx( - &tx, - &["email:alice@example.com".into(), "hashtag:launch-q2".into()], - "summary-1", - 0.84, - 2000, - Some("tree-1"), - )?; - assert_eq!(n, 2); - tx.commit()?; - Ok(()) - }) - .unwrap(); - - // Before the fix: lookup_entity would fail on the summary row - // because entity_kind was "email:alice@example.com" and - // EntityKind::parse rejects it. After the fix, the column stores - // "email" and the lookup succeeds with both rows. - let hits = lookup_entity(&cfg, "email:alice@example.com", None).unwrap(); - assert_eq!(hits.len(), 1, "summary row should be discoverable"); - assert_eq!(hits[0].node_id, "summary-1"); - assert_eq!(hits[0].node_kind, "summary"); - assert_eq!(hits[0].entity_kind, EntityKind::Email); - - // Hashtag row parses as its own kind too. - let hits = lookup_entity(&cfg, "hashtag:launch-q2", None).unwrap(); - assert_eq!(hits.len(), 1); - assert_eq!(hits[0].entity_kind, EntityKind::Hashtag); - - // Mixing leaf + summary entity ids in one lookup also parses cleanly. - let hits = lookup_entity(&cfg, "email:alice", None).unwrap(); - assert_eq!(hits.len(), 1); - assert_eq!(hits[0].entity_kind, EntityKind::Email); -} - -// ---------- get_scores_batch ---------- -// -// Narrow on purpose: returns only chunk_id -> total. `fetch_leaves` is -// the call-site and only needs the float. Missing chunk_ids are absent -// from the map (mirrors per-row `get_score` Ok(None) → caller fallback -// to 0.0 neutral). - -#[test] -fn get_scores_batch_returns_present_chunk_ids() { - let (_tmp, cfg) = test_config(); - let r1 = sample_row("c1", false); - let mut r2 = sample_row("c2", false); - r2.total = 0.3; - upsert_score(&cfg, &r1).unwrap(); - upsert_score(&cfg, &r2).unwrap(); - - let ids = vec!["c1".to_string(), "c2".to_string()]; - let map = get_scores_batch(&cfg, &ids).unwrap(); - assert_eq!(map.len(), 2); - assert!((map.get("c1").copied().unwrap() - 0.7).abs() < 1e-6); - assert!((map.get("c2").copied().unwrap() - 0.3).abs() < 1e-6); -} - -#[test] -fn get_scores_batch_empty_input_and_missing_chunk_ids() { - // Empty input: empty map (no SQL issued). - let (_tmp, cfg) = test_config(); - let empty = get_scores_batch(&cfg, &[]).unwrap(); - assert!(empty.is_empty()); - - // Missing ids: silently absent so `fetch_leaves` can fall back to - // its documented 0.0 neutral without ambient errors. - let r = sample_row("c1", false); - upsert_score(&cfg, &r).unwrap(); - let ids = vec!["c1".to_string(), "ghost:no-such".to_string()]; - 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()); -} diff --git a/src/openhuman/memory_tree/summarise.rs b/src/openhuman/memory_tree/summarise.rs index e7a989b42..0d9495fed 100644 --- a/src/openhuman/memory_tree/summarise.rs +++ b/src/openhuman/memory_tree/summarise.rs @@ -1,443 +1,88 @@ -//! Memory-tree summariser: fold N items into one parent summary via a -//! single LLM call. -//! -//! This module replaces the previous trait-based summariser ladder -//! (`Summariser` + `InertSummariser` + `LlmSummariser`) with one plain -//! `async fn`. Callers pass inputs + context + config and get back -//! either a [`SummaryOutput`] or an error. Resilience (retry, graceful -//! degradation) is the caller's responsibility — see -//! [`fallback_summary`] for the deterministic concat-and-truncate -//! helper used by seal cascades that must never abort. -//! -//! The structured-facet-extraction side-channel that the old summariser -//! carried has been removed from this layer; facet extraction is the -//! `learning` domain's job and runs independently. +//! OpenHuman chat-provider adapter for tinycortex summary preparation. use anyhow::{Context, Result}; -use chrono::{DateTime, Utc}; use crate::openhuman::config::Config; use crate::openhuman::memory::chat::{build_chat_provider, ChatPrompt}; -use crate::openhuman::memory_store::chunks::types::approx_token_count; -use crate::openhuman::memory_store::trees::types::TreeKind; -/// Hard cap on summariser output length (in approximate tokens). Sized -/// to fit the downstream embedder (`nomic-embed-text-v1.5`, 8192-token -/// input ceiling) with headroom for tokenizer drift. -const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 5_000; +pub use tinycortex::memory::tree::{SummaryContext, SummaryInput}; -/// Context window assumed for the model. Sized for the cloud -/// summariser's 120k-token window with headroom; used as the divisor -/// in the per-input clamp so the joined prompt body stays under it -/// at upper-level seals where many children fold together. -const NUM_CTX_TOKENS: u32 = 60_000; - -/// Tokens reserved for system prompt + envelope overhead + tokenizer -/// drift between our 4-chars/token heuristic and the model's tokenizer. -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)] -pub struct SummaryInput { - pub id: String, - 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, -} - -/// Per-seal context — lets logs identify which tree is being sealed -/// without threading config globally. -#[derive(Clone, Debug)] -pub struct SummaryContext<'a> { - pub tree_id: &'a str, - pub tree_kind: TreeKind, - pub target_level: u32, - pub token_budget: u32, -} - -/// Output of a summarise call. +/// Compatibility result carrying provider usage alongside the crate-owned +/// summary output fields. #[derive(Clone, Debug, Default)] pub struct SummaryOutput { pub content: String, pub token_count: u32, - /// Always emitted empty by [`summarise`]. Canonical entity ids are - /// populated separately by the entity extractor; rolling up children's - /// labels mechanically is anti-pattern (see prior `InertSummariser` - /// design note). pub entities: Vec, pub topics: Vec, - /// Provider-reported prompt token count for this summarise call, when - /// the backend returned usage. `0` when usage was unavailable (e.g. - /// the [`fallback_summary`] path, or a provider that doesn't report - /// usage) — callers should fall back to their own estimate in that - /// case. Threaded into the sync audit log (issue #3110). pub input_tokens: u64, - /// Provider-reported completion token count for this summarise call. - /// `0` when usage was unavailable (see [`Self::input_tokens`]). pub output_tokens: u64, - /// Amount billed for this summarise call in USD, from the backend's - /// `openhuman.billing.charged_amount_usd`. `None` when the provider - /// did not report a charge — callers fall back to the hardcoded - /// pricing estimate (issue #3110). pub charged_amount_usd: Option, } -/// Fold `inputs` into a single summary by making one chat-provider call. -/// -/// Returns `Err` on provider build failure, network failure, or empty -/// upstream response. Callers that must not abort (e.g. seal cascades) -/// should match on the error and fall back to [`fallback_summary`]. pub async fn summarise( config: &Config, inputs: &[SummaryInput], - ctx: &SummaryContext<'_>, + context: &SummaryContext<'_>, ) -> Result { - 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 body = build_user_prompt(inputs, per_input_cap); - if body.trim().is_empty() { + let Some(prepared) = tinycortex::memory::tree::prepare_summary_prompt( + inputs, + context, + config.output_language.as_deref(), + ) else { return Ok(SummaryOutput::default()); - } - + }; let provider = build_chat_provider(config).context("memory_tree::summarise: build chat provider")?; - - let prompt = ChatPrompt { - system: system_prompt(effective_budget, config.output_language.as_deref()), - user: body, - temperature: 0.0, - kind: "memory_tree::summarise", - // Left open-ended: the summariser's output budget varies by node - // level (up to ~20k tokens at root) and is enforced prompt-side, so - // a hard wire cap risks truncating a valid large summary. The 402 - // precheck fix is scoped to extraction (TAURI-RUST-C62); revisit if - // the summarise path shows the same low-balance 402 shape. - max_tokens: None, - }; - log::debug!( - "[memory_tree::summarise] provider={} tree_id={} level={} inputs={} budget={}", + "[memory_tree::summarise] provider={} level={} inputs={} budget={}", provider.name(), - ctx.tree_id, - ctx.target_level, + context.target_level, inputs.len(), - ctx.token_budget, + prepared.effective_budget ); - - let (raw, usage) = provider - .chat_for_text_with_usage(&prompt) + let (text, usage) = provider + .chat_for_text_with_usage(&ChatPrompt { + system: prepared.system, + user: prepared.user, + temperature: 0.0, + kind: "memory_tree::summarise", + max_tokens: None, + }) .await .with_context(|| format!("memory_tree::summarise: provider={}", provider.name()))?; - - let (content, token_count) = clamp_to_budget(raw.trim(), effective_budget); - - // Prefer provider-reported usage (real token counts + charged amount) - // over our `body.len() / 4` estimate. `None`/zero means the backend - // didn't surface usage; downstream callers fall back to estimates. - let input_tokens = usage.as_ref().map(|u| u.input_tokens).unwrap_or(0); - let output_tokens = usage.as_ref().map(|u| u.output_tokens).unwrap_or(0); - let charged_amount_usd = usage.as_ref().and_then(|u| { - if u.charged_amount_usd > 0.0 { - Some(u.charged_amount_usd) - } else { - None - } - }); - + let output = + tinycortex::memory::tree::finish_provider_summary(&text, prepared.effective_budget); + let input_tokens = usage.as_ref().map_or(0, |usage| usage.input_tokens); + let output_tokens = usage.as_ref().map_or(0, |usage| usage.output_tokens); + let charged_amount_usd = usage + .as_ref() + .map(|usage| usage.charged_amount_usd) + .filter(|amount| *amount > 0.0); log::debug!( - "[memory_tree::summarise] sealed tree_id={} level={} inputs={} tokens={} usage_input={} usage_output={} charged_usd={:?}", - ctx.tree_id, - ctx.target_level, - inputs.len(), - token_count, + "[memory_tree::summarise] complete tokens={} usage_input={} usage_output={}", + output.token_count, input_tokens, - output_tokens, - charged_amount_usd, + output_tokens ); - Ok(SummaryOutput { - content, - token_count, - entities: Vec::new(), - topics: Vec::new(), + content: output.content, + token_count: output.token_count, + entities: output.entities, + topics: output.topics, input_tokens, output_tokens, charged_amount_usd, }) } -/// Deterministic, dependency-free summary — concatenate inputs with a -/// provenance prefix and truncate to budget. Used by seal cascades when -/// [`summarise`] returns an error and the cascade must still produce a -/// parent row (replaces the old `InertSummariser` soft-fallback role). pub fn fallback_summary(inputs: &[SummaryInput], budget: u32) -> SummaryOutput { - const PROVENANCE_PREFIX: &str = "— "; - let mut parts: Vec = Vec::with_capacity(inputs.len()); - for inp in inputs { - let trimmed = inp.content.trim(); - if trimmed.is_empty() { - continue; - } - parts.push(format!("{PROVENANCE_PREFIX}{trimmed}")); - } - let joined = parts.join("\n\n"); - let (content, token_count) = clamp_to_budget(&joined, budget); - // No provider call happened on the fallback path, so there is no - // real usage to report — leave token counts at 0 and charge at None - // so callers fall back to their estimate. + let output = tinycortex::memory::tree::fallback_summary(inputs, budget); SummaryOutput { - content, - token_count, - entities: Vec::new(), - topics: Vec::new(), - input_tokens: 0, - output_tokens: 0, - charged_amount_usd: None, - } -} - -fn build_user_prompt(inputs: &[SummaryInput], per_input_cap_tokens: u32) -> String { - // Higher-priority inputs (by score) lead the prompt so the most - // important source material — e.g. commit messages and closed/merged - // issues & PRs, which carry a priority boost at ingest — is summarised - // first and is least likely to be truncated under budget pressure. - // `sort_by` is stable, so chronological order is preserved among - // equal-score inputs. - let mut order: Vec<&SummaryInput> = inputs.iter().collect(); - order.sort_by(|a, b| { - b.score - .partial_cmp(&a.score) - .unwrap_or(std::cmp::Ordering::Equal) - }); - - let mut out = String::new(); - for inp in order { - let trimmed = inp.content.trim(); - if trimmed.is_empty() { - continue; - } - let (clamped, _) = clamp_to_budget(trimmed, per_input_cap_tokens); - if !out.is_empty() { - out.push_str("\n\n"); - } - out.push_str(&format!("[{}]\n{clamped}", inp.id)); - } - out -} - -fn system_prompt(budget: u32, output_language: Option<&str>) -> String { - let lang_line = match output_language { - Some(lang) if !lang.trim().is_empty() => { - format!("\nWrite the summary in {lang}.") - } - _ => String::new(), - }; - format!( - "You are folding multiple notes into one compact summary.\n\ - Aim for ~{budget} tokens or fewer. Capture key facts, decisions, and entities.\n\ - Output only the summary prose — no preamble, no JSON, no markdown headings.{lang_line}" - ) -} - -fn clamp_to_budget(text: &str, budget: u32) -> (String, u32) { - let initial = approx_token_count(text); - if initial <= budget { - return (text.to_string(), initial); - } - let char_ceiling = (budget as usize).saturating_mul(4); - let truncated: String = text.chars().take(char_ceiling).collect(); - let tokens = approx_token_count(&truncated); - (truncated, tokens) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn sample_input(id: &str, content: &str) -> SummaryInput { - let ts = Utc::now(); - SummaryInput { - id: id.to_string(), - content: content.to_string(), - token_count: approx_token_count(content), - entities: Vec::new(), - topics: Vec::new(), - time_range_start: ts, - time_range_end: ts, - score: 0.5, - } - } - - #[test] - fn fallback_concatenates_with_provenance_prefix() { - let inputs = vec![sample_input("a", "hello"), sample_input("b", "world")]; - let out = fallback_summary(&inputs, 10_000); - assert!(out.content.contains("hello")); - assert!(out.content.contains("world")); - assert!(out.content.contains("— ")); - assert!(out.entities.is_empty()); - } - - #[test] - fn fallback_truncates_at_budget() { - let inputs = vec![sample_input("a", &"x".repeat(1000))]; - let out = fallback_summary(&inputs, 5); - assert!(out.token_count <= 6); - } - - #[test] - fn fallback_skips_blank_inputs() { - let inputs = vec![sample_input("a", " "), sample_input("b", "kept")]; - let out = fallback_summary(&inputs, 10_000); - assert!(out.content.contains("kept")); - assert_eq!(out.content.matches("— ").count(), 1); - } - - #[test] - fn fallback_reports_no_provider_usage() { - // The fallback path makes no provider call, so it must report - // zero/None usage — github/rebuild then fall back to the estimate. - let inputs = vec![sample_input("a", "hello")]; - let out = fallback_summary(&inputs, 10_000); - assert_eq!(out.input_tokens, 0); - assert_eq!(out.output_tokens, 0); - assert_eq!(out.charged_amount_usd, None); - } - - /// Test `ChatProvider` that reports provider usage (real token counts - /// + a backend charge) so we can prove `summarise` threads it into - /// `SummaryOutput`. - struct UsageReportingProvider { - response: String, - usage: Option, - } - - #[async_trait::async_trait] - impl crate::openhuman::memory::chat::ChatProvider for UsageReportingProvider { - fn name(&self) -> &str { - "test:usage-reporting" - } - - async fn chat_for_json(&self, _prompt: &ChatPrompt) -> Result { - Ok(self.response.clone()) - } - - async fn chat_for_text_with_usage( - &self, - _prompt: &ChatPrompt, - ) -> Result<( - String, - Option, - )> { - Ok((self.response.clone(), self.usage.clone())) - } - } - - fn summary_ctx<'a>(tree_id: &'a str) -> SummaryContext<'a> { - SummaryContext { - tree_id, - tree_kind: TreeKind::Source, - target_level: 1, - token_budget: 5_000, - } - } - - #[tokio::test] - async fn summarise_threads_provider_usage_into_output() { - use crate::openhuman::inference::provider::UsageInfo; - use crate::openhuman::memory::chat::test_override; - - let provider = std::sync::Arc::new(UsageReportingProvider { - response: "a folded summary".to_string(), - usage: Some(UsageInfo { - input_tokens: 1_234, - output_tokens: 56, - charged_amount_usd: 0.0078, - ..Default::default() - }), - }); - - let cfg = Config::default(); - let inputs = vec![sample_input("a", "raw content to fold")]; - let ctx = summary_ctx("tree:test"); - - let out = - test_override::with_provider(provider, async { summarise(&cfg, &inputs, &ctx).await }) - .await - .unwrap(); - - assert_eq!(out.input_tokens, 1_234); - assert_eq!(out.output_tokens, 56); - assert_eq!(out.charged_amount_usd, Some(0.0078)); - assert!(out.content.contains("folded summary")); - } - - #[tokio::test] - async fn summarise_leaves_usage_empty_when_provider_reports_none() { - use crate::openhuman::memory::chat::test_override; - - let provider = std::sync::Arc::new(UsageReportingProvider { - response: "a folded summary".to_string(), - usage: None, - }); - - let cfg = Config::default(); - let inputs = vec![sample_input("a", "raw content to fold")]; - let ctx = summary_ctx("tree:test"); - - let out = - test_override::with_provider(provider, async { summarise(&cfg, &inputs, &ctx).await }) - .await - .unwrap(); - - // No usage → callers must fall back to their estimate. - assert_eq!(out.input_tokens, 0); - assert_eq!(out.output_tokens, 0); - assert_eq!(out.charged_amount_usd, None); - } - - #[tokio::test] - async fn summarise_treats_zero_charge_as_absent() { - use crate::openhuman::inference::provider::UsageInfo; - use crate::openhuman::memory::chat::test_override; - - // A provider that reports token counts but a zero charge (backend - // didn't surface billing) — token counts flow through, but the - // charge must be `None` so callers fall back to the estimate. - let provider = std::sync::Arc::new(UsageReportingProvider { - response: "a folded summary".to_string(), - usage: Some(UsageInfo { - input_tokens: 100, - output_tokens: 10, - charged_amount_usd: 0.0, - ..Default::default() - }), - }); - - let cfg = Config::default(); - let inputs = vec![sample_input("a", "raw content to fold")]; - let ctx = summary_ctx("tree:test"); - - let out = - test_override::with_provider(provider, async { summarise(&cfg, &inputs, &ctx).await }) - .await - .unwrap(); - - assert_eq!(out.input_tokens, 100); - assert_eq!(out.output_tokens, 10); - assert_eq!(out.charged_amount_usd, None); + content: output.content, + token_count: output.token_count, + entities: output.entities, + topics: output.topics, + ..SummaryOutput::default() } } diff --git a/src/openhuman/memory_tree/tree/bucket_seal.rs b/src/openhuman/memory_tree/tree/bucket_seal.rs index 691313ba9..4e196cb39 100644 --- a/src/openhuman/memory_tree/tree/bucket_seal.rs +++ b/src/openhuman/memory_tree/tree/bucket_seal.rs @@ -1,187 +1,23 @@ -//! Append + cascade-seal for summary trees (#709). -//! -//! `append_leaf` pushes a persisted chunk into the L0 buffer of a tree. -//! Seal gates differ by level: -//! -//! - **L0 (leaves → L1)**: seal when `token_sum >= INPUT_TOKEN_BUDGET`. -//! Token-only gating lets small-token items (e.g. commit messages at -//! ~20-50 tokens each) accumulate into large batches so summaries -//! cover meaningful spans of activity. -//! - **L≥1 (summaries → next level)**: seal when `item_ids.len() >= -//! SUMMARY_FANOUT`. Per-summary token size depends on summariser -//! quality, so a token-based gate collapses to a 1:1:1 chain when the -//! summariser is weak. Counting siblings keeps the tree's fan-in -//! stable regardless. -//! -//! When a buffer seals, its items move into the new summary's -//! `child_ids`, the buffer clears, and the new summary id is queued at -//! the next level. The cascade continues upward until a buffer fails its -//! gate. -//! -//! For low-volume sources that never hit the token or sibling-count -//! thresholds, time-based sealing is handled by -//! [`flush_stale_buffers`](super::flush::flush_stale_buffers), which -//! runs on a periodic cadence and force-seals stale L0 buffers. -//! -//! Concurrency: Phase 3a assumes a single-process SQLite workspace. All -//! writes in one seal step run in a single transaction; the async -//! summariser call happens outside any open transaction so a slow LLM -//! doesn't hold DB locks. Callers should serialise `append_leaf` per -//! tree — ingest achieves this by processing one batch's chunks -//! sequentially inside its `persist` task. Blocking SQLite calls inside -//! this async function are acceptable for Phase 3a because the Inert -//! summariser does no real I/O; when a networked summariser lands, wrap -//! DB calls in `tokio::task::spawn_blocking` to keep the runtime healthy. +//! Product adapters for tinycortex-owned bucket and document sealing. -use std::collections::BTreeSet; -use std::sync::Arc; - -use anyhow::{Context, Result}; +use anyhow::Result; use chrono::{DateTime, Utc}; -use futures::stream::{StreamExt, TryStreamExt}; -use rusqlite::Transaction; use crate::openhuman::config::Config; -use crate::openhuman::memory::util::redact::redact; -use crate::openhuman::memory_store::chunks::store::with_connection; -use crate::openhuman::memory_store::content::{ - atomic::stage_summary_with_layout, - paths::{slugify_source_id, SummaryDiskLayout}, - wiki_git::{SummaryCommitBatch, SummaryCommitEntry}, - SummaryComposeInput, -}; -use crate::openhuman::memory_store::trees::types::{ - Buffer, SummaryNode, Tree, INPUT_TOKEN_BUDGET, OUTPUT_TOKEN_BUDGET, SUMMARY_FANOUT, -}; -use crate::openhuman::memory_tree::score::embed::build_write_embedder; -use crate::openhuman::memory_tree::score::extract::EntityExtractor; -use crate::openhuman::memory_tree::score::resolver::canonicalise; -use crate::openhuman::memory_tree::summarise::{ - fallback_summary, summarise, SummaryContext, SummaryInput, SummaryOutput, -}; -use crate::openhuman::memory_tree::tree::factory::TreeFactory; -use crate::openhuman::memory_tree::tree::registry::new_summary_id; -use crate::openhuman::memory_tree::tree::store; +use crate::openhuman::memory_store::trees::types::{Buffer, Tree}; -/// Hard cap on cascade depth — prevents runaway loops if token accounting -/// ever slips. 32 levels at even a 2x fan-in is more than enough for any -/// realistic source. -const MAX_CASCADE_DEPTH: u32 = 32; +pub use tinycortex::memory::tree::{LabelStrategy, LeafRef, MERGE_LEVEL_BASE}; -/// How a sealed summary node's `entities` and `topics` fields get populated. -/// -/// Each tree kind has different correct semantics: -/// - **Source** trees use [`LabelStrategy::ExtractFromContent`] so the -/// summariser's freshly-synthesised text gets its own pass through an -/// extractor. Captures emergent themes that no individual leaf expressed. -/// - **Global** trees use [`LabelStrategy::UnionFromChildren`] — their -/// inputs are already-labeled source-tree summaries; union preserves -/// labels for time-based retrieval ("days that mentioned Alice") -/// without an LLM call. -/// - **Topic** trees use [`LabelStrategy::Empty`] — their scope already -/// pins the dominant theme; inheriting auxiliary entities would -/// cross-pollinate unrelated topic trees and noise the entity index. -#[derive(Clone)] -pub enum LabelStrategy { - /// Run the extractor on the new summary's content; canonicalise the - /// result into `entities` (canonical_ids) and `topics` (labels). - ExtractFromContent(Arc), - /// Dedup-merge each input's `entities` and `topics` into the parent. - UnionFromChildren, - /// Leave both fields empty regardless of inputs. - Empty, +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) } -impl std::fmt::Debug for LabelStrategy { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::ExtractFromContent(ex) => write!(f, "ExtractFromContent({})", ex.name()), - Self::UnionFromChildren => f.write_str("UnionFromChildren"), - Self::Empty => f.write_str("Empty"), - } - } -} - -/// Resolve `entities` and `topics` for a freshly-summarised node according -/// to the chosen strategy. Errors propagate from the extractor (when used). -async fn resolve_labels( - strategy: &LabelStrategy, - inputs: &[SummaryInput], - summary_content: &str, -) -> Result<(Vec, Vec)> { - match strategy { - LabelStrategy::ExtractFromContent(extractor) => { - let extracted = extractor - .extract(summary_content) - .await - .context("seal-time extractor failed")?; - let canonical = canonicalise(&extracted); - let mut entities: Vec = canonical - .into_iter() - .map(|c| c.canonical_id) - .collect::>() - .into_iter() - .collect(); - entities.sort(); - let mut topics: Vec = extracted - .topics - .into_iter() - .map(|t| t.label) - .collect::>() - .into_iter() - .collect(); - topics.sort(); - Ok((entities, topics)) - } - LabelStrategy::UnionFromChildren => { - let mut entities: BTreeSet = BTreeSet::new(); - let mut topics: BTreeSet = BTreeSet::new(); - for inp in inputs { - for e in &inp.entities { - entities.insert(e.clone()); - } - for t in &inp.topics { - topics.insert(t.clone()); - } - } - Ok((entities.into_iter().collect(), topics.into_iter().collect())) - } - LabelStrategy::Empty => Ok((Vec::new(), Vec::new())), - } -} - -/// A single leaf being appended to an L0 buffer. -#[derive(Clone, Debug)] -pub struct LeafRef { - pub chunk_id: String, - pub token_count: u32, - pub timestamp: DateTime, - pub content: String, - pub entities: Vec, - pub topics: Vec, - pub score: f32, -} - -/// Append a leaf to the source tree for `tree`, sealing buffers as they -/// fill. Returns the ids of any summaries that sealed during this call. -/// -/// `strategy` controls how each sealed summary's `entities` and `topics` -/// are populated — see [`LabelStrategy`]. pub async fn append_leaf( config: &Config, tree: &Tree, leaf: &LeafRef, strategy: &LabelStrategy, ) -> Result> { - log::debug!( - "[tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={} strategy={:?}", - tree.id, - leaf.chunk_id, - leaf.token_count, - strategy - ); - - // 1. Push leaf into L0 buffer (transactional). append_to_buffer( config, &tree.id, @@ -190,29 +26,13 @@ pub async fn append_leaf( leaf.token_count as i64, leaf.timestamp, )?; - - // 2. Cascade seals upward until a level stays under budget. - cascade_seals(config, tree, strategy).await + crate::openhuman::tinycortex::cascade_tree(config, tree, 0, false, strategy).await } -/// Queue-oriented variant of [`append_leaf`]. -/// -/// This only appends the leaf to the L0 buffer and returns whether the -/// caller should enqueue a follow-up seal job for level 0. pub fn append_leaf_deferred(config: &Config, tree: &Tree, leaf: &LeafRef) -> Result { - append_to_buffer( - config, - &tree.id, - 0, - &leaf.chunk_id, - leaf.token_count as i64, - leaf.timestamp, - )?; - let buf = store::get_buffer(config, &tree.id, 0)?; - Ok(should_seal(&buf)) + tinycortex::memory::tree::append_leaf_deferred(&engine_config(config), tree, leaf) } -/// Transactionally append a single item to `(tree_id, level)`'s buffer. pub fn append_to_buffer( config: &Config, tree_id: &str, @@ -221,47 +41,16 @@ pub fn append_to_buffer( token_delta: i64, item_ts: DateTime, ) -> Result<()> { - with_connection(config, |conn| { - let tx = conn.unchecked_transaction()?; - let mut buf = store::get_buffer_conn(&tx, tree_id, level)?; - // Idempotent on (tree_id, level, item_id): a retry after a failed - // cascade (step 2 of append_leaf) is a no-op, so duplicated children - // and double-counted tokens can't slip into the buffer. oldest_at - // stays on first-seen. - if buf.item_ids.iter().any(|existing| existing == item_id) { - log::debug!( - "[tree::bucket_seal] append_to_buffer: {item_id} already in buffer \ - tree_id={tree_id} level={level} — no-op" - ); - return Ok(()); - } - buf.item_ids.push(item_id.to_string()); - buf.token_sum = buf.token_sum.saturating_add(token_delta); - buf.oldest_at = match buf.oldest_at { - Some(existing) => Some(existing.min(item_ts)), - None => Some(item_ts), - }; - store::upsert_buffer_tx(&tx, &buf)?; - tx.commit()?; - Ok(()) - }) + tinycortex::memory::tree::append_to_buffer( + &engine_config(config), + tree_id, + level, + item_id, + token_delta, + item_ts, + ) } -async fn cascade_seals( - config: &Config, - tree: &Tree, - strategy: &LabelStrategy, -) -> Result> { - cascade_all_from(config, tree, 0, None, strategy).await -} - -/// Seal buffers starting at `start_level` and cascade upward. When -/// `force_now` is `Some`, the buffer at `start_level` is sealed regardless -/// of token budget (used by time-based flush). Upper levels are sealed -/// only when they cross the budget. -/// -/// `strategy` is forwarded to every sealed level — same semantics as -/// [`append_leaf`]. pub async fn cascade_all_from( config: &Config, tree: &Tree, @@ -269,851 +58,16 @@ pub async fn cascade_all_from( force_now: Option>, strategy: &LabelStrategy, ) -> 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 { - let buf = store::get_buffer(config, &tree.id, level)?; - let forced = first_iteration && force_now.is_some(); - first_iteration = false; - - if !forced && !should_seal(&buf) { - log::debug!( - "[tree::bucket_seal] cascade done tree_id={} stop_level={} token_sum={}", - tree.id, - level, - buf.token_sum - ); - break; - } - if buf.is_empty() { - log::debug!( - "[tree::bucket_seal] cascade hit empty buffer tree_id={} level={} — stopping", - tree.id, - level - ); - break; - } - - // Sync cascade — drives the level walk itself; doesn't need the - // queue follow-ups (we'll hit `seal_one_level` again next iter). - let summary_id = seal_one_level(config, tree, &buf, strategy, false).await?; - sealed_ids.push(summary_id); - level += 1; - } - - Ok(sealed_ids) -} - -/// Level-aware seal gate. -/// -/// L0 buffers gate on `token_sum >= INPUT_TOKEN_BUDGET` **only** — no -/// sibling-count fallback. Token-only gating is deliberate (see the -/// module header): small-token items (commit messages, short emails) -/// accumulate into large batches so each L1 summary covers a -/// meaningful span instead of 10 tiny items. Small buffers that never -/// reach the budget are sealed by time instead: -/// -/// Time-based sealing for low-volume sources is handled separately -/// by `flush_stale_buffers` (see [`crate::openhuman::memory_tree:: -/// tree::flush::flush_stale_buffers`]), which filters buffers -/// by `oldest_at` before calling the cascade (the 3-hourly -/// `memory_queue` scheduler drives it). Keeping the time gate -/// out of `should_seal` avoids prematurely sealing buffers during -/// normal `append_leaf` calls when test/restored data carries older -/// timestamps. -/// -/// L≥1 buffers gate on sibling count alone so the tree's -/// fan-in is independent of per-summary token size — without this, -/// summarisers that emit at the full token budget (e.g. the inert -/// fallback) collapse the cascade into a 1:1:1 chain instead of a -/// real tree. -pub(crate) fn should_seal(buf: &Buffer) -> bool { - if buf.is_empty() { - return false; - } - if buf.level == 0 { - buf.token_sum >= INPUT_TOKEN_BUDGET as i64 - } else { - (buf.item_ids.len() as u32) >= SUMMARY_FANOUT - } -} - -/// Seal `buf` at `level` into one summary at `level + 1`. Returns the new -/// summary id. -/// -/// `strategy` decides how `entities` and `topics` get populated on the new -/// summary node — see [`LabelStrategy`]. -/// -/// When `enqueue_follow_ups` is `true`, the function additionally inserts -/// follow-up job rows **inside the same transaction** that commits the -/// seal: -/// - `seal { tree_id, level: parent_level }` if the parent buffer's gate -/// is now met (parent-cascade enqueue) -/// - `topic_route { NodeRef::Summary { summary_id } }` for source trees -/// (so summary-level entities feed the topic-tree spawn pipeline) -/// -/// Atomic enqueue eliminates the crash window where a seal commits but -/// the post-commit follow-up enqueues are silently lost on a worker -/// crash. The async-pipeline handler (`handle_seal`) passes `true`. The -/// synchronous in-process cascade caller ([`cascade_all_from`]) passes -/// `false` because it drives the cascade itself and topic_route isn't -/// part of the test/flush sync path. -pub(crate) async fn seal_one_level( - config: &Config, - tree: &Tree, - buf: &Buffer, - strategy: &LabelStrategy, - enqueue_follow_ups: bool, -) -> Result { - let level = buf.level; - let target_level = level + 1; - - // Hydrate inputs (synchronous DB reads). - let inputs = hydrate_inputs(config, level, &buf.item_ids)?; - if inputs.is_empty() { - anyhow::bail!( - "[tree::bucket_seal] refused to seal empty buffer tree_id={} level={}", - tree.id, - level - ); - } - - // Compute envelope across children (time range, max score). - let time_range_start = inputs - .iter() - .map(|i| i.time_range_start) - .min() - .unwrap_or_else(Utc::now); - let time_range_end = inputs - .iter() - .map(|i| i.time_range_end) - .max() - .unwrap_or_else(Utc::now); - let score = inputs - .iter() - .map(|i| i.score) - .fold(f32::NEG_INFINITY, f32::max) - .max(0.0); - - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::MemoryTreeBuildProgress { - phase: "seal".to_string(), - step: "summarising".to_string(), - tree_scope: Some(tree.scope.clone()), - level: Some(level), - item_count: Some(inputs.len() as u32), - detail: Some(format!("{} inputs → L{}", inputs.len(), level + 1)), - }, - ); - - let ctx = SummaryContext { - tree_id: &tree.id, - tree_kind: tree.kind, - target_level, - token_budget: OUTPUT_TOKEN_BUDGET, - }; - // #13021: treat a blank summary (LLM returned only whitespace, so - // `summarise()` succeeded with `content = ""`) the same as a hard error — - // fall back to the deterministic `fallback_summary` so we never persist a - // parent node with `content = ""` / `token_count = 0`. Without this, the - // child text is lost: the level buffer is cleared and the parent has no - // recoverable content for the next rollup or for retrieval. - let output = match summarise(config, &inputs, &ctx).await { - Ok(o) if !o.content.trim().is_empty() => o, - Ok(_) => { - log::warn!( - "[memory_tree::seal] summarise returned blank for tree_id={} level={} — using fallback (#13021)", - ctx.tree_id, ctx.target_level, - ); - fallback_summary(&inputs, ctx.token_budget) - } - Err(e) => { - log::warn!( - "[memory_tree::seal] summarise failed for tree_id={} level={}: {e:#} — using fallback", - ctx.tree_id, ctx.target_level, - ); - fallback_summary(&inputs, ctx.token_budget) - } - }; - - // Resolve labels (entities/topics) for the new summary node according - // to the chosen strategy. Done before the write tx so an extractor - // failure aborts the seal cleanly — same shape as the embedder guard - // below. - let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?; - - // Phase 4 (#710): embed the summary BEFORE opening the write tx so an - // embedder failure aborts the seal cleanly — nothing is persisted, - // the buffer stays intact, and a retry re-embeds from scratch. The - // tx below would otherwise commit a summary with no embedding, - // polluting retrieval's semantic rerank. - // - // Embedder context-window guard: `nomic-embed-text-v1.5` accepts - // up to 8192 tokens of input. Summary content is bounded by - // `ctx.token_budget = OUTPUT_TOKEN_BUDGET = 5_000` which fits, but - // we still truncate the input passed to `embed()` to leave - // headroom for tokenizer drift (the persisted summary content - // stays full; only the embedding's "view" of it is clamped). - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::MemoryTreeBuildProgress { - phase: "seal".to_string(), - step: "embedding".to_string(), - tree_scope: Some(tree.scope.clone()), - level: Some(level), - item_count: None, - detail: None, - }, - ); - - // Conservative cap. Slack-style chat content (URLs, mentions, - // emoji) tokenizes 2-4× higher than the 4-chars/token heuristic. - // 1000 approx-tokens (~4000 chars) is comfortably under 8192 - // even at 4× tokenizer ratio. - let embed_input = truncate_for_embed(&output.content, 1_000); - log::info!( - "[tree::bucket_seal] embed input: original_chars={} truncated_chars={}", - output.content.len(), - embed_input.len() - ); - // #002 (FR-002): skip embedding when no usable provider is configured - // (build_write_embedder returns None) rather than writing a fake all-zero - // vector. The summary is sealed embedding-less (re-embeddable later) and - // the semantic-recall degraded flag is already set with a typed cause. - // - // #13021: also skip when `embed_input.trim().is_empty()`. `summarise()` - // returns `SummaryOutput::default()` (content = "") when every input is - // whitespace, and `fallback_summary` joins zero parts the same way. An - // empty `embed_input` is guaranteed to 400 from the upstream embedding - // API — short-circuit to the same "embedding-less, re-embeddable later" - // path as the no-provider case. - let embedding: Option> = if embed_input.trim().is_empty() { - log::warn!( - "[tree::bucket_seal] empty summary content for tree_id={} level={}→{} \ - — sealing without embedding (#13021)", - tree.id, - level, - target_level - ); - None - } else { - match build_write_embedder(config).context("build embedder during seal")? { - None => { - log::warn!( - "[tree::bucket_seal] embeddings unavailable for tree_id={} level={}→{} \ - — sealing summary without embedding (semantic recall degraded)", - tree.id, - level, - target_level - ); - None - } - Some(embedder) => { - let v = match embedder.embed(&embed_input).await { - Ok(v) => v, - Err(e) => { - // #002: classify so the seal job fails fast on - // unrecoverable embed causes (budget/auth/dim) with a - // typed reason instead of retrying; original chain - // preserved as context. - let failure = - crate::openhuman::memory_tree::health::classify_embed_error(&e); - return Err(anyhow::Error::new(failure).context(format!( - "embed summary during seal tree_id={} level={}: {e:#}", - tree.id, level - ))); - } - }; - // Dimension guard: reject wrong-dimensionality vectors before - // they reach the store — same contract as handle_extract's - // pack_checked. Without this a provider returning the wrong - // shape slips into the summary sidecar silently. - crate::openhuman::memory_tree::score::embed::pack_checked(&v).context(format!( - "seal embed dim check tree_id={} level={}", - tree.id, level - ))?; - log::debug!( - "[tree::bucket_seal] embedded summary tree_id={} level={}→{} bytes={} provider={}", - tree.id, - level, - target_level, - output.content.len(), - embedder.name() - ); - crate::openhuman::memory_tree::health::clear_semantic_recall_degraded(); - Some(v) - } - } - }; - - // Build the new summary node. - let now = Utc::now(); - let summary_id = new_summary_id(target_level); - let node = SummaryNode { - id: summary_id.clone(), - tree_id: tree.id.clone(), - // `seal_one_level` runs for source AND topic trees (handle_seal, - // cascade_all_from, flush). Hardcoding Source here would write - // topic-tree summaries with tree_kind='source' in - // mem_tree_summaries, breaking any query filtering on tree_kind. - tree_kind: tree.kind, - level: target_level, - parent_id: None, - child_ids: buf.item_ids.clone(), - content: output.content, - token_count: output.token_count, - entities: node_entities, - topics: node_topics, - time_range_start, - time_range_end, - score, - sealed_at: now, - deleted: false, - embedding, - // Generic seal path (chat/email source trees + the cross-document - // merge tier) is document-agnostic. The per-document subtree seal - // (Notion) sets these via its own seal entrypoint in Task #2. - doc_id: None, - version_ms: None, - }; - - crate::core::event_bus::publish_global( - crate::core::event_bus::DomainEvent::MemoryTreeBuildProgress { - phase: "seal".to_string(), - step: "persisting".to_string(), - tree_scope: Some(tree.scope.clone()), - level: Some(target_level), - item_count: None, - detail: Some(format!( - "summary {} ({} tokens)", - &summary_id[..summary_id.len().min(16)], - output.token_count - )), - }, - ); - - // Phase MD-content: stage the summary .md file BEFORE opening the write - // tx. A staging failure aborts the seal cleanly — nothing is persisted - // and the buffer stays intact for retry. - // - let tree_factory = TreeFactory::from_tree(tree); - let summary_tree_kind = tree_factory.summary_tree_kind(); - let scope_slug = tree_factory.scope_slug(); - // For L1 seals (children are chunks), point each child wikilink at - // the raw archive file the chunk's body lives in — the email - // chunk-store path `email//.md` no longer - // exists, so `[[]]` would be an unresolved Obsidian - // link. We emit the relative path under content_root (with `.md` - // stripped) so the wikilink resolves unambiguously even outside - // Obsidian's unique-basename heuristic — e.g. - // `[[raw/gmail-stevent95-at-gmail-dot-com/_]]`. - // L≥2 children are summary ids whose default `sanitize_filename` - // resolves to existing `wiki/summaries/...md` files — leave - // overrides unset there. - let child_basename_overrides: Option>> = if node.level == 1 { - let overrides: Vec> = node - .child_ids - .iter() - .map(|chunk_id| { - // Surface lookup failures explicitly — silently - // falling back to `[[]]` would commit an - // unresolved Obsidian wikilink without any signal. - // We still yield `None` (so `compose_summary_md` - // takes the sanitised-id fallback) but a warn log - // makes the SQL error visible for diagnosis. - match crate::openhuman::memory_store::chunks::store::get_chunk_raw_refs( - config, chunk_id, - ) { - Ok(Some(refs)) if !refs.is_empty() => { - // RawRef::path is a forward-slash relative path - // under content_root, e.g. - // "raw/gmail-…/1700000_msg-id.md". Strip `.md` - // for Obsidian's extension-less wikilink - // convention. - let r = refs.into_iter().next().expect("non-empty"); - Some(r.path.strip_suffix(".md").unwrap_or(&r.path).to_string()) - } - Ok(_) => { - // No raw_refs persisted for this chunk — most - // commonly slack chunks (we only stage raw - // archive files for gmail today). The wikilink - // falls back to `sanitize_filename(chunk_id)`, - // which produces a deliberately-unresolved - // Obsidian link. Log so the silent-degradation - // path stays visible during diagnosis. - log::debug!( - "[tree::bucket_seal] no raw_refs for chunk_id={chunk_id} \ - — wikilink will fall back to sanitised chunk id" - ); - None - } - Err(e) => { - log::warn!( - "[tree::bucket_seal] get_chunk_raw_refs failed \ - chunk_id={chunk_id} err={e:#} — falling back to \ - chunk_id-based wikilink" - ); - None - } - } - }) - .collect(); - Some(overrides) - } else { - None - }; - let compose_input = SummaryComposeInput { - summary_id: &node.id, - tree_kind: summary_tree_kind, - tree_id: &node.tree_id, - tree_scope: &tree.scope, - level: node.level, - child_ids: &node.child_ids, - child_basenames: child_basename_overrides.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, - }; - // Stage the summary .md file and propagate any error — a staging failure - // aborts the seal entirely so the database never commits a row with - // content_path = NULL. The buffer stays unsealed and the job-retry path - // will re-attempt the file write on next execution. - let content_root = config.memory_tree_content_root(); - // Drop the bundled `.obsidian/` defaults (graph + types) so a user - // opening the vault gets the intended graph-view colour mapping - // without manual configuration. Best-effort and idempotent — never - // overwrites an existing file. - if let Err(err) = - crate::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults(&content_root) - { - log::warn!( - "[tree::bucket_seal] ensure_obsidian_defaults failed: {err:#} — \ - continuing seal without vault defaults" - ); - } - // Merge-tier nodes (document source trees, level ≥ MERGE_LEVEL_BASE) land - // under `source-/merge/`; everything else (chat/email + the - // per-doc subtree is sealed via seal_explicit_children, not here) uses the - // flat layout. - let layout = if node.level >= MERGE_LEVEL_BASE { - SummaryDiskLayout::Merge - } else { - SummaryDiskLayout::Standard - }; - let staged = stage_summary_with_layout(&content_root, &compose_input, &scope_slug, layout) - .with_context(|| { - format!( - "stage_summary failed for {}; seal aborted, buffer stays unsealed for retry", - node.id - ) - })?; - log::debug!( - "[tree::bucket_seal] staged summary {} → {}", - node.id, - staged.content_path - ); - let summary_commit = summary_seal_batch( + crate::openhuman::tinycortex::cascade_tree( + config, tree, - "bucket_seal", - std::slice::from_ref(&node), - std::slice::from_ref(&staged.content_path), + start_level, + force_now.is_some(), + strategy, ) - .with_context(|| { - format!( - "build summary git batch failed for {}; seal aborted, buffer stays unsealed for retry", - node.id - ) - })?; - - // Single write transaction: insert summary, clear this buffer, append - // summary id to parent buffer, bump tree max_level/root if needed, - // and (when `enqueue_follow_ups`) atomically enqueue parent-seal + - // topic_route follow-ups so they can never desync from the commit. - // Re-read `max_level` from inside the tx so cascading seals within - // one call see the updated value from earlier levels. - let summary_id_for_closure = summary_id.clone(); - let target_level_for_closure = target_level; - let tree_id = tree.id.clone(); - with_connection(config, move |conn| { - let tx = conn.unchecked_transaction()?; - - let current_max: u32 = tx - .query_row( - "SELECT max_level FROM mem_tree_trees WHERE id = ?1", - rusqlite::params![&tree_id], - |r| r.get::<_, i64>(0), - ) - .map(|n| n.max(0) as u32) - .context("Failed to read current max_level for tree")?; - - store::insert_summary_tx( - &tx, - &node, - Some(&staged), - &crate::openhuman::memory_store::chunks::store::tree_active_signature(config), - )?; - // Forward-compat: index any entities the summariser emitted into - // `mem_tree_entity_index` so Phase 4 retrieval can resolve - // "summaries mentioning Alice" via the same inverted index as - // leaves. No-op when entities is empty (the current summarise() - // always emits empty — entity extraction is the learning domain's job); - // becomes active once the summariser or a post-seal extractor emits canonical ids. - crate::openhuman::memory_tree::score::store::index_summary_entity_ids_tx( - &tx, - &node.entities, - &node.id, - node.score, - now.timestamp_millis(), - Some(&tree_id), - )?; - // Backlink children → new parent so leaf/parent traversal is a - // single-row lookup in Phase 4. Skipped for levels already bound - // to a parent (shouldn't happen — a child seals at most once). - for child_id in &node.child_ids { - if level == 0 { - tx.execute( - "UPDATE mem_tree_chunks - SET parent_summary_id = ?1 - WHERE id = ?2 AND parent_summary_id IS NULL", - rusqlite::params![&summary_id_for_closure, child_id], - ) - .context("Failed to backlink chunk to parent summary")?; - } else { - tx.execute( - "UPDATE mem_tree_summaries - SET parent_id = ?1 - WHERE id = ?2 AND parent_id IS NULL", - rusqlite::params![&summary_id_for_closure, child_id], - ) - .context("Failed to backlink summary to parent summary")?; - } - } - store::clear_buffer_tx(&tx, &tree_id, level)?; - - // Append to parent buffer. - let mut parent = store::get_buffer_conn(&tx, &tree_id, target_level_for_closure)?; - parent.item_ids.push(summary_id_for_closure.clone()); - parent.token_sum = parent.token_sum.saturating_add(node.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)?; - - // Atomic follow-up enqueues. Done INSIDE this tx — if the commit - // rolls back, the queue rows go with it; if it succeeds, the - // rows are durably visible to the worker pool. Eliminates the - // crash window where the seal commits but post-commit enqueues - // are lost. - if enqueue_follow_ups { - // Parent-cascade: if the new summary made the parent buffer - // cross its gate, enqueue the next level's seal. Dedupe key - // `seal:{tree_id}:{parent_level}` prevents duplicates if a - // parallel path already queued it. - if should_seal(&parent) { - use crate::openhuman::memory_queue::store::enqueue_tx as enqueue_job_tx; - use crate::openhuman::memory_queue::types::{NewJob, SealPayload}; - let parent_seal = SealPayload { - tree_id: tree_id.clone(), - level: target_level_for_closure, - force_now_ms: None, - }; - enqueue_job_tx(&tx, &NewJob::seal(&parent_seal)?)?; - } - // (Topic-tree routing removed: the topic/global trees were - // deleted — source trees plus the entity index are the - // substrate, so a source seal no longer fans out anywhere.) - } - - // Update tree root / max_level if we just climbed. - if target_level_for_closure > current_max { - store::update_tree_after_seal_tx( - &tx, - &tree_id, - &summary_id_for_closure, - target_level_for_closure, - now, - )?; - } else { - // Same max level — still refresh last_sealed_at via same helper - // but keep existing root intact. Root tracking at the same - // level is resolved in Phase 4 retrieval. - refresh_last_sealed_tx(&tx, &tree_id, now)?; - } - - tx.commit()?; - Ok(()) - })?; - - commit_summary_seal(config, &summary_commit).with_context(|| { - format!( - "commit_summary_seal failed for {} after DB seal commit", - summary_id - ) - })?; - - log::info!( - "[tree::bucket_seal] sealed tree_id={} level={}→{} summary_id={} children={}", - tree.id, - level, - target_level, - summary_id, - buf.item_ids.len() - ); - - Ok(summary_id) + .await } -/// Clamp `text` to roughly `max_tokens` tokens before passing to the -/// embedder. Uses the same ~4 chars/token heuristic as -/// `approx_token_count`. Embedders have hard input-size limits (e.g. -/// `nomic-embed-text-v1.5` = 8192 tokens) and an overshoot returns -/// HTTP 500 from Ollama rather than auto-truncating, which would -/// abort the seal transaction. -fn truncate_for_embed(text: &str, max_tokens: u32) -> String { - let approx = crate::openhuman::memory_store::chunks::types::approx_token_count(text); - if approx <= max_tokens { - return text.to_string(); - } - let char_ceiling = (max_tokens as usize).saturating_mul(4); - text.chars().take(char_ceiling).collect() -} - -fn refresh_last_sealed_tx( - tx: &Transaction<'_>, - tree_id: &str, - sealed_at: DateTime, -) -> Result<()> { - tx.execute( - "UPDATE mem_tree_trees SET last_sealed_at_ms = ?1 WHERE id = ?2", - rusqlite::params![sealed_at.timestamp_millis(), tree_id], - ) - .with_context(|| format!("Failed to refresh last_sealed_at for tree {tree_id}"))?; - Ok(()) -} - -fn summary_seal_batch( - tree: &Tree, - reason: &str, - nodes: &[SummaryNode], - content_paths: &[String], -) -> Result { - if nodes.is_empty() { - return Ok(SummaryCommitBatch { - reason: reason.to_string(), - tree_id: tree.id.clone(), - tree_scope: tree.scope.clone(), - entries: Vec::new(), - }); - } - if nodes.len() != content_paths.len() { - anyhow::bail!( - "[tree::bucket_seal] commit_summary_seal: node/path length mismatch nodes={} paths={}", - nodes.len(), - content_paths.len() - ); - } - - let entries = nodes - .iter() - .zip(content_paths.iter()) - .map(|(node, content_path)| SummaryCommitEntry { - summary_id: node.id.clone(), - content_path: content_path.clone(), - level: node.level, - child_count: node.child_ids.len(), - token_count: node.token_count, - time_range_start: node.time_range_start, - time_range_end: node.time_range_end, - }) - .collect(); - - Ok(SummaryCommitBatch { - reason: reason.to_string(), - tree_id: tree.id.clone(), - tree_scope: tree.scope.clone(), - entries, - }) -} - -fn commit_summary_seal(config: &Config, batch: &SummaryCommitBatch) -> Result<()> { - crate::openhuman::memory_store::content::wiki_git::commit_summaries( - &config.memory_tree_content_root(), - batch, - ) -} - -/// Fetch contributions for `item_ids`. At level 0 we pull from -/// `mem_tree_chunks` + `mem_tree_score`; at ≥1 we pull from -/// `mem_tree_summaries`. -fn hydrate_inputs(config: &Config, level: u32, item_ids: &[String]) -> Result> { - if level == 0 { - hydrate_leaf_inputs(config, item_ids) - } else { - hydrate_summary_inputs(config, item_ids) - } -} - -fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result> { - use crate::openhuman::memory_store::chunks::store::get_chunk; - use crate::openhuman::memory_store::content::read as content_read; - use crate::openhuman::memory_tree::score::store::{get_score, list_entity_ids_for_node}; - - let mut out: Vec = Vec::with_capacity(chunk_ids.len()); - for id in chunk_ids { - let chunk = match get_chunk(config, id)? { - Some(c) => c, - None => { - log::warn!( - "[tree::bucket_seal] hydrate_leaf_inputs: missing chunk {id} — skipping" - ); - continue; - } - }; - let score_value = get_score(config, id)?.map(|row| row.total).unwrap_or(0.0); - // Pull canonical entity ids from the inverted index — that's the - // authoritative source for "what entities are attached to this - // chunk." Topics live on the chunk's metadata tags. - // [`LabelStrategy::UnionFromChildren`] reads these fields off - // each `SummaryInput` to roll labels up the tree. - let entities = list_entity_ids_for_node(config, id).unwrap_or_default(); - // Read the full body from disk — the `content` column in SQLite holds - // a ≤500-char preview after the MD-on-disk migration. The summariser - // must receive the complete chunk text so the seal output is not a - // summary of previews. - // - // For pre-MD-migration chunks (no content_path recorded) this call - // returns Err; callers that want to handle legacy rows should check - // content_path presence before calling hydrate_inputs. - let body = content_read::read_chunk_body(config, id).with_context(|| { - format!("[tree::bucket_seal] hydrate_leaf_inputs: read body for chunk {id}") - })?; - out.push(SummaryInput { - id: chunk.id.clone(), - content: body, - token_count: chunk.token_count, - entities, - topics: chunk.metadata.tags.clone(), - time_range_start: chunk.metadata.time_range.0, - time_range_end: chunk.metadata.time_range.1, - score: score_value, - }); - } - Ok(out) -} - -fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result> { - use crate::openhuman::memory_store::content::read as content_read; - use crate::openhuman::memory_store::trees::store::get_summaries_batch; - - // One batched `SELECT … WHERE id IN (?,?,…)` instead of N per-id - // `get_summary` round-trips. Body reads stay per-id because each - // summary's full markdown lives in its own on-disk file — batching - // there would mean concurrent file opens, not a single round-trip. - // Walking the caller's `summary_ids` slice (not the map) preserves - // input order, matching the previous per-id loop's semantics; the - // map lookup keys by id so the order of `HashMap`'s iteration is - // irrelevant. - let node_by_id = get_summaries_batch(config, summary_ids)?; - - let mut out: Vec = Vec::with_capacity(summary_ids.len()); - for id in summary_ids { - let Some(node) = node_by_id.get(id) else { - log::warn!( - "[tree::bucket_seal] hydrate_summary_inputs: missing summary {id} — skipping" - ); - continue; - }; - // Read the full body from disk — `node.content` is a ≤500-char preview - // after the MD-on-disk migration. Higher-level seals (L2+) summarise - // over L1 summary content and need the full text, not a preview. - let body = content_read::read_summary_body(config, id).with_context(|| { - format!("[tree::bucket_seal] hydrate_summary_inputs: read body for summary {id}") - })?; - out.push(SummaryInput { - id: node.id.clone(), - content: body, - token_count: node.token_count, - entities: node.entities.clone(), - topics: node.topics.clone(), - time_range_start: node.time_range_start, - time_range_end: node.time_range_end, - score: node.score, - }); - } - Ok(out) -} - -// ── Document-aware sealing (Notion etc.) ──────────────────────────────── -// -// Document source trees keep ONE physical `mem_tree_trees` row per -// connection (e.g. `notion:{connection_id}`), but inside it each document -// rolls up to its own **doc-root** summary, and those doc-roots merge into -// the connection root. To get that shape without re-keying the shared -// `(tree, level)` buffer — the exact path chat/email seal through — the -// per-document subtree is built as an **isolated side-cascade** -// ([`seal_document_subtree`]) that never touches a shared buffer or the -// tree root. Only the cross-document **merge tier** uses the shared buffer -// + the existing [`cascade_all_from`] engine, starting at -// [`MERGE_LEVEL_BASE`]. -// -// Versioning is forward-only: editing a Notion page calls -// [`seal_document_subtree`] again with a higher `version_ms`, producing a -// *new* doc-root that is appended to the merge buffer alongside the old -// one. Nothing is rewritten or tombstoned; retrieval keeps `max(version_ms)` -// per `doc_id` at read time (see the retrieval layer). - -/// Level offset where the cross-document merge tier lives inside a document -/// source tree. -/// -/// Per-document subtrees occupy small levels (1, 2, …) and are built by -/// [`seal_document_subtree`] as a side-cascade that never enters a shared -/// `(tree, level)` buffer. The merge tier — which summarises *across* -/// documents — uses the shared buffer and the existing cascade engine, -/// starting here. The wide gap guarantees per-doc nodes (small levels) and -/// merge nodes (≥ `MERGE_LEVEL_BASE`) can never collide on `(tree, level)`, -/// and keeps `Tree.root_id` / `max_level` pointing at a merge node. -pub const MERGE_LEVEL_BASE: u32 = 1_000; - -/// Hard cap on how many children one per-document summary fans in, so a -/// single huge document can't produce a doc-root with thousands of direct -/// children. Independent of [`SUMMARY_FANOUT`] (which gates the merge tier). -const DOC_SUBTREE_MAX_FANIN: usize = 32; - -/// Bound on how many sibling batches within a single document-subtree level -/// are sealed concurrently. Each [`seal_explicit_children`] call is one -/// independent summarise + embed round-trip pair, so overlapping them with a -/// small pool collapses the serial `N × (LLM + embed RTT)` wall-time toward -/// `N/CONCURRENCY × RTT` on the ingest path. Kept low: the SQLite write tx -/// serialises on a single writer anyway, and networked summarise/embed -/// providers (Voyage, cloud LLMs) rate-limit — so a large pool buys nothing -/// and risks 429s. `buffered` (ordered) preserves the serial cascade's batch -/// order, so the resulting tree shape is identical. -const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8; - -/// Build (or re-build, for a new version) one document's subtree and merge -/// its doc-root into the connection tree. -/// -/// `doc_id` is the document identity (the chunk `source_id`, e.g. -/// `notion:{conn}:{page_id}`); `version_ms` is the document version -/// (`last_edited_time` epoch-ms). `chunk_ids` are this version's leaf chunk -/// ids, already persisted in `mem_tree_chunks`. -/// -/// Steps: -/// 1. Cascade `chunk_ids` upward (token-budget batches at L0, count batches -/// above) until a **single doc-root** remains — force-sealed even for a -/// one-chunk document so it always surfaces as a doc-root, never loose -/// leaves. Every node is tagged `(doc_id, version_ms)`. -/// 2. Append the doc-root to the connection tree's merge buffer at -/// [`MERGE_LEVEL_BASE`] and run the existing cascade so it folds into the -/// connection root once `SUMMARY_FANOUT` doc-roots accumulate. -/// -/// Returns the doc-root summary id. Idempotent re-runs for the *same* -/// `(doc_id, version_ms, chunk_ids)` produce a new doc-root (new ids); the -/// caller (Notion sync) only invokes this when a new revision is admitted. pub async fn seal_document_subtree( config: &Config, tree: &Tree, @@ -1122,458 +76,25 @@ pub async fn seal_document_subtree( chunk_ids: &[String], strategy: &LabelStrategy, ) -> Result { - if chunk_ids.is_empty() { - anyhow::bail!( - "[tree::bucket_seal] seal_document_subtree: empty chunk set tree_id={} doc_id_hash={}", - tree.id, - redact(doc_id) - ); - } - log::debug!( - "[tree::bucket_seal] seal_document_subtree tree_id={} doc_id_hash={} version_ms={:?} chunks={}", - tree.id, - redact(doc_id), - version_ms, - chunk_ids.len() - ); - - // 1. Per-document side-cascade to a single doc-root. - let mut current_level: u32 = 0; - let mut current_ids: Vec = chunk_ids.to_vec(); - let mut doc_root: Option = None; - - loop { - let batches = if current_level == 0 { - batch_leaves_by_token_budget(config, ¤t_ids)? - } else { - batch_by_count(¤t_ids, DOC_SUBTREE_MAX_FANIN) - }; - - // Within one level the batches are disjoint child sets, and - // `seal_explicit_children` touches no shared `(tree, level)` buffer - // and never advances the tree root (see its docstring) — so the - // sibling seals are independent and their summarise + embed - // round-trips can overlap. Run them with bounded concurrency. - // - // `buffered` (ordered), NOT `buffer_unordered`: the next level groups - // `next_ids` positionally via `batch_by_count`, so the output order - // must match the serial cascade to keep the tree shape identical. - // `try_collect` short-circuits on the first error, preserving the - // old `?`-on-first-failure behaviour. Levels stay sequential — level - // N+1 consumes level N's output — so only siblings overlap, not the - // cross-level fan-in. The blocking SQLite write inside each seal - // still serialises (single writer), which is fine: the win is the - // overlapped network RTTs, not the DB writes. - // Build the per-batch futures eagerly into a `Vec` (they don't run until - // polled) rather than lazily inside a `stream::iter(...).map(closure)`: - // the lazy form makes rustc try to infer a higher-ranked `FnOnce` over - // the borrowed `batch` lifetime, which fails ("implementation of - // `FnOnce` is not general enough") once this async fn is spawned. Eager - // collection ties each future to this stack frame's concrete lifetime. - if batches.len() > 1 { - log::debug!( - "[tree::bucket_seal] doc-subtree concurrent seal tree_id={} doc_id_hash={} level={} batches={} concurrency={}", - tree.id, - redact(doc_id), - current_level, - batches.len(), - DOC_SUBTREE_SEAL_CONCURRENCY - ); - } - let batch_futures: Vec<_> = batches - .iter() - .map(|batch| { - seal_explicit_children( - config, - tree, - current_level, - batch, - Some(doc_id), - version_ms, - strategy, - ) - }) - .collect(); - let nodes: Vec = futures::stream::iter(batch_futures) - .buffered(DOC_SUBTREE_SEAL_CONCURRENCY) - .try_collect() - .await?; - - let next_ids: Vec = nodes.iter().map(|n| n.id.clone()).collect(); - // The terminal level produces exactly one node (the loop breaks once - // `current_ids.len() <= 1`); preserving batch order means the last - // node is the doc-root, identical to the serial assignment. - doc_root = nodes.into_iter().next_back(); - - current_level += 1; - current_ids = next_ids; - if current_ids.len() <= 1 { - break; - } - } - - let doc_root = doc_root.ok_or_else(|| { - anyhow::anyhow!( - "[tree::bucket_seal] seal_document_subtree produced no doc-root tree_id={} doc_id_hash={}", - tree.id, - redact(doc_id) - ) - })?; - log::debug!( - "[tree::bucket_seal] doc-root sealed tree_id={} doc_id_hash={} root_id={} level={}", - tree.id, - redact(doc_id), - doc_root.id, - doc_root.level - ); - - // 2. Feed the doc-root into the cross-document merge tier and cascade - // using the untouched shared engine. - append_to_buffer( - config, - &tree.id, - MERGE_LEVEL_BASE, - &doc_root.id, - doc_root.token_count as i64, - doc_root.time_range_start, - )?; - let merge_sealed = cascade_all_from(config, tree, MERGE_LEVEL_BASE, None, strategy).await?; - log::debug!( - "[tree::bucket_seal] merge cascade tree_id={} doc_id_hash={} merge_sealed={}", - tree.id, - redact(doc_id), - merge_sealed.len() - ); - - Ok(doc_root.id) + crate::openhuman::tinycortex::seal_document_subtree( + config, tree, doc_id, version_ms, chunk_ids, strategy, + ) + .await } -/// Greedily batch leaf chunk ids so each batch stays under -/// [`INPUT_TOKEN_BUDGET`] (and at most [`DOC_SUBTREE_MAX_FANIN`] children). -/// A single oversized chunk forms its own batch. -fn batch_leaves_by_token_budget(config: &Config, chunk_ids: &[String]) -> Result>> { - use crate::openhuman::memory_store::chunks::store::get_chunk; - - let mut batches: Vec> = Vec::new(); - let mut current: Vec = Vec::new(); - let mut token_sum: i64 = 0; - - for id in chunk_ids { - let tokens = match get_chunk(config, id)? { - Some(c) => c.token_count as i64, - None => { - log::warn!( - "[tree::bucket_seal] batch_leaves_by_token_budget: missing chunk {id} — skipping" - ); - continue; - } - }; - let would_exceed = token_sum + tokens > INPUT_TOKEN_BUDGET as i64 - || current.len() >= DOC_SUBTREE_MAX_FANIN; - if would_exceed && !current.is_empty() { - batches.push(std::mem::take(&mut current)); - token_sum = 0; - } - current.push(id.clone()); - token_sum += tokens; - } - if !current.is_empty() { - batches.push(current); - } - if batches.is_empty() { - // All chunks were missing — surface as one empty batch caller rejects. - anyhow::bail!("[tree::bucket_seal] batch_leaves_by_token_budget: no resolvable chunks"); - } - Ok(batches) -} - -/// Split ids into fixed-size batches of at most `max` (used above L0 in the -/// per-document cascade). -fn batch_by_count(ids: &[String], max: usize) -> Vec> { - ids.chunks(max.max(1)).map(|c| c.to_vec()).collect() -} - -/// Seal an **explicit** set of child ids into one summary at `level + 1`, -/// tagging it with `doc_id` / `version_ms`. Unlike [`seal_one_level`] this -/// does NOT touch any shared `(tree, level)` buffer and does NOT advance the -/// tree root/max_level — it is the per-document side-cascade primitive. It -/// reuses the same hydrate → summarise → label → embed → stage → persist -/// pipeline so doc-subtree summaries are indistinguishable from regular -/// summaries except for their `doc_id` / `version_ms` tags. -async fn seal_explicit_children( +pub(crate) async fn seal_one_level( config: &Config, tree: &Tree, - level: u32, - child_ids: &[String], - doc_id: Option<&str>, - version_ms: Option, + buffer: &Buffer, strategy: &LabelStrategy, -) -> Result { - let target_level = level + 1; - let inputs = hydrate_inputs(config, level, child_ids)?; - if inputs.is_empty() { - anyhow::bail!( - "[tree::bucket_seal] seal_explicit_children: empty inputs tree_id={} level={}", - tree.id, - level - ); - } - - let time_range_start = inputs - .iter() - .map(|i| i.time_range_start) - .min() - .unwrap_or_else(Utc::now); - let time_range_end = inputs - .iter() - .map(|i| i.time_range_end) - .max() - .unwrap_or_else(Utc::now); - let score = inputs - .iter() - .map(|i| i.score) - .fold(f32::NEG_INFINITY, f32::max) - .max(0.0); - - let ctx = SummaryContext { - tree_id: &tree.id, - tree_kind: tree.kind, - target_level, - token_budget: OUTPUT_TOKEN_BUDGET, - }; - // Single-input passthrough: if a doc rolls up from exactly one node that - // already fits the summary budget (the common case — a Notion page that is - // a single chunk), there is nothing to summarise. Emit the input verbatim - // as the doc-root content and SKIP the LLM call entirely. The doc-root is - // still a real summary node (so versioning, the merge tier, and the - // read-time latest-version filter all keep working uniformly) — it just - // isn't a redundant paraphrase of one chunk, and costs no inference. - // A single oversized input still goes through the summariser (it genuinely - // needs compression). - let output = if inputs.len() == 1 && inputs[0].token_count <= OUTPUT_TOKEN_BUDGET { - log::debug!( - "[tree::bucket_seal] doc-subtree passthrough (1 input, no LLM) tree_id={} doc_id_hash={} level={}", - tree.id, - doc_id.map(redact).unwrap_or_default(), - level - ); - let only = &inputs[0]; - SummaryOutput { - content: only.content.clone(), - token_count: only.token_count, - entities: Vec::new(), - topics: Vec::new(), - input_tokens: 0, - output_tokens: 0, - charged_amount_usd: None, - } - } else { - // #13021: blank summary (whitespace-only LLM output) → fallback. See - // the matching branch in `seal_one_level` for why we treat blank - // content as a hard fail rather than persisting `content = ""`. - match summarise(config, &inputs, &ctx).await { - Ok(o) if !o.content.trim().is_empty() => o, - Ok(_) => { - log::warn!( - "[tree::bucket_seal] doc-subtree summarise returned blank tree_id={} doc_id_hash={} level={} — fallback (#13021)", - tree.id, doc_id.map(redact).unwrap_or_default(), level, - ); - fallback_summary(&inputs, ctx.token_budget) - } - Err(e) => { - log::warn!( - "[tree::bucket_seal] doc-subtree summarise failed tree_id={} doc_id_hash={} level={}: {e:#} — fallback", - tree.id, doc_id.map(redact).unwrap_or_default(), level, - ); - fallback_summary(&inputs, ctx.token_budget) - } - } - }; - - let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?; - - // Embed before any write so a failure aborts cleanly — same contract as - // seal_one_level. No-provider configs seal embedding-less. - // - // #13021: skip when the embed input is empty/whitespace. `summarise()` - // returns an empty content default when every input is whitespace, which - // would otherwise round-trip a guaranteed 400 from the upstream embedding - // API. The doc-subtree seal is sealed embedding-less, matching the - // no-provider branch. - let embed_input = truncate_for_embed(&output.content, 1_000); - let embedding: Option> = if embed_input.trim().is_empty() { - log::warn!( - "[tree::bucket_seal] doc-subtree: empty summary content for tree_id={} level={} \ - — sealing without embedding (#13021)", - tree.id, - level - ); - None - } else { - match build_write_embedder(config).context("build embedder during doc-subtree seal")? { - None => None, - Some(embedder) => { - let v = embedder.embed(&embed_input).await.map_err(|e| { - let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e); - anyhow::Error::new(failure).context(format!( - "embed doc-subtree summary tree_id={} level={}: {e:#}", - tree.id, level - )) - })?; - crate::openhuman::memory_tree::score::embed::pack_checked(&v).context(format!( - "doc-subtree embed dim check tree_id={} level={}", - tree.id, level - ))?; - Some(v) - } - } - }; - - let now = Utc::now(); - let summary_id = new_summary_id(target_level); - let node = SummaryNode { - id: summary_id.clone(), - 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: node_entities, - topics: node_topics, - time_range_start, - time_range_end, - score, - sealed_at: now, - deleted: false, - embedding, - doc_id: doc_id.map(|s| s.to_string()), - version_ms, - }; - - // Stage the .md file before opening the write tx (same fail-fast as - // seal_one_level). Doc-subtree nodes land under - // `source-/docs//v-/…` so the vault mirrors - // the logical shape. Wikilink overrides are left unset. - let tree_factory = TreeFactory::from_tree(tree); - let summary_tree_kind = tree_factory.summary_tree_kind(); - let scope_slug = tree_factory.scope_slug(); - let compose_input = SummaryComposeInput { - summary_id: &node.id, - tree_kind: summary_tree_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, - }; - let content_root = config.memory_tree_content_root(); - if let Err(err) = - crate::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults(&content_root) - { - log::warn!( - "[tree::bucket_seal] ensure_obsidian_defaults failed during doc-subtree seal: {err:#}" - ); - } - let doc_slug = doc_id.map(slugify_source_id); - let layout = match doc_slug.as_deref() { - Some(slug) => SummaryDiskLayout::DocSubtree { - doc_slug: slug, - version_ms, - }, - None => SummaryDiskLayout::Standard, - }; - let staged = stage_summary_with_layout(&content_root, &compose_input, &scope_slug, layout) - .with_context(|| { - format!( - "stage_summary failed for doc-subtree node {}; seal aborted", - node.id - ) - })?; - let summary_commit = summary_seal_batch( + enqueue_follow_ups: bool, +) -> Result { + crate::openhuman::tinycortex::seal_tree_level( + config, tree, - "document_subtree_seal", - std::slice::from_ref(&node), - std::slice::from_ref(&staged.content_path), + buffer, + strategy, + enqueue_follow_ups, ) - .with_context(|| { - format!( - "build summary git batch failed for doc-subtree node {}; seal aborted", - node.id - ) - })?; - - // Persist the summary row + backlink children — NO buffer / tree-root - // mutation (those belong to the merge tier). - let node_for_tx = node.clone(); - let level_for_tx = level; - let summary_id_for_tx = summary_id.clone(); - let signature = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); - with_connection(config, move |conn| { - let tx = conn.unchecked_transaction()?; - store::insert_summary_tx(&tx, &node_for_tx, Some(&staged), &signature)?; - crate::openhuman::memory_tree::score::store::index_summary_entity_ids_tx( - &tx, - &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_for_tx == 0 { - // Unconditional re-point (no `IS NULL` guard): a byte-identical - // body chunk reused across doc versions upserts to the SAME row - // (content-addressed id), so its single `parent_summary_id` must - // follow the newest version. Doc subtrees seal newest-last, so - // last-write-wins leaves the backlink on the latest doc-root — - // the version retrieval surfaces — instead of stranding it on - // the first (now-superseded) version's summary. - tx.execute( - "UPDATE mem_tree_chunks SET parent_summary_id = ?1 \ - WHERE id = ?2", - rusqlite::params![&summary_id_for_tx, child_id], - ) - .context("backlink chunk to doc-subtree summary")?; - } else { - tx.execute( - "UPDATE mem_tree_summaries SET parent_id = ?1 \ - WHERE id = ?2 AND parent_id IS NULL", - rusqlite::params![&summary_id_for_tx, child_id], - ) - .context("backlink summary to doc-subtree parent")?; - } - } - tx.commit()?; - Ok(()) - })?; - - commit_summary_seal(config, &summary_commit).with_context(|| { - format!( - "commit_summary_seal failed for doc-subtree node {} after DB seal commit", - summary_id - ) - })?; - - log::info!( - "[tree::bucket_seal] doc-subtree sealed tree_id={} doc_id_hash={} level={}→{} summary_id={} children={}", - tree.id, - doc_id.map(redact).unwrap_or_default(), - level, - target_level, - summary_id, - child_ids.len() - ); - - Ok(node) + .await } - -#[cfg(test)] -#[path = "bucket_seal_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs b/src/openhuman/memory_tree/tree/bucket_seal_tests.rs deleted file mode 100644 index d9440b2de..000000000 --- a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs +++ /dev/null @@ -1,1358 +0,0 @@ -//! Unit tests for [`super::bucket_seal`] — append + cascade-seal mechanics -//! for source/topic trees. Covers L0 token gating, L≥1 fanout gating, -//! cascade depth bounds, idempotency on retry, and label-strategy resolution. - -use super::*; -use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider}; -use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree; -use crate::openhuman::memory_store::content as content_store; -use crate::openhuman::memory_store::trees::types::TreeKind; -use std::sync::Arc; -use tempfile::TempDir; - -/// Stage a batch of chunks to the content store so that `read_chunk_body` -/// can find the on-disk file during seals. Tests that call `upsert_chunks` -/// and then trigger a seal MUST also call this helper; otherwise -/// `hydrate_leaf_inputs` will fail with "no content_path for chunk_id". -fn stage_test_chunks( - cfg: &Config, - chunks: &[crate::openhuman::memory_store::chunks::types::Chunk], -) { - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).expect("create content_root for test"); - let staged = - content_store::stage_chunks(&content_root, chunks).expect("stage_chunks for test chunks"); - // Record the content_path + content_sha256 pointers in SQLite so the - // store's `get_chunk_content_pointers` can resolve them later. - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .expect("persist staged chunk pointers"); -} - -fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): seal calls the embedder — force inert so - // tests don't require a running Ollama. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - // #002: opt into the deterministic inert embedder via `provider="none"`. - // This is `Some(inert)` (vector search off by choice) and does NOT set the - // process-global semantic-recall degraded flag — unlike the no-provider - // path, which marks degraded and would leak that signal into parallel - // `pipeline_status` tests. - cfg.embeddings_provider = Some("none".into()); - (tmp, cfg) -} - -fn mk_leaf(id: &str, tokens: u32, ts_ms: i64) -> LeafRef { - use chrono::TimeZone; - LeafRef { - chunk_id: id.to_string(), - token_count: tokens, - timestamp: Utc.timestamp_millis_opt(ts_ms).single().unwrap(), - content: format!("content for {id}"), - entities: vec![], - topics: vec![], - score: 0.5, - } -} - -#[tokio::test] -async fn append_below_budget_does_not_seal() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - // Chunks don't exist in DB — we're only exercising the buffer - // accounting, which doesn't require leaf rows until a seal fires. - let leaf = mk_leaf("leaf-1", 100, 1_700_000_000_000); - let sealed = test_override::with_provider(provider, async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert!(sealed.is_empty()); - - let buf = store::get_buffer(&cfg, &tree.id, 0).unwrap(); - assert_eq!(buf.item_ids, vec!["leaf-1".to_string()]); - assert_eq!(buf.token_sum, 100); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); -} - -/// Build + persist a Notion-style Document chunk (staged to disk so the -/// seal hydrator can read its body). -fn seed_doc_chunk( - cfg: &Config, - doc_id: &str, - seq: u32, - content: &str, -) -> crate::openhuman::memory_store::chunks::types::Chunk { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - let ts = Utc::now(); - let c = Chunk { - id: chunk_id(SourceKind::Document, doc_id, seq, content), - content: content.to_string(), - metadata: Metadata { - source_kind: SourceKind::Document, - source_id: doc_id.to_string(), - owner: "notion:conn1".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["notion".into()], - source_ref: Some(SourceRef::new("notion://page/pageA")), - path_scope: Some("notion:conn1".into()), - }, - token_count: 10, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - upsert_chunks(cfg, &[c.clone()]).unwrap(); - stage_test_chunks(cfg, &[c.clone()]); - c -} - -#[tokio::test] -async fn seal_document_subtree_force_seals_small_doc_to_one_root() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("doc summary")); - - let doc_id = "notion:conn1:pageA"; - let c0 = seed_doc_chunk(&cfg, doc_id, 0, "first chunk body"); - let c1 = seed_doc_chunk(&cfg, doc_id, 1, "second chunk body"); - - let doc_root_id = test_override::with_provider(provider, async { - seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(100), - &[c0.id.clone(), c1.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap() - }) - .await; - - // Two small chunks collapse to exactly ONE doc-root, tagged with the - // document id + version. - let root = store::get_summary(&cfg, &doc_root_id).unwrap().unwrap(); - assert_eq!(root.doc_id.as_deref(), Some(doc_id)); - assert_eq!(root.version_ms, Some(100)); - assert_eq!( - root.child_ids.len(), - 2, - "both chunks roll into the doc-root" - ); - - // The doc-root is fed into the cross-document merge buffer (not the flat - // L0 buffer), so the connection tree can fold it with other documents. - let merge_buf = store::get_buffer(&cfg, &tree.id, MERGE_LEVEL_BASE).unwrap(); - assert!( - merge_buf.item_ids.contains(&doc_root_id), - "doc-root must land in the merge buffer; got {:?}", - merge_buf.item_ids - ); - // Per-doc subtree must NOT pollute the flat L0 buffer. - let l0 = store::get_buffer(&cfg, &tree.id, 0).unwrap(); - assert!( - l0.item_ids.is_empty(), - "L0 buffer stays empty for documents" - ); -} - -#[tokio::test] -async fn seal_document_subtree_new_version_is_additive() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("doc summary")); - - let doc_id = "notion:conn1:pageA"; - - // Version 1. - let v1c = seed_doc_chunk(&cfg, doc_id, 0, "v1 body"); - let v1_root = test_override::with_provider(Arc::clone(&provider), async { - seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(100), - &[v1c.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap() - }) - .await; - - // Version 2 (edited page → new chunk content → new chunk id). - let v2c = seed_doc_chunk(&cfg, doc_id, 0, "v2 body edited"); - let v2_root = test_override::with_provider(Arc::clone(&provider), async { - seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(200), - &[v2c.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap() - }) - .await; - - assert_ne!(v1_root, v2_root, "a new version mints a new doc-root"); - - // Forward-only: BOTH doc-roots persist (nothing tombstoned), and both - // sit in the merge buffer. Read-time latest-wins (drill_down) is what - // surfaces only v2 — the write path never destroys v1. - let r1 = store::get_summary(&cfg, &v1_root).unwrap().unwrap(); - let r2 = store::get_summary(&cfg, &v2_root).unwrap().unwrap(); - assert_eq!(r1.version_ms, Some(100)); - assert_eq!(r2.version_ms, Some(200)); - - let merge_buf = store::get_buffer(&cfg, &tree.id, MERGE_LEVEL_BASE).unwrap(); - assert!(merge_buf.item_ids.contains(&v1_root)); - assert!(merge_buf.item_ids.contains(&v2_root)); -} - -/// A byte-identical body chunk reused across two versions of a multi-chunk doc -/// upserts to the SAME row (content-addressed id). Its single -/// `parent_summary_id` backlink must follow the NEWEST version's doc-root — the -/// one drill_down surfaces — not stay stranded on the first (now-superseded) -/// version. Guards the unconditional re-point in `seal_explicit_children`. -#[tokio::test] -async fn shared_chunk_backlink_repoints_to_latest_doc_version() { - use crate::openhuman::memory_store::chunks::store::with_connection; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("doc summary")); - let doc_id = "notion:conn1:pageA"; - - // seq 0 is byte-identical across versions → one shared row. seq 1 differs, - // so each version is a genuine 2-chunk doc (not the single-chunk passthrough). - let shared = seed_doc_chunk(&cfg, doc_id, 0, "shared body identical across versions"); - - let v1_other = seed_doc_chunk(&cfg, doc_id, 1, "v1 second chunk"); - let v1_root = test_override::with_provider(Arc::clone(&provider), async { - seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(100), - &[shared.id.clone(), v1_other.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap() - }) - .await; - - // After v1, the shared chunk backlinks to v1's doc-root. - let p1: Option = with_connection(&cfg, |conn| { - Ok(conn - .query_row( - "SELECT parent_summary_id FROM mem_tree_chunks WHERE id = ?1", - rusqlite::params![shared.id], - |r| r.get(0), - ) - .unwrap()) - }) - .unwrap(); - assert_eq!(p1.as_deref(), Some(v1_root.as_str())); - - // Re-ingest the same shared chunk (idempotent upsert) and seal version 2. - let _shared_again = seed_doc_chunk(&cfg, doc_id, 0, "shared body identical across versions"); - let v2_other = seed_doc_chunk(&cfg, doc_id, 1, "v2 second chunk edited"); - let v2_root = test_override::with_provider(Arc::clone(&provider), async { - seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(200), - &[shared.id.clone(), v2_other.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap() - }) - .await; - assert_ne!(v1_root, v2_root); - - // The shared chunk's backlink now follows the LATEST version's doc-root. - let p2: Option = with_connection(&cfg, |conn| { - Ok(conn - .query_row( - "SELECT parent_summary_id FROM mem_tree_chunks WHERE id = ?1", - rusqlite::params![shared.id], - |r| r.get(0), - ) - .unwrap()) - }) - .unwrap(); - assert_eq!( - p2.as_deref(), - Some(v2_root.as_str()), - "shared chunk must re-point to the latest doc-root, not stay on v1" - ); -} - -/// Single-chunk passthrough: a doc that rolls up from exactly one -/// budget-fitting chunk must NOT invoke the summariser — the doc-root content -/// is the chunk **verbatim**. Proven two ways: (1) no `ChatProvider` override -/// is installed, so any summarise() call would hit the (unconfigured) cloud -/// path and never reproduce this exact text; (2) the doc-root body is asserted -/// byte-equal to the chunk body. -#[tokio::test] -async fn seal_document_subtree_single_chunk_is_verbatim_passthrough_no_llm() { - use crate::openhuman::memory_store::content::read as content_read; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); - let doc_id = "notion:conn1:pageX"; - // Distinctive content the summariser would never emit verbatim. - let unique = "UNIQUE-PASSTHROUGH-MARKER-7Z\n\n- line one\n- line two"; - let c = seed_doc_chunk(&cfg, doc_id, 0, unique); - - // NOTE: no test_override::with_provider — passthrough must not need the LLM. - let doc_root_id = seal_document_subtree( - &cfg, - &tree, - doc_id, - Some(100), - &[c.id.clone()], - &LabelStrategy::Empty, - ) - .await - .unwrap(); - - let root = store::get_summary(&cfg, &doc_root_id).unwrap().unwrap(); - assert_eq!(root.doc_id.as_deref(), Some(doc_id)); - assert_eq!(root.version_ms, Some(100)); - - // Doc-root body (read full from disk) must be the chunk verbatim. - let body = content_read::read_summary_body(&cfg, &doc_root_id).unwrap(); - assert_eq!( - body.trim(), - unique, - "single-chunk doc-root must be the chunk verbatim (no summarisation / LLM)" - ); -} - -/// Multi-batch document subtree: a doc whose chunks span several level-0 -/// batches must seal its siblings **concurrently** (overlapping the -/// summarise + embed round-trips) while still producing an identical -/// doc-root. A `ChatProvider` probe records the peak number of in-flight -/// summarise calls; with four independent level-0 batches the peak must -/// exceed 1 (proving overlap) yet stay within `DOC_SUBTREE_SEAL_CONCURRENCY` -/// (proving the bound is respected). This is the regression guard for the -/// serial → bounded-concurrency change in `seal_document_subtree`. -#[tokio::test] -async fn seal_document_subtree_seals_sibling_batches_concurrently() { - use crate::openhuman::memory::chat::ChatPrompt; - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use std::sync::atomic::{AtomicUsize, Ordering}; - - /// Records concurrent in-flight `chat_for_json` calls and the peak; a - /// short sleep forces would-be-concurrent calls to actually overlap so - /// the peak reflects real scheduling, not luck. - struct ConcurrencyProbe { - in_flight: AtomicUsize, - peak: AtomicUsize, - calls: AtomicUsize, - } - #[async_trait::async_trait] - impl ChatProvider for ConcurrencyProbe { - fn name(&self) -> &str { - "test:concurrency-probe" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.peak.fetch_max(now, Ordering::SeqCst); - self.calls.fetch_add(1, Ordering::SeqCst); - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - self.in_flight.fetch_sub(1, Ordering::SeqCst); - Ok("doc batch summary".to_string()) - } - } - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); - let doc_id = "notion:conn1:pageM"; - - // Four chunks, each ~60% of the input budget, so the greedy token-budget - // batcher puts each in its OWN level-0 batch → four independent - // summarise calls that should overlap under bounded concurrency. - let per_chunk = INPUT_TOKEN_BUDGET * 6 / 10; - let ts = Utc::now(); - let mut ids = Vec::new(); - let mut chunks = Vec::new(); - for seq in 0..4u32 { - let content = format!("substantive document body number {seq}"); - let c = Chunk { - id: chunk_id(SourceKind::Document, doc_id, seq, &content), - content, - metadata: Metadata { - source_kind: SourceKind::Document, - source_id: doc_id.to_string(), - owner: "notion:conn1".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec!["notion".into()], - source_ref: Some(SourceRef::new("notion://page/pageM")), - path_scope: Some("notion:conn1".into()), - }, - token_count: per_chunk, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - ids.push(c.id.clone()); - chunks.push(c); - } - upsert_chunks(&cfg, &chunks).unwrap(); - stage_test_chunks(&cfg, &chunks); - - let probe = Arc::new(ConcurrencyProbe { - in_flight: AtomicUsize::new(0), - peak: AtomicUsize::new(0), - calls: AtomicUsize::new(0), - }); - let provider: Arc = probe.clone(); - - let doc_root_id = test_override::with_provider(provider, async { - seal_document_subtree(&cfg, &tree, doc_id, Some(7), &ids, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - - // Correctness: the four chunks roll up under a single doc-root tagged - // with the document id + version — tree shape unchanged by batching. - let root = store::get_summary(&cfg, &doc_root_id).unwrap().unwrap(); - assert_eq!(root.doc_id.as_deref(), Some(doc_id)); - assert_eq!(root.version_ms, Some(7)); - assert_eq!( - root.child_ids.len(), - 4, - "all four level-0 siblings must fan into the doc-root" - ); - - // Concurrency: at least two level-0 siblings overlapped, and the pool - // bound was respected. - let peak = probe.peak.load(Ordering::SeqCst); - assert!( - peak >= 2, - "expected overlapping sibling seals (peak in-flight = {peak})" - ); - assert!( - peak <= DOC_SUBTREE_SEAL_CONCURRENCY, - "concurrency must stay within the pool bound (peak = {peak})" - ); -} - -#[tokio::test] -async fn crossing_budget_triggers_seal() { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use chrono::TimeZone; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - - // Persist two chunks that the hydrator can load during seal. - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let mk_chunk = |seq: u32, tokens: u32| Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "test-content"), - content: format!("substantive chunk content {seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: tokens, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - // Budget-relative sizes so the test stays correct as INPUT_TOKEN_BUDGET shifts: - // each leaf is 60% of budget, so the second append crosses the threshold. - let per_leaf = INPUT_TOKEN_BUDGET * 6 / 10; - let c1 = mk_chunk(0, per_leaf); - let c2 = mk_chunk(1, per_leaf); - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - // Stage both chunks to disk so the seal's hydrator can read full bodies. - stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]); - - // Two leaves whose combined token_sum (12k) exceeds the 10k budget. - let leaf1 = LeafRef { - chunk_id: c1.id.clone(), - token_count: per_leaf, - timestamp: ts, - content: c1.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - let leaf2 = LeafRef { - chunk_id: c2.id.clone(), - token_count: per_leaf, - timestamp: ts, - content: c2.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - - let first = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf1, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert!(first.is_empty(), "first append below budget — no seal"); - - let second = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf2, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(second.len(), 1, "second append crosses budget — one seal"); - - let summary_id = &second[0]; - let summary = store::get_summary(&cfg, summary_id).unwrap().unwrap(); - assert_eq!(summary.level, 1); - assert_eq!(summary.child_ids, vec![c1.id.clone(), c2.id.clone()]); - assert!(summary.token_count > 0); - - // L0 buffer cleared, L1 buffer carries the new summary id. - let l0 = store::get_buffer(&cfg, &tree.id, 0).unwrap(); - assert!(l0.is_empty()); - let l1 = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - assert_eq!(l1.item_ids, vec![summary_id.clone()]); - - // Tree metadata updated. - let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); - assert_eq!(t.max_level, 1); - assert_eq!(t.root_id.as_deref(), Some(summary_id.as_str())); - assert!(t.last_sealed_at.is_some()); - - // Leaf → parent backlink populated for both children. - use crate::openhuman::memory_store::chunks::store::with_connection; - let parent: Option = with_connection(&cfg, |conn| { - let p: Option = conn - .query_row( - "SELECT parent_summary_id FROM mem_tree_chunks WHERE id = ?1", - rusqlite::params![c1.id], - |r| r.get(0), - ) - .unwrap(); - Ok(p) - }) - .unwrap(); - assert_eq!(parent.as_deref(), Some(summary_id.as_str())); -} - -#[tokio::test] -async fn fanout_at_l1_triggers_l2_seal() { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::trees::types::SUMMARY_FANOUT; - use chrono::TimeZone; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let mk_chunk = |seq: u32| { - let content = format!("substantive chunk content {seq}"); - Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, &content), - content, - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - // Each leaf alone busts INPUT_TOKEN_BUDGET so the L0→L1 seal - // fires on every append. After SUMMARY_FANOUT seals, the - // L1 buffer's count-based gate trips and cascades to L2. - token_count: INPUT_TOKEN_BUDGET + 1, - seq_in_source: seq, - created_at: ts, - partial_message: false, - } - }; - - let fanout = SUMMARY_FANOUT; - let mut all_sealed: Vec = Vec::new(); - for seq in 0..fanout { - let chunk = mk_chunk(seq); - upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); - // Stage to disk so the seal hydrator can read the full body. - stage_test_chunks(&cfg, &[chunk.clone()]); - let leaf = LeafRef { - chunk_id: chunk.id.clone(), - token_count: chunk.token_count, - timestamp: ts, - content: chunk.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - let sealed = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - all_sealed.extend(sealed); - } - - // First (fanout-1) appends each emit one L1 seal. The final - // append emits an L1 seal AND cascades into one L2 seal. - assert_eq!( - all_sealed.len() as u32, - fanout + 1, - "expected {} L1 seals + 1 L2 seal, got {}", - fanout, - all_sealed.len() - ); - - let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); - assert_eq!(t.max_level, 2, "tree should have climbed to L2"); - - let l1 = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - assert!( - l1.is_empty(), - "L1 buffer should clear when the fanout seal fires" - ); - - let l2 = store::get_buffer(&cfg, &tree.id, 2).unwrap(); - assert_eq!(l2.item_ids.len(), 1, "exactly one L2 summary queued"); - - let l2_summary = store::get_summary(&cfg, &l2.item_ids[0]).unwrap().unwrap(); - assert_eq!(l2_summary.level, 2); - assert_eq!( - l2_summary.child_ids.len() as u32, - fanout, - "L2 summary should fold all {fanout} L1 children" - ); -} - -#[tokio::test] -async fn upper_level_does_not_seal_below_fanout() { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::trees::types::SUMMARY_FANOUT; - use chrono::TimeZone; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - // Emit (fanout - 1) L1 summaries — should leave the L1 buffer - // populated but BELOW the count gate, so no L2 seal. - let stop_before = SUMMARY_FANOUT.saturating_sub(1); - for seq in 0..stop_before { - let content = format!("c{seq}"); - let chunk = Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, &content), - content, - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: INPUT_TOKEN_BUDGET + 1, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); - // Stage to disk so the seal hydrator can read the full body. - stage_test_chunks(&cfg, &[chunk.clone()]); - let leaf = LeafRef { - chunk_id: chunk.id, - token_count: chunk.token_count, - timestamp: ts, - content: chunk.content, - entities: vec![], - topics: vec![], - score: 0.5, - }; - let _ = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - } - - let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); - assert_eq!(t.max_level, 1, "should plateau at L1 below fanout"); - - let l1 = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - assert_eq!( - l1.item_ids.len() as u32, - stop_before, - "L1 buffer should hold the unsealed siblings" - ); - assert_eq!( - store::count_summaries(&cfg, &tree.id).unwrap(), - stop_before as u64 - ); -} - -// ── LabelStrategy tests (#TBD) ──────────────────────────────────────────── -// -// These exercise the three labeling modes seal_one_level supports. We use -// a short token budget so the seal fires on a single leaf — keeps the -// arithmetic of "what entities/topics end up on the parent" obvious. - -/// Helper: persist a substantive chunk and return a `LeafRef` referencing -/// it, with caller-supplied entity/topic labels (used by Union/Empty tests). -/// -/// To match production, entity labels are written into `mem_tree_entity_index` -/// (where seal-time hydration reads them from) and topic labels are stored -/// on `chunk.metadata.tags` (the production source of leaf-level topics). -fn seed_leaf( - cfg: &Config, - seq: u32, - content: &str, - entities: Vec, - topics: Vec, -) -> LeafRef { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_tree::score::extract::EntityKind; - use crate::openhuman::memory_tree::score::resolver::CanonicalEntity; - use crate::openhuman::memory_tree::score::store::index_entity; - use chrono::TimeZone; - let ts = Utc - .timestamp_millis_opt(1_700_000_000_000 + seq as i64) - .unwrap(); - let chunk = Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, content), - content: content.to_string(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: topics.clone(), - source_ref: Some(SourceRef::new(format!("slack://x{seq}"))), - path_scope: None, - }, - // Bust INPUT_TOKEN_BUDGET in one leaf so the seal fires immediately. - token_count: INPUT_TOKEN_BUDGET + 1, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - upsert_chunks(cfg, &[chunk.clone()]).unwrap(); - // Stage the chunk to disk so `hydrate_leaf_inputs` can read the full body - // via `read_chunk_body` during a seal triggered by `append_leaf`. - stage_test_chunks(cfg, &[chunk.clone()]); - // Mirror production indexing: entities go into mem_tree_entity_index - // so the seal hydrator can pull them via list_entity_ids_for_node. - for entity_id in &entities { - let kind = entity_id - .split_once(':') - .map_or(EntityKind::Misc, |(k, _)| { - EntityKind::parse(k).unwrap_or(EntityKind::Misc) - }); - let surface = entity_id - .split_once(':') - .map_or(entity_id.as_str(), |(_, v)| v); - let e = CanonicalEntity { - canonical_id: entity_id.clone(), - kind, - surface: surface.to_string(), - span_start: 0, - span_end: surface.len() as u32, - score: 1.0, - }; - index_entity(cfg, &e, &chunk.id, "leaf", ts.timestamp_millis(), None).unwrap(); - } - LeafRef { - chunk_id: chunk.id.clone(), - token_count: chunk.token_count, - timestamp: ts, - content: chunk.content.clone(), - entities, - topics, - score: 0.5, - } -} - -#[tokio::test] -async fn seal_with_extract_strategy_populates_entities_and_topics() { - use crate::openhuman::memory_tree::score::extract::{CompositeExtractor, EntityExtractor}; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new( - "alice@example.com is leading the #launch sprint this week.", - )); - - // Content the regex extractor can find: an email and a hashtag. The - // StaticChatProvider returns content that the extractor finds. - let leaf = seed_leaf( - &cfg, - 0, - "alice@example.com is leading the #launch sprint this week.", - vec![], - vec![], - ); - - let extractor: Arc = Arc::new(CompositeExtractor::regex_only()); - let strategy = LabelStrategy::ExtractFromContent(extractor); - - let sealed = test_override::with_provider(provider, async { - append_leaf(&cfg, &tree, &leaf, &strategy).await.unwrap() - }) - .await; - assert_eq!(sealed.len(), 1, "single 10k-token leaf should seal L0→L1"); - - let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); - assert!( - summary - .entities - .iter() - .any(|e| e == "email:alice@example.com"), - "ExtractFromContent should surface the email entity from summary text; got entities={:?}", - summary.entities - ); - assert!( - summary.topics.iter().any(|t| t == "launch"), - "ExtractFromContent should surface the hashtag-derived topic; got topics={:?}", - summary.topics - ); -} - -#[tokio::test] -async fn seal_with_union_strategy_inherits_labels_from_children() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - - // Two leaves with overlapping + distinct labels. Union should - // dedup-merge them into the parent. - let leaf1 = seed_leaf( - &cfg, - 0, - "first leaf body", - vec!["email:alice@example.com".into(), "topic:phoenix".into()], - vec!["phoenix".into(), "launch".into()], - ); - let leaf2 = seed_leaf( - &cfg, - 1, - "second leaf body", - vec!["email:alice@example.com".into(), "person:bob".into()], - vec!["launch".into(), "qa".into()], - ); - - // L0 seals when the budget is crossed. With each leaf at 10k tokens, - // the first append triggers a seal containing only leaf1; we want a - // seal containing both, so use UnionFromChildren and a single seal of - // both leaves at once. The simplest way is to lower budget by sealing - // two leaves into one buffer — the second append crosses budget, so - // the seal contains [leaf1, leaf2]. - // - // Adjust by using smaller token counts so both fit in L0 first, then - // a third append triggers a seal containing both. Reuse the helper - // and override the leaf's token_count for this test. - // Each leaf at half the budget so two together hit threshold exactly. - let per_leaf = INPUT_TOKEN_BUDGET / 2; - let leaf1 = LeafRef { - token_count: per_leaf, - ..leaf1 - }; - let leaf2 = LeafRef { - token_count: per_leaf, - ..leaf2 - }; - - // First leaf: under budget, no seal. - let sealed_1 = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf1, &LabelStrategy::UnionFromChildren) - .await - .unwrap() - }) - .await; - assert!(sealed_1.is_empty()); - // Second leaf: crosses budget → one seal covering both leaves. - let sealed_2 = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf2, &LabelStrategy::UnionFromChildren) - .await - .unwrap() - }) - .await; - assert_eq!(sealed_2.len(), 1); - - let summary = store::get_summary(&cfg, &sealed_2[0]).unwrap().unwrap(); - let entities: std::collections::BTreeSet<&str> = - summary.entities.iter().map(String::as_str).collect(); - let topics: std::collections::BTreeSet<&str> = - summary.topics.iter().map(String::as_str).collect(); - assert!(entities.contains("email:alice@example.com")); - assert!(entities.contains("topic:phoenix")); - assert!(entities.contains("person:bob")); - assert_eq!( - entities.len(), - 3, - "expected 3 unique entities; got {entities:?}" - ); - assert!(topics.contains("phoenix")); - assert!(topics.contains("launch")); - assert!(topics.contains("qa")); - assert_eq!(topics.len(), 3, "expected 3 unique topics; got {topics:?}"); -} - -#[tokio::test] -async fn seal_with_empty_strategy_leaves_labels_empty() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - - // Leaf carries labels — Empty strategy should ignore them. - let leaf = seed_leaf( - &cfg, - 0, - "alice@example.com discussing #launch", - vec!["email:alice@example.com".into(), "topic:launch".into()], - vec!["launch".into()], - ); - - let sealed = test_override::with_provider(provider, async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(sealed.len(), 1); - - let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); - assert!( - summary.entities.is_empty(), - "Empty strategy must leave entities empty; got {:?}", - summary.entities - ); - assert!( - summary.topics.is_empty(), - "Empty strategy must leave topics empty; got {:?}", - summary.topics - ); -} - -#[tokio::test] -async fn topic_tree_seal_persists_topic_kind_not_source() { - use crate::openhuman::memory_store::trees::types::TreeStatus; - - let (_tmp, cfg) = test_config(); - // Build a topic tree directly — `seal_one_level` runs for both - // source and topic trees, and previously hardcoded Source on the - // resulting summary regardless of the parent tree's kind. - let tree = Tree { - id: "topic-tree-test-id".to_string(), - kind: TreeKind::Topic, - scope: "topic:launch".to_string(), - root_id: None, - max_level: 0, - status: TreeStatus::Active, - created_at: Utc::now(), - last_sealed_at: None, - }; - store::insert_tree(&cfg, &tree).unwrap(); - - let provider: Arc = Arc::new(StaticChatProvider::new("test summary content")); - let leaf = seed_leaf(&cfg, 0, "topic content", vec![], vec![]); - - let sealed = test_override::with_provider(provider, async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(sealed.len(), 1); - - let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); - assert_eq!( - summary.tree_kind, - TreeKind::Topic, - "topic-tree summary must persist tree_kind=Topic, not Source" - ); -} - -#[test] -fn scope_slug_non_gmail_uses_full_scope() { - // slack:#eng and discord:#eng must NOT produce the same scope slug. - // Previously, stripping everything before ':' made both → "eng". - // After Fix K, only gmail: strips the prefix — others use the full string. - use crate::openhuman::memory_store::content::paths::slugify_source_id; - - // Verify that the slug logic produces distinct values for different platforms. - let slack_slug = slugify_source_id("slack:#eng"); - let discord_slug = slugify_source_id("discord:#eng"); - assert_ne!( - slack_slug, discord_slug, - "slack:#eng and discord:#eng must produce distinct slugs; got slack={slack_slug:?} discord={discord_slug:?}" - ); - // Both must include their platform prefix in the slug. - assert!( - slack_slug.contains("slack"), - "slack slug must include 'slack'; got {slack_slug:?}" - ); - assert!( - discord_slug.contains("discord"), - "discord slug must include 'discord'; got {discord_slug:?}" - ); - - // Confirm gmail: correctly strips the "gmail:" prefix so the participants - // portion (used as the bucket key) matches the chunk path layout. - // scope_slug for a gmail source tree is built by stripping "gmail:" and - // slugifying the remainder; the result must equal slugify of just the - // participants string. - let participants = "alice@x.com|bob@y.com"; - let participants_slug = slugify_source_id(participants); - let gmail_scope = format!("gmail:{participants}"); - // Strip "gmail:" prefix as bucket_seal.rs does. - let gmail_slug = slugify_source_id(&gmail_scope["gmail:".len()..]); - assert_eq!( - participants_slug, gmail_slug, - "gmail scope_slug must equal slugify of participants portion; \ - participants_slug={participants_slug:?} gmail_slug={gmail_slug:?}" - ); - - // Also assert the full-scope slug for gmail is DIFFERENT (shows the bug - // would still exist if we used the full string for gmail). - let gmail_full_slug = slugify_source_id(&gmail_scope); - assert_ne!( - gmail_full_slug, participants_slug, - "slugifying the full 'gmail:...' scope must differ from the participants-only slug" - ); -} - -/// `hydrate_summary_inputs` was rewritten to do one batched -/// `get_summaries_batch` SELECT instead of N per-id `get_summary` -/// round-trips. This test pins three behavioural invariants the per-id -/// loop used to give us for free, and which the HashMap-walk now has to -/// reproduce: -/// -/// 1. **Input order preservation.** We iterate the caller's -/// `summary_ids` slice (not the HashMap) so the `SummaryInput`s come -/// out in the order the caller asked for, even though `HashMap` -/// iteration is not insertion-ordered. -/// 2. **Per-id field propagation by id, not by index.** Distinct field -/// values per summary (content, token_count, score) prove the -/// map.get() is keyed by id — not by enumerate(). -/// 3. **Missing id → silent skip + warn.** Mirrors the per-id -/// `Ok(None)` → continue contract; the request does not error out. -#[tokio::test] -async fn hydrate_summary_inputs_batch_preserves_order_and_skips_missing_ids() { - use crate::openhuman::memory_store::content::atomic::stage_summary; - use crate::openhuman::memory_store::content::SummaryComposeInput; - use crate::openhuman::memory_store::content::SummaryTreeKind; - use crate::openhuman::memory_store::trees::store::insert_tree; - use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind, TreeStatus}; - use crate::openhuman::memory_tree::tree::store::insert_summary_tx; - use chrono::TimeZone; - - let (_tmp, cfg) = test_config(); - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - - let tree = Tree { - id: "tree-hydrate".into(), - kind: TreeKind::Source, - scope: "slack:#eng".into(), - root_id: None, - max_level: 0, - status: TreeStatus::Active, - created_at: ts, - last_sealed_at: None, - }; - insert_tree(&cfg, &tree).unwrap(); - - // Two summaries with distinct content / token_count / score so we - // can prove the per-id HashMap lookup keys by id and not by - // enumerate() over `summary_ids`. - let sum_a = SummaryNode { - id: "sum-a".into(), - tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: vec!["leaf-a".into()], - content: "BODY-A".into(), - token_count: 11, - entities: vec!["entity:alice".into()], - topics: vec!["#a".into()], - time_range_start: ts, - time_range_end: ts, - score: 0.11, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - let sum_b = SummaryNode { - id: "sum-b".into(), - tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, - level: 1, - parent_id: None, - child_ids: vec!["leaf-b".into()], - content: "BODY-B".into(), - token_count: 22, - entities: vec!["entity:bob".into()], - topics: vec!["#b".into()], - time_range_start: ts, - time_range_end: ts, - score: 0.22, - sealed_at: ts, - deleted: false, - embedding: None, - doc_id: None, - version_ms: None, - }; - - // Stage bodies to disk + record content pointers so - // `read_summary_body` (called per-id inside `hydrate_summary_inputs`) - // can resolve the path. Mirrors the production seal write path. - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).expect("create content_root"); - let stage = |n: &SummaryNode| { - stage_summary( - &content_root, - &SummaryComposeInput { - summary_id: &n.id, - tree_kind: SummaryTreeKind::Source, - tree_id: &tree.id, - tree_scope: &tree.scope, - level: n.level, - child_ids: &n.child_ids, - child_basenames: None, - child_count: n.child_ids.len(), - time_range_start: n.time_range_start, - time_range_end: n.time_range_end, - sealed_at: n.sealed_at, - body: &n.content, - }, - "slack-eng", - ) - .unwrap() - }; - let staged_a = stage(&sum_a); - let staged_b = stage(&sum_b); - crate::openhuman::memory_store::chunks::store::with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &sum_a, Some(&staged_a), "test")?; - insert_summary_tx(&tx, &sum_b, Some(&staged_b), "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - // Interleaved order with a ghost id in the middle: if the function - // ever regresses to indexing by position into the batch result or - // returns the HashMap's iteration order, this assertion will catch - // it. `sum-b` deliberately comes first so a naive "iterate the map" - // implementation that happens to land on `sum-a` first would fail. - let ids = vec![ - "sum-b".to_string(), - "ghost:no-such".to_string(), - "sum-a".to_string(), - ]; - let out = hydrate_summary_inputs(&cfg, &ids).unwrap(); - - // Missing id silently skipped → 2 inputs, not 3. - assert_eq!(out.len(), 2, "ghost id must be skipped, not error"); - // Input order preserved across the gap. - assert_eq!(out[0].id, "sum-b"); - assert_eq!(out[1].id, "sum-a"); - // Per-id field propagation: each input's score/token_count/content - // comes from its own row, not from its sibling. - assert_eq!(out[0].token_count, 22); - assert!((out[0].score - 0.22).abs() < 1e-6); - assert_eq!(out[0].content, "BODY-B"); - assert_eq!(out[1].token_count, 11); - assert!((out[1].score - 0.11).abs() < 1e-6); - assert_eq!(out[1].content, "BODY-A"); - // Entities and topics tracked per id too. - assert_eq!(out[0].entities, vec!["entity:bob".to_string()]); - assert_eq!(out[1].entities, vec!["entity:alice".to_string()]); -} - -/// Regression for Sentry #13021 — when the LLM summary collapses to an -/// empty/whitespace string (e.g. the model returns just newlines, which -/// `summarise()` trims to ""), the seal MUST NOT persist a blank parent. -/// -/// Pre-fix, `truncate_for_embed("", 1000)` produced an empty string that -/// got forwarded to the embedding provider. Real providers (cloud / -/// OpenAI-compatible) 400 on that — `"input must be a non-empty string or -/// array of non-empty strings"` — and the failure was captured as a -/// recurring Sentry server fault even though the defect was on the client -/// side. -/// -/// Initial fix (provider guard) short-circuited the embed call but still -/// persisted `content = ""`, losing the child text from the next rollup / -/// retrieval layer. Final fix (this test): when `summarise()` returns -/// blank, fall back to `fallback_summary` — the deterministic concatenation -/// of the inputs — so the parent has recoverable text and the embedding -/// runs on real content. -#[tokio::test] -async fn whitespace_llm_summary_falls_back_to_deterministic_content() { - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use chrono::TimeZone; - - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - // The chat provider returns ONLY whitespace; `summarise()` trims to "" - // and `clamp_to_budget` keeps it empty → `output.content = ""`. - let provider: Arc = Arc::new(StaticChatProvider::new(" \n\t \n")); - - let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let mk_chunk = |seq: u32, tokens: u32| Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "test-content"), - content: format!("non-empty leaf content {seq}"), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: ts, - time_range: (ts, ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: tokens, - seq_in_source: seq, - created_at: ts, - partial_message: false, - }; - let per_leaf = INPUT_TOKEN_BUDGET * 6 / 10; - let c1 = mk_chunk(0, per_leaf); - let c2 = mk_chunk(1, per_leaf); - upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]); - - let leaf1 = LeafRef { - chunk_id: c1.id.clone(), - token_count: per_leaf, - timestamp: ts, - content: c1.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - let leaf2 = LeafRef { - chunk_id: c2.id.clone(), - token_count: per_leaf, - timestamp: ts, - content: c2.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - - let _ = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf1, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - let sealed = test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf2, &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - - assert_eq!(sealed.len(), 1, "second append crosses budget — one seal"); - let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); - - // Content must be the deterministic fallback derived from the inputs, - // not the empty LLM output. `fallback_summary` joins each non-whitespace - // input with a `"— "` provenance prefix. - assert!( - !summary.content.trim().is_empty(), - "expected fallback content when LLM returned blank (#13021), got empty" - ); - assert!( - summary.content.contains("non-empty leaf content 0"), - "fallback must include leaf 0 content; got: {:?}", - summary.content - ); - assert!( - summary.content.contains("non-empty leaf content 1"), - "fallback must include leaf 1 content; got: {:?}", - summary.content - ); - - // Because the persisted content is now non-empty, the embed step runs. - // With `embeddings_provider = "none"` the test wires `InertEmbedder`, - // which returns a zero vector — its presence (not its value) is the - // signal that the new fallback path drove a real embed call. - // - // Read from the per-model sidecar (`mem_tree_summary_embeddings`); the - // legacy `mem_tree_summaries.embedding` column on `SummaryNode` is - // always written as `None` post-#1574 cutover, so checking - // `summary.embedding` alone would silently always pass. - let sidecar_embedding = - crate::openhuman::memory_store::trees::store::get_summary_embedding(&cfg, &sealed[0]) - .unwrap(); - assert!( - sidecar_embedding.is_some(), - "expected sidecar embedding for the fallback-filled summary (#13021)" - ); -} diff --git a/src/openhuman/memory_tree/tree/factory.rs b/src/openhuman/memory_tree/tree/factory.rs index 92ef55227..e0ba2cc18 100644 --- a/src/openhuman/memory_tree/tree/factory.rs +++ b/src/openhuman/memory_tree/tree/factory.rs @@ -21,68 +21,49 @@ use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf, LabelStrateg use crate::openhuman::memory_tree::tree::flush::force_flush_tree; use crate::openhuman::memory_tree::tree::registry::get_or_create_tree; -/// Literal scope used for the singleton global tree. -pub const GLOBAL_SCOPE: &str = "global"; - -/// High-level tree profile. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TreeProfile { - Source, - Topic, - Global, -} +pub use tinycortex::memory::tree::{TreeProfile, GLOBAL_SCOPE}; /// Factory/config object for one tree instance. #[derive(Debug, Clone)] pub struct TreeFactory<'a> { - profile: TreeProfile, - scope: Cow<'a, str>, + inner: tinycortex::memory::tree::TreeFactory<'a>, } impl<'a> TreeFactory<'a> { pub fn source(scope: impl Into>) -> Self { Self { - profile: TreeProfile::Source, - scope: scope.into(), + inner: tinycortex::memory::tree::TreeFactory::source(scope), } } pub fn topic(scope: impl Into>) -> Self { Self { - profile: TreeProfile::Topic, - scope: scope.into(), + inner: tinycortex::memory::tree::TreeFactory::topic(scope), } } pub fn global() -> Self { Self { - profile: TreeProfile::Global, - scope: Cow::Borrowed(GLOBAL_SCOPE), + inner: tinycortex::memory::tree::TreeFactory::global(), } } pub fn from_tree(tree: &'a Tree) -> Self { - match tree.kind { - TreeKind::Source => Self::source(tree.scope.as_str()), - TreeKind::Topic => Self::topic(tree.scope.as_str()), - TreeKind::Global => Self::global(), + Self { + inner: tinycortex::memory::tree::TreeFactory::from_tree(tree), } } pub fn profile(&self) -> TreeProfile { - self.profile + self.inner.profile() } pub fn kind(&self) -> TreeKind { - match self.profile { - TreeProfile::Source => TreeKind::Source, - TreeProfile::Topic => TreeKind::Topic, - TreeProfile::Global => TreeKind::Global, - } + self.inner.kind() } pub fn scope(&self) -> &str { - self.scope.as_ref() + self.inner.scope() } pub fn summary_tree_kind(&self) -> SummaryTreeKind { @@ -90,6 +71,7 @@ impl<'a> TreeFactory<'a> { TreeKind::Source => SummaryTreeKind::Source, TreeKind::Topic => SummaryTreeKind::Topic, TreeKind::Global => SummaryTreeKind::Global, + _ => SummaryTreeKind::Source, } } @@ -104,6 +86,7 @@ impl<'a> TreeFactory<'a> { slugify_source_id(scope) } } + _ => slugify_source_id(scope), } } @@ -111,6 +94,7 @@ impl<'a> TreeFactory<'a> { match self.kind() { TreeKind::Source => LabelStrategy::ExtractFromContent(build_summary_extractor(config)), TreeKind::Topic | TreeKind::Global => LabelStrategy::Empty, + _ => LabelStrategy::ExtractFromContent(build_summary_extractor(config)), } } diff --git a/src/openhuman/memory_tree/tree/flush.rs b/src/openhuman/memory_tree/tree/flush.rs index ac56c5aad..22669033d 100644 --- a/src/openhuman/memory_tree/tree/flush.rs +++ b/src/openhuman/memory_tree/tree/flush.rs @@ -1,15 +1,4 @@ -//! Time-based buffer flush for source trees (#709). -//! -//! The bucket-seal path only fires when a buffer crosses -//! `INPUT_TOKEN_BUDGET` (token volume) or `SUMMARY_FANOUT` (item count -//! — the L0 fallback gate). Low-volume sources (e.g. an email account -//! with two threads a week) can still park a buffer below both -//! thresholds indefinitely, which hurts recall. -//! `flush_stale_buffers` force-seals any buffer whose `oldest_at` is -//! older than `max_age`, regardless of token count or item count. -//! -//! This is meant to run on a cadence (e.g. daily cron). Phase 3a ships -//! the primitive; wiring into a scheduler is not required for merge. +//! Product adapters for tinycortex-owned stale-buffer flushing. use anyhow::Result; use chrono::{DateTime, Duration, Utc}; @@ -17,68 +6,15 @@ use chrono::{DateTime, Duration, Utc}; use crate::openhuman::config::Config; use crate::openhuman::memory_store::trees::types::DEFAULT_FLUSH_AGE_SECS; use crate::openhuman::memory_tree::tree::bucket_seal::{cascade_all_from, LabelStrategy}; -use crate::openhuman::memory_tree::tree::store; -/// Seal every buffer whose oldest item is older than `max_age`. Returns -/// the number of individual seal calls (not trees) that fired. When the -/// same tree has multiple stale levels they're each sealed in order. pub async fn flush_stale_buffers( config: &Config, max_age: Duration, strategy: &LabelStrategy, ) -> Result { - use crate::openhuman::memory_store::trees::store::get_trees_batch; - - let now = Utc::now(); - let cutoff = now - max_age; - let stale = store::list_stale_buffers(config, cutoff)?; - log::info!( - "[tree::flush] found {} stale buffers (max_age={:?})", - stale.len(), - max_age - ); - - // One batched `SELECT … WHERE id IN (?,?,…)` over the distinct - // tree_ids instead of N per-buffer `get_tree` round-trips. `stale` - // is filtered to L0 by `list_stale_buffers` so the same tree_id - // typically appears at most once, but we still de-dup defensively - // before hitting the DB — future widenings of the filter (e.g. - // dropping the `level = 0` clause) would otherwise re-pay the cost. - // Missing rows stay silently absent from the map; the - // `Some(t) => ... / None => warn-and-skip` orphan-buffer path below - // preserves the per-id `get_tree` `Ok(None)` semantics. - // Preserve list_stale_buffers order so the HashMap lookup below stays - // predictable; dedup is purely to avoid redundant DB params (not for - // semantic ordering — the per-buffer loop walks `stale` directly). - let distinct_tree_ids: Vec = { - let mut seen = std::collections::HashSet::new(); - let mut out = Vec::new(); - for buf in &stale { - if seen.insert(buf.tree_id.clone()) { - out.push(buf.tree_id.clone()); - } - } - out - }; - let tree_by_id = get_trees_batch(config, &distinct_tree_ids)?; - - let mut seals: usize = 0; - for buf in stale { - let Some(tree) = tree_by_id.get(&buf.tree_id) else { - log::warn!( - "[tree::flush] orphan buffer tree_id={} level={}", - buf.tree_id, - buf.level - ); - continue; - }; - let sealed = cascade_all_from(config, tree, buf.level, Some(now), strategy).await?; - seals += sealed.len(); - } - Ok(seals) + crate::openhuman::tinycortex::flush_stale_tree_buffers(config, max_age, strategy).await } -/// Convenience wrapper that uses [`DEFAULT_FLUSH_AGE_SECS`]. pub async fn flush_stale_buffers_default( config: &Config, strategy: &LabelStrategy, @@ -86,311 +22,13 @@ pub async fn flush_stale_buffers_default( flush_stale_buffers(config, Duration::seconds(DEFAULT_FLUSH_AGE_SECS), strategy).await } -/// Helper exposed for callers that want a single explicit force-seal (e.g. -/// "user disconnected this account, flush its buffer now"). pub async fn force_flush_tree( config: &Config, tree_id: &str, now: Option>, strategy: &LabelStrategy, ) -> Result> { - let tree = store::get_tree(config, tree_id)? + let tree = crate::openhuman::memory_store::trees::store::get_tree(config, tree_id)? .ok_or_else(|| anyhow::anyhow!("no tree with id {tree_id}"))?; - cascade_all_from(config, &tree, 0, now, strategy).await -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider}; - use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree; - use crate::openhuman::memory_store::chunks::store::upsert_chunks; - use crate::openhuman::memory_store::chunks::types::{ - chunk_id, Chunk, Metadata, SourceKind, SourceRef, - }; - use crate::openhuman::memory_store::content as content_store; - use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf, LeafRef}; - use std::sync::Arc; - use tempfile::TempDir; - - fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) { - let content_root = cfg.memory_tree_content_root(); - std::fs::create_dir_all(&content_root).expect("create content_root for test"); - let staged = content_store::stage_chunks(&content_root, chunks) - .expect("stage_chunks for test chunks"); - crate::openhuman::memory_store::chunks::store::with_connection(cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?; - tx.commit()?; - Ok(()) - }) - .expect("persist staged chunk pointers"); - } - - fn test_config() -> (TempDir, Config) { - let tmp = TempDir::new().unwrap(); - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): flush triggers seals which embed — force inert. - cfg.memory_tree.embedding_endpoint = None; - cfg.memory_tree.embedding_model = None; - cfg.memory_tree.embedding_strict = false; - (tmp, cfg) - } - - #[tokio::test] - async fn flush_seals_old_buffer_even_under_budget() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - - // Persist one chunk with an old timestamp (10 days ago). - let old_ts = Utc::now() - Duration::days(10); - let c = Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "test-content"), - content: "old content that should get sealed".into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: old_ts, - time_range: (old_ts, old_ts), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: 100, - seq_in_source: 0, - created_at: old_ts, - partial_message: false, - }; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c.clone()]); - - let leaf = LeafRef { - chunk_id: c.id.clone(), - token_count: 100, // way under the 10k budget - timestamp: old_ts, - content: c.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap(); - }) - .await; - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); - - let seals = test_override::with_provider(Arc::clone(&provider), async { - flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(seals, 1); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 1); - - let l0 = store::get_buffer(&cfg, &tree.id, 0).unwrap(); - assert!(l0.is_empty()); - } - - #[tokio::test] - async fn flush_does_not_force_seal_under_fanout_upper_buffer() { - // Regression: previously `list_stale_buffers` returned every level, - // and `cascade_all_from` force-sealed the first iteration regardless - // of `should_seal`. A stale L1 buffer with one child would seal into - // a degenerate L2 summary that wraps a single L1 — repeating across - // flush cycles produced the L7→L13 1:1:1 chain in real workspaces. - // Flush must restrict force-seals to L0 and let upper levels gate - // on `SUMMARY_FANOUT` naturally. - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - - // Plant a stale L1 buffer holding a single (synthetic) child id. - // No L0 chunks — the only thing flush could touch is the L1 buffer. - let old_ts = Utc::now() - Duration::days(10); - crate::openhuman::memory_store::chunks::store::with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - crate::openhuman::memory_store::trees::store::upsert_buffer_tx( - &tx, - &crate::openhuman::memory_store::trees::types::Buffer { - tree_id: tree.id.clone(), - level: 1, - item_ids: vec!["fake-l1-child".into()], - token_sum: 100, - oldest_at: Some(old_ts), - }, - )?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - - let seals = test_override::with_provider(provider, async { - flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(seals, 0, "L1 stale buffer must not be force-sealed"); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); - - // The L1 buffer must still be intact — flush cannot touch it. - let l1 = store::get_buffer(&cfg, &tree.id, 1).unwrap(); - assert_eq!(l1.item_ids, vec!["fake-l1-child".to_string()]); - } - - #[tokio::test] - async fn flush_noop_when_buffer_is_recent() { - let (_tmp, cfg) = test_config(); - let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - - // Persist a leaf stamped now so it's NOT stale. - let now = Utc::now(); - let c = Chunk { - id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "test-content"), - content: "fresh".into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - owner: "alice".into(), - timestamp: now, - time_range: (now, now), - tags: vec![], - source_ref: Some(SourceRef::new("slack://x")), - path_scope: None, - }, - token_count: 50, - seq_in_source: 0, - created_at: now, - partial_message: false, - }; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - let leaf = LeafRef { - chunk_id: c.id.clone(), - token_count: 50, - timestamp: now, - content: c.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, &tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap(); - }) - .await; - - let seals = test_override::with_provider(provider, async { - flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - assert_eq!(seals, 0); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); - } - - /// Regression for the `get_trees_batch` refactor: two distinct - /// trees each carry a stale L0 buffer; `flush_stale_buffers` must - /// seal both via a single batched tree-fetch. Pre-refactor this - /// path issued N per-buffer `get_tree(...)` round-trips; post- - /// refactor it's one `SELECT … WHERE id IN (?,?)` and a per-id - /// HashMap lookup. The interesting bit is the HashMap lookup - /// keying by tree_id (not by `enumerate()` position over the input - /// slice) — with two trees, swapping the lookup key would leak one - /// tree's `Tree` into the other tree's seal and either error out - /// or write summaries against the wrong `root_id`. - #[tokio::test] - async fn flush_seals_multiple_distinct_trees_via_batched_lookup() { - let (_tmp, cfg) = test_config(); - let tree_a = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); - let tree_b = get_or_create_source_tree(&cfg, "slack:#design").unwrap(); - assert_ne!( - tree_a.id, tree_b.id, - "test prerequisite: distinct source_ids must yield distinct tree_ids" - ); - let provider: Arc = - Arc::new(StaticChatProvider::new("test summary content")); - - let old_ts = Utc::now() - Duration::days(10); - - // Plant a stale L0 leaf in each tree, backed by real chunks so - // the cascade can seal end-to-end. - for (src_id, content) in [ - ("slack:#eng", "engineering content"), - ("slack:#design", "design content"), - ] { - let c = Chunk { - id: chunk_id(SourceKind::Chat, src_id, 0, content), - content: content.into(), - metadata: Metadata { - source_kind: SourceKind::Chat, - source_id: src_id.into(), - owner: "alice".into(), - timestamp: old_ts, - time_range: (old_ts, old_ts), - tags: vec![], - source_ref: Some(SourceRef::new(&format!("slack://{src_id}"))), - path_scope: None, - }, - token_count: 100, - seq_in_source: 0, - created_at: old_ts, - partial_message: false, - }; - upsert_chunks(&cfg, &[c.clone()]).unwrap(); - stage_test_chunks(&cfg, &[c.clone()]); - let leaf = LeafRef { - chunk_id: c.id.clone(), - token_count: 100, - timestamp: old_ts, - content: c.content.clone(), - entities: vec![], - topics: vec![], - score: 0.5, - }; - let target_tree = if src_id == "slack:#eng" { - &tree_a - } else { - &tree_b - }; - test_override::with_provider(Arc::clone(&provider), async { - append_leaf(&cfg, target_tree, &leaf, &LabelStrategy::Empty) - .await - .unwrap(); - }) - .await; - } - - let seals = test_override::with_provider(provider, async { - flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty) - .await - .unwrap() - }) - .await; - - // Both trees sealed via a single batched fetch — the HashMap - // lookup correctly routed each buffer to its own `Tree`. - assert_eq!(seals, 2, "both stale trees must seal in one flush pass"); - assert_eq!( - store::count_summaries(&cfg, &tree_a.id).unwrap(), - 1, - "tree_a must own its summary (HashMap keyed by id, not position)" - ); - assert_eq!( - store::count_summaries(&cfg, &tree_b.id).unwrap(), - 1, - "tree_b must own its summary (HashMap keyed by id, not position)" - ); - } + cascade_all_from(config, &tree, 0, now.or_else(|| Some(Utc::now())), strategy).await } diff --git a/src/openhuman/memory_tree/tree_runtime/engine.rs b/src/openhuman/memory_tree/tree_runtime/engine.rs index fc1998c6f..e70d4054e 100644 --- a/src/openhuman/memory_tree/tree_runtime/engine.rs +++ b/src/openhuman/memory_tree/tree_runtime/engine.rs @@ -1,668 +1,157 @@ -//! Core summarization engine: ingest raw data, summarize into hour leaves, -//! and propagate summaries upward through the tree. +//! Product adapters around the tinycortex markdown time-tree engine. + +use std::sync::Arc; use anyhow::{Context, Result}; +use async_trait::async_trait; use chrono::{DateTime, Timelike, Utc}; -use std::collections::BTreeMap; +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ChatModel, ModelRequest}; +use tinycortex::memory::tree::runtime::{ + NodeLevel, RuntimeObserver, Summariser, TreeNode, TreeStatus, +}; use crate::core::event_bus::{publish_global, DomainEvent}; use crate::openhuman::config::Config; -use crate::openhuman::memory_tree::tree_runtime::store; -use crate::openhuman::memory_tree::tree_runtime::types::{ - derive_node_ids, derive_parent_id, estimate_tokens, level_from_node_id, NodeLevel, TreeNode, - TreeStatus, -}; -use std::sync::Arc; -use tinyagents::harness::message::Message; -use tinyagents::harness::model::{ChatModel, ModelRequest}; const SUMMARIZATION_TEMP: f64 = 0.3; -/// Maximum characters for a summary response (hard limit enforced after LLM call). -/// Set to 4x the Root token budget as a generous upper bound. -const MAX_SUMMARY_CHARS: usize = 20_000 * 4; +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} -// ── Public API ───────────────────────────────────────────────────────── +struct ChatSummariser<'a>(&'a dyn ChatModel<()>); + +#[async_trait] +impl Summariser for ChatSummariser<'_> { + async fn summarise(&self, system: Option<&str>, content: &str) -> Result { + log::debug!( + "[tree_summarizer] provider call content_chars={} has_system={}", + content.len(), + system.is_some() + ); + let mut messages = Vec::with_capacity(2); + if let Some(system) = system { + messages.push(Message::system(system.to_string())); + } + messages.push(Message::user(content.to_string())); + let response = self + .0 + .invoke( + &(), + ModelRequest::new(messages).with_temperature(SUMMARIZATION_TEMP), + ) + .await + .context("time-tree summarization provider call failed")? + .text(); + log::debug!( + "[tree_summarizer] provider call complete response_chars={}", + response.len() + ); + Ok(response) + } +} + +struct EventObserver; + +impl RuntimeObserver for EventObserver { + fn hour_completed(&self, namespace: &str, node_id: &str, token_count: u32) { + publish_global(DomainEvent::TreeSummarizerHourCompleted { + namespace: namespace.to_string(), + node_id: node_id.to_string(), + token_count, + }); + } + + fn node_propagated(&self, namespace: &str, node_id: &str, level: NodeLevel, token_count: u32) { + publish_global(DomainEvent::TreeSummarizerPropagated { + namespace: namespace.to_string(), + node_id: node_id.to_string(), + level: level.as_str().to_string(), + token_count, + }); + } + + fn rebuild_completed(&self, namespace: &str, total_nodes: u64) { + publish_global(DomainEvent::TreeSummarizerRebuildCompleted { + namespace: namespace.to_string(), + total_nodes, + }); + } +} -/// Run the summarization job for a given namespace. -/// -/// 1. Drains the ingestion buffer. -/// 2. Groups buffered entries by their original hour (from filename timestamps). -/// 3. Summarizes each hour group into its own hour leaf. -/// 4. Propagates summaries upward through day → month → year → root. -/// -/// Returns the last hour leaf node created, or `None` if the buffer was empty. pub async fn run_summarization( config: &Config, provider: &dyn ChatModel<()>, namespace: &str, - _ts: DateTime, + ts: DateTime, ) -> Result> { - // Read buffer entries non-destructively; we only delete after durable writes. - let buffered = store::buffer_read(config, namespace)?; - if buffered.is_empty() { - tracing::debug!("[tree_summarizer] no buffered data for namespace '{namespace}', skipping"); - return Ok(None); - } - - let buffer_filenames: Vec = buffered.iter().map(|(name, _)| name.clone()).collect(); - - tracing::debug!( - "[tree_summarizer] starting summarization for namespace '{}' with {} buffer entries", + log::debug!("[tree_summarizer] tinycortex run namespace={namespace}"); + let result = tinycortex::memory::tree::runtime::run_summarization_observed( + &engine_config(config), + &ChatSummariser(provider), namespace, - buffered.len() + ts, + &EventObserver, + ) + .await; + log::debug!( + "[tree_summarizer] tinycortex run complete namespace={} success={}", + namespace, + result.is_ok() ); - - // Group buffered entries by hour using their buffer filename timestamps. - let hour_groups = group_by_hour(&buffered); - - tracing::debug!( - "[tree_summarizer] grouped into {} distinct hours", - hour_groups.len() - ); - - // Track all ancestor IDs to propagate after all hour leaves are written. - let mut all_propagation_ids: Vec<(String, NodeLevel)> = Vec::new(); - let mut last_hour_node: Option = None; - - for (hour_id, entries) in &hour_groups { - let combined = entries.join("\n\n---\n\n"); - - // Check for an existing hour node and merge content if present - let (existing_summary, existing_created_at) = - match store::read_node(config, namespace, hour_id)? { - Some(existing) => (Some(existing.summary), Some(existing.created_at)), - None => (None, None), - }; - - let to_summarize = if let Some(ref prev) = existing_summary { - format!("{prev}\n\n---\n\n{combined}") - } else { - combined - }; - - let hour_summary = summarize_to_limit( - provider, - &to_summarize, - NodeLevel::Hour.max_tokens(), - "hour", - hour_id, - ) - .await - .context("summarize hour leaf")?; - - let now = Utc::now(); - let hour_node = TreeNode { - node_id: hour_id.clone(), - namespace: namespace.to_string(), - level: NodeLevel::Hour, - parent_id: derive_parent_id(hour_id), - summary: hour_summary.clone(), - token_count: estimate_tokens(&hour_summary), - child_count: 0, - created_at: existing_created_at.unwrap_or(now), - updated_at: now, - metadata: None, - }; - store::write_node(config, &hour_node)?; - - publish_global(DomainEvent::TreeSummarizerHourCompleted { - namespace: namespace.to_string(), - node_id: hour_id.clone(), - token_count: hour_node.token_count, - }); - - tracing::debug!( - "[tree_summarizer] hour leaf {} created ({} tokens)", - hour_id, - hour_node.token_count - ); - - // Derive propagation path for this hour - 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)); - all_propagation_ids.push((month_id, NodeLevel::Month)); - all_propagation_ids.push((year_id, NodeLevel::Year)); - all_propagation_ids.push((root_id, NodeLevel::Root)); - - last_hour_node = Some(hour_node); - } - - // Deduplicate and propagate in bottom-up order (days, months, years, root). - // - // #002 (FR-008): partial success. A single node's summarization failure - // (e.g. the LLM times out on one busy day) must NOT void the entire run — - // the hour leaves are already durably written, and the other ancestors - // should still propagate. We collect per-node failures and continue - // instead of `?`-aborting on the first one; the run still succeeds, the - // failures are logged with their node ids for diagnosis. - let mut seen = std::collections::HashSet::new(); - let mut failed: Vec = Vec::new(); - let mut propagated: u32 = 0; - for level in [ - NodeLevel::Day, - NodeLevel::Month, - NodeLevel::Year, - NodeLevel::Root, - ] { - for (node_id, node_level) in &all_propagation_ids { - if *node_level == level && seen.insert(node_id.clone()) { - match propagate_node(config, provider, namespace, node_id, level).await { - Ok(()) => propagated += 1, - Err(e) => { - log::warn!( - "[tree_summarizer] propagate failed (continuing) namespace='{namespace}' \ - node={node_id} level={}: {e:#}", - level.as_str() - ); - failed.push(node_id.clone()); - } - } - } - } - } - if !failed.is_empty() { - log::warn!( - "[tree_summarizer] partial summarization for '{namespace}': {propagated} node(s) \ - propagated, {} failed ({:?})", - failed.len(), - failed - ); - } - - // Only clear the buffer when propagation was fully successful. If any nodes - // failed, keep the buffer entries so `run_hourly_loop` re-discovers this - // namespace on the next pass and retries the failed levels — otherwise a - // transient day/month/year/root failure becomes sticky degradation. - if failed.is_empty() { - store::buffer_delete(config, namespace, &buffer_filenames) - .context("delete buffer entries after successful summarization")?; - } else { - log::info!( - "[tree_summarizer] keeping buffer for '{namespace}' — {n} failed node(s) \ - will be retried on the next run", - n = failed.len() - ); - } - - Ok(last_hour_node) + result } -/// Rebuild the entire tree from hour leaves upward. -/// Deletes all non-leaf nodes and re-summarizes. -/// Preserves buffered data that hasn't been summarized yet. pub async fn rebuild_tree( config: &Config, provider: &dyn ChatModel<()>, namespace: &str, ) -> Result { - tracing::debug!("[tree_summarizer] rebuilding tree for namespace '{namespace}'"); - - let status = store::get_tree_status(config, namespace)?; - if status.total_nodes == 0 { - return Ok(status); - } - - // Collect all hour leaves first - let base = store::tree_dir(config, namespace); - let mut hour_leaves: Vec = Vec::new(); - collect_hour_leaves_recursive(&base, namespace, "", &mut hour_leaves)?; - - if hour_leaves.is_empty() { - tracing::debug!("[tree_summarizer] no hour leaves found, nothing to rebuild"); - return store::get_tree_status(config, namespace); - } - - // Preserve the buffer directory by moving it to a sibling path *outside* - // the tree directory, so delete_tree() does not destroy it. - let buffer_path = store::buffer_dir(config, namespace); - let tree_base = store::tree_dir(config, namespace); - // Place backup next to the tree dir (e.g. .../tree_buffer_backup) - let buffer_backup = tree_base - .parent() - .unwrap_or(&tree_base) - .join("tree_buffer_backup"); - let buffer_existed = buffer_path.exists(); - if buffer_existed { - if buffer_backup.exists() { - std::fs::remove_dir_all(&buffer_backup)?; - } - std::fs::rename(&buffer_path, &buffer_backup).context("backup buffer before rebuild")?; - tracing::debug!("[tree_summarizer] backed up buffer directory outside tree"); - } - - // Delete and recreate the tree directory - store::delete_tree(config, namespace)?; - - // Restore the buffer directory back inside the tree - if buffer_existed && buffer_backup.exists() { - let restored_buffer = store::buffer_dir(config, namespace); - if let Some(parent) = restored_buffer.parent() { - std::fs::create_dir_all(parent)?; - } - std::fs::rename(&buffer_backup, &restored_buffer) - .context("restore buffer after rebuild")?; - tracing::debug!("[tree_summarizer] restored buffer directory"); - } - - // Re-write all hour leaves - for leaf in &hour_leaves { - store::write_node(config, leaf)?; - } - - // Collect unique ancestor IDs at each level, ordered bottom-up - let mut day_ids = std::collections::BTreeSet::new(); - let mut month_ids = std::collections::BTreeSet::new(); - let mut year_ids = std::collections::BTreeSet::new(); - - for leaf in &hour_leaves { - if let Some(day) = derive_parent_id(&leaf.node_id) { - day_ids.insert(day.clone()); - if let Some(month) = derive_parent_id(&day) { - month_ids.insert(month.clone()); - if let Some(year) = derive_parent_id(&month) { - year_ids.insert(year); - } - } - } - } - - // Propagate bottom-up: days, then months, then years, then root. - // - // #002 (FR-008): partial success — a single node's failure must not abort - // the whole rebuild. Collect-and-continue (the hour leaves are already - // re-written above), logging each failure for diagnosis; the rebuild still - // returns the resulting status rather than erroring out wholesale. - let mut failed: Vec = Vec::new(); - let propagate = |id: &str, level: NodeLevel| { - let node_id = id.to_string(); - async move { - if let Err(e) = propagate_node(config, provider, namespace, &node_id, level).await { - log::warn!( - "[tree_summarizer] rebuild propagate failed (continuing) namespace='{namespace}' \ - node={node_id} level={}: {e:#}", - level.as_str() - ); - Some(node_id) - } else { - None - } - } - }; - for day_id in &day_ids { - if let Some(f) = propagate(day_id, NodeLevel::Day).await { - failed.push(f); - } - } - for month_id in &month_ids { - if let Some(f) = propagate(month_id, NodeLevel::Month).await { - failed.push(f); - } - } - for year_id in &year_ids { - if let Some(f) = propagate(year_id, NodeLevel::Year).await { - failed.push(f); - } - } - if let Some(f) = propagate("root", NodeLevel::Root).await { - failed.push(f); - } - if !failed.is_empty() { - log::warn!( - "[tree_summarizer] partial rebuild for '{namespace}': {} node(s) failed ({:?})", - failed.len(), - failed - ); - } - - let final_status = store::get_tree_status(config, namespace)?; - - publish_global(DomainEvent::TreeSummarizerRebuildCompleted { - namespace: namespace.to_string(), - total_nodes: final_status.total_nodes, - }); - - tracing::debug!( - "[tree_summarizer] rebuild complete for '{}': {} nodes", + log::debug!("[tree_summarizer] tinycortex rebuild namespace={namespace}"); + tinycortex::memory::tree::runtime::rebuild_tree_observed( + &engine_config(config), + &ChatSummariser(provider), namespace, - final_status.total_nodes - ); - Ok(final_status) + &EventObserver, + ) + .await } -// ── Internal ─────────────────────────────────────────────────────────── - -/// Re-summarize a single non-leaf node from its children. -pub(crate) async fn propagate_node( - config: &Config, - provider: &dyn ChatModel<()>, - namespace: &str, - node_id: &str, - level: NodeLevel, -) -> Result<()> { - let children = store::read_children(config, namespace, node_id)?; - if children.is_empty() { - tracing::debug!( - "[tree_summarizer] node {} has no children, skipping propagation", - node_id - ); - return Ok(()); - } - - let child_count = children.len() as u32; - let combined: String = children - .iter() - .map(|c| format!("## {} ({})\n\n{}", c.node_id, c.level.as_str(), c.summary)) - .collect::>() - .join("\n\n---\n\n"); - - let combined_tokens = estimate_tokens(&combined); - let max_tokens = level.max_tokens(); - - let summary = if combined_tokens <= max_tokens { - // Fits within budget — use the combined text directly - tracing::debug!( - "[tree_summarizer] node {} combined children ({} tokens) fits within {} token budget, no LLM needed", - node_id, - combined_tokens, - max_tokens - ); - combined - } else { - // Exceeds budget — summarize with LLM - tracing::debug!( - "[tree_summarizer] node {} combined children ({} tokens) exceeds {} token budget, summarizing", - node_id, - combined_tokens, - max_tokens - ); - summarize_to_limit(provider, &combined, max_tokens, level.as_str(), node_id).await? - }; - - let now = Utc::now(); - let existing = store::read_node(config, namespace, node_id)?; - let created_at = existing.map(|n| n.created_at).unwrap_or(now); - - let node = TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level, - parent_id: derive_parent_id(node_id), - summary: summary.clone(), - token_count: estimate_tokens(&summary), - child_count, - created_at, - updated_at: now, - metadata: None, - }; - store::write_node(config, &node)?; - - publish_global(DomainEvent::TreeSummarizerPropagated { - namespace: namespace.to_string(), - node_id: node_id.to_string(), - level: level.as_str().to_string(), - token_count: node.token_count, - }); - - tracing::debug!( - "[tree_summarizer] propagated node {} (level={}, tokens={}, children={})", - node_id, - level.as_str(), - node.token_count, - child_count - ); - Ok(()) -} - -/// Summarize text to fit within a token limit using the LLM provider. -/// Enforces a hard character limit on the response to prevent runaway output. -async fn summarize_to_limit( - provider: &dyn ChatModel<()>, - content: &str, - max_tokens: u32, - level_name: &str, - node_id: &str, -) -> Result { - let max_chars = (max_tokens as usize) * 4; - let system_prompt = format!( - "You are a hierarchical summarizer. Compress the following content into a concise \ - summary that preserves the most important information.\n\n\ - Rules:\n\ - - The summary MUST be under {max_tokens} tokens (roughly {max_chars} characters).\n\ - - Focus on key events, decisions, facts, patterns, and actionable insights.\n\ - - Preserve names, dates, numbers, and specific details when important.\n\ - - Use clear, dense prose — no filler.\n\n\ - Context: You are summarizing at the {level_name} level for node '{node_id}'.", - ); - - let response = provider - .invoke( - &(), - ModelRequest::new(vec![ - Message::system(system_prompt), - Message::user(content.to_string()), - ]) - .with_temperature(SUMMARIZATION_TEMP), - ) - .await - .with_context(|| { - format!("LLM summarization failed for node {node_id} (level={level_name})") - })? - .text(); - - // Enforce hard character limit on LLM response (use the stricter of the two limits) - let char_limit = max_chars.min(MAX_SUMMARY_CHARS); - let response = if response.len() > char_limit { - tracing::warn!( - "[tree_summarizer] LLM response for node {} (level={}) was {} chars, truncating to {} chars", - node_id, - level_name, - response.len(), - char_limit - ); - // Truncate at a char boundary - let truncated = - &response[..crate::openhuman::util::floor_char_boundary(&response, char_limit)]; - truncated.to_string() - } else { - response - }; - - tracing::debug!( - "[tree_summarizer] LLM summarized {} chars -> {} chars for node {} (level={})", - content.len(), - response.len(), - node_id, - level_name - ); - - Ok(response) -} - -/// Group buffer entries by their hour based on filename timestamps. -/// -/// Buffer filenames are `{timestamp_millis}_{uuid}.md`. We extract the timestamp -/// and derive the hour ID for each entry. -pub(crate) fn group_by_hour(entries: &[(String, String)]) -> BTreeMap> { - let mut groups: BTreeMap> = BTreeMap::new(); - - for (filename, content) in entries { - let hour_id = hour_id_from_buffer_filename(filename).unwrap_or_else(|| { - // Fallback: use current time if filename can't be parsed - let now = Utc::now(); - let (hour, _, _, _, _) = derive_node_ids(&now); - hour - }); - groups.entry(hour_id).or_default().push(content.clone()); - } - - groups -} - -/// Extract the hour node ID from a buffer filename like `1711972800000_abc12345.md`. -fn hour_id_from_buffer_filename(filename: &str) -> Option { - let ts_str = filename.split('_').next()?; - let millis: i64 = ts_str.parse().ok()?; - let dt = DateTime::from_timestamp_millis(millis)?; - let (hour_id, _, _, _, _) = derive_node_ids(&dt); - Some(hour_id) -} - -/// Derive propagation IDs from an hour node_id string like "2024/03/15/14". -fn derive_node_ids_from_hour_id(hour_id: &str) -> (String, String, String, String, String) { - let parts: Vec<&str> = hour_id.split('/').collect(); - if parts.len() == 4 { - let year = parts[0].to_string(); - let month = format!("{}/{}", parts[0], parts[1]); - let day = format!("{}/{}/{}", parts[0], parts[1], parts[2]); - (hour_id.to_string(), day, month, year, "root".to_string()) - } else { - // Fallback - ( - hour_id.to_string(), - "unknown".to_string(), - "unknown".to_string(), - "unknown".to_string(), - "root".to_string(), - ) - } -} - -/// Recursively collect all hour leaf nodes from the tree directory. -fn collect_hour_leaves_recursive( - dir: &std::path::Path, - namespace: &str, - prefix: &str, - leaves: &mut Vec, -) -> Result<()> { - if !dir.exists() { - return Ok(()); - } - for entry in std::fs::read_dir(dir)? { - let entry = entry?; - let name = entry.file_name().to_string_lossy().to_string(); - let ft = entry.file_type()?; - - if ft.is_dir() { - if name == "buffer" || name == "buffer_backup" { - continue; - } - let child_prefix = if prefix.is_empty() { - name.clone() - } else { - format!("{prefix}/{name}") - }; - collect_hour_leaves_recursive(&entry.path(), namespace, &child_prefix, leaves)?; - } else if ft.is_file() && name.ends_with(".md") && name != "summary.md" && name != "root.md" - { - let hour_part = name.trim_end_matches(".md"); - let node_id = if prefix.is_empty() { - hour_part.to_string() - } else { - format!("{prefix}/{hour_part}") - }; - let level = level_from_node_id(&node_id); - if level == NodeLevel::Hour { - let raw = std::fs::read_to_string(entry.path())?; - let node = - crate::openhuman::memory_tree::tree_runtime::store::parse_node_markdown_pub( - &raw, namespace, &node_id, - ) - .with_context(|| format!("failed to parse hour leaf '{node_id}'"))?; - leaves.push(node); - } - } - } - Ok(()) -} - -// ── Hourly background loop ───────────────────────────────────────────── - -/// Start a background task that runs the summarization job every hour. -/// -/// This should be called once at application startup. The task runs -/// indefinitely, sleeping until the next hour boundary. pub async fn run_hourly_loop(config: Config, provider: Arc>) { - tracing::debug!("[tree_summarizer] hourly loop started"); - + log::debug!("[tree_summarizer] hourly loop started"); loop { - // Sleep until the next hour boundary let now = Utc::now(); - let next_hour = { - let base = now - .date_naive() - .and_hms_opt(now.hour(), 0, 0) - .unwrap_or(now.naive_utc()); - let next = base + chrono::Duration::hours(1); - DateTime::::from_naive_utc_and_offset(next, Utc) - }; + let base = now + .date_naive() + .and_hms_opt(now.hour(), 0, 0) + .unwrap_or(now.naive_utc()); + let next_hour = + DateTime::::from_naive_utc_and_offset(base + chrono::Duration::hours(1), Utc); let sleep_duration = (next_hour - now) .to_std() .unwrap_or(std::time::Duration::from_secs(3600)); - - tracing::debug!( - "[tree_summarizer] sleeping {:.0}s until next hour boundary", - sleep_duration.as_secs_f64() + log::debug!( + "[tree_summarizer] sleeping seconds={}", + sleep_duration.as_secs() ); tokio::time::sleep(sleep_duration).await; - // Run summarization for all namespaces that have buffered data let ts = Utc::now(); - let namespaces = discover_active_namespaces(&config); - for ns in &namespaces { - match run_summarization(&config, provider.as_ref(), ns, ts).await { - Ok(Some(node)) => { - tracing::debug!( - "[tree_summarizer] hourly job completed for '{}': node {} ({} tokens)", - ns, - node.node_id, - node.token_count - ); - } - Ok(None) => { - tracing::debug!( - "[tree_summarizer] hourly job skipped for '{}' (no buffered data)", - ns - ); - } - Err(e) => { - tracing::error!("[tree_summarizer] hourly job failed for '{}': {:#}", ns, e); - } + let namespaces = + tinycortex::memory::tree::runtime::discover_active_namespaces(&engine_config(&config)); + log::debug!( + "[tree_summarizer] hourly tick active_namespaces={}", + namespaces.len() + ); + for namespace in namespaces { + if let Err(error) = run_summarization(&config, provider.as_ref(), &namespace, ts).await + { + log::error!( + "[tree_summarizer] hourly run failed namespace={} error={error:#}", + namespace + ); } } } } - -/// Discover namespaces that have pending buffer data by scanning the -/// `memory/namespaces/*/tree/buffer/` directories. -fn discover_active_namespaces(config: &Config) -> Vec { - let namespaces_dir = config.workspace_dir.join("memory").join("namespaces"); - - if !namespaces_dir.exists() { - return vec![]; - } - - let mut active = Vec::new(); - if let Ok(entries) = std::fs::read_dir(&namespaces_dir) { - for entry in entries.flatten() { - let name = entry.file_name().to_string_lossy().to_string(); - let buffer_dir = entry.path().join("tree").join("buffer"); - if buffer_dir.exists() { - // Check if buffer has any .md files - if let Ok(buffer_entries) = std::fs::read_dir(&buffer_dir) { - let has_entries = buffer_entries - .flatten() - .any(|e| e.path().extension().map(|ext| ext == "md").unwrap_or(false)); - if has_entries { - active.push(name); - } - } - } - } - } - active -} - -#[cfg(test)] -#[path = "engine_tests.rs"] -mod tests; diff --git a/src/openhuman/memory_tree/tree_runtime/engine_tests.rs b/src/openhuman/memory_tree/tree_runtime/engine_tests.rs deleted file mode 100644 index 866eb41f6..000000000 --- a/src/openhuman/memory_tree/tree_runtime/engine_tests.rs +++ /dev/null @@ -1,707 +0,0 @@ -use super::*; -use crate::openhuman::config::Config; -use crate::openhuman::inference::provider::traits::Provider; -use async_trait::async_trait; -use chrono::{TimeZone, Utc}; -use tempfile::TempDir; - -// ── Shared helpers ──────────────────────────────────────────────────────── - -fn test_config(tmp: &TempDir) -> Config { - Config { - workspace_dir: tmp.path().join("workspace"), - action_dir: tmp.path().join("workspace"), - config_path: tmp.path().join("config.toml"), - ..Config::default() - } -} - -/// A stub Provider whose `chat_with_system` returns a fixed string. -struct StubProvider { - reply: String, -} - -impl StubProvider { - fn with_reply(reply: impl Into) -> Self { - Self { - reply: reply.into(), - } - } -} - -#[async_trait] -impl Provider for StubProvider { - async fn chat_with_system( - &self, - _system: Option<&str>, - _message: &str, - _model: &str, - _temperature: f64, - ) -> anyhow::Result { - log::debug!("[memory_tree_test] StubProvider::chat_with_system called"); - Ok(self.reply.clone()) - } -} - -/// #002 (FR-008): a Provider that errors only when the system prompt names a -/// specific summarization level (`level_name` is embedded by -/// `summarize_to_limit`), succeeding otherwise. Lets a test force a single -/// propagation node to fail while the rest of the run proceeds. -struct FailAtLevelProvider { - fail_level: &'static str, - reply: String, -} - -#[async_trait] -impl Provider for FailAtLevelProvider { - async fn chat_with_system( - &self, - system: Option<&str>, - _message: &str, - _model: &str, - _temperature: f64, - ) -> anyhow::Result { - let sys = system.unwrap_or(""); - if sys.contains(&format!("at the {} level", self.fail_level)) { - anyhow::bail!("simulated {} summarization failure", self.fail_level); - } - Ok(self.reply.clone()) - } -} - -/// Wrap a stub [`Provider`] as the crate `ChatModel` the engine now consumes. -/// The baked model/temperature are inert for the stubs (they ignore both). -fn as_model( - provider: impl Provider + 'static, -) -> std::sync::Arc> { - crate::openhuman::inference::provider::chat_model_from_provider( - Box::new(provider), - "test-model".to_string(), - 0.3, - ) -} - -// ── group_by_hour ──────────────────────────────────────────────────────── - -#[test] -fn group_by_hour_empty_input_returns_empty_map() { - log::debug!("[memory_tree_test] group_by_hour: empty input"); - let groups = group_by_hour(&[]); - assert!(groups.is_empty()); -} - -#[test] -fn group_by_hour_single_entry_maps_to_correct_hour() { - // Timestamp 1_711_958_400_000 = 2024-04-01T08:00:00Z - let filename = "1711958400000_abc12345.md".to_string(); - let content = "hello".to_string(); - let entries = vec![(filename, content)]; - let groups = group_by_hour(&entries); - log::debug!("[memory_tree_test] group_by_hour single: groups={groups:?}"); - assert_eq!(groups.len(), 1); - let key = groups.keys().next().unwrap(); - // Hour node id should be "YYYY/MM/DD/HH" - assert_eq!(key.matches('/').count(), 3, "hour id must have 3 slashes"); - assert!( - key.starts_with("2024/04/01/"), - "expected 2024-04-01 hour key; got: {key}" - ); - assert!(key.ends_with("/08"), "expected hour /08; got: {key}"); -} - -#[test] -fn group_by_hour_same_hour_entries_are_merged() { - // Two timestamps within the same hour. - let ts_a = 1_711_958_400_000_i64; // 2024-04-01T08:00:00Z - let ts_b = ts_a + 1_800_000; // +30 min, still same hour - - let entries = vec![ - (format!("{ts_a}_uuid1.md"), "msg-a".to_string()), - (format!("{ts_b}_uuid2.md"), "msg-b".to_string()), - ]; - let groups = group_by_hour(&entries); - log::debug!("[memory_tree_test] group_by_hour same-hour: groups={groups:?}"); - assert_eq!( - groups.len(), - 1, - "same-hour entries must collapse into one group" - ); - let contents = groups.values().next().unwrap(); - assert_eq!(contents.len(), 2); - assert!(contents.contains(&"msg-a".to_string())); - assert!(contents.contains(&"msg-b".to_string())); -} - -#[test] -fn group_by_hour_different_hours_produce_separate_groups() { - // 2024-04-01T08:00:00Z and 2024-04-01T09:00:00Z - let ts_h8 = 1_711_958_400_000_i64; - let ts_h9 = 1_711_962_000_000_i64; - - let entries = vec![ - (format!("{ts_h8}_uuid1.md"), "hour-8-msg".to_string()), - (format!("{ts_h9}_uuid2.md"), "hour-9-msg".to_string()), - ]; - let groups = group_by_hour(&entries); - log::debug!("[memory_tree_test] group_by_hour diff-hours: groups={groups:?}"); - assert_eq!(groups.len(), 2); - let keys: Vec<&String> = groups.keys().collect(); - // BTreeMap returns sorted keys. - assert!( - keys[0].ends_with("/08"), - "first key should end in /08; got {}", - keys[0] - ); - assert!( - keys[1].ends_with("/09"), - "second key should end in /09; got {}", - keys[1] - ); -} - -#[test] -fn group_by_hour_unparseable_filename_falls_back_to_current_hour() { - // A filename with no timestamp prefix should fall back without panic. - let entries = vec![("bad-filename.md".to_string(), "content".to_string())]; - let groups = group_by_hour(&entries); - // Must produce exactly one group (fallback to now). - assert_eq!(groups.len(), 1); - let key = groups.keys().next().unwrap(); - assert_eq!( - key.matches('/').count(), - 3, - "fallback key must still be a valid hour id" - ); -} - -#[test] -fn group_by_hour_output_is_ordered_by_hour_id() { - // BTreeMap guarantees sorted iteration — verify the API honours that. - // Timestamps verified: 2024-04-01T{08,10,12}:00:00Z - let ts_h08 = 1_711_958_400_000_i64; // 2024-04-01T08:00:00Z - let ts_h10 = 1_711_965_600_000_i64; // 2024-04-01T10:00:00Z - let ts_h12 = 1_711_972_800_000_i64; // 2024-04-01T12:00:00Z - - let entries = vec![ - (format!("{ts_h10}_a.md"), "10h".to_string()), - (format!("{ts_h08}_b.md"), "8h".to_string()), - (format!("{ts_h12}_c.md"), "12h".to_string()), - ]; - let groups = group_by_hour(&entries); - let keys: Vec<&String> = groups.keys().collect(); - log::debug!("[memory_tree_test] group_by_hour ordering: keys={keys:?}"); - assert_eq!(keys.len(), 3); - // Sorted lexicographically: .../08 < .../10 < .../12 - assert!(keys[0] < keys[1]); - assert!(keys[1] < keys[2]); -} - -// ── propagate_node ──────────────────────────────────────────────────────── - -#[tokio::test] -async fn propagate_node_with_no_children_is_noop() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let provider = as_model(StubProvider::with_reply("should not be called")); - // No children exist for "2024/03/15" — propagate must succeed silently. - let result = propagate_node( - &cfg, - provider.as_ref(), - "test-ns", - "2024/03/15", - NodeLevel::Day, - ) - .await; - assert!(result.is_ok(), "empty children must not error: {result:?}"); - // No node should have been written. - let node = store::read_node(&cfg, "test-ns", "2024/03/15").unwrap(); - assert!( - node.is_none(), - "propagate with no children must not write a node" - ); -} - -#[tokio::test] -async fn propagate_node_day_from_hour_children_fits_budget() { - // Write two small hour leaves; their combined tokens are well within - // Day::max_tokens() so the LLM is NOT called — combined text is used directly. - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let now = Utc.with_ymd_and_hms(2024, 3, 15, 8, 0, 0).unwrap(); - let make_hour_node = |ns: &str, hour_id: &str, summary: &str| TreeNode { - node_id: hour_id.to_string(), - namespace: ns.to_string(), - level: NodeLevel::Hour, - parent_id: derive_parent_id(hour_id), - summary: summary.to_string(), - token_count: estimate_tokens(summary), - child_count: 0, - created_at: now, - updated_at: now, - metadata: None, - }; - - let ns = "test-ns"; - store::write_node( - &cfg, - &make_hour_node(ns, "2024/03/15/08", "Meeting at 8am."), - ) - .unwrap(); - store::write_node( - &cfg, - &make_hour_node(ns, "2024/03/15/09", "Stand-up at 9am."), - ) - .unwrap(); - - // StubProvider reply should NOT be used when content fits budget. - let provider = as_model(StubProvider::with_reply("SHOULD_NOT_APPEAR")); - propagate_node(&cfg, provider.as_ref(), ns, "2024/03/15", NodeLevel::Day) - .await - .unwrap(); - - let day_node = store::read_node(&cfg, ns, "2024/03/15").unwrap().unwrap(); - log::debug!( - "[memory_tree_test] propagate day: summary={}", - day_node.summary - ); - assert_eq!(day_node.level, NodeLevel::Day); - assert!(day_node.summary.contains("Meeting at 8am.")); - assert!(day_node.summary.contains("Stand-up at 9am.")); - // LLM reply must not appear when content fits the budget. - assert!(!day_node.summary.contains("SHOULD_NOT_APPEAR")); - assert!(day_node.child_count >= 2); -} - -#[tokio::test] -async fn propagate_node_month_from_day_children() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let now = Utc::now(); - let ns = "test-ns"; - - let make_day = |id: &str, text: &str| TreeNode { - node_id: id.to_string(), - namespace: ns.to_string(), - level: NodeLevel::Day, - parent_id: derive_parent_id(id), - summary: text.to_string(), - token_count: estimate_tokens(text), - child_count: 1, - created_at: now, - updated_at: now, - metadata: None, - }; - - store::write_node(&cfg, &make_day("2024/03/14", "Day 14 recap.")).unwrap(); - store::write_node(&cfg, &make_day("2024/03/15", "Day 15 recap.")).unwrap(); - - let provider = as_model(StubProvider::with_reply("Month summary from LLM.")); - propagate_node(&cfg, provider.as_ref(), ns, "2024/03", NodeLevel::Month) - .await - .unwrap(); - - let month = store::read_node(&cfg, ns, "2024/03").unwrap().unwrap(); - assert_eq!(month.level, NodeLevel::Month); - // Either both day summaries appear (budget fit) or the stub reply is used. - let has_day_content = month.summary.contains("Day 14") || month.summary.contains("Day 15"); - let has_stub = month.summary.contains("Month summary from LLM."); - assert!( - has_day_content || has_stub, - "month summary must contain day content or stub reply; got: {}", - month.summary - ); -} - -#[tokio::test] -async fn propagate_node_preserves_created_at_on_update() { - // If a node already exists, propagate must NOT overwrite `created_at`. - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "test-ns"; - let now = Utc::now(); - - // Write an existing day node. - let existing = TreeNode { - node_id: "2024/03/15".to_string(), - namespace: ns.to_string(), - level: NodeLevel::Day, - parent_id: Some("2024/03".to_string()), - summary: "old summary".to_string(), - token_count: 10, - child_count: 0, - created_at: now - chrono::Duration::hours(5), - updated_at: now - chrono::Duration::hours(5), - metadata: None, - }; - store::write_node(&cfg, &existing).unwrap(); - let original_created_at = existing.created_at; - - // Write a child so propagation has something to do. - let child = TreeNode { - node_id: "2024/03/15/10".to_string(), - namespace: ns.to_string(), - level: NodeLevel::Hour, - parent_id: Some("2024/03/15".to_string()), - summary: "hour content".to_string(), - token_count: 5, - child_count: 0, - created_at: now, - updated_at: now, - metadata: None, - }; - store::write_node(&cfg, &child).unwrap(); - - let provider = as_model(StubProvider::with_reply("updated summary")); - propagate_node(&cfg, provider.as_ref(), ns, "2024/03/15", NodeLevel::Day) - .await - .unwrap(); - - let updated = store::read_node(&cfg, ns, "2024/03/15").unwrap().unwrap(); - assert_eq!( - updated.created_at, original_created_at, - "created_at must be preserved across propagation updates" - ); - assert!( - updated.updated_at >= now, - "updated_at must be refreshed on re-propagation" - ); -} - -// ── run_summarization end-to-end ────────────────────────────────────────── - -#[tokio::test] -async fn run_summarization_empty_buffer_returns_none() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let provider = as_model(StubProvider::with_reply("should not be called")); - let ts = Utc::now(); - let result = run_summarization(&cfg, provider.as_ref(), "test-ns", ts) - .await - .unwrap(); - assert!(result.is_none(), "empty buffer must return None"); -} - -#[tokio::test] -async fn run_summarization_drains_buffer_and_writes_hour_node() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "test-ns"; - - // Write two buffer entries at the same hour so they merge into one hour leaf. - let ts = Utc.with_ymd_and_hms(2024, 3, 15, 10, 0, 0).unwrap(); - store::buffer_write(&cfg, ns, "entry one", &ts, None).unwrap(); - store::buffer_write(&cfg, ns, "entry two", &ts, None).unwrap(); - - let provider = as_model(StubProvider::with_reply("hour leaf summary from LLM")); - let last_node = run_summarization(&cfg, provider.as_ref(), ns, ts) - .await - .unwrap(); - - let node = last_node.expect("non-empty buffer must return an hour node"); - log::debug!( - "[memory_tree_test] run_summarization: hour node_id={}", - node.node_id - ); - assert_eq!(node.level, NodeLevel::Hour); - assert_eq!(node.namespace, ns); - - // Buffer must be drained after successful run. - let remaining = store::buffer_read(&cfg, ns).unwrap(); - assert!( - remaining.is_empty(), - "buffer must be empty after summarization" - ); - - // The hour node must be persisted. - let stored = store::read_node(&cfg, ns, &node.node_id).unwrap(); - assert!(stored.is_some(), "hour node must be written to disk"); -} - -#[tokio::test] -async fn run_summarization_builds_ancestor_chain() { - // After a successful run the day/month/year/root ancestor chain must be written. - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "ancestor-test"; - let ts = Utc.with_ymd_and_hms(2024, 3, 15, 10, 0, 0).unwrap(); - - store::buffer_write(&cfg, ns, "test content", &ts, None).unwrap(); - - let provider = as_model(StubProvider::with_reply("summary text")); - run_summarization(&cfg, provider.as_ref(), ns, ts) - .await - .unwrap(); - - // Day, month, year, and root must all be present. - assert!( - store::read_node(&cfg, ns, "2024/03/15").unwrap().is_some(), - "day node must be propagated" - ); - assert!( - store::read_node(&cfg, ns, "2024/03").unwrap().is_some(), - "month node must be propagated" - ); - assert!( - store::read_node(&cfg, ns, "2024").unwrap().is_some(), - "year node must be propagated" - ); - assert!( - store::read_node(&cfg, ns, "root").unwrap().is_some(), - "root node must be propagated" - ); -} - -#[tokio::test] -async fn run_summarization_multi_hour_groups_produce_multiple_hour_leaves() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "multi-hour-test"; - - let ts_h08 = Utc.with_ymd_and_hms(2024, 3, 15, 8, 0, 0).unwrap(); - let ts_h14 = Utc.with_ymd_and_hms(2024, 3, 15, 14, 0, 0).unwrap(); - - // Write entries for two different hours. - store::buffer_write(&cfg, ns, "morning entry", &ts_h08, None).unwrap(); - store::buffer_write(&cfg, ns, "afternoon entry", &ts_h14, None).unwrap(); - - let provider = as_model(StubProvider::with_reply("grouped summary")); - run_summarization(&cfg, provider.as_ref(), ns, ts_h14) - .await - .unwrap(); - - // Both hour leaves must be written. - let hour_08 = store::read_node(&cfg, ns, "2024/03/15/08").unwrap(); - let hour_14 = store::read_node(&cfg, ns, "2024/03/15/14").unwrap(); - assert!(hour_08.is_some(), "hour-08 leaf must be written"); - assert!(hour_14.is_some(), "hour-14 leaf must be written"); - - // Buffer must be empty. - assert!(store::buffer_read(&cfg, ns).unwrap().is_empty()); -} - -#[tokio::test] -async fn rebuild_tree_restores_buffer_and_rewrites_ancestors() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "rebuild-test"; - let ts = Utc.with_ymd_and_hms(2024, 3, 15, 10, 0, 0).unwrap(); - - let make_hour = |id: &str, text: &str| TreeNode { - node_id: id.to_string(), - namespace: ns.to_string(), - level: NodeLevel::Hour, - parent_id: derive_parent_id(id), - summary: text.to_string(), - token_count: estimate_tokens(text), - child_count: 0, - created_at: ts, - updated_at: ts, - metadata: None, - }; - - // Seed the tree with hour leaves and an outdated ancestor/root. - store::write_node(&cfg, &make_hour("2024/03/15/10", "hour ten")).unwrap(); - store::write_node(&cfg, &make_hour("2024/03/15/11", "hour eleven")).unwrap(); - store::write_node( - &cfg, - &TreeNode { - node_id: "2024/03/15".into(), - namespace: ns.into(), - level: NodeLevel::Day, - parent_id: Some("2024/03".into()), - summary: "stale day".into(), - token_count: 2, - child_count: 1, - created_at: ts, - updated_at: ts, - metadata: None, - }, - ) - .unwrap(); - store::write_node( - &cfg, - &TreeNode { - node_id: "root".into(), - namespace: ns.into(), - level: NodeLevel::Root, - parent_id: None, - summary: "stale root".into(), - token_count: 2, - child_count: 1, - created_at: ts, - updated_at: ts, - metadata: None, - }, - ) - .unwrap(); - - // Preserve unsummarized buffer content across rebuild. - store::buffer_write(&cfg, ns, "pending buffer item", &ts, None).unwrap(); - let provider = as_model(StubProvider::with_reply("rebuilt summary")); - - let status = rebuild_tree(&cfg, provider.as_ref(), ns).await.unwrap(); - assert!(status.total_nodes >= 5, "expected leaf + ancestor chain"); - - let restored_buffer = store::buffer_read(&cfg, ns).unwrap(); - assert_eq!( - restored_buffer.len(), - 1, - "buffer entries must survive rebuild" - ); - - let day = store::read_node(&cfg, ns, "2024/03/15").unwrap().unwrap(); - assert!( - day.summary.contains("hour ten") || day.summary.contains("rebuilt summary"), - "day node should be regenerated from hour leaves" - ); - - let root = store::read_node(&cfg, ns, "root").unwrap().unwrap(); - assert!( - root.summary.contains("rebuilt summary") || root.summary.contains("hour ten"), - "root node should be regenerated during rebuild" - ); -} - -#[tokio::test] -async fn rebuild_tree_partial_success_when_one_level_fails() { - // #002 (FR-008): a single propagation node failing must NOT abort the - // whole rebuild. Seed two LARGE hour leaves so the day-level combine - // exceeds the 2000-token Day budget and forces an LLM call there; a - // provider that fails only at the "day" level makes that one node fail. - // The rebuild must still return Ok and the hour leaves must survive. - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - let ns = "partial-rebuild"; - let ts = Utc.with_ymd_and_hms(2024, 3, 15, 10, 0, 0).unwrap(); - - // ~5000 chars each (~1250 tokens) → combined ~2500 tokens > Day budget 2000. - let big = "word ".repeat(1000); - let make_hour = |id: &str| TreeNode { - node_id: id.to_string(), - namespace: ns.to_string(), - level: NodeLevel::Hour, - parent_id: derive_parent_id(id), - summary: big.clone(), - token_count: estimate_tokens(&big), - child_count: 0, - created_at: ts, - updated_at: ts, - metadata: None, - }; - store::write_node(&cfg, &make_hour("2024/03/15/10")).unwrap(); - store::write_node(&cfg, &make_hour("2024/03/15/11")).unwrap(); - - let provider = as_model(FailAtLevelProvider { - fail_level: "day", - reply: "ok summary".to_string(), - }); - - // Must NOT error despite the day-level summarization failing. - let status = rebuild_tree(&cfg, provider.as_ref(), ns) - .await - .expect("partial failure must not abort the rebuild"); - - // The hour leaves (written before propagation) survive. - assert!( - status.total_nodes >= 2, - "hour leaves must survive a partial rebuild" - ); - assert!(store::read_node(&cfg, ns, "2024/03/15/10") - .unwrap() - .is_some()); - assert!(store::read_node(&cfg, ns, "2024/03/15/11") - .unwrap() - .is_some()); -} - -#[tokio::test] -async fn rebuild_tree_on_empty_namespace_is_noop() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let provider = as_model(StubProvider::with_reply("unused")); - let status = rebuild_tree(&cfg, provider.as_ref(), "empty-rebuild") - .await - .unwrap(); - assert_eq!(status.total_nodes, 0); - assert_eq!(status.depth, 0); -} - -#[tokio::test] -async fn summarize_to_limit_truncates_overlong_provider_output() { - let provider = as_model(StubProvider::with_reply("x".repeat(MAX_SUMMARY_CHARS + 50))); - let summary = summarize_to_limit(provider.as_ref(), "short input", 10, "day", "2024/03/15") - .await - .unwrap(); - - assert_eq!(summary.len(), 40, "max_tokens=10 should clamp to 40 chars"); - assert!(summary.chars().all(|c| c == 'x')); -} - -#[test] -fn hour_id_from_buffer_filename_parses_and_rejects_invalid_inputs() { - let parsed = hour_id_from_buffer_filename("1711958400000_uuid.md").unwrap(); - assert_eq!(parsed, "2024/04/01/08"); - - assert!(hour_id_from_buffer_filename("not-a-timestamp.md").is_none()); - assert!(hour_id_from_buffer_filename("abc_123.md").is_none()); -} - -#[test] -fn derive_node_ids_from_hour_id_falls_back_for_non_hour_ids() { - let ids = derive_node_ids_from_hour_id("2024/03/15"); - assert_eq!( - ids, - ( - "2024/03/15".to_string(), - "unknown".to_string(), - "unknown".to_string(), - "unknown".to_string(), - "root".to_string(), - ) - ); -} - -#[test] -fn discover_active_namespaces_requires_markdown_entries_in_buffer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - - let base = cfg.workspace_dir.join("memory").join("namespaces"); - std::fs::create_dir_all(base.join("alpha").join("tree").join("buffer")).unwrap(); - std::fs::create_dir_all(base.join("beta").join("tree").join("buffer")).unwrap(); - std::fs::create_dir_all(base.join("gamma").join("tree")).unwrap(); - - std::fs::write( - base.join("alpha") - .join("tree") - .join("buffer") - .join("entry.md"), - "alpha", - ) - .unwrap(); - std::fs::write( - base.join("beta") - .join("tree") - .join("buffer") - .join("entry.txt"), - "beta", - ) - .unwrap(); - - let active = discover_active_namespaces(&cfg); - assert_eq!(active, vec!["alpha".to_string()]); -} diff --git a/src/openhuman/memory_tree/tree_runtime/store.rs b/src/openhuman/memory_tree/tree_runtime/store.rs index 8a8018f83..8531e8238 100644 --- a/src/openhuman/memory_tree/tree_runtime/store.rs +++ b/src/openhuman/memory_tree/tree_runtime/store.rs @@ -1,522 +1,88 @@ -//! Markdown file-based persistence for the summary tree. -//! -//! Each tree node is stored as a markdown file with YAML frontmatter in the -//! memory namespaces directory: -//! `{workspace}/memory/namespaces/{namespace}/tree/` -//! -//! The folder hierarchy mirrors the time hierarchy: -//! root.md, 2024/summary.md, 2024/03/summary.md, 2024/03/15/summary.md, 2024/03/15/14.md +//! `Config` adapters for tinycortex-owned markdown tree persistence. -use anyhow::{Context, Result}; -use chrono::{DateTime, TimeZone, Utc}; -use serde_json::Value; use std::path::{Path, PathBuf}; +use anyhow::Result; +use chrono::{DateTime, Utc}; +use serde_json::Value; + use crate::openhuman::config::Config; -use crate::openhuman::memory_tree::tree_runtime::types::{ - derive_parent_id, estimate_tokens, level_from_node_id, node_id_to_path, NodeLevel, TreeNode, - TreeStatus, -}; +use crate::openhuman::memory_tree::tree_runtime::types::{TreeNode, TreeStatus}; -// ── Path helpers ─────────────────────────────────────────────────────── +fn engine_config(config: &Config) -> tinycortex::memory::MemoryConfig { + crate::openhuman::tinycortex::memory_config_from(config, config.workspace_dir.clone()) +} -/// Base tree directory for a namespace. pub fn tree_dir(config: &Config, namespace: &str) -> PathBuf { - config - .workspace_dir - .join("memory") - .join("namespaces") - .join(sanitize(namespace)) - .join("tree") + tinycortex::memory::tree::runtime::store::tree_dir(&engine_config(config), namespace) } -/// Buffer directory where raw ingested content is staged before summarization. pub fn buffer_dir(config: &Config, namespace: &str) -> PathBuf { - tree_dir(config, namespace).join("buffer") + tinycortex::memory::tree::runtime::store::buffer_dir(&engine_config(config), namespace) } -/// Absolute file path for a given node. pub fn node_file_path(config: &Config, namespace: &str, node_id: &str) -> PathBuf { - tree_dir(config, namespace).join(node_id_to_path(node_id)) + tinycortex::memory::tree::runtime::store::node_file_path( + &engine_config(config), + namespace, + node_id, + ) } -/// Sanitize a namespace string for use as a directory name. -/// Rejects namespaces containing path-traversal or reserved characters. -fn sanitize(namespace: &str) -> String { - let trimmed = namespace.trim(); - // Replace characters that are unsafe for directory names - trimmed - .replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|', '.'], "_") - .replace("__", "_") -} +pub use tinycortex::memory::tree::runtime::store::{validate_namespace, validate_node_id}; -/// Validate a namespace string, returning an error for empty or dangerous input. -pub fn validate_namespace(namespace: &str) -> Result<(), String> { - let trimmed = namespace.trim(); - if trimmed.is_empty() { - return Err("namespace must not be empty".to_string()); - } - if trimmed.contains("..") { - return Err("namespace must not contain '..'".to_string()); - } - if trimmed.starts_with('/') || trimmed.starts_with('\\') { - return Err("namespace must not start with a path separator".to_string()); - } - Ok(()) -} - -/// Validate a node_id against the allowed canonical formats. -/// Accepts: "root", "YYYY", "YYYY/MM", "YYYY/MM/DD", "YYYY/MM/DD/HH". -/// Rejects path traversal, empty segments, and non-numeric components. -pub fn validate_node_id(node_id: &str) -> Result<(), String> { - if node_id == "root" { - return Ok(()); - } - - // Reject path traversal and dangerous characters - if node_id.contains("..") || node_id.starts_with('/') || node_id.ends_with('/') { - return Err(format!( - "invalid node_id '{node_id}': contains path traversal or leading/trailing slashes" - )); - } - - let parts: Vec<&str> = node_id.split('/').collect(); - if parts.is_empty() || parts.len() > 4 { - return Err(format!( - "invalid node_id '{node_id}': expected 1-4 segments (YYYY[/MM[/DD[/HH]]])" - )); - } - - // All parts must be non-empty numeric strings - for (i, part) in parts.iter().enumerate() { - if part.is_empty() { - return Err(format!( - "invalid node_id '{node_id}': empty segment at position {i}" - )); - } - if !part.chars().all(|c| c.is_ascii_digit()) { - return Err(format!( - "invalid node_id '{node_id}': non-numeric segment '{part}' at position {i}" - )); - } - } - - // Basic range validation - if parts.len() >= 2 { - let month: u32 = parts[1].parse().unwrap_or(0); - if !(1..=12).contains(&month) { - return Err(format!( - "invalid node_id '{node_id}': month {month} out of range 1-12" - )); - } - } - if parts.len() >= 3 { - let day: u32 = parts[2].parse().unwrap_or(0); - if !(1..=31).contains(&day) { - return Err(format!( - "invalid node_id '{node_id}': day {day} out of range 1-31" - )); - } - } - if parts.len() >= 4 { - let hour: u32 = parts[3].parse().unwrap_or(99); - if hour > 23 { - return Err(format!( - "invalid node_id '{node_id}': hour {hour} out of range 0-23" - )); - } - } - - Ok(()) -} - -// ── Write ────────────────────────────────────────────────────────────── - -/// Write a tree node to disk as a markdown file with YAML frontmatter. pub fn write_node(config: &Config, node: &TreeNode) -> Result<()> { - let path = node_file_path(config, &node.namespace, &node.node_id); - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent) - .with_context(|| format!("create dirs for {}", parent.display()))?; - } - - let metadata_line = match &node.metadata { - Some(m) => format!("metadata: {m}\n"), - None => String::new(), - }; - - let frontmatter = format!( - "---\n\ - node_id: \"{}\"\n\ - namespace: \"{}\"\n\ - level: {}\n\ - parent_id: {}\n\ - token_count: {}\n\ - child_count: {}\n\ - created_at: {}\n\ - updated_at: {}\n\ - {}\ - ---\n\n", - node.node_id, - node.namespace, - node.level.as_str(), - match &node.parent_id { - Some(pid) => format!("\"{pid}\""), - None => "~".to_string(), - }, - node.token_count, - node.child_count, - node.created_at.to_rfc3339(), - node.updated_at.to_rfc3339(), - metadata_line, - ); - - let content = format!("{frontmatter}{}\n", node.summary); - std::fs::write(&path, content) - .with_context(|| format!("write tree node {}", path.display()))?; - - tracing::debug!( - "[tree_summarizer] wrote node {} (level={}, tokens={}) -> {}", - node.node_id, - node.level.as_str(), - node.token_count, - path.display() - ); - Ok(()) + tinycortex::memory::tree::runtime::store::write_node(&engine_config(config), node) } -// ── Read ─────────────────────────────────────────────────────────────── - -/// Read a single tree node from its markdown file. Returns `None` if the file -/// does not exist. pub fn read_node(config: &Config, namespace: &str, node_id: &str) -> Result> { - let path = node_file_path(config, namespace, node_id); - if !path.exists() { - return Ok(None); - } - let raw = std::fs::read_to_string(&path) - .with_context(|| format!("read tree node {}", path.display()))?; - parse_node_markdown(&raw, namespace, node_id).map(Some) + tinycortex::memory::tree::runtime::store::read_node(&engine_config(config), namespace, node_id) } -/// Read all direct children of a node. pub fn read_children(config: &Config, namespace: &str, parent_id: &str) -> Result> { - let parent_level = level_from_node_id(parent_id); - let base = tree_dir(config, namespace); - - match parent_level { - NodeLevel::Root => read_subdirectory_summaries(&base, namespace, ""), - NodeLevel::Year | NodeLevel::Month => { - read_subdirectory_summaries(&base, namespace, parent_id) - } - NodeLevel::Day => read_hour_leaves(&base, namespace, parent_id), - NodeLevel::Hour => Ok(vec![]), // leaves have no children - } + tinycortex::memory::tree::runtime::store::read_children( + &engine_config(config), + namespace, + parent_id, + ) } -/// Walk up from a node to the root, returning all ancestors (excluding the node itself). pub fn read_ancestors(config: &Config, namespace: &str, node_id: &str) -> Result> { - let mut ancestors = Vec::new(); - let mut current = derive_parent_id(node_id); - while let Some(pid) = current { - if let Some(node) = read_node(config, namespace, &pid)? { - ancestors.push(node); - } - current = derive_parent_id(&pid); - } - Ok(ancestors) + tinycortex::memory::tree::runtime::store::read_ancestors( + &engine_config(config), + namespace, + node_id, + ) } -/// Recursively count all `.md` files in the tree directory. pub fn count_nodes(config: &Config, namespace: &str) -> Result { - let base = tree_dir(config, namespace); - if !base.exists() { - return Ok(0); - } - count_md_files(&base) + tinycortex::memory::tree::runtime::store::count_nodes(&engine_config(config), namespace) } -/// Scan the tree to produce a status summary. pub fn get_tree_status(config: &Config, namespace: &str) -> Result { - let base = tree_dir(config, namespace); - let total_nodes = if base.exists() { - count_md_files(&base)? - } else { - 0 - }; - - // Determine depth by checking which levels exist. - let mut depth = 0u32; - let root_path = base.join("root.md"); - if root_path.exists() { - depth = 1; - } - - // Scan for years/months/days/hours to figure out actual depth and date range. - let mut oldest: Option> = None; - let mut newest: Option> = None; - - if base.exists() { - for entry in std::fs::read_dir(&base).into_iter().flatten().flatten() { - let name = entry.file_name().to_string_lossy().to_string(); - if entry.file_type().map(|t| t.is_dir()).unwrap_or(false) && name.len() == 4 { - if depth < 2 { - depth = 2; - } - // Scan months, days, hours inside - let year_dir = entry.path(); - for month_entry in std::fs::read_dir(&year_dir).into_iter().flatten().flatten() { - if month_entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { - if depth < 3 { - depth = 3; - } - let month_dir = month_entry.path(); - for day_entry in std::fs::read_dir(&month_dir) - .into_iter() - .flatten() - .flatten() - { - if day_entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { - if depth < 4 { - depth = 4; - } - // Check for hour .md files - let day_dir = day_entry.path(); - for hour_entry in - std::fs::read_dir(&day_dir).into_iter().flatten().flatten() - { - let hname = - hour_entry.file_name().to_string_lossy().to_string(); - if hname.ends_with(".md") && hname != "summary.md" { - if depth < 5 { - depth = 5; - } - // Try to parse timestamp from path - if let Some(ts) = timestamp_from_hour_path( - &name, - month_entry.file_name().to_string_lossy().as_ref(), - day_entry.file_name().to_string_lossy().as_ref(), - &hname, - ) { - match &oldest { - None => oldest = Some(ts), - Some(o) if ts < *o => oldest = Some(ts), - _ => {} - } - match &newest { - None => newest = Some(ts), - Some(n) if ts > *n => newest = Some(ts), - _ => {} - } - } - } - } - } - } - } - } - } - } - } - - Ok(TreeStatus { - namespace: namespace.to_string(), - total_nodes, - depth, - oldest_entry: oldest, - newest_entry: newest, - last_run_at: None, // filled by caller if needed - }) + tinycortex::memory::tree::runtime::store::get_tree_status(&engine_config(config), namespace) } -/// Pull the root-level summary out of every tree summarizer namespace -/// that has been written to the given workspace. -/// -/// Each namespace's `root.md` body is truncated to `per_namespace_cap` -/// chars so a single huge namespace can't dominate the prompt; we then -/// stop accumulating once the running total crosses `total_cap` so -/// workspaces with dozens of namespaces can't blow the context window. -/// -/// Failures (missing files, parse errors) are logged at debug level -/// and silently dropped — user memory is best-effort context, never a -/// hard requirement for running a turn or rendering a prompt dump. -/// -/// Returns a stable-ordered `Vec<(namespace, body, updated_at)>` so -/// byte-identical inputs produce byte-identical output across process -/// restarts (the renderer downstream relies on this for KV-cache prefix -/// reuse). The `updated_at` is the namespace root node's timestamp, so -/// the prompt renderer can stamp how current each summary is and the -/// model never serves stale memory as today's (#2944). A plain tuple -/// (rather than a prompt-layer struct) keeps this module free of any -/// `agent::prompts` dependency. pub fn collect_root_summaries_with_caps( workspace_dir: &Path, per_namespace_cap: usize, total_cap: usize, ) -> Vec<(String, String, DateTime)> { - // The store functions all read `config.workspace_dir` and nothing - // else, so we shim a tiny `Config` from the caller's path. Cheap - // (a few allocations) and avoids forcing every call site to thread - // a real `Config` through just for two read calls. - let config = Config { - workspace_dir: workspace_dir.to_path_buf(), - ..Config::default() - }; - - let namespaces = match list_namespaces_with_root(&config) { - Ok(list) => list, - Err(e) => { - tracing::debug!( - error = %e, - workspace = %workspace_dir.display(), - "[tree_summarizer] could not enumerate namespaces" - ); - return Vec::new(); - } - }; - - if namespaces.is_empty() { - tracing::debug!( - workspace = %workspace_dir.display(), - "[tree_summarizer] no namespaces with a root summary" - ); - return Vec::new(); - } - - let mut out = Vec::new(); - let mut total_chars: usize = 0; - - for ns in namespaces { - if total_chars >= total_cap { - tracing::debug!( - namespace = %ns, - total_chars, - "[tree_summarizer] stopping at namespace — total budget exhausted" - ); - break; - } - - let node = match read_node(&config, &ns, "root") { - Ok(Some(node)) => node, - Ok(None) => continue, - Err(e) => { - tracing::debug!( - namespace = %ns, - error = %e, - "[tree_summarizer] failed to read root summary" - ); - continue; - } - }; - - let body = node.summary.trim(); - if body.is_empty() { - continue; - } - - // Per-namespace cap (char count, not byte length, so non-ASCII - // text doesn't silently overshoot). - let body_chars = body.chars().count(); - let truncated: String = if body_chars > per_namespace_cap { - body.chars().take(per_namespace_cap).collect::() + "\n\n[... truncated]" - } else { - body.to_string() - }; - let truncated_chars = truncated.chars().count(); - - // Total cap — use char counts consistently. If this entry - // would push us over, clip to the remaining budget so we - // still get something for the namespace instead of dropping - // it entirely. - let remaining = total_cap.saturating_sub(total_chars); - let final_body = if truncated_chars > remaining { - let mut clipped: String = truncated.chars().take(remaining).collect(); - clipped.push_str("\n\n[... truncated]"); - clipped - } else { - truncated - }; - - total_chars += final_body.chars().count(); - let final_chars = final_body.chars().count(); - tracing::debug!( - namespace = %ns, - chars = final_chars, - running_total = total_chars, - "[tree_summarizer] including namespace in root-summary collection" - ); - out.push((ns, final_body, node.updated_at)); - } - - tracing::info!( - included = out.len(), - total_chars, - "[tree_summarizer] collected root summaries" - ); - out + tinycortex::memory::tree::runtime::store::collect_root_summaries_with_caps( + workspace_dir, + per_namespace_cap, + total_cap, + ) } -/// Enumerate every namespace under the workspace that has a `root.md` -/// summary written. Returns the on-disk directory names (already -/// sanitised) — these are the keys callers should pass back into -/// [`read_node`] / [`tree_dir`] when reading content. -/// -/// Used by the orchestrator's prompt builder to inject "user memory" -/// into the system prompt: each namespace's root summary is the -/// densest/highest-quality artefact we can hand the model, capped by -/// `NodeLevel::Root::max_tokens()` (currently 20 000 tokens). -/// -/// Skips namespaces that exist on disk but have not yet been -/// summarised (no `root.md`) — those would render as empty headings -/// and only burn cache space. pub fn list_namespaces_with_root(config: &Config) -> Result> { - let base = config.workspace_dir.join("memory").join("namespaces"); - if !base.exists() { - return Ok(Vec::new()); - } - - let mut out = Vec::new(); - for entry in std::fs::read_dir(&base) - .with_context(|| format!("scan namespaces dir {}", base.display()))? - { - let entry = entry?; - if !entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { - continue; - } - let ns_name = entry.file_name().to_string_lossy().to_string(); - let root_path = entry.path().join("tree").join("root.md"); - if root_path.exists() { - out.push(ns_name); - } - } - - // Stable order so the prompt body stays cache-friendly across - // process restarts. Without this, `read_dir` ordering is - // filesystem-dependent and would shuffle the cache prefix bytes. - out.sort(); - Ok(out) + tinycortex::memory::tree::runtime::store::list_namespaces_with_root(&engine_config(config)) } -/// Remove the entire tree directory for a namespace. pub fn delete_tree(config: &Config, namespace: &str) -> Result { - let base = tree_dir(config, namespace); - if !base.exists() { - return Ok(0); - } - let count = count_md_files(&base)?; - std::fs::remove_dir_all(&base).with_context(|| format!("delete tree at {}", base.display()))?; - tracing::debug!( - "[tree_summarizer] deleted tree for namespace '{}' ({} nodes)", - namespace, - count - ); - Ok(count) + tinycortex::memory::tree::runtime::store::delete_tree(&engine_config(config), namespace) } -// ── Buffer operations ────────────────────────────────────────────────── - -/// Append raw content to the ingestion buffer as a timestamped file. -/// Optionally includes metadata as a JSON object stored alongside the content. pub fn buffer_write( config: &Config, namespace: &str, @@ -524,355 +90,31 @@ pub fn buffer_write( ts: &DateTime, metadata: Option<&Value>, ) -> Result { - let dir = buffer_dir(config, namespace); - std::fs::create_dir_all(&dir) - .with_context(|| format!("create buffer dir {}", dir.display()))?; - - let filename = format!( - "{}_{}.md", - ts.timestamp_millis(), - &uuid::Uuid::new_v4().to_string()[..8] - ); - let path = dir.join(&filename); - - // If metadata is provided, write it as a YAML frontmatter block - let file_content = if let Some(meta) = metadata { - let meta_str = serde_json::to_string(meta).unwrap_or_default(); - format!("---\nmetadata: {meta_str}\n---\n\n{content}") - } else { - content.to_string() - }; - - std::fs::write(&path, file_content) - .with_context(|| format!("write buffer entry {}", path.display()))?; - - tracing::debug!( - "[tree_summarizer] buffered {} bytes for namespace '{}' -> {}", - content.len(), + tinycortex::memory::tree::runtime::store::buffer_write( + &engine_config(config), namespace, - filename - ); - Ok(path) -} - -/// Read all buffered entries non-destructively, returning `(filename, content)` pairs -/// sorted by filename (chronological). Files remain on disk until explicitly deleted -/// via [`buffer_delete`]. -pub fn buffer_read(config: &Config, namespace: &str) -> Result> { - let dir = buffer_dir(config, namespace); - if !dir.exists() { - return Ok(vec![]); - } - - let mut entries: Vec<(String, PathBuf)> = Vec::new(); - for entry in std::fs::read_dir(&dir)? { - let entry = entry?; - let path = entry.path(); - if path.extension().map(|e| e == "md").unwrap_or(false) { - let name = entry.file_name().to_string_lossy().to_string(); - entries.push((name, path)); - } - } - - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - let mut contents = Vec::with_capacity(entries.len()); - for (name, path) in &entries { - let raw = std::fs::read_to_string(path) - .with_context(|| format!("read buffer entry {}", path.display()))?; - // Strip metadata frontmatter if present, pass raw content - let text = strip_buffer_frontmatter(&raw); - contents.push((name.clone(), text)); - } - - tracing::debug!( - "[tree_summarizer] read {} buffer entries for namespace '{}'", - contents.len(), - namespace - ); - Ok(contents) -} - -/// Delete specific buffer entries by filename after they have been successfully -/// processed and durably written as hour leaves. -pub fn buffer_delete(config: &Config, namespace: &str, filenames: &[String]) -> Result<()> { - let dir = buffer_dir(config, namespace); - for name in filenames { - let path = dir.join(name); - if path.exists() { - std::fs::remove_file(&path).with_context(|| { - format!( - "failed to remove buffer entry '{}' at {}", - name, - path.display() - ) - })?; - } - } - tracing::debug!( - "[tree_summarizer] deleted {} buffer entries for namespace '{}'", - filenames.len(), - namespace - ); - Ok(()) -} - -/// Read and drain all buffered entries. Convenience wrapper that calls -/// [`buffer_read`] then [`buffer_delete`]. Use the split API when you need -/// to defer deletion until after durable writes complete. -pub fn buffer_drain(config: &Config, namespace: &str) -> Result> { - let entries = buffer_read(config, namespace)?; - if entries.is_empty() { - return Ok(entries); - } - let filenames: Vec = entries.iter().map(|(name, _)| name.clone()).collect(); - buffer_delete(config, namespace, &filenames)?; - tracing::debug!( - "[tree_summarizer] drained {} buffer entries for namespace '{}'", - entries.len(), - namespace - ); - Ok(entries) -} - -/// Strip the optional metadata frontmatter from a buffer entry, -/// returning only the content body. -fn strip_buffer_frontmatter(raw: &str) -> String { - let trimmed = raw.trim_start(); - if !trimmed.starts_with("---") { - return raw.to_string(); - } - let after_open = &trimmed[3..]; - if let Some(close_pos) = after_open.find("\n---") { - let body_start = close_pos + 4; - after_open[body_start..] - .trim_start_matches('\n') - .to_string() - } else { - raw.to_string() - } -} - -// ── Internal helpers ─────────────────────────────────────────────────── - -/// Read summary.md files from subdirectories of a given parent path. -fn read_subdirectory_summaries( - base: &Path, - namespace: &str, - parent_id: &str, -) -> Result> { - let scan_dir = if parent_id.is_empty() { - base.to_path_buf() - } else { - base.join(parent_id) - }; - if !scan_dir.exists() { - return Ok(vec![]); - } - - let mut children = Vec::new(); - for entry in std::fs::read_dir(&scan_dir)? { - let entry = entry?; - if !entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { - continue; - } - let child_name = entry.file_name().to_string_lossy().to_string(); - // Skip non-numeric directories and the buffer directory - if child_name == "buffer" - || child_name == "buffer_backup" - || child_name.chars().any(|c| !c.is_ascii_digit()) - { - continue; - } - let child_id = if parent_id.is_empty() { - child_name - } else { - format!("{parent_id}/{child_name}") - }; - let summary_path = entry.path().join("summary.md"); - if summary_path.exists() { - let raw = std::fs::read_to_string(&summary_path)?; - if let Ok(node) = parse_node_markdown(&raw, namespace, &child_id) { - children.push(node); - } - } - } - - children.sort_by(|a, b| a.node_id.cmp(&b.node_id)); - Ok(children) -} - -/// Read hour leaf .md files (excluding summary.md) from a day directory. -fn read_hour_leaves(base: &Path, namespace: &str, day_id: &str) -> Result> { - let day_dir = base.join(day_id); - if !day_dir.exists() { - return Ok(vec![]); - } - - let mut leaves = Vec::new(); - for entry in std::fs::read_dir(&day_dir)? { - let entry = entry?; - let name = entry.file_name().to_string_lossy().to_string(); - if !name.ends_with(".md") || name == "summary.md" { - continue; - } - let hour_part = name.trim_end_matches(".md"); - let node_id = format!("{day_id}/{hour_part}"); - let raw = std::fs::read_to_string(entry.path())?; - if let Ok(node) = parse_node_markdown(&raw, namespace, &node_id) { - leaves.push(node); - } - } - - leaves.sort_by(|a, b| a.node_id.cmp(&b.node_id)); - Ok(leaves) -} - -/// Public entry point for parsing a markdown node (used by engine rebuild). -pub fn parse_node_markdown_pub(raw: &str, namespace: &str, node_id: &str) -> Result { - parse_node_markdown(raw, namespace, node_id) -} - -/// Parse a markdown file with YAML frontmatter into a `TreeNode`. -fn parse_node_markdown(raw: &str, namespace: &str, node_id: &str) -> Result { - let (frontmatter, body_raw) = split_frontmatter(raw); - let body = body_raw.trim_end().to_string(); - - let level = frontmatter - .get("level") - .and_then(|v| NodeLevel::from_str_label(v)) - .unwrap_or_else(|| level_from_node_id(node_id)); - - let parent_id = frontmatter - .get("parent_id") - .and_then(|v| { - let trimmed = v.trim().trim_matches('"'); - if trimmed == "~" || trimmed.is_empty() { - None - } else { - Some(trimmed.to_string()) - } - }) - .or_else(|| derive_parent_id(node_id)); - - let token_count = frontmatter - .get("token_count") - .and_then(|v| v.parse::().ok()) - .unwrap_or_else(|| estimate_tokens(&body)); - - let child_count = frontmatter - .get("child_count") - .and_then(|v| v.parse::().ok()) - .unwrap_or(0); - - // Deterministic fallbacks. `updated_at` now flows into the system - // prompt as a freshness stamp (#2944), and the renderer relies on - // byte-identical output across process restarts for KV-cache reuse — - // `Utc::now()` would change on every read and break that, *and* would - // misrepresent an undated legacy node as "fresh today" (the exact - // stale-as-current bug this work fixes). For nodes that predate the - // `created_at`/`updated_at` frontmatter we fall back to UNIX_EPOCH so - // the date is stable and conservatively renders as ancient rather - // than current. `updated_at` falls back to `created_at` (its best - // available estimate of when the node was last touched). - let created_at = frontmatter - .get("created_at") - .and_then(|v| DateTime::parse_from_rfc3339(v).ok()) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or(DateTime::::UNIX_EPOCH); - - let updated_at = frontmatter - .get("updated_at") - .and_then(|v| DateTime::parse_from_rfc3339(v).ok()) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or(created_at); - - let metadata = frontmatter.get("metadata").map(|v| v.to_string()); - - Ok(TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level, - parent_id, - summary: body, - token_count, - child_count, - created_at, - updated_at, + content, + ts, metadata, - }) + ) } -/// Split markdown into (frontmatter key-value map, body text). -fn split_frontmatter(raw: &str) -> (std::collections::HashMap, String) { - let mut map = std::collections::HashMap::new(); - let trimmed = raw.trim_start(); - - if !trimmed.starts_with("---") { - return (map, raw.to_string()); - } - - // Find the closing --- - let after_open = &trimmed[3..]; - if let Some(close_pos) = after_open.find("\n---") { - let fm_block = &after_open[..close_pos]; - let body_start = close_pos + 4; // skip "\n---" - let body = after_open[body_start..] - .trim_start_matches('\n') - .to_string(); - - for line in fm_block.lines() { - let line = line.trim(); - if line.is_empty() { - continue; - } - if let Some(colon_pos) = line.find(':') { - let key = line[..colon_pos].trim().to_string(); - let value = line[colon_pos + 1..].trim().trim_matches('"').to_string(); - map.insert(key, value); - } - } - - (map, body) - } else { - (map, raw.to_string()) - } +pub fn buffer_read(config: &Config, namespace: &str) -> Result> { + tinycortex::memory::tree::runtime::store::buffer_read(&engine_config(config), namespace) } -fn count_md_files(dir: &Path) -> Result { - let mut count = 0u64; - for entry in std::fs::read_dir(dir)? { - let entry = entry?; - let ft = entry.file_type()?; - if ft.is_dir() { - let name = entry.file_name().to_string_lossy().to_string(); - if name == "buffer" || name == "buffer_backup" { - continue; // skip buffer directories - } - count += count_md_files(&entry.path())?; - } else if ft.is_file() && entry.path().extension().map(|e| e == "md").unwrap_or(false) { - count += 1; - } - } - Ok(count) +pub fn buffer_delete(config: &Config, namespace: &str, filenames: &[String]) -> Result<()> { + tinycortex::memory::tree::runtime::store::buffer_delete( + &engine_config(config), + namespace, + filenames, + ) } -fn timestamp_from_hour_path( - year: &str, - month: &str, - day: &str, - hour_file: &str, -) -> Option> { - let hour = hour_file.trim_end_matches(".md"); - let y: i32 = year.parse().ok()?; - let m: u32 = month.parse().ok()?; - let d: u32 = day.parse().ok()?; - let h: u32 = hour.parse().ok()?; - chrono::Utc.with_ymd_and_hms(y, m, d, h, 0, 0).single() +pub fn buffer_drain(config: &Config, namespace: &str) -> Result> { + tinycortex::memory::tree::runtime::store::buffer_drain(&engine_config(config), namespace) } -// ── Tests ────────────────────────────────────────────────────────────── - -#[cfg(test)] -#[path = "store_tests.rs"] -mod tests; +pub fn parse_node_markdown_pub(raw: &str, namespace: &str, node_id: &str) -> Result { + tinycortex::memory::tree::runtime::store::parse_node_markdown_pub(raw, namespace, node_id) +} diff --git a/src/openhuman/memory_tree/tree_runtime/store_tests.rs b/src/openhuman/memory_tree/tree_runtime/store_tests.rs deleted file mode 100644 index f10729692..000000000 --- a/src/openhuman/memory_tree/tree_runtime/store_tests.rs +++ /dev/null @@ -1,361 +0,0 @@ -use super::*; -use tempfile::TempDir; - -fn test_config(tmp: &TempDir) -> Config { - let config = Config { - workspace_dir: tmp.path().join("workspace"), - action_dir: tmp.path().join("workspace"), - config_path: tmp.path().join("config.toml"), - ..Config::default() - }; - std::fs::create_dir_all(&config.workspace_dir).unwrap(); - config -} - -fn make_node(namespace: &str, node_id: &str, summary: &str) -> TreeNode { - let level = level_from_node_id(node_id); - TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level, - parent_id: derive_parent_id(node_id), - summary: summary.to_string(), - token_count: estimate_tokens(summary), - child_count: 0, - created_at: Utc::now(), - updated_at: Utc::now(), - metadata: None, - } -} - -#[test] -fn write_and_read_node_roundtrip() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - let node = make_node(ns, "root", "All-time summary of events."); - write_node(&config, &node).unwrap(); - - let read_back = read_node(&config, ns, "root").unwrap().unwrap(); - assert_eq!(read_back.node_id, "root"); - assert_eq!(read_back.level, NodeLevel::Root); - assert_eq!(read_back.summary, "All-time summary of events."); - assert!(read_back.parent_id.is_none()); -} - -#[test] -fn write_and_read_hour_leaf() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - let node = make_node(ns, "2024/03/15/14", "Hour 14 summary."); - write_node(&config, &node).unwrap(); - - let read_back = read_node(&config, ns, "2024/03/15/14").unwrap().unwrap(); - assert_eq!(read_back.level, NodeLevel::Hour); - assert_eq!(read_back.parent_id.as_deref(), Some("2024/03/15")); - assert_eq!(read_back.summary, "Hour 14 summary."); -} - -#[test] -fn read_children_of_day() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - // Write some hour leaves - for hour in [10, 11, 14] { - let node = make_node( - ns, - &format!("2024/03/15/{hour:02}"), - &format!("Hour {hour}."), - ); - write_node(&config, &node).unwrap(); - } - // Write the day summary (should not appear as a child) - let day = make_node(ns, "2024/03/15", "Day summary."); - write_node(&config, &day).unwrap(); - - let children = read_children(&config, ns, "2024/03/15").unwrap(); - assert_eq!(children.len(), 3); - assert_eq!(children[0].node_id, "2024/03/15/10"); - assert_eq!(children[1].node_id, "2024/03/15/11"); - assert_eq!(children[2].node_id, "2024/03/15/14"); -} - -#[test] -fn read_children_of_root() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - for year in ["2023", "2024"] { - let node = make_node(ns, year, &format!("Year {year} summary.")); - write_node(&config, &node).unwrap(); - } - - let children = read_children(&config, ns, "root").unwrap(); - assert_eq!(children.len(), 2); - assert_eq!(children[0].node_id, "2023"); - assert_eq!(children[1].node_id, "2024"); -} - -#[test] -fn read_node_missing_returns_none() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - assert!(read_node(&config, "ns", "root").unwrap().is_none()); -} - -#[test] -fn count_nodes_and_status() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - write_node(&config, &make_node(ns, "root", "root")).unwrap(); - write_node(&config, &make_node(ns, "2024", "year")).unwrap(); - write_node(&config, &make_node(ns, "2024/03", "month")).unwrap(); - write_node(&config, &make_node(ns, "2024/03/15", "day")).unwrap(); - write_node(&config, &make_node(ns, "2024/03/15/14", "hour")).unwrap(); - - assert_eq!(count_nodes(&config, ns).unwrap(), 5); - - let status = get_tree_status(&config, ns).unwrap(); - assert_eq!(status.total_nodes, 5); - assert_eq!(status.depth, 5); -} - -#[test] -fn delete_tree_removes_all() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - write_node(&config, &make_node(ns, "root", "root")).unwrap(); - write_node(&config, &make_node(ns, "2024/03/15/14", "hour")).unwrap(); - - let deleted = delete_tree(&config, ns).unwrap(); - assert!(deleted >= 2); - assert_eq!(count_nodes(&config, ns).unwrap(), 0); -} - -#[test] -fn buffer_write_and_drain() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - let ts1 = Utc.with_ymd_and_hms(2024, 3, 15, 10, 0, 0).unwrap(); - let ts2 = Utc.with_ymd_and_hms(2024, 3, 15, 11, 0, 0).unwrap(); - - buffer_write(&config, ns, "entry one", &ts1, None).unwrap(); - buffer_write(&config, ns, "entry two", &ts2, None).unwrap(); - - let drained = buffer_drain(&config, ns).unwrap(); - assert_eq!(drained.len(), 2); - // Sorted by filename (timestamp prefix), so ts1 < ts2 - assert_eq!(drained[0].1, "entry one"); - assert_eq!(drained[1].1, "entry two"); - - // Buffer should be empty now - let again = buffer_drain(&config, ns).unwrap(); - assert!(again.is_empty()); -} - -#[test] -fn buffer_write_with_metadata() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - let now = Utc::now(); - - let meta = serde_json::json!({"source": "test", "priority": 1}); - buffer_write(&config, ns, "entry with meta", &now, Some(&meta)).unwrap(); - - let drained = buffer_drain(&config, ns).unwrap(); - assert_eq!(drained.len(), 1); - // Content should be stripped of frontmatter - assert_eq!(drained[0].1, "entry with meta"); -} - -#[test] -fn ancestors_walk_to_root() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - let ns = "test-ns"; - - write_node(&config, &make_node(ns, "root", "root")).unwrap(); - write_node(&config, &make_node(ns, "2024", "year")).unwrap(); - write_node(&config, &make_node(ns, "2024/03", "month")).unwrap(); - write_node(&config, &make_node(ns, "2024/03/15", "day")).unwrap(); - - let ancestors = read_ancestors(&config, ns, "2024/03/15/14").unwrap(); - let ids: Vec<&str> = ancestors.iter().map(|n| n.node_id.as_str()).collect(); - assert_eq!(ids, vec!["2024/03/15", "2024/03", "2024", "root"]); -} - -#[test] -fn frontmatter_parsing() { - let raw = "---\nnode_id: \"root\"\nlevel: root\ntoken_count: 42\n---\n\nHello world."; - let (fm, body) = split_frontmatter(raw); - assert_eq!(fm.get("level").unwrap(), "root"); - assert_eq!(fm.get("token_count").unwrap(), "42"); - assert_eq!(body, "Hello world."); -} - -#[test] -fn parse_node_markdown_uses_deterministic_fallback_timestamps() { - // #2944: `updated_at` now flows into the system prompt as a freshness - // stamp, and the renderer relies on byte-stable output for KV-cache - // reuse. Legacy nodes that predate the `created_at`/`updated_at` - // frontmatter must therefore fall back deterministically — never to - // `Utc::now()`, which would change on every read and masquerade an - // undated node as "fresh today". - - // No timestamps at all → both fall back to UNIX_EPOCH (renders as an - // unmistakably ancient date rather than today's). - let raw = "---\nnode_id: \"root\"\nlevel: root\n---\n\nUndated summary."; - let node = parse_node_markdown_pub(raw, "ns", "root").unwrap(); - assert_eq!(node.created_at, DateTime::::UNIX_EPOCH); - assert_eq!(node.updated_at, DateTime::::UNIX_EPOCH); - - // `created_at` present but `updated_at` missing → `updated_at` falls - // back to `created_at`, its best estimate of last-touched. - let raw = - "---\nnode_id: \"root\"\nlevel: root\ncreated_at: 2026-05-25T09:00:00Z\n---\n\nSummary."; - let node = parse_node_markdown_pub(raw, "ns", "root").unwrap(); - let created = DateTime::parse_from_rfc3339("2026-05-25T09:00:00Z") - .unwrap() - .with_timezone(&Utc); - assert_eq!(node.created_at, created); - assert_eq!(node.updated_at, created); -} - -#[test] -fn validate_node_id_accepts_valid() { - assert!(validate_node_id("root").is_ok()); - assert!(validate_node_id("2024").is_ok()); - assert!(validate_node_id("2024/03").is_ok()); - assert!(validate_node_id("2024/03/15").is_ok()); - assert!(validate_node_id("2024/03/15/14").is_ok()); -} - -#[test] -fn validate_node_id_rejects_traversal() { - assert!(validate_node_id("..").is_err()); - assert!(validate_node_id("../etc").is_err()); - assert!(validate_node_id("2024/../etc").is_err()); - assert!(validate_node_id("/2024").is_err()); - assert!(validate_node_id("2024/").is_err()); -} - -#[test] -fn validate_node_id_rejects_non_numeric() { - assert!(validate_node_id("abc").is_err()); - assert!(validate_node_id("2024/abc").is_err()); - assert!(validate_node_id("2024/03/15/foo").is_err()); -} - -#[test] -fn validate_node_id_rejects_out_of_range() { - assert!(validate_node_id("2024/13").is_err()); // month 13 - assert!(validate_node_id("2024/03/32").is_err()); // day 32 - assert!(validate_node_id("2024/03/15/24").is_err()); // hour 24 -} - -#[test] -fn validate_namespace_rejects_dangerous() { - assert!(validate_namespace("").is_err()); - assert!(validate_namespace(" ").is_err()); - assert!(validate_namespace("../etc").is_err()); - assert!(validate_namespace("/absolute").is_err()); -} - -#[test] -fn validate_namespace_accepts_valid() { - assert!(validate_namespace("my-namespace").is_ok()); - assert!(validate_namespace("skill:gmail:user@example.com").is_ok()); -} - -#[test] -fn list_namespaces_with_root_returns_only_summarised() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - - // ns_a has a root node — should be returned. - write_node(&config, &make_node("ns_a", "root", "alpha summary")).unwrap(); - // ns_b has only an hour leaf, no root — should be filtered out. - write_node(&config, &make_node("ns_b", "2024/03/15/14", "hour")).unwrap(); - // ns_c has a root. - write_node(&config, &make_node("ns_c", "root", "gamma summary")).unwrap(); - - let listed = list_namespaces_with_root(&config).unwrap(); - // Sorted alphabetically for cache stability — see fn docs. - assert_eq!(listed, vec!["ns_a".to_string(), "ns_c".to_string()]); -} - -#[test] -fn collect_root_summaries_respects_per_namespace_cap() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - - let big = "x".repeat(50); - write_node(&config, &make_node("ns", "root", &big)).unwrap(); - - // Per-namespace cap of 10 should clip the body. - let result = collect_root_summaries_with_caps(&config.workspace_dir, 10, 10_000); - assert_eq!(result.len(), 1); - let (ns, body, _updated_at) = &result[0]; - assert_eq!(ns, "ns"); - assert!( - body.starts_with("xxxxxxxxxx"), - "expected the first 10 x's, got: {body}" - ); - assert!(body.contains("[... truncated]")); -} - -#[test] -fn collect_root_summaries_carries_root_updated_at() { - // #2944: the namespace root's `updated_at` must survive the - // disk round-trip so the prompt renderer can stamp memory freshness. - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - - let fixed = chrono::DateTime::parse_from_rfc3339("2026-05-25T09:00:00Z") - .unwrap() - .with_timezone(&Utc); - let mut node = make_node("ns", "root", "summary body"); - node.updated_at = fixed; - write_node(&config, &node).unwrap(); - - let result = collect_root_summaries_with_caps(&config.workspace_dir, 1000, 10_000); - assert_eq!(result.len(), 1); - let (ns, _body, updated_at) = &result[0]; - assert_eq!(ns, "ns"); - assert_eq!(*updated_at, fixed, "root updated_at must round-trip"); -} - -#[test] -fn collect_root_summaries_stops_at_total_cap() { - let tmp = TempDir::new().unwrap(); - let config = test_config(&tmp); - - write_node(&config, &make_node("aaa", "root", "first")).unwrap(); - write_node(&config, &make_node("bbb", "root", "second")).unwrap(); - write_node(&config, &make_node("ccc", "root", "third")).unwrap(); - - // Total cap of 5 chars — should accept aaa ("first" = 5), - // then break before reading bbb because total >= cap. - let result = collect_root_summaries_with_caps(&config.workspace_dir, 100, 5); - assert_eq!(result.len(), 1); - assert_eq!(result[0].0, "aaa"); -} - -#[test] -fn collect_root_summaries_returns_empty_for_unknown_workspace() { - let tmp = TempDir::new().unwrap(); - let result = collect_root_summaries_with_caps(&tmp.path().join("nope"), 100, 1000); - assert!(result.is_empty()); -} diff --git a/src/openhuman/memory_tree/tree_runtime/types.rs b/src/openhuman/memory_tree/tree_runtime/types.rs index c5c621f4d..dbf693ba8 100644 --- a/src/openhuman/memory_tree/tree_runtime/types.rs +++ b/src/openhuman/memory_tree/tree_runtime/types.rs @@ -1,294 +1,6 @@ -//! Domain types for the tree summarizer. +//! Stable host path for tinycortex-owned markdown tree runtime types. -use chrono::{DateTime, Datelike, Timelike, Utc}; -use serde::{Deserialize, Serialize}; -use std::path::PathBuf; - -// ── Node level ───────────────────────────────────────────────────────── - -/// Hierarchical level of a tree node. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum NodeLevel { - Root, - Year, - Month, - Day, - Hour, -} - -impl NodeLevel { - /// Maximum number of tokens allowed at this level. - pub fn max_tokens(&self) -> u32 { - match self { - Self::Hour => 1_000, - Self::Day => 2_000, - Self::Month => 4_000, - Self::Year => 8_000, - Self::Root => 20_000, - } - } - - /// The level above this one in the hierarchy (`None` for root). - pub fn parent_level(&self) -> Option { - match self { - Self::Hour => Some(Self::Day), - Self::Day => Some(Self::Month), - Self::Month => Some(Self::Year), - Self::Year => Some(Self::Root), - Self::Root => None, - } - } - - /// True only for the leaf level (hour). - pub fn is_leaf(&self) -> bool { - matches!(self, Self::Hour) - } - - /// Parse a level string from YAML frontmatter. - pub fn from_str_label(s: &str) -> Option { - match s.trim().to_ascii_lowercase().as_str() { - "root" => Some(Self::Root), - "year" => Some(Self::Year), - "month" => Some(Self::Month), - "day" => Some(Self::Day), - "hour" => Some(Self::Hour), - _ => None, - } - } - - /// Label for display / frontmatter. - pub fn as_str(&self) -> &'static str { - match self { - Self::Root => "root", - Self::Year => "year", - Self::Month => "month", - Self::Day => "day", - Self::Hour => "hour", - } - } -} - -// ── Tree node ────────────────────────────────────────────────────────── - -/// A single node in the summary tree. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TreeNode { - pub node_id: String, - pub namespace: String, - pub level: NodeLevel, - pub parent_id: Option, - pub summary: String, - pub token_count: u32, - pub child_count: u32, - pub created_at: DateTime, - pub updated_at: DateTime, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -// ── Status ───────────────────────────────────────────────────────────── - -/// Metadata about an entire tree within a namespace. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TreeStatus { - pub namespace: String, - pub total_nodes: u64, - pub depth: u32, - pub oldest_entry: Option>, - pub newest_entry: Option>, - pub last_run_at: Option>, -} - -// ── Ingest request ───────────────────────────────────────────────────── - -/// Input for appending raw content to the ingestion buffer. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct IngestRequest { - pub namespace: String, - pub content: String, - #[serde(default)] - pub timestamp: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -// ── Query result ─────────────────────────────────────────────────────── - -/// Result of a tree query at a specific node. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct QueryResult { - pub node: TreeNode, - pub children: Vec, -} - -// ── Helpers ──────────────────────────────────────────────────────────── - -/// Rough token estimate: ~4 characters per token. -pub fn estimate_tokens(text: &str) -> u32 { - (text.len() as u32).div_ceil(4) -} - -/// Derive the parent node ID from a node ID. -/// -/// - `"2024/03/15/14"` → `Some("2024/03/15")` -/// - `"2024/03/15"` → `Some("2024/03")` -/// - `"2024/03"` → `Some("2024")` -/// - `"2024"` → `Some("root")` -/// - `"root"` → `None` -pub fn derive_parent_id(node_id: &str) -> Option { - if node_id == "root" { - return None; - } - match node_id.rfind('/') { - Some(pos) => Some(node_id[..pos].to_string()), - None => Some("root".to_string()), - } -} - -/// Determine the `NodeLevel` from a node ID string. -pub fn level_from_node_id(node_id: &str) -> NodeLevel { - if node_id == "root" { - return NodeLevel::Root; - } - match node_id.matches('/').count() { - 0 => NodeLevel::Year, // "2024" - 1 => NodeLevel::Month, // "2024/03" - 2 => NodeLevel::Day, // "2024/03/15" - _ => NodeLevel::Hour, // "2024/03/15/14" - } -} - -/// Derive all ancestor node IDs from a timestamp (hour through root). -/// -/// Returns `(hour_id, day_id, month_id, year_id, root_id)`. -pub fn derive_node_ids(ts: &DateTime) -> (String, String, String, String, String) { - let year = format!("{}", ts.year()); - let month = format!("{}/{:02}", ts.year(), ts.month()); - let day = format!("{}/{:02}/{:02}", ts.year(), ts.month(), ts.day()); - let hour = format!( - "{}/{:02}/{:02}/{:02}", - ts.year(), - ts.month(), - ts.day(), - ts.hour() - ); - (hour, day, month, year, "root".to_string()) -} - -/// Convert a node ID to a relative file path within the tree directory. -/// -/// - `"root"` → `root.md` -/// - `"2024"` → `2024/summary.md` -/// - `"2024/03"` → `2024/03/summary.md` -/// - `"2024/03/15"` → `2024/03/15/summary.md` -/// - `"2024/03/15/14"` → `2024/03/15/14.md` (hour leaf — file, not folder) -pub fn node_id_to_path(node_id: &str) -> PathBuf { - if node_id == "root" { - return PathBuf::from("root.md"); - } - let level = level_from_node_id(node_id); - if level.is_leaf() { - // Hour leaf: "2024/03/15/14" → "2024/03/15/14.md" - PathBuf::from(format!("{node_id}.md")) - } else { - // Non-leaf: "2024/03" → "2024/03/summary.md" - PathBuf::from(node_id).join("summary.md") - } -} - -// ── Tests ────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - #[test] - fn node_level_max_tokens() { - assert_eq!(NodeLevel::Hour.max_tokens(), 1_000); - assert_eq!(NodeLevel::Day.max_tokens(), 2_000); - assert_eq!(NodeLevel::Month.max_tokens(), 4_000); - assert_eq!(NodeLevel::Year.max_tokens(), 8_000); - assert_eq!(NodeLevel::Root.max_tokens(), 20_000); - } - - #[test] - fn node_level_parent_chain() { - assert_eq!(NodeLevel::Hour.parent_level(), Some(NodeLevel::Day)); - assert_eq!(NodeLevel::Day.parent_level(), Some(NodeLevel::Month)); - assert_eq!(NodeLevel::Month.parent_level(), Some(NodeLevel::Year)); - assert_eq!(NodeLevel::Year.parent_level(), Some(NodeLevel::Root)); - assert_eq!(NodeLevel::Root.parent_level(), None); - } - - #[test] - fn derive_parent_id_chain() { - assert_eq!(derive_parent_id("2024/03/15/14"), Some("2024/03/15".into())); - assert_eq!(derive_parent_id("2024/03/15"), Some("2024/03".into())); - assert_eq!(derive_parent_id("2024/03"), Some("2024".into())); - assert_eq!(derive_parent_id("2024"), Some("root".into())); - assert_eq!(derive_parent_id("root"), None); - } - - #[test] - fn level_from_node_id_all_levels() { - assert_eq!(level_from_node_id("root"), NodeLevel::Root); - assert_eq!(level_from_node_id("2024"), NodeLevel::Year); - assert_eq!(level_from_node_id("2024/03"), NodeLevel::Month); - assert_eq!(level_from_node_id("2024/03/15"), NodeLevel::Day); - assert_eq!(level_from_node_id("2024/03/15/14"), NodeLevel::Hour); - } - - #[test] - fn derive_node_ids_from_timestamp() { - let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 0).unwrap(); - let (hour, day, month, year, root) = derive_node_ids(&ts); - assert_eq!(hour, "2024/03/15/14"); - assert_eq!(day, "2024/03/15"); - assert_eq!(month, "2024/03"); - assert_eq!(year, "2024"); - assert_eq!(root, "root"); - } - - #[test] - fn node_id_to_path_mapping() { - assert_eq!(node_id_to_path("root"), PathBuf::from("root.md")); - assert_eq!(node_id_to_path("2024"), PathBuf::from("2024/summary.md")); - assert_eq!( - node_id_to_path("2024/03"), - PathBuf::from("2024/03/summary.md") - ); - assert_eq!( - node_id_to_path("2024/03/15"), - PathBuf::from("2024/03/15/summary.md") - ); - assert_eq!( - node_id_to_path("2024/03/15/14"), - PathBuf::from("2024/03/15/14.md") - ); - } - - #[test] - fn estimate_tokens_rough() { - assert_eq!(estimate_tokens(""), 0); - assert_eq!(estimate_tokens("abcd"), 1); - assert_eq!(estimate_tokens("abcdefgh"), 2); - // Roughly 4 chars per token - let text = "a".repeat(4000); - assert_eq!(estimate_tokens(&text), 1000); - } - - #[test] - fn node_level_roundtrip() { - for level in [ - NodeLevel::Root, - NodeLevel::Year, - NodeLevel::Month, - NodeLevel::Day, - NodeLevel::Hour, - ] { - assert_eq!(NodeLevel::from_str_label(level.as_str()), Some(level)); - } - } -} +pub use tinycortex::memory::tree::runtime::{ + 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/openhuman/task_sources/pipeline_tests.rs b/src/openhuman/task_sources/pipeline_tests.rs index 863eb1d48..385ad190e 100644 --- a/src/openhuman/task_sources/pipeline_tests.rs +++ b/src/openhuman/task_sources/pipeline_tests.rs @@ -2,7 +2,7 @@ use super::*; use crate::openhuman::config::Config; use crate::openhuman::memory_sync::composio::providers::{ register_provider, ComposioProvider, NormalizedTask, ProviderContext, ProviderUserProfile, - SyncOutcome, SyncReason, TaskFetchFilter, + TaskFetchFilter, }; use crate::openhuman::task_sources::store; use crate::openhuman::task_sources::types::{FilterSpec, ProviderSlug, SourceTarget}; @@ -35,13 +35,6 @@ impl ComposioProvider for StubProvider { ) -> Result { Ok(ProviderUserProfile::default()) } - async fn sync( - &self, - _ctx: &ProviderContext, - _reason: SyncReason, - ) -> Result { - Ok(SyncOutcome::default()) - } async fn fetch_tasks( &self, _ctx: &ProviderContext, diff --git a/src/openhuman/tinycortex/config.rs b/src/openhuman/tinycortex/config.rs index d45f3c3eb..7e30d40bc 100644 --- a/src/openhuman/tinycortex/config.rs +++ b/src/openhuman/tinycortex/config.rs @@ -35,6 +35,7 @@ use crate::openhuman::config::Config; /// is gated by the W3 golden-workspace harness). pub fn memory_config_from(config: &Config, workspace: PathBuf) -> MemoryConfig { let mut mc = MemoryConfig::new(workspace); + mc.content_root = Some(config.memory_tree_content_root()); mc.embedding = EmbeddingConfig { dim: config.memory.embedding_dimensions, model: config.memory.embedding_model.clone(), diff --git a/src/openhuman/tinycortex/ingest.rs b/src/openhuman/tinycortex/ingest.rs new file mode 100644 index 000000000..0420013e8 --- /dev/null +++ b/src/openhuman/tinycortex/ingest.rs @@ -0,0 +1,58 @@ +//! Host adapters for tinycortex on-demand ingestion. + +use tinycortex::memory::ingest::TreeJobSink; +use tinycortex::memory::score::extract::{LlmEntityExtractor, LlmExtractorConfig}; +use tinycortex::memory::score::ScoringConfig; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_queue::{self, ExtractChunkPayload, NewJob}; + +pub struct HostTreeJobSink { + config: Config, +} + +impl HostTreeJobSink { + pub fn new(config: Config) -> Self { + Self { config } + } +} + +impl TreeJobSink for HostTreeJobSink { + fn enqueue_extract(&self, chunk_id: &str) -> anyhow::Result<()> { + let job = NewJob::extract_chunk(&ExtractChunkPayload { + chunk_id: chunk_id.into(), + })?; + memory_queue::enqueue(&self.config, &job)?; + Ok(()) + } +} + +fn scoring_config(config: &Config) -> ScoringConfig { + match super::build_chat_provider(config) { + Ok(provider) => { + let mut extractor = LlmExtractorConfig::default(); + extractor.output_language = config.output_language.clone(); + ScoringConfig::with_llm_extractor(std::sync::Arc::new(LlmEntityExtractor::new( + extractor, provider, + ))) + } + Err(error) => { + tracing::warn!(%error, "[memory:ingest] chat provider unavailable; using regex scoring"); + ScoringConfig::default_regex_only() + } + } +} + +pub fn context( + config: &Config, +) -> ( + tinycortex::memory::MemoryConfig, + HostTreeJobSink, + ScoringConfig, +) { + ( + super::memory_config_from(config, config.workspace_dir.clone()), + HostTreeJobSink::new(config.clone()), + scoring_config(config), + ) +} diff --git a/src/openhuman/tinycortex/mod.rs b/src/openhuman/tinycortex/mod.rs index 45b890464..b013e42d9 100644 --- a/src/openhuman/tinycortex/mod.rs +++ b/src/openhuman/tinycortex/mod.rs @@ -36,21 +36,31 @@ mod chat; mod config; mod embeddings; +mod ingest; #[cfg(test)] mod parity; mod queue_driver; +mod seal; +mod summariser; +mod sync; pub use chat::{build_chat_provider, SeamChatProvider}; pub use config::memory_config_from; pub use embeddings::SeamEmbedder; +pub use ingest::{context as ingest_context, HostTreeJobSink}; pub use queue_driver::{ classify_worker_error, HostQueueDelegates, WorkerErrorAction, WorkerReport, }; -// `chunk_tree_scope` is a crate-internal pure helper (the single canonical host -// copy after the W4 handlers deletion); read paths like -// `memory_tree::retrieval::cover` import it through this seam. -pub(crate) use queue_driver::chunk_tree_scope; - +pub use seal::{ + cascade_tree, flush_stale_tree_buffers, seal_document_subtree, + seal_one_level as seal_tree_level, +}; +pub use summariser::HostSummariser; +pub use sync::{ + load_composio_sync_state, run_composio_connection, run_composio_connection_with_budgets, + run_gmail_backfill, run_slack_search_backfill, run_source_pipeline, sync_context, + HostSyncAdapter, SourcePipelineFailure, HOST_SYNC_STATE_NAMESPACE, +}; // Facade re-exports — the rest of the host imports memory-engine types through // this one seam so consumer import paths stay stable as the internals flip to // the crate (the type-unification decision, spec §0.5). `MemoryTaint` is the diff --git a/src/openhuman/tinycortex/queue_driver.rs b/src/openhuman/tinycortex/queue_driver.rs index b69c9b813..0498122e4 100644 --- a/src/openhuman/tinycortex/queue_driver.rs +++ b/src/openhuman/tinycortex/queue_driver.rs @@ -54,7 +54,7 @@ use crate::openhuman::memory_tree::health; use crate::openhuman::memory_tree::score; use crate::openhuman::memory_tree::score::embed::{build_write_embedder, pack_checked, Embedder}; use crate::openhuman::memory_tree::score::store as score_store; -use crate::openhuman::memory_tree::tree::{bucket_seal, TreeFactory}; +use crate::openhuman::memory_tree::tree::TreeFactory; // ── Pure scope helpers (ported verbatim from `memory_queue::handlers`) ──────── // These pin the SAME source→tree mapping the append-buffer path uses, so reads @@ -358,7 +358,7 @@ impl QueueDelegates for HostQueueDelegates { let body = content_read::read_chunk_body(config, &chunk.id) .with_context(|| format!("read full body for extract chunk_id={}", chunk.id))?; let preview = std::mem::replace(&mut chunk.content, body); - let scoring_cfg = score::ScoringConfig::from_config(config); + let scoring_cfg = score::scoring_config_from(config); let result = score::score_chunk(&chunk, &scoring_cfg).await?; chunk.content = preview; @@ -493,7 +493,9 @@ impl QueueDelegates for HostQueueDelegates { }; trees_store::upsert_buffer_tx(&tx, &buf)?; } - let should_seal = bucket_seal::should_seal(&buf); + let memory_config = + super::memory_config_from(&self.config, self.config.workspace_dir.clone()); + let should_seal = tinycortex::memory::tree::should_seal(&memory_config, &buf); if is_source_target { if let Some(cid) = lifecycle_chunk_id.as_deref() { chunk_store::set_chunk_lifecycle_status_tx( @@ -532,12 +534,15 @@ impl QueueDelegates for HostQueueDelegates { }; let buf = trees_store::get_buffer(&self.config, &tree.id, payload.level)?; let forced = payload.force_now_ms.is_some(); - if buf.is_empty() || (!forced && !bucket_seal::should_seal(&buf)) { + let memory_config = + super::memory_config_from(&self.config, self.config.workspace_dir.clone()); + if buf.is_empty() + || (!forced && !tinycortex::memory::tree::should_seal(&memory_config, &buf)) + { return Ok(None); } let strategy = TreeFactory::from_tree(&tree).label_strategy(&self.config); - let summary_id = - bucket_seal::seal_one_level(&self.config, &tree, &buf, &strategy, true).await?; + let summary_id = super::seal_tree_level(&self.config, &tree, &buf, &strategy, true).await?; // Best-effort: rewrite the sealed summary's on-disk obsidian tags. Entity // rows were committed inside seal_one_level, so they are visible here. if let Err(e) = content_store::update_summary_tags(&self.config, &summary_id) { @@ -579,7 +584,7 @@ impl QueueDelegates for HostQueueDelegates { // One physical tree per connection scope (e.g. notion:{connection_id}). let tree = get_or_create_source_tree(&self.config, &payload.tree_scope)?; let strategy = TreeFactory::from_tree(&tree).label_strategy(&self.config); - bucket_seal::seal_document_subtree( + super::seal_document_subtree( &self.config, &tree, &payload.doc_id, diff --git a/src/openhuman/tinycortex/seal.rs b/src/openhuman/tinycortex/seal.rs new file mode 100644 index 000000000..e014c9299 --- /dev/null +++ b/src/openhuman/tinycortex/seal.rs @@ -0,0 +1,234 @@ +//! Product compute and notification adapters for tinycortex sealing. + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use chrono::Duration; + +use crate::core::event_bus::{publish_global, DomainEvent}; +use crate::openhuman::config::Config; +use crate::openhuman::memory_store::content::wiki_git::{SummaryCommitBatch, SummaryCommitEntry}; +use crate::openhuman::memory_store::trees::types::{Buffer, SummaryNode, Tree}; +use crate::openhuman::memory_tree::score::embed::{build_write_embedder, Embedder as HostEmbedder}; +use crate::openhuman::memory_tree::tree::bucket_seal::LabelStrategy; + +use super::{memory_config_from, HostSummariser}; + +struct EmbedderBridge<'a>(&'a dyn HostEmbedder); + +#[async_trait] +impl tinycortex::memory::score::embed::Embedder for EmbedderBridge<'_> { + fn name(&self) -> &'static str { + self.0.name() + } + + async fn embed(&self, text: &str) -> Result> { + let vector = self.0.embed(text).await.map_err(|error| { + let failure = crate::openhuman::memory_tree::health::classify_embed_error(&error); + anyhow::Error::new(failure).context(format!("seal embedding failed: {error:#}")) + })?; + crate::openhuman::memory_tree::score::embed::pack_checked(&vector) + .context("seal embedding dimension check")?; + crate::openhuman::memory_tree::health::clear_semantic_recall_degraded(); + Ok(vector) + } +} + +struct Observer<'a> { + config: &'a Config, +} + +impl tinycortex::memory::tree::SealObserver for Observer<'_> { + fn progress(&self, tree: &Tree, step: &str, level: u32, item_count: Option) { + publish_global(DomainEvent::MemoryTreeBuildProgress { + phase: "seal".to_string(), + step: step.to_string(), + tree_scope: Some(tree.scope.clone()), + level: Some(level), + item_count, + detail: None, + }); + } + + fn summary_committed( + &self, + tree: &Tree, + node: &SummaryNode, + content_path: &str, + reason: &str, + ) -> Result<()> { + crate::openhuman::memory_store::content::wiki_git::commit_summaries( + &self.config.memory_tree_content_root(), + &SummaryCommitBatch { + reason: reason.to_string(), + tree_id: tree.id.clone(), + tree_scope: tree.scope.clone(), + entries: vec![SummaryCommitEntry { + summary_id: node.id.clone(), + content_path: content_path.to_string(), + level: node.level, + child_count: node.child_ids.len(), + token_count: node.token_count, + time_range_start: node.time_range_start, + time_range_end: node.time_range_end, + }], + }, + ) + } +} + +pub async fn seal_one_level( + config: &Config, + tree: &Tree, + buffer: &Buffer, + strategy: &LabelStrategy, + enqueue_follow_ups: bool, +) -> Result { + if let Err(error) = crate::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults( + &config.memory_tree_content_root(), + ) { + log::warn!("[tree::bucket_seal] obsidian defaults failed: {error:#}"); + } + let host_embedder = build_write_embedder(config)?; + let embedder_bridge = host_embedder.as_deref().map(EmbedderBridge); + let summariser = HostSummariser::new(config.clone()); + let observer = Observer { config }; + let strategy = match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + tinycortex::memory::tree::LabelStrategy::ExtractFromContent(extractor.clone()) + } + LabelStrategy::UnionFromChildren => { + tinycortex::memory::tree::LabelStrategy::UnionFromChildren + } + LabelStrategy::Empty => tinycortex::memory::tree::LabelStrategy::Empty, + }; + tinycortex::memory::tree::seal_one_level_with_services( + &memory_config_from(config, config.workspace_dir.clone()), + tree, + buffer, + &tinycortex::memory::tree::SealServices { + summariser: &summariser, + embedder: embedder_bridge + .as_ref() + .map(|bridge| bridge as &dyn tinycortex::memory::score::embed::Embedder), + observer: &observer, + }, + &strategy, + enqueue_follow_ups, + ) + .await +} + +pub async fn seal_document_subtree( + config: &Config, + tree: &Tree, + doc_id: &str, + version_ms: Option, + chunk_ids: &[String], + strategy: &LabelStrategy, +) -> Result { + if let Err(error) = crate::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults( + &config.memory_tree_content_root(), + ) { + log::warn!("[tree::bucket_seal] obsidian defaults failed: {error:#}"); + } + let host_embedder = build_write_embedder(config)?; + let embedder_bridge = host_embedder.as_deref().map(EmbedderBridge); + let summariser = HostSummariser::new(config.clone()); + let observer = Observer { config }; + let strategy = match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + tinycortex::memory::tree::LabelStrategy::ExtractFromContent(extractor.clone()) + } + LabelStrategy::UnionFromChildren => { + tinycortex::memory::tree::LabelStrategy::UnionFromChildren + } + LabelStrategy::Empty => tinycortex::memory::tree::LabelStrategy::Empty, + }; + tinycortex::memory::tree::seal_document_subtree_with_services( + &memory_config_from(config, config.workspace_dir.clone()), + tree, + doc_id, + version_ms, + chunk_ids, + &tinycortex::memory::tree::SealServices { + summariser: &summariser, + embedder: embedder_bridge + .as_ref() + .map(|bridge| bridge as &dyn tinycortex::memory::score::embed::Embedder), + observer: &observer, + }, + &strategy, + ) + .await +} + +pub async fn cascade_tree( + config: &Config, + tree: &Tree, + start_level: u32, + force: bool, + strategy: &LabelStrategy, +) -> Result> { + let host_embedder = build_write_embedder(config)?; + let embedder_bridge = host_embedder.as_deref().map(EmbedderBridge); + let summariser = HostSummariser::new(config.clone()); + let observer = Observer { config }; + let strategy = match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + tinycortex::memory::tree::LabelStrategy::ExtractFromContent(extractor.clone()) + } + LabelStrategy::UnionFromChildren => { + tinycortex::memory::tree::LabelStrategy::UnionFromChildren + } + LabelStrategy::Empty => tinycortex::memory::tree::LabelStrategy::Empty, + }; + tinycortex::memory::tree::cascade_all_from_with_services( + &memory_config_from(config, config.workspace_dir.clone()), + tree, + start_level, + force, + &tinycortex::memory::tree::SealServices { + summariser: &summariser, + embedder: embedder_bridge + .as_ref() + .map(|bridge| bridge as &dyn tinycortex::memory::score::embed::Embedder), + observer: &observer, + }, + &strategy, + false, + ) + .await +} + +pub async fn flush_stale_tree_buffers( + config: &Config, + max_age: Duration, + strategy: &LabelStrategy, +) -> Result { + let host_embedder = build_write_embedder(config)?; + let embedder_bridge = host_embedder.as_deref().map(EmbedderBridge); + let summariser = HostSummariser::new(config.clone()); + let observer = Observer { config }; + let strategy = match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + tinycortex::memory::tree::LabelStrategy::ExtractFromContent(extractor.clone()) + } + LabelStrategy::UnionFromChildren => { + tinycortex::memory::tree::LabelStrategy::UnionFromChildren + } + LabelStrategy::Empty => tinycortex::memory::tree::LabelStrategy::Empty, + }; + tinycortex::memory::tree::flush_stale_buffers_with_services( + &memory_config_from(config, config.workspace_dir.clone()), + max_age, + &tinycortex::memory::tree::SealServices { + summariser: &summariser, + embedder: embedder_bridge + .as_ref() + .map(|bridge| bridge as &dyn tinycortex::memory::score::embed::Embedder), + observer: &observer, + }, + &strategy, + ) + .await +} diff --git a/src/openhuman/tinycortex/summariser.rs b/src/openhuman/tinycortex/summariser.rs new file mode 100644 index 000000000..28cc0ba8e --- /dev/null +++ b/src/openhuman/tinycortex/summariser.rs @@ -0,0 +1,63 @@ +//! OpenHuman LLM adapter for tinycortex tree summarization. + +use async_trait::async_trait; +use tinycortex::memory::tree::{ + Summariser, SummaryCall, SummaryContext, SummaryInput, SummaryOutput, +}; + +use crate::openhuman::config::Config; + +#[derive(Clone)] +pub struct HostSummariser { + config: Config, +} + +impl HostSummariser { + pub fn new(config: Config) -> Self { + Self { config } + } + + async fn call( + &self, + inputs: &[SummaryInput], + context: &SummaryContext<'_>, + ) -> anyhow::Result { + let output = + crate::openhuman::memory_tree::summarise::summarise(&self.config, inputs, context) + .await?; + Ok(SummaryCall { + output: SummaryOutput { + content: output.content, + token_count: output.token_count, + entities: output.entities, + topics: output.topics, + }, + input_tokens: output.input_tokens, + output_tokens: output.output_tokens, + charged_amount_usd: output.charged_amount_usd, + }) + } +} + +#[async_trait] +impl Summariser for HostSummariser { + fn name(&self) -> &str { + "openhuman" + } + + async fn summarise( + &self, + inputs: &[SummaryInput], + context: &SummaryContext<'_>, + ) -> anyhow::Result { + Ok(self.call(inputs, context).await?.output) + } + + async fn summarise_with_usage( + &self, + inputs: &[SummaryInput], + context: &SummaryContext<'_>, + ) -> anyhow::Result { + self.call(inputs, context).await + } +} diff --git a/src/openhuman/tinycortex/sync.rs b/src/openhuman/tinycortex/sync.rs new file mode 100644 index 000000000..37967d7b8 --- /dev/null +++ b/src/openhuman/tinycortex/sync.rs @@ -0,0 +1,524 @@ +//! OpenHuman service adapters for tinycortex live synchronization. + +use async_trait::async_trait; +use tinycortex::memory::sync::{ + ClickUpSyncPipeline, ComposioClient, ExternalSourceReader, GitHubSyncPipeline, + GithubRepoSyncPipeline, GmailSyncPipeline, LinearSyncPipeline, LocalDocument, + LocalDocumentSink, NotionSyncPipeline, SkillDocSink, SkillDocument, + SlackSearchBackfillPipeline, SlackSyncPipeline, SyncContext, SyncDispatcher, SyncEvent, + SyncEventSink, SyncOutcome, SyncPipeline, SyncStage, SyncStateStore, WorkspaceSourcePipeline, +}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_sources::{MemorySourceEntry, SourceKind}; +use crate::openhuman::memory_store::MemoryClientRef; + +pub const HOST_SYNC_STATE_NAMESPACE: &str = "composio-sync-state"; + +pub struct HostSyncAdapter { + memory: MemoryClientRef, + config: Option, +} + +#[derive(Debug)] +pub struct SourcePipelineFailure { + pub message: String, + pub actions_called: u32, + pub provider_cost_usd: f64, +} + +impl std::fmt::Display for SourcePipelineFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.message) + } +} + +impl SourcePipelineFailure { + fn without_usage(message: impl Into) -> Self { + Self { + message: message.into(), + actions_called: 0, + provider_cost_usd: 0.0, + } + } +} + +impl HostSyncAdapter { + pub fn new(memory: MemoryClientRef) -> Self { + Self { + memory, + config: None, + } + } + + fn with_config(memory: MemoryClientRef, config: Config) -> Self { + Self { + memory, + config: Some(config), + } + } +} + +#[async_trait] +impl ExternalSourceReader for HostSyncAdapter { + async fn list_items( + &self, + source: &tinycortex::memory::sources::MemorySourceEntry, + ) -> anyhow::Result> { + let config = self + .config + .as_ref() + .ok_or_else(|| anyhow::anyhow!("external source reader requires host config"))?; + let host_source: MemorySourceEntry = serde_json::from_value(serde_json::to_value(source)?)?; + let reader = crate::openhuman::memory_sources::readers::reader_for(&host_source.kind); + let items = reader + .list_items(&host_source, config) + .await + .map_err(anyhow::Error::msg)?; + serde_json::from_value(serde_json::to_value(items)?).map_err(Into::into) + } + + async fn read_item( + &self, + source: &tinycortex::memory::sources::MemorySourceEntry, + item_id: &str, + ) -> anyhow::Result { + let config = self + .config + .as_ref() + .ok_or_else(|| anyhow::anyhow!("external source reader requires host config"))?; + let host_source: MemorySourceEntry = serde_json::from_value(serde_json::to_value(source)?)?; + let reader = crate::openhuman::memory_sources::readers::reader_for(&host_source.kind); + let content = reader + .read_item(&host_source, item_id, config) + .await + .map_err(anyhow::Error::msg)?; + serde_json::from_value(serde_json::to_value(content)?).map_err(Into::into) + } +} + +pub fn sync_context(memory: MemoryClientRef) -> SyncContext { + let adapter = std::sync::Arc::new(HostSyncAdapter::new(memory)); + SyncContext { + events: adapter.clone(), + documents: adapter.clone(), + state: adapter, + local_documents: None, + external_sources: None, + summariser: None, + } +} + +fn source_sync_context(memory: MemoryClientRef, config: &Config, local: bool) -> SyncContext { + let adapter = std::sync::Arc::new(HostSyncAdapter::with_config(memory, config.clone())); + SyncContext { + events: adapter.clone(), + documents: adapter.clone(), + state: adapter.clone(), + local_documents: local.then(|| adapter.clone() as std::sync::Arc), + external_sources: local.then(|| adapter as std::sync::Arc), + summariser: local.then(|| { + std::sync::Arc::new(super::HostSummariser::new(config.clone())) + as std::sync::Arc + }), + } +} + +pub async fn run_source_pipeline( + source: &MemorySourceEntry, + config: &Config, +) -> Result { + let memory = crate::openhuman::memory::global::client_if_ready() + .ok_or_else(|| SourcePipelineFailure::without_usage("memory client is not ready"))?; + let mut memory_config = super::memory_config_from(config, config.workspace_dir.clone()); + memory_config.sync.interval_secs = config.memory_sync_interval_secs; + memory_config.sync.budget.max_items = source.max_items; + memory_config.sync.budget.max_tokens_per_sync = source.max_tokens_per_sync; + memory_config.sync.budget.max_cost_per_sync_usd = source.max_cost_per_sync_usd; + memory_config.sync.budget.sync_depth_days = source.sync_depth_days; + + let pipeline = build_pipeline(source, config, &mut memory_config) + .map_err(SourcePipelineFailure::without_usage)?; + let pipeline_id = pipeline.id().to_owned(); + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(pipeline) + .map_err(|error| SourcePipelineFailure::without_usage(error.to_string()))?; + dispatcher + .tick( + &pipeline_id, + &memory_config, + &source_sync_context(memory, config, source.kind != SourceKind::Composio), + ) + .await + .map_err(|error| { + let usage = error.downcast_ref::(); + SourcePipelineFailure { + message: error.to_string(), + actions_called: usage.map_or(0, |error| error.actions_called), + provider_cost_usd: usage.map_or(0.0, |error| error.provider_cost_usd), + } + }) +} + +/// Run a Composio connection through tinycortex, preserving any source-level +/// budgets already configured in OpenHuman's registry. +pub async fn run_composio_connection( + toolkit: &str, + connection_id: &str, + config: &Config, +) -> Result { + run_composio_connection_with_budgets(toolkit, connection_id, config, None, None).await +} + +/// Run a Composio connection with request-scoped budget overrides. +/// +/// Provider RPCs carry these values in `ProviderContext`, before a source has +/// necessarily been persisted in the registry. Explicit values therefore take +/// precedence, while `None` preserves the registered/default source budget. +pub async fn run_composio_connection_with_budgets( + toolkit: &str, + connection_id: &str, + config: &Config, + max_items: Option, + sync_depth_days: Option, +) -> Result { + let mut source = config + .memory_sources + .iter() + .find(|source| { + source.kind == SourceKind::Composio + && source.connection_id.as_deref() == Some(connection_id) + }) + .cloned() + .unwrap_or_else(|| { + let (max_items, sync_depth_days) = + crate::openhuman::memory_sources::memory_sync_defaults_for_toolkit(toolkit); + MemorySourceEntry { + id: format!("composio:{toolkit}:{connection_id}"), + kind: SourceKind::Composio, + label: format!("{toolkit} connection"), + enabled: true, + toolkit: Some(toolkit.to_ascii_lowercase()), + connection_id: Some(connection_id.to_string()), + path: None, + glob: None, + url: None, + branch: None, + paths: Vec::new(), + max_commits: None, + max_issues: None, + max_prs: None, + query: None, + since_days: None, + max_items, + selector: None, + max_tokens_per_sync: None, + max_cost_per_sync_usd: None, + sync_depth_days, + } + }); + + source.max_items = max_items; + source.sync_depth_days = sync_depth_days; + + tracing::debug!( + toolkit, + connection_id, + source_id = %source.id, + max_items = ?source.max_items, + sync_depth_days = ?source.sync_depth_days, + "[tinycortex:sync] dispatching Composio connection" + ); + run_source_pipeline(&source, config).await +} + +pub async fn load_composio_sync_state( + toolkit: &str, + connection_id: &str, +) -> anyhow::Result { + let memory = crate::openhuman::memory::global::client_if_ready() + .ok_or_else(|| anyhow::anyhow!("memory client is not ready"))?; + let adapter = HostSyncAdapter::new(memory); + tinycortex::memory::sync::SyncState::load(&adapter, toolkit, connection_id).await +} + +pub async fn run_slack_search_backfill( + connection_id: &str, + backfill_days: i64, + config: &Config, +) -> Result { + let memory = crate::openhuman::memory::global::client_if_ready() + .ok_or_else(|| SourcePipelineFailure::without_usage("memory client is not ready"))?; + let mut memory_config = super::memory_config_from(config, config.workspace_dir.clone()); + let composio = composio_config(config).map_err(SourcePipelineFailure::without_usage)?; + memory_config.sync.composio = Some(composio.clone()); + let pipeline = std::sync::Arc::new(SlackSearchBackfillPipeline::new( + ComposioClient::new(composio), + connection_id, + backfill_days, + )); + let pipeline_id = pipeline.id().to_owned(); + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(pipeline) + .map_err(|error| SourcePipelineFailure::without_usage(error.to_string()))?; + dispatcher + .tick( + &pipeline_id, + &memory_config, + &source_sync_context(memory, config, false), + ) + .await + .map_err(|error| SourcePipelineFailure::without_usage(error.to_string())) +} + +pub async fn run_gmail_backfill( + connection_id: &str, + query: &str, + max_pages: usize, + page_size: usize, + config: &Config, +) -> Result { + let memory = crate::openhuman::memory::global::client_if_ready() + .ok_or_else(|| SourcePipelineFailure::without_usage("memory client is not ready"))?; + let mut memory_config = super::memory_config_from(config, config.workspace_dir.clone()); + let composio = composio_config(config).map_err(SourcePipelineFailure::without_usage)?; + memory_config.sync.composio = Some(composio.clone()); + let pipeline = std::sync::Arc::new( + GmailSyncPipeline::new(ComposioClient::new(composio), connection_id) + .with_limits(max_pages, page_size) + .with_query(query), + ); + let pipeline_id = pipeline.id().to_owned(); + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(pipeline) + .map_err(|error| SourcePipelineFailure::without_usage(error.to_string()))?; + dispatcher + .tick( + &pipeline_id, + &memory_config, + &source_sync_context(memory, config, false), + ) + .await + .map_err(|error| SourcePipelineFailure::without_usage(error.to_string())) +} + +fn build_pipeline( + source: &MemorySourceEntry, + config: &Config, + memory_config: &mut tinycortex::memory::config::MemoryConfig, +) -> Result, String> { + if source.kind != SourceKind::Composio { + let crate_source: tinycortex::memory::sources::MemorySourceEntry = serde_json::from_value( + serde_json::to_value(source).map_err(|error| error.to_string())?, + ) + .map_err(|error| error.to_string())?; + if source.kind == SourceKind::GithubRepo { + return GithubRepoSyncPipeline::new(crate_source) + .map(|pipeline| std::sync::Arc::new(pipeline) as std::sync::Arc) + .map_err(|error| error.to_string()); + } + return WorkspaceSourcePipeline::new(crate_source) + .map(|pipeline| std::sync::Arc::new(pipeline) as std::sync::Arc) + .map_err(|error| error.to_string()); + } + + let toolkit = source + .toolkit + .as_deref() + .map(str::trim) + .filter(|toolkit| !toolkit.is_empty()) + .ok_or_else(|| "composio source missing toolkit".to_string())? + .to_ascii_lowercase(); + let connection_id = source + .connection_id + .as_deref() + .map(str::trim) + .filter(|connection_id| !connection_id.is_empty()) + .ok_or_else(|| "composio source missing connection_id".to_string())?; + let composio = composio_config(config)?; + memory_config.sync.composio = Some(composio.clone()); + let client = ComposioClient::new(composio); + let pipeline: std::sync::Arc = match toolkit.as_str() { + "gmail" => std::sync::Arc::new(GmailSyncPipeline::new(client, connection_id)), + "github" => std::sync::Arc::new(GitHubSyncPipeline::new(client, connection_id)), + "notion" => std::sync::Arc::new(NotionSyncPipeline::new(client, connection_id)), + "linear" => std::sync::Arc::new(LinearSyncPipeline::new(client, connection_id)), + "clickup" => std::sync::Arc::new(ClickUpSyncPipeline::new(client, connection_id)), + "slack" => std::sync::Arc::new(SlackSyncPipeline::new(client, connection_id)), + _ => { + return Err(format!( + "tinycortex sync does not support toolkit '{toolkit}'" + )) + } + }; + Ok(pipeline) +} + +fn composio_config( + config: &Config, +) -> Result { + use tinycortex::memory::config::{ComposioMode, ComposioSyncConfig, SecretString}; + + if config.composio.mode.eq_ignore_ascii_case("direct") { + let api_key = crate::openhuman::credentials::get_composio_api_key(config)? + .or_else(|| config.composio.api_key.clone()) + .ok_or_else(|| "Composio direct API key is not configured".to_string())?; + Ok(ComposioSyncConfig { + mode: ComposioMode::Direct, + base_url: "https://backend.composio.dev/api/v3".into(), + api_key: Some(SecretString::new(api_key)), + bearer_token: None, + entity_id: Some(config.composio.entity_id.clone()), + }) + } else { + let bearer = crate::api::jwt::get_session_token(config)? + .ok_or_else(|| "OpenHuman backend bearer token is not configured".to_string())?; + Ok(ComposioSyncConfig { + mode: ComposioMode::Proxied, + base_url: crate::api::config::effective_backend_api_url(&config.api_url), + api_key: None, + bearer_token: Some(SecretString::new(bearer)), + entity_id: Some(config.composio.entity_id.clone()), + }) + } +} + +#[async_trait] +impl SkillDocSink for HostSyncAdapter { + async fn store(&self, document: SkillDocument) -> anyhow::Result<()> { + tracing::debug!( + toolkit = %document.toolkit, + connection_id = %document.connection_id, + document_id = %document.document_id, + "[tinycortex:sync] storing synchronized document" + ); + self.memory + .store_skill_sync( + &document.namespace_skill_id, + &document.connection_id, + &document.title, + &document.content, + Some("tinycortex-sync".into()), + Some(document.metadata), + Some("medium".into()), + None, + None, + Some(document.document_id), + ) + .await + .map_err(anyhow::Error::msg) + } + + async fn delete(&self, namespace_skill_id: &str, document_id: &str) -> anyhow::Result<()> { + let namespace = format!("skill-{}", namespace_skill_id.trim()); + tracing::debug!( + namespace, + document_id, + "[tinycortex:sync] deleting synchronized document" + ); + self.memory + .delete_document(&namespace, document_id) + .await + .map(|_| ()) + .map_err(anyhow::Error::msg) + } +} + +#[async_trait] +impl LocalDocumentSink for HostSyncAdapter { + async fn upsert(&self, document: LocalDocument) -> anyhow::Result<()> { + let config = self + .config + .as_ref() + .ok_or_else(|| anyhow::anyhow!("local document sink missing host config"))?; + let input = crate::openhuman::memory_sync::canonicalize::document::DocumentInput { + provider: "memory_sources:local".into(), + title: document.title, + body: document.body, + modified_at: document.modified_at, + source_ref: document.source_ref, + }; + crate::openhuman::memory::ingest_pipeline::ingest_document_with_scope( + config, + &document.source_id, + &document.owner, + document.tags, + input, + document.path_scope, + ) + .await + .map(|_| ()) + .map_err(anyhow::Error::msg) + } + + async fn delete(&self, source_id: &str) -> anyhow::Result<()> { + let config = self + .config + .clone() + .ok_or_else(|| anyhow::anyhow!("local document sink missing host config"))?; + let source_id = source_id.to_owned(); + tokio::task::spawn_blocking(move || { + crate::openhuman::memory_store::chunks::store::delete_chunks_by_source( + &config, + crate::openhuman::memory_store::chunks::types::SourceKind::Document, + &source_id, + ) + }) + .await + .map_err(|error| anyhow::anyhow!("local delete task failed: {error}"))??; + Ok(()) + } +} + +#[async_trait] +impl SyncStateStore for HostSyncAdapter { + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { + self.memory + .kv_get(Some(namespace), key) + .await + .map_err(anyhow::Error::msg) + } + + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.memory + .kv_set(Some(namespace), key, value) + .await + .map_err(anyhow::Error::msg) + } +} + +#[async_trait] +impl SyncEventSink for HostSyncAdapter { + async fn emit(&self, event: SyncEvent) -> anyhow::Result<()> { + crate::core::event_bus::publish_global( + crate::core::event_bus::DomainEvent::MemorySyncStageChanged { + trigger: "tinycortex".into(), + stage: stage_name(event.stage).into(), + provider: Some(event.toolkit), + connection_id: event.connection_id, + detail: event.message, + source_id: Some(event.source_id), + }, + ); + Ok(()) + } +} + +fn stage_name(stage: SyncStage) -> &'static str { + match stage { + SyncStage::Requested => "requested", + SyncStage::Fetching => "fetching", + SyncStage::Stored => "stored", + SyncStage::Ingesting => "ingesting", + SyncStage::Completed => "completed", + SyncStage::Failed => "failed", + } +} diff --git a/tests/embeddings_rpc_e2e.rs b/tests/embeddings_rpc_e2e.rs index fb26c2a48..444e72658 100644 --- a/tests/embeddings_rpc_e2e.rs +++ b/tests/embeddings_rpc_e2e.rs @@ -579,11 +579,11 @@ async fn embeddings_embed_with_none_returns_empty_vectors() { let result = assert_no_rpc_error(&resp, "embeddings_embed none"); let inner = result.get("result").unwrap_or(result); - // NoopEmbedding.embed() returns an empty vec — count and dimensions should both be 0. + // The inert TinyAgents adapter preserves input cardinality with empty vectors. assert_eq!( inner.get("count").and_then(Value::as_u64), - Some(0), - "noop provider should return count=0: {inner}" + Some(2), + "noop provider should preserve input count: {inner}" ); assert_eq!( inner.get("dimensions").and_then(Value::as_u64), @@ -594,10 +594,7 @@ async fn embeddings_embed_with_none_returns_empty_vectors() { .get("vectors") .and_then(Value::as_array) .expect("embed result must include 'vectors' array: {inner}"); - assert!( - vectors.is_empty(), - "noop provider must return empty vectors array: {vectors:?}" - ); + assert_eq!(vectors, &vec![json!([]), json!([])]); } #[tokio::test(flavor = "multi_thread")] diff --git a/tests/json_rpc_e2e.rs b/tests/json_rpc_e2e.rs index e25ecc7db..90668c326 100644 --- a/tests/json_rpc_e2e.rs +++ b/tests/json_rpc_e2e.rs @@ -12188,13 +12188,6 @@ mod task_sources_stub { ) -> Result { Ok(ProviderUserProfile::default()) } - async fn sync( - &self, - _ctx: &ProviderContext, - _reason: SyncReason, - ) -> Result { - Ok(SyncOutcome::default()) - } async fn fetch_tasks( &self, _ctx: &ProviderContext, diff --git a/tests/memory_artifacts_e2e.rs b/tests/memory_artifacts_e2e.rs index a0a5e456f..3543cd875 100644 --- a/tests/memory_artifacts_e2e.rs +++ b/tests/memory_artifacts_e2e.rs @@ -9,16 +9,17 @@ use tempfile::tempdir; use chrono::{TimeZone, Utc}; use openhuman_core::openhuman::config::Config; +use openhuman_core::openhuman::memory::ingest_pipeline::ingest_chat; use openhuman_core::openhuman::memory::tree_source::registry::get_or_create_source_tree; use openhuman_core::openhuman::memory_queue::drain_until_idle; use openhuman_core::openhuman::memory_store::content::atomic::stage_summary; use openhuman_core::openhuman::memory_store::content::obsidian::ensure_obsidian_defaults; +use openhuman_core::openhuman::memory_store::content::raw::{write_raw_items, RawItem, RawKind}; use openhuman_core::openhuman::memory_store::content::wiki_git::{ get_read_pointer_tag, set_read_pointer_tag, }; use openhuman_core::openhuman::memory_store::content::{SummaryComposeInput, SummaryTreeKind}; -use openhuman_core::openhuman::memory_sync::composio::providers::slack::ingest::ingest_page_into_memory_tree; -use openhuman_core::openhuman::memory_sync::composio::providers::slack::SlackMessage; +use openhuman_core::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; use openhuman_core::openhuman::memory_tree::ingest::{ingest_summary, SummaryIngestInput}; fn make_config(workspace_dir: &std::path::Path) -> Config { @@ -35,21 +36,36 @@ async fn sync_raw_artifacts_and_mocked_summary_match_obsidian_contract() { let config = make_config(&workspace_dir); let ts = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let msg = SlackMessage { - channel_id: "C123".into(), - channel_name: "engineering".into(), - is_private: false, - author: "alice".into(), - author_id: "U123".into(), - text: "Phoenix migration launch window is Friday at 22:00 UTC.".into(), - timestamp: ts, - ts_raw: "1700000000.000100".into(), - thread_ts: None, - permalink: Some("https://slack.example.test/archives/C123/p1700000000000100".into()), + write_raw_items( + &config.memory_tree_content_root(), + "slack:conn-slack-1", + &[RawItem { + uid: "1700000000.000100", + created_at_ms: ts.timestamp_millis(), + markdown: "**Channel:** #engineering\n**Author:** alice\n\nPhoenix migration launch window is Friday at 22:00 UTC.", + kind: RawKind::Chat, + }], + ) + .expect("seed raw Slack artifact"); + let batch = ChatBatch { + platform: "slack".into(), + channel_label: "#engineering".into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp: ts, + text: "Phoenix migration launch window is Friday at 22:00 UTC.".into(), + source_ref: Some("https://slack.example.test/archives/C123/p1700000000000100".into()), + }], }; - ingest_page_into_memory_tree(&config, "alice", "conn-slack-1", &[msg]) - .await - .expect("seed slack sync"); + ingest_chat( + &config, + "slack:conn-slack-1", + "alice", + vec!["slack".into(), "ingested".into()], + batch, + ) + .await + .expect("seed slack sync"); drain_until_idle(&config).await.expect("drain sync jobs"); let content_root = config.memory_tree_content_root(); diff --git a/tests/raw_coverage/embeddings_ollama_raw_coverage_e2e.rs b/tests/raw_coverage/embeddings_ollama_raw_coverage_e2e.rs index 442535735..e0e9840af 100644 --- a/tests/raw_coverage/embeddings_ollama_raw_coverage_e2e.rs +++ b/tests/raw_coverage/embeddings_ollama_raw_coverage_e2e.rs @@ -270,7 +270,7 @@ async fn openai_embed_reports_response_validation_and_http_errors() { .await .expect_err("missing embedding") .to_string() - .contains("missing 'embedding'")); + .contains("missing `embedding`")); let (dim_url, _) = serve_mock_openai(OpenAiMockBehavior::DimensionMismatch).await; let dim_provider = OpenAiEmbedding::new(&dim_url, "k", "m", 2); @@ -288,7 +288,7 @@ async fn openai_embed_reports_response_validation_and_http_errors() { .await .expect_err("missing data") .to_string() - .contains("missing 'data'")); + .contains("missing `data`")); let (non_2xx_url, _) = serve_mock_openai(OpenAiMockBehavior::Non2xx).await; let non_2xx_provider = OpenAiEmbedding::new(&non_2xx_url, "k", "m", 2); @@ -297,7 +297,7 @@ async fn openai_embed_reports_response_validation_and_http_errors() { .await .expect_err("http error") .to_string() - .contains("Embedding API error")); + .contains("returned HTTP")); } #[tokio::test] @@ -716,14 +716,17 @@ async fn ollama_embed_preserves_positions_and_validates_request_and_response() { let base_url = serve_mock_ollama(app).await; let provider = OllamaEmbedding::try_new(&base_url, "mock-ollama", 2).expect("provider"); - let vectors = provider + let mixed_error = provider .embed(&[" alpha ", "", "beta", " "]) .await + .expect_err("mixed blank batch should be rejected"); + assert!(mixed_error.to_string().contains("must not mix blank")); + + let vectors = provider + .embed(&[" alpha ", "beta"]) + .await .expect("ollama embed"); - assert_eq!( - vectors, - vec![vec![1.0, 2.0], vec![], vec![3.0, 4.0], vec![]] - ); + assert_eq!(vectors, vec![vec![1.0, 2.0], vec![3.0, 4.0]]); let all_blank = provider.embed(&["", " \n\t "]).await.expect("blank embed"); assert_eq!(all_blank, vec![Vec::::new(), Vec::::new()]); @@ -811,12 +814,15 @@ async fn ollama_embed_recovers_nan_batch_with_per_text_fallback() { let base_url = serve_mock_ollama(app).await; let provider = OllamaEmbedding::try_new(&base_url, "mock-ollama", 2).expect("provider"); - let vectors = provider - .embed(&["good", "bad", " "]) + let batch_error = provider + .embed(&["good", "bad"]) .await - .expect("nan batch recovery"); - assert_eq!(vectors, vec![vec![9.0, 8.0], vec![], vec![]]); + .expect_err("a NaN-producing item should fail the recovered batch"); + assert!(batch_error.to_string().contains("without NaN values")); - let single_nan = provider.embed(&["bad"]).await.expect("single nan recovery"); - assert_eq!(single_nan, vec![Vec::::new()]); + let single_error = provider + .embed(&["bad"]) + .await + .expect_err("single NaN-producing input should fail"); + assert!(single_error.to_string().contains("without NaN values")); } diff --git a/tests/raw_coverage/memory_core_threads_raw_coverage_e2e.rs b/tests/raw_coverage/memory_core_threads_raw_coverage_e2e.rs index fabd05020..76976bd9b 100644 --- a/tests/raw_coverage/memory_core_threads_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_core_threads_raw_coverage_e2e.rs @@ -367,7 +367,10 @@ async fn memory_read_rpc_filters_graphs_scores_reset_and_wipe_seeded_rows() { .value .unwrap(); assert!(score.kept); - assert!(score.llm_consulted); + assert!( + !score.llm_consulted, + "TinyCortex does not persist admission-time LLM importance" + ); let tree_graph = read_rpc::graph_export_rpc(&cfg, GraphMode::Tree) .await diff --git a/tests/raw_coverage/memory_sync_providers_raw_coverage_e2e.rs b/tests/raw_coverage/memory_sync_providers_raw_coverage_e2e.rs index f49881ae7..fe893a08b 100644 --- a/tests/raw_coverage/memory_sync_providers_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_sync_providers_raw_coverage_e2e.rs @@ -20,32 +20,19 @@ use openhuman_core::openhuman::credentials::{ }; use openhuman_core::openhuman::memory::global as memory_global; use openhuman_core::openhuman::memory::jobs::drain_until_idle; -use openhuman_core::openhuman::memory::tree_source::get_or_create_source_tree; -use openhuman_core::openhuman::memory_store::chunks::store::{ - count_chunks_by_lifecycle_status, get_chunk_raw_refs, list_chunks, ListChunksQuery, - CHUNK_STATUS_BUFFERED, -}; -use openhuman_core::openhuman::memory_store::chunks::types::SourceKind; -use openhuman_core::openhuman::memory_store::content::read::read_chunk_body; -use openhuman_core::openhuman::memory_store::trees::store as tree_store; use openhuman_core::openhuman::memory_sync::composio::bus::{ ComposioConfigChangedSubscriber, ComposioConnectionCreatedSubscriber, ComposioTriggerSubscriber, }; use openhuman_core::openhuman::memory_sync::composio::providers::clickup::ClickUpProvider; use openhuman_core::openhuman::memory_sync::composio::providers::github::GitHubProvider; -use openhuman_core::openhuman::memory_sync::composio::providers::gmail::ingest as gmail_ingest; use openhuman_core::openhuman::memory_sync::composio::providers::gmail::GmailProvider; use openhuman_core::openhuman::memory_sync::composio::providers::linear::LinearProvider; use openhuman_core::openhuman::memory_sync::composio::providers::notion::NotionProvider; -use openhuman_core::openhuman::memory_sync::composio::providers::slack::ingest as slack_ingest; -use openhuman_core::openhuman::memory_sync::composio::providers::slack::{ - SlackMessage, SlackProvider, -}; +use openhuman_core::openhuman::memory_sync::composio::providers::slack::SlackProvider; use openhuman_core::openhuman::memory_sync::composio::providers::sync_state::SyncState; use openhuman_core::openhuman::memory_sync::composio::providers::{ ComposioProvider, ProviderContext, SyncReason, TaskFetchFilter, }; -use openhuman_core::openhuman::memory_tree::tree::bucket_seal::LabelStrategy; static ENV_LOCK: &OnceLock> = &crate::SHARED_ENV_LOCK; @@ -320,329 +307,6 @@ async fn configured_loopback_context( (config, ctx, server) } -#[tokio::test] -async fn gmail_ingest_archives_account_messages_and_legacy_participant_buckets() { - let _guard = env_lock(); - let tmp = TempDir::new().expect("tempdir"); - let config = config_in(&tmp); - let _workspace = EnvGuard::set_path("OPENHUMAN_WORKSPACE", tmp.path()); - let _home = EnvGuard::set_path("HOME", tmp.path()); - let _backend = EnvGuard::unset("BACKEND_URL"); - persist_config(&config).await; - - let page = vec![ - json!({ - "id": "gmail-round17-a", - "from": "Ava ", - "to": ["Ben ", "Casey "], - "cc": "ignored@example.test", - "subject": "Re: Coverage thread", - "date": "2026-05-29T10:00:00Z", - "markdown": "First useful message body." - }), - json!({ - "id": "gmail-round17-b", - "from": "ben@example.test", - "to": "ava@example.test, casey@example.test", - "subject": "Fwd: Coverage thread", - "internalDate": "1780052400000", - "markdown": "Second useful message body." - }), - json!({ - "id": "gmail-round17-empty", - "from": "nobody@example.test", - "to": "ava@example.test", - "subject": "No archive body", - "date": "2026-05-29T12:00:00Z", - "markdown": " " - }), - json!({ - "from": "missing-id@example.test", - "to": "ava@example.test", - "subject": "No id", - "date": "2026-05-29T13:00:00Z", - "markdown": "missing id skips per-account ingest" - }), - ]; - - let chunks = gmail_ingest::ingest_page_into_memory_tree( - &config, - "owner-round17", - Some("round17@example.test"), - &page, - ) - .await - .expect("per-account gmail ingest"); - assert!(chunks >= 2, "expected useful account messages to chunk"); - - let raw_root = config.memory_tree_content_root().join("raw"); - let archived: Vec<_> = walk_files(&raw_root) - .into_iter() - .filter(|p| { - p.file_name() - .and_then(|name| name.to_str()) - .is_some_and(|name| name.contains("gmail-round17-a")) - }) - .collect(); - assert_eq!(archived.len(), 1, "raw archive should include message a"); - let archived_body = std::fs::read_to_string(&archived[0]).expect("archived body"); - assert!(archived_body.contains("**From:** Ava")); - assert!(archived_body.contains("First useful message body.")); - - let legacy = gmail_ingest::ingest_page_into_memory_tree( - &config, - "owner-round17", - None, - &[ - json!({ - "id": "legacy-orphan", - "from": "not an address", - "to": [], - "subject": "Fw: ", - "date": "2026-05-29T14:00:00Z", - "markdown": "orphan fallback body" - }), - json!({ - "from": "", - "to": [], - "subject": "Skipped", - "date": "2026-05-29T15:00:00Z", - "markdown": "no id and no participants" - }), - ], - ) - .await - .expect("legacy gmail ingest"); - assert!(legacy >= 1, "orphan fallback bucket should ingest"); -} - -#[tokio::test] -async fn gmail_raw_backed_messages_drain_into_source_tree_summary() { - let _guard = env_lock(); - let tmp = TempDir::new().expect("tempdir"); - let _workspace = EnvGuard::set_path("OPENHUMAN_WORKSPACE", tmp.path()); - let _home = EnvGuard::set_path("HOME", tmp.path()); - let _backend = EnvGuard::unset("BACKEND_URL"); - let config = config_in(&tmp); - persist_config(&config).await; - - let page = vec![ - json!({ - "id": "gmail-tree-a", - "from": "Ava ", - "to": ["Ben "], - "subject": "Phoenix launch plan", - "date": "2026-05-29T10:00:00Z", - "markdown": "Phoenix migration launches Friday. Ava owns rollout validation and Ben owns customer notices." - }), - json!({ - "id": "gmail-tree-b", - "from": "Ben ", - "to": ["Ava "], - "subject": "Re: Phoenix launch plan", - "date": "2026-05-29T10:05:00Z", - "markdown": "Confirmed. Customer notices go out after staging checks and the rollback doc is reviewed." - }), - ]; - - let outcome = gmail_ingest::ingest_page_into_memory_tree_with_outcome( - &config, - "owner-gmail-tree", - Some("flow@example.test"), - &page, - ) - .await - .expect("gmail ingest outcome"); - - assert_eq!( - outcome.item_ids_ingested, - vec!["gmail-tree-a".to_string(), "gmail-tree-b".to_string()] - ); - assert!( - outcome.chunks_written >= 2, - "expected one or more chunks per message" - ); - - let source_id = "gmail:flow-at-example-dot-test"; - let chunks = list_chunks( - &config, - &ListChunksQuery { - source_kind: Some(SourceKind::Email), - source_id: Some(source_id.to_string()), - limit: Some(10), - ..Default::default() - }, - ) - .expect("list gmail chunks"); - assert_eq!(chunks.len(), outcome.chunks_written); - - for chunk in &chunks { - let refs = get_chunk_raw_refs(&config, &chunk.id) - .expect("raw refs lookup") - .expect("raw refs must be set before extract can run"); - assert_eq!(refs.len(), 1); - assert!( - refs[0].path.contains("gmail-tree-"), - "raw ref should point at the source message file: {:?}", - refs[0].path - ); - let full_body = read_chunk_body(&config, &chunk.id).expect("read raw-backed chunk body"); - assert!( - full_body.contains("Phoenix") || full_body.contains("Customer notices"), - "chunk body should hydrate from raw archive, got: {full_body}" - ); - } - - drain_until_idle(&config) - .await - .expect("extract and append jobs should drain"); - let buffered = - count_chunks_by_lifecycle_status(&config, CHUNK_STATUS_BUFFERED).expect("buffered count"); - assert_eq!(buffered, outcome.chunks_written as u64); - - let tree = get_or_create_source_tree(&config, source_id).expect("source tree"); - let l0 = tree_store::get_buffer(&config, &tree.id, 0).expect("source L0 buffer"); - assert_eq!( - l0.item_ids.len(), - outcome.chunks_written, - "all Gmail chunks should reach the source tree buffer" - ); - - let sealed = openhuman_core::openhuman::memory_tree::tree::flush::flush_stale_buffers( - &config, - chrono::Duration::zero(), - &LabelStrategy::Empty, - ) - .await - .expect("force flush gmail source tree"); - assert!(sealed > 0, "low-volume Gmail source should seal on flush"); - - let l1 = - tree_store::list_summaries_at_level(&config, &tree.id, 1).expect("list source summaries"); - assert!( - !l1.is_empty(), - "Gmail source tree should have a sealed summary after flush" - ); - let summary_body = openhuman_core::openhuman::memory_store::content::read::read_summary_body( - &config, &l1[0].id, - ) - .expect("read gmail summary body"); - assert!( - summary_body.contains("Phoenix") || summary_body.contains("Customer"), - "summary should preserve Gmail content, got: {summary_body}" - ); -} - -#[tokio::test] -async fn slack_provider_profile_postprocess_trigger_and_ingest_use_loopback_composio() { - let _guard = env_lock(); - let tmp = TempDir::new().expect("tempdir"); - let _workspace = EnvGuard::set_path("OPENHUMAN_WORKSPACE", tmp.path()); - let _home = EnvGuard::set_path("HOME", tmp.path()); - let _backend = EnvGuard::unset("BACKEND_URL"); - let requests: Arc>> = Arc::new(Mutex::new(Vec::new())); - let (config, ctx, server) = - configured_loopback_context(&tmp, "slack", "conn-slack-round17", Arc::clone(&requests)) - .await; - - let provider = SlackProvider::new(); - let profile = provider - .fetch_user_profile(&ctx) - .await - .expect("slack profile"); - assert_eq!(profile.username.as_deref(), Some("U17A")); - assert_eq!(profile.email.as_deref(), Some("round17@example.test")); - - let mut channels = json!({ - "data": { - "channels": [ - { "id": "C17", "name": "coverage", "is_private": false }, - { "id": "", "name": "dropped" } - ] - } - }); - provider.post_process_action_result("SLACK_LIST_CONVERSATIONS", None, &mut channels); - assert_eq!(channels["channels"].as_array().unwrap().len(), 1); - - let mut history = json!({ - "data": { - "messages": [ - { - "ts": "1714003200.000100", - "user": "U17A", - "text": "shipping coverage with <@U17B>", - "permalink": "https://coverage.slack.com/archives/C17/p1714003200000100" - }, - { "ts": "1714003300.000200", "user": "U17B", "text": " " } - ] - } - }); - provider.post_process_action_result("SLACK_FETCH_CONVERSATION_HISTORY", None, &mut history); - assert_eq!(history["messages"].as_array().unwrap().len(), 1); - - provider - .on_trigger( - &ctx, - "SLACK_CHANNEL_ARCHIVE", - &json!({ "event": "channel" }), - ) - .await - .expect("slack non-message trigger"); - - let messages = vec![ - SlackMessage { - channel_id: "C17".to_string(), - channel_name: "coverage".to_string(), - is_private: false, - author: "Ava".to_string(), - author_id: "U17A".to_string(), - text: "Slack raw archive body".to_string(), - timestamp: chrono::DateTime::parse_from_rfc3339("2026-05-29T10:00:00Z") - .unwrap() - .with_timezone(&chrono::Utc), - ts_raw: "1714003200.000100".to_string(), - thread_ts: Some("1714003200.000100".to_string()), - permalink: Some( - "https://coverage.slack.com/archives/C17/p1714003200000100".to_string(), - ), - }, - SlackMessage { - channel_id: "G17".to_string(), - channel_name: "private-coverage".to_string(), - is_private: true, - author: String::new(), - author_id: String::new(), - text: " ".to_string(), - timestamp: chrono::DateTime::parse_from_rfc3339("2026-05-29T10:01:00Z") - .unwrap() - .with_timezone(&chrono::Utc), - ts_raw: "1714003260.000200".to_string(), - thread_ts: None, - permalink: None, - }, - ]; - let chunks = slack_ingest::ingest_page_into_memory_tree( - &config, - "owner-round17", - "conn-slack-round17", - &messages, - ) - .await - .expect("slack ingest"); - assert!(chunks >= 1); - - let called_tools: Vec = requests - .lock() - .unwrap() - .iter() - .filter_map(|b| b.get("tool").and_then(Value::as_str).map(str::to_string)) - .collect(); - assert!(called_tools.contains(&"SLACK_TEST_AUTH".to_string())); - assert!(called_tools.contains(&"SLACK_RETRIEVE_DETAILED_USER_INFORMATION".to_string())); - - server.abort(); -} - #[tokio::test] async fn github_clickup_and_composio_bus_cover_provider_branches() { let _guard = env_lock(); @@ -902,7 +566,7 @@ async fn slack_sync_max_items_caps_ingest_to_exact_count() { /// Build M distinct, valid Gmail messages in the upstream (pre-post-process) /// Composio shape. Uses `messageId` / `sender` / `messageText` / /// `messageTimestamp` so `reshape_message` maps them correctly into the slim -/// envelope that `ingest_page_into_memory_tree` expects. +/// envelope consumed by the tinycortex Gmail sync pipeline. fn gmail_cap_messages(m: usize) -> Vec { (1..=m) .map(|i| { @@ -1110,8 +774,10 @@ async fn gmail_sync_stops_after_an_all_already_synced_page() { for i in 1..=3 { state.mark_synced(format!("gmail-cap-msg-{i}")); } + let state_adapter = + openhuman_core::openhuman::tinycortex::HostSyncAdapter::new(memory.clone()); state - .save(&memory) + .save(&state_adapter) .await .expect("save pre-seeded sync state"); diff --git a/tests/raw_coverage/memory_sync_round23_raw_coverage_e2e.rs b/tests/raw_coverage/memory_sync_round23_raw_coverage_e2e.rs index 8f0e2dfc5..567b0e316 100644 --- a/tests/raw_coverage/memory_sync_round23_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_sync_round23_raw_coverage_e2e.rs @@ -349,9 +349,9 @@ async fn notion_profile_prefers_bot_owner_and_sync_paginates_into_memory_tree() .await .expect("notion sync"); assert_eq!(outcome.items_ingested, 2); - assert!(outcome.summary.contains("fetched 3, persisted 2")); - assert_eq!(outcome.details["results_fetched"], 3); - assert_eq!(outcome.details["results_persisted"], 2); + assert_eq!(outcome.summary, "sync completed"); + assert_eq!(outcome.details["more_pending"], false); + assert_eq!(outcome.details["actions_called"], 4); let calls = requests.lock().unwrap().clone(); let fetch_calls: Vec = calls @@ -360,7 +360,7 @@ async fn notion_profile_prefers_bot_owner_and_sync_paginates_into_memory_tree() .cloned() .collect(); assert_eq!(fetch_calls.len(), 2); - assert_eq!(fetch_calls[0]["arguments"]["page_size"], 50); + assert_eq!(fetch_calls[0]["arguments"]["page_size"], 25); assert_eq!(fetch_calls[1]["arguments"]["start_cursor"], "page-2"); server.abort(); } diff --git a/tests/raw_coverage/memory_sync_slack_bus_raw_coverage_e2e.rs b/tests/raw_coverage/memory_sync_slack_bus_raw_coverage_e2e.rs index 0e8df522a..efaeeff8c 100644 --- a/tests/raw_coverage/memory_sync_slack_bus_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_sync_slack_bus_raw_coverage_e2e.rs @@ -6,7 +6,7 @@ //! OPENHUMAN_WORKSPACE, and config loading are process globals. use std::collections::HashMap; -use std::path::{Path, PathBuf}; +use std::path::Path; use std::sync::{Arc, Mutex, OnceLock}; use axum::routing::any; @@ -270,17 +270,15 @@ async fn configured_loopback_context( async fn slack_full_sync_search_backfill_and_bus_use_loopback_composio() { let _guard = env_lock(); let tmp = TempDir::new().expect("tempdir"); - let dump_dir = tmp.path().join("slack-dumps"); let _workspace = EnvGuard::set_path("OPENHUMAN_WORKSPACE", tmp.path()); let _home = EnvGuard::set_path("HOME", tmp.path()); let _backend = EnvGuard::unset("BACKEND_URL"); let _triage_off = EnvGuard::set("OPENHUMAN_TRIGGER_TRIAGE_DISABLED", "1"); let _pacing = EnvGuard::set("OPENHUMAN_SLACK_INTER_CALL_PACING_MS", "0"); let _backfill = EnvGuard::set("OPENHUMAN_SLACK_BACKFILL_DAYS", "1"); - let _dump = EnvGuard::set_path("OPENHUMAN_SLACK_DUMP_DIR", &dump_dir); let requests: Arc>> = Arc::new(Mutex::new(Vec::new())); - let (config, ctx, server) = configured_loopback_context(&tmp, Arc::clone(&requests)).await; + let (_config, ctx, server) = configured_loopback_context(&tmp, Arc::clone(&requests)).await; let provider = SlackProvider::new(); let profile = provider @@ -298,43 +296,26 @@ async fn slack_full_sync_search_backfill_and_bus_use_loopback_composio() { .expect("slack full sync"); assert_eq!(outcome.toolkit, "slack"); assert_eq!(outcome.connection_id.as_deref(), Some("conn-slack-round19")); - assert_eq!(outcome.items_ingested, 4); - // Slack now rides the generic orchestrator: two channels synced cleanly, - // none errored. (`channels_processed` → orchestrator's `scopes_synced`.) - assert_eq!(outcome.details["scopes_synced"], 2); - assert_eq!(outcome.details["scopes_errored"], 0); + assert_eq!(outcome.items_ingested, 3); + assert_eq!(outcome.details["more_pending"], false); + assert_eq!(outcome.details["actions_called"], 6); let search = run_backfill_via_search(&ctx, 2) .await .expect("slack search backfill"); assert_eq!(search.items_ingested, 1); - assert_eq!(search.details["channels_flushed"], 1); - assert_eq!(search.details["channels_failed"], 0); + assert_eq!(search.details["more_pending"], false); + assert_eq!(search.details["actions_called"], 5); - let raw_root = config.memory_tree_content_root().join("raw"); - let raw_files = walk_files(&raw_root); - let raw_bodies = raw_files - .iter() - .filter_map(|path| std::fs::read_to_string(path).ok()) - .collect::>() - .join("\n"); - assert!(raw_bodies.contains("shipping Slack sync coverage with @Ben Round19")); - assert!(raw_bodies.contains("private sync note for @Ben Round19")); - assert!(raw_bodies.contains("search backfill hit for @Ava Round19")); - assert!(raw_bodies.contains("**Channel:** #coverage")); - assert!(raw_bodies.contains("**Channel:** private:private-coverage")); - - let dumped = walk_files(&dump_dir); - assert!( - dumped.iter().any(|p| p.to_string_lossy().contains("users")), - "user directory response should be dumped" - ); - assert!( - dumped - .iter() - .any(|p| p.to_string_lossy().contains("history")), - "history responses should be dumped" - ); + let documents = openhuman_core::openhuman::memory::ops::doc_list(Some( + openhuman_core::openhuman::memory::ops::NamespaceOnlyParams { + namespace: "skill-slack".into(), + }, + )) + .await + .expect("list synchronized Slack documents") + .value; + assert_eq!(documents["documents"].as_array().unwrap().len(), 4); let trigger_sub = ComposioTriggerSubscriber::new(); assert_eq!(trigger_sub.name(), "composio::trigger"); @@ -456,26 +437,3 @@ async fn gmail_post_process_reshapes_nested_messages_and_honors_raw_html_flag() assert_eq!(raw_passthrough["messages"][0]["messageId"], "raw-round19"); assert!(raw_passthrough["messages"][0].get("markdown").is_none()); } - -fn walk_files(root: &Path) -> Vec { - let mut out = Vec::new(); - if !root.exists() { - return out; - } - let mut stack = vec![root.to_path_buf()]; - while let Some(path) = stack.pop() { - let entries = match std::fs::read_dir(&path) { - Ok(entries) => entries, - Err(_) => continue, - }; - for entry in entries.flatten() { - let child = entry.path(); - if child.is_dir() { - stack.push(child); - } else { - out.push(child); - } - } - } - out -} diff --git a/tests/raw_coverage/memory_sync_tree_round21_raw_coverage_e2e.rs b/tests/raw_coverage/memory_sync_tree_round21_raw_coverage_e2e.rs index 4f1a23434..538d652b9 100644 --- a/tests/raw_coverage/memory_sync_tree_round21_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_sync_tree_round21_raw_coverage_e2e.rs @@ -35,7 +35,7 @@ use openhuman_core::openhuman::memory_sync::composio::providers::{ ComposioProvider, ProviderContext, SyncReason, TaskFetchFilter, }; use openhuman_core::openhuman::memory_tree::retrieval::source::query_source; -use openhuman_core::openhuman::memory_tree::score::embed::{pack_embedding, EMBEDDING_DIM}; +use tinycortex::memory::score::embed::{pack_embedding, EMBEDDING_DIM}; use openhuman_core::openhuman::memory_tree::tree::store as tree_store; use openhuman_core::openhuman::memory_tree::tree::TreeStatus; @@ -358,16 +358,16 @@ async fn linear_provider_profile_tasks_sync_and_periodic_bookkeeping_use_loopbac .await .expect("linear sync"); assert_eq!(sync.items_ingested, 4); - assert_eq!(sync.details["issues_fetched"], 4); - assert_eq!(sync.details["issues_persisted"], 4); - assert_eq!(sync.details["cursor"], "2026-05-30T10:00:00.000Z"); + assert_eq!(sync.details["more_pending"], false); + assert_eq!(sync.details["actions_called"], 3); let second = provider .sync(&ctx, SyncReason::Manual) .await .expect("second sync"); assert_eq!(second.items_ingested, 0); - assert_eq!(second.details["issues_persisted"], 0); + assert_eq!(second.details["more_pending"], false); + assert_eq!(second.details["actions_called"], 3); record_sync_success("linear", "conn-linear-round21"); record_sync_success("linear", "conn-linear-round21"); @@ -422,7 +422,12 @@ async fn slack_sync_status_rpc_reads_mock_connections_and_persisted_state() { state.advance_cursor(r#"{"C21":"1714003200.000100"}"#); state.mark_synced("C21:1714003200.000100"); state.record_requests(7); - state.save(&memory).await.expect("save slack sync state"); + let state_adapter = + openhuman_core::openhuman::tinycortex::HostSyncAdapter::new(memory.clone()); + state + .save(&state_adapter) + .await + .expect("save slack sync state"); let outcome = sync_status_rpc(&config, SyncStatusRequest::default()) .await diff --git a/tests/raw_coverage/memory_threads_raw_coverage_e2e.rs b/tests/raw_coverage/memory_threads_raw_coverage_e2e.rs index 42d838c5c..ce3f47801 100644 --- a/tests/raw_coverage/memory_threads_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_threads_raw_coverage_e2e.rs @@ -1876,8 +1876,8 @@ async fn memory_read_rpc_score_index_and_summary_helpers_cover_dashboard_paths() .value .expect("score breakdown"); assert!(breakdown.kept); - assert!(breakdown.llm_consulted); - assert!(breakdown + assert!(!breakdown.llm_consulted); + assert!(!breakdown .signals .iter() .any(|signal| signal.name == "llm_importance" && signal.weight == 2.0)); @@ -2999,22 +2999,6 @@ impl ComposioProvider for RawCoverageProvider { } } - async fn sync( - &self, - _ctx: &ProviderContext, - reason: SyncReason, - ) -> Result { - Ok(ComposioSyncOutcome { - toolkit: "raw_coverage".into(), - connection_id: Some("conn-1".into()), - reason: reason.as_str().into(), - items_ingested: 1, - started_at_ms: 10, - finished_at_ms: 25, - summary: "synced".into(), - details: json!({ "reason": reason.as_str() }), - }) - } } struct EmptySlugProvider; @@ -3032,22 +3016,6 @@ impl ComposioProvider for EmptySlugProvider { Ok(ProviderUserProfile::default()) } - async fn sync( - &self, - _ctx: &ProviderContext, - reason: SyncReason, - ) -> Result { - Ok(ComposioSyncOutcome { - toolkit: String::new(), - connection_id: None, - reason: reason.as_str().into(), - items_ingested: 0, - started_at_ms: 0, - finished_at_ms: 0, - summary: String::new(), - details: Value::Null, - }) - } } #[tokio::test] @@ -4767,7 +4735,8 @@ async fn memory_sources_types_registry_and_sync_state_cover_public_persistence_e MemoryClient::from_workspace_dir(tmp.path().join("memory-sync-state")) .expect("memory client"), ); - let fresh = SyncState::load(&memory, "gmail", "conn-raw") + let adapter = openhuman_core::openhuman::tinycortex::HostSyncAdapter::new(memory.clone()); + let fresh = SyncState::load(&adapter, "gmail", "conn-raw") .await .expect("fresh state"); assert_eq!(fresh.toolkit, "gmail"); @@ -4778,9 +4747,9 @@ async fn memory_sources_types_registry_and_sync_state_cover_public_persistence_e saved.mark_synced("msg-1"); saved.daily_budget.date = "2000-01-01".into(); saved.daily_budget.requests_used = DEFAULT_DAILY_REQUEST_LIMIT; - saved.save(&memory).await.expect("save state"); + saved.save(&adapter).await.expect("save state"); - let loaded = SyncState::load(&memory, "gmail", "conn-raw") + let loaded = SyncState::load(&adapter, "gmail", "conn-raw") .await .expect("load saved state"); assert_eq!(loaded.cursor.as_deref(), Some("cursor-raw")); @@ -4790,16 +4759,18 @@ async fn memory_sources_types_registry_and_sync_state_cover_public_persistence_e memory .kv_set( - Some("composio-sync-state"), - "gmail:bad-json", + Some(openhuman_core::openhuman::tinycortex::HOST_SYNC_STATE_NAMESPACE), + "composio-sync-state:gmail:bad-json", &json!("not a sync state"), ) .await .expect("write bad state"); - assert!(SyncState::load(&memory, "gmail", "bad-json") + let recovered = SyncState::load(&adapter, "gmail", "bad-json") .await - .unwrap_err() - .contains("deserialize failed")); + .expect("malformed state recovers to defaults"); + assert_eq!(recovered.toolkit, "gmail"); + assert_eq!(recovered.connection_id, "bad-json"); + assert!(recovered.cursor.is_none()); } #[test] diff --git a/tests/raw_coverage/memory_tree_sync_raw_coverage_e2e.rs b/tests/raw_coverage/memory_tree_sync_raw_coverage_e2e.rs index 2f20896cb..f956414a1 100644 --- a/tests/raw_coverage/memory_tree_sync_raw_coverage_e2e.rs +++ b/tests/raw_coverage/memory_tree_sync_raw_coverage_e2e.rs @@ -411,7 +411,7 @@ async fn composio_providers_sync_state_and_bus_surfaces_cover_read_write_edges() } #[tokio::test] -async fn default_composio_provider_hooks_return_expected_noop_shapes() { +async fn default_composio_provider_hooks_cover_defaults_and_sync_preconditions() { struct MinimalProvider; #[async_trait] @@ -447,22 +447,6 @@ async fn default_composio_provider_hooks_return_expected_noop_shapes() { }) } - async fn sync( - &self, - ctx: &ProviderContext, - reason: SyncReason, - ) -> Result { - Ok(SyncOutcome { - toolkit: ctx.toolkit.clone(), - connection_id: ctx.connection_id.clone(), - reason: reason.as_str().into(), - items_ingested: 3, - started_at_ms: 10, - finished_at_ms: 25, - summary: "synced".into(), - details: json!({"reason": reason.as_str()}), - }) - } } let tmp = TempDir::new().expect("tempdir"); @@ -499,11 +483,11 @@ async fn default_composio_provider_hooks_return_expected_noop_shapes() { let profile = provider.fetch_user_profile(&ctx).await.expect("profile"); assert_eq!(profile.email.as_deref(), Some("round14@example.com")); - let sync = provider + let sync_error = provider .sync(&ctx, SyncReason::Manual) .await - .expect("sync outcome"); - assert_eq!(sync.elapsed_ms(), 15); + .expect_err("sync requires an initialized memory client"); + assert!(sync_error.contains("memory client is not ready")); let extracted = ExtractedEntities { entities: vec![ diff --git a/vendor/tinyagents b/vendor/tinyagents index 19dc2c438..b34b2f83d 160000 --- a/vendor/tinyagents +++ b/vendor/tinyagents @@ -1 +1 @@ -Subproject commit 19dc2c438e1f8b2715a45ba7030bc455e611dcb2 +Subproject commit b34b2f83d7af492c9b7ccaf558c38bece4182dde diff --git a/vendor/tinycortex b/vendor/tinycortex index a8e10f7dd..671e78a01 160000 --- a/vendor/tinycortex +++ b/vendor/tinycortex @@ -1 +1 @@ -Subproject commit a8e10f7dd8ebdb9b0905e1380fefcc6bf5a65207 +Subproject commit 671e78a01411de5bcda8f3d1816ac6c67485d694