fix: Updated the startup and shutdown events, now they are using life…

sdk/python/agentfield/agent_server.py

@@ -3,14 +3,15 @@
 import os
 import signal
 import urllib.parse
+from contextlib import AsyncExitStack, asynccontextmanager
 from datetime import datetime
 from typing import Optional
 
 import uvicorn
 from agentfield.agent_utils import AgentUtils
 from agentfield.logger import log_debug, log_error, log_info, log_success, log_warn
 from agentfield.utils import get_free_port
-from fastapi import Request
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.routing import APIRoute
 
@@ -48,15 +49,21 @@ async def agentfield_process_logs(request: Request):
             if not node_logs.logs_enabled():
                 return JSONResponse(
                     status_code=404,
-                    content={"error": "logs_disabled", "message": "Process logs API is disabled"},
+                    content={
+                        "error": "logs_disabled",
+                        "message": "Process logs API is disabled",
+                    },
                 )
             auth = request.headers.get("authorization") or request.headers.get(
                 "Authorization"
             )
             if not node_logs.verify_internal_bearer(auth):
                 return JSONResponse(
                     status_code=401,
-                    content={"error": "unauthorized", "message": "Valid Authorization Bearer required"},
+                    content={
+                        "error": "unauthorized",
+                        "message": "Valid Authorization Bearer required",
+                    },
                 )
             qp = request.query_params
             try:
@@ -107,7 +114,9 @@ async def debug_tasks():
                     name = t.get_name()
                 except Exception:
                     name = "?"
-                buf.write(f"=== Task {name} done={t.done()} cancelled={t.cancelled()} ===\n")
+                buf.write(
+                    f"=== Task {name} done={t.done()} cancelled={t.cancelled()} ===\n"
+                )
                 try:
                     coro = t.get_coro()
                     buf.write(f"coro: {coro!r}\n")
@@ -121,7 +130,9 @@ async def debug_tasks():
                                 f"  {frame.f_code.co_filename}:{frame.f_lineno} in {frame.f_code.co_name}\n"
                             )
                     else:
-                        buf.write("  <no stack — task is suspended on a Future/awaitable>\n")
+                        buf.write(
+                            "  <no stack — task is suspended on a Future/awaitable>\n"
+                        )
                 except Exception as e:
                     buf.write(f"  <stack error: {e}>\n")
                 out.append(buf.getvalue())
@@ -315,7 +326,10 @@ async def approval_webhook(request: Request):
             approval_request_id = body.get("approval_request_id", "")
 
             if not execution_id or not decision:
-                return {"error": "execution_id and decision are required", "status": 400}
+                return {
+                    "error": "execution_id and decision are required",
+                    "status": 400,
+                }
 
             # Parse the raw response field (may be a JSON string or dict)
             raw_response = None
@@ -340,9 +354,13 @@ async def approval_webhook(request: Request):
             # Try to resolve by approval_request_id first, then by execution_id
             resolved = False
             if approval_request_id:
-                resolved = await self.agent._pause_manager.resolve(approval_request_id, result)
+                resolved = await self.agent._pause_manager.resolve(
+                    approval_request_id, result
+                )
             if not resolved and execution_id:
-                resolved = await self.agent._pause_manager.resolve_by_execution_id(execution_id, result)
+                resolved = await self.agent._pause_manager.resolve_by_execution_id(
+                    execution_id, result
+                )
 
             if self.agent.dev_mode:
                 log_debug(
@@ -766,8 +784,27 @@ def serve(
         # Setup fast lifecycle signal handlers
         self.agent.agentfield_handler.setup_fast_lifecycle_signal_handlers()
 
-        # Add startup event handler for resilient lifecycle
-        @self.agent.on_event("startup")
+        @asynccontextmanager
+        async def internal_lifespan(app: FastAPI):
+            # Add startup event handler for resilient lifecycle
+            await startup_resilient_lifecycle()
+            try:
+                yield
+            finally:
+                # Add shutdown event handler for cleanup
+                await shutdown_cleanup()
+
+        existing_lifespan = self.agent.router.lifespan_context
+
+        @asynccontextmanager
+        async def merged_lifespan(app: FastAPI):
+            async with AsyncExitStack() as stack:
+                await stack.enter_async_context(internal_lifespan(app))
+                await stack.enter_async_context(existing_lifespan(app))
+                yield
+
+        self.agent.router.lifespan_context = merged_lifespan
+
         async def startup_resilient_lifecycle():
             """Resilient lifecycle startup: connection manager handles AgentField server connectivity"""
 
@@ -848,8 +885,6 @@ def on_disconnected():
                         "Agent started in local mode - will connect to AgentField server when available"
                     )
 
-        # Add shutdown event handler for cleanup
-        @self.agent.on_event("shutdown")
         async def shutdown_cleanup():
             """Cleanup all resources when FastAPI shuts down"""
 

sdk/python/tests/test_agent_server.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import sys
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -66,6 +67,21 @@ async def _post(app, path, **kwargs):
         return await client.post(path, **kwargs)
 
 
+class _FakeConnectionManager:
+    def __init__(self, agent, config):
+        self.agent = agent
+        self.config = config
+        self.on_connected = None
+        self.on_disconnected = None
+
+    async def start(self):
+        self.agent.lifecycle_events.append("agentfield-start")
+        return False
+
+    async def stop(self):
+        self.agent.lifecycle_events.append("agentfield-stop")
+
+
 # ---------------------------------------------------------------------------
 # Route registration and health endpoint
 # ---------------------------------------------------------------------------
@@ -120,6 +136,58 @@ async def test_info_endpoint():
     assert "registered_at" in data
 
 
+def test_serve_preserves_existing_lifespan_until_shutdown(monkeypatch):
+    events = []
+
+    @asynccontextmanager
+    async def existing_lifespan(app):
+        app.lifecycle_events.append("caller-start")
+        try:
+            yield
+        finally:
+            app.lifecycle_events.append("caller-stop")
+
+    app = make_agent_app()
+    app.router.lifespan_context = existing_lifespan
+    app.lifecycle_events = events
+    app.agentfield_handler = SimpleNamespace(
+        start_heartbeat=lambda interval: events.append("heartbeat-start"),
+        setup_fast_lifecycle_signal_handlers=lambda: events.append(
+            "signal-handlers"
+        ),
+        stop_heartbeat=lambda: events.append("heartbeat-stop"),
+        send_enhanced_heartbeat=AsyncMock(return_value=True),
+        enhanced_heartbeat_loop=AsyncMock(),
+    )
+    app.connection_manager = None
+    app.memory_event_client = None
+    app.client = SimpleNamespace(aclose=AsyncMock())
+
+    def fake_uvicorn_run(served_app, **config):
+        async def exercise_lifespan():
+            async with served_app.router.lifespan_context(served_app):
+                events.append("serving")
+
+        asyncio.run(exercise_lifespan())
+
+    monkeypatch.setattr("agentfield.connection_manager.ConnectionManager", _FakeConnectionManager)
+    monkeypatch.setattr("agentfield.agent_server.uvicorn.run", fake_uvicorn_run)
+
+    AgentServer(app).serve(port=8001)
+
+    assert events == [
+        "heartbeat-start",
+        "signal-handlers",
+        "agentfield-start",
+        "caller-start",
+        "serving",
+        "caller-stop",
+        "agentfield-stop",
+        "heartbeat-stop",
+    ]
+    app.client.aclose.assert_awaited_once()
+
+
 # ---------------------------------------------------------------------------
 # Shutdown endpoint
 # ---------------------------------------------------------------------------