harness-sdk-python 0.25.1__tar.gz → 0.26.0__tar.gz

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: harness-sdk-python
-Version: 0.25.1
+Version: 0.26.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

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "harness-sdk-python"
-version = "0.25.1"
+version = "0.26.0"
 description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
 readme = "README.md"
 license = "MIT"

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/src/harness_sdk/__init__.py

@@ -11,6 +11,7 @@ from .checkpoints import (
     langgraph_snapshot,
     langgraph_switch_values,
 )
+from .deepagents import DeepagentsMessages
 from .linear_thread import linear_thread
 from .run_manager import RunManager, RunManagerDescriptor, run_manager
 
@@ -18,6 +19,7 @@ __all__ = [
     "CheckpointIndex",
     "CheckpointStore",
     "CheckpointThread",
+    "DeepagentsMessages",
     "MemoryCheckpointIndex",
     "Plan",
     "PostgresCheckpointIndex",

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/src/harness_sdk/checkpoints.py

@@ -11,6 +11,7 @@ Rows and entries are insert-only; ``last_checkpoint_id`` is the only update.
 """
 
 import asyncio
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import Any, Callable, Protocol
 
@@ -158,50 +159,69 @@ class MemoryCheckpointIndex:
 
 
 class PostgresCheckpointIndex:
-    """psycopg-backed CheckpointIndex; ``setup()`` creates the two tables."""
+    """psycopg-backed CheckpointIndex over an ``AsyncConnection`` or ``AsyncConnectionPool``; ``setup()`` creates the two tables."""
 
     def __init__(self, conn: Any) -> None:
         try:
-            import psycopg  # noqa: F401
+            import psycopg
+            import psycopg_pool
         except ImportError as exc:
             raise ImportError(
                 "harness_sdk.checkpoints.PostgresCheckpointIndex requires the "
                 "'postgres' extra: pip install 'harness-sdk-python[postgres]'"
             ) from exc
+        if not isinstance(
+            conn, (psycopg.AsyncConnection, psycopg_pool.AsyncConnectionPool)
+        ):
+            raise TypeError(
+                f"conn must be an AsyncConnection or AsyncConnectionPool, got {conn!r}"
+            )
         self._conn = conn
 
+    @asynccontextmanager
+    async def _connection(self):
+        import psycopg
+
+        if isinstance(self._conn, psycopg.AsyncConnection):
+            yield self._conn
+        else:
+            async with self._conn.connection() as conn:
+                yield 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 with self._connection() as conn:
+            await 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 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()
+        async with self._connection() as conn:
+            cur = await 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 _seed_root(self, thread_id: str, checkpoint_ns: str) -> None:
+    async def _seed_root(self, conn: Any, thread_id: str, checkpoint_ns: str) -> None:
         from psycopg.types.json import Jsonb
 
-        await self._conn.execute(
+        await conn.execute(
             "INSERT INTO checkpoint_threads (thread_id, checkpoint_ns, branches)"
             " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
             (thread_id, checkpoint_ns, Jsonb([_ROOT_ENTRY])),
@@ -223,9 +243,9 @@ class PostgresCheckpointIndex:
             "fork_branch": fork_branch,
             "last_checkpoint_id": last_checkpoint_id,
         }
-        async with self._conn.transaction():
-            await self._seed_root(thread_id, checkpoint_ns)
-            cur = await self._conn.execute(
+        async with self._connection() as conn, conn.transaction():
+            await self._seed_root(conn, thread_id, checkpoint_ns)
+            cur = await conn.execute(
                 "UPDATE checkpoint_threads SET branches = branches || %s"
                 " WHERE thread_id = %s AND checkpoint_ns = %s"
                 " RETURNING jsonb_array_length(branches)",
@@ -240,26 +260,27 @@ class PostgresCheckpointIndex:
         branch_ids: list[int],
         checkpoint_id: str,
     ) -> None:
-        await self._seed_root(thread_id, checkpoint_ns)
-        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 - 1 = ANY(%s)"
-            " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
-            " ELSE elem END ORDER BY idx)"
-            " 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 with self._connection() as conn:
+            await self._seed_root(conn, thread_id, checkpoint_ns)
+            cur = await 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 conn.execute(
+                "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
+                " CASE WHEN idx::int - 1 = ANY(%s)"
+                " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
+                " ELSE elem END ORDER BY idx)"
+                " 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:
+        async with self._connection() as conn, conn.cursor() as cur:
             await cur.executemany(
                 "INSERT INTO checkpoint_messages"
                 " (thread_id, checkpoint_ns, message_id, anchor_checkpoint_id, branch_id)"
@@ -272,12 +293,13 @@ class PostgresCheckpointIndex:
     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()
+        async with self._connection() as conn:
+            cur = await 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]}

harness_sdk_python-0.26.0/src/harness_sdk/deepagents.py

@@ -0,0 +1,210 @@
+"""DeepagentsMessages: the ``messages`` slice of a Statewire host, projected
+from a langgraph checkpointer in the deepagents dialect. The checkpointer is
+the source of truth and the wire state is its projection, restamped with
+``branch: {id, siblings}`` metadata after every ingest. A rewind restores the
+projection from the target message's anchor checkpoint; a branch switch clones
+the target branch's tip as the checkpointer's latest checkpoint, so resumption
+and reboot both land on the switched world.
+"""
+
+from datetime import datetime, timezone
+from typing import Any
+
+from statepatch import plain
+
+from .checkpoints import (
+    CheckpointIndex,
+    CheckpointStore,
+    Snapshot,
+    _values_snapshot,
+    checkpoint_thread,
+    langgraph_snapshot,
+)
+
+
+def _role(message: dict[str, Any]) -> str:
+    return "user" if message["type"] == "human" else "assistant"
+
+
+def _user_text(message: dict[str, Any]) -> str:
+    return "\n\n".join(
+        part["text"] for part in message["parts"] if part["type"] == "text"
+    )
+
+
+def _checkpointer_store(
+    checkpointer: Any, thread_id: str, checkpoint_ns: str
+) -> CheckpointStore:
+    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 checkpointer.aget_tuple({"configurable": configurable})
+            if saved is None:
+                return None
+            return _values_snapshot(
+                saved.checkpoint["id"],
+                saved.checkpoint["channel_values"],
+                "messages",
+                "lineage",
+            )
+
+    return _Store()
+
+
+class DeepagentsMessages:
+    """A host's ``state["messages"]`` over a langgraph checkpointer: boot, turn streaming, rewind, and branch switching for the deepagents dialect."""
+
+    def __init__(self, checkpointer: Any) -> None:
+        if not hasattr(checkpointer, "aget_tuple") or not hasattr(checkpointer, "aput"):
+            raise TypeError(
+                f"checkpointer must be a langgraph checkpoint saver, got {checkpointer!r}"
+            )
+        self._checkpointer = checkpointer
+        self._captured: dict[str, Any] | None = None
+
+    async def boot(
+        self, state: Any, *, index: CheckpointIndex, thread_id: str
+    ) -> None:
+        """Adopt ``state["messages"]`` and load the thread's history into it with branch stamps."""
+        self._state = state
+        self._thread_id = thread_id
+        self._store = _checkpointer_store(self._checkpointer, thread_id, "")
+        self._thread = await checkpoint_thread(
+            store=self._store, index=index, thread_id=thread_id, role=_role
+        )
+        if self._thread.transcript:
+            # the graph carries no lineage channel: re-derive the world path from the tail row
+            row = await index.get_message(
+                thread_id, "", self._thread.transcript[-1]["id"]
+            )
+            if row is not None and row["branch_id"] != 0:
+                await self._thread.switch_state(row["branch_id"])
+        self._restamp()
+
+    async def get_message_meta(self, message_id: str | None) -> dict[str, Any] | None:
+        """The thread projection's message meta, for RunManager."""
+        return await self._thread.get_message_meta(message_id)
+
+    async def get_message_child_id(self, parent_id: str) -> str | None:
+        """The thread projection's child id, for RunManager."""
+        return await self._thread.get_message_child_id(parent_id)
+
+    def config(self) -> dict[str, Any]:
+        """The langgraph config addressing this thread."""
+        return {"configurable": {"thread_id": self._thread_id}}
+
+    async def plan_turn(
+        self, ctx: Any
+    ) -> tuple[dict[str, Any] | None, dict[str, Any]]:
+        """The dispatch's langgraph (payload, config): plans the resumption, restores a rewind's window, appends the batch to the wire state, and confirms it applied."""
+        config = self.config()
+        if ctx.resume:
+            return None, config
+        plan = await self._thread.plan_run(ctx)
+        if ctx.has_rollback:
+            base = (
+                (await self._snapshot(plan.checkpoint_id)).messages
+                if plan.checkpoint_id is not None
+                else []
+            )
+            self._replace([*base, *plan.replay])
+        turn: list[Any] = [*plan.replay]
+        for message in ctx.messages:
+            turn.append(self._append_user(message))
+        ctx.applied()
+        if plan.fork:
+            if plan.checkpoint_id is not None:
+                config["configurable"]["checkpoint_id"] = plan.checkpoint_id
+            else:
+                from langchain_core.messages import RemoveMessage
+                from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+                turn = [RemoveMessage(id=REMOVE_ALL_MESSAGES), *turn]
+        return ({"messages": turn} if turn else None), config
+
+    def append_event(self, namespace: tuple, mode: str, event: Any) -> None:
+        """Fold one ``astream`` event into the wire state; the root checkpoint stream is captured for the turn's ingest."""
+        if mode == "checkpoints":
+            if namespace == ():
+                self._captured = event
+            return
+        from statewire.langgraph import append_langgraph_event
+
+        append_langgraph_event(self._state, namespace, mode, event)
+
+    async def ingest_turn(self, *, refetch: bool = False) -> None:
+        """Ingest the turn's captured checkpoint and restamp; ``refetch`` discards the capture and reads the checkpointer's latest instead."""
+        captured, self._captured = self._captured, None
+        if refetch or captured is None:
+            await self._thread.ingest_latest()
+        else:
+            await self._thread.ingest(langgraph_snapshot(captured))
+        self._restamp()
+
+    def append_users(self, messages: Any) -> list[dict[str, Any]]:
+        """Append harness user messages to the wire state as human entries, deduplicated by id; returns the langgraph turn entries."""
+        return [self._append_user(message) for message in messages]
+
+    async def switch_branch(self, branch_id: int) -> None:
+        """Switch to ``branch_id``: its tip becomes the checkpointer's latest checkpoint and the wire state rebaselines."""
+        transcript, lineage = await self._thread.switch_state(branch_id)
+        checkpoint_id = await self._clone_tip(self._thread.checkpoint_id)
+        await self._thread.ingest(Snapshot(checkpoint_id, transcript, lineage))
+        self._restamp()
+
+    async def _clone_tip(self, checkpoint_id: str) -> str:
+        from langgraph.checkpoint.base.id import uuid6
+
+        parent = {
+            "configurable": {
+                "thread_id": self._thread_id,
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoint_id,
+            }
+        }
+        saved = await self._checkpointer.aget_tuple(parent)
+        if saved is None:
+            raise LookupError(f"checkpoint {checkpoint_id!r} is not in the checkpointer")
+        clone = {
+            **saved.checkpoint,
+            "id": str(uuid6(clock_seq=-2)),
+            "ts": datetime.now(timezone.utc).isoformat(),
+        }
+        metadata = {
+            "source": "update",
+            "step": (saved.metadata or {}).get("step", -1) + 1,
+            "parents": {},
+        }
+        written = await self._checkpointer.aput(parent, clone, metadata, {})
+        return written["configurable"]["checkpoint_id"]
+
+    async def _snapshot(self, checkpoint_id: str) -> Snapshot:
+        snapshot = await self._store.get(checkpoint_id)
+        if snapshot is None:
+            raise LookupError(
+                f"checkpoint {checkpoint_id!r} is not in the checkpointer"
+            )
+        return snapshot
+
+    def _append_user(self, message: dict[str, Any]) -> dict[str, Any]:
+        entry = {"id": message["id"], "type": "human", "content": _user_text(message)}
+        if not any(m["id"] == entry["id"] for m in plain(self._state["messages"])):
+            self._state["messages"].append(entry)
+        return entry
+
+    def _replace(self, messages: list[dict[str, Any]]) -> None:
+        from statewire.langgraph import replace_messages
+
+        replace_messages(self._state, messages)
+
+    def _restamp(self) -> None:
+        stamped = []
+        for message in self._thread.transcript:
+            branch, siblings = self._thread.sibling_info(message["id"])
+            stamped.append({**message, "branch": {"id": branch, "siblings": siblings}})
+        self._replace(stamped)

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/src/harness_sdk/run_manager.py

@@ -48,6 +48,7 @@ window.
 """
 
 import asyncio
+import contextlib
 import logging
 from dataclasses import dataclass, field
 from typing import Any, AsyncIterator, Callable, Coroutine, Iterable, Protocol
@@ -1610,7 +1611,7 @@ class RunManagerNamespace(TrayMapNamespace):
                 f"{type(api).__name__} must implement get_message_meta "
                 "(RunManager.Thread) to host @run_manager"
             )
-        self._mint_lock = asyncio.Lock()
+        self._mint: "asyncio.Task[Any] | None" = None
         self._mounted: list[RunManager] = []
         api.state["runs"] = []
         api.state["status"] = "ready"
@@ -1628,14 +1629,26 @@ class RunManagerNamespace(TrayMapNamespace):
         return tray
 
     async def _machine(self, run_id: str) -> tuple[RunManager, bool]:
-        tray = self.live()
-        if tray is not None:
-            return tray, False
-        async with self._mint_lock:
+        # No await between the live() check and the mint's publication, and no
+        # lock held across storage I/O: same-tick callers await the shared mint
+        # task and merge into the machine it mounts; the creator raises its
+        # mint's failure while followers re-mint and progress.
+        while True:
             tray = self.live()
             if tray is not None:
                 return tray, False
-            return await self.get_or_create(run_id), True  # type: ignore[return-value]
+            pending = self._mint
+            if pending is None:
+                pending = self._api.create_task(self.get_or_create(run_id))
+                self._mint = pending
+                pending.add_done_callback(self._mint_settled)
+                return await pending, True
+            with contextlib.suppress(Exception):
+                await pending
+
+    def _mint_settled(self, task: "asyncio.Task[Any]") -> None:
+        if self._mint is task:
+            self._mint = None
 
     # ─── Mounted view ───────────────────────────────────────
 

{harness_sdk_python-0.25.1 → harness_sdk_python-0.26.0}/tests/test_checkpoints.py

@@ -663,11 +663,12 @@ def postgres_dsn():
         subprocess.run([docker, "rm", "-f", container], capture_output=True)
 
 
-@pytest.fixture
-async def pg_index(postgres_dsn):
+@pytest.fixture(params=["conn", "pool"])
+async def pg_index(request, postgres_dsn):
     import uuid
 
     import psycopg
+    from psycopg_pool import AsyncConnectionPool
 
     from harness_sdk import PostgresCheckpointIndex
 
@@ -675,15 +676,29 @@ async def pg_index(postgres_dsn):
     with psycopg.connect(postgres_dsn, autocommit=True) as admin:
         admin.execute(f'CREATE DATABASE "{dbname}"')
     dsn = postgres_dsn.rsplit("/", 1)[0] + f"/{dbname}"
-    conn = await psycopg.AsyncConnection.connect(dsn, autocommit=True)
-    index = PostgresCheckpointIndex(conn)
-    await index.setup()
-    yield index
-    await conn.close()
+    if request.param == "conn":
+        conn = await psycopg.AsyncConnection.connect(dsn, autocommit=True)
+        index = PostgresCheckpointIndex(conn)
+        await index.setup()
+        yield index
+        await conn.close()
+    else:
+        async with AsyncConnectionPool(dsn, open=False) as pool:
+            index = PostgresCheckpointIndex(pool)
+            await index.setup()
+            yield index
     with psycopg.connect(postgres_dsn, autocommit=True) as admin:
         admin.execute(f'DROP DATABASE "{dbname}" WITH (FORCE)')
 
 
+def test_postgres_index_rejects_other_conn_types():
+    pytest.importorskip("psycopg")
+    from harness_sdk import PostgresCheckpointIndex
+
+    with pytest.raises(TypeError, match="AsyncConnection or AsyncConnectionPool"):
+        PostgresCheckpointIndex(object())
+
+
 async def test_postgres_branches_round_trip(pg_index):
     assert await pg_index.read_thread(THREAD, "") == []
     first = await pg_index.append_branch(

harness_sdk_python-0.26.0/tests/test_deepagents_messages.py

@@ -0,0 +1,238 @@
+"""Contract: DeepagentsMessages at its public seam — a langgraph checkpointer
+projected onto ``state["messages"]``, driven the way the host's stream loop
+drives it (plan_turn, append_event per stream event, ingest_turn per turn)."""
+
+import itertools
+from dataclasses import dataclass
+
+import pytest
+
+from harness_sdk import DeepagentsMessages, MemoryCheckpointIndex
+
+pytest.importorskip("langgraph.graph")
+
+from langchain_core.messages import AIMessage, RemoveMessage  # noqa: E402
+from langgraph.checkpoint.memory import InMemorySaver  # noqa: E402
+from langgraph.graph import START, StateGraph  # noqa: E402
+from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages  # noqa: E402
+
+THREAD = "t1"
+
+
+def _graph(saver):
+    from typing import Annotated, TypedDict
+
+    counter = itertools.count()
+
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    def node(state):
+        n = next(counter)
+        return {"messages": [AIMessage(id=f"a{n}", content=f"reply {n}")]}
+
+    builder = StateGraph(State)
+    builder.add_node("node", node)
+    builder.add_edge(START, "node")
+    return builder.compile(checkpointer=saver)
+
+
+@dataclass
+class Ctx:
+    trigger: str
+    messages: tuple = ()
+    resume: bool = False
+    has_rollback: bool = False
+    rollback_to: str | None = None
+    applied_count: int = 0
+
+    def applied(self):
+        self.applied_count += 1
+
+
+def user(id, text):
+    return {"id": id, "role": "user", "parts": [{"type": "text", "text": text}]}
+
+
+async def booted(saver, index):
+    state = {"messages": []}
+    messages = DeepagentsMessages(saver)
+    await messages.boot(state, index=index, thread_id=THREAD)
+    return messages, state
+
+
+async def stream(graph, messages, payload, config):
+    async for namespace, mode, event in graph.astream(
+        payload,
+        config=config,
+        stream_mode=["messages", "updates", "checkpoints"],
+        subgraphs=True,
+    ):
+        messages.append_event(namespace, mode, event)
+
+
+async def turn(graph, messages, ctx):
+    payload, config = await messages.plan_turn(ctx)
+    await stream(graph, messages, payload, config)
+    await messages.ingest_turn()
+
+
+def ids(state):
+    return [m["id"] for m in state["messages"]]
+
+
+def branch(state, message_id):
+    return next(m["branch"] for m in state["messages"] if m["id"] == message_id)
+
+
+async def seeded():
+    """Two turns: [u1, a0, u2, a1] on branch 0."""
+    saver, index = InMemorySaver(), MemoryCheckpointIndex()
+    graph = _graph(saver)
+    messages, state = await booted(saver, index)
+    await turn(graph, messages, Ctx("message-send", (user("u1", "one"),)))
+    await turn(graph, messages, Ctx("message-send", (user("u2", "two"),)))
+    return saver, index, graph, messages, state
+
+
+# ─── Boot ───────────────────────────────────────────────────
+
+
+async def test_boot_fresh_thread_projects_empty():
+    _, state = await booted(InMemorySaver(), MemoryCheckpointIndex())
+    assert state["messages"] == []
+
+
+async def test_boot_rejects_a_non_saver():
+    with pytest.raises(TypeError, match="checkpoint saver"):
+        DeepagentsMessages(object())
+
+
+async def test_reboot_restores_history_with_stamps():
+    saver, index, _, _, state = await seeded()
+    _, rebooted = await booted(saver, index)
+    assert rebooted["messages"] == state["messages"]
+    assert ids(rebooted) == ["u1", "a0", "u2", "a1"]
+
+
+# ─── Turn streaming ─────────────────────────────────────────
+
+
+async def test_send_turn_projects_and_stamps():
+    saver, index = InMemorySaver(), MemoryCheckpointIndex()
+    graph = _graph(saver)
+    messages, state = await booted(saver, index)
+    ctx = Ctx("message-send", (user("u1", "one"),))
+    await turn(graph, messages, ctx)
+    assert ctx.applied_count == 1
+    assert ids(state) == ["u1", "a0"]
+    assert [m["type"] for m in state["messages"]] == ["human", "ai"]
+    assert branch(state, "a0") == {"id": 0, "siblings": [0]}
+
+
+async def test_plan_turn_resume_carries_no_payload():
+    _, _, _, messages, _ = await seeded()
+    ctx = Ctx("message-send", resume=True)
+    payload, config = await messages.plan_turn(ctx)
+    assert payload is None
+    assert config == {"configurable": {"thread_id": THREAD}}
+    assert ctx.applied_count == 0
+
+
+async def test_append_users_dedupes_by_id():
+    _, _, _, messages, state = await seeded()
+    entries = messages.append_users([user("u3", "three")])
+    assert entries == [{"id": "u3", "type": "human", "content": "three"}]
+    messages.append_users([user("u3", "three")])
+    assert ids(state).count("u3") == 1
+
+
+async def test_ingest_turn_refetch_discards_the_capture():
+    saver, index = InMemorySaver(), MemoryCheckpointIndex()
+    graph = _graph(saver)
+    messages, state = await booted(saver, index)
+    payload, config = await messages.plan_turn(
+        Ctx("message-send", (user("u1", "one"),))
+    )
+    await stream(graph, messages, payload, config)
+    # the stopped run persisted a further exchange after the capture
+    await graph.ainvoke(
+        {"messages": [{"id": "u2", "type": "human", "content": "two"}]},
+        {"configurable": {"thread_id": THREAD}},
+    )
+    await messages.ingest_turn(refetch=True)
+    assert ids(state) == ["u1", "a0", "u2", "a1"]
+
+
+# ─── Rewind ─────────────────────────────────────────────────
+
+
+async def test_reload_restores_the_window_from_the_anchor_checkpoint():
+    _, _, graph, messages, state = await seeded()
+    ctx = Ctx("message-reload", has_rollback=True, rollback_to="u2")
+    payload, config = await messages.plan_turn(ctx)
+    assert ids(state) == ["u1", "a0", "u2"]
+    assert config["configurable"]["checkpoint_id"] is not None
+    assert [m["id"] for m in payload["messages"]] == ["u2"]
+    await stream(graph, messages, payload, config)
+    await messages.ingest_turn()
+    assert ids(state) == ["u1", "a0", "u2", "a2"]
+    assert branch(state, "a2") == {"id": 1, "siblings": [0, 1]}
+    assert branch(state, "u2") == {"id": 0, "siblings": [0]}
+
+
+async def test_edit_of_the_first_message_starts_fresh():
+    _, _, graph, messages, state = await seeded()
+    ctx = Ctx(
+        "message-edit", (user("u1b", "redo"),), has_rollback=True, rollback_to=None
+    )
+    payload, config = await messages.plan_turn(ctx)
+    assert ids(state) == ["u1b"]
+    assert payload["messages"][0] == RemoveMessage(id=REMOVE_ALL_MESSAGES)
+    assert "checkpoint_id" not in config["configurable"]
+    await stream(graph, messages, payload, config)
+    await messages.ingest_turn()
+    assert ids(state) == ["u1b", "a2"]
+    assert branch(state, "u1b") == {"id": 1, "siblings": [0, 1]}
+
+
+# ─── Branch switch ──────────────────────────────────────────
+
+
+async def forked():
+    """seeded() plus a reload at u2: branch 1 replaces a1 with a2."""
+    saver, index, graph, messages, state = await seeded()
+    await turn(
+        graph, messages, Ctx("message-reload", has_rollback=True, rollback_to="u2")
+    )
+    assert ids(state) == ["u1", "a0", "u2", "a2"]
+    return saver, index, graph, messages, state
+
+
+async def test_switch_branch_rebaselines_and_stamps():
+    _, _, _, messages, state = await forked()
+    await messages.switch_branch(0)
+    assert ids(state) == ["u1", "a0", "u2", "a1"]
+    assert branch(state, "a1") == {"id": 0, "siblings": [0, 1]}
+
+
+async def test_switch_branch_survives_a_reboot():
+    saver, index, _, messages, _ = await forked()
+    await messages.switch_branch(0)
+    _, state = await booted(saver, index)
+    assert ids(state) == ["u1", "a0", "u2", "a1"]
+    assert branch(state, "a1") == {"id": 0, "siblings": [0, 1]}
+
+
+async def test_send_after_a_switch_continues_the_switched_world():
+    _, _, graph, messages, state = await forked()
+    await messages.switch_branch(0)
+    await turn(graph, messages, Ctx("message-send", (user("u3", "three"),)))
+    assert ids(state) == ["u1", "a0", "u2", "a1", "u3", "a3"]
+    assert branch(state, "u3") == {"id": 0, "siblings": [0]}
+
+
+async def test_switch_branch_unknown_raises():
+    _, _, _, messages, _ = await seeded()
+    with pytest.raises(KeyError, match="unknown branch_id"):
+        await messages.switch_branch(7)

harness_sdk_python-0.26.0/tests/test_mint_concurrency.py

@@ -0,0 +1,59 @@
+import asyncio
+
+import pytest
+from pinned.testing import MemoryStateStore
+from run_helpers import Script, run_host
+
+
+class WedgeStore(MemoryStateStore):
+    """get_tray parks for flagged tray ids until ``fail()`` fires the store's
+    op deadline, the way a half-open conn parks a postgres transaction."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.wedged: set[str] = set()
+        self.parked = asyncio.Event()
+        self._deadline = asyncio.Event()
+
+    def fail(self) -> None:
+        self._deadline.set()
+
+    async def get_tray(self, ns, id, tray_ns, tray_id, epoch):
+        if tray_id in self.wedged:
+            self.wedged.discard(tray_id)
+            self.parked.set()
+            await self._deadline.wait()
+            raise TimeoutError("store op deadline")
+        return await super().get_tray(ns, id, tray_ns, tray_id, epoch)
+
+
+async def test_a_parked_mint_fails_at_the_deadline_and_followers_progress():
+    """No lock is held across mint storage I/O: a wedged mint fails loudly at
+    the store deadline, the waiting command re-mints its own machine, and the
+    wedged command surfaces the store error."""
+    script = Script()
+    store = WedgeStore()
+    async with run_host(script, store=store) as (drv, host):
+        api = host.live
+        store.wedged.add("run-a")
+        parked = asyncio.create_task(api.runs.for_run("run-a"))
+        await asyncio.wait_for(store.parked.wait(), 5)
+        follower = asyncio.create_task(api.runs.for_run("run-b"))
+        await asyncio.sleep(0.05)
+        assert not follower.done()  # same-tick mints merge: followers await the pending mint
+        store.fail()
+        tray = await asyncio.wait_for(follower, 5)
+        assert tray.id == "run-b"
+        with pytest.raises(TimeoutError):
+            await parked
+
+
+async def test_concurrent_same_run_id_mints_resolve_to_one_machine():
+    script = Script()
+    async with run_host(script) as (drv, host):
+        api = host.live
+        trays = await asyncio.gather(
+            *(asyncio.create_task(api.runs.for_run("run-x")) for _ in range(3))
+        )
+        assert all(tray is trays[0] for tray in trays)
+        assert list(api.runs.instances) == ["run-x"]