harness-sdk-python 0.21.0__tar.gz → 0.23.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 (37) hide show
  1. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/PKG-INFO +1 -1
  2. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/pyproject.toml +1 -1
  3. harness_sdk_python-0.23.0/src/harness_sdk/__init__.py +31 -0
  4. harness_sdk_python-0.23.0/src/harness_sdk/checkpoints.py +620 -0
  5. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/run_manager.py +5 -5
  6. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/run_helpers.py +52 -49
  7. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_admit_steer.py +1 -1
  8. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_applied_visibility.py +3 -2
  9. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_batches.py +1 -1
  10. harness_sdk_python-0.23.0/tests/test_checkpoints.py +607 -0
  11. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_durability.py +1 -1
  12. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_fenced_postgres.py +3 -2
  13. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_run_crash_sweep.py +10 -5
  14. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_stopping_window.py +1 -1
  15. harness_sdk_python-0.21.0/src/harness_sdk/__init__.py +0 -9
  16. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/.gitignore +0 -0
  17. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/README.md +0 -0
  18. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/examples/__init__.py +0 -0
  19. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/examples/runs_app.py +0 -0
  20. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/fenced_postgres.py +0 -0
  21. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/linear_thread.py +0 -0
  22. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_branch_anchor.py +0 -0
  23. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_dispatching.py +0 -0
  24. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_edit_dispatched.py +0 -0
  25. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_edit_reload.py +0 -0
  26. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_enqueue.py +0 -0
  27. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_input_required.py +0 -0
  28. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_linear_thread.py +0 -0
  29. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_meta.py +0 -0
  30. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_outcomes.py +0 -0
  31. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_placement.py +0 -0
  32. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_prepare_hooks.py +0 -0
  33. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_restore.py +0 -0
  34. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_rewind_during_run.py +0 -0
  35. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_settle.py +0 -0
  36. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_steer.py +0 -0
  37. {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_stop_continue.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: harness-sdk-python
3
- Version: 0.21.0
3
+ Version: 0.23.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.21.0"
3
+ version = "0.23.0"
4
4
  description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -0,0 +1,31 @@
1
+ from .checkpoints import (
2
+ CheckpointIndex,
3
+ CheckpointStore,
4
+ CheckpointThread,
5
+ MemoryCheckpointIndex,
6
+ Plan,
7
+ PostgresCheckpointIndex,
8
+ Snapshot,
9
+ checkpoint_thread,
10
+ langgraph_checkpoints,
11
+ langgraph_switch_values,
12
+ )
13
+ from .linear_thread import linear_thread
14
+ from .run_manager import RunManager, RunManagerDescriptor, run_manager
15
+
16
+ __all__ = [
17
+ "CheckpointIndex",
18
+ "CheckpointStore",
19
+ "CheckpointThread",
20
+ "MemoryCheckpointIndex",
21
+ "Plan",
22
+ "PostgresCheckpointIndex",
23
+ "RunManager",
24
+ "RunManagerDescriptor",
25
+ "Snapshot",
26
+ "checkpoint_thread",
27
+ "langgraph_checkpoints",
28
+ "langgraph_switch_values",
29
+ "linear_thread",
30
+ "run_manager",
31
+ ]
@@ -0,0 +1,620 @@
1
+ """Checkpoint subsystem: a sidecar branch/message index over an external
2
+ checkpoint store (``get(checkpoint_id)`` only), designed for langgraph.
3
+
4
+ A branch is a resumption point plus its continuation; edits, regenerations,
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.
11
+ """
12
+
13
+ import asyncio
14
+ from dataclasses import dataclass, field
15
+ from typing import Any, Callable, Protocol
16
+
17
+ from .linear_thread import linear_thread
18
+
19
+ _FORK_TRIGGERS = ("message-edit", "message-reload")
20
+ _TIP_TRIGGERS = (
21
+ "message-send",
22
+ "input-resume",
23
+ "error-continue",
24
+ "stop-continue",
25
+ "steer",
26
+ )
27
+
28
+
29
+ @dataclass(frozen=True)
30
+ class Snapshot:
31
+ """One checkpoint's projection: its id, message list, and lineage channel."""
32
+
33
+ checkpoint_id: str
34
+ messages: list[dict[str, Any]]
35
+ lineage: dict[str, Any] | None = None
36
+
37
+
38
+ @dataclass(frozen=True)
39
+ class Plan:
40
+ """A resumption: the checkpoint to resume from, survivors to re-append, and whether a branch forks."""
41
+
42
+ checkpoint_id: str | None
43
+ replay: list[dict[str, Any]]
44
+ fork: bool
45
+
46
+
47
+ class CheckpointStore(Protocol):
48
+ """The whole read seam over the external framework's checkpoint store."""
49
+
50
+ async def get(self, checkpoint_id: str | None) -> Snapshot | None:
51
+ """The snapshot named by ``checkpoint_id``, the latest for None, None when absent."""
52
+ ...
53
+
54
+
55
+ class CheckpointIndex(Protocol):
56
+ """The sidecar index: branch entries per thread row, one row per message."""
57
+
58
+ async def read_thread(
59
+ self, thread_id: str, checkpoint_ns: str
60
+ ) -> list[dict[str, Any]]: ...
61
+
62
+ async def append_branch(
63
+ self,
64
+ thread_id: str,
65
+ checkpoint_ns: str,
66
+ *,
67
+ fork_message_id: str | None,
68
+ last_checkpoint_id: str,
69
+ ) -> int: ...
70
+
71
+ async def bump_tips(
72
+ self,
73
+ thread_id: str,
74
+ checkpoint_ns: str,
75
+ branch_ids: list[int],
76
+ checkpoint_id: str,
77
+ ) -> None: ...
78
+
79
+ async def upsert_messages(self, rows: list[dict[str, Any]]) -> None: ...
80
+
81
+ async def get_message(
82
+ self, thread_id: str, checkpoint_ns: str, message_id: str
83
+ ) -> dict[str, Any] | None: ...
84
+
85
+
86
+ def _check_branch_ids(branch_ids: list[int], count: int) -> None:
87
+ for branch_id in branch_ids:
88
+ if not isinstance(branch_id, int) or not 1 <= branch_id <= count:
89
+ raise KeyError(f"unknown branch_id {branch_id!r} (index holds {count})")
90
+
91
+
92
+ @dataclass
93
+ class MemoryCheckpointIndex:
94
+ """Dict-backed CheckpointIndex for tests."""
95
+
96
+ _threads: dict[tuple[str, str], list[dict[str, Any]]] = field(default_factory=dict)
97
+ _messages: dict[tuple[str, str, str], dict[str, Any]] = field(default_factory=dict)
98
+
99
+ async def read_thread(
100
+ self, thread_id: str, checkpoint_ns: str
101
+ ) -> list[dict[str, Any]]:
102
+ return [dict(e) for e in self._threads.get((thread_id, checkpoint_ns), [])]
103
+
104
+ async def append_branch(
105
+ self,
106
+ thread_id: str,
107
+ checkpoint_ns: str,
108
+ *,
109
+ fork_message_id: str | None,
110
+ last_checkpoint_id: str,
111
+ ) -> int:
112
+ branches = self._threads.setdefault((thread_id, checkpoint_ns), [])
113
+ branches.append(
114
+ {
115
+ "fork_message_id": fork_message_id,
116
+ "last_checkpoint_id": last_checkpoint_id,
117
+ }
118
+ )
119
+ return len(branches)
120
+
121
+ async def bump_tips(
122
+ self,
123
+ thread_id: str,
124
+ checkpoint_ns: str,
125
+ branch_ids: list[int],
126
+ checkpoint_id: str,
127
+ ) -> None:
128
+ branches = self._threads.get((thread_id, checkpoint_ns), [])
129
+ _check_branch_ids(branch_ids, len(branches))
130
+ for branch_id in branch_ids:
131
+ branches[branch_id - 1]["last_checkpoint_id"] = checkpoint_id
132
+
133
+ async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
134
+ for row in rows:
135
+ key = (row["thread_id"], row["checkpoint_ns"], row["message_id"])
136
+ self._messages.setdefault(
137
+ key,
138
+ {
139
+ "anchor_checkpoint_id": row["anchor_checkpoint_id"],
140
+ "branch_id": row["branch_id"],
141
+ },
142
+ )
143
+
144
+ async def get_message(
145
+ self, thread_id: str, checkpoint_ns: str, message_id: str
146
+ ) -> dict[str, Any] | None:
147
+ row = self._messages.get((thread_id, checkpoint_ns, message_id))
148
+ return dict(row) if row is not None else None
149
+
150
+
151
+ class PostgresCheckpointIndex:
152
+ """psycopg-backed CheckpointIndex; ``setup()`` creates the two tables."""
153
+
154
+ def __init__(self, conn: Any) -> None:
155
+ try:
156
+ import psycopg # noqa: F401
157
+ except ImportError as exc:
158
+ raise ImportError(
159
+ "harness_sdk.checkpoints.PostgresCheckpointIndex requires the "
160
+ "'postgres' extra: pip install 'harness-sdk-python[postgres]'"
161
+ ) from exc
162
+ self._conn = conn
163
+
164
+ async def setup(self) -> None:
165
+ await self._conn.execute(
166
+ "CREATE TABLE IF NOT EXISTS checkpoint_threads ("
167
+ " thread_id TEXT NOT NULL,"
168
+ " checkpoint_ns TEXT NOT NULL,"
169
+ " branches JSONB NOT NULL,"
170
+ " PRIMARY KEY (thread_id, checkpoint_ns))"
171
+ )
172
+ await self._conn.execute(
173
+ "CREATE TABLE IF NOT EXISTS checkpoint_messages ("
174
+ " thread_id TEXT NOT NULL,"
175
+ " checkpoint_ns TEXT NOT NULL,"
176
+ " message_id TEXT NOT NULL,"
177
+ " anchor_checkpoint_id TEXT,"
178
+ " branch_id INTEGER NOT NULL,"
179
+ " PRIMARY KEY (thread_id, checkpoint_ns, message_id))"
180
+ )
181
+
182
+ async def read_thread(
183
+ self, thread_id: str, checkpoint_ns: str
184
+ ) -> list[dict[str, Any]]:
185
+ cur = await self._conn.execute(
186
+ "SELECT branches FROM checkpoint_threads WHERE thread_id = %s AND checkpoint_ns = %s",
187
+ (thread_id, checkpoint_ns),
188
+ )
189
+ row = await cur.fetchone()
190
+ return row[0] if row is not None else []
191
+
192
+ async def append_branch(
193
+ self,
194
+ thread_id: str,
195
+ checkpoint_ns: str,
196
+ *,
197
+ fork_message_id: str | None,
198
+ last_checkpoint_id: str,
199
+ ) -> int:
200
+ from psycopg.types.json import Jsonb
201
+
202
+ entry = {
203
+ "fork_message_id": fork_message_id,
204
+ "last_checkpoint_id": last_checkpoint_id,
205
+ }
206
+ 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
+ )
212
+ cur = await self._conn.execute(
213
+ "UPDATE checkpoint_threads SET branches = branches || %s"
214
+ " WHERE thread_id = %s AND checkpoint_ns = %s"
215
+ " RETURNING jsonb_array_length(branches)",
216
+ (Jsonb([entry]), thread_id, checkpoint_ns),
217
+ )
218
+ return (await cur.fetchone())[0]
219
+
220
+ async def bump_tips(
221
+ self,
222
+ thread_id: str,
223
+ checkpoint_ns: str,
224
+ branch_ids: list[int],
225
+ checkpoint_id: str,
226
+ ) -> None:
227
+ cur = await self._conn.execute(
228
+ "SELECT jsonb_array_length(branches) FROM checkpoint_threads"
229
+ " WHERE thread_id = %s AND checkpoint_ns = %s",
230
+ (thread_id, checkpoint_ns),
231
+ )
232
+ row = await cur.fetchone()
233
+ _check_branch_ids(branch_ids, row[0] if row is not None else 0)
234
+ await self._conn.execute(
235
+ "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
236
+ " CASE WHEN idx::int = ANY(%s)"
237
+ " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
238
+ " ELSE elem END)"
239
+ " FROM jsonb_array_elements(branches) WITH ORDINALITY AS t(elem, idx))"
240
+ " WHERE thread_id = %s AND checkpoint_ns = %s",
241
+ (branch_ids, checkpoint_id, thread_id, checkpoint_ns),
242
+ )
243
+
244
+ async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
245
+ async with self._conn.cursor() as cur:
246
+ await cur.executemany(
247
+ "INSERT INTO checkpoint_messages"
248
+ " (thread_id, checkpoint_ns, message_id, anchor_checkpoint_id, branch_id)"
249
+ " VALUES (%(thread_id)s, %(checkpoint_ns)s, %(message_id)s,"
250
+ " %(anchor_checkpoint_id)s, %(branch_id)s)"
251
+ " ON CONFLICT (thread_id, checkpoint_ns, message_id) DO NOTHING",
252
+ rows,
253
+ )
254
+
255
+ async def get_message(
256
+ self, thread_id: str, checkpoint_ns: str, message_id: str
257
+ ) -> dict[str, Any] | None:
258
+ cur = await self._conn.execute(
259
+ "SELECT anchor_checkpoint_id, branch_id FROM checkpoint_messages"
260
+ " WHERE thread_id = %s AND checkpoint_ns = %s AND message_id = %s",
261
+ (thread_id, checkpoint_ns, message_id),
262
+ )
263
+ row = await cur.fetchone()
264
+ if row is None:
265
+ return None
266
+ return {"anchor_checkpoint_id": row[0], "branch_id": row[1]}
267
+
268
+
269
+ def _message_dict(message: Any) -> dict[str, Any]:
270
+ if not isinstance(message, dict):
271
+ dump = getattr(message, "model_dump", None)
272
+ if dump is None:
273
+ raise TypeError(f"message must be a dict or a BaseMessage, got {message!r}")
274
+ message = dump()
275
+ message_id = message.get("id")
276
+ if not isinstance(message_id, str) or message_id == "":
277
+ raise ValueError(f"message without an id: {message!r}")
278
+ return message
279
+
280
+
281
+ def langgraph_checkpoints(
282
+ saver: Any,
283
+ *,
284
+ thread_id: str,
285
+ checkpoint_ns: str = "",
286
+ messages_key: str = "messages",
287
+ lineage_key: str = "lineage",
288
+ ) -> CheckpointStore:
289
+ """A CheckpointStore over a langgraph BaseCheckpointSaver's ``aget_tuple``."""
290
+
291
+ class _Store:
292
+ async def get(self, checkpoint_id: str | None) -> Snapshot | None:
293
+ configurable: dict[str, Any] = {
294
+ "thread_id": thread_id,
295
+ "checkpoint_ns": checkpoint_ns,
296
+ }
297
+ if checkpoint_id is not None:
298
+ configurable["checkpoint_id"] = checkpoint_id
299
+ saved = await saver.aget_tuple({"configurable": configurable})
300
+ if saved is None:
301
+ 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,
310
+ )
311
+
312
+ return _Store()
313
+
314
+
315
+ def langgraph_switch_values(
316
+ transcript: list[dict[str, Any]],
317
+ lineage: dict[str, Any],
318
+ *,
319
+ messages_key: str = "messages",
320
+ lineage_key: str = "lineage",
321
+ ) -> dict[str, Any]:
322
+ """``update_state`` channel values that replace the messages channel with ``transcript`` and set ``lineage``."""
323
+ from langchain_core.messages import RemoveMessage, convert_to_messages
324
+ from langgraph.graph.message import REMOVE_ALL_MESSAGES
325
+
326
+ return {
327
+ messages_key: [
328
+ RemoveMessage(id=REMOVE_ALL_MESSAGES),
329
+ *convert_to_messages([_message_dict(m) for m in transcript]),
330
+ ],
331
+ lineage_key: lineage,
332
+ }
333
+
334
+
335
+ @dataclass
336
+ class CheckpointThread:
337
+ """One thread's current window over the store, with the sidecar index kept in step."""
338
+
339
+ transcript: list[dict[str, Any]]
340
+ branches: list[dict[str, Any]]
341
+ lineage: dict[str, Any] | None
342
+ checkpoint_id: str | None
343
+ _store: CheckpointStore
344
+ _index: CheckpointIndex
345
+ _thread_id: str
346
+ _ns: str
347
+ _role: Callable[[dict[str, Any]], str]
348
+ _path: list[int] = field(default_factory=lambda: [0])
349
+ _pending_fork: dict[str, Any] | None = None
350
+
351
+ def _projection(self):
352
+ return linear_thread(messages=lambda: self.transcript, role=self._role)
353
+
354
+ async def get_message_meta(self, message_id: str | None) -> dict[str, Any] | None:
355
+ return await self._projection().get_message_meta(message_id)
356
+
357
+ async def get_message_child_id(self, parent_id: str) -> str | None:
358
+ return await self._projection().get_message_child_id(parent_id)
359
+
360
+ 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."""
362
+ index = next(
363
+ (i for i, m in enumerate(self.transcript) if m["id"] == message_id), None
364
+ )
365
+ if index is None:
366
+ raise KeyError(f"message {message_id!r} is not in the current window")
367
+ parent_id = self.transcript[index - 1]["id"] if index > 0 else None
368
+ return [
369
+ {"branch_id": i + 1, **entry}
370
+ for i, entry in enumerate(self.branches)
371
+ if entry["fork_message_id"] == parent_id
372
+ ]
373
+
374
+ def _branch_at(self, index: int) -> int:
375
+ branch = 0
376
+ for b in self._path:
377
+ fork_id = self.branches[b - 1]["fork_message_id"] if b != 0 else None
378
+ start = next(
379
+ (i + 1 for i, m in enumerate(self.transcript) if m["id"] == fork_id), 0
380
+ )
381
+ if start <= index:
382
+ branch = b
383
+ return branch
384
+
385
+ def sibling_info(self, message_id: str) -> tuple[int, int]:
386
+ """1-based (position, count) of ``message_id`` within its sibling fan-out."""
387
+ entries = self.sibling_branches(message_id)
388
+ index = next(i for i, m in enumerate(self.transcript) if m["id"] == message_id)
389
+ branch = self._branch_at(index)
390
+ group = sorted({e["branch_id"] for e in entries} | {branch})
391
+ return group.index(branch) + 1, len(group)
392
+
393
+ async def plan(self, trigger: str, *, rollback_to: str | None = None) -> Plan:
394
+ """The resumption for ``trigger``; a fork trigger names its last surviving message via ``rollback_to``."""
395
+ if trigger in _TIP_TRIGGERS:
396
+ if rollback_to is not None:
397
+ raise ValueError(f"rollback_to is invalid for trigger {trigger!r}")
398
+ return Plan(self.checkpoint_id, [], False)
399
+ if trigger not in _FORK_TRIGGERS:
400
+ raise ValueError(f"unknown trigger {trigger!r}")
401
+ if rollback_to is None:
402
+ self._pending_fork = {
403
+ "fork_message_id": None,
404
+ "anchor_checkpoint_id": None,
405
+ "fork_branch": 0,
406
+ }
407
+ return Plan(None, [], True)
408
+ row = await self._index.get_message(self._thread_id, self._ns, rollback_to)
409
+ if row is None:
410
+ raise KeyError(f"message {rollback_to!r} is not indexed")
411
+ anchor = row["anchor_checkpoint_id"]
412
+ index = next(
413
+ (i for i, m in enumerate(self.transcript) if m["id"] == rollback_to), None
414
+ )
415
+ if index is None:
416
+ raise KeyError(f"message {rollback_to!r} is not in the current window")
417
+ survivors: list[dict[str, Any]] = []
418
+ for message in reversed(self.transcript[: index + 1]):
419
+ meta = await self._index.get_message(self._thread_id, self._ns, message["id"])
420
+ if meta is None or meta["anchor_checkpoint_id"] != anchor:
421
+ break
422
+ survivors.append(message)
423
+ survivors.reverse()
424
+ self._pending_fork = {
425
+ "fork_message_id": rollback_to,
426
+ "anchor_checkpoint_id": anchor,
427
+ "fork_branch": row["branch_id"],
428
+ }
429
+ return Plan(anchor, survivors, True)
430
+
431
+ async def ingest(self, snapshot: Snapshot) -> None:
432
+ """Index one new checkpoint: bump the path's tips, upsert the message diff, and append the pending fork's branch entry; idempotent."""
433
+ pending, self._pending_fork = self._pending_fork, None
434
+ path = (
435
+ list(snapshot.lineage["branch_ids"])
436
+ if snapshot.lineage is not None
437
+ else list(self._path)
438
+ )
439
+ anchor = self.checkpoint_id
440
+ if pending is not None:
441
+ # the new world descends from the fork message's branch, not the window's leaf
442
+ if pending["fork_branch"] in path:
443
+ path = path[: path.index(pending["fork_branch"]) + 1]
444
+ branch_id = await self._index.append_branch(
445
+ self._thread_id,
446
+ self._ns,
447
+ fork_message_id=pending["fork_message_id"],
448
+ last_checkpoint_id=snapshot.checkpoint_id,
449
+ )
450
+ self.branches = [
451
+ *self.branches,
452
+ {
453
+ "fork_message_id": pending["fork_message_id"],
454
+ "last_checkpoint_id": snapshot.checkpoint_id,
455
+ },
456
+ ]
457
+ path = [*path, branch_id]
458
+ anchor = pending["anchor_checkpoint_id"]
459
+ tips = [b for b in path if b != 0]
460
+ if tips:
461
+ await self._index.bump_tips(
462
+ self._thread_id, self._ns, tips, snapshot.checkpoint_id
463
+ )
464
+ for tip in tips:
465
+ if tip <= len(self.branches):
466
+ self.branches[tip - 1]["last_checkpoint_id"] = snapshot.checkpoint_id
467
+ known = {m["id"] for m in self.transcript}
468
+ rows = [
469
+ {
470
+ "thread_id": self._thread_id,
471
+ "checkpoint_ns": self._ns,
472
+ "message_id": message["id"],
473
+ "anchor_checkpoint_id": anchor,
474
+ "branch_id": path[-1] if path else 0,
475
+ }
476
+ for message in snapshot.messages
477
+ if message["id"] not in known
478
+ ]
479
+ if rows:
480
+ await self._index.upsert_messages(rows)
481
+ self.transcript = list(snapshot.messages)
482
+ self.checkpoint_id = snapshot.checkpoint_id
483
+ self.lineage = snapshot.lineage
484
+ self._path = path
485
+
486
+ def _load(self, snapshot: Snapshot) -> "CheckpointThread":
487
+ self.transcript = list(snapshot.messages)
488
+ self.checkpoint_id = snapshot.checkpoint_id
489
+ self.lineage = snapshot.lineage
490
+ self._path = (
491
+ list(snapshot.lineage["branch_ids"])
492
+ if snapshot.lineage is not None
493
+ else [0]
494
+ )
495
+ return self
496
+
497
+ def _tip_of(self, branch_id: int) -> str | None:
498
+ if branch_id == 0:
499
+ return None
500
+ _check_branch_ids([branch_id], len(self.branches))
501
+ return self.branches[branch_id - 1]["last_checkpoint_id"]
502
+
503
+ async def _get(self, checkpoint_id: str | None) -> Snapshot:
504
+ snapshot = await self._store.get(checkpoint_id)
505
+ if snapshot is None:
506
+ raise LookupError(f"checkpoint {checkpoint_id!r} is not in the store")
507
+ return snapshot
508
+
509
+ async def switch_branch(self, branch_id: int) -> "CheckpointThread":
510
+ """Load the branch's subtree tip as the current window."""
511
+ return self._load(await self._get(self._tip_of(branch_id)))
512
+
513
+ async def _world_path(self, branch_id: int) -> list[int]:
514
+ path = [0]
515
+ while branch_id != 0:
516
+ path.insert(1, branch_id)
517
+ fork_id = self.branches[branch_id - 1]["fork_message_id"]
518
+ if fork_id is None:
519
+ break
520
+ row = await self._index.get_message(self._thread_id, self._ns, fork_id)
521
+ if row is None:
522
+ raise KeyError(f"fork message {fork_id!r} is not indexed")
523
+ branch_id = row["branch_id"]
524
+ return path
525
+
526
+ async def switch_state(
527
+ self, branch_id: int
528
+ ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
529
+ """Switch like ``switch_branch`` and return (transcript, lineage) for the host to write back into graph state."""
530
+ await self.switch_branch(branch_id)
531
+ lineage = self.lineage or {
532
+ "branch_ids": await self._world_path(branch_id),
533
+ "parent_checkpoint_id": None,
534
+ }
535
+ self._path = list(lineage["branch_ids"])
536
+ return list(self.transcript), lineage
537
+
538
+ async def page_up(self) -> "CheckpointThread":
539
+ """Load the parent world's tail checkpoint as the current window."""
540
+ parent = (self.lineage or {}).get("parent_checkpoint_id")
541
+ if parent is None:
542
+ raise ValueError("the current window has no parent world")
543
+ return self._load(await self._get(parent))
544
+
545
+ async def page_down(self) -> "CheckpointThread":
546
+ """Load the newest world forked at the current window's tail message."""
547
+ if not self.transcript:
548
+ raise ValueError("the current window has no tail message")
549
+ tail = self.transcript[-1]["id"]
550
+ candidates = [
551
+ i + 1
552
+ for i, entry in enumerate(self.branches)
553
+ if entry["fork_message_id"] == tail
554
+ ]
555
+ if not candidates:
556
+ raise LookupError(f"no world forks at tail message {tail!r}")
557
+ return self._load(await self._get(self._tip_of(candidates[-1])))
558
+
559
+ async def locate(self, message_id: str) -> "CheckpointThread":
560
+ """Load a window containing ``message_id``: its branch's tip, or a fork's tip below it that took the pointer."""
561
+ row = await self._index.get_message(self._thread_id, self._ns, message_id)
562
+ if row is None:
563
+ raise KeyError(f"message {message_id!r} is not indexed")
564
+ queue = [row["branch_id"]]
565
+ checked: set[int] = set()
566
+ while queue:
567
+ branch_id = queue.pop(0)
568
+ checked.add(branch_id)
569
+ snapshot = await self._get(self._tip_of(branch_id))
570
+ if any(m["id"] == message_id for m in snapshot.messages):
571
+ return self._load(snapshot)
572
+ for i, entry in enumerate(self.branches):
573
+ candidate = i + 1
574
+ if candidate in checked or candidate in queue:
575
+ continue
576
+ fork_id = entry["fork_message_id"]
577
+ if fork_id is None:
578
+ fork_branch = 0
579
+ else:
580
+ fork_row = await self._index.get_message(
581
+ self._thread_id, self._ns, fork_id
582
+ )
583
+ if fork_row is None:
584
+ continue
585
+ fork_branch = fork_row["branch_id"]
586
+ if fork_branch in checked:
587
+ queue.append(candidate)
588
+ raise LookupError(f"no branch tip contains message {message_id!r}")
589
+
590
+
591
+ async def checkpoint_thread(
592
+ *,
593
+ store: CheckpointStore,
594
+ index: CheckpointIndex,
595
+ thread_id: str,
596
+ checkpoint_ns: str = "",
597
+ role: Callable[[dict[str, Any]], str],
598
+ ) -> CheckpointThread:
599
+ """Boot a CheckpointThread: the thread row and the latest snapshot load concurrently."""
600
+ if not isinstance(thread_id, str) or not thread_id:
601
+ raise ValueError(f"thread_id must be a non-empty string, got {thread_id!r}")
602
+ if not callable(role):
603
+ raise TypeError("role must be callable")
604
+ branches, snapshot = await asyncio.gather(
605
+ index.read_thread(thread_id, checkpoint_ns), store.get(None)
606
+ )
607
+ thread = CheckpointThread(
608
+ transcript=[],
609
+ branches=branches,
610
+ lineage=None,
611
+ checkpoint_id=None,
612
+ _store=store,
613
+ _index=index,
614
+ _thread_id=thread_id,
615
+ _ns=checkpoint_ns,
616
+ _role=role,
617
+ )
618
+ if snapshot is not None:
619
+ thread._load(snapshot)
620
+ return thread