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.
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/PKG-INFO +1 -1
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/pyproject.toml +1 -1
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/__init__.py +2 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/checkpoints.py +128 -87
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/run_helpers.py +4 -4
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_checkpoints.py +184 -65
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/.gitignore +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/README.md +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/examples/__init__.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/examples/runs_app.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/fenced_postgres.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/linear_thread.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/src/harness_sdk/run_manager.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_admit_steer.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_applied_visibility.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_batches.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_branch_anchor.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_dispatching.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_durability.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_edit_dispatched.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_edit_reload.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_enqueue.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_fenced_postgres.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_input_required.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_linear_thread.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_meta.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_outcomes.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_placement.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_prepare_hooks.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_restore.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_rewind_during_run.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_run_crash_sweep.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_settle.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_steer.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_stop_continue.py +0 -0
- {harness_sdk_python-0.23.1 → harness_sdk_python-0.25.0}/tests/test_stopping_window.py +0 -0
|
@@ -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
|
|
7
|
-
|
|
8
|
-
the state's ``branch_ids`` path
|
|
9
|
-
|
|
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
|
|
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.
|
|
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
|
|
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.
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
300
|
-
if
|
|
332
|
+
state = await graph.aget_state({"configurable": configurable})
|
|
333
|
+
if state.metadata is None:
|
|
301
334
|
return None
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
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
|
-
"""
|
|
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
|
|
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
|
|
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
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
459
|
-
|
|
460
|
-
|
|
461
|
-
|
|
462
|
-
|
|
463
|
-
|
|
464
|
-
|
|
465
|
-
|
|
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]
|
|
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
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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":
|
|
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
|
|
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
|
|
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
|
-
|
|
564
|
-
|
|
565
|
-
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
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
|
-
|
|
195
|
-
|
|
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, "")) ==
|
|
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, ""))[
|
|
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[
|
|
187
|
-
assert branches[
|
|
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, "
|
|
239
|
-
{"branch_id": 2, "
|
|
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
|
|
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
|
|
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"] == "
|
|
333
|
-
assert branches[1]["last_checkpoint_id"] == "
|
|
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, ""))[
|
|
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
|
|
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
|
|
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 == "
|
|
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
|
-
|
|
427
|
-
from langgraph.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
|
-
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
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
|
-
|
|
440
|
-
from
|
|
535
|
+
pytest.importorskip("langgraph.graph")
|
|
536
|
+
from langchain_core.messages import AIMessage, HumanMessage
|
|
441
537
|
|
|
442
|
-
|
|
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
|
|
445
|
-
|
|
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
|
|
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
|
-
|
|
459
|
-
assert
|
|
460
|
-
assert
|
|
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
|
-
|
|
468
|
-
from
|
|
558
|
+
pytest.importorskip("langgraph.graph")
|
|
559
|
+
from typing import TypedDict
|
|
560
|
+
|
|
561
|
+
class State(TypedDict):
|
|
562
|
+
messages: list
|
|
469
563
|
|
|
470
|
-
|
|
471
|
-
await
|
|
472
|
-
store = langgraph_checkpoints(
|
|
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
|
-
|
|
583
|
-
|
|
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):
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|