Lab 13.5: Extending ADK with Custom Firestore Persistence

Goal

In this lab, you will learn how to extend the ADK's core functionality by implementing a Custom Session Service. You will take a specialized FirestoreSessionService implementation and "plug it in" to an agent's runtime.

By the end of this lab, you will have an agent that persists its memory to Google Cloud Firestore, surviving restarts of the Python process.

Step 1: Pre-requisites

Enable Firestore: In the GCP Console, ensure you have a Firestore database in Native mode.
Auth: gcloud auth application-default login
Dependencies:
Environment setup:
- Local machine: If you're continuing the course in sequence, cd into the adk-training project you created in Module 2 and skip the block below -- your environment and dependencies already carry forward. If you're starting fresh at this module, create that project first:
  uv init adk-training --python 3.10 && cd adk-training uv add "google-adk>=2.1.0" python-dotenv
- DevContainer / Codespaces (Alternative): If you opened this repository in a DevContainer or GitHub Codespaces, Python, uv, and the core ADK dependencies are already pre-installed in your container environment.
```
uv add google-cloud-firestore
```

Step 2: The Custom Provider

Create a file named firestore_provider.py. This contains the implementation that inherits from BaseSessionService. Study the code to see how it uses firestore.AsyncClient to map ADK concepts (Apps, Users, Sessions) to document paths.

# firestore_provider.py
from typing import Any, Optional
import uuid
import time
from google.cloud import firestore
from google.adk.sessions.base_session_service import BaseSessionService, GetSessionConfig, ListSessionsResponse
from google.adk.sessions.session import Session
from google.adk.events.event import Event

class FirestoreSessionService(BaseSessionService):
    def __init__(self, project_id: str):
        self._client = firestore.AsyncClient(project=project_id)

    async def create_session(self, *, app_name: str, user_id: str, state: Optional[dict] = None, session_id: Optional[str] = None) -> Session:
        sid = session_id or str(uuid.uuid4())
        # Reference: apps/{app}/users/{user}/sessions/{sid}
        session_ref = self._client.collection("apps").document(app_name)\
            .collection("users").document(user_id)\
            .collection("sessions").document(sid)

        # Check if exists, else create
        doc = await session_ref.get()
        if not doc.exists:
            await session_ref.set({
                "created_at": time.time(),
                "state": state or {}
            })

        return Session(id=sid, app_name=app_name, user_id=user_id, state=doc.to_dict().get("state", {}) if doc.exists else (state or {}))

    async def get_session(self, *, app_name: str, user_id: str, session_id: str, config: Optional[GetSessionConfig] = None) -> Optional[Session]:
        session_ref = self._client.collection("apps").document(app_name)\
            .collection("users").document(user_id)\
            .collection("sessions").document(session_id)
        doc = await session_ref.get()
        if not doc.exists:
            return None

        # Rebuild the event history from the events sub-collection -- this is
        # what actually gives the agent its memory across separate script
        # runs. Returning state alone (without events) isn't enough: the
        # model recalls prior turns from conversation history, not just the
        # state dict.
        events = []
        async for event_doc in session_ref.collection("events").order_by("timestamp").stream():
            events.append(Event.model_validate(event_doc.to_dict()))

        return Session(
            id=session_id,
            app_name=app_name,
            user_id=user_id,
            state=doc.to_dict().get("state", {}),
            events=events,
        )

    async def append_event(self, session: Session, event: Event) -> Event:
        # Let the base implementation apply event.actions.state_delta to
        # session.state in-memory first (it also handles temp-scoped state
        # and trims it before we persist) -- without this, state changes are
        # silently dropped: they'd never make it past this method.
        event = await super().append_event(session=session, event=event)

        # Save event to sub-collection
        event_ref = self._client.collection("apps").document(session.app_name)\
            .collection("users").document(session.user_id)\
            .collection("sessions").document(session.id)\
            .collection("events").document(event.invocation_id)

        await event_ref.set(event.model_dump())
        print(f"🔥 [Firestore] Persisted event from {event.author}")

        # Persist the now-updated state too, so a fresh get_session() call
        # (e.g. from a new process) sees it.
        await self.update_session_state(session)
        return event

    async def update_session_state(self, session: Session) -> None:
        # Update the main session document's state
        session_ref = self._client.collection("apps").document(session.app_name)\
            .collection("users").document(session.user_id)\
            .collection("sessions").document(session.id)

        await session_ref.update({"state": session.state})
        print(f"🔥 [Firestore] Updated session state in cloud.")

    # Mandatory stubs for BaseSessionService
    async def list_sessions(self, *, app_name: str, user_id: Optional[str] = None) -> ListSessionsResponse:
        return ListSessionsResponse(sessions=[])

    async def delete_session(self, *, app_name: str, user_id: str, session_id: str) -> None:
        pass

Step 3: Integrate the Provider

Exercise: Open agent.py. Your task is to modify the main() function to use your new FirestoreSessionService instead of the default in-memory runner.

# agent.py
import asyncio
import os
from google.adk import Agent, Runner
from google.adk.apps import App
from firestore_provider import FirestoreSessionService # Import your custom class
from dotenv import load_dotenv

load_dotenv()

# --- 1. Define a simple agent ---
agent = Agent(
    model="gemini-3.5-flash",
    name="PersistentAgent",
    instruction="You are a helpful assistant that remembers the user's favorite color."
)

async def main():
    # TODO: 1. Setup metadata
    project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
    app = App(name="extensibility_demo", root_agent=agent)
    
    # TODO: 2. Instantiate your custom service
    # custom_fs = ...
    
    # TODO: 3. Create a base Runner (NOT InMemoryRunner)
    # Inject your custom_fs into the session_service parameter
    # runner = Runner(app=app, session_service=...)
    
    # 4. Test it
    print("Agent is now powered by Firestore persistence!")
    await runner.run_debug("My favorite color is blue.", user_id="student_1")

if __name__ == "__main__":
    asyncio.run(main())

Step 4: Verify Persistence

Run the script once to set the favorite color.
Stop the script.
Modify the last line of main() to ask: "What is my favorite color?".
Run it again. If the agent answers "Blue", your custom provider is working!

Self-Reflection Questions

How does the use of an Abstract Base Class (BaseSessionService) make the ADK more flexible for large companies?
If you wanted to use Redis instead of Firestore, what parts of the firestore_provider.py would you need to rewrite?
Why is it better to inject the session service into the Runner rather than hardcoding it inside the Agent?

🔐 Reference Solution Vault

🔐

Module 13_5 Reference Solution Vault

Active Comprehension Gate • Zero Guessable URLs • Guided Verification

Step 1/3: Question 1 of 3

Goal​

Step 1: Pre-requisites​

Step 2: The Custom Provider​

Step 3: Integrate the Provider​

Step 4: Verify Persistence​

Self-Reflection Questions​

🔐 Reference Solution Vault​