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.
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/PKG-INFO +1 -1
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/pyproject.toml +1 -1
- harness_sdk_python-0.23.0/src/harness_sdk/__init__.py +31 -0
- harness_sdk_python-0.23.0/src/harness_sdk/checkpoints.py +620 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/run_manager.py +5 -5
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/run_helpers.py +52 -49
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_admit_steer.py +1 -1
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_applied_visibility.py +3 -2
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_batches.py +1 -1
- harness_sdk_python-0.23.0/tests/test_checkpoints.py +607 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_durability.py +1 -1
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_fenced_postgres.py +3 -2
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_run_crash_sweep.py +10 -5
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_stopping_window.py +1 -1
- harness_sdk_python-0.21.0/src/harness_sdk/__init__.py +0 -9
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/.gitignore +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/README.md +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/examples/__init__.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/examples/runs_app.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/fenced_postgres.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/src/harness_sdk/linear_thread.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_branch_anchor.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_dispatching.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_edit_dispatched.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_edit_reload.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_enqueue.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_input_required.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_linear_thread.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_meta.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_outcomes.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_placement.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_prepare_hooks.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_restore.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_rewind_during_run.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_settle.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_steer.py +0 -0
- {harness_sdk_python-0.21.0 → harness_sdk_python-0.23.0}/tests/test_stop_continue.py +0 -0
|
@@ -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
|