yycode 0.3.2__py3-none-any.whl

agent/session_replay.py

@@ -0,0 +1,95 @@
+"""Replay view model derived from canonical session messages."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
+
+from agent.task_memory import is_task_summary_memory
+
+
+ReplayRole = Literal["user", "assistant", "tool", "system"]
+ReplayKind = Literal["message", "summary", "tool", "context"]
+
+
+@dataclass(frozen=True)
+class ReplayEvent:
+    """One UI/protocol replay event derived from a stored message."""
+
+    role: ReplayRole
+    kind: ReplayKind
+    content: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+def build_session_replay(messages: list[BaseMessage]) -> list[ReplayEvent]:
+    """Build a display-friendly replay model from canonical session messages."""
+    events: list[ReplayEvent] = []
+    for index, message in enumerate(messages):
+        metadata = {
+            "message_index": index,
+            "message_type": type(message).__name__,
+        }
+        metadata.update(getattr(message, "additional_kwargs", {}) or {})
+        if is_task_summary_memory(message):
+            events.append(
+                ReplayEvent(
+                    role="system",
+                    kind="summary",
+                    content=_message_text(message),
+                    metadata=metadata,
+                )
+            )
+        elif isinstance(message, HumanMessage):
+            events.append(
+                ReplayEvent(
+                    role="user",
+                    kind="message",
+                    content=_message_text(message),
+                    metadata=metadata,
+                )
+            )
+        elif isinstance(message, AIMessage):
+            content = _message_text(message)
+            if not content.strip():
+                continue
+            events.append(
+                ReplayEvent(
+                    role="assistant",
+                    kind="message",
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+        elif isinstance(message, ToolMessage):
+            content = _message_text(message)
+            if not _should_replay_tool_message(content, metadata):
+                continue
+            metadata["tool_name"] = message.name or ""
+            metadata["tool_call_id"] = message.tool_call_id
+            events.append(
+                ReplayEvent(
+                    role="tool",
+                    kind="tool",
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+    return events
+
+
+def _should_replay_tool_message(content: str, metadata: dict[str, Any]) -> bool:
+    if metadata.get("context_compressed"):
+        return True
+    return len(content) <= 2_000
+
+
+def _message_text(message: BaseMessage) -> str:
+    content = getattr(message, "content", "")
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        return "\n".join(str(item) for item in content)
+    return str(content)

agent/session_store.py

@@ -0,0 +1,186 @@
+"""Persistence for canonical session message history."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict
+
+
+SESSION_VERSION = 1
+
+
+class SessionStoreError(RuntimeError):
+    """Raised when persisted session data cannot be safely used."""
+
+
+@dataclass(frozen=True)
+class SessionRecord:
+    """Small listing record for a persisted session."""
+
+    session_id: str
+    path: Path
+    updated_at: str
+    workdir: str
+
+
+class SessionStore:
+    """Abstract session message store."""
+
+    def load(self, session_id: str) -> list[BaseMessage]:
+        raise NotImplementedError
+
+    def save(
+        self,
+        session_id: str,
+        messages: list[BaseMessage],
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def delete(self, session_id: str) -> None:
+        raise NotImplementedError
+
+    def list_sessions(self) -> list[SessionRecord]:
+        raise NotImplementedError
+
+
+class FileSessionStore(SessionStore):
+    """File-backed message store grouped by workspace hash."""
+
+    def __init__(
+        self,
+        app_root: Path,
+        workdir: Path,
+        root: Path | str | None = None,
+    ) -> None:
+        self.app_root = app_root.expanduser().resolve()
+        self.workdir = workdir.expanduser().resolve()
+        raw_root = root or os.environ.get("YOYO_SESSION_DIR")
+        self.root = Path(raw_root).expanduser().resolve() if raw_root else self.app_root / "sessions"
+        self.workspace_hash = workspace_hash(self.workdir)
+
+    @property
+    def workspace_dir(self) -> Path:
+        return self.root / self.workspace_hash
+
+    def load(self, session_id: str) -> list[BaseMessage]:
+        path = self._session_path(session_id)
+        if not path.exists():
+            return []
+        try:
+            payload = json.loads(path.read_text(encoding="utf-8"))
+        except json.JSONDecodeError as exc:
+            raise SessionStoreError(f"Session file is not valid JSON: {path}") from exc
+
+        if payload.get("version") != SESSION_VERSION:
+            raise SessionStoreError(f"Unsupported session file version: {payload.get('version')}")
+
+        saved_workdir = Path(str(payload.get("workdir", ""))).expanduser().resolve()
+        if saved_workdir != self.workdir:
+            raise SessionStoreError(
+                f"Session workdir mismatch: saved {saved_workdir}, current {self.workdir}"
+            )
+
+        try:
+            return list(messages_from_dict(payload.get("messages", [])))
+        except Exception as exc:
+            raise SessionStoreError(f"Could not deserialize session messages: {path}") from exc
+
+    def save(
+        self,
+        session_id: str,
+        messages: list[BaseMessage],
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
+        path = self._session_path(session_id)
+        now = _utc_now()
+        existing = self._read_existing(path)
+        payload = {
+            "version": SESSION_VERSION,
+            "session_id": session_id,
+            "created_at": existing.get("created_at") or now,
+            "updated_at": now,
+            "workdir": str(self.workdir),
+            "workspace_hash": self.workspace_hash,
+            "app_root": str(self.app_root),
+            "model": (metadata or {}).get("model"),
+            "messages": messages_to_dict(messages),
+        }
+        path.parent.mkdir(parents=True, exist_ok=True)
+        _atomic_write_json(path, payload)
+
+    def delete(self, session_id: str) -> None:
+        path = self._session_path(session_id)
+        if path.exists():
+            path.unlink()
+
+    def list_sessions(self) -> list[SessionRecord]:
+        if not self.workspace_dir.exists():
+            return []
+        records: list[SessionRecord] = []
+        for path in sorted(self.workspace_dir.glob("*.json")):
+            payload = self._read_existing(path)
+            records.append(
+                SessionRecord(
+                    session_id=str(payload.get("session_id") or path.stem),
+                    path=path,
+                    updated_at=str(payload.get("updated_at") or ""),
+                    workdir=str(payload.get("workdir") or ""),
+                )
+            )
+        return records
+
+    def _session_path(self, session_id: str) -> Path:
+        safe_id = _safe_session_id(session_id)
+        return self.workspace_dir / f"{safe_id}.json"
+
+    def _read_existing(self, path: Path) -> dict[str, Any]:
+        if not path.exists():
+            return {}
+        try:
+            payload = json.loads(path.read_text(encoding="utf-8"))
+        except json.JSONDecodeError:
+            return {}
+        return payload if isinstance(payload, dict) else {}
+
+
+def workspace_hash(workdir: Path) -> str:
+    """Return a short stable hash for a resolved workspace path."""
+    normalized = str(workdir.expanduser().resolve())
+    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()[:16]
+
+
+def _safe_session_id(session_id: str) -> str:
+    if not session_id or session_id in {".", ".."}:
+        raise SessionStoreError("Session id must not be empty")
+    if any(separator in session_id for separator in ("/", "\\")):
+        raise SessionStoreError("Session id must not contain path separators")
+    return session_id
+
+
+def _utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
+    data = json.dumps(payload, ensure_ascii=False, indent=2)
+    with tempfile.NamedTemporaryFile(
+        "w",
+        encoding="utf-8",
+        dir=path.parent,
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+        delete=False,
+    ) as tmp:
+        tmp.write(data)
+        tmp.write("\n")
+        tmp_path = Path(tmp.name)
+    tmp_path.replace(path)

agent/skills.py

@@ -0,0 +1,254 @@
+"""Skill discovery and loading utilities."""
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Optional
+
+
+MAX_SKILL_CHARS = 8_000
+ALL_SKILLS_TOKENS = {"all", "*"}
+DEFAULT_SKILL_DIRS = ["skills"]
+
+
+@dataclass(frozen=True)
+class LoadedSkill:
+    """A loaded skill document."""
+
+    name: str
+    description: str
+    path: Path
+    content: str
+
+
+def parse_skill_paths(raw_paths: Optional[str]) -> list[str]:
+    """Parse a comma, newline, or os.pathsep separated string."""
+    if not raw_paths:
+        return []
+
+    paths = []
+    for chunk in raw_paths.replace("\n", ",").replace(os.pathsep, ",").split(","):
+        path = chunk.strip()
+        if path:
+            paths.append(path)
+    return paths
+
+
+class SkillRegistry:
+    """Discover and load skills for a workspace."""
+
+    def __init__(self, workdir: Path, skill_dirs: Optional[Iterable[str]] = None):
+        self.workdir = workdir
+        dirs = list(skill_dirs) if skill_dirs is not None else DEFAULT_SKILL_DIRS
+        self.search_dirs = [_resolve_path(path, workdir) for path in dirs]
+
+    def list_skills(self) -> list[LoadedSkill]:
+        """Return metadata for all discoverable skills."""
+        skills: list[LoadedSkill] = []
+        seen_names: set[str] = set()
+        seen_paths: set[Path] = set()
+
+        for skill_file in _iter_all_skill_files(self.search_dirs):
+            if skill_file in seen_paths:
+                continue
+            skill = _load_skill_file(skill_file)
+            if not skill or skill.name in seen_names:
+                continue
+            skills.append(skill)
+            seen_names.add(skill.name)
+            seen_paths.add(skill_file)
+
+        return sorted(skills, key=lambda skill: skill.name)
+
+    def load_skills(self, refs: Iterable[str]) -> list[LoadedSkill]:
+        """Load skills by name, path, or all token."""
+        loaded: list[LoadedSkill] = []
+        seen_paths: set[Path] = set()
+
+        for ref in refs:
+            for skill_file in _resolve_skill_files(ref, self.workdir, self.search_dirs):
+                if skill_file in seen_paths:
+                    continue
+                skill = _load_skill_file(skill_file)
+                if skill:
+                    loaded.append(skill)
+                    seen_paths.add(skill_file)
+
+        return loaded
+
+    def format_skill_list(self) -> str:
+        """Format skill metadata for tool output."""
+        skills = self.list_skills()
+        if not skills:
+            return "No skills found."
+
+        lines = ["Available skills:"]
+        for skill in skills:
+            description = skill.description or "(no description)"
+            lines.append(f"- {skill.name}: {description}")
+        return "\n".join(lines)
+
+    def format_skill_catalog_prompt(self) -> str:
+        """Format skill metadata for prompt injection without full content."""
+        skills = self.list_skills()
+        if not skills:
+            return ""
+
+        lines = [
+            "Available local skills:",
+            "You only have skill names and descriptions by default.",
+            "Call load_skill to load the full instructions for any skill you want to use.",
+        ]
+        for skill in skills:
+            description = skill.description or "(no description)"
+            lines.append(f"- {skill.name}: {description}")
+        return "\n".join(lines)
+
+    def format_loaded_skills(self, refs: Iterable[str]) -> str:
+        """Format full loaded skills for tool output."""
+        skills = self.load_skills(refs)
+        if not skills:
+            return "No matching skills found."
+
+        sections = ["Loaded skills:"]
+        for skill in skills:
+            content = skill.content
+            if len(content) > MAX_SKILL_CHARS:
+                content = content[:MAX_SKILL_CHARS] + "\n\n[Skill content truncated.]"
+            sections.append(f"## {skill.name}")
+            sections.append(f"Source: {skill.path}")
+            if skill.description:
+                sections.append(f"Description: {skill.description}")
+            sections.append("")
+            sections.append(content)
+            sections.append("")
+        return "\n".join(sections).rstrip()
+
+
+def load_skills(
+    skill_paths: Iterable[str],
+    workdir: Path,
+    skill_dirs: Optional[Iterable[str]] = None,
+) -> list[LoadedSkill]:
+    """Backward-compatible wrapper for loading skills."""
+    return SkillRegistry(workdir, skill_dirs).load_skills(skill_paths)
+
+
+def discover_skills(workdir: Path, skill_dirs: Optional[Iterable[str]] = None) -> list[LoadedSkill]:
+    """Discover all skills visible to the workspace."""
+    return SkillRegistry(workdir, skill_dirs).list_skills()
+
+
+def _resolve_path(raw_path: str, workdir: Path) -> Path:
+    path = Path(raw_path).expanduser()
+    if not path.is_absolute():
+        path = workdir / path
+    return path.resolve()
+
+
+def _resolve_skill_files(raw_ref: str, workdir: Path, search_dirs: list[Path]) -> list[Path]:
+    ref = raw_ref.strip()
+    if ref.lower() in ALL_SKILLS_TOKENS:
+        return _iter_all_skill_files(search_dirs)
+
+    explicit_path = _resolve_path(ref, workdir)
+    if explicit_path.exists():
+        return _iter_skill_files(explicit_path)
+
+    for search_dir in search_dirs:
+        for candidate in _named_skill_candidates(ref, search_dir):
+            if candidate.exists():
+                return _iter_skill_files(candidate)
+    return []
+
+
+def _named_skill_candidates(name: str, search_dir: Path) -> list[Path]:
+    return [
+        search_dir / name,
+        search_dir / f"{name}.md",
+        search_dir / name / "SKILL.md",
+    ]
+
+
+def _iter_all_skill_files(search_dirs: list[Path]) -> list[Path]:
+    skill_files: list[Path] = []
+    seen_paths: set[Path] = set()
+
+    for search_dir in search_dirs:
+        if not search_dir.is_dir():
+            continue
+        for skill_file in _iter_skill_files(search_dir):
+            if skill_file not in seen_paths:
+                skill_files.append(skill_file)
+                seen_paths.add(skill_file)
+        for child in sorted(search_dir.iterdir()):
+            if not child.is_dir():
+                continue
+            for skill_file in _iter_skill_files(child):
+                if skill_file not in seen_paths:
+                    skill_files.append(skill_file)
+                    seen_paths.add(skill_file)
+    return skill_files
+
+
+def _iter_skill_files(path: Path) -> list[Path]:
+    if path.is_file():
+        return [path]
+    if not path.is_dir():
+        return []
+
+    skill_md = path / "SKILL.md"
+    if skill_md.is_file():
+        return [skill_md]
+    return sorted(p for p in path.glob("*.md") if p.is_file())
+
+
+def _load_skill_file(path: Path) -> Optional[LoadedSkill]:
+    content = _read_skill_file(path)
+    if not content:
+        return None
+    metadata, body = _parse_skill_content(content)
+    return LoadedSkill(
+        name=metadata.get("name") or _skill_name(path),
+        description=metadata.get("description", ""),
+        path=path,
+        content=body,
+    )
+
+
+def _read_skill_file(path: Path) -> str:
+    try:
+        return path.read_text(encoding="utf-8", errors="backslashreplace").strip()
+    except OSError:
+        return ""
+
+
+def _skill_name(path: Path) -> str:
+    if path.name == "SKILL.md":
+        return path.parent.name
+    return path.stem
+
+
+def _parse_skill_content(content: str) -> tuple[dict[str, str], str]:
+    if not content.startswith("---\n"):
+        return {}, content
+
+    match = re.match(r"^---\n(.*?)\n---\n?(.*)$", content, re.DOTALL)
+    if not match:
+        return {}, content
+
+    raw_metadata, body = match.groups()
+    metadata: dict[str, str] = {}
+    for line in raw_metadata.splitlines():
+        if ":" not in line:
+            continue
+        key, value = line.split(":", 1)
+        metadata[key.strip()] = _strip_quotes(value.strip())
+    return metadata, body.strip()
+
+
+def _strip_quotes(value: str) -> str:
+    if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
+        return value[1:-1]
+    return value