From 451e09b8fbf63ca0cca71b8d01ed420e7d05c8c9 Mon Sep 17 00:00:00 2001 From: Susheem Koul <105606472+susheem-k@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:30:45 +0530 Subject: [PATCH 1/5] Add CHRONICLE_ENABLED env flag to disable LIVE recording. Lets agents flip recording on/off per process for with/without Chronicle runs without code changes; replay stays unaffected. Co-authored-by: Cursor --- CHANGELOG.md | 6 ++ README.md | 3 +- chronicle/__init__.py | 2 + chronicle/api.py | 8 +++ chronicle/boundary.py | 18 +++++- chronicle/config.py | 23 +++++++ chronicle/envelope/capture.py | 5 ++ chronicle/session.py | 5 ++ chronicle/wrap.py | 17 +++++- tests/test_enabled.py | 109 ++++++++++++++++++++++++++++++++++ 10 files changed, 190 insertions(+), 6 deletions(-) create mode 100644 chronicle/config.py create mode 100644 tests/test_enabled.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 21332aa..29fe6eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **`CHRONICLE_ENABLED`**: set to `0` / `false` / `off` / `no` to turn off LIVE + recording. `@boundary`, `wrap`, `wrap_llm`, `record()`, and `EnvelopeRecorder` + become passthrough so an agent can be run with and without Chronicle. Replay is + unaffected. Check with `chronicle.is_enabled()`. + ## [0.3.0] - 2026-07-24 ### Added diff --git a/README.md b/README.md index 674ff72..f794b7f 100644 --- a/README.md +++ b/README.md @@ -464,6 +464,7 @@ See `examples/langgraph_demo/agent.py`. | Variable | Purpose | |---|---| +| `CHRONICLE_ENABLED` | Set to `0` / `false` / `off` / `no` to disable LIVE recording (`@boundary`, `wrap`, `record()`, `EnvelopeRecorder` become passthrough). Default on. Replay is unaffected. | | `CHRONICLE_BUILD_ID` | Pin runtime build ID in envelope metadata | | `CHRONICLE_STORE` | Default envelope store path | | `PHOENIX_COLLECTOR_ENDPOINT` | Phoenix OTLP endpoint (default `http://localhost:4317`) | @@ -578,7 +579,7 @@ during the live run, so tests are deterministic, free, and fast. Only the option `@boundary` and `wrap()` are **transparent**: they never change what your function returns or raises. Recording adds one wrapper call and one JSON append per boundary crossing, so the cost scales with how many boundaries you mark, not with anything in a -hot loop. Enable it where you want a record; leave it off elsewhere. +hot loop. Set `CHRONICLE_ENABLED=0` to make LIVE recording a no-op for A/B runs.
diff --git a/chronicle/__init__.py b/chronicle/__init__.py index 38899c0..37160f2 100644 --- a/chronicle/__init__.py +++ b/chronicle/__init__.py @@ -2,6 +2,7 @@ from chronicle.api import record, replay_trace from chronicle.boundary import boundary, wrap_llm +from chronicle.config import is_enabled from chronicle.envelope.schema import ( ActionResult, ContextMetadata, @@ -65,6 +66,7 @@ def __getattr__(name: str): "get_session", "instrument_langgraph", "instrument_otel", + "is_enabled", "open_store", "record", "redact_secrets", diff --git a/chronicle/api.py b/chronicle/api.py index 3a3b35d..cbf17f7 100644 --- a/chronicle/api.py +++ b/chronicle/api.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from pathlib import Path +from chronicle.config import is_enabled from chronicle.envelope.backends import Store, open_store from chronicle.replay.plan import ReplayPlan from chronicle.session import ChronicleSession, reset_session @@ -39,8 +40,15 @@ def record( export="fixtures/traces/incident-001/", ) as session: run_agent(...) + + 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. """ session = reset_session() + if not is_enabled(): + yield session + return if store is not None: # A Store instance (has append) is used directly; a path/URL string is routed # through open_store, so store="sqlite:///runs.db" or an http control-plane URL diff --git a/chronicle/boundary.py b/chronicle/boundary.py index 86f5c53..236b502 100644 --- a/chronicle/boundary.py +++ b/chronicle/boundary.py @@ -21,11 +21,13 @@ from collections.abc import Callable, Mapping from typing import Any, TypeVar +from chronicle.config import is_enabled from chronicle.envelope.schema import ActionResult, InputState from chronicle.session import ( SessionMode, get_session, model_version_from, + peek_session, result_to_action_result, sampling_params_from, ) @@ -123,7 +125,13 @@ def _bind_boundary( @functools.wraps(fn) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - session = get_session() + if not is_enabled(): + # Passthrough for LIVE; still honor an active REPLAY session. + session = peek_session() + if session is None or session.mode == SessionMode.LIVE: + return await fn(*args, **kwargs) + else: + session = get_session() if session.mode == SessionMode.LIVE: return await _record_call_async( session, fn, boundary_id, kind, args, kwargs, @@ -141,7 +149,13 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @functools.wraps(fn) def wrapper(*args: Any, **kwargs: Any) -> Any: - session = get_session() + if not is_enabled(): + # Passthrough for LIVE; still honor an active REPLAY session. + session = peek_session() + if session is None or session.mode == SessionMode.LIVE: + return fn(*args, **kwargs) + else: + session = get_session() if session.mode == SessionMode.LIVE: return _record_call( session, fn, boundary_id, kind, args, kwargs, diff --git a/chronicle/config.py b/chronicle/config.py new file mode 100644 index 0000000..61092b7 --- /dev/null +++ b/chronicle/config.py @@ -0,0 +1,23 @@ +"""Runtime configuration from the environment.""" + +from __future__ import annotations + +import os + +# Explicit falsy tokens. Unset means enabled (backward compatible). +_DISABLED = frozenset({"0", "false", "no", "off"}) + + +def is_enabled() -> bool: + """Whether Chronicle LIVE recording / instrumentation is active. + + Controlled by ``CHRONICLE_ENABLED`` (default on). Set to ``0``, ``false``, + ``off``, or ``no`` to make ``@boundary``, ``wrap``, ``wrap_llm``, + ``record()``, and ``EnvelopeRecorder`` no-ops for live runs so an agent can + be timed with and without Chronicle. Replay is unaffected so cut-point + fixtures keep working. + """ + raw = os.environ.get("CHRONICLE_ENABLED") + if raw is None: + return True + return raw.strip().lower() not in _DISABLED diff --git a/chronicle/envelope/capture.py b/chronicle/envelope/capture.py index 135876c..c1b4425 100644 --- a/chronicle/envelope/capture.py +++ b/chronicle/envelope/capture.py @@ -9,6 +9,7 @@ from collections.abc import Callable from typing import Any, ParamSpec, TypeVar +from chronicle.config import is_enabled from chronicle.envelope.schema import ( ActionResult, ContextMetadata, @@ -141,6 +142,8 @@ def _on_error(input_state, exc): if inspect.iscoroutinefunction(fn): @functools.wraps(fn) async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + if not is_enabled(): + return await fn(*args, **kwargs) state, input_state = _prepare(args, kwargs) try: result = await fn(*args, **kwargs) @@ -154,6 +157,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R: @functools.wraps(fn) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + if not is_enabled(): + return fn(*args, **kwargs) state, input_state = _prepare(args, kwargs) try: result = fn(*args, **kwargs) diff --git a/chronicle/session.py b/chronicle/session.py index 831c025..bd546ec 100644 --- a/chronicle/session.py +++ b/chronicle/session.py @@ -256,6 +256,11 @@ def get_session() -> ChronicleSession: return session +def peek_session() -> ChronicleSession | None: + """Return the context session if one exists, without creating one.""" + return _session.get() + + def reset_session() -> ChronicleSession: session = ChronicleSession() _session.set(session) diff --git a/chronicle/wrap.py b/chronicle/wrap.py index 2bd0fda..fd45d7e 100644 --- a/chronicle/wrap.py +++ b/chronicle/wrap.py @@ -18,8 +18,9 @@ from typing import Any from chronicle.boundary import boundary +from chronicle.config import is_enabled from chronicle.envelope.schema import ActionResult, InputState -from chronicle.session import SessionMode, get_session, sampling_params_from +from chronicle.session import SessionMode, get_session, peek_session, sampling_params_from def instrument_langgraph(nodes: Mapping[str, Callable], *, kind: str = "custom") -> dict[str, Callable]: @@ -72,7 +73,12 @@ def _wrap_completion(create: Callable, boundary_id: str) -> Callable: if inspect.iscoroutinefunction(create): @functools.wraps(create) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - session = get_session() + if not is_enabled(): + session = peek_session() + if session is None or session.mode is SessionMode.LIVE: + return await create(*args, **kwargs) + else: + session = get_session() input_state = _input_state(kwargs) if session.mode is SessionMode.REPLAY and _should_stub(session, boundary_id): return _stub(session, boundary_id) @@ -84,7 +90,12 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @functools.wraps(create) def wrapper(*args: Any, **kwargs: Any) -> Any: - session = get_session() + if not is_enabled(): + session = peek_session() + if session is None or session.mode is SessionMode.LIVE: + return create(*args, **kwargs) + else: + session = get_session() input_state = _input_state(kwargs) if session.mode is SessionMode.REPLAY and _should_stub(session, boundary_id): return _stub(session, boundary_id) diff --git a/tests/test_enabled.py b/tests/test_enabled.py new file mode 100644 index 0000000..2f62411 --- /dev/null +++ b/tests/test_enabled.py @@ -0,0 +1,109 @@ +"""CHRONICLE_ENABLED kill switch for LIVE recording.""" + +from __future__ import annotations + +import chronicle +from chronicle import ReplayPlan, boundary +from chronicle.envelope.capture import EnvelopeRecorder +from chronicle.envelope.store import EnvelopeStore +from chronicle.session import reset_session + + +def test_is_enabled_defaults_on(monkeypatch): + monkeypatch.delenv("CHRONICLE_ENABLED", raising=False) + assert chronicle.is_enabled() is True + + +def test_is_enabled_falsy_tokens(monkeypatch): + for value in ("0", "false", "FALSE", "off", "no", " No "): + monkeypatch.setenv("CHRONICLE_ENABLED", value) + assert chronicle.is_enabled() is False, value + + +def test_is_enabled_truthy_tokens(monkeypatch): + for value in ("1", "true", "yes", "on"): + monkeypatch.setenv("CHRONICLE_ENABLED", value) + assert chronicle.is_enabled() is True, value + + +def test_boundary_passthrough_when_disabled(monkeypatch): + monkeypatch.setenv("CHRONICLE_ENABLED", "0") + reset_session() + + @boundary("tool", kind="tool") + def tool(x: str) -> dict: + return {"value": x} + + with chronicle.record("t", store="unused.jsonl") as session: + result = tool("hi") + + assert result == {"value": "hi"} + assert session._recorded_envelopes == [] + assert session.store is None + + +def test_boundary_records_when_enabled(monkeypatch, tmp_path): + monkeypatch.setenv("CHRONICLE_ENABLED", "1") + reset_session() + + @boundary("tool", kind="tool") + def tool(x: str) -> dict: + return {"value": x} + + with chronicle.record("t", store=str(tmp_path / "runs.jsonl")) as session: + tool("hi") + + assert len(session._recorded_envelopes) == 1 + + +def test_replay_still_works_when_disabled(monkeypatch, tmp_path): + """Fixtures and cut-point replay keep working even if CHRONICLE_ENABLED=0.""" + monkeypatch.setenv("CHRONICLE_ENABLED", "1") + + @boundary("tool", kind="tool") + def tool(x: str) -> dict: + return {"value": x} + + trace_dir = tmp_path / "trace" + with chronicle.record("t", export=str(trace_dir)): + tool("recorded") + + monkeypatch.setenv("CHRONICLE_ENABLED", "0") + with chronicle.replay_trace(str(trace_dir), ReplayPlan().stub("tool", 1)): + assert tool("ignored") == {"value": "recorded"} + + +def test_wrap_passthrough_when_disabled(monkeypatch): + monkeypatch.setenv("CHRONICLE_ENABLED", "0") + reset_session() + + class Completions: + def create(self, **kwargs): + return {"choices": [{"message": {"content": "ok"}}], "model": "m"} + + class Chat: + completions = Completions() + + class Client: + chat = Chat() + + client = chronicle.wrap(Client()) + with chronicle.record("t") as session: + resp = client.chat.completions.create(model="m", messages=[]) + + assert resp["choices"][0]["message"]["content"] == "ok" + assert session._recorded_envelopes == [] + + +def test_envelope_recorder_passthrough_when_disabled(monkeypatch, tmp_path): + monkeypatch.setenv("CHRONICLE_ENABLED", "0") + store = EnvelopeStore(tmp_path / "runs.jsonl") + recorder = EnvelopeRecorder(store=store, model_version="m") + + @recorder.wrap_node("agent") + def agent(state: dict) -> dict: + return {**state, "completion": "done"} + + result = agent({"messages": []}) + assert result["completion"] == "done" + assert store.read_all() == [] From fef8dfea1b3146087e611ae2ba7da5613837880e Mon Sep 17 00:00:00 2001 From: Susheem Koul <105606472+susheem-k@users.noreply.github.com> Date: Tue, 4 Aug 2026 23:33:59 +0530 Subject: [PATCH 2/5] Add BufferedStore with batched JSONL flush. Buffers envelopes in memory and flushes in batches (JsonlStore.append_many) so recording can avoid a sync open/write per crossing. Co-authored-by: Cursor --- CHANGELOG.md | 3 ++ chronicle/__init__.py | 2 + chronicle/envelope/__init__.py | 2 + chronicle/envelope/backends.py | 85 ++++++++++++++++++++++++++++++++++ chronicle/envelope/store.py | 8 ++++ tests/test_stores.py | 18 +++++++ 6 files changed, 118 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 29fe6eb..3767e11 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 recording. `@boundary`, `wrap`, `wrap_llm`, `record()`, and `EnvelopeRecorder` become passthrough so an agent can be run with and without Chronicle. Replay is unaffected. Check with `chronicle.is_enabled()`. +- **`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")`. ## [0.3.0] - 2026-07-24 diff --git a/chronicle/__init__.py b/chronicle/__init__.py index 37160f2..e18a93d 100644 --- a/chronicle/__init__.py +++ b/chronicle/__init__.py @@ -14,6 +14,7 @@ ToolSchema, ) from chronicle.envelope.backends import ( + BufferedStore, JsonlStore, RemoteStore, SqliteStore, @@ -43,6 +44,7 @@ def __getattr__(name: str): __all__ = [ "ActionResult", "BoundaryMode", + "BufferedStore", "ChronicleSession", "ContextMetadata", "Envelope", diff --git a/chronicle/envelope/__init__.py b/chronicle/envelope/__init__.py index 89db763..62b9722 100644 --- a/chronicle/envelope/__init__.py +++ b/chronicle/envelope/__init__.py @@ -1,4 +1,5 @@ from chronicle.envelope.backends import ( + BufferedStore, JsonlStore, RemoteStore, SqliteStore, @@ -10,6 +11,7 @@ from chronicle.envelope.store import EnvelopeStore __all__ = [ + "BufferedStore", "Envelope", "EnvelopeRecorder", "EnvelopeStore", diff --git a/chronicle/envelope/backends.py b/chronicle/envelope/backends.py index 75046f4..f5a40e9 100644 --- a/chronicle/envelope/backends.py +++ b/chronicle/envelope/backends.py @@ -5,6 +5,8 @@ - ``JsonlStore`` (the default ``EnvelopeStore``): append-only JSONL on local disk. Zero config, perfect for local development and CI fixtures. +- ``BufferedStore``: in-memory buffer + batched flush over any inner store. Cuts + per-crossing disk/network cost; use ``buffered:32:runs.jsonl`` via ``open_store``. - ``SqliteStore``: durable, queryable SQLite. Zero dependency (stdlib ``sqlite3``). A good fit for a single deployed agent instance. - ``RemoteStore``: ships envelopes to a Chronicle control plane over HTTP. Point many @@ -47,6 +49,75 @@ def find_by_trace_id(self, trace_id: str) -> list[Envelope]: ... def find_by_envelope_id(self, envelope_id: str) -> Envelope | None: ... +class BufferedStore: + """In-memory buffer in front of any ``Store``, with batched flush to disk/network. + + ``append`` is cheap (list append under a lock). When the buffer reaches + ``batch_size``, or when ``flush()`` / context-exit is called, envelopes are + written to the inner store in one batch. Prefer an inner store that implements + ``append_many`` (``JsonlStore`` does) so a flush is a single open/write. + """ + + def __init__( + self, + inner: Store, + *, + batch_size: int = 32, + ) -> None: + if batch_size < 1: + raise ValueError("batch_size must be >= 1") + self.inner = inner + self.batch_size = batch_size + self._buf: list[Envelope] = [] + self._lock = threading.Lock() + + def append(self, envelope: Envelope) -> None: + with self._lock: + self._buf.append(envelope) + if len(self._buf) >= self.batch_size: + self._flush_unlocked() + + def flush(self) -> None: + with self._lock: + self._flush_unlocked() + + def _flush_unlocked(self) -> None: + if not self._buf: + return + batch = self._buf + self._buf = [] + append_many = getattr(self.inner, "append_many", None) + if callable(append_many): + append_many(batch) + else: + for envelope in batch: + self.inner.append(envelope) + + def read_all(self) -> list[Envelope]: + self.flush() + return self.inner.read_all() + + def find_by_trace_id(self, trace_id: str) -> list[Envelope]: + self.flush() + return self.inner.find_by_trace_id(trace_id) + + def find_by_envelope_id(self, envelope_id: str) -> Envelope | None: + self.flush() + return self.inner.find_by_envelope_id(envelope_id) + + def __enter__(self) -> BufferedStore: + return self + + def __exit__(self, *exc: object) -> None: + self.flush() + + def close(self) -> None: + self.flush() + close = getattr(self.inner, "close", None) + if callable(close): + close() + + class SqliteStore: """Append-only envelope store backed by SQLite (stdlib, zero dependency). @@ -160,9 +231,23 @@ def open_store(target: str | Path, **kwargs) -> Store: - ``http(s)://...`` -> RemoteStore (control plane), accepts api_key/timeout - ``sqlite:///path`` or ``*.db`` / ``*.sqlite`` -> SqliteStore + - ``buffered:N:inner`` -> BufferedStore(batch_size=N) over open_store(inner) + e.g. ``buffered:32:runs.jsonl`` or ``buffered:64:sqlite:///runs.db`` - anything else -> JsonlStore (local file, the default) """ text = str(target) + if text.startswith("buffered:"): + # buffered:: + rest = text[len("buffered:") :] + size_str, _, inner = rest.partition(":") + if not size_str or not inner: + raise ValueError( + "buffered store target must look like 'buffered:32:runs.jsonl'" + ) + 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)) + return BufferedStore(open_store(inner, **inner_kwargs), batch_size=batch_size) if text.startswith(("http://", "https://")): return RemoteStore(text, **kwargs) if text.startswith("sqlite:///"): diff --git a/chronicle/envelope/store.py b/chronicle/envelope/store.py index d41144f..12cfd0f 100644 --- a/chronicle/envelope/store.py +++ b/chronicle/envelope/store.py @@ -21,6 +21,14 @@ def append(self, envelope: Envelope) -> None: with open(self.path, "a", encoding="utf-8") as f: f.write(line + "\n") + def append_many(self, envelopes: list[Envelope]) -> None: + """Append a batch in one open/write — used by ``BufferedStore`` flushes.""" + if not envelopes: + return + payload = "".join(e.model_dump_json() + "\n" for e in envelopes) + with open(self.path, "a", encoding="utf-8") as f: + f.write(payload) + def read_all(self) -> list[Envelope]: envelopes: list[Envelope] = [] with open(self.path, encoding="utf-8") as f: diff --git a/tests/test_stores.py b/tests/test_stores.py index 71f4efb..b178c92 100644 --- a/tests/test_stores.py +++ b/tests/test_stores.py @@ -10,6 +10,7 @@ import chronicle from chronicle import ( + BufferedStore, EnvelopeStore, JsonlStore, RemoteStore, @@ -50,6 +51,7 @@ def test_backends_satisfy_store_protocol(tmp_path): assert isinstance(SqliteStore(":memory:"), Store) assert isinstance(JsonlStore(tmp_path / "r.jsonl"), Store) assert isinstance(RemoteStore("http://localhost:1"), Store) + assert isinstance(BufferedStore(JsonlStore(tmp_path / "b.jsonl")), Store) def test_open_store_dispatch(tmp_path): @@ -57,6 +59,22 @@ def test_open_store_dispatch(tmp_path): assert isinstance(open_store(str(tmp_path / "r.db")), SqliteStore) assert isinstance(open_store("sqlite:///" + str(tmp_path / "x.db")), SqliteStore) assert isinstance(open_store("https://cp.example"), RemoteStore) + buffered = open_store(f"buffered:8:{tmp_path / 'buf.jsonl'}") + assert isinstance(buffered, BufferedStore) + assert buffered.batch_size == 8 + + +def test_buffered_store_batches_jsonl_flush(tmp_path): + path = tmp_path / "runs.jsonl" + store = BufferedStore(JsonlStore(path), batch_size=3) + store.append(_env("t", 1)) + store.append(_env("t", 2)) + assert path.read_text(encoding="utf-8").strip() == "" # still buffered + store.append(_env("t", 3)) # hits batch_size -> flush + assert len(JsonlStore(path).read_all()) == 3 + store.append(_env("t", 4)) + store.flush() + assert len(store.read_all()) == 4 def test_record_into_sqlite(tmp_path): From dd529a1b716e6e2e4d267af0da1b2c2ccbecb408 Mon Sep 17 00:00:00 2001 From: Susheem Koul <105606472+susheem-k@users.noreply.github.com> Date: Tue, 4 Aug 2026 23:39:10 +0530 Subject: [PATCH 3/5] Speed up LIVE recording hot path for closer-to-baseline benches. Cache boundary signatures, use model_construct on envelope build, add JsonlStore keep_open and retain_envelopes=False for store-only buffered recording. Co-authored-by: Cursor --- CHANGELOG.md | 4 ++ chronicle/api.py | 7 ++- chronicle/boundary.py | 70 +++++++++++++++++++++-------- chronicle/envelope/backends.py | 10 ++++- chronicle/envelope/store.py | 42 ++++++++++++++--- chronicle/session.py | 45 +++++++++++++++---- docs/testbench-bench-baseline.json | 62 +++++++++++++++++++++++++ docs/testbench-bench-optimized.json | 62 +++++++++++++++++++++++++ 8 files changed, 268 insertions(+), 34 deletions(-) create mode 100644 docs/testbench-bench-baseline.json create mode 100644 docs/testbench-bench-optimized.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 3767e11..d91afa5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,10 @@ 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). ## [0.3.0] - 2026-07-24 diff --git a/chronicle/api.py b/chronicle/api.py index cbf17f7..035845f 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,12 @@ 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 # 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..b72184e 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 dataclass rows. + if messages and not isinstance(messages[0], Mapping): + messages = [_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,6 +431,16 @@ 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] + # Fast path for chat-message shaped dataclasses (role/content) used by most agents. + role = getattr(value, "role", None) + content = getattr(value, "content", None) + if isinstance(role, str) and isinstance(content, str): + return {"role": role, "content": content} + 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() diff --git a/chronicle/envelope/backends.py b/chronicle/envelope/backends.py index f5a40e9..5bda7d8 100644 --- a/chronicle/envelope/backends.py +++ b/chronicle/envelope/backends.py @@ -247,6 +247,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 +259,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 + } +] From 879a5d67417b73f270bb50ad1fd2c9bc60bfdd95 Mon Sep 17 00:00:00 2001 From: Susheem Koul <105606472+susheem-k@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:38:02 +0530 Subject: [PATCH 4/5] Fix BufferedStore flush/durability and message capture gaps. Flush stores on record() exit, restore batches after failed writes, and keep full message metadata when coercing inputs to JSON-safe dicts. Co-authored-by: Cursor --- CHANGELOG.md | 4 ++ chronicle/api.py | 11 +++++- chronicle/boundary.py | 24 ++++++++---- chronicle/envelope/backends.py | 16 ++++++-- tests/test_json_safe_messages.py | 66 ++++++++++++++++++++++++++++++++ tests/test_stores.py | 43 +++++++++++++++++++++ 6 files changed, 152 insertions(+), 12 deletions(-) create mode 100644 tests/test_json_safe_messages.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d91afa5..0d2cac1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 035845f..df4ee2d 100644 --- a/chronicle/api.py +++ b/chronicle/api.py @@ -66,7 +66,16 @@ def record( 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 and retain_envelopes: diff --git a/chronicle/boundary.py b/chronicle/boundary.py index b72184e..37238db 100644 --- a/chronicle/boundary.py +++ b/chronicle/boundary.py @@ -354,9 +354,9 @@ def _bind_input_state(fn, args, kwargs, cached_sig=None) -> InputState: messages = source.get("messages") or [] if not messages and "user_message" in source: messages = [{"role": "user", "content": source["user_message"]}] - # Messages must be dicts for the envelope schema; coerce dataclass rows. - if messages and not isinstance(messages[0], Mapping): - messages = [_json_safe(m) for m in messages] + # 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"), @@ -431,11 +431,7 @@ 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] - # Fast path for chat-message shaped dataclasses (role/content) used by most agents. - role = getattr(value, "role", None) - content = getattr(value, "content", None) - if isinstance(role, str) and isinstance(content, str): - return {"role": role, "content": content} + # 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) @@ -446,4 +442,16 @@ def _json_safe(value: Any, _depth: int = 0) -> Any: 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 5bda7d8..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() 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") From 6d15c4fec478d51c13c4f9d43ec9351145eb9088 Mon Sep 17 00:00:00 2001 From: Susheem Koul <105606472+susheem-k@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:47:43 +0530 Subject: [PATCH 5/5] Remove unused typing.Text import from JsonlStore. Fixes the ruff F401 that was failing CI lint on PR #37. Co-authored-by: Cursor --- chronicle/envelope/store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chronicle/envelope/store.py b/chronicle/envelope/store.py index b0adc6d..43e1c03 100644 --- a/chronicle/envelope/store.py +++ b/chronicle/envelope/store.py @@ -3,7 +3,7 @@ from __future__ import annotations from pathlib import Path -from typing import IO, Text +from typing import IO from chronicle.envelope.schema import Envelope