diff --git a/python/pyproject.toml b/python/pyproject.toml index 3b2539f..80c0f13 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "rewind-agent" -version = "0.16.1" +version = "0.17.0" 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 0d95131..05815d3 100644 --- a/python/rewind_agent/__init__.py +++ b/python/rewind_agent/__init__.py @@ -34,8 +34,18 @@ def search(query): wrap_crew, ) from .cached_call import cached_llm_call -from .explicit import ExplicitClient, RewindReplayDivergenceError +from .explicit import ( + ExplicitClient, + RewindReplayDivergenceError, + RewindServerError, + StepNotFoundError, + StepResponse, + cached_tool, + get_default_client, + set_default_client, +) from . import connector +from .intercept import DefaultPredicates, Predicates from .assertions import Assertions, AssertionResult from .openai_agents import openai_agents_hooks from .pydantic_ai import pydantic_ai_hooks @@ -99,6 +109,17 @@ def search(query): "cached_llm_call", # One-call connector for any agent (see docs/hdk.md) "connector", + # Intercept predicate types (for `connector.setup(predicates=…)`) + "Predicates", + "DefaultPredicates", + # Default-client discovery + module-level cached_tool decorator + "set_default_client", + "get_default_client", + "cached_tool", + # Public step-fetch helper (Phase 0 commit 3) + "StepResponse", + "StepNotFoundError", + "RewindServerError", ] @@ -120,4 +141,4 @@ def import_from_langfuse(trace_id: str, **kwargs) -> str: return _import(trace_id, **kwargs) -__version__ = "0.16.1" +__version__ = "0.17.0" diff --git a/python/rewind_agent/connector.py b/python/rewind_agent/connector.py index ceab305..7288cd9 100644 --- a/python/rewind_agent/connector.py +++ b/python/rewind_agent/connector.py @@ -70,9 +70,14 @@ from contextlib import contextmanager from typing import Iterator, Sequence -from rewind_agent.explicit import ExplicitClient +from rewind_agent.explicit import ( + ExplicitClient, + get_default_client, + set_default_client, +) from rewind_agent.intercept import ( DefaultPredicates, + Predicates, install, is_installed, uninstall, @@ -146,6 +151,7 @@ def setup( *, base_url: str | None = None, llm_hosts: Sequence[str] | None = None, + predicates: Predicates | None = None, enabled: bool | None = None, thread_id: str | None = None, metadata: dict | None = None, @@ -153,8 +159,9 @@ def setup( """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. + when ``llm_hosts`` is set or ``predicates`` is provided), yields the + :class:`ExplicitClient` for use inside the block, and tears both + down on exit. Parameters ---------- @@ -168,19 +175,57 @@ def setup( Sequence of hostnames to treat as LLM gateways. ``None`` (default) reads ``$REWIND_LLM_HOSTS``; empty / unset falls through to intercept's strict-by-default provider list. + Mutually exclusive with ``predicates``. + predicates: + Fully custom :class:`~rewind_agent.intercept.Predicates` + instance forwarded directly to :func:`intercept.install`. Use + when hostname-substring matching is not enough (e.g. matching + only specific path prefixes, or routing decisions that depend + on headers). Mutually exclusive with ``llm_hosts`` — passing + both raises :class:`ValueError` rather than silently picking + a winner. enabled: ``None`` (default) reads ``$REWIND_ENABLED`` (any value other than ``"0"`` is on); ``True`` forces on; ``False`` forces off. When off, ``setup()`` is a true no-op — yields ``None``, no HTTP, no install. thread_id, metadata: - Forwarded to :meth:`ExplicitClient.session`. + Forwarded to :meth:`ExplicitClient.session`. Ignored on the + replay-dispatch path — when ``REWIND_SESSION_ID`` etc. are set, + ``setup()`` attaches to the existing session instead of starting + a new one, so per-session metadata has no effect. Yields ------ ExplicitClient | None The recording client, or ``None`` when disabled. + + Raises + ------ + ValueError + When both ``predicates=`` and ``llm_hosts=`` are provided. + TypeError + When ``predicates=`` is not a :class:`~rewind_agent.intercept.Predicates` + instance (catches the common typo of passing a callable, a string, + or a list of hostnames). """ + if predicates is not None and llm_hosts is not None: + raise ValueError( + "setup() accepts either `predicates=` or `llm_hosts=`, not both. " + "Use `predicates=` for fully custom matching; use `llm_hosts=` " + "for the hostname-substring shortcut." + ) + if predicates is not None and not isinstance(predicates, Predicates): + # Boundary check parity with set_default_client(): catches the + # common typos of passing a callable, a string, or a list of + # hostnames where a Predicates instance was expected. Predicates + # is a runtime_checkable Protocol, so duck-typed instances are + # accepted. + raise TypeError( + f"setup(predicates=...) expected a Predicates instance, " + f"got {type(predicates).__name__}" + ) + if not _enabled(enabled): yield None return @@ -188,27 +233,39 @@ def setup( # base_url resolution lives in ExplicitClient.__init__ so all callers # share a single source of truth (kwarg > $REWIND_URL > localhost). client = ExplicitClient(base_url=base_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() + if predicates is None: + hosts = _resolve_hosts(llm_hosts) + predicates = _HostPredicates(hosts) if hosts else None + + # Stack-semantics for the default-client binding: save the previous + # value (which may be a different client or None), bind ours for the + # duration of the block, restore on exit — even if install() or + # client.session().__enter__ raises. The outer try/finally below + # guarantees the module-global never stays polluted across a setup() + # failure. + previous_default = get_default_client() + set_default_client(client) + try: + 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() + finally: + set_default_client(previous_default) diff --git a/python/rewind_agent/explicit.py b/python/rewind_agent/explicit.py index 8c7c14e..8f9f3ce 100644 --- a/python/rewind_agent/explicit.py +++ b/python/rewind_agent/explicit.py @@ -34,10 +34,14 @@ def get_pods(cluster: str) -> str: import json import logging import os +import threading import time import urllib.error +import urllib.parse import urllib.request +import weakref from contextlib import asynccontextmanager, contextmanager +from dataclasses import dataclass, field from typing import Any, Callable logger = logging.getLogger("rewind.explicit") @@ -58,6 +62,46 @@ def get_pods(cluster: str) -> str: _SESSION_CACHE_TTL = 7200 # 2 hours +class RewindServerError(RuntimeError): + """Raised by :meth:`ExplicitClient.get_step` when a transport or + server-side failure prevents step retrieval. + + Distinct from :class:`StepNotFoundError` so callers (e.g. replay + handlers) can decide whether to retry transient infra failures vs + treat the step as genuinely absent. The existing ``_post`` / ``_get`` + helpers swallow such failures to ``None`` for legacy reasons; the + new public step-fetch path opts into explicit propagation. + """ + + +class StepNotFoundError(LookupError): + """Raised by :meth:`ExplicitClient.get_step` ONLY when the + requested ``step_number`` does not exist on the resolved timeline. + + Transport / server failures raise :class:`RewindServerError` instead, + so the two cases are distinguishable by exception type. + """ + + +@dataclass(frozen=True) +class StepResponse: + """Typed view of a single recorded step (LLM call, tool call, etc.). + + Returned by :meth:`ExplicitClient.get_step` / + :meth:`ExplicitClient.get_step_sync`. Replay handlers use this to + extract the recorded request/response without parsing untyped JSON + blobs by hand. + """ + + step_number: int + step_type: str + request_body: Any | None = None + response_body: Any | None = None + model: str | None = None + tool_name: str | None = None + raw: dict = field(default_factory=dict) + + class RewindReplayDivergenceError(RuntimeError): """Phase 1 (Santa #4): strict-match replay lookup returned HTTP 409. @@ -151,6 +195,36 @@ def _get(self, path: str) -> dict | list | None: logger.debug("Rewind GET %s failed: %s", path, e) return None + def _get_or_raise(self, path: str) -> dict | list: + """Variant of ``_get`` for paths where the caller wants to + distinguish transport failure from a successful empty response. + + Raises :class:`RewindServerError` on network errors, non-2xx + responses, or invalid JSON. Returns the parsed body on success. + + The existing ``_get`` helper swallows everything to ``None`` for + legacy reasons (recording paths must never crash the agent). + Public read APIs that need crisp error semantics opt in here + rather than reverse-engineering the silent-failure path. + """ + url = f"{self.base_url}/api{path}" + req = urllib.request.Request(url, method="GET") + try: + with urllib.request.urlopen(req, timeout=_TIMEOUT) as resp: + return json.loads(resp.read()) + except urllib.error.HTTPError as e: + raise RewindServerError( + f"Rewind GET {path} returned {e.code}: {e.reason}" + ) from e + except urllib.error.URLError as e: + raise RewindServerError( + f"Rewind GET {path} transport error: {e.reason}" + ) from e + except (json.JSONDecodeError, OSError, TimeoutError) as e: + raise RewindServerError( + f"Rewind GET {path} failed: {e}" + ) from e + # ── Session lifecycle ────────────────────────────────────── @contextmanager @@ -677,6 +751,111 @@ def fork(self, session_id: str, *, at_step: int, label: str, result = self._post(f"/sessions/{session_id}/fork", body) return result["fork_timeline_id"] if result else None + # ── Step fetch ──────────────────────────────────────────── + + def get_step_sync( + self, + session_id: str, + *, + step_number: int, + timeline_id: str | None = None, + ) -> StepResponse: + """Fetch a single recorded step by number (sync). + + When ``timeline_id`` is omitted, resolves the session's root + timeline. Replay handlers typically pass the explicit timeline + from the dispatch payload so forks resolve correctly. + + Raises: + StepNotFoundError: when the requested ``step_number`` does + not exist on the resolved timeline, OR when ``session_id`` + resolves to no timelines at all (unknown / empty session). + These are the two true "absence" cases. + RewindServerError: when the rewind server is unreachable, + returns a non-2xx response, returns malformed JSON, or + returns a non-empty timelines list with no root timeline + (server data inconsistency). Replay handlers should retry + on this; downstream sf-rewind style consumers should NOT + swallow it. + """ + # Quote both path components: session_id is opaque caller input + # and timeline_id can in principle contain reserved characters. + # Existing legacy paths leave them raw — that's an SDK-wide + # consistency drift we're starting to walk back; this new public + # helper is built right. + quoted_sid = urllib.parse.quote(session_id, safe="") + + tid = timeline_id + if tid is None: + timelines = self._get_or_raise(f"/sessions/{quoted_sid}/timelines") + if not isinstance(timelines, list): + raise RewindServerError( + f"Rewind GET /sessions/{session_id}/timelines returned " + f"non-list body: {type(timelines).__name__}" + ) + if not timelines: + # Empty list = unknown / freshly-empty session. True absence. + raise StepNotFoundError( + f"Session {session_id} has no timelines" + ) + root = next( + (t for t in timelines if t.get("parent_timeline_id") is None), + None, + ) + if root is None: + # Non-empty list with no root entry = server data + # inconsistency, NOT a "step doesn't exist" case. + raise RewindServerError( + f"Server returned {len(timelines)} timelines for session " + f"{session_id} but none with parent_timeline_id=None" + ) + tid = root["id"] + + quoted_tid = urllib.parse.quote(tid, safe="") + path = ( + f"/sessions/{quoted_sid}/steps?" + f"timeline={quoted_tid}&include_blobs=1" + ) + steps = self._get_or_raise(path) + if not isinstance(steps, list): + raise RewindServerError( + f"Rewind GET {path} returned non-list body: " + f"{type(steps).__name__}" + ) + for s in steps: + if s.get("step_number") == step_number: + return StepResponse( + step_number=s["step_number"], + step_type=s.get("step_type", ""), + request_body=s.get("request_body"), + response_body=s.get("response_body"), + model=s.get("model"), + tool_name=s.get("tool_name"), + raw=s, + ) + raise StepNotFoundError( + f"Step {step_number} not found on timeline {tid} of session {session_id}" + ) + + async def get_step( + self, + session_id: str, + *, + step_number: int, + timeline_id: str | None = None, + ) -> StepResponse: + """Async variant of :meth:`get_step_sync` — runs the HTTP call + in a thread executor to avoid blocking the event loop.""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self.get_step_sync( + session_id, + step_number=step_number, + timeline_id=timeline_id, + ), + ) + # ── Cached tool decorator ───────────────────────────────── def cached_tool(self, name: str | None = None): @@ -759,6 +938,152 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: return decorator +# --------------------------------------------------------------------------- # +# Default-client discovery (Phase 0 commit 2). +# +# A blessed module-level handle so wrappers (e.g. sf-rewind) can let +# decorators find an active recording client at call time without inventing +# their own module global. Accepted trade-off documented in the design plan: +# this is a plain module attribute (not a ContextVar) for simplicity. Nested +# `connector.setup()` blocks use stack-semantics (entry saves prior, exit +# restores) but per-asyncio-task multi-client topologies are NOT supported — +# real consumers run a single bootstrap. Revisit if a real workload needs +# multi-sidecar in one process. +# --------------------------------------------------------------------------- # + +_default_client: "ExplicitClient | None" = None +_default_client_lock = threading.Lock() + + +def set_default_client(client: "ExplicitClient | None") -> None: + """Bind the process-wide default :class:`ExplicitClient`. + + Call this once at app startup (after constructing the client) so that + the module-level :func:`cached_tool` decorator can find it at call + time. Pass ``None`` to clear. + + Threading / async semantics + --------------------------- + The binding is **process-global**, NOT per-thread or per-asyncio task. + Reads and writes are serialized by an internal :class:`threading.Lock`, + so individual ``set_default_client`` / ``get_default_client`` calls are + atomic. However, the read-modify-write pattern in + :func:`rewind_agent.connector.setup` (save previous → bind ours → + restore on exit) is NOT atomic across concurrent ``setup()`` blocks + in different threads: thread B can observe thread A's client as + its "previous" and then on exit restore A's client even though A + has already exited. Single-threaded or single-event-loop consumers + are safe; multi-threaded multi-bootstrap consumers should use + :meth:`ExplicitClient.cached_tool` with an explicit client instance. + + Raises: + TypeError: when ``client`` is neither an :class:`ExplicitClient` + nor ``None``. Catches the common typo of passing a base-URL + string instead of a client instance. + """ + global _default_client + if client is not None and not isinstance(client, ExplicitClient): + raise TypeError( + f"set_default_client expected ExplicitClient or None, got {type(client).__name__}" + ) + with _default_client_lock: + _default_client = client + + +def get_default_client() -> "ExplicitClient | None": + """Return the currently-bound default client, or ``None``. + + See :func:`set_default_client` for the threading contract. + """ + with _default_client_lock: + return _default_client + + +# Per-(client, func) wrapper cache for module-level cached_tool. Keyed +# weakly on the client so wrappers don't keep the client alive past its +# normal lifetime. The inner dict is keyed by id(func) (functions don't +# have a meaningful weak-ref story for module-level decorators); since +# decorated functions live for the process lifetime in practice, the +# inner dict not collecting is fine. +_module_cached_tool_wrappers: "weakref.WeakKeyDictionary[ExplicitClient, dict[int, Callable[..., Any]]]" = weakref.WeakKeyDictionary() +_module_cached_tool_lock = threading.Lock() + + +def _resolve_module_cached_wrapper( + client: "ExplicitClient", + func: Callable, + tool_name: str, +) -> Callable[..., Any]: + """Return a stable wrapper for (client, func, tool_name); built once + per (client, func) pair on first call, reused thereafter.""" + func_key = id(func) + with _module_cached_tool_lock: + per_client = _module_cached_tool_wrappers.get(client) + if per_client is None: + per_client = {} + _module_cached_tool_wrappers[client] = per_client + wrapper = per_client.get(func_key) + if wrapper is None: + wrapper = client.cached_tool(tool_name)(func) + per_client[func_key] = wrapper + return wrapper + + +def cached_tool(name: str | None = None): + """Module-level :func:`ExplicitClient.cached_tool` that lazy-resolves + the default client at call time. + + Decorate at import time: + + from rewind_agent import cached_tool + + @cached_tool("list_clusters") + async def list_clusters(...): ... + + Resolution happens *each call* via :func:`get_default_client`. If no + default client is bound when the function is called, the decorated + function still runs — it just isn't recorded. This keeps imports safe + at module load before app startup has bound a client. + + Per-call cost is a single dict lookup: the underlying recording + wrapper is built once per ``(client, func)`` pair and cached in a + :class:`weakref.WeakKeyDictionary` so it dies with the client. + + Threading / async note + ---------------------- + The default-client binding is **process-global**, not per-thread or + per-asyncio task. See :func:`set_default_client` for the full + contract; in particular, concurrent + :func:`rewind_agent.connector.setup` blocks in different threads can + corrupt the stack-restore. If multiple threads or tasks need + different clients, use :meth:`ExplicitClient.cached_tool` directly + with an explicit client instance. + """ + def decorator(func: Callable) -> Callable: + tool_name = name or func.__name__ + + if inspect.iscoroutinefunction(func): + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + client = get_default_client() + if client is None: + return await func(*args, **kwargs) + inner = _resolve_module_cached_wrapper(client, func, tool_name) + return await inner(*args, **kwargs) + return async_wrapper + else: + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + client = get_default_client() + if client is None: + return func(*args, **kwargs) + inner = _resolve_module_cached_wrapper(client, func, tool_name) + return inner(*args, **kwargs) + return sync_wrapper + + return decorator + + def _serialize_args(args: tuple, kwargs: dict) -> dict: """Convert function args to a JSON-serializable dict.""" result: dict[str, Any] = {} diff --git a/python/rewind_agent/testing/__init__.py b/python/rewind_agent/testing/__init__.py new file mode 100644 index 0000000..7566d70 --- /dev/null +++ b/python/rewind_agent/testing/__init__.py @@ -0,0 +1,33 @@ +"""Public testing utilities for downstream packages. + +This module exposes the helpers downstream packages (e.g. the planned +``sf-rewind`` connector) use to write tests against Rewind's recording +and replay APIs without standing up a real server. + +Stability: symbols listed in ``__all__`` follow normal SDK semver — +breaking changes require a minor bump in 0.x. Symbols under +``rewind_agent.testing._unstable`` are explicitly NOT covered by that +commitment and may break in any release. + +Example +------- + +>>> from rewind_agent import ExplicitClient +>>> from rewind_agent.testing import StubRewindServer +>>> +>>> with StubRewindServer() as server: +... client = ExplicitClient(server.base_url) +... with client.session("my-test"): +... client.record_tool_call("ping", {}, "ok", duration_ms=1) +... assert len(server.recorded_steps) == 1 +""" + +from rewind_agent.testing._stub_server import StubRewindServer +from rewind_agent.testing._dispatch import make_dispatch_payload +from rewind_agent.testing._wait import wait_for_session + +__all__ = [ + "StubRewindServer", + "make_dispatch_payload", + "wait_for_session", +] diff --git a/python/rewind_agent/testing/_dispatch.py b/python/rewind_agent/testing/_dispatch.py new file mode 100644 index 0000000..a1641db --- /dev/null +++ b/python/rewind_agent/testing/_dispatch.py @@ -0,0 +1,44 @@ +"""``make_dispatch_payload`` — build a runner-compatible +:class:`DispatchPayload` for tests. + +Public API: imported via ``from rewind_agent.testing import make_dispatch_payload``. +""" + +from __future__ import annotations + +from rewind_agent.runner import DispatchPayload + + +def make_dispatch_payload( + *, + session_id: str, + job_id: str = "job-test", + replay_context_id: str = "ctx-test", + replay_context_timeline_id: str = "tl-fork", + source_timeline_id: str | None = None, + base_url: str = "http://127.0.0.1:4800", + at_step: int = 1, + dispatch_token: str = "tok-test", +) -> DispatchPayload: + """Build a :class:`DispatchPayload` with sensible test defaults. + + All fields default to non-empty placeholders so a single + ``make_dispatch_payload(session_id="s")`` call produces a valid + payload. Override individually when a test needs to pin a specific + value (e.g. ``at_step=5`` to exercise mid-turn replay). + + By convention, ``source_timeline_id`` defaults to + ``replay_context_timeline_id`` (the ``ReuseContext`` shape — the + common case in tests). Set them differently to exercise the + ``CreateAndDispatch`` shape where edits live on the source timeline. + """ + return DispatchPayload( + job_id=job_id, + session_id=session_id, + replay_context_id=replay_context_id, + replay_context_timeline_id=replay_context_timeline_id, + source_timeline_id=source_timeline_id or replay_context_timeline_id, + base_url=base_url, + at_step=at_step, + dispatch_token=dispatch_token, + ) diff --git a/python/rewind_agent/testing/_stub_server.py b/python/rewind_agent/testing/_stub_server.py new file mode 100644 index 0000000..223b730 --- /dev/null +++ b/python/rewind_agent/testing/_stub_server.py @@ -0,0 +1,225 @@ +"""In-process stub of the Rewind HTTP API. + +Public API: :class:`StubRewindServer`. Imported via +``from rewind_agent.testing import StubRewindServer``. + +This stub serves the subset of routes that recording + replay handlers +hit during unit tests: + +* ``POST /api/sessions/start`` — creates a session, returns + ``{session_id, root_timeline_id}``. +* ``POST /api/sessions/{id}/end`` — finalizes a session. +* ``POST /api/sessions/{id}/llm-calls`` — records an LLM call. +* ``POST /api/sessions/{id}/tool-calls`` — records a tool call. +* ``GET /api/sessions`` — lists all started sessions. +* ``GET /api/sessions/{id}/timelines`` — returns a single root timeline. +* ``GET /api/sessions/{id}/steps`` — returns recorded steps. + +Anything beyond this minimal surface belongs in +:mod:`rewind_agent.testing._unstable`. +""" + +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any + + +class _StubHandler(BaseHTTPRequestHandler): + server: "_StubServer" # type: ignore[assignment] + + def do_POST(self) -> None: # noqa: N802 + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length)) if length else {} + owner = self.server # type: ignore[attr-defined] + + if self.path == "/api/sessions/start": + # Decide the response inside the lock, then write the socket + # outside it. Both branches now follow the same pattern, so a + # slow socket write never blocks other handlers waiting on + # the lock. + status = 201 + with owner.lock: + key = body.get("client_session_key") + if key and key in owner.sessions_by_client_key: + sid, tid = owner.sessions_by_client_key[key] + status = 200 + else: + idx = len(owner.sessions) + 1 + sid = f"s-{idx}" + tid = f"tl-{idx}" + owner.sessions[sid] = { + "session_id": sid, + "name": body.get("name", ""), + "root_timeline_id": tid, + "metadata": body.get("metadata", {}), + "thread_id": body.get("thread_id"), + "ended": False, + } + if key: + owner.sessions_by_client_key[key] = (sid, tid) + self._json(status, {"session_id": sid, "root_timeline_id": tid}) + return + + if self.path.endswith("/end"): + sid = self.path.split("/")[3] + with owner.lock: + session = owner.sessions.get(sid) + if session is not None: + session["ended"] = True + self._json(200, {"session_id": sid}) + return + + if self.path.endswith("/llm-calls") and "replay-lookup" not in self.path: + sid = self.path.split("/")[3] + with owner.lock: + # Match real-server semantics: step_number is per-session, + # not global. Tests authored against this stub can rely on + # `step_number == 1` being the first call within a session + # regardless of how many other sessions have recorded. + step_number = sum( + 1 for s in owner.recorded_steps + if s.get("session_id") == sid + ) + 1 + owner.recorded_steps.append({ + "session_id": sid, + "step_number": step_number, + "step_type": "llm_call", + "model": body.get("model"), + "request_body": body.get("request_body"), + "response_body": body.get("response_body"), + "tool_name": None, + }) + self._json(201, {"step_number": step_number}) + return + + if self.path.endswith("/tool-calls") and "replay-lookup" not in self.path: + sid = self.path.split("/")[3] + with owner.lock: + step_number = sum( + 1 for s in owner.recorded_steps + if s.get("session_id") == sid + ) + 1 + owner.recorded_steps.append({ + "session_id": sid, + "step_number": step_number, + "step_type": "tool_call", + "tool_name": body.get("tool_name"), + "request_body": body.get("request_body"), + "response_body": body.get("response_body"), + "model": None, + }) + self._json(201, {"step_number": step_number}) + return + + if "replay-lookup" in self.path: + self._json(200, {"hit": False}) + return + + self._json(404, {"error": f"unhandled POST {self.path}"}) + + def do_GET(self) -> None: # noqa: N802 + owner = self.server # type: ignore[attr-defined] + + if self.path == "/api/sessions": + with owner.lock: + listing = list(owner.sessions.values()) + self._json(200, listing) + return + + if "/timelines" in self.path: + sid = self.path.split("/")[3] + with owner.lock: + session = owner.sessions.get(sid) + if session is None: + self._json(404, {"error": "session not found"}) + return + self._json(200, [ + { + "id": session["root_timeline_id"], + "session_id": sid, + "parent_timeline_id": None, + } + ]) + return + + if "/steps" in self.path: + sid = self.path.split("/")[3] + with owner.lock: + steps = [ + s for s in owner.recorded_steps + if s.get("session_id") == sid + ] + self._json(200, steps) + return + + self._json(404, {"error": f"unhandled GET {self.path}"}) + + def _json(self, status: int, payload: Any) -> None: + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *_args, **_kwargs) -> None: # silence + pass + + +class _StubServer(HTTPServer): + """:class:`HTTPServer` subclass that carries the recording state.""" + + def __init__(self, addr) -> None: + super().__init__(addr, _StubHandler) + self.lock = threading.Lock() + self.sessions: dict[str, dict[str, Any]] = {} + self.sessions_by_client_key: dict[str, tuple[str, str]] = {} + self.recorded_steps: list[dict[str, Any]] = [] + + +class StubRewindServer: + """In-process stub of the Rewind HTTP API for unit tests. + + Bind on a random localhost port. Use as a context manager: + + with StubRewindServer() as server: + client = ExplicitClient(server.base_url) + ... + + Inspectable state: + * ``server.base_url`` — pass to :class:`ExplicitClient`. + * ``server.recorded_steps`` — list of step dicts as recorded. + * ``server.sessions`` — mapping of ``session_id`` to session dict. + """ + + def __init__(self) -> None: + self._server = _StubServer(("127.0.0.1", 0)) + self.base_url = f"http://127.0.0.1:{self._server.server_address[1]}" + self._thread = threading.Thread( + target=self._server.serve_forever, + daemon=True, + ) + + def __enter__(self) -> "StubRewindServer": + self._thread.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self._server.shutdown() + self._server.server_close() + # Join the worker thread so a test asserting "port is free / + # serve_forever has returned" right after exit doesn't race. + # serve_forever has already returned by the time shutdown() is + # done, so the join is bounded; the timeout is defensive. + self._thread.join(timeout=2.0) + + @property + def recorded_steps(self) -> list[dict[str, Any]]: + return self._server.recorded_steps + + @property + def sessions(self) -> dict[str, dict[str, Any]]: + return self._server.sessions diff --git a/python/rewind_agent/testing/_unstable/__init__.py b/python/rewind_agent/testing/_unstable/__init__.py new file mode 100644 index 0000000..498f92f --- /dev/null +++ b/python/rewind_agent/testing/_unstable/__init__.py @@ -0,0 +1,17 @@ +"""Explicitly unstable testing helpers. + +Symbols under :mod:`rewind_agent.testing._unstable` are NOT covered by +the SDK's semver commitment. They may be renamed, restructured, or +removed in any 0.x release without warning. Use them only inside the +SDK's own tests, or accept that downstream tests pinning to them will +break on upgrade. + +If a downstream consumer needs a stable version of something in here, +file an issue or a PR — promotion to :mod:`rewind_agent.testing` is a +deliberate decision, not an accident. + +This package is intentionally empty in v0.17 — it exists as the boundary +marker so the stability commitment of the parent module is unambiguous. +""" + +__all__: list[str] = [] diff --git a/python/rewind_agent/testing/_wait.py b/python/rewind_agent/testing/_wait.py new file mode 100644 index 0000000..1203bc1 --- /dev/null +++ b/python/rewind_agent/testing/_wait.py @@ -0,0 +1,46 @@ +"""``wait_for_session`` — poll a :class:`StubRewindServer` for a session +matching a name, used by integration-style tests. + +Public API: imported via ``from rewind_agent.testing import wait_for_session``. +""" + +from __future__ import annotations + +import time +from typing import Any + +from rewind_agent.testing._stub_server import StubRewindServer + + +def wait_for_session( + server: StubRewindServer, + *, + name: str, + timeout: float = 5.0, + poll_interval: float = 0.02, +) -> dict[str, Any]: + """Poll ``server`` until a session with ``name`` appears, or raise. + + Returns the session dict (the same shape :class:`StubRewindServer` records + in ``server.sessions``). + + Raises: + TimeoutError: when no matching session appears within ``timeout`` + seconds. Caller is responsible for naming sessions distinctly. + """ + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + # Snapshot under the stub's lock for symmetry with the rest of + # the stub. CPython's GIL would make the bare iteration safe in + # practice, but explicit locking keeps the test infrastructure + # honest about its concurrency contract. + with server._server.lock: # type: ignore[attr-defined] + snapshot = list(server.sessions.values()) + for session in snapshot: + if session.get("name") == name: + return session + time.sleep(poll_interval) + raise TimeoutError( + f"No session named {name!r} appeared within {timeout}s; " + f"recorded sessions: {[s.get('name') for s in server.sessions.values()]}" + ) diff --git a/python/tests/test_connector.py b/python/tests/test_connector.py index 099681e..01d15ff 100644 --- a/python/tests/test_connector.py +++ b/python/tests/test_connector.py @@ -14,7 +14,7 @@ _session_id, _timeline_id, ) -from rewind_agent.intercept import is_installed, uninstall +from rewind_agent.intercept import DefaultPredicates, is_installed, uninstall class _MockHandler(BaseHTTPRequestHandler): @@ -281,5 +281,177 @@ def test_module_attribute(self): self.assertTrue(callable(rewind_agent.connector.setup)) +class TestPredicatesKwarg(_ConnectorTestBase): + """`predicates=` is the structured alternative to the `llm_hosts=` shortcut. + + Lets callers pass a fully-custom Predicates instance (e.g. an SF gateway + predicate) without going through the hostname-substring path. Phase 0 + commit 1 of the public-helpers SDK PR. + """ + + def _capture_predicates(self, **setup_kwargs): + 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="capture", base_url=self.base_url, **setup_kwargs): + pass + return captured["predicates"] + + def test_passthrough_to_intercept_install(self): + from rewind_agent.intercept import DefaultPredicates + + sentinel = DefaultPredicates() + forwarded = self._capture_predicates(predicates=sentinel) + # The exact instance the caller passed must reach intercept.install, + # untouched — no wrapping in _HostPredicates. + self.assertIs(forwarded, sentinel) + + def test_custom_predicate_subclass_passes_through(self): + from rewind_agent.intercept import DefaultPredicates + + class _MyPreds(DefaultPredicates): + def is_llm_call(self, req) -> bool: # noqa: D401 + return "my-custom-gw" in req.url_parts.netloc.lower() + + custom = _MyPreds() + forwarded = self._capture_predicates(predicates=custom) + self.assertIs(forwarded, custom) + # The forwarded predicate is fully functional — caller's behavior is preserved. + self.assertTrue(forwarded.is_llm_call(_fake_req("my-custom-gw.internal"))) + self.assertFalse(forwarded.is_llm_call(_fake_req("unrelated.example"))) + + def test_predicates_and_llm_hosts_together_raises(self): + # Mutually exclusive: passing both is operator confusion and the + # silent-precedence outcome ("predicates wins, llm_hosts ignored") + # would mask the misconfiguration. Refuse explicitly. + with self.assertRaises(ValueError) as ctx: + with setup( + name="conflict", + base_url=self.base_url, + llm_hosts=("a.example",), + predicates=DefaultPredicates(), + ): + pass + self.assertIn("predicates", str(ctx.exception)) + self.assertIn("llm_hosts", str(ctx.exception)) + + def test_predicates_none_falls_through_to_llm_hosts_path(self): + # Sanity: omitting `predicates=` doesn't disturb the existing + # llm_hosts path. Hosts kwarg still produces a _HostPredicates. + forwarded = self._capture_predicates(llm_hosts=("a.example",)) + self.assertIsInstance(forwarded, _HostPredicates) + + def test_predicates_wrong_type_raises_type_error(self): + # Boundary check parity with set_default_client(): catches typos + # like passing a callable or a string. Predicates is a + # runtime_checkable Protocol, so a duck-typed object with the + # right methods would be accepted — but a string definitely + # shouldn't be. + with self.assertRaises(TypeError): + with setup(name="bad", base_url=self.base_url, predicates="oops"): # type: ignore[arg-type] + pass + with self.assertRaises(TypeError): + with setup(name="bad", base_url=self.base_url, predicates=lambda r: True): # type: ignore[arg-type] + pass + + +class TestDefaultClientLeakOnFailure(_ConnectorTestBase): + """Regression: setup() must restore the previous default client even + when install() or session().__enter__ raises. Without the outer + try/finally, a failure mid-setup leaves the module-global polluted + across the failure, poisoning all subsequent cached_tool() calls in + the process.""" + + def test_install_failure_restores_previous_default(self): + from rewind_agent.explicit import ( + ExplicitClient, + get_default_client, + set_default_client, + ) + + # Outer client representing an "always-on" baseline. + outer = ExplicitClient(self.base_url) + set_default_client(outer) + try: + self.assertIs(get_default_client(), outer) + + def boom(predicates=None): + raise RuntimeError("install blew up") + + with mock.patch("rewind_agent.connector.install", side_effect=boom): + with self.assertRaises(RuntimeError): + with setup(name="will-fail", base_url=self.base_url): + self.fail("setup() body must not be entered when install fails") + + # Default client is restored to the outer baseline, not left + # pointing at the half-initialized inner client. + self.assertIs(get_default_client(), outer) + finally: + set_default_client(None) + + def test_session_enter_failure_restores_previous_default(self): + """Round-2 santa-review gap: the previous test only covered + install() failure. session().__enter__ raises (e.g., POST + /sessions/start hangs / 5xx) is the other path the outer + try/finally must protect.""" + from rewind_agent.explicit import ( + ExplicitClient, + get_default_client, + set_default_client, + ) + + outer = ExplicitClient(self.base_url) + set_default_client(outer) + try: + self.assertIs(get_default_client(), outer) + + class _BoomSession: + def __enter__(self_inner): + raise RuntimeError("session start blew up") + + def __exit__(self_inner, *_): + return False + + # Make any newly-constructed ExplicitClient.session() raise on + # __enter__, but leave session_async (etc.) alone. This is + # the exact failure shape replay/recording would surface from + # an unreachable rewind sidecar at session-start time. + with mock.patch.object( + ExplicitClient, "session", lambda *a, **kw: _BoomSession() + ): + with self.assertRaises(RuntimeError): + with setup(name="will-fail-session", base_url=self.base_url): + self.fail("setup() body must not be entered when session.__enter__ fails") + + # Default client must be restored even though the failure + # happened inside `with client.session(...):` and the inner + # try/finally never reached its restore. + self.assertIs(get_default_client(), outer) + finally: + set_default_client(None) + + +class TestPredicatesPackageReExport(unittest.TestCase): + """The package root re-exports Predicates / DefaultPredicates so callers + that pass `predicates=` don't have to reach into the private intercept + package.""" + + def test_predicates_classes_at_package_root(self): + self.assertTrue(hasattr(rewind_agent, "Predicates")) + self.assertTrue(hasattr(rewind_agent, "DefaultPredicates")) + # The re-exports must be the *same* objects intercept exposes. + from rewind_agent.intercept import ( + DefaultPredicates as _DP, + Predicates as _P, + ) + self.assertIs(rewind_agent.Predicates, _P) + self.assertIs(rewind_agent.DefaultPredicates, _DP) + + if __name__ == "__main__": unittest.main() diff --git a/python/tests/test_explicit.py b/python/tests/test_explicit.py index ac7cb9d..617e25d 100644 --- a/python/tests/test_explicit.py +++ b/python/tests/test_explicit.py @@ -538,5 +538,530 @@ def test_kwarg_wins_over_env_var(self): self.assertEqual(client.base_url, "http://rewind.from-kwarg:7777") +class _StepAwareHandler(BaseHTTPRequestHandler): + """Mock that serves per-step lookups via the existing list endpoint. + + Mirrors the real Rust server: GET /api/sessions/{id}/steps?timeline=… + returns the full step list; with ``&include_blobs=1`` the response + bodies include ``request_body`` and ``response_body``. Per-step + fetch in commit 3 is built on top of this list endpoint (no new + Rust route — Python-only filter by step_number). + """ + + fixture_steps: list[dict] = [] + # Test instrumentation: lets tests assert which paths were touched. + paths_called: list[str] = [] + timelines_path_called: bool = False + # Override the /timelines response for tests that exercise edge + # cases (empty list / non-empty list with no root / etc.). None + # = use the default 2-timeline fixture. + timelines_override: list[dict] | None = None + + def do_GET(self): # noqa: N802 — stdlib API + _StepAwareHandler.paths_called.append(self.path) + if "/timelines" in self.path: + _StepAwareHandler.timelines_path_called = True + if _StepAwareHandler.timelines_override is not None: + self._respond(200, _StepAwareHandler.timelines_override) + return + self._respond(200, [ + {"id": "tl-root", "parent_timeline_id": None, "session_id": "s1"}, + {"id": "tl-fork", "parent_timeline_id": "tl-root", "session_id": "s1"}, + ]) + return + if "/steps" in self.path: + include_blobs = "include_blobs=1" in self.path + steps = [] + for s in _StepAwareHandler.fixture_steps: + copy = { + "step_number": s["step_number"], + "step_type": s["step_type"], + "model": s.get("model"), + "tool_name": s.get("tool_name"), + } + if include_blobs: + copy["request_body"] = s.get("request_body") + copy["response_body"] = s.get("response_body") + steps.append(copy) + self._respond(200, steps) + return + self._respond(404, {"error": "not found"}) + + def _respond(self, status, body): + 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, *_): + pass + + +class TestGetStep(unittest.TestCase): + """Tests for ExplicitClient.get_step / get_step_sync (Phase 0 commit 3). + + Replay handlers today reach into private SDK helpers or hit raw HTTP + to fetch step content. A typed public helper closes that gap. + """ + + @classmethod + def setUpClass(cls): + cls.server = HTTPServer(("127.0.0.1", 0), _StepAwareHandler) + cls.port = cls.server.server_address[1] + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + cls.client = ExplicitClient(f"http://127.0.0.1:{cls.port}") + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + + def setUp(self): + _StepAwareHandler.paths_called = [] + _StepAwareHandler.timelines_path_called = False + _StepAwareHandler.timelines_override = None + _StepAwareHandler.fixture_steps = [ + { + "step_number": 1, + "step_type": "llm_call", + "model": "gpt-4o", + "request_body": {"messages": [{"role": "user", "content": "hi"}]}, + "response_body": {"choices": [{"message": {"content": "hello"}}]}, + }, + { + "step_number": 2, + "step_type": "tool_call", + "tool_name": "get_pods", + "request_body": {"args": ["dev"]}, + "response_body": "pod-1\npod-2", + }, + { + "step_number": 3, + "step_type": "llm_call", + "model": "gpt-4o", + "request_body": {"messages": [{"role": "user", "content": "next"}]}, + "response_body": {"choices": [{"message": {"content": "ok"}}]}, + }, + ] + + def test_returns_typed_step_response(self): + from rewind_agent.explicit import StepResponse + + step = self.client.get_step_sync("s1", step_number=2) + + self.assertIsInstance(step, StepResponse) + self.assertEqual(step.step_number, 2) + self.assertEqual(step.step_type, "tool_call") + self.assertEqual(step.request_body, {"args": ["dev"]}) + self.assertEqual(step.response_body, "pod-1\npod-2") + + def test_returns_llm_call_with_request_body(self): + step = self.client.get_step_sync("s1", step_number=1) + self.assertEqual(step.step_type, "llm_call") + self.assertEqual(step.model, "gpt-4o") + self.assertEqual( + step.request_body, + {"messages": [{"role": "user", "content": "hi"}]}, + ) + + def test_unknown_step_raises_step_not_found(self): + from rewind_agent.explicit import StepNotFoundError + + with self.assertRaises(StepNotFoundError): + self.client.get_step_sync("s1", step_number=999) + + def test_explicit_timeline_id_passes_through(self): + # When `timeline_id` is provided we must NOT auto-resolve via + # /timelines — the caller's choice wins. Verify two things: + # (1) the /timelines endpoint was NOT hit; (2) the caller-supplied + # timeline_id reached the server in the /steps query string. + step = self.client.get_step_sync("s1", timeline_id="tl-fork", step_number=1) + self.assertEqual(step.step_number, 1) + self.assertFalse( + _StepAwareHandler.timelines_path_called, + "auto-resolve via /timelines must be skipped when caller passes timeline_id", + ) + steps_paths = [p for p in _StepAwareHandler.paths_called if "/steps" in p] + self.assertEqual(len(steps_paths), 1, f"expected one /steps call, got {steps_paths}") + self.assertIn( + "timeline=tl-fork", + steps_paths[0], + f"caller-supplied timeline_id missing from server URL: {steps_paths[0]}", + ) + + def test_default_timeline_resolution_hits_timelines_endpoint(self): + # Counterpart to the above: when timeline_id is omitted we DO hit + # /timelines once to find the root timeline. Locks the contract so + # a refactor that loses auto-resolve is caught. + self.client.get_step_sync("s1", step_number=1) + self.assertTrue( + _StepAwareHandler.timelines_path_called, + "auto-resolve via /timelines must run when caller omits timeline_id", + ) + steps_paths = [p for p in _StepAwareHandler.paths_called if "/steps" in p] + self.assertEqual(len(steps_paths), 1) + # Auto-resolved root timeline reached the server URL. + self.assertIn("timeline=tl-root", steps_paths[0]) + + def test_async_get_step(self): + from rewind_agent.explicit import StepResponse + + async def run(): + step = await self.client.get_step("s1", step_number=3) + self.assertIsInstance(step, StepResponse) + self.assertEqual(step.step_number, 3) + self.assertEqual(step.step_type, "llm_call") + + asyncio.run(run()) + + def test_step_response_is_immutable(self): + import dataclasses + + step = self.client.get_step_sync("s1", step_number=1) + # Frozen dataclass: assignment raises FrozenInstanceError specifically. + # Don't use bare Exception here — that would let unrelated regressions + # (e.g. AttributeError because StepResponse stops being a dataclass) + # mask the actual contract. + with self.assertRaises(dataclasses.FrozenInstanceError): + step.step_number = 999 # type: ignore[misc] + # And StepResponse is the public type used in __all__. + self.assertTrue(hasattr(rewind_agent_module, "StepResponse")) + + def test_package_root_re_exports(self): + import rewind_agent + self.assertTrue(hasattr(rewind_agent, "StepResponse")) + self.assertTrue(hasattr(rewind_agent, "StepNotFoundError")) + self.assertTrue(hasattr(rewind_agent, "RewindServerError")) + + def test_empty_timelines_raises_step_not_found(self): + """An empty timelines list = unknown / freshly-empty session. + That's true 'absence', not server data inconsistency, so it raises + StepNotFoundError (NOT RewindServerError).""" + from rewind_agent.explicit import StepNotFoundError + + _StepAwareHandler.timelines_override = [] + with self.assertRaises(StepNotFoundError): + self.client.get_step_sync("s1", step_number=1) + + def test_non_empty_timelines_with_no_root_raises_server_error(self): + """Non-empty list with NO entry having parent_timeline_id=None is + server data inconsistency. Per round-3 fix, this raises + RewindServerError so callers can distinguish it from real + absences and decide to retry / log / page.""" + from rewind_agent.explicit import RewindServerError + + _StepAwareHandler.timelines_override = [ + {"id": "tl-orphan-1", "parent_timeline_id": "tl-missing", "session_id": "s1"}, + {"id": "tl-orphan-2", "parent_timeline_id": "tl-missing", "session_id": "s1"}, + ] + with self.assertRaises(RewindServerError): + self.client.get_step_sync("s1", step_number=1) + + def test_session_id_is_url_quoted(self): + """Round-3 fix: session_id, like timeline_id, is URL-quoted on + the way out so reserved characters in opaque IDs don't break + the URL or open a path-traversal-shaped surface.""" + # Use a session id that contains URL-reserved characters. The + # mock handler ignores the session id for routing — it serves + # the same fixture regardless — so the test verifies what hit + # the wire, not server-side behavior. + weird_sid = "s/with?reserved#chars" + try: + self.client.get_step_sync(weird_sid, step_number=1) + except Exception: + pass # The mock 404s for unrecognized paths; we only care about the URL. + + # Both /timelines and /steps requests should have the session_id + # percent-encoded. + self.assertTrue( + _StepAwareHandler.paths_called, + "expected at least one GET to reach the server", + ) + for p in _StepAwareHandler.paths_called: + self.assertNotIn( + "s/with", + p, + f"raw session_id leaked into URL: {p}", + ) + self.assertIn( + "s%2Fwith%3Freserved%23chars", + p, + f"session_id was not properly URL-quoted in: {p}", + ) + + +class TestGetStepServerErrors(unittest.TestCase): + """Round-2 santa-review fix: get_step distinguishes 'step absent' + (StepNotFoundError) from 'transport / server failure' + (RewindServerError). Replay handlers can decide whether to retry + transient infra failures vs treat the step as genuinely missing.""" + + def test_unreachable_server_raises_rewind_server_error(self): + from rewind_agent.explicit import ( + ExplicitClient, + RewindServerError, + StepNotFoundError, + ) + + # Port 1 is reserved/closed on standard hosts — guaranteed + # connection refused. The exception type we want is + # RewindServerError, NOT StepNotFoundError. + client = ExplicitClient("http://127.0.0.1:1") + with self.assertRaises(RewindServerError): + client.get_step_sync("any-session", step_number=1) + + # Defensive: verify it's NOT StepNotFoundError (which would let + # callers silently swallow infra failures). + try: + client.get_step_sync("any-session", step_number=1) + except StepNotFoundError: + self.fail( + "get_step_sync masked transport failure as StepNotFoundError" + ) + except RewindServerError: + pass + + +class TestModuleCachedToolStableWrapper(unittest.TestCase): + """Round-2 santa-review fix: module-level cached_tool used to call + `client.cached_tool(name)(func)` on EVERY invocation, allocating a + fresh wrapper per call. After the fix, the inner wrapper is built + once per (client, func) pair and reused. Locks the contract so a + refactor that loses the cache is caught.""" + + def setUp(self): + from rewind_agent import explicit as mod + mod.set_default_client(None) + # Also clear the wrapper cache so test order doesn't matter. + mod._module_cached_tool_wrappers.clear() + _session_id.set(None) + _timeline_id.set(None) + _replay_context_id.set(None) + + def tearDown(self): + from rewind_agent import explicit as mod + mod.set_default_client(None) + + def test_call_path_uses_cached_wrapper_not_per_call_rebuild(self): + """Genuinely lock the perf claim: instrument ExplicitClient.cached_tool + to count invocations and assert it ran exactly ONCE across N calls + of the decorated function. A regression that reverts sync_wrapper / + async_wrapper to `client.cached_tool(name)(func)(*args)` per call + would show N invocations instead of 1.""" + from rewind_agent.explicit import ( + ExplicitClient, + cached_tool as module_cached_tool, + set_default_client, + ) + + @module_cached_tool("count_check") + def add(a: int, b: int) -> int: + return a + b + + server = HTTPServer(("127.0.0.1", 0), MockRewindHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + client = ExplicitClient(f"http://127.0.0.1:{port}") + set_default_client(client) + + # Wrap the real cached_tool method to count invocations. + real_cached_tool = client.cached_tool + invocations = {"count": 0} + + def counting_cached_tool(name): + invocations["count"] += 1 + return real_cached_tool(name) + + client.cached_tool = counting_cached_tool # type: ignore[method-assign] + + with client.session("call-path-cache-test"): + # Five calls of the decorated function. + for i in range(5): + self.assertEqual(add(i, i), 2 * i) + + # If sync_wrapper rebuilds the wrapper each call, + # invocations["count"] would be 5. With caching, exactly 1. + self.assertEqual( + invocations["count"], + 1, + f"client.cached_tool was invoked {invocations['count']} times " + f"across 5 add() calls; module-level cached_tool must build " + f"the wrapper once per (client, func) pair and reuse it.", + ) + finally: + server.shutdown() + + +# Module reference used by TestGetStep.test_step_response_is_immutable +import rewind_agent as rewind_agent_module # noqa: E402 + + +class TestDefaultClient(unittest.TestCase): + """Tests for module-level default-client discovery (Phase 0 commit 2). + + The blessed singleton lets wrapper libraries (like the planned sf-rewind) + bind a recording client at app startup and have a module-level + `cached_tool` decorator find it at call time, without re-implementing + the module-global pattern in every consumer. + """ + + def setUp(self): + # The default-client module attr is process-global. Reset before each + # test so leakage is impossible. + from rewind_agent import explicit as mod + mod.set_default_client(None) + _session_id.set(None) + _timeline_id.set(None) + _replay_context_id.set(None) + + def tearDown(self): + from rewind_agent import explicit as mod + mod.set_default_client(None) + + def test_get_returns_none_when_unset(self): + from rewind_agent.explicit import get_default_client + self.assertIsNone(get_default_client()) + + def test_set_and_get_round_trip(self): + from rewind_agent.explicit import ( + get_default_client, + set_default_client, + ) + client = ExplicitClient("http://127.0.0.1:1") + set_default_client(client) + self.assertIs(get_default_client(), client) + + def test_set_none_clears(self): + from rewind_agent.explicit import ( + get_default_client, + set_default_client, + ) + client = ExplicitClient("http://127.0.0.1:1") + set_default_client(client) + set_default_client(None) + self.assertIsNone(get_default_client()) + + def test_module_cached_tool_records_when_default_client_set(self): + """`@cached_tool` at module level decorates at import time but + resolves the active client at *call* time. This is the load-bearing + contract for sf-rewind's tools.py pattern.""" + # Decorate BEFORE setting the default client to lock the lazy-resolve + # contract. + from rewind_agent.explicit import cached_tool as module_cached_tool + + @module_cached_tool("noop_add") + def add(a: int, b: int) -> int: + return a + b + + server = HTTPServer(("127.0.0.1", 0), MockRewindHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + from rewind_agent.explicit import set_default_client + client = ExplicitClient(f"http://127.0.0.1:{port}") + set_default_client(client) + with client.session("default-client-test"): + result = add(2, 3) + self.assertEqual(result, 5) + # Recording happened against the bound default client. + self.assertTrue( + any(s.get("tool_name") == "noop_add" for s in MockRewindHandler.recorded_steps), + f"expected a noop_add tool_call step, got {MockRewindHandler.recorded_steps}", + ) + finally: + server.shutdown() + + def test_module_cached_tool_runs_unrecorded_when_no_default_client(self): + """If `@cached_tool` is decorated and called before + `set_default_client(...)` is invoked, the function still runs — it + just doesn't record. This keeps imports safe at module load.""" + from rewind_agent.explicit import cached_tool as module_cached_tool + + @module_cached_tool("unbound_add") + def add(a: int, b: int) -> int: + return a + b + + # No default client set; no session active. + result = add(2, 3) + self.assertEqual(result, 5) + + def test_module_cached_tool_async(self): + from rewind_agent.explicit import cached_tool as module_cached_tool + + @module_cached_tool("noop_aadd") + async def aadd(a: int, b: int) -> int: + return a + b + + server = HTTPServer(("127.0.0.1", 0), MockRewindHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + async def run(): + from rewind_agent.explicit import set_default_client + client = ExplicitClient(f"http://127.0.0.1:{port}") + set_default_client(client) + async with client.session_async("default-client-async-test"): + result = await aadd(2, 3) + self.assertEqual(result, 5) + + try: + asyncio.run(run()) + self.assertTrue( + any(s.get("tool_name") == "noop_aadd" for s in MockRewindHandler.recorded_steps), + f"expected a noop_aadd tool_call step, got {MockRewindHandler.recorded_steps}", + ) + finally: + server.shutdown() + + def test_setup_binds_default_client_inside_block(self): + """connector.setup() binds the default client on entry and restores + the previous value on exit. Stack semantics, not full ContextVar + isolation — accepted trade-off documented in the design plan.""" + from rewind_agent import connector + from rewind_agent.explicit import ( + get_default_client, + set_default_client, + ) + + server = HTTPServer(("127.0.0.1", 0), MockRewindHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + base_url = f"http://127.0.0.1:{port}" + + # Stack semantics: inner block restores outer client on exit. + outer = ExplicitClient(base_url) + set_default_client(outer) + self.assertIs(get_default_client(), outer) + + try: + with connector.setup(name="bind-test", base_url=base_url) as inner: + # Inside the block, the default client is the connector's + # client, NOT the outer one we set manually. + self.assertIsNotNone(inner) + self.assertIs(get_default_client(), inner) + # On exit, the outer client is restored. + self.assertIs(get_default_client(), outer) + finally: + server.shutdown() + + def test_set_default_client_rejects_wrong_type(self): + """Cheap type check at the API boundary. Catches the common typo of + passing a string URL instead of a constructed client.""" + from rewind_agent.explicit import set_default_client + with self.assertRaises(TypeError): + set_default_client("http://127.0.0.1:1") # type: ignore[arg-type] + + def test_package_root_re_exports(self): + import rewind_agent + self.assertTrue(hasattr(rewind_agent, "set_default_client")) + self.assertTrue(hasattr(rewind_agent, "get_default_client")) + self.assertTrue(hasattr(rewind_agent, "cached_tool")) + + if __name__ == "__main__": unittest.main() diff --git a/python/tests/test_testing_module.py b/python/tests/test_testing_module.py new file mode 100644 index 0000000..90c60f6 --- /dev/null +++ b/python/tests/test_testing_module.py @@ -0,0 +1,210 @@ +"""Tests for the public ``rewind_agent.testing`` module (Phase 0 commit 4). + +Promoting test infrastructure from ``python/tests/`` (private) to a public +module is a real semver commitment. These tests lock the public API contract: +import-by-name (catches accidental rename), lifecycle, payload shape, +``wait_for_session`` timeout semantics, and the ``_unstable`` boundary. +""" + +import threading +import time +import unittest +import urllib.error +import urllib.request + + +class TestPublicSurface(unittest.TestCase): + """Lock the names of the stable public symbols. + + If a refactor accidentally renames or removes one of these, this test + fires before downstream consumers (sf-rewind, integration smoke tests) + break. New stable symbols should be ADDED here, not silently introduced. + """ + + def test_stable_symbols_importable(self): + # Import via the documented module path. + from rewind_agent.testing import ( + StubRewindServer, + make_dispatch_payload, + wait_for_session, + ) + self.assertTrue(callable(StubRewindServer)) + self.assertTrue(callable(make_dispatch_payload)) + self.assertTrue(callable(wait_for_session)) + + def test_dunder_all_lists_stable_symbols_only(self): + import rewind_agent.testing as testing + self.assertIn("StubRewindServer", testing.__all__) + self.assertIn("make_dispatch_payload", testing.__all__) + self.assertIn("wait_for_session", testing.__all__) + # Internal/unstable surfaces must NOT appear in __all__. + for name in testing.__all__: + self.assertFalse( + name.startswith("_"), + f"{name} starts with '_'; private symbols don't belong in __all__", + ) + + +class TestStubRewindServer(unittest.TestCase): + """``StubRewindServer`` is the in-process Rewind server that downstream + tests bind ``ExplicitClient`` against. The contract: start it, point a + client at ``server.base_url``, run normal recording calls, stop it on + teardown. + """ + + def test_start_stop_lifecycle_via_context_manager(self): + from rewind_agent.testing import StubRewindServer + + with StubRewindServer() as server: + self.assertTrue(server.base_url.startswith("http://127.0.0.1:")) + # Server is responsive while inside the block. + req = urllib.request.Request( + f"{server.base_url}/api/sessions/start", + data=b'{"name": "lifecycle"}', + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=2.0) as resp: + self.assertIn(resp.status, (200, 201)) + + # After exit, the server is shut down — connections fail with a + # transport error specifically, not just any Exception. Tighten + # the assertion so a regression that surfaces as a different + # exception class (e.g. a hang masked by KeyboardInterrupt) + # actually fails the test. + with self.assertRaises(urllib.error.URLError): + with urllib.request.urlopen( + f"{server.base_url}/api/sessions/start", timeout=0.5 + ): + pass + + def test_explicit_client_can_record_against_server(self): + from rewind_agent import ExplicitClient + from rewind_agent.testing import StubRewindServer + + with StubRewindServer() as server: + client = ExplicitClient(server.base_url) + with client.session("stub-test"): + step = client.record_tool_call( + "ping", {"x": 1}, "ok", duration_ms=5, + ) + self.assertIsNotNone(step) + self.assertGreaterEqual(len(server.recorded_steps), 1) + + def test_recorded_steps_observable_for_assertions(self): + """Tests need to assert what was recorded. Stub exposes a list of + captured step dicts.""" + from rewind_agent import ExplicitClient + from rewind_agent.testing import StubRewindServer + + with StubRewindServer() as server: + client = ExplicitClient(server.base_url) + with client.session("assert-test"): + client.record_tool_call("foo", {}, "out", duration_ms=1) + client.record_tool_call("bar", {}, "out", duration_ms=1) + + tool_names = [s.get("tool_name") for s in server.recorded_steps] + self.assertEqual(tool_names, ["foo", "bar"]) + + +class TestMakeDispatchPayload(unittest.TestCase): + """``make_dispatch_payload`` builds a runner-compatible + :class:`DispatchPayload` for tests that exercise replay handlers. + + The contract: required fields have sensible defaults so a single + ``make_dispatch_payload(session_id="s")`` call produces a valid payload; + callers override individually. + """ + + def test_returns_dispatch_payload_instance(self): + from rewind_agent.runner import DispatchPayload + from rewind_agent.testing import make_dispatch_payload + + payload = make_dispatch_payload(session_id="s-1") + self.assertIsInstance(payload, DispatchPayload) + self.assertEqual(payload.session_id, "s-1") + + def test_defaults_populate_required_fields(self): + from rewind_agent.testing import make_dispatch_payload + + payload = make_dispatch_payload(session_id="s-1") + # Every required DispatchPayload field is non-empty. + self.assertTrue(payload.job_id) + self.assertTrue(payload.replay_context_id) + self.assertTrue(payload.replay_context_timeline_id) + self.assertTrue(payload.source_timeline_id) + self.assertTrue(payload.base_url) + self.assertTrue(payload.dispatch_token) + # at_step defaults to 1 (replay-from-start). + self.assertEqual(payload.at_step, 1) + + def test_overrides_apply(self): + from rewind_agent.testing import make_dispatch_payload + + payload = make_dispatch_payload( + session_id="sess", + job_id="job-42", + at_step=7, + base_url="http://my.test:1234", + ) + self.assertEqual(payload.job_id, "job-42") + self.assertEqual(payload.at_step, 7) + self.assertEqual(payload.base_url, "http://my.test:1234") + + +class TestWaitForSession(unittest.TestCase): + """``wait_for_session`` polls the stub for a session and returns when + it appears. Used by integration-style tests that bring up an agent in + a thread and need to assert "the agent recorded its session" without + relying on internal sync primitives. + """ + + def test_returns_session_when_recorded(self): + from rewind_agent import ExplicitClient + from rewind_agent.testing import StubRewindServer, wait_for_session + + with StubRewindServer() as server: + client = ExplicitClient(server.base_url) + + def record_in_thread(): + # Tiny delay to exercise the polling path. + time.sleep(0.05) + with client.session("background"): + client.record_tool_call( + "ping", {}, "ok", duration_ms=1, + ) + + t = threading.Thread(target=record_in_thread) + t.start() + try: + session = wait_for_session(server, name="background", timeout=2.0) + self.assertIsNotNone(session) + self.assertEqual(session["name"], "background") + finally: + t.join() + + def test_raises_timeout_when_no_session_recorded(self): + from rewind_agent.testing import StubRewindServer, wait_for_session + + with StubRewindServer() as server: + with self.assertRaises(TimeoutError): + wait_for_session(server, name="never-happens", timeout=0.2) + + +class TestUnstableBoundary(unittest.TestCase): + """``rewind_agent.testing._unstable`` exists as an explicit signal that + APIs underneath it may break in any 0.x release. Promoting a symbol + OUT of ``_unstable`` is a deliberate decision; this test makes sure + the boundary is real (the submodule importable) and not just + documentation.""" + + def test_unstable_submodule_importable(self): + import rewind_agent.testing._unstable # noqa: F401 + + def test_unstable_not_in_dunder_all(self): + import rewind_agent.testing as testing + self.assertNotIn("_unstable", testing.__all__) + + +if __name__ == "__main__": + unittest.main()