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
This commit is contained in:
lwangverizon
2026-01-13 21:23:22 -08:00
committed by Copybara-Service
parent 1bedffe457
commit 8e69a58df4
2 changed files with 167 additions and 18 deletions
+46 -18
View File
@@ -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,
+121
View File
@@ -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"