From 8e69a58df4eadeccbb100b7264bb518a46b61fd7 Mon Sep 17 00:00:00 2001 From: lwangverizon Date: Tue, 13 Jan 2026 21:23:22 -0800 Subject: [PATCH] feat: Add support to automatically create a session if one does not exist feature/auto-create-new-session Merge https://github.com/google/adk-python/pull/4072 **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change **2. Or, if no issue exists, describe the change:** **Problem:** When building frontend applications with ADK, there's a limitation where frontends cannot always guarantee that `create_session` is called before initiating a conversation. This creates friction in the user experience because: - Users may refresh the page or navigate directly to a conversation URL with a specific session_id - Frontend state management may lose track of whether a session was already created - Mobile apps or single-page applications have complex lifecycle management where ensuring `create_session` is called first adds unnecessary complexity - This forces developers to implement additional logic to check session existence before every conversation Currently, if `get_session` is called with a non-existent session_id, it returns `None`, requiring the frontend to explicitly handle this case and call `create_session` separately. **Solution:** Modified the `get_session` method in `DatabaseSessionService` to automatically create a session if it doesn't exist in the database. This "get or create" pattern is common in many frameworks and provides a more developer-friendly API. The implementation: 1. Attempts to fetch the session from the database 2. If the session doesn't exist (returns `None`), automatically calls `create_session` with the provided parameters 3. Retrieves and returns the newly created session 4. Maintains backward compatibility - existing code continues to work without changes This allows frontends to simply call `get_session` with a session_id and be confident that the session will be available, regardless of whether it was previously created. **Benefits:** - Simplifies frontend integration by removing the need to track session creation state - Reduces API calls (no need to check existence before calling get_session) - Follows the principle of least surprise - getting a session with an ID should work reliably - No breaking changes to existing code that checks for `None` return values ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. **pytest results:** COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4072 from lwangverizon:feature/auto-create-new-session 5475c6ae91d12332598d521302736eb1db79a8be PiperOrigin-RevId: 856019482 --- src/google/adk/runners.py | 64 ++++++++++++----- tests/unittests/test_runners.py | 121 ++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+), 18 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index cbf2c595..97eb85df 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -149,6 +149,7 @@ class Runner: memory_service: Optional[BaseMemoryService] = None, credential_service: Optional[BaseCredentialService] = None, plugin_close_timeout: float = 5.0, + auto_create_session: bool = False, ): """Initializes the Runner. @@ -175,6 +176,9 @@ class Runner: memory_service: The memory service for the runner. credential_service: The credential service for the runner. plugin_close_timeout: The timeout in seconds for plugin close methods. + auto_create_session: Whether to automatically create a session when + not found. Defaults to False. If False, a missing session raises + ValueError with a helpful message. Raises: ValueError: If `app` is provided along with `agent` or `plugins`, or if @@ -195,6 +199,7 @@ class Runner: self.plugin_manager = PluginManager( plugins=plugins, close_timeout=plugin_close_timeout ) + self.auto_create_session = auto_create_session ( self._agent_origin_app_name, self._agent_origin_dir, @@ -343,9 +348,43 @@ class Runner: return message return ( f'{message}. {self._app_name_alignment_hint} ' - 'The mismatch prevents the runner from locating the session.' + 'The mismatch prevents the runner from locating the session. ' + 'To automatically create a session when missing, set ' + 'auto_create_session=True when constructing the runner.' ) + async def _get_or_create_session( + self, *, user_id: str, session_id: str + ) -> Session: + """Gets the session or creates it if auto-creation is enabled. + + This helper first attempts to retrieve the session. If not found and + auto_create_session is True, it creates a new session with the provided + identifiers. Otherwise, it raises a ValueError with a helpful message. + + Args: + user_id: The user ID of the session. + session_id: The session ID of the session. + + Returns: + The existing or newly created `Session`. + + Raises: + ValueError: If the session is not found and auto_create_session is False. + """ + session = await self.session_service.get_session( + app_name=self.app_name, user_id=user_id, session_id=session_id + ) + if not session: + if self.auto_create_session: + session = await self.session_service.create_session( + app_name=self.app_name, user_id=user_id, session_id=session_id + ) + else: + message = self._format_session_not_found_message(session_id) + raise ValueError(message) + return session + def run( self, *, @@ -455,12 +494,9 @@ class Runner: invocation_id: Optional[str] = None, ) -> AsyncGenerator[Event, None]: with tracer.start_as_current_span('invocation'): - session = await self.session_service.get_session( - app_name=self.app_name, user_id=user_id, session_id=session_id + session = await self._get_or_create_session( + user_id=user_id, session_id=session_id ) - if not session: - message = self._format_session_not_found_message(session_id) - raise ValueError(message) if not invocation_id and not new_message: raise ValueError( 'Running an agent requires either a new_message or an ' @@ -534,12 +570,9 @@ class Runner: rewind_before_invocation_id: str, ) -> None: """Rewinds the session to before the specified invocation.""" - session = await self.session_service.get_session( - app_name=self.app_name, user_id=user_id, session_id=session_id + session = await self._get_or_create_session( + user_id=user_id, session_id=session_id ) - if not session: - raise ValueError(f'Session not found: {session_id}') - rewind_event_index = -1 for i, event in enumerate(session.events): if event.invocation_id == rewind_before_invocation_id: @@ -967,14 +1000,9 @@ class Runner: stacklevel=2, ) if not session: - session = await self.session_service.get_session( - app_name=self.app_name, user_id=user_id, session_id=session_id + session = await self._get_or_create_session( + user_id=user_id, session_id=session_id ) - if not session: - raise ValueError( - f'Session not found for user id: {user_id} and session id:' - f' {session_id}' - ) invocation_context = self._new_invocation_context_for_live( session, live_request_queue=live_request_queue, diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 710ca90d..bb44ce73 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -68,6 +68,24 @@ class MockAgent(BaseAgent): ) +class MockLiveAgent(BaseAgent): + """Mock live agent for unit testing.""" + + def __init__(self, name: str): + super().__init__(name=name, sub_agents=[]) + + async def _run_live_impl( + self, invocation_context: InvocationContext + ) -> AsyncGenerator[Event, None]: + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="live hello")] + ), + ) + + class MockLlmAgent(LlmAgent): """Mock LLM agent for unit testing.""" @@ -237,6 +255,109 @@ async def test_session_not_found_message_includes_alignment_hint(): assert "Ensure the runner app_name matches" in message +@pytest.mark.asyncio +async def test_session_auto_creation(): + + class RunnerWithMismatch(Runner): + + def _infer_agent_origin( + self, agent: BaseAgent + ) -> tuple[Optional[str], Optional[Path]]: + del agent + return "expected_app", Path("/workspace/agents/expected_app") + + session_service = InMemorySessionService() + runner = RunnerWithMismatch( + app_name="expected_app", + agent=MockLlmAgent("test_agent"), + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + + agen = runner.run_async( + user_id="user", + session_id="missing", + new_message=types.Content(role="user", parts=[types.Part(text="hi")]), + ) + + event = await agen.__anext__() + await agen.aclose() + + # Verify that session_id="missing" doesn't error out - session is auto-created + assert event.author == "test_agent" + assert event.content.parts[0].text == "Test LLM response" + + +@pytest.mark.asyncio +async def test_rewind_auto_create_session_on_missing_session(): + """When auto_create_session=True, rewind should create session if missing. + + The newly created session won't contain the target invocation, so + `rewind_async` should raise an Invocation ID not found error (rather than + a session not found error), demonstrating auto-creation occurred. + """ + session_service = InMemorySessionService() + runner = Runner( + app_name="auto_create_app", + agent=MockLlmAgent("agent_for_rewind"), + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + + with pytest.raises(ValueError, match=r"Invocation ID not found: inv_missing"): + await runner.rewind_async( + user_id="user", + session_id="missing", + rewind_before_invocation_id="inv_missing", + ) + + # Verify the session actually exists now due to auto-creation. + session = await session_service.get_session( + app_name="auto_create_app", user_id="user", session_id="missing" + ) + assert session is not None + assert session.app_name == "auto_create_app" + + +@pytest.mark.asyncio +async def test_run_live_auto_create_session(): + """run_live should auto-create session when missing and yield events.""" + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + runner = Runner( + app_name="live_app", + agent=MockLiveAgent("live_agent"), + session_service=session_service, + artifact_service=artifact_service, + auto_create_session=True, + ) + + # An empty LiveRequestQueue is sufficient for our mock agent. + from google.adk.agents.live_request_queue import LiveRequestQueue + + live_queue = LiveRequestQueue() + + agen = runner.run_live( + user_id="user", + session_id="missing", + live_request_queue=live_queue, + ) + + event = await agen.__anext__() + await agen.aclose() + + assert event.author == "live_agent" + assert event.content.parts[0].text == "live hello" + + # Session should have been created automatically. + session = await session_service.get_session( + app_name="live_app", user_id="user", session_id="missing" + ) + assert session is not None + + @pytest.mark.asyncio async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch): project_root = tmp_path / "workspace"