Add tracking system to support optional PostHog tracking in AI engine (#6040)

Co-authored-by: ConnorYoh <[email protected]>

engine/src/stirling/api/app.py

@@ -4,8 +4,11 @@ from contextlib import asynccontextmanager
 from typing import Annotated
 
 from fastapi import Depends, FastAPI
+from pydantic_ai import Agent
+from pydantic_ai.models.instrumented import InstrumentationSettings
 
 from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
+from stirling.api.middleware import UserIdMiddleware
 from stirling.api.routes import (
     agent_draft_router,
     execution_router,
@@ -15,7 +18,7 @@ from stirling.api.routes import (
 )
 from stirling.config import AppSettings, load_settings
 from stirling.contracts import HealthResponse
-from stirling.services import build_runtime
+from stirling.services import build_runtime, setup_posthog_tracking
 
 
 def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
@@ -37,10 +40,16 @@ async def lifespan(fast_api: FastAPI):
     fast_api.state.pdf_question_agent = PdfQuestionAgent(runtime)
     fast_api.state.user_spec_agent = UserSpecAgent(runtime)
     fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime)
+    tracer_provider = setup_posthog_tracking(settings)
+    if tracer_provider:
+        Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
     yield
+    if tracer_provider:
+        tracer_provider.shutdown()
 
 
 app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
+app.add_middleware(UserIdMiddleware)
 app.include_router(orchestrator_router)
 app.include_router(pdf_edit_router)
 app.include_router(pdf_question_router)

engine/src/stirling/api/middleware.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.responses import Response
+
+from stirling.services.tracking import current_user_id
+
+_USER_ID_HEADER = "X-User-Id"
+
+
+class UserIdMiddleware(BaseHTTPMiddleware):
+    """Extract X-User-Id header and set it as the current user for PostHog tracking."""
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        user_id = request.headers.get(_USER_ID_HEADER)
+        if user_id:
+            token = current_user_id.set(user_id)
+            try:
+                return await call_next(request)
+            finally:
+                current_user_id.reset(token)
+        return await call_next(request)