Merge pull request #55 from luis5tb/bugfix-adk-sessions

Bugfix adk sessions and configuration

Author: luis5tb

.env.example

@@ -100,6 +100,18 @@ DATABASE_URL=sqlite+aiosqlite:///./lightspeed_agent.db
 # For PostgreSQL in production:
 # DATABASE_URL=postgresql+asyncpg://user:password@localhost:5432/lightspeed_agent
 
+# -----------------------------------------------------------------------------
+# Session Configuration
+# -----------------------------------------------------------------------------
+# Session storage backend: "memory" (default) or "database"
+# "memory": Sessions stored in-memory (lost on restart, suitable for dev)
+# "database": Sessions persisted to PostgreSQL (requires SESSION_DATABASE_URL)
+SESSION_BACKEND=memory
+
+# Session database URL (required when SESSION_BACKEND=database)
+# Use a SEPARATE database from DATABASE_URL for security isolation.
+# SESSION_DATABASE_URL=postgresql+asyncpg://sessions:password@localhost:5433/agent_sessions
+
 # -----------------------------------------------------------------------------
 # Google Cloud Service Control (for usage reporting)
 # -----------------------------------------------------------------------------

deploy/cloudrun/service.yaml

@@ -126,9 +126,11 @@ spec:
                 secretKeyRef:
                   name: database-url
                   key: latest
-            # Session Database (required for session persistence)
-            # If not set, uses in-memory storage (sessions lost on restart)
-            # For production, always configure this for session persistence
+            # Session backend: "database" for production persistence,
+            # "memory" for in-memory (sessions lost on restart)
+            - name: SESSION_BACKEND
+              value: "database"
+            # Session Database (required when SESSION_BACKEND=database)
             - name: SESSION_DATABASE_URL
               valueFrom:
                 secretKeyRef:

deploy/podman/lightspeed-agent-configmap.yaml

@@ -40,6 +40,11 @@ data:
   AGENT_HOST: "0.0.0.0"
   AGENT_PORT: "8000"
 
+  # Session Configuration
+  # Session backend: "memory" for dev (no persistence), "database" for persistent sessions
+  # To use "database", ensure SESSION_DATABASE_URL is configured in the secret
+  SESSION_BACKEND: "memory"
+
   # Database Configuration
   # DATABASE_URL and passwords are stored in secrets (contain credentials)
   # Marketplace PostgreSQL (in marketplace-handler pod)

docs/configuration.md

@@ -101,7 +101,12 @@ AGENT_PORT=8000
 | Variable | Default | Description |
 |----------|---------|-------------|
 | `DATABASE_URL` | `sqlite+aiosqlite:///./lightspeed_agent.db` | Marketplace database connection URL (orders, DCR clients, auth) |
-| `SESSION_DATABASE_URL` | (uses DATABASE_URL) | Session database URL for ADK sessions. Optional - for security isolation. |
+| `SESSION_BACKEND` | `memory` | Session storage backend: `memory` (in-memory, no persistence) or `database` (PostgreSQL, persistent) |
+| `SESSION_DATABASE_URL` | *(empty)* | Session database URL for ADK sessions. Required when `SESSION_BACKEND=database`. |
+
+> **Note:** Setting `SESSION_BACKEND=database` without providing `SESSION_DATABASE_URL`
+> will cause a startup validation error. This is intentional to prevent running
+> production workloads without session persistence.
 
 **SQLite (Development):**
 
@@ -121,11 +126,14 @@ DATABASE_URL=postgresql+asyncpg://user:password@localhost:5432/lightspeed_agent
 DATABASE_URL=postgresql+asyncpg://user:password@/lightspeed_agent?host=/cloudsql/project:region:instance
 ```
 
-**Security Isolation (Optional):**
+**Security Isolation (Recommended for Production):**
 
-For production deployments, you can use separate databases for marketplace data and agent sessions:
+For production deployments, use `SESSION_BACKEND=database` with a separate database for agent sessions:
 
 ```bash
+# Explicit database backend for sessions
+SESSION_BACKEND=database
+
 # Shared marketplace database (orders, DCR clients, auth data)
 DATABASE_URL=postgresql+asyncpg://marketplace:pass@db:5432/marketplace
 
@@ -138,6 +146,14 @@ This separation ensures:
 - Compromised agents can't access DCR credentials or order information
 - Different retention policies can be applied to sessions vs. marketplace data
 
