From 24b7f200f1f65ec38386dc5fa5b975e7aa9d435c Mon Sep 17 00:00:00 2001
From: Markus <markusleben@users.noreply.github.com>
Date: Thu, 9 Apr 2026 12:23:33 +0200
Subject: [PATCH] fix: improve session_search lineage recall

- merge matched child/root hits within the same session lineage
- anchor summaries with matched snippets and better fallback previews
- add regression tests for lineage recall, metadata, and truncation
---
 tests/tools/test_session_search.py | 337 +++++++++++++++++++++++++++++
 tools/session_search_tool.py       | 140 +++++++++---
 2 files changed, 452 insertions(+), 25 deletions(-)

diff --git a/tests/tools/test_session_search.py b/tests/tools/test_session_search.py
index acb64d62fbb42..0b3dee3d0ed2b 100644
--- a/tests/tools/test_session_search.py
+++ b/tests/tools/test_session_search.py
@@ -146,6 +146,15 @@ def test_match_at_beginning(self):
         result = _truncate_around_matches(text, "KEYWORD")
         assert "KEYWORD" in result
 
+    def test_quoted_phrase_prefers_exact_match_over_common_terms(self):
+        text = (
+            ("commonterm " * 20000)
+            + "SPECIAL UNIQUE TARGET PHRASE "
+            + ("tail " * 20000)
+        )
+        result = _truncate_around_matches(text, '"SPECIAL UNIQUE TARGET PHRASE" commonterm')
+        assert "SPECIAL UNIQUE TARGET PHRASE" in result
+
 
 # =========================================================================
 # session_search (dispatcher)
@@ -284,3 +293,331 @@ def _get_session(session_id):
         assert result["count"] == 0
         assert result["results"] == []
         assert result["sessions_searched"] == 0
