From 7e64515aa14fb492352633b3e1ecc9150f02f82a Mon Sep 17 00:00:00 2001 From: shuaixr <1025sxr@gmail.com> Date: Mon, 13 Apr 2026 00:09:17 +0800 Subject: [PATCH] feat: implement Agent.stop() in python sdk --- CHANGELOG.md | 33 +++++++ sdk/python/agentfield/agent.py | 99 +++++++++++++++++++ .../tests/test_agent_graceful_shutdown.py | 45 ++++++++- 3 files changed, 173 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d466843a..d628bf25 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,39 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) +## [Unreleased] + +### Added + +- Feat: implement Agent.stop() in python sdk + +Implements `Agent.stop()` method that performs a clean async shutdown of an agent instance: +- Marks agent as shutting down and transitions status to OFFLINE +- Stops heartbeat background worker +- Notifies AgentField control plane of graceful shutdown (best effort) +- Cleans up async execution resources, memory event clients, and connection managers +- Idempotent: repeated calls have no additional effect after the first + +Useful for applications that manage agent lifecycle programmatically (e.g., +context managers, signal handlers, test teardown). Uses try/except around each +cleanup step so failures in one subsystem don't prevent cleanup of others. + +### Testing + +- Test(sdk-python): strengthen Agent.stop() idempotency and branch coverage + +Expanded `test_agent_stop_is_idempotent` with mock assertions verifying that all +cleanup side effects (heartbeat stop, shutdown notification, connection manager +stop, memory client close, async resource cleanup) are invoked exactly once across +two consecutive stop() calls. + +Added `test_agent_stop_skips_shutdown_notification_when_not_connected` to verify +graceful degradation: when `agentfield_connected=False`, the shutdown notification +is skipped but local cleanup still runs. + +Removed obsolete TODO and dead implementation guard (`pytest.skip`); Agent.stop() +is now fully implemented. + ## [0.1.67-rc.1] - 2026-04-11 diff --git a/sdk/python/agentfield/agent.py b/sdk/python/agentfield/agent.py index 33158c40..6bdd846f 100644 --- a/sdk/python/agentfield/agent.py +++ b/sdk/python/agentfield/agent.py @@ -4070,6 +4070,105 @@ def _clear_current(self) -> None: # Also clear from thread-local storage clear_current_agent() + async def stop(self) -> None: + """ + Programmatically stop the agent and clean up resources. + + This method performs a graceful shutdown by: + 1. Marking the agent as shutting down and its status as OFFLINE. + 2. Stopping the heartbeat background worker. + 3. Notifying the AgentField control plane that the agent is shutting down. + 4. Cleaning up resources and event subscriptions. + + The method is idempotent; calling it multiple times has no additional effect. + + Example: + ```python + app = Agent("my_agent") + # ... start agent in a background task or loop ... + + # Later, shut down cleanly + await app.stop() + ``` + """ + if getattr(self, "_shutdown_requested", False): + # Already shutting down or stopped + return + + self._shutdown_requested = True + + from agentfield.types import AgentStatus + + self._current_status = AgentStatus.OFFLINE + + if hasattr(self, "agentfield_handler") and self.agentfield_handler: + try: + self.agentfield_handler.stop_heartbeat() + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Heartbeat stop error during stop(): {e}") + + try: + if ( + getattr(self, "agentfield_connected", False) + and hasattr(self, "client") + and self.client + ): + success = await self.client.notify_graceful_shutdown(self.node_id) + if self.dev_mode: + from agentfield.logger import log_info + + state = "sent" if success else "failed" + log_info(f"Shutdown notification {state}") + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Shutdown notification error during stop(): {e}") + + try: + if getattr(self, "connection_manager", None): + await self.connection_manager.stop() + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Connection manager stop error during stop(): {e}") + + try: + if getattr(self, "memory_event_client", None): + await self.memory_event_client.close() + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Memory event client close error during stop(): {e}") + + try: + await self._cleanup_async_resources() + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Resource cleanup error during stop(): {e}") + + try: + from agentfield.agent_registry import clear_current_agent + + clear_current_agent() + except Exception as e: + if self.dev_mode: + from agentfield.logger import log_error + + log_error(f"Registry clear error during stop(): {e}") + + if self.dev_mode: + from agentfield.logger import log_success + + log_success("Agent programmatically stopped") + def _emit_workflow_event_sync( self, context: ExecutionContext, diff --git a/sdk/python/tests/test_agent_graceful_shutdown.py b/sdk/python/tests/test_agent_graceful_shutdown.py index 624cc6ad..ecfccf3f 100644 --- a/sdk/python/tests/test_agent_graceful_shutdown.py +++ b/sdk/python/tests/test_agent_graceful_shutdown.py @@ -1,4 +1,3 @@ -# TODO: source bug — see test_agent_stop_is_idempotent # TODO: source bug — see test_graceful_shutdown_cancels_in_flight_tasks_within_deadline # TODO: source bug — see test_graceful_shutdown_force_cancels_tasks_after_timeout @@ -6,7 +5,7 @@ import os import signal from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest @@ -38,12 +37,50 @@ async def test_agent_stop_is_idempotent(): enable_did=False, ) - if not hasattr(agent, "stop"): - pytest.skip("source bug: Agent.stop() is not implemented") + heartbeat_stop = Mock() + notify_shutdown = AsyncMock(return_value=True) + stop_connection_manager = AsyncMock() + close_memory_event_client = AsyncMock() + + agent.agentfield_handler = SimpleNamespace(stop_heartbeat=heartbeat_stop) + agent.agentfield_connected = True + agent.client = SimpleNamespace(notify_graceful_shutdown=notify_shutdown) + agent.connection_manager = SimpleNamespace(stop=stop_connection_manager) + agent.memory_event_client = SimpleNamespace(close=close_memory_event_client) + agent._cleanup_async_resources = AsyncMock() await agent.stop() await agent.stop() + assert agent._shutdown_requested is True + assert agent._current_status == AgentStatus.OFFLINE + heartbeat_stop.assert_called_once() + notify_shutdown.assert_awaited_once_with(agent.node_id) + stop_connection_manager.assert_awaited_once() + close_memory_event_client.assert_awaited_once() + agent._cleanup_async_resources.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_agent_stop_skips_shutdown_notification_when_not_connected(): + agent = Agent( + node_id="shutdown-agent-disconnected", + agentfield_server="http://agentfield", + auto_register=False, + enable_mcp=False, + enable_did=False, + ) + + notify_shutdown = AsyncMock(return_value=True) + agent.agentfield_connected = False + agent.client = SimpleNamespace(notify_graceful_shutdown=notify_shutdown) + agent._cleanup_async_resources = AsyncMock() + + await agent.stop() + + notify_shutdown.assert_not_awaited() + agent._cleanup_async_resources.assert_awaited_once() + def test_fast_lifecycle_signal_handler_marks_shutdown_and_notifies(monkeypatch): agent = make_shutdown_agent()