harness-sdk-python 0.32.2__py3-none-any.whl → 0.34.0__py3-none-any.whl

harness_sdk/__init__.py

@@ -8,16 +8,26 @@ from .checkpoints import (
     Snapshot,
     checkpoint_thread,
     langgraph_checkpoints,
+    langgraph_checkpointer,
     langgraph_snapshot,
+    langgraph_state_snapshot,
     langgraph_switch_values,
 )
-from .deepagents import DeepagentsMessages
+from .deepagents import (
+    DeepagentsMessages,
+    deepagents_drop_ephemeral,
+    deepagents_ingest,
+    deepagents_restamp,
+    deepagents_role,
+    deepagents_upsert_ephemeral,
+)
 from .linear_thread import linear_thread
 from .run_manager import RunManager, RunManagerDescriptor, run_manager
 from .voice import (
     VoiceManager,
     VoiceManagerDescriptor,
     VoiceProvider,
+    VoiceTurnCommitDeferred,
     end_voice,
     start_voice,
     temp_frontend_tool_result,
@@ -39,10 +49,18 @@ __all__ = [
     "VoiceManager",
     "VoiceManagerDescriptor",
     "VoiceProvider",
+    "VoiceTurnCommitDeferred",
     "checkpoint_thread",
+    "deepagents_drop_ephemeral",
+    "deepagents_ingest",
+    "deepagents_restamp",
+    "deepagents_role",
+    "deepagents_upsert_ephemeral",
     "end_voice",
     "langgraph_checkpoints",
+    "langgraph_checkpointer",
     "langgraph_snapshot",
+    "langgraph_state_snapshot",
     "langgraph_switch_values",
     "linear_thread",
     "run_manager",

harness_sdk/checkpoints.py

@@ -27,6 +27,7 @@ _TIP_TRIGGERS = (
     "stop-continue",
     "steer",
 )
+_UNREAD: Any = object()
 
 
 @dataclass(frozen=True)
@@ -326,37 +327,66 @@ def _values_snapshot(
     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}")
+    # an unwritten channel is absent from the checkpoint
+    messages = values.get(messages_key, [])
+    if not isinstance(messages, list):
+        raise TypeError(f"{messages_key} channel must be a list, got {messages!r}")
     return Snapshot(
         checkpoint_id=checkpoint_id,
-        messages=[_message_dict(m) for m in values.get(messages_key, [])],
+        messages=[_message_dict(m) for m in messages],
         lineage=lineage,
     )
 
 
-def langgraph_checkpoints(
-    graph: Any,
+def langgraph_state_snapshot(
+    state: Any,
+    *,
+    messages_key: str = "messages",
+    lineage_key: str = "lineage",
+) -> Snapshot | None:
+    """A Snapshot from a compiled graph's ``aget_state`` result; None for a thread without a checkpoint."""
+    if state.metadata is None:
+        return None
+    return _values_snapshot(
+        state.config["configurable"]["checkpoint_id"],
+        state.values,
+        messages_key,
+        lineage_key,
+    )
+
+
+def _thread_config(
+    thread_id: str, checkpoint_ns: str, checkpoint_id: str | None
+) -> dict[str, Any]:
+    configurable: dict[str, Any] = {
+        "thread_id": thread_id,
+        "checkpoint_ns": checkpoint_ns,
+    }
+    if checkpoint_id is not None:
+        configurable["checkpoint_id"] = checkpoint_id
+    return {"configurable": configurable}
+
+
+def langgraph_checkpointer(
+    checkpointer: Any,
     *,
     thread_id: str,
     checkpoint_ns: str = "",
     messages_key: str = "messages",
     lineage_key: str = "lineage",
 ) -> CheckpointStore:
-    """A CheckpointStore over a compiled langgraph graph's ``aget_state``."""
+    """A CheckpointStore over a langgraph checkpoint saver'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
-            state = await graph.aget_state({"configurable": configurable})
-            if state.metadata is None:
+            saved = await checkpointer.aget_tuple(
+                _thread_config(thread_id, checkpoint_ns, checkpoint_id)
+            )
+            if saved is None:
                 return None
             return _values_snapshot(
-                state.config["configurable"]["checkpoint_id"],
-                state.values,
+                saved.checkpoint["id"],
+                saved.checkpoint["channel_values"],
                 messages_key,
                 lineage_key,
             )
@@ -364,6 +394,29 @@ def langgraph_checkpoints(
     return _Store()
 
 
+def langgraph_checkpoints(
+    graph: Any,
+    *,
+    thread_id: str,
+    checkpoint_ns: str = "",
+    messages_key: str = "messages",
+    lineage_key: str = "lineage",
+) -> CheckpointStore:
+    """A CheckpointStore over a compiled langgraph graph's ``aget_state``."""
+
+    class _Store:
+        async def get(self, checkpoint_id: str | None) -> Snapshot | None:
+            return langgraph_state_snapshot(
+                await graph.aget_state(
+                    _thread_config(thread_id, checkpoint_ns, checkpoint_id)
+                ),
+                messages_key=messages_key,
+                lineage_key=lineage_key,
+            )
+
+    return _Store()
+
+
 def langgraph_snapshot(
     payload: dict[str, Any],
     *,
@@ -415,6 +468,10 @@ class CheckpointThread:
     _path: list[int] = field(default_factory=lambda: [0])
     _pending_fork: dict[str, Any] | None = None
 
+    def abandon_plan(self) -> None:
+        """Discard a planned fork whose run will not be ingested."""
+        self._pending_fork = None
+
     def _projection(self):
         return linear_thread(messages=lambda: self.transcript, role=self._role)
 
@@ -463,6 +520,8 @@ class CheckpointThread:
 
     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``."""
+        # an abandoned plan's fork must not leak into the next turn's ingest
+        self.abandon_plan()
         if trigger in _TIP_TRIGGERS:
             if rollback_to is not None:
                 raise ValueError(f"rollback_to is invalid for trigger {trigger!r}")
@@ -657,15 +716,21 @@ async def checkpoint_thread(
     thread_id: str,
     checkpoint_ns: str = "",
     role: Callable[[dict[str, Any]], str],
+    snapshot: Snapshot | None = _UNREAD,
 ) -> CheckpointThread:
-    """Boot a CheckpointThread: the thread row and the latest snapshot load concurrently."""
+    """Boot a CheckpointThread: the thread row and the latest snapshot load concurrently; ``snapshot`` hands in a latest snapshot already read (``None`` for a thread without one) and skips the store read."""
     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)
-    )
+    if snapshot is _UNREAD:
+        branches, snapshot = await asyncio.gather(
+            index.read_thread(thread_id, checkpoint_ns), store.get(None)
+        )
+    else:
+        if snapshot is not None and not isinstance(snapshot, Snapshot):
+            raise TypeError(f"snapshot must be a Snapshot or None, got {snapshot!r}")
+        branches = await index.read_thread(thread_id, checkpoint_ns)
     thread = CheckpointThread(
         transcript=[],
         branches=branches,

harness_sdk/deepagents.py

@@ -8,6 +8,7 @@ the target branch's tip as the checkpointer's latest checkpoint, so resumption
 and reboot both land on the switched world.
 """
 
+import asyncio
 from datetime import datetime, timezone
 from typing import Any
 
@@ -15,16 +16,20 @@ from statepatch import plain
 
 from .checkpoints import (
     CheckpointIndex,
-    CheckpointStore,
+    CheckpointThread,
     Snapshot,
-    _values_snapshot,
     checkpoint_thread,
+    langgraph_checkpointer,
     langgraph_snapshot,
 )
 
 
-def _role(message: dict[str, Any]) -> str:
-    return "user" if message["type"] == "human" else "assistant"
+def deepagents_role(message: dict[str, Any]) -> str:
+    """The harness role of a deepagents message: ``human`` is the user, ``ai``/``tool``/``system`` the assistant."""
+    try:
+        return _TYPE_TO_ROLE[message["type"]]
+    except KeyError as exc:
+        raise ValueError(f"unsupported message type: {message['type']!r}") from exc
 
 
 def _user_text(message: dict[str, Any]) -> str:
@@ -33,28 +38,102 @@ def _user_text(message: dict[str, Any]) -> str:
     )
 
 
-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()
+_ROLE_TO_TYPE = {"user": "human", "assistant": "ai"}
+_TYPE_TO_ROLE = {"human": "user", "ai": "assistant", "tool": "assistant", "system": "assistant"}
+
+
+def _message_type(role: str) -> str:
+    try:
+        return _ROLE_TO_TYPE[role]
+    except KeyError as exc:
+        raise ValueError(f"unsupported message role: {role!r}") from exc
+
+
+def _next_metadata(saved: Any | None) -> dict[str, Any]:
+    if saved is None:
+        return {"source": "update", "step": 0, "parents": {}}
+    metadata = saved.metadata
+    if not isinstance(metadata, dict):
+        raise TypeError(f"checkpoint metadata must be an object, got {metadata!r}")
+    step = metadata["step"]
+    if isinstance(step, bool) or not isinstance(step, int):
+        raise TypeError(f"checkpoint metadata step must be an integer, got {step!r}")
+    return {**metadata, "source": "update", "step": step + 1}
+
+
+def _checkpoint_messages(checkpoint: dict[str, Any]) -> list[Any]:
+    # an unwritten channel is absent from the checkpoint
+    messages = checkpoint["channel_values"].get("messages", [])
+    if not isinstance(messages, list):
+        raise TypeError(f"messages channel must be a list, got {messages!r}")
+    return messages
+
+
+def deepagents_restamp(
+    state: Any,
+    thread: CheckpointThread,
+    ephemeral: dict[str, dict[str, Any]],
+) -> None:
+    """Replace a Statewire messages projection with a thread and ephemeral entries."""
+    from statewire.langgraph import replace_messages
+
+    stamped = []
+    for message in thread.transcript:
+        branch, siblings = thread.sibling_info(message["id"])
+        stamped.append({**message, "branch": {"id": branch, "siblings": siblings}})
+    durable = {message["id"] for message in stamped}
+    stamped.extend(
+        {**entry, "ephemeral": True}
+        for id_, entry in ephemeral.items()
+        if id_ not in durable
+    )
+    replace_messages(state, stamped)
+
+
+def deepagents_upsert_ephemeral(
+    state: Any,
+    thread: CheckpointThread,
+    ephemeral: dict[str, dict[str, Any]],
+    message: dict[str, Any],
+) -> None:
+    """Add or replace one ephemeral deepagents message and restamp the projection."""
+    message_type = _message_type(message["role"])
+    if message["id"] not in ephemeral and any(
+        m["id"] == message["id"] for m in thread.transcript
+    ):
+        raise ValueError(f"message {message['id']!r} is already durable")
+    ephemeral[message["id"]] = {
+        "id": message["id"],
+        "type": message_type,
+        "content": message["text"],
+    }
+    deepagents_restamp(state, thread, ephemeral)
+
+
+def deepagents_drop_ephemeral(
+    state: Any,
+    thread: CheckpointThread,
+    ephemeral: dict[str, dict[str, Any]],
+    message_id: str,
+) -> None:
+    """Remove one ephemeral deepagents message and restamp the projection."""
+    del ephemeral[message_id]
+    deepagents_restamp(state, thread, ephemeral)
+
+
+async def deepagents_ingest(
+    thread: CheckpointThread,
+    captured: dict[str, Any] | None,
+    *,
+    refetch: bool = False,
+) -> None:
+    """Ingest a captured root checkpoint or explicitly refetch the checkpointer tip."""
+    if refetch:
+        await thread.ingest_latest()
+    elif captured is None:
+        raise RuntimeError("turn did not produce a root checkpoint")
+    else:
+        await thread.ingest(langgraph_snapshot(captured))
 
 
 class DeepagentsMessages:
@@ -68,6 +147,7 @@ class DeepagentsMessages:
         self._checkpointer = checkpointer
         self._captured: dict[str, Any] | None = None
         self._ephemeral: dict[str, dict[str, Any]] = {}
+        self._write_lock = asyncio.Lock()
 
     async def boot(
         self, state: Any, *, index: CheckpointIndex, thread_id: str
@@ -75,16 +155,22 @@ class DeepagentsMessages:
         """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._store = langgraph_checkpointer(
+            self._checkpointer, thread_id=thread_id, checkpoint_ns=""
+        )
         self._thread = await checkpoint_thread(
-            store=self._store, index=index, thread_id=thread_id, role=_role
+            store=self._store, index=index, thread_id=thread_id, role=deepagents_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:
+            if row is None:
+                raise KeyError(
+                    f"message {self._thread.transcript[-1]['id']!r} is not indexed"
+                )
+            if row["branch_id"] != 0:
                 await self._thread.switch_state(row["branch_id"])
         self._restamp()
 
@@ -104,6 +190,8 @@ class DeepagentsMessages:
         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."""
+        self._captured = None
+        self._thread.abandon_plan()
         config = self.config()
         if ctx.resume:
             return None, config
@@ -142,29 +230,32 @@ class DeepagentsMessages:
     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))
+        await deepagents_ingest(self._thread, captured, refetch=refetch)
         self._restamp()
 
+    def abandon_turn(self) -> None:
+        """Discard a stream capture and any fork planned for the abandoned turn."""
+        self._captured = None
+        self._thread.abandon_plan()
+
     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 append_messages(self, messages: list[dict[str, Any]]) -> None:
         """Append ``{id, role, text}`` entries as a new checkpoint on the current tip and restamp; the voice commit path."""
+        async with self._write_lock:
+            await self._append_messages(messages)
+
+    async def _append_messages(self, messages: list[dict[str, Any]]) -> None:
         from langchain_core.messages import AIMessage, HumanMessage
         from langgraph.checkpoint.base import empty_checkpoint
         from langgraph.checkpoint.base.id import uuid6
 
         if not messages:
             raise ValueError("messages must be non-empty")
-        for message in messages:
-            if message["role"] not in ("user", "assistant"):
-                raise ValueError(f"unsupported message role: {message['role']!r}")
         entries = [
-            (HumanMessage if m["role"] == "user" else AIMessage)(
+            (HumanMessage if _message_type(m["role"]) == "human" else AIMessage)(
                 id=m["id"], content=m["text"]
             )
             for m in messages
@@ -173,10 +264,10 @@ class DeepagentsMessages:
         saved = await self._checkpointer.aget_tuple(parent)
         if saved is None:
             checkpoint = empty_checkpoint()
-            step = -1
+            previous_messages = []
         else:
             checkpoint = {**saved.checkpoint}
-            step = (saved.metadata or {}).get("step", -1)
+            previous_messages = _checkpoint_messages(checkpoint)
             parent["configurable"]["checkpoint_id"] = saved.checkpoint["id"]
         version = self._checkpointer.get_next_version(
             checkpoint["channel_versions"].get("messages"), None
@@ -188,7 +279,7 @@ class DeepagentsMessages:
             "channel_values": {
                 **checkpoint["channel_values"],
                 "messages": [
-                    *checkpoint["channel_values"].get("messages", []),
+                    *previous_messages,
                     *entries,
                 ],
             },
@@ -197,17 +288,18 @@ class DeepagentsMessages:
                 "messages": version,
             },
         }
-        metadata = {"source": "update", "step": step + 1, "parents": {}}
+        metadata = _next_metadata(saved)
         await self._checkpointer.aput(parent, checkpoint, metadata, {"messages": version})
         await self._thread.ingest_latest()
         self._restamp()
 
     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 with self._write_lock:
+            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
@@ -227,11 +319,7 @@ class DeepagentsMessages:
             "id": str(uuid6(clock_seq=-2)),
             "ts": datetime.now(timezone.utc).isoformat(),
         }
-        metadata = {
-            "source": "update",
-            "step": (saved.metadata or {}).get("step", -1) + 1,
-            "parents": {},
-        }
+        metadata = _next_metadata(saved)
         written = await self._checkpointer.aput(parent, clone, metadata, {})
         return written["configurable"]["checkpoint_id"]
 
@@ -245,23 +333,15 @@ class DeepagentsMessages:
 
     def upsert_ephemeral_message(self, message: dict[str, Any]) -> None:
         """Render an in-flight ``{id, role, text}`` entry in ``state["messages"]`` tagged ``ephemeral: true``; an ``append_messages`` of the same id replaces it durably."""
-        if message["role"] not in ("user", "assistant"):
-            raise ValueError(f"unsupported message role: {message['role']!r}")
-        if message["id"] not in self._ephemeral and any(
-            m["id"] == message["id"] for m in self._thread.transcript
-        ):
-            raise ValueError(f"message {message['id']!r} is already durable")
-        self._ephemeral[message["id"]] = {
-            "id": message["id"],
-            "type": "human" if message["role"] == "user" else "ai",
-            "content": message["text"],
-        }
-        self._restamp()
+        deepagents_upsert_ephemeral(
+            self._state, self._thread, self._ephemeral, message
+        )
 
     def drop_ephemeral_message(self, message_id: str) -> None:
         """Remove an in-flight entry from the ``state["messages"]`` projection."""
-        del self._ephemeral[message_id]
-        self._restamp()
+        deepagents_drop_ephemeral(
+            self._state, self._thread, self._ephemeral, message_id
+        )
 
     def _append_user(self, message: dict[str, Any]) -> dict[str, Any]:
         entry = {"id": message["id"], "type": "human", "content": _user_text(message)}
@@ -275,14 +355,4 @@ class DeepagentsMessages:
         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}})
-        durable = {message["id"] for message in stamped}
-        stamped.extend(
-            {**entry, "ephemeral": True}
-            for id_, entry in self._ephemeral.items()
-            if id_ not in durable
-        )
-        self._replace(stamped)
+        deepagents_restamp(self._state, self._thread, self._ephemeral)