my-pi-agent 0.1.0

package/src/my_agent_core/registry.py

@@ -0,0 +1,186 @@
+"""工具注册表 —— 持有工具集合，按名字查表与执行。"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from my_agent_core.tools import Tool, ToolResult, tool
+
+__all__ = ["Tool", "ToolRegistry", "ToolResult", "tool"]
+
+
+class ToolRegistry:
+    """工具注册表：持有工具集合，按名字查表与执行。"""
+
+    def __init__(self):
+        self._tools: dict[str, Tool] = {}
+
+    def register(self, tool: Tool) -> None:
+        self._tools[tool.name] = tool
+
+    def unregister(self, name: str) -> None:
+        self._tools.pop(name, None)
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def list(self) -> list[Tool]:
+        """当前全部工具（发现顺序）。"""
+        return list(self._tools.values())
+
+    def get_schemas(self) -> list[dict[str, Any]]:
+        """生成全部工具的 OpenAI tools 参数。"""
+        return [t.to_openai_schema() for t in self._tools.values()]
+
+    async def execute_tool(
+        self,
+        tool_call: Any = None,
+        *,
+        name: str | None = None,
+        args: dict[str, Any] | str | None = None,
+        signal: Any | None = None,
+        on_update: Callable[[Any], None] | None = None,
+        tool_call_id: str | None = None,
+    ) -> ToolResult:
+        """执行单个 tool 调用。
+
+        支持多种入参形式：
+        1. 原生结构化传参: execute_tool(name="add", args={"a": 1})
+        2. 元组传参: execute_tool(("add", {"a": 1})) 或 execute_tool(("add", {"a": 1}, on_update))
+        3. 实体传参: execute_tool(ToolCall(id=..., name="add", args={"a": 1}))
+        4. 字典传参: execute_tool({"name": "add", "args": {"a": 1}})
+        5. 协议传参: execute_tool({"function": {"name": "add", "arguments": "..."}})
+        任何错误都转成 ToolResult，永不抛。
+        """
+        target_name = name
+        target_args = args
+
+        if target_name is None:
+            if isinstance(tool_call, tuple):
+                if len(tool_call) >= 2:
+                    target_name, target_args = str(tool_call[0]), tool_call[1]
+                if len(tool_call) >= 3 and on_update is None and callable(tool_call[2]):
+                    on_update = tool_call[2]  # pyright: ignore[reportAssignmentType]
+                if len(tool_call) >= 4 and tool_call_id is None:
+                    tool_call_id = str(tool_call[3])
+            elif isinstance(tool_call, dict):
+                if "function" in tool_call:
+                    fn = tool_call.get("function") or {}
+                    target_name = fn.get("name", "")
+                    target_args = fn.get("arguments", "{}")
+                else:
+                    target_name = tool_call.get("name", "")
+                    target_args = tool_call.get("args", {})
+                if tool_call_id is None:
+                    tool_call_id = tool_call.get("id")
+            elif isinstance(tool_call, str):
+                target_name = tool_call
+            elif tool_call is not None:
+                target_name = getattr(tool_call, "name", "")  # pyright: ignore[reportAttributeAccessIssue]
+                target_args = getattr(tool_call, "args", {})  # pyright: ignore[reportAttributeAccessIssue]
+                if tool_call_id is None:
+                    tool_call_id = getattr(tool_call, "id", None)  # pyright: ignore[reportAttributeAccessIssue]
+
+        final_args: dict[str, Any]
+        if isinstance(target_args, str):
+            try:
+                parsed = json.loads(target_args) if target_args.strip() else {}
+                if not isinstance(parsed, dict):
+                    return ToolResult(
+                        ok=False,
+                        error=f"Invalid JSON arguments for tool '{target_name}': expected object dict",
+                    )
+                final_args = parsed
+            except (json.JSONDecodeError, TypeError) as exc:
+                return ToolResult(
+                    ok=False,
+                    error=f"Invalid JSON arguments for tool '{target_name}': {exc}",
+                )
+        elif isinstance(target_args, dict):
+            final_args = target_args
+        elif target_args is None:
+            final_args = {}
+        else:
+            return ToolResult(
+                ok=False,
+                error=f"Invalid arguments for tool '{target_name}': expected dict, got {type(target_args).__name__}",
+            )
+
+        target = self._tools.get(target_name or "")
+        if target is None:
+            available = ", ".join(sorted(self._tools))
+            return ToolResult(
+                ok=False, error=f"Unknown tool '{target_name}'. Available: {available}"
+            )
+        return await target.execute(
+            final_args,
+            signal=signal,
+            on_update=on_update,
+            tool_call_id=tool_call_id,
+        )
+
+    async def execute_batch(
+        self,
+        tool_calls: Sequence[Any],
+        signal: Any | None = None,
+        on_updates: Sequence[Callable[[Any], None] | None] | None = None,
+        on_update_factory: Callable[[int, Any], Callable[[Any], None] | None]
+        | None = None,
+    ) -> list[ToolResult]:
+        """批量执行工具调用（全员只读并发；只要包含一个写入则整批保序串行，防止因果时序倒置）。"""
+        if not tool_calls:
+            return []
+
+        def _get_name(tc: Any) -> str:
+            if isinstance(tc, tuple) and len(tc) >= 2:
+                return str(tc[0])
+            if isinstance(tc, dict):
+                if "function" in tc:
+                    return str((tc.get("function") or {}).get("name", ""))
+                return str(tc.get("name", ""))
+            if tc is not None:
+                name_attr = getattr(tc, "name", None)  # pyright: ignore[reportAttributeAccessIssue]
+                if name_attr is not None:
+                    return str(name_attr)
+            return ""
+
+        def _get_update_cb(idx: int, tc: Any) -> Callable[[Any], None] | None:
+            if isinstance(tc, tuple) and len(tc) >= 3 and callable(tc[2]):
+                return tc[2]  # pyright: ignore[reportReturnType]
+            if on_updates is not None and idx < len(on_updates):
+                return on_updates[idx]
+            if on_update_factory is not None:
+                return on_update_factory(idx, tc)
+            return None
+
+        # 检查这批工具中是否包含任何不安全的写工具（或未知工具）
+        has_sequential = any(
+            (t := self._tools.get(_get_name(tc))) is None or not t.is_parallel_safe
+            for tc in tool_calls
+        )
+
+        if has_sequential:
+            # 只要包含一个写操作，整批严格按大模型输出的原始顺序串行执行，确保因果顺序绝对正确
+            results = []
+            for i, tc in enumerate(tool_calls):
+                results.append(
+                    await self.execute_tool(
+                        tc, signal=signal, on_update=_get_update_cb(i, tc)
+                    )
+                )
+            return results
+
+        # 全部都是只读安全工具时，安全并发执行
+        return list(
+            await asyncio.gather(
+                *(
+                    self.execute_tool(
+                        tc, signal=signal, on_update=_get_update_cb(i, tc)
+                    )
+                    for i, tc in enumerate(tool_calls)
+                )
+            )
+        )

package/src/my_agent_core/session/__init__.py

@@ -0,0 +1,79 @@
+# pyright: reportAttributeAccessIssue=false
+"""Session package: entries, tree, state projection, storage abstractions, and Session manager."""
+
+from __future__ import annotations
+
+from .entries import (
+    BaseSessionEntry,
+    BranchSummaryEntry,
+    CompactionEntry,
+    CustomEntry,
+    LabelEntry,
+    LeafEntry,
+    MessageEntry,
+    ModelChangeEntry,
+    SessionEntry,
+    SessionInfoEntry,
+    ThinkingLevelChangeEntry,
+)
+from .jsonl import (
+    SessionJsonlError,
+    entry_from_json_line,
+    entry_to_json_line,
+)
+from .memory import (
+    SessionState,
+)
+from .session import (
+    Session,
+    SessionTree,
+)
+from .storage import (
+    InMemorySessionStorage,
+    JsonlSessionStorage,
+    SessionStorage,
+)
+from .store import (
+    SessionMeta,
+    SessionStore,
+)
+from .tree import (
+    SessionTreeError,
+    entries_by_id,
+    lowest_common_ancestor,
+    path_to_entry,
+)
+
+__all__ = [
+    # Session manager & tree
+    "Session",
+    "SessionTree",
+    "SessionStore",
+    "SessionMeta",
+    # Modular entries
+    "BaseSessionEntry",
+    "SessionInfoEntry",
+    "MessageEntry",
+    "ModelChangeEntry",
+    "ThinkingLevelChangeEntry",
+    "CompactionEntry",
+    "BranchSummaryEntry",
+    "LabelEntry",
+    "LeafEntry",
+    "CustomEntry",
+    "SessionEntry",
+    # Tree algorithms
+    "SessionTreeError",
+    "entries_by_id",
+    "path_to_entry",
+    "lowest_common_ancestor",
+    # Memory projection
+    "SessionState",
+    # Storage protocol & driver
+    "SessionStorage",
+    "InMemorySessionStorage",
+    "JsonlSessionStorage",
+    "SessionJsonlError",
+    "entry_to_json_line",
+    "entry_from_json_line",
+]

package/src/my_agent_core/session/entries.py

@@ -0,0 +1,197 @@
+"""Session entry definitions for the modular session subsystem.
+
+Provides 9 polymorphic discriminated entry types backed by Pydantic v2,
+supporting camelCase and snake_case aliasing and strict extra field rejection.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Annotated, Any, Literal
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic.alias_generators import to_camel
+
+from my_agent_llm.models import Message
+
+
+class BaseSessionEntry(BaseModel):
+    """会话树节点基类。
+
+    提供自增 ID、父节点指针、浮点时间戳，并开启 camelCase 互转与未知字段禁止。
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",
+        alias_generator=to_camel,
+        populate_by_name=True,
+    )
+
+    id: str = Field(default_factory=lambda: uuid4().hex)
+    parent_id: str | None = None
+    timestamp: float = Field(default_factory=time.time)
+
+
+class SessionInfoEntry(BaseSessionEntry):
+    """记录会话元数据（工作目录、标题、创建时间），通常作为流首项。"""
+
+    type: Literal["session_info", "sessionInfo"] = "session_info"
+    cwd: str | None = None
+    title: str | None = None
+    name: str | None = None
+    created_at: float | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def _sync_title_name(self) -> SessionInfoEntry:
+        if self.name and not self.title:
+            self.title = self.name
+        elif self.title and not self.name:
+            self.name = self.title
+        return self
+
+
+class MessageEntry(BaseSessionEntry):
+    """包装 Agent 对话消息（User / Assistant / Tool 交互）。"""
+
+    type: Literal["message"] = "message"
+    message: Message
+
+    @property
+    def role(self) -> str:
+        return self.message.role
+
+    @property
+    def content(self) -> str:
+        return self.message.content
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        return self.message.metadata or {}
+
+
+class ModelChangeEntry(BaseSessionEntry):
+    """记录运行时大模型变更（兼容 Pi 原厂 modelId 字段）。"""
+
+    type: Literal["model_change", "modelChange"] = "model_change"
+    model: str | None = None
+    model_id: str | None = Field(default=None, alias="modelId")
+    provider: str | None = None
+
+    @model_validator(mode="after")
+    def _sync_model_fields(self) -> ModelChangeEntry:
+        if self.model_id and not self.model:
+            self.model = self.model_id
+        elif self.model and not self.model_id:
+            self.model_id = self.model
+        return self
+
+
+class ThinkingLevelChangeEntry(BaseSessionEntry):
+    """记录推理思考等级调整。"""
+
+    type: Literal["thinking_level_change", "thinkingLevelChange"] = "thinking_level_change"
+    thinking_level: str
+
+
+class CompactionEntry(BaseSessionEntry):
+    """记录上下文压缩覆盖的条目 ID 清单与折叠摘要（兼容 Pi 原厂 firstKeptEntryId 与 tokensBefore）。"""
+
+    type: Literal["compaction"] = "compaction"
+    summary: str
+    first_kept_entry_id: str | None = Field(default=None, alias="firstKeptEntryId")
+    tokens_before: int | None = Field(default=None, alias="tokensBefore")
+    replaces_entry_ids: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    @property
+    def role(self) -> str:
+        return "system"
+
+    @property
+    def content(self) -> str:
+        return self.summary
+
+
+class BranchSummaryEntry(BaseSessionEntry):
+    """记录分支折叠探索摘要。"""
+
+    type: Literal["branch_summary", "branchSummary"] = "branch_summary"
+    summary: str
+    details: dict[str, Any] = Field(default_factory=dict)
+
+
+class LabelEntry(BaseSessionEntry):
+    """用户书签/检查点。"""
+
+    type: Literal["label"] = "label"
+    label: str
+    target_id: str | None = None
+
+
+class LeafEntry(BaseSessionEntry):
+    """指向当前分支活动叶节点的指针（分支切换仅需追加一条 LeafEntry，零文件重写）。"""
+
+    type: Literal["leaf"] = "leaf"
+    leaf_id: str
+
+
+class CustomEntry(BaseSessionEntry):
+    """扩展与遥测隔离槽位。"""
+
+    type: Literal["custom"] = "custom"
+    namespace: str | None = None
+    custom_type: str | None = Field(default=None, alias="customType")
+    data: dict[str, Any] = Field(default_factory=dict)
+
+
+class CustomMessageEntry(BaseSessionEntry):
+    """扩展上下文注入消息条目（对标 Pi 原厂 CustomMessageEntry）。"""
+
+    type: Literal["custom_message", "customMessage"] = "custom_message"
+    custom_type: str = Field(default="custom", alias="customType")
+    content: str | list[Any] = ""
+    display: bool = True
+    details: dict[str, Any] = Field(default_factory=dict)
+
+
+class SessionHeaderEntry(BaseSessionEntry):
+    """Pi 原厂会话头节点。"""
+
+    type: Literal["session"] = "session"
+    version: int = 3
+    cwd: str = ""
+    parent_session: str | None = Field(default=None, alias="parentSession")
+
+
+SessionEntry = Annotated[
+    SessionInfoEntry
+    | SessionHeaderEntry
+    | MessageEntry
+    | ModelChangeEntry
+    | ThinkingLevelChangeEntry
+    | CompactionEntry
+    | BranchSummaryEntry
+    | LabelEntry
+    | LeafEntry
+    | CustomEntry
+    | CustomMessageEntry,
+    Field(discriminator="type"),
+]
+
+__all__ = [
+    "BaseSessionEntry",
+    "SessionInfoEntry",
+    "SessionHeaderEntry",
+    "MessageEntry",
+    "ModelChangeEntry",
+    "ThinkingLevelChangeEntry",
+    "CompactionEntry",
+    "BranchSummaryEntry",
+    "LabelEntry",
+    "LeafEntry",
+    "CustomEntry",
+    "CustomMessageEntry",
+    "SessionEntry",
+]

package/src/my_agent_core/session/jsonl.py

@@ -0,0 +1,60 @@
+"""JSONL format serialization, parsing, and error definitions for session entries.
+
+对齐 Tau (tau_agent.session.jsonl) 架构设计：
+本模块纯粹负责 SessionEntry 与单行 JSONL 文本之间的双向序列化与反序列化，
+以及行格式异常定义。持久化驱动与文件锁已归位至 storage.py。
+"""
+
+from __future__ import annotations
+
+import json
+
+from pydantic import TypeAdapter
+
+from .entries import SessionEntry
+
+
+class SessionJsonlError(ValueError):
+    """JSONL 序列化、反序列化或格式损坏异常。"""
+
+
+_ENTRY_ADAPTER: TypeAdapter[SessionEntry] = TypeAdapter(SessionEntry)
+
+
+def entry_to_json_line(entry: SessionEntry) -> str:
+    """将 SessionEntry 序列化为单行 JSON 字符串（不带尾随换行符）。"""
+    try:
+        return entry.model_dump_json(by_alias=True, exclude_none=True)
+    except Exception as exc:
+        raise SessionJsonlError(
+            f"Failed to serialize entry to JSON line: {exc}"
+        ) from exc
+
+
+def entry_from_json_line(line: str) -> SessionEntry:
+    """从单行 JSON 字符串解析 SessionEntry。"""
+    stripped = line.strip()
+    if not stripped:
+        raise SessionJsonlError("JSONL line is empty or blank")
+
+    try:
+        raw = json.loads(stripped)
+    except json.JSONDecodeError as exc:
+        raise SessionJsonlError(f"Invalid JSON line: {exc}") from exc
+
+    if not isinstance(raw, dict):
+        raise SessionJsonlError(f"Expected JSON object, got {type(raw).__name__}")
+
+    try:
+        return _ENTRY_ADAPTER.validate_python(raw)
+    except Exception as exc:
+        raise SessionJsonlError(
+            f"Cannot parse JSON object as SessionEntry: {exc}"
+        ) from exc
+
+
+__all__ = [
+    "SessionJsonlError",
+    "entry_to_json_line",
+    "entry_from_json_line",
+]

package/src/my_agent_core/session/memory.py

@@ -0,0 +1,137 @@
+"""Pure in-memory state fold projection for session entries.
+
+Provides the immutable SessionState dataclass and pure function folding over
+path entries from root to leaf, applying model changes, thinking levels,
+labels, leaf pointers, and compaction history replacements.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from my_agent_llm.models import Message
+
+from .entries import (
+    CompactionEntry,
+    LabelEntry,
+    LeafEntry,
+    MessageEntry,
+    ModelChangeEntry,
+    SessionEntry,
+    ThinkingLevelChangeEntry,
+)
+from .tree import path_to_entry
+
+
+def _apply_compaction(
+    items: list[tuple[str, Message]],
+    entry: CompactionEntry,
+) -> list[tuple[str, Message]]:
+    """将 replaces_entry_ids 范围内的消息条目折叠为一条 UserMessage 摘要。"""
+    summary_text = entry.summary
+    prefix = "Previous conversation summary:\n"
+    if summary_text.startswith(prefix):
+        content = summary_text
+    else:
+        content = f"{prefix}{summary_text}"
+
+    summary_msg = Message(role="user", content=content)
+    replaces_set = set(entry.replaces_entry_ids)
+
+    if not replaces_set:
+        return [*items, (entry.id, summary_msg)]
+
+    new_items: list[tuple[str, Message]] = []
+    inserted = False
+    for eid, msg in items:
+        if eid in replaces_set:
+            if not inserted:
+                new_items.append((entry.id, summary_msg))
+                inserted = True
+        else:
+            new_items.append((eid, msg))
+
+    if not inserted:
+        new_items.append((entry.id, summary_msg))
+
+    return new_items
+
+
+@dataclass(frozen=True, slots=True)
+class SessionState:
+    """不可变运行时会话状态快照。
+
+    通过事件溯源沿树状 DAG 路径纯函数折叠聚合生成，无锁且线程安全。
+    """
+
+    messages: tuple[Message, ...] = ()
+    model: str | None = None
+    provider: str | None = None
+    thinking_level: str | None = None
+    label: str | None = None
+    active_leaf_id: str | None = None
+
+    @classmethod
+    def from_entries(
+        cls,
+        entries: Sequence[SessionEntry],
+        leaf_id: str | None = None,
+    ) -> SessionState:
+        """纯函数折叠投影：沿根至 leaf_id 的路径条目无锁计算出最新运行时状态。"""
+        if not entries:
+            return cls()
+
+        # 确定目标叶节点 ID
+        if leaf_id is not None:
+            target_leaf_id = leaf_id
+        else:
+            # 若存在 LeafEntry 指针，取最后一条 LeafEntry 指定的叶节点
+            target_leaf_id = None
+            for entry in reversed(entries):
+                if isinstance(entry, LeafEntry):
+                    target_leaf_id = entry.leaf_id
+                    break
+            if target_leaf_id is None:
+                target_leaf_id = entries[-1].id
+
+        # 提取从根至目标叶节点的有序单链路径
+        path = path_to_entry(entries, target_leaf_id)
+
+        # 沿路径纯函数折叠投影
+        items: list[tuple[str, Message]] = []
+        model: str | None = None
+        provider: str | None = None
+        thinking_level: str | None = None
+        label: str | None = None
+        active_leaf_id: str | None = target_leaf_id
+
+        for entry in path:
+            if isinstance(entry, MessageEntry):
+                items.append((entry.id, entry.message))
+            elif isinstance(entry, CompactionEntry):
+                items = _apply_compaction(items, entry)
+            elif isinstance(entry, ModelChangeEntry):
+                model = entry.model
+                if entry.provider is not None:
+                    provider = entry.provider
+            elif isinstance(entry, ThinkingLevelChangeEntry):
+                thinking_level = entry.thinking_level
+            elif isinstance(entry, LabelEntry):
+                label = entry.label
+            elif isinstance(entry, LeafEntry):
+                active_leaf_id = entry.leaf_id
+
+        return cls(
+            messages=tuple(msg for _, msg in items),
+            model=model,
+            provider=provider,
+            thinking_level=thinking_level,
+            label=label,
+            active_leaf_id=active_leaf_id,
+        )
+
+
+__all__ = [
+    "SessionState",
+]