solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4from __future__ import annotations
5
6import json
7import multiprocessing
8import os
9import time
10import traceback
11from pathlib import Path
12from queue import Empty
13from typing import Any
14
15DAY = "20240101"
16STREAM = "test"
17SEGMENT = "143022_300"
18SOURCE_ID = "merge_source"
19TARGET_ID = "merge_target"
20
21
22def _segment_dir(journal: Path) -> Path:
23 return journal / "chronicle" / DAY / STREAM / SEGMENT
24
25
26def _merge_rewrite_worker(
27 journal_path: str,
28 barrier: Any,
29 append_done: Any,
30 errors: Any,
31) -> None:
32 os.environ["SOLSTONE_JOURNAL"] = journal_path
33 try:
34 from solstone.think.entities import merge as merge_mod
35
36 plan = merge_mod._plan_segment_rewrites(SOURCE_ID, TARGET_ID)
37 barrier.wait(timeout=5)
38 if not append_done.wait(timeout=5):
39 raise TimeoutError("append worker did not finish")
40 time.sleep(0.05)
41 merge_mod._apply_segment_plan(plan["operations"], SOURCE_ID, TARGET_ID)
42 except BaseException:
43 errors.put(traceback.format_exc())
44 raise
45
46
47def _append_worker(
48 journal_path: str,
49 seg_dir_path: str,
50 barrier: Any,
51 append_done: Any,
52 errors: Any,
53) -> None:
54 os.environ["SOLSTONE_JOURNAL"] = journal_path
55 try:
56 from solstone.apps.speakers.attribution import append_speaker_correction
57
58 barrier.wait(timeout=5)
59 append_speaker_correction(
60 Path(seg_dir_path),
61 {
62 "sentence_id": 2,
63 "original_speaker": "other_source",
64 "corrected_speaker": "other_target",
65 "original_method": "user_corrected",
66 "timestamp": 1,
67 },
68 )
69 append_done.set()
70 except BaseException:
71 errors.put(traceback.format_exc())
72 append_done.set()
73 raise
74
75
76def _drain_errors(errors: Any) -> list[str]:
77 found = []
78 while True:
79 try:
80 found.append(errors.get_nowait())
81 except Empty:
82 return found
83
84
85def test_entity_merge_corrections_rewrite_preserves_concurrent_append(
86 tmp_path: Path,
87) -> None:
88 ctx = multiprocessing.get_context("spawn")
89 journal = tmp_path / "journal"
90 seg_dir = _segment_dir(journal)
91 talents_dir = seg_dir / "talents"
92 talents_dir.mkdir(parents=True)
93 corrections_path = talents_dir / "speaker_corrections.json"
94 corrections_path.write_text(
95 json.dumps(
96 {
97 "corrections": [
98 {
99 "sentence_id": 1,
100 "original_speaker": SOURCE_ID,
101 "corrected_speaker": SOURCE_ID,
102 "original_method": "acoustic",
103 "timestamp": 0,
104 }
105 ]
106 },
107 indent=2,
108 ),
109 encoding="utf-8",
110 )
111
112 barrier = ctx.Barrier(2)
113 append_done = ctx.Event()
114 errors = ctx.Queue()
115 processes = [
116 ctx.Process(
117 target=_merge_rewrite_worker,
118 args=(str(journal), barrier, append_done, errors),
119 ),
120 ctx.Process(
121 target=_append_worker,
122 args=(str(journal), str(seg_dir), barrier, append_done, errors),
123 ),
124 ]
125
126 for process in processes:
127 process.start()
128 for process in processes:
129 process.join(timeout=10)
130 for process in processes:
131 if process.is_alive():
132 process.terminate()
133 process.join(timeout=2)
134
135 error_text = "\n".join(_drain_errors(errors))
136 assert all(not process.is_alive() for process in processes), error_text
137 assert all(process.exitcode == 0 for process in processes), error_text
138
139 data = json.loads(corrections_path.read_text(encoding="utf-8"))
140 by_sentence = {entry["sentence_id"]: entry for entry in data["corrections"]}
141 assert by_sentence[1]["original_speaker"] == TARGET_ID
142 assert by_sentence[1]["corrected_speaker"] == TARGET_ID
143 assert by_sentence[2]["original_speaker"] == "other_source"
144 assert by_sentence[2]["corrected_speaker"] == "other_target"