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: