fix: flush un-retained turns in SessionEnd hook before daemon stop

hindsight-integrations/claude-code/scripts/session_end.py

@@ -1,20 +1,161 @@
 #!/usr/bin/env python3
-"""SessionEnd hook: daemon cleanup.
+"""SessionEnd hook: final retain flush + daemon cleanup.
 
-Fires once when a Claude Code session terminates. If the plugin
-auto-started a hindsight-embed daemon, this is where we stop it.
+Fires once when a Claude Code session terminates. Two responsibilities:
+1. Flush any un-retained turns that were skipped by retainEveryNTurns throttle.
+2. Stop the auto-started hindsight-embed daemon (if any).
 
 Port of: Openclaw's service.stop() in index.js
 """
 
 import json
 import os
 import sys
+import time
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
+from lib.bank import derive_bank_id, ensure_bank_mission
+from lib.client import HindsightClient
 from lib.config import debug_log, load_config
-from lib.daemon import stop_daemon
+from lib.content import prepare_retention_transcript
+from lib.daemon import get_api_url, stop_daemon
+from lib.state import get_turn_count
+
+
+def _read_transcript(transcript_path: str) -> list:
+    """Read a JSONL transcript file and return list of message dicts."""
+    if not transcript_path or not os.path.isfile(transcript_path):
+        return []
+    messages = []
+    try:
+        with open(transcript_path, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    entry = json.loads(line)
+                    if entry.get("type") in ("user", "assistant"):
+                        msg = entry.get("message", {})
+                        if isinstance(msg, dict) and msg.get("role"):
+                            messages.append(msg)
+                    elif "role" in entry and "content" in entry:
+                        messages.append(entry)
+                except json.JSONDecodeError:
+                    continue
+    except OSError:
+        pass
+    return messages
+
+
+def _flush_remaining_turns(config: dict, hook_input: dict) -> None:
+    """Retain the full session transcript, bypassing retainEveryNTurns throttle.
+
+    Called by SessionEnd to ensure turns between the last throttled retain
+    and session end are not silently dropped.
+    """
+    session_id = hook_input.get("session_id", "unknown")
+    transcript_path = hook_input.get("transcript_path", "")
+
+    retain_every_n = max(1, config.get("retainEveryNTurns", 1))
+
+    # If retainEveryNTurns == 1, the Stop hook retains every turn — no flush needed.
+    # Also skip if turn_count is already a multiple (last Stop hook already flushed).
+    if retain_every_n <= 1:
+        debug_log(config, "SessionEnd flush: retainEveryNTurns=1, Stop hook covers all turns — skipping")
+        return
+    turn_count = get_turn_count(session_id)
+    if turn_count % retain_every_n == 0:
+        debug_log(config, f"SessionEnd flush: turn {turn_count} is a multiple of {retain_every_n} — already retained by Stop hook")
+        return
+
+    debug_log(config, f"SessionEnd flush: turn {turn_count} not a multiple of {retain_every_n} — flushing remaining turns")
+
+    all_messages = _read_transcript(transcript_path)
+    if not all_messages:
+        debug_log(config, "SessionEnd flush: no messages in transcript, skipping")
+        return
+
+    retain_roles = config.get("retainRoles", ["user", "assistant"])
+    include_tool_calls = config.get("retainToolCalls", True)
+    transcript, message_count = prepare_retention_transcript(
+        all_messages, retain_roles, True, include_tool_calls=include_tool_calls
+    )
+    if not transcript:
+        debug_log(config, "SessionEnd flush: empty transcript after formatting, skipping")
+        return
+
+    def _dbg(*a):
+        debug_log(config, *a)
+
+    try:
+        api_url = get_api_url(config, debug_fn=_dbg, allow_daemon_start=False)
+    except RuntimeError as e:
+        print(f"[Hindsight] SessionEnd flush: cannot reach API — {e}", file=sys.stderr)
+        return
+
+    api_token = config.get("hindsightApiToken")
+    try:
+        client = HindsightClient(api_url, api_token)
+    except ValueError as e:
+        print(f"[Hindsight] SessionEnd flush: invalid API URL — {e}", file=sys.stderr)
+        return
+
+    bank_id = derive_bank_id(hook_input, config)
+    ensure_bank_mission(client, bank_id, config, debug_fn=_dbg)
+
+    # Use session_id as document_id so this upserts the same document as prior
+    # full-session retains (deduplicates; last write wins).
+    document_id = session_id
+
+    template_vars = {
+        "session_id": session_id,
+        "bank_id": bank_id,
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "user_id": os.environ.get("HINDSIGHT_USER_ID", ""),
+    }
+
+    def _resolve_template(value: str) -> str:
+        for k, v in template_vars.items():
+            value = value.replace(f"{{{k}}}", v)
+        return value
+
+    raw_tags = config.get("retainTags", [])
+    tags = None
+    if raw_tags:
+        resolved = []
+        for original in raw_tags:
+            tag = _resolve_template(original)
+            if ":" in tag and tag.split(":", 1)[1] == "":
+                continue
+            resolved.append(tag)
+        tags = resolved or None
+
+    metadata = {
+        "retained_at": template_vars["timestamp"],
+        "message_count": str(message_count),
+        "session_id": session_id,
+        "flush_reason": "session_end",
+    }
+    for k, v in config.get("retainMetadata", {}).items():
+        metadata[k] = _resolve_template(str(v))
+
+    debug_log(config, f"SessionEnd flush: retaining bank '{bank_id}', {message_count} messages, {len(transcript)} chars")
+
+    try:
+        response = client.retain(
+            bank_id=bank_id,
+            content=transcript,
+            document_id=document_id,
+            context=config.get("retainContext", "claude-code"),
+            metadata=metadata,
+            tags=tags,
+            timeout=15,
+        )
+        debug_log(config, f"SessionEnd flush: retain response {json.dumps(response)[:200]}")
+    except Exception as e:
+        print(f"[Hindsight] SessionEnd flush failed: {e}", file=sys.stderr)
 
 
 def main():
@@ -28,6 +169,13 @@ def main():
 
     debug_log(config, f"SessionEnd hook, reason: {hook_input.get('reason', 'unknown')}")
 
+    # Flush any turns not captured by the Stop-hook throttle before stopping daemon
+    if config.get("autoRetain"):
+        try:
+            _flush_remaining_turns(config, hook_input)
+        except Exception as e:
+            print(f"[Hindsight] SessionEnd flush error: {e}", file=sys.stderr)
+
     # Stop daemon if we started it
     def _dbg(*a):
         debug_log(config, *a)

hindsight-integrations/claude-code/tests/test_hooks.py

@@ -12,14 +12,11 @@
 import json
 import os
 import sys
-import time
 from unittest.mock import MagicMock, patch
 
 import pytest
-
 from conftest import FakeHTTPResponse, make_hook_input, make_memory, make_transcript_file
 
-
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -560,7 +557,7 @@ def capture(req, timeout=None):
         if "body" in captured:
             assert captured["body"]["items"][0]["context"] == "claude-code"
 
-    def test_disabled_auto_retain_does_not_call_api(self, monkeypatch, tmp_path):
+    def test_disabled_auto_retain_does_not_call_api_retain(self, monkeypatch, tmp_path):
         (tmp_path / "plugin_root").mkdir(exist_ok=True)
         (tmp_path / "plugin_data").mkdir(exist_ok=True)
         settings = {"autoRetain": False, "autoRecall": False, "hindsightApiUrl": "http://fake:9077"}
@@ -595,3 +592,114 @@ def capture(req, timeout=None):
             mod.main()
 
         assert "called" not in captured
+
+
+# ---------------------------------------------------------------------------
+# session_end hook
+# ---------------------------------------------------------------------------
+
+
+class TestSessionEndHook:
+    """Tests for session_end.py flush-on-exit behaviour."""
+
+    def test_flushes_remaining_turns_on_session_end(self, monkeypatch, tmp_path):
+        """SessionEnd should retain the full transcript when the last Stop hook was skipped."""
+        messages = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "world"}]
+        transcript = make_transcript_file(tmp_path, messages)
+        hook_input = make_hook_input(transcript_path=transcript, session_id="sess-flush")
+        captured = {}
+
+        def capture(req, timeout=None):
+            if "/memories" in req.full_url and "/recall" not in req.full_url:
+                captured["called"] = True
+                captured["body"] = json.loads(req.data.decode())
+            return FakeHTTPResponse({})
+
+        # Fire retain on turn 1 of 3 — should be skipped
+        _run_hook(
+            "retain", hook_input, monkeypatch, tmp_path,
+            urlopen_side_effect=capture,
+            extra_settings={"retainEveryNTurns": 3},
+        )
+        assert "called" not in captured, "turn 1/3 should not trigger retain"
+
+        # SessionEnd should detect turn_count=1 % 3 != 0 and flush
+        captured.clear()
+        session_end_input = {"reason": "normal", "session_id": "sess-flush", "transcript_path": transcript}
+        _run_hook(
+            "session_end", session_end_input, monkeypatch, tmp_path,
+            urlopen_side_effect=capture,
+            extra_settings={"retainEveryNTurns": 3},
+        )
+        assert "called" in captured, "SessionEnd should flush remaining turns"
+        item = captured["body"]["items"][0]
+        assert "hello" in item["content"]
+        assert item["document_id"] == "sess-flush"
+
+    def test_skips_flush_when_retain_every_n_is_one(self, monkeypatch, tmp_path):
+        """When retainEveryNTurns=1, every Stop hook retains — no extra flush needed."""
+        messages = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "world"}]
+        transcript = make_transcript_file(tmp_path, messages)
+        session_end_input = {"reason": "normal", "session_id": "sess-no-flush", "transcript_path": transcript}
+        captured = {}
+
+        def capture(req, timeout=None):
+            if "/memories" in req.full_url and "/recall" not in req.full_url:
+                captured["called"] = True
+            return FakeHTTPResponse({})
+
+        _run_hook(
+            "session_end", session_end_input, monkeypatch, tmp_path,
+            urlopen_side_effect=capture,
+            extra_settings={"retainEveryNTurns": 1},
+        )
+        assert "called" not in captured, "No flush when retainEveryNTurns=1"
+
+    def test_skips_flush_when_last_stop_already_retained(self, monkeypatch, tmp_path):
+        """When turn_count is a multiple of retainEveryNTurns, the Stop hook already retained."""
+        messages = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "world"}]
+        transcript = make_transcript_file(tmp_path, messages)
+        hook_input = make_hook_input(transcript_path=transcript, session_id="sess-already-retained")
+        captured = {}
+
+        def capture(req, timeout=None):
+            if "/memories" in req.full_url and "/recall" not in req.full_url:
+                captured["called"] = True
+            return FakeHTTPResponse({})
+
+        # Fire retain exactly retainEveryNTurns times so turn_count % n == 0
+        for _ in range(3):
+            _run_hook(
+                "retain", hook_input, monkeypatch, tmp_path,
+                urlopen_side_effect=capture,
+                extra_settings={"retainEveryNTurns": 3},
+            )
+
+        # The 3rd call should have triggered retain; now SessionEnd should NOT add another
+        captured.clear()
+        session_end_input = {"reason": "normal", "session_id": "sess-already-retained", "transcript_path": transcript}
+        _run_hook(
+            "session_end", session_end_input, monkeypatch, tmp_path,
+            urlopen_side_effect=capture,
+            extra_settings={"retainEveryNTurns": 3},
+        )
+        assert "called" not in captured, "No extra flush when turn_count % retainEveryNTurns == 0"
+
+    def test_skips_flush_when_auto_retain_disabled(self, monkeypatch, tmp_path):
+        """SessionEnd should not retain when autoRetain is disabled."""
+        messages = [{"role": "user", "content": "hello"}]
+        transcript = make_transcript_file(tmp_path, messages)
+        captured = {}
+
+        def capture(req, timeout=None):
+            if "/memories" in req.full_url and "/recall" not in req.full_url:
+                captured["called"] = True
+            return FakeHTTPResponse({})
+
+        session_end_input = {"reason": "normal", "session_id": "sess-disabled", "transcript_path": transcript}
+        _run_hook(
+            "session_end", session_end_input, monkeypatch, tmp_path,
+            urlopen_side_effect=capture,
+            extra_settings={"autoRetain": False, "retainEveryNTurns": 3},
+        )
+        assert "called" not in captured, "No flush when autoRetain=false"