+**Switching to In-Memory Sessions:**
+
+To disable database session persistence on a running deployment (e.g., for debugging),
+set `SESSION_BACKEND=memory` and redeploy:
+
+- **Cloud Run:** Update the `SESSION_BACKEND` env var to `memory` and redeploy the service.
+- **Podman:** Update `SESSION_BACKEND` in the ConfigMap to `memory` and restart the pod.
+
 ### Dynamic Client Registration (DCR)
 
 DCR allows Google Cloud Marketplace customers to automatically register as OAuth clients.
@@ -365,6 +381,7 @@ LOG_LEVEL=DEBUG
 LOG_FORMAT=text
 AGENT_LOGGING_DETAIL=detailed
 DATABASE_URL=sqlite+aiosqlite:///./dev.db
+SESSION_BACKEND=memory
 ```
 
 ### Staging
@@ -376,6 +393,8 @@ SKIP_JWT_VALIDATION=false
 LOG_LEVEL=INFO
 LOG_FORMAT=json
 DATABASE_URL=postgresql+asyncpg://user:pass@staging-db:5432/insights
+SESSION_BACKEND=database
+SESSION_DATABASE_URL=postgresql+asyncpg://sessions:pass@staging-db:5432/sessions
 ```
 
 ### Production
@@ -386,5 +405,6 @@ DEBUG=false
 SKIP_JWT_VALIDATION=false
 LOG_LEVEL=INFO
 LOG_FORMAT=json
