pico-cli-core 0.1.1__py3-none-any.whl

pico_core/session.py

@@ -0,0 +1,192 @@
+"""Append-only, tree-based session model (ADR-0002).
+
+A session is a tree of immutable nodes. Each node carries an id, a parent
+pointer, a timestamp, and a payload. Nodes are never edited or deleted — only
+built upon. Branches are root-to-leaf timelines; forking rewinds to an earlier
+node and starts a new branch.
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from pico_ai.types import ToolCall, Usage
+
+
+def new_id() -> str:
+    """Return a fresh, unique node/session id."""
+    return uuid.uuid4().hex
+
+
+def utc_now() -> str:
+    """Return the current UTC timestamp as an ISO-8601 string."""
+    return datetime.now(timezone.utc).isoformat()
+
+
+class UserPayload(BaseModel):
+    """A user message."""
+
+    kind: Literal["user"] = "user"
+    content: str
+
+
+class AssistantBlock(BaseModel):
+    """A block of an assistant response: prose or reasoning."""
+
+    kind: Literal["text", "thinking"]
+    text: str = ""
+    thinking: str = ""
+
+
+class AssistantPayload(BaseModel):
+    """An assistant response, block-granular (text / thinking)."""
+
+    kind: Literal["assistant"] = "assistant"
+    blocks: list[AssistantBlock] = Field(default_factory=list)
+    usage: Usage | None = None
+    #: Provider-stream wall-time in milliseconds (None = unknown, e.g.
+    #: sessions persisted before trace timing existed).
+    duration_ms: float | None = None
+
+    @property
+    def text(self) -> str:
+        return "".join(b.text for b in self.blocks if b.kind == "text")
+
+
+class ToolRequestPayload(BaseModel):
+    """A request to run a tool."""
+
+    kind: Literal["tool_request"] = "tool_request"
+    tool_call: ToolCall
+
+
+class ToolResultPayload(BaseModel):
+    """The output of running a tool."""
+
+    kind: Literal["tool_result"] = "tool_result"
+    tool_call_id: str
+    name: str
+    content: str
+    is_error: bool = False
+
+
+class CompactionSummaryPayload(BaseModel):
+    """A summary of older turns produced by compaction."""
+
+    kind: Literal["compaction_summary"] = "compaction_summary"
+    summary: str
+
+
+Payload = Annotated[
+    UserPayload
+    | AssistantPayload
+    | ToolRequestPayload
+    | ToolResultPayload
+    | CompactionSummaryPayload,
+    Field(discriminator="kind"),
+]
+
+
+class Node(BaseModel):
+    """An immutable, append-only unit of event data in a session."""
+
+    id: str = Field(default_factory=new_id)
+    parent_id: str | None = None
+    timestamp: str = Field(default_factory=utc_now)
+    payload: Payload
+
+
+class Session(BaseModel):
+    """A tree of nodes with a single active branch."""
+
+    id: str = Field(default_factory=new_id)
+    nodes: dict[str, Node] = Field(default_factory=dict)
+    active_leaf_id: str | None = None
+
+    # -- construction -------------------------------------------------------
+
+    def append(self, parent_id: str | None, payload: Payload) -> Node:
+        """Append a node as a child of ``parent_id`` and make it the active leaf."""
+        if parent_id is not None and parent_id not in self.nodes:
+            raise KeyError(f"unknown parent node id: {parent_id}")
+        node = Node(parent_id=parent_id, payload=payload)
+        self.nodes[node.id] = node
+        self.active_leaf_id = node.id
+        return node
+
+    def fork(self, node_id: str) -> None:
+        """Rewind to ``node_id``, starting a new branch from it."""
+        if node_id not in self.nodes:
+            raise KeyError(f"unknown node id: {node_id}")
+        self.active_leaf_id = node_id
+
+    # -- traversal ----------------------------------------------------------
+
+    def active_branch(self) -> list[Node]:
+        """Return the root-to-leaf timeline of the active branch."""
+        if self.active_leaf_id is None:
+            return []
+        branch: list[Node] = []
+        current: Node | None = self.nodes[self.active_leaf_id]
+        while current is not None:
+            branch.append(current)
+            current = self.nodes.get(current.parent_id) if current.parent_id else None
+        branch.reverse()
+        return branch
+
+    def leaves(self) -> list[Node]:
+        """Return every leaf node (nodes with no children)."""
+        child_ids = {n.parent_id for n in self.nodes.values() if n.parent_id}
+        return [n for n in self.nodes.values() if n.id not in child_ids]
+
+    # -- persistence --------------------------------------------------------
+
+    def to_jsonl(self) -> str:
+        """Serialize the session as newline-delimited JSON."""
+        meta = json.dumps(
+            {"session_id": self.id, "active_leaf_id": self.active_leaf_id}
+        )
+        lines = [meta]
+        lines.extend(node.model_dump_json() for node in self.nodes.values())
+        return "\n".join(lines) + "\n"
+
+    @classmethod
+    def from_jsonl(cls, text: str) -> "Session":
+        """Rebuild a session from newline-delimited JSON."""
+        session: Session | None = None
+        for line in text.splitlines():
+            if not line.strip():
+                continue
+            data = json.loads(line)
+            if "session_id" in data and "parent_id" not in data:
+                session = cls(
+                    id=data["session_id"], active_leaf_id=data["active_leaf_id"]
+                )
+                continue
+            if session is None:
+                session = cls()
+            node = Node.model_validate(data)
+            session.nodes[node.id] = node
+        if session is None:
+            session = cls()
+        if session.active_leaf_id is None and session.nodes:
+            # No meta header present: the last-appended node is the active leaf.
+            session.active_leaf_id = next(reversed(session.nodes))
+        return session
+
+    def save(self, path: Path) -> None:
+        """Persist the session to ``path`` (parent dirs are created)."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(self.to_jsonl(), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: Path) -> "Session":
+        """Load a session persisted by :meth:`save`."""
+        return cls.from_jsonl(Path(path).read_text(encoding="utf-8"))

