From 0969acea959aed0450d0994544ec9765f5c3e0b8 Mon Sep 17 00:00:00 2001
From: Susheem Koul <105606472+susheem-k@users.noreply.github.com>
Date: Fri, 24 Jul 2026 19:42:24 +0530
Subject: [PATCH] Add session on_enter/on_leave for pre-call governance on
 boundaries.

Governors can abort or MUTATE kwargs before the wrapped call; on_leave pairs cleanup. Bumps to 0.3.0.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 CHANGELOG.md           |  10 ++++
 chronicle/__init__.py  |   2 +-
 chronicle/boundary.py  |  89 ++++++++++++++++++++++---------
 chronicle/session.py   |   9 ++++
 pyproject.toml         |   2 +-
 tests/test_on_enter.py | 115 +++++++++++++++++++++++++++++++++++++++++
 6 files changed, 200 insertions(+), 27 deletions(-)
 create mode 100644 tests/test_on_enter.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0a68e85..21332aa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.3.0] - 2026-07-24
+
+### Added
+- **`session.on_enter` / `session.on_leave`**: pre-call and paired cleanup hooks on
+  `@boundary` / `wrap_llm` (LIVE + live cut-point). `on_enter` runs after input
+  capture and before the wrapped function; it may raise to abort the call, or
+  return a kwargs mapping to merge (MUTATE). `on_leave` runs after the attempt
+  when `on_enter` completed, including when the function raises. Governors
+  (e.g. TokenOps) use this for LLM-kind pre_call without a separate wrap.
+
 ## [0.2.0] - 2026-07-23
 
 ### Changed
