Skip to content

Commit d029bce

Browse files
nicolasmotaDeanChensj
authored andcommitted
feat(sessions): add get_user_state(app_name, user_id) to BaseSessionService
Closes: #5592 Problem: BaseSessionService has no public method to read user-scoped state without an active session_id. Callers that need to bootstrap user context before a new session exists are forced to call the expensive list_sessions or maintain a separate process-level cache as a workaround. Solution: Add get_user_state(app_name, user_id) -> dict[str, Any] to BaseSessionService. Implemented in InMemorySessionService, DatabaseSessionService, and SqliteSessionService. VertexAiSessionService raises NotImplementedError because the Vertex AI Agent Engine API does not expose user state independently of a session. The default in BaseSessionService also raises NotImplementedError to preserve backward compatibility for existing custom subclasses. Keys are returned without the user: prefix, consistent with how user state is stored internally (the prefix is applied by the state-merging layer). Merge #5596 Change-Id: Id1015034c62810aafd2b2411a0376742bb80e8c8
1 parent 7eb9b3d commit d029bce

8 files changed

Lines changed: 246 additions & 1 deletion

File tree

src/google/adk/cli/utils/local_storage.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import asyncio
1919
import logging
2020
from pathlib import Path
21+
from typing import Any
2122
from typing import Mapping
2223
from typing import Optional
2324

@@ -205,6 +206,13 @@ async def delete_session(
205206
app_name=app_name, user_id=user_id, session_id=session_id
206207
)
207208

209+
@override
210+
async def get_user_state(
211+
self, *, app_name: str, user_id: str
212+
) -> dict[str, Any]:
213+
service = await self._get_service(app_name)
214+
return await service.get_user_state(app_name=app_name, user_id=user_id)
215+
208216
@override
209217
async def append_event(self, session: Session, event: Event) -> Event:
210218
service = await self._get_service(session.app_name)

src/google/adk/sessions/base_session_service.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,46 @@ async def delete_session(
111111
) -> None:
112112
"""Deletes a session."""
113113

114+
async def get_user_state(
115+
self, *, app_name: str, user_id: str
116+
) -> dict[str, Any]:
117+
"""Returns the user-scoped state for the given app and user.
118+
119+
User state is keyed by ``(app_name, user_id)`` and shared across all
120+
sessions of the same user within the same app. The returned dictionary
121+
uses raw keys **without** the ``user:`` prefix (e.g. ``"my_key"`` rather
122+
than ``"user:my_key"``).
123+
124+
This method exists so that callers can read user state without holding an
125+
active ``session_id``. A common use case is bootstrapping context at the
126+
start of a new session before calling ``create_session``, which would
127+
otherwise require an expensive ``list_sessions`` call just to access
128+
user-scoped data.
129+
130+
Returns an empty dict when no user state has been stored for this
131+
``(app_name, user_id)`` combination.
132+
133+
Args:
134+
app_name: The name of the app.
135+
user_id: The ID of the user.
136+
137+
Returns:
138+
A dictionary of raw (un-prefixed) user-scoped key/value pairs, or an
139+
empty dict when no user state exists.
140+
141+
Raises:
142+
NotImplementedError: When the concrete ``BaseSessionService``
143+
implementation does not support reading user state independently of a
144+
session. Callers should catch this, then enumerate sessions via
145+
``list_sessions`` and call ``get_session`` on each result to access
146+
the merged state, or accept that user state is unavailable.
147+
"""
148+
raise NotImplementedError(
149+
f'{type(self).__name__} does not support get_user_state. '
150+
'To read user state, enumerate sessions via list_sessions and '
151+
'call get_session on each result to access the merged state.'
152+
)
153+
114154
async def append_event(self, session: Session, event: Event) -> Event:
115155
"""Appends an event to a session object."""
116156
if event.partial:

