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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: harness-sdk-python
-Version: 0.22.0
+Version: 0.23.1
 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.22.0 → harness_sdk_python-0.23.1}/pyproject.toml

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

{harness_sdk_python-0.22.0 → harness_sdk_python-0.23.1}/src/harness_sdk/__init__.py

@@ -8,6 +8,7 @@ from .checkpoints import (
     Snapshot,
     checkpoint_thread,
     langgraph_checkpoints,
+    langgraph_switch_values,
 )
 from .linear_thread import linear_thread
 from .run_manager import RunManager, RunManagerDescriptor, run_manager
@@ -24,6 +25,7 @@ __all__ = [
     "Snapshot",
     "checkpoint_thread",
     "langgraph_checkpoints",
+    "langgraph_switch_values",
     "linear_thread",
     "run_manager",
 ]

{harness_sdk_python-0.22.0 → harness_sdk_python-0.23.1}/src/harness_sdk/checkpoints.py

@@ -312,6 +312,26 @@ def langgraph_checkpoints(
     return _Store()
 
 
+def langgraph_switch_values(
+    transcript: list[dict[str, Any]],
+    lineage: dict[str, Any],
+    *,
+    messages_key: str = "messages",
+    lineage_key: str = "lineage",
+) -> dict[str, Any]:
+    """``update_state`` channel values that replace the messages channel with ``transcript`` and set ``lineage``."""
+    from langchain_core.messages import RemoveMessage, convert_to_messages
+    from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+    return {
+        messages_key: [
+            RemoveMessage(id=REMOVE_ALL_MESSAGES),
+            *convert_to_messages([_message_dict(m) for m in transcript]),
+        ],
+        lineage_key: lineage,
+    }
+
+
 @dataclass
 class CheckpointThread:
     """One thread's current window over the store, with the sidecar index kept in step."""
@@ -351,6 +371,24 @@ class CheckpointThread:
             if entry["fork_message_id"] == parent_id
         ]
 
+    def _branch_at(self, index: int) -> int:
+        branch = 0
+        for b in self._path:
+            fork_id = self.branches[b - 1]["fork_message_id"] if b != 0 else None
+            start = next(
+                (i + 1 for i, m in enumerate(self.transcript) if m["id"] == fork_id), 0
+            )
+            if start <= index:
+                branch = b
+        return branch
+
+    def sibling_info(self, message_id: str) -> tuple[int, list[int]]:
+        """``message_id``'s own branch id and its sibling fan-out's branch ids in creation order."""
+        entries = self.sibling_branches(message_id)
+        index = next(i for i, m in enumerate(self.transcript) if m["id"] == message_id)
+        branch = self._branch_at(index)
+        return branch, sorted({e["branch_id"] for e in entries} | {branch})
+
     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:
@@ -471,6 +509,31 @@ class CheckpointThread:
         """Load the branch's subtree tip as the current window."""
         return self._load(await self._get(self._tip_of(branch_id)))
 
