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

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

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

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

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

@@ -3,11 +3,11 @@ 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.
+position; entry 0 is the root branch and every fork entry records the branch
+it forked from as ``fork_branch``. Every ingested checkpoint advances
+``last_checkpoint_id`` for the state's ``branch_ids`` path's leaf branch, so
+a branch's pointer is always its own continuation's most recent checkpoint.
+Rows and entries are insert-only; ``last_checkpoint_id`` is the only update.
 """
 
 import asyncio
@@ -16,6 +16,8 @@ from typing import Any, Callable, Protocol
 
 from .linear_thread import linear_thread
 
+_ROOT_ENTRY = {"fork_message_id": None, "fork_branch": None, "last_checkpoint_id": None}
+
 _FORK_TRIGGERS = ("message-edit", "message-reload")
 _TIP_TRIGGERS = (
     "message-send",
@@ -65,6 +67,7 @@ class CheckpointIndex(Protocol):
         checkpoint_ns: str,
         *,
         fork_message_id: str | None,
+        fork_branch: int,
         last_checkpoint_id: str,
     ) -> int: ...
 
@@ -85,7 +88,7 @@ class CheckpointIndex(Protocol):
 
 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:
+        if not isinstance(branch_id, int) or not 0 <= branch_id < count:
             raise KeyError(f"unknown branch_id {branch_id!r} (index holds {count})")
 
 
@@ -107,16 +110,20 @@ class MemoryCheckpointIndex:
         checkpoint_ns: str,
         *,
         fork_message_id: str | None,
+        fork_branch: int,
         last_checkpoint_id: str,
     ) -> int:
         branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
+        if not branches:
+            branches.append(dict(_ROOT_ENTRY))
         branches.append(
             {
                 "fork_message_id": fork_message_id,
+                "fork_branch": fork_branch,
                 "last_checkpoint_id": last_checkpoint_id,
             }
         )
-        return len(branches)
+        return len(branches) - 1
 
     async def bump_tips(
         self,
@@ -125,10 +132,12 @@ class MemoryCheckpointIndex:
         branch_ids: list[int],
         checkpoint_id: str,
     ) -> None:
-        branches = self._threads.get((thread_id, checkpoint_ns), [])
+        branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
+        if not branches:
+            branches.append(dict(_ROOT_ENTRY))
         _check_branch_ids(branch_ids, len(branches))
         for branch_id in branch_ids:
-            branches[branch_id - 1]["last_checkpoint_id"] = checkpoint_id
+            branches[branch_id]["last_checkpoint_id"] = checkpoint_id
 
     async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
         for row in rows:
@@ -189,33 +198,40 @@ class PostgresCheckpointIndex:
         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:
+        from psycopg.types.json import Jsonb
+
+        await self._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])),
+        )
+
     async def append_branch(
         self,
         thread_id: str,
         checkpoint_ns: str,
         *,
         fork_message_id: str | None,
+        fork_branch: int,
         last_checkpoint_id: str,
     ) -> int:
         from psycopg.types.json import Jsonb
 
         entry = {
             "fork_message_id": fork_message_id,
+            "fork_branch": fork_branch,
             "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),
-            )
+            await self._seed_root(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]
+            return (await cur.fetchone())[0] - 1
 
     async def bump_tips(
         self,
@@ -224,6 +240,7 @@ 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",
@@ -233,9 +250,9 @@ class PostgresCheckpointIndex:
         _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)"
+            " CASE WHEN idx::int - 1 = ANY(%s)"
             " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
-            " ELSE elem END)"
+            " 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),
@@ -358,7 +375,7 @@ class CheckpointThread:
         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."""
+        """Fork entries at ``message_id``'s parent, as {branch_id, **entry} in allocation order; the source branch's continuation has no entry."""
         index = next(
             (i for i, m in enumerate(self.transcript) if m["id"] == message_id), None
         )
@@ -366,15 +383,15 @@ class CheckpointThread:
             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}
+            {"branch_id": i, **entry}
             for i, entry in enumerate(self.branches)
-            if entry["fork_message_id"] == parent_id
+            if i and 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
+            fork_id = self.branches[b]["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
             )
