tracefork-agent 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (58) hide show
  1. tracefork_agent-0.1.0/.gitignore +26 -0
  2. tracefork_agent-0.1.0/PKG-INFO +24 -0
  3. tracefork_agent-0.1.0/README.md +8 -0
  4. tracefork_agent-0.1.0/pyproject.toml +22 -0
  5. tracefork_agent-0.1.0/src/tracefork/__init__.py +9 -0
  6. tracefork_agent-0.1.0/src/tracefork/adapters/__init__.py +17 -0
  7. tracefork_agent-0.1.0/src/tracefork/adapters/httpx.py +210 -0
  8. tracefork_agent-0.1.0/src/tracefork/adapters/openai.py +179 -0
  9. tracefork_agent-0.1.0/src/tracefork/adapters/python_tools.py +190 -0
  10. tracefork_agent-0.1.0/src/tracefork/assertions.py +206 -0
  11. tracefork_agent-0.1.0/src/tracefork/bootstrap.py +57 -0
  12. tracefork_agent-0.1.0/src/tracefork/boundaries/__init__.py +20 -0
  13. tracefork_agent-0.1.0/src/tracefork/boundaries/base.py +50 -0
  14. tracefork_agent-0.1.0/src/tracefork/boundaries/errors.py +7 -0
  15. tracefork_agent-0.1.0/src/tracefork/boundaries/registry.py +34 -0
  16. tracefork_agent-0.1.0/src/tracefork/boundaries/runtime.py +238 -0
  17. tracefork_agent-0.1.0/src/tracefork/canonicalization/__init__.py +13 -0
  18. tracefork_agent-0.1.0/src/tracefork/canonicalization/canonicalizer.py +31 -0
  19. tracefork_agent-0.1.0/src/tracefork/canonicalization/json.py +77 -0
  20. tracefork_agent-0.1.0/src/tracefork/canonicalization/rules.py +50 -0
  21. tracefork_agent-0.1.0/src/tracefork/cli/__init__.py +1 -0
  22. tracefork_agent-0.1.0/src/tracefork/cli/bootstrap.py +51 -0
  23. tracefork_agent-0.1.0/src/tracefork/cli/commands/diff.py +51 -0
  24. tracefork_agent-0.1.0/src/tracefork/cli/commands/eval.py +34 -0
  25. tracefork_agent-0.1.0/src/tracefork/cli/commands/init.py +34 -0
  26. tracefork_agent-0.1.0/src/tracefork/cli/commands/inspect.py +39 -0
  27. tracefork_agent-0.1.0/src/tracefork/cli/commands/record.py +35 -0
  28. tracefork_agent-0.1.0/src/tracefork/cli/commands/replay.py +127 -0
  29. tracefork_agent-0.1.0/src/tracefork/cli/main.py +46 -0
  30. tracefork_agent-0.1.0/src/tracefork/cli/output/__init__.py +5 -0
  31. tracefork_agent-0.1.0/src/tracefork/cli/output/console.py +6 -0
  32. tracefork_agent-0.1.0/src/tracefork/cli/suites.py +270 -0
  33. tracefork_agent-0.1.0/src/tracefork/diff.py +212 -0
  34. tracefork_agent-0.1.0/src/tracefork/errors.py +70 -0
  35. tracefork_agent-0.1.0/src/tracefork/metrics.py +121 -0
  36. tracefork_agent-0.1.0/src/tracefork/models/__init__.py +31 -0
  37. tracefork_agent-0.1.0/src/tracefork/models/provenance.py +51 -0
  38. tracefork_agent-0.1.0/src/tracefork/models/replay.py +100 -0
  39. tracefork_agent-0.1.0/src/tracefork/models/span.py +79 -0
  40. tracefork_agent-0.1.0/src/tracefork/models/trace.py +60 -0
  41. tracefork_agent-0.1.0/src/tracefork/py.typed +0 -0
  42. tracefork_agent-0.1.0/src/tracefork/recording/__init__.py +13 -0
  43. tracefork_agent-0.1.0/src/tracefork/recording/context.py +40 -0
  44. tracefork_agent-0.1.0/src/tracefork/recording/lifecycle.py +87 -0
  45. tracefork_agent-0.1.0/src/tracefork/recording/recorder.py +237 -0
  46. tracefork_agent-0.1.0/src/tracefork/redaction.py +54 -0
  47. tracefork_agent-0.1.0/src/tracefork/replay/__init__.py +23 -0
  48. tracefork_agent-0.1.0/src/tracefork/replay/matcher.py +206 -0
  49. tracefork_agent-0.1.0/src/tracefork/replay/session.py +295 -0
  50. tracefork_agent-0.1.0/src/tracefork/serialization/__init__.py +17 -0
  51. tracefork_agent-0.1.0/src/tracefork/serialization/fixture.py +51 -0
  52. tracefork_agent-0.1.0/src/tracefork/serialization/reader.py +69 -0
  53. tracefork_agent-0.1.0/src/tracefork/serialization/writer.py +22 -0
  54. tracefork_agent-0.1.0/src/tracefork/storage/__init__.py +12 -0
  55. tracefork_agent-0.1.0/src/tracefork/storage/base.py +36 -0
  56. tracefork_agent-0.1.0/src/tracefork/storage/filesystem.py +69 -0
  57. tracefork_agent-0.1.0/src/tracefork/storage/sqlite.py +69 -0
  58. tracefork_agent-0.1.0/src/tracefork/trajectory.py +93 -0
@@ -0,0 +1,26 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+
8
+ # Environments
9
+ .venv/
10
+
11
+ # Tooling caches
12
+ .pytest_cache/
13
+ .mypy_cache/
14
+ .ruff_cache/
15
+ .hypothesis/
16
+ .coverage
17
+ coverage.xml
18
+ htmlcov/
19
+
20
+ # TraceFork local state
21
+ .tracefork/
22
+ site/
23
+
24
+ # Editors
25
+ .idea/
26
+ .vscode/
@@ -0,0 +1,24 @@
1
+ Metadata-Version: 2.5
2
+ Name: tracefork-agent
3
+ Version: 0.1.0
4
+ Summary: Replay AI agent failures locally and turn them into regression tests.
5
+ License-Expression: Apache-2.0
6
+ Requires-Python: >=3.12
7
+ Requires-Dist: pydantic>=2.7
8
+ Requires-Dist: pyyaml>=6.0
9
+ Requires-Dist: rich>=13.7
10
+ Requires-Dist: typer>=0.12
11
+ Provides-Extra: httpx
12
+ Requires-Dist: httpx>=0.27; extra == 'httpx'
13
+ Provides-Extra: openai
14
+ Requires-Dist: openai>=1.40; extra == 'openai'
15
+ Description-Content-Type: text/markdown
16
+
17
+ # tracefork-core
18
+
19
+ Framework-independent core engine for TraceFork: trace domain models, the
20
+ recording engine, the boundary abstraction, canonicalization and matching, and
21
+ the hermetic replay engine.
22
+
23
+ This package must never depend on an agent framework. Framework support lives
24
+ in adapter packages (`tracefork-openai`, `tracefork-langgraph`, ...).
@@ -0,0 +1,8 @@
1
+ # tracefork-core
2
+
3
+ Framework-independent core engine for TraceFork: trace domain models, the
4
+ recording engine, the boundary abstraction, canonicalization and matching, and
5
+ the hermetic replay engine.
6
+
7
+ This package must never depend on an agent framework. Framework support lives
8
+ in adapter packages (`tracefork-openai`, `tracefork-langgraph`, ...).
@@ -0,0 +1,22 @@
1
+ [project]
2
+ name = "tracefork-agent"
3
+ version = "0.1.0"
4
+ description = "Replay AI agent failures locally and turn them into regression tests."
5
+ readme = "README.md"
6
+ requires-python = ">=3.12"
7
+ license = "Apache-2.0"
8
+ dependencies = ["pydantic>=2.7", "pyyaml>=6.0", "rich>=13.7", "typer>=0.12"]
9
+
10
+ [project.optional-dependencies]
11
+ openai = ["openai>=1.40"]
12
+ httpx = ["httpx>=0.27"]
13
+
14
+ [project.scripts]
15
+ tracefork = "tracefork.cli.main:main"
16
+
17
+ [build-system]
18
+ requires = ["hatchling"]
19
+ build-backend = "hatchling.build"
20
+
21
+ [tool.hatch.build.targets.wheel]
22
+ packages = ["src/tracefork"]
@@ -0,0 +1,9 @@
1
+ """TraceFork core: record, replay and diff AI agent executions."""
2
+
3
+ from tracefork.adapters import ToolBox
4
+ from tracefork.recording import Recording, record, span
5
+ from tracefork.redaction import RedactionEngine, redactor
6
+
7
+ __version__ = "0.1.0"
8
+
9
+ __all__ = ["Recording", "RedactionEngine", "ToolBox", "__version__", "record", "redactor", "span"]
@@ -0,0 +1,17 @@
1
+ """Adapters connecting the boundary runtime to concrete technologies."""
2
+
3
+ from tracefork.adapters.python_tools import (
4
+ TOOL_BOUNDARY_TYPE,
5
+ PythonToolHandler,
6
+ ToolBox,
7
+ WrappedTool,
8
+ canonicalize_argument,
9
+ )
10
+
11
+ __all__ = [
12
+ "TOOL_BOUNDARY_TYPE",
13
+ "PythonToolHandler",
14
+ "ToolBox",
15
+ "WrappedTool",
16
+ "canonicalize_argument",
17
+ ]
@@ -0,0 +1,210 @@
1
+ """httpx adapter (TF-080..082).
2
+
3
+ An ``httpx.AsyncClient`` transport that routes every request through the
4
+ boundary runtime, so HTTP traffic is recorded and replayed hermetically.
5
+ Secret headers are redacted before anything is persisted (TF-081). Async-only
6
+ in v0.1: the boundary runtime is async-first.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import asyncio
12
+ import base64
13
+ import json
14
+ from typing import Any
15
+
16
+ import httpx
17
+
18
+ from tracefork.boundaries import BoundaryRuntime
19
+ from tracefork.boundaries.base import LiveCall
20
+ from tracefork.errors import AdapterError
21
+ from tracefork.models import BoundaryRequest, BoundaryResponse
22
+
23
+ HTTPX_BOUNDARY_TYPE = "http.httpx"
24
+
25
+ REDACTED = "[REDACTED]"
26
+
27
+ # Secret-safe defaults (TF-081): matched case-insensitively.
28
+ _SECRET_HEADERS = {"authorization", "cookie", "set-cookie", "x-api-key", "proxy-authorization"}
29
+
30
+ # Content-coding headers describe the WIRE encoding, but transports hand us
31
+ # decoded payloads. Recording them makes the replayed httpx.Response try to
32
+ # decompress already-decompressed content (found by the live smoke), so they
33
+ # are stripped on both directions.
34
+ _CONTENT_CODING_HEADERS = {"content-encoding", "content-length", "transfer-encoding"}
35
+
36
+
37
+ class HTTPXHandler:
38
+ """Translates httpx request/response objects to and from canonical data."""
39
+
40
+ async def execute(self, request: BoundaryRequest, call_live: LiveCall) -> BoundaryResponse:
41
+ response: httpx.Response = await call_live()
42
+ await response.aread() # live responses stream; reading is a no-op otherwise
43
+ return BoundaryResponse(
44
+ response=_response_payload(request.request, response),
45
+ metadata={"status_code": response.status_code},
46
+ )
47
+
48
+ def restore(self, response: dict[str, Any], metadata: dict[str, Any]) -> httpx.Response:
49
+ recorded_request = response["request"]
50
+ request = httpx.Request(
51
+ method=recorded_request["method"],
52
+ url=recorded_request["url"],
53
+ headers=_plain_headers(recorded_request["headers"]),
54
+ content=_request_content(recorded_request),
55
+ )
56
+ status: int = response["status_code"]
57
+ return httpx.Response(
58
+ status_code=status,
59
+ headers=_plain_headers(response["headers"]),
60
+ content=_response_content(response),
61
+ request=request,
62
+ )
63
+
64
+
65
+ class TraceForkAsyncTransport(httpx.AsyncBaseTransport):
66
+ """httpx transport that routes requests through the boundary runtime.
67
+
68
+ Wrap the real transport::
69
+
70
+ client = httpx.AsyncClient(
71
+ transport=TraceForkAsyncTransport(inner=httpx.AsyncTransport(), runtime=runtime)
72
+ )
73
+ """
74
+
75
+ def __init__(self, inner: httpx.AsyncBaseTransport, runtime: BoundaryRuntime) -> None:
76
+ self._inner = inner
77
+ self._runtime = runtime
78
+ runtime.registry.register(HTTPX_BOUNDARY_TYPE, HTTPXHandler())
79
+
80
+ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
81
+ payload = _request_payload(request)
82
+ name = f"{request.method} {request.url.path}"
83
+
84
+ async def call_live() -> httpx.Response:
85
+ return await self._inner.handle_async_request(request)
86
+
87
+ response: httpx.Response = await self._runtime.invoke(
88
+ HTTPX_BOUNDARY_TYPE, name, payload, call_live
89
+ )
90
+ return response
91
+
92
+
93
+ class TraceForkTransport(httpx.BaseTransport):
94
+ """Sync httpx transport that routes requests through the boundary runtime.
95
+
96
+ Each call is executed in a fresh event loop (the sync SDK path cannot
97
+ await); using it inside a running loop raises ``AdapterError`` — use
98
+ :class:`TraceForkAsyncTransport` from async code.
99
+ """
100
+
101
+ def __init__(self, inner: httpx.BaseTransport, runtime: BoundaryRuntime) -> None:
102
+ self._inner = inner
103
+ self._runtime = runtime
104
+ runtime.registry.register(HTTPX_BOUNDARY_TYPE, HTTPXHandler())
105
+
106
+ def handle_request(self, request: httpx.Request) -> httpx.Response:
107
+ try:
108
+ asyncio.get_running_loop()
109
+ except RuntimeError:
110
+ pass
111
+ else:
112
+ msg = "sync httpx transport used inside a running event loop; use AsyncClient"
113
+ raise AdapterError(msg)
114
+
115
+ async def call_live() -> httpx.Response:
116
+ response = self._inner.handle_request(request)
117
+ response.read() # buffer now; aread() later is a no-op
118
+ return response
119
+
120
+ response: httpx.Response = asyncio.run(
121
+ self._runtime.invoke(
122
+ HTTPX_BOUNDARY_TYPE,
123
+ f"{request.method} {request.url.path}",
124
+ _request_payload(request),
125
+ call_live,
126
+ )
127
+ )
128
+ return response
129
+
130
+
131
+ def _request_payload(request: httpx.Request) -> dict[str, Any]:
132
+ headers = {
133
+ key: value
134
+ for key, value in _safe_headers(request.headers).items()
135
+ if key not in _CONTENT_CODING_HEADERS
136
+ }
137
+ payload: dict[str, Any] = {
138
+ "method": request.method,
139
+ "url": str(request.url),
140
+ "headers": headers,
141
+ }
142
+ body = _decode_body(request.headers.get("content-type"), request.content)
143
+ if body is not None:
144
+ payload.update(body)
145
+ return payload
146
+
147
+
148
+ def _response_payload(request_payload: dict[str, Any], response: httpx.Response) -> dict[str, Any]:
149
+ headers = {
150
+ key: value
151
+ for key, value in _safe_headers(response.headers).items()
152
+ if key not in _CONTENT_CODING_HEADERS
153
+ }
154
+ payload: dict[str, Any] = {
155
+ "request": request_payload,
156
+ "status_code": response.status_code,
157
+ "headers": headers,
158
+ }
159
+ body = _decode_body(response.headers.get("content-type"), response.content)
160
+ if body is not None:
161
+ payload.update(body)
162
+ return payload
163
+
164
+
165
+ def _decode_body(content_type: str | None, content: bytes) -> dict[str, Any] | None:
166
+ if not content:
167
+ return None
168
+ if content_type and "application/json" in content_type:
169
+ try:
170
+ return {"json": json.loads(content)}
171
+ except json.JSONDecodeError:
172
+ pass
173
+ try:
174
+ return {"content": content.decode("utf-8")}
175
+ except UnicodeDecodeError:
176
+ # Binary bodies are stored base64 so replays are byte-identical
177
+ # (errors="replace" would silently corrupt them).
178
+ return {"content_base64": base64.b64encode(content).decode("ascii")}
179
+
180
+
181
+ def _request_content(recorded_request: dict[str, Any]) -> bytes:
182
+ if "json" in recorded_request:
183
+ return str(json.dumps(recorded_request["json"])).encode("utf-8")
184
+ if "content" in recorded_request:
185
+ return str(recorded_request["content"]).encode("utf-8")
186
+ if "content_base64" in recorded_request:
187
+ return base64.b64decode(recorded_request["content_base64"])
188
+ return b""
189
+
190
+
191
+ def _response_content(recorded_response: dict[str, Any]) -> bytes:
192
+ if "json" in recorded_response:
193
+ return str(json.dumps(recorded_response["json"])).encode("utf-8")
194
+ if "content" in recorded_response:
195
+ return str(recorded_response["content"]).encode("utf-8")
196
+ if "content_base64" in recorded_response:
197
+ return base64.b64decode(recorded_response["content_base64"])
198
+ return b""
199
+
200
+
201
+ def _safe_headers(headers: Any) -> dict[str, str]:
202
+ """Lower-case header map with secret headers redacted (TF-081)."""
203
+ return {
204
+ str(key).lower(): REDACTED if str(key).lower() in _SECRET_HEADERS else str(value)
205
+ for key, value in headers.items()
206
+ }
207
+
208
+
209
+ def _plain_headers(headers: dict[str, str]) -> list[tuple[str, str]]:
210
+ return list(headers.items())
@@ -0,0 +1,179 @@
1
+ """OpenAI Responses API adapter (TF-070..074).
2
+
3
+ Routes ``client.responses.create`` (sync and async, streaming and not) through
4
+ the boundary runtime. The adapter only translates:
5
+
6
+ - native kwargs → canonical request payload,
7
+ - SDK response/stream objects → canonical payloads (with usage metadata),
8
+ - recorded payloads → SDK objects rebuilt via the SDK's own construction
9
+ machinery, so application code keeps working unchanged in replay.
10
+
11
+ The live callable is supplied by the instrumentation and is structurally
12
+ unreachable on the replay path (ADR 0004, ADR 0003).
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import asyncio
18
+ import inspect
19
+ import time
20
+ from collections.abc import AsyncIterator, Callable, Iterator
21
+ from typing import Any
22
+
23
+ from openai import AsyncOpenAI, OpenAI
24
+ from openai._models import construct_type
25
+ from openai.types.responses import Response
26
+ from openai.types.responses.response_stream_event import ResponseStreamEvent
27
+
28
+ from tracefork.adapters import canonicalize_argument
29
+ from tracefork.boundaries import BoundaryRuntime
30
+ from tracefork.boundaries.base import LiveCall
31
+ from tracefork.errors import AdapterError
32
+ from tracefork.models import BoundaryRequest, BoundaryResponse
33
+
34
+ LLM_OPENAI_BOUNDARY_TYPE = "llm.openai"
35
+
36
+
37
+ class OpenAIHandler:
38
+ """Translates Responses API traffic between native and canonical shapes."""
39
+
40
+ async def execute(self, request: BoundaryRequest, call_live: LiveCall) -> BoundaryResponse:
41
+ started = time.perf_counter()
42
+ if request.metadata.get("stream"):
43
+ stream = await _call(call_live)
44
+ if hasattr(stream, "__aiter__"):
45
+ events = [event.model_dump(mode="json") async for event in stream]
46
+ else:
47
+ events = [event.model_dump(mode="json") for event in stream]
48
+ terminal_types = ("response.completed", "response.incomplete")
49
+ terminal = next((e for e in reversed(events) if e.get("type") in terminal_types), None)
50
+ usage = _usage_metadata((terminal or {}).get("response", {}).get("usage"))
51
+ return BoundaryResponse(
52
+ response={"__stream__": True, "events": events},
53
+ metadata={
54
+ "stream": True,
55
+ "usage": usage,
56
+ "finish_reason": terminal.get("type").removeprefix("response.")
57
+ if terminal
58
+ else None,
59
+ "latency_ms": _latency_ms(started),
60
+ },
61
+ )
62
+ response = await _call(call_live)
63
+ payload = response.model_dump(mode="json")
64
+ return BoundaryResponse(
65
+ response=payload,
66
+ metadata={
67
+ "stream": False,
68
+ "usage": _usage_metadata(payload.get("usage")),
69
+ "finish_reason": payload.get("status"),
70
+ "latency_ms": _latency_ms(started),
71
+ },
72
+ )
73
+
74
+ def restore(self, response: Any, metadata: dict[str, Any]) -> Any:
75
+ if metadata.get("stream"):
76
+ events = [
77
+ construct_type(type_=ResponseStreamEvent, value=event)
78
+ for event in response["events"]
79
+ ]
80
+ return _ReplayedStream(events)
81
+ return construct_type(type_=Response, value=response)
82
+
83
+
84
+ class _ReplayedStream:
85
+ """In-memory stream replays a recorded event sequence.
86
+
87
+ Buffers the full event list (also done at record time); supports both the
88
+ sync and async iteration protocols, matching whichever client is used.
89
+ """
90
+
91
+ def __init__(self, events: list[Any]) -> None:
92
+ self._events = events
93
+
94
+ def __iter__(self) -> Iterator[Any]:
95
+ return iter(self._events)
96
+
97
+ def __aiter__(self) -> AsyncIterator[Any]:
98
+ async def _gen() -> AsyncIterator[Any]:
99
+ for event in self._events:
100
+ yield event
101
+
102
+ return _gen()
103
+
104
+
105
+ def instrument_openai(client: OpenAI | AsyncOpenAI, runtime: BoundaryRuntime) -> None:
106
+ """Route ``client.responses.create`` through the boundary runtime, in place.
107
+
108
+ Registers the shared OpenAI handler on the runtime's registry. Recording,
109
+ matching and replay semantics remain in the core runtime.
110
+ """
111
+ runtime.registry.register(LLM_OPENAI_BOUNDARY_TYPE, OpenAIHandler())
112
+ original = client.responses.create
113
+ is_async = isinstance(client, AsyncOpenAI)
114
+
115
+ if is_async:
116
+
117
+ async def create_async(**kwargs: Any) -> Any:
118
+ return await _invoke(runtime, original, kwargs, sync_client=False)
119
+
120
+ client.responses.create = create_async # type: ignore[method-assign]
121
+ return
122
+
123
+ def create_sync(**kwargs: Any) -> Any:
124
+ try:
125
+ asyncio.get_running_loop()
126
+ except RuntimeError:
127
+ pass
128
+ else:
129
+ msg = "sync OpenAI client used inside a running event loop; use AsyncOpenAI"
130
+ raise AdapterError(msg)
131
+ return asyncio.run(_invoke(runtime, original, kwargs, sync_client=True))
132
+
133
+ client.responses.create = create_sync # type: ignore[method-assign]
134
+
135
+
136
+ async def _invoke(
137
+ runtime: BoundaryRuntime,
138
+ original: Callable[..., Any],
139
+ kwargs: dict[str, Any],
140
+ *,
141
+ sync_client: bool,
142
+ ) -> Any:
143
+ model = kwargs.get("model") or "responses.create"
144
+ stream = bool(kwargs.get("stream"))
145
+ request_payload = {str(key): canonicalize_argument(value) for key, value in kwargs.items()}
146
+
147
+ async def call_live() -> Any:
148
+ result = original(**kwargs)
149
+ if inspect.isawaitable(result):
150
+ result = await result
151
+ return result
152
+
153
+ return await runtime.invoke(
154
+ LLM_OPENAI_BOUNDARY_TYPE,
155
+ model,
156
+ request_payload,
157
+ call_live,
158
+ metadata={"stream": stream, "sync": sync_client},
159
+ )
160
+
161
+
162
+ async def _call(call_live: LiveCall) -> Any:
163
+ return await call_live()
164
+
165
+
166
+ def _usage_metadata(usage: dict[str, Any] | None) -> dict[str, Any]:
167
+ if not usage:
168
+ return {}
169
+ input_details = usage.get("input_tokens_details") or {}
170
+ return {
171
+ "input_tokens": usage.get("input_tokens"),
172
+ "output_tokens": usage.get("output_tokens"),
173
+ "total_tokens": usage.get("total_tokens"),
174
+ "cached_tokens": input_details.get("cached_tokens"),
175
+ }
176
+
177
+
178
+ def _latency_ms(started: float) -> int:
179
+ return int((time.perf_counter() - started) * 1000)
@@ -0,0 +1,190 @@
1
+ """Python function tool adapter (TF-060..063).
2
+
3
+ Python tools are ordinary functions routed through the boundary runtime, so
4
+ recording, matching and replay semantics live entirely in the core (ADR 0004).
5
+ The adapter only translates: argument canonicalization on the way in, the
6
+ recorded response on the way out.
7
+
8
+ Usage::
9
+
10
+ runtime = BoundaryRuntime(registry=registry)
11
+ tools = ToolBox(runtime)
12
+
13
+ @tools.tool()
14
+ async def search_orders(customer_id: int) -> dict:
15
+ ...
16
+
17
+ result = await search_orders(customer_id=912) # recorded or replayed
18
+ """
19
+
20
+ import dataclasses
21
+ import functools
22
+ import inspect
23
+ from collections.abc import Callable
24
+ from datetime import UTC, date, datetime
25
+ from enum import Enum
26
+ from typing import Any
27
+ from uuid import UUID
28
+
29
+ from pydantic import BaseModel
30
+
31
+ from tracefork.boundaries import BoundaryRuntime
32
+ from tracefork.boundaries.base import LiveCall
33
+ from tracefork.canonicalization import canonical_json
34
+ from tracefork.errors import AdapterError
35
+ from tracefork.models import BoundaryRequest, BoundaryResponse
36
+
37
+ TOOL_BOUNDARY_TYPE = "tool.python"
38
+
39
+
40
+ class PythonToolHandler:
41
+ """Generic boundary handler for Python tools: run live, restore verbatim."""
42
+
43
+ async def execute(self, request: BoundaryRequest, call_live: LiveCall) -> BoundaryResponse:
44
+ return BoundaryResponse(response=await call_live(), metadata={})
45
+
46
+ def restore(self, response: Any, metadata: dict[str, Any]) -> Any:
47
+ return response
48
+
49
+
50
+ def canonicalize_argument(value: Any) -> Any:
51
+ """Convert a tool argument to a canonical, JSON-ready value (TF-063).
52
+
53
+ Supports primitives, datetimes, UUIDs, enums, dataclasses, Pydantic
54
+ models, mappings and sequences (including deterministically ordered sets).
55
+ Anything else raises ``AdapterError`` — objects are never silently
56
+ stringified.
57
+ """
58
+ if value is None or isinstance(value, str | bool):
59
+ return value
60
+ if isinstance(value, int | float):
61
+ return value # finite check happens in canonicalization.json
62
+ if isinstance(value, datetime):
63
+ if value.tzinfo is None:
64
+ value = value.replace(tzinfo=UTC)
65
+ return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
66
+ if isinstance(value, date):
67
+ return value.isoformat()
68
+ if isinstance(value, UUID):
69
+ return str(value)
70
+ if isinstance(value, Enum):
71
+ return canonicalize_argument(value.value)
72
+ if dataclasses.is_dataclass(value) and not isinstance(value, type):
73
+ return {
74
+ field.name: canonicalize_argument(getattr(value, field.name))
75
+ for field in dataclasses.fields(value)
76
+ }
77
+ if isinstance(value, BaseModel):
78
+ return canonicalize_argument(value.model_dump(mode="python"))
79
+ if isinstance(value, dict):
80
+ return {str(key): canonicalize_argument(item) for key, item in value.items()}
81
+ if isinstance(value, list | tuple):
82
+ return [canonicalize_argument(item) for item in value]
83
+ if isinstance(value, set | frozenset):
84
+ items = [canonicalize_argument(item) for item in value]
85
+ return sorted(items, key=canonical_json)
86
+ msg = (
87
+ f"tool argument of type {type(value).__name__} cannot be canonicalized; "
88
+ "pass JSON-compatible data, dataclasses, Pydantic models, enums, "
89
+ "datetimes, UUIDs or collections thereof"
90
+ )
91
+ raise AdapterError(msg)
92
+
93
+
94
+ class WrappedTool:
95
+ """A Python function routed through the boundary runtime.
96
+
97
+ Calling the wrapper records the call (RECORD mode), replays it
98
+ (REPLAY + policy REPLAY), or executes it live per the active context. The
99
+ real function body is unreachable on the replay path (TF-062).
100
+ """
101
+
102
+ def __init__(
103
+ self,
104
+ func: Callable[..., Any],
105
+ runtime: BoundaryRuntime,
106
+ name: str | None = None,
107
+ boundary_type: str = TOOL_BOUNDARY_TYPE,
108
+ ) -> None:
109
+ self._func = func
110
+ self._runtime = runtime
111
+ self.tool_name = name if name is not None else f"{func.__module__}.{func.__qualname__}"
112
+ self._boundary_type = boundary_type
113
+ functools.update_wrapper(self, func, updated=())
114
+
115
+ @property
116
+ def boundary_type(self) -> str:
117
+ return self._boundary_type
118
+
119
+ @property
120
+ def func(self) -> Callable[..., Any]:
121
+ return self._func
122
+
123
+ async def __call__(self, *args: Any, **kwargs: Any) -> Any:
124
+ request = {
125
+ "args": [canonicalize_argument(arg) for arg in args],
126
+ "kwargs": {str(key): canonicalize_argument(value) for key, value in kwargs.items()},
127
+ }
128
+
129
+ async def call_live() -> Any:
130
+ result = self._func(*args, **kwargs)
131
+ if inspect.isawaitable(result):
132
+ result = await result
133
+ return result
134
+
135
+ return await self._runtime.invoke(self.boundary_type, self.tool_name, request, call_live)
136
+
137
+
138
+ class ToolBox:
139
+ """Binds tool decorators to a boundary runtime and registers the handler.
140
+
141
+ Dependency-injected by design: there is no global registry and no hidden
142
+ mutation. The shared :class:`PythonToolHandler` is registered once.
143
+ """
144
+
145
+ def __init__(self, runtime: BoundaryRuntime) -> None:
146
+ self._runtime = runtime
147
+ self._handler_registered = False
148
+
149
+ def tool(
150
+ self, *, name: str | None = None, boundary_type: str = TOOL_BOUNDARY_TYPE
151
+ ) -> Callable[[Callable[..., Any]], WrappedTool]:
152
+ """Decorate an async or sync function as a TraceFork tool.
153
+
154
+ ``boundary_type`` re-labels the boundary family (e.g. ``llm.demo`` for
155
+ a scripted LLM); it decides the span kind and replay family.
156
+ """
157
+
158
+ def decorator(func: Callable[..., Any]) -> WrappedTool:
159
+ return self.wrap(func, name=name, boundary_type=boundary_type)
160
+
161
+ return decorator
162
+
163
+ def wrap(
164
+ self,
165
+ func: Callable[..., Any],
166
+ *,
167
+ name: str | None = None,
168
+ boundary_type: str = TOOL_BOUNDARY_TYPE,
169
+ ) -> WrappedTool:
170
+ """Wrap an existing function as a TraceFork tool."""
171
+ self._ensure_handler(boundary_type)
172
+ return WrappedTool(func, self._runtime, name=name, boundary_type=boundary_type)
173
+
174
+ def _ensure_handler(self, boundary_type: str) -> None:
175
+ if boundary_type == TOOL_BOUNDARY_TYPE:
176
+ if not self._handler_registered:
177
+ self._runtime.registry.register(boundary_type, PythonToolHandler())
178
+ self._handler_registered = True
179
+ return
180
+ # Custom boundary types get their own handler registration.
181
+ self._runtime.registry.register(boundary_type, PythonToolHandler())
182
+
183
+
184
+ __all__ = [
185
+ "TOOL_BOUNDARY_TYPE",
186
+ "PythonToolHandler",
187
+ "ToolBox",
188
+ "WrappedTool",
189
+ "canonicalize_argument",
190
+ ]