+    async def _world_path(self, branch_id: int) -> list[int]:
+        path = [0]
+        while branch_id != 0:
+            path.insert(1, branch_id)
+            fork_id = self.branches[branch_id - 1]["fork_message_id"]
+            if fork_id is None:
+                break
+            row = await self._index.get_message(self._thread_id, self._ns, fork_id)
+            if row is None:
+                raise KeyError(f"fork message {fork_id!r} is not indexed")
+            branch_id = row["branch_id"]
+        return path
+
+    async def switch_state(
+        self, branch_id: int
+    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+        """Switch like ``switch_branch`` and return (transcript, lineage) for the host to write back into graph state."""
+        await self.switch_branch(branch_id)
+        lineage = self.lineage or {
+            "branch_ids": await self._world_path(branch_id),
+            "parent_checkpoint_id": None,
+        }
+        self._path = list(lineage["branch_ids"])
+        return list(self.transcript), lineage
+
     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")

{harness_sdk_python-0.22.0 → harness_sdk_python-0.23.1}/tests/test_checkpoints.py

@@ -13,6 +13,7 @@ from harness_sdk import (
     Snapshot,
     checkpoint_thread,
     langgraph_checkpoints,
+    langgraph_switch_values,
 )
 
 THREAD = "t1"
@@ -242,6 +243,29 @@ async def test_sibling_branches_order_by_allocation():
         thread.sibling_branches("a2")
 
 
+async def test_sibling_info_lists_the_fan_out_ids_in_creation_order():
+    store, _, thread = await seeded()
+    for cid, reply in [("cp3", "a2b"), ("cp4", "a2c"), ("cp5", "a2d")]:
+        await thread.plan("message-reload", rollback_to="u2")
+        await thread.ingest(store.put(cid, [U1, A1, U2, msg(reply, "ai")]))
+    assert thread.sibling_info("a2d") == (3, [1, 2, 3])
+    assert thread.sibling_info("u2") == (0, [0])
+    lineage = {"branch_ids": [0, 1], "parent_checkpoint_id": None}
+    await thread.ingest(store.put("cp6", [U1, A1, U2, msg("a2b", "ai")], lineage))
+    assert thread.sibling_info("a2b") == (1, [1, 2, 3])
+
+
+async def test_sibling_info_counts_the_window_timeline_among_the_forks():
+    store, _, thread = await seeded()
+    await thread.plan("message-edit", rollback_to="a1")
+    await thread.ingest(store.put("cp3", [U1, A1, msg("u2c"), msg("a2c", "ai")]))
+    assert thread.sibling_info("u2c") == (1, [1])
+    lineage = {"branch_ids": [0], "parent_checkpoint_id": None}
+    await thread.ingest(store.put("cp4", [U1, A1, U2, A2], lineage))
+    assert thread.sibling_info("u2") == (0, [0, 1])  # branch 0 before the edit fork
+    assert thread.sibling_info("u1") == (0, [0])
+
+
 # ─── Switch branch ──────────────────────────────────────────
 
 
@@ -267,6 +291,51 @@ async def test_switch_branch_unknown_raises():
         await thread.switch_branch(7)
 
 
+# ─── Switch as a state write ────────────────────────────────
+
+
+async def forked():
+    """Two reloads at u2: branch 1 tips cp3 with a2b, branch 2 tips cp4 with a2c."""
+    store, index, thread = await seeded()
+    await thread.plan("message-reload", rollback_to="u2")
+    a2b = msg("a2b", "ai")
+    await thread.ingest(store.put("cp3", [U1, A1, U2, a2b]))
+    await thread.plan("message-reload", rollback_to="u2")
+    await thread.ingest(store.put("cp4", [U1, A1, U2, msg("a2c", "ai")]))
+    return store, index, thread, [U1, A1, U2, a2b]
+
+
+async def test_switch_state_returns_the_transcript_and_world_lineage():
+    _, _, thread, world = await forked()
+    transcript, lineage = await thread.switch_state(1)
+    assert transcript == world
+    assert lineage == {"branch_ids": [0, 1], "parent_checkpoint_id": None}
+    assert thread.checkpoint_id == "cp3"
+    assert thread.transcript == world
+
+
+async def test_switch_state_returns_a_compacted_world_lineage():
+    _, _, thread, window = await compacted()
+    await thread.page_up()
+    transcript, lineage = await thread.switch_state(1)
+    assert transcript == window
+    assert lineage == {"branch_ids": [1], "parent_checkpoint_id": "cp2"}
+
+
+async def test_switch_state_ingest_round_trip_keeps_the_index_consistent():
+    store, index, thread, _ = await forked()
+    transcript, lineage = await thread.switch_state(1)
+    rows_before = len(index._messages)
+    await thread.ingest(store.put("cp5", transcript, lineage))
+    assert len(index._messages) == rows_before
+    branches = await index.read_thread(THREAD, "")
+    assert branches[0]["last_checkpoint_id"] == "cp5"
+    assert branches[1]["last_checkpoint_id"] == "cp4"
+    assert await row(index, "a2b") == {"anchor_checkpoint_id": "cp1", "branch_id": 1}
+    assert thread.checkpoint_id == "cp5"
+    assert thread.sibling_info("a2b") == (1, [1, 2])
+
+
 # ─── Paging across a compaction ─────────────────────────────
 
 
@@ -405,6 +474,28 @@ async def test_langgraph_adapter_rejects_a_message_without_an_id():
         await store.get(None)
 
 
+async def test_langgraph_switch_values_replaces_the_messages_channel():
+    pytest.importorskip("langgraph.graph")
+    from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+    from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+    lineage = {"branch_ids": [0, 1], "parent_checkpoint_id": None}
+    values = langgraph_switch_values(
+        [
+            {"id": "u1", "type": "human", "content": "hi"},
+            {"id": "a1", "type": "ai", "content": "yo"},
+        ],
+        lineage,
+    )
+    remove, u1, a1 = values["messages"]
+    assert remove == RemoveMessage(id=REMOVE_ALL_MESSAGES)
+    assert isinstance(u1, HumanMessage) and u1.id == "u1" and u1.content == "hi"
+    assert isinstance(a1, AIMessage) and a1.id == "a1" and a1.content == "yo"
+    assert values["lineage"] == lineage
+    with pytest.raises(ValueError, match="without an id"):
+        langgraph_switch_values([{"type": "human", "content": "x"}], lineage)
+
+
 # ─── PostgresCheckpointIndex ────────────────────────────────