stepfork 0.1.0a2__py3-none-any.whl

stepfork/recorder/session.py

@@ -0,0 +1,500 @@
+"""Recording sessions that capture agent executions as `.sftrace` traces."""
+
+from __future__ import annotations
+
+import re
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from datetime import UTC, datetime
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Literal, TypeVar
+
+from stepfork.replay.session import active_replay
+from stepfork.trace import (
+    ErrorEvent,
+    EventBase,
+    EventStatus,
+    EventType,
+    FailureInfo,
+    JsonValue,
+    LLMRequest,
+    LLMResponse,
+    ReplayPolicy,
+    RunEnd,
+    RunStart,
+    RunStatus,
+    StateChange,
+    ToolCall,
+    ToolResult,
+    Trace,
+    TraceSerializationError,
+    save_trace,
+)
+from stepfork.trace.jsonable import Serializer, to_json_value
+from stepfork.trace.manifest import EnvironmentInfo, TraceTotals
+from stepfork.version import __version__
+
+DEFAULT_TRACE_DIR = Path(".stepfork") / "traces"
+
+E = TypeVar("E", bound=EventBase)
+
+_current_recorder: ContextVar[RecordingSession | None] = ContextVar(
+    "stepfork_recording",
+    default=None,
+)
+
+
+def active_recorder() -> RecordingSession | None:
+    """Return the recording session active in the current context, if any."""
+    return _current_recorder.get()
+
+
+class RecordingSession:
+    """Capture LLM, tool, and state events for one agent run.
+
+    Enter the session as a context manager; the trace is persisted when the
+    block exits and any application exception is always preserved::
+
+        with record("booking-agent", output="failure.sftrace") as session:
+            result = run_agent()
+            session.set_output(result)
+    """
+
+    def __init__(
+        self,
+        agent_name: str,
+        *,
+        output: str | Path | None = None,
+        overwrite: bool = False,
+        serializer: Serializer | None = None,
+        run_input: Any = None,
+    ) -> None:
+        self.trace = Trace(
+            agent_name=agent_name,
+            environment=EnvironmentInfo(
+                python=f"{sys.version_info.major}.{sys.version_info.minor}",
+                os=sys.platform,
+                stepfork_version=__version__,
+            ),
+        )
+        self._serializer = serializer
+        self._run_input = run_input
+        self._output_path = (
+            Path(output)
+            if output is not None
+            else _default_output_path(agent_name, self.trace.run_id)
+        )
+        self._overwrite = overwrite
+        self._stack: list[str] = []
+        self._token: Token[RecordingSession | None] | None = None
+        self._started = False
+        self._finished = False
+        self._inert = False
+        self._pending_output: JsonValue | None = None
+        self._has_output = False
+        self.path: Path | None = None
+
+    @property
+    def agent_name(self) -> str:
+        """Name of the recorded agent."""
+        return self.trace.agent_name
+
+    @property
+    def run_id(self) -> str:
+        """Opaque run identifier for this recording."""
+        return self.trace.run_id
+
+    @property
+    def output_path(self) -> Path:
+        """Where the trace bundle will be (or was) written."""
+        return self._output_path
+
+    def to_json(self, value: Any) -> JsonValue:
+        """Convert a runtime value into a JSON trace payload."""
+        return to_json_value(value, serializer=self._serializer)
+
+    def __enter__(self) -> RecordingSession:
+        if self._started:
+            raise RuntimeError("RecordingSession cannot be re-entered")
+        self._started = True
+        token = _current_recorder.set(self)
+        try:
+            if active_replay() is not None:
+                # Frozen replay: agent entrypoints keep their record() blocks.
+                # Stay inert so replay never writes a new trace bundle.
+                self._inert = True
+                self._stack = []
+                self._token = token
+                return self
+            if self._run_input is None:
+                start = self.trace.add(RunStart())
+            else:
+                start = self.trace.add(RunStart(input=self.to_json(self._run_input)))
+        except BaseException:
+            _current_recorder.reset(token)
+            raise
+        self._stack = [start.id]
+        self._token = token
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> Literal[False]:
+        try:
+            if self._inert:
+                return False
+            if not self._finished:
+                self._finished = True
+                self._finalize(exc)
+                self.path = save_trace(
+                    self.trace,
+                    self._output_path,
+                    overwrite=self._overwrite,
+                )
+        except Exception as save_error:
+            if exc is None:
+                raise
+            exc.add_note(f"stepfork: trace was not saved: {save_error}")
+        finally:
+            if self._token is not None:
+                _current_recorder.reset(self._token)
+                self._token = None
+        return False
+
+    def push_scope(self, event_id: str) -> None:
+        """Make ``event_id`` the parent scope for subsequently recorded events."""
+        if self._inert:
+            return
+        self._stack.append(event_id)
+
+    def pop_scope(self) -> None:
+        """Leave the current dependency scope."""
+        if self._inert:
+            return
+        if len(self._stack) <= 1:
+            raise RuntimeError("cannot leave the root recording scope")
+        self._stack.pop()
+
+    def record_llm_request(
+        self,
+        *,
+        model: str,
+        input: Any,
+        provider: str | None = None,
+        replay_policy: ReplayPolicy = ReplayPolicy.FROZEN,
+    ) -> LLMRequest:
+        """Record an outgoing LLM request."""
+        event = LLMRequest(
+            model=model,
+            input=self.to_json(input),
+            provider=provider,
+            replay_policy=replay_policy,
+        )
+        return self._append(event)
+
+    def record_llm_response(
+        self,
+        request: LLMRequest,
+        output: Any,
+        *,
+        model: str | None = None,
+        prompt_tokens: int | None = None,
+        completion_tokens: int | None = None,
+        total_tokens: int | None = None,
+        cost_usd: float | None = None,
+    ) -> LLMResponse:
+        """Record the LLM response paired with ``request``."""
+        event = LLMResponse(
+            model=model if model is not None else request.model,
+            output=self.to_json(output),
+            duration_ms=_elapsed_ms(request.timestamp),
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            cost_usd=cost_usd,
+            parent_id=request.id,
+        )
+        return self._append(event)
+
+    def record_tool_call(
+        self,
+        *,
+        name: str,
+        input: Any,
+        replay_policy: ReplayPolicy = ReplayPolicy.FROZEN,
+    ) -> ToolCall:
+        """Record an outgoing tool call."""
+        event = ToolCall(
+            name=name,
+            input=self.to_json(input),
+            replay_policy=replay_policy,
+        )
+        return self._append(event)
+
+    def record_tool_result(
+        self,
+        call: ToolCall,
+        output: Any,
+        *,
+        status: EventStatus = EventStatus.OK,
+    ) -> ToolResult:
+        """Record the result paired with ``call``."""
+        event = ToolResult(
+            name=call.name,
+            output=self.to_json(output),
+            duration_ms=_elapsed_ms(call.timestamp),
+            status=status,
+            parent_id=call.id,
+        )
+        return self._append(event)
+
+    def record_tool_failure(self, call: ToolCall, exc: BaseException) -> ToolResult:
+        """Record a failed tool result and a linked error event for ``exc``."""
+        error_type = type(exc).__name__
+        message = str(exc) or error_type
+        result = self._append(
+            ToolResult(
+                name=call.name,
+                output={"error_type": error_type, "message": message},
+                duration_ms=_elapsed_ms(call.timestamp),
+                status=EventStatus.ERROR,
+                parent_id=call.id,
+            )
+        )
+        self._append(
+            ErrorEvent(
+                error_type=error_type,
+                message=message,
+                parent_id=result.id,
+            )
+        )
+        return result
+
+    def record_llm_failure(self, request: LLMRequest, exc: BaseException) -> ErrorEvent:
+        """Record a failed LLM call linked to ``request``."""
+        error_type = type(exc).__name__
+        message = str(exc) or error_type
+        return self._append(
+            ErrorEvent(
+                error_type=error_type,
+                message=message,
+                parent_id=request.id,
+            )
+        )
+
+    def record_state_change(
+        self,
+        *,
+        key: str | None = None,
+        before: Any = None,
+        after: Any = None,
+    ) -> StateChange:
+        """Record a change to agent state."""
+        event = StateChange(
+            key=key,
+            before=self.to_json(before),
+            after=self.to_json(after),
+        )
+        return self._append(event)
+
+    def record_error(
+        self,
+        error_type: str,
+        message: str,
+        *,
+        details: Any = None,
+        parent_id: str | None = None,
+    ) -> ErrorEvent:
+        """Record an error event in the current scope."""
+        event = ErrorEvent(
+            error_type=error_type,
+            message=message or error_type,
+            details=self.to_json(details) if details is not None else None,
+            parent_id=parent_id,
+        )
+        return self._append(event)
+
+    def set_output(self, value: Any) -> None:
+        """Set the final run output recorded in ``run_end``."""
+        self._pending_output = self.to_json(value)
+        self._has_output = True
+
+    @contextmanager
+    def llm_call(
+        self,
+        *,
+        model: str,
+        input: Any,
+        provider: str | None = None,
+        replay_policy: ReplayPolicy = ReplayPolicy.FROZEN,
+    ) -> Iterator[LLMCallHandle]:
+        """Manually record an LLM call around a ``with`` body.
+
+        This API records during :func:`record` sessions. Replayable agents
+        should prefer :func:`stepfork.llm_request`, which substitutes captured
+        responses during frozen replay instead of executing the body.
+        """
+        request = self.record_llm_request(
+            model=model,
+            input=input,
+            provider=provider,
+            replay_policy=replay_policy,
+        )
+        handle = LLMCallHandle(request)
+        self.push_scope(request.id)
+        try:
+            yield handle
+        except BaseException as exc:
+            self.record_llm_failure(request, exc)
+            raise
+        finally:
+            self.pop_scope()
+        self.record_llm_response(
+            request,
+            handle.output,
+            model=handle.model,
+            prompt_tokens=handle.prompt_tokens,
+            completion_tokens=handle.completion_tokens,
+            total_tokens=handle.total_tokens,
+            cost_usd=handle.cost_usd,
+        )
+
+    def _append(self, event: E) -> E:
+        if event.parent_id is None and self._stack:
+            event = event.model_copy(update={"parent_id": self._stack[-1]})
+        return self.trace.add(event)
+
+    def _finalize(self, exc: BaseException | None) -> None:
+        output: JsonValue | None = None
+        failure: FailureInfo | None = None
+
+        if exc is not None:
+            error_type = type(exc).__name__
+            message = str(exc) or error_type
+            error_event = self.record_error(error_type, message)
+            failure = FailureInfo(
+                type=error_type,
+                message=message,
+                step=error_event.step,
+            )
+            status = RunStatus.FAILED
+        else:
+            status = RunStatus.COMPLETED
+            if self._has_output:
+                output = self._pending_output
+
+        run_end = self.trace.add(RunEnd(run_status=status, output=output))
+        duration_ms = max(
+            0,
+            int((run_end.timestamp - self.trace.created_at).total_seconds() * 1000),
+        )
+        self.trace.totals = TraceTotals(
+            events=len(self.trace.events),
+            llm_calls=sum(
+                event.type.value == "llm_request" for event in self.trace.events
+            ),
+            tool_calls=sum(
+                event.type.value == "tool_call" for event in self.trace.events
+            ),
+            cost_usd=_total_cost(self.trace),
+            duration_ms=duration_ms,
+        )
+        self.trace.status = status
+        self.trace.failure = failure
+
+
+class LLMCallHandle:
+    """Mutable handle yielded by :meth:`RecordingSession.llm_call`."""
+
+    def __init__(self, request: LLMRequest) -> None:
+        self.request = request
+        self.model: str | None = request.model
+        self.output: Any = None
+        self.prompt_tokens: int | None = None
+        self.completion_tokens: int | None = None
+        self.total_tokens: int | None = None
+        self.cost_usd: float | None = None
+
+    def set_output(
+        self,
+        value: Any,
+        *,
+        model: str | None = None,
+        prompt_tokens: int | None = None,
+        completion_tokens: int | None = None,
+        total_tokens: int | None = None,
+        cost_usd: float | None = None,
+    ) -> None:
+        """Attach the response payload and optional usage metadata."""
+        self.output = value
+        if model is not None:
+            self.model = model
+        self.prompt_tokens = prompt_tokens
+        self.completion_tokens = completion_tokens
+        self.total_tokens = total_tokens
+        self.cost_usd = cost_usd
+
+
+def record(
+    agent_name: str,
+    *,
+    output: str | Path | None = None,
+    overwrite: bool = False,
+    serializer: Serializer | None = None,
+    run_input: Any = None,
+) -> RecordingSession:
+    """Create a recording session for one agent run.
+
+    When ``output`` is omitted the trace is written under
+    ``.stepfork/traces/``. The session persists the trace on exit, including
+    when the recorded block raises, and always re-raises the original error.
+
+    Inside an active replay the session is inert: the block runs, but no
+    trace is recorded or written, so replaying an entrypoint never creates
+    new bundles.
+    """
+    return RecordingSession(
+        agent_name,
+        output=output,
+        overwrite=overwrite,
+        serializer=serializer,
+        run_input=run_input,
+    )
+
+
+def _elapsed_ms(since: datetime) -> int:
+    return max(0, int((datetime.now(UTC) - since).total_seconds() * 1000))
+
+
+def _total_cost(trace: Trace) -> float | None:
+    costs = [
+        event.cost_usd
+        for event in trace.events
+        if event.type is EventType.LLM_RESPONSE
+        and isinstance(event, LLMResponse)
+        and event.cost_usd is not None
+    ]
+    if not costs:
+        return None
+    return sum(costs)
+
+
+def _default_output_path(agent_name: str, run_id: str) -> Path:
+    safe_agent = re.sub(r"[^A-Za-z0-9._-]+", "-", agent_name).strip("-") or "agent"
+    return DEFAULT_TRACE_DIR / f"{safe_agent}-{run_id}.sftrace"
+
+
+__all__ = [
+    "DEFAULT_TRACE_DIR",
+    "LLMCallHandle",
+    "RecordingSession",
+    "TraceSerializationError",
+    "active_recorder",
+    "record",
+]

