diff --git a/.gitignore b/.gitignore index 9933adf6a..2a185ac8b 100644 --- a/.gitignore +++ b/.gitignore @@ -66,6 +66,7 @@ tauri.key tauri.key.pub /target/ src-tauri/target/ +.target-codex/ workflow .fastembed_cache diff --git a/Cargo.lock b/Cargo.lock index 170f8578a..d48d5a324 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4541,7 +4541,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "openhuman" -version = "0.53.3" +version = "0.53.4" dependencies = [ "aes-gcm", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index 41af288b8..2f7f976bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,10 @@ path = "src/main.rs" name = "slack-backfill" path = "src/bin/slack_backfill.rs" +[[bin]] +name = "gmail-backfill-3d" +path = "src/bin/gmail_backfill_3d.rs" + [lib] name = "openhuman_core" crate-type = ["rlib"] diff --git a/docs/memory-tree-async-pipeline.excalidraw b/docs/memory-tree-async-pipeline.excalidraw new file mode 100644 index 000000000..29ee1a1ce --- /dev/null +++ b/docs/memory-tree-async-pipeline.excalidraw @@ -0,0 +1,1353 @@ +{ + "type": "excalidraw", + "version": 2, + "source": "openhuman-memory-tree-async-pipeline", + "elements": [ + { + "id": "title", + "type": "text", + "x": 355, + "y": 20, + "width": 740, + "height": 40, + "text": "OpenHuman Memory Tree Async Pipeline", + "fontSize": 30, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 32, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 1, + "version": 1, + "versionNonce": 1, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "subtitle", + "type": "text", + "x": 235, + "y": 64, + "width": 980, + "height": 24, + "text": "Leaf ingestion -> jobs queue -> workers -> source/topic/global tree building", + "fontSize": 18, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 18, + "strokeColor": "#495057", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 2, + "version": 1, + "versionNonce": 2, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane1", + "type": "rectangle", + "x": 40, + "y": 120, + "width": 310, + "height": 340, + "strokeColor": "#1971c2", + "backgroundColor": "#e7f5ff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 3, + "version": 1, + "versionNonce": 3, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane2", + "type": "rectangle", + "x": 390, + "y": 120, + "width": 340, + "height": 340, + "strokeColor": "#2b8a3e", + "backgroundColor": "#ebfbee", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 4, + "version": 1, + "versionNonce": 4, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane3", + "type": "rectangle", + "x": 770, + "y": 120, + "width": 390, + "height": 340, + "strokeColor": "#e67700", + "backgroundColor": "#fff4e6", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 5, + "version": 1, + "versionNonce": 5, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane4", + "type": "rectangle", + "x": 1200, + "y": 120, + "width": 440, + "height": 340, + "strokeColor": "#9c36b5", + "backgroundColor": "#f8f0fc", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 6, + "version": 1, + "versionNonce": 6, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane5", + "type": "rectangle", + "x": 40, + "y": 500, + "width": 760, + "height": 220, + "strokeColor": "#0b7285", + "backgroundColor": "#e3fafc", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 7, + "version": 1, + "versionNonce": 7, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane6", + "type": "rectangle", + "x": 840, + "y": 500, + "width": 800, + "height": 220, + "strokeColor": "#495057", + "backgroundColor": "#f1f3f5", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 8, + "version": 1, + "versionNonce": 8, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h1", + "type": "text", + "x": 135, + "y": 135, + "width": 120, + "height": 28, + "text": "1. Ingest", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 9, + "version": 1, + "versionNonce": 9, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h2", + "type": "text", + "x": 505, + "y": 135, + "width": 110, + "height": 28, + "text": "2. Queue", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#2b8a3e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 10, + "version": 1, + "versionNonce": 10, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h3", + "type": "text", + "x": 890, + "y": 135, + "width": 150, + "height": 28, + "text": "3. Workers", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#e67700", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 11, + "version": 1, + "versionNonce": 11, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h4", + "type": "text", + "x": 1320, + "y": 135, + "width": 200, + "height": 28, + "text": "4. Tree State", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 12, + "version": 1, + "versionNonce": 12, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h5", + "type": "text", + "x": 275, + "y": 515, + "width": 290, + "height": 28, + "text": "5. Scheduler / Background", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#0b7285", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 13, + "version": 1, + "versionNonce": 13, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h6", + "type": "text", + "x": 1100, + "y": 515, + "width": 280, + "height": 28, + "text": "6. Leaf Lifecycle", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#343a40", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 14, + "version": 1, + "versionNonce": 14, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b1", + "type": "rectangle", + "x": 75, + "y": 185, + "width": 240, + "height": 240, + "strokeColor": "#1971c2", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 15, + "version": 1, + "versionNonce": 15, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t1", + "type": "text", + "x": 93, + "y": 205, + "width": 204, + "height": 198, + "text": "JSON-RPC / source ingestion\n\nchat | email | document\n\ncanonicalise\n-> chunk_markdown\n-> score_chunks_fast\n-> upsert_chunks_tx\n-> lifecycle_status = pending_extraction\n-> persist fast score rows\n-> enqueue extract_chunk per chunk\n-> wake_workers()", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 16, + "version": 1, + "versionNonce": 16, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b2", + "type": "rectangle", + "x": 435, + "y": 185, + "width": 250, + "height": 240, + "strokeColor": "#2b8a3e", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 17, + "version": 1, + "versionNonce": 17, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t2", + "type": "text", + "x": 453, + "y": 205, + "width": 214, + "height": 198, + "text": "SQLite: memory_tree/chunks.db\n\nmem_tree_chunks\nmem_tree_score\nmem_tree_entity_index\nmem_tree_jobs\nmem_tree_trees\nmem_tree_buffers\nmem_tree_summaries\n\njobs fields\nkind | payload_json | dedupe_key\nstatus | attempts | available_at_ms\nlocked_until_ms | last_error", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 18, + "version": 1, + "versionNonce": 18, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b3", + "type": "rectangle", + "x": 815, + "y": 185, + "width": 300, + "height": 135, + "strokeColor": "#e67700", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 19, + "version": 1, + "versionNonce": 19, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t3", + "type": "text", + "x": 833, + "y": 205, + "width": 264, + "height": 108, + "text": "jobs::start(workspace_dir)\n\nrecover_stale_locks()\nspawn 3 worker tasks\nNotify wakeup + 5s polling fallback\nshared Semaphore(3) for LLM-bound work", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 104, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 20, + "version": 1, + "versionNonce": 20, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b4", + "type": "rectangle", + "x": 815, + "y": 335, + "width": 300, + "height": 90, + "strokeColor": "#d9480f", + "backgroundColor": "#fff8f0", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 21, + "version": 1, + "versionNonce": 21, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t4", + "type": "text", + "x": 833, + "y": 355, + "width": 264, + "height": 54, + "text": "Handlers\nextract_chunk | append_buffer | seal\ntopic_route | digest_daily | flush_stale", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 50, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 22, + "version": 1, + "versionNonce": 22, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b5", + "type": "rectangle", + "x": 1240, + "y": 185, + "width": 360, + "height": 240, + "strokeColor": "#9c36b5", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 23, + "version": 1, + "versionNonce": 23, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t5", + "type": "text", + "x": 1258, + "y": 205, + "width": 324, + "height": 198, + "text": "Tree building outputs\n\nsource tree\nL0 buffer -> seal -> L1/L2/... summaries\n\ntopic tree\ncurator hotness gate\noptional append_buffer(topic)\n\nglobal tree\ndigest_daily -> daily node\nappend_daily_and_cascade", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 24, + "version": 1, + "versionNonce": 24, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b6", + "type": "rectangle", + "x": 85, + "y": 575, + "width": 670, + "height": 105, + "strokeColor": "#0b7285", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 25, + "version": 1, + "versionNonce": 25, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t6", + "type": "text", + "x": 103, + "y": 597, + "width": 634, + "height": 72, + "text": "Scheduler loop\n\nUTC daily tick -> enqueue digest_daily(yesterday) + flush_stale(today)\nflush_stale scans stale buffers and enqueues force seal jobs\nworkers consume these through the same mem_tree_jobs pipeline", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 68, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 26, + "version": 1, + "versionNonce": 26, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s1", + "type": "rectangle", + "x": 875, + "y": 585, + "width": 130, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#fff3bf", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 27, + "version": 1, + "versionNonce": 27, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st1", + "type": "text", + "x": 891, + "y": 607, + "width": 98, + "height": 24, + "text": "pending_extraction", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 28, + "version": 1, + "versionNonce": 28, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s2", + "type": "rectangle", + "x": 1045, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#d3f9d8", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 29, + "version": 1, + "versionNonce": 29, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st2", + "type": "text", + "x": 1069, + "y": 607, + "width": 62, + "height": 24, + "text": "admitted", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 30, + "version": 1, + "versionNonce": 30, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s3", + "type": "rectangle", + "x": 1195, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#d0ebff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 31, + "version": 1, + "versionNonce": 31, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st3", + "type": "text", + "x": 1223, + "y": 607, + "width": 54, + "height": 24, + "text": "buffered", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 32, + "version": 1, + "versionNonce": 32, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s4", + "type": "rectangle", + "x": 1345, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#e5dbff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 33, + "version": 1, + "versionNonce": 33, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st4", + "type": "text", + "x": 1375, + "y": 607, + "width": 50, + "height": 24, + "text": "sealed", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 34, + "version": 1, + "versionNonce": 34, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s5", + "type": "rectangle", + "x": 1045, + "y": 665, + "width": 110, + "height": 36, + "strokeColor": "#c92a2a", + "backgroundColor": "#ffe3e3", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 35, + "version": 1, + "versionNonce": 35, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st5", + "type": "text", + "x": 1074, + "y": 672, + "width": 52, + "height": 20, + "text": "dropped", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 16, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 36, + "version": 1, + "versionNonce": 36, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "life-note", + "type": "text", + "x": 1185, + "y": 665, + "width": 390, + "height": 36, + "text": "extract_chunk decides admitted vs dropped. append_buffer moves admitted leaves into buffers. seal creates summaries and parent links.", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 32, + "strokeColor": "#495057", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 37, + "version": 1, + "versionNonce": 37, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "a1", + "type": "arrow", + "x": 315, + "y": 305, + "width": 115, + "height": 0, + "points": [[0, 0], [115, 0]], + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 38, + "version": 1, + "versionNonce": 38, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [115, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a2", + "type": "arrow", + "x": 685, + "y": 305, + "width": 125, + "height": 0, + "points": [[0, 0], [125, 0]], + "strokeColor": "#e67700", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 39, + "version": 1, + "versionNonce": 39, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [125, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a3", + "type": "arrow", + "x": 1115, + "y": 305, + "width": 125, + "height": 0, + "points": [[0, 0], [125, 0]], + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 40, + "version": 1, + "versionNonce": 40, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [125, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a4", + "type": "arrow", + "x": 430, + "y": 575, + "width": 70, + "height": 120, + "points": [[0, 0], [70, -120]], + "strokeColor": "#0b7285", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 41, + "version": 1, + "versionNonce": 41, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [70, -120], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a5", + "type": "arrow", + "x": 1005, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#2b8a3e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 42, + "version": 1, + "versionNonce": 42, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a6", + "type": "arrow", + "x": 1155, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 43, + "version": 1, + "versionNonce": 43, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a7", + "type": "arrow", + "x": 1305, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 44, + "version": 1, + "versionNonce": 44, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a8", + "type": "arrow", + "x": 1100, + "y": 655, + "width": 0, + "height": 10, + "points": [[0, 0], [0, 10]], + "strokeColor": "#c92a2a", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 45, + "version": 1, + "versionNonce": 45, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [0, 10], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "foot", + "type": "text", + "x": 40, + "y": 750, + "width": 1540, + "height": 60, + "text": "Job kinds in play: extract_chunk -> append_buffer(source) -> optional seal -> topic_route -> optional append_buffer(topic). Independent background flow: scheduler -> digest_daily / flush_stale -> seal. All paths go through mem_tree_jobs, so retries, dedupe, stale lock recovery, and worker wakeups stay centralized.", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 56, + "strokeColor": "#343a40", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 46, + "version": 1, + "versionNonce": 46, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + } + ], + "appState": { + "gridSize": null, + "viewBackgroundColor": "#ffffff" + }, + "files": {} +} diff --git a/src/bin/gmail_backfill_3d.rs b/src/bin/gmail_backfill_3d.rs new file mode 100644 index 000000000..717090d15 --- /dev/null +++ b/src/bin/gmail_backfill_3d.rs @@ -0,0 +1,474 @@ +//! Backfill the last N days of Gmail into the memory-tree content store. +//! +//! 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. +//! +//! 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`. +//! +//! # Prerequisites +//! +//! - Signed-in openhuman session JWT in the same workspace the desktop app +//! uses (stored at `/auth-profiles.json`). +//! - Active Gmail connection on Composio for that user. +//! +//! # Usage +//! +//! ```sh +//! cargo run --bin gmail-backfill-3d +//! cargo run --bin gmail-backfill-3d -- --days 7 +//! cargo run --bin gmail-backfill-3d -- --days 14 --page-size 100 +//! cargo run --bin gmail-backfill-3d -- --skip-drain +//! cargo run --bin gmail-backfill-3d -- --skip-verify +//! cargo run --bin gmail-backfill-3d -- --wipe +//! ``` +//! +//! Set `RUST_LOG=info` (or `debug`) for detailed output. + +use anyhow::{Context, Result}; +use clap::Parser; +use serde_json::{json, Value}; + +use openhuman_core::openhuman::composio::client::build_composio_client; +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::tree::content_store::read::{ + verify_chunk_file, verify_summary_file, VerifyResult, +}; +use openhuman_core::openhuman::memory::tree::jobs::drain_until_idle; +use openhuman_core::openhuman::memory::tree::store::{ + get_chunk_content_pointers, list_chunks, list_summaries_with_content_path, ListChunksQuery, +}; + +#[derive(Parser, Debug)] +#[command( + name = "gmail-backfill-3d", + about = "Backfill last N days of Gmail into the memory-tree content store (.md files + SQLite)." +)] +struct Cli { + /// Lookback window in days. Default 3. + #[arg(long, default_value_t = 3)] + days: u32, + + /// Page size per `GMAIL_FETCH_EMAILS` call (1–500). + #[arg(long, default_value_t = 50)] + page_size: u32, + + /// Cap on pages we will request. Guards against runaway pagination. + #[arg(long, default_value_t = 40)] + max_pages: u32, + + /// Include SPAM and TRASH messages in the fetch. + #[arg(long, default_value_t = false)] + include_spam_trash: bool, + + /// Extra Gmail search query AND-ed with the default scope. + #[arg(long)] + query: Option, + + /// Skip draining the async worker pool after ingest (useful for quick + /// smoke-test of file writes only). + #[arg(long, default_value_t = false)] + skip_drain: bool, + + /// Skip the post-drain integrity check (SHA-256 file verification). + #[arg(long, default_value_t = false)] + skip_verify: bool, + + /// Override the owner string embedded in chunk metadata. Defaults to + /// `"gmail-backfill"`. + #[arg(long)] + owner: Option, + + /// Wipe `chunks.db` (+ wal/shm) AND `/` before running. + /// Useful after a chunker change that invalidates existing chunk IDs. + #[arg(long, default_value_t = false)] + wipe: bool, +} + +#[tokio::main] +async fn main() -> Result<()> { + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")) + .format_timestamp_secs() + .try_init() + .ok(); + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .with_target(true) + .try_init() + .ok(); + + let cli = Cli::parse(); + if cli.days == 0 { + anyhow::bail!("--days must be >= 1"); + } + + let config = Config::load_or_init() + .await + .context("[gmail_backfill_3d] Config::load_or_init failed")?; + + if cli.wipe { + wipe_memory_tree_state(&config)?; + } + + let client = build_composio_client(&config).ok_or_else(|| { + anyhow::anyhow!( + "No Composio client — user not signed in (no JWT). \ + Sign in via the desktop app first, then re-run this binary." + ) + })?; + + 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 mut query = format!("in:inbox newer_than:{}d", cli.days); + if !cli.include_spam_trash { + query.push_str(" -in:spam -in:trash"); + } + if let Some(extra) = cli + .query + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + query.push(' '); + query.push_str(extra); + } + + log::info!( + "[gmail_backfill_3d] start days={} page_size={} max_pages={} query={:?}", + cli.days, + cli.page_size, + cli.max_pages, + query, + ); + + let content_root = config.memory_tree_content_root(); + log::info!( + "[gmail_backfill_3d] content_root={}", + content_root.display() + ); + + // ─── 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 = client + .execute_tool("GMAIL_FETCH_EMAILS", Some(args.clone())) + .await + .map_err(|e| anyhow::anyhow!("GMAIL_FETCH_EMAILS 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; + } + + let chunks_this_page = ingest_page_into_memory_tree(&config, &owner, &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, + } + } + + log::info!( + "[gmail_backfill_3d] fetch+ingest done pages={} total_chunks={} cost=~${:.4}", + total_pages, + total_chunks, + total_cost, + ); + + // ─── Drain async worker pool ──────────────────────────────────────────── + + if cli.skip_drain { + log::info!("[gmail_backfill_3d] skipping worker pool drain (--skip-drain)"); + } else { + log::info!("[gmail_backfill_3d] draining async worker pool…"); + drain_until_idle(&config).await?; + log::info!("[gmail_backfill_3d] worker pool idle"); + } + + // ─── Integrity check ──────────────────────────────────────────────────── + + 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, + ); + + // 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, + ); + + 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, + ); + 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) +} diff --git a/src/bin/slack_backfill.rs b/src/bin/slack_backfill.rs index fffeffd78..d6529f6e8 100644 --- a/src/bin/slack_backfill.rs +++ b/src/bin/slack_backfill.rs @@ -96,6 +96,21 @@ struct Cli { /// re-attempts cascade on the next append. #[arg(long = "seal-probe", default_value_t = false)] seal_probe: bool, + + /// Fire N back-to-back `SLACK_FETCH_CONVERSATION_HISTORY` calls + /// against the first listed channel and report a per-call tally + /// of {success, ratelimit, other-failure} + total duration. No + /// pacing by default (see --probe-pacing-ms), no ingestion. Used + /// to characterise Composio/Slack quota behaviour without + /// touching the memory tree. + #[arg(long = "probe-ratelimit")] + probe_ratelimit: Option, + + /// Sleep this many milliseconds between probe calls. 0 = fire + /// back-to-back (default). Use to find the threshold at which + /// rate-limits stop firing. + #[arg(long = "probe-pacing-ms", default_value_t = 0)] + probe_pacing_ms: u64, } #[tokio::main] @@ -243,6 +258,130 @@ async fn main() -> Result<()> { return Ok(()); } + if let Some(n) = cli.probe_ratelimit { + // Pure quota probe: fire N back-to-back + // SLACK_FETCH_CONVERSATION_HISTORY calls against the first + // discoverable channel. No pacing, no retry, no ingest. Reports + // a per-call status table + summary so we can characterise + // Composio/Slack rate-limit behaviour without contaminating the + // memory tree or burning extra quota on retries. + log::info!("[probe-ratelimit] requesting one channel via SLACK_LIST_CONVERSATIONS"); + let list_resp = client + .execute_tool( + "SLACK_LIST_CONVERSATIONS", + Some(serde_json::json!({ "exclude_archived": true, "limit": 1 })), + ) + .await + .map_err(|e| anyhow::anyhow!("SLACK_LIST_CONVERSATIONS failed: {e:#}"))?; + if !list_resp.successful { + anyhow::bail!( + "SLACK_LIST_CONVERSATIONS returned non-success: {:?}", + list_resp.error + ); + } + let channel_id = ["/data/channels/0/id", "/channels/0/id", "/data/0/id"] + .iter() + .find_map(|p| list_resp.data.pointer(p).and_then(|v| v.as_str())) + .map(str::to_string) + .ok_or_else(|| { + anyhow::anyhow!( + "could not find a channel id in SLACK_LIST_CONVERSATIONS response: {}", + serde_json::to_string(&list_resp.data).unwrap_or_default() + ) + })?; + log::info!("[probe-ratelimit] firing {n} calls against channel={channel_id}"); + + #[derive(Debug)] + enum Outcome { + Ok, + Ratelimit, + OtherFail(String), + Transport(String), + } + let mut outcomes: Vec<(u32, std::time::Duration, Outcome)> = Vec::with_capacity(n as usize); + let probe_started = Instant::now(); + for i in 1..=n { + if i > 1 && cli.probe_pacing_ms > 0 { + tokio::time::sleep(std::time::Duration::from_millis(cli.probe_pacing_ms)).await; + } + let t0 = Instant::now(); + let resp = client + .execute_tool( + "SLACK_FETCH_CONVERSATION_HISTORY", + Some(serde_json::json!({ "channel": channel_id, "limit": 1000 })), + ) + .await; + let dt = t0.elapsed(); + let outcome = match resp { + Err(e) => Outcome::Transport(format!("{e:#}")), + Ok(r) if r.successful => Outcome::Ok, + Ok(r) => { + let err = r.error.as_deref().unwrap_or("provider failure"); + if err.contains("ratelimited") + || err.contains("rate_limit") + || err.contains("rate limit") + { + log::warn!( + "[probe-ratelimit] call {i} ratelimited; body: {}", + serde_json::to_string(&r.data).unwrap_or_default() + ); + Outcome::Ratelimit + } else { + Outcome::OtherFail(err.to_string()) + } + } + }; + log::info!( + "[probe-ratelimit] call {i}/{n} took {:.2}s -> {:?}", + dt.as_secs_f64(), + outcome + ); + outcomes.push((i, dt, outcome)); + } + let total = probe_started.elapsed(); + + let ok = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Ok)) + .count(); + let rl = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Ratelimit)) + .count(); + let other = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::OtherFail(_))) + .count(); + let transport = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Transport(_))) + .count(); + let avg_ms = if !outcomes.is_empty() { + outcomes.iter().map(|(_, d, _)| d.as_millis()).sum::() / outcomes.len() as u128 + } else { + 0 + }; + + println!("=== probe-ratelimit summary ==="); + println!("channel: {channel_id}"); + println!("calls fired: {n}"); + println!("total duration: {:.2}s", total.as_secs_f64()); + println!("avg per call: {avg_ms} ms"); + println!("successful: {ok}"); + println!("ratelimited: {rl}"); + println!("other failures: {other}"); + println!("transport errors: {transport}"); + if rl > 0 { + let first_rl = outcomes + .iter() + .find(|(_, _, o)| matches!(o, Outcome::Ratelimit)) + .map(|(i, _, _)| *i) + .unwrap_or(0); + println!("first ratelimit: call #{first_rl}"); + } + return Ok(()); + } + let connections = client .list_connections() .await diff --git a/src/core/jsonrpc.rs b/src/core/jsonrpc.rs index 90d5712b6..21fcb45cf 100644 --- a/src/core/jsonrpc.rs +++ b/src/core/jsonrpc.rs @@ -805,7 +805,10 @@ async fn run_server_inner( /// /// Guarded by `std::sync::Once` so repeated calls to `bootstrap_skill_runtime` /// are safe and idempotent. -fn register_domain_subscribers(workspace_dir: std::path::PathBuf) { +fn register_domain_subscribers( + workspace_dir: std::path::PathBuf, + config: crate::openhuman::config::Config, +) { use std::sync::{Arc, Once}; static REGISTERED: Once = Once::new(); @@ -840,6 +843,7 @@ fn register_domain_subscribers(workspace_dir: std::path::PathBuf) { } crate::openhuman::composio::register_composio_trigger_subscriber(); crate::openhuman::composio::start_periodic_sync(); + crate::openhuman::memory::tree::jobs::start(config.clone()); // Restart requests go through a subscriber so every trigger path shares // the same respawn logic. @@ -880,7 +884,7 @@ pub async fn bootstrap_skill_runtime() { // Register domain subscribers for cross-module event handling. // Uses a Once guard so repeated calls to bootstrap_skill_runtime() // cannot double-subscribe. - register_domain_subscribers(workspace_dir.clone()); + register_domain_subscribers(workspace_dir.clone(), cfg.clone()); // --- Sub-agent definition registry bootstrap --- // Loads built-in archetype definitions plus any custom TOML files diff --git a/src/openhuman/composio/providers/gmail/ingest.rs b/src/openhuman/composio/providers/gmail/ingest.rs new file mode 100644 index 000000000..09ccb1c66 --- /dev/null +++ b/src/openhuman/composio/providers/gmail/ingest.rs @@ -0,0 +1,642 @@ +//! 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::tree::canonicalize::email::{EmailMessage, EmailThread}; +use crate::openhuman::memory::tree::canonicalize::email_clean::{ + extract_email, parse_message_date, +}; +use crate::openhuman::memory::tree::ingest::{ingest_email, IngestResult}; +use crate::openhuman::memory::tree::util::redact::redact; + +/// 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"]; + +/// 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 source_ref = if id.is_empty() { + None + } else { + Some(format!("gmail://msg/{id}")) + }; + Some(EmailMessage { + from, + to, + cc, + subject, + sent_at, + body, + source_ref, + }) +} + +/// 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`] which fans out +/// to the chunker + scorer + source tree downstream. +/// +/// `source_id` = `"gmail:{participants}"` where participants is +/// `addr1|addr2|...` (sorted, deduped, lowercased). This groups all +/// correspondence between the same people into one path subtree. +/// +/// 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, + page_messages: &[Value], +) -> Result { + if page_messages.is_empty() { + return Ok(0); + } + let buckets = bucket_by_participants(page_messages); + let mut total_chunks = 0usize; + let mut total_buckets = 0usize; + for (participants, raw_msgs) in &buckets { + 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; + } + // source_id encodes participants so every unique conversation set + // lands in its own path subtree. + let source_id = format!("gmail:{}", participants); + let thread_subject = pick_thread_subject(&messages); + log::info!( + "[composio:gmail][ingest] bucket participants_hash={} messages={} source_id_hash={}", + redact(participants), + messages.len(), + redact(&source_id) + ); + 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; + } + 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}", + redact(owner) + ); + Ok(total_chunks) +} + +/// 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, + }]; + 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, + }]; + assert_eq!(pick_thread_subject(&messages), "(no subject)"); + } +} diff --git a/src/openhuman/composio/providers/gmail/mod.rs b/src/openhuman/composio/providers/gmail/mod.rs index bdc05ebe2..a20f787d9 100644 --- a/src/openhuman/composio/providers/gmail/mod.rs +++ b/src/openhuman/composio/providers/gmail/mod.rs @@ -1,3 +1,4 @@ +pub mod ingest; mod post_process; mod provider; mod sync; diff --git a/src/openhuman/config/schema/load.rs b/src/openhuman/config/schema/load.rs index badc8c2e6..ffc204e1d 100644 --- a/src/openhuman/config/schema/load.rs +++ b/src/openhuman/config/schema/load.rs @@ -980,6 +980,17 @@ impl Config { } } + // Phase MD-content: chunk body directory override. Empty string means + // "fall back to default", consistent with other memory_tree env vars. + if let Ok(dir) = std::env::var("OPENHUMAN_MEMORY_TREE_CONTENT_DIR") { + let trimmed = dir.trim(); + self.memory_tree.content_dir = if trimmed.is_empty() { + None + } else { + Some(std::path::PathBuf::from(trimmed)) + }; + } + // Auto-update overrides if let Some(flag) = env.get("OPENHUMAN_AUTO_UPDATE_ENABLED") { let normalized = flag.trim().to_ascii_lowercase(); diff --git a/src/openhuman/config/schema/storage_memory.rs b/src/openhuman/config/schema/storage_memory.rs index 790bccd92..574806bd0 100644 --- a/src/openhuman/config/schema/storage_memory.rs +++ b/src/openhuman/config/schema/storage_memory.rs @@ -2,6 +2,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; #[derive(Debug, Clone, Serialize, Deserialize, Default, JsonSchema)] pub struct StorageConfig { @@ -96,6 +97,7 @@ impl Default for MemoryConfig { /// - `OPENHUMAN_MEMORY_SUMMARISE_ENDPOINT` /// - `OPENHUMAN_MEMORY_SUMMARISE_MODEL` /// - `OPENHUMAN_MEMORY_SUMMARISE_TIMEOUT_MS` +/// - `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (Phase MD-content) #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] pub struct MemoryTreeConfig { /// Ollama endpoint for the embedder (e.g. `http://localhost:11434`). @@ -154,6 +156,17 @@ pub struct MemoryTreeConfig { /// tokens and therefore takes longer to generate. #[serde(default = "default_memory_tree_llm_summariser_timeout_ms")] pub llm_summariser_timeout_ms: Option, + + /// Phase MD-content: root directory where chunk `.md` files are stored. + /// + /// Resolved at runtime via [`super::types::Config::memory_tree_content_root`]: + /// - `Some(path)` → use that path verbatim. + /// - `None` → default `/memory_tree/content/`. + /// + /// Env override: `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (empty string = fall + /// back to default, consistent with other memory_tree env vars). + #[serde(default = "default_memory_tree_content_dir")] + pub content_dir: Option, } /// Returns `None` so that existing installs that never opted into Phase 4 @@ -199,6 +212,12 @@ fn default_memory_tree_llm_summariser_timeout_ms() -> Option { Some(120_000) } +/// Returns `None` so the default `/memory_tree/content/` path is +/// used unless explicitly overridden via TOML or env var. +fn default_memory_tree_content_dir() -> Option { + None +} + impl Default for MemoryTreeConfig { fn default() -> Self { Self { @@ -212,6 +231,43 @@ impl Default for MemoryTreeConfig { llm_summariser_endpoint: default_memory_tree_llm_endpoint(), llm_summariser_model: default_memory_tree_llm_endpoint(), llm_summariser_timeout_ms: default_memory_tree_llm_summariser_timeout_ms(), + content_dir: default_memory_tree_content_dir(), } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_tree_config_default_content_dir_is_none() { + let cfg = MemoryTreeConfig::default(); + assert!( + cfg.content_dir.is_none(), + "default content_dir must be None so workspace default path is used" + ); + } + + /// Verify that the env-var override logic correctly maps non-empty strings + /// to `Some(PathBuf)` and empty/blank strings to `None`. We test the + /// logic inline (not via `apply_env_overrides`) to avoid mutating the + /// process environment in a way that could race with parallel tests. + #[test] + fn content_dir_env_override_logic() { + // Simulate the load.rs overlay logic. + let apply = |raw: &str| -> Option { + let trimmed = raw.trim(); + if trimmed.is_empty() { + None + } else { + Some(PathBuf::from(trimmed)) + } + }; + + assert_eq!(apply("/tmp/foo"), Some(PathBuf::from("/tmp/foo"))); + assert_eq!(apply(" /tmp/foo "), Some(PathBuf::from("/tmp/foo"))); + assert_eq!(apply(""), None); + assert_eq!(apply(" "), None); + } +} diff --git a/src/openhuman/config/schema/types.rs b/src/openhuman/config/schema/types.rs index e834d0ebc..866b4569b 100644 --- a/src/openhuman/config/schema/types.rs +++ b/src/openhuman/config/schema/types.rs @@ -203,6 +203,23 @@ pub struct Config { pub chat_onboarding_completed: bool, } +impl Config { + /// Resolve the root directory where chunk `.md` files are stored. + /// + /// Resolution order: + /// 1. `memory_tree.content_dir` if `Some`. + /// 2. Default: `/memory_tree/content/`. + /// + /// This is the only place in the codebase that should compute the content + /// root — all code that needs the path should call this method. + pub fn memory_tree_content_root(&self) -> PathBuf { + self.memory_tree + .content_dir + .clone() + .unwrap_or_else(|| self.workspace_dir.join("memory_tree").join("content")) + } +} + impl Default for Config { fn default() -> Self { let openhuman_dir = diff --git a/src/openhuman/memory/tree/canonicalize/chat.rs b/src/openhuman/memory/tree/canonicalize/chat.rs index 1a1b91ed9..43e180c2d 100644 --- a/src/openhuman/memory/tree/canonicalize/chat.rs +++ b/src/openhuman/memory/tree/canonicalize/chat.rs @@ -4,10 +4,9 @@ //! from the same channel becomes one [`CanonicalisedSource`]; the chunker //! slices it by token budget downstream. //! -//! Output format: +//! Output format (no leading `# ...` header — that info lives in front-matter +//! once Phase MD-content lands; the chunker splits at `## ` boundaries): //! ```md -//! # Chat transcript — {platform} / {channel} -//! //! ## 2026-04-21T10:12:00Z — Alice //! Message body here. //! @@ -67,10 +66,9 @@ pub fn canonicalise( let last_ts = messages.last().map(|m| m.timestamp).unwrap(); let mut md = String::new(); - md.push_str(&format!( - "# Chat transcript — {} / {}\n\n", - batch.platform, batch.channel_label - )); + // 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", @@ -150,7 +148,7 @@ mod tests { } #[test] - fn includes_header_and_per_message_sections() { + fn includes_per_message_sections_without_header() { let b = ChatBatch { platform: "slack".into(), channel_label: "#eng".into(), @@ -159,10 +157,15 @@ mod tests { let out = canonicalise("slack:#eng", "alice", &[], b) .unwrap() .unwrap(); - assert!(out - .markdown - .starts_with("# Chat transcript — slack / #eng\n\n")); - assert!(out.markdown.contains("## ")); + // 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")); } diff --git a/src/openhuman/memory/tree/canonicalize/document.rs b/src/openhuman/memory/tree/canonicalize/document.rs index 1fa6e9e8f..a3e171dd4 100644 --- a/src/openhuman/memory/tree/canonicalize/document.rs +++ b/src/openhuman/memory/tree/canonicalize/document.rs @@ -39,7 +39,8 @@ pub fn canonicalise( } let mut md = String::new(); - md.push_str(&format!("# {} — {}\n\n", doc.provider, doc.title)); + // 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'); @@ -85,7 +86,7 @@ mod tests { } #[test] - fn renders_title_and_body() { + fn renders_body_without_header() { let out = canonicalise( "d1", "alice", @@ -94,7 +95,11 @@ mod tests { ) .unwrap() .unwrap(); - assert!(out.markdown.starts_with("# notion — Launch plan\n\n")); + // 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")); } diff --git a/src/openhuman/memory/tree/canonicalize/email.rs b/src/openhuman/memory/tree/canonicalize/email.rs index f7d12df54..925026ac7 100644 --- a/src/openhuman/memory/tree/canonicalize/email.rs +++ b/src/openhuman/memory/tree/canonicalize/email.rs @@ -1,13 +1,16 @@ //! Email threads → canonical Markdown. //! -//! Email sources are scoped by **thread**. One thread becomes one -//! [`CanonicalisedSource`]. Headers (From, To, Subject, Date) surface in a -//! small frontmatter-style block per message; the body follows as 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::{normalize_source_ref, CanonicalisedSource}; +use super::{email_clean, normalize_source_ref, CanonicalisedSource}; use crate::openhuman::memory::tree::types::{Metadata, SourceKind}; /// One email in a thread. @@ -55,10 +58,9 @@ pub fn canonicalise( let last_ts = messages.last().map(|m| m.sent_at).unwrap(); let mut md = String::new(); - md.push_str(&format!( - "# Email thread — {} — {}\n\n", - thread.provider, thread.thread_subject - )); + // 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"); @@ -71,7 +73,12 @@ pub fn canonicalise( } md.push_str(&format!("Subject: {}\n", msg.subject)); md.push_str(&format!("Date: {}\n\n", msg.sent_at.to_rfc3339())); - md.push_str(msg.body.trim()); + 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"); } @@ -128,10 +135,19 @@ mod tests { email(2000, "alice@example.com", "Re: Launch", "agreed"), ], }; - let out = canonicalise("gmail:t1", "alice@example.com", &[], t) - .unwrap() - .unwrap(); - assert!(out.markdown.contains("# Email thread — gmail — Launch")); + 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")); @@ -139,6 +155,50 @@ mod tests { 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, + }], + }; + 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 { diff --git a/src/openhuman/memory/tree/canonicalize/email_clean.rs b/src/openhuman/memory/tree/canonicalize/email_clean.rs new file mode 100644 index 000000000..cac33160f --- /dev/null +++ b/src/openhuman/memory/tree/canonicalize/email_clean.rs @@ -0,0 +1,382 @@ +//! 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 { + 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); + 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)) { + 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 { + 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> { + let raw = m.get("date")?; + if let Some(s) = raw.as_str() { + let s = s.trim(); + if s.is_empty() { + return None; + } + // Epoch millis as a string? + 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()); + } + } + if let Some(ms) = raw.as_i64() { + return DateTime::from_timestamp_millis(ms); + } + None +} + +#[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 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(&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()); + } +} diff --git a/src/openhuman/memory/tree/canonicalize/mod.rs b/src/openhuman/memory/tree/canonicalize/mod.rs index 2e57bb29a..168fafaec 100644 --- a/src/openhuman/memory/tree/canonicalize/mod.rs +++ b/src/openhuman/memory/tree/canonicalize/mod.rs @@ -12,6 +12,7 @@ pub mod chat; pub mod document; pub mod email; +pub mod email_clean; use serde::{Deserialize, Serialize}; diff --git a/src/openhuman/memory/tree/chunker.rs b/src/openhuman/memory/tree/chunker.rs index 4e0c13b38..62ffc04d0 100644 --- a/src/openhuman/memory/tree/chunker.rs +++ b/src/openhuman/memory/tree/chunker.rs @@ -5,16 +5,26 @@ //! so later phases (#709 seal budget = 10k tokens) can ingest them without //! blowing past the summariser ceiling. //! -//! Splitting strategy: prefer paragraph boundaries (`\n\n`), then single -//! newlines, then hard character-count cut. We never split a character in half. +//! ## 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/line/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**: original paragraph-based greedy packing (unchanged). use crate::openhuman::memory::tree::types::{approx_token_count, Chunk, Metadata, SourceKind}; +use crate::openhuman::memory::tree::util::redact::redact; /// Default upper bound on per-chunk tokens. /// -/// Aligned with the LLD's summariser 10k budget so a single chunk never blows -/// a seal on its own. -pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 10_000; +/// Sized below the L0 seal budget (`source_tree::types::TOKEN_BUDGET = 4_500`) +/// so each seal accumulates roughly 1–3 chunks before firing — natural pacing +/// for the local 1B summariser, which produces noticeably better summaries +/// with smaller (≤4–5k) inputs than at the previous 10k cap. +pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 3_000; /// Tunable settings for the chunker. #[derive(Clone, Debug)] @@ -49,34 +59,255 @@ pub struct ChunkerInput { /// 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 the paragraph/line/char splitter and emit each piece with +/// `partial_message = true`. +/// - **Document**: original paragraph-based greedy packing (unchanged). pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec { let now = chrono::Utc::now(); - let pieces = split_by_token_budget(&input.markdown, opts.max_tokens); + 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 = super::types::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={} len={} pieces={}", + "[memory_tree::chunker] source_kind={} source_id_hash={} len={} units={}", input.source_kind.as_str(), - input.source_id, + redact(&input.source_id), input.markdown.len(), - pieces.len() + units.len() ); - pieces - .into_iter() - .enumerate() - .map(|(idx, content)| { - let seq = idx as u32; - let token_count = approx_token_count(&content); - let id = super::types::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, + // 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 = super::types::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 = super::types::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, + }); } - }) - .collect() + 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 = super::types::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].iter().copied().collect(); + let piece = piece_lines.join("\n").trim_end().to_string(); + if !piece.is_empty() { + pieces.push(piece); + } + } + + pieces } /// Split `text` into pieces each ≤ `max_tokens` tokens. @@ -85,7 +316,7 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec /// 1. Paragraph (`\n\n`) /// 2. Line (`\n`) /// 3. Hard character cut (last resort; preserves UTF-8 code points) -fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec { +pub(crate) fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec { let max_tokens = max_tokens.max(1); if text.is_empty() { return vec![String::new()]; @@ -187,22 +418,33 @@ mod tests { 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: "hello world".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_eq!(chunks[0].content, "hello world"); + assert!(chunks[0].content.contains("hello world")); assert_eq!(chunks[0].seq_in_source, 0); + assert!(!chunks[0].partial_message); } #[test] - fn empty_input_produces_one_empty_chunk() { + fn empty_chat_input_produces_one_empty_chunk() { let input = ChunkerInput { source_kind: SourceKind::Chat, source_id: "x".into(), @@ -212,10 +454,275 @@ mod tests { let chunks = chunk_markdown(&input, &ChunkerOptions::default()); assert_eq!(chunks.len(), 1); assert_eq!(chunks[0].content, ""); + assert!(!chunks[0].partial_message); } #[test] - fn paragraph_boundaries_preferred() { + 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); @@ -225,7 +732,7 @@ mod tests { source_kind: SourceKind::Document, source_id: "doc1".into(), markdown: text, - metadata: meta(), + metadata: meta_doc(), }; let chunks = chunk_markdown( &input, @@ -234,8 +741,6 @@ mod tests { }, ); assert!(chunks.len() >= 2); - // Every chunk should be a multiple of complete paragraphs — so starts - // with 'a', 'b', or 'c' and doesn't mix them arbitrarily. for c in &chunks { let first = c.content.chars().next().unwrap(); assert!( @@ -247,16 +752,16 @@ mod tests { } #[test] - fn falls_back_to_line_split_when_no_paragraphs() { + 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::Chat, + source_kind: SourceKind::Document, source_id: "x".into(), markdown: text, - metadata: meta(), + metadata: meta_doc(), }; let chunks = chunk_markdown( &input, @@ -271,54 +776,14 @@ mod tests { } #[test] - fn chunk_ids_are_stable_across_runs() { - let input = ChunkerInput { - source_kind: SourceKind::Chat, - source_id: "slack:#eng".into(), - markdown: "hello".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 text = (0..10) - .map(|i| format!("p{i}")) - .collect::>() - .join("\n\n"); - let input = ChunkerInput { - source_kind: SourceKind::Document, - source_id: "d".into(), - markdown: text, - metadata: meta(), - }; - let chunks = chunk_markdown( - &input, - &ChunkerOptions { - max_tokens: 2, // forces one chunk per paragraph basically - }, - ); - for (idx, c) in chunks.iter().enumerate() { - assert_eq!(c.seq_in_source, idx as u32); - } - } - - #[test] - fn utf8_boundaries_preserved_on_hard_split() { + fn utf8_boundaries_preserved_on_hard_split_document() { // Single long line with no paragraph/line splits → falls to hard cut. - // Use a 3-byte-per-char Unicode codepoint (Chinese) and force a cut. let text = "中".repeat(400); let input = ChunkerInput { source_kind: SourceKind::Document, source_id: "d".into(), markdown: text.clone(), - metadata: meta(), + metadata: meta_doc(), }; let chunks = chunk_markdown( &input, @@ -332,12 +797,12 @@ mod tests { } #[test] - fn zero_token_budget_is_clamped_without_empty_leading_chunk() { + 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(), + metadata: meta_doc(), }; let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 0 }); assert!(!chunks.is_empty()); diff --git a/src/openhuman/memory/tree/content_store/atomic.rs b/src/openhuman/memory/tree/content_store/atomic.rs new file mode 100644 index 000000000..ef1798bb9 --- /dev/null +++ b/src/openhuman/memory/tree/content_store/atomic.rs @@ -0,0 +1,416 @@ +//! 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_abs_path, summary_rel_path}; + +/// 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))?; + f.write_all(bytes) + .map_err(|e| anyhow::anyhow!("write tempfile {:?}: {e}", tmp_path))?; + f.sync_all() + .map_err(|e| anyhow::anyhow!("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 + )) + } + } + } +} + +/// A summary that has been written to disk and is ready for SQLite upsert. +#[derive(Debug, Clone)] +pub struct StagedSummary { + pub summary_id: String, + /// Relative content path (forward-slash, e.g. `"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 +/// `date_for_global` argument is required when `input.tree_kind == +/// SummaryTreeKind::Global`. The `scope_slug` must already be slugified by the +/// caller. +/// +/// 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, + date_for_global: Option>, +) -> anyhow::Result { + let rel_path = summary_rel_path( + input.tree_kind, + scope_slug, + input.level, + input.summary_id, + date_for_global, + ); + let abs_path = summary_abs_path( + content_root, + input.tree_kind, + scope_slug, + input.level, + input.summary_id, + date_for_global, + ); + + let composed = compose_summary_md(input); + let body_bytes = composed.body.as_bytes(); + let sha256 = sha256_hex(body_bytes); + + // Idempotent re-stage: if the file already exists, read and hash its + // body bytes. If the on-disk hash matches the new body's hash, return + // the StagedSummary unchanged (true idempotency). If the hashes differ + // the on-disk file is stale/corrupted — re-write it atomically with the + // new content so the db row and disk file are always consistent. + // + // Not re-writing would leave SQLite storing a content_sha256 that + // doesn't match the actual on-disk bytes, breaking integrity checks. + if abs_path.exists() { + let disk_sha = read_body_sha256(&abs_path).unwrap_or_default(); + if disk_sha == sha256 { + log::debug!( + "[content_store::atomic] summary already on disk with matching sha: {}", + input.summary_id + ); + return Ok(StagedSummary { + summary_id: input.summary_id.to_string(), + content_path: rel_path, + content_sha256: sha256, + }); + } + // Hash mismatch — overwrite atomically. + log::debug!( + "[content_store::atomic] summary on-disk sha mismatch for {} — re-staging", + input.summary_id + ); + // Remove the stale file first; write_if_new's fast-path would skip it. + let _ = std::fs::remove_file(&abs_path); + } + + let full_bytes = composed.full.as_bytes(); + write_if_new(&abs_path, full_bytes)?; + + 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::tree::content_store::compose::SummaryComposeInput; + use crate::openhuman::memory::tree::content_store::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 + } + + 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_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", None).unwrap(); + assert_eq!(staged.summary_id, "summary:L1:test1"); + assert!(staged.content_path.starts_with("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", None).unwrap(); + let second = stage_summary(dir.path(), &input, "person-alex", None).unwrap(); + assert_eq!(first.content_sha256, second.content_sha256); + assert_eq!(first.content_path, second.content_path); + } + + #[test] + fn stage_summary_global_uses_date_in_path() { + use chrono::TimeZone; + let dir = TempDir::new().unwrap(); + let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).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", Some(date)).unwrap(); + assert!( + staged.content_path.contains("2026-04-28"), + "global summary path must contain date; 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", None).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", None).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", None).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/tree/content_store/compose.rs b/src/openhuman/memory/tree/content_store/compose.rs new file mode 100644 index 000000000..75e98969f --- /dev/null +++ b/src/openhuman/memory/tree/content_store/compose.rs @@ -0,0 +1,878 @@ +//! 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. + +use chrono::{DateTime, Utc}; + +use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind; +use crate::openhuman::memory::tree::types::{Chunk, SourceKind}; + +/// 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))); + 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))); + } + + if meta.tags.is_empty() { + fm.push_str("tags: []\n"); + } else { + fm.push_str("tags:\n"); + for tag in &meta.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) +} + +// ── Summary composition ────────────────────────────────────────────────────── + +/// Input data required to compose a summary `.md` file. +pub struct SummaryComposeInput<'a> { + pub summary_id: &'a str, + pub tree_kind: SummaryTreeKind, + 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, + pub level: u32, + /// Child ids (chunk_ids at L0 → L1, summary_ids for cascades). + pub child_ids: &'a [String], + /// Total child count (== child_ids.len() unless truncated). + pub child_count: usize, + pub time_range_start: DateTime, + pub time_range_end: DateTime, + 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. +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 + if r.child_ids.is_empty() { + fm.push_str("children: []\n"); + } else { + fm.push_str("children:\n"); + for id in r.child_ids { + fm.push_str(&format!(" - {}\n", yaml_scalar(id))); + } + } + 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")); + + // aliases: human-readable title + let alias = build_summary_alias(r); + fm.push_str("aliases:\n"); + fm.push_str(&format!(" - {}\n", yaml_scalar(&alias))); + + 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"`) +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> { + rewrite_tags(file_bytes, new_tags) +} + +/// 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` + 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. +fn yaml_scalar(s: &str) -> String { + 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::*; + use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind; + use crate::openhuman::memory::tree::types::{Metadata, SourceKind, SourceRef}; + 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())), + }, + 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 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, + }, + 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, + }, + 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, + }, + 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_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(" - child-1"), "must list child ids"); + assert!(fm.contains(" - child-2"), "must list child ids"); + assert!(fm.contains("tags: []"), "must start with empty tags"); + // 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 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_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_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_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: []")); + // Body must be unchanged + assert!(rewritten_str.ends_with("summary body text")); + } +} diff --git a/src/openhuman/memory/tree/content_store/mod.rs b/src/openhuman/memory/tree/content_store/mod.rs new file mode 100644 index 000000000..0555d721b --- /dev/null +++ b/src/openhuman/memory/tree/content_store/mod.rs @@ -0,0 +1,169 @@ +//! Content store for memory-tree chunk and summary `.md` files (Phase MD-content). +//! +//! Bodies are stored on disk as `.md` files with YAML front-matter. +//! SQLite holds `content_path` (relative, forward-slash) and `content_sha256` +//! (over body bytes only) as pointers + integrity tokens. +//! +//! ## Module layout +//! +//! - [`paths`] — path generation + `slugify_source_id` + summary path builders +//! - [`compose`] — YAML front-matter + body composition; tag rewriting +//! - [`atomic`] — tempfile+fsync+rename writes; SHA-256; `stage_summary` +//! - [`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 paths; +pub mod read; +pub mod tags; + +use std::path::Path; + +use crate::openhuman::memory::tree::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, +} + +/// Update the `tags:` block in a summary's on-disk `.md` file after an +/// extraction job runs. +/// +/// Delegates to [`tags::update_summary_tags`]. +pub fn update_summary_tags( + config: &crate::openhuman::config::Config, + summary_id: &str, +) -> 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. +/// +/// `content_root` — absolute path to the root of the content store. +pub fn stage_chunks(content_root: &Path, chunks: &[Chunk]) -> anyhow::Result> { + let mut staged = Vec::with_capacity(chunks.len()); + + for chunk in chunks { + let source_kind = chunk.metadata.source_kind.as_str(); + let source_id = &chunk.metadata.source_id; + + let rel_path = paths::chunk_rel_path(source_kind, source_id, &chunk.id); + let abs_path = paths::chunk_abs_path(content_root, source_kind, source_id, &chunk.id); + + let (full_bytes, body_bytes) = compose::compose_chunk_file(chunk); + let sha256 = atomic::sha256_hex(&body_bytes); + + match atomic::write_if_new(&abs_path, &full_bytes) { + Ok(written) => { + if written { + log::debug!("[content_store] wrote chunk {} → {}", chunk.id, rel_path); + } else { + log::debug!( + "[content_store] chunk {} already on disk at {}", + chunk.id, + rel_path + ); + } + } + Err(e) => { + log::error!( + "[content_store] failed to write chunk {} to {}: {e}", + chunk.id, + rel_path + ); + return Err(e); + } + } + + 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::tree::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, + }, + 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); + } +} diff --git a/src/openhuman/memory/tree/content_store/paths.rs b/src/openhuman/memory/tree/content_store/paths.rs new file mode 100644 index 000000000..8046bf71b --- /dev/null +++ b/src/openhuman/memory/tree/content_store/paths.rs @@ -0,0 +1,473 @@ +//! 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 chrono::{DateTime, Utc}; + +use crate::openhuman::memory::tree::util::redact::redact; + +/// Which kind of summary tree a summary belongs to. Determines the top-level +/// directory under `/summaries/`. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum SummaryTreeKind { + /// Per-source-tree summary. Layout: `summaries/source//L/.md` + Source, + /// Global digest tree. Layout: `summaries/global//L/.md` + Global, + /// Per-topic (entity) tree. Layout: `summaries/topic//L/.md` + Topic, +} + +/// Build the relative content path for a summary, using forward slashes. +/// +/// Path layout depends on tree_kind: +/// - Source: `"summaries/source//L/.md"` +/// - Global: `"summaries/global//L/.md"` +/// Falls back to `unknown-date` (with a warn log) if `date_for_global` is +/// `None` — preferable to panicking inside a path utility. +/// - Topic: `"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. +/// +/// The `summary_id` is sanitized into a filesystem-safe filename by replacing +/// characters illegal on Windows (`:`, `\`, `*`, `?`, `"`, `<`, `>`, `|`) with `-`. +pub fn summary_rel_path( + tree_kind: SummaryTreeKind, + scope_slug: &str, + level: u32, + summary_id: &str, + date_for_global: Option>, +) -> String { + // Strip a trailing `.md` from summary_id if accidentally included. + let id = summary_id.strip_suffix(".md").unwrap_or(summary_id); + // Sanitize to a cross-platform filename (colons are illegal on Windows NTFS). + let filename = sanitize_filename(id); + + match tree_kind { + SummaryTreeKind::Source => { + format!("summaries/source/{}/L{}/{}.md", scope_slug, level, filename) + } + SummaryTreeKind::Global => { + // Fall back to a sentinel date rather than panic — a path-utility + // panic would propagate up through seal/digest/janitor codepaths + // and abort otherwise-recoverable work. Callers should always + // pass a date for Global; the warn log surfaces the contract + // violation without taking the process down. + let date_str = match date_for_global { + Some(d) => d.format("%Y-%m-%d").to_string(), + None => { + log::warn!( + "[content_store::paths] summary_rel_path called for Global \ + without date_for_global; using sentinel 'unknown-date'. \ + Caller bug — please pass a date." + ); + "unknown-date".to_string() + } + }; + format!("summaries/global/{}/L{}/{}.md", date_str, level, filename) + } + SummaryTreeKind::Topic => { + format!("summaries/topic/{}/L{}/{}.md", scope_slug, level, filename) + } + } +} + +/// 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.) +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, + date_for_global: Option>, +) -> PathBuf { + let rel = summary_rel_path(tree_kind, scope_slug, level, summary_id, date_for_global); + 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::*; + + // ─── 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", + None, + ); + // Colons in summary_id are replaced with '-' for cross-platform filenames. + assert_eq!( + p, + "summaries/source/gmail-alice-x-com-bob-y-com/L1/summary-L1-abc.md" + ); + } + + #[test] + fn summary_rel_path_global() { + use chrono::TimeZone; + let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap(); + let p = summary_rel_path( + SummaryTreeKind::Global, + "global", + 0, + "summary:L0:daily", + Some(date), + ); + assert_eq!(p, "summaries/global/2026-04-28/L0/summary-L0-daily.md"); + } + + #[test] + fn summary_rel_path_topic() { + let p = summary_rel_path( + SummaryTreeKind::Topic, + "person-alex-johnson", + 1, + "summary:L1:xyz", + None, + ); + assert_eq!( + p, + "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", + None, + ); + assert_eq!(p, "summaries/topic/entity-slug/L2/summary-L2-foo.md"); + } + + #[test] + fn summary_rel_path_global_falls_back_to_sentinel_without_date() { + // Caller bug to omit date for Global, but a path utility shouldn't + // panic — fall back to a sentinel `unknown-date` segment so the + // file lands somewhere predictable rather than aborting the seal. + let p = summary_rel_path(SummaryTreeKind::Global, "global", 0, "summary:L0:x", None); + assert_eq!(p, "summaries/global/unknown-date/L0/summary-L0-x.md"); + } + + #[test] + fn summary_abs_path_rooted_under_content_root() { + use chrono::TimeZone; + use std::path::Path; + let root = Path::new("/workspace/content"); + let date = chrono::Utc.with_ymd_and_hms(2026, 1, 15, 0, 0, 0).unwrap(); + let abs = summary_abs_path( + root, + SummaryTreeKind::Global, + "global", + 0, + "daily-123", + Some(date), + ); + assert!(abs.starts_with(root)); + assert!(abs.ends_with("daily-123.md")); + } +} diff --git a/src/openhuman/memory/tree/content_store/read.rs b/src/openhuman/memory/tree/content_store/read.rs new file mode 100644 index 000000000..068f8f4f0 --- /dev/null +++ b/src/openhuman/memory/tree/content_store/read.rs @@ -0,0 +1,405 @@ +//! Read and verify chunk and summary `.md` files from the content store. + +use std::path::Path; + +use super::atomic::sha256_hex; +use super::compose::split_front_matter; +use crate::openhuman::memory::tree::util::redact::redact; + +/// 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::tree::store::get_chunk_content_pointers; + + let pointers = get_chunk_content_pointers(config, chunk_id)?.ok_or_else(|| { + anyhow::anyhow!( + "[content_store::read] no content_path for chunk_id={} (pre-MD-migration row?)", + chunk_id + ) + })?; + let (rel_path, expected_sha256) = pointers; + + let content_root = config.memory_tree_content_root(); + // Reconstruct the absolute path from the stored relative forward-slash path. + let abs_path = { + let mut p = content_root.clone(); + for component in rel_path.split('/') { + p.push(component); + } + p + }; + + 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), + ) + })?; + + // Verify the on-disk body matches the SHA stored at write time. A mismatch + // means the file was tampered with, the tx that committed the pointer + // raced with a separate writer, or the disk corrupted — all unsafe to + // hand back to a consumer. Fail loudly rather than serve stale/corrupt + // bytes into the LLM extractor / summariser pipeline. + if result.sha256 != expected_sha256 { + return Err(anyhow::anyhow!( + "[content_store::read] sha256 mismatch for chunk_id={} \ + expected={} actual={} path_hash={}", + chunk_id, + expected_sha256, + result.sha256, + redact(&rel_path), + )); + } + + Ok(result.body) +} + +use anyhow::Context as _; + +/// 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::tree::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 = { + let mut p = content_root.clone(); + for component in rel_path.split('/') { + p.push(component); + } + p + }; + + 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), + ) + })?; + + // Verify the on-disk body matches the SHA stored at seal time. See the + // matching guard in `read_chunk_body` for rationale. + if result.sha256 != expected_sha256 { + return Err(anyhow::anyhow!( + "[content_store::read] sha256 mismatch for summary_id={} \ + expected={} actual={} path_hash={}", + summary_id, + expected_sha256, + result.sha256, + redact(&rel_path), + )); + } + + Ok(result.body) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new}; + use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file; + use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind}; + 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, + }, + token_count: 8, + seq_in_source: 0, + created_at: ts, + partial_message: false, + } + } + + #[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::tree::content_store::atomic::{sha256_hex, write_if_new}; + use crate::openhuman::memory::tree::content_store::compose::{ + compose_summary_md, SummaryComposeInput, + }; + use crate::openhuman::memory::tree::content_store::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_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 + ); + } +} diff --git a/src/openhuman/memory/tree/content_store/tags.rs b/src/openhuman/memory/tree/content_store/tags.rs new file mode 100644 index 000000000..f0283e21c --- /dev/null +++ b/src/openhuman/memory/tree/content_store/tags.rs @@ -0,0 +1,420 @@ +//! Post-extraction tag rewriting for chunk and summary `.md` files. +//! +//! After the LLM extraction job runs, it produces a list of entities. Each +//! entity is converted to an Obsidian-style hierarchical tag (`kind/Value`) +//! and written into the `tags:` block in the file's front-matter. +//! +//! The body bytes (and therefore the SHA-256) are never changed — only the +//! front-matter is rewritten. + +use std::path::Path; + +use super::compose::{rewrite_summary_tags as compose_rewrite_summary_tags, rewrite_tags}; +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::score::store::list_entity_ids_for_node; +use crate::openhuman::memory::tree::store::get_summary_content_pointers; + +/// Rewrite the `tags:` block in a chunk's on-disk `.md` file. +/// +/// `abs_path` — absolute path to the chunk file. +/// `tags` — new list of tag strings (Obsidian `kind/Value` format). +/// +/// The operation is atomic: the new file is written to a sibling temp path and +/// then renamed over the original. If the file does not exist, the call is a +/// no-op (returns `Ok(())`). +/// +/// Note: unlike the initial chunk write, tag rewrites MAY overwrite an +/// existing file. The immutability contract covers the **body** only; tags are +/// explicitly designed to be updated post-extraction. +pub fn update_chunk_tags(abs_path: &Path, tags: &[String]) -> anyhow::Result<()> { + if !abs_path.exists() { + log::debug!( + "[content_store::tags] skipping tag update — file not found: {}", + abs_path.display() + ); + return Ok(()); + } + + let old_bytes = + std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?; + + let new_bytes = rewrite_tags(&old_bytes, tags) + .map_err(|e| anyhow::anyhow!("rewrite_tags {:?}: {e}", abs_path))?; + + // Write the new content atomically via a sibling temp file. + let parent = abs_path.parent().unwrap_or_else(|| Path::new(".")); + let tmp_name = format!(".tmp_tags_{}.md", crate_temp_id()); + let tmp_path = parent.join(&tmp_name); + + { + use std::io::Write; + let mut f = std::fs::File::create(&tmp_path) + .map_err(|e| anyhow::anyhow!("create tag-rewrite tempfile {:?}: {e}", tmp_path))?; + f.write_all(&new_bytes) + .map_err(|e| anyhow::anyhow!("write tag-rewrite tempfile {:?}: {e}", tmp_path))?; + f.sync_all() + .map_err(|e| anyhow::anyhow!("fsync tag-rewrite tempfile {:?}: {e}", tmp_path))?; + } + + std::fs::rename(&tmp_path, abs_path).map_err(|e| { + let _ = std::fs::remove_file(&tmp_path); + anyhow::anyhow!("rename tag-rewrite {:?} -> {:?}: {e}", tmp_path, abs_path) + })?; + + log::debug!( + "[content_store::tags] updated tags in {}", + abs_path.display() + ); + Ok(()) +} + +/// Rewrite the `tags:` block in a summary's on-disk `.md` file. +/// +/// Reads entity rows from `mem_tree_entity_index` for `summary_id`, converts +/// them to `kind/Value` Obsidian tags, rewrites the YAML `tags:` block +/// atomically (tempfile + fsync + rename), and verifies the body SHA-256 is +/// unchanged afterwards. +/// +/// Best-effort: tag-rewrite failures should not fail the extraction job. Callers +/// should log a warning and continue — the entity index is the authoritative source. +pub fn update_summary_tags(config: &Config, summary_id: &str) -> anyhow::Result<()> { + // 1. Fetch content_path from SQLite. + let pointers = get_summary_content_pointers(config, summary_id)?; + let (rel_path, expected_sha) = match pointers { + Some(p) => p, + None => { + log::debug!( + "[content_store::tags] update_summary_tags: no content_path for summary {summary_id} — skipping" + ); + return Ok(()); + } + }; + + let content_root = config.memory_tree_content_root(); + let abs_path = { + let mut p = content_root; + for component in rel_path.split('/') { + p.push(component); + } + p + }; + + if !abs_path.exists() { + log::debug!( + "[content_store::tags] update_summary_tags: file missing for summary {summary_id} \ + at {} — skipping", + abs_path.display() + ); + return Ok(()); + } + + // 2. Fetch entity_index rows and build the merged tag list. + let entity_ids = list_entity_ids_for_node(config, summary_id)?; + let tags: Vec = entity_ids + .iter() + .filter_map(|eid| { + // entity_id format: "kind:surface" + let (kind, surface) = eid.split_once(':')?; + Some(entity_tag(kind, surface)) + }) + .collect(); + + // Sort + dedup for stability. + let mut tags = tags; + tags.sort(); + tags.dedup(); + + // 3. Read + atomic rewrite of the front-matter `tags:` block. + let old_bytes = std::fs::read(&abs_path) + .map_err(|e| anyhow::anyhow!("read summary {:?}: {e}", abs_path))?; + + let new_bytes = compose_rewrite_summary_tags(&old_bytes, &tags) + .map_err(|e| anyhow::anyhow!("rewrite_summary_tags {:?}: {e}", abs_path))?; + + let parent = abs_path.parent().unwrap_or_else(|| Path::new(".")); + let tmp_name = format!(".tmp_sum_tags_{}.md", crate_temp_id()); + let tmp_path = parent.join(&tmp_name); + + { + use std::io::Write; + let mut f = std::fs::File::create(&tmp_path).map_err(|e| { + anyhow::anyhow!("create summary tag-rewrite tempfile {:?}: {e}", tmp_path) + })?; + f.write_all(&new_bytes).map_err(|e| { + anyhow::anyhow!("write summary tag-rewrite tempfile {:?}: {e}", tmp_path) + })?; + f.sync_all().map_err(|e| { + anyhow::anyhow!("fsync summary tag-rewrite tempfile {:?}: {e}", tmp_path) + })?; + } + + std::fs::rename(&tmp_path, &abs_path).map_err(|e| { + let _ = std::fs::remove_file(&tmp_path); + anyhow::anyhow!( + "rename summary tag-rewrite {:?} -> {:?}: {e}", + tmp_path, + abs_path + ) + })?; + + // 4. Sanity check: body sha must still match after the rewrite. + let verify_bytes = std::fs::read(&abs_path) + .map_err(|e| anyhow::anyhow!("re-read after tag rewrite {:?}: {e}", abs_path))?; + let content = std::str::from_utf8(&verify_bytes) + .map_err(|e| anyhow::anyhow!("UTF-8 after tag rewrite {:?}: {e}", abs_path))?; + let body_after = super::compose::split_front_matter(content) + .ok_or_else(|| anyhow::anyhow!("no front-matter after tag rewrite {:?}", abs_path))? + .1; + let actual_sha = super::atomic::sha256_hex(body_after.as_bytes()); + if actual_sha != expected_sha { + return Err(anyhow::anyhow!( + "[content_store::tags] update_summary_tags body mutated after rewrite \ + summary_id={summary_id} expected_sha={expected_sha} actual_sha={actual_sha}" + )); + } + + log::debug!( + "[content_store::tags] updated {} tags in summary file summary_id={summary_id} n_tags={}", + tags.len(), + tags.len() + ); + Ok(()) +} + +/// Slugify an entity kind string for use in an Obsidian hierarchical tag. +/// +/// Output: lowercase, spaces and non-alphanumeric chars replaced with `-`, +/// consecutive dashes collapsed, leading/trailing dashes stripped. +/// +/// Example: `"Person"` → `"person"`, `"GitHub Repo"` → `"github-repo"` +pub fn slugify_tag_kind(kind: &str) -> String { + slugify_tag_component(kind) +} + +/// Slugify an entity value string for use in an Obsidian hierarchical tag. +/// +/// Like `slugify_tag_kind`, but capitalises the first letter of each word +/// so values are visually distinct from kinds: +/// +/// `"alice johnson"` → `"Alice-Johnson"`, +/// `"project Phoenix"` → `"Project-Phoenix"` +pub fn slugify_tag_value(value: &str) -> String { + // Split on non-alphanumeric boundaries, capitalise first letter of each word. + let mut parts: Vec = Vec::new(); + let mut current = String::new(); + + for ch in value.chars() { + if ch.is_alphanumeric() || ch == '_' { + current.push(ch); + } else if !current.is_empty() { + parts.push(capitalise(¤t)); + current.clear(); + } + } + if !current.is_empty() { + parts.push(capitalise(¤t)); + } + + let joined = parts.join("-"); + if joined.is_empty() { + "unknown".to_string() + } else { + joined + } +} + +/// Build an Obsidian-style `kind/Value` tag string from raw entity kind + surface. +pub fn entity_tag(kind: &str, surface: &str) -> String { + format!("{}/{}", slugify_tag_kind(kind), slugify_tag_value(surface)) +} + +fn slugify_tag_component(s: &str) -> String { + let lower = s.to_lowercase(); + let mut out = String::new(); + let mut last_dash = true; + for ch in lower.chars() { + if ch.is_ascii_alphanumeric() || ch == '_' { + out.push(ch); + last_dash = false; + } else if !last_dash { + out.push('-'); + last_dash = true; + } + } + let trimmed = out.trim_end_matches('-'); + if trimmed.is_empty() { + "unknown".to_string() + } else { + trimmed.to_string() + } +} + +fn capitalise(s: &str) -> String { + let mut chars = s.chars(); + match chars.next() { + None => String::new(), + Some(first) => { + let upper: String = first.to_uppercase().collect(); + upper + chars.as_str() + } + } +} + +fn crate_temp_id() -> String { + use std::time::{SystemTime, UNIX_EPOCH}; + let ns = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos(); + format!("{ns:08x}") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new}; + use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file; + use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind}; + 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: "tags_test".into(), + content: "hello from tags test".into(), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: "slack:#eng".into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: vec!["old/Tag".into()], + source_ref: None, + }, + token_count: 4, + seq_in_source: 0, + created_at: ts, + partial_message: false, + } + } + + #[test] + fn update_chunk_tags_replaces_tag_block() { + let dir = TempDir::new().unwrap(); + let chunk = sample_chunk(); + let (full, _) = compose_chunk_file(&chunk); + let path = dir.path().join("0.md"); + write_if_new(&path, &full).unwrap(); + + update_chunk_tags( + &path, + &["person/Alice-Smith".into(), "project/Phoenix".into()], + ) + .unwrap(); + + let updated = std::fs::read_to_string(&path).unwrap(); + assert!(updated.contains(" - person/Alice-Smith")); + assert!(updated.contains(" - project/Phoenix")); + assert!(!updated.contains(" - old/Tag")); + // Body unchanged. + assert!(updated.ends_with("hello from tags test")); + } + + #[test] + fn update_chunk_tags_is_noop_for_missing_file() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("nonexistent.md"); + assert!(update_chunk_tags(&path, &["p/X".into()]).is_ok()); + } + + #[test] + fn slugify_tag_kind_examples() { + assert_eq!(slugify_tag_kind("Person"), "person"); + assert_eq!(slugify_tag_kind("GitHub Repo"), "github-repo"); + assert_eq!(slugify_tag_kind("EMAIL"), "email"); + } + + #[test] + fn slugify_tag_value_capitalises_words() { + assert_eq!(slugify_tag_value("alice johnson"), "Alice-Johnson"); + assert_eq!(slugify_tag_value("project Phoenix"), "Project-Phoenix"); + assert_eq!(slugify_tag_value("OPENAI"), "OPENAI"); + } + + #[test] + fn entity_tag_builds_obsidian_tag() { + assert_eq!( + entity_tag("person", "Alice Johnson"), + "person/Alice-Johnson" + ); + assert_eq!(entity_tag("ORG", "Tinyhumans AI"), "org/Tinyhumans-AI"); + } + + // ─── update_summary_tags tests ──────────────────────────────────────────── + + /// Write a summary .md file to disk with empty tags and verify rewriting works. + #[test] + fn rewrite_summary_tags_preserves_body_and_replaces_tags() { + use crate::openhuman::memory::tree::content_store::compose::{ + compose_summary_md, SummaryComposeInput, + }; + use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind; + + let dir = TempDir::new().unwrap(); + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let body = "summary body for tag test\n"; + let children = vec!["c1".to_string()]; + let input = SummaryComposeInput { + summary_id: "sum:L1:tagtest", + tree_kind: SummaryTreeKind::Source, + tree_id: "t1", + tree_scope: "gmail:alice@x.com", + level: 1, + child_ids: &children, + 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"); + write_if_new(&path, composed.full.as_bytes()).unwrap(); + + // Original must have `tags: []` + let original = std::fs::read_to_string(&path).unwrap(); + assert!(original.contains("tags: []")); + + // Rewrite the tags block + let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()]; + let file_bytes = std::fs::read(&path).unwrap(); + let rewritten = super::compose_rewrite_summary_tags(&file_bytes, &new_tags).unwrap(); + + // Write rewritten bytes back (simulating atomic rewrite) + let tmp = dir.path().join("sum.tmp.md"); + { + use std::io::Write; + let mut f = std::fs::File::create(&tmp).unwrap(); + f.write_all(&rewritten).unwrap(); + } + std::fs::rename(&tmp, &path).unwrap(); + + let updated = std::fs::read_to_string(&path).unwrap(); + assert!(updated.contains(" - person/Alice-Smith")); + assert!(updated.contains(" - topic/Memory")); + assert!(!updated.contains("tags: []")); + // Body unchanged + assert!(updated.ends_with(body)); + + // Body sha unchanged + use crate::openhuman::memory::tree::content_store::compose::split_front_matter; + let (_, body_after) = split_front_matter(&updated).unwrap(); + let sha = sha256_hex(body_after.as_bytes()); + let expected_sha = sha256_hex(body.as_bytes()); + assert_eq!( + sha, expected_sha, + "body sha must be stable after tag rewrite" + ); + } +} diff --git a/src/openhuman/memory/tree/global_tree/digest.rs b/src/openhuman/memory/tree/global_tree/digest.rs index e55187ee4..e82ae130e 100644 --- a/src/openhuman/memory/tree/global_tree/digest.rs +++ b/src/openhuman/memory/tree/global_tree/digest.rs @@ -19,11 +19,17 @@ //! - Idempotency: if an L0 daily node already exists for the target day, //! return `DigestOutcome::Skipped` rather than emitting a duplicate. +use std::collections::BTreeSet; + use anyhow::{Context, Result}; use chrono::{DateTime, Duration, NaiveDate, TimeZone, Utc}; use rusqlite::OptionalExtension; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::{ + atomic::stage_summary, paths::slugify_source_id, read as content_read, SummaryComposeInput, + SummaryTreeKind, +}; use crate::openhuman::memory::tree::global_tree::registry::get_or_create_global_tree; use crate::openhuman::memory::tree::global_tree::seal::append_daily_and_cascade; use crate::openhuman::memory::tree::global_tree::GLOBAL_TOKEN_BUDGET; @@ -154,6 +160,25 @@ pub async fn end_of_day_digest( .await .context("embed daily summary during end_of_day_digest")?; + // L0 daily node inherits entities/topics by union of contributing + // source-tree summaries. Each input was already labeled at source-tree + // seal time, so emergent themes don't need another extractor pass + // here — global is a sink; union preserves "days that mentioned X" + // retrieval without an extra LLM call. See LabelStrategy in + // source_tree::bucket_seal for the full design. + let mut entities_set: BTreeSet = BTreeSet::new(); + let mut topics_set: BTreeSet = BTreeSet::new(); + for inp in &inputs { + for e in &inp.entities { + entities_set.insert(e.clone()); + } + for t in &inp.topics { + topics_set.insert(t.clone()); + } + } + let daily_entities: Vec = entities_set.into_iter().collect(); + let daily_topics: Vec = topics_set.into_iter().collect(); + let now = Utc::now(); let daily_id = new_summary_id(0); let daily = SummaryNode { @@ -165,8 +190,8 @@ pub async fn end_of_day_digest( child_ids: inputs.iter().map(|i| i.id.clone()).collect(), content: output.content, token_count: output.token_count, - entities: output.entities, - topics: output.topics, + entities: daily_entities, + topics: daily_topics, time_range_start: day_start, time_range_end: day_end, score, @@ -175,6 +200,44 @@ pub async fn end_of_day_digest( embedding: Some(embedding), }; + // Phase MD-content: stage the L0 daily .md file before the write tx. + // `date_for_global` = day_start (the calendar day this digest covers). + let daily_compose_input = SummaryComposeInput { + summary_id: &daily.id, + tree_kind: SummaryTreeKind::Global, + tree_id: &daily.tree_id, + tree_scope: &global.scope, + level: daily.level, + child_ids: &daily.child_ids, + child_count: daily.child_ids.len(), + time_range_start: daily.time_range_start, + time_range_end: daily.time_range_end, + sealed_at: daily.sealed_at, + body: &daily.content, + }; + // Stage the summary .md file — abort the digest on failure so the database + // never commits a row with content_path = NULL. The digest job is retried + // via the normal job-retry path. + let content_root_daily = config.memory_tree_content_root(); + let global_scope_slug = slugify_source_id(&global.scope); + let staged_daily = stage_summary( + &content_root_daily, + &daily_compose_input, + &global_scope_slug, + Some(day_start), + ) + .with_context(|| { + format!( + "stage_summary failed for daily {}; digest aborted for retry", + daily.id + ) + })?; + log::debug!( + "[global_tree::digest] staged daily summary {} → {}", + daily.id, + staged_daily.content_path + ); + // Persist the daily node. Note: we do NOT backlink parent_id on the // child summaries here — their parents are their own source trees, not // the global tree. The global-tree child_ids are cross-source @@ -183,7 +246,7 @@ pub async fn end_of_day_digest( let tree_id_clone = global.id.clone(); with_connection(config, move |conn| { let tx = conn.unchecked_transaction()?; - store::insert_summary_tx(&tx, &daily_clone)?; + store::insert_summary_tx(&tx, &daily_clone, Some(&staged_daily))?; // Index any entities the summariser emitted (no-op under inert). crate::openhuman::memory::tree::score::store::index_summary_entity_ids_tx( &tx, @@ -315,9 +378,23 @@ fn pick_source_contribution( } }; + // Read the full body from disk — `node.content` is a ≤500-char preview + // after the MD-on-disk migration. The digest summariser must receive the + // complete summary text so the daily recap is not assembled from previews. + let body = match content_read::read_summary_body(config, &node.id) { + Ok(b) => b, + Err(e) => { + log::warn!( + "[global_tree::digest] read_summary_body failed for {} — using preview: {e:#}", + node.id + ); + // Non-fatal: fall back to preview for pre-MD-migration rows. + node.content.clone() + } + }; Ok(Some(SummaryInput { id: node.id, - content: format!("[{}]\n{}", source_tree.scope, node.content), + content: format!("[{}]\n{}", source_tree.scope, body), token_count: node.token_count, entities: node.entities, topics: node.topics, diff --git a/src/openhuman/memory/tree/global_tree/digest_tests.rs b/src/openhuman/memory/tree/global_tree/digest_tests.rs index 0db1ca750..39822e8b3 100644 --- a/src/openhuman/memory/tree/global_tree/digest_tests.rs +++ b/src/openhuman/memory/tree/global_tree/digest_tests.rs @@ -1,5 +1,8 @@ use super::*; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::content_store; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::source_tree::types::TreeStatus; @@ -7,6 +10,24 @@ use crate::openhuman::memory::tree::store::upsert_chunks; use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; 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: &[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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::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(); @@ -39,6 +60,7 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime token_count: 6_000, seq_in_source: 0, created_at: ts, + partial_message: false, }; let c2 = Chunk { id: chunk_id(SourceKind::Chat, scope, 1, "test-content"), @@ -55,8 +77,10 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime token_count: 6_000, seq_in_source: 1, created_at: ts, + partial_message: false, }; upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap(); + stage_test_chunks(cfg, &[c1.clone(), c2.clone()]); let leaf1 = LeafRef { chunk_id: c1.id.clone(), @@ -76,8 +100,12 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime topics: vec![], score: 0.5, }; - append_leaf(cfg, &tree, &leaf1, &summariser).await.unwrap(); - append_leaf(cfg, &tree, &leaf2, &summariser).await.unwrap(); + append_leaf(cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + append_leaf(cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); // 12k tokens > 10k budget → one L1 summary covering `ts`. } @@ -216,3 +244,157 @@ async fn seven_days_cascade_to_weekly_seal() { assert_eq!(t.max_level, 1); assert_eq!(t.status, TreeStatus::Active); } + +/// Seed a source tree whose sealed L1 summary carries the given entities +/// and topics. Entities are written into `mem_tree_entity_index` (where +/// seal-time hydration reads them); topics are stored on chunk metadata +/// tags. The seal then unions both into the L1 summary. +async fn seed_source_tree_with_labeled_l1( + cfg: &Config, + scope: &str, + ts: DateTime, + entities: Vec, + topics: Vec, +) { + 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; + + let tree = get_or_create_source_tree(cfg, scope).unwrap(); + let summariser = InertSummariser::new(); + + let mut chunks: Vec = Vec::new(); + for seq in 0..2u32 { + chunks.push(Chunk { + id: chunk_id(SourceKind::Chat, scope, seq, "labeled-test"), + content: format!("labeled chunk {seq} in {scope}"), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: scope.into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: topics.clone(), + source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))), + }, + token_count: 6_000, + seq_in_source: seq, + created_at: ts, + partial_message: false, + }); + } + upsert_chunks(cfg, &chunks).unwrap(); + stage_test_chunks(cfg, &chunks); + + for chunk in &chunks { + 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(), + Some(scope), + ) + .unwrap(); + } + } + + // Two 6k-token leaves total 12k → exceeds L0 budget → seal fires on + // the second append, producing one L1 summary that unions all leaf + // labels (every leaf has the same set, so dedup yields the input set). + for chunk in &chunks { + let leaf = LeafRef { + chunk_id: chunk.id.clone(), + token_count: 6_000, + timestamp: ts, + content: chunk.content.clone(), + entities: entities.clone(), + topics: topics.clone(), + score: 0.5, + }; + append_leaf( + cfg, + &tree, + &leaf, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn daily_digest_unions_labels_from_source_summaries() { + let (_tmp, cfg) = test_config(); + let summariser = InertSummariser::new(); + + let day = NaiveDate::from_ymd_opt(2025, 5, 1).unwrap(); + let ts = day.and_hms_opt(10, 0, 0).unwrap().and_utc(); + + // Source A's L1 carries (alice, phoenix-migration). Source B's L1 + // carries (bob, phoenix-migration, qa). The daily L0 should union to + // (alice, bob, phoenix-migration) for entities and (phoenix-migration, + // qa) for topics — overlap dedup'd. + seed_source_tree_with_labeled_l1( + &cfg, + "slack:#a", + ts, + vec!["email:alice@example.com".into(), "topic:phoenix".into()], + vec!["phoenix-migration".into()], + ) + .await; + seed_source_tree_with_labeled_l1( + &cfg, + "slack:#b", + ts, + vec!["person:bob".into(), "topic:phoenix".into()], + vec!["phoenix-migration".into(), "qa".into()], + ) + .await; + + let outcome = end_of_day_digest(&cfg, day, &summariser).await.unwrap(); + let daily_id = match outcome { + DigestOutcome::Emitted { daily_id, .. } => daily_id, + other => panic!("expected Emitted, got {other:?}"), + }; + + let daily = store::get_summary(&cfg, &daily_id).unwrap().unwrap(); + let entities: std::collections::BTreeSet<&str> = + daily.entities.iter().map(String::as_str).collect(); + let topics: std::collections::BTreeSet<&str> = + daily.topics.iter().map(String::as_str).collect(); + + assert!(entities.contains("email:alice@example.com")); + assert!(entities.contains("person:bob")); + assert!(entities.contains("topic:phoenix")); + assert_eq!( + entities.len(), + 3, + "expected 3 unique entities (deduped); got {entities:?}" + ); + assert!(topics.contains("phoenix-migration")); + assert!(topics.contains("qa")); + assert_eq!( + topics.len(), + 2, + "expected 2 unique topics (deduped); got {topics:?}" + ); +} diff --git a/src/openhuman/memory/tree/global_tree/recap.rs b/src/openhuman/memory/tree/global_tree/recap.rs index ddbccfde7..1fbf9e089 100644 --- a/src/openhuman/memory/tree/global_tree/recap.rs +++ b/src/openhuman/memory/tree/global_tree/recap.rs @@ -158,14 +158,31 @@ fn assemble_recap(covering: &[&SummaryNode], level: u32) -> RecapOutput { #[cfg(test)] mod tests { use super::*; + use crate::openhuman::memory::tree::content_store; use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome}; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::store::upsert_chunks; use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; 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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::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(); @@ -216,6 +233,7 @@ mod tests { token_count: 6_000, seq_in_source: 0, created_at: ts, + partial_message: false, }; let c2 = Chunk { id: chunk_id(SourceKind::Chat, scope, 1, "test-content"), @@ -232,8 +250,10 @@ mod tests { token_count: 6_000, seq_in_source: 1, created_at: ts, + partial_message: false, }; upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap(); + stage_test_chunks(cfg, &[c1.clone(), c2.clone()]); append_leaf( cfg, &tree, @@ -247,6 +267,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); @@ -263,6 +284,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); diff --git a/src/openhuman/memory/tree/global_tree/seal.rs b/src/openhuman/memory/tree/global_tree/seal.rs index efe77b658..3af02a58a 100644 --- a/src/openhuman/memory/tree/global_tree/seal.rs +++ b/src/openhuman/memory/tree/global_tree/seal.rs @@ -11,10 +11,15 @@ //! `mem_tree_summaries` on both sides (children and output), since even L0 //! is a sealed summary node rather than a raw chunk. +use std::collections::BTreeSet; + use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::{ + atomic::stage_summary, SummaryComposeInput, SummaryTreeKind, +}; use crate::openhuman::memory::tree::global_tree::{ GLOBAL_TOKEN_BUDGET, MONTHLY_SEAL_THRESHOLD, WEEKLY_SEAL_THRESHOLD, YEARLY_SEAL_THRESHOLD, }; @@ -186,6 +191,26 @@ async fn seal_one_level( .await .context("summariser failed during global seal")?; + // Global-tree summaries inherit their entity/topic labels via union + // from their already-labeled inputs (source-tree summaries carry + // labels from the source-tree seal extractor; global L1+ inputs + // carry labels from this same union path one level down). We + // deliberately do NOT run an extractor on the daily/weekly/monthly + // synthesis: the inputs already cover what the summary represents, + // and global is a sink — no second-pass labeling earns its keep. + let mut entities_set: BTreeSet = BTreeSet::new(); + let mut topics_set: BTreeSet = BTreeSet::new(); + for inp in &inputs { + for e in &inp.entities { + entities_set.insert(e.clone()); + } + for t in &inp.topics { + topics_set.insert(t.clone()); + } + } + let node_entities: Vec = entities_set.into_iter().collect(); + let node_topics: Vec = topics_set.into_iter().collect(); + // Phase 4 (#710): embed BEFORE opening the write tx so an embedder // error aborts the cascade without half-committing the summary. let embedder = @@ -208,8 +233,8 @@ async fn seal_one_level( child_ids: buf.item_ids.clone(), content: output.content, token_count: output.token_count, - entities: output.entities, - topics: output.topics, + entities: node_entities, + topics: node_topics, time_range_start, time_range_end, score, @@ -218,6 +243,49 @@ async fn seal_one_level( embedding: Some(embedding), }; + // Phase MD-content: stage the global summary .md file before opening the + // write tx. date_for_global = time_range_start date (daily for L0, or + // the start of the range for higher levels). + let global_date = Some(time_range_start); + let compose_input_global = SummaryComposeInput { + summary_id: &node.id, + tree_kind: SummaryTreeKind::Global, + tree_id: &node.tree_id, + tree_scope: &tree.scope, + level: node.level, + child_ids: &node.child_ids, + 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 — abort the seal on failure so the database + // never commits a row with content_path = NULL. The job-retry path will + // re-attempt the file write on next execution. + let content_root_global = config.memory_tree_content_root(); + // Global tree scope is typically the literal "global" string. + // Use it as-is for the path (slugify passes through short ascii strings unchanged). + let global_scope_slug = + crate::openhuman::memory::tree::content_store::paths::slugify_source_id(&tree.scope); + let staged_global = stage_summary( + &content_root_global, + &compose_input_global, + &global_scope_slug, + global_date, + ) + .with_context(|| { + format!( + "stage_summary failed for {}; global-tree seal aborted for retry", + node.id + ) + })?; + log::debug!( + "[global_tree::seal] staged summary {} → {}", + node.id, + staged_global.content_path + ); + // Single write transaction: insert the new summary, clear this level's // buffer, append the new id to the parent buffer, and bump the tree's // max_level/root_id if we just climbed. Re-read `max_level` inside the @@ -238,7 +306,7 @@ async fn seal_one_level( .map(|n| n.max(0) as u32) .context("Failed to read current max_level for global tree")?; - store::insert_summary_tx(&tx, &node)?; + store::insert_summary_tx(&tx, &node, Some(&staged_global))?; // Index any entities the summariser emitted. No-op under // InertSummariser (entities stays empty by design — see // summariser/inert.rs). Becomes active when the Ollama summariser @@ -383,7 +451,7 @@ mod tests { fn insert_daily(cfg: &Config, node: &SummaryNode) { with_connection(cfg, |conn| { let tx = conn.unchecked_transaction()?; - store::insert_summary_tx(&tx, node)?; + store::insert_summary_tx(&tx, node, None)?; tx.commit()?; Ok(()) }) diff --git a/src/openhuman/memory/tree/ingest.rs b/src/openhuman/memory/tree/ingest.rs index 91386ee3d..07b3c54bb 100644 --- a/src/openhuman/memory/tree/ingest.rs +++ b/src/openhuman/memory/tree/ingest.rs @@ -1,14 +1,12 @@ -//! Ingest orchestrator (Phase 1 + Phase 2): +//! Ingest orchestrator for the async memory-tree pipeline. //! -//! canonicalise → chunk → score → admission gate → persist (chunks + scores + entity index) +//! The hot path now does: +//! `canonicalise -> chunk -> fast score -> persist chunks/score rows -> enqueue extract jobs` //! -//! Phase 2 inserts scoring between chunker and persistence. Low-scoring -//! chunks are dropped (their rationale is still persisted to -//! `mem_tree_score` for diagnostics); surviving chunks get their entities -//! indexed so later phases can resolve "which chunks mention Alice?" in -//! O(lookup). +//! The slower work (full extraction, admission, tree buffering, sealing, +//! topic routing, daily digests) runs out of the SQLite-backed jobs queue. -use anyhow::{Context, Result}; +use anyhow::Result; use serde::{Deserialize, Serialize}; use crate::openhuman::config::Config; @@ -19,25 +17,21 @@ use crate::openhuman::memory::tree::canonicalize::{ CanonicalisedSource, }; use crate::openhuman::memory::tree::chunker::{chunk_markdown, ChunkerInput, ChunkerOptions}; -use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked}; +use crate::openhuman::memory::tree::content_store; +use crate::openhuman::memory::tree::jobs::{self, ExtractChunkPayload, NewJob}; use crate::openhuman::memory::tree::score::{self, ScoreResult, ScoringConfig}; -use crate::openhuman::memory::tree::source_tree::{ - append_leaf, get_or_create_source_tree, LeafRef, -}; use crate::openhuman::memory::tree::store; -use crate::openhuman::memory::tree::topic_tree::route_leaf_to_topic_trees; -use crate::openhuman::memory::tree::types::Chunk; -/// Outcome of one ingest call — extended with per-chunk admission info. +/// Outcome of one ingest call. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct IngestResult { pub source_id: String, - /// Number of chunks that passed the admission gate and were persisted. + /// Number of chunks persisted and queued for async extraction. pub chunks_written: usize, - /// Number of chunks that failed the admission gate and were NOT persisted - /// (their score rationale IS persisted for diagnostics). + /// 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 that were persisted (in source order). + /// IDs of all chunks written and queued. pub chunk_ids: Vec, } @@ -52,7 +46,6 @@ impl IngestResult { } } -/// Ingest a batch of chat messages scoped to one channel/group. pub async fn ingest_chat( config: &Config, source_id: &str, @@ -60,11 +53,6 @@ pub async fn ingest_chat( tags: Vec, batch: ChatBatch, ) -> Result { - log::debug!( - "[memory_tree::ingest] chat source_id={} msg_count={}", - source_id, - batch.messages.len() - ); let canonical = match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -73,7 +61,6 @@ pub async fn ingest_chat( persist(config, source_id, canonical).await } -/// Ingest a single email thread. pub async fn ingest_email( config: &Config, source_id: &str, @@ -81,11 +68,6 @@ pub async fn ingest_email( tags: Vec, thread: EmailThread, ) -> Result { - log::debug!( - "[memory_tree::ingest] email source_id={} msg_count={}", - source_id, - thread.messages.len() - ); let canonical = match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -94,7 +76,6 @@ pub async fn ingest_email( persist(config, source_id, canonical).await } -/// Ingest a single standalone document. pub async fn ingest_document( config: &Config, source_id: &str, @@ -102,13 +83,6 @@ pub async fn ingest_document( tags: Vec, doc: DocumentInput, ) -> Result { - let title_len = doc.title.chars().count(); - log::debug!( - "[memory_tree::ingest] document source_id={} has_title={} title_len={}", - source_id, - !doc.title.trim().is_empty(), - title_len - ); let canonical = match document::canonicalise(source_id, owner, &tags, doc).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -122,7 +96,6 @@ async fn persist( source_id: &str, canonical: CanonicalisedSource, ) -> Result { - // 1. Chunk let input = ChunkerInput { source_kind: canonical.metadata.source_kind, source_id: source_id.to_string(), @@ -134,15 +107,15 @@ async fn persist( return Ok(IngestResult::empty(source_id)); } - // 2. Score (async; extractor chain driven by config.memory_tree — - // regex-only unless llm_extractor_endpoint + model are set, in - // which case LlmEntityExtractor runs as second-pass on borderline - // chunks with soft-fallback to regex on LLM failure). - let scoring_cfg = ScoringConfig::from_config(config); - let scores = score::score_chunks(&chunks, &scoring_cfg).await?; + // 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_tree::ingest] stage_chunks failed: {e}"))?; - // Fail fast on scorer length mismatch — silently truncating via zip would - // drop chunks (or their score rationale) without trace. + 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_tree::ingest] scorer length mismatch: chunks={} scores={}", @@ -151,78 +124,79 @@ async fn persist( ); } - // 3. Partition kept vs dropped - let mut kept_chunks: Vec = Vec::new(); - let mut all_results: Vec<(ScoreResult, i64)> = Vec::new(); - for (chunk, result) in chunks.iter().zip(scores.into_iter()) { - let ts_ms = chunk.metadata.timestamp.timestamp_millis(); - if result.kept { - kept_chunks.push(chunk.clone()); - } - all_results.push((result, ts_ms)); - } - + let all_results: Vec<(ScoreResult, i64)> = chunks + .iter() + .zip(scores.into_iter()) + .map(|(chunk, result)| (result, chunk.metadata.timestamp.timestamp_millis())) + .collect(); let dropped = all_results.iter().filter(|(r, _)| !r.kept).count(); - log::debug!( - "[memory_tree::ingest] scoring source_id={} kept={} dropped={}", - source_id, - kept_chunks.len(), - dropped - ); - // 3.5. Phase 4 (#710) — embed every kept chunk BEFORE writing. A failed - // embed aborts the whole batch so a retry stays idempotent on - // `chunk_id` (upsert_chunks_tx ON CONFLICT REPLACE). Serial per - // chunk for now; parallelism is an explicit follow-up. We pack - // each vector into its SQLite BLOB representation up front so the - // write tx below doesn't have to juggle floats. - let embedder = build_embedder_from_config(config).context("build embedder during ingest")?; - log::debug!( - "[memory_tree::ingest] embedding source_id={} provider={} kept={}", - source_id, - embedder.name(), - kept_chunks.len() - ); - let mut chunk_embedding_blobs: Vec<(String, Vec)> = Vec::with_capacity(kept_chunks.len()); - for chunk in &kept_chunks { - let vector = embedder - .embed(&chunk.content) - .await - .with_context(|| format!("embed chunk_id={} during ingest", chunk.id))?; - let packed = pack_checked(&vector) - .with_context(|| format!("pack embedding for chunk_id={} during ingest", chunk.id))?; - chunk_embedding_blobs.push((chunk.id.clone(), packed)); - } - - // 4. Persist (blocking SQLite — isolate on a dedicated thread). - // Chunks + scores + embeddings all commit in one tx so retries are - // idempotent. The chunk upsert preserves any prior embedding on - // conflict (see store::upsert_chunks_tx), so we follow with a - // deliberate UPDATE that writes the fresh blob. let config_owned = config.clone(); - let kept_for_store = kept_chunks.clone(); + let staged_for_store = staged.clone(); let results_for_store = all_results.clone(); - let embeddings_for_store = chunk_embedding_blobs.clone(); let written = tokio::task::spawn_blocking(move || -> Result { + use std::collections::{HashMap, HashSet}; store::with_connection(&config_owned, |conn| { let tx = conn.unchecked_transaction()?; - let n = store::upsert_chunks_tx(&tx, &kept_for_store)?; - for (chunk_id, blob) in &embeddings_for_store { - let changed = tx.execute( - "UPDATE mem_tree_chunks SET embedding = ?1 WHERE id = ?2", - rusqlite::params![blob, chunk_id], - )?; - if changed == 0 { - log::warn!( - "[memory_tree::ingest] embedding update affected 0 rows chunk_id={chunk_id} — \ - upsert missed the row?" - ); + + // 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 = store::get_chunk_lifecycle_status_tx(&tx, &s.chunk.id)?; + prior.insert(s.chunk.id.clone(), status); + } + + let n = store::upsert_staged_chunks_tx(&tx, &staged_for_store)?; + + // 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(store::CHUNK_STATUS_PENDING_EXTRACTION), + ); + if needs_processing { + store::set_chunk_lifecycle_status_tx( + &tx, + &s.chunk.id, + store::CHUNK_STATUS_PENDING_EXTRACTION, + )?; + to_schedule.insert(s.chunk.id.clone()); } } + for (result, ts_ms) in &results_for_store { - // Persist rationale for EVERY chunk (kept or dropped). - // Index entities only for kept chunks (handled inside persist_score_tx). + 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(n) @@ -231,98 +205,26 @@ async fn persist( .await .map_err(|e| anyhow::anyhow!("persist join error: {e}"))??; - // 5. Source-tree append (Phase 3a #709). Each kept leaf pushes into - // the tree's L0 buffer and cascades upward when token_sum crosses - // the budget. Entities/topics from the scorer are threaded in so - // sealed summaries inherit the child signal set. Failures here - // log at warn level but don't fail the ingest — leaves are already - // persisted, and a later flush/retry can still rebuild the tree. - if let Err(e) = append_leaves_to_tree(config, source_id, &kept_chunks, &all_results).await { - log::warn!( - "[memory_tree::ingest] source_tree append failed source_id={} err={:#}", - source_id, - e - ); - } + jobs::wake_workers(); Ok(IngestResult { source_id: source_id.to_string(), chunks_written: written, chunks_dropped: dropped, - chunk_ids: kept_chunks.iter().map(|c| c.id.clone()).collect(), + chunk_ids: staged.iter().map(|s| s.chunk.id.clone()).collect(), }) } -/// Push every kept chunk into its source tree. Scoped to Phase 3a — all -/// chunks from one ingest batch share the same `source_id`, so they share -/// one tree lookup. The summariser comes from `build_summariser(config)`: -/// Ollama-backed [`LlmSummariser`] when `memory_tree.llm_summariser_*` is -/// set, else deterministic [`InertSummariser`]. Both share the same trait, -/// so the seal cascade doesn't notice the difference. -async fn append_leaves_to_tree( - config: &Config, - source_id: &str, - kept_chunks: &[Chunk], - all_results: &[(ScoreResult, i64)], -) -> Result<()> { - if kept_chunks.is_empty() { - return Ok(()); - } - let tree = get_or_create_source_tree(config, source_id)?; - let summariser = crate::openhuman::memory::tree::source_tree::build_summariser(config); - - // Build a chunk_id → (score, entities, topics) map for quick lookup. - use std::collections::HashMap; - let mut score_by_id: HashMap = HashMap::new(); - for (r, _) in all_results { - score_by_id.insert(r.chunk_id.clone(), r); - } - - for chunk in kept_chunks { - let (score_value, entities, topics) = match score_by_id.get(&chunk.id) { - Some(r) => ( - r.total, - r.canonical_entities - .iter() - .map(|e| e.canonical_id.clone()) - .collect::>(), - chunk.metadata.tags.clone(), - ), - None => (0.0, Vec::new(), chunk.metadata.tags.clone()), - }; - let leaf = LeafRef { - chunk_id: chunk.id.clone(), - token_count: chunk.token_count, - timestamp: chunk.metadata.timestamp, - content: chunk.content.clone(), - entities: entities.clone(), - topics, - score: score_value, - }; - append_leaf(config, &tree, &leaf, summariser.as_ref()).await?; - - // Phase 3c (#709): route the leaf to every matching topic tree - // and tick the curator for each entity. Non-fatal on error — - // the source-tree append has already succeeded above. - if let Err(e) = - route_leaf_to_topic_trees(config, &leaf, &entities, summariser.as_ref()).await - { - log::warn!( - "[memory_tree::ingest] topic_tree routing failed chunk_id={} err={:#}", - chunk.id, - e - ); - } - } - Ok(()) -} - #[cfg(test)] mod tests { use super::*; use crate::openhuman::memory::tree::canonicalize::chat::ChatMessage; + use crate::openhuman::memory::tree::jobs::drain_until_idle; use crate::openhuman::memory::tree::score::store::{count_scores, lookup_entity}; - use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery}; + use crate::openhuman::memory::tree::store::{ + count_chunks, count_chunks_by_lifecycle_status, get_chunk_embedding, list_chunks, + ListChunksQuery, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, + }; use crate::openhuman::memory::tree::types::SourceKind; use chrono::{TimeZone, Utc}; use tempfile::TempDir; @@ -331,16 +233,12 @@ mod tests { let tmp = TempDir::new().unwrap(); let mut cfg = Config::default(); cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): disable Ollama-backed embeddings in tests. - // Ingest/seal call `build_embedder_from_config`; falling back to - // inert (zero vectors) keeps tests deterministic and network-free. cfg.memory_tree.embedding_endpoint = None; cfg.memory_tree.embedding_model = None; cfg.memory_tree.embedding_strict = false; (tmp, cfg) } - /// Build a substantive batch that reliably passes the admission gate. fn substantive_batch() -> ChatBatch { ChatBatch { platform: "slack".into(), @@ -349,17 +247,14 @@ mod tests { 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. Please \ - confirm availability by replying here. alice@example.com" + 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'll handle the coordination and cut a release \ - candidate tonight. #launch-q2 will be tracked in Notion." + text: "Confirmed, I will handle the coordination and launch tracking tonight." .into(), source_ref: None, }, @@ -368,25 +263,40 @@ mod tests { } #[tokio::test] - async fn ingest_chat_writes_substantive_chunks() { + 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!(out.chunks_dropped, 0); assert_eq!(count_chunks(&cfg).unwrap(), 1); - // Score row persisted for the kept chunk - assert_eq!(count_scores(&cfg).unwrap(), 1); - // Entity index populated from regex extraction (alice@example.com + hashtag) - let alice_hits = lookup_entity(&cfg, "email:alice@example.com", None).unwrap(); - assert_eq!(alice_hits.len(), 1); + + drain_until_idle(&cfg).await.unwrap(); + + // Final lifecycle is `buffered`: extract → admitted → append_buffer → buffered. + // The single packed chunk does not cross 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); + assert!(get_chunk_embedding(&cfg, &out.chunk_ids[0]) + .unwrap() + .is_some()); } #[tokio::test] - async fn low_signal_chunks_are_dropped_but_score_persists() { + async fn low_signal_chunks_end_up_dropped_after_queue_processing() { let (_tmp, cfg) = test_config(); let batch = ChatBatch { platform: "slack".into(), @@ -394,18 +304,22 @@ mod tests { messages: vec![ChatMessage { author: "alice".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "+1".into(), // extremely low-signal + text: "+1".into(), source_ref: None, }], }; let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch) .await .unwrap(); - assert_eq!(out.chunks_written, 0); - assert_eq!(out.chunks_dropped, 1); - // Chunk NOT in chunks table - assert_eq!(count_chunks(&cfg).unwrap(), 0); - // Score row IS persisted for diagnostics + 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); } @@ -421,7 +335,6 @@ mod tests { .await .unwrap(); assert_eq!(out.chunks_written, 0); - assert_eq!(out.chunks_dropped, 0); assert_eq!(count_chunks(&cfg).unwrap(), 0); assert_eq!(count_scores(&cfg).unwrap(), 0); } @@ -432,9 +345,7 @@ mod tests { let doc = DocumentInput { provider: "notion".into(), title: "Launch plan".into(), - body: "We are planning to ship Phoenix on Friday after review. \ - Coordination is via email and the launch thread tracks \ - the relevant decisions. alice@example.com owns this." + 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()), @@ -445,74 +356,8 @@ mod tests { 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); } - - #[tokio::test] - async fn chunks_preserve_source_ref_when_kept() { - let (_tmp, cfg) = test_config(); - let doc = DocumentInput { - provider: "notion".into(), - title: "t".into(), - body: "Phoenix launch plan with enough substance to pass the admission \ - gate: we are reviewing the migration runbook alice@example.com \ - on Friday evening." - .into(), - modified_at: Utc::now(), - source_ref: Some("notion://x".into()), - }; - ingest_document(&cfg, "notion:x", "alice", vec![], doc) - .await - .unwrap(); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows.len(), 1); - assert_eq!( - rows[0].metadata.source_ref.as_ref().unwrap().value, - "notion://x" - ); - } - - // ── Phase 4 (#710) ────────────────────────────────────────────── - - /// Ingesting with an inert embedder must still write embeddings for - /// every kept chunk. Use a trimmed test_config so we go through the - /// full embed → persist pipeline. - #[tokio::test] - async fn ingest_writes_chunk_embeddings() { - use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM; - use crate::openhuman::memory::tree::store::get_chunk_embedding; - - let (_tmp, cfg) = test_config(); - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()) - .await - .unwrap(); - assert!(out.chunks_written >= 1); - for id in &out.chunk_ids { - let v = get_chunk_embedding(&cfg, id) - .unwrap() - .unwrap_or_else(|| panic!("missing embedding for {id}")); - assert_eq!(v.len(), EMBEDDING_DIM); - } - } - - /// When the embedder errors (here: Ollama required but unavailable), - /// ingest must fail and persist nothing. Point embedding at a dead - /// port so the HTTP call refuses, and flip `strict=true` so the - /// factory returns Ollama (not the inert fallback). - #[tokio::test] - async fn ingest_fails_when_embedder_fails_and_persists_nothing() { - let (_tmp, mut cfg) = test_config(); - cfg.memory_tree.embedding_endpoint = Some("http://127.0.0.1:1".into()); - cfg.memory_tree.embedding_model = Some("nomic-embed-text".into()); - cfg.memory_tree.embedding_timeout_ms = Some(500); - cfg.memory_tree.embedding_strict = true; - - let result = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()).await; - assert!(result.is_err(), "expected ingest to bail when embed fails"); - - // No chunks, no scores persisted — retry stays clean. - assert_eq!(count_chunks(&cfg).unwrap(), 0); - assert_eq!(count_scores(&cfg).unwrap(), 0); - } } diff --git a/src/openhuman/memory/tree/jobs/handlers/mod.rs b/src/openhuman/memory/tree/jobs/handlers/mod.rs new file mode 100644 index 000000000..2de3d5e1b --- /dev/null +++ b/src/openhuman/memory/tree/jobs/handlers/mod.rs @@ -0,0 +1,861 @@ +use anyhow::{Context, Result}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::{ + self as content_store, read as content_read, tags as content_tags, +}; +use crate::openhuman::memory::tree::global_tree::digest::{self, DigestOutcome}; +use crate::openhuman::memory::tree::jobs::store; +use crate::openhuman::memory::tree::jobs::types::{ + AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload, + Job, JobKind, NewJob, NodeRef, SealPayload, TopicRoutePayload, +}; +use crate::openhuman::memory::tree::score; +use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked}; +use crate::openhuman::memory::tree::score::extract::build_summary_extractor; +use crate::openhuman::memory::tree::score::store as score_store; +use crate::openhuman::memory::tree::source_tree::{ + build_summariser, get_or_create_source_tree, LabelStrategy, LeafRef, +}; +use crate::openhuman::memory::tree::store as chunk_store; +use crate::openhuman::memory::tree::topic_tree::curator; + +pub async fn handle_job(config: &Config, job: &Job) -> Result<()> { + match job.kind { + JobKind::ExtractChunk => handle_extract(config, job).await, + JobKind::AppendBuffer => handle_append_buffer(config, job).await, + JobKind::Seal => handle_seal(config, job).await, + JobKind::TopicRoute => handle_topic_route(config, job).await, + JobKind::DigestDaily => handle_digest_daily(config, job).await, + JobKind::FlushStale => handle_flush_stale(config, job).await, + } +} + +async fn handle_extract(config: &Config, job: &Job) -> Result<()> { + let payload: ExtractChunkPayload = + serde_json::from_str(&job.payload_json).context("parse ExtractChunk payload")?; + let Some(chunk) = chunk_store::get_chunk(config, &payload.chunk_id)? else { + log::warn!( + "[memory_tree::jobs] extract chunk missing chunk_id={}", + payload.chunk_id + ); + return Ok(()); + }; + + // Read the full body from disk (the `content` column in SQLite holds a + // ≤500-char preview after the MD-on-disk migration). Both the scorer and + // the embedder need the complete text so extraction and semantic indexing + // operate over the full chunk body, not a truncated preview. + let body = content_read::read_chunk_body(config, &chunk.id) + .with_context(|| format!("read full body for extract chunk_id={}", chunk.id))?; + // Score a clone of the chunk with the full body swapped in. + let chunk_with_body = { + let mut c = chunk.clone(); + c.content = body.clone(); + c + }; + + let scoring_cfg = score::ScoringConfig::from_config(config); + let result = score::score_chunk(&chunk_with_body, &scoring_cfg).await?; + let packed_embedding = if result.kept { + let embedder = + build_embedder_from_config(config).context("build embedder in extract handler")?; + // Reuse the body already read — avoid a second disk read. + let vector = embedder + .embed(&body) + .await + .with_context(|| format!("embed chunk_id={} in extract handler", chunk.id))?; + Some( + pack_checked(&vector) + .with_context(|| format!("pack embedding for chunk_id={}", chunk.id))?, + ) + } else { + None + }; + + // Build follow-up job payloads before opening the tx — construction is + // cheap and doesn't require a database connection. The two jobs are + // enqueued inside the SAME transaction that commits the lifecycle update, + // so a crash anywhere rolls everything back together and prevents the + // "lifecycle committed but job lost" crash window. + let source_job = if result.kept { + Some(NewJob::append_buffer(&AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: chunk.id.clone(), + }, + target: AppendTarget::Source { + source_id: chunk.metadata.source_id.clone(), + }, + })?) + } else { + None + }; + let route_job = if result.kept { + Some(NewJob::topic_route(&TopicRoutePayload { + node: NodeRef::Leaf { + chunk_id: chunk.id.clone(), + }, + })?) + } else { + None + }; + + let (did_enqueue_source, did_enqueue_route) = chunk_store::with_connection(config, |conn| { + let tx = conn.unchecked_transaction()?; + score::persist_score_tx( + &tx, + &result, + chunk.metadata.timestamp.timestamp_millis(), + None, + )?; + + if result.kept { + tx.execute( + "UPDATE mem_tree_chunks + SET embedding = ?1, + lifecycle_status = ?2 + WHERE id = ?3", + rusqlite::params![ + packed_embedding, + chunk_store::CHUNK_STATUS_ADMITTED, + chunk.id, + ], + )?; + } else { + tx.execute( + "UPDATE mem_tree_chunks + SET lifecycle_status = ?1 + WHERE id = ?2", + rusqlite::params![chunk_store::CHUNK_STATUS_DROPPED, chunk.id], + )?; + } + + // Enqueue follow-up jobs inside the SAME transaction so they are + // atomically visible with the lifecycle update. + let mut eq_src = false; + let mut eq_route = false; + if let Some(ref j) = source_job { + eq_src = store::enqueue_tx(&tx, j)?.is_some(); + } + if let Some(ref j) = route_job { + eq_route = store::enqueue_tx(&tx, j)?.is_some(); + } + + tx.commit()?; + Ok((eq_src, eq_route)) + })?; + + // Phase MD-content: rewrite the `tags:` block in the on-disk chunk file + // with Obsidian-style hierarchical tags derived from the extracted entities. + // This runs after the tx commits so the entity index is visible to readers. + // It is a filesystem op and therefore lives outside the SQL tx — best-effort. + if result.kept { + if let Some(content_path) = chunk_store::get_chunk_content_path(config, &chunk.id)? { + let content_root = config.memory_tree_content_root(); + let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?; + let obsidian_tags: Vec = entity_ids + .iter() + .filter_map(|eid| { + // entity_id format: "kind:surface" + let (kind, surface) = eid.split_once(':')?; + Some(content_tags::entity_tag(kind, surface)) + }) + .collect(); + + // Build the absolute path from the stored relative path. + let abs_path = { + let mut p = content_root.clone(); + for component in content_path.split('/') { + p.push(component); + } + p + }; + + if let Err(e) = content_tags::update_chunk_tags(&abs_path, &obsidian_tags) { + log::warn!( + "[memory_tree::jobs] failed to update tags in chunk file chunk_id={} path_hash={}: {e}", + chunk.id, + crate::openhuman::memory::tree::util::redact::redact(&content_path), + ); + // Non-fatal: tag rewrite failure does not block the pipeline. + } else { + log::debug!( + "[memory_tree::jobs] updated {} obsidian tags in chunk file chunk_id={}", + obsidian_tags.len(), + chunk.id, + ); + } + } + } + + // Signal workers after the tx commits (no atomicity requirement on signaling). + if did_enqueue_source { + super::worker::wake_workers(); + } + if did_enqueue_route { + super::worker::wake_workers(); + } + + Ok(()) +} + +async fn handle_append_buffer(config: &Config, job: &Job) -> Result<()> { + use crate::openhuman::memory::tree::source_tree::bucket_seal::should_seal; + use crate::openhuman::memory::tree::source_tree::store as src_store; + + let payload: AppendBufferPayload = + serde_json::from_str(&job.payload_json).context("parse AppendBuffer payload")?; + + // Hydrate the leaf-shaped record from either a chunk row or a summary + // row. The downstream buffer-push doesn't care which kind produced + // the LeafRef. + let (leaf, chunk_id_for_lifecycle): (LeafRef, Option) = match &payload.node { + NodeRef::Leaf { chunk_id } => { + let Some(chunk) = chunk_store::get_chunk(config, chunk_id)? else { + log::warn!("[memory_tree::jobs] append_buffer chunk missing chunk_id={chunk_id}"); + return Ok(()); + }; + let score_row = score_store::get_score(config, &chunk.id)? + .ok_or_else(|| anyhow::anyhow!("missing score row for chunk {}", chunk.id))?; + let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?; + // Read the full body from disk — the `content` column in SQLite + // is a ≤500-char preview after the MD-on-disk migration. The + // summariser receives this LeafRef and must see the complete text. + let body = content_read::read_chunk_body(config, chunk_id) + .with_context(|| format!("read chunk body in append_buffer chunk_id={chunk_id}"))?; + let leaf = LeafRef { + chunk_id: chunk.id.clone(), + token_count: chunk.token_count, + timestamp: chunk.metadata.timestamp, + content: body, + entities: entity_ids, + topics: chunk.metadata.tags.clone(), + score: score_row.total, + }; + (leaf, Some(chunk.id)) + } + NodeRef::Summary { summary_id } => { + let Some(summary) = src_store::get_summary(config, summary_id)? else { + log::warn!( + "[memory_tree::jobs] append_buffer summary missing summary_id={summary_id}" + ); + return Ok(()); + }; + // Read the full body from disk — `summary.content` is a ≤500-char + // preview after the MD-on-disk migration. The summariser receives + // this LeafRef when sealing higher-level nodes and must see the + // complete summary text. + let body = content_read::read_summary_body(config, summary_id).with_context(|| { + format!("read summary body in append_buffer summary_id={summary_id}") + })?; + // Build a LeafRef from the summary's already-populated fields. + // `chunk_id` carries the source-node id (any string); buffer + // accounting uses it as the item id only. + let leaf = LeafRef { + chunk_id: summary.id.clone(), + token_count: summary.token_count, + timestamp: summary.time_range_start, + content: body, + entities: summary.entities.clone(), + topics: summary.topics.clone(), + score: summary.score, + }; + (leaf, None) // summaries have no chunk lifecycle to update + } + }; + + // Resolve target tree (no tx open yet — this can create a row). + let tree = match &payload.target { + AppendTarget::Source { source_id } => Some(get_or_create_source_tree(config, source_id)?), + AppendTarget::Topic { tree_id } => src_store::get_tree(config, tree_id)?, + }; + let Some(tree) = tree else { + // Target topic tree doesn't exist (e.g. archived between + // topic_route and this append). Drop on the floor — the + // topic_route was advisory and the source-tree path already + // ran for this leaf. + return Ok(()); + }; + + let is_source_target = matches!(payload.target, AppendTarget::Source { .. }); + let leaf_for_tx = leaf.clone(); + let tree_for_tx = tree.clone(); + let lifecycle_chunk_id = chunk_id_for_lifecycle.clone(); + + // ATOMIC: buffer push + seal enqueue (if gate met) + lifecycle update + // happen in a single SQLite transaction. Eliminates the crash window + // where the buffer commits but the seal job is lost — which can + // duplicate the leaf into two summaries on retry-after-seal-cleared. + let did_enqueue_seal = chunk_store::with_connection(config, move |conn| { + let tx = conn.unchecked_transaction()?; + + // 1. Push leaf into L0 buffer (idempotent on (tree, level, item_id)). + let mut buf = src_store::get_buffer_conn(&tx, &tree_for_tx.id, 0)?; + if !buf.item_ids.iter().any(|x| x == &leaf_for_tx.chunk_id) { + buf.item_ids.push(leaf_for_tx.chunk_id.clone()); + buf.token_sum = buf.token_sum.saturating_add(leaf_for_tx.token_count as i64); + buf.oldest_at = match buf.oldest_at { + Some(existing) => Some(existing.min(leaf_for_tx.timestamp)), + None => Some(leaf_for_tx.timestamp), + }; + src_store::upsert_buffer_tx(&tx, &buf)?; + } + + // 2. If the gate is met, enqueue a seal job atomically. + let did_enqueue = if should_seal(&buf) { + let seal = SealPayload { + tree_id: tree_for_tx.id.clone(), + level: 0, + force_now_ms: None, + }; + store::enqueue_tx(&tx, &NewJob::seal(&seal)?)?.is_some() + } else { + false + }; + + // 3. Lifecycle transition (Source target with a leaf chunk). + // Last step in the tx — its presence is the "this handler + // finished" marker. Same tx as the push + seal-enqueue, so a + // crash anywhere rolls everything back together. + if is_source_target { + if let Some(chunk_id) = lifecycle_chunk_id.as_deref() { + chunk_store::set_chunk_lifecycle_status_tx( + &tx, + chunk_id, + chunk_store::CHUNK_STATUS_BUFFERED, + )?; + } + } + + tx.commit()?; + Ok(did_enqueue) + })?; + + if did_enqueue_seal { + super::worker::wake_workers(); + } + Ok(()) +} + +async fn handle_seal(config: &Config, job: &Job) -> Result<()> { + use crate::openhuman::memory::tree::source_tree::bucket_seal::{seal_one_level, should_seal}; + use crate::openhuman::memory::tree::source_tree::store as src_store; + use crate::openhuman::memory::tree::source_tree::types::TreeKind; + + let payload: SealPayload = + serde_json::from_str(&job.payload_json).context("parse Seal payload")?; + let Some(tree) = src_store::get_tree(config, &payload.tree_id)? else { + log::warn!( + "[memory_tree::jobs] seal tree missing tree_id={}", + payload.tree_id + ); + return Ok(()); + }; + + // Seal exactly one level. Parents only get sealed via a follow-up job + // so each level is its own crash-recovery checkpoint and each LLM + // summariser call competes for a fresh slot from the global semaphore. + let buf = src_store::get_buffer(config, &tree.id, payload.level)?; + let forced = payload.force_now_ms.is_some(); + if buf.is_empty() { + log::debug!( + "[memory_tree::jobs] seal skipped — empty buffer tree_id={} level={}", + tree.id, + payload.level + ); + return Ok(()); + } + if !forced && !should_seal(&buf) { + // Another job sealed this level out from under us (or the buffer + // hasn't crossed the gate yet); idempotent no-op. + log::debug!( + "[memory_tree::jobs] seal gate not met tree_id={} level={} token_sum={}", + tree.id, + payload.level, + buf.token_sum + ); + return Ok(()); + } + + // Pick the labeling strategy for this tree kind. Source trees mint + // emergent themes via the seal-time extractor; topic trees stay empty + // by design (scope already pins the canonical id). Global trees never + // reach here — `digest_daily` handles them — but Empty is a safe + // defensive default. + let strategy = match tree.kind { + TreeKind::Source => LabelStrategy::ExtractFromContent(build_summary_extractor(config)), + TreeKind::Topic => LabelStrategy::Empty, + TreeKind::Global => LabelStrategy::Empty, + }; + + let summariser = build_summariser(config); + // `seal_one_level` with `enqueue_follow_ups: true` atomically inserts + // the parent-cascade seal (if the parent buffer now meets its gate) + // and the summary-side `topic_route` (for source trees) inside the + // same SQLite transaction that commits the seal. This eliminates the + // crash window where the seal succeeds but the follow-up enqueues + // are silently lost. + let summary_id = + seal_one_level(config, &tree, &buf, summariser.as_ref(), &strategy, true).await?; + + // Phase MD-content: rewrite the `tags:` block in the sealed summary's + // on-disk .md file. Entity index rows were committed inside + // `seal_one_level` (via `index_summary_entity_ids_tx`), so they are + // visible here. Best-effort: failure does not abort the seal. + if let Err(e) = content_store::update_summary_tags(config, &summary_id) { + log::warn!( + "[memory_tree::jobs] update_summary_tags failed for summary_id={summary_id}: {e:#}" + ); + } + + super::worker::wake_workers(); + Ok(()) +} + +async fn handle_topic_route(config: &Config, job: &Job) -> Result<()> { + let payload: TopicRoutePayload = + serde_json::from_str(&job.payload_json).context("parse TopicRoute payload")?; + + // Resolve the source node id and verify it exists. `mem_tree_entity_index` + // already indexes both chunks and summaries via `node_kind`, so the + // canonical-id loop below is identical for either case. + let node_id: String = match &payload.node { + NodeRef::Leaf { chunk_id } => { + if chunk_store::get_chunk(config, chunk_id)?.is_none() { + log::warn!("[memory_tree::jobs] topic_route chunk missing chunk_id={chunk_id}"); + return Ok(()); + } + chunk_id.clone() + } + NodeRef::Summary { summary_id } => { + if crate::openhuman::memory::tree::source_tree::store::get_summary(config, summary_id)? + .is_none() + { + log::warn!( + "[memory_tree::jobs] topic_route summary missing summary_id={summary_id}" + ); + return Ok(()); + } + summary_id.clone() + } + }; + + let entity_ids = score_store::list_entity_ids_for_node(config, &node_id)?; + if entity_ids.is_empty() { + log::debug!("[memory_tree::jobs] topic_route no entities for node_id={node_id} — skipping"); + return Ok(()); + } + + let summariser = build_summariser(config); + for entity_id in entity_ids { + let _ = curator::maybe_spawn_topic_tree(config, &entity_id, summariser.as_ref()).await?; + if let Some(tree) = crate::openhuman::memory::tree::source_tree::store::get_tree_by_scope( + config, + crate::openhuman::memory::tree::source_tree::types::TreeKind::Topic, + &entity_id, + )? { + let job = NewJob::append_buffer(&AppendBufferPayload { + node: payload.node.clone(), + target: AppendTarget::Topic { + tree_id: tree.id.clone(), + }, + })?; + if store::enqueue(config, &job)?.is_some() { + super::worker::wake_workers(); + } + } + } + Ok(()) +} + +async fn handle_digest_daily(config: &Config, job: &Job) -> Result<()> { + let payload: DigestDailyPayload = + serde_json::from_str(&job.payload_json).context("parse DigestDaily payload")?; + let day = chrono::NaiveDate::parse_from_str(&payload.date_iso, "%Y-%m-%d") + .with_context(|| format!("invalid digest date {}", payload.date_iso))?; + let summariser = build_summariser(config); + match digest::end_of_day_digest(config, day, summariser.as_ref()).await? { + DigestOutcome::Emitted { daily_id, .. } => { + log::info!("[memory_tree::jobs] emitted digest daily_id={daily_id}"); + } + DigestOutcome::EmptyDay => {} + DigestOutcome::Skipped { existing_id } => { + log::debug!("[memory_tree::jobs] digest skipped existing_id={existing_id}"); + } + } + Ok(()) +} + +async fn handle_flush_stale(config: &Config, job: &Job) -> Result<()> { + let payload: FlushStalePayload = + serde_json::from_str(&job.payload_json).context("parse FlushStale payload")?; + let age_secs = payload + .max_age_secs + .unwrap_or(crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS); + let cutoff = chrono::Utc::now() - chrono::Duration::seconds(age_secs); + let buffers = + crate::openhuman::memory::tree::source_tree::store::list_stale_buffers(config, cutoff)?; + for buf in buffers { + let seal = SealPayload { + tree_id: buf.tree_id.clone(), + level: buf.level, + force_now_ms: Some(chrono::Utc::now().timestamp_millis()), + }; + if store::enqueue(config, &NewJob::seal(&seal)?)?.is_some() { + super::worker::wake_workers(); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::content_store; + use crate::openhuman::memory::tree::jobs::store::{count_by_status, count_total}; + use crate::openhuman::memory::tree::jobs::types::JobStatus; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf_deferred, LeafRef}; + use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; + use crate::openhuman::memory::tree::source_tree::store as src_store; + use crate::openhuman::memory::tree::store::with_connection; + use chrono::TimeZone; + use rusqlite::params; + 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) + } + + /// Build a minimal `Job` row for direct handler invocation. Mirrors + /// what `claim_next` would produce for a freshly-claimed row. + fn mk_running_job(kind: JobKind, payload_json: String) -> Job { + let now_ms = chrono::Utc::now().timestamp_millis(); + Job { + id: "test-job-id".into(), + kind, + payload_json, + dedupe_key: None, + status: JobStatus::Running, + attempts: 1, + max_attempts: 5, + available_at_ms: now_ms, + locked_until_ms: Some(now_ms + 60_000), + last_error: None, + created_at_ms: now_ms, + started_at_ms: Some(now_ms), + completed_at_ms: None, + } + } + + /// Count rows in `mem_tree_jobs` matching a specific kind. + fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 { + with_connection(cfg, |conn| { + let n: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1", + params![kind], + |r| r.get(0), + )?; + Ok(n.max(0) as u64) + }) + .unwrap() + } + + /// Seed a source tree and push enough labeled leaves into its L0 buffer + /// to cross `TOKEN_BUDGET`, returning the tree. The caller can then + /// fire `handle_seal` and inspect the result. + async fn seed_source_tree_ready_to_seal( + cfg: &Config, + ) -> crate::openhuman::memory::tree::source_tree::types::Tree { + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap(); + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"), + content: "alice@example.com leading the rollout".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")), + }, + // Bust budget so the L0 buffer is "ready" for seal. + token_count: 10_000, + seq_in_source: 0, + created_at: ts, + partial_message: false, + }; + upsert_chunks(cfg, &[chunk.clone()]).unwrap(); + // Stage to disk so `hydrate_leaf_inputs` can read the full body via + // `read_chunk_body` when `handle_seal` fires and calls `seal_one_level`. + let content_root = cfg.memory_tree_content_root(); + std::fs::create_dir_all(&content_root).unwrap(); + let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap(); + with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 10_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + // append_leaf_deferred only buffers; doesn't seal. handle_seal will. + let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap(); + tree + } + + #[tokio::test] + async fn source_tree_seal_handler_enqueues_summary_topic_route() { + let (_tmp, cfg) = test_config(); + let tree = seed_source_tree_ready_to_seal(&cfg).await; + + let payload = SealPayload { + tree_id: tree.id.clone(), + level: 0, + force_now_ms: None, + }; + let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap()); + + // Pre-condition: queue has no topic_route jobs. + assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0); + + super::handle_seal(&cfg, &job).await.unwrap(); + + // Post-condition: source-tree seal must enqueue exactly one + // topic_route job carrying NodeRef::Summary { summary_id: }. + assert_eq!( + count_jobs_of_kind(&cfg, "topic_route"), + 1, + "source-tree seal must enqueue summary-side topic_route" + ); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1); + + // Inspect the enqueued payload to confirm it's a Summary variant. + let payload_json: String = with_connection(&cfg, |conn| { + let s: String = conn + .query_row( + "SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'", + [], + |r| r.get(0), + ) + .unwrap(); + Ok(s) + }) + .unwrap(); + let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap(); + match p.node { + NodeRef::Summary { summary_id } => { + assert!(summary_id.starts_with("summary:L1:")); + } + other => panic!("expected NodeRef::Summary, got {other:?}"), + } + } + + #[tokio::test] + async fn topic_tree_seal_handler_does_not_enqueue_topic_route() { + let (_tmp, cfg) = test_config(); + // Spawn a topic tree directly via the registry (skipping curator's + // hotness gate — we just need a TreeKind::Topic with leaves). + let topic_tree = + crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree( + &cfg, + "topic:phoenix-migration", + ) + .unwrap(); + // Push a single 10k-token leaf so L0 is gate-ready. + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"), + content: "topic content".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")), + }, + token_count: 10_000, + seq_in_source: 0, + created_at: ts, + partial_message: false, + }; + upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); + // Stage to disk so `hydrate_leaf_inputs` can read the full body + // when `handle_seal` fires. + let content_root = cfg.memory_tree_content_root(); + std::fs::create_dir_all(&content_root).unwrap(); + let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 10_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap(); + + let payload = SealPayload { + tree_id: topic_tree.id.clone(), + level: 0, + force_now_ms: None, + }; + let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap()); + + super::handle_seal(&cfg, &job).await.unwrap(); + + // Topic-tree seals are sinks: must not enqueue any topic_route. + assert_eq!( + count_jobs_of_kind(&cfg, "topic_route"), + 0, + "topic-tree seal must NOT enqueue topic_route (trees are sinks)" + ); + // The seal itself should still have produced a summary node. + assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1); + } + + #[tokio::test] + async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() { + let (_tmp, cfg) = test_config(); + + // 1. Create a target topic tree with a clean L0 buffer. + let topic_tree = + crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree( + &cfg, + "email:alice@example.com", + ) + .unwrap(); + let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap(); + assert!(l0_before.is_empty()); + + // 2. Manually insert a summary node we can route. The simplest way + // is to create a separate source tree, push two 6k leaves into + // it, and let the seal produce a summary we can address. + let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + use crate::openhuman::memory::tree::source_tree::bucket_seal::seal_one_level; + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let content_root = cfg.memory_tree_content_root(); + std::fs::create_dir_all(&content_root).unwrap(); + for seq in 0..2 { + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"), + content: format!("source 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")), + }, + token_count: 6_000, + seq_in_source: seq, + created_at: ts, + partial_message: false, + }; + upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); + // Stage to disk so `hydrate_leaf_inputs` can read the full body + // during `seal_one_level`. + let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 6_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap(); + } + // Force-seal the source tree's L0 to mint the summary. + let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap(); + let summariser = build_summariser(&cfg); + let summary_id = seal_one_level( + &cfg, + &source_tree, + &buf, + summariser.as_ref(), + &crate::openhuman::memory::tree::source_tree::bucket_seal::LabelStrategy::Empty, + // No follow-up enqueues — the test scopes assertions to the + // append_buffer handler, not seal-side fan-out. + false, + ) + .await + .unwrap(); + + // 3. Build an append_buffer payload routing the summary into the + // topic tree. + let payload = AppendBufferPayload { + node: NodeRef::Summary { + summary_id: summary_id.clone(), + }, + target: AppendTarget::Topic { + tree_id: topic_tree.id.clone(), + }, + }; + let job = mk_running_job( + JobKind::AppendBuffer, + serde_json::to_string(&payload).unwrap(), + ); + + // Clear out any pending append_buffer jobs minted upstream so the + // post-condition assertion below is unambiguous. + let pre = count_total(&cfg).unwrap(); + + super::handle_append_buffer(&cfg, &job).await.unwrap(); + + // 4. Topic tree's L0 buffer should now hold the summary id. + let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap(); + assert_eq!(l0_after.item_ids, vec![summary_id]); + assert!(l0_after.token_sum > 0); + + // No new jobs should have been enqueued (buffer didn't cross gate). + assert_eq!(count_total(&cfg).unwrap(), pre); + } +} diff --git a/src/openhuman/memory/tree/jobs/mod.rs b/src/openhuman/memory/tree/jobs/mod.rs new file mode 100644 index 000000000..bcc5aec95 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/mod.rs @@ -0,0 +1,45 @@ +//! Async job pipeline for memory-tree work. +//! +//! Replaces the previous synchronous `append_leaf → cascade_seal → LLM +//! summarise` chain on the ingest hot path with a SQLite-backed job queue +//! and a worker pool. The shape is: +//! +//! ```text +//! ingest::persist +//! └── writes chunk row (lifecycle = pending_extraction) +//! enqueues `extract_chunk` +//! +//! worker pool (3 tasks) ──► claims jobs by kind: +//! extract_chunk → LLM extraction → admission decision → enqueue append_buffer +//! append_buffer → push to L0 → enqueue seal if gate met → enqueue topic_route +//! seal → seal one level → enqueue parent seal if cascading +//! topic_route → match topics → enqueue per-topic append_buffer +//! digest_daily → call global_tree::digest::end_of_day_digest +//! flush_stale → enqueue seals for time-stale buffers +//! +//! scheduler (1 task) ──► daily wall-clock tick: +//! enqueues digest_daily(yesterday) + flush_stale(today) +//! ``` +//! +//! All persistence lives in the same `chunks.db` as `mem_tree_chunks` so a +//! producer can insert its side-effect and its follow-up job in one tx. +//! See [`store::enqueue_tx`] for the in-tx producer entry point. + +mod handlers; +pub mod scheduler; +pub mod store; +pub mod testing; +pub mod types; +mod worker; + +pub use scheduler::{backfill_missing_digests, trigger_digest}; +pub use store::{ + claim_next, count_by_status, count_total, enqueue, enqueue_tx, get_job, mark_done, mark_failed, + recover_stale_locks, DEFAULT_LOCK_DURATION_MS, +}; +pub use testing::drain_until_idle; +pub use types::{ + AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload, + Job, JobKind, JobStatus, NewJob, NodeRef, SealPayload, TopicRoutePayload, +}; +pub use worker::{start, wake_workers}; diff --git a/src/openhuman/memory/tree/jobs/scheduler.rs b/src/openhuman/memory/tree/jobs/scheduler.rs new file mode 100644 index 000000000..ea77035d2 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/scheduler.rs @@ -0,0 +1,216 @@ +use std::time::Duration; + +use anyhow::Result; +use chrono::{Datelike, Duration as ChronoDuration, NaiveDate, TimeZone, Utc}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::store; +use crate::openhuman::memory::tree::jobs::types::{DigestDailyPayload, FlushStalePayload, NewJob}; + +static STARTED: std::sync::Once = std::sync::Once::new(); + +/// Start the daily wall-clock scheduler. Takes the full `Config` so the +/// digest enqueues match the same workspace + LLM settings the workers +/// see — not `Config::default()`. +pub fn start(config: Config) { + STARTED.call_once(|| { + tokio::spawn(async move { + loop { + if let Err(err) = enqueue_daily_jobs(&config) { + log::warn!("[memory_tree::jobs] scheduler enqueue failed: {err:#}"); + } + tokio::time::sleep(next_sleep_duration()).await; + } + }); + }); +} + +fn enqueue_daily_jobs(config: &Config) -> anyhow::Result<()> { + let now = Utc::now(); + let yesterday = now.date_naive() - ChronoDuration::days(1); + let date_iso = yesterday.format("%Y-%m-%d").to_string(); + + if store::enqueue( + config, + &NewJob::digest_daily(&DigestDailyPayload { + date_iso: date_iso.clone(), + })?, + )? + .is_some() + { + super::worker::wake_workers(); + } + + let today_iso = now.date_naive().format("%Y-%m-%d").to_string(); + if store::enqueue( + config, + &NewJob::flush_stale(&FlushStalePayload::default(), &today_iso)?, + )? + .is_some() + { + super::worker::wake_workers(); + } + + Ok(()) +} + +/// Manually enqueue a `digest_daily` job for `date`. Idempotent — if a +/// digest already ran for that day, the handler's `find_existing_daily` +/// check will return `Skipped` without doing any work; if a job for the +/// same date is already queued or running, the partial unique index on +/// `dedupe_key` suppresses the duplicate. +/// +/// Useful for catch-up after the process was down across midnight, or +/// to force a re-run for testing / debugging. +pub fn trigger_digest(config: &Config, date: NaiveDate) -> Result> { + let payload = DigestDailyPayload { + date_iso: date.format("%Y-%m-%d").to_string(), + }; + let job_id = store::enqueue(config, &NewJob::digest_daily(&payload)?)?; + if job_id.is_some() { + log::info!( + "[memory_tree::jobs] manual digest trigger enqueued date={} id={:?}", + payload.date_iso, + job_id.as_deref() + ); + super::worker::wake_workers(); + } else { + log::debug!( + "[memory_tree::jobs] manual digest trigger dedupe-suppressed date={} \ + (an active job for this date already exists)", + payload.date_iso + ); + } + Ok(job_id) +} + +/// Enqueue `digest_daily` jobs for the last `days_back` calendar days +/// (excluding today). Catch-up helper for cases where the scheduler +/// missed days because the process was down. +/// +/// Returns the number of jobs newly enqueued. Days that already have a +/// completed digest are still re-enqueued — the handler is idempotent +/// and skips them — so this is safe to call repeatedly. +pub fn backfill_missing_digests(config: &Config, days_back: i64) -> Result { + if days_back <= 0 { + return Ok(0); + } + let today = Utc::now().date_naive(); + let mut enqueued = 0usize; + for offset in 1..=days_back { + let date = today - ChronoDuration::days(offset); + if trigger_digest(config, date)?.is_some() { + enqueued += 1; + } + } + log::info!( + "[memory_tree::jobs] backfill_missing_digests window={}d enqueued={}", + days_back, + enqueued + ); + Ok(enqueued) +} + +fn next_sleep_duration() -> Duration { + let now = Utc::now(); + let tomorrow = now.date_naive() + ChronoDuration::days(1); + let next = Utc + .with_ymd_and_hms(tomorrow.year(), tomorrow.month(), tomorrow.day(), 0, 5, 0) + // UTC has no DST gaps/overlaps, so `single()` always returns + // `Some` for any valid (Y, M, D, h, m, s). Fallback retained + // only as a defensive belt-and-braces against future API churn. + .single() + .unwrap_or_else(|| now + ChronoDuration::hours(24)); + (next - now) + .to_std() + .unwrap_or_else(|_| Duration::from_secs(60)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::store::{ + claim_next, count_by_status, count_total, mark_done, DEFAULT_LOCK_DURATION_MS, + }; + use crate::openhuman::memory::tree::jobs::types::JobStatus; + 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 trigger_digest_enqueues_a_job() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let id = trigger_digest(&cfg, date).unwrap(); + assert!(id.is_some(), "first trigger must enqueue"); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1); + } + + #[test] + fn trigger_digest_dedupes_active_jobs() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let first = trigger_digest(&cfg, date).unwrap(); + let second = trigger_digest(&cfg, date).unwrap(); + assert!(first.is_some()); + assert!( + second.is_none(), + "duplicate trigger must be dedupe-suppressed while active" + ); + assert_eq!(count_total(&cfg).unwrap(), 1); + } + + #[test] + fn trigger_digest_after_done_creates_fresh_row() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let id1 = trigger_digest(&cfg, date).unwrap().unwrap(); + // Simulate a worker finishing the job — claim it first so we have a + // Job snapshot for the claim-token-gated mark_done. + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + assert_eq!(claimed.id, id1); + mark_done(&cfg, &claimed).unwrap(); + + let id2 = trigger_digest(&cfg, date).unwrap(); + assert!( + id2.is_some(), + "after the prior job completes, a fresh trigger must enqueue" + ); + assert_ne!(id2.unwrap(), id1); + assert_eq!(count_total(&cfg).unwrap(), 2); + } + + #[test] + fn backfill_missing_digests_enqueues_one_per_day() { + let (_tmp, cfg) = test_config(); + let n = backfill_missing_digests(&cfg, 5).unwrap(); + assert_eq!(n, 5, "expected one job per day in the 5-day window"); + assert_eq!(count_total(&cfg).unwrap(), 5); + } + + #[test] + fn backfill_missing_digests_zero_window_is_noop() { + let (_tmp, cfg) = test_config(); + let n = backfill_missing_digests(&cfg, 0).unwrap(); + assert_eq!(n, 0); + assert_eq!(count_total(&cfg).unwrap(), 0); + } + + #[test] + fn backfill_missing_digests_is_idempotent_while_active() { + let (_tmp, cfg) = test_config(); + let n1 = backfill_missing_digests(&cfg, 3).unwrap(); + let n2 = backfill_missing_digests(&cfg, 3).unwrap(); + assert_eq!(n1, 3); + assert_eq!(n2, 0, "second call must be fully dedupe-suppressed"); + assert_eq!(count_total(&cfg).unwrap(), 3); + } +} diff --git a/src/openhuman/memory/tree/jobs/store.rs b/src/openhuman/memory/tree/jobs/store.rs new file mode 100644 index 000000000..53931c587 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/store.rs @@ -0,0 +1,618 @@ +//! 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. + +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::{params, Connection, OptionalExtension, Transaction}; +use uuid::Uuid; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::types::{Job, JobKind, JobStatus, NewJob}; +use crate::openhuman::memory::tree::store::with_connection; + +/// 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; + +/// 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; + +/// 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)) +} + +/// 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) +} + +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_tree::jobs] enqueue suppressed by dedupe kind={} key={:?}", + job.kind.as_str(), + job.dedupe_key + ); + return Ok(None); + } + log::debug!( + "[memory_tree::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( + "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 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", + 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_tree::jobs] claimed id={} kind={} attempt={}/{}", + j.id, + j.kind.as_str(), + j.attempts, + j.max_attempts + ); + } + Ok(row) + }) +} + +/// 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_tree::jobs] mark_done id={job_id} was a no-op \ + (stale lease: attempts={claim_attempts} started_at_ms={claim_started_at:?})" + ); + } + Ok(()) + }) +} + +/// 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<()> { + 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; + with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + + if attempts >= max_attempts { + log::warn!( + "[memory_tree::jobs] terminal failure id={job_id} \ + attempts={attempts}/{max_attempts} err={error}" + ); + let n = conn.execute( + "UPDATE mem_tree_jobs + SET status = 'failed', + completed_at_ms = ?1, + locked_until_ms = NULL, + last_error = ?2 + WHERE id = ?3 + AND attempts = ?4 + AND started_at_ms IS ?5", + params![now_ms, error, job_id, attempts, claim_started_at], + )?; + if n == 0 { + log::warn!( + "[memory_tree::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_tree::jobs] retry id={job_id} attempt={attempts}/{max_attempts} \ + next_at_ms={next_at} err={error}" + ); + 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_tree::jobs] mark_failed(retry) id={job_id} was a no-op \ + (stale lease: attempts={attempts} started_at_ms={claim_started_at:?})" + ); + } + } + Ok(()) + }) +} + +/// 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 { + log::warn!("[memory_tree::jobs] recovered {n} stale-locked job(s) at startup"); + } + Ok(n) + }) +} + +/// 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) + }) +} + +/// 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) + }) +} + +/// 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 + 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 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, + 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::tree::jobs::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()); + } + + #[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()); + } + + #[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); + } + + /// 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); + } +} diff --git a/src/openhuman/memory/tree/jobs/testing.rs b/src/openhuman/memory/tree/jobs/testing.rs new file mode 100644 index 000000000..e8211fb78 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/testing.rs @@ -0,0 +1,15 @@ +use anyhow::Result; + +use crate::openhuman::config::Config; + +/// Deterministically run queued memory-tree jobs until no immediately +/// claimable work remains. Intended for tests that need the async pipeline +/// to settle without spawning background tasks. +pub async fn drain_until_idle(config: &Config) -> Result<()> { + loop { + if !super::worker::run_once(config).await? { + break; + } + } + Ok(()) +} diff --git a/src/openhuman/memory/tree/jobs/types.rs b/src/openhuman/memory/tree/jobs/types.rs new file mode 100644 index 000000000..e6bf044f2 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/types.rs @@ -0,0 +1,439 @@ +//! 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`. + +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, + /// Match a chunk's entities against active topic trees and enqueue + /// per-topic `AppendBuffer` jobs. + TopicRoute, + /// Build the global tree's daily digest for a given UTC date. + DigestDaily, + /// Walk stale buffers and enqueue `Seal` jobs for any over the age cap. + FlushStale, +} + +impl JobKind { + pub fn as_str(&self) -> &'static str { + match self { + JobKind::ExtractChunk => "extract_chunk", + JobKind::AppendBuffer => "append_buffer", + JobKind::Seal => "seal", + JobKind::TopicRoute => "topic_route", + JobKind::DigestDaily => "digest_daily", + JobKind::FlushStale => "flush_stale", + } + } + + pub fn parse(s: &str) -> Result { + Ok(match s { + "extract_chunk" => JobKind::ExtractChunk, + "append_buffer" => JobKind::AppendBuffer, + "seal" => JobKind::Seal, + "topic_route" => JobKind::TopicRoute, + "digest_daily" => JobKind::DigestDaily, + "flush_stale" => JobKind::FlushStale, + other => return Err(anyhow!("unknown JobKind '{other}'")), + }) + } + + /// True when handling this kind should hold a slot from the global + /// LLM concurrency semaphore. `TopicRoute` is bound because + /// `maybe_spawn_topic_tree → backfill_topic_tree` can transitively + /// trigger summariser LLM calls when an entity first crosses the + /// hotness threshold. + pub fn is_llm_bound(&self) -> bool { + matches!( + self, + JobKind::ExtractChunk | JobKind::Seal | JobKind::DigestDaily | JobKind::TopicRoute + ) + } +} + +/// 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 { + pub fn as_str(&self) -> &'static str { + match self { + JobStatus::Ready => "ready", + JobStatus::Running => "running", + JobStatus::Done => "done", + JobStatus::Failed => "failed", + JobStatus::Cancelled => "cancelled", + } + } + + 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}'")), + }) + } + + 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, 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 { + 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 { + 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 { + 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)] +pub struct TopicRoutePayload { + pub node: NodeRef, +} + +impl TopicRoutePayload { + pub fn dedupe_key(&self) -> String { + format!("topic_route:{}", self.node.dedupe_fragment()) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct DigestDailyPayload { + /// UTC calendar date in `YYYY-MM-DD` form. Stored as a string so the + /// dedupe key doesn't need to know about chrono. + pub date_iso: String, +} + +impl DigestDailyPayload { + pub fn dedupe_key(&self) -> String { + format!("digest_daily:{}", self.date_iso) + } +} + +#[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 { + pub fn dedupe_key(&self, date_iso: &str) -> String { + format!("flush_stale:{date_iso}") + } +} + +/// 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, + 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 { + 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, + }) + } + + 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, + }) + } + + 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, + }) + } + + pub fn topic_route(p: &TopicRoutePayload) -> Result { + Ok(Self { + kind: JobKind::TopicRoute, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn digest_daily(p: &DigestDailyPayload) -> Result { + Ok(Self { + kind: JobKind::DigestDaily, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn flush_stale(p: &FlushStalePayload, date_iso: &str) -> Result { + Ok(Self { + kind: JobKind::FlushStale, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key(date_iso)), + 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::TopicRoute, + JobKind::DigestDaily, + JobKind::FlushStale, + ] { + assert_eq!(JobKind::parse(k.as_str()).unwrap(), k); + } + } + + #[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()); + + let r_leaf = TopicRoutePayload { + node: NodeRef::Leaf { + chunk_id: "x".into(), + }, + }; + let r_summary = TopicRoutePayload { + node: NodeRef::Summary { + summary_id: "x".into(), + }, + }; + assert_ne!(r_leaf.dedupe_key(), r_summary.dedupe_key()); + } + + #[test] + fn llm_bound_kinds() { + assert!(JobKind::ExtractChunk.is_llm_bound()); + assert!(JobKind::Seal.is_llm_bound()); + assert!(JobKind::DigestDaily.is_llm_bound()); + assert!(JobKind::TopicRoute.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"), + } + } +} diff --git a/src/openhuman/memory/tree/jobs/worker.rs b/src/openhuman/memory/tree/jobs/worker.rs new file mode 100644 index 000000000..76b101db0 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/worker.rs @@ -0,0 +1,109 @@ +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use anyhow::Result; +use tokio::sync::{Notify, Semaphore}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::handlers; +use crate::openhuman::memory::tree::jobs::store::{ + claim_next, mark_done, mark_failed, recover_stale_locks, DEFAULT_LOCK_DURATION_MS, +}; + +const WORKER_COUNT: usize = 3; +const POLL_INTERVAL: Duration = Duration::from_secs(5); + +static WORKER_NOTIFY: OnceLock> = OnceLock::new(); +static STARTED: std::sync::Once = std::sync::Once::new(); + +pub fn wake_workers() { + if let Some(notify) = WORKER_NOTIFY.get() { + notify.notify_waiters(); + } +} + +/// Start the worker pool + daily scheduler. Takes the full `Config` so +/// each spawned task sees the user's actual settings (LLM endpoints, +/// embedder model, timeouts) — not `Config::default()`. Without this, +/// workers fall back to inert/regex-only behavior regardless of what's +/// in `config.toml`, defeating the entire async pipeline. +/// +/// Idempotent (`Once`-guarded) so repeat calls during bootstrap are +/// safe no-ops after the first. +pub fn start(config: Config) { + STARTED.call_once(|| { + let notify = WORKER_NOTIFY + .get_or_init(|| Arc::new(Notify::new())) + .clone(); + let llm_slots = Arc::new(Semaphore::new(3)); + if let Err(err) = recover_stale_locks(&config) { + log::warn!("[memory_tree::jobs] recover_stale_locks failed at startup: {err:#}"); + } + + for idx in 0..WORKER_COUNT { + let notify = notify.clone(); + let llm_slots = llm_slots.clone(); + let cfg = config.clone(); + tokio::spawn(async move { + loop { + match run_once_with_semaphore(&cfg, llm_slots.clone()).await { + Ok(true) => continue, + Ok(false) => { + tokio::select! { + _ = notify.notified() => {} + _ = tokio::time::sleep(POLL_INTERVAL) => {} + } + } + Err(err) => { + log::error!("[memory_tree::jobs] worker={} loop error: {:#}", idx, err); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } + } + }); + } + + super::scheduler::start(config); + }); +} + +pub async fn run_once(config: &Config) -> Result { + let llm_slots = Arc::new(Semaphore::new(1)); + run_once_with_semaphore(config, llm_slots).await +} + +async fn run_once_with_semaphore(config: &Config, llm_slots: Arc) -> Result { + let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else { + return Ok(false); + }; + + let permit = if job.kind.is_llm_bound() { + Some(llm_slots.acquire().await?) + } else { + None + }; + let result = handlers::handle_job(config, &job).await; + drop(permit); + + match result { + Ok(()) => { + mark_done(config, &job)?; + } + Err(err) => { + // Preserve the full anyhow cause chain in the persisted + // last_error so a reader of mem_tree_jobs can see the root + // cause, not just the top-level message. Mirrors the {:#} + // log format used right above. + let message = format!("{err:#}"); + log::warn!( + "[memory_tree::jobs] job failed id={} kind={} err={:#}", + job.id, + job.kind.as_str(), + err + ); + mark_failed(config, &job, &message)?; + } + } + + Ok(true) +} diff --git a/src/openhuman/memory/tree/mod.rs b/src/openhuman/memory/tree/mod.rs index 038b352d3..7a10d546c 100644 --- a/src/openhuman/memory/tree/mod.rs +++ b/src/openhuman/memory/tree/mod.rs @@ -24,8 +24,10 @@ pub mod canonicalize; pub mod chunker; +pub mod content_store; pub mod global_tree; pub mod ingest; +pub mod jobs; pub mod retrieval; pub mod rpc; pub mod schemas; @@ -34,6 +36,7 @@ pub mod source_tree; pub mod store; pub mod topic_tree; pub mod types; +pub mod util; pub use retrieval::{all_retrieval_controller_schemas, all_retrieval_registered_controllers}; pub use schemas::{ diff --git a/src/openhuman/memory/tree/retrieval/drill_down.rs b/src/openhuman/memory/tree/retrieval/drill_down.rs index 9a307bfb1..8b30ad257 100644 --- a/src/openhuman/memory/tree/retrieval/drill_down.rs +++ b/src/openhuman/memory/tree/retrieval/drill_down.rs @@ -23,6 +23,7 @@ use std::collections::VecDeque; use anyhow::Result; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::read as content_read; use crate::openhuman::memory::tree::retrieval::types::{ hit_from_chunk, hit_from_summary, RetrievalHit, }; @@ -179,15 +180,27 @@ fn walk_with_embeddings( continue; } // Is it a summary? - if let Some(summary) = store::get_summary(config, &id)? { + if let Some(mut summary) = store::get_summary(config, &id)? { let scope = store::get_tree(config, &summary.tree_id)? .map(|t| t.scope) .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 { - for next in summary.child_ids { + for next in child_ids { frontier.push_back((next, depth + 1)); } } @@ -195,11 +208,21 @@ fn walk_with_embeddings( } // Else try as a chunk (leaf). Chunk embeddings live in a separate // blob column — fetch via the existing accessor. - if let Some(chunk) = get_chunk(config, &id)? { + if let Some(mut chunk) = get_chunk(config, &id)? { // Propagate DB errors rather than silently treating them as // "no embedding" — the caller should know if the store is broken. let emb = get_chunk_embedding(config, &chunk.id)?; 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; @@ -215,7 +238,10 @@ fn walk_with_embeddings( #[cfg(test)] mod tests { use super::*; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::content_store; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::source_tree::types::TreeKind; @@ -240,6 +266,8 @@ mod tests { let ts = Utc::now(); let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap(); let summariser = InertSummariser::new(); + 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 { @@ -254,18 +282,32 @@ mod tests { tags: vec![], source_ref: Some(SourceRef::new("slack://x")), }, - token_count: 6_000, + token_count: crate::openhuman::memory::tree::source_tree::types::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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); leaf_ids.push(c.id.clone()); append_leaf( cfg, &tree, &LeafRef { chunk_id: c.id.clone(), - token_count: 6_000, + token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET + * 6 + / 10, timestamp: ts, content: c.content.clone(), entities: vec![], @@ -273,6 +315,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); @@ -401,6 +444,7 @@ mod tests { token_count: 10, seq_in_source: 0, created_at: ts, + partial_message: false, }; let leaf_a_2 = Chunk { id: "chat:slack:#eng:1".into(), @@ -423,15 +467,24 @@ mod tests { seq_in_source: 3, ..leaf_a_1.clone() }; - upsert_chunks( - cfg, - &[ - leaf_a_1.clone(), - leaf_a_2.clone(), - leaf_b_1.clone(), - leaf_b_2.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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) .unwrap(); let l1_a = SummaryNode { @@ -471,9 +524,9 @@ mod tests { with_connection(cfg, |conn| { let tx = conn.unchecked_transaction()?; tree_store::insert_tree_conn(&tx, &tree)?; - tree_store::insert_summary_tx(&tx, &l1_a)?; - tree_store::insert_summary_tx(&tx, &l1_b)?; - tree_store::insert_summary_tx(&tx, &root)?; + tree_store::insert_summary_tx(&tx, &l1_a, None)?; + tree_store::insert_summary_tx(&tx, &l1_b, None)?; + tree_store::insert_summary_tx(&tx, &root, None)?; tx.commit()?; Ok(()) }) diff --git a/src/openhuman/memory/tree/retrieval/fetch.rs b/src/openhuman/memory/tree/retrieval/fetch.rs index 3eced1209..df9522493 100644 --- a/src/openhuman/memory/tree/retrieval/fetch.rs +++ b/src/openhuman/memory/tree/retrieval/fetch.rs @@ -13,6 +13,7 @@ use anyhow::Result; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::read as content_read; use crate::openhuman::memory::tree::retrieval::types::{hit_from_chunk, RetrievalHit}; use crate::openhuman::memory::tree::score::store::get_score; use crate::openhuman::memory::tree::store::get_chunk; @@ -66,7 +67,22 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result 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) }) @@ -80,11 +96,26 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result (TempDir, Config) { let tmp = TempDir::new().unwrap(); let mut cfg = Config::default(); @@ -113,6 +144,7 @@ mod tests { token_count: 20, seq_in_source: seq, created_at: ts, + partial_message: false, } } @@ -129,6 +161,7 @@ mod tests { 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(); @@ -142,6 +175,7 @@ mod tests { 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()], @@ -159,6 +193,7 @@ mod tests { 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(); @@ -170,6 +205,7 @@ mod tests { 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")); diff --git a/src/openhuman/memory/tree/retrieval/global.rs b/src/openhuman/memory/tree/retrieval/global.rs index fd5199f29..9d6e22bf4 100644 --- a/src/openhuman/memory/tree/retrieval/global.rs +++ b/src/openhuman/memory/tree/retrieval/global.rs @@ -87,8 +87,11 @@ fn recap_to_hits(recap: RecapOutput, tree_id: &str, tree_scope: &str) -> Vec (TempDir, Config) { let tmp = TempDir::new().unwrap(); let mut cfg = Config::default(); @@ -134,8 +151,10 @@ mod tests { token_count: 6_000, seq_in_source: seq, created_at: ts, + partial_message: false, }; upsert_chunks(cfg, &[c.clone()]).unwrap(); + stage_test_chunks(cfg, &[c.clone()]); append_leaf( cfg, &tree, @@ -149,6 +168,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); diff --git a/src/openhuman/memory/tree/retrieval/integration_test.rs b/src/openhuman/memory/tree/retrieval/integration_test.rs index 95cfe3c6b..736f56d95 100644 --- a/src/openhuman/memory/tree/retrieval/integration_test.rs +++ b/src/openhuman/memory/tree/retrieval/integration_test.rs @@ -160,11 +160,11 @@ async fn topic_entity_surfaces_after_ingest() { // ── Phase 4 (#710): embedding + semantic rerank tests ─────────────────── /// Ingest with an inert embedder must populate every kept chunk's -/// `embedding` column. This guards against regressions where the embed -/// step is silently skipped (e.g. future refactors threading embeddings -/// through a different code path). +/// `embedding` column. Embeddings are written by the async `extract_chunk` +/// handler, so the test drains the queue before inspecting. #[tokio::test] async fn ingest_populates_chunk_embeddings() { + use crate::openhuman::memory::tree::jobs::drain_until_idle; use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM; use crate::openhuman::memory::tree::store::get_chunk_embedding; @@ -172,7 +172,11 @@ async fn ingest_populates_chunk_embeddings() { let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], chat_about_phoenix(0)) .await .unwrap(); - assert!(out.chunks_written >= 1, "expected at least one kept chunk"); + assert!( + out.chunks_written >= 1, + "expected at least one persisted chunk" + ); + drain_until_idle(&cfg).await.unwrap(); 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}")); @@ -188,8 +192,11 @@ async fn ingest_populates_chunk_embeddings() { /// the seal from firing on short batches. #[tokio::test] async fn seal_populates_summary_embedding() { + use crate::openhuman::memory::tree::content_store; use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::store as src_store; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; @@ -216,10 +223,24 @@ async fn seal_populates_summary_embedding() { token_count: tokens, seq_in_source: seq, created_at: ts, + partial_message: false, }; let c1 = mk_chunk(0, 6_000); let c2 = mk_chunk(1, 6_000); upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); + { + 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, &[c1.clone(), c2.clone()]) + .expect("stage_chunks for test chunks"); + crate::openhuman::memory::tree::store::with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .expect("persist staged chunk pointers"); + } let leaf_of = |c: &Chunk| LeafRef { chunk_id: c.id.clone(), @@ -230,12 +251,24 @@ async fn seal_populates_summary_embedding() { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf_of(&c1), &summariser) - .await - .unwrap(); - let sealed = append_leaf(&cfg, &tree, &leaf_of(&c2), &summariser) - .await - .unwrap(); + append_leaf( + &cfg, + &tree, + &leaf_of(&c1), + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap(); + let sealed = append_leaf( + &cfg, + &tree, + &leaf_of(&c2), + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap(); assert_eq!(sealed.len(), 1, "expected one seal at the budget crossing"); let summary = src_store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); diff --git a/src/openhuman/memory/tree/retrieval/rpc.rs b/src/openhuman/memory/tree/retrieval/rpc.rs index 171ff0419..5f1b8948c 100644 --- a/src/openhuman/memory/tree/retrieval/rpc.rs +++ b/src/openhuman/memory/tree/retrieval/rpc.rs @@ -282,11 +282,26 @@ mod tests { //! initialises the schema idempotently on first access, so read-only //! calls return empty responses rather than erroring. use super::*; + use crate::openhuman::memory::tree::content_store; use crate::openhuman::memory::tree::store::upsert_chunks; use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceRef}; 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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::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(); @@ -316,6 +331,7 @@ mod tests { token_count: 20, seq_in_source: seq, created_at: ts, + partial_message: false, } } @@ -539,6 +555,7 @@ mod tests { 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 req = FetchLeavesRequest { chunk_ids: vec![c1.id.clone(), c2.id.clone()], }; @@ -552,6 +569,7 @@ mod tests { 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 req = FetchLeavesRequest { chunk_ids: vec![c1.id.clone(), "ghost:nonexistent".into()], }; diff --git a/src/openhuman/memory/tree/retrieval/source.rs b/src/openhuman/memory/tree/retrieval/source.rs index 46d5bc473..8580bd9e1 100644 --- a/src/openhuman/memory/tree/retrieval/source.rs +++ b/src/openhuman/memory/tree/retrieval/source.rs @@ -21,6 +21,7 @@ use anyhow::Result; use chrono::{Duration, Utc}; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::read as content_read; use crate::openhuman::memory::tree::retrieval::types::{ hit_from_summary, QueryResponse, RetrievalHit, }; @@ -124,7 +125,19 @@ fn collect_hits_and_nodes( // 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 node in level_nodes { + 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())); } @@ -293,7 +306,10 @@ fn filter_by_window(hits: Vec, window_days: u32) -> Vec) { let tree = get_or_create_source_tree(cfg, scope).unwrap(); let summariser = InertSummariser::new(); + 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"), @@ -328,17 +346,32 @@ mod tests { tags: vec!["eng".into()], source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))), }, - token_count: 6_000, + token_count: crate::openhuman::memory::tree::source_tree::types::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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); append_leaf( cfg, &tree, &LeafRef { chunk_id: c.id.clone(), - token_count: 6_000, + token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET + * 6 + / 10, timestamp: ts, content: c.content.clone(), entities: vec![], @@ -346,6 +379,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); diff --git a/src/openhuman/memory/tree/retrieval/topic.rs b/src/openhuman/memory/tree/retrieval/topic.rs index a91ef3756..6e7b031a9 100644 --- a/src/openhuman/memory/tree/retrieval/topic.rs +++ b/src/openhuman/memory/tree/retrieval/topic.rs @@ -17,6 +17,7 @@ use anyhow::Result; use chrono::{Duration, TimeZone, Utc}; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::read as content_read; use crate::openhuman::memory::tree::retrieval::types::{ hit_from_summary, QueryResponse, RetrievalHit, }; @@ -229,7 +230,7 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result id.clone(), None => return Ok(None), }; - let summary = match store::get_summary(config, &root_id)? { + let mut summary = match store::get_summary(config, &root_id)? { Some(s) => s, None => { log::warn!( @@ -238,6 +239,17 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result summary.content = body, + Err(e) => { + log::warn!( + "[retrieval::topic] read_summary_body failed for topic root — serving preview: {e:#}" + ); + } + } Ok(Some(hit_from_summary(&summary, &tree.scope))) } @@ -258,13 +270,23 @@ async fn entity_hit_to_retrieval_hit( tokio::task::spawn_blocking(move || -> Result> { if node_kind == "summary" { - let summary = match store::get_summary(&config_owned, &node_id)? { + let mut summary = match store::get_summary(&config_owned, &node_id)? { Some(s) => s, None => { log::warn!("[retrieval::topic] entity index points at missing summary row"); return Ok(None); } }; + // Hydrate the full body from disk — `summary.content` is a + // ≤500-char preview after the MD-on-disk migration. + match content_read::read_summary_body(&config_owned, &node_id) { + Ok(body) => summary.content = body, + Err(e) => { + log::warn!( + "[retrieval::topic] read_summary_body failed — serving preview: {e:#}" + ); + } + } // Prefer tree scope from the summary's parent tree if resolvable. let scope = if let Some(tid) = &tree_id_opt { store::get_tree(&config_owned, tid)? @@ -283,13 +305,21 @@ async fn entity_hit_to_retrieval_hit( // Leaf: fetch chunk and hydrate. use crate::openhuman::memory::tree::retrieval::types::hit_from_chunk; use crate::openhuman::memory::tree::store::get_chunk; - let chunk = match get_chunk(&config_owned, &node_id)? { + let mut chunk = match get_chunk(&config_owned, &node_id)? { Some(c) => c, None => { log::warn!("[retrieval::topic] entity index points at missing chunk row"); return Ok(None); } }; + // 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_owned, &node_id) { + Ok(body) => chunk.content = body, + Err(e) => { + log::warn!("[retrieval::topic] read_chunk_body failed — serving preview: {e:#}"); + } + } let scope = if let Some(tid) = &tree_id_opt { store::get_tree(&config_owned, tid)? .map(|t: Tree| t.scope) @@ -545,7 +575,7 @@ mod tests { with_connection(&cfg, |conn| { let tx = conn.unchecked_transaction()?; tree_store::insert_tree_conn(&tx, &tree)?; - tree_store::insert_summary_tx(&tx, &summary)?; + tree_store::insert_summary_tx(&tx, &summary, None)?; score_store::index_entities_tx( &tx, &[entity], diff --git a/src/openhuman/memory/tree/rpc.rs b/src/openhuman/memory/tree/rpc.rs index 2c9132d4f..2d370f5cf 100644 --- a/src/openhuman/memory/tree/rpc.rs +++ b/src/openhuman/memory/tree/rpc.rs @@ -178,3 +178,143 @@ pub async fn get_chunk_rpc( format!("memory_tree: get_chunk id={}", req.id), )) } + +/// Manual-trigger surface for the global tree's daily digest. Default +/// behavior (no `date_iso`) targets yesterday in UTC, matching the +/// scheduler's autonomous behavior. Pass an explicit `YYYY-MM-DD` to +/// re-run a specific date (idempotent — the handler skips if a daily +/// node already exists for that day). +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct TriggerDigestRequest { + /// UTC calendar date in `YYYY-MM-DD` form. When omitted, defaults to + /// `yesterday` (today minus one day, UTC). + #[serde(default)] + pub date_iso: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TriggerDigestResponse { + /// True when the job was newly enqueued; false when an active job for + /// the same date was suppressed by the dedupe partial unique index. + pub enqueued: bool, + /// ID of the freshly-inserted job row (None when dedupe-suppressed). + pub job_id: Option, + /// The actual date the digest will run for, echoed back as + /// `YYYY-MM-DD`. Useful when the caller didn't pass `date_iso` and + /// wants to know what default got chosen. + pub date_iso: String, +} + +pub async fn trigger_digest_rpc( + config: &Config, + req: TriggerDigestRequest, +) -> Result, String> { + use crate::openhuman::memory::tree::jobs; + use chrono::{Duration as ChronoDuration, NaiveDate, Utc}; + + let date = match req + .date_iso + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + Some(s) => NaiveDate::parse_from_str(s, "%Y-%m-%d") + .map_err(|e| format!("invalid date_iso (expected YYYY-MM-DD): {e}"))?, + None => Utc::now().date_naive() - ChronoDuration::days(1), + }; + let date_iso = date.format("%Y-%m-%d").to_string(); + + // Run the synchronous enqueue on a blocking thread — `trigger_digest` + // touches SQLite and we don't want to block the async runtime even + // for the few-microsecond INSERT. + let cfg_clone = config.clone(); + let date_for_blocking = date; + let job_id = + tokio::task::spawn_blocking(move || jobs::trigger_digest(&cfg_clone, date_for_blocking)) + .await + .map_err(|e| format!("trigger_digest join error: {e}"))? + .map_err(|e| format!("trigger_digest: {e}"))?; + + let enqueued = job_id.is_some(); + Ok(RpcOutcome::single_log( + TriggerDigestResponse { + enqueued, + job_id, + date_iso: date_iso.clone(), + }, + format!("memory_tree: trigger_digest date={date_iso} enqueued={enqueued}"), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::store::count_total; + use chrono::{Duration as ChronoDuration, 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; + (tmp, cfg) + } + + #[tokio::test] + async fn trigger_digest_with_explicit_date_enqueues() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("2026-04-27".into()), + }; + let outcome = trigger_digest_rpc(&cfg, req).await.unwrap(); + let resp = outcome.value; + assert!(resp.enqueued); + assert!(resp.job_id.is_some()); + assert_eq!(resp.date_iso, "2026-04-27"); + assert_eq!(count_total(&cfg).unwrap(), 1); + } + + #[tokio::test] + async fn trigger_digest_with_no_date_defaults_to_yesterday() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest::default(); + let outcome = trigger_digest_rpc(&cfg, req).await.unwrap(); + let resp = outcome.value; + assert!(resp.enqueued); + let expected = (Utc::now().date_naive() - ChronoDuration::days(1)) + .format("%Y-%m-%d") + .to_string(); + assert_eq!(resp.date_iso, expected); + } + + #[tokio::test] + async fn trigger_digest_rejects_malformed_date() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("not-a-date".into()), + }; + let err = trigger_digest_rpc(&cfg, req).await.unwrap_err(); + assert!( + err.contains("invalid date_iso"), + "expected schema-shaped error message, got: {err}" + ); + assert_eq!(count_total(&cfg).unwrap(), 0); + } + + #[tokio::test] + async fn trigger_digest_dedupes_active_jobs() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("2026-04-27".into()), + }; + let first = trigger_digest_rpc(&cfg, req.clone()).await.unwrap().value; + let second = trigger_digest_rpc(&cfg, req).await.unwrap().value; + assert!(first.enqueued); + assert!(!second.enqueued, "duplicate must be dedupe-suppressed"); + assert!(second.job_id.is_none()); + assert_eq!(count_total(&cfg).unwrap(), 1); + } +} diff --git a/src/openhuman/memory/tree/schemas.rs b/src/openhuman/memory/tree/schemas.rs index 6e5d59480..c558a85f6 100644 --- a/src/openhuman/memory/tree/schemas.rs +++ b/src/openhuman/memory/tree/schemas.rs @@ -23,6 +23,7 @@ pub fn all_controller_schemas() -> Vec { schemas("ingest"), schemas("list_chunks"), schemas("get_chunk"), + schemas("trigger_digest"), ] } @@ -40,6 +41,10 @@ pub fn all_registered_controllers() -> Vec { schema: schemas("get_chunk"), handler: handle_get_chunk, }, + RegisteredController { + schema: schemas("trigger_digest"), + handler: handle_trigger_digest, + }, ] } @@ -184,6 +189,44 @@ pub fn schemas(function: &str) -> ControllerSchema { required: false, }], }, + "trigger_digest" => ControllerSchema { + namespace: NAMESPACE, + function: "trigger_digest", + description: "Manually enqueue a daily-digest job for the global \ + tree. Idempotent — re-running for a day that already has a \ + digest is a no-op (the handler skips). When no date is \ + supplied, defaults to yesterday in UTC, matching the \ + scheduler's autonomous behavior.", + inputs: vec![FieldSchema { + name: "date_iso", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "UTC calendar date as `YYYY-MM-DD`. Optional; \ + defaults to yesterday when omitted.", + required: false, + }], + outputs: vec![ + FieldSchema { + name: "enqueued", + ty: TypeSchema::Bool, + comment: "True when a fresh job row was inserted; false \ + when an active job for the same date suppressed it.", + required: true, + }, + FieldSchema { + name: "job_id", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "ID of the newly enqueued job row, when enqueued.", + required: false, + }, + FieldSchema { + name: "date_iso", + ty: TypeSchema::String, + comment: "The date the digest will cover, echoed back \ + as `YYYY-MM-DD`.", + required: true, + }, + ], + }, _ => ControllerSchema { namespace: NAMESPACE, function: "unknown", @@ -228,6 +271,14 @@ fn handle_get_chunk(params: Map) -> ControllerFuture { }) } +fn handle_trigger_digest(params: Map) -> ControllerFuture { + Box::pin(async move { + let config = config_rpc::load_config_with_timeout().await?; + let req = parse_value::(Value::Object(params))?; + to_json(tree_rpc::trigger_digest_rpc(&config, req).await?) + }) +} + fn parse_value(v: Value) -> Result { serde_json::from_value(v).map_err(|e| format!("invalid params: {e}")) } diff --git a/src/openhuman/memory/tree/score/extract/llm.rs b/src/openhuman/memory/tree/score/extract/llm.rs index 6ff895618..20072aac6 100644 --- a/src/openhuman/memory/tree/score/extract/llm.rs +++ b/src/openhuman/memory/tree/score/extract/llm.rs @@ -36,7 +36,7 @@ use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; -use super::types::{EntityKind, ExtractedEntities, ExtractedEntity}; +use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; use super::EntityExtractor; // ── Configuration ──────────────────────────────────────────────────────── @@ -59,6 +59,15 @@ pub struct LlmExtractorConfig { /// 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, } impl Default for LlmExtractorConfig { @@ -73,8 +82,13 @@ impl Default for LlmExtractorConfig { EntityKind::Location, EntityKind::Event, EntityKind::Product, + EntityKind::Datetime, + EntityKind::Technology, + EntityKind::Artifact, + EntityKind::Quantity, ], strict_kinds: false, + emit_topics: false, } } } @@ -103,7 +117,7 @@ impl LlmEntityExtractor { messages: vec![ OllamaMessage { role: "system".to_string(), - content: SYSTEM_PROMPT.to_string(), + content: build_system_prompt(self.cfg.emit_topics), }, OllamaMessage { role: "user".to_string(), @@ -128,19 +142,54 @@ impl EntityExtractor for LlmEntityExtractor { // 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). A caller - // distinguishes "LLM had nothing to say" from "LLM ran and returned - // zero entities" by inspecting `llm_importance` — `None` means the - // call didn't complete successfully. - Ok(self.extract_or_empty(text).await) + // 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 { + Some(extracted) => return Ok(extracted), + None => { + // Transport 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; + } + } + } + } + + log::warn!( + "[memory_tree::extract::llm] transport failed after {} attempts — \ + returning empty extraction", + MAX_ATTEMPTS + ); + Ok(ExtractedEntities::default()) } } impl LlmEntityExtractor { - /// Internal: wraps the actual HTTP call and returns `ExtractedEntities` - /// for every failure mode via soft-fallback. Split out of `extract` so - /// the error branches can share logging without `?`-propagation. - async fn extract_or_empty(&self, text: &str) -> ExtractedEntities { + /// Internal: one attempt at calling Ollama. + /// + /// Returns: + /// - `Some(extracted)` — call completed (HTTP returned). Includes the + /// "HTTP non-success" and "malformed JSON" cases, which return + /// `Some(empty)` because retrying the same input won't help. + /// - `None` — transport-level failure (DNS, connect refused, timeout + /// before any HTTP response). Caller may retry. + async fn try_extract(&self, text: &str) -> Option { let url = format!("{}/api/chat", self.cfg.endpoint.trim_end_matches('/')); let body = self.build_request(text); log::debug!( @@ -152,11 +201,8 @@ impl LlmEntityExtractor { let resp = match self.http.post(&url).json(&body).send().await { Ok(r) => r, Err(e) => { - log::warn!( - "[memory_tree::extract::llm] transport failure to {url}: {e} — \ - returning empty extraction" - ); - return ExtractedEntities::default(); + log::warn!("[memory_tree::extract::llm] transport failure to {url}: {e}"); + return None; } }; @@ -168,7 +214,7 @@ impl LlmEntityExtractor { returning empty extraction", truncate_for_log(&body, 200) ); - return ExtractedEntities::default(); + return Some(ExtractedEntities::default()); } let envelope: OllamaChatResponse = match resp.json().await { @@ -178,7 +224,7 @@ impl LlmEntityExtractor { "[memory_tree::extract::llm] response body not Ollama-shaped JSON: {e} — \ returning empty extraction" ); - return ExtractedEntities::default(); + return Some(ExtractedEntities::default()); } }; log::debug!( @@ -194,38 +240,97 @@ impl LlmEntityExtractor { response: {e}; content was: {} — returning empty extraction", truncate_for_log(&envelope.message.content, 400) ); - return ExtractedEntities::default(); + return Some(ExtractedEntities::default()); } }; - parsed.into_extracted_entities(text, &self.cfg) + Some(parsed.into_extracted_entities(text, &self.cfg)) } } // ── Prompt ─────────────────────────────────────────────────────────────── -const SYSTEM_PROMPT: &str = "\ -You are a named-entity extractor and importance rater. Return JSON only — \ +/// 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) -> 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 { + "" + }; + + format!( + "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\", - \"text\": \"\" } + {{ \"kind\": \"person|organization|location|event|product|datetime|technology|artifact|quantity\", + \"text\": \"\" }} ], - \"importance\": 0.0, +{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 -"; +" + ) +} // ── Wire types (Ollama API) ────────────────────────────────────────────── @@ -265,6 +370,11 @@ struct OllamaResponseMessage { 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)] @@ -355,9 +465,26 @@ impl LlmExtractionOutput { 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: Vec::new(), + topics, llm_importance, llm_importance_reason: self.importance_reason, } @@ -373,6 +500,14 @@ fn parse_kind(s: &str) -> Option { "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, } diff --git a/src/openhuman/memory/tree/score/extract/llm_tests.rs b/src/openhuman/memory/tree/score/extract/llm_tests.rs index a1f61c5b2..0ee9bf364 100644 --- a/src/openhuman/memory/tree/score/extract/llm_tests.rs +++ b/src/openhuman/memory/tree/score/extract/llm_tests.rs @@ -1,5 +1,39 @@ use super::*; +#[test] +fn build_system_prompt_default_omits_topics() { + let p = build_system_prompt(false); + 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); + 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 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)); @@ -8,6 +42,42 @@ fn parse_kind_normalisation() { 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"; @@ -71,6 +141,7 @@ fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() { text: "Alice".into(), }, ], + topics: vec![], importance: None, importance_reason: None, }; @@ -100,6 +171,7 @@ fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() { text: "Alice".into(), }, ], + topics: vec![], importance: None, importance_reason: None, }; @@ -138,6 +210,7 @@ fn into_extracted_entities_drops_hallucinations() { text: "ImaginaryPerson".into(), }, ], + topics: vec![], importance: Some(0.7), importance_reason: Some("substantive".into()), }; @@ -154,6 +227,7 @@ fn into_extracted_entities_drops_hallucinations() { fn into_extracted_entities_clamps_importance() { let out = LlmExtractionOutput { entities: vec![], + topics: vec![], importance: Some(1.5), importance_reason: None, }; @@ -169,6 +243,7 @@ fn into_extracted_entities_strict_drops_unknown_kinds() { kind: "spaceship".into(), text: "Enterprise".into(), }], + topics: vec![], importance: None, importance_reason: None, }; @@ -187,6 +262,7 @@ fn into_extracted_entities_lenient_falls_back_to_misc() { kind: "spaceship".into(), text: "Enterprise".into(), }], + topics: vec![], importance: None, importance_reason: None, }; @@ -204,6 +280,7 @@ fn into_extracted_entities_disallowed_known_kind_falls_back_to_misc() { kind: "person".into(), text: "Alice".into(), }], + topics: vec![], importance: None, importance_reason: None, }; diff --git a/src/openhuman/memory/tree/score/extract/mod.rs b/src/openhuman/memory/tree/score/extract/mod.rs index f79a97b19..9fd914aaa 100644 --- a/src/openhuman/memory/tree/score/extract/mod.rs +++ b/src/openhuman/memory/tree/score/extract/mod.rs @@ -10,6 +10,83 @@ pub mod llm; pub mod regex; pub mod types; +use std::sync::Arc; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::util::redact::redact_endpoint; + pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor}; pub use llm::{LlmEntityExtractor, LlmExtractorConfig}; pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; + +/// 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 +/// `memory_tree.llm_extractor_endpoint` and `..._model` are both set +/// +/// 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 endpoint = config + .memory_tree + .llm_extractor_endpoint + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()); + let model = config + .memory_tree + .llm_extractor_model + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()); + + let (Some(endpoint), Some(model)) = (endpoint, model) else { + log::debug!( + "[memory_tree::extract] summary extractor: LLM not configured — using regex-only" + ); + return Arc::new(CompositeExtractor::regex_only()); + }; + + let timeout_ms = config + .memory_tree + .llm_extractor_timeout_ms + .unwrap_or(15_000); + + let cfg = LlmExtractorConfig { + endpoint: endpoint.to_string(), + model: model.to_string(), + timeout: std::time::Duration::from_millis(timeout_ms), + emit_topics: true, + ..LlmExtractorConfig::default() + }; + + match LlmEntityExtractor::new(cfg) { + Ok(llm) => { + // Drop to debug (diagnostic, not always-on) and redact the endpoint + // so embedded credentials (e.g. api keys in URL) don't leak. + log::debug!( + "[memory_tree::extract] summary extractor: regex + LLM endpoint={} model={} \ + timeout_ms={} emit_topics=true", + redact_endpoint(endpoint), + model, + timeout_ms + ); + Arc::new(CompositeExtractor::new(vec![ + Box::new(RegexEntityExtractor), + Box::new(llm), + ])) + } + Err(err) => { + log::warn!( + "[memory_tree::extract] summary extractor: LlmEntityExtractor construction \ + failed: {err:#} — falling back to regex-only" + ); + Arc::new(CompositeExtractor::regex_only()) + } + } +} diff --git a/src/openhuman/memory/tree/score/extract/types.rs b/src/openhuman/memory/tree/score/extract/types.rs index 8e91c8437..89b4fd672 100644 --- a/src/openhuman/memory/tree/score/extract/types.rs +++ b/src/openhuman/memory/tree/score/extract/types.rs @@ -26,12 +26,22 @@ pub enum EntityKind { Url, Handle, Hashtag, - // Semantic (reserved — not emitted in Phase 2) + // 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 @@ -54,6 +64,10 @@ impl EntityKind { 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", } @@ -70,6 +84,10 @@ impl EntityKind { "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}")), @@ -207,7 +225,12 @@ mod tests { 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); } diff --git a/src/openhuman/memory/tree/score/mod.rs b/src/openhuman/memory/tree/score/mod.rs index 15f17be76..7ae5ade51 100644 --- a/src/openhuman/memory/tree/score/mod.rs +++ b/src/openhuman/memory/tree/score/mod.rs @@ -336,6 +336,22 @@ pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result 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. diff --git a/src/openhuman/memory/tree/score/mod_tests.rs b/src/openhuman/memory/tree/score/mod_tests.rs index ac46d0083..2ce45d981 100644 --- a/src/openhuman/memory/tree/score/mod_tests.rs +++ b/src/openhuman/memory/tree/score/mod_tests.rs @@ -11,6 +11,7 @@ fn test_chunk(content: &str) -> Chunk { metadata: meta, seq_in_source: 0, created_at: Utc::now(), + partial_message: false, } } diff --git a/src/openhuman/memory/tree/score/store.rs b/src/openhuman/memory/tree/score/store.rs index ecb314b8d..e8670e57b 100644 --- a/src/openhuman/memory/tree/score/store.rs +++ b/src/openhuman/memory/tree/score/store.rs @@ -371,6 +371,21 @@ pub fn lookup_entity( }) } +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| { diff --git a/src/openhuman/memory/tree/source_tree/bucket_seal.rs b/src/openhuman/memory/tree/source_tree/bucket_seal.rs index 0f9f2b2c0..10377fb3b 100644 --- a/src/openhuman/memory/tree/source_tree/bucket_seal.rs +++ b/src/openhuman/memory/tree/source_tree/bucket_seal.rs @@ -26,12 +26,20 @@ //! summariser does no real I/O; when a networked summariser lands, wrap //! DB calls in `tokio::task::spawn_blocking` to keep the runtime healthy. +use std::collections::BTreeSet; +use std::sync::Arc; + use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use rusqlite::Transaction; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::{ + atomic::stage_summary, paths::slugify_source_id, SummaryComposeInput, SummaryTreeKind, +}; use crate::openhuman::memory::tree::score::embed::build_embedder_from_config; +use crate::openhuman::memory::tree::score::extract::EntityExtractor; +use crate::openhuman::memory::tree::score::resolver::canonicalise; use crate::openhuman::memory::tree::source_tree::registry::new_summary_id; use crate::openhuman::memory::tree::source_tree::store; use crate::openhuman::memory::tree::source_tree::summariser::{ @@ -47,6 +55,88 @@ use crate::openhuman::memory::tree::store::with_connection; /// realistic source. const MAX_CASCADE_DEPTH: u32 = 32; +/// 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, +} + +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 { @@ -61,17 +151,22 @@ pub struct LeafRef { /// 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, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result> { log::debug!( - "[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={}", + "[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={} strategy={:?}", tree.id, leaf.chunk_id, - leaf.token_count + leaf.token_count, + strategy ); // 1. Push leaf into L0 buffer (transactional). @@ -85,7 +180,24 @@ pub async fn append_leaf( )?; // 2. Cascade seals upward until a level stays under budget. - cascade_seals(config, tree, summariser).await + cascade_seals(config, tree, summariser, 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)) } /// Transactionally append a single item to `(tree_id, level)`'s buffer. @@ -127,20 +239,25 @@ async fn cascade_seals( config: &Config, tree: &Tree, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result> { - cascade_all_from(config, tree, 0, summariser, None).await + cascade_all_from(config, tree, 0, summariser, 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, start_level: u32, summariser: &dyn Summariser, force_now: Option>, + strategy: &LabelStrategy, ) -> Result> { let mut sealed_ids: Vec = Vec::new(); let mut level: u32 = start_level; @@ -169,7 +286,9 @@ pub async fn cascade_all_from( break; } - let summary_id = seal_one_level(config, tree, &buf, summariser).await?; + // 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, summariser, strategy, false).await?; sealed_ids.push(summary_id); level += 1; } @@ -185,7 +304,7 @@ pub async fn cascade_all_from( /// 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. -fn should_seal(buf: &Buffer) -> bool { +pub(crate) fn should_seal(buf: &Buffer) -> bool { if buf.is_empty() { return false; } @@ -198,11 +317,31 @@ fn should_seal(buf: &Buffer) -> bool { /// Seal `buf` at `level` into one summary at `level + 1`. Returns the new /// summary id. -async fn seal_one_level( +/// +/// `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, summariser: &dyn Summariser, + strategy: &LabelStrategy, + enqueue_follow_ups: bool, ) -> Result { let level = buf.level; let target_level = level + 1; @@ -237,7 +376,7 @@ async fn seal_one_level( // Run summariser — async, OUTSIDE any DB transaction. let ctx = SummaryContext { tree_id: &tree.id, - tree_kind: TreeKind::Source, + tree_kind: tree.kind, target_level, token_budget: TOKEN_BUDGET, }; @@ -246,6 +385,12 @@ async fn seal_one_level( .await .context("summariser failed during seal")?; + // 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 @@ -291,14 +436,18 @@ async fn seal_one_level( let node = SummaryNode { id: summary_id.clone(), tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, + // `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: output.entities, - topics: output.topics, + entities: node_entities, + topics: node_topics, time_range_start, time_range_end, score, @@ -307,13 +456,88 @@ async fn seal_one_level( embedding: Some(embedding), }; + // 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. + // + // `bucket_seal.rs` handles both Source and Topic tree seals (Topic trees + // use the same cascade machinery via `handle_seal` in the job handler). + // Map TreeKind to SummaryTreeKind accordingly. + let summary_tree_kind = match tree.kind { + TreeKind::Topic => SummaryTreeKind::Topic, + _ => SummaryTreeKind::Source, + }; + let scope_slug = { + // Path slug semantics per source kind: + // + // - Gmail source trees: scope is `"gmail:"` where + // participants is `addr1|addr2|...`. Strip the `gmail:` prefix so the + // path is `summaries/source//...` and mirrors the + // chunk layout under `email//`. + // + // - Topic trees: scope is the canonical entity_id (e.g. + // `"email:alice@example.com"`). Slugify the FULL string so topic-tree + // summaries and source-tree summaries don't share a path prefix. + // + // - All other source kinds (slack:, discord:, document:, …): slugify the + // FULL scope string. Stripping the prefix for non-Gmail sources was a + // bug — `"slack:#eng"` and `"discord:#eng"` would both produce slug + // `"eng"` and collide in `summaries/source/eng/`. + let s = &tree.scope; + match tree.kind { + TreeKind::Topic => slugify_source_id(s), + _ => { + if s.starts_with("gmail:") { + // Strip "gmail:" prefix; slugify the participants portion. + slugify_source_id(&s["gmail:".len()..]) + } else { + // All other source kinds: slugify the full scope string. + slugify_source_id(s) + } + } + } + }; + 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_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(); + let staged = + stage_summary(&content_root, &compose_input, &scope_slug, None).with_context(|| { + format!( + "stage_summary failed for {}; seal aborted, buffer stays unsealed for retry", + node.id + ) + })?; + log::debug!( + "[source_tree::bucket_seal] staged summary {} → {}", + node.id, + staged.content_path + ); + // Single write transaction: insert summary, clear this buffer, append - // summary id to parent buffer, bump tree max_level/root if needed. + // 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(); + let tree_kind = tree.kind; with_connection(config, move |conn| { let tx = conn.unchecked_transaction()?; @@ -326,7 +550,7 @@ async fn seal_one_level( .map(|n| n.max(0) as u32) .context("Failed to read current max_level for tree")?; - store::insert_summary_tx(&tx, &node)?; + store::insert_summary_tx(&tx, &node, Some(&staged))?; // 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 @@ -375,6 +599,43 @@ async fn seal_one_level( }; 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::tree::jobs::store::enqueue_tx as enqueue_job_tx; + use crate::openhuman::memory::tree::jobs::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)?)?; + } + // Source-tree summary routing: feed the new summary's + // entities back into the topic-tree spawn pipeline. Topic + // and global trees are sinks — no fan-out from their seals. + if matches!(tree_kind, TreeKind::Source) { + use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx; + use crate::openhuman::memory::tree::jobs::types::{ + NewJob, NodeRef, TopicRoutePayload, + }; + let route = TopicRoutePayload { + node: NodeRef::Summary { + summary_id: summary_id_for_closure.clone(), + }, + }; + enqueue_job_tx(&tx, &NewJob::topic_route(&route)?)?; + } + } + // Update tree root / max_level if we just climbed. if target_level_for_closure > current_max { store::update_tree_after_seal_tx( @@ -447,7 +708,8 @@ fn hydrate_inputs(config: &Config, level: u32, item_ids: &[String]) -> Result Result> { - use crate::openhuman::memory::tree::score::store::get_score; + use crate::openhuman::memory::tree::content_store::read as content_read; + use crate::openhuman::memory::tree::score::store::{get_score, list_entity_ids_for_node}; use crate::openhuman::memory::tree::store::get_chunk; let mut out: Vec = Vec::with_capacity(chunk_ids.len()); @@ -461,17 +723,30 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result (row.total, Vec::new(), chunk.metadata.tags.clone()), - None => (0.0, Vec::new(), chunk.metadata.tags.clone()), - }; + 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!("[source_tree::bucket_seal] hydrate_leaf_inputs: read body for chunk {id}") + })?; out.push(SummaryInput { id: chunk.id.clone(), - content: chunk.content.clone(), + content: body, token_count: chunk.token_count, entities, - topics, + topics: chunk.metadata.tags.clone(), time_range_start: chunk.metadata.time_range.0, time_range_end: chunk.metadata.time_range.1, score: score_value, @@ -481,6 +756,8 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result Result> { + use crate::openhuman::memory::tree::content_store::read as content_read; + let mut out: Vec = Vec::with_capacity(summary_ids.len()); for id in summary_ids { let node = match store::get_summary(config, id)? { @@ -492,9 +769,15 @@ fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result (TempDir, Config) { let tmp = TempDir::new().unwrap(); let mut cfg = Config::default(); @@ -36,7 +57,9 @@ async fn append_below_budget_does_not_seal() { // 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 = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert!(sealed.is_empty()); let buf = store::get_buffer(&cfg, &tree.id, 0).unwrap(); @@ -72,15 +95,21 @@ async fn crossing_budget_triggers_seal() { token_count: tokens, seq_in_source: seq, created_at: ts, + partial_message: false, }; - let c1 = mk_chunk(0, 6_000); - let c2 = mk_chunk(1, 6_000); + // Budget-relative sizes so the test stays correct as TOKEN_BUDGET shifts: + // each leaf is 60% of budget, so the second append crosses the threshold. + let per_leaf = 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: 6_000, + token_count: per_leaf, timestamp: ts, content: c1.content.clone(), entities: vec![], @@ -89,7 +118,7 @@ async fn crossing_budget_triggers_seal() { }; let leaf2 = LeafRef { chunk_id: c2.id.clone(), - token_count: 6_000, + token_count: per_leaf, timestamp: ts, content: c2.content.clone(), entities: vec![], @@ -97,10 +126,14 @@ async fn crossing_budget_triggers_seal() { score: 0.5, }; - let first = append_leaf(&cfg, &tree, &leaf1, &summariser).await.unwrap(); + let first = append_leaf(&cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert!(first.is_empty(), "first append below budget — no seal"); - let second = append_leaf(&cfg, &tree, &leaf2, &summariser).await.unwrap(); + let second = append_leaf(&cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(second.len(), 1, "second append crosses budget — one seal"); let summary_id = &second[0]; @@ -169,6 +202,7 @@ async fn fanout_at_l1_triggers_l2_seal() { token_count: 10_000, seq_in_source: seq, created_at: ts, + partial_message: false, } }; @@ -177,6 +211,8 @@ async fn fanout_at_l1_triggers_l2_seal() { 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, @@ -186,7 +222,9 @@ async fn fanout_at_l1_triggers_l2_seal() { topics: vec![], score: 0.5, }; - let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); all_sealed.extend(sealed); } @@ -253,8 +291,11 @@ async fn upper_level_does_not_seal_below_fanout() { token_count: 10_000, 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, @@ -264,7 +305,9 @@ async fn upper_level_does_not_seal_below_fanout() { topics: vec![], score: 0.5, }; - let _ = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let _ = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); } let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); @@ -281,3 +324,334 @@ async fn upper_level_does_not_seal_below_fanout() { 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::tree::score::extract::EntityKind; + use crate::openhuman::memory::tree::score::resolver::CanonicalEntity; + use crate::openhuman::memory::tree::score::store::index_entity; + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; + 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}"))), + }, + // Bust TOKEN_BUDGET in one leaf so the seal fires immediately. + token_count: 10_000, + 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}; + use std::sync::Arc; + + let (_tmp, cfg) = test_config(); + let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + let summariser = InertSummariser::new(); + + // Content the regex extractor can find: an email and a hashtag. The + // inert summariser concatenates leaf content into the L1 summary, so + // these tokens survive into the summary text and the extractor finds + // them when run on the summary content. + 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 = append_leaf(&cfg, &tree, &leaf, &summariser, &strategy) + .await + .unwrap(); + 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 summariser = InertSummariser::new(); + + // 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 = 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 = append_leaf( + &cfg, + &tree, + &leaf1, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + assert!(sealed_1.is_empty()); + // Second leaf: crosses budget → one seal covering both leaves. + let sealed_2 = append_leaf( + &cfg, + &tree, + &leaf2, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + 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 summariser = InertSummariser::new(); + + // 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 = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + 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::tree::source_tree::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 summariser = InertSummariser::new(); + let leaf = seed_leaf(&cfg, 0, "topic content", vec![], vec![]); + + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + 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::tree::content_store::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" + ); +} diff --git a/src/openhuman/memory/tree/source_tree/flush.rs b/src/openhuman/memory/tree/source_tree/flush.rs index 920c47061..24e8a1314 100644 --- a/src/openhuman/memory/tree/source_tree/flush.rs +++ b/src/openhuman/memory/tree/source_tree/flush.rs @@ -13,7 +13,7 @@ use anyhow::Result; use chrono::{DateTime, Duration, Utc}; use crate::openhuman::config::Config; -use crate::openhuman::memory::tree::source_tree::bucket_seal::cascade_all_from; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{cascade_all_from, LabelStrategy}; use crate::openhuman::memory::tree::source_tree::store; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS; @@ -25,6 +25,7 @@ pub async fn flush_stale_buffers( config: &Config, max_age: Duration, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result { let now = Utc::now(); let cutoff = now - max_age; @@ -48,7 +49,8 @@ pub async fn flush_stale_buffers( continue; } }; - let sealed = cascade_all_from(config, &tree, buf.level, summariser, Some(now)).await?; + let sealed = + cascade_all_from(config, &tree, buf.level, summariser, Some(now), strategy).await?; seals += sealed.len(); } Ok(seals) @@ -58,11 +60,13 @@ pub async fn flush_stale_buffers( pub async fn flush_stale_buffers_default( config: &Config, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result { flush_stale_buffers( config, Duration::seconds(DEFAULT_FLUSH_AGE_SECS), summariser, + strategy, ) .await } @@ -74,15 +78,17 @@ pub async fn force_flush_tree( tree_id: &str, summariser: &dyn Summariser, now: Option>, + strategy: &LabelStrategy, ) -> Result> { let tree = store::get_tree(config, tree_id)? .ok_or_else(|| anyhow::anyhow!("no tree with id {tree_id}"))?; - cascade_all_from(config, &tree, 0, summariser, now).await + cascade_all_from(config, &tree, 0, summariser, now, strategy).await } #[cfg(test)] mod tests { use super::*; + use crate::openhuman::memory::tree::content_store; use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; @@ -90,6 +96,20 @@ mod tests { use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; 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::tree::store::with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + crate::openhuman::memory::tree::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(); @@ -124,8 +144,10 @@ mod tests { 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(), @@ -136,12 +158,15 @@ mod tests { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); - - let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser) + append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) .await .unwrap(); + assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); + + let seals = + flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(seals, 1); assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 1); @@ -172,6 +197,7 @@ mod tests { token_count: 50, seq_in_source: 0, created_at: now, + partial_message: false, }; upsert_chunks(&cfg, &[c.clone()]).unwrap(); let leaf = LeafRef { @@ -183,11 +209,14 @@ mod tests { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); - - let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser) + append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) .await .unwrap(); + + let seals = + flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(seals, 0); assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); } diff --git a/src/openhuman/memory/tree/source_tree/mod.rs b/src/openhuman/memory/tree/source_tree/mod.rs index f97d07995..1f9694468 100644 --- a/src/openhuman/memory/tree/source_tree/mod.rs +++ b/src/openhuman/memory/tree/source_tree/mod.rs @@ -22,7 +22,7 @@ pub mod store; pub mod summariser; pub mod types; -pub use bucket_seal::{append_leaf, LeafRef}; +pub use bucket_seal::{append_leaf, append_leaf_deferred, LabelStrategy, LeafRef}; pub use registry::get_or_create_source_tree; pub use store::{get_summary_embedding, set_summary_embedding}; pub use summariser::{build_summariser, inert::InertSummariser, llm::LlmSummariser, Summariser}; diff --git a/src/openhuman/memory/tree/source_tree/store.rs b/src/openhuman/memory/tree/source_tree/store.rs index ca2ac4398..36fd9f126 100644 --- a/src/openhuman/memory/tree/source_tree/store.rs +++ b/src/openhuman/memory/tree/source_tree/store.rs @@ -21,6 +21,7 @@ use chrono::{DateTime, TimeZone, Utc}; use rusqlite::{params, Connection, OptionalExtension, Transaction}; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::StagedSummary; use crate::openhuman::memory::tree::score::embed::{decode_optional_blob, pack_checked}; use crate::openhuman::memory::tree::source_tree::types::{ Buffer, SummaryNode, Tree, TreeKind, TreeStatus, @@ -180,7 +181,16 @@ fn row_to_tree(row: &rusqlite::Row<'_>) -> rusqlite::Result { /// 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. -pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Result<()> { +/// +/// 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>, +) -> Result<()> { let embedding_blob: Option> = match node.embedding.as_deref() { Some(v) => Some( pack_checked(v) @@ -188,14 +198,30 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res ), None => 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 - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)", + score, sealed_at_ms, deleted, embedding, + content_path, content_sha256 + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)", params![ node.id, node.tree_id, @@ -203,7 +229,7 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res node.level, node.parent_id, serde_json::to_string(&node.child_ids)?, - node.content, + content_preview, node.token_count, serde_json::to_string(&node.entities)?, serde_json::to_string(&node.topics)?, @@ -213,6 +239,8 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res node.sealed_at.timestamp_millis(), node.deleted as i64, embedding_blob, + content_path, + content_sha256, ], ) .with_context(|| format!("Failed to insert summary id={}", node.id))?; diff --git a/src/openhuman/memory/tree/source_tree/store_tests.rs b/src/openhuman/memory/tree/source_tree/store_tests.rs index 248716d7f..6d98ddbd6 100644 --- a/src/openhuman/memory/tree/source_tree/store_tests.rs +++ b/src/openhuman/memory/tree/source_tree/store_tests.rs @@ -71,7 +71,7 @@ fn summary_insert_and_fetch() { let node = sample_summary("sum-1", "tree-1", 1); with_connection(&cfg, |conn| { let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &node)?; + insert_summary_tx(&tx, &node, None)?; tx.commit()?; Ok(()) }) @@ -90,8 +90,8 @@ fn summary_insert_is_idempotent_on_id() { let node = sample_summary("sum-1", "tree-1", 1); with_connection(&cfg, |conn| { let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &node)?; - insert_summary_tx(&tx, &node)?; + insert_summary_tx(&tx, &node, None)?; + insert_summary_tx(&tx, &node, None)?; tx.commit()?; Ok(()) }) diff --git a/src/openhuman/memory/tree/source_tree/summariser/llm.rs b/src/openhuman/memory/tree/source_tree/summariser/llm.rs index 09c8e6ba4..8f9f31135 100644 --- a/src/openhuman/memory/tree/source_tree/summariser/llm.rs +++ b/src/openhuman/memory/tree/source_tree/summariser/llm.rs @@ -6,10 +6,12 @@ //! When the source / topic / global tree's bucket-seal cascade decides to //! fold N contributions (raw leaves at L0→L1, or lower-level summaries at //! L_n→L_{n+1}), this summariser is asked to produce the parent node's -//! `content` + derived `entities` + `topics`. The seal machinery itself -//! (bucket budgeting, level promotion, `mem_tree_summaries` persistence) -//! is unchanged — only the text inside the summary row differs from -//! [`super::inert::InertSummariser`]. +//! `content`. The seal machinery itself (bucket budgeting, level +//! promotion, `mem_tree_summaries` persistence) is unchanged — only the +//! text inside the summary row differs from [`super::inert::InertSummariser`]. +//! Entities and topics on `SummaryOutput` are always emitted empty by +//! this summariser; canonical entity ids are populated separately by the +//! entity extractor. //! //! ## Soft-fallback contract //! @@ -18,14 +20,12 @@ //! parent row. We therefore promise **never** to return `Err`: every //! failure (transport, HTTP status, JSON shape) falls back to the same //! deterministic concat-and-truncate behaviour as `InertSummariser` and -//! logs a warn. Callers distinguish "LLM ran fine" from "we fell back" -//! only by observing whether the returned `entities`/`topics` are -//! populated — the inert branch emits empty vecs. +//! logs a warn. //! //! ## Prompt shape //! //! The system prompt commits the model to returning JSON with the shape -//! `{ summary, entities, topics }`. We use Ollama's `format: "json"` + +//! `{ summary }`. We use Ollama's `format: "json"` + //! `temperature: 0.0` to maximise determinism — same knobs the entity //! extractor already uses with success. @@ -40,15 +40,24 @@ use super::inert::InertSummariser; use super::{Summariser, SummaryContext, SummaryInput, SummaryOutput}; use crate::openhuman::memory::tree::types::approx_token_count; -/// Hard cap on summary OUTPUT tokens, regardless of the seal's -/// `ctx.token_budget`. Driven by the embedder's context window — -/// `nomic-embed-text-v1.5` accepts up to 8192 tokens and Phase 4 -/// (`source_tree::bucket_seal`) embeds the summary right after we -/// produce it. If our summary overshoots, the embedder returns 500 -/// and the whole seal transaction rolls back → no summary persists. -/// 6000 leaves a safety margin for tokenizer differences and the -/// JSON wrapper. -const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 6_000; +/// Hard cap on summariser output length (in approximate tokens). +/// +/// Two constraints set this: +/// +/// 1. The downstream embedder (`nomic-embed-text-v1.5`) accepts up to +/// 8192 tokens, and Phase 4 (`source_tree::bucket_seal`) embeds the +/// summary right after we produce it. An overshoot returns HTTP 500 +/// and rolls back the whole seal transaction. +/// 2. Empirically, small instruction-tuned models running locally +/// degrade quickly past ~3500 tokens — they drift, hallucinate, or +/// produce repetitive boilerplate as they extend toward longer +/// targets. Keeping the cap below that breakeven keeps output +/// quality stable on local Ollama deployments. +/// +/// 3500 sits comfortably under the embedder ceiling AND below the local +/// LLM quality cliff. The post-generation [`clamp_to_budget`] enforces +/// this regardless of what the model produces. +const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 3_500; /// Context window we ask Ollama for. Must match the value below in /// [`OllamaOptions::num_ctx`] so the per-input clamp computed in @@ -261,31 +270,18 @@ impl Summariser for LlmSummariser { let (content, token_count) = clamp_to_budget(&parsed.summary, effective_budget); log::debug!( - "[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={} \ - surface_entities_dropped={} topics={}", + "[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={}", ctx.tree_id, ctx.target_level, inputs.len(), - token_count, - parsed.entities.len(), - parsed.topics.len() + token_count ); - // Drop LLM-emitted entities. The model returns surface forms - // ("Alice", "she"), but `SummaryNode.entities` is indexed via - // `index_summary_entity_ids_tx` as canonical ids. Surface forms - // would silently corrupt that index — searches by canonical id - // would not find these summaries. Canonicalisation is the - // entity extractor's job, not the summariser's. Topics stay - // because they're free-form labels, not indexed as canonical - // ids. The `entities` field stays in the prompt to nudge the - // model toward entity-aware summarisation; we just don't - // persist its output. Ok(SummaryOutput { content, token_count, entities: Vec::new(), - topics: dedupe_sorted(parsed.topics), + topics: Vec::new(), }) } } @@ -314,20 +310,21 @@ fn build_user_prompt(inputs: &[SummaryInput], per_input_cap_tokens: u32) -> Stri out } -/// System prompt. Token budget is templated in so the model aims under it. -fn system_prompt(budget: u32) -> String { - format!( - "You are a precise summariser. Summarise the user-provided contributions into a \ - single cohesive passage that preserves concrete facts, decisions, named entities, \ - and temporal ordering. Do not invent facts. Stay well under {budget} tokens.\n\ - \n\ - Return JSON only — no prose, no markdown, no commentary. Schema:\n\ - {{\n\ - \x20 \"summary\": \"\",\n\ - \x20 \"entities\": [\"\"],\n\ - \x20 \"topics\": [\"\"]\n\ - }}" - ) +/// System prompt. Length isn't templated in — empirically, telling small +/// instruction-tuned models "stay under N tokens" makes them produce +/// curt, generic output even when the input has plenty of substance. +/// Output is clamped post-generation by [`clamp_to_budget`] in the +/// caller, so we don't need the model to self-police length. +fn system_prompt(_budget: u32) -> String { + "You are a precise summariser. Summarise the user-provided contributions into a \ + single cohesive passage that preserves concrete facts, decisions, \ + and temporal ordering. Do not invent facts.\n\ + \n\ + Return JSON only — no prose, no markdown, no commentary. Schema:\n\ + {\n\ + \x20 \"summary\": \"\"\n\ + }" + .to_string() } /// Truncate to the caller's token budget using the same ~4 chars/token @@ -343,16 +340,6 @@ fn clamp_to_budget(text: &str, budget: u32) -> (String, u32) { (truncated, tokens) } -fn dedupe_sorted(mut items: Vec) -> Vec { - for item in items.iter_mut() { - *item = item.trim().to_string(); - } - items.retain(|s| !s.is_empty()); - items.sort(); - items.dedup(); - items -} - fn truncate_for_log(s: &str, max_chars: usize) -> String { if s.chars().count() <= max_chars { return s.to_string(); @@ -405,10 +392,6 @@ struct OllamaResponseMessage { struct LlmSummaryOutput { #[serde(default)] summary: String, - #[serde(default)] - entities: Vec, - #[serde(default)] - topics: Vec, } #[cfg(test)] @@ -493,12 +476,16 @@ mod tests { } #[test] - fn system_prompt_templates_budget() { + fn system_prompt_describes_schema() { + // Budget is no longer templated into the prompt — small models + // produced overly curt output when told to "stay under N tokens". + // The clamp in `clamp_to_budget` handles enforcement instead. let p = system_prompt(4096); - assert!(p.contains("4096")); + assert!(!p.contains("4096")); + assert!(!p.contains("Stay well under")); assert!(p.contains("\"summary\"")); - assert!(p.contains("\"entities\"")); - assert!(p.contains("\"topics\"")); + assert!(!p.contains("\"entities\"")); + assert!(!p.contains("\"topics\"")); } #[test] @@ -516,19 +503,6 @@ mod tests { assert!(t <= 6); } - #[test] - fn dedupe_sorted_trims_and_dedupes() { - let out = dedupe_sorted(vec![ - "Bob".into(), - " Alice ".into(), - "Bob".into(), - "".into(), - " ".into(), - "Alice".into(), - ]); - assert_eq!(out, vec!["Alice", "Bob"]); - } - #[test] fn truncate_for_log_short_input_unchanged() { assert_eq!(truncate_for_log("hi", 10), "hi"); @@ -592,16 +566,24 @@ mod tests { assert!(!req.stream); assert_eq!(req.options.temperature, 0.0); assert_eq!(req.messages[0].role, "system"); - assert!(req.messages[0].content.contains("2048")); + assert!(req.messages[0].content.contains("\"summary\"")); assert_eq!(req.messages[1].role, "user"); assert_eq!(req.messages[1].content, "body"); } #[test] - fn llm_output_deserialises_with_missing_fields() { + fn llm_output_deserialises_with_only_summary() { let v: LlmSummaryOutput = serde_json::from_str(r#"{"summary":"hi"}"#).unwrap(); assert_eq!(v.summary, "hi"); - assert!(v.entities.is_empty()); - assert!(v.topics.is_empty()); + } + + #[test] + fn llm_output_ignores_extraneous_fields() { + // Prompt no longer asks for entities/topics, but if the model + // emits them anyway we should still parse `summary` cleanly. + let v: LlmSummaryOutput = + serde_json::from_str(r#"{"summary":"hi","entities":["Alice"],"topics":["x"]}"#) + .unwrap(); + assert_eq!(v.summary, "hi"); } } diff --git a/src/openhuman/memory/tree/source_tree/types.rs b/src/openhuman/memory/tree/source_tree/types.rs index 4fded7db8..2a4d23082 100644 --- a/src/openhuman/memory/tree/source_tree/types.rs +++ b/src/openhuman/memory/tree/source_tree/types.rs @@ -172,13 +172,17 @@ impl Buffer { } } -/// Token ceiling for one summariser invocation — aligned with the Phase 1 -/// chunker ceiling so a single leaf never busts a seal on its own. +/// Token ceiling for one summariser invocation. +/// +/// Sized for the local 1B summariser (`gemma3:1b-it-qat`), which produces +/// noticeably better summaries with ≤4-5k input than at higher caps. The +/// chunker's `DEFAULT_CHUNK_MAX_TOKENS` (3_000) sits below this so each +/// L0 buffer accumulates roughly 1-3 chunks before sealing. /// /// Gates only the L0 → L1 seal: leaves are fan-in by raw token volume so /// the summariser input stays bounded. Summaries above L0 use /// [`SUMMARY_FANOUT`] instead — see `bucket_seal::should_seal`. -pub const TOKEN_BUDGET: u32 = 10_000; +pub const TOKEN_BUDGET: u32 = 4_500; /// Sibling count that triggers a seal at level ≥ 1 (summaries → next level). /// diff --git a/src/openhuman/memory/tree/store.rs b/src/openhuman/memory/tree/store.rs index 9da6da299..9f50ea891 100644 --- a/src/openhuman/memory/tree/store.rs +++ b/src/openhuman/memory/tree/store.rs @@ -13,6 +13,7 @@ use rusqlite::{params, Connection, OptionalExtension, Transaction}; use std::time::Duration; use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::content_store::StagedChunk; use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind, SourceRef}; const DB_DIR: &str = "memory_tree"; @@ -21,6 +22,12 @@ const DEFAULT_LIST_LIMIT: usize = 100; const MAX_LIST_LIMIT: usize = 10_000; const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); +pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction"; +pub const CHUNK_STATUS_ADMITTED: &str = "admitted"; +pub const CHUNK_STATUS_BUFFERED: &str = "buffered"; +pub const CHUNK_STATUS_SEALED: &str = "sealed"; +pub const CHUNK_STATUS_DROPPED: &str = "dropped"; + const SCHEMA: &str = " PRAGMA foreign_keys = ON; @@ -174,6 +181,35 @@ CREATE TABLE IF NOT EXISTS mem_tree_entity_hotness ( 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 +); + +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'); "; /// Upsert a batch of chunks atomically. @@ -249,6 +285,66 @@ pub(crate) fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result Ok(chunks.len()) } +/// 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], +) -> Result { + if staged.is_empty() { + return Ok(0); + } + let mut stmt = tx.prepare( + "INSERT INTO mem_tree_chunks ( + id, source_kind, source_id, 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) + ON CONFLICT(id) DO UPDATE SET + source_kind = excluded.source_kind, + source_id = excluded.source_id, + 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; + // Store a ≤500-char preview in the `content` column; full body is on disk. + 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.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()) +} + fn upsert_chunks_with_statement( stmt: &mut rusqlite::Statement<'_>, chunks: &[Chunk], @@ -361,6 +457,70 @@ pub fn count_chunks(config: &Config) -> Result { }) } +pub fn set_chunk_lifecycle_status(config: &Config, chunk_id: &str, status: &str) -> Result<()> { + with_connection(config, |conn| { + set_chunk_lifecycle_status_conn(conn, chunk_id, status) + }) +} + +pub(crate) fn set_chunk_lifecycle_status_tx( + tx: &Transaction<'_>, + chunk_id: &str, + status: &str, +) -> Result<()> { + set_chunk_lifecycle_status_conn(tx, chunk_id, status) +} + +pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result> { + with_connection(config, |conn| { + get_chunk_lifecycle_status_conn(conn, chunk_id) + }) +} + +pub(crate) fn get_chunk_lifecycle_status_tx( + tx: &Transaction<'_>, + chunk_id: &str, +) -> Result> { + get_chunk_lifecycle_status_conn(tx, chunk_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) +} + +pub fn count_chunks_by_lifecycle_status(config: &Config, status: &str) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_chunks WHERE lifecycle_status = ?1", + params![status], + |r| r.get(0), + )?; + Ok(n.max(0) as u64) + }) +} + +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_tree::store] lifecycle update affected 0 rows chunk_id={} status={}", + chunk_id, + status + ); + } + Ok(()) +} + fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result { let id: String = row.get(0)?; let source_kind_s: String = row.get(1)?; @@ -401,6 +561,10 @@ fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result { token_count: token_count.max(0) as u32, seq_in_source: seq.max(0) as u32, created_at, + // partial_message is not stored in SQLite — it's a transient chunker + // signal. Chunks read back from DB always get false (the column doesn't + // exist; callers that need this flag hold the Chunk in memory). + partial_message: false, }) } @@ -451,9 +615,124 @@ pub(crate) fn with_connection( // legacy summaries from Phases 1-3 read back as None; retrieval // tolerates NULL by dropping the row to the bottom of a 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. New writes start at + // 'pending_extraction'; the extract handler advances them to 'admitted' + // (then 'buffered' / 'sealed') or 'dropped'. + 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")?; + // Phase MD-content (#TBD): pointer + integrity hash. Body lives at + // / as a .md file. Both nullable so chunks + // ingested before this migration read back with NULL (body still in + // `content`). New writes populate both columns. The `content` column + // stores a 500-char plain-text preview instead of the full body. + 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): same pointer pattern for summary nodes. + // `content_path` is the relative path to the .md file under + // `/summaries/...`. `content_sha256` is the SHA-256 hex + // of the body bytes only (front-matter excluded). Both nullable so + // legacy rows (from before this migration) read back with NULL — callers + // fall back to the `content` column for those rows. + add_column_if_missing(&conn, "mem_tree_summaries", "content_path", "TEXT")?; + add_column_if_missing(&conn, "mem_tree_summaries", "content_sha256", "TEXT")?; f(&conn) } +/// Return both `content_path` and `content_sha256` stored in SQLite for `chunk_id`. +/// +/// Returns `Ok(None)` if the chunk does not exist or has no content_path recorded yet. +pub fn get_chunk_content_pointers( + config: &Config, + chunk_id: &str, +) -> Result> { + with_connection(config, |conn| { + let row = conn + .query_row( + "SELECT content_path, content_sha256 FROM mem_tree_chunks WHERE id = ?1", + params![chunk_id], + |r| { + let path: Option = r.get(0)?; + let sha: Option = r.get(1)?; + Ok((path, sha)) + }, + ) + .optional()?; + Ok(row.and_then(|(p, s)| p.zip(s))) + }) +} + +/// Return the `content_path` stored in SQLite for `chunk_id`, if any. +pub fn get_chunk_content_path(config: &Config, chunk_id: &str) -> Result> { + with_connection(config, |conn| { + let row = conn + .query_row( + "SELECT content_path FROM mem_tree_chunks WHERE id = ?1", + params![chunk_id], + |r| r.get::<_, Option>(0), + ) + .optional()? + .flatten(); + Ok(row) + }) +} + +/// Return both `content_path` and `content_sha256` stored in SQLite for `summary_id`. +/// +/// Returns `Ok(None)` if the summary does not exist or has no content_path recorded yet +/// (legacy rows pre-MD-content migration). +pub fn get_summary_content_pointers( + config: &Config, + summary_id: &str, +) -> Result> { + with_connection(config, |conn| { + let row = conn + .query_row( + "SELECT content_path, content_sha256 FROM mem_tree_summaries WHERE id = ?1", + params![summary_id], + |r| { + let path: Option = r.get(0)?; + let sha: Option = r.get(1)?; + Ok((path, sha)) + }, + ) + .optional()?; + Ok(row.and_then(|(p, s)| p.zip(s))) + }) +} + +/// List all summary rows that have a non-NULL `content_path`. Used by the +/// bin integrity checker. +pub fn list_summaries_with_content_path(config: &Config) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT id, content_path, content_sha256 + FROM mem_tree_summaries + WHERE content_path IS NOT NULL AND content_sha256 IS NOT NULL + AND deleted = 0", + )?; + let rows = stmt + .query_map([], |r| { + let id: String = r.get(0)?; + let path: String = r.get(1)?; + let sha: String = r.get(2)?; + Ok((id, path, sha)) + })? + .collect::>>() + .context("Failed to list summaries with content_path")?; + Ok(rows) + }) +} + fn normalized_limit(requested: Option) -> i64 { let clamped = requested .unwrap_or(DEFAULT_LIST_LIMIT) diff --git a/src/openhuman/memory/tree/store_tests.rs b/src/openhuman/memory/tree/store_tests.rs index 698c9e243..9c12a1493 100644 --- a/src/openhuman/memory/tree/store_tests.rs +++ b/src/openhuman/memory/tree/store_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::openhuman::memory::tree::types::chunk_id; use chrono::TimeZone; +use rusqlite::params; use tempfile::TempDir; fn test_config() -> (TempDir, Config) { @@ -27,6 +28,7 @@ fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk { token_count: 12, seq_in_source: seq, created_at: ts, + partial_message: false, } } @@ -164,3 +166,37 @@ fn empty_batch_is_noop() { assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0); assert_eq!(count_chunks(&cfg).unwrap(), 0); } + +#[test] +fn schema_has_content_path_and_content_sha256_columns() { + // Phase MD-content: verify that with_connection applies the additive + // migrations for the new pointer + hash columns on a fresh DB. + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + 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 == "content_path" { + has_content_path = true; + } + if name == "content_sha256" { + has_content_sha256 = true; + } + } + 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(); +} diff --git a/src/openhuman/memory/tree/topic_tree/backfill.rs b/src/openhuman/memory/tree/topic_tree/backfill.rs index e98a53e53..5e241347e 100644 --- a/src/openhuman/memory/tree/topic_tree/backfill.rs +++ b/src/openhuman/memory/tree/topic_tree/backfill.rs @@ -1,30 +1,52 @@ //! Topic-tree backfill — hydrate a freshly-materialised topic tree with -//! every historical leaf mentioning the entity (#709 Phase 3c). +//! recent leaves mentioning the entity (#709 Phase 3c). //! //! When the curator decides an entity has crossed the hotness threshold //! for the first time, we create a fresh topic tree AND walk the -//! `mem_tree_entity_index` inverted index to append every prior leaf into +//! `mem_tree_entity_index` inverted index to append matching leaves into //! its L0 buffer. Reusing `bucket_seal::append_leaf` means the cascade -//! fires automatically — a well-established entity may seal several -//! levels as soon as the tree is spawned. +//! fires automatically. +//! +//! ## Why bounded by hotness window +//! +//! Hotness uses a 30-day recency decay (see `topic_tree::hotness`). Leaves +//! older than 30 days contribute zero to current hotness, so by definition +//! they cannot be the reason a tree is spawning *now*. Including them +//! bloats the spawn latency, wastes summariser LLM calls, and amplifies +//! ancient signal that has already decayed away. We cap the backfill +//! window at [`BACKFILL_WINDOW_DAYS`] to align with the hotness math. +//! +//! Older content is still queryable through source-tree retrieval and the +//! entity index — it just doesn't get its own slot in the topic tree. //! //! Backfill is intentionally best-effort: missing chunks are skipped with //! a warn log rather than failing the whole spawn, because Phase 3c is //! additive — a partial topic tree is still useful. use anyhow::{Context, Result}; +use chrono::Utc; use crate::openhuman::config::Config; use crate::openhuman::memory::tree::score::store::lookup_entity; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::Tree; use crate::openhuman::memory::tree::store::get_chunk; +use crate::openhuman::memory::tree::util::redact::redact; /// Max leaves to pull from the entity index during backfill. A hard cap /// keeps initial spawn latency bounded even for very active entities. const BACKFILL_LIMIT: usize = 500; +/// Backfill window in days — matches `topic_tree::hotness::recency_decay`'s +/// hard cliff. Leaves older than this contribute zero to current hotness +/// so they cannot have driven the spawn decision. +pub const BACKFILL_WINDOW_DAYS: i64 = 30; + +const DAY_MS: i64 = 24 * 60 * 60 * 1_000; + /// Walk the entity index for `entity_id` and append every discovered leaf /// to `tree`. Returns the number of leaves appended (NOT the number of /// summaries sealed). Idempotent: `append_leaf` itself is a no-op when a @@ -35,19 +57,65 @@ pub async fn backfill_topic_tree( entity_id: &str, summariser: &dyn Summariser, ) -> Result { - log::info!( - "[topic_tree::backfill] start entity_id={} tree_id={}", + backfill_topic_tree_at( + config, + tree, entity_id, - tree.id + summariser, + Utc::now().timestamp_millis(), + ) + .await +} + +/// Deterministic variant — backfill against a caller-supplied `now_ms` +/// for the recency window. Used by tests so the 30-day cutoff doesn't +/// depend on the wall clock. +pub async fn backfill_topic_tree_at( + config: &Config, + tree: &Tree, + entity_id: &str, + summariser: &dyn Summariser, + now_ms: i64, +) -> Result { + let cutoff_ms = now_ms.saturating_sub(BACKFILL_WINDOW_DAYS.saturating_mul(DAY_MS)); + log::info!( + "[topic_tree::backfill] start entity_id_hash={} tree_id={} window_days={} cutoff_ms={}", + redact(entity_id), + tree.id, + BACKFILL_WINDOW_DAYS, + cutoff_ms ); let hits = lookup_entity(config, entity_id, Some(BACKFILL_LIMIT)) - .with_context(|| format!("failed to lookup entity {entity_id}"))?; + .with_context(|| format!("failed to lookup entity {}", redact(entity_id)))?; if hits.is_empty() { log::debug!( - "[topic_tree::backfill] no entity-index hits for entity_id={} — empty backfill", - entity_id + "[topic_tree::backfill] no entity-index hits for entity_id_hash={} — empty backfill", + redact(entity_id) + ); + return Ok(0); + } + + // Drop hits older than the hotness recency window — see module docs. + let total_hits = hits.len(); + let mut hits: Vec<_> = hits + .into_iter() + .filter(|h| h.timestamp_ms >= cutoff_ms) + .collect(); + let dropped = total_hits - hits.len(); + if dropped > 0 { + log::debug!( + "[topic_tree::backfill] dropped {dropped} hits older than {BACKFILL_WINDOW_DAYS}d \ + for entity_id_hash={}", + redact(entity_id) + ); + } + if hits.is_empty() { + log::debug!( + "[topic_tree::backfill] all entity-index hits fell outside the {BACKFILL_WINDOW_DAYS}d \ + window for entity_id_hash={} — empty backfill", + redact(entity_id) ); return Ok(0); } @@ -55,7 +123,6 @@ pub async fn backfill_topic_tree( // Sort by timestamp ASC so the buffer's `oldest_at` and the sealed // summary's `time_range_start` reflect the true historical order, not // the DESC ordering `lookup_entity` returns. - let mut hits = hits; hits.sort_by_key(|h| h.timestamp_ms); let mut appended = 0usize; @@ -77,9 +144,9 @@ pub async fn backfill_topic_tree( Some(c) => c, None => { log::warn!( - "[topic_tree::backfill] missing chunk {} for entity {} — skipping", + "[topic_tree::backfill] missing chunk {} for entity_id_hash={} — skipping", hit.node_id, - entity_id + redact(entity_id) ); continue; } @@ -95,20 +162,25 @@ pub async fn backfill_topic_tree( score: hit.score, }; - append_leaf(config, tree, &leaf, summariser) + // Topic-tree backfill: empty labels for sealed summaries — the + // tree's scope already pins the canonical id, so cross-pollinating + // descendants' entities would noise the index. See LabelStrategy. + append_leaf(config, tree, &leaf, summariser, &LabelStrategy::Empty) .await .with_context(|| { format!( - "backfill append_leaf failed tree_id={} chunk_id={}", - tree.id, chunk.id + "backfill append_leaf failed entity_id_hash={} tree_id={} chunk_id={}", + redact(entity_id), + tree.id, + chunk.id ) })?; appended += 1; } log::info!( - "[topic_tree::backfill] done entity_id={} tree_id={} appended={}", - entity_id, + "[topic_tree::backfill] done entity_id_hash={} tree_id={} appended={}", + redact(entity_id), tree.id, appended ); @@ -158,6 +230,7 @@ mod tests { token_count: tokens, seq_in_source: seq, created_at: ts, + partial_message: false, } } @@ -172,6 +245,11 @@ mod tests { } } + /// Deterministic "now" used by the windowed-backfill tests: 1 hour + /// after the latest seeded leaf so all three sit inside the 30-day + /// cutoff. Lets us keep the legacy 2023-era timestamps unchanged. + const TEST_NOW_MS: i64 = 1_700_000_020_000 + 3_600_000; + #[tokio::test] async fn backfill_appends_all_entity_leaves() { let (_tmp, cfg) = test_config(); @@ -212,9 +290,15 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 3); // L0 buffer should hold all three leaves (combined tokens well @@ -226,6 +310,38 @@ mod tests { assert_eq!(buf.oldest_at.unwrap().timestamp_millis(), 1_700_000_000_000); } + #[tokio::test] + async fn backfill_drops_leaves_older_than_window() { + let (_tmp, cfg) = test_config(); + // c_old is 60d before TEST_NOW_MS — outside the 30d cutoff. + // c_new is 5d before TEST_NOW_MS — inside the window. + let old_ts = TEST_NOW_MS - 60 * DAY_MS; + let new_ts = TEST_NOW_MS - 5 * DAY_MS; + let c_old = mk_chunk("slack:#eng", 0, old_ts, 100); + let c_new = mk_chunk("slack:#eng", 1, new_ts, 100); + upsert_chunks(&cfg, &[c_old.clone(), c_new.clone()]).unwrap(); + + let e = sample_entity("email:alice@example.com", "alice@example.com"); + index_entity(&cfg, &e, &c_old.id, "leaf", old_ts, Some("source:slack")).unwrap(); + index_entity(&cfg, &e, &c_new.id, "leaf", new_ts, Some("source:slack")).unwrap(); + + let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); + let summariser = InertSummariser::new(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); + assert_eq!(n, 1, "only the in-window leaf should be appended"); + let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); + assert_eq!(buf.item_ids.len(), 1); + assert_eq!(buf.item_ids[0], c_new.id); + } + #[tokio::test] async fn backfill_skips_missing_chunks_without_failing() { let (_tmp, cfg) = test_config(); @@ -247,9 +363,15 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 1, "only the existing chunk should be appended"); let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); assert_eq!(buf.item_ids.len(), 1); @@ -273,12 +395,24 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); - backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); + backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); // append_leaf is idempotent so the buffer still has exactly one row. let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); assert_eq!(buf.item_ids.len(), 1); @@ -300,9 +434,15 @@ mod tests { .unwrap(); let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 0); } } diff --git a/src/openhuman/memory/tree/topic_tree/curator.rs b/src/openhuman/memory/tree/topic_tree/curator.rs index ae4006272..85cba1845 100644 --- a/src/openhuman/memory/tree/topic_tree/curator.rs +++ b/src/openhuman/memory/tree/topic_tree/curator.rs @@ -186,7 +186,11 @@ mod tests { } fn seed_leaf_for_entity(cfg: &Config, entity_id: &str, source_tree: &str, seq: u32) { - let ts_ms = 1_700_000_000_000 + (seq as i64) * 1_000; + // Use a "now-anchored" timestamp so backfill's 30-day window + // (see topic_tree::backfill::BACKFILL_WINDOW_DAYS) always + // includes these seeded leaves. Spread by seq to keep ordering + // deterministic. + let ts_ms = Utc::now().timestamp_millis() - (seq as i64) * 1_000; let ts = Utc.timestamp_millis_opt(ts_ms).unwrap(); let c = Chunk { id: chunk_id(SourceKind::Chat, source_tree, seq, "test-content"), @@ -203,6 +207,7 @@ mod tests { token_count: 50, seq_in_source: seq, created_at: ts, + partial_message: false, }; upsert_chunks(cfg, &[c.clone()]).unwrap(); let e = CanonicalEntity { diff --git a/src/openhuman/memory/tree/topic_tree/routing.rs b/src/openhuman/memory/tree/topic_tree/routing.rs index 686756125..1bd77092a 100644 --- a/src/openhuman/memory/tree/topic_tree/routing.rs +++ b/src/openhuman/memory/tree/topic_tree/routing.rs @@ -21,7 +21,9 @@ use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::store as src_store; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::{TreeKind, TreeStatus}; @@ -90,7 +92,16 @@ async fn route_one_entity( entities: vec![entity_id.to_string()], ..leaf.clone() }; - append_leaf(config, &tree, &topic_leaf, summariser).await?; + // Topic-tree seals leave entities/topics empty: the tree's + // scope already pins the canonical id this tree represents. + append_leaf( + config, + &tree, + &topic_leaf, + summariser, + &LabelStrategy::Empty, + ) + .await?; } else { log::debug!( "[topic_tree::routing] skip archived topic tree id={} entity={}", @@ -162,6 +173,7 @@ mod tests { token_count: tokens, seq_in_source: seq, created_at: ts, + partial_message: false, }; let id = c.id.clone(); upsert_chunks(cfg, &[c]).unwrap(); @@ -295,9 +307,15 @@ mod tests { crate::openhuman::memory::tree::topic_tree::types::TOPIC_RECHECK_EVERY - 1; crate::openhuman::memory::tree::topic_tree::store::upsert(&cfg, &counters).unwrap(); - // Seed a leaf in slack and gmail referencing Alice. - let c1 = persist_chunk(&cfg, "slack:#eng", 0, 1_700_000_000_000, 100); - let c2 = persist_chunk(&cfg, "gmail:alice", 0, 1_700_000_010_000, 100); + // Seed leaves in slack and gmail referencing Alice. Anchor the + // timestamps to "now" so the 30-day backfill window + // (topic_tree::backfill::BACKFILL_WINDOW_DAYS) covers them. + let now_ms = Utc::now().timestamp_millis(); + let ts_c1 = now_ms - 20_000; + let ts_c2 = now_ms - 10_000; + let ts_c3 = now_ms; + let c1 = persist_chunk(&cfg, "slack:#eng", 0, ts_c1, 100); + let c2 = persist_chunk(&cfg, "gmail:alice", 0, ts_c2, 100); let e = CanonicalEntity { canonical_id: entity_id.into(), kind: EntityKind::Email, @@ -306,24 +324,16 @@ mod tests { span_end: entity_id.len() as u32, score: 1.0, }; - index_entity(&cfg, &e, &c1, "leaf", 1_700_000_000_000, Some("slack:#eng")).unwrap(); - index_entity( - &cfg, - &e, - &c2, - "leaf", - 1_700_000_010_000, - Some("gmail:alice"), - ) - .unwrap(); + index_entity(&cfg, &e, &c1, "leaf", ts_c1, Some("slack:#eng")).unwrap(); + index_entity(&cfg, &e, &c2, "leaf", ts_c2, Some("gmail:alice")).unwrap(); // A third leaf arrives — should both fan out to (future) topic tree // and push the curator over the recheck cadence, materialising it. - let c3 = persist_chunk(&cfg, "slack:#eng", 1, 1_700_000_020_000, 100); + let c3 = persist_chunk(&cfg, "slack:#eng", 1, ts_c3, 100); let leaf = LeafRef { chunk_id: c3.clone(), token_count: 100, - timestamp: Utc.timestamp_millis_opt(1_700_000_020_000).unwrap(), + timestamp: Utc.timestamp_millis_opt(ts_c3).unwrap(), content: "new mention".into(), entities: vec![entity_id.into()], topics: vec![], diff --git a/src/openhuman/memory/tree/types.rs b/src/openhuman/memory/tree/types.rs index 1d266da82..a7e05a045 100644 --- a/src/openhuman/memory/tree/types.rs +++ b/src/openhuman/memory/tree/types.rs @@ -232,6 +232,12 @@ pub struct Chunk { /// When this chunk was persisted to the local store. #[serde(with = "chrono::serde::ts_milliseconds")] pub created_at: DateTime, + /// True when this chunk is a sub-split of a single logical unit (e.g. a + /// chat message or email body that exceeded `max_tokens`). The full logical + /// unit was split into multiple pieces; each piece carries this flag so + /// downstream scorers can lower its weight relative to whole-unit chunks. + #[serde(default)] + pub partial_message: bool, } /// Deterministic chunk id. diff --git a/src/openhuman/memory/tree/util/mod.rs b/src/openhuman/memory/tree/util/mod.rs new file mode 100644 index 000000000..0c32b1c7f --- /dev/null +++ b/src/openhuman/memory/tree/util/mod.rs @@ -0,0 +1,3 @@ +//! Shared utility helpers for the memory-tree subsystem. + +pub mod redact; diff --git a/src/openhuman/memory/tree/util/redact.rs b/src/openhuman/memory/tree/util/redact.rs new file mode 100644 index 000000000..a9b2ab4ba --- /dev/null +++ b/src/openhuman/memory/tree/util/redact.rs @@ -0,0 +1,136 @@ +//! PII redaction helpers for log output. +//! +//! Per project rule (CLAUDE.md): "Never log secrets or full PII." +//! After the participant-bucketing change introduced in the MD-content PR, +//! source_ids and content_paths can embed full email addresses, so any log +//! line that prints them needs to redact. + +use sha2::{Digest, Sha256}; + +/// Redact a string by hashing it to 8 hex chars. Stable across runs for the +/// same input — safe to grep for in logs when debugging with the raw value +/// available externally. +/// +/// Use for source_ids, entity_ids, content_paths and similar PII-bearing +/// strings in log output. +pub fn redact(s: &str) -> String { + let mut h = Sha256::new(); + h.update(s.as_bytes()); + let d = h.finalize(); + format!("{:08x}", u32::from_be_bytes([d[0], d[1], d[2], d[3]])) +} + +/// Redact a URL/endpoint by stripping path, query, fragment and credentials, +/// keeping only the host (and port if present). +/// +/// Examples: +/// - `"http://localhost:11434/api/chat"` → `"localhost:11434"` +/// - `"https://user:pass@example.com/foo?q=1"` → `"example.com"` +/// - `"ollama://host:1234"` → `"host:1234"` +/// +/// Does not pull in a URL-parsing crate; uses cheap string splitting which is +/// sufficient for the endpoint-config strings this codebase passes around. +pub fn redact_endpoint(url: &str) -> String { + // Strip scheme (everything before "://"). + let after_scheme = url.split_once("://").map(|(_, r)| r).unwrap_or(url); + // Take only the authority (everything up to the first '/', '?', or '#') so + // any '@' in the path / query (e.g. `?email=foo@bar`) doesn't get treated + // as a userinfo separator. + let authority = after_scheme + .split(['/', '?', '#']) + .next() + .unwrap_or(after_scheme); + // Within the authority, the LAST '@' separates userinfo from host:port. + // (RFC 3986: userinfo may itself contain '@' — split-on-first would + // truncate the host. Use rsplit so `user:p@ss@example.com` extracts + // `example.com` correctly.) + let host_port = authority + .rsplit_once('@') + .map(|(_, r)| r) + .unwrap_or(authority); + host_port.to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + // ── redact ─────────────────────────────────────────────────────────────── + + #[test] + fn redact_returns_eight_hex_chars() { + let r = redact("alice@example.com"); + assert_eq!(r.len(), 8, "must be 8 hex chars; got {r:?}"); + assert!(r.chars().all(|c| c.is_ascii_hexdigit()), "must be hex"); + } + + #[test] + fn redact_is_stable_across_calls() { + assert_eq!(redact("alice@example.com"), redact("alice@example.com")); + } + + #[test] + fn redact_is_different_for_different_inputs() { + assert_ne!(redact("alice@example.com"), redact("bob@example.com")); + } + + #[test] + fn redact_empty_string_does_not_panic() { + let r = redact(""); + assert_eq!(r.len(), 8); + } + + // ── redact_endpoint ───────────────────────────────────────────────────── + + #[test] + fn redact_endpoint_strips_path_and_query() { + assert_eq!( + redact_endpoint("http://localhost:11434/api/chat"), + "localhost:11434" + ); + } + + #[test] + fn redact_endpoint_strips_credentials() { + assert_eq!( + redact_endpoint("https://user:pass@example.com/foo"), + "example.com" + ); + } + + #[test] + fn redact_endpoint_no_scheme_passthrough() { + // No "://" present — treat the whole string as host/path; still strip path. + assert_eq!(redact_endpoint("localhost:11434/api"), "localhost:11434"); + } + + #[test] + fn redact_endpoint_just_host() { + assert_eq!(redact_endpoint("https://example.com"), "example.com"); + } + + #[test] + fn redact_endpoint_strips_fragment() { + assert_eq!(redact_endpoint("http://host:9090/path#frag"), "host:9090"); + } + + #[test] + fn redact_endpoint_strips_query() { + assert_eq!(redact_endpoint("http://host/path?q=1"), "host"); + } + + #[test] + fn redact_endpoint_empty_does_not_panic() { + let r = redact_endpoint(""); + // Empty input: no scheme, no host — returns empty string. + assert_eq!(r, ""); + } + + #[test] + fn redact_endpoint_ollama_style() { + assert_eq!( + redact_endpoint("http://127.0.0.1:11434/v1/chat/completions"), + "127.0.0.1:11434" + ); + } +} diff --git a/tests/json_rpc_e2e.rs b/tests/json_rpc_e2e.rs index eb15b8794..54cd8c27a 100644 --- a/tests/json_rpc_e2e.rs +++ b/tests/json_rpc_e2e.rs @@ -825,6 +825,7 @@ async fn json_rpc_memory_tree_end_to_end() { "openhuman.memory_tree_ingest".to_string(), "openhuman.memory_tree_list_chunks".to_string(), "openhuman.memory_tree_get_chunk".to_string(), + "openhuman.memory_tree_trigger_digest".to_string(), ]; assert_eq!(controllers.len(), expected_methods.len()); for method in &expected_methods {