Commit Graph
70 Commits
Author SHA1 Message Date
George WealeandCopybara-Service f50847460f fix: Per-session locking and row-level locking in DatabaseSessionService.append_event
This change introduces an in-process `asyncio.Lock` per session to serialize `append_event` calls for the same session ID within a single process. For supported database dialects (MySQL, PostgreSQL, MariaDB), it also uses `SELECT ... FOR UPDATE` to acquire row-level locks on the session, app state, and user state records, preventing race conditions across different processes or database connections. A new test case verifies that concurrent updates to stale session objects correctly merge all state changes.

Close #1049

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 867752676
2026-02-09 13:38:38 -08:00
George WealeandCopybara-Service da73e718ef fix: Enable pool_pre_ping by default for non-SQLite database engines
This change sets `pool_pre_ping=True` in SQLAlchemy engine kwargs for database backends other than SQLite. This helps ensure that connections from the pool are still valid before being used, preventing issues with stale or disconnected connections. Tests are added to verify the default behavior and that explicit overrides are respected

Close #4211

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864886767
2026-02-03 08:16:55 -08:00
George WealeandCopybara-Service 666cebe369 fix: Add update_timestamp_tz property to StorageSession
This property is a compatibility alias that returns the update timestamp as a POSIX timestamp. It infers whether the database is SQLite using sqlalchemy.inspect to call get_update_timestamp correctly

Close #4334

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864595914
2026-02-02 17:39:42 -08:00
Liang WuandCopybara-Service 63a8eba53f fix: Ensure database sessions are always rolled back on errors
Fixes issue#3328

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 863314939
2026-01-30 11:10:02 -08:00
Google Team MemberandCopybara-Service 00aba2d884 feat: Improve asyncio loop handling and test cleanup
This CL enhances asyncio event loop management and test isolation.

-   **BigQuery Analytics Plugin:** Ensure the asyncio event loop is consistently closed within the BigQuery analytics plugin. This prevents potential resource leaks. Add checks to handle potential deadlocks in Python 3.13+ when creating loops during interpreter shutdown.
-   **Test Thread Pool Cleanup:** Introduce a pytest fixture (`cleanup_thread_pools`) to automatically shut down and clear all tool-related thread pools after each test run in `test_functions_thread_pool.py`. This improves test isolation and prevents order-dependent test failures.
-   **Streaming Test Loop Restoration:** Refactor event loop handling in `test_streaming.py`. A new `_run_with_loop` method is introduced in the custom test runners to create a temporary event loop for each test execution, run the coroutine, and crucially, restore the original event loop afterwards. This prevents tests from interfering with each other's loop state.
-   **Resource Closure:** Ensure services are closed properly in tests by adding `await service.close()` in `test_service_factory.py` and using `async with session_service` in `test_session_service.py`.

PiperOrigin-RevId: 863305565
2026-01-30 10:47:09 -08:00
Shangjie ChenandCopybara-Service 1063fa532c fix: Reload the stale session in database_session_service when the update_time of storage is later than the current in memory session object
This changes the behavior of how we handle stale session, instead of reject the transaction entirely, we aggresively refresh the session.

Removed synchronous inspect(self).session calls within StorageSession.update_timestamp_tz and _dialect_name. This introspection was causing deadlocks/hangs when used with sqlalchemy.ext.asyncio in Python 3.13.

