Skip to content

feat(byo): add lifespan hook for BYO agents #1064

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+50 −4
Open

feat(byo): add lifespan hook for BYO agents #1064

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 23 additions & 3 deletions python/packages/kagent-adk/src/kagent/adk/_a2a.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import faulthandler
import logging
import os
from typing import Callable, List
from typing import Callable, List, Any, Optional

import httpx
from a2a.server.apps import A2AFastAPIApplication
Expand Down Expand Up @@ -54,12 +54,14 @@ def __init__(
agent_card: AgentCard,
kagent_url: str,
app_name: str,
lifespan: Optional[Callable[[Any], Any]] = None,
plugins: List[BasePlugin] = None,
):
self.root_agent = root_agent
self.kagent_url = kagent_url
self.app_name = app_name
self.agent_card = agent_card
self._lifespan = lifespan
self.plugins = plugins if plugins is not None else []

def build(self) -> FastAPI:
Expand Down Expand Up @@ -101,7 +103,7 @@ def create_runner() -> Runner:
)

faulthandler.enable()
app = FastAPI(lifespan=token_service.lifespan())
app = FastAPI(lifespan=self._compose_lifespan(token_service.lifespan()))

# Health check/readiness probe
app.add_route("/health", methods=["GET"], route=health_check)
Expand Down Expand Up @@ -139,7 +141,10 @@ def create_runner() -> Runner:
)

faulthandler.enable()
app = FastAPI()
if self._lifespan is not None:
app = FastAPI(lifespan=self._lifespan)
else:
app = FastAPI()
Comment on lines +144 to +147
Copy link

Copilot AI Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The build_local() method doesn't compose the user-defined lifespan with any base lifespan like build() does with token_service.lifespan(). While this may be intentional for local mode, it creates inconsistency. If users define a lifespan expecting composition behavior (e.g., wrapping around a base lifespan), it won't work in local mode. Consider either using _compose_lifespan() with a no-op base lifespan or documenting this behavioral difference.

Suggested change
if self._lifespan is not None:
app = FastAPI(lifespan=self._lifespan)
else:
app = FastAPI()
from contextlib import asynccontextmanager
@asynccontextmanager
async def noop_lifespan(app):
yield
composed_lifespan = self._compose_lifespan(noop_lifespan)
app = FastAPI(lifespan=composed_lifespan)

Copilot uses AI. Check for mistakes.

app.add_route("/health", methods=["GET"], route=health_check)
app.add_route("/thread_dump", methods=["GET"], route=thread_dump)
Expand Down Expand Up @@ -186,3 +191,18 @@ async def test(self, task: str):
# Key Concept: is_final_response() marks the concluding message for the turn.
jsn = event.model_dump_json()
logger.info(f" [Event] {jsn}")

def _compose_lifespan(self, base_lifespan: Callable[[Any], Any]) -> Callable[[Any], Any]:
if self._lifespan is None:
return base_lifespan

# Compose base lifespan with optional user-provided lifespan
from contextlib import asynccontextmanager

@asynccontextmanager
async def composed(app):
async with base_lifespan(app):
async with self._lifespan(app):
yield

return composed
13 changes: 12 additions & 1 deletion python/packages/kagent-adk/src/kagent/adk/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import asyncio
import importlib
import sys
Copy link

Copilot AI Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import of 'sys' is not used.

Suggested change
import sys

Copilot uses AI. Check for mistakes.
import json
import logging
import os
Expand Down Expand Up @@ -93,7 +95,16 @@ def run(
agent_card = json.load(f)
agent_card = AgentCard.model_validate(agent_card)

kagent_app = KAgentApp(root_agent, agent_card, app_cfg.url, app_cfg.app_name)
# Attempt to import optional user-defined lifespan(app) from the agent package
lifespan = None
try:
Copy link

Copilot AI Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The dynamic import uses the agent name directly without adjusting sys.path to include working_dir. If the agent module is in a subdirectory (e.g., working_dir/name), this import may fail. Consider using sys.path.insert(0, working_dir) before the import to ensure the module can be found, similar to how AgentLoader handles the working directory.

Suggested change
try:
try:
sys.path.insert(0, working_dir)

Copilot uses AI. Check for mistakes.
module_candidate = importlib.import_module(name)
if hasattr(module_candidate, "lifespan"):
lifespan = module_candidate.lifespan
except Exception:
logger.exception("Failed to load agent module for lifespan")
Copy link

Copilot AI Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error message 'Failed to load agent module for lifespan' is too generic and doesn't indicate which module failed to load. Consider including the module name in the error message: f\"Failed to load agent module '{name}' for lifespan\" to aid in debugging.

Suggested change
logger.exception("Failed to load agent module for lifespan")
logger.exception(f"Failed to load agent module '{name}' for lifespan")

Copilot uses AI. Check for mistakes.

kagent_app = KAgentApp(root_agent, agent_card, app_cfg.url, app_cfg.app_name, lifespan=lifespan)

if local:
logger.info("Running in local mode with InMemorySessionService")
Expand Down
3 changes: 3 additions & 0 deletions python/samples/adk/basic/basic/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1,4 @@
__all__ = ["agent", "lifespan"]

from . import agent
from .lifespan import lifespan
12 changes: 12 additions & 0 deletions python/samples/adk/basic/basic/lifespan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import logging
from contextlib import asynccontextmanager
from typing import Any


@asynccontextmanager
async def lifespan(app: Any):
logging.info("Lifespan: setup")
try:
yield
finally:
logging.info("Lifespan: teardown")
Loading