stepfork/recorder/tooling.py

@@ -0,0 +1,208 @@
+"""`@trace_tool`: capture tool boundaries during recording and replay."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Any, ParamSpec, TypeVar, cast, overload
+
+from stepfork.recorder.session import RecordingSession, active_recorder
+from stepfork.replay.session import ReplayDecision, active_replay
+from stepfork.trace import JsonValue, ReplayPolicy, ToolCall
+from stepfork.trace.jsonable import Serializer, to_json_value
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+@overload
+def trace_tool(func: Callable[P, R], /) -> Callable[P, R]: ...
+
+
+@overload
+def trace_tool(
+    func: None = None,
+    *,
+    name: str | None = None,
+    replay_policy: ReplayPolicy = ReplayPolicy.FROZEN,
+    serializer: Serializer | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def trace_tool(
+    func: Callable[P, R] | None = None,
+    *,
+    name: str | None = None,
+    replay_policy: ReplayPolicy = ReplayPolicy.FROZEN,
+    serializer: Serializer | None = None,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    """Record a tool call and its result inside an active recording.
+
+    Outside a recording (and outside replay) the decorated function runs
+    unchanged. During frozen replay the function body is not executed and the
+    recorded result is returned instead.
+
+    Supported forms::
+
+        @trace_tool
+        def search(destination: str) -> dict: ...
+
+        @trace_tool(name="flight_search")
+        def search(destination: str) -> dict: ...
+    """
+
+    def decorate(target: Callable[P, R]) -> Callable[P, R]:
+        tool_name = name if name is not None else target.__name__
+        if inspect.iscoroutinefunction(target):
+
+            async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
+                run = _ToolRun(
+                    target=target,
+                    tool_name=tool_name,
+                    args=args,
+                    kwargs=kwargs,
+                    replay_policy=replay_policy,
+                    serializer=serializer,
+                )
+                if run.passthrough:
+                    return await cast(Awaitable[Any], target(*args, **kwargs))
+                if run.substituted:
+                    return run.substitute()
+                try:
+                    result = await cast(Awaitable[Any], target(*args, **kwargs))
+                except BaseException as exc:
+                    run.record_failure(exc)
+                    raise
+                else:
+                    run.record_success(result)
+                    return result
+                finally:
+                    run.close()
+
+            return cast(
+                Callable[P, R],
+                functools.wraps(target)(async_wrapper),
+            )
+
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            run = _ToolRun(
+                target=target,
+                tool_name=tool_name,
+                args=args,
+                kwargs=kwargs,
+                replay_policy=replay_policy,
+                serializer=serializer,
+            )
+            if run.passthrough:
+                return target(*args, **kwargs)
+            if run.substituted:
+                return cast(R, run.substitute())
+            try:
+                result = target(*args, **kwargs)
+            except BaseException as exc:
+                run.record_failure(exc)
+                raise
+            else:
+                run.record_success(result)
+                return result
+            finally:
+                run.close()
+
+        return functools.wraps(target)(wrapper)
+
+    if func is None:
+        return decorate
+    return decorate(func)
+
+
+class _ToolRun:
+    """One instrumented tool invocation shared by sync and async wrappers."""
+
+    def __init__(
+        self,
+        *,
+        target: Callable[..., Any],
+        tool_name: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        replay_policy: ReplayPolicy,
+        serializer: Serializer | None,
+    ) -> None:
+        self.replay = active_replay()
+        self.session: RecordingSession | None = None
+        self.call: ToolCall | None = None
+        self.decision: ReplayDecision | None = None
+        self.input_json: JsonValue | None = None
+
+        if self.replay is None and active_recorder() is None:
+            return
+
+        bound = _bind_arguments(target, args, kwargs)
+        if self.replay is not None:
+            self.input_json = to_json_value(bound, serializer=serializer)
+            self.decision = self.replay.before_tool(
+                name=tool_name,
+                input=self.input_json,
+            )
+            return
+
+        session = active_recorder()
+        if session is None:  # pragma: no cover - defensive
+            return
+        self.session = session
+        self.call = session.record_tool_call(
+            name=tool_name,
+            input=bound,
+            replay_policy=replay_policy,
+        )
+        session.push_scope(self.call.id)
+
+    @property
+    def passthrough(self) -> bool:
+        """True when neither recording nor replay is active."""
+        return self.replay is None and self.session is None
+
+    @property
+    def substituted(self) -> bool:
+        """True when replay returns a captured value instead of executing."""
+        return self.decision is not None and not self.decision.execute
+
+    def substitute(self) -> JsonValue | None:
+        """Return the captured output for frozen replay."""
+        if self.decision is None:  # pragma: no cover - defensive
+            raise RuntimeError("no replay decision to substitute")
+        return self.decision.value
+
+    def record_success(self, result: Any) -> None:
+        """Record the successful tool result when a session is active."""
+        if self.session is None or self.call is None:
+            return
+        self.session.record_tool_result(self.call, result)
+
+    def record_failure(self, exc: BaseException) -> None:
+        """Record the failed tool result when a session is active."""
+        if self.session is None or self.call is None:
+            return
+        self.session.record_tool_failure(self.call, exc)
+
+    def close(self) -> None:
+        """Leave the recording scope opened for this tool call."""
+        if self.session is None or self.call is None:
+            return
+        self.session.pop_scope()
+
+
+def _bind_arguments(
+    target: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    try:
+        bound = inspect.signature(target).bind(*args, **kwargs)
+    except (TypeError, ValueError):
+        return {"args": list(args), "kwargs": dict(kwargs)}
+    arguments: dict[str, Any] = dict(bound.arguments)
+    arguments.pop("self", None)
+    arguments.pop("cls", None)
+    return arguments