From 09d4c5456ddd535f40fe927b3d4fd1fc50bf1695 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 30 Jul 2026 18:36:39 +0800 Subject: [PATCH 1/3] feat(memory): add global and session search --- flocks/config/config_writer.py | 66 +- flocks/hooks/builtin/slug_generator.py | 2 +- flocks/memory/__init__.py | 2 + flocks/memory/config.py | 33 +- flocks/memory/manager.py | 399 ++++++++--- flocks/memory/search/hybrid.py | 69 +- flocks/memory/sync/indexer.py | 333 +++++---- flocks/memory/types.py | 8 +- flocks/session/features/memory.py | 7 +- flocks/session/message.py | 214 +++++- flocks/session/session.py | 8 + flocks/storage/__init__.py | 2 + flocks/storage/session_search.py | 671 ++++++++++++++++++ flocks/storage/storage.py | 39 +- flocks/storage/vector.py | 205 +++++- tests/config/test_config_init.py | 40 +- tests/memory/test_memory_scope.py | 311 ++++++++ .../memory/test_session_transcript_search.py | 622 ++++++++++++++++ .../session/test_message_parts_persistence.py | 19 +- tests/storage/test_storage.py | 44 ++ 20 files changed, 2779 insertions(+), 315 deletions(-) create mode 100644 flocks/storage/session_search.py create mode 100644 tests/memory/test_memory_scope.py create mode 100644 tests/memory/test_session_transcript_search.py diff --git a/flocks/config/config_writer.py b/flocks/config/config_writer.py index 7251c6d17..9ba939133 100644 --- a/flocks/config/config_writer.py +++ b/flocks/config/config_writer.py @@ -77,6 +77,8 @@ def ensure_config_files() -> None: "error": str(e), }) + ConfigWriter.ensure_memory_config() + class ConfigWriter: """Atomic read-modify-write operations on the provider section of flocks.json.""" @@ -106,9 +108,13 @@ def _read_raw(cls) -> Dict[str, Any]: return {} @classmethod - def _write_raw(cls, data: Dict[str, Any]) -> None: + def _write_raw( + cls, + data: Dict[str, Any], + path: Optional[Path] = None, + ) -> None: """Atomic write: write to tmp file then rename, then clear Config cache.""" - path = cls._get_config_path() + path = path or cls._get_config_path() path.parent.mkdir(parents=True, exist_ok=True) # Atomic write via temp file in same directory @@ -136,6 +142,62 @@ def _write_raw(cls, data: Dict[str, Any]) -> None: log.debug("config_writer.written", {"path": str(path)}) + @classmethod + def ensure_memory_config(cls) -> bool: + """Persist the editable Memory Search config when absent.""" + path = Config.get_config_file() + try: + text = path.read_text(encoding="utf-8") if path.exists() else "" + data = json.loads(text) if text.strip() else {} + except (json.JSONDecodeError, OSError) as exc: + log.error( + "config_writer.memory_config_init_failed", + {"path": str(path), "error": str(exc)}, + ) + return False + + if not isinstance(data, dict): + log.error( + "config_writer.memory_config_init_failed", + {"path": str(path), "error": "top-level config must be an object"}, + ) + return False + if "memory" in data: + return False + + from flocks.memory.config import MemoryConfig + + default_config = MemoryConfig() + data["memory"] = { + "search": { + "embedding": default_config.search.embedding.model_dump( + mode="json", + exclude_none=True, + ), + }, + } + cls._write_raw(data, path=path) + log.info("config_writer.memory_config_initialized", {"path": str(path)}) + return True + + @classmethod + def enable_memory_source(cls, source: str) -> bool: + """Persist a Memory source without rewriting unrelated config.""" + data = cls._read_raw() + memory = data.get("memory") + if not isinstance(memory, dict): + memory = {} + sources = memory.get("sources") + if not isinstance(sources, list): + sources = ["memory"] + if source in sources: + return False + memory["sources"] = [*sources, source] + data["memory"] = memory + cls._write_raw(data) + log.info("config_writer.memory_source_enabled", {"source": source}) + return True + # ------------------------------------------------------------------ # Provider-level CRUD # ------------------------------------------------------------------ diff --git a/flocks/hooks/builtin/slug_generator.py b/flocks/hooks/builtin/slug_generator.py index 2ff78a3dd..2bbc2073a 100644 --- a/flocks/hooks/builtin/slug_generator.py +++ b/flocks/hooks/builtin/slug_generator.py @@ -46,7 +46,7 @@ async def generate_slug_via_llm( """ # Get provider configuration - provider_id = getattr(config.memory.embedding, 'provider', 'openai') + provider_id = getattr(config.memory.search.embedding, 'provider', 'openai') if provider_id == "auto": provider_id = "openai" diff --git a/flocks/memory/__init__.py b/flocks/memory/__init__.py index d13797cb3..36c70deec 100644 --- a/flocks/memory/__init__.py +++ b/flocks/memory/__init__.py @@ -28,6 +28,7 @@ from flocks.memory.config import ( MemoryConfig, MemoryEmbeddingConfig, + MemorySearchConfig, MemoryChunkingConfig, MemorySyncConfig, MemoryQueryConfig, @@ -68,6 +69,7 @@ # Config "MemoryConfig", "MemoryEmbeddingConfig", + "MemorySearchConfig", "MemoryChunkingConfig", "MemorySyncConfig", "MemoryQueryConfig", diff --git a/flocks/memory/config.py b/flocks/memory/config.py index 7d349862a..21dea8f3d 100644 --- a/flocks/memory/config.py +++ b/flocks/memory/config.py @@ -10,6 +10,10 @@ class MemoryEmbeddingConfig(BaseModel): """Embedding provider configuration""" + enabled: bool = Field( + False, + description="Enable vector embeddings for Memory search", + ) provider: Literal["auto", "openai", "google", "local"] = Field( "auto", description="Embedding provider (auto=try openai then google)" @@ -32,6 +36,15 @@ class MemoryEmbeddingConfig(BaseModel): ) +class MemorySearchConfig(BaseModel): + """Memory search configuration.""" + + embedding: MemoryEmbeddingConfig = Field( + default_factory=MemoryEmbeddingConfig, + description="Embedding configuration", + ) + + class MemoryChunkingConfig(BaseModel): """Text chunking configuration""" tokens: int = Field( @@ -52,7 +65,7 @@ class MemorySyncSessionConfig(BaseModel): ) delta_messages: int = Field( 50, - description="Number of new messages to trigger sync" + description="Batch size for session transcript reconciliation" ) @@ -64,7 +77,7 @@ class MemorySyncConfig(BaseModel): ) on_search: bool = Field( True, - description="Sync before search if dirty" + description="Reconcile filesystem Memory before every search" ) watch: bool = Field( True, @@ -295,9 +308,9 @@ class MemoryConfig(BaseModel): ) # Sub-configurations - embedding: MemoryEmbeddingConfig = Field( - default_factory=MemoryEmbeddingConfig, - description="Embedding configuration" + search: MemorySearchConfig = Field( + default_factory=MemorySearchConfig, + description="Memory search configuration", ) chunking: MemoryChunkingConfig = Field( default_factory=MemoryChunkingConfig, @@ -330,7 +343,15 @@ class MemoryConfig(BaseModel): def resolve_memory_config(app_config: object) -> MemoryConfig: - """Resolve runtime Memory config, using defaults when absent.""" + """Resolve runtime Memory config, using defaults when absent. + + Args: + app_config: Loaded application configuration. + + Returns: + Configured Memory settings, or defaults when the application has no + Memory section. + """ memory_config = getattr(app_config, "memory", None) if isinstance(memory_config, MemoryConfig): return memory_config diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index 2f1857faa..99978a79e 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -35,6 +35,90 @@ def _safe_resolve_memory_path(memory_root: Path, rel_path: str) -> Path: return resolved +class _MemoryIndexCoordinator: + """Own the process-wide Memory file indexer for one SQLite database.""" + + def __init__(self) -> None: + self.indexer: Optional[MemoryIndexer] = None + self.signature: Optional[tuple[Any, ...]] = None + self.initialized = False + self.sync_lock = asyncio.Lock() + self.write_lock = asyncio.Lock() + + def configure( + self, + *, + workspace_dir: Path, + provider_id: Optional[str], + embedding_model: str, + config: MemoryConfig, + ) -> MemoryIndexer: + """Create or reuse the one global file indexer.""" + signature = ( + provider_id, + embedding_model, + config.chunking.model_dump_json(), + config.batch.model_dump_json(), + config.cache.model_dump_json(), + tuple(config.extra_paths), + ) + if self.indexer is not None and self.signature == signature: + return self.indexer + + self.indexer = MemoryIndexer( + project_id="global", + workspace_dir=workspace_dir, + provider_id=provider_id, + embedding_model=embedding_model, + config=config, + ) + self.signature = signature + self.initialized = False + return self.indexer + + async def _sync_locked( + self, + *, + force: bool, + progress_callback: Optional[Callable[[MemorySyncProgress], None]], + ) -> Dict[str, Any]: + """Reconcile while the coordinator sync lock is held.""" + if self.indexer is None: + raise RuntimeError("Memory file indexer is not configured") + async with self.write_lock: + stats = await self.indexer.sync( + force=force, + progress_callback=progress_callback, + ) + self.initialized = True + return stats + + async def sync_on_start(self) -> Optional[Dict[str, Any]]: + """Run the initial reconciliation once for a shared indexer.""" + async with self.sync_lock: + if self.initialized: + return None + return await self._sync_locked( + force=False, + progress_callback=None, + ) + + async def sync( + self, + *, + force: bool = False, + progress_callback: Optional[ + Callable[[MemorySyncProgress], None] + ] = None, + ) -> Dict[str, Any]: + """Serialize global Memory index reconciliation.""" + async with self.sync_lock: + return await self._sync_locked( + force=force, + progress_callback=progress_callback, + ) + + class MemoryManager: """ Memory manager - orchestrates memory system @@ -47,6 +131,7 @@ class MemoryManager: # Singleton cache by project_id _instances: Dict[str, "MemoryManager"] = {} + _index_coordinators: Dict[str, _MemoryIndexCoordinator] = {} def __init__( self, @@ -67,11 +152,17 @@ def __init__( self.config = config # Provider configuration - self.provider_id = config.embedding.provider - if self.provider_id == "auto": + self._embedding_enabled = config.search.embedding.enabled + self._requested_provider = config.search.embedding.provider + self.provider_id: Optional[str] = ( + config.search.embedding.provider + if self._embedding_enabled + else None + ) + if self._embedding_enabled and self.provider_id == "auto": self.provider_id = "openai" # Default fallback - self.embedding_model = config.embedding.model + self.embedding_model = config.search.embedding.model # Components (lazy initialization) self.search_engine: Optional[HybridSearch] = None @@ -79,11 +170,19 @@ def __init__( # State self._initialized = False - self._dirty = False - self._sync_lock = asyncio.Lock() self._init_lock = asyncio.Lock() - self._write_lock = asyncio.Lock() - + self._index_coordinator: Optional[_MemoryIndexCoordinator] = None + + @classmethod + def _coordinator_for_active_db(cls) -> _MemoryIndexCoordinator: + """Return the global Memory index owner for the active database.""" + key = str(Storage.get_db_path().resolve()) + coordinator = cls._index_coordinators.get(key) + if coordinator is None: + coordinator = _MemoryIndexCoordinator() + cls._index_coordinators[key] = coordinator + return coordinator + @classmethod def get_instance( cls, @@ -111,25 +210,37 @@ def get_instance( if project_id in cls._instances: instance = cls._instances[project_id] - old_provider = instance.provider_id + old_enabled = instance._embedding_enabled + old_provider = instance._requested_provider old_model = instance.embedding_model instance.config = config instance.workspace_dir = Path(workspace_dir) - new_provider = config.embedding.provider - if new_provider == "auto": - new_provider = "openai" - new_model = config.embedding.model + new_enabled = config.search.embedding.enabled + new_provider = config.search.embedding.provider + new_model = config.search.embedding.model - if new_provider != old_provider or new_model != old_model: - instance.provider_id = new_provider + if ( + new_enabled != old_enabled + or new_provider != old_provider + or new_model != old_model + ): + instance._embedding_enabled = new_enabled + instance._requested_provider = new_provider + instance.provider_id = ( + ("openai" if new_provider == "auto" else new_provider) + if new_enabled + else None + ) instance.embedding_model = new_model instance._initialized = False instance.search_engine = None instance.indexer = None log.info("manager.config_changed", { "project_id": project_id, + "old_enabled": old_enabled, + "new_enabled": new_enabled, "old_provider": old_provider, "new_provider": new_provider, "old_model": old_model, @@ -157,25 +268,27 @@ async def initialize(self) -> None: log.info("manager.init.start", {"project_id": self.project_id}) try: - await Storage.init() - await Provider.init() - - provider = Provider.get(self.provider_id) - if not provider: - raise ValueError(f"Provider {self.provider_id} not found") + await Storage._ensure_init() - if not provider.supports_embeddings(): - for fallback_id in ["openai", "google"]: - fallback = Provider.get(fallback_id) - if fallback and fallback.supports_embeddings(): - log.warn("manager.provider.fallback", { - "from": self.provider_id, - "to": fallback_id, - }) - self.provider_id = fallback_id - break - else: - raise ValueError("No provider with embeddings support available") + if self._embedding_enabled: + await Provider.init() + provider = Provider.get(self.provider_id) if self.provider_id else None + if not provider or not provider.supports_embeddings(): + for fallback_id in ["openai", "google"]: + fallback = Provider.get(fallback_id) + if fallback and fallback.supports_embeddings(): + log.warn("manager.provider.fallback", { + "from": self.provider_id, + "to": fallback_id, + }) + self.provider_id = fallback_id + break + else: + log.info( + "manager.embedding.unavailable", + {"project_id": self.project_id}, + ) + self.provider_id = None self.search_engine = HybridSearch( project_id=self.project_id, @@ -184,22 +297,47 @@ async def initialize(self) -> None: config=self.config.query, ) - self.indexer = MemoryIndexer( - project_id=self.project_id, + coordinator = self._coordinator_for_active_db() + self._index_coordinator = coordinator + previous_indexer = coordinator.indexer + self.indexer = coordinator.configure( workspace_dir=self.workspace_dir, provider_id=self.provider_id, embedding_model=self.embedding_model, config=self.config, ) + if self.indexer is not previous_indexer: + for manager in self._instances.values(): + if manager._index_coordinator is coordinator: + manager.indexer = self.indexer self._initialized = True + if self.config.sync.on_session_start: + await coordinator.sync_on_start() + if ( + "session" in self.config.sources + and self.config.sync.sessions.enabled + ): + if Storage.session_search_available(): + await self._ensure_session_index_ready() + else: + log.warn( + "manager.session_search.disabled", + { + "project_id": self.project_id, + "reason": ( + "SQLite runtime does not support FTS5" + ), + }, + ) log.info("manager.init.complete", { "project_id": self.project_id, - "provider": self.provider_id, + "provider": self.provider_id or "fts", "model": self.embedding_model, }) except Exception as e: + self._initialized = False log.error("manager.init.failed", {"error": str(e)}) raise @@ -225,23 +363,125 @@ async def search( if not self._initialized: await self.initialize() - # Trigger sync if configured and dirty - if self.config.sync.on_search and self._dirty: - await self.sync(reason="search") - - # Execute search - results = await self.search_engine.search( - query=query, - max_results=max_results or self.config.query.max_results, - min_score=min_score or self.config.query.min_score, - sources=sources or [MemorySource(s) for s in self.config.sources], + selected_sources = ( + list(sources) + if sources is not None + else [MemorySource(source) for source in self.config.sources] ) + limit = ( + max_results + if max_results is not None + else self.config.query.max_results + ) + threshold = ( + min_score + if min_score is not None + else self.config.query.min_score + ) + + if sources is not None and MemorySource.SESSION in selected_sources: + await self._persist_session_source() + + # Filesystem tools and external editors can update Memory without going + # through MemoryManager. Reconcile on every search and let the indexer + # skip files whose content hash is unchanged. + if self.config.sync.on_search: + await self.sync(reason="search") + + results: List[MemorySearchResult] = [] + errors: List[Exception] = [] + successful_sources = 0 + + if MemorySource.MEMORY in selected_sources: + try: + results.extend( + await self.search_engine.search( + query=query, + max_results=limit, + min_score=threshold, + sources=[MemorySource.MEMORY], + ) + ) + successful_sources += 1 + except Exception as exc: + errors.append(exc) + log.warn("manager.search.memory_failed", {"error": str(exc)}) + + if MemorySource.SESSION in selected_sources: + try: + await self._ensure_session_index_ready() + from flocks.storage.session_search import session_fts_search + + raw_results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=self.project_id, + query=query, + max_results=limit + * self.config.query.hybrid.candidate_multiplier, + ) + results.extend( + MemorySearchResult( + path=result["path"], + start_line=result["start_line"], + end_line=result["end_line"], + score=result["score"], + snippet=result["text"][:700], + source=MemorySource.SESSION, + citation=result["citation"], + ) + for result in raw_results + if result["score"] >= threshold + ) + successful_sources += 1 + except Exception as exc: + errors.append(exc) + log.warn("manager.search.session_failed", {"error": str(exc)}) + + if successful_sources == 0 and errors: + from flocks.storage.session_search import ( + SessionSearchUnavailableError, + ) + + if len(errors) == 1 and isinstance( + errors[0], + SessionSearchUnavailableError, + ): + raise errors[0] + raise RuntimeError( + "All requested memory sources failed: " + + "; ".join(str(error) for error in errors) + ) + + deduplicated: Dict[str, MemorySearchResult] = {} + for result in sorted(results, key=lambda item: item.score, reverse=True): + key = f"{result.source.value}:{result.path}" + deduplicated.setdefault(key, result) + results = list(deduplicated.values())[:limit] # Decorate citations if enabled if self.config.citations != "off": results = decorate_citations(results, mode=self.config.citations) return results + + async def _persist_session_source(self) -> None: + """Persist explicit Session search opt-in without touching other config.""" + if "session" in self.config.sources: + return + from flocks.config.config_writer import ConfigWriter + + await asyncio.to_thread(ConfigWriter.enable_memory_source, "session") + self.config.sources.append("session") + + async def _ensure_session_index_ready(self) -> None: + """Run the one-time historical Session backfill when required.""" + if not self.config.sync.sessions.enabled: + return + from flocks.storage.session_search import ensure_session_index_ready + + await ensure_session_index_ready( + batch_size=self.config.sync.sessions.delta_messages, + ) async def read_file( self, @@ -315,7 +555,8 @@ async def write_memory( file_path = _safe_resolve_memory_path(memory_root, path) file_path.parent.mkdir(parents=True, exist_ok=True) - async with self._write_lock: + coordinator = self._index_coordinator or self._coordinator_for_active_db() + async with coordinator.write_lock: if append: needs_separator = file_path.exists() and file_path.stat().st_size > 0 with open(file_path, "a", encoding="utf-8") as f: @@ -325,14 +566,10 @@ async def write_memory( else: with open(file_path, "w", encoding="utf-8") as f: f.write(content) - - # Mark as dirty for next sync - self._dirty = True - log.info("manager.write", {"path": path, "append": append, "length": len(content)}) return path - + async def sync( self, reason: Optional[str] = None, @@ -353,27 +590,25 @@ async def sync( if not self._initialized: await self.initialize() - async with self._sync_lock: - log.info("manager.sync.start", { - "project_id": self.project_id, - "reason": reason, - "force": force, - }) - - try: - stats = await self.indexer.sync( - force=force, - progress_callback=progress_callback, - ) - - self._dirty = False - - log.info("manager.sync.complete", stats) - return stats - - except Exception as e: - log.error("manager.sync.failed", {"error": str(e)}) - raise + coordinator = self._index_coordinator or self._coordinator_for_active_db() + log.info("manager.sync.start", { + "project_id": self.project_id, + "reason": reason, + "force": force, + }) + + try: + stats = await coordinator.sync( + force=force, + progress_callback=progress_callback, + ) + + log.info("manager.sync.complete", stats) + return stats + + except Exception as e: + log.error("manager.sync.failed", {"error": str(e)}) + raise def status(self) -> MemoryProviderStatus: """ @@ -385,25 +620,29 @@ def status(self) -> MemoryProviderStatus: # TODO: Implement comprehensive status collection return MemoryProviderStatus( enabled=True, - provider=self.provider_id, + provider=self.provider_id or "fts", model=self.embedding_model, - requested_provider=self.config.embedding.provider, + requested_provider=self.config.search.embedding.provider, workspace_dir=str(self.workspace_dir), sources=[MemorySource(s) for s in self.config.sources], - dirty=self._dirty, cache={"enabled": self.config.cache.enabled}, fts={"enabled": True}, # Always available - vector={"enabled": True}, # Always available + vector={"enabled": self.provider_id is not None}, ) async def close(self) -> None: """Close and cleanup manager""" + coordinator = self._index_coordinator self._initialized = False self.search_engine = None self.indexer = None + self._index_coordinator = None self._instances.pop(self.project_id, None) + if coordinator is not None and not any( + manager._index_coordinator is coordinator + for manager in self._instances.values() + ): + for key, candidate in list(self._index_coordinators.items()): + if candidate is coordinator: + self._index_coordinators.pop(key, None) log.info("manager.closed", {"project_id": self.project_id}) - - def mark_dirty(self) -> None: - """Mark as needing sync""" - self._dirty = True diff --git a/flocks/memory/search/hybrid.py b/flocks/memory/search/hybrid.py index b4707e3f6..4dd9dd48d 100644 --- a/flocks/memory/search/hybrid.py +++ b/flocks/memory/search/hybrid.py @@ -24,7 +24,7 @@ class HybridSearch: def __init__( self, project_id: str, - provider_id: str, + provider_id: Optional[str], embedding_model: str, config: MemoryQueryConfig, ): @@ -69,14 +69,41 @@ async def search( }) try: - if not self.config.hybrid.enabled: - # Vector-only search - return await self._vector_search( + if self.provider_id is None: + results = await self._keyword_search( query=query, max_results=max_results, - min_score=min_score, sources=sources, ) + return [ + result + for result in results + if result.score >= min_score + ][:max_results] + + if not self.config.hybrid.enabled: + try: + return await self._vector_search( + query=query, + max_results=max_results, + min_score=min_score, + sources=sources, + ) + except Exception as exc: + log.warn( + "search.vector.failed_fts_fallback", + {"error": str(exc)}, + ) + results = await self._keyword_search( + query=query, + max_results=max_results, + sources=sources, + ) + return [ + result + for result in results + if result.score >= min_score + ][:max_results] # Hybrid search: parallel vector + keyword search candidate_limit = max_results * self.config.hybrid.candidate_multiplier @@ -97,13 +124,31 @@ async def search( ) # Handle exceptions - if isinstance(vector_results, Exception): + vector_failed = isinstance(vector_results, Exception) + keyword_failed = isinstance(keyword_results, Exception) + if vector_failed: log.warn("search.vector.failed", {"error": str(vector_results)}) vector_results = [] - if isinstance(keyword_results, Exception): + if keyword_failed: log.warn("search.keyword.failed", {"error": str(keyword_results)}) keyword_results = [] + + if vector_failed and keyword_failed: + raise RuntimeError("Both vector and keyword search failed") + + if keyword_results and not vector_results: + return [ + result + for result in keyword_results + if result.score >= min_score + ][:max_results] + if vector_results and not keyword_results: + return [ + result + for result in vector_results + if result.score >= min_score + ][:max_results] # Merge results merged = self._merge_results( @@ -140,6 +185,9 @@ async def _vector_search( ) -> List[MemorySearchResult]: """Execute vector similarity search""" try: + if self.provider_id is None: + raise RuntimeError("Embedding provider is not configured") + # Generate query embedding query_embedding = await Provider.embed( text=query, @@ -230,10 +278,7 @@ def _merge_results( kw_range = kw_max - kw_min normalised_keyword: List[tuple[MemorySearchResult, float]] = [] for r in keyword_results: - # When all scores are identical (including single-result case), - # use 0.5 as a neutral midpoint instead of 1.0 to avoid - # inflating keyword importance in the weighted combination. - norm = (r.score - kw_min) / kw_range if kw_range > 0 else 0.5 + norm = (r.score - kw_min) / kw_range if kw_range > 0 else 1.0 normalised_keyword.append((r, norm)) else: normalised_keyword = [] @@ -301,7 +346,7 @@ def decorate_citations( decorated = [] for result in results: - citation = format_citation(result) + citation = result.citation or format_citation(result) decorated.append(result.model_copy(update={"citation": citation})) return decorated diff --git a/flocks/memory/sync/indexer.py b/flocks/memory/sync/indexer.py index 614ba8176..c784436a7 100644 --- a/flocks/memory/sync/indexer.py +++ b/flocks/memory/sync/indexer.py @@ -6,16 +6,21 @@ from typing import List, Optional, Callable, Dict, Any from pathlib import Path -from datetime import datetime import asyncio +import math import uuid from flocks.provider import Provider -from flocks.storage import Storage, insert_chunks, get_embedding_from_cache, put_embedding_to_cache +from flocks.storage import ( + Storage, + get_embedding_from_cache, + put_embedding_to_cache, + replace_memory_file_index, +) from flocks.memory.types import MemoryFileEntry, MemoryChunk, MemorySyncProgress from flocks.memory.config import MemoryConfig -from flocks.memory.utils.hash import compute_hash, compute_text_hash -from flocks.memory.utils.text import is_memory_path +from flocks.memory.paths import classify_memory_path +from flocks.memory.utils.hash import compute_text_hash from flocks.memory.sync.chunking import TextChunker from flocks.utils.log import Log @@ -29,7 +34,7 @@ def __init__( self, project_id: str, workspace_dir: Path, - provider_id: str, + provider_id: Optional[str], embedding_model: str, config: MemoryConfig, ): @@ -78,7 +83,11 @@ async def sync( try: content_cache: Dict[str, str] = {} - memory_files = await self._scan_memory_files(_content_cache=content_cache) + indexed_files = await self._get_indexed_files() + memory_files = await self._scan_memory_files( + _content_cache=content_cache, + _indexed_files=None if force else indexed_files, + ) stats["files_scanned"] = len(memory_files) if progress_callback: @@ -88,12 +97,18 @@ async def sync( label="Scanning files" )) - indexed_files = await self._get_indexed_files() - for idx, file_entry in enumerate(memory_files): if not force: - indexed = indexed_files.get(file_entry.path) + indexed = indexed_files.get( + ( + file_entry.scope.value, + file_entry.scope_id, + file_entry.path, + ) + ) if indexed and indexed["hash"] == file_entry.hash: + if not self._metadata_matches(file_entry, indexed): + await self._update_file_metadata(file_entry) stats["files_skipped"] += 1 log.debug("indexer.file.skipped", {"path": file_entry.path}) content_cache.pop(file_entry.abs_path, None) @@ -113,7 +128,10 @@ async def sync( )) deleted_count = await self._clean_deleted_files( - current_files=[f.path for f in memory_files] + current_files=[ + (file.scope.value, file.scope_id, file.path) + for file in memory_files + ] ) if deleted_count > 0: log.info("indexer.cleaned", {"deleted": deleted_count}) @@ -126,15 +144,21 @@ async def sync( raise async def _scan_memory_files( - self, *, _content_cache: Optional[Dict[str, str]] = None, + self, + *, + _content_cache: Optional[Dict[str, str]] = None, + _indexed_files: Optional[ + Dict[tuple[str, str, str], Dict[str, Any]] + ] = None, ) -> List[MemoryFileEntry]: """ Scan workspace for memory files. - Filesystem I/O (glob, stat, read) is offloaded to a thread to avoid - blocking the event loop. When *_content_cache* is passed, file - contents read during hash calculation are stored there for later - reuse in ``_index_file``. + Filesystem I/O is offloaded to a thread to avoid blocking the event + loop. When *_indexed_files* is passed, files whose mtime and size still + match the stored manifest reuse the stored content hash without being + read. New or changed files are read once, hashed, and cached for + ``_index_file``. """ from flocks.config import Config @@ -153,8 +177,24 @@ def _add(fp: Path) -> None: resolved = str(fp.resolve()) if resolved in seen: return + classified = classify_memory_path(memory_root, fp) + if classified is None: + return seen.add(resolved) - files.append(self._create_file_entry(fp, memory_root, _content_cache=_content_cache)) + scope, scope_id, rel_path = classified + indexed_file = ( + _indexed_files.get((scope.value, scope_id, rel_path)) + if _indexed_files is not None + else None + ) + files.append( + self._create_file_entry( + fp, + memory_root, + _content_cache=_content_cache, + _indexed_file=indexed_file, + ) + ) for fp in memory_root.glob("**/*.md"): if fp.is_file(): @@ -180,58 +220,112 @@ def _add(fp: Path) -> None: return files def _create_file_entry( - self, file_path: Path, memory_root: Path, *, _content_cache: Optional[Dict[str, str]] = None, + self, + file_path: Path, + memory_root: Path, + *, + _content_cache: Optional[Dict[str, str]] = None, + _indexed_file: Optional[Dict[str, Any]] = None, ) -> MemoryFileEntry: """ Create file entry from path. - When *_content_cache* is provided the raw text is stored there keyed - by absolute path so that ``_index_file`` can reuse it without a second - disk read (fixes the TOCTOU + double-I/O issue). + An unchanged indexed file reuses its stored hash based on mtime and + size. Otherwise the raw text is read and optionally cached by absolute + path so that ``_index_file`` can reuse it without a second disk read. """ - try: - rel_path = str(file_path.relative_to(memory_root)) - except ValueError: - rel_path = file_path.name + classified = classify_memory_path(memory_root, file_path) + if classified is None: + raise ValueError(f"Unsupported Memory index path: {file_path}") + scope, scope_id, rel_path = classified stat = file_path.stat() - - content = file_path.read_text(encoding="utf-8") - content_hash = compute_text_hash(content) - - if _content_cache is not None: - _content_cache[str(file_path)] = content + if ( + _indexed_file is not None + and math.isclose( + _indexed_file["mtime"], + stat.st_mtime, + rel_tol=0, + abs_tol=1e-6, + ) + and _indexed_file["size"] == stat.st_size + ): + content_hash = str(_indexed_file["hash"]) + else: + content = file_path.read_text(encoding="utf-8") + content_hash = compute_text_hash(content) + if _content_cache is not None: + _content_cache[str(file_path)] = content return MemoryFileEntry( + scope=scope, + scope_id=scope_id, path=rel_path, abs_path=str(file_path), mtime_ms=stat.st_mtime * 1000, size=stat.st_size, hash=content_hash, ) + + @staticmethod + def _metadata_matches( + file_entry: MemoryFileEntry, + indexed_file: Dict[str, Any], + ) -> bool: + """Return whether current file metadata matches the stored manifest.""" + return ( + math.isclose( + indexed_file["mtime"], + file_entry.mtime_ms / 1000, + rel_tol=0, + abs_tol=1e-6, + ) + and indexed_file["size"] == file_entry.size + ) + + async def _update_file_metadata(self, file_entry: MemoryFileEntry) -> None: + """Refresh metadata after a content-preserving filesystem change.""" + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + UPDATE memory_files + SET mtime = ?, size = ? + WHERE scope = ? AND scope_id = ? AND path = ? + AND source = 'memory' AND hash = ? + """, + ( + file_entry.mtime_ms / 1000, + file_entry.size, + file_entry.scope.value, + file_entry.scope_id, + file_entry.path, + file_entry.hash, + ), + ) + await db.commit() - async def _get_indexed_files(self) -> Dict[str, Dict[str, Any]]: + async def _get_indexed_files( + self, + ) -> Dict[tuple[str, str, str], Dict[str, Any]]: """ Get indexed files from database Returns: Dict mapping path to file info """ - import aiosqlite - - indexed = {} + indexed: Dict[tuple[str, str, str], Dict[str, Any]] = {} try: async with Storage.connect(Storage.get_db_path()) as db: cursor = await db.execute(""" - SELECT path, hash, mtime, size + SELECT scope, scope_id, path, hash, mtime, size FROM memory_files - WHERE project_id = ? - """, (self.project_id,)) + WHERE source = 'memory' + """) rows = await cursor.fetchall() - for path, hash_val, mtime, size in rows: - indexed[path] = { + for scope, scope_id, path, hash_val, mtime, size in rows: + indexed[(scope, scope_id, path)] = { "hash": hash_val, "mtime": mtime, "size": size, @@ -263,29 +357,48 @@ async def _index_file( # Chunk text chunks = self.chunker.chunk_text(content, file_entry.path) stats["chunks"] = len(chunks) - + + # FTS indexing is always available. Embeddings are optional. + chunk_records: List[Dict[str, Any]] = [] if not chunks: - log.warn("indexer.no_chunks", {"path": file_entry.path}) - return stats - - # Generate embeddings for chunks - chunk_records = [] - - # Use batch processing if enabled - if self.config.batch.enabled and len(chunks) > 1: - chunk_records = await self._generate_embeddings_batch( - chunks, file_entry, stats - ) + log.debug("indexer.no_chunks", {"path": file_entry.path}) + elif self.provider_id is None: + chunk_records = [ + self._create_chunk_record(chunk, file_entry, None, None) + for chunk in chunks + ] else: - chunk_records = await self._generate_embeddings_sequential( - chunks, file_entry, stats - ) + try: + if self.config.batch.enabled and len(chunks) > 1: + chunk_records = await self._generate_embeddings_batch( + chunks, file_entry, stats + ) + else: + chunk_records = await self._generate_embeddings_sequential( + chunks, file_entry, stats + ) + except Exception as exc: + log.warn( + "indexer.embedding.failed_fts_fallback", + {"path": file_entry.path, "error": str(exc)}, + ) + chunk_records = [ + self._create_chunk_record(chunk, file_entry, None, None) + for chunk in chunks + ] - await self._delete_file_chunks(file_entry.path) - await insert_chunks(Storage.get_db_path(), chunk_records) - - # Update file entry in database - await self._update_file_entry(file_entry) + await replace_memory_file_index( + Storage.get_db_path(), + file_entry={ + "scope": file_entry.scope.value, + "scope_id": file_entry.scope_id, + "path": file_entry.path, + "hash": file_entry.hash, + "mtime": file_entry.mtime_ms / 1000, + "size": file_entry.size, + }, + chunks=chunk_records, + ) log.info("indexer.file.indexed", { "path": file_entry.path, @@ -392,14 +505,15 @@ def _create_chunk_record( self, chunk: MemoryChunk, file_entry: MemoryFileEntry, - embedding: List[float], - dims: int, + embedding: Optional[List[float]], + dims: Optional[int], ) -> Dict[str, Any]: """Create chunk record for database""" return { "id": str(uuid.uuid4()), + "scope": file_entry.scope.value, + "scope_id": file_entry.scope_id, "path": file_entry.path, - "project_id": self.project_id, "source": "memory", "start_line": chunk.start_line, "end_line": chunk.end_line, @@ -415,7 +529,7 @@ async def _get_cached_embedding( text_hash: str, ) -> Optional[tuple[List[float], int]]: """Get embedding from cache""" - if not self.config.cache.enabled: + if not self.config.cache.enabled or self.provider_id is None: return None return await get_embedding_from_cache( @@ -432,7 +546,7 @@ async def _put_cached_embedding( dims: int, ) -> None: """Put embedding to cache""" - if not self.config.cache.enabled: + if not self.config.cache.enabled or self.provider_id is None: return await put_embedding_to_cache( @@ -444,46 +558,10 @@ async def _put_cached_embedding( dims=dims, ) - async def _delete_file_chunks(self, path: str) -> None: - """Delete all existing chunks for a file before re-indexing.""" - import aiosqlite - - try: - async with Storage.connect(Storage.get_db_path()) as db: - await db.execute( - "DELETE FROM memory_chunks WHERE project_id = ? AND path = ?", - (self.project_id, path), - ) - await db.commit() - except Exception as e: - log.error("indexer.delete_chunks.failed", {"path": path, "error": str(e)}) - - async def _update_file_entry(self, file_entry: MemoryFileEntry) -> None: - """Update file entry in database""" - import aiosqlite - - now = datetime.now().timestamp() - - try: - async with Storage.connect(Storage.get_db_path()) as db: - await db.execute(""" - INSERT OR REPLACE INTO memory_files - (path, project_id, source, hash, mtime, size, indexed_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, ( - file_entry.path, - self.project_id, - "memory", - file_entry.hash, - file_entry.mtime_ms / 1000, - file_entry.size, - now, - )) - await db.commit() - except Exception as e: - log.error("indexer.update_file.failed", {"path": file_entry.path, "error": str(e)}) - - async def _clean_deleted_files(self, current_files: List[str]) -> int: + async def _clean_deleted_files( + self, + current_files: List[tuple[str, str, str]], + ) -> int: """ Clean up deleted files from database @@ -493,15 +571,14 @@ async def _clean_deleted_files(self, current_files: List[str]) -> int: Returns: Number of deleted files """ - import aiosqlite - try: async with Storage.connect(Storage.get_db_path()) as db: cursor = await db.execute(""" - SELECT path FROM memory_files WHERE project_id = ? - """, (self.project_id,)) + SELECT scope, scope_id, path FROM memory_files + WHERE source = 'memory' + """) - indexed_paths = [row[0] for row in await cursor.fetchall()] + indexed_paths = [tuple(row) for row in await cursor.fetchall()] # Find deleted files deleted = [p for p in indexed_paths if p not in current_files] @@ -509,18 +586,32 @@ async def _clean_deleted_files(self, current_files: List[str]) -> int: if not deleted: return 0 - placeholders = ",".join("?" * len(deleted)) - params = (self.project_id, *deleted) - - await db.execute(f""" - DELETE FROM memory_chunks - WHERE project_id = ? AND path IN ({placeholders}) - """, params) - - await db.execute(f""" - DELETE FROM memory_files - WHERE project_id = ? AND path IN ({placeholders}) - """, params) + for scope, scope_id, path in deleted: + params = (scope, scope_id, path) + await db.execute( + """ + DELETE FROM memory_fts + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) + await db.execute( + """ + DELETE FROM memory_chunks + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) + await db.execute( + """ + DELETE FROM memory_files + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + params, + ) await db.commit() diff --git a/flocks/memory/types.py b/flocks/memory/types.py index 4e7211887..8d19387c1 100644 --- a/flocks/memory/types.py +++ b/flocks/memory/types.py @@ -52,7 +52,6 @@ class MemoryProviderStatus(BaseModel): # Statistics files: int = Field(0, description="Number of indexed files") chunks: int = Field(0, description="Number of indexed chunks") - dirty: bool = Field(False, description="Whether sync is needed") # Configuration workspace_dir: Optional[str] = Field(None, description="Workspace directory") @@ -68,11 +67,8 @@ class MemoryProviderStatus(BaseModel): class MemoryFileEntry(BaseModel): """File entry for indexing""" - scope: MemoryScope = Field( - MemoryScope.GLOBAL, - description="Memory visibility scope", - ) - scope_id: str = Field("global", description="Scope identifier") + scope: MemoryScope = Field(..., description="Memory visibility scope") + scope_id: str = Field(..., description="Scope identifier") path: str = Field(..., description="Relative path") abs_path: str = Field(..., description="Absolute path") mtime_ms: float = Field(..., description="Modification time (milliseconds)") diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index b20db5bbf..cc415b482 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -65,7 +65,10 @@ async def initialize(self) -> bool: try: config = await Config.get() if getattr(config, "memory", None) is None: - log.info("session.memory.no_config", {"session_id": self.session_id}) + log.info( + "session.memory.no_config", + {"session_id": self.session_id}, + ) memory_config = resolve_memory_config(config) self._manager = MemoryManager.get_instance( @@ -139,7 +142,7 @@ async def search( "session_id": self.session_id, "error": str(e), }) - return [] + raise async def write( self, diff --git a/flocks/session/message.py b/flocks/session/message.py index 0ca768a1a..9f3e3a627 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -484,11 +484,12 @@ async def quiesce_parts(cls, session_id: str, *, persist: bool) -> None: async with _session_locks.get(session_id): if persist and session_id in cls._parts_cache: - if cls._parts_storage_format.get(session_id) == "legacy": - await cls._persist_parts(session_id) - else: - for message_id in list(cls._parts_cache[session_id]): - await cls._persist_parts(session_id, message_id=message_id) + for message_id in list(cls._parts_cache[session_id]): + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) @classmethod def _cache_token(cls, session_id: str) -> tuple[int, int]: @@ -647,7 +648,10 @@ async def _flush_later() -> None: try: await asyncio.sleep(cls._PARTS_PERSIST_DEBOUNCE_MS / 1000) async with _session_locks.get(session_id): - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_parts( + session_id, + message_id=message_id, + ) except asyncio.CancelledError: pass except Exception as exc: @@ -1369,6 +1373,117 @@ async def _persist_parts(cls, session_id: str, *, message_id: Optional[str] = No await Storage.delete(cls._parts_item_key(session_id, stale_mid)) serialized.pop(stale_mid, None) persisted_mids.discard(stale_mid) + + @classmethod + async def _persist_indexed_state( + cls, + session_id: str, + message_id: str, + *, + include_messages: bool = False, + include_parts: bool = False, + delete_message: bool = False, + ) -> None: + """Persist canonical message data and its derived FTS row atomically.""" + set_entries = [] + delete_keys = [] + + if include_messages: + messages = cls._messages_cache.get(session_id, []) + serialized_messages = [] + for index, message in enumerate(messages): + normalized = cls._normalize_assistant_message(message) + if normalized is not message: + messages[index] = normalized + serialized_messages.append(normalized.model_dump()) + set_entries.append( + (f"{cls._MESSAGE_PREFIX}:{session_id}", serialized_messages, "json") + ) + + storage_format = cls._parts_storage_format.setdefault( + session_id, + "per_message", + ) + if include_parts: + all_parts = cls._parts_cache.get(session_id, {}) + serialized = cls._parts_serialized_cache.setdefault(session_id, {}) + if storage_format == "legacy": + serialized = { + mid: cls._serialize_message_parts(message_parts) + for mid, message_parts in all_parts.items() + } + cls._parts_serialized_cache[session_id] = serialized + set_entries.append( + ( + cls._parts_blob_key(session_id), + serialized, + "message_parts", + ) + ) + elif delete_message: + delete_keys.append(cls._parts_item_key(session_id, message_id)) + else: + serialized_one = cls._serialize_message_parts( + all_parts.get(message_id, []) + ) + serialized[message_id] = serialized_one + set_entries.append( + ( + cls._parts_item_key(session_id, message_id), + serialized_one, + "message_part", + ) + ) + + from flocks.session.session import Session + from flocks.storage.session_search import ( + delete_message_document, + upsert_session_document, + ) + + session = await Session.get_by_id_unfiltered(session_id) + message = next( + ( + item + for item in cls._messages_cache.get(session_id, []) + if item.id == message_id + ), + None, + ) + parts = list( + cls._parts_cache.get(session_id, {}).get(message_id, []) + ) + + async def _sync_search_index(db) -> None: + if not Storage.session_search_available(): + return + if delete_message or message is None: + await delete_message_document(db, message_id) + return + if session is None: + return + await upsert_session_document( + db, + project_id=session.project_id, + message=message, + parts=parts, + ) + + await Storage.mutate_many( + set_entries=set_entries, + delete_keys=delete_keys, + transaction_hook=_sync_search_index, + ) + + if include_parts: + persisted_mids = cls._parts_persisted_mids.setdefault( + session_id, + set(), + ) + if delete_message: + persisted_mids.discard(message_id) + else: + persisted_mids.add(message_id) @classmethod async def create( @@ -1473,8 +1588,12 @@ async def create( cls._parts_cache[session_id][message.id].append(part) # Persist to storage - await cls._persist_messages(session_id) - await cls._persist_parts(session_id) + await cls._persist_indexed_state( + session_id, + message.id, + include_messages=True, + include_parts=True, + ) log.info("message.created", { "id": message.id, @@ -1658,7 +1777,11 @@ async def store_part(cls, session_id: str, message_id: str, part: PartType) -> P cls._schedule_parts_flush(session_id, message_id=message_id) else: cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.debug("message.part.stored" if not updated else "message.part.updated", { "session_id": session_id, @@ -1711,7 +1834,11 @@ async def upsert_message_info(cls, session_id: str, message_info: MessageInfo) - messages.append(message_info) cls._rebuild_id_index(session_id) - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_info.id, + include_messages=True, + ) log.debug("message.upserted", { "id": message_info.id, @@ -1916,7 +2043,13 @@ async def delete(cls, session_id: str, message_id: str) -> bool: had_pending_parts_flush = session_id in cls._parts_flush_tasks cls._cancel_parts_flush_task(session_id) try: - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_messages=True, + include_parts=True, + delete_message=True, + ) except BaseException: # Message metadata is the deletion commit point. Restore # every in-memory index/cache if it was not persisted so a @@ -1942,23 +2075,6 @@ async def delete(cls, session_id: str, message_id: str) -> bool: ) raise - try: - if cls._parts_storage_format.get(session_id) == "legacy": - await cls._persist_parts(session_id) - else: - await Storage.delete(cls._parts_item_key(session_id, message_id)) - cls._parts_persisted_mids.setdefault(session_id, set()).discard( - message_id - ) - except Exception as exc: - # Metadata deletion has committed. Orphaned parts are not - # user-visible and can be cleaned later; restoring the - # message here would make cache and durable metadata diverge. - log.warn("message.delete.parts_cleanup_failed", { - "session_id": session_id, - "message_id": message_id, - "error": str(exc), - }) log.info("message.deleted", {"id": message_id, "session_id": session_id}) return True return False @@ -2019,9 +2135,21 @@ async def clear(cls, session_id: str) -> int: cls._parts_fully_loaded.add(session_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_messages(session_id) - await Storage.clear(prefix=cls._parts_item_prefix(session_id)) - await Storage.delete(cls._parts_blob_key(session_id)) + from flocks.storage.session_search import delete_session_documents + + async def _delete_search_index(db) -> None: + if not Storage.session_search_available(): + return + await delete_session_documents(db, [session_id]) + + await Storage.mutate_many( + set_entries=[ + (f"{cls._MESSAGE_PREFIX}:{session_id}", [], "json"), + ], + delete_keys=[cls._parts_blob_key(session_id)], + delete_prefixes=[cls._parts_item_prefix(session_id)], + transaction_hook=_delete_search_index, + ) log.info("messages.cleared", { "session_id": session_id, @@ -2352,7 +2480,11 @@ async def update(cls, session_id: str, message_id: str, **updates) -> Optional[M updated = message.model_copy(update=patch) messages[msg_index] = updated - await cls._persist_messages(session_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_messages=True, + ) log.info("message.updated", { "id": message_id, @@ -2394,7 +2526,11 @@ async def add_part(cls, session_id: str, message_id: str, part: PartType) -> Opt cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_added", { "message_id": message_id, @@ -2437,7 +2573,11 @@ async def update_part(cls, session_id: str, message_id: str, part_id: str, **upd cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_updated", { "message_id": message_id, @@ -2472,7 +2612,11 @@ async def remove_part(cls, session_id: str, message_id: str, part_id: str) -> bo cls._touch_parts_revision(session_id, message_id) cls._cancel_parts_flush_task(session_id) - await cls._persist_parts(session_id, message_id=message_id) + await cls._persist_indexed_state( + session_id, + message_id, + include_parts=True, + ) log.info("message.part_removed", { "message_id": message_id, diff --git a/flocks/session/session.py b/flocks/session/session.py index 14a56e740..6a904412a 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -769,8 +769,15 @@ async def _delete_locked(cls, project_id: str, session_id: str) -> bool: session_ids = [session.id for session in sessions] from flocks.permission.next import PermissionNext + from flocks.storage.session_search import delete_session_documents permission_keys = await PermissionNext.deletion_storage_keys(session_ids) + + async def _delete_search_index(db) -> None: + if not Storage.session_search_available(): + return + await delete_session_documents(db, session_ids) + await Storage.mutate_many( delete_keys=[ key @@ -794,6 +801,7 @@ async def _delete_locked(cls, project_id: str, session_id: str) -> bool: f"system_prompts:{session.id}:", ) ], + transaction_hook=_delete_search_index, ) PermissionNext.clear_session_runtime(session_ids) diff --git a/flocks/storage/__init__.py b/flocks/storage/__init__.py index 91b269afa..865756804 100644 --- a/flocks/storage/__init__.py +++ b/flocks/storage/__init__.py @@ -8,6 +8,7 @@ vector_search, fts_search, insert_chunks, + replace_memory_file_index, get_embedding_from_cache, put_embedding_to_cache, cosine_similarity, @@ -22,6 +23,7 @@ "vector_search", "fts_search", "insert_chunks", + "replace_memory_file_index", "get_embedding_from_cache", "put_embedding_to_cache", "cosine_similarity", diff --git a/flocks/storage/session_search.py b/flocks/storage/session_search.py new file mode 100644 index 000000000..6668f9841 --- /dev/null +++ b/flocks/storage/session_search.py @@ -0,0 +1,671 @@ +"""Derived FTS5 index for persisted session transcripts.""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime +import hashlib +from pathlib import Path +import sqlite3 +from typing import Any, Iterable, Optional, Sequence + +import aiosqlite + +from flocks.storage.storage import Storage +from flocks.utils.log import Log + +log = Log.create(service="storage.session_search") + +_SESSION_BACKFILL_KEY = "history-v1" +_reconcile_locks: dict[str, asyncio.Lock] = {} + +_SESSION_SEARCH_UNAVAILABLE_MESSAGE = ( + "Session search is unavailable because this SQLite runtime does not " + "support FTS5. Session messages will continue to be stored normally." +) + + +class SessionSearchUnavailableError(RuntimeError): + """Raised when the active SQLite runtime cannot provide Session FTS.""" + + +def _is_fts5_unavailable_error(error: BaseException) -> bool: + """Return whether an SQLite failure specifically means FTS5 is absent.""" + current: Optional[BaseException] = error + while current is not None: + if isinstance(current, sqlite3.OperationalError): + message = str(current).casefold() + if "fts5" in message and ( + "no such module" in message or "unknown module" in message + ): + return True + current = current.__cause__ or current.__context__ + return False + + +def require_session_search_available() -> None: + """Raise a stable, user-facing error when Session FTS is disabled.""" + if not Storage.session_search_available(): + raise SessionSearchUnavailableError(_SESSION_SEARCH_UNAVAILABLE_MESSAGE) + + +SESSION_SEARCH_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS session_transcript_index_state ( + id INTEGER PRIMARY KEY, + message_id TEXT NOT NULL UNIQUE, + session_id TEXT NOT NULL, + project_id TEXT, + role TEXT NOT NULL, + created_at INTEGER NOT NULL, + source_updated_at INTEGER NOT NULL, + content_hash TEXT NOT NULL, + indexed_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_session_transcript_state_session + ON session_transcript_index_state(session_id); + +CREATE INDEX IF NOT EXISTS idx_session_transcript_state_project + ON session_transcript_index_state(project_id); + +CREATE VIRTUAL TABLE IF NOT EXISTS session_transcript_fts USING fts5( + text, + tokenize = 'unicode61 remove_diacritics 2' +); +""" + + +SESSION_SEARCH_META_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS session_transcript_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at INTEGER NOT NULL +); +""" + + +def _value(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, dict): + return value.get(name, default) + return getattr(value, name, default) + + +def _message_timestamp(message: Any, name: str) -> int: + time_value = _value(message, "time", {}) + raw = _value(time_value, name, 0) + try: + return int(raw or 0) + except (TypeError, ValueError): + return 0 + + +def build_session_document( + message: Any, + parts: Iterable[Any], +) -> Optional[dict[str, Any]]: + """Build a searchable user/assistant document from authoritative parts.""" + role_value = _value(message, "role", "") + role = getattr(role_value, "value", role_value) + if role not in {"user", "assistant"}: + return None + + text_parts: list[str] = [] + for part in parts: + if _value(part, "type") != "text": + continue + if bool(_value(part, "synthetic", False)) or bool( + _value(part, "ignored", False) + ): + continue + text = str(_value(part, "text", "") or "").strip() + if text: + text_parts.append(text) + + text = "\n".join(text_parts).strip() + if not text: + return None + + created_at = _message_timestamp(message, "created") + updated_at = ( + _message_timestamp(message, "updated") + or _message_timestamp(message, "completed") + or created_at + ) + return { + "message_id": str(_value(message, "id")), + "session_id": str(_value(message, "sessionID")), + "role": role, + "created_at": created_at, + "source_updated_at": updated_at, + "text": text, + "content_hash": hashlib.sha256(text.encode("utf-8")).hexdigest(), + } + + +async def ensure_session_search_tables(db_path: Path) -> bool: + """Create Session search tables if the SQLite runtime supports FTS5. + + Returns: + ``True`` when Session FTS is available and the schema is ready. + ``False`` only when SQLite explicitly reports that the FTS5 module is + unavailable. All other database errors are propagated. + """ + async with Storage.connect(db_path) as db: + await db.executescript(SESSION_SEARCH_META_SCHEMA_SQL) + try: + await db.execute( + """ + CREATE VIRTUAL TABLE temp._flocks_session_fts5_probe + USING fts5(text) + """ + ) + await db.execute( + "DROP TABLE temp._flocks_session_fts5_probe" + ) + except sqlite3.OperationalError as exc: + if _is_fts5_unavailable_error(exc): + # Messages may be created, updated, or deleted while Session + # indexing is disabled. Force a complete reconciliation if a + # future runtime restores FTS5 support. + await db.execute( + "DELETE FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + await db.commit() + return False + raise + + cursor = await db.execute( + """ + SELECT name + FROM sqlite_master + WHERE name IN ( + 'session_transcript_index_state', + 'session_transcript_fts' + ) + """ + ) + existing_tables = {row[0] for row in await cursor.fetchall()} + await db.executescript(SESSION_SEARCH_SCHEMA_SQL) + if existing_tables != { + "session_transcript_index_state", + "session_transcript_fts", + }: + await db.execute( + "DELETE FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + await db.commit() + return True + + +def _reconcile_lock(db_path: Path) -> asyncio.Lock: + """Return the process-local reconcile owner for one SQLite database.""" + key = str(db_path.resolve()) + lock = _reconcile_locks.get(key) + if lock is None: + lock = asyncio.Lock() + _reconcile_locks[key] = lock + return lock + + +async def _session_index_is_ready(db_path: Path) -> bool: + async with Storage.connect(db_path) as db: + cursor = await db.execute( + "SELECT 1 FROM session_transcript_meta WHERE key = ?", + (_SESSION_BACKFILL_KEY,), + ) + return await cursor.fetchone() is not None + + +async def _mark_session_index_ready(db_path: Path) -> None: + now = int(datetime.now(UTC).timestamp() * 1000) + async with Storage.connect(db_path) as db: + await db.execute( + """ + INSERT INTO session_transcript_meta (key, value, updated_at) + VALUES (?, 'complete', ?) + ON CONFLICT(key) DO UPDATE SET + value = excluded.value, + updated_at = excluded.updated_at + """, + (_SESSION_BACKFILL_KEY, now), + ) + await db.commit() + + +async def ensure_session_index_ready(*, batch_size: int = 50) -> bool: + """Backfill legacy transcripts once, then use realtime message indexing. + + Returns ``True`` when this call performed the historical backfill and + ``False`` when a previous successful pass already made the index ready. + """ + require_session_search_available() + db_path = Storage.get_db_path() + if await _session_index_is_ready(db_path): + return False + + async with _reconcile_lock(db_path): + if await _session_index_is_ready(db_path): + return False + stats = await _reconcile_session_index_unlocked(batch_size=batch_size) + await _mark_session_index_ready(db_path) + log.info("session_search.backfill.complete", stats) + return True + + +async def upsert_session_document( + db: aiosqlite.Connection, + *, + project_id: str, + message: Any, + parts: Sequence[Any], +) -> bool: + """Synchronize one message into the transcript FTS index.""" + message_id = str(_value(message, "id")) + document = build_session_document(message, parts) + cursor = await db.execute( + """ + SELECT id, project_id, role, created_at, source_updated_at, content_hash + FROM session_transcript_index_state + WHERE message_id = ? + """, + (message_id,), + ) + existing = await cursor.fetchone() + + if document is None: + if existing is None: + return False + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (existing[0],), + ) + await db.execute( + "DELETE FROM session_transcript_index_state WHERE id = ?", + (existing[0],), + ) + return True + + unchanged = existing is not None and ( + existing[1], + existing[2], + existing[3], + existing[4], + existing[5], + ) == ( + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + ) + if unchanged: + cursor = await db.execute( + "SELECT text FROM session_transcript_fts WHERE rowid = ?", + (existing[0],), + ) + indexed = await cursor.fetchone() + if indexed is not None and indexed[0] == document["text"]: + return False + + indexed_at = int(datetime.now(UTC).timestamp() * 1000) + if existing is None: + cursor = await db.execute( + """ + INSERT INTO session_transcript_index_state ( + message_id, session_id, project_id, role, created_at, + source_updated_at, content_hash, indexed_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + document["message_id"], + document["session_id"], + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + indexed_at, + ), + ) + rowid = cursor.lastrowid + else: + rowid = existing[0] + await db.execute( + """ + UPDATE session_transcript_index_state + SET session_id = ?, project_id = ?, role = ?, created_at = ?, + source_updated_at = ?, content_hash = ?, indexed_at = ? + WHERE id = ? + """, + ( + document["session_id"], + project_id, + document["role"], + document["created_at"], + document["source_updated_at"], + document["content_hash"], + indexed_at, + rowid, + ), + ) + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (rowid,), + ) + + await db.execute( + "INSERT INTO session_transcript_fts(rowid, text) VALUES (?, ?)", + (rowid, document["text"]), + ) + return True + + +async def delete_session_documents( + db: aiosqlite.Connection, + session_ids: Sequence[str], +) -> int: + """Delete all derived transcript rows for the supplied sessions.""" + if not session_ids: + return 0 + placeholders = ",".join("?" for _ in session_ids) + cursor = await db.execute( + f""" + SELECT id FROM session_transcript_index_state + WHERE session_id IN ({placeholders}) + """, + tuple(session_ids), + ) + rowids = [row[0] for row in await cursor.fetchall()] + if rowids: + rowid_placeholders = ",".join("?" for _ in rowids) + await db.execute( + f"DELETE FROM session_transcript_fts WHERE rowid IN ({rowid_placeholders})", + tuple(rowids), + ) + cursor = await db.execute( + f""" + DELETE FROM session_transcript_index_state + WHERE session_id IN ({placeholders}) + """, + tuple(session_ids), + ) + return max(cursor.rowcount, 0) + + +async def delete_message_document( + db: aiosqlite.Connection, + message_id: str, +) -> bool: + """Delete one derived transcript row.""" + cursor = await db.execute( + "SELECT id FROM session_transcript_index_state WHERE message_id = ?", + (message_id,), + ) + row = await cursor.fetchone() + if row is None: + return False + await db.execute( + "DELETE FROM session_transcript_fts WHERE rowid = ?", + (row[0],), + ) + await db.execute( + "DELETE FROM session_transcript_index_state WHERE id = ?", + (row[0],), + ) + return True + + +async def reconcile_session_index( + *, + project_id: Optional[str] = None, + batch_size: int = 50, +) -> dict[str, int]: + """Rebuild missing/stale rows and remove orphaned derived rows.""" + require_session_search_available() + db_path = Storage.get_db_path() + async with _reconcile_lock(db_path): + stats = await _reconcile_session_index_unlocked( + project_id=project_id, + batch_size=batch_size, + ) + if project_id is None: + await _mark_session_index_ready(db_path) + return stats + + +async def _reconcile_session_index_unlocked( + *, + project_id: Optional[str] = None, + batch_size: int = 50, +) -> dict[str, int]: + """Repair Session FTS while bounding loaded TextParts to one batch.""" + from flocks.session.message import Message + from flocks.session.session import Session + + sessions = [ + session + for session in await Session.list_all_unfiltered() + if session.status != "deleted" + and (project_id is None or session.project_id == project_id) + ] + stats = {"scanned": 0, "updated": 0, "deleted": 0} + effective_batch_size = max(1, batch_size) + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_seen ( + message_id TEXT PRIMARY KEY + ) + """ + ) + if project_id is None: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_candidates AS + SELECT id, message_id + FROM session_transcript_index_state + """ + ) + else: + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_candidates AS + SELECT id, message_id + FROM session_transcript_index_state + WHERE project_id = ? + """, + (project_id,), + ) + await db.execute( + """ + CREATE TEMP TABLE session_transcript_reconcile_fts_candidates AS + SELECT rowid + FROM session_transcript_fts + """ + ) + + for session in sessions: + # Session deletion owns this same lifecycle lock. Holding it while + # rebuilding prevents a deleted transcript from being reinserted + # after its transactional FTS cleanup. + async with Session.lifecycle_lock(session.id): + current = await Session.get_by_id_unfiltered(session.id) + if ( + current is None + or current.status == "deleted" + or Session.is_lifecycle_transitioning(session.id) + ): + continue + + messages = await Message.list( + session.id, + include_archived=True, + ) + for offset in range(0, len(messages), effective_batch_size): + batch_messages = messages[ + offset : offset + effective_batch_size + ] + batch = [] + for message in batch_messages: + item = await Message.get_with_parts_lazy( + session.id, + message.id, + ) + if item is not None: + batch.append(item) + + stats["scanned"] += len(batch) + try: + await db.execute("BEGIN IMMEDIATE") + for item in batch: + document = build_session_document( + item.info, + item.parts, + ) + if document is not None: + await db.execute( + """ + INSERT OR IGNORE INTO + session_transcript_reconcile_seen ( + message_id + ) + VALUES (?) + """, + (document["message_id"],), + ) + if await upsert_session_document( + db, + project_id=session.project_id, + message=item.info, + parts=item.parts, + ): + stats["updated"] += 1 + await db.commit() + except BaseException: + await db.rollback() + raise + + try: + await db.execute("BEGIN IMMEDIATE") + cursor = await db.execute( + """ + SELECT count(*) + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + """ + ) + stale_count = int((await cursor.fetchone())[0]) + if stale_count: + await db.execute( + """ + DELETE FROM session_transcript_fts + WHERE rowid IN ( + SELECT candidate.id + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + ) + """ + ) + await db.execute( + """ + DELETE FROM session_transcript_index_state + WHERE id IN ( + SELECT candidate.id + FROM session_transcript_reconcile_candidates AS candidate + LEFT JOIN session_transcript_reconcile_seen AS seen + ON seen.message_id = candidate.message_id + WHERE seen.message_id IS NULL + ) + """ + ) + stats["deleted"] += stale_count + + cursor = await db.execute( + """ + DELETE FROM session_transcript_fts + WHERE rowid IN ( + SELECT candidate.rowid + FROM session_transcript_reconcile_fts_candidates AS candidate + WHERE candidate.rowid NOT IN ( + SELECT id FROM session_transcript_index_state + ) + ) + """ + ) + stats["deleted"] += max(cursor.rowcount, 0) + await db.commit() + except BaseException: + await db.rollback() + raise + + log.info( + "session_search.reconciled", + {"project_id": project_id or "*", **stats}, + ) + return stats + + +async def session_fts_search( + *, + db_path: Path, + project_id: str, + query: str, + max_results: int, +) -> list[dict[str, Any]]: + """Search all indexed session messages using FTS5 BM25 ranking.""" + from flocks.storage.vector import build_fts_query + + require_session_search_available() + del project_id # Retained for API compatibility; Session search is global. + fts_query = build_fts_query(query) + if not fts_query: + return [] + + async with Storage.connect(db_path) as db: + cursor = await db.execute( + """ + SELECT + s.message_id, + s.session_id, + s.role, + s.created_at, + snippet(session_transcript_fts, 0, '', '', ' … ', 24), + bm25(session_transcript_fts) + FROM session_transcript_fts + JOIN session_transcript_index_state s + ON s.id = session_transcript_fts.rowid + WHERE session_transcript_fts MATCH ? + ORDER BY bm25(session_transcript_fts) + LIMIT ? + """, + (fts_query, max_results), + ) + rows = await cursor.fetchall() + + count = len(rows) + results: list[dict[str, Any]] = [] + for index, row in enumerate(rows): + message_id, session_id, role, created_at, snippet, _rank = row + score = 1.0 if count == 1 else 1.0 - (index / (2 * count)) + results.append( + { + "path": f"sessions/{session_id}/messages/{message_id}", + "source": "session", + "start_line": 1, + "end_line": 1, + "text": snippet, + "score": score, + "citation": ( + f"session:{session_id} message:{message_id} " + f"role:{role} created_at:{created_at}" + ), + } + ) + return results diff --git a/flocks/storage/storage.py b/flocks/storage/storage.py index 43a8724cc..b1330e127 100644 --- a/flocks/storage/storage.py +++ b/flocks/storage/storage.py @@ -25,6 +25,7 @@ T = TypeVar("T", bound=BaseModel) DDLScript = str | Callable[[aiosqlite.Connection], Awaitable[None]] +TransactionHook = Callable[[aiosqlite.Connection], Awaitable[None]] R = TypeVar("R") @@ -72,6 +73,7 @@ class Storage: _log = Log.create(service="storage") _db_path: Optional[Path] = None _initialized = False + _session_search_available = True _db_identity: Optional[Tuple[int, int]] = None # PID of the process that called ``init()``. Used by ``_ensure_init`` to # detect ``fork()`` (uvicorn ``--reload`` / multiprocessing workers) and @@ -151,6 +153,11 @@ def get_db_path(cls) -> Path: data_dir = Config.get_data_path() return data_dir / "flocks.db" + @classmethod + def session_search_available(cls) -> bool: + """Return whether Session FTS is supported by the active SQLite runtime.""" + return cls._session_search_available + @staticmethod def _file_identity(db_path: Path) -> Optional[Tuple[int, int]]: """Return the filesystem identity used to detect an online DB replacement.""" @@ -1354,6 +1361,20 @@ async def _bootstrap_schema(cls) -> None: except Exception as e: cls._log.warn("storage.vector.init.failed", {"error": str(e)}) + from flocks.storage.session_search import ensure_session_search_tables + + cls._session_search_available = await ensure_session_search_tables( + cls._db_path + ) + if not cls._session_search_available: + cls._log.warn( + "storage.session_search.disabled", + { + "reason": "SQLite runtime does not support FTS5", + "db_path": str(cls._db_path), + }, + ) + # Create model management tables await cls._create_model_management_tables() @@ -1555,12 +1576,22 @@ async def mutate_many( set_entries: Sequence[Tuple[str, Any, str]] = (), delete_keys: Sequence[str] = (), delete_prefixes: Sequence[str] = (), + transaction_hook: Optional[TransactionHook] = None, ) -> int: - """Apply related set/delete operations in one SQLite transaction.""" + """Apply related set/delete operations in one SQLite transaction. + + ``transaction_hook`` is reserved for derived relational indexes that + must commit atomically with their canonical KV records. + """ entries = list(set_entries) keys_to_delete = list(delete_keys) prefixes_to_delete = list(delete_prefixes) - if not entries and not keys_to_delete and not prefixes_to_delete: + if ( + not entries + and not keys_to_delete + and not prefixes_to_delete + and transaction_hook is None + ): return 0 routing_paths = { @@ -1568,6 +1599,8 @@ async def mutate_many( *(cls.route_db_path_for_key(key) for key in keys_to_delete), *(cls.route_db_path_for_prefix(prefix) for prefix in prefixes_to_delete), } + if not routing_paths: + routing_paths = {cls.get_db_path()} if len(routing_paths) != 1: raise ValueError("Storage.mutate_many operations must target the same database") @@ -1614,6 +1647,8 @@ async def _write() -> int: (cls._like_prefix_pattern(prefix),), ) deleted += max(cursor.rowcount, 0) + if transaction_hook is not None: + await transaction_hook(db) await db.commit() return deleted except BaseException: diff --git a/flocks/storage/vector.py b/flocks/storage/vector.py index 18401f90a..7344812a4 100644 --- a/flocks/storage/vector.py +++ b/flocks/storage/vector.py @@ -7,7 +7,6 @@ from typing import List, Optional, Dict, Any, Tuple from pathlib import Path -import aiosqlite import json import math from datetime import datetime @@ -22,20 +21,23 @@ VECTOR_SCHEMA_SQL = """ -- Memory files index table CREATE TABLE IF NOT EXISTS memory_files ( - path TEXT PRIMARY KEY, - project_id TEXT NOT NULL, + scope TEXT NOT NULL, + scope_id TEXT NOT NULL, + path TEXT NOT NULL, source TEXT NOT NULL, -- 'memory' | 'session' hash TEXT NOT NULL, mtime REAL NOT NULL, size INTEGER NOT NULL, - indexed_at REAL NOT NULL + indexed_at REAL NOT NULL, + PRIMARY KEY (scope, scope_id, path) ); -- Memory chunks table CREATE TABLE IF NOT EXISTS memory_chunks ( id TEXT PRIMARY KEY, + scope TEXT NOT NULL, + scope_id TEXT NOT NULL, path TEXT NOT NULL, - project_id TEXT NOT NULL, source TEXT NOT NULL, start_line INTEGER NOT NULL, end_line INTEGER NOT NULL, @@ -46,7 +48,8 @@ embedding_dims INTEGER, created_at REAL NOT NULL, updated_at REAL NOT NULL, - FOREIGN KEY (path) REFERENCES memory_files(path) ON DELETE CASCADE + FOREIGN KEY (scope, scope_id, path) + REFERENCES memory_files(scope, scope_id, path) ON DELETE CASCADE ); -- Embedding cache table (shared across projects) @@ -62,10 +65,13 @@ ); -- Indexes for performance -CREATE INDEX IF NOT EXISTS idx_memory_files_project ON memory_files(project_id); +CREATE INDEX IF NOT EXISTS idx_memory_files_scope + ON memory_files(scope, scope_id); CREATE INDEX IF NOT EXISTS idx_memory_files_source ON memory_files(source); -CREATE INDEX IF NOT EXISTS idx_memory_chunks_project ON memory_chunks(project_id); -CREATE INDEX IF NOT EXISTS idx_memory_chunks_path ON memory_chunks(path); +CREATE INDEX IF NOT EXISTS idx_memory_chunks_scope + ON memory_chunks(scope, scope_id); +CREATE INDEX IF NOT EXISTS idx_memory_chunks_path + ON memory_chunks(scope, scope_id, path); CREATE INDEX IF NOT EXISTS idx_memory_chunks_source ON memory_chunks(source); CREATE INDEX IF NOT EXISTS idx_memory_embedding_cache_accessed ON memory_embedding_cache(accessed_at); """ @@ -78,7 +84,8 @@ chunk_id UNINDEXED, path UNINDEXED, source UNINDEXED, - project_id UNINDEXED, + scope UNINDEXED, + scope_id UNINDEXED, start_line UNINDEXED, end_line UNINDEXED, tokenize = 'porter unicode61' @@ -93,7 +100,7 @@ async def ensure_vector_tables(db_path: Path) -> Dict[str, Any]: Returns: Status dict with table availability info """ - status = { + status: Dict[str, Any] = { "vector_tables": False, "fts5": False, "fts5_error": None, @@ -101,6 +108,28 @@ async def ensure_vector_tables(db_path: Path) -> Dict[str, Any]: try: async with Storage.connect(db_path) as db: + cursor = await db.execute("PRAGMA table_info(memory_files)") + file_columns = {row[1] for row in await cursor.fetchall()} + cursor = await db.execute("PRAGMA table_info(memory_chunks)") + chunk_columns = {row[1] for row in await cursor.fetchall()} + cursor = await db.execute("PRAGMA table_info(memory_fts)") + fts_columns = {row[1] for row in await cursor.fetchall()} + old_scope_schema = any( + columns + and not {"scope", "scope_id"}.issubset(columns) + for columns in (file_columns, chunk_columns, fts_columns) + ) + if old_scope_schema: + await db.execute("BEGIN IMMEDIATE") + try: + await db.execute("DROP TABLE IF EXISTS memory_fts") + await db.execute("DROP TABLE IF EXISTS memory_chunks") + await db.execute("DROP TABLE IF EXISTS memory_files") + await db.commit() + log.info("vector.memory_scope_schema.rebuilt") + except BaseException: + await db.rollback() + raise # Create vector tables await db.executescript(VECTOR_SCHEMA_SQL) await db.commit() @@ -157,7 +186,7 @@ def bm25_rank_to_score(rank: float) -> float: Returns: Normalized score """ - normalized = max(0, rank) if math.isfinite(rank) else 999 + normalized = abs(rank) if math.isfinite(rank) else 999 return 1 / (1 + normalized) @@ -177,7 +206,8 @@ async def vector_search( Args: db_path: Database path - project_id: Project ID to filter + project_id: Current Session project ID (retained for API compatibility; + Memory file search is global) embedding: Query embedding vector max_results: Maximum results to return min_score: Minimum similarity score @@ -187,16 +217,16 @@ async def vector_search( List of search results """ results = [] + del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: - # Build query query = """ SELECT id, path, source, start_line, end_line, text, embedding FROM memory_chunks - WHERE project_id = ? AND embedding IS NOT NULL + WHERE embedding IS NOT NULL """ - params = [project_id] + params: list[Any] = [] if sources: placeholders = ",".join("?" * len(sources)) @@ -254,7 +284,7 @@ def build_fts_query(raw: str) -> Optional[str]: """ Build FTS5 query string from raw text - Extracts alphanumeric tokens and combines with AND. + Extracts Unicode word tokens and combines them with AND. Args: raw: Raw query text @@ -264,8 +294,9 @@ def build_fts_query(raw: str) -> Optional[str]: """ import re - # Extract alphanumeric tokens - tokens = re.findall(r'[A-Za-z0-9_]+', raw) + # Python's Unicode-aware ``\w`` keeps CJK and other scripts intact while + # quoting prevents user input from becoming FTS5 syntax. + tokens = re.findall(r"\w+", raw, flags=re.UNICODE) tokens = [t.strip() for t in tokens if t.strip()] if not tokens: @@ -288,7 +319,8 @@ async def fts_search( Args: db_path: Database path - project_id: Project ID to filter + project_id: Current Session project ID (retained for API compatibility; + Memory file search is global) query: Search query (FTS5 format) max_results: Maximum results to return sources: Optional list of sources to filter @@ -297,6 +329,7 @@ async def fts_search( List of search results with BM25 scores """ results = [] + del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -317,9 +350,8 @@ async def fts_search( rank FROM memory_fts f WHERE f.text MATCH ? - AND f.project_id = ? """ - params = [fts_query, project_id] + params = [fts_query] if sources: placeholders = ",".join("?" * len(sources)) @@ -364,7 +396,7 @@ async def insert_chunks( Args: db_path: Database path chunks: List of chunk dicts with keys: - - id, path, project_id, source, start_line, end_line, + - id, scope, scope_id, path, source, start_line, end_line, hash, text, embedding, embedding_model, embedding_dims Returns: @@ -377,14 +409,15 @@ async def insert_chunks( # Insert into chunks table await db.executemany(""" INSERT OR REPLACE INTO memory_chunks - (id, path, project_id, source, start_line, end_line, hash, text, + (id, scope, scope_id, path, source, start_line, end_line, hash, text, embedding, embedding_model, embedding_dims, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, [ ( chunk["id"], + chunk["scope"], + chunk["scope_id"], chunk["path"], - chunk["project_id"], chunk["source"], chunk["start_line"], chunk["end_line"], @@ -401,16 +434,24 @@ async def insert_chunks( # Insert into FTS5 table (if exists) try: + chunk_ids = [chunk["id"] for chunk in chunks] + if chunk_ids: + placeholders = ",".join("?" for _ in chunk_ids) + await db.execute( + f"DELETE FROM memory_fts WHERE chunk_id IN ({placeholders})", + tuple(chunk_ids), + ) await db.executemany(""" INSERT OR REPLACE INTO memory_fts - (chunk_id, path, source, project_id, start_line, end_line, text) - VALUES (?, ?, ?, ?, ?, ?, ?) + (chunk_id, path, source, scope, scope_id, start_line, end_line, text) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, [ ( chunk["id"], chunk["path"], chunk["source"], - chunk["project_id"], + chunk["scope"], + chunk["scope_id"], chunk["start_line"], chunk["end_line"], chunk["text"], @@ -430,6 +471,112 @@ async def insert_chunks( raise +async def replace_memory_file_index( + db_path: Path, + *, + file_entry: Dict[str, Any], + chunks: List[Dict[str, Any]], +) -> int: + """Atomically replace one Memory file's metadata, chunks, and FTS rows.""" + scope = file_entry["scope"] + scope_id = file_entry["scope_id"] + path = file_entry["path"] + now = datetime.now().timestamp() + async with Storage.connect(db_path) as db: + try: + await db.execute("BEGIN IMMEDIATE") + await db.execute( + """ + DELETE FROM memory_fts + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + (scope, scope_id, path), + ) + await db.execute( + """ + DELETE FROM memory_chunks + WHERE scope = ? AND scope_id = ? + AND source = 'memory' AND path = ? + """, + (scope, scope_id, path), + ) + await db.execute( + """ + INSERT OR REPLACE INTO memory_files ( + scope, scope_id, path, source, hash, mtime, size, indexed_at + ) VALUES (?, ?, ?, 'memory', ?, ?, ?, ?) + """, + ( + scope, + scope_id, + path, + file_entry["hash"], + file_entry["mtime"], + file_entry["size"], + now, + ), + ) + await db.executemany( + """ + INSERT INTO memory_chunks ( + id, scope, scope_id, path, source, start_line, end_line, hash, + text, embedding, embedding_model, embedding_dims, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + chunk["id"], + chunk["scope"], + chunk["scope_id"], + chunk["path"], + chunk["source"], + chunk["start_line"], + chunk["end_line"], + chunk["hash"], + chunk["text"], + ( + json.dumps(chunk["embedding"]) + if chunk.get("embedding") + else None + ), + chunk.get("embedding_model"), + chunk.get("embedding_dims"), + now, + now, + ) + for chunk in chunks + ], + ) + await db.executemany( + """ + INSERT INTO memory_fts ( + chunk_id, path, source, scope, scope_id, start_line, end_line, + text + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + chunk["id"], + chunk["path"], + chunk["source"], + chunk["scope"], + chunk["scope_id"], + chunk["start_line"], + chunk["end_line"], + chunk["text"], + ) + for chunk in chunks + ], + ) + await db.commit() + return len(chunks) + except BaseException: + await db.rollback() + raise + + async def get_embedding_from_cache( db_path: Path, text_hash: str, diff --git a/tests/config/test_config_init.py b/tests/config/test_config_init.py index 665cd6368..981e2496d 100644 --- a/tests/config/test_config_init.py +++ b/tests/config/test_config_init.py @@ -2,6 +2,8 @@ Tests for config file initialization from examples. """ +import json + import pytest @@ -44,8 +46,12 @@ def test_ensure_config_files_creates_from_examples(tmp_path, monkeypatch): assert mcp_file.exists() assert secret_file.exists() - # Content should match examples - assert config_file.read_text(encoding="utf-8") == '{"test": "config"}' + # Existing example content is preserved and Memory defaults are persisted. + config_data = json.loads(config_file.read_text(encoding="utf-8")) + assert config_data["test"] == "config" + assert set(config_data["memory"]) == {"search"} + assert config_data["memory"]["search"]["embedding"]["provider"] == "auto" + assert config_data["memory"]["search"]["embedding"]["enabled"] is False assert mcp_file.read_text(encoding="utf-8") == '{"test": "mcp"}' assert secret_file.read_text(encoding="utf-8") == '{"test": "secret"}' @@ -81,11 +87,37 @@ def test_ensure_config_files_skips_if_exists(tmp_path, monkeypatch): ensure_config_files = config_writer.ensure_config_files ensure_config_files() - # File should still have original content - assert config_file.read_text() == '{"test": "existing"}' + # Existing fields are preserved while the missing Memory config is added. + config_data = json.loads(config_file.read_text(encoding="utf-8")) + assert config_data["test"] == "existing" + assert set(config_data["memory"]) == {"search"} assert mcp_file.read_text() == '{"test": "mcp-existing"}' +def test_ensure_memory_config_is_written_to_flocks_json( + tmp_path, + monkeypatch, +): + """The generated Memory section belongs to the primary flocks.json.""" + config_dir = tmp_path / "home" / ".flocks" / "config" + config_dir.mkdir(parents=True) + monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(config_dir)) + flocks_json = config_dir / "flocks.json" + flocks_json.write_text("{}", encoding="utf-8") + flocks_jsonc = config_dir / "flocks.jsonc" + flocks_jsonc.write_text('{"test": "jsonc"}', encoding="utf-8") + + from flocks.config.config import Config + from flocks.config.config_writer import ConfigWriter + + Config._global_config = None + Config._cached_config = None + assert ConfigWriter.ensure_memory_config() is True + memory_config = json.loads(flocks_json.read_text(encoding="utf-8"))["memory"] + assert set(memory_config) == {"search"} + assert flocks_jsonc.read_text(encoding="utf-8") == '{"test": "jsonc"}' + + def test_ensure_config_files_handles_missing_examples(tmp_path, monkeypatch): """Test that ensure_config_files handles missing example files gracefully.""" config_dir = tmp_path / "home" / ".flocks" / "config" diff --git a/tests/memory/test_memory_scope.py b/tests/memory/test_memory_scope.py new file mode 100644 index 000000000..63d0c20a8 --- /dev/null +++ b/tests/memory/test_memory_scope.py @@ -0,0 +1,311 @@ +"""Tests for Global and Project Memory scope isolation.""" + +import os +from pathlib import Path +import sqlite3 +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.memory.config import MemoryConfig +from flocks.memory.manager import MemoryManager +from flocks.memory.sync.indexer import MemoryIndexer +from flocks.memory.types import MemoryScope +from flocks.storage import ( + Storage, + ensure_vector_tables, + fts_search, + replace_memory_file_index, +) + + +def _file_entry( + scope: str, + scope_id: str, + path: str, +) -> dict[str, object]: + return { + "scope": scope, + "scope_id": scope_id, + "path": path, + "hash": f"hash:{scope}:{scope_id}:{path}", + "mtime": 1, + "size": 10, + } + + +def _chunk( + scope: str, + scope_id: str, + path: str, + text: str, +) -> dict[str, object]: + return { + "id": f"chunk:{scope}:{scope_id}:{path}", + "scope": scope, + "scope_id": scope_id, + "path": path, + "source": "memory", + "start_line": 1, + "end_line": 1, + "hash": f"hash:{text}", + "text": text, + "embedding": None, + "embedding_model": None, + "embedding_dims": None, + } + + +@pytest.mark.asyncio +async def test_search_reconciles_filesystem_before_every_search( + tmp_path: Path, +) -> None: + manager = MemoryManager( + project_id="prj_alpha", + workspace_dir=str(tmp_path), + config=MemoryConfig(), + ) + manager._initialized = True + manager.sync = AsyncMock(return_value={}) + manager.search_engine = SimpleNamespace( + search=AsyncMock(return_value=[]), + ) + + await manager.search("new filesystem memory") + + manager.sync.assert_awaited_once_with(reason="search") + + +@pytest.mark.asyncio +async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: + db_path = tmp_path / "scope.db" + await Storage.init(db_path) + records = [ + ("global", "", "MEMORY.md", "scopeword global"), + ( + "project", + "prj_alpha", + "projects/prj_alpha/MEMORY.md", + "scopeword alpha", + ), + ( + "project", + "prj_beta", + "projects/prj_beta/MEMORY.md", + "scopeword beta", + ), + ] + for scope, scope_id, path, text in records: + await replace_memory_file_index( + db_path, + file_entry=_file_entry(scope, scope_id, path), + chunks=[_chunk(scope, scope_id, path, text)], + ) + + alpha = await fts_search(db_path, "prj_alpha", "scopeword") + default = await fts_search(db_path, "default", "scopeword") + + expected_paths = { + "MEMORY.md", + "projects/prj_alpha/MEMORY.md", + "projects/prj_beta/MEMORY.md", + } + assert {result["path"] for result in alpha} == expected_paths + assert {result["path"] for result in default} == expected_paths + + +@pytest.mark.asyncio +async def test_indexer_scans_global_and_all_projects( + tmp_path: Path, +) -> None: + memory_root = tmp_path / "memory" + (memory_root / "daily").mkdir(parents=True) + (memory_root / "projects" / "prj_alpha").mkdir(parents=True) + (memory_root / "projects" / "prj_beta").mkdir(parents=True) + (memory_root / "MEMORY.md").write_text("global", encoding="utf-8") + (memory_root / "daily" / "2026-01-01.md").write_text( + "daily", + encoding="utf-8", + ) + (memory_root / "projects" / "prj_alpha" / "MEMORY.md").write_text( + "alpha", + encoding="utf-8", + ) + (memory_root / "projects" / "prj_beta" / "MEMORY.md").write_text( + "beta", + encoding="utf-8", + ) + indexer = MemoryIndexer( + project_id="prj_alpha", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + files = await indexer._scan_memory_files() + + identities = {(entry.scope.value, entry.scope_id, entry.path) for entry in files} + assert ("global", "", "MEMORY.md") in identities + assert ("global", "", "daily/2026-01-01.md") in identities + assert ( + "project", + "prj_alpha", + "projects/prj_alpha/MEMORY.md", + ) in identities + assert ( + "project", + "prj_beta", + "projects/prj_beta/MEMORY.md", + ) in identities + + +@pytest.mark.asyncio +async def test_indexer_does_not_read_unchanged_files( + tmp_path: Path, +) -> None: + db_path = tmp_path / "metadata-scan.db" + await Storage.init(db_path) + memory_root = tmp_path / "memory" + memory_root.mkdir() + memory_file = memory_root / "MEMORY.md" + memory_file.write_text("stable memory", encoding="utf-8") + indexer = MemoryIndexer( + project_id="global", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + initial = await indexer.sync() + with patch.object( + Path, + "read_text", + side_effect=AssertionError("unchanged Memory file was read"), + ): + unchanged = await indexer.sync() + + assert initial["files_indexed"] == 1 + assert unchanged["files_indexed"] == 0 + assert unchanged["files_skipped"] == 1 + + +@pytest.mark.asyncio +async def test_indexer_refreshes_metadata_without_reindexing_unchanged_content( + tmp_path: Path, +) -> None: + db_path = tmp_path / "metadata-refresh.db" + await Storage.init(db_path) + memory_root = tmp_path / "memory" + memory_root.mkdir() + memory_file = memory_root / "MEMORY.md" + memory_file.write_text("stable memory", encoding="utf-8") + indexer = MemoryIndexer( + project_id="global", + workspace_dir=tmp_path, + provider_id=None, + embedding_model="unused", + config=MemoryConfig(), + ) + + with patch("flocks.config.Config.get_data_path", return_value=tmp_path): + await indexer.sync() + before = memory_file.stat() + os.utime( + memory_file, + ns=(before.st_atime_ns, before.st_mtime_ns + 2_000_000_000), + ) + with patch.object( + indexer, + "_index_file", + wraps=indexer._index_file, + ) as index_file: + touched = await indexer.sync() + indexed_files = await indexer._get_indexed_files() + with patch.object( + Path, + "read_text", + side_effect=AssertionError("refreshed Memory file was read again"), + ): + unchanged = await indexer.sync() + + indexed = indexed_files[("global", "", "MEMORY.md")] + assert touched["files_indexed"] == 0 + assert touched["files_skipped"] == 1 + index_file.assert_not_awaited() + assert indexed["mtime"] == memory_file.stat().st_mtime + assert unchanged["files_indexed"] == 0 + assert unchanged["files_skipped"] == 1 + + +@pytest.mark.asyncio +async def test_old_memory_index_schema_is_rebuilt_without_other_data( + tmp_path: Path, +) -> None: + db_path = tmp_path / "legacy.db" + connection = sqlite3.connect(db_path) + connection.executescript( + """ + CREATE TABLE memory_files ( + path TEXT PRIMARY KEY, + project_id TEXT NOT NULL, + source TEXT NOT NULL, + hash TEXT NOT NULL, + mtime REAL NOT NULL, + size INTEGER NOT NULL, + indexed_at REAL NOT NULL + ); + CREATE TABLE memory_chunks ( + id TEXT PRIMARY KEY, + path TEXT NOT NULL, + project_id TEXT NOT NULL, + source TEXT NOT NULL, + start_line INTEGER NOT NULL, + end_line INTEGER NOT NULL, + hash TEXT NOT NULL, + text TEXT NOT NULL, + embedding BLOB, + embedding_model TEXT, + embedding_dims INTEGER, + created_at REAL NOT NULL, + updated_at REAL NOT NULL + ); + CREATE VIRTUAL TABLE memory_fts USING fts5( + text, chunk_id, path, source, project_id, start_line, end_line + ); + CREATE TABLE memory_embedding_cache ( + text_hash TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + embedding BLOB NOT NULL, + dims INTEGER NOT NULL, + created_at REAL NOT NULL, + accessed_at REAL NOT NULL, + PRIMARY KEY (text_hash, provider, model) + ); + INSERT INTO memory_embedding_cache + VALUES ('hash', 'provider', 'model', '[1.0]', 1, 1, 1); + CREATE VIRTUAL TABLE session_transcript_fts USING fts5(text); + INSERT INTO session_transcript_fts VALUES ('preserved'); + """ + ) + connection.commit() + connection.close() + + await ensure_vector_tables(db_path) + + connection = sqlite3.connect(db_path) + columns = {row[1] for row in connection.execute("PRAGMA table_info(memory_files)")} + cache_row = connection.execute("SELECT text_hash FROM memory_embedding_cache").fetchone() + marker_row = connection.execute( + "SELECT text FROM session_transcript_fts" + ).fetchone() + connection.close() + + assert {"scope", "scope_id", "path"}.issubset(columns) + assert cache_row == ("hash",) + assert marker_row == ("preserved",) diff --git a/tests/memory/test_session_transcript_search.py b/tests/memory/test_session_transcript_search.py new file mode 100644 index 000000000..aa2c05c51 --- /dev/null +++ b/tests/memory/test_session_transcript_search.py @@ -0,0 +1,622 @@ +"""Session transcript FTS lifecycle tests.""" + +from pathlib import Path +import sqlite3 +from unittest.mock import AsyncMock, Mock +import uuid + +import pytest + +from flocks.config.config import Config +from flocks.memory.config import MemoryConfig +from flocks.memory.manager import MemoryManager +from flocks.memory.search.hybrid import HybridSearch +from flocks.memory.types import MemorySearchResult +from flocks.memory.types import MemorySource +from flocks.provider import Provider +from flocks.session.message import Message, MessageRole +from flocks.session.session import Session, SessionInfo +from flocks.storage.session_search import ( + SessionSearchUnavailableError, + _is_fts5_unavailable_error, + ensure_session_index_ready, + ensure_session_search_tables, + reconcile_session_index, + session_fts_search, +) +from flocks.storage import session_search as session_search_module +from flocks.storage.storage import Storage + + +@pytest.fixture(autouse=True) +async def isolate_transcript_search( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + flocks_root = tmp_path / "flocks-home" + data_dir = flocks_root / "data" + monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) + monkeypatch.setenv("FLOCKS_DATA_DIR", str(data_dir)) + monkeypatch.setenv("FLOCKS_LOG_DIR", str(flocks_root / "logs")) + monkeypatch.setenv("FLOCKS_RECORD_DIR", str(data_dir / "records")) + + Config._global_config = None + Config.clear_cache() + Storage._initialized = False + Storage._db_path = None + Session.invalidate_cache() + Message.invalidate_cache() + MemoryManager._instances.clear() + await Storage.init() + + yield + + Session.invalidate_cache() + Message.invalidate_cache() + MemoryManager._instances.clear() + Config._global_config = None + Config.clear_cache() + Storage._initialized = False + Storage._db_path = None + + +async def _create_session(tmp_path: Path, project_id: str = "project-search"): + session = SessionInfo( + id=f"session-{uuid.uuid4().hex}", + project_id=project_id, + directory=str(tmp_path), + agent="rex", + memory_enabled=True, + ) + await Storage.set( + f"session:{project_id}:{session.id}", + session, + "session", + ) + Session.invalidate_cache() + return session + + +def test_only_explicit_missing_fts5_errors_are_classified() -> None: + assert _is_fts5_unavailable_error( + sqlite3.OperationalError("no such module: fts5") + ) + assert _is_fts5_unavailable_error( + sqlite3.OperationalError("unknown module: fts5") + ) + assert not _is_fts5_unavailable_error( + sqlite3.OperationalError("database is locked") + ) + assert not _is_fts5_unavailable_error( + RuntimeError("no such module: fts5") + ) + + +@pytest.mark.asyncio +async def test_session_schema_probe_degrades_when_fts5_module_is_missing( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + class MissingFtsConnection: + def __init__(self): + self.statements: list[str] = [] + self.committed = False + + async def executescript(self, sql: str): + self.statements.append(sql) + + async def execute(self, sql: str, _parameters=()): + self.statements.append(sql) + if "_flocks_session_fts5_probe" in sql: + raise sqlite3.OperationalError("no such module: fts5") + return None + + async def commit(self): + self.committed = True + + class MissingFtsContext: + def __init__(self): + self.connection = MissingFtsConnection() + + async def __aenter__(self): + return self.connection + + async def __aexit__(self, *_args): + return None + + context = MissingFtsContext() + monkeypatch.setattr( + Storage, + "connect", + classmethod(lambda _cls, _path=None: context), + ) + + assert not await ensure_session_search_tables(tmp_path / "missing-fts.db") + assert context.connection.committed + assert any( + "DELETE FROM session_transcript_meta" in statement + for statement in context.connection.statements + ) + + +@pytest.mark.asyncio +async def test_messages_persist_when_session_search_is_disabled( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + index_message = AsyncMock( + side_effect=AssertionError("Session FTS hook must be disabled") + ) + monkeypatch.setattr( + session_search_module, + "upsert_session_document", + index_message, + ) + monkeypatch.setattr(Storage, "_session_search_available", False) + + message = await Message.create( + session.id, + MessageRole.USER, + "canonical message survives without FTS5", + ) + + stored = await Message.get(session.id, message.id) + assert stored is not None + parts = await Message.parts(message.id, session.id) + assert [part.text for part in parts if part.type == "text"] == [ + "canonical message survives without FTS5" + ] + index_message.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_search_reports_fts5_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(Storage, "_session_search_available", False) + + with pytest.raises( + SessionSearchUnavailableError, + match="SQLite runtime does not support FTS5", + ): + await session_fts_search( + db_path=Storage.get_db_path(), + project_id="default", + query="anything", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_memory_manager_starts_without_fts5_and_session_search_fails_clearly( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + monkeypatch.setattr(Storage, "_session_search_available", False) + + manager = MemoryManager( + project_id="default", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["session"]), + ) + + await manager.initialize() + + assert manager._initialized + with pytest.raises( + SessionSearchUnavailableError, + match="SQLite runtime does not support FTS5", + ): + await manager.search( + query="anything", + sources=[MemorySource.SESSION], + ) + + +@pytest.mark.asyncio +async def test_text_part_updates_and_message_delete_update_fts( + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.USER, + "initial searchable phrase", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="searchable", + max_results=10, + ) + assert [result["path"] for result in results] == [ + f"sessions/{session.id}/messages/{message.id}" + ] + + part = (await Message.parts(message.id, session.id))[0] + await Message.update_part( + session.id, + message.id, + part.id, + text="replacement transcript text", + ) + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="initial", + max_results=10, + ) + assert await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="replacement", + max_results=10, + ) + + assert await Message.delete(session.id, message.id) + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="replacement", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_session_search_is_global_across_projects( + tmp_path: Path, +) -> None: + alpha = await _create_session(tmp_path, project_id="prj_alpha") + beta = await _create_session(tmp_path, project_id="prj_beta") + alpha_message = await Message.create( + alpha.id, + MessageRole.USER, + "cross project session marker alpha", + ) + beta_message = await Message.create( + beta.id, + MessageRole.ASSISTANT, + "cross project session marker beta", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=alpha.project_id, + query="cross project session marker", + max_results=10, + ) + + assert {result["path"] for result in results} == { + f"sessions/{alpha.id}/messages/{alpha_message.id}", + f"sessions/{beta.id}/messages/{beta_message.id}", + } + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute("DELETE FROM session_transcript_index_state") + await db.commit() + stats = await reconcile_session_index(batch_size=1) + rebuilt = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=alpha.project_id, + query="cross project session marker", + max_results=10, + ) + assert stats["updated"] == 2 + assert {result["path"] for result in rebuilt} == { + f"sessions/{alpha.id}/messages/{alpha_message.id}", + f"sessions/{beta.id}/messages/{beta_message.id}", + } + + +@pytest.mark.asyncio +async def test_reconciliation_restores_history_and_removes_orphans( + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.ASSISTANT, + "historical reconciliation marker", + ) + + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute( + "INSERT INTO session_transcript_fts(rowid, text) VALUES (999, 'orphan')" + ) + await db.commit() + + stats = await reconcile_session_index( + project_id=session.project_id, + batch_size=1, + ) + assert stats["updated"] == 1 + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="reconciliation", + max_results=10, + ) + assert results[0]["path"].endswith(message.id) + + async with Storage.connect(Storage.get_db_path()) as db: + cursor = await db.execute( + "SELECT count(*) FROM session_transcript_fts WHERE rowid = 999" + ) + assert (await cursor.fetchone())[0] == 0 + + +@pytest.mark.asyncio +async def test_session_history_backfill_runs_only_once( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + message = await Message.create( + session.id, + MessageRole.USER, + "legacy transcript backfill marker", + ) + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute("DELETE FROM session_transcript_fts") + await db.execute("DELETE FROM session_transcript_index_state") + await db.commit() + + reconcile = AsyncMock( + wraps=session_search_module._reconcile_session_index_unlocked + ) + monkeypatch.setattr( + session_search_module, + "_reconcile_session_index_unlocked", + reconcile, + ) + + assert await ensure_session_index_ready(batch_size=1) + assert not await ensure_session_index_ready(batch_size=1) + assert reconcile.await_count == 1 + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="legacy transcript", + max_results=10, + ) + assert [result["path"] for result in results] == [ + f"sessions/{session.id}/messages/{message.id}" + ] + + +@pytest.mark.asyncio +async def test_failed_session_backfill_is_retried( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_reconcile = session_search_module._reconcile_session_index_unlocked + attempts = 0 + + async def fail_once(*, project_id=None, batch_size=50): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("injected backfill failure") + return await original_reconcile( + project_id=project_id, + batch_size=batch_size, + ) + + monkeypatch.setattr( + session_search_module, + "_reconcile_session_index_unlocked", + fail_once, + ) + + with pytest.raises(RuntimeError, match="injected backfill failure"): + await ensure_session_index_ready(batch_size=1) + + assert await ensure_session_index_ready(batch_size=1) + assert attempts == 2 + + +@pytest.mark.asyncio +async def test_memory_managers_share_one_global_file_indexer( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + sync = AsyncMock( + return_value={ + "files_scanned": 0, + "files_indexed": 0, + "files_skipped": 0, + "chunks_created": 0, + "embeddings_generated": 0, + "cache_hits": 0, + } + ) + monkeypatch.setattr("flocks.memory.sync.indexer.MemoryIndexer.sync", sync) + + config = MemoryConfig(sources=["memory"]) + alpha = MemoryManager.get_instance( + project_id="prj_alpha", + workspace_dir=str(tmp_path / "alpha"), + config=config, + ) + beta = MemoryManager.get_instance( + project_id="prj_beta", + workspace_dir=str(tmp_path / "beta"), + config=config, + ) + + await alpha.initialize() + await beta.initialize() + + assert alpha.indexer is beta.indexer + assert sync.await_count == 1 + + +@pytest.mark.asyncio +async def test_synthetic_text_is_not_indexed(tmp_path: Path) -> None: + session = await _create_session(tmp_path) + await Message.create( + session.id, + MessageRole.ASSISTANT, + "synthetic compaction marker", + synthetic=True, + ) + + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id=session.project_id, + query="compaction", + max_results=10, + ) + + +@pytest.mark.asyncio +async def test_explicit_session_search_persists_opt_in_without_embeddings( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + session = await _create_session(tmp_path) + await Message.create( + session.id, + MessageRole.USER, + "session source opt in marker", + ) + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + + manager = MemoryManager( + project_id=session.project_id, + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + results = await manager.search( + query="marker", + sources=[MemorySource.SESSION], + ) + + assert results + assert results[0].source is MemorySource.SESSION + assert manager.provider_id is None + assert "session" in manager.config.sources + + config_path = Config.get_config_file() + persisted = config_path.read_text(encoding="utf-8") + assert '"sources": [' in persisted + assert '"session"' in persisted + + +@pytest.mark.asyncio +async def test_memory_search_uses_fts_without_embedding_provider( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + provider_init = AsyncMock() + provider_get = Mock() + monkeypatch.setattr(Provider, "init", provider_init) + monkeypatch.setattr(Provider, "get", provider_get) + async with Storage.connect(Storage.get_db_path()) as db: + await db.execute( + """ + INSERT INTO memory_fts ( + text, chunk_id, path, source, scope, scope_id, + start_line, end_line + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + "durable keyword memory", + "chunk-1", + "MEMORY.md", + "memory", + "global", + "", + 1, + 1, + ), + ) + await db.commit() + + manager = MemoryManager( + project_id="project-memory", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + results = await manager.search("durable") + + assert manager.provider_id is None + provider_init.assert_not_awaited() + provider_get.assert_not_called() + assert [result.path for result in results] == ["MEMORY.md"] + + +@pytest.mark.asyncio +async def test_memory_sync_indexes_fts_when_embeddings_are_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(Provider, "init", AsyncMock()) + monkeypatch.setattr(Provider, "get", lambda _provider_id: None) + memory_root = Config.get_data_path() / "memory" + memory_root.mkdir(parents=True, exist_ok=True) + (memory_root / "notes.md").write_text( + "fts fallback indexing marker", + encoding="utf-8", + ) + + manager = MemoryManager( + project_id="project-sync", + workspace_dir=str(tmp_path), + config=MemoryConfig(sources=["memory"]), + ) + await manager.sync(force=True) + results = await manager.search("fallback") + + assert results + assert results[0].path == "notes.md" + + +@pytest.mark.asyncio +async def test_embedding_failure_falls_back_to_keyword_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = HybridSearch( + project_id="project-fallback", + provider_id="openai", + embedding_model="embedding-model", + config=MemoryConfig().query, + ) + monkeypatch.setattr( + engine, + "_vector_search", + AsyncMock(side_effect=RuntimeError("embedding unavailable")), + ) + monkeypatch.setattr( + engine, + "_keyword_search", + AsyncMock( + return_value=[ + MemorySearchResult( + path="MEMORY.md", + start_line=1, + end_line=1, + score=1.0, + snippet="keyword fallback", + source=MemorySource.MEMORY, + ) + ] + ), + ) + + results = await engine.search( + query="fallback", + max_results=6, + min_score=0.35, + sources=[MemorySource.MEMORY], + ) + assert [result.path for result in results] == ["MEMORY.md"] diff --git a/tests/session/test_message_parts_persistence.py b/tests/session/test_message_parts_persistence.py index 61dafc1ff..a827828e2 100644 --- a/tests/session/test_message_parts_persistence.py +++ b/tests/session/test_message_parts_persistence.py @@ -193,10 +193,10 @@ async def test_delete_restores_caches_when_message_persistence_fails( id="msg_a", part_id="part_a", ) - persist_messages = AsyncMock( + mutate_many = AsyncMock( side_effect=RuntimeError("message storage unavailable") ) - monkeypatch.setattr(Message, "_persist_messages", persist_messages) + monkeypatch.setattr(Storage, "mutate_many", mutate_many) with pytest.raises(RuntimeError, match="message storage unavailable"): await Message.delete(session_id, "msg_a") @@ -210,9 +210,7 @@ async def test_delete_restores_caches_when_message_persistence_fails( @pytest.mark.asyncio -async def test_delete_commits_when_parts_cleanup_fails( - monkeypatch: pytest.MonkeyPatch, -) -> None: +async def test_delete_removes_message_and_parts_atomically() -> None: session_id = "ses_parts_delete_parts_failure" await Message.create( session_id, @@ -221,19 +219,10 @@ async def test_delete_commits_when_parts_cleanup_fails( id="msg_a", part_id="part_a", ) - original_delete = Storage.delete - - async def fail_parts_delete(key: str) -> None: - if key == f"message_parts:{session_id}:msg_a": - raise RuntimeError("parts storage unavailable") - await original_delete(key) - - monkeypatch.setattr(Storage, "delete", fail_parts_delete) - assert await Message.delete(session_id, "msg_a") is True assert await Message.get(session_id, "msg_a") is None assert await Storage.get(f"message:{session_id}") == [] - assert await Storage.get(f"message_parts:{session_id}:msg_a") is not None + assert await Storage.get(f"message_parts:{session_id}:msg_a") is None @pytest.mark.asyncio diff --git a/tests/storage/test_storage.py b/tests/storage/test_storage.py index 36f644d2d..c15d92a22 100644 --- a/tests/storage/test_storage.py +++ b/tests/storage/test_storage.py @@ -16,6 +16,7 @@ from pydantic import BaseModel from flocks.project.instance import Instance +from flocks.storage import session_search as session_search_module from flocks.storage.storage import Storage from flocks.task.store import TaskStore from flocks.workflow.store import WorkflowStore @@ -28,6 +29,49 @@ class StorageTestModel(BaseModel): value: int +@pytest.mark.asyncio +async def test_storage_init_continues_when_session_fts_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + session_search_module, + "ensure_session_search_tables", + AsyncMock(return_value=False), + ) + + with patch.object(Storage, "_initialized", False), patch.object( + Storage, + "_db_path", + None, + ), patch.object(Storage, "_session_search_available", True): + await Storage.init(tmp_path / "fts-unavailable.db") + + assert Storage._initialized + assert not Storage.session_search_available() + + +@pytest.mark.asyncio +async def test_storage_init_propagates_unexpected_session_schema_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + failure = RuntimeError("unexpected schema failure") + monkeypatch.setattr( + session_search_module, + "ensure_session_search_tables", + AsyncMock(side_effect=failure), + ) + + with patch.object(Storage, "_initialized", False), patch.object( + Storage, + "_db_path", + None, + ), patch.object(Storage, "_session_search_available", True): + with pytest.raises(RuntimeError, match="unexpected schema failure"): + await Storage.init(tmp_path / "broken-session-schema.db") + + def _require_sqlite_recover() -> None: sqlite_bin = shutil.which("sqlite3") if sqlite_bin is None: From c3b5e10f763c4b4d3cf3eec058e2f5b05f5d0bd7 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 6 Aug 2026 10:01:40 +0800 Subject: [PATCH 2/3] fix(memory): isolate memory and session search --- flocks/memory/bootstrap.py | 4 +- flocks/memory/manager.py | 9 +- flocks/session/features/memory.py | 117 ++++++- flocks/storage/session_search.py | 48 +-- flocks/storage/vector.py | 20 +- flocks/tool/system/memory.py | 4 +- tests/memory/test_memory_scope.py | 46 ++- .../memory/test_session_transcript_search.py | 311 +++++++++++++++++- 8 files changed, 508 insertions(+), 51 deletions(-) diff --git a/flocks/memory/bootstrap.py b/flocks/memory/bootstrap.py index 88e0b4992..6bf032faa 100644 --- a/flocks/memory/bootstrap.py +++ b/flocks/memory/bootstrap.py @@ -57,7 +57,7 @@ ### Managing Memory Files: - The injected USER, Global, and Project files are a snapshot for this run. Read the file again before changing it. -- Use `read`, `glob`, and `grep` to inspect Memory explicitly, and `memory_search` for indexed recall across all projects. +- Use `read`, `glob`, and `grep` to inspect Memory explicitly, and `memory_search` for indexed recall across USER, Global, Daily, and the current Project. - Use `write` only to create a missing curated Memory file. Use `edit` for precise entry-level changes to an existing curated file. - Never write or edit `daily/`; only the Session lifecycle may append Daily entries. - **User profile**: Maintain `{memory_root}/USER.md` only for facts about the user. @@ -88,7 +88,7 @@ - Verify stale or conflicting Memory against current authoritative evidence before replacing or removing it. ### Available Tools: -- `memory_search` - Reconcile and search indexed Memory across all projects +- `memory_search` - Reconcile and search USER, Global, Daily, and current Project Memory - `read`, `glob`, `grep` - Inspect Memory files - `write` - Create a missing Memory file - `edit` - Precisely update an existing Memory file diff --git a/flocks/memory/manager.py b/flocks/memory/manager.py index 99978a79e..9978c8d4c 100644 --- a/flocks/memory/manager.py +++ b/flocks/memory/manager.py @@ -4,7 +4,7 @@ Coordinates all memory system components: indexing, search, and sync. """ -from typing import Optional, List, Dict, Any, Callable +from typing import Optional, List, Dict, Any, Callable, Set from pathlib import Path import asyncio import os @@ -347,6 +347,7 @@ async def search( max_results: Optional[int] = None, min_score: Optional[float] = None, sources: Optional[List[MemorySource]] = None, + readable_session_ids: Optional[Set[str]] = None, ) -> List[MemorySearchResult]: """ Search memory @@ -356,6 +357,7 @@ async def search( max_results: Maximum results (default from config) min_score: Minimum similarity score (default from config) sources: Sources to search (default from config) + readable_session_ids: Session IDs the caller may read Returns: List of search results @@ -418,6 +420,11 @@ async def search( query=query, max_results=limit * self.config.query.hybrid.candidate_multiplier, + readable_session_ids=( + readable_session_ids + if readable_session_ids is not None + else set() + ), ) results.extend( MemorySearchResult( diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index cc415b482..9279d3247 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -4,7 +4,7 @@ Bridges Session and MemoryManager for seamless memory access within sessions. """ -from typing import Optional, List, Dict, Any, Set +from typing import Optional, List, Dict, Any, Set, TYPE_CHECKING from pathlib import Path import asyncio @@ -13,6 +13,10 @@ from flocks.config import Config from flocks.utils.log import Log +if TYPE_CHECKING: + from flocks.auth.context import AuthUser + from flocks.session.session import SessionInfo + log = Log.create(service="session.memory") @@ -94,6 +98,93 @@ async def initialize(self) -> bool: "error": str(e), }) return False + + async def _resolve_search_caller( + self, + session: "SessionInfo", + ) -> Optional["AuthUser"]: + """Resolve the authenticated caller, falling back to Session owner.""" + from flocks.auth.context import ( + API_TOKEN_SERVICE_USER_ID, + AuthUser, + get_current_auth_user, + ) + + caller = get_current_auth_user() + if caller is not None: + return caller + + owner_id = getattr(session, "owner_user_id", None) + if not owner_id: + return None + if owner_id == API_TOKEN_SERVICE_USER_ID: + return AuthUser( + id=API_TOKEN_SERVICE_USER_ID, + username=API_TOKEN_SERVICE_USER_ID, + role="admin", + ) + + from flocks.auth.service import AuthService + + owner = await AuthService.get_user_by_id(owner_id) + if owner is None: + return None + to_auth_user = getattr(owner, "to_auth_user", None) + if callable(to_auth_user): + return to_auth_user() + return AuthUser( + id=str(owner.id), + username=str(owner.username), + role=str(owner.role), + status=str(getattr(owner, "status", "active")), + ) + + async def _search_access_context( + self, + ) -> tuple["SessionInfo", Optional["AuthUser"], Set[str]]: + """Validate the current Session and resolve its effective caller.""" + from flocks.project.project import Project + from flocks.session.policy import SessionPolicy + from flocks.session.session import Session + + session = await Session.get_by_id_unfiltered(self.session_id) + if session is None: + raise PermissionError("Session not found") + + caller = await self._resolve_search_caller(session) + shared_project_ids = Project.shared_project_ids() + if caller is not None and not SessionPolicy.can_read( + session, + caller, + shared_project_ids=shared_project_ids, + ): + raise PermissionError("Session access denied") + return session, caller, shared_project_ids + + async def _readable_session_ids( + self, + current_session: "SessionInfo", + caller: Optional["AuthUser"], + shared_project_ids: Set[str], + ) -> Set[str]: + """Return readable, non-deleted Session IDs in the current project.""" + if caller is None: + return {current_session.id} + + from flocks.session.policy import SessionPolicy + from flocks.session.session import Session + + return { + session.id + for session in await Session.list_all_unfiltered() + if session.project_id == self.project_id + and session.status != "deleted" + and SessionPolicy.can_read( + session, + caller, + shared_project_ids=shared_project_ids, + ) + } async def search( self, @@ -122,11 +213,33 @@ async def search( return [] try: - results = await self._manager.search( + manager = self._manager + if manager is None: + raise RuntimeError("Memory manager is not initialized") + current_session, caller, shared_project_ids = ( + await self._search_access_context() + ) + selected_sources = ( + list(sources) + if sources is not None + else [ + MemorySource(source) + for source in manager.config.sources + ] + ) + readable_session_ids = None + if MemorySource.SESSION in selected_sources: + readable_session_ids = await self._readable_session_ids( + current_session, + caller, + shared_project_ids, + ) + results = await manager.search( query=query, max_results=max_results, min_score=min_score, sources=sources, + readable_session_ids=readable_session_ids, ) log.debug("session.memory.search", { diff --git a/flocks/storage/session_search.py b/flocks/storage/session_search.py index 6668f9841..c9193d844 100644 --- a/flocks/storage/session_search.py +++ b/flocks/storage/session_search.py @@ -618,35 +618,43 @@ async def session_fts_search( project_id: str, query: str, max_results: int, + readable_session_ids: Optional[set[str]] = None, ) -> list[dict[str, Any]]: - """Search all indexed session messages using FTS5 BM25 ranking.""" + """Search readable Session messages in the current project.""" from flocks.storage.vector import build_fts_query require_session_search_available() - del project_id # Retained for API compatibility; Session search is global. fts_query = build_fts_query(query) if not fts_query: return [] + if readable_session_ids is not None and not readable_session_ids: + return [] + + sql = """ + SELECT + s.message_id, + s.session_id, + s.role, + s.created_at, + snippet(session_transcript_fts, 0, '', '', ' … ', 24), + bm25(session_transcript_fts) + FROM session_transcript_fts + JOIN session_transcript_index_state s + ON s.id = session_transcript_fts.rowid + WHERE session_transcript_fts MATCH ? + AND s.project_id = ? + """ + params: list[Any] = [fts_query, project_id] + if readable_session_ids is not None: + ordered_ids = sorted(readable_session_ids) + placeholders = ",".join("?" for _ in ordered_ids) + sql += f" AND s.session_id IN ({placeholders})" + params.extend(ordered_ids) + sql += " ORDER BY bm25(session_transcript_fts) LIMIT ?" + params.append(max_results) async with Storage.connect(db_path) as db: - cursor = await db.execute( - """ - SELECT - s.message_id, - s.session_id, - s.role, - s.created_at, - snippet(session_transcript_fts, 0, '', '', ' … ', 24), - bm25(session_transcript_fts) - FROM session_transcript_fts - JOIN session_transcript_index_state s - ON s.id = session_transcript_fts.rowid - WHERE session_transcript_fts MATCH ? - ORDER BY bm25(session_transcript_fts) - LIMIT ? - """, - (fts_query, max_results), - ) + cursor = await db.execute(sql, params) rows = await cursor.fetchall() count = len(rows) diff --git a/flocks/storage/vector.py b/flocks/storage/vector.py index 7344812a4..c10be74c0 100644 --- a/flocks/storage/vector.py +++ b/flocks/storage/vector.py @@ -206,8 +206,7 @@ async def vector_search( Args: db_path: Database path - project_id: Current Session project ID (retained for API compatibility; - Memory file search is global) + project_id: Current Session project ID embedding: Query embedding vector max_results: Maximum results to return min_score: Minimum similarity score @@ -217,7 +216,6 @@ async def vector_search( List of search results """ results = [] - del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -225,8 +223,12 @@ async def vector_search( SELECT id, path, source, start_line, end_line, text, embedding FROM memory_chunks WHERE embedding IS NOT NULL + AND ( + scope = 'global' + OR (scope = 'project' AND scope_id = ?) + ) """ - params: list[Any] = [] + params: list[Any] = [project_id] if sources: placeholders = ",".join("?" * len(sources)) @@ -319,8 +321,7 @@ async def fts_search( Args: db_path: Database path - project_id: Current Session project ID (retained for API compatibility; - Memory file search is global) + project_id: Current Session project ID query: Search query (FTS5 format) max_results: Maximum results to return sources: Optional list of sources to filter @@ -329,7 +330,6 @@ async def fts_search( List of search results with BM25 scores """ results = [] - del project_id # Memory file search is intentionally global across scopes. try: async with Storage.connect(db_path) as db: @@ -350,8 +350,12 @@ async def fts_search( rank FROM memory_fts f WHERE f.text MATCH ? + AND ( + f.scope = 'global' + OR (f.scope = 'project' AND f.scope_id = ?) + ) """ - params = [fts_query] + params = [fts_query, project_id] if sources: placeholders = ",".join("?" * len(sources)) diff --git a/flocks/tool/system/memory.py b/flocks/tool/system/memory.py index 8fff9e8b5..3d9b65a1e 100644 --- a/flocks/tool/system/memory.py +++ b/flocks/tool/system/memory.py @@ -67,8 +67,8 @@ def evict_session_memory(session_id: str) -> None: @ToolRegistry.register_function( name="memory_search", description=( - "Search persistent memory globally across Global, Daily, all Project " - "Memory files, and optional Session History sources." + "Search USER, Global, Daily, and current Project Memory, plus optional " + "readable Session History from the current project." ), category=ToolCategory.SEARCH, parameters=[ diff --git a/tests/memory/test_memory_scope.py b/tests/memory/test_memory_scope.py index 63d0c20a8..cf2633f6e 100644 --- a/tests/memory/test_memory_scope.py +++ b/tests/memory/test_memory_scope.py @@ -17,6 +17,7 @@ ensure_vector_tables, fts_search, replace_memory_file_index, + vector_search, ) @@ -40,6 +41,7 @@ def _chunk( scope_id: str, path: str, text: str, + embedding: list[float] | None = None, ) -> dict[str, object]: return { "id": f"chunk:{scope}:{scope_id}:{path}", @@ -51,9 +53,9 @@ def _chunk( "end_line": 1, "hash": f"hash:{text}", "text": text, - "embedding": None, - "embedding_model": None, - "embedding_dims": None, + "embedding": embedding, + "embedding_model": "test" if embedding else None, + "embedding_dims": len(embedding) if embedding else None, } @@ -78,11 +80,15 @@ async def test_search_reconciles_filesystem_before_every_search( @pytest.mark.asyncio -async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: +async def test_memory_search_uses_global_and_current_project_scopes( + tmp_path: Path, +) -> None: db_path = tmp_path / "scope.db" await Storage.init(db_path) records = [ + ("global", "", "USER.md", "scopeword user"), ("global", "", "MEMORY.md", "scopeword global"), + ("global", "", "daily/2026-08-03.md", "scopeword daily"), ( "project", "prj_alpha", @@ -100,19 +106,35 @@ async def test_memory_search_is_global_across_scopes(tmp_path: Path) -> None: await replace_memory_file_index( db_path, file_entry=_file_entry(scope, scope_id, path), - chunks=[_chunk(scope, scope_id, path, text)], + chunks=[_chunk(scope, scope_id, path, text, [1.0, 0.0])], ) - alpha = await fts_search(db_path, "prj_alpha", "scopeword") - default = await fts_search(db_path, "default", "scopeword") - - expected_paths = { + expected_global_paths = { + "USER.md", "MEMORY.md", + "daily/2026-08-03.md", + } + expected_alpha_paths = expected_global_paths | { "projects/prj_alpha/MEMORY.md", - "projects/prj_beta/MEMORY.md", } - assert {result["path"] for result in alpha} == expected_paths - assert {result["path"] for result in default} == expected_paths + + alpha_fts = await fts_search(db_path, "prj_alpha", "scopeword") + default_fts = await fts_search(db_path, "default", "scopeword") + alpha_vector = await vector_search( + db_path, + "prj_alpha", + [1.0, 0.0], + ) + default_vector = await vector_search( + db_path, + "default", + [1.0, 0.0], + ) + + assert {result["path"] for result in alpha_fts} == expected_alpha_paths + assert {result["path"] for result in alpha_vector} == expected_alpha_paths + assert {result["path"] for result in default_fts} == expected_global_paths + assert {result["path"] for result in default_vector} == expected_global_paths @pytest.mark.asyncio diff --git a/tests/memory/test_session_transcript_search.py b/tests/memory/test_session_transcript_search.py index aa2c05c51..70de19940 100644 --- a/tests/memory/test_session_transcript_search.py +++ b/tests/memory/test_session_transcript_search.py @@ -7,6 +7,7 @@ import pytest +from flocks.auth.context import AuthUser, reset_current_auth_user, set_current_auth_user from flocks.config.config import Config from flocks.memory.config import MemoryConfig from flocks.memory.manager import MemoryManager @@ -14,6 +15,7 @@ from flocks.memory.types import MemorySearchResult from flocks.memory.types import MemorySource from flocks.provider import Provider +from flocks.session.features.memory import SessionMemory from flocks.session.message import Message, MessageRole from flocks.session.session import Session, SessionInfo from flocks.storage.session_search import ( @@ -60,13 +62,25 @@ async def isolate_transcript_search( Storage._db_path = None -async def _create_session(tmp_path: Path, project_id: str = "project-search"): +async def _create_session( + tmp_path: Path, + project_id: str = "project-search", + *, + owner_user_id: str | None = None, + owner_username: str | None = None, + metadata: dict | None = None, + status: str = "active", +): session = SessionInfo( id=f"session-{uuid.uuid4().hex}", project_id=project_id, directory=str(tmp_path), agent="rex", memory_enabled=True, + owner_user_id=owner_user_id, + owner_username=owner_username, + metadata=metadata or {}, + status=status, ) await Storage.set( f"session:{project_id}:{session.id}", @@ -267,7 +281,7 @@ async def test_text_part_updates_and_message_delete_update_fts( @pytest.mark.asyncio -async def test_session_search_is_global_across_projects( +async def test_session_search_is_limited_to_current_project( tmp_path: Path, ) -> None: alpha = await _create_session(tmp_path, project_id="prj_alpha") @@ -292,7 +306,6 @@ async def test_session_search_is_global_across_projects( assert {result["path"] for result in results} == { f"sessions/{alpha.id}/messages/{alpha_message.id}", - f"sessions/{beta.id}/messages/{beta_message.id}", } async with Storage.connect(Storage.get_db_path()) as db: @@ -309,10 +322,299 @@ async def test_session_search_is_global_across_projects( assert stats["updated"] == 2 assert {result["path"] for result in rebuilt} == { f"sessions/{alpha.id}/messages/{alpha_message.id}", - f"sessions/{beta.id}/messages/{beta_message.id}", } +@pytest.mark.asyncio +async def test_session_search_filters_readable_ids_within_project( + tmp_path: Path, +) -> None: + readable = await _create_session(tmp_path, project_id="prj_alpha") + private = await _create_session(tmp_path, project_id="prj_alpha") + readable_message = await Message.create( + readable.id, + MessageRole.USER, + "same project permission marker readable", + ) + await Message.create( + private.id, + MessageRole.USER, + "same project permission marker private", + ) + + results = await session_fts_search( + db_path=Storage.get_db_path(), + project_id="prj_alpha", + query="same project permission marker", + max_results=10, + readable_session_ids={readable.id}, + ) + + assert [result["path"] for result in results] == [ + f"sessions/{readable.id}/messages/{readable_message.id}" + ] + assert not await session_fts_search( + db_path=Storage.get_db_path(), + project_id="prj_alpha", + query="same project permission marker", + max_results=10, + readable_session_ids=set(), + ) + + +@pytest.mark.asyncio +async def test_session_memory_uses_session_read_policy( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + caller = AuthUser(id="user-a", username="alice", role="member") + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + ) + owned = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + ) + archived = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + status="archived", + ) + private = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + shared = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + metadata={"shared_read_access_user_ids": [caller.id]}, + ) + deleted = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=caller.id, + owner_username=caller.username, + status="deleted", + ) + other_project = await _create_session( + tmp_path, + project_id="prj_beta", + owner_user_id=caller.id, + owner_username=caller.username, + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + resolved_session, resolved_caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + resolved_caller, + shared_projects, + ) + finally: + reset_current_auth_user(token) + + assert readable_ids == {current.id, owned.id, archived.id, shared.id} + assert private.id not in readable_ids + assert deleted.id not in readable_ids + assert other_project.id not in readable_ids + + +@pytest.mark.asyncio +async def test_session_memory_honors_shared_project_access( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + caller = AuthUser(id="user-a", username="alice", role="member") + current = await _create_session( + tmp_path, + project_id="prj_shared", + owner_user_id="user-b", + owner_username="bob", + ) + sibling = await _create_session( + tmp_path, + project_id="prj_shared", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr( + Project, + "shared_project_ids", + lambda: {"prj_shared"}, + ) + + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + resolved_session, resolved_caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + resolved_caller, + shared_projects, + ) + finally: + reset_current_auth_user(token) + + assert readable_ids == {current.id, sibling.id} + + +@pytest.mark.asyncio +async def test_session_memory_without_caller_only_reads_current_session( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.auth.service import AuthService + from flocks.project.project import Project + + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="missing-user", + ) + await _create_session(tmp_path, project_id="prj_alpha") + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + monkeypatch.setattr( + AuthService, + "get_user_by_id", + AsyncMock(return_value=None), + ) + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + + resolved_session, caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + caller, + shared_projects, + ) + + assert caller is None + assert readable_ids == {current.id} + + +@pytest.mark.asyncio +async def test_session_memory_falls_back_to_session_owner( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.auth.service import AuthService + from flocks.project.project import Project + + owner_auth = AuthUser(id="user-a", username="alice", role="member") + owner = Mock() + owner.to_auth_user.return_value = owner_auth + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=owner_auth.id, + owner_username=owner_auth.username, + ) + sibling = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id=owner_auth.id, + owner_username=owner_auth.username, + ) + private = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + monkeypatch.setattr( + AuthService, + "get_user_by_id", + AsyncMock(return_value=owner), + ) + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + + resolved_session, caller, shared_projects = ( + await memory._search_access_context() + ) + readable_ids = await memory._readable_session_ids( + resolved_session, + caller, + shared_projects, + ) + + assert caller == owner_auth + assert readable_ids == {current.id, sibling.id} + assert private.id not in readable_ids + + +@pytest.mark.asyncio +async def test_session_memory_rejects_unreadable_current_session( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.project.project import Project + + current = await _create_session( + tmp_path, + project_id="prj_alpha", + owner_user_id="user-b", + owner_username="bob", + ) + monkeypatch.setattr(Project, "shared_project_ids", lambda: set()) + caller = AuthUser(id="user-a", username="alice", role="member") + token = set_current_auth_user(caller) + try: + memory = SessionMemory( + session_id=current.id, + project_id=current.project_id, + workspace_dir=str(tmp_path), + enabled=True, + ) + with pytest.raises(PermissionError, match="Session access denied"): + await memory._search_access_context() + finally: + reset_current_auth_user(token) + + @pytest.mark.asyncio async def test_reconciliation_restores_history_and_removes_orphans( tmp_path: Path, @@ -499,6 +801,7 @@ async def test_explicit_session_search_persists_opt_in_without_embeddings( results = await manager.search( query="marker", sources=[MemorySource.SESSION], + readable_session_ids={session.id}, ) assert results From 1e0864c3c9a470132d968fa1ad446111be8abb9d Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 6 Aug 2026 10:58:53 +0800 Subject: [PATCH 3/3] refactor(hooks): remove unused slug generator --- flocks/hooks/builtin/slug_generator.py | 94 -------------------------- 1 file changed, 94 deletions(-) delete mode 100644 flocks/hooks/builtin/slug_generator.py diff --git a/flocks/hooks/builtin/slug_generator.py b/flocks/hooks/builtin/slug_generator.py deleted file mode 100644 index 2bbc2073a..000000000 --- a/flocks/hooks/builtin/slug_generator.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -LLM Slug Generator - Generate descriptive filenames using LLM - -Uses LLM to generate a 1-2 word slug for session memory filenames. -""" - -from typing import Optional -import re - -from flocks.provider import Provider -from flocks.utils.log import Log - -log = Log.create(service="hooks.slug_generator") - - -async def generate_slug_via_llm( - conversation: str, - config: any, - session_id: str, - timeout_seconds: int = 15, -) -> Optional[str]: - """ - Generate a slug using LLM - - Args: - conversation: Conversation summary - config: Configuration object - session_id: Session ID (for logging) - timeout_seconds: Timeout in seconds - - Returns: - slug string or None (on failure) - - Examples: - >>> await generate_slug_via_llm("user: Design API\\nassistant: Sure...") - "api-design" - """ - try: - # Construct prompt - prompt = f"""Based on this conversation, generate a short 1-2 word filename slug (lowercase, hyphen-separated, no file extension). - -Conversation summary: -{conversation[:2000]} - -Reply with ONLY the slug, nothing else. Examples: "vendor-pitch", "api-design", "bug-fix" -""" - - # Get provider configuration - provider_id = getattr(config.memory.search.embedding, 'provider', 'openai') - if provider_id == "auto": - provider_id = "openai" - - # Call LLM (use lightweight model) - response = await Provider.chat( - messages=[{"role": "user", "content": prompt}], - provider_id=provider_id, - model="gpt-3.5-turbo", # Fast lightweight model - max_tokens=50, - temperature=0.7, - ) - - # Extract and clean slug - if response and response.get('content'): - text = response['content'].strip() - - # Clean format - slug = text.lower().replace(" ", "-").replace("_", "-") - - # Remove invalid characters - slug = re.sub(r'[^a-z0-9-]', '', slug) - slug = re.sub(r'-+', '-', slug) - slug = slug.strip('-') - - # Limit length - slug = slug[:30] - - if slug: - log.debug("slug_generator.success", { - "session_id": session_id, - "slug": slug, - }) - return slug - - log.warn("slug_generator.no_result", { - "session_id": session_id, - }) - return None - - except Exception as e: - log.error("slug_generator.error", { - "session_id": session_id, - "error": str(e), - }) - return None