diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0656d83..ba4a9c4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,3 +17,13 @@ repos: - id: ruff args: ["--fix"] - id: ruff-format + + # Match CI / `make lint`: project mypy config from pyproject.toml. + - repo: local + hooks: + - id: mypy + name: mypy + entry: python -m mypy + language: system + pass_filenames: false + types: [python] diff --git a/CHANGELOG.md b/CHANGELOG.md index 2532f9f..bcb446f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.1.2] - 2026-07-24 + +### Added + +- LLM-kind Chronicle `@boundary` / `wrap_llm` runs **pre_call** via `session.on_enter` + (and ledger admit/complete via `on_leave`). A bare `@boundary(..., kind="llm")` under + `tokenops_run` is enough for worst-case halt / output-cap MUTATE — no `wrap_complete` + required. `wrap_complete` still works and opts out of the double pre_call. + +### Changed + +- Require `agent-chronicle>=0.3.0` (on_enter / on_leave hooks). + ## [0.1.1] - 2026-07-24 ### Changed diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1da2af7..5da085e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -19,7 +19,7 @@ cp .env.example .env # optional API keys for live demos Install [Git LFS](https://git-lfs.com) if you need local copies of demo videos / the product deck (tracked as LFS pointers). -Optional local hooks (same checks as CI): +Optional local hooks (same checks as CI — ruff + mypy): ```bash pre-commit install diff --git a/docs/guides/onboarding.md b/docs/guides/onboarding.md index db7a165..3c99f44 100644 --- a/docs/guides/onboarding.md +++ b/docs/guides/onboarding.md @@ -34,7 +34,7 @@ Chronicle records decision boundaries; TokenOps observes them for cost/governanc | **Or embedded Store** *(single-process / tests)* | Omit `TOKENOPS_URL` or set `TOKENOPS_EMBEDDED=1`. | | **LLM API keys** | Only for real model calls — not required for TokenOps itself or offline tests. | | **FastAPI** | Only if you use `instrument_app`. Non-FastAPI: pass kwargs / `RequestContext` to `tokenops_run`. | -| **Chronicle `@boundary`** | Only if tools should be governed; LLM-only stacks can stop at `wrap_complete`. | +| **Chronicle `@boundary`** | Tools *and* LLM: `kind="llm"` runs pre_call via `on_enter`; tools stay observe-only. LLM-only stacks can use bare `@boundary(..., kind="llm")` under `tokenops_run` instead of `wrap_complete`. | You do **not** need A2A, `create_a2a_app`, LangChain, or the Admin UI for governance to work. @@ -132,7 +132,8 @@ The **agent** (via `instrument_app` / `tokenops_run` kwargs), not the UI. Client ### When do I need Chronicle `@boundary`? -When tools (search, fetch, etc.) should appear on the ledger / be governed. Without `@boundary` + the crossing hook, TokenOps only sees LLM calls you put through `wrap_complete`. +- **LLM calls:** `@boundary(..., kind="llm")` under `tokenops_run` runs **pre_call** (via `on_enter`) and observe — enough for LLM-only stacks without `wrap_complete`. +- **Tools** (search, fetch, etc.): still need `@boundary` + the crossing hook so they appear on the ledger / are governed. Without that, TokenOps only sees LLM crossings you decorate (or put through `wrap_complete`). ### Embedded Store vs `TOKENOPS_URL`? diff --git a/pyproject.toml b/pyproject.toml index 80f5cf3..4bfdf5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "agent-tokenops" -version = "0.1.1" +version = "0.1.2" description = "TokenOps control plane and SDK for run-aware agent token governance" readme = "README.md" requires-python = ">=3.10" @@ -34,7 +34,7 @@ classifiers = [ "Topic :: System :: Systems Administration", ] dependencies = [ - "agent-chronicle>=0.2.0", + "agent-chronicle @ git+https://github.com/theagentplane/chronicle.git@feat/boundary-on-enter-precall", "openai>=1.0", "anthropic>=0.40", "pyyaml>=6.0", diff --git a/src/tokenops/__init__.py b/src/tokenops/__init__.py index c21eb09..0d35c1c 100644 --- a/src/tokenops/__init__.py +++ b/src/tokenops/__init__.py @@ -4,7 +4,7 @@ load_env() -__version__ = "0.1.1" +__version__ = "0.1.2" def init() -> None: diff --git a/src/tokenops/control/crossing.py b/src/tokenops/control/crossing.py index 8b3fa3f..8185a30 100644 --- a/src/tokenops/control/crossing.py +++ b/src/tokenops/control/crossing.py @@ -1,17 +1,37 @@ -"""Chronicle ``on_crossing`` → TokenOps govern ingest. +"""Chronicle boundary hooks → TokenOps govern (pre_call + observe). -Chronicle's ``@boundary`` invokes ``session.on_crossing`` after LIVE record and -LIVE cut-point. ``reset_session()`` creates a fresh session with ``on_crossing`` -cleared, so ``install_crossing_hook`` wraps ``reset_session`` to re-attach. +Chronicle ``@boundary`` / ``wrap_llm`` invoke: + +* ``session.on_enter`` after input capture, before the wrapped call (LLM pre_call) +* ``session.on_crossing`` after LIVE record / cut-point (observe) +* ``session.on_leave`` after the attempt when ``on_enter`` completed (ledger complete) + +``reset_session()`` clears hooks, so ``install_crossing_hook`` wraps ``reset_session`` +to re-attach all three. """ from __future__ import annotations +from contextvars import ContextVar from typing import Any from chronicle.envelope.schema import Envelope, InputState from chronicle.session import ChronicleSession, get_session +# When wrap_complete / wrap_stream already run pre_call + admit, skip boundary pre_call +# so LLM calls are not governed twice. +_wrap_owns_precall: ContextVar[bool] = ContextVar("tokenops_wrap_owns_precall", default=False) +_inflight_seg: ContextVar[str | None] = ContextVar("tokenops_inflight_seg", default=None) + + +def wrap_owns_precall(active: bool) -> Any: + """Context manager token: ``wrap_complete`` sets True while it owns pre_call.""" + return _wrap_owns_precall.set(active) + + +def reset_wrap_owns_precall(token: Any) -> None: + _wrap_owns_precall.reset(token) + def _service_name() -> str: from tokenops.control.context import current_span @@ -29,6 +49,90 @@ def _input_state_to_dict(input_state: InputState) -> dict[str, object]: return d +def _estimate_input_tokens(state: dict[str, object]) -> int: + messages = state.get("messages") + if messages is None: + messages = state + return max(1, len(str(messages)) // 4) + + +def on_enter( + boundary_id: str, + kind: str, + input_state: InputState, +) -> dict[str, Any] | None: + """LLM-kind pre_call: may Halt, or return kwargs patches (max_output / model).""" + del boundary_id # used for tracing only; attribution comes from run scope + if kind != "llm" or _wrap_owns_precall.get(): + return None + + from tokenops.control.context import current_governance, current_registration + from tokenops.control.core import CallRequest + + gov_ctx = current_governance() + if gov_ctx is None or current_registration() is None: + return None + + governor = gov_ctx.governor + controls = governor.controls + begin = getattr(controls, "begin_call", None) + if callable(begin): + begin() + + state = _input_state_to_dict(input_state) + provider = str(state.get("provider") or gov_ctx.provider or "") + model = str(state.get("model") or gov_ctx.model or "") + raw_cap = state.get("max_output_tokens") + max_out: int | None = None + if isinstance(raw_cap, bool): + max_out = None + elif isinstance(raw_cap, int): + max_out = raw_cap + elif isinstance(raw_cap, str): + try: + max_out = int(raw_cap) + except ValueError: + max_out = None + + request = CallRequest( + attr=gov_ctx.attr, + provider=provider, + model=model, + estimated_input_tokens=_estimate_input_tokens(state), + max_output_tokens=max_out, + ) + governor.pre_call(request) + + seg = f"run:{gov_ctx.attr.run_id}" + governor.ledger.admit(seg) + _inflight_seg.set(seg) + + patch: dict[str, Any] = {} + call = getattr(controls, "call", None) + if call is not None: + if getattr(call, "max_output_tokens", None) is not None: + patch["max_output_tokens"] = call.max_output_tokens + if getattr(call, "model_override", None): + patch["model"] = call.model_override + return patch or None + + +def on_leave(boundary_id: str, kind: str, input_state: InputState) -> None: + """Release inflight admit from :func:`on_enter` (success or failure).""" + del boundary_id, kind, input_state + from tokenops.control.context import current_governance + + seg = _inflight_seg.get() + if not seg: + return + gov_ctx = current_governance() + try: + if gov_ctx is not None: + gov_ctx.governor.ledger.complete(seg) + finally: + _inflight_seg.set(None) + + def on_crossing( boundary_id: str, kind: str, @@ -73,11 +177,13 @@ def recorded_envelopes(self) -> list[Envelope]: def _attach(session: ChronicleSession) -> ChronicleSession: session.on_crossing = on_crossing + session.on_enter = on_enter + session.on_leave = on_leave return session def install_crossing_hook() -> None: - """Attach ``on_crossing`` to the current session and every ``reset_session()``. + """Attach enter/crossing/leave hooks to the current session and every ``reset_session()``. Idempotent. Patches ``chronicle.session.reset_session`` (and rebinds the ``chronicle`` package export when present) so callers keep the wrapped entrypoint. diff --git a/src/tokenops/control/integration.py b/src/tokenops/control/integration.py index 0de485e..4a28e5d 100644 --- a/src/tokenops/control/integration.py +++ b/src/tokenops/control/integration.py @@ -233,38 +233,44 @@ def wrap_complete( traced = wrap_llm(boundary_id, dispatch) def governed(p: str, m: str, messages) -> object: - controls.begin_call() - request = CallRequest( - attr=attr, - provider=provider, - model=m, - estimated_input_tokens=estimate(messages), - max_output_tokens=controls.call.max_output_tokens, - ) - governor.pre_call(request) + from tokenops.control.crossing import reset_wrap_owns_precall, wrap_owns_precall - use_model = controls.call.model_override or m - messages = consume_carry(controls, messages) - if controls.call.compact: # deep prompt compaction - messages = _compact_messages(messages) - - governor.ledger.admit(seg) + owns = wrap_owns_precall(True) try: - cap = controls.call.max_output_tokens - penalties: dict = {} - attempt = 0 - while True: - controls.retry = False - # Chronicle records + on_crossing → Governor.observe (when bound) - response = traced(p, use_model, messages, max_output_tokens=cap, **penalties) - if controls.retry and attempt < max_call_retries: - attempt += 1 - cap = _tighten_cap(cap) - penalties = {"frequency_penalty": 1.0, "presence_penalty": 0.6} - continue - return response + controls.begin_call() + request = CallRequest( + attr=attr, + provider=provider, + model=m, + estimated_input_tokens=estimate(messages), + max_output_tokens=controls.call.max_output_tokens, + ) + governor.pre_call(request) + + use_model = controls.call.model_override or m + messages = consume_carry(controls, messages) + if controls.call.compact: # deep prompt compaction + messages = _compact_messages(messages) + + governor.ledger.admit(seg) + try: + cap = controls.call.max_output_tokens + penalties: dict = {} + attempt = 0 + while True: + controls.retry = False + # Chronicle records + on_crossing → Governor.observe (when bound) + response = traced(p, use_model, messages, max_output_tokens=cap, **penalties) + if controls.retry and attempt < max_call_retries: + attempt += 1 + cap = _tighten_cap(cap) + penalties = {"frequency_penalty": 1.0, "presence_penalty": 0.6} + continue + return response + finally: + governor.ledger.complete(seg) finally: - governor.ledger.complete(seg) + reset_wrap_owns_precall(owns) return governed @@ -366,39 +372,45 @@ def _stream_once(p, use_model, messages, max_output_tokens=None, **penalties): traced = wrap_llm(boundary_id, _stream_once) def governed(p: str, m: str, messages) -> object: - controls.begin_call() - governor.pre_call( - CallRequest( - attr=attr, - provider=provider, - model=m, - estimated_input_tokens=estimate(messages), - max_output_tokens=controls.call.max_output_tokens, - ) - ) - use_model = controls.call.model_override or m - messages = consume_carry(controls, messages) - if controls.call.compact: # deep prompt compaction - messages = _compact_messages(messages) + from tokenops.control.crossing import reset_wrap_owns_precall, wrap_owns_precall - governor.ledger.admit(seg) + owns = wrap_owns_precall(True) try: - cap = controls.call.max_output_tokens - penalties: dict = {} - attempt = 0 - while True: - controls.retry = False - cancel_flag["cancelled"] = False - resp = traced(p, use_model, messages, max_output_tokens=cap, **penalties) - if cancel_flag["cancelled"] and on_cancel: - on_cancel() - if (cancel_flag["cancelled"] or controls.retry) and attempt < max_call_retries: - attempt += 1 - cap = _tighten_cap(cap) - penalties = {"frequency_penalty": 1.0, "presence_penalty": 0.6} - continue - return resp + controls.begin_call() + governor.pre_call( + CallRequest( + attr=attr, + provider=provider, + model=m, + estimated_input_tokens=estimate(messages), + max_output_tokens=controls.call.max_output_tokens, + ) + ) + use_model = controls.call.model_override or m + messages = consume_carry(controls, messages) + if controls.call.compact: # deep prompt compaction + messages = _compact_messages(messages) + + governor.ledger.admit(seg) + try: + cap = controls.call.max_output_tokens + penalties: dict = {} + attempt = 0 + while True: + controls.retry = False + cancel_flag["cancelled"] = False + resp = traced(p, use_model, messages, max_output_tokens=cap, **penalties) + if cancel_flag["cancelled"] and on_cancel: + on_cancel() + if (cancel_flag["cancelled"] or controls.retry) and attempt < max_call_retries: + attempt += 1 + cap = _tighten_cap(cap) + penalties = {"frequency_penalty": 1.0, "presence_penalty": 0.6} + continue + return resp + finally: + governor.ledger.complete(seg) finally: - governor.ledger.complete(seg) + reset_wrap_owns_precall(owns) return governed diff --git a/tests/test_llm_boundary_precall.py b/tests/test_llm_boundary_precall.py new file mode 100644 index 0000000..7018025 --- /dev/null +++ b/tests/test_llm_boundary_precall.py @@ -0,0 +1,94 @@ +"""LLM-kind @boundary runs pre_call via Chronicle on_enter (no wrap_complete).""" + +from __future__ import annotations + +import chronicle.session as chronicle_session +import pytest +from chronicle import boundary + +from conftest import toy_price +from tokenops.control import ApplyControls, Halt, build_governor, install_crossing_hook +from tokenops.control.attribution import _build_attribution +from tokenops.control.context import SpanContext, _governance_scope, clear, run_scope +from tokenops.control.models import RunRegistration +from tokenops.control.store import Store +from tokenops.providers.types import ModelResponse + + +@pytest.fixture +def store(tmp_path): + s = Store(str(tmp_path / "b.db")) + yield s + s.close() + + +@pytest.fixture(autouse=True) +def _hooks(): + install_crossing_hook() + + +def test_llm_boundary_pre_call_halt(store): + clear() + reg = store.register_run(RunRegistration(run_id="r-halt")) + config = { + "governance": { + "budgets": [{"id": "run_llm_cap", "limit_micros": 1, "dimension": "run"}], + "policies": { + "cost_budget": {"budget": "run_llm_cap"}, + "pre_call_worst_case": {"budget": "run_llm_cap", "default_max_output": 1024}, + }, + } + } + controls = ApplyControls() + gov = build_governor(config, toy_price, controls) + attr = _build_attribution(reg, service="agent") + gov.ledger.open_run("r-halt") + + called = {"n": 0} + + @boundary("agent.chat", kind="llm") + def chat(model: str, messages: list, *, max_output_tokens: int | None = None) -> ModelResponse: + called["n"] += 1 + return ModelResponse(content="x", input_tokens=1, output_tokens=1) + + chronicle_session.reset_session().begin_trace("r-halt") + with run_scope(reg, SpanContext(span_id="s1", service="agent")): + with _governance_scope(gov, attr, provider="openai", model="gpt-4o-mini"): + with pytest.raises(Halt): + chat("gpt-4o-mini", [{"role": "user", "content": "hi"}]) + + assert called["n"] == 0 + chronicle_session.reset_session() + + +def test_llm_boundary_mutate_max_output(store): + clear() + reg = store.register_run(RunRegistration(run_id="r-mut")) + config = { + "governance": { + "budgets": [{"id": "run_llm_cap", "limit_micros": 10_000_000, "dimension": "run"}], + "policies": { + "pre_call_worst_case": {"budget": "run_llm_cap", "default_max_output": 256}, + }, + } + } + controls = ApplyControls() + gov = build_governor(config, toy_price, controls) + attr = _build_attribution(reg, service="agent") + gov.ledger.open_run("r-mut") + + seen: dict = {} + + @boundary("agent.chat", kind="llm") + def chat(model: str, messages: list, *, max_output_tokens: int | None = None) -> ModelResponse: + seen["max_output_tokens"] = max_output_tokens + return ModelResponse(content="ok", input_tokens=2, output_tokens=1) + + chronicle_session.reset_session().begin_trace("r-mut") + with run_scope(reg, SpanContext(span_id="s1", service="agent")): + with _governance_scope(gov, attr, provider="openai", model="gpt-4o-mini"): + out = chat("gpt-4o-mini", [{"role": "user", "content": "hi"}]) + + assert out.content == "ok" + assert seen["max_output_tokens"] == 256 + chronicle_session.reset_session()