diff --git a/chronicle/__init__.py b/chronicle/__init__.py
index fe1c7ee..c599bcd 100644
--- a/chronicle/__init__.py
+++ b/chronicle/__init__.py
@@ -18,7 +18,7 @@
 from chronicle.session import ChronicleSession, SessionMode, get_session, reset_session
 from chronicle.wrap import instrument_langgraph, wrap
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "ActionResult",
diff --git a/chronicle/boundary.py b/chronicle/boundary.py
index bd67319..86f5c53 100644
--- a/chronicle/boundary.py
+++ b/chronicle/boundary.py
@@ -44,9 +44,11 @@ def boundary(
     """
     Annotate a decision boundary for Chronicle record and replay.
 
-    LIVE mode:     execute the function, record an envelope, return its real value
+    LIVE mode:     optional ``session.on_enter`` (may abort / patch kwargs), execute
+                   the function, record an envelope, ``on_crossing``, then ``on_leave``
     REPLAY + STUB: return the recorded fixture without executing
     REPLAY + LIVE: execute the function (cut-point), capture input/result for asserts
+                   (``on_enter`` / ``on_leave`` / ``on_crossing`` still apply)
 
     Works on sync functions and ``async def`` coroutines. The wrapper is
     transparent: the caller always gets exactly what the function returned (or the
@@ -160,26 +162,55 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 # Recording (LIVE mode)
 # --------------------------------------------------------------------------- #
 
+def _apply_on_enter(session, boundary_id, kind, input_state, kwargs) -> tuple[dict, bool]:
+    """Run ``on_enter`` if set. Returns ``(call_kwargs, entered)``.
+
+    ``entered`` is True only when ``on_enter`` returned (so ``on_leave`` can pair).
+    A raise from ``on_enter`` (e.g. TokenOps ``Halt``, a ``BaseException``) aborts
+    before the wrapped function and does not mark entered.
+    """
+    call_kwargs = dict(kwargs)
+    if session.on_enter is None:
+        return call_kwargs, False
+    patch = session.on_enter(boundary_id, kind, input_state)
+    if patch:
+        call_kwargs.update(dict(patch))
+    return call_kwargs, True
+
+
+def _run_on_leave(session, boundary_id, kind, input_state, entered: bool) -> None:
+    if entered and session.on_leave is not None:
+        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)
+    call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs)
     try:
-        result = fn(*args, **kwargs)
-    except Exception as exc:
-        _record_failure(session, boundary_id, kind, input_state, exc)
-        raise
-    _record_success(session, boundary_id, kind, input_state, result, extract_result, extract_metadata)
-    return result
+        try:
+            result = fn(*args, **call_kwargs)
+        except Exception as exc:
+            _record_failure(session, boundary_id, kind, input_state, exc)
+            raise
+        _record_success(session, boundary_id, kind, input_state, result, extract_result, extract_metadata)
+        return result
+    finally:
+        _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)
+    call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs)
     try:
-        result = await fn(*args, **kwargs)
-    except Exception as exc:
-        _record_failure(session, boundary_id, kind, input_state, exc)
-        raise
-    _record_success(session, boundary_id, kind, input_state, result, extract_result, extract_metadata)
-    return result
+        try:
+            result = await fn(*args, **call_kwargs)
+        except Exception as exc:
+            _record_failure(session, boundary_id, kind, input_state, exc)
+            raise
+        _record_success(session, boundary_id, kind, input_state, result, extract_result, extract_metadata)
+        return result
+    finally:
+        _run_on_leave(session, boundary_id, kind, input_state, entered)
 
 
 def _record_success(session, boundary_id, kind, input_state, result, extract_result, extract_metadata):
@@ -228,25 +259,33 @@ def _call_metadata(result, kind, extract_metadata):
 def _live_cutpoint_call(session, fn, boundary_id, kind, args, kwargs, extract_input, invocation_index):
     input_state = _capture_input(fn, args, kwargs, extract_input)
     session.capture_live_input(boundary_id, invocation_index, input_state)
+    call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs)
     try:
-        result = fn(*args, **kwargs)
-    except Exception:
-        _advance_cutpoint(session, boundary_id, invocation_index)
-        raise
-    _finish_cutpoint(session, boundary_id, kind, input_state, result, invocation_index)
-    return result
+        try:
+            result = fn(*args, **call_kwargs)
+        except Exception:
+            _advance_cutpoint(session, boundary_id, invocation_index)
+            raise
+        _finish_cutpoint(session, boundary_id, kind, input_state, result, invocation_index)
+        return result
+    finally:
+        _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)
     session.capture_live_input(boundary_id, invocation_index, input_state)
+    call_kwargs, entered = _apply_on_enter(session, boundary_id, kind, input_state, kwargs)
     try:
-        result = await fn(*args, **kwargs)
-    except Exception:
-        _advance_cutpoint(session, boundary_id, invocation_index)
-        raise
-    _finish_cutpoint(session, boundary_id, kind, input_state, result, invocation_index)
-    return result
+        try:
+            result = await fn(*args, **call_kwargs)
+        except Exception:
+            _advance_cutpoint(session, boundary_id, invocation_index)
+            raise
+        _finish_cutpoint(session, boundary_id, kind, input_state, result, invocation_index)
+        return result
+    finally:
+        _run_on_leave(session, boundary_id, kind, input_state, entered)
 
 
 def _finish_cutpoint(session, boundary_id, kind, input_state, result, invocation_index):
diff --git a/chronicle/session.py b/chronicle/session.py
index d683307..3210a6b 100644
--- a/chronicle/session.py
+++ b/chronicle/session.py
@@ -54,6 +54,15 @@ class ChronicleSession:
     # Optional observer for boundary crossings (LIVE record + LIVE cut-point).
     # Signature: (boundary_id, kind, input_state, result) -> None
     on_crossing: Callable[[str, str, InputState, Any], None] | None = None
+    # Optional pre-call hook (LIVE record + LIVE cut-point), after input capture
+    # and before the wrapped function runs. May raise to abort (e.g. a governor
+    # Halt). May return a mapping of kwargs to merge into the call (MUTATE).
+    # Signature: (boundary_id, kind, input_state) -> Mapping[str, Any] | None
+    on_enter: Callable[[str, str, InputState], Mapping[str, Any] | None] | None = None
+    # Optional post-call cleanup (LIVE), always run after a successful on_enter
+    # whether the function returned or raised. Signature:
+    # (boundary_id, kind, input_state) -> None
+    on_leave: Callable[[str, str, InputState], None] | None = None
     # Applied to each envelope before it is retained or stored, so secrets never
     # reach a committed fixture. Empty by default; set to default_redactors() or
     # your own. Signature: (str) -> str. See chronicle.redaction.
diff --git a/pyproject.toml b/pyproject.toml
index 11d68f7..e646544 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "agent-chronicle"
-version = "0.2.0"
+version = "0.3.0"
 description = "Record-and-replay for agent decision graphs: reproduce a prod agent failure as a committed regression test — and re-run your fix without live LLM calls."
 readme = "README.md"
 requires-python = ">=3.10"
diff --git a/tests/test_on_enter.py b/tests/test_on_enter.py
new file mode 100644
index 0000000..a56a915
--- /dev/null
+++ b/tests/test_on_enter.py
@@ -0,0 +1,115 @@
+"""Tests for ChronicleSession.on_enter / on_leave pre-call hooks."""
+
+from __future__ import annotations
+
+import pytest
+
+from chronicle.boundary import boundary
+from chronicle.session import reset_session
+
+
+@boundary("chat", kind="llm")
+def chat(model: str, messages: list, *, max_output_tokens: int | None = None) -> dict:
+    return {
+        "content": "ok",
+        "model": model,
+        "max_output_tokens": max_output_tokens,
+        "n_messages": len(messages),
+    }
+
+
+@pytest.mark.layer1
+def test_on_enter_runs_before_function():
+    order: list[str] = []
+
+    @boundary("step", kind="llm")
+    def step(x: int) -> int:
+        order.append("fn")
+        return x
+
+    session = reset_session()
+    session.enable_live()
+
+    def on_enter(boundary_id, kind, input_state):
+        order.append(f"enter:{boundary_id}:{kind}")
+        return None
+
+    session.on_enter = on_enter
+    assert step(1) == 1
+    assert order == ["enter:step:llm", "fn"]
+
+
+@pytest.mark.layer1
+def test_on_enter_can_patch_kwargs():
+    session = reset_session()
+    session.enable_live()
+    session.on_enter = lambda *_: {"max_output_tokens": 128}
+
+    out = chat("gpt-4o-mini", [{"role": "user", "content": "hi"}])
+    assert out["max_output_tokens"] == 128
+
+
+@pytest.mark.layer1
+def test_on_enter_raise_skips_function_and_on_leave():
+    called = {"fn": False, "leave": False}
+
+    @boundary("blocked", kind="llm")
+    def blocked() -> str:
+        called["fn"] = True
+        return "nope"
+
+    session = reset_session()
+    session.enable_live()
+
+    class PreCallAbort(BaseException):
+        pass
+
+    def on_enter(*_):
+        raise PreCallAbort("halt")
+
+    def on_leave(*_):
+        called["leave"] = True
+
+    session.on_enter = on_enter
+    session.on_leave = on_leave
+
+    with pytest.raises(PreCallAbort):
+        blocked()
+    assert called == {"fn": False, "leave": False}
+
+
+@pytest.mark.layer1
+def test_on_leave_pairs_with_successful_on_enter():
+    events: list[str] = []
+
+    @boundary("t", kind="tool")
+    def tool(x: int) -> int:
+        events.append("fn")
+        return x
+
+    session = reset_session()
+    session.enable_live()
+    session.on_enter = lambda *_: (events.append("enter") or None)
+    session.on_leave = lambda *_: events.append("leave")
+
+    assert tool(3) == 3
+    assert events == ["enter", "fn", "leave"]
+
+
+@pytest.mark.layer1
+def test_on_leave_runs_when_function_raises():
+    events: list[str] = []
+
+    @boundary("t", kind="tool")
+    def boom() -> None:
+        events.append("fn")
+        raise ValueError("x")
+
+    session = reset_session()
+    session.enable_live()
+    session.on_enter = lambda *_: (events.append("enter") or None)
+    session.on_leave = lambda *_: events.append("leave")
+
+    with pytest.raises(ValueError):
+        boom()
+    assert events == ["enter", "fn", "leave"]