my-pi-agent 0.1.0

package/src/my_agent_core/session/session.py

@@ -0,0 +1,400 @@
+# pyright: reportUnusedCallResult=false, reportAttributeAccessIssue=false
+"""Session 管理：树结构会话 + JSONL 文件持久化（纯正 Tau 只追加多态体系，配 my-agent-llm Message）。
+
+一个会话 = 一棵树（entry 带 id/parent_id）+ current_id 指针。rewind = 移动指针、
+旧分支保留。文件格式为纯粹的 SessionEntry 流：第 1 行是 SessionInfoEntry，后续每行为 MessageEntry/CompactionEntry 等。
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from collections.abc import Iterable
+from datetime import datetime
+from pathlib import Path
+from typing import Any, cast
+from uuid import uuid4
+
+from my_agent_llm import Message
+
+from .entries import (
+    CompactionEntry,
+    MessageEntry,
+    SessionEntry,
+    SessionHeaderEntry,
+    SessionInfoEntry,
+)
+from .jsonl import entry_from_json_line, entry_to_json_line
+from .tree import path_to_entry
+
+
+class SessionTree:
+    """树：entries（id→entry）+ current_id 指针 + root_id。"""
+
+    def __init__(self) -> None:
+        self.entries: dict[str, SessionEntry] = {}
+        self.current_id: str | None = None
+        self.root_id: str | None = None
+
+    def add_entry(self, role: str, content: str, parent_id: str | None = None, **metadata: Any) -> MessageEntry:
+        """追加到 current 下（或指定 parent）。首个 entry 成为根。"""
+        if parent_id is None:
+            parent_id = self.current_id
+        entry = MessageEntry(
+            id=uuid4().hex[:8],
+            parent_id=parent_id,
+            message=Message(
+                role=cast(Any, role),
+                content=content,
+                metadata=metadata or None,
+            ),
+        )
+        self.entries[entry.id] = entry
+        self.current_id = entry.id
+        if self.root_id is None:
+            self.root_id = entry.id
+        return entry
+
+    def get_current_path(self) -> list[SessionEntry]:
+        """根 → current 的路径（Agent 上下文用）。空树返回 []。"""
+        if self.current_id is None:
+            return []
+        return self.get_path_to_entry(self.current_id)
+
+    def get_path_to_entry(self, entry_id: str) -> list[SessionEntry]:
+        """根 → entry_id 的路径（fork 用）。不存在抛 ValueError。"""
+        if entry_id not in self.entries:
+            raise ValueError(f"Entry {entry_id} not found")
+        try:
+            return path_to_entry(self.entries, entry_id)
+        except Exception as exc:
+            raise ValueError(str(exc)) from exc
+
+    def rewind(self, entry_id: str) -> None:
+        """移动 current 指针到已有节点；新 entry 将在其下追加（长新枝）。不存在抛 ValueError。"""
+        if entry_id not in self.entries:
+            raise ValueError(f"Entry {entry_id} not found")
+        self.current_id = entry_id
+
+    @classmethod
+    def from_jsonl_iter(cls, lines: Iterable[str]) -> SessionTree:
+        """从迭代器恢复树。中途某行损坏抛 ValueError（带行号），尾行由 load 处理。"""
+        tree = cls()
+        for idx, line in enumerate(lines, start=2):  # start=2 因为 line 1 是 session_info
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entry = entry_from_json_line(line)
+            except Exception as exc:
+                raise ValueError(f"Corrupted entry at line {idx}: {exc}") from exc
+            tree.entries[entry.id] = entry
+            tree.current_id = entry.id
+            if tree.root_id is None:
+                tree.root_id = entry.id
+        return tree
+
+
+class Session:
+    """一个会话 = 树 + 路径 + JSONL 文件（纯正 Tau 多态 SessionEntry 流）。
+
+    文件第 1 行是 SessionInfoEntry，后续每行是一个多态 SessionEntry。
+    """
+
+    def __init__(self, *, path: Path, cwd: str | None = None, metadata: dict | None = None):
+        """新建会话（纯对话，不含 system）。不立即写文件。"""
+        self.path = Path(path)
+        self.id = f"{datetime.now().strftime('%Y%m%d-%H%M%S')}-{uuid4().hex[:8]}"
+        self.created_at = datetime.now().isoformat()
+        self.cwd = cwd or str(Path.cwd())
+        self.tree = SessionTree()
+        self.compaction_floor: str | None = None
+        self.metadata = metadata or {}  # 额外元数据（子代理：agent_type/parent_session_id），进 SessionInfoEntry
+        self._storage: Any | None = None
+
+    @classmethod
+    def load(cls, path: Path) -> Session:
+        """从 JSONL 文件恢复整棵树。首行非 SessionInfoEntry → ValueError；
+        非尾行损坏 → ValueError（带行号）；尾行撕裂 → 丢弃该行（宽容兜底）。"""
+        path = Path(path).resolve()
+        try:
+            with open(path, encoding="utf-8") as f:
+                lines = list(f)
+        except OSError as exc:
+            raise ValueError(f"Failed to read session file {path}: {exc}") from exc
+        if not lines:
+            raise ValueError(f"Session file {path} is empty")
+
+        try:
+            first_entry = entry_from_json_line(lines[0])
+        except Exception as exc:
+            raise ValueError(f"Session file {path}: invalid header/info: {exc}") from exc
+
+        if not isinstance(first_entry, (SessionInfoEntry, SessionHeaderEntry)):
+            raise ValueError(f"Session file {path}: line 1 must be a valid SessionHeaderEntry or SessionInfoEntry")
+
+        tree_lines = lines[1:]
+        if tree_lines and tree_lines[-1].strip():
+            try:
+                entry_from_json_line(tree_lines[-1])
+            except Exception:
+                tree_lines = tree_lines[:-1]
+
+        tree = SessionTree.from_jsonl_iter(tree_lines)
+        meta = dict(getattr(first_entry, "metadata", {}) or {})
+        if getattr(first_entry, "name", None):
+            meta["name"] = first_entry.name
+        if getattr(first_entry, "title", None):
+            meta["title"] = first_entry.title
+        if getattr(first_entry, "parent_session", None):
+            meta["parent_session"] = first_entry.parent_session
+            meta["parent_session_path"] = first_entry.parent_session
+        for e in tree.entries.values():
+            if isinstance(e, SessionInfoEntry):
+                sub_name = getattr(e, "name", None) or getattr(e, "title", None)
+                if sub_name:
+                    meta["name"] = sub_name
+                    meta["title"] = sub_name
+
+        # 优先恢复 current_id / root_id
+        cur = meta.get("current_id")
+        if isinstance(cur, str) and cur in tree.entries:
+            tree.current_id = cur
+        root = meta.get("root_id")
+        if isinstance(root, str) and root in tree.entries:
+            tree.root_id = root
+
+        session = cls(path=path, cwd=first_entry.cwd)
+        session.id = first_entry.id
+        session.created_at = (
+            datetime.fromtimestamp(first_entry.created_at).isoformat()
+            if first_entry.created_at
+            else datetime.now().isoformat()
+        )
+        session.tree = tree
+        session.compaction_floor = meta.get("compaction_floor")
+        session.metadata = {k: v for k, v in meta.items() if k not in ("current_id", "root_id", "compaction_floor")}
+        return session
+
+    def add_message(self, role: str, content: str, parent_id: str | None = None, **metadata: Any) -> MessageEntry:
+        """加到树 + save()。"""
+        entry = self.tree.add_entry(role, content, parent_id, **metadata)
+        self.save()
+        return entry
+
+    def get_current_path_messages(self) -> list[Message]:
+        """当前路径 → list[Message]（Agent 上下文用）。"""
+        return [
+            e.message
+            if isinstance(e, MessageEntry)
+            else Message(
+                role=cast(Any, getattr(e, "role", "system")),
+                content=getattr(e, "content", ""),
+                metadata=getattr(e, "metadata", None),
+            )
+            for e in self.tree.get_current_path()
+        ]
+
+    def add_summary_cache(
+        self,
+        summary: str,
+        *,
+        covered_count: int,
+        retained_tail: list[dict],
+        tokens_before: int,
+        summary_usage: dict | None = None,
+        summary_model: str | None = None,
+    ) -> None:
+        """写一条 CompactionEntry（不动 current_id）+ 更新 compaction_floor。"""
+        metadata: dict[str, Any] = {
+            "retained_tail": retained_tail,
+            "covered_count": covered_count,
+            "tokens_before": tokens_before,
+        }
+        if summary_usage is not None:
+            metadata["summary_usage"] = summary_usage
+        if summary_model is not None:
+            metadata["summary_model"] = summary_model
+        entry = CompactionEntry(
+            parent_id=self.tree.current_id,
+            summary=summary,
+            replaces_entry_ids=[],
+            metadata=metadata,
+        )
+        self.tree.entries[entry.id] = entry
+        self.compaction_floor = self.tree.current_id
+        self.save()
+
+    def get_full_history_messages(self) -> list[Message]:
+        """完整对话历史（排除 CompactionEntry 节点）——宿主看历史、Agent 恢复上下文用。"""
+        return [
+            e.message
+            if isinstance(e, MessageEntry)
+            else Message(
+                role=cast(Any, getattr(e, "role", "system")),
+                content=getattr(e, "content", ""),
+                metadata=getattr(e, "metadata", None),
+            )
+            for e in self.tree.get_current_path()
+            if getattr(e, "type", "message") == "message"
+        ]
+
+    def get_latest_compaction_cache(self) -> dict | None:
+        """最新一条 CompactionEntry → {summary, covered_count, retained_tail}；无则 None。"""
+        cache_entries = [
+            e
+            for e in self.tree.entries.values()
+            if isinstance(e, CompactionEntry) or getattr(e, "type", None) == "compaction"
+        ]
+        if not cache_entries:
+            return None
+        latest = max(cache_entries, key=lambda e: len(self.tree.get_path_to_entry(e.id)))
+        md = getattr(latest, "metadata", {}) or {}
+        try:
+            covered_count = int(md.get("covered_count", 0))
+        except (ValueError, TypeError):
+            covered_count = 0
+        return {
+            "summary": getattr(latest, "summary", getattr(latest, "content", "")),
+            "covered_count": covered_count,
+            "retained_tail": list(md.get("retained_tail", [])),
+        }
+
+    def rewind(self, entry_id: str) -> None:
+        """移动 current 指针（旧分支保留）+ save。压缩后只能回 floor（含）之后，否则 ValueError。"""
+        if self.compaction_floor is not None and not self._after_floor(entry_id):
+            raise ValueError(
+                f"Cannot rewind past compaction floor {self.compaction_floor}: "
+                f"entry {entry_id} is prior to compacted history"
+            )
+        self.tree.rewind(entry_id)
+        self.save()
+
+    def _after_floor(self, entry_id: str) -> bool:
+        """entry_id 是否在 compaction_floor（含）之后长出的节点：等于 floor，或沿 parent 链能走到 floor。"""
+        if entry_id == self.compaction_floor:
+            return True
+        cur = self.tree.entries.get(entry_id)
+        while cur is not None and cur.parent_id is not None:
+            if cur.parent_id == self.compaction_floor:
+                return True
+            cur = self.tree.entries.get(cur.parent_id)
+        return False
+
+    def save(self) -> None:
+        """原子落盘：临时文件 + fsync + os.replace。写入首条 SessionInfoEntry 与所有多态 SessionEntry。"""
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        meta = dict(self.metadata)
+        meta["current_id"] = self.tree.current_id
+        meta["root_id"] = self.tree.root_id
+        meta["compaction_floor"] = self.compaction_floor
+
+        created_at_val: float | None = None
+        try:
+            if isinstance(self.created_at, str):
+                created_at_val = datetime.fromisoformat(self.created_at).timestamp()
+            elif isinstance(self.created_at, (int, float)):
+                created_at_val = float(self.created_at)
+        except Exception:
+            created_at_val = None
+
+        name_val = meta.get("name") or meta.get("title")
+        info_entry = SessionInfoEntry(
+            id=self.id,
+            cwd=self.cwd,
+            title=name_val,
+            name=name_val,
+            created_at=created_at_val,
+            metadata=meta,
+        )
+
+        temp_path: Path | None = None
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=self.path.parent,
+                prefix=f".{self.path.name}.",
+                suffix=".tmp",
+                delete=False,
+            ) as f:
+                temp_path = Path(f.name)
+                f.write(entry_to_json_line(info_entry) + "\n")
+                for e in self.tree.entries.values():
+                    f.write(entry_to_json_line(e) + "\n")
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(temp_path, self.path)
+        except Exception:
+            if temp_path is not None:
+                temp_path.unlink(missing_ok=True)
+            raise
+
+    def append_entry(self, entry: SessionEntry) -> None:
+        """追加一条任意多态 SessionEntry 并原子持久化。"""
+        if isinstance(entry, SessionInfoEntry):
+            name_val = getattr(entry, "name", None) or getattr(entry, "title", None)
+            if name_val:
+                self.metadata["name"] = name_val
+                self.metadata["title"] = name_val
+            if entry.cwd:
+                self.cwd = entry.cwd
+        self.tree.entries[entry.id] = entry
+        if self.tree.root_id is None:
+            self.tree.root_id = entry.id
+        self.tree.current_id = entry.id
+        self.save()
+
+    @property
+    def store(self):
+        """兼容层：暴露包含 append_entry 的 storage/store 接口。"""
+
+        class _SessionStoreProxy:
+            def __init__(self, session: Session):
+                self._session = session
+
+            def append_entry(self, entry: SessionEntry) -> None:
+                self._session.append_entry(entry)
+
+        return _SessionStoreProxy(self)
+
+    def reset(self) -> None:
+        """清空树 + 原子重写（纯对话，不含 system）。唯一破坏性操作。"""
+        self.tree = SessionTree()
+        self.compaction_floor = None
+        self.save()
+
+    def fork(self, entry_id: str, new_path: Path | None = None) -> Session:
+        """从某 entry 分叉为新会话：复制根到 entry 的路径为新会话（新 id/路径，独立演化）。"""
+        if entry_id not in self.tree.entries:
+            raise ValueError(f"Entry {entry_id} not found")
+        if new_path is None:
+            sid = f"{datetime.now().strftime('%Y%m%d-%H%M%S')}-{uuid4().hex[:8]}"
+            new_path = self.path.parent / f"{sid}.jsonl"
+        new_session = Session(path=new_path, cwd=self.cwd, metadata=dict(self.metadata))
+        for entry in self.tree.get_path_to_entry(entry_id):
+            if isinstance(entry, MessageEntry):
+                new_session.add_message(entry.role, entry.content, **entry.metadata)
+        return new_session
+
+    @property
+    def storage(self) -> Any:
+        """底层只追加存储对象。"""
+        if self._storage is None:
+            from my_agent_core.session.storage import JsonlSessionStorage
+
+            self._storage = JsonlSessionStorage(self.path)
+        return self._storage
+
+    def get_state(self) -> Any:
+        """获取从根到当前 current_id 的 SessionState 状态快照。"""
+        from my_agent_core.session.memory import SessionState
+
+        return SessionState(
+            messages=tuple(self.get_current_path_messages()),
+            active_leaf_id=self.tree.current_id,
+        )
+
+    def get_history(self) -> list[Message]:
+        """获取当前路径的对话历史列表。"""
+        return self.get_current_path_messages()

package/src/my_agent_core/session/storage.py

@@ -0,0 +1,245 @@
+# pyright: reportUnusedCallResult=false
+"""Locked, append-only session storage implementations and protocols.
+
+对齐 Tau (tau_agent.session.storage) 架构设计：
+集中放置 SessionStorage 协议、InMemorySessionStorage 以及 JsonlSessionStorage。
+文件锁底层调用以内聚辅助函数 _lock_file / _unlock_file 放置于文件末尾，
+彻底杜绝模块顶层的条件分支定义，具备工业级跨进程防踩踏与超时保护能力。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import sys
+import time
+from collections.abc import AsyncIterator, Sequence
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from .entries import SessionEntry
+from .jsonl import (
+    SessionJsonlError,
+    entry_from_json_line,
+    entry_to_json_line,
+)
+
+
+@runtime_checkable
+class SessionStorage(Protocol):
+    """纯异步只追加会话存储协议。
+
+    彻底废除全量重写 rewrite_history，保证历史记录发生即不可变。
+    """
+
+    async def append(self, entry: SessionEntry) -> None:
+        """追加单个条目到存储末尾。"""
+        ...
+
+    async def append_batch(self, entries: Sequence[SessionEntry]) -> None:
+        """原子追加一批条目到存储末尾。"""
+        ...
+
+    async def read_all(self) -> list[SessionEntry]:
+        """读取存储中所有按追加时序排列的历史条目。"""
+        ...
+
+
+class InMemorySessionStorage:
+    """纯内存只追加会话存储实现。
+
+    用于极速单元测试和无持久化需求的会话场景，协程安全。
+    """
+
+    def __init__(self, initial_entries: Sequence[SessionEntry] | None = None) -> None:
+        self._entries: list[SessionEntry] = (
+            list(initial_entries) if initial_entries else []
+        )
+        self._lock = asyncio.Lock()
+
+    async def append(self, entry: SessionEntry) -> None:
+        async with self._lock:
+            self._entries.append(entry)
+
+    async def append_batch(self, entries: Sequence[SessionEntry]) -> None:
+        async with self._lock:
+            self._entries.extend(entries)
+
+    async def read_all(self) -> list[SessionEntry]:
+        async with self._lock:
+            return list(self._entries)
+
+
+class JsonlSessionStorage:
+    """基于单文件行级追加的 SessionStorage 实现。
+
+    支持 `.{name}.lock` 跨进程文件锁，以及异常退出遗留 `.tmp` 碎片的自愈清理。
+    """
+
+    def __init__(self, path: Path | str) -> None:
+        self.path = Path(path)
+        self.lock_path = self.path.parent / f".{self.path.name}.lock"
+        self._async_lock = asyncio.Lock()
+        self._remove_incomplete_temp()
+
+    def _remove_incomplete_temp(self) -> list[Path]:
+        """清理异常中断遗留的临时碎片文件。"""
+        cleaned: list[Path] = []
+        parent = self.path.parent
+        if not parent.exists():
+            return cleaned
+
+        patterns = [
+            f".{self.path.name}*.tmp",
+            f"{self.path.name}*.tmp",
+        ]
+        seen: set[Path] = set()
+        for pattern in patterns:
+            for tmp_file in parent.glob(pattern):
+                if tmp_file in seen or not tmp_file.is_file():
+                    continue
+                seen.add(tmp_file)
+                with contextlib.suppress(OSError):
+                    tmp_file.unlink(missing_ok=True)
+                    cleaned.append(tmp_file)
+        return cleaned
+
+    @contextlib.asynccontextmanager
+    async def _process_lock(self, timeout: float = 10.0) -> AsyncIterator[None]:
+        """基于 `.{name}.lock` 的跨进程文件锁。"""
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        start_time = time.monotonic()
+        fd = os.open(str(self.lock_path), os.O_RDWR | os.O_CREAT)
+        try:
+            acquired = False
+            while not acquired:
+                try:
+                    _lock_file(fd)
+                    acquired = True
+                except (BlockingIOError, OSError, PermissionError):
+                    if time.monotonic() - start_time > timeout:
+                        raise SessionJsonlError(
+                            f"Timeout ({timeout}s) waiting for lock on {self.lock_path}"
+                        ) from None
+                    await asyncio.sleep(0.01)
+            try:
+                yield
+            finally:
+                _unlock_file(fd)
+        finally:
+            os.close(fd)
+
+    @contextlib.asynccontextmanager
+    async def _lock(self, timeout: float = 10.0) -> AsyncIterator[None]:
+        """协程锁与跨进程锁的双重上下文。"""
+        async with self._async_lock, self._process_lock(timeout=timeout):
+            yield
+
+    async def append(self, entry: SessionEntry) -> None:
+        """追加单个条目到存储末尾。"""
+        async with self._lock():
+            self._remove_incomplete_temp()
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+            line = entry_to_json_line(entry)
+            try:
+                with open(self.path, "a", encoding="utf-8") as f:
+                    f.write(line + "\n")
+                    f.flush()
+                    os.fsync(f.fileno())
+            except OSError as exc:
+                raise SessionJsonlError(
+                    f"Failed to append to {self.path}: {exc}"
+                ) from exc
+
+    async def append_batch(self, entries: Sequence[SessionEntry]) -> None:
+        """原子追加一批条目到存储末尾。"""
+        if not entries:
+            return
+        async with self._lock():
+            self._remove_incomplete_temp()
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+            lines = [entry_to_json_line(e) for e in entries]
+            try:
+                with open(self.path, "a", encoding="utf-8") as f:
+                    for line in lines:
+                        f.write(line + "\n")
+                    f.flush()
+                    os.fsync(f.fileno())
+            except OSError as exc:
+                raise SessionJsonlError(
+                    f"Failed to append batch to {self.path}: {exc}"
+                ) from exc
+
+    async def read_all(self) -> list[SessionEntry]:
+        """读取存储中所有按追加时序排列的历史条目。"""
+        async with self._lock():
+            self._remove_incomplete_temp()
+            if not self.path.exists():
+                return []
+
+            try:
+                with open(self.path, encoding="utf-8") as f:
+                    raw_lines = f.readlines()
+            except OSError as exc:
+                raise SessionJsonlError(f"Failed to read {self.path}: {exc}") from exc
+
+            lines = [ln.strip() for ln in raw_lines if ln.strip()]
+            if not lines:
+                return []
+
+            entries: list[SessionEntry] = []
+            for i, line in enumerate(lines):
+                is_last = i == len(lines) - 1
+                try:
+                    entry = entry_from_json_line(line)
+                    entries.append(entry)
+                except SessionJsonlError as exc:
+                    # 尾行撕裂容忍：如果是最后一行且前面已有有效条目，丢弃该损坏尾行（宽容兜底）
+                    if _should_tolerate_tail_error(is_last, len(entries)):
+                        break
+                    raise exc
+            return entries
+
+
+# ─── 私有辅助函数（对齐 Tau 内聚风格，置于模块末尾） ──────────────────────────
+
+
+def _should_tolerate_tail_error(is_last: bool, count: int) -> bool:
+    """尾行撕裂容忍：最后一行且前面已有有效条目时为 True。"""
+    return is_last and count > 0
+
+
+def _lock_file(fd: int) -> None:
+    """底层文件锁适配（非阻塞排他锁），内部按需判断操作系统。"""
+    if sys.platform == "win32":
+        import msvcrt
+
+        msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)
+        return
+
+    import fcntl
+
+    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+
+
+def _unlock_file(fd: int) -> None:
+    """底层文件解锁适配。"""
+    if sys.platform == "win32":
+        import msvcrt
+
+        with contextlib.suppress(OSError):
+            msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)
+        return
+
+    import fcntl
+
+    with contextlib.suppress(OSError):
+        fcntl.flock(fd, fcntl.LOCK_UN)
+
+
+__all__ = [
+    "SessionStorage",
+    "InMemorySessionStorage",
+    "JsonlSessionStorage",
+]