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

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

engine/tests/conftest.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import pytest
+
+from stirling.config import AppSettings, load_settings
+from stirling.services import build_runtime
+from stirling.services.runtime import AppRuntime
+
+
+@pytest.fixture(autouse=True)
+def clear_settings_cache() -> Iterator[None]:
+    load_settings.cache_clear()
+    yield
+    load_settings.cache_clear()
+
+
+def build_app_settings() -> AppSettings:
+    return AppSettings(
+        smart_model_name="test",
+        fast_model_name="test",
+        smart_model_max_tokens=8192,
+        fast_model_max_tokens=2048,
+        posthog_enabled=False,
+        posthog_api_key="",
+        posthog_host="https://eu.i.posthog.com",
+    )
+
+
+@pytest.fixture
+def app_settings() -> AppSettings:
+    return build_app_settings()
+
+
+@pytest.fixture
+def runtime(app_settings: AppSettings) -> AppRuntime:
+    return build_runtime(app_settings)

engine/tests/test_pdf_edit_agent.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass
 import pytest
 
 from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
-from stirling.config import AppSettings
 from stirling.contracts import (
     EditCannotDoResponse,
     EditClarificationRequest,
@@ -14,16 +13,7 @@ from stirling.contracts import (
     ToolOperationStep,
 )
 from stirling.models.tool_models import CompressParams, OperationId, RotateParams
-from stirling.services import build_runtime
-
-
-def build_test_settings() -> AppSettings:
-    return AppSettings(
-        smart_model_name="test",
-        fast_model_name="test",
-        smart_model_max_tokens=8192,
-        fast_model_max_tokens=2048,
-    )
+from stirling.services.runtime import AppRuntime
 
 
 @dataclass(frozen=True)
@@ -61,10 +51,11 @@ class RecordingParameterSelector:
 class StubPdfEditAgent(PdfEditAgent):
     def __init__(
         self,
+        runtime: AppRuntime,
         selection: PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse,
         parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None,
     ) -> None:
-        super().__init__(build_runtime(build_test_settings()))
+        super().__init__(runtime)
         self.selection = selection
         if parameter_selector is not None:
             self.parameter_selector = parameter_selector
@@ -77,9 +68,10 @@ class StubPdfEditAgent(PdfEditAgent):
 
 
 @pytest.mark.anyio
-async def test_pdf_edit_agent_builds_multi_step_plan() -> None:
+async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> None:
     parameter_selector = RecordingParameterSelector()
     agent = StubPdfEditAgent(
+        runtime,
         PdfEditPlanSelection(
             operations=[OperationId.ROTATE, OperationId.COMPRESS],
             summary="Rotate the PDF, then compress it.",
@@ -104,9 +96,10 @@ async def test_pdf_edit_agent_builds_multi_step_plan() -> None:
 
 
 @pytest.mark.anyio
-async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector() -> None:
+async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtime: AppRuntime) -> None:
     parameter_selector = RecordingParameterSelector()
     agent = StubPdfEditAgent(
+        runtime,
         PdfEditPlanSelection(
             operations=[OperationId.ROTATE, OperationId.COMPRESS],
             summary="Rotate the PDF, then compress it.",
@@ -134,12 +127,13 @@ async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector() -> N
 
 
 @pytest.mark.anyio
-async def test_pdf_edit_agent_returns_clarification_without_partial_plan() -> None:
+async def test_pdf_edit_agent_returns_clarification_without_partial_plan(runtime: AppRuntime) -> None:
     agent = StubPdfEditAgent(
+        runtime,
         EditClarificationRequest(
             question="Which pages should be rotated?",
             reason="The request does not say which pages to change.",
-        )
+        ),
     )
 
     response = await agent.handle(PdfEditRequest(user_message="Rotate some pages."))
@@ -148,11 +142,12 @@ async def test_pdf_edit_agent_returns_clarification_without_partial_plan() -> No
 
 
 @pytest.mark.anyio
-async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan() -> None:
+async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan(runtime: AppRuntime) -> None:
     agent = StubPdfEditAgent(
+        runtime,
         EditCannotDoResponse(
             reason="This request requires OCR, which is not part of PDF edit planning.",
-        )
+        ),
     )
 
     response = await agent.handle(PdfEditRequest(user_message="Read this scan and summarize it."))

engine/tests/test_pdf_question_agent.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import pytest
 
 from stirling.agents import PdfQuestionAgent
-from stirling.config import AppSettings
 from stirling.contracts import (
     ExtractedFileText,
     PdfQuestionAnswerResponse,
@@ -12,12 +11,12 @@ from stirling.contracts import (
     PdfQuestionRequest,
     PdfTextSelection,
 )
-from stirling.services import build_runtime
+from stirling.services.runtime import AppRuntime
 
 
 class StubPdfQuestionAgent(PdfQuestionAgent):
-    def __init__(self, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
-        super().__init__(build_runtime(build_test_settings()))
+    def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
+        super().__init__(runtime)
         self.response = response
 
     async def _run_answer_agent(
@@ -27,15 +26,6 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
         return self.response
 
 
-def build_test_settings() -> AppSettings:
-    return AppSettings(
-        smart_model_name="test",
-        fast_model_name="test",
-        smart_model_max_tokens=8192,
-        fast_model_max_tokens=2048,
-    )
-
-
 def invoice_page() -> ExtractedFileText:
     return ExtractedFileText(
         file_name="invoice.pdf",
@@ -44,8 +34,8 @@ def invoice_page() -> ExtractedFileText:
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_requires_extracted_text() -> None:
-    agent = PdfQuestionAgent(build_runtime(build_test_settings()))
+async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None:
+    agent = PdfQuestionAgent(runtime)
 
     response = await agent.handle(
         PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
@@ -55,12 +45,13 @@ async def test_pdf_question_agent_requires_extracted_text() -> None:
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_returns_grounded_answer() -> None:
+async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None:
     agent = StubPdfQuestionAgent(
+        runtime,
         PdfQuestionAnswerResponse(
             answer="The invoice total is 120.00.",
             evidence=[invoice_page()],
-        )
+        ),
     )
 
     response = await agent.handle(
@@ -76,8 +67,8 @@ async def test_pdf_question_agent_returns_grounded_answer() -> None:
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient() -> None:
-    agent = StubPdfQuestionAgent(PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
+async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None:
+    agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
 
     response = await agent.handle(
         PdfQuestionRequest(

engine/tests/test_stirling_api.py

@@ -1,3 +1,4 @@
+from conftest import build_app_settings
 from fastapi.testclient import TestClient
 
 from stirling.api import app
@@ -8,7 +9,7 @@ from stirling.api.dependencies import (
     get_pdf_question_agent,
     get_user_spec_agent,
 )
-from stirling.config import AppSettings, load_settings
+from stirling.config import load_settings
 from stirling.contracts import (
     AgentDraft,
     AgentDraftRequest,
@@ -27,16 +28,6 @@ from stirling.contracts import (
 from stirling.models.tool_models import RotateParams
 
 
-class StubSettingsProvider:
-    def __call__(self) -> AppSettings:
-        return AppSettings(
-            smart_model_name="test",
-            fast_model_name="test",
-            smart_model_max_tokens=8192,
-            fast_model_max_tokens=2048,
-        )
-
-
 class StubOrchestratorAgent:
     async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse:
         return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000)
@@ -72,41 +63,16 @@ class StubExecutionPlanningAgent:
         return CannotContinueExecutionAction(reason=str(request.current_step_index))
 
 
+app.dependency_overrides[load_settings] = build_app_settings
+app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
+app.dependency_overrides[get_pdf_edit_agent] = lambda: StubPdfEditAgent()
+app.dependency_overrides[get_pdf_question_agent] = lambda: StubPdfQuestionAgent()
+app.dependency_overrides[get_user_spec_agent] = lambda: StubUserSpecAgent()
+app.dependency_overrides[get_execution_planning_agent] = lambda: StubExecutionPlanningAgent()
+
 client: TestClient = TestClient(app)
 
 
-def override_settings() -> AppSettings:
-    return StubSettingsProvider()()
-
-
-def override_orchestrator_agent() -> StubOrchestratorAgent:
-    return StubOrchestratorAgent()
-
-
-def override_pdf_edit_agent() -> StubPdfEditAgent:
-    return StubPdfEditAgent()
-
-
-def override_pdf_question_agent() -> StubPdfQuestionAgent:
-    return StubPdfQuestionAgent()
-
-
-def override_user_spec_agent() -> StubUserSpecAgent:
-    return StubUserSpecAgent()
-
-
-def override_execution_agent() -> StubExecutionPlanningAgent:
-    return StubExecutionPlanningAgent()
-
-
-app.dependency_overrides[load_settings] = override_settings
-app.dependency_overrides[get_orchestrator_agent] = override_orchestrator_agent
-app.dependency_overrides[get_pdf_edit_agent] = override_pdf_edit_agent
-app.dependency_overrides[get_pdf_question_agent] = override_pdf_question_agent
-app.dependency_overrides[get_user_spec_agent] = override_user_spec_agent
-app.dependency_overrides[get_execution_planning_agent] = override_execution_agent
-
-
 def test_health_route() -> None:
     response = client.get("/health")
 

engine/tests/test_stirling_contracts.py

@@ -1,8 +1,4 @@
-from collections.abc import Iterator
-
-import pytest
-
-from stirling.config import AppSettings, load_settings
+from stirling.config import AppSettings
 from stirling.contracts import (
     AgentExecutionRequest,
     AgentSpec,
@@ -76,19 +72,15 @@ def test_pdf_question_answer_defaults_evidence_list() -> None:
     assert response.evidence == []
 
 
-@pytest.fixture(autouse=True)
-def clear_settings_cache() -> Iterator[None]:
-    load_settings.cache_clear()
-    yield
-    load_settings.cache_clear()
-
-
 def test_app_settings_accepts_model_configuration() -> None:
     settings = AppSettings(
         smart_model_name="claude-sonnet-4-5-20250929",
         fast_model_name="claude-haiku-4-5-20251001",
         smart_model_max_tokens=8192,
         fast_model_max_tokens=2048,
+        posthog_enabled=False,
+        posthog_api_key="",
+        posthog_host="https://eu.i.posthog.com",
     )
 
     assert settings.smart_model_name

engine/tests/test_user_spec_agent.py

@@ -4,7 +4,6 @@ import pytest
 from pydantic import ValidationError
 
 from stirling.agents import UserSpecAgent
-from stirling.config import AppSettings
 from stirling.contracts import (
     AgentDraft,
     AgentDraftRequest,
@@ -15,21 +14,12 @@ from stirling.contracts import (
     ToolOperationStep,
 )
 from stirling.models.tool_models import CompressParams, OperationId, RotateParams
-from stirling.services import build_runtime
-
-
-def build_test_settings() -> AppSettings:
-    return AppSettings(
-        smart_model_name="test",
-        fast_model_name="test",
-        smart_model_max_tokens=8192,
-        fast_model_max_tokens=2048,
-    )
+from stirling.services.runtime import AppRuntime
 
 
 class StubUserSpecAgent(UserSpecAgent):
-    def __init__(self, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
-        super().__init__(build_runtime(build_test_settings()))
+    def __init__(self, runtime: AppRuntime, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
+        super().__init__(runtime)
         self.draft_result = draft_result
         self.revision_result = revision_result
         self.edit_plan = EditPlanResponse(
@@ -53,8 +43,8 @@ class StubUserSpecAgent(UserSpecAgent):
 
 
 class ClarifyingUserSpecAgent(UserSpecAgent):
-    def __init__(self) -> None:
-        super().__init__(build_runtime(build_test_settings()))
+    def __init__(self, runtime: AppRuntime) -> None:
+        super().__init__(runtime)
 
     async def _build_edit_plan(self, user_message: str) -> EditClarificationRequest:
         return EditClarificationRequest(
@@ -64,8 +54,9 @@ class ClarifyingUserSpecAgent(UserSpecAgent):
 
 
 @pytest.mark.anyio
-async def test_user_spec_agent_drafts_agent_spec() -> None:
+async def test_user_spec_agent_drafts_agent_spec(runtime: AppRuntime) -> None:
     agent = StubUserSpecAgent(
+        runtime,
         AgentDraft(
             name="Invoice Cleanup",
             description="Prepare invoices for review.",
@@ -100,7 +91,7 @@ async def test_user_spec_agent_drafts_agent_spec() -> None:
 
 
 @pytest.mark.anyio
-async def test_user_spec_agent_revises_existing_draft() -> None:
+async def test_user_spec_agent_revises_existing_draft(runtime: AppRuntime) -> None:
     current_draft = AgentDraft(
         name="Invoice Cleanup",
         description="Prepare invoices for review.",
@@ -113,6 +104,7 @@ async def test_user_spec_agent_revises_existing_draft() -> None:
         ],
     )
     agent = StubUserSpecAgent(
+        runtime,
         draft_result=current_draft,
         revision_result=AgentDraft(
             name="Invoice Cleanup",
@@ -152,8 +144,8 @@ def test_tool_operation_step_rejects_mismatched_parameters() -> None:
 
 
 @pytest.mark.anyio
-async def test_user_spec_agent_propagates_edit_clarification() -> None:
-    agent = ClarifyingUserSpecAgent()
+async def test_user_spec_agent_propagates_edit_clarification(runtime: AppRuntime) -> None:
+    agent = ClarifyingUserSpecAgent(runtime)
 
     response = await agent.draft(AgentDraftRequest(user_message="Build an agent to rotate some pages."))