diff --git a/CHANGELOG.md b/CHANGELOG.md index 3767e11..0d2cac1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **`BufferedStore`**: in-memory buffer with batched flush over any inner store (`JsonlStore.append_many` for one open/write). Also `open_store("buffered:32:runs.jsonl")`. +- **Recording hot-path speedups**: cache `inspect.signature` per boundary, + dataclass-aware `_json_safe`, `Envelope.model_construct` on LIVE record, + `JsonlStore(keep_open=True)`, and `retain_envelopes=` on `record()` / + session (skip in-memory list when only the store write is needed). +- **BufferedStore durability**: `record()` flushes the store on context exit; + failed flushes restore the in-memory batch instead of dropping it. +- **Message capture**: `_json_safe` keeps full dataclass / duck-typed message + fields (not just `role`/`content`); every messages entry is coerced to a dict. ## [0.3.0] - 2026-07-24 diff --git a/chronicle/api.py b/chronicle/api.py index cbf17f7..df4ee2d 100644 --- a/chronicle/api.py +++ b/chronicle/api.py @@ -27,6 +27,7 @@ def record( build_id: str | None = None, redactors: list[Callable[[str], str]] | None = None, export: str | Path | None = None, + retain_envelopes: bool = True, ) -> Iterator[ChronicleSession]: """Record a run in one block. @@ -44,6 +45,9 @@ def record( When ``CHRONICLE_ENABLED`` is off, this is a no-op: yields a fresh session with no store and does not export. Boundaries inside the block also skip LIVE recording. + + Set ``retain_envelopes=False`` when you only need the store write (skips the + in-session list; ``export_trace`` will be empty). """ session = reset_session() if not is_enabled(): @@ -60,11 +64,21 @@ def record( session.build_id = build_id if redactors is not None: session.redactors = redactors + session.retain_envelopes = retain_envelopes session.begin_trace(trace_id) - yield session + try: + yield session + finally: + # Buffered (and keep-open) stores must flush so a short run or remainder + # batch is not left only in memory when the context exits. + store_obj = session.store + if store_obj is not None: + flush = getattr(store_obj, "flush", None) + if callable(flush): + flush() # Export only on a clean exit, so a crash mid-run doesn't overwrite a fixture # with a partial trace. Call session.export_trace(...) yourself if you need it. - if export is not None: + if export is not None and retain_envelopes: session.export_trace(export) diff --git a/chronicle/boundary.py b/chronicle/boundary.py index 236b502..37238db 100644 --- a/chronicle/boundary.py +++ b/chronicle/boundary.py @@ -16,6 +16,7 @@ from __future__ import annotations +import dataclasses import functools import inspect from collections.abc import Callable, Mapping @@ -121,6 +122,12 @@ def _bind_boundary( ) -> Callable[..., Any]: """Shared LIVE / replay wrapper for ``@boundary`` and ``wrap_llm`` (sync + async).""" + # Resolve once — inspect.signature dominates per-crossing cost otherwise. + try: + cached_sig: inspect.Signature | None = inspect.signature(fn) + except (TypeError, ValueError): + cached_sig = None + if inspect.iscoroutinefunction(fn): @functools.wraps(fn) @@ -135,14 +142,14 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: if session.mode == SessionMode.LIVE: return await _record_call_async( session, fn, boundary_id, kind, args, kwargs, - extract_input, extract_result, extract_metadata, + extract_input, extract_result, extract_metadata, cached_sig, ) invocation_index = session._replay_cursor.get(boundary_id, 0) + 1 if session.replay_plan.should_stub(boundary_id, invocation_index): return session.stub_result(boundary_id, kind) return await _live_cutpoint_call_async( session, fn, boundary_id, kind, args, kwargs, - extract_input, invocation_index, + extract_input, invocation_index, cached_sig, ) return async_wrapper @@ -159,14 +166,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: if session.mode == SessionMode.LIVE: return _record_call( session, fn, boundary_id, kind, args, kwargs, - extract_input, extract_result, extract_metadata, + extract_input, extract_result, extract_metadata, cached_sig, ) invocation_index = session._replay_cursor.get(boundary_id, 0) + 1 if session.replay_plan.should_stub(boundary_id, invocation_index): return session.stub_result(boundary_id, kind) return _live_cutpoint_call( session, fn, boundary_id, kind, args, kwargs, - extract_input, invocation_index, + extract_input, invocation_index, cached_sig, ) return wrapper @@ -197,8 +204,11 @@ def _run_on_leave(session, boundary_id, kind, input_state, entered: bool) -> Non session.on_leave(boundary_id, kind, input_state) -def _record_call(session, fn, boundary_id, kind, args, kwargs, extract_input, extract_result, extract_metadata): - input_state = _capture_input(fn, args, kwargs, extract_input) +def _record_call( + session, fn, boundary_id, kind, args, kwargs, + extract_input, extract_result, extract_metadata, cached_sig=None, +): + input_state = _capture_input(fn, args, kwargs, extract_input, cached_sig) call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs) try: try: @@ -212,8 +222,11 @@ def _record_call(session, fn, boundary_id, kind, args, kwargs, extract_input, ex _run_on_leave(session, boundary_id, kind, input_state, entered) -async def _record_call_async(session, fn, boundary_id, kind, args, kwargs, extract_input, extract_result, extract_metadata): - input_state = _capture_input(fn, args, kwargs, extract_input) +async def _record_call_async( + session, fn, boundary_id, kind, args, kwargs, + extract_input, extract_result, extract_metadata, cached_sig=None, +): + input_state = _capture_input(fn, args, kwargs, extract_input, cached_sig) call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs) try: try: @@ -270,8 +283,10 @@ def _call_metadata(result, kind, extract_metadata): # Cut-point (REPLAY mode, live boundary). No envelope; capture for assertions. # --------------------------------------------------------------------------- # -def _live_cutpoint_call(session, fn, boundary_id, kind, args, kwargs, extract_input, invocation_index): - input_state = _capture_input(fn, args, kwargs, extract_input) +def _live_cutpoint_call( + session, fn, boundary_id, kind, args, kwargs, extract_input, invocation_index, cached_sig=None, +): + input_state = _capture_input(fn, args, kwargs, extract_input, cached_sig) session.capture_live_input(boundary_id, invocation_index, input_state) call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs) try: @@ -286,8 +301,10 @@ def _live_cutpoint_call(session, fn, boundary_id, kind, args, kwargs, extract_in _run_on_leave(session, boundary_id, kind, input_state, entered) -async def _live_cutpoint_call_async(session, fn, boundary_id, kind, args, kwargs, extract_input, invocation_index): - input_state = _capture_input(fn, args, kwargs, extract_input) +async def _live_cutpoint_call_async( + session, fn, boundary_id, kind, args, kwargs, extract_input, invocation_index, cached_sig=None, +): + input_state = _capture_input(fn, args, kwargs, extract_input, cached_sig) session.capture_live_input(boundary_id, invocation_index, input_state) call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs) try: @@ -321,23 +338,26 @@ def _advance_cutpoint(session, boundary_id, invocation_index): _IO_KEYS = ("messages", "system_prompt", "rag_chunks") -def _capture_input(fn, args, kwargs, extract_input) -> InputState: +def _capture_input(fn, args, kwargs, extract_input, cached_sig=None) -> InputState: if extract_input is not None: return extract_input(*args, **kwargs) # The default capture must never break the wrapped call. try: - return _bind_input_state(fn, args, kwargs) + return _bind_input_state(fn, args, kwargs, cached_sig) except Exception: return InputState(messages=[], graph_state={"args": _json_safe(list(args)), "kwargs": _json_safe(dict(kwargs))}) -def _bind_input_state(fn, args, kwargs) -> InputState: - graph_state = _bound_arguments(fn, args, kwargs) +def _bind_input_state(fn, args, kwargs, cached_sig=None) -> InputState: + graph_state = _bound_arguments(fn, args, kwargs, cached_sig) source = _io_source(graph_state) messages = source.get("messages") or [] if not messages and "user_message" in source: messages = [{"role": "user", "content": source["user_message"]}] - return InputState( + # Messages must be dicts for the envelope schema; coerce each non-mapping row. + if messages: + messages = [m if isinstance(m, Mapping) else _json_safe(m) for m in messages] + return InputState.model_construct( messages=messages, system_prompt=source.get("system_prompt"), rag_chunks=_coerce_rag_chunks(source.get("rag_chunks")), @@ -345,14 +365,16 @@ def _bind_input_state(fn, args, kwargs) -> InputState: ) -def _bound_arguments(fn, args, kwargs) -> dict[str, Any]: +def _bound_arguments(fn, args, kwargs, cached_sig=None) -> dict[str, Any]: """Record the call by real parameter names. Skip self/cls, flatten **kwargs. Falls back to positional capture when the callable has no introspectable signature (some builtins / C callables). """ + sig = cached_sig try: - sig = inspect.signature(fn) + if sig is None: + sig = inspect.signature(fn) bound = sig.bind_partial(*args, **kwargs) except (TypeError, ValueError): return {"args": _json_safe(list(args)), "kwargs": _json_safe(dict(kwargs))} @@ -409,9 +431,27 @@ def _json_safe(value: Any, _depth: int = 0) -> Any: return {str(k): _json_safe(v, _depth + 1) for k, v in value.items()} if isinstance(value, (list, tuple, set)): return [_json_safe(v, _depth + 1) for v in value] + # Prefer full structured dumps so tool_calls / name / id are not stripped. + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return { + f.name: _json_safe(getattr(value, f.name), _depth + 1) + for f in dataclasses.fields(value) + } if hasattr(value, "model_dump"): try: return value.model_dump() except Exception: return repr(value) + # Duck-typed chat message: keep every public attribute, not just role/content. + role = getattr(value, "role", None) + content = getattr(value, "content", None) + if isinstance(role, str) and isinstance(content, str): + data = getattr(value, "__dict__", None) + if isinstance(data, dict) and data: + return { + str(k): _json_safe(v, _depth + 1) + for k, v in data.items() + if not str(k).startswith("_") + } + return {"role": role, "content": content} return repr(value) diff --git a/chronicle/envelope/backends.py b/chronicle/envelope/backends.py index f5a40e9..2b3f64e 100644 --- a/chronicle/envelope/backends.py +++ b/chronicle/envelope/backends.py @@ -88,10 +88,20 @@ def _flush_unlocked(self) -> None: self._buf = [] append_many = getattr(self.inner, "append_many", None) if callable(append_many): - append_many(batch) - else: - for envelope in batch: + try: + append_many(batch) + except Exception: + # Nothing committed — put the batch back ahead of any newer appends. + self._buf = batch + self._buf + raise + return + for i, envelope in enumerate(batch): + try: self.inner.append(envelope) + except Exception: + # Keep the failed envelope and everything after it. + self._buf = batch[i:] + self._buf + raise def read_all(self) -> list[Envelope]: self.flush() @@ -247,6 +257,11 @@ def open_store(target: str | Path, **kwargs) -> Store: batch_kwargs = {k: v for k, v in kwargs.items() if k == "batch_size"} inner_kwargs = {k: v for k, v in kwargs.items() if k != "batch_size"} batch_size = int(batch_kwargs.get("batch_size", size_str)) + # Prefer a kept-open JSONL handle under the buffer — one fd for the run. + if "keep_open" not in inner_kwargs and not str(inner).startswith( + ("http://", "https://", "sqlite:///", "buffered:") + ) and not str(inner).endswith((".db", ".sqlite")): + inner_kwargs["keep_open"] = True return BufferedStore(open_store(inner, **inner_kwargs), batch_size=batch_size) if text.startswith(("http://", "https://")): return RemoteStore(text, **kwargs) @@ -254,4 +269,7 @@ def open_store(target: str | Path, **kwargs) -> Store: return SqliteStore(text[len("sqlite:///"):], **kwargs) if text.endswith((".db", ".sqlite")): return SqliteStore(text, **kwargs) - return JsonlStore(text) + keep_open = bool(kwargs.pop("keep_open", False)) + if kwargs: + raise TypeError(f"unexpected open_store kwargs for JsonlStore: {sorted(kwargs)}") + return JsonlStore(text, keep_open=keep_open) diff --git a/chronicle/envelope/store.py b/chronicle/envelope/store.py index 12cfd0f..b0adc6d 100644 --- a/chronicle/envelope/store.py +++ b/chronicle/envelope/store.py @@ -3,33 +3,49 @@ from __future__ import annotations from pathlib import Path +from typing import IO, Text from chronicle.envelope.schema import Envelope class EnvelopeStore: - """Append-only JSONL store for immutable envelope records.""" + """Append-only JSONL store for immutable envelope records. - def __init__(self, path: str | Path) -> None: + Set ``keep_open=True`` to hold one append file handle across writes (and + ``append_many`` flushes). Call ``close()`` when done — ``BufferedStore`` does + this on context exit. + """ + + def __init__(self, path: str | Path, *, keep_open: bool = False) -> None: self.path = Path(path) self.path.parent.mkdir(parents=True, exist_ok=True) if not self.path.exists(): self.path.touch() + self._fp: IO[str] | None = None + if keep_open: + self._fp = open(self.path, "a", encoding="utf-8") def append(self, envelope: Envelope) -> None: - line = envelope.model_dump_json() + line = envelope.model_dump_json() + "\n" + if self._fp is not None: + self._fp.write(line) + return with open(self.path, "a", encoding="utf-8") as f: - f.write(line + "\n") + f.write(line) def append_many(self, envelopes: list[Envelope]) -> None: - """Append a batch in one open/write — used by ``BufferedStore`` flushes.""" + """Append a batch in one write — used by ``BufferedStore`` flushes.""" if not envelopes: return payload = "".join(e.model_dump_json() + "\n" for e in envelopes) + if self._fp is not None: + self._fp.write(payload) + return with open(self.path, "a", encoding="utf-8") as f: f.write(payload) def read_all(self) -> list[Envelope]: + self.flush() envelopes: list[Envelope] = [] with open(self.path, encoding="utf-8") as f: for line in f: @@ -47,6 +63,22 @@ def find_by_envelope_id(self, envelope_id: str) -> Envelope | None: return envelope return None + def flush(self) -> None: + if self._fp is not None: + self._fp.flush() + + def close(self) -> None: + if self._fp is not None: + self._fp.flush() + self._fp.close() + self._fp = None + + def __enter__(self) -> EnvelopeStore: + return self + + def __exit__(self, *exc: object) -> None: + self.close() + @staticmethod def load_fixture(path: str | Path) -> Envelope: return Envelope.from_file(str(path)) diff --git a/chronicle/session.py b/chronicle/session.py index bd546ec..8e1ef96 100644 --- a/chronicle/session.py +++ b/chronicle/session.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from contextvars import ContextVar from dataclasses import dataclass, field +from datetime import datetime, timezone from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Any @@ -71,6 +72,9 @@ class ChronicleSession: # reach a committed fixture. Empty by default; set to default_redactors() or # your own. Signature: (str) -> str. See chronicle.redaction. redactors: list[Callable[[str], str]] = field(default_factory=list) + # When False, envelopes are written to ``store`` only and not kept on the + # session (``export_trace`` will be empty). Cuts memory traffic on hot paths. + retain_envelopes: bool = True _sequence: int = 0 _invocation_counts: dict[str, int] = field(default_factory=dict) @@ -153,23 +157,31 @@ def record_envelope( sequence = self.next_sequence() parent_id = self._last_envelope_id - envelope = Envelope( + # model_construct: fields are produced by Chronicle itself; skip pydantic + # validation on the hot LIVE path. + envelope = Envelope.model_construct( + schema_version="1.0", + envelope_id=str(uuid.uuid4()), trace_id=self.trace_id, node_id=boundary_id, boundary_kind=kind, parent_envelope_id=parent_id, sequence=sequence, invocation_index=invocation_index, - metadata=ContextMetadata( + timestamp=datetime.now(timezone.utc), + metadata=ContextMetadata.model_construct( # Prefer what the call actually used; fall back to the session # default only when the boundary surfaced no real metadata. model_version=model_version or self.model_version, build_id=self.build_id, - sampling_params=sampling_params or SamplingParams(), + sampling_params=sampling_params or SamplingParams.model_construct( + temperature=None, top_p=None, max_tokens=None, seed=None, extra={}, + ), tool_schemas=tool_schemas or [], framework="chronicle.boundary", node_id=boundary_id, trace_id=self.trace_id, + extra={}, ), input_state=input_state, action_result=action_result, @@ -182,7 +194,8 @@ def record_envelope( self._push_envelope(envelope.envelope_id) try: - self._recorded_envelopes.append(envelope) + if self.retain_envelopes: + self._recorded_envelopes.append(envelope) self._last_envelope_id = envelope.envelope_id if self.store is not None: self.store.append(envelope) @@ -290,26 +303,42 @@ def envelope_to_return_value(envelope: Envelope, kind: str) -> Any: def result_to_action_result(result: Any, kind: str) -> ActionResult: if kind == "tool" and isinstance(result, dict): - return ActionResult( + return ActionResult.model_construct( + tool_calls=[], completion=result.get("status", str(result)), + finish_reason=None, + token_usage={}, raw_response=result, + error=None, + error_type=None, ) if kind == "llm" and isinstance(result, dict): tool_calls = [ - ToolCall( + ToolCall.model_construct( id=tc.get("id"), name=tc.get("name", ""), arguments=tc.get("arguments", {}), ) for tc in result.get("tool_calls", []) ] - return ActionResult( + return ActionResult.model_construct( tool_calls=tool_calls, completion=result.get("completion"), finish_reason=result.get("finish_reason"), token_usage=_as_token_usage(result.get("token_usage") or result.get("usage")), + raw_response=None, + error=None, + error_type=None, ) - return ActionResult(completion=str(result), raw_response=result if isinstance(result, dict) else None) + return ActionResult.model_construct( + tool_calls=[], + completion=str(result), + finish_reason=None, + token_usage={}, + raw_response=result if isinstance(result, dict) else None, + error=None, + error_type=None, + ) def _as_token_usage(source: Any) -> dict[str, int]: diff --git a/docs/testbench-bench-baseline.json b/docs/testbench-bench-baseline.json new file mode 100644 index 0000000..48ee1cd --- /dev/null +++ b/docs/testbench-bench-baseline.json @@ -0,0 +1,62 @@ +[ + { + "mode": "off", + "workload": "mas1", + "reps": 560, + "best_ms": 8.71962446251473, + "mean_ms": 9.836970051794328, + "median_ms": 9.594138962506804, + "p95_ms": 10.779918612502115, + "vs_off_x": 1.0 + }, + { + "mode": "on", + "workload": "mas1", + "reps": 560, + "best_ms": 10.395640100000492, + "mean_ms": 11.269004414290457, + "median_ms": 11.038737200010473, + "p95_ms": 11.371199950008304, + "vs_off_x": 1.1922119059933 + }, + { + "mode": "buffered", + "workload": "mas1", + "reps": 560, + "best_ms": 10.096395062510055, + "mean_ms": 14.089789878572187, + "median_ms": 11.694157687497864, + "p95_ms": 17.407402000003458, + "vs_off_x": 1.1578933365667294 + }, + { + "mode": "off", + "workload": "mas2", + "reps": 560, + "best_ms": 0.199022599986165, + "mean_ms": 0.21913233213451477, + "median_ms": 0.21450883748457272, + "p95_ms": 0.23961912499999016, + "vs_off_x": 1.0 + }, + { + "mode": "on", + "workload": "mas2", + "reps": 560, + "best_ms": 3.2445309125023414, + "mean_ms": 4.38107248392693, + "median_ms": 4.0147967374878135, + "p95_ms": 4.734604474992921, + "vs_off_x": 16.302324021130683 + }, + { + "mode": "buffered", + "workload": "mas2", + "reps": 560, + "best_ms": 1.7535519875082173, + "mean_ms": 2.360719767856218, + "median_ms": 2.376095287490898, + "p95_ms": 2.680181137498039, + "vs_off_x": 8.810818407709048 + } +] diff --git a/docs/testbench-bench-optimized.json b/docs/testbench-bench-optimized.json new file mode 100644 index 0000000..2013107 --- /dev/null +++ b/docs/testbench-bench-optimized.json @@ -0,0 +1,62 @@ +[ + { + "mode": "off", + "workload": "mas1", + "reps": 900, + "best_ms": 8.04147713999555, + "mean_ms": 8.746666014441063, + "median_ms": 8.383836030006933, + "p95_ms": 10.076082429986855, + "vs_off_x": 1.0 + }, + { + "mode": "on", + "workload": "mas1", + "reps": 900, + "best_ms": 9.25757988999976, + "mean_ms": 9.615125595557604, + "median_ms": 9.658935989991733, + "p95_ms": 9.869413420001365, + "vs_off_x": 1.1512287766081848 + }, + { + "mode": "buffered", + "workload": "mas1", + "reps": 900, + "best_ms": 8.309829709996848, + "mean_ms": 9.160616334443652, + "median_ms": 8.892162369993457, + "p95_ms": 10.113950100003422, + "vs_off_x": 1.033371054263974 + }, + { + "mode": "off", + "workload": "mas2", + "reps": 900, + "best_ms": 0.20378561999677913, + "mean_ms": 0.20900474111234266, + "median_ms": 0.209779119995801, + "p95_ms": 0.2142730100058543, + "vs_off_x": 1.0 + }, + { + "mode": "on", + "workload": "mas2", + "reps": 900, + "best_ms": 3.1873719199938932, + "mean_ms": 3.432608691113678, + "median_ms": 3.478748710003856, + "p95_ms": 3.5867101499934506, + "vs_off_x": 15.640808807040802 + }, + { + "mode": "buffered", + "workload": "mas2", + "reps": 900, + "best_ms": 1.7306703100075538, + "mean_ms": 1.8545640499986702, + "median_ms": 1.863409790003061, + "p95_ms": 1.9177167800080497, + "vs_off_x": 8.492602716692707 + } +] diff --git a/tests/test_json_safe_messages.py b/tests/test_json_safe_messages.py new file mode 100644 index 0000000..31314cd --- /dev/null +++ b/tests/test_json_safe_messages.py @@ -0,0 +1,66 @@ +"""Input capture / _json_safe keep message metadata intact.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import chronicle +from chronicle import boundary +from chronicle.boundary import _json_safe +from chronicle.session import reset_session + + +@dataclass +class RichMessage: + role: str + content: str + name: str | None = None + tool_calls: list | None = None + + +def test_json_safe_dataclass_keeps_extra_message_fields(): + msg = RichMessage( + role="assistant", + content="ok", + name="planner", + tool_calls=[{"id": "1", "name": "search", "arguments": {}}], + ) + out = _json_safe(msg) + assert out["role"] == "assistant" + assert out["content"] == "ok" + assert out["name"] == "planner" + assert out["tool_calls"][0]["name"] == "search" + + +def test_json_safe_duck_typed_message_keeps_public_attrs(): + msg = SimpleNamespace( + role="assistant", + content="done", + tool_call_id="call_9", + name="tool", + ) + out = _json_safe(msg) + assert out["tool_call_id"] == "call_9" + assert out["name"] == "tool" + + +def test_bind_input_coerces_every_message_not_just_first(): + reset_session() + + @boundary("agent", kind="llm") + def agent(messages): + return {"completion": "ok", "finish_reason": "stop"} + + mixed = [ + {"role": "user", "content": "hi"}, + RichMessage(role="assistant", content="yo", name="bot", tool_calls=[]), + ] + with chronicle.record("t") as session: + agent(mixed) + + recorded = session._recorded_envelopes[-1].input_state.messages + assert recorded[0] == {"role": "user", "content": "hi"} + assert recorded[1]["role"] == "assistant" + assert recorded[1]["name"] == "bot" + assert recorded[1]["tool_calls"] == [] diff --git a/tests/test_stores.py b/tests/test_stores.py index b178c92..cd13eeb 100644 --- a/tests/test_stores.py +++ b/tests/test_stores.py @@ -8,6 +8,8 @@ import threading import warnings +import pytest + import chronicle from chronicle import ( BufferedStore, @@ -77,6 +79,47 @@ def test_buffered_store_batches_jsonl_flush(tmp_path): assert len(store.read_all()) == 4 +def test_record_flushes_buffered_store_on_exit(tmp_path): + """Short runs below batch_size must still land on disk when record() exits.""" + path = tmp_path / "runs.jsonl" + store = BufferedStore(JsonlStore(path), batch_size=32) + + with chronicle.record("t-buf", store=store): + + @boundary("agent", kind="tool") + def do(x): + return {"ok": x} + + do(1) + + assert len(JsonlStore(path).read_all()) == 1 + + +def test_buffered_store_restores_batch_when_append_many_fails(tmp_path): + class BoomStore: + def append(self, envelope): + raise AssertionError("should use append_many") + + def append_many(self, envelopes): + raise OSError("disk full") + + def read_all(self): + return [] + + def find_by_trace_id(self, trace_id): + return [] + + def find_by_envelope_id(self, envelope_id): + return None + + store = BufferedStore(BoomStore(), batch_size=2) + store.append(_env("t", 1)) + with pytest.raises(OSError, match="disk full"): + store.append(_env("t", 2)) # triggers flush + # Failed batch is back in the buffer, not silently dropped. + assert len(store._buf) == 2 + + def test_record_into_sqlite(tmp_path): store = SqliteStore(tmp_path / "runs.db")