44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
"""Shared, bounded initialization for async agent and tool entry points."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from deerflow.utils.file_io import FileIOPool
|
|
|
|
# Serialize cold imports/client setup without occupying the request loop or
|
|
# authentication's executor. Queued initializers consume no worker threads.
|
|
_INITIALIZATION_IO = FileIOPool(1, 1, "agent-init")
|
|
|
|
|
|
def _initialize_in_worker[Result](
|
|
function: Callable[..., Result], args: tuple[Any, ...], kwargs: dict[str, Any], loop: asyncio.AbstractEventLoop
|
|
) -> tuple[Result, contextvars.Context]:
|
|
from deerflow.mcp.cache import mcp_initialization_loop
|
|
|
|
# Synchronous construction may wait for MCP discovery, while its owning
|
|
# request loop remains free to run discovery and other requests.
|
|
with mcp_initialization_loop(loop):
|
|
result = function(*args, **kwargs)
|
|
return result, contextvars.copy_context()
|
|
|
|
|
|
async def run_background_init[Result](function: Callable[..., Result], /, *args: Any, **kwargs: Any) -> Result:
|
|
"""Run thread-safe initialization with the caller's MCP loop and context.
|
|
|
|
Cancellation drains already-submitted work before propagating and never
|
|
commits its context changes. Successful initialization writes ContextVars
|
|
back to this task so tool registries remain available during execution.
|
|
"""
|
|
result, initialized_context = await _INITIALIZATION_IO.run_until_complete(
|
|
_initialize_in_worker, function, args, kwargs, asyncio.get_running_loop()
|
|
)
|
|
# Ordinary to_thread only copies context into a worker. Tool discovery also
|
|
# sets or clears the deferred registry, which the calling task must inherit.
|
|
for variable, value in initialized_context.items():
|
|
variable.set(value)
|
|
return result
|