src/google/adk/sessions/database_session_service.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,7 @@ async def _with_session_lock(
329329
else:
330330
self._session_lock_ref_count[lock_key] = remaining
331331

332-
async def _prepare_tables(self):
332+
async def _prepare_tables(self) -> None:
333333
"""Ensure database tables are ready for use.
334334
335335
This method is called lazily before each database operation. It checks the
@@ -627,6 +627,22 @@ async def delete_session(
627627
await sql_session.execute(stmt)
628628
await sql_session.commit()
629629

630+
@override
631+
async def get_user_state(
632+
self, *, app_name: str, user_id: str
633+
) -> dict[str, Any]:
634+
await self._prepare_tables()
635+
schema = self._get_schema_classes()
636+
async with self._rollback_on_exception_session(
637+
read_only=True
638+
) as sql_session:
639+
storage_user_state = await sql_session.get(
640+
schema.StorageUserState, (app_name, user_id)
641+
)
642+
if storage_user_state is None:
643+
return {}
644+
return dict(storage_user_state.state or {})
645+
630646
@override
631647
async def append_event(self, session: Session, event: Event) -> Event:
632648
await self._prepare_tables()

src/google/adk/sessions/in_memory_session_service.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,12 @@ def _delete_session_impl(
312312

313313
self.sessions[app_name][user_id].pop(session_id)
314314

315+
@override
316+
async def get_user_state(
317+
self, *, app_name: str, user_id: str
318+
) -> dict[str, Any]:
319+
return dict(self.user_state.get(app_name, {}).get(user_id, {}))
320+
315321
@override
316322
async def append_event(self, session: Session, event: Event) -> Event:
317323
if event.partial:

src/google/adk/sessions/sqlite_session_service.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,13 @@ async def delete_session(
359359
)
360360
await db.commit()
361361

362+
@override
363+
async def get_user_state(
364+
self, *, app_name: str, user_id: str
365+
) -> dict[str, Any]:
366+
async with self._get_db_connection() as db:
367+
return await self._get_user_state(db, app_name, user_id)
368+
362369
@override
363370
async def append_event(self, session: Session, event: Event) -> Event:
364371
if event.partial:

src/google/adk/sessions/vertex_ai_session_service.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,27 @@ async def delete_session(
276276
logger.error('Error deleting session %s: %s', session_id, e)
277277
raise
278278

279+
@override
280+
async def get_user_state(
281+
self, *, app_name: str, user_id: str
282+
) -> dict[str, Any]:
283+
"""Not supported by the Vertex AI Agent Engine backend.
284+
285+
The Vertex AI Agent Engine API does not expose user state independently of
286+
a session. To read user state, enumerate sessions via ``list_sessions``
287+
and call ``get_session`` on each result to access the merged state.
288+
289+
Raises:
290+
NotImplementedError: Always, because the Vertex AI Agent Engine API does
291+
not provide a way to query user state without a session.
292+
"""
293+
raise NotImplementedError(
294+
'VertexAiSessionService does not support get_user_state. '
295+
'The Vertex AI Agent Engine API does not expose user state '
296+
'independently of a session. To read user state, enumerate sessions '
297+
'via list_sessions and call get_session on each result.'
298+
)
299+
279300
@override
280301
async def append_event(self, session: Session, event: Event) -> Event:
281302
# Update the in-memory session.

tests/unittests/cli/utils/test_local_storage.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
from google.adk.cli.utils.local_storage import create_local_database_session_service
2020
from google.adk.cli.utils.local_storage import create_local_session_service
2121
from google.adk.cli.utils.local_storage import PerAgentDatabaseSessionService
22+
from google.adk.events.event import Event
23+
from google.adk.events.event_actions import EventActions
2224
from google.adk.sessions.sqlite_session_service import SqliteSessionService
2325
import pytest
2426

@@ -90,3 +92,28 @@ def test_create_local_database_session_service_returns_sqlite(
9092
service = create_local_database_session_service(base_dir=tmp_path)
9193

9294
assert isinstance(service, SqliteSessionService)
95+
96+
97+
@pytest.mark.asyncio
98+
async def test_per_agent_session_service_get_user_state(tmp_path: Path) -> None:
99+
agent_a = tmp_path / "agent_a"
100+
agent_b = tmp_path / "agent_b"
101+
agent_a.mkdir()
102+
agent_b.mkdir()
103+
104+
service = PerAgentDatabaseSessionService(agents_root=tmp_path)
105+
106+
session_a = await service.create_session(app_name="agent_a", user_id="user_a")
107+
await service.append_event(
108+
session_a,
109+
Event(
110+
author="system",
111+
actions=EventActions(state_delta={"user:profile": {"name": "Alice"}}),
112+
),
113+
)
114+
115+
state_a = await service.get_user_state(app_name="agent_a", user_id="user_a")
116+
state_b = await service.get_user_state(app_name="agent_b", user_id="user_b")
117+
118+
assert state_a == {"profile": {"name": "Alice"}}
119+
assert state_b == {}

tests/unittests/sessions/test_session_service.py

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from google.adk.sessions.database_session_service import DatabaseSessionService
3131
from google.adk.sessions.in_memory_session_service import InMemorySessionService
3232
from google.adk.sessions.sqlite_session_service import SqliteSessionService
33+
from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
3334
from google.genai import types
3435
import pytest
3536
from sqlalchemy import delete
@@ -1650,3 +1651,122 @@ async def tracking_fn(**kwargs):
16501651
finally:
16511652
database_session_service._select_required_state = original_fn
16521653
await service.close()
1654+
1655+
1656+
@pytest.mark.asyncio
1657+
async def test_get_user_state_returns_empty_dict_when_no_state_exists(
1658+
session_service,
1659+
):
1660+
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
1661+
assert state == {}
1662+
1663+
1664+
@pytest.mark.asyncio
1665+
async def test_get_user_state_returns_state_written_via_append_event(
1666+
session_service,
1667+
):
1668+
session = await session_service.create_session(
1669+
app_name='my_app', user_id='u1'
1670+
)
1671+
await session_service.append_event(
1672+
session,
1673+
Event(
1674+
author='system',
1675+
actions=EventActions(
1676+
state_delta={'user:profile': {'name': 'Alice'}, 'session_key': 1}
1677+
),
1678+
),
1679+
)
1680+
1681+
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
1682+
1683+
assert state == {'profile': {'name': 'Alice'}}
1684+
assert 'session_key' not in state
1685+
1686+
1687+
@pytest.mark.asyncio
1688+
async def test_get_user_state_is_not_visible_across_users(session_service):
1689+
session = await session_service.create_session(
1690+
app_name='my_app', user_id='u1'
1691+
)
1692+
await session_service.append_event(
1693+
session,
1694+
Event(
1695+
author='system',
1696+
actions=EventActions(state_delta={'user:secret': 'only-for-u1'}),
1697+
),
1698+
)
1699+
1700+
other_state = await session_service.get_user_state(
1701+
app_name='my_app', user_id='u2'
1702+
)
1703+
assert other_state == {}
1704+
1705+
1706+
@pytest.mark.asyncio
1707+
async def test_get_user_state_is_not_visible_across_apps(session_service):
1708+
session = await session_service.create_session(
1709+
app_name='my_app', user_id='u1'
1710+
)
1711+
await session_service.append_event(
1712+
session,
1713+
Event(
1714+
author='system',
1715+
actions=EventActions(state_delta={'user:data': 'only-app-a'}),
1716+
),
1717+
)
1718+
1719+
other_state = await session_service.get_user_state(
1720+
app_name='other_app', user_id='u1'
1721+
)
1722+
assert other_state == {}
1723+
1724+
1725+
@pytest.mark.asyncio
1726+
async def test_get_user_state_available_before_session_is_created(
1727+
session_service,
1728+
):
1729+
first_session = await session_service.create_session(
1730+
app_name='my_app', user_id='u1'
1731+
)
1732+
await session_service.append_event(
1733+
first_session,
1734+
Event(
1735+
author='system',
1736+
actions=EventActions(state_delta={'user:ctx': {'v': 1}}),
1737+
),
1738+
)
1739+
1740+
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
1741+
assert state == {'ctx': {'v': 1}}
1742+
1743+
1744+
@pytest.mark.asyncio
1745+
async def test_get_user_state_reflects_latest_write(session_service):
1746+
session = await session_service.create_session(
1747+
app_name='my_app', user_id='u1'
1748+
)
1749+
await session_service.append_event(
1750+
session,
1751+
Event(
1752+
author='system',
1753+
actions=EventActions(state_delta={'user:counter': 1}),
1754+
),
1755+
)
1756+
await session_service.append_event(
1757+
session,
1758+
Event(
1759+
author='system',
1760+
actions=EventActions(state_delta={'user:counter': 2}),
1761+
),
1762+
)
1763+
1764+
state = await session_service.get_user_state(app_name='my_app', user_id='u1')
1765+
assert state['counter'] == 2
1766+
1767+
1768+
@pytest.mark.asyncio
1769+
async def test_vertex_ai_session_service_raises_not_implemented_for_get_user_state():
1770+
service = VertexAiSessionService(project='proj', location='us-central1')
1771+
with pytest.raises(NotImplementedError):
1772+
await service.get_user_state(app_name='my_app', user_id='u1')

0 commit comments

Comments
 (0)