@@ -383,11 +400,16 @@ class CheckpointThread:
         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."""
+        """``message_id``'s own branch id and its full fan-out's branch ids (source branch included) 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})
+        sources = {e["fork_branch"] for e in entries}
+        if len(sources) > 1:
+            raise ValueError(
+                f"fork entries at one fan-out disagree on fork_branch: {sorted(sources)}"
+            )
+        return branch, sorted({e["branch_id"] for e in entries} | sources | {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``."""
@@ -427,14 +449,25 @@ class CheckpointThread:
         }
         return Plan(anchor, survivors, True)
 
+    async def plan_run(self, ctx: Any) -> Plan:
+        """The resumption for a RunManager dispatch: trigger and rollback surface read off the run context."""
+        rollback = ctx.rollback_to if ctx.has_rollback else None
+        return await self.plan(ctx.trigger, rollback_to=rollback)
+
+    async def ingest_latest(self) -> None:
+        """Ingest the store's latest checkpoint."""
+        await self.ingest(await self._get(None))
+
     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."""
+        """Index one new checkpoint: bump the path leaf's tip, 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)
         )
+        if not path:
+            raise ValueError("lineage branch_ids must not be empty")
         anchor = self.checkpoint_id
         if pending is not None:
             # the new world descends from the fork message's branch, not the window's leaf
@@ -444,25 +477,29 @@ class CheckpointThread:
                 self._thread_id,
                 self._ns,
                 fork_message_id=pending["fork_message_id"],
+                fork_branch=pending["fork_branch"],
                 last_checkpoint_id=snapshot.checkpoint_id,
             )
+            if not self.branches:
+                self.branches = [dict(_ROOT_ENTRY)]
             self.branches = [
                 *self.branches,
                 {
                     "fork_message_id": pending["fork_message_id"],
+                    "fork_branch": pending["fork_branch"],
                     "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
+        tip = path[-1]
+        await self._index.bump_tips(
+            self._thread_id, self._ns, [tip], snapshot.checkpoint_id
+        )
+        if not self.branches:
+            self.branches = [dict(_ROOT_ENTRY)]
+        if tip < len(self.branches):
+            self.branches[tip]["last_checkpoint_id"] = snapshot.checkpoint_id
         known = {m["id"] for m in self.transcript}
         rows = [
             {
@@ -470,7 +507,7 @@ class CheckpointThread:
                 "checkpoint_ns": self._ns,
                 "message_id": message["id"],
                 "anchor_checkpoint_id": anchor,
-                "branch_id": path[-1] if path else 0,
+                "branch_id": path[-1],
             }
             for message in snapshot.messages
             if message["id"] not in known
@@ -493,11 +530,12 @@ class CheckpointThread:
         )
         return self
 
-    def _tip_of(self, branch_id: int) -> str | None:
-        if branch_id == 0:
-            return None
+    def _tip_of(self, branch_id: int) -> str:
         _check_branch_ids([branch_id], len(self.branches))
-        return self.branches[branch_id - 1]["last_checkpoint_id"]
+        tip = self.branches[branch_id]["last_checkpoint_id"]
+        if tip is None:
+            raise LookupError(f"branch {branch_id} has no recorded tip")
+        return tip
 
     async def _get(self, checkpoint_id: str | None) -> Snapshot:
         snapshot = await self._store.get(checkpoint_id)
@@ -506,20 +544,14 @@ class CheckpointThread:
         return snapshot
 
     async def switch_branch(self, branch_id: int) -> "CheckpointThread":
-        """Load the branch's subtree tip as the current window."""
+        """Load the branch's own continuation 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]:
+    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"]
+            branch_id = self.branches[branch_id]["fork_branch"]
         return path
 
     async def switch_state(
@@ -528,7 +560,7 @@ class CheckpointThread:
         """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),
+            "branch_ids": self._world_path(branch_id),
             "parent_checkpoint_id": None,
         }
         self._path = list(lineage["branch_ids"])
