diff --git a/sdk/python/agentfield/agent.py b/sdk/python/agentfield/agent.py index 5f120ea61..857330e23 100644 --- a/sdk/python/agentfield/agent.py +++ b/sdk/python/agentfield/agent.py @@ -1,3 +1,4 @@ +from typing import Coroutine import asyncio import inspect import os @@ -500,6 +501,71 @@ def _bind_trigger_payload( return (), {} +class _NotificationDispatcher: + _SHUTDOWN = object() + + def __init__(self, dev_mode: bool): + self._queue: asyncio.Queue[Any] | None = None + self._dispatcher_task: asyncio.Task[Any] | None = None + self._dev_mode = dev_mode + + def start(self): + if self._dispatcher_task is not None: + return + self._queue = asyncio.Queue() + self._dispatcher_task = asyncio.create_task(self._run()) + + def is_start(self): + return self._dispatcher_task is not None + + def submit(self, coro_factory: Callable[[], Coroutine[Any, Any, None]]): + if self._queue is None: + if self._dev_mode: + log_error( + "Coroutine factory submitted before _NotoficationDispatcher even start" + ) + return + self._queue.put_nowait(coro_factory) + + async def _run(self): + if self._queue is None or self._dispatcher_task is None: + return + while True: + coro_factory = await self._queue.get() + if coro_factory is _NotificationDispatcher._SHUTDOWN: + self._queue.task_done() + break + try: + await coro_factory() + except Exception as e: + if self._dev_mode: + log_error(f"Notification dilivery failed: {e}") + finally: + self._queue.task_done() + + async def shutdown(self, timeout: int = 5): + if self._dispatcher_task is None or self._queue is None: + return + self._queue.put_nowait(_NotificationDispatcher._SHUTDOWN) + try: + await asyncio.wait_for(self._dispatcher_task, timeout=timeout) + except asyncio.TimeoutError: + self._dispatcher_task.cancel() + try: + await self._dispatcher_task + except asyncio.CancelledError: + pass + except Exception as e: + if self._dev_mode: + log_error(f"Notification dispatcher shutdown failed: {e}") + finally: + self._dispatcher_task = None + + +def _create_coro_factory(coro): + return lambda: coro + + class Agent(FastAPI): """ AgentField Agent - FastAPI subclass for creating AI agent nodes. @@ -711,6 +777,9 @@ def __init__( # prevent GC of fire-and-forget async execution tasks self._background_tasks: set[asyncio.Task] = set() + # Manage background notifications in order + self._notification_dispatcher = _NotificationDispatcher(dev_mode=self.dev_mode) + # Cooperative cancel registry. The control plane's cancel dispatcher # POSTs /_internal/executions/{id}/cancel to signal that the user's # reasoner code should stop. We track the active asyncio.Task per @@ -2245,12 +2314,15 @@ async def _execute_reasoner_endpoint( if hasattr(self, "workflow_handler") and self.workflow_handler: execution_context.reasoner_name = reasoner_id - await self.workflow_handler.notify_call_start( - execution_context.execution_id, - execution_context, - reasoner_id, - payload_dict, - parent_execution_id=execution_context.parent_execution_id, + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_start( + execution_context.execution_id, + execution_context, + reasoner_id, + payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) start_time = time.time() @@ -2358,29 +2430,36 @@ async def _execute_reasoner_endpoint( if hasattr(self, "workflow_handler") and self.workflow_handler: end_time = time.time() - await self.workflow_handler.notify_call_complete( - execution_context.execution_id, - execution_context.workflow_id, - result, - int((end_time - start_time) * 1000), - execution_context, - input_data=payload_dict, - parent_execution_id=execution_context.parent_execution_id, + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_complete( + execution_context.execution_id, + execution_context.workflow_id, + result, + int((end_time - start_time) * 1000), + execution_context, + input_data=payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) return result except asyncio.CancelledError as cancel_err: if hasattr(self, "workflow_handler") and self.workflow_handler: end_time = time.time() - await self.workflow_handler.notify_call_error( - execution_context.execution_id, - execution_context.workflow_id, - "Execution cancelled by upstream client", - int((end_time - start_time) * 1000), - execution_context, - input_data=payload_dict, - parent_execution_id=execution_context.parent_execution_id, + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_error( + execution_context.execution_id, + execution_context.workflow_id, + "Execution cancelled by upstream client", + int((end_time - start_time) * 1000), + execution_context, + input_data=payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) + raise cancel_err except ExecuteError as exec_err: # Propagate upstream HTTP status codes from cross-agent calls. @@ -2388,15 +2467,20 @@ async def _execute_reasoner_endpoint( # (unhandled exception) and then 502 at the outer control plane. if hasattr(self, "workflow_handler") and self.workflow_handler: end_time = time.time() - await self.workflow_handler.notify_call_error( - execution_context.execution_id, - execution_context.workflow_id, - str(exec_err), - int((end_time - start_time) * 1000), - execution_context, - input_data=payload_dict, - parent_execution_id=execution_context.parent_execution_id, + error_msg = str(exec_err) + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_error( + execution_context.execution_id, + execution_context.workflow_id, + error_msg, + int((end_time - start_time) * 1000), + execution_context, + input_data=payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) + detail = {"error": str(exec_err)} if exec_err.error_details: detail["error_details"] = exec_err.error_details @@ -2408,28 +2492,37 @@ async def _execute_reasoner_endpoint( if hasattr(self, "workflow_handler") and self.workflow_handler: end_time = time.time() detail = getattr(http_exc, "detail", None) or str(http_exc) - await self.workflow_handler.notify_call_error( - execution_context.execution_id, - execution_context.workflow_id, - detail, - int((end_time - start_time) * 1000), - execution_context, - input_data=payload_dict, - parent_execution_id=execution_context.parent_execution_id, + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_error( + execution_context.execution_id, + execution_context.workflow_id, + detail, + int((end_time - start_time) * 1000), + execution_context, + input_data=payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) + raise except Exception as exc: if hasattr(self, "workflow_handler") and self.workflow_handler: end_time = time.time() - await self.workflow_handler.notify_call_error( - execution_context.execution_id, - execution_context.workflow_id, - str(exc), - int((end_time - start_time) * 1000), - execution_context, - input_data=payload_dict, - parent_execution_id=execution_context.parent_execution_id, + error_msg = str(exc) + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_error( + execution_context.execution_id, + execution_context.workflow_id, + error_msg, + int((end_time - start_time) * 1000), + execution_context, + input_data=payload_dict, + parent_execution_id=execution_context.parent_execution_id, + ) ) + raise finally: reset_execution_context(context_token) @@ -3111,39 +3204,50 @@ async def _run_async_skill(*args, **kwargs): self._current_execution_context = child_context input_payload = _build_invocation_payload(args, kwargs) - await self.workflow_handler.notify_call_start( - child_context.execution_id, - child_context, - skill_id, - input_payload, - parent_execution_id=current_context.execution_id, + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_start( + child_context.execution_id, + child_context, + skill_id, + input_payload, + parent_execution_id=current_context.execution_id, + ) ) start_time = time.time() try: result = await original_func(*args, **kwargs) duration_ms = int((time.time() - start_time) * 1000) - await self.workflow_handler.notify_call_complete( - child_context.execution_id, - child_context.workflow_id, - result, - duration_ms, - child_context, - input_data=input_payload, - parent_execution_id=current_context.execution_id, + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_complete( + child_context.execution_id, + child_context.workflow_id, + result, + duration_ms, + child_context, + input_data=input_payload, + parent_execution_id=current_context.execution_id, + ) ) + return result except Exception as exc: duration_ms = int((time.time() - start_time) * 1000) - await self.workflow_handler.notify_call_error( - child_context.execution_id, - child_context.workflow_id, - str(exc), - duration_ms, - child_context, - input_data=input_payload, - parent_execution_id=current_context.execution_id, + error_msg = str(exc) + + self._notification_dispatcher.submit( + lambda: self.workflow_handler.notify_call_error( + child_context.execution_id, + child_context.workflow_id, + error_msg, + duration_ms, + child_context, + input_data=input_payload, + parent_execution_id=current_context.execution_id, + ) ) + raise finally: reset_execution_context(token) @@ -4373,6 +4477,29 @@ async def _cleanup_async_resources(self) -> None: if self.dev_mode: log_debug(f"Error cleaning up AsyncExecutionManager: {e}") + if self._background_tasks: + try: + await asyncio.wait_for( + asyncio.gather(*self._background_tasks, return_exceptions=True), + timeout = 5 + ) + if self.dev_mode: + log_debug("Background tasks are cleaned up") + except Exception as e: + if self.dev_mode: + log_error(f"Error cleaning up background tasks: {e}") + finally: + self._background_tasks.clear() + try: + await self._notification_dispatcher.shutdown() + if self.dev_mode: + log_debug( + "Notification dispatcher queue cleared and dispatcher shutdown" + ) + except Exception as e: + if self.dev_mode: + log_error(f"Error while shutdown notification dispatcher: {e}") + if getattr(self, "client", None) is not None: try: await self.client.aclose() diff --git a/sdk/python/agentfield/agent_server.py b/sdk/python/agentfield/agent_server.py index c438c377b..e4e4ed2e2 100644 --- a/sdk/python/agentfield/agent_server.py +++ b/sdk/python/agentfield/agent_server.py @@ -887,6 +887,9 @@ def on_disconnected(): # Start connection manager (non-blocking) connected = await self.agent.connection_manager.start() + # Start notification dispatcher + self.agent._notification_dispatcher.start() + # Connect memory event client - works independently of AgentField server connection if self.agent.memory_event_client: try: diff --git a/sdk/python/agentfield/agent_workflow.py b/sdk/python/agentfield/agent_workflow.py index f99efcdec..2d9a45563 100644 --- a/sdk/python/agentfield/agent_workflow.py +++ b/sdk/python/agentfield/agent_workflow.py @@ -71,12 +71,14 @@ async def execute_with_tracking( start_time = time.time() parent_execution_id = parent_context.execution_id if parent_context else None - await self.notify_call_start( - execution_context.execution_id, - execution_context, - reasoner_name, - input_data, - parent_execution_id=parent_execution_id, + self.agent._notification_dispatcher.submit( + lambda: self.notify_call_start( + execution_context.execution_id, + execution_context, + reasoner_name, + input_data, + parent_execution_id=parent_execution_id, + ) ) try: @@ -84,27 +86,36 @@ async def execute_with_tracking( if inspect.isawaitable(result): result = await result duration_ms = int((time.time() - start_time) * 1000) - await self.notify_call_complete( - execution_context.execution_id, - execution_context.workflow_id, - result, - duration_ms, - execution_context, - input_data=input_data, - parent_execution_id=parent_execution_id, + + self.agent._notification_dispatcher.submit( + lambda: self.notify_call_complete( + execution_context.execution_id, + execution_context.workflow_id, + result, + duration_ms, + execution_context, + input_data=input_data, + parent_execution_id=parent_execution_id, + ) ) + return result except Exception as exc: # pragma: no cover - re-raised duration_ms = int((time.time() - start_time) * 1000) - await self.notify_call_error( - execution_context.execution_id, - execution_context.workflow_id, - str(exc), - duration_ms, - execution_context, - input_data=input_data, - parent_execution_id=parent_execution_id, + error_msg = str(exc) + + self.agent._notification_dispatcher.submit( + lambda: self.notify_call_error( + execution_context.execution_id, + execution_context.workflow_id, + error_msg, + duration_ms, + execution_context, + input_data=input_data, + parent_execution_id=parent_execution_id, + ) ) + raise finally: reset_execution_context(token) @@ -121,7 +132,7 @@ async def notify_call_start( *, parent_execution_id: Optional[str] = None, ) -> None: - await self._emit_execution_transition_log( + self._emit_execution_transition_log( context, reasoner_name, event_type="reasoner.started", @@ -131,6 +142,7 @@ async def notify_call_start( input_data=input_data, parent_execution_id=parent_execution_id, ) + payload = self._build_event_payload( context, reasoner_name, @@ -151,7 +163,7 @@ async def notify_call_complete( input_data: Optional[Dict[str, Any]] = None, parent_execution_id: Optional[str] = None, ) -> None: - await self._emit_execution_transition_log( + self._emit_execution_transition_log( context, context.reasoner_name, event_type="reasoner.completed", @@ -163,6 +175,7 @@ async def notify_call_complete( input_data=input_data, parent_execution_id=parent_execution_id, ) + payload = self._build_event_payload( context, context.reasoner_name, @@ -185,7 +198,7 @@ async def notify_call_error( input_data: Optional[Dict[str, Any]] = None, parent_execution_id: Optional[str] = None, ) -> None: - await self._emit_execution_transition_log( + self._emit_execution_transition_log( context, context.reasoner_name, event_type="reasoner.failed", @@ -291,7 +304,7 @@ async def _ensure_execution_registered( context.execution_id = body.get("execution_id", context.execution_id) context.workflow_id = body.get("workflow_id", context.workflow_id) context.run_id = body.get("run_id", context.run_id) - await self._emit_execution_transition_log( + self._emit_execution_transition_log( context, reasoner_name, event_type="execution.registered", @@ -303,7 +316,7 @@ async def _ensure_execution_registered( except Exception as exc: # pragma: no cover - network failure path if getattr(self.agent, "dev_mode", False): log_warn(f"Workflow registration failed: {exc}") - await self._emit_execution_transition_log( + self._emit_execution_transition_log( context, reasoner_name, event_type="execution.registration.failed", @@ -348,7 +361,7 @@ def _build_event_payload( payload["input_data"] = input_data return payload - async def _emit_execution_transition_log( + def _emit_execution_transition_log( self, context: ExecutionContext, reasoner_name: str, diff --git a/sdk/python/tests/helpers.py b/sdk/python/tests/helpers.py index 71feaf035..99848fc81 100644 --- a/sdk/python/tests/helpers.py +++ b/sdk/python/tests/helpers.py @@ -1,6 +1,9 @@ """Shared testing utilities for AgentField SDK unit tests.""" from __future__ import annotations +from typing import Callable +from typing import Coroutine + import asyncio import threading @@ -92,6 +95,29 @@ def notify_graceful_shutdown_sync(self, node_id: str) -> bool: return True +class InlineTestingDispatcher: + def __init__(self): + self._started = True + + def start(self): + self._started = True + + def is_start(self): + return self._started + + def submit(self, coro_factory: Callable[[], Coroutine[Any, Any, None]]): + coro = coro_factory() + try: + coro.send(None) + except StopIteration: + pass + except RuntimeWarning: + pass + + async def shutdown(self, timeout: int = 5): + self._started = False + + @dataclass class StubAgent: """Light-weight stand-in for Agent used across module tests.""" @@ -113,6 +139,9 @@ class StubAgent: agentfield_connected: bool = True _current_status: AgentStatus = AgentStatus.STARTING callback_candidates: List[str] = field(default_factory=list) + _notification_dispatcher: InlineTestingDispatcher = field( + default_factory=InlineTestingDispatcher + ) def _build_vc_metadata(self): return {"agent_default": True} diff --git a/sdk/python/tests/test_agent_core.py b/sdk/python/tests/test_agent_core.py index 2eaa22384..cb1775485 100644 --- a/sdk/python/tests/test_agent_core.py +++ b/sdk/python/tests/test_agent_core.py @@ -27,6 +27,7 @@ def make_agent_stub(): api_base="http://agentfield/api/v1", _get_auth_headers=lambda: {}, ) + agent._background_tasks = set() return agent @@ -74,11 +75,31 @@ def __init__(self): async def stop(self): self.stopped = True + class DummyNotificationDispatcher: + def __init__(self): + self.stopped = False + + async def shutdown(self): + self.stopped = True + + async def dummy_async_task(sleep_delay: int): + await asyncio.sleep(sleep_delay) + manager = DummyManager() + notification_dispatcher = DummyNotificationDispatcher() agent._async_execution_manager = manager + agent._notification_dispatcher = notification_dispatcher + + for i in range(1,6): + task = asyncio.create_task(dummy_async_task(i)) + agent._background_tasks.add(task) + task.add_done_callback(agent._background_tasks.discard) + await agent._cleanup_async_resources() assert manager.stopped is True assert agent._async_execution_manager is None + assert len(agent._background_tasks) == 0 + assert notification_dispatcher.stopped is True @pytest.mark.asyncio diff --git a/sdk/python/tests/test_agent_integration.py b/sdk/python/tests/test_agent_integration.py index 20e59e4f6..374d0897b 100644 --- a/sdk/python/tests/test_agent_integration.py +++ b/sdk/python/tests/test_agent_integration.py @@ -15,6 +15,8 @@ async def test_agent_reasoner_routing_and_workflow(monkeypatch): agent, agentfield_client = create_test_agent( monkeypatch, callback_url="https://callback.example.com" ) + # Manually start notification dispatcher + agent._notification_dispatcher.start() # Disable async execution for this test to get synchronous 200 responses agent.async_config.enable_async_execution = False # Disable agentfield_server to prevent async callback execution diff --git a/sdk/python/tests/test_agent_server.py b/sdk/python/tests/test_agent_server.py index 7e7dd847e..5c93510f5 100644 --- a/sdk/python/tests/test_agent_server.py +++ b/sdk/python/tests/test_agent_server.py @@ -42,6 +42,9 @@ def make_agent_app(**overrides): resolve_by_execution_id=AsyncMock(return_value=False), ), ) + app._notification_dispatcher = overrides.get( + "_notification_dispatcher", SimpleNamespace(start=MagicMock()) + ) return app diff --git a/sdk/python/tests/test_did_auth_invariants.py b/sdk/python/tests/test_did_auth_invariants.py index dc95f0f69..9e3f4f408 100644 --- a/sdk/python/tests/test_did_auth_invariants.py +++ b/sdk/python/tests/test_did_auth_invariants.py @@ -142,7 +142,9 @@ def test_invariant_signature_is_deterministic_for_same_payload(self, ed25519_key same signature out (Ed25519 determinism property). """ try: - from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey # noqa: F401 + from cryptography.hazmat.primitives.asymmetric.ed25519 import ( + Ed25519PrivateKey, # noqa: F401 + ) except ImportError: pytest.skip("cryptography library not available")