Update rft_flightrecorder.py

rft_flightrecorder.py

@@ -6,6 +6,8 @@ import zipfile
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Tuple
 
+from filelock import FileLock, Timeout
+
 from cryptography.hazmat.primitives.asymmetric.ed25519 import (
     Ed25519PrivateKey,
     Ed25519PublicKey,
@@ -17,6 +19,8 @@ EVENT_SPEC = "rft-flight-event-v0"
 ROOT_SPEC = "rft-flight-session-root-v0"
 DEFAULT_LOG_PATH = "flightlog.jsonl"
 
+LOCK_TIMEOUT_S = 10  # keep small; we only lock for quick file IO
+
 
 # ============================================================
 # Canonical JSON + hashing
@@ -45,6 +49,20 @@ def short(h: str, n: int = 12) -> str:
     return h[:n] if len(h) > n else h
 
 
+# ============================================================
+# File locking helpers
+# ============================================================
+
+def _lock_for(path: str) -> FileLock:
+    # one lock file per jsonl
+    return FileLock(path + ".lock", timeout=LOCK_TIMEOUT_S)
+
+
+def _read_jsonl_locked(path: str) -> Tuple[List[Dict[str, Any]], int]:
+    with _lock_for(path):
+        return read_jsonl(path)
+
+
 # ============================================================
 # Key management (Ed25519)
 # ============================================================
@@ -196,49 +214,59 @@ def append_event(
     if not sid:
         return None, "Missing session_id."
 
-    all_events, _corrupt = read_jsonl(log_path)
+    # IMPORTANT: Atomic critical section:
+    # read -> compute seq/prev -> build event -> append
+    try:
+        with _lock_for(log_path):
+            all_events, _corrupt = read_jsonl(log_path)
 
-    payload = parse_payload_text(payload_text)
-    meta = {
-        "model_id": (model_id or "unknown").strip(),
-        "run_mode": (run_mode or "unknown").strip(),
-    }
+            payload = parse_payload_text(payload_text)
+            meta = {
+                "model_id": (model_id or "unknown").strip(),
+                "run_mode": (run_mode or "unknown").strip(),
+            }
 
-    last = last_event_for_session(all_events, sid)
-    prev_hash = last.get("event_hash_sha256") if last else ("0" * 64)
-    seq = next_seq(all_events, sid)
+            last = last_event_for_session(all_events, sid)
+            prev_hash = last.get("event_hash_sha256") if last else ("0" * 64)
+            seq = next_seq(all_events, sid)
+
+            # If parent hash not provided, default to prev hash (simple linear causality)
+            parent = (parent_event_hash or "").strip()
+            if not parent and seq > 1:
+                parent = prev_hash
+
+            core = make_event_core(
+                session_id=sid,
+                seq=seq,
+                event_type=event_type,
+                payload=payload,
+                parent_event_hash=parent,
+                prev_event_hash=prev_hash,
+                meta=meta,
+            )
 
-    # If parent hash not provided, default to prev hash (simple linear causality)
-    parent = (parent_event_hash or "").strip()
-    if not parent and seq > 1:
-        parent = prev_hash
+            event_hash = sha256_hex(canon(core))
+            event = dict(core)
+            event["event_hash_sha256"] = event_hash
 
-    core = make_event_core(
-        session_id=sid,
-        seq=seq,
-        event_type=event_type,
-        payload=payload,
-        parent_event_hash=parent,
-        prev_event_hash=prev_hash,
-        meta=meta,
-    )
+            if sign_event:
+                if not sk_hex or not sk_hex.strip():
+                    return None, "Signing enabled but private key is missing."
+                try:
+                    sk = load_sk(sk_hex)
+                    sig = sk.sign(bytes.fromhex(event_hash))
+                    event["signature_ed25519"] = sig.hex()
+                except Exception:
+                    return None, "Failed to sign event (invalid private key?)."
 
-    event_hash = sha256_hex(canon(core))
-    event = dict(core)
-    event["event_hash_sha256"] = event_hash
+            append_jsonl(log_path, event)
 
-    if sign_event:
-        if not sk_hex or not sk_hex.strip():
-            return None, "Signing enabled but private key is missing."
-        try:
-            sk = load_sk(sk_hex)
-            sig = sk.sign(bytes.fromhex(event_hash))
-            event["signature_ed25519"] = sig.hex()
-        except Exception:
-            return None, "Failed to sign event (invalid private key?)."
+        return event, f"OK. Appended event #{seq} ({event_type}). Hash: {event_hash}"
 
-    append_jsonl(log_path, event)
-    return event, f"OK. Appended event #{seq} ({event_type}). Hash: {event_hash}"
+    except Timeout:
+        return None, "Busy: log is locked (try again)."
+    except Exception:
+        return None, "Failed to append event (unexpected error)."
 
 
 def start_session(
@@ -280,56 +308,91 @@ def finalise_session(
     if not sid:
         return None, "Missing session_id."
 
-    all_events, _corrupt = read_jsonl(log_path)
-    evs = events_for_session(all_events, sid)
-    if not evs:
-        return None, "No events found for this session."
-
-    first_hash = evs[0]["event_hash_sha256"]
-    last_hash = evs[-1]["event_hash_sha256"]
-    count = len(evs)
-
-    root_core = {
-        "spec": ROOT_SPEC,
-        "session_id": sid,
-        "first_event_hash_sha256": first_hash,
-        "last_event_hash_sha256": last_hash,
-        "event_count": count,
-    }
-    root_hash = sha256_hex(canon(root_core))
+    try:
+        # Lock while reading events + appending session_end (append_event is locked too,
+        # but we want a consistent anchor over a stable snapshot of the file)
+        with _lock_for(log_path):
+            all_events, _corrupt = read_jsonl(log_path)
+            evs = events_for_session(all_events, sid)
+            if not evs:
+                return None, "No events found for this session."
 
-    anchor = dict(root_core)
-    anchor["root_hash_sha256"] = root_hash
-    anchor["created_utc"] = now_utc_iso()
-    anchor["model_id"] = (model_id or "unknown").strip()
-    anchor["run_mode"] = (run_mode or "unknown").strip()
+            # Ensure correct order
+            evs.sort(key=lambda x: int(x.get("seq", 0)))
 
-    if sign_anchor:
-        if not sk_hex or not sk_hex.strip():
-            return None, "Anchor signing enabled but private key is missing."
-        try:
-            sk = load_sk(sk_hex)
-            sig = sk.sign(bytes.fromhex(root_hash))
-            anchor["signature_ed25519"] = sig.hex()
-        except Exception:
-            return None, "Failed to sign anchor (invalid private key?)."
+            first_hash = evs[0]["event_hash_sha256"]
+            last_hash = evs[-1]["event_hash_sha256"]
+            count = len(evs)
 
-    payload_text = json.dumps({"anchor": anchor}, ensure_ascii=False)
-    ev, msg = append_event(
-        log_path=log_path,
-        session_id=sid,
-        event_type="session_end",
-        payload_text=payload_text,
-        parent_event_hash=last_hash,
-        sign_event=sign_anchor,
-        sk_hex=sk_hex,
-        model_id=model_id,
-        run_mode=run_mode,
-    )
-    if not ev:
-        return None, msg
+            root_core = {
+                "spec": ROOT_SPEC,
+                "session_id": sid,
+                "first_event_hash_sha256": first_hash,
+                "last_event_hash_sha256": last_hash,
+                "event_count": count,
+            }
+            root_hash = sha256_hex(canon(root_core))
+
+            anchor = dict(root_core)
+            anchor["root_hash_sha256"] = root_hash
+            anchor["created_utc"] = now_utc_iso()
+            anchor["model_id"] = (model_id or "unknown").strip()
+            anchor["run_mode"] = (run_mode or "unknown").strip()
+
+            if sign_anchor:
+                if not sk_hex or not sk_hex.strip():
+                    return None, "Anchor signing enabled but private key is missing."
+                try:
+                    sk = load_sk(sk_hex)
+                    sig = sk.sign(bytes.fromhex(root_hash))
+                    anchor["signature_ed25519"] = sig.hex()
+                except Exception:
+                    return None, "Failed to sign anchor (invalid private key?)."
+
+            # Append session_end inside the same lock to prevent new events slipping in
+            payload_text = json.dumps({"anchor": anchor}, ensure_ascii=False)
+
+            # Build the session_end event manually so it stays inside our lock
+            # (we avoid calling append_event here to prevent nested lock complexity)
+            meta = {
+                "model_id": (model_id or "unknown").strip(),
+                "run_mode": (run_mode or "unknown").strip(),
+            }
+            seq = int(evs[-1].get("seq", 0)) + 1
+            prev_hash = evs[-1]["event_hash_sha256"]
+            parent = prev_hash
+
+            core = make_event_core(
+                session_id=sid,
+                seq=seq,
+                event_type="session_end",
+                payload=parse_payload_text(payload_text),
+                parent_event_hash=parent,
+                prev_event_hash=prev_hash,
+                meta=meta,
+            )
+            event_hash = sha256_hex(canon(core))
+            event = dict(core)
+            event["event_hash_sha256"] = event_hash
+
+            if sign_anchor:
+                if not sk_hex or not sk_hex.strip():
+                    return None, "Anchor signing enabled but private key is missing."
+                try:
+                    sk = load_sk(sk_hex)
+                    sig = sk.sign(bytes.fromhex(event_hash))
+                    event["signature_ed25519"] = sig.hex()
+                except Exception:
+                    return None, "Failed to sign session_end event."
+
+            append_jsonl(log_path, event)
 
-    return anchor, f"OK. Session finalised. Root hash: {root_hash} (last event hash: {ev['event_hash_sha256']})"
+        return anchor, f"OK. Session finalised. Root hash: {root_hash} (last event hash: {event_hash})"
+
+    except Timeout:
+        return None, "Busy: log is locked (try again)."
+    except Exception:
+        return None, "Failed to finalise session (unexpected error)."
 
 
 # ============================================================
@@ -451,9 +514,11 @@ def verify_session_from_events(
 
 
 def verify_session(log_path: str, session_id: str, pk_hex: str, require_signatures: bool) -> Tuple[str, bool, str]:
-    all_events, _corrupt = read_jsonl(log_path)
+    try:
+        all_events, _corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return "FAIL", False, "Busy: log is locked (try again)."
     evs = events_for_session(all_events, session_id)
-    # Ensure correct order even if file was rearranged
     evs.sort(key=lambda x: int(x.get("seq", 0)))
     return verify_session_from_events(evs, session_id, pk_hex=pk_hex, require_signatures=require_signatures)
 
@@ -467,7 +532,11 @@ def session_timeline_rows(log_path: str, session_id: str) -> Tuple[List[List[Any
     if not sid:
         return [], "Missing session_id."
 
-    all_events, _corrupt = read_jsonl(log_path)
+    try:
+        all_events, _corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return [], "Busy: log is locked (try again)."
+
     evs = events_for_session(all_events, sid)
     if not evs:
         return [], "No events found."
@@ -496,7 +565,11 @@ def get_event_by_hash(log_path: str, session_id: str, event_hash: str) -> Tuple[
     if not sid or not h:
         return None, "Missing session_id or event hash."
 
-    all_events, _corrupt = read_jsonl(log_path)
+    try:
+        all_events, _corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return None, "Busy: log is locked (try again)."
+
     evs = events_for_session(all_events, sid)
     for e in evs:
         if e.get("event_hash_sha256") == h:
@@ -505,7 +578,11 @@ def get_event_by_hash(log_path: str, session_id: str, event_hash: str) -> Tuple[
 
 
 def list_sessions(log_path: str) -> Tuple[List[str], str]:
-    all_events, corrupt = read_jsonl(log_path)
+    try:
+        all_events, corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return [], "Busy: log is locked (try again)."
+
     counts: Dict[str, int] = {}
     for e in all_events:
         if isinstance(e, dict) and e.get("spec") == EVENT_SPEC:
@@ -518,7 +595,20 @@ def list_sessions(log_path: str) -> Tuple[List[str], str]:
 
 
 def diagnostics(log_path: str) -> Dict[str, Any]:
-    all_events, corrupt = read_jsonl(log_path)
+    try:
+        all_events, corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return {
+            "log_path": log_path,
+            "exists": os.path.exists(log_path),
+            "bytes": os.path.getsize(log_path) if os.path.exists(log_path) else 0,
+            "total_events": 0,
+            "sessions": 0,
+            "signed_events": 0,
+            "corrupt_lines_ignored": 0,
+            "error": "Busy: log is locked (try again).",
+        }
+
     size = os.path.getsize(log_path) if os.path.exists(log_path) else 0
     sessions = set()
     signed = 0
@@ -550,7 +640,11 @@ def export_session_bundle(log_path: str, session_id: str) -> Tuple[Optional[str]
     if not sid:
         return None, "Missing session_id."
 
-    all_events, _corrupt = read_jsonl(log_path)
+    try:
+        all_events, _corrupt = _read_jsonl_locked(log_path)
+    except Timeout:
+        return None, "Busy: log is locked (try again)."
+
     evs = events_for_session(all_events, sid)
     if not evs:
         return None, "No events found."
@@ -650,8 +744,13 @@ def import_bundle_verify(
     )
 
     if store_into_log and ok:
-        for e in evs:
-            append_jsonl(log_path, e)
+        # Store under a single lock to avoid interleaving with live writers.
+        try:
+            with _lock_for(log_path):
+                for e in evs:
+                    append_jsonl(log_path, e)
+        except Timeout:
+            return "FAIL", False, "Busy: log is locked (try again).", None
 
     stored_msg = "Stored into local flightlog.jsonl." if (store_into_log and ok) else None
     return status, ok, report, stored_msg