feat(byo): add lifespan hook for BYO agents #1064
Conversation
f5d469f to
1febeb3
Compare
There was a problem hiding this comment.
Pull Request Overview
This PR adds support for optional user-defined lifespan context managers in the KAgent ADK framework. Users can now define a lifespan function in their agent module that will be executed during the FastAPI application lifecycle, enabling custom setup and teardown logic.
Key Changes:
- Added a sample
lifespan.pymodule demonstrating the lifespan pattern with setup and teardown logging - Extended
KAgentAppto accept an optionallifespanparameter and compose it with the built-in token service lifespan - Modified the CLI to dynamically import and use user-defined lifespan functions from agent modules
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| python/samples/adk/basic/basic/lifespan.py | Sample implementation of user-defined lifespan with logging |
| python/samples/adk/basic/basic/init.py | Exports the new lifespan function for discoverability |
| python/packages/kagent-adk/src/kagent/adk/cli.py | Dynamically imports lifespan from agent modules and passes it to KAgentApp |
| python/packages/kagent-adk/src/kagent/adk/_a2a.py | Adds lifespan parameter and composition logic to integrate user-defined and built-in lifespans |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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: |
There was a problem hiding this comment.
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.
| try: | |
| try: | |
| sys.path.insert(0, working_dir) |
There was a problem hiding this comment.
AgentLoader.load_agent() already handles adding the working directory to sys.path internally. Since our import_module call happens after that call, we inherit the same sys.path setup.
| if self._lifespan is not None: | ||
| app = FastAPI(lifespan=self._lifespan) | ||
| else: | ||
| app = FastAPI() |
There was a problem hiding this comment.
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.
| 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) |
c725353 to
3cd7e32
Compare
|
@EItanya Any thoughts on this one? |
f0a677c to
59cc016
Compare
EItanya
left a comment
There was a problem hiding this comment.
Hey there, sorry for taking so long to review this, I was at KubeCon, and then I looked at this and the async context Python logic scared me 😆 . I think at a glance it makes sense to me. You do a magic import here because of how the agent is already magically imported?
Also, have you seen this? https://github.com/uriyyo/fastapi-lifespan-manager. Not saying we should use it directly but could get inspired.
I got it from the following issue: fastapi/fastapi#9397
59cc016 to
6330b8d
Compare
I'm looking for a way to be able to run background code in a BYO agent. Whilst this is possible by setting up your own event loop, it would be nice if we could hook into the existing [FastAPI lifespan](https://fastapi.tiangolo.com/advanced/events/#lifespan) instead. So I figured I'd propose doing just that... That being said, I'm completely open to alternatives if anyone knows a better way of doing this 😄 Signed-off-by: Brian Fox <878612+onematchfox@users.noreply.github.com>
6330b8d to
54f7c24
Compare
Signed-off-by: Brian Fox <878612+onematchfox@users.noreply.github.com>
54f7c24 to
bcf8f28
Compare
No worries. Yes, ultimately I'm just trying to replicate what the ADK
I hadn't seen this before and I like the approach. Used the inspiration to refactor in bcf8f28. |
I'm looking for a way to be able to run background process(es) in a BYO agent. Whilst this is possible by setting up your own event loop, I thought it might be nicer(?) if we could hook into the existing FastAPI lifespan instead. So I figured I'd propose doing just that... That being said, I'm completely open to alternatives if anyone knows a better way of doing this 😄