diff --git a/python/pyproject.toml b/python/pyproject.toml index 171ce9d..3b2539f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "rewind-agent" -version = "0.16.0" +version = "0.16.1" description = "Chrome DevTools for AI agents — record, inspect, fork, replay, diff." readme = "README.md" license = "MIT" diff --git a/python/rewind_agent/__init__.py b/python/rewind_agent/__init__.py index 0d0af4f..0d95131 100644 --- a/python/rewind_agent/__init__.py +++ b/python/rewind_agent/__init__.py @@ -15,6 +15,10 @@ def search(query): # Wrap LangGraph / CrewAI for automatic instrumentation graph = rewind_agent.wrap_langgraph(graph) crew = rewind_agent.wrap_crew(crew) + + # One-call connector for custom agents (any HTTP transport, any framework) + with rewind_agent.connector.setup(name="my-agent"): + run_agent_loop() """ from .patch import init, uninit, session, replay, thread @@ -31,6 +35,7 @@ def search(query): ) from .cached_call import cached_llm_call from .explicit import ExplicitClient, RewindReplayDivergenceError +from . import connector from .assertions import Assertions, AssertionResult from .openai_agents import openai_agents_hooks from .pydantic_ai import pydantic_ai_hooks @@ -92,6 +97,8 @@ def search(query): "RewindReplayDivergenceError", # Cached LLM call decorator (Phase 2 / Tier 2) "cached_llm_call", + # One-call connector for any agent (see docs/hdk.md) + "connector", ] @@ -113,4 +120,4 @@ def import_from_langfuse(trace_id: str, **kwargs) -> str: return _import(trace_id, **kwargs) -__version__ = "0.16.0" +__version__ = "0.16.1" diff --git a/python/rewind_agent/connector.py b/python/rewind_agent/connector.py new file mode 100644 index 0000000..e554b7f --- /dev/null +++ b/python/rewind_agent/connector.py @@ -0,0 +1,175 @@ +"""One-call connector for any agent that wants Rewind recording. + +The Tier-1 ergonomic wrapper described in [docs/hdk.md](docs/hdk.md). +Composes :class:`ExplicitClient` + :func:`intercept.install` in the +correct order so callers don't have to remember the dance — and +don't hit the silent-no-op trap where intercept records nothing +because no session is active. + +Usage +----- + +>>> import rewind_agent +>>> with rewind_agent.connector.setup(name="my-agent"): +... run_agent_loop() + +For a custom LLM gateway hostname, pass ``llm_hosts`` (or set +``REWIND_LLM_HOSTS`` in the env): + +>>> with rewind_agent.connector.setup( +... name="my-agent", +... llm_hosts=("llm-gateway.corp.example",), +... ): +... ... + +The yielded value is the underlying :class:`ExplicitClient`, available +for non-HTTP record paths inside the block: + +>>> with rewind_agent.connector.setup(name="my-agent") as client: +... resp = my_grpc_llm.chat(req) # not HTTP — intercept can't see it +... client.record_llm_call(req, resp.dict(), model="...", duration_ms=...) + +Environment variables (env > default; kwargs override env) +---------------------------------------------------------- + +* ``REWIND_ENABLED`` — set to ``0`` to make ``setup()`` a no-op with + zero overhead. Yields ``None`` instead of a client; the ``with`` + block runs unmodified. +* ``REWIND_URL`` — Rewind server URL. Default ``http://127.0.0.1:4800``. +* ``REWIND_LLM_HOSTS`` — comma-separated hostnames to treat as LLM + gateways in addition to the strict-by-default provider list. + +Replay-context interaction +-------------------------- + +If ``REWIND_SESSION_ID`` and ``REWIND_REPLAY_CONTEXT_ID`` are set in +the environment (the runner subprocess pattern documented in +docs/runners.md), ``setup()`` skips creating a fresh session and lets +``intercept.install()`` attach to the existing replay context. This +makes the connector safe to drop into runner-driven replay handlers +without phantom sessions. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager +from typing import Iterable, Iterator + +from rewind_agent.explicit import ExplicitClient +from rewind_agent.intercept import ( + DefaultPredicates, + install, + is_installed, + uninstall, +) + +logger = logging.getLogger(__name__) + + +class _HostPredicates(DefaultPredicates): + """DefaultPredicates extended with caller-provided LLM gateway hostnames.""" + + def __init__(self, hosts: tuple[str, ...]) -> None: + super().__init__() + self._hosts = tuple(h for h in (s.strip().lower() for s in hosts) if h) + + def is_llm_call(self, req) -> bool: # type: ignore[override] + netloc = req.url_parts.netloc.lower() + if any(h in netloc for h in self._hosts): + return True + return super().is_llm_call(req) + + +def _resolve_hosts(llm_hosts: Iterable[str] | None) -> tuple[str, ...]: + if llm_hosts is not None: + return tuple(llm_hosts) + env = os.environ.get("REWIND_LLM_HOSTS", "") + return tuple(h for h in env.split(",") if h.strip()) if env else () + + +def _enabled(enabled: bool | None) -> bool: + if enabled is not None: + return enabled + return os.environ.get("REWIND_ENABLED", "1") != "0" + + +def _is_replay_dispatch() -> bool: + """Detect runner-subprocess replay env vars. + + When these are set, intercept.install() will attach to the existing + replay context — we must NOT start a fresh session in that case. + """ + return bool( + os.environ.get("REWIND_SESSION_ID") + and os.environ.get("REWIND_REPLAY_CONTEXT_ID") + ) + + +@contextmanager +def setup( + name: str, + *, + base_url: str | None = None, + llm_hosts: Iterable[str] | None = None, + enabled: bool | None = None, + thread_id: str | None = None, + metadata: dict | None = None, +) -> Iterator[ExplicitClient | None]: + """Connect any agent to Rewind for the duration of a ``with`` block. + + Starts a session, installs HTTP intercept (with custom predicates + when ``llm_hosts`` is set), yields the :class:`ExplicitClient` for + use inside the block, and tears both down on exit. + + Parameters + ---------- + name: + Session name shown in ``rewind show`` and the dashboard. + base_url: + Rewind server URL. Defaults to ``$REWIND_URL`` or + ``http://127.0.0.1:4800``. + llm_hosts: + Iterable of hostnames to treat as LLM gateways. Defaults to + the parsed value of ``$REWIND_LLM_HOSTS`` or no extras. + enabled: + Override the ``$REWIND_ENABLED`` kill switch. + thread_id, metadata: + Forwarded to :meth:`ExplicitClient.session`. + + Yields + ------ + ExplicitClient | None + The recording client, or ``None`` when disabled. + """ + if not _enabled(enabled): + yield None + return + + url = base_url or os.environ.get("REWIND_URL", "http://127.0.0.1:4800") + client = ExplicitClient(base_url=url) + hosts = _resolve_hosts(llm_hosts) + predicates = _HostPredicates(hosts) if hosts else None + + if _is_replay_dispatch(): + # Runner-driven replay: intercept.install() will attach to the + # existing replay context via env vars. Don't create a phantom + # session. + already_installed = is_installed() + install(predicates=predicates) + try: + yield client + finally: + if not already_installed: + uninstall() + return + + with client.session(name, thread_id=thread_id, metadata=metadata): + already_installed = is_installed() + install(predicates=predicates) + try: + yield client + finally: + if not already_installed: + uninstall() diff --git a/python/tests/test_connector.py b/python/tests/test_connector.py new file mode 100644 index 0000000..ddc2c5c --- /dev/null +++ b/python/tests/test_connector.py @@ -0,0 +1,246 @@ +"""Tests for the one-call connector (rewind_agent.connector.setup).""" + +import os +import threading +import unittest +from http.server import HTTPServer, BaseHTTPRequestHandler +from unittest import mock + +import rewind_agent +from rewind_agent.connector import _HostPredicates, _is_replay_dispatch, setup +from rewind_agent.explicit import ( + ExplicitClient, + _replay_context_id, + _session_id, + _timeline_id, +) +from rewind_agent.intercept import is_installed, uninstall + + +class _MockHandler(BaseHTTPRequestHandler): + """Minimal mock — only what the connector needs.""" + + sessions_started: list = [] + sessions_ended: list = [] + + def do_POST(self): # noqa: N802 — stdlib API + import json + + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length)) if length else {} + + if self.path == "/api/sessions/start": + _MockHandler.sessions_started.append(body) + self._respond(201, { + "session_id": f"s-{len(_MockHandler.sessions_started)}", + "root_timeline_id": f"tl-{len(_MockHandler.sessions_started)}", + }) + elif self.path.endswith("/end"): + _MockHandler.sessions_ended.append(self.path) + self._respond(200, {"session_id": self.path.split("/")[3]}) + else: + self._respond(404, {"error": "unhandled"}) + + def _respond(self, status, body): + import json + + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(body).encode()) + + def log_message(self, *_): # silence + pass + + +class _ConnectorTestBase(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.server = HTTPServer(("127.0.0.1", 0), _MockHandler) + cls.port = cls.server.server_address[1] + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + cls.base_url = f"http://127.0.0.1:{cls.port}" + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + + def setUp(self): + _MockHandler.sessions_started = [] + _MockHandler.sessions_ended = [] + _session_id.set(None) + _timeline_id.set(None) + _replay_context_id.set(None) + # Make sure intercept is not lingering from a prior test. + if is_installed(): + uninstall() + + +class TestKillSwitch(_ConnectorTestBase): + def test_disabled_via_env_yields_none_and_no_http(self): + with mock.patch.dict(os.environ, {"REWIND_ENABLED": "0"}, clear=False): + with setup(name="off", base_url=self.base_url) as client: + self.assertIsNone(client) + self.assertFalse(is_installed()) + self.assertEqual(_MockHandler.sessions_started, []) + + def test_disabled_via_kwarg(self): + with setup(name="off", base_url=self.base_url, enabled=False) as client: + self.assertIsNone(client) + self.assertEqual(_MockHandler.sessions_started, []) + + +class TestSessionLifecycle(_ConnectorTestBase): + def test_starts_session_and_installs_intercept(self): + with setup(name="my-agent", base_url=self.base_url) as client: + self.assertIsInstance(client, ExplicitClient) + self.assertEqual(_session_id.get(), "s-1") + self.assertEqual(_timeline_id.get(), "tl-1") + self.assertTrue(is_installed()) + + # Cleanup happened. + self.assertIsNone(_session_id.get()) + self.assertFalse(is_installed()) + self.assertEqual(len(_MockHandler.sessions_started), 1) + self.assertEqual(_MockHandler.sessions_started[0]["name"], "my-agent") + self.assertEqual(len(_MockHandler.sessions_ended), 1) + + def test_does_not_uninstall_intercept_if_already_installed(self): + # Simulate: operator already called intercept.install() at startup. + from rewind_agent.intercept import install as intercept_install + intercept_install() + try: + with setup(name="reentrant", base_url=self.base_url): + self.assertTrue(is_installed()) + # We were not the installer, so we must NOT have uninstalled. + self.assertTrue(is_installed()) + finally: + uninstall() + + def test_propagates_thread_id_and_metadata(self): + with setup( + name="threaded", + base_url=self.base_url, + thread_id="conv-42", + metadata={"app": "test"}, + ): + pass + body = _MockHandler.sessions_started[0] + self.assertEqual(body.get("thread_id"), "conv-42") + self.assertEqual(body.get("metadata"), {"app": "test"}) + + +class TestReplayDispatch(_ConnectorTestBase): + def test_replay_env_skips_session_start(self): + with mock.patch.dict( + os.environ, + { + "REWIND_SESSION_ID": "s-replay", + "REWIND_REPLAY_CONTEXT_ID": "ctx-replay", + "REWIND_REPLAY_CONTEXT_TIMELINE_ID": "tl-fork", + "REWIND_URL": self.base_url, + }, + clear=False, + ): + self.assertTrue(_is_replay_dispatch()) + with setup(name="replay-handler", base_url=self.base_url) as client: + self.assertIsInstance(client, ExplicitClient) + self.assertTrue(is_installed()) + # No phantom session was created on /api/sessions/start. + self.assertEqual(_MockHandler.sessions_started, []) + self.assertEqual(_MockHandler.sessions_ended, []) + + def test_partial_replay_env_does_not_trigger_replay_path(self): + # Only one of the two required vars set → not a replay dispatch. + with mock.patch.dict( + os.environ, + {"REWIND_SESSION_ID": "s-only"}, + clear=False, + ), mock.patch.dict(os.environ, {"REWIND_REPLAY_CONTEXT_ID": ""}, clear=False): + os.environ.pop("REWIND_REPLAY_CONTEXT_ID", None) + self.assertFalse(_is_replay_dispatch()) + + +class TestHostPredicates(_ConnectorTestBase): + def test_hosts_from_kwarg(self): + captured = {} + + def fake_install(predicates=None): + captured["predicates"] = predicates + + with mock.patch("rewind_agent.connector.install", side_effect=fake_install), \ + mock.patch("rewind_agent.connector.uninstall"), \ + mock.patch("rewind_agent.connector.is_installed", return_value=False): + with setup( + name="custom", + base_url=self.base_url, + llm_hosts=("llm-gateway.example.com",), + ): + pass + + preds = captured["predicates"] + self.assertIsInstance(preds, _HostPredicates) + self.assertEqual(preds._hosts, ("llm-gateway.example.com",)) + + def test_hosts_from_env(self): + captured = {} + + def fake_install(predicates=None): + captured["predicates"] = predicates + + with mock.patch.dict( + os.environ, + {"REWIND_LLM_HOSTS": "a.example,b.example , ,c.example"}, + clear=False, + ), mock.patch("rewind_agent.connector.install", side_effect=fake_install), \ + mock.patch("rewind_agent.connector.uninstall"), \ + mock.patch("rewind_agent.connector.is_installed", return_value=False): + with setup(name="env-hosts", base_url=self.base_url): + pass + + preds = captured["predicates"] + self.assertIsInstance(preds, _HostPredicates) + # Empty entries dropped, surrounding whitespace stripped. + self.assertEqual( + preds._hosts, + ("a.example", "b.example", "c.example"), + ) + + def test_no_hosts_uses_default_predicates(self): + captured = {} + + def fake_install(predicates=None): + captured["predicates"] = predicates + + with mock.patch.dict(os.environ, {"REWIND_LLM_HOSTS": ""}, clear=False), \ + mock.patch("rewind_agent.connector.install", side_effect=fake_install), \ + mock.patch("rewind_agent.connector.uninstall"), \ + mock.patch("rewind_agent.connector.is_installed", return_value=False): + with setup(name="defaults", base_url=self.base_url): + pass + + # No custom hosts → no _HostPredicates wrapper; install gets None + # (DefaultPredicates is applied by intercept itself). + self.assertIsNone(captured["predicates"]) + + def test_predicate_matches_substring(self): + preds = _HostPredicates(("internal-gateway",)) + + class FakeReq: + class _Parts: + netloc = "private-internal-gateway.example.com" + + url_parts = _Parts() + + self.assertTrue(preds.is_llm_call(FakeReq())) + + +class TestPublicExport(unittest.TestCase): + def test_module_attribute(self): + self.assertTrue(hasattr(rewind_agent, "connector")) + self.assertTrue(callable(rewind_agent.connector.setup)) + + +if __name__ == "__main__": + unittest.main()