fix: Fix event loop closed and thread leak errors in unit tests

- Wrap run_live generator in aclosing/Aclosing in InMemoryRunner and test runners to ensure they are closed when tests exit early.
- Refactor test_streaming.py to use a single, robust TestRunner helper, eliminating 10 duplicated CustomTestRunner definitions.
- Implement close() and async context manager support directly in PerAgentDatabaseSessionService in local_storage.py to clean up SQLAlchemy engines and connection threads without changing BaseSessionService.
- Wrap PerAgentDatabaseSessionService in async with blocks (or try-finally) in test_local_storage.py to ensure they are closed.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 936950444

Author: DeanChensj

src/google/adk/cli/utils/local_storage.py

@@ -236,6 +236,21 @@ async def append_event(self, session: Session, event: Event) -> Event:
     service = await self._get_service(session.app_name)
     return await service.append_event(session, event)
 
+  async def close(self) -> None:
+    """Closes all underlying session services."""
+    for service in self._services.values():
+      if hasattr(service, "close"):
+        await service.close()
+    self._services.clear()
+
+  async def __aenter__(self) -> PerAgentDatabaseSessionService:
+    """Enters the async context manager."""
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+    """Exits the async context manager and closes the service."""
+    await self.close()
+
 
 class PerAgentFileArtifactService(BaseArtifactService):
   """Routes artifact storage to per-agent `.adk/artifacts` folders."""

tests/unittests/cli/utils/test_local_storage.py

@@ -38,21 +38,20 @@ async def test_per_agent_session_service_creates_scoped_dot_adk(
   agent_a.mkdir()
   agent_b.mkdir()
 
-  service = PerAgentDatabaseSessionService(agents_root=tmp_path)
+  async with PerAgentDatabaseSessionService(agents_root=tmp_path) as service:
+    await service.create_session(app_name="agent_a", user_id="user_a")
+    await service.create_session(app_name="agent_b", user_id="user_b")
 
-  await service.create_session(app_name="agent_a", user_id="user_a")
-  await service.create_session(app_name="agent_b", user_id="user_b")
+    assert (agent_a / ".adk" / "session.db").exists()
+    assert (agent_b / ".adk" / "session.db").exists()
 
-  assert (agent_a / ".adk" / "session.db").exists()
-  assert (agent_b / ".adk" / "session.db").exists()
+    agent_a_sessions = await service.list_sessions(app_name="agent_a")
+    agent_b_sessions = await service.list_sessions(app_name="agent_b")
 
-  agent_a_sessions = await service.list_sessions(app_name="agent_a")
-  agent_b_sessions = await service.list_sessions(app_name="agent_b")
-
-  assert len(agent_a_sessions.sessions) == 1
-  assert agent_a_sessions.sessions[0].app_name == "agent_a"
-  assert len(agent_b_sessions.sessions) == 1
-  assert agent_b_sessions.sessions[0].app_name == "agent_b"
+    assert len(agent_a_sessions.sessions) == 1
+    assert agent_a_sessions.sessions[0].app_name == "agent_a"
+    assert len(agent_b_sessions.sessions) == 1
+    assert agent_b_sessions.sessions[0].app_name == "agent_b"
 
 
 @pytest.mark.asyncio
@@ -68,26 +67,28 @@ async def test_per_agent_session_service_respects_app_name_alias(
       per_agent=True,
       app_name_to_dir={logical_name: folder_name},
   )
+  try:
+    session = await service.create_session(
+        app_name=logical_name,
+        user_id="user",
+    )
 
-  session = await service.create_session(
-      app_name=logical_name,
-      user_id="user",
-  )
-
-  assert session.app_name == logical_name
-  assert (tmp_path / folder_name / ".adk" / "session.db").exists()
+    assert session.app_name == logical_name
+    assert (tmp_path / folder_name / ".adk" / "session.db").exists()
+  finally:
+    if isinstance(service, PerAgentDatabaseSessionService):
+      await service.close()
 
 
 @pytest.mark.asyncio
 async def test_per_agent_session_service_routes_built_in_agents_to_root_dot_adk(
     tmp_path: Path,
 ) -> None:
-  service = PerAgentDatabaseSessionService(agents_root=tmp_path)
+  async with PerAgentDatabaseSessionService(agents_root=tmp_path) as service:
+    await service.create_session(app_name="__helper", user_id="user")
 
-  await service.create_session(app_name="__helper", user_id="user")
-
-  assert not (tmp_path / "__helper").exists()
-  assert (tmp_path / ".adk" / "session.db").exists()
+    assert not (tmp_path / "__helper").exists()
+    assert (tmp_path / ".adk" / "session.db").exists()
 
 
 def test_create_local_database_session_service_returns_sqlite(
@@ -106,22 +107,25 @@ async def test_per_agent_session_service_get_user_state(tmp_path: Path) -> None:
   agent_a.mkdir()
   agent_b.mkdir()
 
-  service = PerAgentDatabaseSessionService(agents_root=tmp_path)
-
-  session_a = await service.create_session(app_name="agent_a", user_id="user_a")
-  await service.append_event(
-      session_a,
-      Event(
-          author="system",
-          actions=EventActions(state_delta={"user:profile": {"name": "Alice"}}),
-      ),
-  )
-
-  state_a = await service.get_user_state(app_name="agent_a", user_id="user_a")
-  state_b = await service.get_user_state(app_name="agent_b", user_id="user_b")
-
-  assert state_a == {"profile": {"name": "Alice"}}
-  assert not state_b
+  async with PerAgentDatabaseSessionService(agents_root=tmp_path) as service:
+    session_a = await service.create_session(
+        app_name="agent_a", user_id="user_a"
+    )
+    await service.append_event(
+        session_a,
+        Event(
+            author="system",
+            actions=EventActions(
+                state_delta={"user:profile": {"name": "Alice"}}
+            ),
+        ),
+    )
+
+    state_a = await service.get_user_state(app_name="agent_a", user_id="user_a")
+    state_b = await service.get_user_state(app_name="agent_b", user_id="user_b")
+
+    assert state_a == {"profile": {"name": "Alice"}}
+    assert not state_b
 
 
 @pytest.mark.asyncio

tests/unittests/streaming/test_multi_agent_streaming.py

@@ -94,10 +94,13 @@ async def consume_responses(session: testing_utils.Session):
             live_request_queue=live_request_queue,
             run_config=run_config or testing_utils.RunConfig(),
         )
-        async for response in run_res:
-          collected_responses.append(response)
-          if len(collected_responses) >= 5:
-            return
+        from contextlib import aclosing
+
+        async with aclosing(run_res) as agen:
+          async for response in agen:
+            collected_responses.append(response)
+            if len(collected_responses) >= 5:
+              return
 
       try:
         session = self.session