Skip to content

feat(sessions): introduce pluggable SessionDataTransformer hooks for …#5280

Open
sgmoorthy wants to merge 1 commit intogoogle:mainfrom
sgmoorthy:feat/session-data-masking
Open

feat(sessions): introduce pluggable SessionDataTransformer hooks for …#5280
sgmoorthy wants to merge 1 commit intogoogle:mainfrom
sgmoorthy:feat/session-data-masking

Conversation

@sgmoorthy
Copy link
Copy Markdown

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):

  • Closes: #drafted_issue_number
  • Related: #drafted_issue_number

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

Problem:
Saving sensitive strings, tokens, and identifiers (e.g. LLM inputs) inherently places unencrypted PII straight into the database table JSON schema. The current DatabaseSessionService persists bare event payloads and state natively without an option to selectively redact or encrypt content prior to database ingestion.

Solution:
Exposes a transformer: Optional[SessionDataTransformer] = None argument in DatabaseSessionService. It adopts a formal Protocol containing 4 hook strategies ([before|after]_[persist|load]_[event|state]), allowing customized tokenizations, field-level encryption, or redactions natively. This process reliably isolates raw database serialization blocks and keeps caller memory safely un-mutated, protecting local state access while securely scrubbing the storage records.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Unit Test Summary:

tests\unittests\sessions\test_session_service.py .... [100%]
======================= 4 passed, 2 warnings in 23.60s ========================

4 additional tests successfully verified StorageEvent behaviors using a MockPIIMaskerTransformer. Assertions validated JSON mappings cleanly tokenizing strings (e.g. foo -> foo_masked) strictly in SQL layers while raising programmatic failures correctly on bad schemas. All existing legacy unittests successfully passed locally alongside it.

Manual End-to-End (E2E) Tests:

I verified tokenizations correctly function within live applications via the standard Python initialization methodology:

Setup Script:

import asyncio
from google.adk.sessions import DatabaseSessionService, SessionDataTransformer
from google.adk.events import Event, EventActions

class CustomRedactor(SessionDataTransformer):
    def before_persist_state(self, state):
        return {k: v.replace("John Doe", "[REDACTED_NAME]") if isinstance(v, str) else v for k, v in state.items()}
    def after_load_state(self, state): return state
    def before_persist_event(self, event): return event
    def after_load_event(self, event): return event

async def main():
    service = DatabaseSessionService("sqlite+aiosqlite:///demo_db.sqlite", transformer=CustomRedactor())
    session = await service.create_session(app_name="banking", user_id="u123", session_id="flow_1")
    event = Event(author="system", invocation_id="1", actions=EventActions(state_delta={"customer_name": "John Doe"}))
    await service.append_event(session, event)
    
    print(f"[Live Session State]: {session.state}") 

    async with service.db_engine.connect() as conn:
        from sqlalchemy import text
        row = await conn.execute(text("SELECT state FROM user_states WHERE user_id = 'u123'"))
        print(f"[Raw DB Storage Payload]: {row.fetchone()[0]}")

asyncio.run(main())

Logs/Evidence:

[Live Session State]: {'customer_name': 'John Doe'}
[Raw DB Storage Payload]: {"customer_name": "[REDACTED_NAME]"}

(The generated SQLite verification successfully proved backend storage intercepts [REDACTED_NAME], while the caller runtime session objects simultaneously retrieved standard John Doe strings).

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Built and installed the local configuration to a clean isolated environment using uv build and uv venv to guarantee python wheel module imports operated effectively. All legacy instantiations gracefully default to backwards-compatible structures.

@adk-bot adk-bot added the services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc label Apr 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants