"""Session-owned Python tool environments with a shared invocation boundary."""
from __future__ import annotations
import asyncio
from collections.abc import Callable, Mapping, Sequence
from contextvars import ContextVar, copy_context
from copy import deepcopy
from dataclasses import dataclass, replace
from pathlib import Path
from threading import RLock, Thread
from typing import Any, Literal
from ..agent import AgentRunContext
from ..ids import new_uuid7
from ..messages import ImageContent, TextContent, ToolCall, ToolMessage
from ..tools.base import (
_SYNC_TOOL_RUNNER,
AgentTool,
ToolResult,
_bind_tool_call,
_reset_tool_call,
current_tool_call_id,
)
from ..tools.invocation import invoke_tool
from .enums import ToolCallStatus, ToolKind, ToolOrigin
from .errors import ToolErrorReport
from .profiling import MAX_RETAINED_CALLS, ActiveToolCall, JITHarnessProfile, ToolProfiler, ToolProfileStats
from .python import _load_python_tool
from .results import ToolCatalogEntry, ToolInspection, ToolMutationResult
from .snapshot import (
_ENTRYPOINT_FILE,
_MANAGEMENT_NAMES,
JITHarnessSnapshot,
PythonToolSource,
TextToolSource,
_validate_file_path,
_validate_name,
)
from .text import load_text_tool
@dataclass(frozen=True)
class HarnessBinding:
"""One request's harness binding and loop, published as a contextvar.
Internal: composition entry points read it to decide whether a call belongs
to the running request.
"""
harness: JITHarness
context: AgentRunContext | None
loop: asyncio.AbstractEventLoop
_BINDING: ContextVar[HarnessBinding | None] = ContextVar("zett_jit_binding", default=None)
_CALL_STACK: ContextVar[tuple[str, ...]] = ContextVar("zett_jit_stack", default=())
_PROFILE_CALL: ContextVar[tuple[JITHarness, ActiveToolCall] | None] = ContextVar("zett_jit_profile_call", default=None)
async def _run_sync_tool(function: Callable[..., Any], **arguments: Any) -> ToolResult:
"""Use one worker per active sync invocation, never a pool parents can block.
Cancelling the await cannot stop Python's worker. It is daemonized, receives
an independent Context, and can no longer dispatch after request detachment.
"""
loop = asyncio.get_running_loop()
completed = loop.create_future()
copied = copy_context()
def settle(value: Any, error: BaseException | None) -> None:
if not completed.done():
if error is None:
completed.set_result(value)
else:
completed.set_exception(error)
def run() -> None:
try:
value, error = copied.run(function, **arguments), None
except BaseException as failure:
value, error = None, failure
try:
loop.call_soon_threadsafe(settle, value, error)
except RuntimeError:
# An uninterruptible worker may finish after its loop has closed.
pass
Thread(target=run, name="zett-jit-tool", daemon=True).start()
return await completed
[docs]
class JITHarness:
"""An independently evolving toolbox of Python and Markdown capabilities.
Args:
snapshot: Definitions to activate. Activation executes Python module
bodies, including imports, but does not call tool entry points.
tools: Application-owned base tools; they are never serialized as code.
manage_base_tools: Wrap base tools for unloading and profiling. False
leaves their request registrations intact and skips their profiling.
max_call_depth: Bound nested composition, including recursive calls.
error_capacity: Maximum distinct recent errors across this harness.
error_message_chars: Maximum characters retained in each error message.
call_capacity: Maximum finished call records retained, oldest dropped
first; lifetime aggregates survive eviction.
Note:
This is not a Python sandbox. Source has the process's filesystem,
network, and import privileges. Unload controls harness dispatch, not
what arbitrary Python can do. Module globals are not persisted.
Session-local profiles retain a bounded window of managed dispatches
without payloads, independently of executable snapshots; inspect them
with profile().
Examples:
Save definitions for another session::
harness = JITHarness(tools=[read_file])
harness.create_tool(source, entrypoint="preview_file")
harness.save("harness.json")
"""
def __init__(
self,
snapshot: JITHarnessSnapshot | None = None,
*,
tools: Sequence[AgentTool] = (),
manage_base_tools: bool = True,
max_call_depth: int = 32,
error_capacity: int = 100,
error_message_chars: int = 4000,
call_capacity: int = MAX_RETAINED_CALLS,
) -> None:
if max_call_depth < 1:
raise ValueError("max_call_depth must be positive")
if not isinstance(manage_base_tools, bool):
raise ValueError("manage_base_tools must be a boolean")
self.manage_base_tools = manage_base_tools
self.max_call_depth = max_call_depth
# Guards every registry mutation and snapshot read below. It is a plain
# re-entrant lock because helpers such as create_tool call each other.
self._lock = RLock()
# True only while a module body executes. It stops that trusted code from
# mutating the registry or dispatching tools through the harness.
self._compiling = False
# The request currently bound by attach(). Every wrapper checks this
# identity, so a stale wrapper cannot run against a later request.
self._context: AgentRunContext | None = None
# Tasks performing composed calls, mapped to the request that owns them.
# detach() cancels them so nested work cannot outlive its request.
self._nested_tasks: dict[asyncio.Task[Any], AgentRunContext] = {}
# Serializable Python or Markdown source for dynamic tools. Unloading keeps
# these entries; this mapping is what snapshots and inspect_tool expose.
self._sources: dict[str, PythonToolSource | TextToolSource] = {}
# Executable Python tools and non-executable Markdown readers. Updates
# replace the entry; unload only hides the name, without reloading source.
self._dynamic: dict[str, AgentTool] = {}
# Base tools bound for the current request: constructor tools plus any a
# tool-registering extension added. Rebuilt on attach, dropped on detach.
self._bases: dict[str, AgentTool] = {}
# Tools passed to the constructor. They are never serialized and are the
# ones that stay known across requests when a request adds no binding.
self._provided: dict[str, AgentTool] = {}
# Control tools registered for the current request; cleared on detach
# because they close over that request's harness state.
self._management: dict[str, AgentTool] = {}
# Wrapper handlers created in this request. Registry synchronization
# recognizes old wrappers (for example a restored policy baseline)
# without mistaking them for a new base implementation or wrapping twice.
self._wrapper_handlers: dict[int, Callable[..., Any]] = {}
# Session-local counters, call records, and the recent-error cache. Kept
# for the lifetime of this harness and never written into snapshots.
self._profiler = ToolProfiler(
error_capacity=error_capacity, error_message_chars=error_message_chars, call_capacity=call_capacity
)
initial = snapshot or JITHarnessSnapshot()
# Every base name this harness has seen, accumulated across requests. An
# unbound name is reported as "missing" instead of unknown.
self._base_names = set(initial.base_tools)
# Names hidden from dispatch while their source is retained. Unload only
# edits this set; loading removes the name again.
self._unloaded = set(initial.unloaded_tools)
records: tuple[PythonToolSource | TextToolSource, ...] = (*initial.tools, *initial.text_tools)
dynamic_names = {record.name for record in records}
for candidate in tools:
_validate_name(candidate.name)
if candidate.name in dynamic_names or candidate.name in self._provided:
raise ValueError(f"Duplicate tool: {candidate.name}")
if candidate.local_tool_search:
raise ValueError("JIT harness does not manage client-side tool search")
self._provided[candidate.name] = replace(candidate, parameters=deepcopy(candidate.parameters))
self._base_names.update(self._provided)
self._bases = dict(self._provided)
for record in initial.tools:
registered, rebuilt = self._compile(record.source, record.entrypoint, record.files)
if registered.name != record.name:
raise ValueError(f"Source entrypoint does not match stored tool name {record.name!r}")
self._sources[record.name] = rebuilt
self._dynamic[record.name] = registered
for text_record in initial.text_tools:
self._sources[text_record.name] = text_record.model_copy(deep=True)
self._dynamic[text_record.name] = load_text_tool(text_record)
def _compile(
self, source: str, entrypoint: str, files: Mapping[str, str] | None = None
) -> tuple[AgentTool, PythonToolSource]:
self._compiling = True
try:
return _load_python_tool(source, entrypoint, self, files)
finally:
self._compiling = False
def _check_mutation(self) -> None:
if self._compiling:
raise RuntimeError("Cannot change the harness while loading Python source")
binding = _BINDING.get()
if binding is not None and binding.harness is self and binding.context is not self._context:
raise RuntimeError("This tool's request has already closed")
def _known(self, name: str) -> None:
if name in _MANAGEMENT_NAMES:
raise ValueError("JIT management tools cannot be changed or composed")
if name not in self._sources and name not in self._base_names:
raise ValueError(f"Unknown tool: {name}")
[docs]
def origin(self, name: str) -> ToolOrigin:
"""Return which registry currently serves one tool name.
Labels observability events and profile records: ``BASE`` for application
tools, ``PYTHON`` for executable dynamic tools, ``TEXT`` for Markdown
workflows, ``MANAGEMENT`` for
the harness control functions, and ``UNKNOWN`` for a name this harness has
never seen. The answer does not depend on the request binding.
"""
with self._lock:
return self._resolve_origin(name)
def _resolve_origin(self, name: str) -> ToolOrigin:
if name in _MANAGEMENT_NAMES:
return ToolOrigin.MANAGEMENT
if name in self._sources:
return ToolOrigin.TEXT if isinstance(self._sources[name], TextToolSource) else ToolOrigin.PYTHON
if name in self._base_names:
return ToolOrigin.BASE
return ToolOrigin.UNKNOWN
def _raw_tool(self, name: str, context: AgentRunContext | None) -> AgentTool:
if self._compiling:
raise RuntimeError("Cannot call tools while loading Python source")
self._known(name)
if context is not self._context:
raise RuntimeError("This tool's request has already closed")
if self._is_unloaded(name):
raise ValueError(f"Tool {name!r} is unloaded")
if context is not None and name not in context.tools:
raise ValueError(f"Tool {name!r} is not registered in the current request")
registered = self._dynamic.get(name) or self._bases.get(name)
if registered is None:
raise ValueError(f"Missing base tool: {name}")
return registered
def _is_unloaded(self, name: str) -> bool:
return name in self._unloaded and (self.manage_base_tools or name in self._sources)
def _wrapper(self, registered: AgentTool, context: AgentRunContext | None) -> AgentTool:
name = registered.name
async def guarded(**arguments: Any) -> ToolResult:
# Resolve on every invocation: an old wrapper cannot call a tool
# after unload or retain a superseded implementation.
return await self._profiled_call(name, context, arguments)
if context is not None:
self._wrapper_handlers[id(guarded)] = guarded
# Preserve provider deferral: a base tool marked deferred must keep
# reaching the provider as such, and the manager never defers its own.
return replace(registered, handler=guarded, parameters=deepcopy(registered.parameters))
async def _execute_handler(
self, name: str, target: AgentTool, context: AgentRunContext | None, arguments: dict[str, Any]
) -> ToolResult:
binding_token = _BINDING.set(HarnessBinding(self, context, asyncio.get_running_loop()))
stack_token = _CALL_STACK.set((*_CALL_STACK.get(), name))
worker_token = _SYNC_TOOL_RUNNER.set(_run_sync_tool)
try:
return await target(arguments)
finally:
_SYNC_TOOL_RUNNER.reset(worker_token)
_CALL_STACK.reset(stack_token)
_BINDING.reset(binding_token)
async def _profiled_call(self, name: str, context: AgentRunContext | None, arguments: dict[str, Any]) -> ToolResult:
# Invalid source-loading calls and stale references cannot append to a
# new request's profile. Every actual dispatch attempt is tracked once.
with self._lock:
if self._compiling:
raise RuntimeError("Cannot call tools while loading Python source")
if context is not self._context:
raise RuntimeError("This tool's request has already closed")
origin = self._resolve_origin(name)
if origin is ToolOrigin.BASE and not self.manage_base_tools:
target = self._raw_tool(name, context)
if len(_CALL_STACK.get()) >= self.max_call_depth:
raise RuntimeError(f"JIT tool call depth exceeded {self.max_call_depth}")
else:
target = None
if target is not None:
return await self._execute_handler(name, target, context, arguments)
previous = _PROFILE_CALL.get()
parent = previous[1] if previous is not None and previous[0] is self else None
call_id = current_tool_call_id()
# Direct Python calls do not bind an ID; do not reuse their parent's ID.
if call_id is None or (parent is not None and call_id == parent.call_id):
call_id = new_uuid7()
active = self._profiler.start(
name,
arguments,
call_id=call_id,
parent=parent,
origin=origin,
)
token = _PROFILE_CALL.set((self, active))
output = None
error: BaseException | None = None
status = ToolCallStatus.REJECTED
try:
with self._lock:
if name in _MANAGEMENT_NAMES:
if context is None or name not in context.tools:
raise ValueError(f"Tool {name!r} is not registered in the current request")
target = self._management[name]
else:
target = self._raw_tool(name, context)
stack = _CALL_STACK.get()
if len(stack) >= self.max_call_depth:
raise RuntimeError(f"JIT tool call depth exceeded {self.max_call_depth}")
status = ToolCallStatus.FAILED
output = await self._execute_handler(name, target, context, arguments)
status = ToolCallStatus.COMPLETED
return output
except asyncio.CancelledError as failure:
status, error = ToolCallStatus.CANCELLED, failure
raise
except BaseException as failure:
error = failure
raise
finally:
try:
self._profiler.finish(active, status, output, error)
finally:
_PROFILE_CALL.reset(token)
def _publish(self, name: str) -> None:
context = self._context
if context is not None:
if self._is_unloaded(name):
context.tools.pop(name, None)
else:
registered = self._dynamic.get(name) or self._bases.get(name)
if registered is not None:
context.tools[name] = (
registered
if name in self._bases and not self.manage_base_tools
else self._wrapper(registered, context)
)
@property
def active(self) -> bool:
"""Return whether this harness is currently bound to one request.
The extension checks this before releasing a session, so it never drops
a harness whose request is still running.
"""
with self._lock:
return self._context is not None
[docs]
def attach(self, context: AgentRunContext, management_tools: Sequence[AgentTool] = ()) -> None:
"""Bind this harness to one request and mount its callable tools.
Rebuilds the base registry from the request's current tools, installs
harness-owned wrappers, hides unloaded tools, and registers the
management tools the model may call. Called by
:class:`~zett_agent.jit_harness.extension.JITHarnessExtension` once per
request before the model runs; attaching twice without :meth:`detach`
raises instead of silently replacing the binding.
Args:
context: Request-scoped tool registry and session identity to bind.
management_tools: Harness control tools to expose to the model.
"""
with self._lock:
if self._context is not None:
raise RuntimeError("This harness already has an active request")
bases = dict(self._provided)
for registered in context.tools.values():
_validate_name(registered.name)
if registered.name in self._dynamic:
raise ValueError(f"Base tool conflicts with dynamic tool: {registered.name}")
if registered.local_tool_search:
raise ValueError("JIT harness does not manage client-side tool search")
bases[registered.name] = registered
self._context = context
self._bases = bases
self._base_names.update(bases)
for name in self._base_names | self._sources.keys():
self._publish(name)
for registered in management_tools:
self._management[registered.name] = registered
context.register_tool(self._wrapper(registered, context))
[docs]
def detach(self, context: AgentRunContext) -> None:
"""Release one request binding and stop tracking its nested calls.
Deliberately tolerant of repeated or late cleanup: a context that is not
the currently bound one is ignored, so releasing an older request cannot
dismantle the request that is running now.
"""
with self._lock:
if self._context is context:
self._context = None
# Extension tools often close over a request. Do not retain
# those callables or allow them to run after their request.
self._bases = dict(self._provided)
self._management.clear()
self._wrapper_handlers.clear()
for task, owner in tuple(self._nested_tasks.items()):
if owner is context:
task.get_loop().call_soon_threadsafe(task.cancel)
def _prepare_tool(
self,
source: str,
*,
kind: ToolKind,
entrypoint: str | None,
name: str | None,
description: str | None,
files: Mapping[str, str] | None = None,
) -> tuple[AgentTool, PythonToolSource | TextToolSource]:
if kind is ToolKind.TEXT:
if entrypoint is not None:
raise ValueError("Text tools do not have a Python entrypoint")
if files:
raise ValueError("Text tools have a single document, not a package")
if name is None or description is None:
raise ValueError("Text tools require name and description")
record = TextToolSource(name=name, source=source, description=description)
return load_text_tool(record), record
if entrypoint is None:
raise ValueError("Python tools require entrypoint")
if name is not None or description is not None:
raise ValueError("Python tool name and description come from the @tool entrypoint")
return self._compile(source, entrypoint, files)
[docs]
def describe(self, name: str) -> ToolCatalogEntry:
"""Return one tool's compact catalog entry.
Holds the name, origin, status, and description, without measurements.
Read inspect_tool() for schema and source, or profile() for statistics.
"""
with self._lock:
self._known(name)
registered = self._dynamic.get(name) or self._bases.get(name)
status: Literal["loaded", "unloaded", "missing", "unavailable"] = (
"unloaded" if self._is_unloaded(name) else "loaded" if registered is not None else "missing"
)
if self._context is not None and status == "loaded" and name not in self._context.tools:
status = "unavailable"
return {
"name": name,
"origin": self._resolve_origin(name),
"status": status,
"description": registered.description if registered is not None else "",
}
[docs]
def definition(self, name: str) -> ToolInspection:
"""Read metadata and the argument schema, without source or measurements."""
with self._lock:
info: ToolInspection = {**self.describe(name)}
registered = self._dynamic.get(name) or self._bases.get(name)
if registered is not None:
info["parameters"] = deepcopy(registered.parameters)
info["guidelines"] = list(registered.guidelines)
source = self._sources.get(name)
if isinstance(source, PythonToolSource):
info["entrypoint"] = source.entrypoint
if source.files:
# A Python tool is a package whose entry point is __init__.
sizes = {path: len(text) for path, text in sorted(source.files.items())}
info["files"] = {_ENTRYPOINT_FILE: len(source.source), **sizes}
return info
[docs]
def source(self, name: str, path: str | None = None) -> str:
"""Read exact retained source of one tool file; base tools have none.
``path`` selects an extra module of a Python tool package, and the
default or ``__init__.py`` reads the entry point module. Text tools,
which have a single document, reject a path.
"""
with self._lock:
self._known(name)
record = self._sources.get(name)
if record is None:
raise ValueError(f"Base tool {name!r} has no retained Python source")
if isinstance(record, TextToolSource):
if path is not None:
raise ValueError("Text tools hold a single document, not a package")
return record.source
if path is None or path == _ENTRYPOINT_FILE:
return record.source
_validate_file_path(path)
try:
return record.files[path]
except KeyError:
raise ValueError(f"Unknown package file: {path}") from None
[docs]
def profile(self, name: str | None = None, *, offset: int = 0, limit: int = 100) -> JITHarnessProfile:
"""Read aggregate counters and a page of finished, payload-free calls.
All managed dispatch attempts through harness wrappers are counted, including
nested calls, validation failures, and cancellations. Middleware and
lifecycle hooks run outside this measured boundary. Histories remain
in this harness across requests/updates/unload, bounded by call_capacity,
but are not restored from executable snapshots. Base calls are excluded when
manage_base_tools=False. Export with profile().model_dump_json().
"""
return self._profiler.report(name, offset, limit)
[docs]
def profile_stats(self, name: str) -> ToolProfileStats:
"""Read a tool's full aggregates without fetching its call records."""
return self._profiler.summary(name)
[docs]
def snapshot(self) -> JITHarnessSnapshot:
"""Copy Python/text definitions and unload state, not runtime globals."""
with self._lock:
tools = tuple(
record.model_copy(deep=True)
for _, record in sorted(self._sources.items())
if isinstance(record, PythonToolSource)
)
return JITHarnessSnapshot(
# Only a package needs the newer reader; single-file snapshots stay
# byte-compatible with earlier releases.
format_version=2 if any(record.files for record in tools) else 1,
tools=tools,
text_tools=tuple(
record.model_copy(deep=True)
for _, record in sorted(self._sources.items())
if isinstance(record, TextToolSource)
),
base_tools=tuple(sorted(self._base_names)),
unloaded_tools=tuple(sorted(self._unloaded)),
)
[docs]
def save(self, path: str | Path) -> None:
"""Export the current snapshot to an explicit, atomically replaced file."""
self.snapshot().save(path)
@staticmethod
def _result_content(registered: AgentTool, output: ToolResult) -> str | list[TextContent | ImageContent]:
if isinstance(output, (TextContent, ImageContent)):
return [output]
if (
isinstance(output, list)
and output
and all(isinstance(part, (TextContent, ImageContent)) for part in output)
):
return output
return registered.serialize_result(output)
