feat! Update session service interface to be async.

Also keep the sync version in the InMemorySessionService as create_session_sync() as a temporary migration option.

PiperOrigin-RevId: 759224250
This commit is contained in:
Shangjie Chen
2025-05-15 11:16:43 -07:00
committed by Copybara-Service
parent d161a2c3f7
commit 5b3204c356
23 changed files with 264 additions and 268 deletions

View File

@@ -70,11 +70,11 @@ class _TestingAgentWithEscalateAction(BaseAgent):
)
def _create_parent_invocation_context(
async def _create_parent_invocation_context(
test_name: str, agent: BaseAgent
) -> InvocationContext:
session_service = InMemorySessionService()
session = session_service.create_session(
session = await session_service.create_session(
app_name='test_app', user_id='test_user'
)
return InvocationContext(
@@ -95,7 +95,7 @@ async def test_run_async(request: pytest.FixtureRequest):
agent,
],
)
parent_ctx = _create_parent_invocation_context(
parent_ctx = await _create_parent_invocation_context(
request.function.__name__, loop_agent
)
events = [e async for e in loop_agent.run_async(parent_ctx)]
@@ -119,7 +119,7 @@ async def test_run_async_with_escalate_action(request: pytest.FixtureRequest):
name=f'{request.function.__name__}_test_loop_agent',
sub_agents=[non_escalating_agent, escalating_agent],
)
parent_ctx = _create_parent_invocation_context(
parent_ctx = await _create_parent_invocation_context(
request.function.__name__, loop_agent
)
events = [e async for e in loop_agent.run_async(parent_ctx)]