harness-sdk-python 0.20.0__tar.gz → 0.22.0__tar.gz

{harness_sdk_python-0.20.0 → harness_sdk_python-0.22.0}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.5
 Name: harness-sdk-python
-Version: 0.20.0
+Version: 0.22.0
 Summary: RunManager: the harness-sdk runs subsystem for Python Statewire hosts
 Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
 License-Expression: MIT
 Requires-Python: <4.0,>=3.12
-Requires-Dist: pinned<0.14,>=0.13.0
-Requires-Dist: statewire<0.11,>=0.10.0
+Requires-Dist: pinned<0.15,>=0.14.0
+Requires-Dist: statewire<0.12,>=0.11.0
 Provides-Extra: deepagents
 Requires-Dist: deepagents>=0.6.12; extra == 'deepagents'
 Requires-Dist: langchain-core>=0.3; extra == 'deepagents'

{harness_sdk_python-0.20.0 → harness_sdk_python-0.22.0}/pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "harness-sdk-python"
-version = "0.20.0"
+version = "0.22.0"
 description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
 readme = "README.md"
 license = "MIT"
 requires-python = ">=3.12,<4.0"
-dependencies = ["statewire>=0.10.0,<0.11", "pinned>=0.13.0,<0.14"]
+dependencies = ["statewire>=0.11.0,<0.12", "pinned>=0.14.0,<0.15"]
 
 [project.optional-dependencies]
 postgres = ["langgraph-checkpoint-postgres>=2.0.0"]

harness_sdk_python-0.22.0/src/harness_sdk/__init__.py

@@ -0,0 +1,29 @@
+from .checkpoints import (
+    CheckpointIndex,
+    CheckpointStore,
+    CheckpointThread,
+    MemoryCheckpointIndex,
+    Plan,
+    PostgresCheckpointIndex,
+    Snapshot,
+    checkpoint_thread,
+    langgraph_checkpoints,
+)
+from .linear_thread import linear_thread
+from .run_manager import RunManager, RunManagerDescriptor, run_manager
+
+__all__ = [
+    "CheckpointIndex",
+    "CheckpointStore",
+    "CheckpointThread",
+    "MemoryCheckpointIndex",
+    "Plan",
+    "PostgresCheckpointIndex",
+    "RunManager",
+    "RunManagerDescriptor",
+    "Snapshot",
+    "checkpoint_thread",
+    "langgraph_checkpoints",
+    "linear_thread",
+    "run_manager",
+]

harness_sdk_python-0.22.0/src/harness_sdk/checkpoints.py

