Improve restore flow and seed data validation

Author: acailic

agent_debugger_sdk/core/context/session_manager.py

@@ -5,7 +5,7 @@
 import uuid
 from collections.abc import Awaitable, Callable
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import httpx
 
@@ -21,6 +21,77 @@ class _CheckpointRestoreError(Exception):
     pass
 
 
+def _resolve_restore_server_url(server_url: str | None) -> str:
+    """Resolve the checkpoint restore server URL."""
+    if server_url is not None:
+        return server_url
+
+    from agent_debugger_sdk.config import get_config
+
+    config = get_config()
+    return config.endpoint or "http://localhost:8000"
+
+
+async def _fetch_checkpoint_payload(
+    client: httpx.AsyncClient,
+    checkpoint_id: str,
+    server_url: str,
+) -> dict[str, Any]:
+    """Fetch and decode checkpoint payload data from the server."""
+    try:
+        response = await client.get(f"{server_url}/api/checkpoints/{checkpoint_id}")
+        response.raise_for_status()
+        payload = response.json()
+    except httpx.HTTPStatusError as e:
+        raise _CheckpointRestoreError(
+            f"Failed to restore checkpoint {checkpoint_id!r} from {server_url}: "
+            f"{e.response.status_code} {e.response.reason_phrase}"
+        ) from e
+    except httpx.RequestError as e:
+        raise _CheckpointRestoreError(
+            f"Network error while restoring checkpoint {checkpoint_id!r} from {server_url}: {e}"
+        ) from e
+    except Exception as e:
+        raise _CheckpointRestoreError(
+            f"Unexpected error while restoring checkpoint {checkpoint_id!r} from {server_url}: {e}"
+        ) from e
+
+    if not isinstance(payload, dict):
+        raise _CheckpointRestoreError(
+            f"Unexpected checkpoint payload type for {checkpoint_id!r} from {server_url}: "
+            f"{type(payload).__name__}"
+        )
+
+    return payload
+
+
+def _build_restored_session(
+    checkpoint_id: str,
+    checkpoint_data: dict[str, Any],
+    *,
+    session_id: str | None = None,
+    label: str = "",
+) -> tuple[Session, BaseCheckpointState | None]:
+    """Build a restored Session and validated checkpoint state from payload data."""
+    from agent_debugger_sdk.checkpoints import validate_checkpoint_state
+
+    state_dict = checkpoint_data.get("state", {})
+    original_session_id = checkpoint_data.get("session_id", "")
+
+    session = Session(
+        id=session_id or str(uuid.uuid4()),
+        agent_name=label or f"restored from {checkpoint_id[:8]}",
+        framework=state_dict.get("framework", "custom"),
+        config={
+            "restored_from_checkpoint": checkpoint_id,
+            "original_session_id": original_session_id,
+        },
+    )
+
+    restored_state = validate_checkpoint_state(state_dict)
+    return session, restored_state
+
+
 class SessionManager:
     """Manage session lifecycle for TraceContext.
 
@@ -79,48 +150,9 @@ async def restore_from_checkpoint(
         Example:
             >>> session, state = await SessionManager.restore_from_checkpoint("ckpt_123")
         """
-        from agent_debugger_sdk.checkpoints import validate_checkpoint_state
-        from agent_debugger_sdk.config import get_config
+        resolved_server_url = _resolve_restore_server_url(server_url)
 
-        if server_url is None:
-            config = get_config()
-            server_url = config.endpoint or "http://localhost:8000"
-
-        # Fetch checkpoint data using a temporary client (avoids connection leaks)
-        # All processing happens inside the context manager to ensure checkpoint_data is in scope
         async with httpx.AsyncClient() as client:
-            try:
-                response = await client.get(f"{server_url}/api/checkpoints/{checkpoint_id}")
-                response.raise_for_status()
-                checkpoint_data = response.json()
-            except httpx.HTTPStatusError as e:
-                raise _CheckpointRestoreError(
-                    f"Failed to restore checkpoint {checkpoint_id!r} from {server_url}: "
-                    f"{e.response.status_code} {e.response.reason_phrase}"
-                ) from e
-            except httpx.RequestError as e:
-                raise _CheckpointRestoreError(
-                    f"Network error while restoring checkpoint {checkpoint_id!r} from {server_url}: {e}"
-                ) from e
-            except Exception as e:
-                # Catch any other unexpected errors and wrap them
-                raise _CheckpointRestoreError(
-                    f"Unexpected error while restoring checkpoint {checkpoint_id!r} from {server_url}: {e}"
-                ) from e
-
-            # Process the checkpoint data inside the context manager where checkpoint_data is in scope
-            state_dict = checkpoint_data.get("state", {})
-            original_session_id = checkpoint_data.get("session_id", "")
-
-            session = Session(
-                id=session_id or str(uuid.uuid4()),
-                agent_name=label or f"restored from {checkpoint_id[:8]}",
-                framework=state_dict.get("framework", "custom"),
-                config={
-                    "restored_from_checkpoint": checkpoint_id,
-                    "original_session_id": original_session_id,
-                },
-            )
-
-            restored_state = validate_checkpoint_state(state_dict)
-            return session, restored_state
+            checkpoint_data = await _fetch_checkpoint_payload(client, checkpoint_id, resolved_server_url)
+
+        return _build_restored_session(checkpoint_id, checkpoint_data, session_id=session_id, label=label)

