Feature/clone session #4191
Feature/clone session #4191
Conversation
Summary of ChangesHello @lwangverizon, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the ADK session services by adding a Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
services
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable clone_session feature across all session service backends, complete with a comprehensive set of unit tests. The implementation is well-structured and addresses a key use case for transitioning user states. My review highlights several opportunities for improvement, primarily focusing on performance and correctness. I've identified a recurring N+1 query pattern that could be optimized, as well as redundant data fetching that can be avoided. Additionally, there are a couple of potential correctness issues in the VertexAiSessionService related to missing deep copies that could lead to unexpected state mutations. Addressing these points will enhance the robustness and efficiency of the new feature.
| # Merge states from all source sessions | ||
| merged_state = {} | ||
| for session in source_sessions: | ||
| merged_state.update(session.state) |
There was a problem hiding this comment.
The session state is being merged using a shallow copy (.update()). This is inconsistent with the other SessionService implementations, which use copy.deepcopy(). If session.state contains mutable objects (e.g., lists, dicts), changes in the cloned session could unintentionally alter the state of the source sessions. To prevent these side effects and ensure correctness, you should use a deep copy.
Note: You will need to add import copy at the top of the file.
| merged_state.update(session.state) | |
| merged_state.update(copy.deepcopy(session.state)) |
| for event in all_events: | ||
| await self.append_event(new_session, event) |
There was a problem hiding this comment.
The event object from a source session is passed directly to self.append_event(). The base implementation of append_event can mutate this object (e.g., by calling _trim_temp_delta_state). This creates an unintended side effect, modifying the original event from the source session. To prevent this, you should pass a deep copy of the event.
Note: This also requires import copy at the top of the file.
| for event in all_events: | |
| await self.append_event(new_session, event) | |
| for event in all_events: | |
| await self.append_event(new_session, copy.deepcopy(event)) |
| return await self.get_session( | ||
| app_name=new_session.app_name, | ||
| user_id=new_session.user_id, | ||
| session_id=new_session.id, | ||
| ) |
There was a problem hiding this comment.
This final get_session call is redundant and inefficient as it performs an extra database query to fetch data you already have. The new_session object returned by create_session contains the correct state, and you have the deduplicated all_events list. You can construct the final Session object in memory.
# Return the new session with events
new_session.events = all_events
return new_session| for sess in list_response.sessions: | ||
| full_session = self._get_session_impl( | ||
| app_name=app_name, | ||
| user_id=src_user_id, | ||
| session_id=sess.id, | ||
| ) | ||
| if full_session: | ||
| source_sessions.append(full_session) |
There was a problem hiding this comment.
This loop exhibits an N+1-like pattern. You call _list_sessions_impl and then loop to call _get_session_impl for each session. While this is an in-memory operation, it's less efficient than necessary and follows a pattern that is problematic in database-backed services. You could optimize this by directly accessing the sessions for the user from self.sessions and collecting the events without repeated get_session calls.
| return self._get_session_impl( | ||
| app_name=new_session.app_name, | ||
| user_id=new_session.user_id, | ||
| session_id=new_session.id, | ||
| ) |
There was a problem hiding this comment.
This final call to _get_session_impl is unnecessary. You have already created the new_session object and populated the storage_session with the correct events and latest_update_time. Instead of re-fetching, you can update the new_session object you already have in memory and return it directly.
new_session.events = all_events
new_session.last_update_time = latest_update_time
return new_session| for sess in list_response.sessions: | ||
| full_session = await self.get_session( | ||
| app_name=app_name, | ||
| user_id=src_user_id, | ||
| session_id=sess.id, | ||
| ) | ||
| if full_session: | ||
| source_sessions.append(full_session) |
There was a problem hiding this comment.
This loop introduces an N+1 query problem. You call list_sessions once, and then get_session is called inside a loop for every session found. This can lead to a large number of database queries and impact performance, especially for users with many sessions.
You could optimize this by:
- Getting all session objects from
list_sessions. - Collecting all
session_ids. - Fetching all events for those session IDs in a single query using
WHERE session_id IN (...). - Mapping the events back to their respective sessions in memory.
| return await self.get_session( | ||
| app_name=new_session.app_name, | ||
| user_id=new_session.user_id, | ||
| session_id=new_session.id, | ||
| ) |
There was a problem hiding this comment.
This get_session call to fetch the newly created session is redundant. You have all the necessary information to construct the Session object in memory and return it directly. The new_session object returned from create_session already contains the correct state, and you have the all_events list. You can simply assign the events to the session and return it.
# Return the new session with events
new_session.events = all_events
return new_session| for sess in list_response.sessions: | ||
| full_session = await self.get_session( | ||
| app_name=app_name, | ||
| user_id=src_user_id, | ||
| session_id=sess.id, | ||
| ) | ||
| if full_session: | ||
| source_sessions.append(full_session) |
There was a problem hiding this comment.
This loop creates an N+1 query problem. list_sessions is called once, and then get_session is called for each session found. This results in multiple API calls to Vertex AI, which can be inefficient for users with many sessions. If the API supports batch retrieval of sessions with their events, that would be a more performant approach.
| return await self.get_session( | ||
| app_name=new_session.app_name, | ||
| user_id=new_session.user_id, | ||
| session_id=new_session.id, | ||
| ) |
There was a problem hiding this comment.
This final get_session call is redundant. You already have all the necessary information to construct the final Session object in memory. The new_session object from create_session has the correct state, and you have the all_events list. You can assign the events to the session and return it directly, avoiding an unnecessary API call.
# Return the new session with events
new_session.events = all_events
return new_session2 participants
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
In e-commerce and customer service scenarios, agentic chatbots commonly need to handle both non-signed-in (guest) and signed-in (authenticated) user states. A typical user journey looks like:
session_id,cookie_id, ordevice_id)customer_id,account_number)Without this capability, the agent loses all prior context when the user authenticates, forcing them to repeat their questions and preferences. This creates a poor user experience and reduces conversion rates.
Current limitation: ADK session services do not provide a way to clone or transfer session history from one user identity to another.
Solution:
Add a new
clone_sessionmethod toBaseSessionServicethat enables applications to seamlessly transfer conversation context when users transition between identity states.Primary use case - Guest to Authenticated transition:
The method is implemented across all session service backends:
InMemorySessionServiceDatabaseSessionService(SQLAlchemy)SqliteSessionService(aiosqlite)VertexAiSessionServiceAdditional supported scenarios:
Key features:
new_session_idis not providedValueErrorfor missing sources,AlreadyExistsErrorfor duplicate IDs)API signature:
Testing Plan
Unit Tests:
Added 10 new test cases that run across all 3 testable session services (InMemory, Database, SQLite) for a total of 30 new tests:
test_clone_session_basictest_clone_session_with_different_user_idtest_clone_session_with_custom_session_idtest_clone_session_with_existing_id_raises_errortest_clone_session_preserves_event_contenttest_clone_session_does_not_affect_sourcetest_clone_all_user_sessionstest_clone_session_no_source_raises_errortest_clone_all_sessions_no_sessions_raises_errortest_clone_session_deduplicates_eventspytest results:
Manual End-to-End (E2E) Tests:
Checklist
Additional context
Files changed:
src/google/adk/sessions/base_session_service.pyclone_sessionmethod with comprehensive docstringsrc/google/adk/sessions/in_memory_session_service.pyclone_sessionsrc/google/adk/sessions/database_session_service.pyclone_sessionsrc/google/adk/sessions/sqlite_session_service.pyclone_sessionsrc/google/adk/sessions/vertex_ai_session_service.pyclone_session(note:new_session_idnot supported due to Vertex AI API limitation)tests/unittests/sessions/test_session_service.pyDesign decisions:
new_session_idparameter raisesValueErrorforVertexAiSessionServicesince the backend auto-generates session IDs