Skip to content

Commit 8e69a58

Browse files
lwangverizoncopybara-github
authored andcommitted
feat: Add support to automatically create a session if one does not exist
feature/auto-create-new-session Merge #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=#4072 from lwangverizon:feature/auto-create-new-session 5475c6a PiperOrigin-RevId: 856019482
1 parent 1bedffe commit 8e69a58

File tree

2 files changed

+167
-18
lines changed

2 files changed

+167
-18
lines changed

src/google/adk/runners.py

Lines changed: 46 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ def __init__(
149149
memory_service: Optional[BaseMemoryService] = None,
150150
credential_service: Optional[BaseCredentialService] = None,
151151
plugin_close_timeout: float = 5.0,
152+
auto_create_session: bool = False,
152153
):
153154
"""Initializes the Runner.
154155
@@ -175,6 +176,9 @@ def __init__(
175176
memory_service: The memory service for the runner.
176177
credential_service: The credential service for the runner.
177178
plugin_close_timeout: The timeout in seconds for plugin close methods.
179+
auto_create_session: Whether to automatically create a session when
180+
not found. Defaults to False. If False, a missing session raises
181+
ValueError with a helpful message.
178182
179183
Raises:
180184
ValueError: If `app` is provided along with `agent` or `plugins`, or if
@@ -195,6 +199,7 @@ def __init__(
195199
self.plugin_manager = PluginManager(
196200
plugins=plugins, close_timeout=plugin_close_timeout
197201
)
202+
self.auto_create_session = auto_create_session
198203
(
199204
self._agent_origin_app_name,
200205
self._agent_origin_dir,
@@ -343,9 +348,43 @@ def _format_session_not_found_message(self, session_id: str) -> str:
343348
return message
344349
return (
345350
f'{message}. {self._app_name_alignment_hint} '
346-
'The mismatch prevents the runner from locating the session.'
351+
'The mismatch prevents the runner from locating the session. '
352+
'To automatically create a session when missing, set '
353+
'auto_create_session=True when constructing the runner.'
347354
)
348355

356+
async def _get_or_create_session(
357+
self, *, user_id: str, session_id: str
358+
) -> Session:
359+
"""Gets the session or creates it if auto-creation is enabled.
360+
361+
This helper first attempts to retrieve the session. If not found and
362+
auto_create_session is True, it creates a new session with the provided
363+
identifiers. Otherwise, it raises a ValueError with a helpful message.
364+
365+
Args:
366+
user_id: The user ID of the session.
367+
session_id: The session ID of the session.
368+
369+
Returns:
370+
The existing or newly created `Session`.
371+
372+
Raises:
373+
ValueError: If the session is not found and auto_create_session is False.
374+
"""
375+
session = await self.session_service.get_session(
376+
app_name=self.app_name, user_id=user_id, session_id=session_id
377+
)
378+
if not session:
379+
if self.auto_create_session:
380+
session = await self.session_service.create_session(
381+
app_name=self.app_name, user_id=user_id, session_id=session_id
382+
)
383+
else:
384+
message = self._format_session_not_found_message(session_id)
385+
raise ValueError(message)
386+
return session
387+
349388
def run(
350389
self,
351390
*,
@@ -455,12 +494,9 @@ async def _run_with_trace(
455494
invocation_id: Optional[str] = None,
456495
) -> AsyncGenerator[Event, None]:
457496
with tracer.start_as_current_span('invocation'):
458-
session = await self.session_service.get_session(
459-
app_name=self.app_name, user_id=user_id, session_id=session_id
497+
session = await self._get_or_create_session(
498+
user_id=user_id, session_id=session_id
460499
)
461-
if not session:
462-
message = self._format_session_not_found_message(session_id)
463-
raise ValueError(message)
464500
if not invocation_id and not new_message:
465501
raise ValueError(
466502
'Running an agent requires either a new_message or an '
@@ -534,12 +570,9 @@ async def rewind_async(
534570
rewind_before_invocation_id: str,
535571
) -> None:
536572
"""Rewinds the session to before the specified invocation."""
537-
session = await self.session_service.get_session(
538-
app_name=self.app_name, user_id=user_id, session_id=session_id
573+
session = await self._get_or_create_session(
574+
user_id=user_id, session_id=session_id
539575
)
540-
if not session:
541-
raise ValueError(f'Session not found: {session_id}')
542-
543576
rewind_event_index = -1
544577
for i, event in enumerate(session.events):
545578
if event.invocation_id == rewind_before_invocation_id:
@@ -967,14 +1000,9 @@ async def run_live(
9671000
stacklevel=2,
9681001
)
9691002
if not session:
970-
session = await self.session_service.get_session(
971-
app_name=self.app_name, user_id=user_id, session_id=session_id
1003+
session = await self._get_or_create_session(
1004+
user_id=user_id, session_id=session_id
9721005
)
973-
if not session:
974-
raise ValueError(
975-
f'Session not found for user id: {user_id} and session id:'
976-
f' {session_id}'
977-
)
9781006
invocation_context = self._new_invocation_context_for_live(
9791007
session,
9801008
live_request_queue=live_request_queue,

tests/unittests/test_runners.py

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,24 @@ async def _run_async_impl(
6868
)
6969

7070

71+
class MockLiveAgent(BaseAgent):
72+
"""Mock live agent for unit testing."""
73+
74+
def __init__(self, name: str):
75+
super().__init__(name=name, sub_agents=[])
76+
77+
async def _run_live_impl(
78+
self, invocation_context: InvocationContext
79+
) -> AsyncGenerator[Event, None]:
80+
yield Event(
81+
invocation_id=invocation_context.invocation_id,
82+
author=self.name,
83+
content=types.Content(
84+
role="model", parts=[types.Part(text="live hello")]
85+
),
86+
)
87+
88+
7189
class MockLlmAgent(LlmAgent):
7290
"""Mock LLM agent for unit testing."""
7391

@@ -237,6 +255,109 @@ def _infer_agent_origin(
237255
assert "Ensure the runner app_name matches" in message
238256

239257

258+
@pytest.mark.asyncio
259+
async def test_session_auto_creation():
260+
261+
class RunnerWithMismatch(Runner):
262+
263+
def _infer_agent_origin(
264+
self, agent: BaseAgent
265+
) -> tuple[Optional[str], Optional[Path]]:
266+
del agent
267+
return "expected_app", Path("/workspace/agents/expected_app")
268+
269+
session_service = InMemorySessionService()
270+
runner = RunnerWithMismatch(
271+
app_name="expected_app",
272+
agent=MockLlmAgent("test_agent"),
273+
session_service=session_service,
274+
artifact_service=InMemoryArtifactService(),
275+
auto_create_session=True,
276+
)
277+
278+
agen = runner.run_async(
279+
user_id="user",
280+
session_id="missing",
281+
new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
282+
)
283+
284+
event = await agen.__anext__()
285+
await agen.aclose()
286+
287+
# Verify that session_id="missing" doesn't error out - session is auto-created
288+
assert event.author == "test_agent"
289+
assert event.content.parts[0].text == "Test LLM response"
290+
291+
292+
@pytest.mark.asyncio
293+
async def test_rewind_auto_create_session_on_missing_session():
294+
"""When auto_create_session=True, rewind should create session if missing.
295+
296+
The newly created session won't contain the target invocation, so
297+
`rewind_async` should raise an Invocation ID not found error (rather than
298+
a session not found error), demonstrating auto-creation occurred.
299+
"""
300+
session_service = InMemorySessionService()
301+
runner = Runner(
302+
app_name="auto_create_app",
303+
agent=MockLlmAgent("agent_for_rewind"),
304+
session_service=session_service,
305+
artifact_service=InMemoryArtifactService(),
306+
auto_create_session=True,
307+
)
308+
309+
with pytest.raises(ValueError, match=r"Invocation ID not found: inv_missing"):
310+
await runner.rewind_async(
311+
user_id="user",
312+
session_id="missing",
313+
rewind_before_invocation_id="inv_missing",
314+
)
315+
316+
# Verify the session actually exists now due to auto-creation.
317+
session = await session_service.get_session(
318+
app_name="auto_create_app", user_id="user", session_id="missing"
319+
)
320+
assert session is not None
321+
assert session.app_name == "auto_create_app"
322+
323+
324+
@pytest.mark.asyncio
325+
async def test_run_live_auto_create_session():
326+
"""run_live should auto-create session when missing and yield events."""
327+
session_service = InMemorySessionService()
328+
artifact_service = InMemoryArtifactService()
329+
runner = Runner(
330+
app_name="live_app",
331+
agent=MockLiveAgent("live_agent"),
332+
session_service=session_service,
333+
artifact_service=artifact_service,
334+
auto_create_session=True,
335+
)
336+
337+
# An empty LiveRequestQueue is sufficient for our mock agent.
338+
from google.adk.agents.live_request_queue import LiveRequestQueue
339+
340+
live_queue = LiveRequestQueue()
341+
342+
agen = runner.run_live(
343+
user_id="user",
344+
session_id="missing",
345+
live_request_queue=live_queue,
346+
)
347+
348+
event = await agen.__anext__()
349+
await agen.aclose()
350+
351+
assert event.author == "live_agent"
352+
assert event.content.parts[0].text == "live hello"
353+
354+
# Session should have been created automatically.
355+
session = await session_service.get_session(
356+
app_name="live_app", user_id="user", session_id="missing"
357+
)
358+
assert session is not None
359+
360+
240361
@pytest.mark.asyncio
241362
async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch):
242363
project_root = tmp_path / "workspace"

0 commit comments

Comments
 (0)