api/app_context.py

@@ -28,29 +28,35 @@ def init_app_context() -> None:
     _redaction_pipeline = RedactionPipeline.from_config()
 
 
+def _ensure_initialized() -> None:
+    """Initialize app context lazily when first accessed."""
+    if engine is None or async_session_maker is None or trace_intelligence is None or _redaction_pipeline is None:
+        init_app_context()
+
+
 def require_engine() -> AsyncEngine:
     """Return the initialized database engine."""
-    if engine is None:
-        raise RuntimeError("API app context has not been initialized")
+    _ensure_initialized()
+    assert engine is not None
     return engine
 
 
 def require_session_maker() -> async_sessionmaker[AsyncSession]:
     """Return the initialized async session maker."""
-    if async_session_maker is None:
-        raise RuntimeError("API app context has not been initialized")
+    _ensure_initialized()
+    assert async_session_maker is not None
     return async_session_maker
 
 
 def require_trace_intelligence() -> TraceIntelligence:
     """Return the initialized trace intelligence service."""
-    if trace_intelligence is None:
-        raise RuntimeError("API app context has not been initialized")
+    _ensure_initialized()
+    assert trace_intelligence is not None
     return trace_intelligence
 
 
 def _get_redaction_pipeline() -> RedactionPipeline:
     """Return the configured redaction pipeline."""
-    if _redaction_pipeline is None:
-        raise RuntimeError("API app context has not been initialized")
+    _ensure_initialized()
+    assert _redaction_pipeline is not None
     return _redaction_pipeline

scripts/seed_demo_sessions.py

@@ -27,6 +27,18 @@
 # Session enrichment data: realistic values for demo sessions
 # Note: failure_count is computed in API layer (services.py) as errors count
 # behavior_alert_count is computed in API layer from AnomalyAlertModel records
+def validate_session_enrichment(session_id: str, enrichment: dict[str, object]) -> None:
+    """Validate curated enrichment metrics for demo seed sessions."""
+    total_tokens = enrichment.get("total_tokens")
+    total_cost_usd = enrichment.get("total_cost_usd")
+
+    if not isinstance(total_tokens, int) or total_tokens <= 0:
+        raise ValueError(f"Seed enrichment for {session_id} must define positive total_tokens")
+
+    if not isinstance(total_cost_usd, (int, float)) or float(total_cost_usd) <= 0:
+        raise ValueError(f"Seed enrichment for {session_id} must define positive total_cost_usd")
+
+
 SESSION_ENRICHMENT = {
     "seed-prompt-injection": {
         "total_tokens": 856,
@@ -95,19 +107,41 @@
 }
 
 
+def validate_session_metrics(total_tokens: int, total_cost_usd: float, *, context: str) -> None:
+    """Validate curated session metrics before persisting demo seed data."""
+    if total_tokens < 0:
+        raise ValueError(f"{context}: total_tokens must be non-negative, got {total_tokens}")
+    if total_cost_usd < 0:
+        raise ValueError(f"{context}: total_cost_usd must be non-negative, got {total_cost_usd}")
+
+    has_tokens = total_tokens > 0
+    has_cost = total_cost_usd > 0
+    if has_tokens != has_cost:
+        raise ValueError(
+            f"{context}: total_tokens and total_cost_usd must either both be zero or both be positive "
+            f"(got total_tokens={total_tokens}, total_cost_usd={total_cost_usd})"
+        )
+
+
 async def enrich_session(session_id: str, session_maker: async_sessionmaker[AsyncSession]) -> None:
     """Enrich a session with realistic data fields and behavior alerts."""
     enrichment = SESSION_ENRICHMENT.get(session_id, {})
     if not enrichment:
         return
 
+    validate_session_enrichment(session_id, enrichment)
+
+    total_tokens = enrichment.get("total_tokens", 0)
+    total_cost_usd = enrichment.get("total_cost_usd", 0.0)
+    validate_session_metrics(total_tokens, total_cost_usd, context=f"seed enrichment for {session_id}")
+
     async with session_maker() as db_session:
         repo = TraceRepository(db_session)
 
         # Update session fields
         update_data = {
-            "total_tokens": enrichment.get("total_tokens", 0),
-            "total_cost_usd": enrichment.get("total_cost_usd", 0.0),
+            "total_tokens": total_tokens,
+            "total_cost_usd": total_cost_usd,
             "errors": enrichment.get("errors", 0),
         }
 

tests/conftest.py

@@ -56,13 +56,17 @@ def setup_test_db():
     from storage.engine import create_db_engine
 
     async def _setup():
-        # Reset app context globals so each xdist worker uses its own engine
-        app_context.engine = None
-        app_context.async_session_maker = None
+        # Build schema with a short-lived engine bound to this temporary loop.
+        # The app context is initialized lazily inside each test loop.
         engine = create_db_engine()
         async with engine.begin() as conn:
             await conn.run_sync(Base.metadata.create_all)
-        app_context.init_app_context()
+        await engine.dispose()
+
+        app_context.engine = None
+        app_context.async_session_maker = None
+        app_context.trace_intelligence = None
+        app_context._redaction_pipeline = None
 
     asyncio.run(_setup())
     yield