feat: Update support for segment analytics

src/__tests__/segment.test.ts

@@ -9,20 +9,16 @@
 
 import { beforeEach, describe, expect, it, jest } from '@jest/globals'
 
-import posthog from '../loader-module'
 import { PostHog } from '../posthog-core'
-import { uuidv7 } from '../uuidv7'
 
-describe(`Module-based loader in Node env`, () => {
+describe(`Segment integration`, () => {
     let segment: any
     let segmentIntegration: any
     let posthogName: string
 
-    beforeEach(() => {
-        jest.spyOn(posthog, '_send_request').mockReturnValue()
-        jest.spyOn(console, 'log').mockReturnValue()
-        posthogName = uuidv7()
+    jest.setTimeout(500)
 
+    beforeEach(() => {
         // Create something that looks like the Segment Analytics 2.0 API. We
         // could use the actual client, but it's a little more tricky and we'd
         // want to mock out the network requests, for which we don't have a good
@@ -39,6 +35,7 @@ describe(`Module-based loader in Node env`, () => {
                 // To ensure the Promise isn't resolved instantly, we use a
                 // setTimeout with a delay of 0 to ensure it happens as a
                 // microtask in the future.
+
                 return new Promise((resolve) => {
                     setTimeout(() => {
                         segmentIntegration = integration
@@ -49,36 +46,66 @@ describe(`Module-based loader in Node env`, () => {
         }
     })
 
-    it('should call loaded after the segment integration has been set up', (done) => {
-        // This test is to ensure that, by the time the `loaded` callback is
-        // called, the PostHog Segment integration have completed registration.
-        // If this is not the case, then we end up in odd situations where we
-        // try to send segment events but we do not get the enriched event
-        // information as the integration provides.
-        const posthog = new PostHog()
-        jest.spyOn(posthog, 'capture')
+    it('should call loaded after the segment integration has been set up', async () => {
+        const loadPromise = new Promise((resolve) => {
+            return new PostHog().init(
+                `test-token`,
+                {
+                    debug: true,
+                    persistence: `localStorage`,
+                    api_host: `https://test.com`,
+                    segment: segment,
+                    loaded: resolve,
+                },
+                posthogName
+            )
+        })
+        expect(segmentIntegration).toBeUndefined()
+        await loadPromise
+        expect(segmentIntegration).toBeDefined()
+    })
 
-        posthog.init(
-            `test-token`,
-            {
-                debug: true,
-                persistence: `localStorage`,
-                api_host: `https://test.com`,
-                segment: segment,
-                loaded: () => {
-                    expect(segmentIntegration).toBeDefined()
-                    done()
+    it('should set properties from the segment user', async () => {
+        const posthog = await new Promise<PostHog>((resolve) => {
+            return new PostHog().init(
+                `test-token`,
+                {
+                    debug: true,
+                    persistence: `localStorage`,
+                    api_host: `https://test.com`,
+                    segment: segment,
+                    loaded: resolve,
                 },
-            },
-            posthogName
-        )
+                posthogName
+            )
+        })
 
-        // Assuming we've set up our mocks correctly, the segmentIntegration
-        // shouldn't have been set by now, but just to be sure we're actually
-        // checking that loaded callback handles async code, we explicitly check
-        // this first.
-        expect(segmentIntegration).toBeUndefined()
+        expect(posthog.get_distinct_id()).toBe('test-id')
+        expect(posthog.get_property('$device_id')).toBe('test-anonymous-id')
     })
 
-    // TODO: add tests for distinct id setting and event enrichment.
+    it('should handle the segment user being a promise', async () => {
+        segment.user = () =>
+            Promise.resolve({
+                anonymousId: () => 'test-anonymous-id',
+                id: () => 'test-id',
+            })
+
+        const posthog = await new Promise<PostHog>((resolve) => {
+            return new PostHog().init(
+                `test-token`,
+                {
+                    debug: true,
+                    persistence: `localStorage`,
+                    api_host: `https://test.com`,
+                    segment: segment,
+                    loaded: resolve,
+                },
+                posthogName
+            )
+        })
+
+        expect(posthog.get_distinct_id()).toBe('test-id')
+        expect(posthog.get_property('$device_id')).toBe('test-anonymous-id')
+    })
 })

src/extensions/segment-integration.ts

@@ -19,8 +19,21 @@
 import { PostHog } from '../posthog-core'
 import { logger } from '../utils/logger'
 
+import { uuidv7 } from '../uuidv7'
+import { _isFunction } from '../utils/type-utils'
+
+export type SegmentUser = {
+    anonymousId(): string | undefined
+    id(): string | undefined
+}
+
+export type SegmentAnalytics = {
+    user: () => SegmentUser | Promise<SegmentUser>
+    register: (integration: SegmentPlugin) => Promise<void>
+}
+
 // Loosely based on https://github.com/segmentio/analytics-next/blob/master/packages/core/src/plugins/index.ts
-interface SegmentPluginContext {
+interface SegmentContext {
     event: {
         event: string
         userId?: string
@@ -29,28 +42,33 @@ interface SegmentPluginContext {
     }
 }
 
+type SegmentFunction = (ctx: SegmentContext) => Promise<SegmentContext> | SegmentContext
+
 interface SegmentPlugin {
     name: string
     version: string
     type: 'enrichment'
     isLoaded: () => boolean
-    load: (ctx: SegmentPluginContext, instance: any, config?: any) => Promise<unknown>
-    unload?: (ctx: SegmentPluginContext, instance: any) => Promise<unknown> | unknown
+    load: (ctx: SegmentContext, instance: any, config?: any) => Promise<unknown>
+    unload?: (ctx: SegmentContext, instance: any) => Promise<unknown> | unknown
     ready?: () => Promise<unknown>
-    track?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
-    identify?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
-    page?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
-    group?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
-    alias?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
-    screen?: (ctx: SegmentPluginContext) => Promise<SegmentPluginContext> | SegmentPluginContext
+    track?: SegmentFunction
+    identify?: SegmentFunction
+    page?: SegmentFunction
+    group?: SegmentFunction
+    alias?: SegmentFunction
+    screen?: SegmentFunction
 }
 
-export const createSegmentIntegration = (posthog: PostHog): SegmentPlugin => {
+const createSegmentIntegration = (posthog: PostHog): SegmentPlugin => {
     if (!Promise || !Promise.resolve) {
         logger.warn('This browser does not have Promise support, and can not use the segment integration')
     }
 
-    const enrichEvent = (ctx: SegmentPluginContext, eventName: string) => {
+    const enrichEvent = (ctx: SegmentContext, eventName: string | undefined) => {
+        if (!eventName) {
+            return ctx
+        }
         if (!ctx.event.userId && ctx.event.anonymousId !== posthog.get_distinct_id()) {
             // This is our only way of detecting that segment's analytics.reset() has been called so we also call it
             posthog.reset()
@@ -62,7 +80,7 @@ export const createSegmentIntegration = (posthog: PostHog): SegmentPlugin => {
             posthog.reloadFeatureFlags()
         }
 
-        const additionalProperties = posthog._calculate_event_properties(eventName, ctx.event.properties)
+        const additionalProperties = posthog._calculate_event_properties(eventName, ctx.event.properties ?? {})
         ctx.event.properties = Object.assign({}, additionalProperties, ctx.event.properties)
         return ctx
     }
@@ -81,3 +99,49 @@ export const createSegmentIntegration = (posthog: PostHog): SegmentPlugin => {
         screen: (ctx) => enrichEvent(ctx, '$screen'),
     }
 }
+
+function setupPostHogFromSegment(posthog: PostHog, done: () => void) {
+    const segment = posthog.config.segment
+    if (!segment) {
+        return done()
+    }
+
+    const bootstrapUser = (user: SegmentUser) => {
+        // Use segments anonymousId instead
+        const getSegmentAnonymousId = () => user.anonymousId() || uuidv7()
+        posthog.config.get_device_id = getSegmentAnonymousId
+
+        // If a segment user ID exists, set it as the distinct_id
+        if (user.id()) {
+            posthog.register({
+                distinct_id: user.id(),
+                $device_id: getSegmentAnonymousId(),
+            })
+            posthog.persistence!.set_user_state('identified')
+        }
+
+        done()
+    }
+
+    const segmentUser = segment.user()
+
+    // If segmentUser is a promise then we need to wait for it to resolve
+    if ('then' in segmentUser && _isFunction(segmentUser.then)) {
+        segmentUser.then((user) => bootstrapUser(user))
+    } else {
+        bootstrapUser(segmentUser as SegmentUser)
+    }
+}
+
+export function setupSegmentIntegration(posthog: PostHog, done: () => void) {
+    const segment = posthog.config.segment
+    if (!segment) {
+        return done()
+    }
+
+    setupPostHogFromSegment(posthog, () => {
+        segment.register(createSegmentIntegration(posthog)).then(() => {
+            done()
+        })
+    })
+}

src/posthog-core.ts

@@ -49,7 +49,7 @@ import {
     ToolbarParams,
 } from './types'
 import { SentryIntegration } from './extensions/sentry-integration'
-import { createSegmentIntegration } from './extensions/segment-integration'
+import { setupSegmentIntegration } from './extensions/segment-integration'
 import { PageViewManager } from './page-view'
 import { PostHogSurveys } from './posthog-surveys'
 import { RateLimiter } from './rate-limiter'
@@ -228,7 +228,6 @@ export class PostHog {
     analyticsDefaultEndpoint: string
 
     SentryIntegration: typeof SentryIntegration
-    segmentIntegration: () => any
 
     private _debugEventEmitter = new SimpleEventEmitter()
 
@@ -242,7 +241,6 @@ export class PostHog {
         this.config = defaultConfig()
         this.decideEndpointWasHit = false
         this.SentryIntegration = SentryIntegration
-        this.segmentIntegration = () => createSegmentIntegration(this)
         this.__request_queue = []
         this.__loaded = false
         this.analyticsDefaultEndpoint = '/e/'
@@ -377,19 +375,6 @@ export class PostHog {
 
         this._gdpr_init()
 
-        if (config.segment) {
-            // Use segments anonymousId instead
-            this.config.get_device_id = () => config.segment.user().anonymousId()
-
-            // If a segment user ID exists, set it as the distinct_id
-            if (config.segment.user().id()) {
-                this.register({
-                    distinct_id: config.segment.user().id(),
-                })
-                this.persistence.set_user_state('identified')
-            }
-        }
-
         // isUndefined doesn't provide typehint here so wouldn't reduce bundle as we'd need to assign
         // eslint-disable-next-line posthog-js/no-direct-undefined-check
         if (config.bootstrap?.distinctID !== undefined) {
@@ -447,9 +432,7 @@ export class PostHog {
 
         // We wan't to avoid promises for IE11 compatibility, so we use callbacks here
         if (config.segment) {
-            config.segment.register(this.segmentIntegration()).then(() => {
-                this._loaded()
-            })
+            setupSegmentIntegration(this, () => this._loaded())
         } else {
             this._loaded()
         }

src/types.ts

@@ -1,5 +1,6 @@
 import type { MaskInputOptions, SlimDOMOptions } from 'rrweb-snapshot'
 import { PostHog } from './posthog-core'
+import type { SegmentAnalytics } from './extensions/segment-integration'
 
 export type Property = any
 export type Properties = Record<string, Property>
@@ -141,7 +142,7 @@ export interface PostHogConfig {
     // Should only be used for testing. Could negatively impact performance.
     disable_compression: boolean
     bootstrap: BootstrapConfig
-    segment?: any
+    segment?: SegmentAnalytics
     __preview_send_client_session_params?: boolean
     disable_scroll_properties?: boolean
     // Let the pageview scroll stats use a custom css selector for the root element, e.g. `main`