feat(kagent-adk): Use new mcp session for a2a requests (#1220)

Updates `kagent-adk` to allow for a new agent to be used for each runner
instance during an A2A request. This allows each a2a request to use a
new mcp session instead of reusing one single mcp session during the
entire lifetime of the agent.

---------

Signed-off-by: JM Huibonhoa <[email protected]>

Author: jmhbh

python/packages/kagent-adk/src/kagent/adk/_a2a.py

@@ -53,15 +53,26 @@ def thread_dump(request: Request) -> PlainTextResponse:
 class KAgentApp:
     def __init__(
         self,
-        root_agent: BaseAgent,
+        root_agent_factory: Callable[[], BaseAgent],
         agent_card: AgentCard,
         kagent_url: str,
         app_name: str,
         lifespan: Optional[Callable[[Any], Any]] = None,
         plugins: List[BasePlugin] = None,
         stream: bool = False,
     ):
-        self.root_agent = root_agent
+        """Initialize the KAgent application.
+
+        Args:
+            root_agent_factory: Root agent factory function that returns a new agent instance
+            agent_card: Agent card configuration for A2A protocol
+            kagent_url: URL of the KAgent backend server
+            app_name: Application name for identification
+            lifespan: Optional lifespan function
+            plugins: Optional list of plugins
+            stream: Whether to stream the response
+        """
+        self.root_agent_factory = root_agent_factory
         self.kagent_url = kagent_url
         self.app_name = app_name
         self.agent_card = agent_card
@@ -81,9 +92,10 @@ def build(self, local=False) -> FastAPI:
             )
             session_service = KAgentSessionService(http_client)
 
-        adk_app = App(name=self.app_name, root_agent=self.root_agent, plugins=self.plugins)
-
         def create_runner() -> Runner:
+            root_agent = self.root_agent_factory()
+            adk_app = App(name=self.app_name, root_agent=root_agent, plugins=self.plugins)
+
             return Runner(
                 app=adk_app,
                 session_service=session_service,
@@ -138,12 +150,8 @@ async def test(self, task: str):
             session_id=SESSION_ID,
             user_id=USER_ID,
         )
-        if isinstance(self.root_agent, Callable):
-            agent_factory = self.root_agent
-            root_agent = agent_factory()
-        else:
-            root_agent = self.root_agent
 
+        root_agent = self.root_agent_factory()
         runner = Runner(
             agent=root_agent,
             app_name=self.app_name,

python/packages/kagent-adk/src/kagent/adk/_agent_executor.py

@@ -185,6 +185,11 @@ async def execute(
                     logger.error("Failed to publish failure event: %s", enqueue_error, exc_info=True)
         finally:
             clear_kagent_span_attributes(context_token)
+            # close the runner which cleans up the mcptoolsets 
+            # since the runner is created for each a2a request
+            # and the mcptoolsets are not shared between requests
+            # this is necessary to gracefully handle mcp toolset connections
+            await runner.close()
 
     async def _handle_request(
         self,

python/packages/kagent-adk/src/kagent/adk/cli.py

@@ -64,11 +64,16 @@ def static(
     sts_integration = create_sts_integration()
     if sts_integration:
         plugins = [sts_integration]
-    root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
-    maybe_add_skills(root_agent)
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
+
+        maybe_add_skills(root_agent)
+
+        return root_agent
 
     kagent_app = KAgentApp(
-        root_agent,
+        root_agent_factory,
         agent_card,
         app_cfg.url,
         app_cfg.app_name,
@@ -136,16 +141,22 @@ def run(
 ):
     app_cfg = KAgentConfig()
 
-    agent_loader = AgentLoader(agents_dir=working_dir)
-    root_agent = agent_loader.load_agent(name)
-
     plugins = None
     sts_integration = create_sts_integration()
     if sts_integration:
         plugins = [sts_integration]
-        add_to_agent(sts_integration, root_agent)
 
-    maybe_add_skills(root_agent)
+    agent_loader = AgentLoader(agents_dir=working_dir)
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_loader.load_agent(name)
+
+        if sts_integration:
+            add_to_agent(sts_integration, root_agent)
+
+        maybe_add_skills(root_agent)
+
+        return root_agent
 
     # Load agent config to get stream setting
     agent_config = None
@@ -171,7 +182,7 @@ def run(
         logger.exception(f"Failed to load agent module '{name}' for lifespan")
 
     kagent_app = KAgentApp(
-        root_agent,
+        root_agent_factory,
         agent_card,
         app_cfg.url,
         app_cfg.app_name,
@@ -202,9 +213,13 @@ async def test_agent(agent_config: AgentConfig, agent_card: AgentCard, task: str
     sts_integration = create_sts_integration()
     if sts_integration:
         plugins = [sts_integration]
-    root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
-    maybe_add_skills(root_agent)
-    app = KAgentApp(root_agent, agent_card, app_cfg.url, app_cfg.app_name, plugins=plugins)
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
+        maybe_add_skills(root_agent)
+        return root_agent
+
+    app = KAgentApp(root_agent_factory, agent_card, app_cfg.url, app_cfg.app_name, plugins=plugins)
     await app.test(task)