@@ -0,0 +1,556 @@
+"""Checkpoint subsystem: a sidecar branch/message index over an external
+checkpoint store (``get(checkpoint_id)`` only), designed for langgraph.
+
+A branch is a resumption point plus its continuation; edits, regenerations,
+reloads, and compactions all create one. ``branch_id`` is the entry's array
+position + 1; the root branch is 0 and the array holds fork entries only.
+Every ingested checkpoint advances ``last_checkpoint_id`` for every branch in
+the state's ``branch_ids`` path, so a branch's pointer is always its
+subtree's most recent checkpoint. Rows and entries are insert-only;
+``last_checkpoint_id`` is the only update.
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any, Callable, Protocol
+
+from .linear_thread import linear_thread
+
+_FORK_TRIGGERS = ("message-edit", "message-reload")
+_TIP_TRIGGERS = (
+    "message-send",
+    "input-resume",
+    "error-continue",
+    "stop-continue",
+    "steer",
+)
+
+
+@dataclass(frozen=True)
+class Snapshot:
+    """One checkpoint's projection: its id, message list, and lineage channel."""
+
+    checkpoint_id: str
+    messages: list[dict[str, Any]]
+    lineage: dict[str, Any] | None = None
+
+
+@dataclass(frozen=True)
+class Plan:
+    """A resumption: the checkpoint to resume from, survivors to re-append, and whether a branch forks."""
+
+    checkpoint_id: str | None
+    replay: list[dict[str, Any]]
+    fork: bool
+
+
+class CheckpointStore(Protocol):
+    """The whole read seam over the external framework's checkpoint store."""
+
+    async def get(self, checkpoint_id: str | None) -> Snapshot | None:
+        """The snapshot named by ``checkpoint_id``, the latest for None, None when absent."""
+        ...
+
+
+class CheckpointIndex(Protocol):
+    """The sidecar index: branch entries per thread row, one row per message."""
+
+    async def read_thread(
+        self, thread_id: str, checkpoint_ns: str
+    ) -> list[dict[str, Any]]: ...
+
+    async def append_branch(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        *,
+        fork_message_id: str | None,
+        last_checkpoint_id: str,
+    ) -> int: ...
+
+    async def bump_tips(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        branch_ids: list[int],
+        checkpoint_id: str,
+    ) -> None: ...
+
+    async def upsert_messages(self, rows: list[dict[str, Any]]) -> None: ...
+
+    async def get_message(
+        self, thread_id: str, checkpoint_ns: str, message_id: str
+    ) -> dict[str, Any] | None: ...
+
+
+def _check_branch_ids(branch_ids: list[int], count: int) -> None:
+    for branch_id in branch_ids:
+        if not isinstance(branch_id, int) or not 1 <= branch_id <= count:
+            raise KeyError(f"unknown branch_id {branch_id!r} (index holds {count})")
+
+
+@dataclass
+class MemoryCheckpointIndex:
+    """Dict-backed CheckpointIndex for tests."""
+
+    _threads: dict[tuple[str, str], list[dict[str, Any]]] = field(default_factory=dict)
+    _messages: dict[tuple[str, str, str], dict[str, Any]] = field(default_factory=dict)
+
+    async def read_thread(
+        self, thread_id: str, checkpoint_ns: str
+    ) -> list[dict[str, Any]]:
+        return [dict(e) for e in self._threads.get((thread_id, checkpoint_ns), [])]
+
+    async def append_branch(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        *,
+        fork_message_id: str | None,
+        last_checkpoint_id: str,
+    ) -> int:
+        branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
+        branches.append(
+            {
+                "fork_message_id": fork_message_id,
+                "last_checkpoint_id": last_checkpoint_id,
+            }
+        )
+        return len(branches)
+
+    async def bump_tips(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        branch_ids: list[int],
+        checkpoint_id: str,
+    ) -> None:
+        branches = self._threads.get((thread_id, checkpoint_ns), [])
+        _check_branch_ids(branch_ids, len(branches))
+        for branch_id in branch_ids:
+            branches[branch_id - 1]["last_checkpoint_id"] = checkpoint_id
+
+    async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
+        for row in rows:
+            key = (row["thread_id"], row["checkpoint_ns"], row["message_id"])
+            self._messages.setdefault(
+                key,
+                {
+                    "anchor_checkpoint_id": row["anchor_checkpoint_id"],
+                    "branch_id": row["branch_id"],
+                },
+            )
+
+    async def get_message(
+        self, thread_id: str, checkpoint_ns: str, message_id: str
+    ) -> dict[str, Any] | None:
+        row = self._messages.get((thread_id, checkpoint_ns, message_id))
+        return dict(row) if row is not None else None
+
+
+class PostgresCheckpointIndex:
+    """psycopg-backed CheckpointIndex; ``setup()`` creates the two tables."""
+
+    def __init__(self, conn: Any) -> None:
+        try:
+            import psycopg  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "harness_sdk.checkpoints.PostgresCheckpointIndex requires the "
+                "'postgres' extra: pip install 'harness-sdk-python[postgres]'"
+            ) from exc
+        self._conn = conn
+
+    async def setup(self) -> None:
+        await self._conn.execute(
+            "CREATE TABLE IF NOT EXISTS checkpoint_threads ("
+            " thread_id TEXT NOT NULL,"
+            " checkpoint_ns TEXT NOT NULL,"
+            " branches JSONB NOT NULL,"
+            " PRIMARY KEY (thread_id, checkpoint_ns))"
+        )
+        await self._conn.execute(
+            "CREATE TABLE IF NOT EXISTS checkpoint_messages ("
+            " thread_id TEXT NOT NULL,"
+            " checkpoint_ns TEXT NOT NULL,"
+            " message_id TEXT NOT NULL,"
+            " anchor_checkpoint_id TEXT,"
+            " branch_id INTEGER NOT NULL,"
+            " PRIMARY KEY (thread_id, checkpoint_ns, message_id))"
+        )
+
+    async def read_thread(
+        self, thread_id: str, checkpoint_ns: str
+    ) -> list[dict[str, Any]]:
+        cur = await self._conn.execute(
+            "SELECT branches FROM checkpoint_threads WHERE thread_id = %s AND checkpoint_ns = %s",
+            (thread_id, checkpoint_ns),
+        )
+        row = await cur.fetchone()
+        return row[0] if row is not None else []
+
+    async def append_branch(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        *,
+        fork_message_id: str | None,
+        last_checkpoint_id: str,
+    ) -> int:
+        from psycopg.types.json import Jsonb
+
+        entry = {
+            "fork_message_id": fork_message_id,
+            "last_checkpoint_id": last_checkpoint_id,
+        }
+        async with self._conn.transaction():
+            await self._conn.execute(
+                "INSERT INTO checkpoint_threads (thread_id, checkpoint_ns, branches)"
+                " VALUES (%s, %s, '[]'::jsonb) ON CONFLICT DO NOTHING",
+                (thread_id, checkpoint_ns),
+            )
+            cur = await self._conn.execute(
+                "UPDATE checkpoint_threads SET branches = branches || %s"
+                " WHERE thread_id = %s AND checkpoint_ns = %s"
+                " RETURNING jsonb_array_length(branches)",
+                (Jsonb([entry]), thread_id, checkpoint_ns),
+            )
+            return (await cur.fetchone())[0]
+
+    async def bump_tips(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        branch_ids: list[int],
+        checkpoint_id: str,
+    ) -> None:
+        cur = await self._conn.execute(
+            "SELECT jsonb_array_length(branches) FROM checkpoint_threads"
+            " WHERE thread_id = %s AND checkpoint_ns = %s",
+            (thread_id, checkpoint_ns),
+        )
+        row = await cur.fetchone()
+        _check_branch_ids(branch_ids, row[0] if row is not None else 0)
+        await self._conn.execute(
+            "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
+            " CASE WHEN idx::int = ANY(%s)"
+            " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
+            " ELSE elem END)"
+            " FROM jsonb_array_elements(branches) WITH ORDINALITY AS t(elem, idx))"
+            " WHERE thread_id = %s AND checkpoint_ns = %s",
+            (branch_ids, checkpoint_id, thread_id, checkpoint_ns),
+        )
+
+    async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
+        async with self._conn.cursor() as cur:
+            await cur.executemany(
+                "INSERT INTO checkpoint_messages"
+                " (thread_id, checkpoint_ns, message_id, anchor_checkpoint_id, branch_id)"
+                " VALUES (%(thread_id)s, %(checkpoint_ns)s, %(message_id)s,"
+                " %(anchor_checkpoint_id)s, %(branch_id)s)"
+                " ON CONFLICT (thread_id, checkpoint_ns, message_id) DO NOTHING",
+                rows,
+            )
+
+    async def get_message(
+        self, thread_id: str, checkpoint_ns: str, message_id: str
+    ) -> dict[str, Any] | None:
+        cur = await self._conn.execute(
+            "SELECT anchor_checkpoint_id, branch_id FROM checkpoint_messages"
+            " WHERE thread_id = %s AND checkpoint_ns = %s AND message_id = %s",
+            (thread_id, checkpoint_ns, message_id),
+        )
+        row = await cur.fetchone()
+        if row is None:
+            return None
+        return {"anchor_checkpoint_id": row[0], "branch_id": row[1]}
+
+
+def _message_dict(message: Any) -> dict[str, Any]:
+    if not isinstance(message, dict):
+        dump = getattr(message, "model_dump", None)
+        if dump is None:
+            raise TypeError(f"message must be a dict or a BaseMessage, got {message!r}")
+        message = dump()
+    message_id = message.get("id")
+    if not isinstance(message_id, str) or message_id == "":
+        raise ValueError(f"message without an id: {message!r}")
+    return message
+
+
+def langgraph_checkpoints(
+    saver: Any,
+    *,
+    thread_id: str,
+    checkpoint_ns: str = "",
+    messages_key: str = "messages",
+    lineage_key: str = "lineage",
+) -> CheckpointStore:
+    """A CheckpointStore over a langgraph BaseCheckpointSaver's ``aget_tuple``."""
+
+    class _Store:
+        async def get(self, checkpoint_id: str | None) -> Snapshot | None:
+            configurable: dict[str, Any] = {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+            }
+            if checkpoint_id is not None:
+                configurable["checkpoint_id"] = checkpoint_id
+            saved = await saver.aget_tuple({"configurable": configurable})
+            if saved is None:
+                return None
+            values = saved.checkpoint["channel_values"]
+            lineage = values.get(lineage_key)
+            if lineage is not None and not isinstance(lineage, dict):
+                raise ValueError(f"lineage channel must be an object, got {lineage!r}")
+            return Snapshot(
+                checkpoint_id=saved.checkpoint["id"],
+                messages=[_message_dict(m) for m in values.get(messages_key, [])],
+                lineage=lineage,
+            )
+
+    return _Store()
+
+
+@dataclass
+class CheckpointThread:
+    """One thread's current window over the store, with the sidecar index kept in step."""
+
+    transcript: list[dict[str, Any]]
+    branches: list[dict[str, Any]]
+    lineage: dict[str, Any] | None
+    checkpoint_id: str | None
+    _store: CheckpointStore
+    _index: CheckpointIndex
+    _thread_id: str
+    _ns: str
+    _role: Callable[[dict[str, Any]], str]
+    _path: list[int] = field(default_factory=lambda: [0])
+    _pending_fork: dict[str, Any] | None = None
+
+    def _projection(self):
+        return linear_thread(messages=lambda: self.transcript, role=self._role)
+
+    async def get_message_meta(self, message_id: str | None) -> dict[str, Any] | None:
+        return await self._projection().get_message_meta(message_id)
+
+    async def get_message_child_id(self, parent_id: str) -> str | None:
+        return await self._projection().get_message_child_id(parent_id)
+
+    def sibling_branches(self, message_id: str) -> list[dict[str, Any]]:
+        """Branch entries forked at ``message_id``'s parent, as {branch_id, **entry} in allocation order; the transcript in hand is the extra continuation."""
+        index = next(
+            (i for i, m in enumerate(self.transcript) if m["id"] == message_id), None
+        )
+        if index is None:
+            raise KeyError(f"message {message_id!r} is not in the current window")
+        parent_id = self.transcript[index - 1]["id"] if index > 0 else None
+        return [
+            {"branch_id": i + 1, **entry}
+            for i, entry in enumerate(self.branches)
+            if entry["fork_message_id"] == parent_id
+        ]
+
+    async def plan(self, trigger: str, *, rollback_to: str | None = None) -> Plan:
+        """The resumption for ``trigger``; a fork trigger names its last surviving message via ``rollback_to``."""
+        if trigger in _TIP_TRIGGERS:
+            if rollback_to is not None:
+                raise ValueError(f"rollback_to is invalid for trigger {trigger!r}")
+            return Plan(self.checkpoint_id, [], False)
+        if trigger not in _FORK_TRIGGERS:
+            raise ValueError(f"unknown trigger {trigger!r}")
+        if rollback_to is None:
+            self._pending_fork = {
+                "fork_message_id": None,
+                "anchor_checkpoint_id": None,
+                "fork_branch": 0,
+            }
+            return Plan(None, [], True)
+        row = await self._index.get_message(self._thread_id, self._ns, rollback_to)
+        if row is None:
+            raise KeyError(f"message {rollback_to!r} is not indexed")
+        anchor = row["anchor_checkpoint_id"]
+        index = next(
+            (i for i, m in enumerate(self.transcript) if m["id"] == rollback_to), None
+        )
+        if index is None:
+            raise KeyError(f"message {rollback_to!r} is not in the current window")
+        survivors: list[dict[str, Any]] = []
+        for message in reversed(self.transcript[: index + 1]):
+            meta = await self._index.get_message(self._thread_id, self._ns, message["id"])
+            if meta is None or meta["anchor_checkpoint_id"] != anchor:
+                break
+            survivors.append(message)
+        survivors.reverse()
+        self._pending_fork = {
+            "fork_message_id": rollback_to,
+            "anchor_checkpoint_id": anchor,
+            "fork_branch": row["branch_id"],
+        }
+        return Plan(anchor, survivors, True)
+
+    async def ingest(self, snapshot: Snapshot) -> None:
+        """Index one new checkpoint: bump the path's tips, upsert the message diff, and append the pending fork's branch entry; idempotent."""
+        pending, self._pending_fork = self._pending_fork, None
+        path = (
+            list(snapshot.lineage["branch_ids"])
+            if snapshot.lineage is not None
+            else list(self._path)
+        )
+        anchor = self.checkpoint_id
+        if pending is not None:
+            # the new world descends from the fork message's branch, not the window's leaf
+            if pending["fork_branch"] in path:
+                path = path[: path.index(pending["fork_branch"]) + 1]
+            branch_id = await self._index.append_branch(
+                self._thread_id,
+                self._ns,
+                fork_message_id=pending["fork_message_id"],
+                last_checkpoint_id=snapshot.checkpoint_id,
+            )
+            self.branches = [
+                *self.branches,
+                {
+                    "fork_message_id": pending["fork_message_id"],
+                    "last_checkpoint_id": snapshot.checkpoint_id,
+                },
+            ]
+            path = [*path, branch_id]
+            anchor = pending["anchor_checkpoint_id"]
+        tips = [b for b in path if b != 0]
+        if tips:
+            await self._index.bump_tips(
+                self._thread_id, self._ns, tips, snapshot.checkpoint_id
+            )
+            for tip in tips:
+                if tip <= len(self.branches):
+                    self.branches[tip - 1]["last_checkpoint_id"] = snapshot.checkpoint_id
+        known = {m["id"] for m in self.transcript}
+        rows = [
+            {
+                "thread_id": self._thread_id,
+                "checkpoint_ns": self._ns,
+                "message_id": message["id"],
+                "anchor_checkpoint_id": anchor,
+                "branch_id": path[-1] if path else 0,
+            }
+            for message in snapshot.messages
+            if message["id"] not in known
+        ]
+        if rows:
+            await self._index.upsert_messages(rows)
+        self.transcript = list(snapshot.messages)
+        self.checkpoint_id = snapshot.checkpoint_id
+        self.lineage = snapshot.lineage
+        self._path = path
+
+    def _load(self, snapshot: Snapshot) -> "CheckpointThread":
+        self.transcript = list(snapshot.messages)
+        self.checkpoint_id = snapshot.checkpoint_id
+        self.lineage = snapshot.lineage
+        self._path = (
+            list(snapshot.lineage["branch_ids"])
+            if snapshot.lineage is not None
+            else [0]
+        )
+        return self
+
+    def _tip_of(self, branch_id: int) -> str | None:
+        if branch_id == 0:
+            return None
+        _check_branch_ids([branch_id], len(self.branches))
+        return self.branches[branch_id - 1]["last_checkpoint_id"]
+
+    async def _get(self, checkpoint_id: str | None) -> Snapshot:
+        snapshot = await self._store.get(checkpoint_id)
+        if snapshot is None:
+            raise LookupError(f"checkpoint {checkpoint_id!r} is not in the store")
+        return snapshot
+
+    async def switch_branch(self, branch_id: int) -> "CheckpointThread":
+        """Load the branch's subtree tip as the current window."""
+        return self._load(await self._get(self._tip_of(branch_id)))
+
+    async def page_up(self) -> "CheckpointThread":
+        """Load the parent world's tail checkpoint as the current window."""
+        parent = (self.lineage or {}).get("parent_checkpoint_id")
+        if parent is None:
+            raise ValueError("the current window has no parent world")
+        return self._load(await self._get(parent))
+
+    async def page_down(self) -> "CheckpointThread":
+        """Load the newest world forked at the current window's tail message."""
+        if not self.transcript:
+            raise ValueError("the current window has no tail message")
+        tail = self.transcript[-1]["id"]
+        candidates = [
+            i + 1
+            for i, entry in enumerate(self.branches)
+            if entry["fork_message_id"] == tail
+        ]
+        if not candidates:
+            raise LookupError(f"no world forks at tail message {tail!r}")
+        return self._load(await self._get(self._tip_of(candidates[-1])))
+
+    async def locate(self, message_id: str) -> "CheckpointThread":
+        """Load a window containing ``message_id``: its branch's tip, or a fork's tip below it that took the pointer."""
+        row = await self._index.get_message(self._thread_id, self._ns, message_id)
+        if row is None:
+            raise KeyError(f"message {message_id!r} is not indexed")
+        queue = [row["branch_id"]]
+        checked: set[int] = set()
+        while queue:
+            branch_id = queue.pop(0)
+            checked.add(branch_id)
+            snapshot = await self._get(self._tip_of(branch_id))
+            if any(m["id"] == message_id for m in snapshot.messages):
+                return self._load(snapshot)
+            for i, entry in enumerate(self.branches):
+                candidate = i + 1
+                if candidate in checked or candidate in queue:
+                    continue
+                fork_id = entry["fork_message_id"]
+                if fork_id is None:
+                    fork_branch = 0
+                else:
+                    fork_row = await self._index.get_message(
+                        self._thread_id, self._ns, fork_id
+                    )
+                    if fork_row is None:
+                        continue
+                    fork_branch = fork_row["branch_id"]
+                if fork_branch in checked:
+                    queue.append(candidate)
+        raise LookupError(f"no branch tip contains message {message_id!r}")
+
+
+async def checkpoint_thread(
+    *,
+    store: CheckpointStore,
+    index: CheckpointIndex,
+    thread_id: str,
+    checkpoint_ns: str = "",
+    role: Callable[[dict[str, Any]], str],
+) -> CheckpointThread:
+    """Boot a CheckpointThread: the thread row and the latest snapshot load concurrently."""
+    if not isinstance(thread_id, str) or not thread_id:
+        raise ValueError(f"thread_id must be a non-empty string, got {thread_id!r}")
+    if not callable(role):
+        raise TypeError("role must be callable")
+    branches, snapshot = await asyncio.gather(
+        index.read_thread(thread_id, checkpoint_ns), store.get(None)
+    )
+    thread = CheckpointThread(
+        transcript=[],
+        branches=branches,
+        lineage=None,
+        checkpoint_id=None,
+        _store=store,
+        _index=index,
+        _thread_id=thread_id,
+        _ns=checkpoint_ns,
+        _role=role,
+    )
+    if snapshot is not None:
+        thread._load(snapshot)
+    return thread