Closes issue: https://github.com/google/adk-python/issues/1733

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 861353058
2026-01-26 14:19:14 -08:00
George WealeandCopybara-Service 2367901ec5 chore: Upgrade to headers to 2026
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 858763407
2026-01-20 14:50:09 -08:00
Liang WuandCopybara-Service 4b29d15b3e fix: Handle async driver URLs in migration tool
The migration tool uses synchronous SQLAlchemy engines but users often provide async driver URLs (e.g., postgresql+asyncpg://) since that's what ADK requires at runtime.

This fix:
- Makes `to_sync_url()` public in `_schema_check_utils.py` for reuse
- Updates `migrate_from_sqlalchemy_pickle.py` to convert async URLs
- Updates `migrate_from_sqlalchemy_sqlite.py` to convert async URLs
- Adds comprehensive unit tests for `to_sync_url()` function
- Adds integration test for migration with async driver URLs

Fixes #4176

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 858359061
2026-01-19 19:39:20 -08:00
Liang WuandCopybara-Service 75231a30f1 fix: Handle NOT_FOUND error when fetching Vertex AI sessions
Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 855464349
2026-01-12 18:17:41 -08:00
Alexis MarasiganandCopybara-Service b725045e5a fix: fix httpx client closure during event pagination
Merge https://github.com/google/adk-python/pull/3756

move event iteration inside api_client context in get_session

Move event iteration inside the api_client context manager in VertexAiSessionService.get_session() to prevent client closure during multi-page event fetching.

**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

**1. Link to an existing issue (if applicable):**

- Closes: #3757

**2. Or, if no issue exists, describe the change:**

**Problem:**

When a session contains more than 100 events (requiring pagination), `VertexAiSessionService.get_session()` fails with:

```
RuntimeError: Cannot send a request, as the client has been closed.
```

The root cause is that the `events_iterator` is consumed **outside** the `async with self._get_api_client() as api_client:` context block. When the iterator needs to fetch page 2, 3, etc., the API client has already been closed because the `async with` block has exited.

```python
# Current buggy flow:
async with self._get_api_client() as api_client:
    get_session_response, events_iterator = await asyncio.gather(...)
# ← Client closed here

async for event in events_iterator:  # ← Fails on page 2+ (client closed)
    session.events.append(...)
```

**Solution:**

Move the session creation, user validation, and event iteration **inside** the `async with` block so the API client remains open during the entire pagination process:

```python
async with self._get_api_client() as api_client:
    get_session_response, events_iterator = await asyncio.gather(...)
    # Validation and session creation...
    async for event in events_iterator:  # ← Now works for all pages
        session.events.append(...)
# Client closed after all events are fetched
```

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

```bash
pytest tests/unittests/sessions/test_vertex_ai_session_service.py -v
```

**Added regression test:** `test_get_session_pagination_keeps_client_open`
- Creates a `MockAsyncClientWithPagination` that tracks whether it's inside the `async with` context
- Raises `RuntimeError` if iteration happens outside the context (matching real httpx behavior)
- Simulates 3 pages of events (100 + 100 + 50 = 250 events)
- Verifies all 250 events are successfully retrieved

**Manual End-to-End (E2E) Tests:**

1. Deploy an ADK agent to Vertex AI Agent Engine
2. Create a session and send 100+ messages to accumulate >100 events
3. Verify `get_session()` successfully retrieves all events without error

**Before fix:**
```
RuntimeError: Cannot send a request, as the client has been closed.
```

**After fix:**
- Session with 201 events (3 pages) loads successfully
- All events are retrieved and appended to the session

### Checklist

- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective or that my feature works.
- [x] New and existing unit tests pass locally with my changes.
- [x] I have manually tested my changes end-to-end.
- [x] Any dependent changes have been merged and published in downstream modules.

### Additional context

This bug affects any production deployment where users have extended conversations. Sessions accumulating >100 events (which triggers pagination) become completely unusable as the agent cannot load the session to process new messages.

The fix is minimal and maintains backward compatibility - it only changes the scope of the `async with` block without altering any logic or return values.

**Affected versions:** Tested on google-adk 1.19.0, but the bug exists in earlier versions as well.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3756 from AlexisMarasigan:fix/vertex-ai-session-service-paginatio 01fbafa6524312f24f7c9feaffb07bff0ad49b77
PiperOrigin-RevId: 855451813
2026-01-12 17:27:29 -08:00
Dinesh ThummaandCopybara-Service 94d48fce32 fix: Database reserved keyword issue
The fix will use quotes to escape "key", which is column name in the metadata table. Should work for different database types.

Merge https://github.com/google/adk-python/pull/4106

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4106 from DineshThumma9:fix/mysql-reserved-keyword-issue e39d0d02f3695d6890bc3267417b5dad58f7e8ee
PiperOrigin-RevId: 854411915
2026-01-09 18:20:46 -08:00
George WealeandCopybara-Service b8917bc80e fix: Handle SQLite URLs in SqliteSessionService
The SqliteSessionService now accepts database paths in the form of SQLite URLs (e.g., "sqlite:///./sessions.db", "sqlite+aiosqlite:////absolute.db")

Close #4077

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 853922433
2026-01-08 16:14:07 -08:00
Liang WuandCopybara-Service ce64787c3e feat: Add database schema migration command and script
Final part of https://github.com/google/adk-python/discussions/3605.

This change introduces:
-   A new `adk migrate session` CLI command to run database schema upgrades.
-   A migration script to upgrade from the old Pickle-based session schema (v0) to the new JSON-based schema (v1).
-   A migration runner that orchestrates the upgrade process, handling sequential migrations and using temporary SQLite databases for intermediate steps if needed.
-   Unit tests for the v0 to v1 migration.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 852983323
2026-01-06 16:36:45 -08:00
George WealeandCopybara-Service 4ddb2cb2a8 chore: Close database engines to avoid aiosqlite pytest hangs
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 852428755
2026-01-05 13:18:08 -08:00
George WealeandCopybara-Service 838530ebe0 fix: rehydration of EventActions in StorageEvent.to_event
The change updates the `StorageEvent.to_event` method to use `EventActions.model_validate` when rehydrating the `actions` field. This ensures that nested models within `EventActions`, such as `EventCompaction`, are correctly reconstructed from the stored data

Close #4047

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 852408683
2026-01-05 12:26:10 -08:00
Liang WuandCopybara-Service ba91fea541 feat: Use new JSON-based schema for newly-created databases
Part 2 of https://github.com/google/adk-python/discussions/3605.

The DatabaseSessionService now checks for the usage of a V1 schema based on the "adk_internal_metadata" table. Table creation and subsequent operations use either the V0 or V1 SQLAlchemy models accordingly. New databases will default to V1.

Migration script and CLI command will be provided in the next change.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 845443406
2025-12-16 14:49:31 -08:00
Liang WuandCopybara-Service 7e6ef71eec feat: Introduce new database schema for DatabaseSessionService
Part 1 of https://github.com/google/adk-python/discussions/3605.

This change adds a new schema that uses JSON serialization to store Events data in the database. A new "adk_internal_metadata" table is also added to store information like schema version. Since we want to keep supporting existing DB, we fork from the original schema and call it "v0", while the new one is called "v1".

The change is no-op for existing users. In later change, the new schema will be used for new databases, and migration scripts will be provided for existing databases.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 844986248
2025-12-15 17:32:48 -08:00
Shangjie ChenandCopybara-Service 9d918d45df feat!: Rollback the DB migration as it is breaking
Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 839818479
2025-12-03 10:39:43 -08:00
Shangjie ChenandCopybara-Service 0094eea3ca feat!: Migrate DatabaseSessionService to use JSON serialization schema
Also provide a command line tool `adk migrate session` for DB migration

Addresses https://github.com/google/adk-python/discussions/3605
Addresses https://github.com/google/adk-python/issues/3681

To verify:
```
# Start one postgres DB
docker run --name my-postgres -d -e POSTGRES_DB=agent  -e POSTGRES_USER=agent  -e POSTGRES_PASSWORD=agent  -e PGDATA=/var/lib/postgresql/data/pgdata -v pgvolume:/var/lib/postgresql/data -p 5532:5432 postgres

# Connect to an old version of ADK and produce some query data
adk web --session_service_uri=postgresql://agent:agent@localhost:5532/agent

# Check out to the latest branch and restart ADK web
# You should see error log ask you to migrate the DB

# Start a new DB
docker run --name migration-test-db \
  -d \ --rm \ -e POSTGRES_DB=agent \ -e POSTGRES_USER=agent \ -e POSTGRES_PASSWORD=agent -e PGDATA=/var/lib/postgresql/data/pgdata -v migration_test_vol:/var/lib/postgresql/data -p 5533:5432 postgres

# DB Migration
adk migrate session \
  --source_db_url="postgresql://agent:agent@localhost:5532/agent" \
  --dest_db_url="postgresql://agent:agent@localhost:5533/agent"

# Run ADK web with the new DB
adk web --session_service_uri=postgresql+asyncpg://agent:agent@localhost:5533/agent

# You should see the data from old DB is migrated
```

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 837341139
2025-11-26 19:49:08 -08:00
George WealeandCopybara-Service a3e4ad3cd1 fix: Update session last update time when appending events
The `database_session_service` now updates the `update_time` of a session to the event's timestamp when an event is appended

Close #2721

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 834994070
2025-11-20 18:27:02 -08:00
Hangfei LinandCopybara-Service 3ad30a58f9 fix: Add transcription fields to session events
This change introduces `input_transcription` and `output_transcription` fields to session events, enabling the storage and retrieval of transcription data in both the database and Vertex AI session services.

Closes #3172

Co-authored-by: Hangfei Lin <hangfei@google.com>
PiperOrigin-RevId: 834366848
2025-11-19 11:08:38 -08:00
George WealeandCopybara-Service 0ec01956e8 fix: keep vertex session event history intact
Close #3504

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 833417574
2025-11-17 11:13:08 -08:00
Shangjie ChenandCopybara-Service 9211f4ce8c fix: Use async for to loop through event iterator to get all events in vertex_ai_session_service
Fix https://github.com/google/adk-python/issues/3559

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 832476367
2025-11-14 15:22:27 -08:00
GitMarco27andCopybara-Service 74959414d8 feat: full async implementation of DatabaseSessionService
Merge https://github.com/google/adk-python/pull/2889

# Implement Full async DatabaseSessionService

**Target Issue:** #1005

## Overview

This PR introduces an asynchronous implementation of the `DatabaseSessionService` with minimal breaking changes. The primary goal is to enable effective use of ADK in fully async environments and API endpoints while avoiding event loop blocking during database I/O operations.

## Changes

- Converted `DatabaseSessionService` to use async/await patterns throughout

## Testing Plan

The implementation has been tested following the project's contribution guidelines:

### Unit Tests
- All existing unit tests pass successfully
- Minor update to test requirements added to support `aiosqlite`

### Manual End-to-End Testing
- E2E tests performed using:
  - **LLM Provider:** LiteLLM
  - **Database:** PostgreSQL with `asyncpg` driver

 ```python
from google.adk.sessions.database_session_service import DatabaseSessionService

connection_string: str = (
    "postgresql+asyncpg://PG_USER:PG_PSWD@PG_HOST:5432/PG_DB"
)
session_service: DatabaseSessionService = DatabaseSessionService(
    db_url=connection_string
)

session = await session_service.create_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

sessions = await session_service.list_sessions(app_name="test_app", user_id="test_user")
assert len(sessions.sessions) > 0

session = await session_service.get_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

await session_service.delete_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert (
    await session_service.get_session(
        app_name="test_app", session_id="test_session", user_id="test_user"
    )
    is None
)
```

The implementation have been also tested using the following configurations for llm provider and Runner:

```python
def get_azure_openai_model(deployment_id: str | None = None) -> LiteLlm:
    ...

    if not deployment_id:
        deployment_id = os.getenv("AZURE_OPENAI_DEPLOYMENT_ID")

    logger.info(f"Using Azure OpenAI deployment ID: {deployment_id}")

    return LiteLlm(
        model=f"azure/{os.getenv('AZURE_OPENAI_DEPLOYMENT_ID')}",
        stream=True,
    )

...

    @staticmethod
    def _get_runner(agent: Agent) -> Runner:
        storage=DatabaseSessionService(db_url=get_pg_connection_string())
        return Runner(
            agent=agent,
            app_name=APP_NAME,
            session_service=storage,
        )

...

async for event in self.runner.run_async(
    user_id=user_id,
    session_id=session_id,
    new_message=content,
    run_config=(
        RunConfig(
            streaming_mode=StreamingMode.SSE, response_modalities=["TEXT"]
        )
        if stream
        else RunConfig()
    ),
):
    last_event = event
    if stream:
        yield event

...

```

## Breaking Changes

- Database connection string format may need updates for async drivers

Co-authored-by: Shangjie Chen <deanchen@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2889 from GitMarco27:feature/async_database_session_service e1b1b14934c1fb7975a6832cdd1549e94acab985
PiperOrigin-RevId: 830525148
2025-11-10 11:18:02 -08:00
Amy WuandCopybara-Service 7937e9ed40 fix: Use async with for API client in Vertex session management to ensure proper resource management
PiperOrigin-RevId: 830485639
2025-11-10 09:49:10 -08:00