style(session): sort imports in atomic agent-history test

Co-authored-by: Mervin Praison <[email protected]>

src/praisonai-agents/praisonaiagents/session/api.py

@@ -502,6 +502,10 @@ class Session:
     ) -> bool:
         """Persist one sub-agent transcript onto this session's own record.
 
+        Built-in stores merge the entry under their session file lock. Older
+        custom stores retain the bounded read/write verification fallback below;
+        that fallback cannot guarantee atomic preservation of concurrent entries.
+
         The store replaces the whole ``AGENT_HISTORY_KEY`` map per write, so a
         naive read-then-replace would let a concurrent writer (another agent
         under the same parent) clobber this agent's entry, or vice versa. We
@@ -509,6 +513,14 @@ class Session:
         *other* key changed underneath us, retrying a bounded number of times.
         Legacy per-agent records are always merged in (they stay on disk).
         """
+        merge_map = getattr(session_store, "merge_session_metadata_map", None)
+        if callable(merge_map):
+            # Built-in stores merge inside their locked read/modify/write path.
+            # A post-write comparison cannot detect entries lost to a stale map.
+            return merge_map(
+                self.session_id, AGENT_HISTORY_KEY, {agent_key: messages},
+                defaults=self._legacy_agent_histories(session_store),
+            )
         update_meta = getattr(session_store, "update_session_metadata", None)
         if not callable(update_meta):
             return False

src/praisonai-agents/praisonaiagents/session/store.py

@@ -1622,6 +1622,27 @@ class DefaultSessionStore:
             session_id, _apply, error_label="update session metadata"
         )
 
+    def merge_session_metadata_map(
+        self, session_id: str, key: str, updates: Dict[str, Any],
+        *, defaults: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """Merge nested metadata entries using the freshest map under FileLock.
+
+        Defaults fill missing entries; persisted entries win over defaults,
+        and the supplied updates win over both.
+        """
+        def _apply(session: SessionData) -> None:
+            merged = dict(defaults or {})
+            current = session.metadata.get(key)
+            if isinstance(current, dict):
+                merged.update(current)
+            merged.update(updates)
+            session.metadata[key] = merged
+
+        return self._modify_session_locked(
+            session_id, _apply, error_label="merge session metadata map"
+        )
+
     def rename_session(self, session_id: str, title: str) -> bool:
         """Give a session a human-readable title (Issue #3737).
 

src/praisonai-agents/tests/unit/session/test_session_atomic_agent_history.py

@@ -0,0 +1,91 @@
+"""Concurrent Session closes must merge different agent histories under lock."""
+
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from threading import Event
+
+import pytest
+
+from praisonaiagents.session.api import AGENT_HISTORY_KEY, Session
+from praisonaiagents.session.store import DefaultSessionStore
+
+
+@pytest.mark.parametrize("existing", [{}, {"peer:Assistant": [
+    {"role": "user", "content": "peer"},
+]}])
+def test_two_closes_preserve_both_agent_histories(tmp_path, monkeypatch, existing):
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("PRAISONAI_HOME", str(tmp_path / "home"))
+    directory = str(tmp_path / "sessions")
+    first = DefaultSessionStore(session_dir=directory)
+    second = DefaultSessionStore(session_dir=directory)
+    first.update_session_metadata("parent", **{AGENT_HISTORY_KEY: existing})
+    a, b = Session(session_id="parent"), Session(session_id="parent")
+    monkeypatch.setattr(a, "_get_session_store", lambda: first)
+    monkeypatch.setattr(b, "_get_session_store", lambda: second)
+    ah = [{"role": "user", "content": "a"}]
+    bh = [{"role": "user", "content": "b"}]
+    a._agents["a:Assistant"] = {"agent": None, "chat_history": ah}
+    b._agents["b:Assistant"] = {"agent": None, "chat_history": bh}
+    ready, completed = Event(), Event()
+    original = first._modify_session_locked
+
+    def pause_before_lock(*args, **kwargs):
+        ready.set()
+        assert completed.wait(10), "second writer did not complete"
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(first, "_modify_session_locked", pause_before_lock)
+
+    def close_second():
+        assert ready.wait(10), "first writer did not reach write boundary"
+        try:
+            b.close()
+        finally:
+            completed.set()
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        fa, fb = pool.submit(a.close), pool.submit(close_second)
+        fa.result(timeout=15)
+        fb.result(timeout=15)
+    fresh = DefaultSessionStore(session_dir=directory)
+    assert fresh.get_session("parent").metadata[AGENT_HISTORY_KEY] == {
+        **existing, "a:Assistant": ah, "b:Assistant": bh,
+    }
+
+
+def test_metadata_map_defaults_do_not_overwrite_current_entries(tmp_path):
+    store = DefaultSessionStore(session_dir=str(tmp_path))
+    store.update_session_metadata("parent", marker="keep", **{AGENT_HISTORY_KEY: {
+        "current": [], "peer": ["present"],
+    }})
+    assert store.merge_session_metadata_map(
+        "parent", AGENT_HISTORY_KEY, {"new": ["new"]},
+        defaults={"current": ["legacy"], "legacy-only": ["old"]},
+    )
+    data = DefaultSessionStore(session_dir=str(tmp_path)).get_session("parent")
+    assert data.metadata[AGENT_HISTORY_KEY] == {
+        "current": [], "peer": ["present"], "new": ["new"],
+        "legacy-only": ["old"],
+    }
+    assert data.metadata["marker"] == "keep"
+
+
+@pytest.mark.parametrize("failure", ["read", "write"])
+def test_metadata_map_failure_preserves_disk(tmp_path, monkeypatch, failure):
+    import praisonaiagents.session.store as store_module
+
+    store = DefaultSessionStore(session_dir=str(tmp_path))
+    store.update_session_metadata("parent", **{AGENT_HISTORY_KEY: {"peer": ["old"]}})
+    path = Path(store._get_session_path("parent"))
+    before = path.read_bytes()
+
+    def fail(*args, **kwargs):
+        raise OSError("injected I/O failure")
+
+    if failure == "read":
+        monkeypatch.setattr(store, "_load_session_from_disk", fail)
+    else:
+        monkeypatch.setattr(store_module.os, "replace", fail)
+    assert not store.merge_session_metadata_map("parent", AGENT_HISTORY_KEY, {"new": []})
+    assert path.read_bytes() == before