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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/PKG-INFO +1 -1
  2. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/pyproject.toml +1 -1
  3. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/__init__.py +2 -0
  4. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/checkpoints.py +128 -87
  5. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/run_helpers.py +4 -4
  6. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_checkpoints.py +184 -65
  7. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/.gitignore +0 -0
  8. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/README.md +0 -0
  9. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/examples/__init__.py +0 -0
  10. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/examples/runs_app.py +0 -0
  11. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/fenced_postgres.py +0 -0
  12. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/linear_thread.py +0 -0
  13. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/run_manager.py +0 -0
  14. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_admit_steer.py +0 -0
  15. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_applied_visibility.py +0 -0
  16. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_batches.py +0 -0
  17. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_branch_anchor.py +0 -0
  18. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_dispatching.py +0 -0
  19. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_durability.py +0 -0
  20. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_edit_dispatched.py +0 -0
  21. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_edit_reload.py +0 -0
  22. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_enqueue.py +0 -0
  23. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_fenced_postgres.py +0 -0
  24. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_input_required.py +0 -0
  25. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_linear_thread.py +0 -0
  26. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_meta.py +0 -0
  27. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_outcomes.py +0 -0
  28. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_placement.py +0 -0
  29. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_prepare_hooks.py +0 -0
  30. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_restore.py +0 -0
  31. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_rewind_during_run.py +0 -0
  32. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_run_crash_sweep.py +0 -0
  33. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_settle.py +0 -0
  34. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_steer.py +0 -0
  35. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_stop_continue.py +0 -0
  36. {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_stopping_window.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: harness-sdk-python
3
- Version: 0.23.1
3
+ Version: 0.25.0
4
4
  Summary: RunManager: the harness-sdk runs subsystem for Python Statewire hosts
5
5
  Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
6
6
  License-Expression: MIT
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "harness-sdk-python"
3
- version = "0.23.1"
3
+ version = "0.25.0"
4
4
  description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -8,6 +8,7 @@ from .checkpoints import (
8
8
  Snapshot,
9
9
  checkpoint_thread,
10
10
  langgraph_checkpoints,
11
+ langgraph_snapshot,
11
12
  langgraph_switch_values,
12
13
  )
13
14
  from .linear_thread import linear_thread
@@ -25,6 +26,7 @@ __all__ = [
25
26
  "Snapshot",
26
27
  "checkpoint_thread",
27
28
  "langgraph_checkpoints",
29
+ "langgraph_snapshot",
28
30
  "langgraph_switch_values",
29
31
  "linear_thread",
30
32
  "run_manager",
@@ -3,11 +3,11 @@ checkpoint store (``get(checkpoint_id)`` only), designed for langgraph.
3
3
 
4
4
  A branch is a resumption point plus its continuation; edits, regenerations,
5
5
  reloads, and compactions all create one. ``branch_id`` is the entry's array
6
- position + 1; the root branch is 0 and the array holds fork entries only.
7
- Every ingested checkpoint advances ``last_checkpoint_id`` for every branch in
8
- the state's ``branch_ids`` path, so a branch's pointer is always its
9
- subtree's most recent checkpoint. Rows and entries are insert-only;
10
- ``last_checkpoint_id`` is the only update.
6
+ position; entry 0 is the root branch and every fork entry records the branch
7
+ it forked from as ``fork_branch``. Every ingested checkpoint advances
8
+ ``last_checkpoint_id`` for the state's ``branch_ids`` path's leaf branch, so
9
+ a branch's pointer is always its own continuation's most recent checkpoint.
10
+ Rows and entries are insert-only; ``last_checkpoint_id`` is the only update.
11
11
  """
12
12
 
13
13
  import asyncio
@@ -16,6 +16,8 @@ from typing import Any, Callable, Protocol
16
16
 
17
17
  from .linear_thread import linear_thread
18
18
 
19
+ _ROOT_ENTRY = {"fork_message_id": None, "fork_branch": None, "last_checkpoint_id": None}
20
+
19
21
  _FORK_TRIGGERS = ("message-edit", "message-reload")
20
22
  _TIP_TRIGGERS = (
21
23
  "message-send",
@@ -65,6 +67,7 @@ class CheckpointIndex(Protocol):
65
67
  checkpoint_ns: str,
66
68
  *,
67
69
  fork_message_id: str | None,
70
+ fork_branch: int,
68
71
  last_checkpoint_id: str,
69
72
  ) -> int: ...
70
73
 
@@ -85,7 +88,7 @@ class CheckpointIndex(Protocol):
85
88
 
86
89
  def _check_branch_ids(branch_ids: list[int], count: int) -> None:
87
90
  for branch_id in branch_ids:
88
- if not isinstance(branch_id, int) or not 1 <= branch_id <= count:
91
+ if not isinstance(branch_id, int) or not 0 <= branch_id < count:
89
92
  raise KeyError(f"unknown branch_id {branch_id!r} (index holds {count})")
90
93
 
91
94
 
@@ -107,16 +110,20 @@ class MemoryCheckpointIndex:
107
110
  checkpoint_ns: str,
108
111
  *,
109
112
  fork_message_id: str | None,
113
+ fork_branch: int,
110
114
  last_checkpoint_id: str,
111
115
  ) -> int:
112
116
  branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
117
+ if not branches:
118
+ branches.append(dict(_ROOT_ENTRY))
113
119
  branches.append(
114
120
  {
115
121
  "fork_message_id": fork_message_id,
122
+ "fork_branch": fork_branch,
116
123
  "last_checkpoint_id": last_checkpoint_id,
117
124
  }
118
125
  )
119
- return len(branches)
126
+ return len(branches) - 1
120
127
 
121
128
  async def bump_tips(
122
129
  self,
@@ -125,10 +132,12 @@ class MemoryCheckpointIndex:
125
132
  branch_ids: list[int],
126
133
  checkpoint_id: str,
127
134
  ) -> None:
128
- branches = self._threads.get((thread_id, checkpoint_ns), [])
135
+ branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
136
+ if not branches:
137
+ branches.append(dict(_ROOT_ENTRY))
129
138
  _check_branch_ids(branch_ids, len(branches))
130
139
  for branch_id in branch_ids:
131
- branches[branch_id - 1]["last_checkpoint_id"] = checkpoint_id
140
+ branches[branch_id]["last_checkpoint_id"] = checkpoint_id
132
141
 
133
142
  async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
134
143
  for row in rows:
@@ -189,33 +198,40 @@ class PostgresCheckpointIndex:
189
198
  row = await cur.fetchone()
190
199
  return row[0] if row is not None else []
191
200
 
201
+ async def _seed_root(self, thread_id: str, checkpoint_ns: str) -> None:
202
+ from psycopg.types.json import Jsonb
203
+
204
+ await self._conn.execute(
205
+ "INSERT INTO checkpoint_threads (thread_id, checkpoint_ns, branches)"
206
+ " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
207
+ (thread_id, checkpoint_ns, Jsonb([_ROOT_ENTRY])),
208
+ )
209
+
192
210
  async def append_branch(
193
211
  self,
194
212
  thread_id: str,
195
213
  checkpoint_ns: str,
196
214
  *,
197
215
  fork_message_id: str | None,
216
+ fork_branch: int,
198
217
  last_checkpoint_id: str,
199
218
  ) -> int:
200
219
  from psycopg.types.json import Jsonb
201
220
 
202
221
  entry = {
203
222
  "fork_message_id": fork_message_id,
223
+ "fork_branch": fork_branch,
204
224
  "last_checkpoint_id": last_checkpoint_id,
205
225
  }
206
226
  async with self._conn.transaction():
207
- await self._conn.execute(
208
- "INSERT INTO checkpoint_threads (thread_id, checkpoint_ns, branches)"
209
- " VALUES (%s, %s, '[]'::jsonb) ON CONFLICT DO NOTHING",
210
- (thread_id, checkpoint_ns),
211
- )
227
+ await self._seed_root(thread_id, checkpoint_ns)
212
228
  cur = await self._conn.execute(
213
229
  "UPDATE checkpoint_threads SET branches = branches || %s"
214
230
  " WHERE thread_id = %s AND checkpoint_ns = %s"
215
231
  " RETURNING jsonb_array_length(branches)",
216
232
  (Jsonb([entry]), thread_id, checkpoint_ns),
217
233
  )
218
- return (await cur.fetchone())[0]
234
+ return (await cur.fetchone())[0] - 1
219
235
 
220
236
  async def bump_tips(
221
237
  self,
@@ -224,6 +240,7 @@ class PostgresCheckpointIndex:
224
240
  branch_ids: list[int],
225
241
  checkpoint_id: str,
226
242
  ) -> None:
243
+ await self._seed_root(thread_id, checkpoint_ns)
227
244
  cur = await self._conn.execute(
228
245
  "SELECT jsonb_array_length(branches) FROM checkpoint_threads"
229
246
  " WHERE thread_id = %s AND checkpoint_ns = %s",
@@ -233,9 +250,9 @@ class PostgresCheckpointIndex:
233
250
  _check_branch_ids(branch_ids, row[0] if row is not None else 0)
234
251
  await self._conn.execute(
235
252
  "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
236
- " CASE WHEN idx::int = ANY(%s)"
253
+ " CASE WHEN idx::int - 1 = ANY(%s)"
237
254
  " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
238
- " ELSE elem END)"
255
+ " ELSE elem END ORDER BY idx)"
239
256
  " FROM jsonb_array_elements(branches) WITH ORDINALITY AS t(elem, idx))"
240
257
  " WHERE thread_id = %s AND checkpoint_ns = %s",
241
258
  (branch_ids, checkpoint_id, thread_id, checkpoint_ns),
@@ -278,15 +295,31 @@ def _message_dict(message: Any) -> dict[str, Any]:
278
295
  return message
279
296
 
280
297
 
298
+ def _values_snapshot(
299
+ checkpoint_id: str,
300
+ values: dict[str, Any],
301
+ messages_key: str,
302
+ lineage_key: str,
303
+ ) -> Snapshot:
304
+ lineage = values.get(lineage_key)
305
+ if lineage is not None and not isinstance(lineage, dict):
306
+ raise ValueError(f"lineage channel must be an object, got {lineage!r}")
307
+ return Snapshot(
308
+ checkpoint_id=checkpoint_id,
309
+ messages=[_message_dict(m) for m in values.get(messages_key, [])],
310
+ lineage=lineage,
311
+ )
312
+
313
+
281
314
  def langgraph_checkpoints(
282
- saver: Any,
315
+ graph: Any,
283
316
  *,
284
317
  thread_id: str,
285
318
  checkpoint_ns: str = "",
286
319
  messages_key: str = "messages",
287
320
  lineage_key: str = "lineage",
288
321
  ) -> CheckpointStore:
289
- """A CheckpointStore over a langgraph BaseCheckpointSaver's ``aget_tuple``."""
322
+ """A CheckpointStore over a compiled langgraph graph's ``aget_state``."""
290
323
 
291
324
  class _Store:
292
325
  async def get(self, checkpoint_id: str | None) -> Snapshot | None:
@@ -296,22 +329,34 @@ def langgraph_checkpoints(
296
329
  }
297
330
  if checkpoint_id is not None:
298
331
  configurable["checkpoint_id"] = checkpoint_id
299
- saved = await saver.aget_tuple({"configurable": configurable})
300
- if saved is None:
332
+ state = await graph.aget_state({"configurable": configurable})
333
+ if state.metadata is None:
301
334
  return None
302
- values = saved.checkpoint["channel_values"]
303
- lineage = values.get(lineage_key)
304
- if lineage is not None and not isinstance(lineage, dict):
305
- raise ValueError(f"lineage channel must be an object, got {lineage!r}")
306
- return Snapshot(
307
- checkpoint_id=saved.checkpoint["id"],
308
- messages=[_message_dict(m) for m in values.get(messages_key, [])],
309
- lineage=lineage,
335
+ return _values_snapshot(
336
+ state.config["configurable"]["checkpoint_id"],
337
+ state.values,
338
+ messages_key,
339
+ lineage_key,
310
340
  )
311
341
 
312
342
  return _Store()
313
343
 
314
344
 
345
+ def langgraph_snapshot(
346
+ payload: dict[str, Any],
347
+ *,
348
+ messages_key: str = "messages",
349
+ lineage_key: str = "lineage",
350
+ ) -> Snapshot:
351
+ """A Snapshot from one ``stream_mode="checkpoints"`` event payload, for ingesting a run's own output."""
352
+ return _values_snapshot(
353
+ payload["config"]["configurable"]["checkpoint_id"],
354
+ payload["values"],
355
+ messages_key,
356
+ lineage_key,
357
+ )
358
+
359
+
315
360
  def langgraph_switch_values(
316
361
  transcript: list[dict[str, Any]],
317
362
  lineage: dict[str, Any],
@@ -358,7 +403,7 @@ class CheckpointThread:
358
403
  return await self._projection().get_message_child_id(parent_id)
359
404
 
360
405
  def sibling_branches(self, message_id: str) -> list[dict[str, Any]]:
361
- """Branch entries forked at ``message_id``'s parent, as {branch_id, **entry} in allocation order; the transcript in hand is the extra continuation."""
406
+ """Fork entries at ``message_id``'s parent, as {branch_id, **entry} in allocation order; the source branch's continuation has no entry."""
362
407
  index = next(
363
408
  (i for i, m in enumerate(self.transcript) if m["id"] == message_id), None
364
409
  )
@@ -366,15 +411,15 @@ class CheckpointThread:
366
411
  raise KeyError(f"message {message_id!r} is not in the current window")
367
412
  parent_id = self.transcript[index - 1]["id"] if index > 0 else None
368
413
  return [
369
- {"branch_id": i + 1, **entry}
414
+ {"branch_id": i, **entry}
370
415
  for i, entry in enumerate(self.branches)
371
- if entry["fork_message_id"] == parent_id
416
+ if i and entry["fork_message_id"] == parent_id
372
417
  ]
373
418
 
374
419
  def _branch_at(self, index: int) -> int:
375
420
  branch = 0
376
421
  for b in self._path:
377
- fork_id = self.branches[b - 1]["fork_message_id"] if b != 0 else None
422
+ fork_id = self.branches[b]["fork_message_id"] if b != 0 else None
378
423
  start = next(
379
424
  (i + 1 for i, m in enumerate(self.transcript) if m["id"] == fork_id), 0
380
425
  )
@@ -383,11 +428,16 @@ class CheckpointThread:
383
428
  return branch
384
429
 
385
430
  def sibling_info(self, message_id: str) -> tuple[int, list[int]]:
386
- """``message_id``'s own branch id and its sibling fan-out's branch ids in creation order."""
431
+ """``message_id``'s own branch id and its full fan-out's branch ids (source branch included) in creation order."""
387
432
  entries = self.sibling_branches(message_id)
388
433
  index = next(i for i, m in enumerate(self.transcript) if m["id"] == message_id)
389
434
  branch = self._branch_at(index)
390
- return branch, sorted({e["branch_id"] for e in entries} | {branch})
435
+ sources = {e["fork_branch"] for e in entries}
436
+ if len(sources) > 1:
437
+ raise ValueError(
438
+ f"fork entries at one fan-out disagree on fork_branch: {sorted(sources)}"
439
+ )
440
+ return branch, sorted({e["branch_id"] for e in entries} | sources | {branch})
391
441
 
392
442
  async def plan(self, trigger: str, *, rollback_to: str | None = None) -> Plan:
393
443
  """The resumption for ``trigger``; a fork trigger names its last surviving message via ``rollback_to``."""
@@ -427,14 +477,25 @@ class CheckpointThread:
427
477
  }
428
478
  return Plan(anchor, survivors, True)
429
479
 
480
+ async def plan_run(self, ctx: Any) -> Plan:
481
+ """The resumption for a RunManager dispatch: trigger and rollback surface read off the run context."""
482
+ rollback = ctx.rollback_to if ctx.has_rollback else None
483
+ return await self.plan(ctx.trigger, rollback_to=rollback)
484
+
485
+ async def ingest_latest(self) -> None:
486
+ """Ingest the store's latest checkpoint."""
487
+ await self.ingest(await self._get(None))
488
+
430
489
  async def ingest(self, snapshot: Snapshot) -> None:
431
- """Index one new checkpoint: bump the path's tips, upsert the message diff, and append the pending fork's branch entry; idempotent."""
490
+ """Index one new checkpoint: bump the path leaf's tip, upsert the message diff, and append the pending fork's branch entry; idempotent."""
432
491
  pending, self._pending_fork = self._pending_fork, None
433
492
  path = (
434
493
  list(snapshot.lineage["branch_ids"])
435
494
  if snapshot.lineage is not None
436
495
  else list(self._path)
437
496
  )
497
+ if not path:
498
+ raise ValueError("lineage branch_ids must not be empty")
438
499
  anchor = self.checkpoint_id
439
500
  if pending is not None:
440
501
  # the new world descends from the fork message's branch, not the window's leaf
@@ -444,25 +505,29 @@ class CheckpointThread:
444
505
  self._thread_id,
445
506
  self._ns,
446
507
  fork_message_id=pending["fork_message_id"],
508
+ fork_branch=pending["fork_branch"],
447
509
  last_checkpoint_id=snapshot.checkpoint_id,
448
510
  )
511
+ if not self.branches:
512
+ self.branches = [dict(_ROOT_ENTRY)]
449
513
  self.branches = [
450
514
  *self.branches,
451
515
  {
452
516
  "fork_message_id": pending["fork_message_id"],
517
+ "fork_branch": pending["fork_branch"],
453
518
  "last_checkpoint_id": snapshot.checkpoint_id,
454
519
  },
455
520
  ]
456
521
  path = [*path, branch_id]
457
522
  anchor = pending["anchor_checkpoint_id"]
458
- tips = [b for b in path if b != 0]
459
- if tips:
460
- await self._index.bump_tips(
461
- self._thread_id, self._ns, tips, snapshot.checkpoint_id
462
- )
463
- for tip in tips:
464
- if tip <= len(self.branches):
465
- self.branches[tip - 1]["last_checkpoint_id"] = snapshot.checkpoint_id
523
+ tip = path[-1]
524
+ await self._index.bump_tips(
525
+ self._thread_id, self._ns, [tip], snapshot.checkpoint_id
526
+ )
527
+ if not self.branches:
528
+ self.branches = [dict(_ROOT_ENTRY)]
529
+ if tip < len(self.branches):
530
+ self.branches[tip]["last_checkpoint_id"] = snapshot.checkpoint_id
466
531
  known = {m["id"] for m in self.transcript}
467
532
  rows = [
468
533
  {
@@ -470,7 +535,7 @@ class CheckpointThread:
470
535
  "checkpoint_ns": self._ns,
471
536
  "message_id": message["id"],
472
537
  "anchor_checkpoint_id": anchor,
473
- "branch_id": path[-1] if path else 0,
538
+ "branch_id": path[-1],
474
539
  }
475
540
  for message in snapshot.messages
476
541
  if message["id"] not in known
@@ -493,11 +558,12 @@ class CheckpointThread:
493
558
  )
494
559
  return self
495
560
 
496
- def _tip_of(self, branch_id: int) -> str | None:
497
- if branch_id == 0:
498
- return None
561
+ def _tip_of(self, branch_id: int) -> str:
499
562
  _check_branch_ids([branch_id], len(self.branches))
500
- return self.branches[branch_id - 1]["last_checkpoint_id"]
563
+ tip = self.branches[branch_id]["last_checkpoint_id"]
564
+ if tip is None:
565
+ raise LookupError(f"branch {branch_id} has no recorded tip")
566
+ return tip
501
567
 
502
568
  async def _get(self, checkpoint_id: str | None) -> Snapshot:
503
569
  snapshot = await self._store.get(checkpoint_id)
@@ -506,20 +572,14 @@ class CheckpointThread:
506
572
  return snapshot
507
573
 
508
574
  async def switch_branch(self, branch_id: int) -> "CheckpointThread":
509
- """Load the branch's subtree tip as the current window."""
575
+ """Load the branch's own continuation tip as the current window."""
510
576
  return self._load(await self._get(self._tip_of(branch_id)))
511
577
 
512
- async def _world_path(self, branch_id: int) -> list[int]:
578
+ def _world_path(self, branch_id: int) -> list[int]:
513
579
  path = [0]
514
580
  while branch_id != 0:
515
581
  path.insert(1, branch_id)
516
- fork_id = self.branches[branch_id - 1]["fork_message_id"]
517
- if fork_id is None:
518
- break
519
- row = await self._index.get_message(self._thread_id, self._ns, fork_id)
520
- if row is None:
521
- raise KeyError(f"fork message {fork_id!r} is not indexed")
522
- branch_id = row["branch_id"]
582
+ branch_id = self.branches[branch_id]["fork_branch"]
523
583
  return path
524
584
 
525
585
  async def switch_state(
@@ -528,7 +588,7 @@ class CheckpointThread:
528
588
  """Switch like ``switch_branch`` and return (transcript, lineage) for the host to write back into graph state."""
529
589
  await self.switch_branch(branch_id)
530
590
  lineage = self.lineage or {
531
- "branch_ids": await self._world_path(branch_id),
591
+ "branch_ids": self._world_path(branch_id),
532
592
  "parent_checkpoint_id": None,
533
593
  }
534
594
  self._path = list(lineage["branch_ids"])
@@ -547,44 +607,25 @@ class CheckpointThread:
547
607
  raise ValueError("the current window has no tail message")
548
608
  tail = self.transcript[-1]["id"]
549
609
  candidates = [
550
- i + 1
610
+ i
551
611
  for i, entry in enumerate(self.branches)
552
- if entry["fork_message_id"] == tail
612
+ if i and entry["fork_message_id"] == tail
553
613
  ]
554
614
  if not candidates:
555
615
  raise LookupError(f"no world forks at tail message {tail!r}")
556
616
  return self._load(await self._get(self._tip_of(candidates[-1])))
557
617
 
558
618
  async def locate(self, message_id: str) -> "CheckpointThread":
559
- """Load a window containing ``message_id``: its branch's tip, or a fork's tip below it that took the pointer."""
619
+ """Load ``message_id``'s own branch's tip as the current window."""
560
620
  row = await self._index.get_message(self._thread_id, self._ns, message_id)
561
621
  if row is None:
562
622
  raise KeyError(f"message {message_id!r} is not indexed")
563
- queue = [row["branch_id"]]
564
- checked: set[int] = set()
565
- while queue:
566
- branch_id = queue.pop(0)
567
- checked.add(branch_id)
568
- snapshot = await self._get(self._tip_of(branch_id))
569
- if any(m["id"] == message_id for m in snapshot.messages):
570
- return self._load(snapshot)
571
- for i, entry in enumerate(self.branches):
572
- candidate = i + 1
573
- if candidate in checked or candidate in queue:
574
- continue
575
- fork_id = entry["fork_message_id"]
576
- if fork_id is None:
577
- fork_branch = 0
578
- else:
579
- fork_row = await self._index.get_message(
580
- self._thread_id, self._ns, fork_id
581
- )
582
- if fork_row is None:
583
- continue
584
- fork_branch = fork_row["branch_id"]
585
- if fork_branch in checked:
586
- queue.append(candidate)
587
- raise LookupError(f"no branch tip contains message {message_id!r}")
623
+ snapshot = await self._get(self._tip_of(row["branch_id"]))
624
+ if not any(m["id"] == message_id for m in snapshot.messages):
625
+ raise LookupError(
626
+ f"branch {row['branch_id']} tip does not contain message {message_id!r}"
627
+ )
628
+ return self._load(snapshot)
588
629
 
589
630
 
590
631
  async def checkpoint_thread(
@@ -189,11 +189,11 @@ class RunDriver:
189
189
  async def side(self, method: str, params: Any = None, *, client_id: str = "c2"):
190
190
  """A second client's command: its own seq stream, its verdicts riding
191
191
  the other stream."""
192
- await attach(self._app, client_id=client_id)
193
192
  seq = self._side_seqs[client_id] = self._side_seqs.get(client_id, 0) + 1
194
- await self._submit(
195
- {"cmd": [_statement(method, params, seq)]}, client_id=client_id
196
- )
193
+ async with attach(self._app, client_id=client_id):
194
+ await self._submit(
195
+ {"cmd": [_statement(method, params, seq)]}, client_id=client_id
196
+ )
197
197
 
198
198
  async def pump(self, timeout: float = 5) -> None:
199
199
  env = await self._stream.next_event(timeout=timeout)
@@ -30,6 +30,22 @@ def msg(id, type="human"):
30
30
  U1, A1, U2, A2 = msg("u1"), msg("a1", "ai"), msg("u2"), msg("a2", "ai")
31
31
 
32
32
 
33
+ def root(last_checkpoint_id=None):
34
+ return {
35
+ "fork_message_id": None,
36
+ "fork_branch": None,
37
+ "last_checkpoint_id": last_checkpoint_id,
38
+ }
39
+
40
+
41
+ def fork(fork_message_id, fork_branch, last_checkpoint_id):
42
+ return {
43
+ "fork_message_id": fork_message_id,
44
+ "fork_branch": fork_branch,
45
+ "last_checkpoint_id": last_checkpoint_id,
46
+ }
47
+
48
+
33
49
  @dataclass
34
50
  class FakeStore:
35
51
  snapshots: dict[str, Snapshot] = field(default_factory=dict)
@@ -79,7 +95,7 @@ async def test_boot_fresh_thread_is_empty():
79
95
  async def test_boot_loads_latest_snapshot_and_branches():
80
96
  store, index = FakeStore(), MemoryCheckpointIndex()
81
97
  await index.append_branch(
82
- THREAD, "", fork_message_id="u1", last_checkpoint_id="cp2"
98
+ THREAD, "", fork_message_id="u1", fork_branch=0, last_checkpoint_id="cp2"
83
99
  )
84
100
  lineage = {"branch_ids": [1], "parent_checkpoint_id": None}
85
101
  store.put("cp2", [U1, A1], lineage)
@@ -87,9 +103,7 @@ async def test_boot_loads_latest_snapshot_and_branches():
87
103
  assert thread.transcript == [U1, A1]
88
104
  assert thread.checkpoint_id == "cp2"
89
105
  assert thread.lineage == lineage
90
- assert thread.branches == [
91
- {"fork_message_id": "u1", "last_checkpoint_id": "cp2"}
92
- ]
106
+ assert thread.branches == [root(), fork("u1", 0, "cp2")]
93
107
 
94
108
 
95
109
  async def test_boot_rejects_empty_thread_id():
@@ -143,7 +157,7 @@ async def test_ingest_is_idempotent():
143
157
  store, index, thread = await seeded()
144
158
  await thread.ingest(store.snapshots["cp2"])
145
159
  assert await row(index, "u2") == {"anchor_checkpoint_id": "cp1", "branch_id": 0}
146
- assert thread.branches == []
160
+ assert thread.branches == [root("cp2")]
147
161
  assert thread.transcript == [U1, A1, U2, A2]
148
162
 
149
163
 
@@ -153,13 +167,11 @@ async def test_fork_ingest_appends_the_branch_entry_once():
153
167
  a2b = msg("a2b", "ai")
154
168
  snapshot = store.put("cp3", [U1, A1, U2, a2b])
155
169
  await thread.ingest(snapshot)
156
- assert thread.branches == [
157
- {"fork_message_id": "u2", "last_checkpoint_id": "cp3"}
158
- ]
170
+ assert thread.branches == [root("cp2"), fork("u2", 0, "cp3")]
159
171
  assert await index.read_thread(THREAD, "") == thread.branches
160
172
  assert await row(index, "a2b") == {"anchor_checkpoint_id": "cp1", "branch_id": 1}
161
173
  await thread.ingest(snapshot) # redelivery: the pending fork is consumed
162
- assert len(await index.read_thread(THREAD, "")) == 1
174
+ assert len(await index.read_thread(THREAD, "")) == 2
163
175
 
164
176
 
165
177
  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():
169
181
  await thread.ingest(store.put("cp3", [U1, A1, U2, a2b]))
170
182
  u3 = msg("u3")
171
183
  await thread.ingest(store.put("cp4", [U1, A1, U2, a2b, u3]))
172
- assert (await index.read_thread(THREAD, ""))[0]["last_checkpoint_id"] == "cp4"
184
+ assert (await index.read_thread(THREAD, ""))[1]["last_checkpoint_id"] == "cp4"
173
185
  assert await row(index, "u3") == {"anchor_checkpoint_id": "cp3", "branch_id": 1}
174
186
 
175
187
 
@@ -183,8 +195,8 @@ async def test_fork_from_a_shared_message_descends_from_its_branch():
183
195
  u2c, a2c = msg("u2c"), msg("a2c", "ai")
184
196
  await thread.ingest(store.put("cp4", [U1, A1, u2c, a2c]))
185
197
  branches = await index.read_thread(THREAD, "")
186
- assert branches[0]["last_checkpoint_id"] == "cp3" # branch 1 untouched
187
- assert branches[1] == {"fork_message_id": "a1", "last_checkpoint_id": "cp4"}
198
+ assert branches[1]["last_checkpoint_id"] == "cp3" # branch 1 untouched
199
+ assert branches[2] == fork("a1", 0, "cp4")
188
200
  assert await row(index, "u2c") == {"anchor_checkpoint_id": None, "branch_id": 2}
189
201
 
190
202
 
@@ -225,6 +237,60 @@ async def test_plan_rejects_unknown_rollback_and_trigger():
225
237
  await thread.plan("bogus")
226
238
 
227
239
 
240
+ # ─── Run resumption seam ────────────────────────────────────
241
+
242
+
243
+ @dataclass(frozen=True)
244
+ class Ctx:
245
+ trigger: str
246
+ has_rollback: bool = False
247
+ rollback_to: str | None = None
248
+
249
+
250
+ async def test_plan_run_edit_allocates_a_fork_entry():
251
+ store, index, thread = await seeded()
252
+ plan = await thread.plan_run(Ctx("message-edit", True, "a1"))
253
+ assert plan == Plan(None, [U1, A1], True)
254
+ await thread.ingest(store.put("cp3", [U1, A1, msg("u2b"), msg("a2b", "ai")]))
255
+ assert await index.read_thread(THREAD, "") == [
256
+ root("cp2"),
257
+ {"fork_message_id": "a1", "fork_branch": 0, "last_checkpoint_id": "cp3"},
258
+ ]
259
+
260
+
261
+ async def test_plan_run_reload_allocates_a_fork_entry():
262
+ store, index, thread = await seeded()
263
+ plan = await thread.plan_run(Ctx("message-reload", True, "u2"))
264
+ assert plan == Plan("cp1", [U2], True)
265
+ await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
266
+ assert await index.read_thread(THREAD, "") == [
267
+ root("cp2"),
268
+ {"fork_message_id": "u2", "fork_branch": 0, "last_checkpoint_id": "cp3"},
269
+ ]
270
+
271
+
272
+ async def test_plan_run_tip_triggers_allocate_nothing():
273
+ store, index, thread = await seeded()
274
+ for trigger in (
275
+ "message-send",
276
+ "input-resume",
277
+ "error-continue",
278
+ "stop-continue",
279
+ "steer",
280
+ ):
281
+ assert await thread.plan_run(Ctx(trigger)) == Plan("cp2", [], False)
282
+ await thread.ingest(store.put("cp3", [U1, A1, U2, A2, msg("u3")]))
283
+ assert await index.read_thread(THREAD, "") == [root("cp3")]
284
+
285
+
286
+ async def test_ingest_latest_reads_the_store_tip():
287
+ store, _, thread = await seeded()
288
+ store.put("cp3", [U1, A1, U2, A2, msg("u3")])
289
+ await thread.ingest_latest()
290
+ assert thread.checkpoint_id == "cp3"
291
+ assert thread.transcript[-1] == msg("u3")
292
+
293
+
228
294
  # ─── Sibling fan-out ────────────────────────────────────────
229
295
 
230
296
 
@@ -235,8 +301,8 @@ async def test_sibling_branches_order_by_allocation():
235
301
  await thread.plan("message-reload", rollback_to="u2")
236
302
  await thread.ingest(store.put("cp4", [U1, A1, U2, msg("a2c", "ai")]))
237
303
  assert thread.sibling_branches("a2c") == [
238
- {"branch_id": 1, "fork_message_id": "u2", "last_checkpoint_id": "cp3"},
239
- {"branch_id": 2, "fork_message_id": "u2", "last_checkpoint_id": "cp4"},
304
+ {"branch_id": 1, **fork("u2", 0, "cp3")},
305
+ {"branch_id": 2, **fork("u2", 0, "cp4")},
240
306
  ]
241
307
  assert thread.sibling_branches("u1") == []
242
308
  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():
248
314
  for cid, reply in [("cp3", "a2b"), ("cp4", "a2c"), ("cp5", "a2d")]:
249
315
  await thread.plan("message-reload", rollback_to="u2")
250
316
  await thread.ingest(store.put(cid, [U1, A1, U2, msg(reply, "ai")]))
251
- assert thread.sibling_info("a2d") == (3, [1, 2, 3])
317
+ assert thread.sibling_info("a2d") == (3, [0, 1, 2, 3])
252
318
  assert thread.sibling_info("u2") == (0, [0])
253
319
  lineage = {"branch_ids": [0, 1], "parent_checkpoint_id": None}
254
320
  await thread.ingest(store.put("cp6", [U1, A1, U2, msg("a2b", "ai")], lineage))
255
- assert thread.sibling_info("a2b") == (1, [1, 2, 3])
321
+ assert thread.sibling_info("a2b") == (1, [0, 1, 2, 3])
322
+
323
+
324
+ async def test_sibling_info_is_branch_invariant():
325
+ store, _, thread = await seeded()
326
+ await thread.plan("message-reload", rollback_to="u2")
327
+ await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
328
+ from_fork = thread.sibling_info("a2b")[1]
329
+ await thread.switch_state(0)
330
+ from_source = thread.sibling_info("a2")[1]
331
+ assert from_fork == from_source == [0, 1]
256
332
 
257
333
 
258
334
  async def test_sibling_info_counts_the_window_timeline_among_the_forks():
259
335
  store, _, thread = await seeded()
260
336
  await thread.plan("message-edit", rollback_to="a1")
261
337
  await thread.ingest(store.put("cp3", [U1, A1, msg("u2c"), msg("a2c", "ai")]))
262
- assert thread.sibling_info("u2c") == (1, [1])
338
+ assert thread.sibling_info("u2c") == (1, [0, 1])
263
339
  lineage = {"branch_ids": [0], "parent_checkpoint_id": None}
264
340
  await thread.ingest(store.put("cp4", [U1, A1, U2, A2], lineage))
265
341
  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():
269
345
  # ─── Switch branch ──────────────────────────────────────────
270
346
 
271
347
 
272
- async def test_switch_branch_loads_the_subtree_tip():
348
+ async def test_switch_branch_loads_the_branch_tip():
273
349
  store, index, thread = await seeded()
274
350
  await thread.plan("message-reload", rollback_to="u2")
275
351
  a2b = msg("a2b", "ai")
@@ -279,10 +355,14 @@ async def test_switch_branch_loads_the_subtree_tip():
279
355
  assert thread.transcript == [U1, A1, U2, a2b]
280
356
 
281
357
 
282
- async def test_switch_branch_zero_loads_the_thread_latest():
358
+ async def test_switch_branch_round_trips_to_the_source_branch():
283
359
  store, _, thread = await seeded()
360
+ await thread.plan("message-reload", rollback_to="u2")
361
+ await thread.ingest(store.put("cp3", [U1, A1, U2, msg("a2b", "ai")]))
362
+ assert thread.checkpoint_id == "cp3"
284
363
  await thread.switch_branch(0)
285
364
  assert thread.checkpoint_id == "cp2"
365
+ assert thread.transcript == [U1, A1, U2, A2]
286
366
 
287
367
 
288
368
  async def test_switch_branch_unknown_raises():
@@ -329,11 +409,12 @@ async def test_switch_state_ingest_round_trip_keeps_the_index_consistent():
329
409
  await thread.ingest(store.put("cp5", transcript, lineage))
330
410
  assert len(index._messages) == rows_before
331
411
  branches = await index.read_thread(THREAD, "")
332
- assert branches[0]["last_checkpoint_id"] == "cp5"
333
- assert branches[1]["last_checkpoint_id"] == "cp4"
412
+ assert branches[0]["last_checkpoint_id"] == "cp2"
413
+ assert branches[1]["last_checkpoint_id"] == "cp5"
414
+ assert branches[2]["last_checkpoint_id"] == "cp4"
334
415
  assert await row(index, "a2b") == {"anchor_checkpoint_id": "cp1", "branch_id": 1}
335
416
  assert thread.checkpoint_id == "cp5"
336
- assert thread.sibling_info("a2b") == (1, [1, 2])
417
+ assert thread.sibling_info("a2b") == (1, [0, 1, 2])
337
418
 
338
419
 
339
420
  # ─── Paging across a compaction ─────────────────────────────
@@ -342,7 +423,7 @@ async def test_switch_state_ingest_round_trip_keeps_the_index_consistent():
342
423
  async def compacted():
343
424
  store, index, old = await seeded()
344
425
  branch_id = await index.append_branch(
345
- THREAD, "", fork_message_id="a2", last_checkpoint_id="cp3"
426
+ THREAD, "", fork_message_id="a2", fork_branch=0, last_checkpoint_id="cp3"
346
427
  )
347
428
  summary = msg("sum", "ai")
348
429
  store.put(
@@ -365,7 +446,7 @@ async def compacted():
365
446
  async def test_compaction_window_restarts_the_branch_path():
366
447
  _, index, thread, window = await compacted()
367
448
  assert thread.transcript == window
368
- assert (await index.read_thread(THREAD, ""))[0]["last_checkpoint_id"] == "cp4"
449
+ assert (await index.read_thread(THREAD, ""))[1]["last_checkpoint_id"] == "cp4"
369
450
  assert await row(index, "u3") == {"anchor_checkpoint_id": "cp3", "branch_id": 1}
370
451
 
371
452
 
@@ -407,73 +488,108 @@ async def test_locate_unknown_message_raises():
407
488
  await thread.locate("nope")
408
489
 
409
490
 
410
- async def test_locate_follows_a_fork_that_took_the_pointer():
491
+ async def test_locate_loads_the_own_branch_tip_across_forks():
411
492
  store, index, thread = await seeded()
412
493
  await thread.plan("message-reload", rollback_to="u2")
413
494
  a2b = msg("a2b", "ai")
414
495
  await thread.ingest(store.put("cp3", [U1, A1, U2, a2b]))
415
496
  await thread.plan("message-edit", rollback_to="a1")
416
497
  await thread.ingest(store.put("cp4", [U1, A1, msg("u2c"), msg("a2c", "ai")]))
417
- # u2 rides branch 0, whose tip (the thread latest, cp4) dropped it
498
+ # u2 rides branch 0, whose own tip stays cp2 while forks advance the thread
418
499
  await thread.locate("u2")
419
- assert thread.checkpoint_id == "cp3"
500
+ assert thread.checkpoint_id == "cp2"
420
501
  assert any(m["id"] == "u2" for m in thread.transcript)
502
+ await thread.locate("a2b")
503
+ assert thread.checkpoint_id == "cp3"
421
504
 
422
505
 
423
506
  # ─── langgraph adapter ──────────────────────────────────────
424
507
 
425
508
 
426
- async def _saved(saver, checkpoint_id, channel_values):
427
- from langgraph.checkpoint.base import empty_checkpoint
509
+ def _graph(state_type):
510
+ from langgraph.checkpoint.memory import InMemorySaver
511
+ from langgraph.graph import START, StateGraph
512
+
513
+ builder = StateGraph(state_type)
514
+ builder.add_node("node", lambda state: {})
515
+ builder.add_edge(START, "node")
516
+ return builder.compile(checkpointer=InMemorySaver())
517
+
518
+
519
+ def _messages_graph():
520
+ from typing import Annotated, Any, TypedDict
521
+
522
+ from langgraph.graph.message import add_messages
523
+
524
+ class State(TypedDict):
525
+ messages: Annotated[list, add_messages]
526
+ lineage: dict[str, Any] | None
428
527
 
429
- versions = {key: "1" for key in channel_values}
430
- checkpoint = empty_checkpoint()
431
- checkpoint["id"] = checkpoint_id
432
- checkpoint["channel_values"] = channel_values
433
- checkpoint["channel_versions"] = versions
434
- config = {"configurable": {"thread_id": THREAD, "checkpoint_ns": ""}}
435
- await saver.aput(config, checkpoint, {"source": "loop", "step": 1}, versions)
528
+ return _graph(State)
529
+
530
+
531
+ CONFIG = {"configurable": {"thread_id": THREAD}}
436
532
 
437
533
 
438
534
  async def test_langgraph_adapter_reads_latest_and_by_id():
439
- from langchain_core.messages import HumanMessage
440
- from langgraph.checkpoint.memory import InMemorySaver
535
+ pytest.importorskip("langgraph.graph")
536
+ from langchain_core.messages import AIMessage, HumanMessage
441
537
 
442
- saver = InMemorySaver()
538
+ graph = _messages_graph()
539
+ await graph.ainvoke({"messages": [HumanMessage(id="u1", content="hi")]}, CONFIG)
540
+ store = langgraph_checkpoints(graph, thread_id=THREAD)
541
+ first = await store.get(None)
542
+ assert [m["id"] for m in first.messages] == ["u1"]
543
+ assert first.lineage is None
443
544
  lineage = {"branch_ids": [0], "parent_checkpoint_id": None}
444
- await _saved(saver, "1", {"messages": [HumanMessage(id="u1", content="hi")]})
445
- await _saved(
446
- saver,
447
- "2",
448
- {
449
- "messages": [HumanMessage(id="u1", content="hi"), msg("a1", "ai")],
450
- "lineage": lineage,
451
- },
545
+ await graph.ainvoke(
546
+ {"messages": [AIMessage(id="a1", content="yo")], "lineage": lineage}, CONFIG
452
547
  )
453
- store = langgraph_checkpoints(saver, thread_id=THREAD)
454
548
  latest = await store.get(None)
455
- assert latest.checkpoint_id == "2"
549
+ assert latest.checkpoint_id != first.checkpoint_id
456
550
  assert [m["id"] for m in latest.messages] == ["u1", "a1"]
457
551
  assert latest.lineage == lineage
458
- first = await store.get("1")
459
- assert first.checkpoint_id == "1"
460
- assert first.messages[0]["id"] == "u1"
461
- assert first.lineage is None
462
- other = langgraph_checkpoints(saver, thread_id="absent")
463
- assert await other.get(None) is None
552
+ assert await store.get(first.checkpoint_id) == first
553
+ assert await store.get("absent") is None
554
+ assert await langgraph_checkpoints(graph, thread_id="other").get(None) is None
464
555
 
465
556
 
466
557
  async def test_langgraph_adapter_rejects_a_message_without_an_id():
467
- from langchain_core.messages import HumanMessage
468
- from langgraph.checkpoint.memory import InMemorySaver
558
+ pytest.importorskip("langgraph.graph")
559
+ from typing import TypedDict
560
+
561
+ class State(TypedDict):
562
+ messages: list
469
563
 
470
- saver = InMemorySaver()
471
- await _saved(saver, "1", {"messages": [HumanMessage(content="no id")]})
472
- store = langgraph_checkpoints(saver, thread_id=THREAD)
564
+ graph = _graph(State)
565
+ await graph.ainvoke({"messages": [{"type": "human", "content": "no id"}]}, CONFIG)
566
+ store = langgraph_checkpoints(graph, thread_id=THREAD)
473
567
  with pytest.raises(ValueError, match="without an id"):
474
568
  await store.get(None)
475
569
 
476
570
 
571
+ async def test_langgraph_snapshot_ingests_the_run_output():
572
+ pytest.importorskip("langgraph.graph")
573
+ from langchain_core.messages import HumanMessage
574
+
575
+ from harness_sdk import langgraph_snapshot
576
+
577
+ graph = _messages_graph()
578
+ checkpoints = [
579
+ event
580
+ async for namespace, mode, event in graph.astream(
581
+ {"messages": [HumanMessage(id="u1", content="hi")]},
582
+ CONFIG,
583
+ stream_mode=["checkpoints"],
584
+ subgraphs=True,
585
+ )
586
+ if namespace == ()
587
+ ]
588
+ snapshot = langgraph_snapshot(checkpoints[-1])
589
+ store = langgraph_checkpoints(graph, thread_id=THREAD)
590
+ assert snapshot == await store.get(None)
591
+
592
+
477
593
  async def test_langgraph_switch_values_replaces_the_messages_channel():
478
594
  pytest.importorskip("langgraph.graph")
479
595
  from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
@@ -571,19 +687,22 @@ async def pg_index(postgres_dsn):
571
687
  async def test_postgres_branches_round_trip(pg_index):
572
688
  assert await pg_index.read_thread(THREAD, "") == []
573
689
  first = await pg_index.append_branch(
574
- THREAD, "", fork_message_id="u2", last_checkpoint_id="cp1"
690
+ THREAD, "", fork_message_id="u2", fork_branch=0, last_checkpoint_id="cp1"
575
691
  )
576
692
  second = await pg_index.append_branch(
577
- THREAD, "", fork_message_id=None, last_checkpoint_id="cp2"
693
+ THREAD, "", fork_message_id=None, fork_branch=0, last_checkpoint_id="cp2"
578
694
  )
579
695
  assert (first, second) == (1, 2)
580
- await pg_index.bump_tips(THREAD, "", [2], "cp3")
696
+ await pg_index.bump_tips(THREAD, "", [0, 2], "cp3")
581
697
  assert await pg_index.read_thread(THREAD, "") == [
582
- {"fork_message_id": "u2", "last_checkpoint_id": "cp1"},
583
- {"fork_message_id": None, "last_checkpoint_id": "cp3"},
698
+ root("cp3"),
699
+ fork("u2", 0, "cp1"),
700
+ fork(None, 0, "cp3"),
584
701
  ]
585
702
  with pytest.raises(KeyError, match="unknown branch_id"):
586
703
  await pg_index.bump_tips(THREAD, "", [3], "cp4")
704
+ await pg_index.bump_tips(THREAD, "ns2", [0], "cp1")
705
+ assert await pg_index.read_thread(THREAD, "ns2") == [root("cp1")]
587
706
 
588
707
 
589
708
  async def test_postgres_messages_upsert_keeps_the_first_row(pg_index):