"""Blocking bridges that keep async resources on one persistent event loop."""
from __future__ import annotations
import asyncio
import inspect
import sys
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator
from concurrent.futures import Future
from contextlib import AbstractAsyncContextManager, AbstractContextManager
from contextvars import ContextVar
from functools import wraps
from queue import Queue
from threading import Event, Lock, Thread, get_ident
from typing import Any, Generic, ParamSpec, TypeVar, cast, overload
from ._compat import BaseExceptionGroup, Self, owned_event_loop
_callback_loop: ContextVar[asyncio.AbstractEventLoop | None] = ContextVar("zett_sync_callback_loop", default=None)
# Type variables for generic services: one result, one arbitrary parameter list.
_T = TypeVar("_T")
_P = ParamSpec("_P")
# Method-call signatures preserve both the callee parameters and its result.
_ParametersT = ParamSpec("_ParametersT")
_ResultT = TypeVar("_ResultT")
[docs]
class SyncRuntime(AbstractContextManager):
"""Own one background event loop for all blocking calls in a scope.
Reuse this scope when calling a provider, Agent, or service repeatedly.
Calls preserve context variables and propagate original exceptions. Different
calling threads may share a runtime. Closing cancels streams, awaits their
cleanup, and shuts down the loop and its worker pool. Objects passed into
this runtime remain caller-owned: close their SDK/database resources first.
Construction starts a dedicated thread immediately. A runtime is not a
process-wide singleton: multiple views may share one runtime, or independent
runtimes may own separate threads. Repeated calls reuse the same loop.
Examples:
Call any asynchronous function without writing an async entry point::
with SyncRuntime() as runtime:
reply = runtime.call(agent.run, "Hello")
with runtime.stream(agent.stream, "Continue") as events:
for event in events:
print(event.type)
.. warning::
Do not move already-used asynchronous SDK clients between event loops.
Construct and use them within the same runtime. A blocking call from
the runtime's own loop is rejected instead of deadlocking.
"""
# Persistent loop used for all operations submitted to this runtime.
_loop: asyncio.AbstractEventLoop
# Loop-thread identity used to reject blocking calls from that same thread.
_thread_id: int
# Owned background thread, joined only after async resource cleanup.
_thread: Thread
def __init__(self) -> None:
self._lock = Lock()
self._closed = False
self._ready = Event()
self._streams: set[SyncStream[Any]] = set()
self._thread = Thread(target=self._serve, name="zett-agent-sync", daemon=True)
self._thread.start()
self._ready.wait()
def _serve(self) -> None:
with owned_event_loop() as loop:
self._loop = loop
self._thread_id = get_ident()
self._ready.set()
self._loop.run_forever()
def _check_thread(self) -> None:
if get_ident() == self._thread_id:
raise RuntimeError("A synchronous call cannot block its own event loop; use await instead")
def _check_close(self) -> None:
self._check_thread()
if _callback_loop.get() is self._loop:
raise RuntimeError(
"A synchronous callback cannot close its own runtime or stream; return to the caller first"
)
def _submit(self, function: Callable[[], Coroutine[Any, Any, _T]]) -> Future[_T]:
self._check_thread()
with self._lock:
if self._closed:
raise RuntimeError("SyncRuntime is closed")
coroutine = function()
try:
return asyncio.run_coroutine_threadsafe(coroutine, self._loop)
except BaseException:
coroutine.close()
raise
@overload
def call(self, function: Callable[_P, Awaitable[_T]], /, *args: _P.args, **kwargs: _P.kwargs) -> _T: ...
@overload
def call(self, function: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> _T: ...
[docs]
def call(self, function: Callable[_P, Any], /, *args: _P.args, **kwargs: _P.kwargs) -> Any:
"""Invoke a callable on the owned loop and wait for its final value.
Accept a callable and its arguments, not an already-created coroutine.
If the callable returns an awaitable, await it before returning. Use :meth:`stream`
for async iterators and :meth:`wrap` for a complete service object.
KeyboardInterrupt requests cancellation of the submitted operation.
Ordinary callables execute directly on the loop thread, not in a worker.
Do not submit long-running blocking I/O here; use an async wrapper with
asyncio.to_thread when that operation must run outside the loop.
"""
async def invoke() -> Any:
result = function(*args, **kwargs)
return await result if inspect.isawaitable(result) else result
future = self._submit(invoke)
try:
return future.result()
except BaseException:
future.cancel()
raise
[docs]
def stream(
self, function: Callable[_P, AsyncIterator[_T]], /, *args: _P.args, **kwargs: _P.kwargs
) -> SyncStream[_T]:
"""Create a lazy, closeable iterator with one owner task for its lifetime.
Each next() requests exactly one item; this bridge does not prefetch.
Provider SDKs and transports may maintain their own buffers.
Use a with block when breaking early. Closing cancels an
in-flight next() and waits for the producer's finally blocks.
"""
self._check_thread()
with self._lock:
if self._closed:
raise RuntimeError("SyncRuntime is closed")
stream = SyncStream(self, lambda: function(*args, **kwargs))
self._streams.add(stream)
return stream
[docs]
def wrap(self, value: _T) -> SyncObject[_T]:
"""Expose an object's async methods as blocking methods on this runtime."""
return SyncObject(value, runtime=self)
[docs]
def context(self, value: AbstractAsyncContextManager[_T]) -> SyncContext[_T]:
"""Return a synchronous with adapter; enter and exit share one async task.
The value yielded by the source remains unchanged. Its asynchronous
methods still need runtime.call; they are not automatically wrapped.
"""
return SyncContext(self, value)
[docs]
def close(self) -> None:
"""Cancel owned streams and join the loop thread; repeated calls are safe."""
self._check_close()
with self._lock:
if self._closed:
return
self._closed = True
streams = tuple(self._streams)
errors: list[BaseException] = []
try:
for stream in streams:
try:
stream.close()
except BaseException as error:
errors.append(error)
finally:
self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join()
if errors:
raise BaseExceptionGroup("Synchronous stream cleanup failed", errors)
def __enter__(self) -> Self:
if self._closed:
raise RuntimeError("SyncRuntime is closed")
return self
def __exit__(self, *exc: object) -> None:
self.close()
[docs]
class SyncStream(Iterator[_T], AbstractContextManager, Generic[_T]):
"""Pull async events synchronously, preserving one task and bounded memory.
A stream has one consumer. Another thread may call close() to interrupt a
blocked consumer. Exhaustion, errors, and close release the source iterator.
Do not rely on garbage collection to close an abandoned stream.
"""
def __init__(self, runtime: SyncRuntime, factory: Callable[[], AsyncIterator[_T]]) -> None:
self._runtime = runtime
self._factory = factory
self._lock = Lock()
self._consumer = Lock()
self._closed = False
self._started = False
self._done = Event()
self._error: BaseException | None = None
self._error_delivered = False
self._replies: Queue[tuple[bool, Any]] = Queue()
self._task: asyncio.Task[None] | None = None
self._demand: asyncio.Queue[None] | None = None
async def _produce(self) -> None:
self._task = asyncio.current_task()
self._demand = asyncio.Queue()
terminal: BaseException | None = None
try:
if self._closed:
return
# The first next() starts production; subsequent steps need a permit.
source = self._factory()
try:
while True:
try:
item = await anext(source)
except StopAsyncIteration:
break
self._replies.put((True, item))
await self._demand.get()
finally:
close = getattr(source, "aclose", None)
if close is not None:
await close()
except asyncio.CancelledError:
pass
except BaseException as error:
terminal = error
finally:
self._error = terminal
self._replies.put((False, terminal))
self._done.set()
with self._runtime._lock:
self._runtime._streams.discard(self)
def __next__(self) -> _T:
self._runtime._check_thread()
if not self._consumer.acquire(blocking=False):
raise RuntimeError("A SyncStream supports only one consumer at a time")
try:
with self._lock:
if self._closed:
raise StopIteration
if not self._started:
self._runtime._submit(self._produce)
self._started = True
elif not self._done.is_set():
self._runtime._loop.call_soon_threadsafe(self._request_next)
success, value = self._replies.get()
if not success:
self._closed = True
if value is not None:
self._error_delivered = True
raise value
raise StopIteration
return value
except (KeyboardInterrupt, SystemExit):
self.close()
raise
finally:
self._consumer.release()
def _request_next(self) -> None:
if self._demand is not None and not self._done.is_set():
self._demand.put_nowait(None)
[docs]
def close(self) -> None:
"""Stop consumption and wait until the async source has been closed."""
self._runtime._check_close()
with self._lock:
self._closed = True
started = self._started
if started and not self._done.is_set():
def cancel() -> None:
if self._task is not None and not self._task.done():
self._task.cancel()
self._runtime._loop.call_soon_threadsafe(cancel)
self._done.wait()
with self._runtime._lock:
self._runtime._streams.discard(self)
if self._error is not None and not self._error_delivered:
self._error_delivered = True
raise self._error
def __enter__(self) -> Self:
return self
def __exit__(self, *exc: object) -> None:
self.close()
[docs]
class SyncObject(AbstractContextManager, Generic[_T]):
"""A blocking view of an existing Agent, provider, tool, or storage object.
The original object remains available as ``wrapped``. Attribute reads and
writes access the original object directly on the caller thread; they do
not provide synchronization. Async method results are awaited and async iterators
become :class:`SyncStream`. Methods and tools can be called normally.
Leaving this scope closes only a runtime created by this view. Explicitly
call provider.aclose() on the view to release a caller-owned SDK client.
Dynamic method lookup is typed as Any; this view does not reproduce an
arbitrary object's static interface. Use SyncAgent for explicitly typed
Agent methods, or SyncRuntime.call for a typed callable boundary.
Examples:
Use any provider's ordinary API synchronously::
with provider.sync() as blocking:
try:
with blocking.stream(request) as events:
for event in events:
print(event)
finally:
blocking.aclose()
"""
#: Original service object; its resources remain caller-owned.
wrapped: _T
#: Shared or privately owned bridge used to execute method calls.
runtime: SyncRuntime
# True only when this view must close its runtime on context-manager exit.
_owns_runtime: bool
def __init__(self, wrapped: _T, *, runtime: SyncRuntime | None = None) -> None:
object.__setattr__(self, "wrapped", wrapped)
object.__setattr__(self, "runtime", runtime if runtime is not None else SyncRuntime())
object.__setattr__(self, "_owns_runtime", runtime is None)
@overload
def _invoke(
self,
function: Callable[_ParametersT, AsyncIterator[_ResultT] | Awaitable[AsyncIterator[_ResultT]]],
*args: _ParametersT.args,
**kwargs: _ParametersT.kwargs,
) -> SyncStream[_ResultT]: ...
@overload
def _invoke(
self,
function: Callable[
_ParametersT, AbstractAsyncContextManager[_ResultT] | Awaitable[AbstractAsyncContextManager[_ResultT]]
],
*args: _ParametersT.args,
**kwargs: _ParametersT.kwargs,
) -> SyncContext[_ResultT]: ...
@overload
def _invoke(
self,
function: Callable[_ParametersT, Awaitable[_ResultT]],
*args: _ParametersT.args,
**kwargs: _ParametersT.kwargs,
) -> _ResultT: ...
@overload
def _invoke(
self,
function: Callable[_ParametersT, _ResultT],
*args: _ParametersT.args,
**kwargs: _ParametersT.kwargs,
) -> _ResultT: ...
def _invoke(
self,
function: Callable[_ParametersT, _ResultT],
*args: _ParametersT.args,
**kwargs: _ParametersT.kwargs,
) -> object:
"""Preserve callable parameters and expose its blocking result type.
Awaitables are resolved, asynchronous iterators become SyncStream,
and asynchronous context managers become SyncContext. Specific
overloads precede the ordinary-result fallback to describe these
conversions instead of claiming that the return type is unchanged.
"""
if inspect.isasyncgenfunction(function):
return self.runtime.stream(function, *args, **kwargs)
value = self.runtime.call(function, *args, **kwargs)
if isinstance(value, AsyncIterator):
return self.runtime.stream(lambda: value)
if isinstance(value, AbstractAsyncContextManager):
return self.runtime.context(value)
return value
def __getattr__(self, name: str) -> Any:
value = getattr(self.wrapped, name)
if not callable(value):
return value
@wraps(value)
def invoke(*args: Any, **kwargs: Any) -> Any:
return self._invoke(value, *args, **kwargs)
return invoke
def __setattr__(self, name: str, value: Any) -> None:
setattr(self.wrapped, name, value)
def __call__(self, *args: Any, **kwargs: Any) -> Any:
# Wrapping something calls it, so it is callable whenever this is reached; ``_T`` itself is unbounded.
function = cast("Callable[..., Any]", self.wrapped)
return self._invoke(function, *args, **kwargs)
def __enter__(self) -> Self:
self.runtime.__enter__()
return self
def __exit__(self, *exc: object) -> None:
if self._owns_runtime:
self.runtime.close()
class SyncMethodsMixin:
"""Expose all asynchronous instance methods through a blocking view."""
__slots__ = ()
def sync(self, *, runtime: SyncRuntime | None = None) -> SyncObject[Self]:
"""Return a context-managed synchronous view; optionally share a runtime.
Omitting runtime creates a new background thread for this view. Passing
an existing runtime reuses its loop and leaves its shutdown to the caller.
This API is for service entry points, not Extension lifecycle hooks.
Examples:
Use a model or tool from a plain function::
with read_file.sync() as read:
result = read({"path": "README.md"})
"""
return SyncObject(self, runtime=runtime)
[docs]
class SyncContext(AbstractContextManager, Generic[_T]):
"""Adapt async with to ordinary with without changing the yielded value.
One async task owns both __aenter__ and __aexit__. This matters for task
groups and cancellation scopes that must be exited by their entering task.
An exception from the synchronous with body is passed to __aexit__; its
truthy return suppresses the exception just as async with would.
The supplied runtime remains caller-owned. This is a resource-lifecycle
adapter, not an AgentRunContext and not a synchronous proxy for the yielded
resource's methods.
Examples:
Open a resource and invoke one of its asynchronous methods::
with SyncRuntime() as runtime:
with runtime.context(open_connection()) as connection:
runtime.call(connection.send, "Hello")
"""
def __init__(self, runtime: SyncRuntime, source: AbstractAsyncContextManager[_T]) -> None:
self._source = source
self._exception: tuple[Any, Any, Any] = (None, None, None)
self._events = runtime.stream(self._lifecycle)
async def _lifecycle(self) -> AsyncIterator[Any]:
"""Yield the resource, then the exit decision, from one owning task."""
value = await self._source.__aenter__()
try:
yield value
except BaseException:
if not await self._source.__aexit__(*sys.exc_info()):
raise
else:
yield await self._source.__aexit__(*self._exception)
def __enter__(self) -> _T:
return next(self._events)
def __exit__(self, *exc: Any) -> bool:
self._exception = exc
try:
return bool(next(self._events))
finally:
self._events.close()
async def _worker(function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
"""Run a blocking model operation outside the loop.
Shield the worker from the initial cancellation and wait for it to finish
before releasing resources it may still use. Python cannot forcibly stop
arbitrary thread work; blocking operations need their own I/O timeouts.
Extension lifecycle hooks do not use this helper.
"""
def invoke() -> Any:
token = _callback_loop.set(loop)
try:
return function(*args, **kwargs)
finally:
_callback_loop.reset(token)
loop = asyncio.get_running_loop()
task = asyncio.create_task(asyncio.to_thread(invoke))
try:
return await asyncio.shield(task)
except asyncio.CancelledError:
try:
await task
except Exception:
pass
raise
