Commit Graph
65 Commits
Author SHA1 Message Date
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
Shangjie ChenandCopybara-Service e218254495 feat: Add SqliteSessionService and a migration script to migrate existing DB using DatabaseSessionService to SqliteSessionService
The new Sqlite version has fixed schema and use a single column to store Event data, this should avoid DB migration for future add to the Event object.

- This change introduces `SqliteSessionService`, an asynchronous session service using `aiosqlite` that stores event data as JSON within SQLite.
- A migration script, `migrate_from_sqlalchemy_sqlite.py`, is included to transition data from the older SQLAlchemy-based SQLite schema to this new format.
- The CLI service registry is updated to use SqliteSessionService for sqlite:// URIs.
- Throw error when user trying to access a legacy DB and advice the user to do the migration.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 829971174
2025-11-08 21:19:38 -08:00
hung12ctandCopybara-Service fc15c9a0c3 fix: Update DynamicPickleType to support MySQL dialect
Merge https://github.com/google/adk-python/pull/3282

The `process_bind_param` and `process_result_value` methods in the `DynamicPickleType` class have been modified to handle MySQL dialect in addition to Spanner. This change ensures that pickled values are correctly processed for both database types.

**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: #3283

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

_If applicable, please follow the issue templates to provide as much detail as
possible._

**Problem:**
When using `DatabaseSessionService` with MySQL backend in google-adk v1.17.0, the application crashes with the following error: app.resources.runner:event_generator:260 - Error in event_generator: (builtins.TypeError) 'tuple' object cannot be interpreted as an integer
<img width="1237" height="129" alt="image" src="https://github.com/user-attachments/assets/0a5fc223-600a-4a92-8443-4d37fb1267f6" />

Root cause: The `DynamicPickleType` class in `database_session_service.py` configures MySQL dialect to use `LONGBLOB` for storing pickled data (line 117-118), but the `process_bind_param` and `process_result_value` methods only handle pickle serialization/deserialization for Spanner dialect, not MySQL. This causes MySQL to attempt storing raw Python objects instead of pickled bytes, leading to serialization errors and potential data corruption.

**Solution:**
Added MySQL to the pickle serialization logic in both `process_bind_param` and `process_result_value` methods, treating it the same way as Spanner dialect. This ensures that:
- Data is properly pickled to bytes before being stored in MySQL's LONGBLOB column
- Data is properly unpickled when retrieved from the database
- No breaking changes to existing functionality for other dialects (SQLite, PostgreSQL)

### Testing Plan

_Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes._

**Unit Tests:**

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

**Summary of `pytest` results:**
<img width="929" height="306" alt="image" src="https://github.com/user-attachments/assets/3d548b96-ac49-4101-8405-a289a722293c" />

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

_Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix._

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

_Add any other context or screenshots about the feature request here._

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3282 from hung12ct:fix/mysql-pickle-serialization d9df37adb7dfbbfd8502a0fe65c4f8bca3d0d978
PiperOrigin-RevId: 825834360
2025-10-29 20:39:58 -07:00
Google Team MemberandCopybara-Service 9014a849ea feat: Add api key argument to Vertex Session and Memory services for Express Mode support
We also change VertexAiSessionService and VertexAiMemoryBankService to both use keyword arguments for project, location, agent engine id, and express mode api key

PiperOrigin-RevId: 825719331
2025-10-29 14:54:12 -07:00
Shangjie ChenandCopybara-Service f7e2a7a40e chore: Make VertexAiSessionService fully asynchronous
This CL refactors VertexAiSessionService to use the asynchronous aio client for all Vertex AI API calls. This ensures that the service methods are non-blocking and can be used effectively in an asyncio environment.

PiperOrigin-RevId: 824573356
2025-10-27 10:19:25 -07:00
Shangjie ChenandCopybara-Service 2a901d12f4 chore: Raise AlreadyExistsError when trying to create a resource with same ID
Move the dedupe logic into session service so that the internal error can be surfaced to client

PiperOrigin-RevId: 822294430
2025-10-21 15:11:33 -07:00