pico_core/subagents.py

@@ -0,0 +1,183 @@
+"""Model-invoked sub-agents: the hardcoded ``task`` tool (see ADR-0005).
+
+A parent loop delegates a self-contained chunk of work by calling ``task``
+with a short description and a full prompt. A fresh child ``AgentLoop`` runs
+on its own ``Session`` to completion; only the child's final summary returns
+to the parent as the tool result, so the parent's context stays lean while
+the child's full transcript persists as its own session file.
+
+Children get a restricted toolset by default (read-only research tools),
+fresh todos, and the parent's provider/model unless overridden per-spawn.
+Recursion is bounded: ``task`` is only registered in a child when the
+parent explicitly grants it, and ``max_depth`` caps the nesting — a spawn
+requested beyond it returns an error result instead of running.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+from typing import Any, Protocol
+
+from pydantic import BaseModel, Field
+
+from .fsm import AgentLoop, AgentState
+from .tools import ToolOutcome
+
+
+#: Default child toolset: read-only research tools. No writes, no bash, no
+#: todos (child progress stays out of the parent's todo contract), and no
+#: nested spawning unless the parent explicitly grants ``task``.
+DEFAULT_CHILD_TOOLS = ["read", "grep", "fetch", "websearch"]
+
+#: Maximum spawn nesting (the top-level loop runs at depth 0).
+MAX_DEPTH = 2
+
+#: Default per-child turn cap (streaming iterations of the child loop).
+DEFAULT_MAX_TURNS = 25
+
+#: Default per-child wall-clock budget in seconds.
+DEFAULT_TIMEOUT_S = 300.0
+
+#: Child result text beyond this is cut (with a note) before returning.
+MAX_RESULT_CHARS = 8000
+
+
+class ChildSpec(BaseModel):
+    """What to run in the child loop."""
+
+    description: str = ""
+    prompt: str
+    allowed_tools: list[str] | None = None
+    model: str | None = None
+    max_turns: int = Field(default=DEFAULT_MAX_TURNS, ge=1)
+    timeout_s: float = Field(default=DEFAULT_TIMEOUT_S, gt=0)
+
+
+class ChildFactory(Protocol):
+    """Builds a child loop for ``spec`` at nesting ``depth``."""
+
+    def __call__(self, spec: ChildSpec, depth: int) -> AgentLoop: ...
+
+
+class SpawnTool:
+    """Run a sub-agent to completion and return its summary."""
+
+    name = "task"
+    description = (
+        "Delegate a self-contained task to a sub-agent. Give a short "
+        "description and a full prompt; the sub-agent works autonomously "
+        "with its own tools and context, and its final summary is returned. "
+        "Use for parallelizable research or isolated chunks of work whose "
+        "details the parent does not need."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "description": {"type": "string"},
+            "prompt": {"type": "string"},
+            "allowed_tools": {
+                "type": "array",
+                "items": {"type": "string"},
+            },
+            "model": {"type": "string"},
+            "max_turns": {"type": "integer"},
+            "timeout_s": {"type": "number"},
+        },
+        "required": ["prompt"],
+    }
+
+    def __init__(
+        self,
+        factory: ChildFactory,
+        *,
+        depth: int = 0,
+        session_dir: Path | str = "~/.pico/sessions",
+        default_tools: list[str] | None = None,
+        max_depth: int = MAX_DEPTH,
+        max_result_chars: int = MAX_RESULT_CHARS,
+    ) -> None:
+        self._factory = factory
+        self._depth = depth
+        self._session_dir = Path(session_dir).expanduser()
+        self._default_tools = (
+            list(default_tools) if default_tools is not None else list(DEFAULT_CHILD_TOOLS)
+        )
+        self._max_depth = max_depth
+        self._max_result_chars = max_result_chars
+
+    async def run(self, arguments: dict) -> ToolOutcome:
+        prompt = (arguments.get("prompt") or "")
+        if not isinstance(prompt, str) or not prompt.strip():
+            return ToolOutcome(
+                content="error: task requires a non-empty prompt", is_error=True
+            )
+        if self._depth >= self._max_depth:
+            return ToolOutcome(
+                content=(
+                    f"error: sub-agent nesting limit reached "
+                    f"(depth {self._depth} >= max {self._max_depth})"
+                ),
+                is_error=True,
+            )
+        spec_kwargs: dict[str, Any] = {
+            "description": arguments.get("description", ""),
+            "prompt": prompt.strip(),
+        }
+        for key in ("allowed_tools", "model", "max_turns", "timeout_s"):
+            if arguments.get(key) is not None:
+                spec_kwargs[key] = arguments[key]
+        try:
+            spec = ChildSpec.model_validate(spec_kwargs)
+        except Exception as exc:  # noqa: BLE001 - invalid args are results
+            return ToolOutcome(content=f"error: invalid task: {exc}", is_error=True)
+        if spec.allowed_tools is None:
+            spec = spec.model_copy(update={"allowed_tools": self._default_tools})
+        try:
+            loop = self._factory(spec, self._depth + 1)
+        except Exception as exc:  # noqa: BLE001 - surface as a result
+            return ToolOutcome(
+                content=f"error: could not start sub-agent: {exc}", is_error=True
+            )
+        try:
+            result = await asyncio.wait_for(
+                loop.run(prompt.strip(), max_turns=spec.max_turns),
+                timeout=spec.timeout_s,
+            )
+        except asyncio.TimeoutError:
+            return ToolOutcome(
+                content=(
+                    f"[subagent {loop.session.id} timed out after "
+                    f"{spec.timeout_s:g}s]"
+                ),
+                is_error=True,
+            )
+        except Exception as exc:  # noqa: BLE001 - surface as a result
+            return ToolOutcome(
+                content=f"[subagent {loop.session.id} crashed: {exc}]",
+                is_error=True,
+            )
+        # Best-effort persistence: the child's transcript lives as its own
+        # session file; a save failure must not fail the delegation itself.
+        try:
+            loop.session.save(self._session_dir / f"{loop.session.id}.jsonl")
+        except OSError:
+            pass
+        notes = []
+        if result.truncated:
+            notes.append(f"truncated at max_turns={spec.max_turns}")
+        header = f"[subagent {loop.session.id} done"
+        if notes:
+            header += f" ({'; '.join(notes)})"
+        header += "]"
+        content = f"{header}\n{result.text}" if result.text else header
+        if result.state == AgentState.ERROR:
+            content = f"{content}\n[subagent ended in error: {result.error}]"
+            return ToolOutcome(content=self._cut(content), is_error=True)
+        return ToolOutcome(content=self._cut(content))
+
+    def _cut(self, content: str) -> str:
+        if len(content) <= self._max_result_chars:
+            return content
+        cut = len(content) - self._max_result_chars
+        return content[: self._max_result_chars] + f"\n[... truncated {cut} chars ...]"

pico_core/todos.py

@@ -0,0 +1,150 @@
+"""In-memory todo tracking: a shared list plus the agent-facing ``todo`` tool.
+
+The list lives only for the life of the process (no persistence): the owning
+``AgentSession`` creates one :class:`TodoList`, hands it to the ``TodoTool``,
+and the TUI reads it back to render the read-only side panel.
+"""
+
+from __future__ import annotations
+
+import itertools
+
+from pydantic import BaseModel
+
+from .tools import ToolOutcome
+
+
+class TodoItem(BaseModel):
+    """A single tracked task."""
+
+    id: str
+    text: str
+    status: str = "pending"  # "pending" | "in_progress" | "completed"
+
+
+class TodoList:
+    """An ordered, in-memory list of todos shared by the tool and the UI."""
+
+    def __init__(self) -> None:
+        self._items: list[TodoItem] = []
+        self._ids = itertools.count(1)
+
+    def add(self, text: str) -> TodoItem:
+        """Append a pending todo and return it."""
+        item = TodoItem(id=f"t{next(self._ids)}", text=text)
+        self._items.append(item)
+        return item
+
+    def update(
+        self, todo_id: str, *, status: str | None = None, text: str | None = None
+    ) -> TodoItem | None:
+        """Update a todo's status and/or text; ``None`` when the id is unknown."""
+        item = self.get(todo_id)
+        if item is None:
+            return None
+        if status is not None:
+            item.status = status
+        if text is not None:
+            item.text = text
+        return item
+
+    def get(self, todo_id: str) -> TodoItem | None:
+        """Return the todo with ``todo_id``, or ``None``."""
+        for item in self._items:
+            if item.id == todo_id:
+                return item
+        return None
+
+    def all(self) -> list[TodoItem]:
+        """Return every todo in insertion order."""
+        return list(self._items)
+
+    def clear_completed(self) -> int:
+        """Drop every completed todo; return how many were removed."""
+        before = len(self._items)
+        self._items = [i for i in self._items if i.status != "completed"]
+        return before - len(self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+
+_VALID_STATUSES = ("pending", "in_progress", "completed")
+
+
+def format_todos(items: list[TodoItem]) -> str:
+    """Render todos as one ``[id] status — text`` line each."""
+    if not items:
+        return "(no todos)"
+    return "\n".join(f"[{i.id}] {i.status} — {i.text}" for i in items)
+
+
+class TodoTool:
+    """Add, update, list, or clear the session's in-memory todos."""
+
+    name = "todo"
+    description = (
+        "Track multi-step work: add a todo, update its status "
+        "(pending|in_progress|completed), list all todos, or clear completed ones."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "action": {
+                "type": "string",
+                "enum": ["add", "update", "list", "clear"],
+            },
+            "text": {"type": "string"},
+            "id": {"type": "string"},
+            "status": {"type": "string", "enum": list(_VALID_STATUSES)},
+        },
+        "required": ["action"],
+    }
+
+    def __init__(self, todos: TodoList | None = None) -> None:
+        # NB: `todos or TodoList()` would be wrong — an empty TodoList is
+        # falsy via __len__, so an explicitly shared list would be discarded.
+        self._todos = todos if todos is not None else TodoList()
+
+    @property
+    def todos(self) -> TodoList:
+        """The shared list this tool reads and writes."""
+        return self._todos
+
+    async def run(self, arguments: dict) -> ToolOutcome:
+        action = arguments.get("action", "")
+        if action == "add":
+            text = (arguments.get("text") or "").strip()
+            if not text:
+                return ToolOutcome(
+                    content="error: todo add needs a text", is_error=True
+                )
+            item = self._todos.add(text)
+            return ToolOutcome(content=f"added [{item.id}] pending — {item.text}")
+        if action == "update":
+            todo_id = arguments.get("id", "")
+            status = arguments.get("status")
+            text = arguments.get("text")
+            if status is not None and status not in _VALID_STATUSES:
+                return ToolOutcome(
+                    content=f"error: unknown status: {status}", is_error=True
+                )
+            if status is None and text is None:
+                return ToolOutcome(
+                    content="error: todo update needs a status and/or text",
+                    is_error=True,
+                )
+            updated = self._todos.update(todo_id, status=status, text=text)
+            if updated is None:
+                return ToolOutcome(
+                    content=f"error: unknown todo id: {todo_id}", is_error=True
+                )
+            return ToolOutcome(
+                content=f"updated [{updated.id}] {updated.status} — {updated.text}"
+            )
+        if action == "list":
+            return ToolOutcome(content=format_todos(self._todos.all()))
+        if action == "clear":
+            removed = self._todos.clear_completed()
+            return ToolOutcome(content=f"cleared {removed} completed todo(s)")
+        return ToolOutcome(content=f"error: unknown action: {action}", is_error=True)