diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 413ef800a6e..77742ac9b67 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -495,11 +495,14 @@ class UpdateMemoryRequest(common.BaseModel): class UpdateSessionRequest(common.BaseModel): - """Request to update session state without running the agent.""" + """Request to update session state and/or metadata without running the agent.""" - state_delta: dict[str, Any] + state_delta: Optional[dict[str, Any]] = None """The state changes to apply to the session.""" + title: Optional[str] = None + """A new human-readable title to set on the session.""" + class AppInfo(common.BaseModel): name: str @@ -1301,23 +1304,33 @@ async def update_session( if not session: raise HTTPException(status_code=404, detail="Session not found") - # Create an event to record the state change - import uuid + # Apply state changes (if provided) as an event so they flow through the + # normal state-update path and event history. + if req.state_delta is not None: + import uuid - from ..events.event import Event - from ..events.event import EventActions + from ..events.event import Event + from ..events.event import EventActions - state_update_event = Event( - invocation_id="p-" + str(uuid.uuid4()), - author="user", - actions=EventActions(state_delta=req.state_delta), - ) + state_update_event = Event( + invocation_id="p-" + str(uuid.uuid4()), + author="user", + actions=EventActions(state_delta=req.state_delta), + ) + # This will automatically update the session state through + # __update_session_state. + await self.session_service.append_event( + session=session, event=state_update_event + ) - # Append the event to the session - # This will automatically update the session state through __update_session_state - await self.session_service.append_event( - session=session, event=state_update_event - ) + # Apply a title change (if provided) as out-of-band session metadata. + if req.title is not None: + session = await self.session_service.update_session_title( + app_name=app_name, + user_id=user_id, + session_id=session_id, + title=req.title, + ) return session diff --git a/src/google/adk/sessions/base_session_service.py b/src/google/adk/sessions/base_session_service.py index 06eb6a2534a..2d4c0bda0d0 100644 --- a/src/google/adk/sessions/base_session_service.py +++ b/src/google/adk/sessions/base_session_service.py @@ -151,6 +151,34 @@ async def get_user_state( 'call get_session on each result to access the merged state.' ) + async def update_session_title( + self, *, app_name: str, user_id: str, session_id: str, title: str + ) -> Session: + """Sets a human-readable title on an existing session. + + The title is client-provided metadata (e.g. for rendering a chat-history + list); the service only stores it. It is surfaced as ``Session.title`` by + both ``get_session`` and ``list_sessions``. + + Args: + app_name: The name of the app. + user_id: The ID of the user. + session_id: The ID of the session to update. + title: The new title to store on the session. + + Returns: + The updated Session. + + Raises: + SessionNotFoundError: When the session does not exist. + NotImplementedError: When the concrete ``BaseSessionService`` + implementation does not support session titles (e.g. a managed backend + that does not expose a client-set title). + """ + raise NotImplementedError( + f'{type(self).__name__} does not support update_session_title.' + ) + async def append_event(self, session: Session, event: Event) -> Event: """Appends an event to a session object.""" if event.partial: diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 54755a1fbc3..7e1151913e9 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -32,12 +32,16 @@ try: from sqlalchemy import delete from sqlalchemy import event + from sqlalchemy import inspect from sqlalchemy import MetaData from sqlalchemy import select + from sqlalchemy import text + from sqlalchemy import update from sqlalchemy.engine import Connection from sqlalchemy.engine import make_url from sqlalchemy.exc import ArgumentError from sqlalchemy.exc import IntegrityError + from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import async_sessionmaker from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory @@ -168,9 +172,75 @@ def _ensure_schema_indexes_exist( index.create(bind=connection, checkfirst=True) +def _ensure_schema_columns_exist( + connection: Connection, metadata: MetaData +) -> None: + """Adds missing additive, nullable columns to existing tables in place. + + ``metadata.create_all`` only creates missing *tables*; it never alters an + existing table to add a newly declared column. This reconciles that gap for + additive nullable columns (mirroring ``_ensure_schema_indexes_exist`` for + indexes), so a database created before a column was introduced gains it on + the next connect instead of failing every subsequent read/write. + + Only nullable columns are added: a non-nullable column would require a data + backfill and a real schema migration, so it is skipped with a warning. + """ + logger.debug("Ensuring schema columns exist for metadata tables.") + inspector = inspect(connection) + existing_tables = set(inspector.get_table_names()) + preparer = connection.dialect.identifier_preparer + for table in metadata.sorted_tables: + if table.name not in existing_tables: + # create_all just made it with every declared column; nothing to do. + continue + existing_columns = { + column["name"] for column in inspector.get_columns(table.name) + } + for column in table.columns: + if column.name in existing_columns: + continue + if not column.nullable: + logger.warning( + "Cannot add non-nullable column %s.%s in place; a schema" + " migration is required.", + table.name, + column.name, + ) + continue + column_type = column.type.compile(dialect=connection.dialect) + add_column = text( + f"ALTER TABLE {preparer.quote(table.name)} ADD COLUMN" + f" {preparer.quote(column.name)} {column_type}" + ) + try: + # SAVEPOINT so a concurrent add on another instance does not poison + # the surrounding setup transaction (mirrors _get_or_create_state). + with connection.begin_nested(): + connection.execute(add_column) + logger.info( + "Added missing column %s.%s to existing table.", + table.name, + column.name, + ) + except SQLAlchemyError: + refreshed_columns = { + column_info["name"] + for column_info in inspect(connection).get_columns(table.name) + } + if column.name not in refreshed_columns: + raise + logger.debug( + "Column %s.%s already present after a concurrent add.", + table.name, + column.name, + ) + + def _setup_database_schema(connection: Connection, metadata: MetaData) -> None: - """Ensures tables and indexes declared in metadata exist.""" + """Ensures tables, columns, and indexes declared in metadata exist.""" metadata.create_all(bind=connection) + _ensure_schema_columns_exist(connection, metadata) _ensure_schema_indexes_exist(connection, metadata) @@ -650,6 +720,38 @@ async def get_session( ) return session + @override + async def update_session_title( + self, *, app_name: str, user_id: str, session_id: str, title: str + ) -> Session: + await self.prepare_tables() + schema = self._get_schema_classes() + async with self._rollback_on_exception_session() as sql_session: + storage_session = await sql_session.get( + schema.StorageSession, (app_name, user_id, session_id) + ) + if storage_session is None: + raise SessionNotFoundError(f"Session {session_id} not found.") + # Preserve update_time by passing it explicitly so onupdate=func.now() + # does not fire: a title is display metadata and must not invalidate an + # in-flight copy of the session, which would trip the stale-writer check + # on the next append_event. + await sql_session.execute( + update(schema.StorageSession) + .where(schema.StorageSession.app_name == app_name) + .where(schema.StorageSession.user_id == user_id) + .where(schema.StorageSession.id == session_id) + .values(title=title, update_time=storage_session.update_time) + ) + await sql_session.commit() + + session = await self.get_session( + app_name=app_name, user_id=user_id, session_id=session_id + ) + if session is None: + raise SessionNotFoundError(f"Session {session_id} not found.") + return session + @override async def list_sessions( self, *, app_name: str, user_id: Optional[str] = None diff --git a/src/google/adk/sessions/in_memory_session_service.py b/src/google/adk/sessions/in_memory_session_service.py index 73a54f398b8..bebc24c6729 100644 --- a/src/google/adk/sessions/in_memory_session_service.py +++ b/src/google/adk/sessions/in_memory_session_service.py @@ -24,6 +24,7 @@ from . import _session_util from ..errors.already_exists_error import AlreadyExistsError +from ..errors.session_not_found_error import SessionNotFoundError from ..events.event import Event from ..features import FeatureName from ..features import is_feature_enabled @@ -312,6 +313,21 @@ def _delete_session_impl( self.sessions[app_name][user_id].pop(session_id) + @override + async def update_session_title( + self, *, app_name: str, user_id: str, session_id: str, title: str + ) -> Session: + if ( + app_name not in self.sessions + or user_id not in self.sessions[app_name] + or session_id not in self.sessions[app_name][user_id] + ): + raise SessionNotFoundError(f'Session {session_id} not found.') + session = self.sessions[app_name][user_id][session_id] + session.title = title + copied_session = _copy_session(session) + return self._merge_state(app_name, user_id, copied_session) + @override async def get_user_state( self, *, app_name: str, user_id: str diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py index 9b5862d5610..d5ecc3e7cf6 100644 --- a/src/google/adk/sessions/schemas/v1.py +++ b/src/google/adk/sessions/schemas/v1.py @@ -97,6 +97,13 @@ class StorageSession(Base): PreciseTimestamp, default=func.now(), onupdate=func.now() ) + # Optional client-set human-readable title (e.g. for a chat-history list). + # Nullable and additive: existing databases are reconciled in-place by + # DatabaseSessionService._ensure_schema_columns_exist on connect. + title: Mapped[str | None] = mapped_column( + String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True + ) + storage_events: Mapped[list[StorageEvent]] = relationship( "StorageEvent", back_populates="storage_session", @@ -169,6 +176,7 @@ def to_session( last_update_time=self.get_update_timestamp( is_sqlite=is_sqlite, is_postgresql=is_postgresql ), + title=self.title, ) session._storage_update_marker = self.get_update_marker() return session diff --git a/src/google/adk/sessions/session.py b/src/google/adk/sessions/session.py index dab5476ce31..735b8362b77 100644 --- a/src/google/adk/sessions/session.py +++ b/src/google/adk/sessions/session.py @@ -68,6 +68,15 @@ class Session(BaseModel): ), examples=[1_742_000_000.0], ) + title: str | None = Field( + default=None, + description=( + "Optional human-readable title for the session, e.g. for display in" + " a chat-history list. Client-set via update_session_title; the" + " service only stores it." + ), + examples=["Debugging the payments outage"], + ) _storage_update_marker: str | None = PrivateAttr(default=None) """Internal storage revision marker used for stale-session detection.""" diff --git a/src/google/adk/sessions/sqlite_session_service.py b/src/google/adk/sessions/sqlite_session_service.py index 31b95e375c8..dc96c0960e5 100644 --- a/src/google/adk/sessions/sqlite_session_service.py +++ b/src/google/adk/sessions/sqlite_session_service.py @@ -69,6 +69,7 @@ state TEXT NOT NULL, create_time REAL NOT NULL, update_time REAL NOT NULL, + title TEXT, PRIMARY KEY (app_name, user_id, id) ); """ @@ -240,7 +241,7 @@ async def get_session( ) -> Optional[Session]: async with self._get_db_connection() as db: async with db.execute( - "SELECT state, update_time FROM sessions WHERE app_name=? AND" + "SELECT state, update_time, title FROM sessions WHERE app_name=? AND" " user_id=? AND id=?", (app_name, user_id, session_id), ) as cursor: @@ -249,6 +250,7 @@ async def get_session( return None session_state = json.loads(session_row["state"]) last_update_time = session_row["update_time"] + title = session_row["title"] # Build events query query_parts = [ @@ -293,6 +295,7 @@ async def get_session( state=merged_state, events=events, last_update_time=last_update_time, + title=title, ) @override @@ -304,13 +307,13 @@ async def list_sessions( # Fetch sessions if user_id: session_rows = await db.execute_fetchall( - "SELECT id, user_id, state, update_time FROM sessions WHERE" + "SELECT id, user_id, state, update_time, title FROM sessions WHERE" " app_name=? AND user_id=?", (app_name, user_id), ) else: session_rows = await db.execute_fetchall( - "SELECT id, user_id, state, update_time FROM sessions WHERE" + "SELECT id, user_id, state, update_time, title FROM sessions WHERE" " app_name=?", (app_name,), ) @@ -346,6 +349,7 @@ async def list_sessions( state=merged_state, events=[], last_update_time=row["update_time"], + title=row["title"], ) ) return ListSessionsResponse(sessions=sessions_list) @@ -368,6 +372,30 @@ async def get_user_state( async with self._get_db_connection() as db: return await self._get_user_state(db, app_name, user_id) + @override + async def update_session_title( + self, *, app_name: str, user_id: str, session_id: str, title: str + ) -> Session: + async with self._get_db_connection() as db: + # Preserve update_time: a title is display metadata and must not + # invalidate an in-flight copy of the session (which would trip the + # stale-session check on the next append_event). + async with db.execute( + "UPDATE sessions SET title=? WHERE app_name=? AND user_id=? AND id=?", + (title, app_name, user_id, session_id), + ) as cursor: + updated_rows = cursor.rowcount + if updated_rows == 0: + raise SessionNotFoundError(f"Session {session_id} not found.") + await db.commit() + + session = await self.get_session( + app_name=app_name, user_id=user_id, session_id=session_id + ) + if session is None: + raise SessionNotFoundError(f"Session {session_id} not found.") + return session + @override async def append_event(self, session: Session, event: Event) -> Event: if event.partial: @@ -477,9 +505,24 @@ async def _get_db_connection(self): await db.execute(PRAGMA_FOREIGN_KEYS) if not self._schema_ready: await db.executescript(CREATE_SCHEMA_SQL) + await self._ensure_sessions_columns(db) self._schema_ready = True yield db + async def _ensure_sessions_columns(self, db: aiosqlite.Connection) -> None: + """Adds additive nullable columns missing from an existing sessions table. + + CREATE TABLE IF NOT EXISTS never alters an existing table, so a database + created before the `title` column was introduced would lack it. Add it in + place (nullable, no default) so existing databases keep working instead of + failing on the next read/write. + """ + async with db.execute("PRAGMA table_info(sessions)") as cursor: + columns = {row["name"] for row in await cursor.fetchall()} + if "title" not in columns: + await db.execute("ALTER TABLE sessions ADD COLUMN title TEXT") + await db.commit() + async def _get_state( self, db: aiosqlite.Connection, query: str, params: tuple ) -> dict[str, Any]: diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 60abb9aad1a..e26c5a5e6b1 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -381,6 +381,25 @@ async def get_user_state( 'via list_sessions and call get_session on each result.' ) + @override + async def update_session_title( + self, *, app_name: str, user_id: str, session_id: str, title: str + ) -> Session: + """Not supported by the Vertex AI Agent Engine backend. + + The Vertex AI Agent Engine Session API does not expose a client-set title + field, so sessions managed by this backend cannot store a title. + + Raises: + NotImplementedError: Always, because the Vertex AI Agent Engine API does + not provide a client-set session title. + """ + raise NotImplementedError( + 'VertexAiSessionService does not support update_session_title. ' + 'The Vertex AI Agent Engine API does not expose a client-set session ' + 'title.' + ) + @override async def append_event(self, session: Session, event: Event) -> Event: # Update the in-memory session. diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 83cadef8848..01bc0ce1395 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -1487,6 +1487,40 @@ def test_patch_session_not_found(test_app, test_session_info): logger.info("Patch session not found test passed") +def test_update_session_title(test_app, create_test_session): + """Test patching a session title via the PATCH session endpoint.""" + info = create_test_session + url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}" + + response = test_app.patch(url, json={"title": "My session title"}) + assert response.status_code == 200 + patched_session = response.json() + assert patched_session["id"] == info["session_id"] + assert patched_session["title"] == "My session title" + + # Verify the title persisted. + response = test_app.get(url) + assert response.status_code == 200 + assert response.json()["title"] == "My session title" + + +def test_update_session_title_only_adds_no_event(test_app, create_test_session): + """A title-only patch sets the title without recording a state event.""" + info = create_test_session + url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}" + + original = test_app.get(url).json() + original_event_count = len(original.get("events", [])) + + response = test_app.patch(url, json={"title": "Just a title"}) + assert response.status_code == 200 + + retrieved = test_app.get(url).json() + assert retrieved["title"] == "Just a title" + # A title-only patch must not append a state-change event. + assert len(retrieved.get("events", [])) == original_event_count + + def test_agent_run(test_app, create_test_session): """Test running an agent with a message.""" info = create_test_session diff --git a/tests/unittests/sessions/migration/test_database_schema.py b/tests/unittests/sessions/migration/test_database_schema.py index 5381742097b..9997cc5eb32 100644 --- a/tests/unittests/sessions/migration/test_database_schema.py +++ b/tests/unittests/sessions/migration/test_database_schema.py @@ -216,6 +216,77 @@ async def test_prepare_tables_recreates_missing_latest_events_index(tmp_path): ) +@pytest.mark.asyncio +async def test_prepare_tables_adds_missing_title_column(tmp_path): + """A latest-schema DB predating the title column gains it in place on connect.""" + db_path = tmp_path / 'missing_title.db' + db_url = f'sqlite+aiosqlite:///{db_path}' + + # Create a latest-schema DB with a session. + async with DatabaseSessionService(db_url) as session_service: + await session_service.create_session( + app_name='my_app', user_id='test_user', session_id='s1' + ) + + # Simulate a database created before the title column existed. + engine = create_async_engine(db_url) + async with engine.begin() as conn: + await conn.execute(text('ALTER TABLE sessions DROP COLUMN title')) + await engine.dispose() + + engine = create_async_engine(db_url) + async with engine.connect() as conn: + columns_before = await conn.run_sync( + lambda sync_conn: inspect(sync_conn).get_columns('sessions') + ) + await engine.dispose() + assert 'title' not in {column['name'] for column in columns_before} + + # Reconnecting must heal the column in place and support title updates. + async with DatabaseSessionService(db_url) as session_service: + updated = await session_service.update_session_title( + app_name='my_app', user_id='test_user', session_id='s1', title='Healed' + ) + assert updated.title == 'Healed' + fetched = await session_service.get_session( + app_name='my_app', user_id='test_user', session_id='s1' + ) + assert fetched.title == 'Healed' + + engine = create_async_engine(db_url) + async with engine.connect() as conn: + columns_after = await conn.run_sync( + lambda sync_conn: inspect(sync_conn).get_columns('sessions') + ) + await engine.dispose() + assert 'title' in {column['name'] for column in columns_after} + + +@pytest.mark.asyncio +async def test_prepare_tables_does_not_add_title_to_v0_db(tmp_path): + """The title column reconciliation leaves legacy v0 databases untouched.""" + db_path = tmp_path / 'v0_no_title.db' + await create_v0_db(db_path) + db_url = f'sqlite+aiosqlite:///{db_path}' + + async with DatabaseSessionService(db_url) as session_service: + await session_service.create_session( + app_name='my_app', user_id='test_user', session_id='s1' + ) + assert ( + session_service._db_schema_version + == _schema_check_utils.SCHEMA_VERSION_0_PICKLE + ) + + engine = create_async_engine(db_url) + async with engine.connect() as conn: + columns = await conn.run_sync( + lambda sync_conn: inspect(sync_conn).get_columns('sessions') + ) + await engine.dispose() + assert 'title' not in {column['name'] for column in columns} + + @pytest.mark.asyncio async def test_prepare_tables_recreates_missing_v0_events_index(tmp_path): db_path = tmp_path / 'missing_v0_index.db' diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 9ecb3d7eb7e..da1d7cafbfc 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -2181,3 +2181,108 @@ async def test_database_session_service_sqlite_file_timestamp_read_after_reopen( assert retrieved_session.events[0].timestamp == pytest.approx( raw_epoch_float, abs=1.0 ) + + +async def test_update_session_title(session_service): + """A client-set title round-trips through get_session and list_sessions.""" + app_name = 'my_app' + user_id = 'user' + session = await session_service.create_session( + app_name=app_name, user_id=user_id + ) + # Newly created sessions have no title. + assert session.title is None + + updated = await session_service.update_session_title( + app_name=app_name, + user_id=user_id, + session_id=session.id, + title='First title', + ) + assert updated.title == 'First title' + + # get_session reflects the title. + fetched = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session.id + ) + assert fetched is not None + assert fetched.title == 'First title' + + # list_sessions (which omits events/state) still carries the title. + listed = await session_service.list_sessions( + app_name=app_name, user_id=user_id + ) + assert len(listed.sessions) == 1 + assert listed.sessions[0].title == 'First title' + + # Titles can be overwritten. + updated_again = await session_service.update_session_title( + app_name=app_name, + user_id=user_id, + session_id=session.id, + title='Second title', + ) + assert updated_again.title == 'Second title' + refetched = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session.id + ) + assert refetched is not None + assert refetched.title == 'Second title' + + +async def test_update_session_title_missing_session_raises(session_service): + """Updating the title of a nonexistent session raises SessionNotFoundError.""" + with pytest.raises(SessionNotFoundError): + await session_service.update_session_title( + app_name='my_app', + user_id='user', + session_id='does-not-exist', + title='irrelevant', + ) + + +async def test_update_session_title_preserves_last_update_time(session_service): + """Setting a title is metadata-only and must not bump last_update_time. + + A title update should not invalidate an in-flight copy of the session, so it + must leave the session's update time (and thus the stale-writer marker) alone. + """ + app_name = 'my_app' + user_id = 'user' + session = await session_service.create_session( + app_name=app_name, user_id=user_id + ) + before = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session.id + ) + await session_service.update_session_title( + app_name=app_name, user_id=user_id, session_id=session.id, title='Renamed' + ) + after = await session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session.id + ) + assert after.last_update_time == before.last_update_time + + +async def test_sqlite_session_service_adds_missing_title_column(tmp_path): + """SqliteSessionService heals a sessions table that predates the title column.""" + db_path = tmp_path / 'sqlite.db' + service = SqliteSessionService(str(db_path)) + await service.create_session(app_name='a', user_id='u', session_id='s1') + + # Simulate a database created before the title column existed. + conn = sqlite3.connect(str(db_path)) + try: + conn.execute('ALTER TABLE sessions DROP COLUMN title') + conn.commit() + finally: + conn.close() + + # A fresh service instance must heal the column in place on first connect. + healed = SqliteSessionService(str(db_path)) + updated = await healed.update_session_title( + app_name='a', user_id='u', session_id='s1', title='Healed' + ) + assert updated.title == 'Healed' + fetched = await healed.get_session(app_name='a', user_id='u', session_id='s1') + assert fetched.title == 'Healed' diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 2f112f6de97..95f0eb24522 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -1542,3 +1542,13 @@ async def test_get_session_strips_full_resource_name( mock_api_client_instance.agent_engines.sessions.get.assert_called_once_with( name='reasoningEngines/123/sessions/session-123' ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures('mock_get_api_client') +async def test_update_session_title_not_supported(): + session_service = mock_vertex_ai_session_service() + with pytest.raises(NotImplementedError, match='Vertex AI Agent Engine'): + await session_service.update_session_title( + app_name='123', user_id='user', session_id='1', title='My title' + )