forked from NousResearch/hermes-agent
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhermes_state_search.py
More file actions
2487 lines (2301 loc) · 111 KB
/
Copy pathhermes_state_search.py
File metadata and controls
2487 lines (2301 loc) · 111 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""Full-text / trigram / CJK message search and FTS maintenance for SessionDB.
Mixin contract: this is a plain mixin class consumed by
``hermes_state.SessionDB``. It defines no ``__init__`` and no state of its
own; methods access the host's attributes (``self._conn``, ``self.db_path``,
``self._execute_write`` and other SessionDB methods) established by
``SessionDB.__init__``. It must never import hermes_state (cycle) — shared
module-level constants live in hermes_state_common.
"""
import logging
import json
import os
import re
import sqlite3
import time
from typing import Any, Callable, Collection, Dict, List, Optional, Tuple
from agent.skill_commands import describe_skill_invocation
from hermes_state_common import (
FTS_CJK_STALE_KEY,
FTS_SQL,
FTS_STALE_KEY,
FTS_STORAGE_VERSION,
FTS_TRIGRAM_SQL,
MAX_FTS5_QUERY_CHARS,
SCHEMA_VERSION,
_FTS_CJK_TRIGGERS,
escape_like as _escape_like,
)
# Moved methods logged under the "hermes_state" logger before the split;
# keep that logger identity so log filtering/capture behavior is unchanged.
logger = logging.getLogger("hermes_state")
# Characters FTS5's query grammar rejects outside a quoted phrase. Anything
# missing from this set reaches MATCH raw and raises, which the execute site
# swallows into zero results — the failure this strip step exists to prevent.
# Assembled through re.escape so the backslash cannot be eaten as a regex
# escape inside the class (it was, while the set was written as a literal).
#
# ``%`` is deliberately excluded: a CJK query falls back to a LIKE search that
# needs it preserved as a literal (that path escapes wildcards itself), so
# stripping it here widened those queries onto unrelated rows.
_FTS5_SPECIAL_CHARS = '+{}():"^@/#&|~[]<>,;!?$=\\\''
_FTS5_SPECIAL_RE = re.compile(f"[{re.escape(_FTS5_SPECIAL_CHARS)}]")
class SessionSearchMixin:
"""See module docstring — mixin for SessionDB (Search cluster)."""
_SEARCH_MESSAGE_RESULT_FIELDS = (
"id",
"session_id",
"role",
"snippet",
"timestamp",
"tool_name",
"source",
"model",
"session_started",
"context",
)
@classmethod
def _search_message_fields(
cls, fields: Optional[Collection[str]]
) -> Optional[Tuple[str, ...]]:
"""Validate and canonically order an optional result projection."""
if fields is None:
return None
if isinstance(fields, str):
raise TypeError("search fields must be a collection of field names, not a string")
requested = set(fields)
unknown = requested.difference(cls._SEARCH_MESSAGE_RESULT_FIELDS)
if unknown:
raise ValueError(f"unknown search result field(s): {', '.join(sorted(unknown))}")
return tuple(
field for field in cls._SEARCH_MESSAGE_RESULT_FIELDS if field in requested
)
def _try_incremental_merge_fts(self) -> None:
"""Run one bounded FTS5 merge pass without failing the completed write."""
if not self._fts_enabled:
return
try:
self._merge_fts_incrementally(
max_pages=self._FTS_MERGE_MAX_PAGES_PER_INDEX
)
except sqlite3.Error as exc:
# Routine maintenance is best effort, but unexpected SQLite errors
# must remain visible instead of being silently mistaken for an
# optional missing index.
logger.warning("FTS incremental merge failed: %s", exc)
def fts_rebuild_status(self) -> Optional[Dict[str, Any]]:
"""Return deferred-rebuild progress, or None when no rebuild pending.
Shape: {"pending": True, "total": <rows at drop time>,
"indexed": <rows backfilled>, "percent": <0-100 int>}.
Consumed by search_messages() notes and by status surfaces
(dashboard/desktop can poll this to render a progress indicator).
Reads state_meta directly via _read_ctx instead of calling
get_meta() (which takes self._lock) so search_messages doesn't
block on the writer lock when checking rebuild status.
"""
with self._read_ctx() as conn:
row = conn.execute(
"SELECT key, value FROM state_meta WHERE key IN (?, ?)",
("fts_rebuild_high_water", "fts_rebuild_progress"),
).fetchall()
meta = {r["key"]: r["value"] for r in row}
high_water = meta.get("fts_rebuild_high_water")
if high_water is None:
return None
progress = int(meta.get("fts_rebuild_progress") or 0)
total = int(high_water)
if total <= 0:
return None
pct = min(100, int(100 * progress / total))
return {"pending": True, "total": total, "indexed": progress, "percent": pct}
def _fts_rebuild_finish(self) -> None:
"""Finalize the deferred rebuild: boundary sweep + clear markers.
The sweep is cheap insurance against any write that slipped through
the migration-boundary instant (between high_water capture and
trigger activation): re-index any row near the boundary that the
index is missing. docsize has one row per indexed doc, so the
anti-join is exact and runs on a narrow id range.
The trigram half of the sweep is gated on ``self._trigram_available``
for the same reason ``fts_rebuild_step()`` gates its backfill INSERT:
when the SQLite build has no trigram tokenizer (or the table was
never created), an unconditional INSERT raises ``no such table``
and aborts the whole rebuild — taking ``optimize_fts_storage()``
down with it.
"""
include_trigram = self._trigram_available
def _do(conn):
hw_row = conn.execute(
"SELECT value FROM state_meta WHERE key = 'fts_rebuild_high_water'"
).fetchone()
if hw_row is not None:
hw = int(hw_row[0])
# Sweep a generous window around the boundary.
lo, hi = hw - 1000, hw + 1000
conn.execute(
"INSERT INTO messages_fts(rowid, content, tool_name, tool_calls) "
"SELECT m.id, m.content, m.tool_name, m.tool_calls "
"FROM messages m "
"WHERE m.id > ? AND m.id <= ? "
"AND NOT EXISTS (SELECT 1 FROM messages_fts_docsize d WHERE d.id = m.id)",
(lo, hi),
)
if include_trigram:
conn.execute(
"INSERT INTO messages_fts_trigram(rowid, content, tool_name, tool_calls) "
"SELECT m.id, m.content, m.tool_name, m.tool_calls "
"FROM messages m "
"WHERE m.id > ? AND m.id <= ? AND m.role <> 'tool' "
"AND NOT EXISTS (SELECT 1 FROM messages_fts_trigram_docsize d WHERE d.id = m.id)",
(lo, hi),
)
conn.execute(
"DELETE FROM state_meta WHERE key IN "
"('fts_rebuild_high_water', 'fts_rebuild_progress')"
)
self._execute_write(_do)
logger.info("Deferred FTS rebuild complete — all messages indexed.")
def _fts_teardown_trash_step(self) -> bool:
"""Tear down one chunk of a demoted v22 FTS shadow table.
The trash tables are PLAIN tables (their vtable parent was demoted
away during the migration), so chunked DELETE + final DROP involve
no FTS5 machinery at all. Returns True while teardown work remains.
Single-column-key trash tables (the common shape — FTS shadow
tables carry a rowid/integer PK) are drained with a high-water
marker mirroring :meth:`fts_rebuild_step`: each chunk deletes only
rows after the previously-drained key, so the per-chunk scan is
bounded instead of re-scanning from the start of the table every
chunk (O(n²) total on large trash tables, #79324). Compound-key
trash tables (multi-column PK) cannot use a scalar high-water
comparison, so they keep the legacy chunked ``LIMIT`` delete —
those shadow tables are small by construction.
"""
with self._lock:
trash = [
r[0] for r in self._conn.execute(
"SELECT name FROM sqlite_master WHERE type = 'table' "
"AND name LIKE ? ESCAPE '\\'",
(self._FTS_TRASH_PREFIX.replace("_", "\\_") + "%",),
).fetchall()
]
if not trash:
return False
tbl = trash[0]
def _do(conn):
pk_info = [
(r[1], (r[2] or "").upper())
for r in conn.execute(f"PRAGMA table_info({tbl})")
if r[5] > 0
]
pk_cols = [name for name, _typ in pk_info]
key = ", ".join(pk_cols) if pk_cols else "rowid"
if len(pk_cols) == 1 and (not pk_info or pk_info[0][1] == "INTEGER"):
# High-water drain: delete only rows past the marker key.
# The marker is read/written inside the same BEGIN IMMEDIATE
# transaction as the DELETE, so concurrent callers claim
# disjoint key ranges instead of re-deleting. Only integer
# PKs can anchor a numeric high-water comparison — the FTS
# config shadow table (TEXT pk like 'version') falls back to
# the legacy chunked delete below.
marker_key = f"fts_teardown_{tbl}_progress"
row = conn.execute(
"SELECT value FROM state_meta WHERE key = ?",
(marker_key,),
).fetchone()
high_water = int(row[0]) if row is not None else 0
# Claim the chunk's upper bound: the LAST row of the
# LIMIT window, so a full chunk is deleted per step.
upper_rows = conn.execute(
f"SELECT {key} FROM {tbl} WHERE {key} > ? "
f"ORDER BY {key} LIMIT {self._FTS_REBUILD_CHUNK_ROWS}",
(high_water,),
).fetchall()
if not upper_rows:
# Drained — the DROP is cheap now.
conn.execute(f"DROP TABLE IF EXISTS {tbl}")
conn.execute(
"DELETE FROM state_meta WHERE key = ?", (marker_key,)
)
logger.info("Old FTS shadow table %s torn down.", tbl)
return True
upper = upper_rows[-1][0]
cur = conn.execute(
f"DELETE FROM {tbl} WHERE {key} > ? AND {key} <= ?",
(high_water, upper),
)
if cur.rowcount > 0:
conn.execute(
"INSERT INTO state_meta (key, value) VALUES (?, ?) "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
(marker_key, str(upper)),
)
return True
# Compound-key or rowid trash table: legacy chunked delete.
# These shadow tables are small, so the quadratic re-scan is
# not a concern (#79324 keeps the high-water path for the big
# single-key tables).
cur = conn.execute(
f"DELETE FROM {tbl} WHERE ({key}) IN "
f"(SELECT {key} FROM {tbl} LIMIT {self._FTS_REBUILD_CHUNK_ROWS})"
)
if cur.rowcount == 0:
# Empty — the DROP is cheap now.
conn.execute(f"DROP TABLE IF EXISTS {tbl}")
logger.info("Old FTS shadow table %s torn down.", tbl)
return True # re-check: more trash tables / chunks may remain
try:
return bool(self._execute_write(_do))
except sqlite3.OperationalError as exc:
logger.debug("FTS trash teardown chunk failed (will retry): %s", exc)
return True
def fts_rebuild_step(self) -> bool:
"""Backfill one chunk of the deferred FTS rebuild.
Returns True when more work remains, False when the rebuild is
complete (or none is pending). Safe to call from any process at any
time; chunks are claimed atomically inside the write transaction, so
concurrent callers interleave instead of duplicating rows.
"""
if not self._fts_enabled:
return False
high_water_raw = self.get_meta("fts_rebuild_high_water")
if high_water_raw is None:
return False
high_water = int(high_water_raw)
include_trigram = self._trigram_available
chunk = self._FTS_REBUILD_CHUNK_ROWS
def _do(conn):
# Re-read progress inside the write transaction (BEGIN IMMEDIATE
# is already held by _execute_write) — this is the claim: two
# workers can't read the same progress value concurrently.
row = conn.execute(
"SELECT value FROM state_meta WHERE key = 'fts_rebuild_progress'"
).fetchone()
if row is None:
return False # finished (or cleared) by another process
progress = int(row[0])
if progress >= high_water:
return False
# The chunk upper bound is an id, not a row count, so gaps from
# deleted rows don't shrink chunks below the claimed range.
upper = min(progress + chunk, high_water)
conn.execute(
"INSERT INTO messages_fts(rowid, content, tool_name, tool_calls) "
"SELECT id, content, tool_name, tool_calls FROM messages "
"WHERE id > ? AND id <= ?",
(progress, upper),
)
if include_trigram:
conn.execute(
"INSERT INTO messages_fts_trigram"
"(rowid, content, tool_name, tool_calls) "
"SELECT id, content, tool_name, tool_calls FROM messages "
"WHERE id > ? AND id <= ? AND role <> 'tool'",
(progress, upper),
)
# Publish progress in the same transaction as the rows it
# covers — crash-atomic: either both land or neither does.
conn.execute(
"UPDATE state_meta SET value = ? "
"WHERE key = 'fts_rebuild_progress'",
(str(upper),),
)
return upper < high_water
try:
more = self._execute_write(_do)
except sqlite3.OperationalError as exc:
logger.debug("FTS rebuild chunk failed (will retry): %s", exc)
return True # transient (lock contention) — caller retries
if more is False:
status = self.fts_rebuild_status()
if status is not None and status["indexed"] >= status["total"]:
self._fts_rebuild_finish()
return False
return bool(more)
def fts_cjk_rebuild_status(self) -> Optional[Dict[str, Any]]:
"""CJK-index backfill progress, or None when none is pending."""
with self._read_ctx() as conn:
row = conn.execute(
"SELECT key, value FROM state_meta WHERE key IN (?, ?)",
("fts_cjk_rebuild_high_water", "fts_cjk_rebuild_progress"),
).fetchall()
meta = {r["key"]: r["value"] for r in row}
high_water = meta.get("fts_cjk_rebuild_high_water")
if high_water is None:
return None
progress = int(meta.get("fts_cjk_rebuild_progress") or 0)
total = int(high_water)
if total <= 0:
return None
pct = min(100, int(100 * progress / total))
return {"pending": True, "total": total, "indexed": progress, "percent": pct}
def fts_cjk_rebuild_step(self) -> bool:
"""Backfill one chunk of the CJK index. True while work remains."""
if not self._fts_enabled or not self._fts_cjk_loaded:
return False
high_water_raw = self.get_meta("fts_cjk_rebuild_high_water")
if high_water_raw is None:
return False
high_water = int(high_water_raw)
chunk = self._FTS_REBUILD_CHUNK_ROWS
def _do(conn):
row = conn.execute(
"SELECT value FROM state_meta "
"WHERE key = 'fts_cjk_rebuild_progress'"
).fetchone()
if row is None:
return False # finished (or cleared) by another process
progress = int(row[0])
if progress >= high_water:
return False
upper = min(progress + chunk, high_water)
conn.execute(
"INSERT INTO messages_fts_cjk(rowid, content, tool_name, tool_calls) "
"SELECT id, content, tool_name, tool_calls FROM messages "
"WHERE id > ? AND id <= ? AND role <> 'tool'",
(progress, upper),
)
conn.execute(
"UPDATE state_meta SET value = ? "
"WHERE key = 'fts_cjk_rebuild_progress'",
(str(upper),),
)
return upper < high_water
try:
more = self._execute_write(_do)
except sqlite3.OperationalError as exc:
logger.debug("CJK FTS rebuild chunk failed (will retry): %s", exc)
return True
if more is False:
status = self.fts_cjk_rebuild_status()
if status is not None and status["indexed"] >= status["total"]:
self._fts_cjk_rebuild_finish()
return False
return bool(more)
def _fts_cjk_rebuild_finish(self) -> None:
"""Boundary sweep + clear the cjk markers; index becomes servable."""
def _do(conn):
hw_row = conn.execute(
"SELECT value FROM state_meta "
"WHERE key = 'fts_cjk_rebuild_high_water'"
).fetchone()
if hw_row is not None:
hw = int(hw_row[0])
lo, hi = hw - 1000, hw + 1000
conn.execute(
"INSERT INTO messages_fts_cjk(rowid, content, tool_name, tool_calls) "
"SELECT m.id, m.content, m.tool_name, m.tool_calls "
"FROM messages m "
"WHERE m.id > ? AND m.id <= ? AND m.role <> 'tool' "
"AND NOT EXISTS (SELECT 1 FROM messages_fts_cjk_docsize d WHERE d.id = m.id)",
(lo, hi),
)
conn.execute(
"DELETE FROM state_meta WHERE key IN "
"('fts_cjk_rebuild_high_water', 'fts_cjk_rebuild_progress')"
)
self._execute_write(_do)
self._fts_cjk_available = True
logger.info("CJK FTS index backfill complete — serving CJK search.")
def _fts_cjk_reset_if_stale(self) -> None:
"""Rebuild path for a stale cjk index (triggers were dropped).
The gap's extent is unknown, so the only safe recovery is a from-
scratch rebuild: drop the table + triggers, clear the breadcrumb,
recreate via ``_ensure_fts_cjk_schema`` (which sets fresh backfill
markers on a populated DB). Called from ``optimize_fts_storage`` on
a tokenizer-capable host; no-op when not stale.
"""
if not self._fts_cjk_loaded:
return
def _do(conn):
stale = conn.execute(
"SELECT 1 FROM state_meta WHERE key = ?",
(FTS_CJK_STALE_KEY,),
).fetchone()
if not stale:
return False
for trig in _FTS_CJK_TRIGGERS:
conn.execute(f"DROP TRIGGER IF EXISTS {trig}")
conn.execute("DROP TABLE IF EXISTS messages_fts_cjk")
conn.execute("DROP VIEW IF EXISTS messages_fts_cjk_src")
conn.execute(
"DELETE FROM state_meta WHERE key IN "
f"('{FTS_CJK_STALE_KEY}', 'fts_cjk_rebuild_high_water', "
"'fts_cjk_rebuild_progress')"
)
return True
was_stale = self._execute_write(_do)
if was_stale:
# Recreate outside the write transaction — _ensure_fts_cjk_schema
# uses executescript(), which implicitly commits any pending
# transaction and must not run inside _execute_write's BEGIN
# IMMEDIATE. Sets fresh backfill markers on a populated DB.
with self._lock:
self._ensure_fts_cjk_schema(self._conn)
self._conn.commit()
def _fts_external_index_empty_with_messages(self, conn) -> bool:
"""True when the base FTS table exists but indexes nothing while
``messages`` has rows. Caller must hold ``self._lock``.
This is the post-demote empty-index shape: external-content FTS with
zero ``messages_fts_docsize`` rows against a non-empty messages table.
Healthy installs (and mid-backfill installs that still hold markers)
never match.
"""
try:
has_msg = conn.execute(
"SELECT EXISTS(SELECT 1 FROM messages)"
).fetchone()[0]
if not has_msg:
return False
# docsize is the authoritative "is this rowid indexed" surface for
# external-content FTS5; probing the virtual table itself is
# not reliable across SQLite builds. EXISTS instead of COUNT(*):
# this runs on every writable open via the _init_schema stamp
# condition, and COUNT(*) is a full b-tree scan (~100ms on a
# 2M-row table) while EXISTS is O(1).
has_fts = conn.execute(
"SELECT EXISTS(SELECT 1 FROM messages_fts_docsize)"
).fetchone()[0]
return not has_fts
except sqlite3.OperationalError:
# Table absent / FTS disabled mid-init — not this failure class.
return False
def _fts_index_known_empty(self, conn) -> bool:
"""True when the base external-content index holds no rows.
A missing table counts as empty: the schema ensure that follows
creates it fresh.
"""
try:
n = conn.execute(
"SELECT COUNT(*) FROM messages_fts_docsize"
).fetchone()[0]
return int(n) == 0
except sqlite3.OperationalError:
return True
def _reset_fts_index_to_empty(self, conn) -> None:
"""Delete every indexed row from the v23 external-content tables.
Uses the FTS5 ``'delete-all'`` special command — the documented O(1)
truncate for external-content tables. A plain no-WHERE ``DELETE`` is
O(rows) on external-content FTS5 (each row's delete tokens are
regenerated from the content table; measured ~12µs/row, minutes on a
large index, while holding the write lock) and corrupts the index if
indexed rows have diverged from ``messages`` — precisely the broken-
bookkeeping shape this repair path handles. The backfill chunk worker
replays its whole selected id range with no anti-join, so a replay
from zero is only safe once the index is known empty — this is how a
partially indexed DB gets there.
"""
for tbl in ("messages_fts", "messages_fts_trigram"):
try:
conn.execute(f"INSERT INTO {tbl}({tbl}) VALUES('delete-all')")
except sqlite3.OperationalError:
pass # table absent — already an empty surface
def _seed_fts_rebuild_markers(self, conn, *, force: bool = False) -> int:
"""Write ``fts_rebuild_high_water`` / ``fts_rebuild_progress`` for a
full backfill. Returns the high-water id.
When ``force`` is False and high_water is already set, only repairs a
missing progress key (stuck no-op when high_water exists alone), and
only after the index is known empty: the chunk worker replays its
whole selected id range without an anti-join, so a partially indexed
DB is first reset to a known-empty surface rather than rebuilt from
zero on top of surviving rows. Caller must hold the write
transaction / lock as appropriate.
"""
existing_hw = conn.execute(
"SELECT value FROM state_meta WHERE key = 'fts_rebuild_high_water'"
).fetchone()
if existing_hw is not None and not force:
hw = int(existing_hw[0])
progress = conn.execute(
"SELECT value FROM state_meta WHERE key = 'fts_rebuild_progress'"
).fetchone()
if progress is None:
# high_water without progress: fts_rebuild_step treats missing
# progress as "done by another process" and optimize would
# no-op then stamp. Re-seed progress so the chunk loop runs.
if not self._fts_index_known_empty(conn):
self._reset_fts_index_to_empty(conn)
conn.execute(
"INSERT INTO state_meta (key, value) VALUES "
"('fts_rebuild_progress', '0') "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value"
)
return hw
hw = conn.execute(
"SELECT COALESCE(MAX(id), 0) FROM messages"
).fetchone()[0]
for k, v in (
("fts_rebuild_high_water", str(hw)),
("fts_rebuild_progress", "0"),
):
conn.execute(
"INSERT INTO state_meta (key, value) VALUES (?, ?) "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
(k, v),
)
return int(hw)
def _repair_optimize_bookkeeping(self) -> None:
"""Heal interrupted demote/backfill bookkeeping before optimize runs.
Covers two post-#65798 failure classes:
1. Empty external-content index with messages present and no rebuild
markers (demote crash window after empty v23 tables landed but
before markers, or settle that stamped without backfill). Seed a
full backfill.
2. ``fts_rebuild_high_water`` present without ``fts_rebuild_progress``
(partial meta) — seed progress so the chunk loop is not a no-op,
resetting a partially populated index to a known-empty surface
first so the anti-join-free chunk replay cannot duplicate rows.
Must not invent markers on a still-legacy inline DB: that would make
``optimize_fts_storage`` skip demote (``legacy and not pending``) and
attempt v23-shaped INSERTs against the inline table forever.
"""
def _do(conn):
existing_hw = conn.execute(
"SELECT value FROM state_meta "
"WHERE key = 'fts_rebuild_high_water'"
).fetchone()
if existing_hw is not None:
# Repair orphan high_water-without-progress only. Never
# invent a fresh claim on a healthy complete index.
progress = conn.execute(
"SELECT 1 FROM state_meta "
"WHERE key = 'fts_rebuild_progress'"
).fetchone()
if progress is None:
if not self._fts_index_known_empty(conn):
self._reset_fts_index_to_empty(conn)
conn.execute(
"INSERT INTO state_meta (key, value) VALUES "
"('fts_rebuild_progress', '0') "
"ON CONFLICT(key) DO UPDATE SET value = '0'"
)
return
# No markers. On a still-legacy DB demote owns marker creation.
if self._db_has_legacy_inline_fts(conn):
return
# Non-legacy empty external index (demote crash window / premature
# stamp): seed a full backfill claim.
if self._fts_external_index_empty_with_messages(conn):
conn.execute(
"DELETE FROM state_meta WHERE key = 'fts_storage_version'"
)
self._seed_fts_rebuild_markers(conn, force=True)
self._execute_write(_do)
def fts_optimize_available(self) -> bool:
"""True when `optimize_fts_storage()` has work to do: either this DB
is a legacy inline-FTS install that can be optimized to the v23
external-content schema, or a previous optimize run was interrupted
(legacy vtables already demoted, but backfill markers and/or trash
tables remain) and re-running would resume it, or the CJK-bigram
index needs a backfill/rebuild on this tokenizer-capable host, or
a prior demote left an empty external-content index without markers
(healable on re-run).
False for fresh and fully-optimized installs (and when FTS5 is
unavailable)."""
if not self._fts_enabled or self.read_only:
return False
with self._lock:
if self._db_has_legacy_inline_fts(self._conn):
return True
# Interrupted optimize: demotion already removed the legacy
# vtables (so the check above is False), but the transition is
# unfinished until the backfill markers are cleared and the
# demoted trash tables are torn down. Search stays complete
# through the gap supplement meanwhile; re-running resumes.
if self._conn.execute(
"SELECT 1 FROM state_meta "
"WHERE key = 'fts_rebuild_high_water' LIMIT 1"
).fetchone():
return True
# CJK-bigram index work — only offerable when THIS process can
# tokenize: a pending backfill (markers set at creation on a
# populated DB) or a stale index awaiting a from-scratch rebuild.
if self._fts_cjk_loaded and self._conn.execute(
"SELECT 1 FROM state_meta WHERE key IN "
f"('fts_cjk_rebuild_high_water', '{FTS_CJK_STALE_KEY}') LIMIT 1"
).fetchone():
return True
if self._has_fts_trash(self._conn):
return True
# Pre-fix crash window: empty external-content index with
# messages still present, no markers, no trash (teardown already
# finished or never needed). Re-run seeds markers and backfills.
return self._fts_external_index_empty_with_messages(self._conn)
def _demote_legacy_fts_to_trash(self) -> int:
"""Demote the legacy inline FTS vtables and stage their shadow tables
for chunked teardown. Returns MAX(messages.id) as the rebuild high
water. O(1) schema surgery — the heavy delete is deferred to the
chunked teardown, exactly as the validated auto path did.
Markers are written in the same BEGIN IMMEDIATE as the demote, *before*
the empty v23 schema is created. Schema creation uses
``executescript`` and therefore cannot run inside that transaction
(it issues an implicit COMMIT — see the CJK recreate path). Creating
the empty schema only after markers are durable closes the crash
window where trash + empty v23 tables exist with no backfill claim.
"""
def _stage(conn):
self._drop_fts_triggers(conn)
conn.execute("DROP VIEW IF EXISTS messages_fts_trigram_src")
had = bool(conn.execute(
"SELECT 1 FROM sqlite_master WHERE type = 'table' "
"AND name IN ('messages_fts', 'messages_fts_trigram') "
"AND sql LIKE 'CREATE VIRTUAL TABLE%' LIMIT 1"
).fetchone())
if had:
conn.execute("PRAGMA writable_schema=ON")
conn.execute(
"DELETE FROM sqlite_master WHERE type = 'table' "
"AND name IN ('messages_fts', 'messages_fts_trigram') "
"AND sql LIKE 'CREATE VIRTUAL TABLE%'"
)
conn.execute("PRAGMA writable_schema=RESET")
shadows = [
r[0] for r in conn.execute(
"SELECT name FROM sqlite_master WHERE type = 'table' "
"AND (name LIKE 'messages_fts_%' ESCAPE '\\' "
"OR name LIKE 'messages_fts_trigram_%' ESCAPE '\\')"
).fetchall()
]
for sh in shadows:
conn.execute(f"ALTER TABLE {sh} RENAME TO fts_v22_trash_{sh}")
# Claim the backfill *before* empty v23 tables exist. A crash
# between this commit and schema ensure still leaves markers, so
# optimize-storage resumes instead of tearing down trash and
# stamping an empty index as complete.
hw = self._seed_fts_rebuild_markers(conn, force=True)
conn.execute(
"DELETE FROM state_meta WHERE key = 'fts_optimize_available'"
)
return hw
hw = int(self._execute_write(_stage))
# Create the empty v23 schema outside the write transaction —
# ``_ensure_fts_schema`` uses executescript(), which implicitly
# commits any pending transaction and must not run inside
# ``_execute_write``'s BEGIN IMMEDIATE (same rule as the CJK recreate
# path above). Markers are already durable.
with self._lock:
base_ok = self._ensure_fts_schema(self._conn, "messages_fts", FTS_SQL)
trigram_ok = self._ensure_fts_schema(
self._conn, "messages_fts_trigram", FTS_TRIGRAM_SQL
)
self._trigram_available = bool(trigram_ok)
if not base_ok:
raise sqlite3.OperationalError(
"failed to create v23 messages_fts during optimize-storage demote"
)
self._conn.commit()
return hw
def optimize_fts_storage(
self,
*,
progress_cb: Optional[Callable[[Dict[str, Any]], None]] = None,
vacuum: bool = True,
) -> Dict[str, Any]:
"""Migrate a legacy v22 inline-FTS DB to the v23 external-content
schema, foreground and to completion. Safe to re-run: if a previous
attempt was interrupted it resumes from the progress marker.
``progress_cb`` receives {"phase", "percent", "indexed", "total"}
dicts for a CLI progress bar. Returns a summary dict.
The trigram tokenizer being unavailable is not fatal — the base index
is still rebuilt (CJK falls back to LIKE), mirroring normal startup.
"""
if not self._fts_enabled:
return {"ok": False, "reason": "fts5_unavailable"}
if self.read_only:
return {"ok": False, "reason": "read_only"}
# Heal empty-index / orphan-marker bookkeeping from an interrupted
# demote *before* deciding whether to demote again. This re-seeds
# markers when trash was already staged (or torn down) without a
# backfill claim so the phases below actually run.
self._repair_optimize_bookkeeping()
# Only demote if we're actually still on the legacy shape. If a prior
# run already demoted (markers/trash present), skip straight to
# finishing the backfill + teardown — this is what makes re-running
# after an interruption safe.
with self._lock:
legacy = self._db_has_legacy_inline_fts(self._conn)
pending = self.get_meta("fts_rebuild_high_water") is not None
if legacy and not pending:
self._demote_legacy_fts_to_trash()
elif pending and not legacy:
# Resume mid-demote: markers exist, empty v23 tables may still be
# missing if the process died between the staged demote commit and
# schema ensure. Re-ensure is IF NOT EXISTS and cheap.
with self._lock:
base_ok = self._ensure_fts_schema(
self._conn, "messages_fts", FTS_SQL
)
trigram_ok = self._ensure_fts_schema(
self._conn, "messages_fts_trigram", FTS_TRIGRAM_SQL
)
self._trigram_available = bool(trigram_ok)
if not base_ok:
# Fail fast: without the base table the backfill loop
# below would retry "no such table" errors forever.
raise sqlite3.OperationalError(
"failed to re-create v23 messages_fts "
"on optimize-storage resume"
)
self._conn.commit()
# A stale CJK index (triggers dropped by a tokenizer-less process)
# can only be recovered from scratch — reset it now so the cjk
# backfill phase below rebuilds it. No-op without the tokenizer.
self._fts_cjk_reset_if_stale()
# An optimized v23 DB gaining the cjk index for the first time (no
# legacy work left, tokenizer newly installed): ensure the table +
# markers exist so the backfill phase has work to claim.
if self._fts_cjk_loaded:
with self._lock:
self._ensure_fts_cjk_schema(self._conn)
self._conn.commit()
def _emit(phase: str) -> None:
if progress_cb is None:
return
st = self.fts_rebuild_status()
if st is None:
st = self.fts_cjk_rebuild_status()
progress_cb({
"phase": phase,
"percent": st["percent"] if st else 100,
"indexed": st["indexed"] if st else 0,
"total": st["total"] if st else 0,
})
def _pause(chunk_seconds: float) -> None:
"""Inter-chunk throttle (see the chunk-engine note above).
The chunk methods themselves never sleep, so this loop is the
single place the duty cycle is enforced: without it, back-to-back
BEGIN IMMEDIATE chunks starve any live gateway/CLI process
sharing the DB out of its lock retries (the measured ~85%
write-lock ownership that froze concurrent sessions).
"""
time.sleep(max(
self._FTS_REBUILD_MIN_PAUSE,
chunk_seconds * self._FTS_REBUILD_DUTY_FACTOR,
))
# Phase 1: backfill (foreground, throttled between chunks so a live
# gateway sharing the DB stays responsive).
_emit("backfill")
while True:
_t0 = time.monotonic()
if not self.fts_rebuild_step():
break
_emit("backfill")
_pause(time.monotonic() - _t0)
_emit("backfill")
# Phase 1b: backfill the CJK-bigram index (its own marker pair; a
# no-op when the tokenizer isn't loadable or nothing is pending).
while True:
_t0 = time.monotonic()
if not self.fts_cjk_rebuild_step():
break
_emit("backfill")
_pause(time.monotonic() - _t0)
# Phase 2: tear down the demoted legacy shadow tables in chunks.
_emit("teardown")
while True:
_t0 = time.monotonic()
if not self._fts_teardown_trash_step():
break
_emit("teardown")
_pause(time.monotonic() - _t0)
# Refuse to stamp "optimized" while work remains or the base index is
# still empty against a non-empty messages table. Pre-fix code could
# tear down trash and settle after a no-op backfill when markers were
# missing — permanent search-index loss for historical rows.
with self._lock:
still_pending = self._conn.execute(
"SELECT 1 FROM state_meta "
"WHERE key = 'fts_rebuild_high_water' LIMIT 1"
).fetchone() is not None
still_trash = self._has_fts_trash(self._conn)
empty_index = self._fts_external_index_empty_with_messages(self._conn)
if still_pending or still_trash or empty_index:
reason = (
"backfill_incomplete" if still_pending or empty_index
else "teardown_incomplete"
)
logger.warning(
"FTS storage optimization did not settle (%s): "
"pending=%s trash=%s empty_index=%s",
reason, still_pending, still_trash, empty_index,
)
return {"ok": False, "reason": reason, "vacuumed": None}
# Phase 3: reclaim freed pages to the OS.
vacuum_ok = None
if vacuum:
_emit("vacuum")
try:
with self._lock:
self._conn.execute("VACUUM")
vacuum_ok = True
except sqlite3.OperationalError as exc:
# Most common cause: not enough free disk for VACUUM's temp
# copy. The optimization still succeeded; space just isn't
# reclaimed until a later VACUUM. Non-fatal.
logger.warning("VACUUM after FTS optimize failed: %s", exc)
vacuum_ok = False
# Best-effort: fold the WAL back into the main file so the on-disk
# size settles now rather than at close(). NOTE this is REFUSED
# (SQLITE_BUSY) while any other connection holds a WAL read-mark —
# e.g. a live gateway sharing the DB — so it is not sufficient on
# its own. Callers must therefore NOT size the result by stat()ing
# the file; use :meth:`logical_size_bytes`, which is truthful
# immediately regardless of readers.
try:
with self._lock:
self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
except Exception as exc:
logger.debug(
"WAL checkpoint (TRUNCATE) after optimize VACUUM failed: %s",
exc,
)
# Phase 4: stamp the FTS storage layout as current, clear the "available"
# flag, and advance schema_version if it was somehow still behind (the
# main version normally advances on open now, but bump defensively so a
# DB opened only by pre-decoupling code still settles). The FTS-layout
# marker is the source of truth for "is this DB optimized".
def _settle(conn):
# Re-check inside the write transaction so a concurrent writer
# cannot race a stamp past incomplete work. Returns a refusal
# reason (stamping nothing) or None once the stamp is written.
if conn.execute(
"SELECT 1 FROM state_meta "
"WHERE key = 'fts_rebuild_high_water' LIMIT 1"
).fetchone() is not None:
return "backfill_incomplete"
if self._has_fts_trash(conn):
return "teardown_incomplete"
if self._fts_external_index_empty_with_messages(conn):
return "backfill_incomplete"
conn.execute(
"INSERT INTO state_meta (key, value) VALUES ('fts_storage_version', ?) "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
(str(FTS_STORAGE_VERSION),),
)
conn.execute("DELETE FROM state_meta WHERE key = 'fts_optimize_available'")
conn.execute(
"UPDATE schema_version SET version = ? WHERE version < ?",
(SCHEMA_VERSION, SCHEMA_VERSION),
)
return None
refusal = self._execute_write(_settle)
if refusal is not None:
# A concurrent process re-seeded markers, left trash, or emptied
# the index between the pre-vacuum check above and this write
# transaction. Nothing was stamped. Report the failure instead of
# crashing the CLI with a traceback; a re-run can still settle.
logger.warning(
"FTS storage optimization settle refused (%s)", refusal
)
return {"ok": False, "reason": refusal, "vacuumed": vacuum_ok}
_emit("done")
logger.info(
"FTS storage optimization complete (layout v%d).", FTS_STORAGE_VERSION
)
return {"ok": True, "vacuumed": vacuum_ok}
def get_anchored_view(
self,
session_id: str,
around_message_id: int,
window: int = 5,
bookend: int = 3,
keep_roles: Optional[Tuple[str, ...]] = ("user", "assistant"),
) -> Dict[str, Any]:
"""Return an anchored window plus session bookends.
Built on top of ``get_messages_around``. Three slices:
- ``window``: messages immediately surrounding the anchor. Filtered
to ``keep_roles`` (tool-response noise dropped by default), EXCEPT
the anchor itself is always preserved regardless of role.
- ``bookend_start``: first ``bookend`` user/assistant messages of the
session — but only those whose id is strictly before the window's
first message id. Empty when the window already overlaps the
session head. Empty-content messages (tool-call-only assistant
turns) are skipped so they don't crowd out actual prose openings.
- ``bookend_end``: last ``bookend`` user/assistant messages of the
session, same non-overlap rule at the tail.
Bookends let an FTS5 hit anywhere in a long session yield the goal
(opening) and the resolution (closing) on a single call — without
loading the whole transcript.
Returns ``{"window": [], "messages_before": 0, "messages_after": 0,
"bookend_start": [], "bookend_end": []}`` when the anchor isn't in
the session.