-# DATABASE_URL from Secret Manager
+SESSION_BACKEND=database
+# DATABASE_URL and SESSION_DATABASE_URL from Secret Manager
 ```

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "aiosqlite>=0.20.0",
     # PostgreSQL drivers
     "asyncpg>=0.29.0",  # Async driver for SQLAlchemy async operations
-    "psycopg2-binary>=2.9.0",  # Sync driver for ADK DatabaseSessionService
+    "psycopg2-binary>=2.9.0",  # Sync PostgreSQL driver (used by alembic migrations)
     # Google Cloud authentication (ADC for Procurement API calls)
     "google-auth>=2.0.0",
     "requests>=2.20.0",  # Required by google.auth.transport.requests

src/lightspeed_agent/api/a2a/a2a_setup.py

@@ -6,7 +6,9 @@
 """
 
 import logging
+import re
 from typing import Any
+from urllib.parse import urlparse
 
 from a2a.server.apps import A2AFastAPIApplication
 from a2a.server.request_handlers import DefaultRequestHandler
@@ -28,61 +30,70 @@
 logger = logging.getLogger(__name__)
 
 
+def _normalize_db_url(url: str) -> str:
+    """Ensure a database URL uses an async driver for SQLAlchemy.
+
+    ADK's DatabaseSessionService requires ``create_async_engine``, so any
+    synchronous PostgreSQL scheme must be replaced with ``postgresql+asyncpg``.
+    """
+    scheme, remainder = url.split("://", 1)
+    normalized_scheme = scheme.lower()
+
+    sync_postgres_schemes = {
+        "postgres",
+        "postgresql",
+        "postgresql+psycopg",
+        "postgresql+psycopg2",
+    }
+    if normalized_scheme in sync_postgres_schemes:
+        return f"postgresql+asyncpg://{remainder}"
+
+    return url
+
+
 def _get_session_service() -> Any:
-    """Get the appropriate session service based on configuration.
+    """Get the appropriate session service based on SESSION_BACKEND setting.
+
+    Uses SESSION_BACKEND to determine the session storage:
+    - ``"memory"``: InMemorySessionService (default, no persistence)
+    - ``"database"``: DatabaseSessionService (requires SESSION_DATABASE_URL)
 
-    For production, uses DatabaseSessionService which persists sessions to PostgreSQL.
-    For development, uses InMemorySessionService.
+    When SESSION_BACKEND is ``"database"``, failures are raised immediately
+    rather than silently falling back to in-memory, so misconfigurations are
+    caught at startup.
 
     Security Note:
-        SESSION_DATABASE_URL must be explicitly set to use database persistence.
-        This prevents accidental use of the marketplace database (DATABASE_URL)
-        for session storage, ensuring:
-        - Agents only have access to session data, not marketplace/auth data
-        - Compromised agents can't access DCR credentials or order information
-        - Different retention policies can be applied to sessions vs. marketplace data
+        SESSION_DATABASE_URL should point to a separate database from
+        DATABASE_URL to ensure agents only access session data, not
+        marketplace/auth data.
 
     Returns:
         Session service instance (DatabaseSessionService or InMemorySessionService).
     """
     settings = get_settings()
 
-    # Only use database session service if SESSION_DATABASE_URL is explicitly set
-    # Do NOT fall back to DATABASE_URL to avoid mixing session and marketplace data
-    session_db_url = settings.session_database_url
+    if settings.session_backend == "database":
+        from google.adk.sessions import DatabaseSessionService
+
+        # SESSION_DATABASE_URL is guaranteed non-empty by the model validator
+        db_url = _normalize_db_url(settings.session_database_url)
+
+        # Log which database is being used (without credentials)
+        parsed = urlparse(db_url)
+        db_host = parsed.hostname or parsed.query or "local"
+        logger.info(
+            "Using DatabaseSessionService for session persistence (host=%s)",
+            db_host,
+        )
 
-    # Use database session service for production (non-SQLite databases)
-    if session_db_url and not session_db_url.startswith("sqlite"):
         try:
-            from google.adk.sessions import DatabaseSessionService
-
-            # ADK's DatabaseSessionService uses synchronous SQLAlchemy,
-            # so we need to convert the async URL to sync format
-            db_url = session_db_url
-            if "postgresql+asyncpg" in db_url:
-                # Convert asyncpg URL to sync psycopg2 format
-                db_url = db_url.replace("postgresql+asyncpg", "postgresql+psycopg2")
-            elif "postgresql+aiopg" in db_url:
-                db_url = db_url.replace("postgresql+aiopg", "postgresql+psycopg2")
-
-            # Log which database is being used (without credentials)
-            db_host = db_url.split("@")[-1].split("/")[0] if "@" in db_url else "local"
-            logger.info(
-                "Using DatabaseSessionService for session persistence (host=%s)",
-                db_host,
-            )
             return DatabaseSessionService(db_url=db_url)
-        except ImportError as e:
-            logger.warning(
-                "DatabaseSessionService not available (%s), falling back to InMemorySessionService",
-                e,
-            )
         except Exception as e:
-            logger.warning(
-                "Failed to initialize DatabaseSessionService (%s), "
-                "falling back to InMemorySessionService",
-                e,
-            )
+            # Sanitize error message to avoid leaking credentials from URLs
+            sanitized_msg = re.sub(r"://[^@]+@", "://***@", str(e))
+            raise RuntimeError(
+                f"Failed to initialize DatabaseSessionService: {sanitized_msg}"
+            ) from None
 
     logger.info("Using InMemorySessionService for session management")
     return InMemorySessionService()  # type: ignore[no-untyped-call]

src/lightspeed_agent/config/settings.py

@@ -233,14 +233,22 @@ class Settings(BaseSettings):
         description="Maximum overflow connections beyond pool size",
     )
 
-    # Session database: stores ADK sessions, conversation history, memory
+    # Session configuration: controls ADK session storage backend
     # Separate from marketplace DB for security isolation - each agent can have its own
+    session_backend: Literal["memory", "database"] = Field(
+        default="memory",
+        description=(
+            "Session storage backend. "
+            "'memory' uses in-memory sessions (lost on restart). "
+            "'database' uses DatabaseSessionService and requires SESSION_DATABASE_URL."
+        ),
+    )
     session_database_url: str = Field(
         default="",
         description=(
             "Session database URL for ADK sessions."
-            " If empty, uses DATABASE_URL."
-            " For security isolation, use a separate database."
+            " Required when SESSION_BACKEND=database."
+            " For security isolation, use a separate database from DATABASE_URL."
         ),
     )
 
@@ -307,6 +315,17 @@ def _block_skip_jwt_in_production(self) -> "Settings":
             )
         return self
 
+    @model_validator(mode="after")
+    def _validate_session_backend(self) -> "Settings":
+        """Ensure SESSION_DATABASE_URL is set when SESSION_BACKEND=database."""
+        if self.session_backend == "database" and not self.session_database_url:
+            raise ValueError(
+                "SESSION_BACKEND=database requires SESSION_DATABASE_URL to be set. "
+                "Provide a PostgreSQL connection URL, e.g.: "
+                "SESSION_DATABASE_URL=postgresql+asyncpg://user:pass@host:5432/sessions"
+            )
+        return self
+
     # OpenTelemetry Configuration
     otel_enabled: bool = Field(
         default=False,