+
+    def test_child_match_uses_lineage_conversation_in_fallback_preview(self):
+        """Child hits should preview a transcript that includes the matched child content."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        mock_db = MagicMock()
+        query = "needle phrase"
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "child_sid",
+                "content": query,
+                "source": "telegram",
+                "session_started": 1709400000,
+                "model": "test",
+            },
+        ]
+
+        def _get_session(session_id):
+            if session_id == "child_sid":
+                return {"parent_session_id": "mid_sid", "source": "telegram", "started_at": 1709400000}
+            if session_id == "mid_sid":
+                return {"parent_session_id": "root_sid", "source": "telegram", "started_at": 1709300000}
+            if session_id == "root_sid":
+                return {"parent_session_id": None, "source": "telegram", "started_at": 1709200000}
+            return None
+
+        def _get_messages(session_id):
+            if session_id == "root_sid":
+                return [
+                    {"role": "user", "content": "root question"},
+                    {"role": "assistant", "content": "root answer"},
+                ]
+            if session_id == "mid_sid":
+                return [
+                    {"role": "assistant", "content": "middle context"},
+                ]
+            if session_id == "child_sid":
+                return [
+                    {"role": "assistant", "content": f"here is the {query} we need"},
+                ]
+            return []
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        # Force the raw preview fallback so we can inspect the exact prepared transcript.
+        from unittest.mock import AsyncMock, patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm",
+                     new_callable=AsyncMock,
+                     side_effect=RuntimeError("no provider")):
+            result = json.loads(session_search(query=query, db=mock_db, limit=1))
+
+        assert result["success"] is True
+        assert result["count"] == 1
+        assert result["results"][0]["session_id"] == "root_sid"
+        assert query in result["results"][0]["summary"]
+        # Root + lineage child should both be included in the prepared preview.
+        assert "root answer" in result["results"][0]["summary"]
+        assert "middle context" in result["results"][0]["summary"]
+
+    def test_child_match_keeps_root_identity_but_uses_root_metadata(self):
+        """Lineage-grouped results should not mix root session_id with child metadata."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        mock_db = MagicMock()
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "child_sid",
+                "content": "needle",
+                "source": "telegram",
+                "session_started": 1709400000,
+                "model": "child-model",
+            },
+        ]
+
+        def _get_session(session_id):
+            if session_id == "child_sid":
+                return {"parent_session_id": "root_sid", "source": "telegram", "started_at": 1709400000, "model": "child-model"}
+            if session_id == "root_sid":
+                return {"parent_session_id": None, "source": "cli", "started_at": 1709200000, "model": "root-model"}
+            return None
+
+        def _get_messages(session_id):
+            if session_id == "root_sid":
+                return [{"role": "assistant", "content": "root"}]
+            if session_id == "child_sid":
+                return [{"role": "assistant", "content": "needle child detail"}]
+            return []
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        from unittest.mock import AsyncMock, patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm",
+                     new_callable=AsyncMock,
+                     side_effect=RuntimeError("no provider")):
+            result = json.loads(session_search(query="needle", db=mock_db, limit=1))
+
+        assert result["success"] is True
+        assert result["count"] == 1
+        entry = result["results"][0]
+        assert entry["session_id"] == "root_sid"
+        assert entry["source"] == "cli"
+        assert entry["model"] == "root-model"
+        assert entry["when"] == _format_timestamp(1709200000)
+
+    def test_long_lineage_fallback_preview_is_centered_on_match(self):
+        """Fallback preview should still show the matched child phrase for long lineages."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        query = 'SPECIAL UNIQUE TARGET PHRASE'
+        mock_db = MagicMock()
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "child_sid",
+                "content": query,
+                "source": "telegram",
+                "session_started": 1709400000,
+                "model": "test",
+            },
+        ]
+
+        def _get_session(session_id):
+            if session_id == "child_sid":
+                return {"parent_session_id": "root_sid", "source": "telegram", "started_at": 1709400000}
+            if session_id == "root_sid":
+                return {"parent_session_id": None, "source": "telegram", "started_at": 1709200000}
+            return None
+
+        def _get_messages(session_id):
+            if session_id == "root_sid":
+                return [{"role": "assistant", "content": "commonterm " * 20000}]
+            if session_id == "child_sid":
+                return [{"role": "assistant", "content": f"here is the {query} we need"}]
+            return []
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        from unittest.mock import AsyncMock, patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm",
+                     new_callable=AsyncMock,
+                     side_effect=RuntimeError("no provider")):
+            result = json.loads(session_search(query=f'"{query}" commonterm', db=mock_db, limit=1))
+
+        assert result["success"] is True
+        assert result["count"] == 1
+        assert query in result["results"][0]["summary"]
+
+    def test_multiple_results_keep_per_session_metadata(self):
+        """Each result should use its own root session metadata, not the last task's metadata."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        mock_db = MagicMock()
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "child_a",
+                "content": "alpha",
+                "source": "telegram",
+                "session_started": 1709400000,
+                "model": "child-a-model",
+            },
+            {
+                "session_id": "child_b",
+                "content": "beta",
+                "source": "telegram",
+                "session_started": 1709500000,
+                "model": "child-b-model",
+            },
+        ]
+
+        def _get_session(session_id):
+            mapping = {
+                "child_a": {"parent_session_id": "root_a", "source": "telegram", "started_at": 1709400000, "model": "child-a-model"},
+                "root_a": {"parent_session_id": None, "source": "cli", "started_at": 1709200000, "model": "root-a-model"},
+                "child_b": {"parent_session_id": "root_b", "source": "telegram", "started_at": 1709500000, "model": "child-b-model"},
+                "root_b": {"parent_session_id": None, "source": "discord", "started_at": 1709300000, "model": "root-b-model"},
+            }
+            return mapping.get(session_id)
+
+        def _get_messages(session_id):
+            mapping = {
+                "root_a": [{"role": "assistant", "content": "root a"}],
+                "child_a": [{"role": "assistant", "content": "alpha detail"}],
+                "root_b": [{"role": "assistant", "content": "root b"}],
+                "child_b": [{"role": "assistant", "content": "beta detail"}],
+            }
+            return mapping.get(session_id, [])
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        from unittest.mock import AsyncMock, patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm",
+                     new_callable=AsyncMock,
+                     side_effect=RuntimeError("no provider")):
+            result = json.loads(session_search(query="alpha OR beta", db=mock_db, limit=2))
+
+        assert result["success"] is True
+        assert result["count"] == 2
+        entry_a, entry_b = result["results"]
+        assert entry_a["session_id"] == "root_a"
+        assert entry_a["source"] == "cli"
+        assert entry_a["model"] == "root-a-model"
+        assert entry_a["when"] == _format_timestamp(1709200000)
+        assert entry_b["session_id"] == "root_b"
+        assert entry_b["source"] == "discord"
+        assert entry_b["model"] == "root-b-model"
+        assert entry_b["when"] == _format_timestamp(1709300000)
+
+    def test_same_root_keeps_child_hit_content_even_if_root_hit_ranks_first(self):
+        """Broad queries should include later continuation content for the same root, not only the first root hit."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        mock_db = MagicMock()
+        query = 'ADHS OR Notion'
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "root_sid",
+                "content": "ADHS root mention",
+                "snippet": "ADHS root mention",
+                "source": "telegram",
+                "session_started": 1709200000,
+                "model": "root-model",
+            },
+            {
+                "session_id": "child_sid",
+                "content": "Notion child detail",
+                "snippet": "Notion child detail",
+                "source": "telegram",
+                "session_started": 1709300000,
+                "model": "child-model",
+            },
+        ]
+
+        def _get_session(session_id):
+            if session_id == "root_sid":
+                return {"parent_session_id": None, "source": "telegram", "started_at": 1709200000, "model": "root-model"}
+            if session_id == "child_sid":
+                return {"parent_session_id": "root_sid", "source": "telegram", "started_at": 1709300000, "model": "child-model"}
+            return None
+
+        def _get_messages(session_id):
+            if session_id == "root_sid":
+                return [{"role": "assistant", "content": "ADHS root mention and setup"}]
+            if session_id == "child_sid":
+                return [{"role": "assistant", "content": "Notion child detail with Weekly Actions and Video-Database"}]
+            return []
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        from unittest.mock import AsyncMock, patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm",
+                     new_callable=AsyncMock,
+                     side_effect=RuntimeError("no provider")):
+            result = json.loads(session_search(query=query, db=mock_db, limit=1))
+
+        assert result["success"] is True
+        assert result["count"] == 1
+        summary = result["results"][0]["summary"]
+        assert "ADHS root mention" in summary
+        assert "Notion child detail" in summary
+        assert "Weekly Actions" in summary
+
+    def test_same_root_summary_prompt_includes_all_matched_snippets(self):
+        """Summarizer prompts should receive matched snippets from both root and child hits."""
+        from unittest.mock import MagicMock
+        from tools.session_search_tool import session_search
+
+        mock_db = MagicMock()
+        query = 'ADHS OR Notion'
+        mock_db.search_messages.return_value = [
+            {
+                "session_id": "root_sid",
+                "content": "ADHS root mention",
+                "snippet": "ADHS root mention",
+                "source": "telegram",
+                "session_started": 1709200000,
+                "model": "root-model",
+            },
+            {
+                "session_id": "child_sid",
+                "content": "Notion child detail",
+                "snippet": "Notion child detail",
+                "source": "telegram",
+                "session_started": 1709300000,
+                "model": "child-model",
+            },
+        ]
+
+        def _get_session(session_id):
+            if session_id == "root_sid":
+                return {"parent_session_id": None, "source": "telegram", "started_at": 1709200000, "model": "root-model"}
+            if session_id == "child_sid":
+                return {"parent_session_id": "root_sid", "source": "telegram", "started_at": 1709300000, "model": "child-model"}
+            return None
+
+        def _get_messages(session_id):
+            if session_id == "root_sid":
+                return [{"role": "assistant", "content": "ADHS root mention and setup"}]
+            if session_id == "child_sid":
+                return [{"role": "assistant", "content": "Notion child detail with Weekly Actions and Video-Database"}]
+            return []
+
+        mock_db.get_session.side_effect = _get_session
+        mock_db.get_messages_as_conversation.side_effect = _get_messages
+
+        captured = {}
+
+        async def _fake_async_call_llm(*args, **kwargs):
+            captured['messages'] = kwargs.get('messages')
+            return object()
+
+        from unittest.mock import patch as _patch
+        with _patch("tools.session_search_tool.async_call_llm", side_effect=_fake_async_call_llm), \
+             _patch("tools.session_search_tool.extract_content_or_reasoning", return_value="summary ok"):
+            result = json.loads(session_search(query=query, db=mock_db, limit=1))
+
+        assert result["success"] is True
+        user_prompt = captured['messages'][1]['content']
+        assert 'ADHS root mention' in user_prompt
+        assert 'Notion child detail' in user_prompt
diff --git a/tools/session_search_tool.py b/tools/session_search_tool.py
index 3e9c68af40e2d..96e32d51df8c9 100644
--- a/tools/session_search_tool.py
+++ b/tools/session_search_tool.py
@@ -19,6 +19,7 @@
 import concurrent.futures
 import json
 import logging
+import re
 from typing import Dict, Any, List, Optional, Union
 
 from agent.auxiliary_client import async_call_llm, extract_content_or_reasoning
@@ -96,15 +97,30 @@ def _truncate_around_matches(
     if len(full_text) <= max_chars:
         return full_text
 
-    # Find the first occurrence of any query term
-    query_terms = query.lower().split()
     text_lower = full_text.lower()
+
+    # Prefer exact quoted phrases when present. Falling back to raw whitespace
+    # splitting can anchor on common words like "wurde" long before the actual
+    # phrase match, which makes previews start at unrelated earlier content.
+    quoted_phrases = [p.strip().lower() for p in re.findall(r'"([^"]+)"', query) if p.strip()]
+    remaining_query = re.sub(r'"[^"]+"', ' ', query.lower())
+    query_terms = [
+        token for token in re.findall(r"\b[\w.-]+\b", remaining_query)
+        if token not in {"and", "or", "not"}
+    ]
+
     first_match = len(full_text)
-    for term in query_terms:
-        pos = text_lower.find(term)
+    for phrase in quoted_phrases:
+        pos = text_lower.find(phrase)
         if pos != -1 and pos < first_match:
             first_match = pos
 
+    if first_match == len(full_text):
+        for term in query_terms:
+            pos = text_lower.find(term)
+            if pos != -1 and pos < first_match:
+                first_match = pos
+
     if first_match == len(full_text):
         # No match found, take from the start
         first_match = 0
@@ -123,7 +139,10 @@ def _truncate_around_matches(
 
 
 async def _summarize_session(
-    conversation_text: str, query: str, session_meta: Dict[str, Any]
+    conversation_text: str,
+    query: str,
+    session_meta: Dict[str, Any],
+    matched_snippets: List[str] = None,
 ) -> Optional[str]:
     """Summarize a single session conversation focused on the search query."""
     system_prompt = (
@@ -134,17 +153,21 @@ async def _summarize_session(
         "3. Key decisions, solutions found, or conclusions reached\n"
         "4. Any specific commands, files, URLs, or technical details that were important\n"
         "5. Anything left unresolved or notable\n\n"
+        "Important: matched excerpts are high-confidence evidence from search hits. "
+        "If they mention a topic, do not claim that topic was absent. Use them to anchor the summary, then use the transcript for surrounding context. "
         "Be thorough but concise. Preserve specific details (commands, paths, error messages) "
         "that would be useful to recall. Write in past tense as a factual recap."
     )
 
     source = session_meta.get("source", "unknown")
     started = _format_timestamp(session_meta.get("started_at"))
+    matched_block = "\n".join(f"- {snippet}" for snippet in (matched_snippets or []) if snippet)
 
     user_prompt = (
         f"Search topic: {query}\n"
         f"Session source: {source}\n"
         f"Session date: {started}\n\n"
+        f"MATCHED EXCERPTS:\n{matched_block or '(none)'}\n\n"
         f"CONVERSATION TRANSCRIPT:\n{conversation_text}\n\n"
         f"Summarize this conversation with focus on: {query}"
     )
@@ -293,14 +316,17 @@ def session_search(
                 "message": "No matching sessions found.",
             }, ensure_ascii=False)
 
-        # Resolve child sessions to their parent — delegation stores detailed
-        # content in child sessions, but the user's conversation is the parent.
-        def _resolve_to_parent(session_id: str) -> str:
-            """Walk delegation chain to find the root parent session ID."""
+        # Resolve child sessions to their parent/root for dedup + display,
+        # but keep the raw hit session so we can prepare a transcript that
+        # actually contains the matched content.
+        def _lineage_to_root(session_id: str) -> List[str]:
+            """Return [session, parent, grandparent, ..., root] with cycle protection."""
             visited = set()
+            lineage = []
             sid = session_id
             while sid and sid not in visited:
                 visited.add(sid)
+                lineage.append(sid)
                 try:
                     session = db.get_session(sid)
                     if not session:
@@ -318,7 +344,52 @@ def _resolve_to_parent(session_id: str) -> str:
                         exc_info=True,
                     )
                     break
-            return sid
+            return lineage
+
+        def _resolve_to_parent(session_id: str) -> str:
+            """Walk continuation/delegation chain to find the root parent session ID."""
+            lineage = _lineage_to_root(session_id)
+            return lineage[-1] if lineage else session_id
+
+        def _load_lineage_conversation(hit_session_id: str) -> List[Dict[str, Any]]:
+            """Load root→...→hit conversation so child-only matches stay visible."""
+            lineage = list(reversed(_lineage_to_root(hit_session_id)))
+            messages: List[Dict[str, Any]] = []
+            for sid in lineage:
+                try:
+                    messages.extend(db.get_messages_as_conversation(sid) or [])
+                except Exception as e:
+                    logging.debug(
+                        "Failed to load lineage segment %s: %s",
+                        sid,
+                        e,
+                        exc_info=True,
+                    )
+            return messages
+
+        def _load_lineage_conversation_for_hits(hit_session_ids: List[str]) -> List[Dict[str, Any]]:
+            """Merge root→hit paths for all matched sessions in the same lineage root."""
+            ordered_session_ids: List[str] = []
+            seen_lineage_ids = set()
+            for hit_session_id in hit_session_ids:
+                lineage = list(reversed(_lineage_to_root(hit_session_id)))
+                for sid in lineage:
+                    if sid not in seen_lineage_ids:
+                        seen_lineage_ids.add(sid)
+                        ordered_session_ids.append(sid)
+
+            messages: List[Dict[str, Any]] = []
+            for sid in ordered_session_ids:
+                try:
+                    messages.extend(db.get_messages_as_conversation(sid) or [])
+                except Exception as e:
+                    logging.debug(
+                        "Failed to load merged lineage segment %s: %s",
+                        sid,
+                        e,
+                        exc_info=True,
+                    )
+            return messages
 
         current_lineage_root = (
             _resolve_to_parent(current_session_id) if current_session_id else None
@@ -338,23 +409,36 @@ def _resolve_to_parent(session_id: str) -> str:
             if current_session_id and raw_sid == current_session_id:
                 continue
             if resolved_sid not in seen_sessions:
-                result = dict(result)
-                result["session_id"] = resolved_sid
-                seen_sessions[resolved_sid] = result
-            if len(seen_sessions) >= limit:
-                break
+                if len(seen_sessions) >= limit:
+                    # Ignore additional roots beyond the requested limit, but
+                    # continue collecting more hit sessions for already-selected roots.
+                    continue
+                seen_sessions[resolved_sid] = {
+                    "display_session_id": resolved_sid,
+                    "hit_session_ids": [],
+                    "matched_snippets": [],
+                    "match_info": dict(result),
+                }
+            if raw_sid not in seen_sessions[resolved_sid]["hit_session_ids"]:
+                seen_sessions[resolved_sid]["hit_session_ids"].append(raw_sid)
+            snippet = result.get("snippet") or result.get("content")
+            if snippet and snippet not in seen_sessions[resolved_sid]["matched_snippets"]:
+                seen_sessions[resolved_sid]["matched_snippets"].append(snippet)
 
         # Prepare all sessions for parallel summarization
         tasks = []
-        for session_id, match_info in seen_sessions.items():
+        for session_id, session_info in seen_sessions.items():
             try:
-                messages = db.get_messages_as_conversation(session_id)
+                hit_session_ids = session_info["hit_session_ids"]
+                match_info = session_info["match_info"]
+                matched_snippets = session_info["matched_snippets"]
+                messages = _load_lineage_conversation_for_hits(hit_session_ids)
                 if not messages:
                     continue
                 session_meta = db.get_session(session_id) or {}
                 conversation_text = _format_conversation(messages)
                 conversation_text = _truncate_around_matches(conversation_text, query)
-                tasks.append((session_id, match_info, conversation_text, session_meta))
+                tasks.append((session_id, match_info, conversation_text, session_meta, matched_snippets))
             except Exception as e:
                 logging.warning(
                     "Failed to prepare session %s: %s",
@@ -367,8 +451,8 @@ def _resolve_to_parent(session_id: str) -> str:
         async def _summarize_all() -> List[Union[str, Exception]]:
             """Summarize all sessions in parallel."""
             coros = [
-                _summarize_session(text, query, meta)
-                for _, _, text, meta in tasks
+                _summarize_session(text, query, meta, snippets)
+                for _, _, text, meta, snippets in tasks
             ]
             return await asyncio.gather(*coros, return_exceptions=True)
 
@@ -392,7 +476,7 @@ async def _summarize_all() -> List[Union[str, Exception]]:
             }, ensure_ascii=False)
 
         summaries = []
-        for (session_id, match_info, conversation_text, _), result in zip(tasks, results):
+        for (session_id, match_info, conversation_text, session_meta, _matched_snippets), result in zip(tasks, results):
             if isinstance(result, Exception):
                 logging.warning(
                     "Failed to summarize session %s: %s",
@@ -402,9 +486,9 @@ async def _summarize_all() -> List[Union[str, Exception]]:
 
             entry = {
                 "session_id": session_id,
-                "when": _format_timestamp(match_info.get("session_started")),
-                "source": match_info.get("source", "unknown"),
-                "model": match_info.get("model"),
+                "when": _format_timestamp(session_meta.get("started_at") or match_info.get("session_started")),
+                "source": session_meta.get("source") or match_info.get("source", "unknown"),
+                "model": session_meta.get("model") or match_info.get("model"),
             }
 
             if result:
@@ -412,7 +496,13 @@ async def _summarize_all() -> List[Union[str, Exception]]:
             else:
                 # Fallback: raw preview so matched sessions aren't silently
                 # dropped when the summarizer is unavailable (fixes #3409).
-                preview = (conversation_text[:500] + "\n…[truncated]") if conversation_text else "No preview available."
+                # Re-center a short preview around the match; taking the first
+                # 500 chars of a long lineage transcript can hide the very hit
+                # that justified returning this session.
+                preview = (
+                    _truncate_around_matches(conversation_text, query, max_chars=500)
+                    if conversation_text else "No preview available."
+                )
                 entry["summary"] = f"[Raw preview — summarization unavailable]\n{preview}"
 
             summaries.append(entry)