feat: sync groups in posthog

api/services/auth/depends.py

@@ -13,9 +13,16 @@ from api.schemas.ai_model_configuration import EffectiveAIModelConfiguration
 from api.services.auth.stack_auth import stackauth
 from api.services.configuration.registry import ServiceProviders
 from api.services.mps_billing import ensure_hosted_mps_billing_account_v2
-from api.services.posthog_client import capture_event
+from api.services.posthog_client import (
+    capture_event,
+    group_identify,
+    set_person_properties,
+)
 from api.utils.auth import decode_jwt_token
 
+POSTHOG_ORGANIZATION_GROUP_TYPE = "organization"
+POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY = "uses_mps_billing_v2"
+
 
 async def get_user(
     authorization: Annotated[str | None, Header()] = None,
@@ -94,6 +101,11 @@ async def get_user(
         ) = await db_client.get_or_create_organization_by_provider_id(
             org_provider_id=selected_team_id, user_id=user_model.id
         )
+        if org_was_created:
+            _sync_created_organization_to_posthog(
+                organization=organization,
+                stack_user=stack_user,
+            )
 
         # Check if user's selected organization differs from the current organization
         if user_model.selected_organization_id != organization.id:
@@ -107,6 +119,13 @@ async def get_user(
             # Update the user_model object to reflect the change
             user_model.selected_organization_id = organization.id
 
+            _associate_user_with_posthog_organization(
+                user=user_model,
+                organization=organization,
+                stack_user=stack_user,
+                org_was_created=org_was_created,
+            )
+
             # Only create default configuration if organization was just created
             # This prevents race conditions where multiple concurrent requests
             # might try to create configurations
@@ -156,6 +175,146 @@ async def get_user(
     return user_model
 
 
+def _sync_created_organization_to_posthog(
+    *,
+    organization,
+    stack_user: dict | None = None,
+    created_by_provider_id: str | None = None,
+    uses_mps_billing_v2: bool | None = None,
+) -> None:
+    """Create/update the PostHog organization group for a newly-created org."""
+    try:
+        organization_id = int(organization.id)
+        organization_provider_id = getattr(organization, "provider_id", None)
+        created_by = created_by_provider_id
+        if created_by is None and stack_user and stack_user.get("id"):
+            created_by = str(stack_user["id"])
+        properties = {
+            "organization_id": organization_id,
+            "organization_provider_id": organization_provider_id,
+            "auth_provider": "stack",
+        }
+        if created_by:
+            properties["created_by_provider_id"] = created_by
+        if uses_mps_billing_v2 is not None:
+            properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
+                uses_mps_billing_v2
+            )
+
+        group_identify(
+            POSTHOG_ORGANIZATION_GROUP_TYPE,
+            str(organization_id),
+            properties,
+            distinct_id=created_by,
+        )
+        if created_by:
+            capture_event(
+                distinct_id=created_by,
+                event=PostHogEvent.ORGANIZATION_CREATED,
+                properties=properties,
+                groups={POSTHOG_ORGANIZATION_GROUP_TYPE: str(organization_id)},
+            )
+    except Exception:
+        logger.exception("Failed to sync created organization to PostHog")
+
+
+def _sync_posthog_organization_group_properties(
+    *,
+    organization,
+    uses_mps_billing_v2: bool | None = None,
+) -> None:
+    """Update PostHog organization group properties without creating a person."""
+    try:
+        organization_id = int(organization.id)
+        properties = {
+            "organization_id": organization_id,
+            "organization_provider_id": getattr(organization, "provider_id", None),
+            "auth_provider": "stack",
+        }
+        if uses_mps_billing_v2 is not None:
+            properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
+                uses_mps_billing_v2
+            )
+
+        group_identify(
+            POSTHOG_ORGANIZATION_GROUP_TYPE,
+            str(organization_id),
+            properties,
+        )
+    except Exception:
+        logger.exception("Failed to sync organization group properties to PostHog")
+
+
+def _sync_posthog_organization_mps_billing_v2_status(
+    organization_id: int,
+    *,
+    uses_mps_billing_v2: bool,
+) -> None:
+    """Update the PostHog organization group with current MPS billing status."""
+    try:
+        organization_id = int(organization_id)
+        group_identify(
+            POSTHOG_ORGANIZATION_GROUP_TYPE,
+            str(organization_id),
+            {POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY: uses_mps_billing_v2},
+        )
+    except Exception:
+        logger.exception("Failed to sync organization billing status to PostHog")
+
+
+def _associate_user_with_posthog_organization(
+    *,
+    user: UserModel,
+    organization,
+    stack_user: dict | None = None,
+    user_distinct_id: str | None = None,
+    org_was_created: bool,
+    organization_ids: list[int] | None = None,
+    selected_organization_id: int | None = None,
+    selected_organization_provider_id: str | None = None,
+) -> None:
+    """Attach the Stack user to the PostHog organization group."""
+    try:
+        organization_id = int(organization.id)
+        organization_provider_id = getattr(organization, "provider_id", None)
+        if user_distinct_id is None:
+            if stack_user and stack_user.get("id"):
+                user_distinct_id = str(stack_user["id"])
+            else:
+                user_distinct_id = str(user.provider_id)
+        selected_org_id = selected_organization_id or organization_id
+        selected_org_provider_id = (
+            selected_organization_provider_id or organization_provider_id
+        )
+        person_properties = {
+            "user_id": user.id,
+            "user_provider_id": user_distinct_id,
+            "selected_organization_id": selected_org_id,
+            "selected_organization_provider_id": selected_org_provider_id,
+        }
+        if organization_ids is not None:
+            person_properties["organization_ids"] = organization_ids
+        if user.email:
+            person_properties["email"] = user.email
+        set_person_properties(user_distinct_id, person_properties)
+        event_properties = {
+            "user_id": user.id,
+            "organization_id": organization_id,
+            "organization_provider_id": organization_provider_id,
+            "auth_provider": "stack",
+            "organization_was_created": org_was_created,
+        }
+
+        capture_event(
+            distinct_id=user_distinct_id,
+            event=PostHogEvent.ORGANIZATION_USER_ASSOCIATED,
+            properties=event_properties,
+            groups={POSTHOG_ORGANIZATION_GROUP_TYPE: str(organization_id)},
+        )
+    except Exception:
+        logger.exception("Failed to associate user with PostHog organization")
+
+
 async def get_user_with_selected_organization(
     user: Annotated[UserModel, Depends(get_user)],
 ) -> UserModel:

api/services/posthog_client.py

@@ -1,7 +1,9 @@
+from typing import Any, Optional
+
 from loguru import logger
 from posthog import Posthog
 
-from api.constants import ENABLE_TELEMETRY, POSTHOG_API_KEY, POSTHOG_HOST
+from api.constants import POSTHOG_API_KEY, POSTHOG_HOST
 
 _posthog_client: Posthog | None = None
 
@@ -9,23 +11,84 @@ _posthog_client: Posthog | None = None
 def get_posthog() -> Posthog | None:
     """Return the lazily-initialised PostHog client, or None if not configured."""
     global _posthog_client
-    if _posthog_client is None and POSTHOG_API_KEY and ENABLE_TELEMETRY:
+    if _posthog_client is None and POSTHOG_API_KEY:
         _posthog_client = Posthog(POSTHOG_API_KEY, host=POSTHOG_HOST)
     return _posthog_client
 
 
+def shutdown_posthog() -> None:
+    """Flush queued PostHog messages before a short-lived process exits."""
+    client = get_posthog()
+    if not client:
+        return
+    try:
+        client.shutdown()
+    except Exception:
+        logger.exception("Failed to shut down PostHog client")
+
+
+def flush_posthog() -> None:
+    """Flush queued PostHog messages without shutting down the client."""
+    client = get_posthog()
+    if not client:
+        return
+    try:
+        client.flush()
+    except Exception:
+        logger.exception("Failed to flush PostHog client")
+
+
 def capture_event(
     distinct_id: str,
     event: str,
-    properties: dict | None = None,
+    properties: dict[str, Any] | None = None,
+    groups: Optional[dict[str, str]] = None,
 ) -> None:
     """Fire a PostHog event. Silently no-ops if PostHog is not configured."""
     client = get_posthog()
     if not client:
         return
     try:
-        client.capture(
-            distinct_id=distinct_id, event=event, properties=properties or {}
-        )
+        kwargs: dict[str, Any] = {
+            "distinct_id": distinct_id,
+            "event": event,
+            "properties": properties or {},
+        }
+        if groups:
+            kwargs["groups"] = groups
+        client.capture(**kwargs)
     except Exception:
         logger.exception(f"Failed to send PostHog event '{event}'")
+
+
+def group_identify(
+    group_type: str,
+    group_key: str,
+    properties: dict[str, Any],
+    *,
+    distinct_id: Optional[str] = None,
+) -> None:
+    """Set PostHog group properties. Silently no-ops if PostHog is not configured."""
+    client = get_posthog()
+    if not client:
+        return
+    try:
+        client.group_identify(
+            group_type,
+            group_key,
+            properties,
+            distinct_id=distinct_id,
+        )
+    except Exception:
+        logger.exception("Failed to identify PostHog group")
+
+
+def set_person_properties(distinct_id: str, properties: dict[str, Any]) -> None:
+    """Set PostHog person properties. Silently no-ops if PostHog is not configured."""
+    client = get_posthog()
+    if not client:
+        return
+    try:
+        client.set(distinct_id=distinct_id, properties=properties)
+    except Exception:
+        logger.exception("Failed to set PostHog person properties")