@@ -547,44 +579,25 @@ class CheckpointThread:
             raise ValueError("the current window has no tail message")
         tail = self.transcript[-1]["id"]
         candidates = [
-            i + 1
+            i
             for i, entry in enumerate(self.branches)
-            if entry["fork_message_id"] == tail
+            if i and 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."""
+        """Load ``message_id``'s own branch's tip as the current window."""
         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}")
+        snapshot = await self._get(self._tip_of(row["branch_id"]))
+        if not any(m["id"] == message_id for m in snapshot.messages):
+            raise LookupError(
+                f"branch {row['branch_id']} tip does not contain message {message_id!r}"
+            )
+        return self._load(snapshot)
 
 
 async def checkpoint_thread(

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

@@ -30,6 +30,22 @@ def msg(id, type="human"):
 U1, A1, U2, A2 = msg("u1"), msg("a1", "ai"), msg("u2"), msg("a2", "ai")
 
 
+def root(last_checkpoint_id=None):
+    return {
+        "fork_message_id": None,
+        "fork_branch": None,
+        "last_checkpoint_id": last_checkpoint_id,
+    }
+
+
+def fork(fork_message_id, fork_branch, last_checkpoint_id):
+    return {
+        "fork_message_id": fork_message_id,
+        "fork_branch": fork_branch,
+        "last_checkpoint_id": last_checkpoint_id,
+    }
+
+
 @dataclass
 class FakeStore:
     snapshots: dict[str, Snapshot] = field(default_factory=dict)
@@ -79,7 +95,7 @@ async def test_boot_fresh_thread_is_empty():
 async def test_boot_loads_latest_snapshot_and_branches():
     store, index = FakeStore(), MemoryCheckpointIndex()
     await index.append_branch(
-        THREAD, "", fork_message_id="u1", last_checkpoint_id="cp2"
+        THREAD, "", fork_message_id="u1", fork_branch=0, last_checkpoint_id="cp2"
     )
     lineage = {"branch_ids": [1], "parent_checkpoint_id": None}
     store.put("cp2", [U1, A1], lineage)
@@ -87,9 +103,7 @@ async def test_boot_loads_latest_snapshot_and_branches():
     assert thread.transcript == [U1, A1]
     assert thread.checkpoint_id == "cp2"
     assert thread.lineage == lineage
-    assert thread.branches == [
-        {"fork_message_id": "u1", "last_checkpoint_id": "cp2"}
-    ]
+    assert thread.branches == [root(), fork("u1", 0, "cp2")]
 
 
 async def test_boot_rejects_empty_thread_id():
@@ -143,7 +157,7 @@ async def test_ingest_is_idempotent():
     store, index, thread = await seeded()
     await thread.ingest(store.snapshots["cp2"])
     assert await row(index, "u2") == {"anchor_checkpoint_id": "cp1", "branch_id": 0}
-    assert thread.branches == []
+    assert thread.branches == [root("cp2")]
     assert thread.transcript == [U1, A1, U2, A2]
 
 
@@ -153,13 +167,11 @@ async def test_fork_ingest_appends_the_branch_entry_once():
     a2b = msg("a2b", "ai")
     snapshot = store.put("cp3", [U1, A1, U2, a2b])
     await thread.ingest(snapshot)
-    assert thread.branches == [
-        {"fork_message_id": "u2", "last_checkpoint_id": "cp3"}
-    ]
+    assert thread.branches == [root("cp2"), fork("u2", 0, "cp3")]
     assert await index.read_thread(THREAD, "") == thread.branches
     assert await row(index, "a2b") == {"anchor_checkpoint_id": "cp1", "branch_id": 1}
     await thread.ingest(snapshot)  # redelivery: the pending fork is consumed
-    assert len(await index.read_thread(THREAD, "")) == 1
+    assert len(await index.read_thread(THREAD, "")) == 2
 
 
 async def test_ingest_after_a_fork_bumps_the_forked_branch():
@@ -169,7 +181,7 @@ async def test_ingest_after_a_fork_bumps_the_forked_branch():
     await thread.ingest(store.put("cp3", [U1, A1, U2, a2b]))
     u3 = msg("u3")
     await thread.ingest(store.put("cp4", [U1, A1, U2, a2b, u3]))
-    assert (await index.read_thread(THREAD, ""))[0]["last_checkpoint_id"] == "cp4"
+    assert (await index.read_thread(THREAD, ""))[1]["last_checkpoint_id"] == "cp4"
     assert await row(index, "u3") == {"anchor_checkpoint_id": "cp3", "branch_id": 1}
 
 
@@ -183,8 +195,8 @@ async def test_fork_from_a_shared_message_descends_from_its_branch():
     u2c, a2c = msg("u2c"), msg("a2c", "ai")
     await thread.ingest(store.put("cp4", [U1, A1, u2c, a2c]))
     branches = await index.read_thread(THREAD, "")
-    assert branches[0]["last_checkpoint_id"] == "cp3"  # branch 1 untouched
-    assert branches[1] == {"fork_message_id": "a1", "last_checkpoint_id": "cp4"}
+    assert branches[1]["last_checkpoint_id"] == "cp3"  # branch 1 untouched
+    assert branches[2] == fork("a1", 0, "cp4")
     assert await row(index, "u2c") == {"anchor_checkpoint_id": None, "branch_id": 2}
 
 
@@ -225,6 +237,60 @@ async def test_plan_rejects_unknown_rollback_and_trigger():
         await thread.plan("bogus")
 
 
+# ─── Run resumption seam ────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class Ctx:
+    trigger: str
+    has_rollback: bool = False
+    rollback_to: str | None = None
+
+
+async def test_plan_run_edit_allocates_a_fork_entry():
+    store, index, thread = await seeded()
+    plan = await thread.plan_run(Ctx("message-edit", True, "a1"))
+    assert plan == Plan(None, [U1, A1], True)
+    await thread.ingest(store.put("cp3", [U1, A1, msg("u2b"), msg("a2b", "ai")]))
+    assert await index.read_thread(THREAD, "") == [
+        root("cp2"),
+        {"fork_message_id": "a1", "fork_branch": 0, "last_checkpoint_id": "cp3"},
+    ]
+
+
+async def test_plan_run_reload_allocates_a_fork_entry():
+    store, index, thread = await seeded()
+    plan = await thread.plan_run(Ctx("message-reload", True, "u2"))
+    assert plan == Plan("cp1", [U2], True)
+    await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
+    assert await index.read_thread(THREAD, "") == [
+        root("cp2"),
+        {"fork_message_id": "u2", "fork_branch": 0, "last_checkpoint_id": "cp3"},
+    ]
+
+
+async def test_plan_run_tip_triggers_allocate_nothing():
+    store, index, thread = await seeded()
+    for trigger in (
+        "message-send",
+        "input-resume",
+        "error-continue",
+        "stop-continue",
+        "steer",
+    ):
+        assert await thread.plan_run(Ctx(trigger)) == Plan("cp2", [], False)
+    await thread.ingest(store.put("cp3", [U1, A1, U2, A2, msg("u3")]))
+    assert await index.read_thread(THREAD, "") == [root("cp3")]
+
+
+async def test_ingest_latest_reads_the_store_tip():
+    store, _, thread = await seeded()
+    store.put("cp3", [U1, A1, U2, A2, msg("u3")])
+    await thread.ingest_latest()
+    assert thread.checkpoint_id == "cp3"
+    assert thread.transcript[-1] == msg("u3")
+
+
 # ─── Sibling fan-out ────────────────────────────────────────
 
 
@@ -235,8 +301,8 @@ async def test_sibling_branches_order_by_allocation():
     await thread.plan("message-reload", rollback_to="u2")
     await thread.ingest(store.put("cp4", [U1, A1, U2, msg("a2c", "ai")]))
     assert thread.sibling_branches("a2c") == [
-        {"branch_id": 1, "fork_message_id": "u2", "last_checkpoint_id": "cp3"},
-        {"branch_id": 2, "fork_message_id": "u2", "last_checkpoint_id": "cp4"},
+        {"branch_id": 1, **fork("u2", 0, "cp3")},
+        {"branch_id": 2, **fork("u2", 0, "cp4")},
     ]
     assert thread.sibling_branches("u1") == []
     with pytest.raises(KeyError, match="not in the current window"):
@@ -248,18 +314,28 @@ async def test_sibling_info_lists_the_fan_out_ids_in_creation_order():
     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("a2d") == (3, [0, 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])
+    assert thread.sibling_info("a2b") == (1, [0, 1, 2, 3])
+
+
+async def test_sibling_info_is_branch_invariant():
+    store, _, thread = await seeded()
+    await thread.plan("message-reload", rollback_to="u2")
+    await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
+    from_fork = thread.sibling_info("a2b")[1]
+    await thread.switch_state(0)
+    from_source = thread.sibling_info("a2")[1]
+    assert from_fork == from_source == [0, 1]
 
 
 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])
+    assert thread.sibling_info("u2c") == (1, [0, 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
@@ -269,7 +345,7 @@ async def test_sibling_info_counts_the_window_timeline_among_the_forks():
 # ─── Switch branch ──────────────────────────────────────────
 
 
-async def test_switch_branch_loads_the_subtree_tip():
+async def test_switch_branch_loads_the_branch_tip():
     store, index, thread = await seeded()
     await thread.plan("message-reload", rollback_to="u2")
     a2b = msg("a2b", "ai")
@@ -279,10 +355,14 @@ async def test_switch_branch_loads_the_subtree_tip():
     assert thread.transcript == [U1, A1, U2, a2b]
 
 
-async def test_switch_branch_zero_loads_the_thread_latest():
+async def test_switch_branch_round_trips_to_the_source_branch():
     store, _, thread = await seeded()
+    await thread.plan("message-reload", rollback_to="u2")
+    await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
+    assert thread.checkpoint_id == "cp3"
     await thread.switch_branch(0)
     assert thread.checkpoint_id == "cp2"
+    assert thread.transcript == [U1, A1, U2, A2]
 
 
 async def test_switch_branch_unknown_raises():
@@ -329,11 +409,12 @@ async def test_switch_state_ingest_round_trip_keeps_the_index_consistent():
     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 branches[0]["last_checkpoint_id"] == "cp2"
+    assert branches[1]["last_checkpoint_id"] == "cp5"
+    assert branches[2]["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])
+    assert thread.sibling_info("a2b") == (1, [0, 1, 2])
 
 
 # ─── Paging across a compaction ─────────────────────────────
@@ -342,7 +423,7 @@ async def test_switch_state_ingest_round_trip_keeps_the_index_consistent():
 async def compacted():
     store, index, old = await seeded()
     branch_id = await index.append_branch(
-        THREAD, "", fork_message_id="a2", last_checkpoint_id="cp3"
+        THREAD, "", fork_message_id="a2", fork_branch=0, last_checkpoint_id="cp3"
     )
     summary = msg("sum", "ai")
     store.put(
@@ -365,7 +446,7 @@ async def compacted():
 async def test_compaction_window_restarts_the_branch_path():
     _, index, thread, window = await compacted()
     assert thread.transcript == window
-    assert (await index.read_thread(THREAD, ""))[0]["last_checkpoint_id"] == "cp4"
+    assert (await index.read_thread(THREAD, ""))[1]["last_checkpoint_id"] == "cp4"
     assert await row(index, "u3") == {"anchor_checkpoint_id": "cp3", "branch_id": 1}
 
 
@@ -407,17 +488,19 @@ async def test_locate_unknown_message_raises():
         await thread.locate("nope")
 
 
-async def test_locate_follows_a_fork_that_took_the_pointer():
+async def test_locate_loads_the_own_branch_tip_across_forks():
     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-edit", rollback_to="a1")
     await thread.ingest(store.put("cp4", [U1, A1, msg("u2c"), msg("a2c", "ai")]))
-    # u2 rides branch 0, whose tip (the thread latest, cp4) dropped it
+    # u2 rides branch 0, whose own tip stays cp2 while forks advance the thread
     await thread.locate("u2")
-    assert thread.checkpoint_id == "cp3"
+    assert thread.checkpoint_id == "cp2"
     assert any(m["id"] == "u2" for m in thread.transcript)
+    await thread.locate("a2b")
+    assert thread.checkpoint_id == "cp3"
 
 
 # ─── langgraph adapter ──────────────────────────────────────
@@ -571,19 +654,22 @@ async def pg_index(postgres_dsn):
 async def test_postgres_branches_round_trip(pg_index):
     assert await pg_index.read_thread(THREAD, "") == []
     first = await pg_index.append_branch(
-        THREAD, "", fork_message_id="u2", last_checkpoint_id="cp1"
+        THREAD, "", fork_message_id="u2", fork_branch=0, last_checkpoint_id="cp1"
     )
     second = await pg_index.append_branch(
-        THREAD, "", fork_message_id=None, last_checkpoint_id="cp2"
+        THREAD, "", fork_message_id=None, fork_branch=0, last_checkpoint_id="cp2"
     )
     assert (first, second) == (1, 2)
-    await pg_index.bump_tips(THREAD, "", [2], "cp3")
+    await pg_index.bump_tips(THREAD, "", [0, 2], "cp3")
     assert await pg_index.read_thread(THREAD, "") == [
-        {"fork_message_id": "u2", "last_checkpoint_id": "cp1"},
-        {"fork_message_id": None, "last_checkpoint_id": "cp3"},
+        root("cp3"),
+        fork("u2", 0, "cp1"),
+        fork(None, 0, "cp3"),
     ]
     with pytest.raises(KeyError, match="unknown branch_id"):
         await pg_index.bump_tips(THREAD, "", [3], "cp4")
+    await pg_index.bump_tips(THREAD, "ns2", [0], "cp1")
+    assert await pg_index.read_thread(THREAD, "ns2") == [root("cp1")]
 
 
 async def test_postgres_messages_upsert_keeps_the_first_row(pg_index):