insights/products/wizard/frontend/wizardSessionStreamLogic.ts

import { MakeLogicType, actions, connect, kea, key, listeners, path, props, reducers } from 'kea'
import insights from 'insights-js'

import { FEATURE_FLAGS } from 'lib/constants'
import { featureFlagLogic, getFeatureFlagPayload } from 'lib/logic/featureFlagLogic'
import {
    createPollLoop,
    isPermanentPollError,
    resolvePollingIntervalMs,
    resolveWizardSyncMode,
    sseReconnectDelayMs,
    type WizardSyncMode,
} from 'lib/wizard-sync/pollLoop'
import { logSyncDebug } from 'lib/wizard-sync/wizardSyncDebugLogic'
import { projectLogic } from 'scenes/projectLogic'

import type { FeatureFlagsSet } from '../../../frontend/src/lib/logic/featureFlagLogic'
import { getWizardSessionsStreamRetrieveUrl, wizardSessionsLatestRetrieve } from './generated/api'
import type { WizardSessionDTOApi } from './generated/api.schemas'

export type WizardConnectionStatus = 'idle' | 'connecting' | 'open' | 'closed' | 'error'

export interface WizardSessionStreamLogicProps {
    workflowId: string
    /**
     * Optional. When omitted, the backend pattern-subscribes to all skills under
     * the workflow. Useful when the framework / skill the wizard will run isn't
     * known up front (the common onboarding case).
     */
    skillId?: string
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface wizardSessionStreamLogicValues {
    featureFlags: FeatureFlagsSet // featureFlagLogic
    currentProjectId: number | null // projectLogic
    connectionStatus: WizardConnectionStatus
    lastError: string | null
    latestSession: WizardSessionDTOApi | null
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface wizardSessionStreamLogicActions {
    connect: () => {
        value: true
    }
    connectionErrored: (error: string) => {
        error: string
    }
    connectionOpened: () => {
        value: true
    }
    disconnect: () => {
        value: true
    }
    sessionUpdated: (session: WizardSessionDTOApi) => {
        session: WizardSessionDTOApi
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface wizardSessionStreamLogicMeta {
    key: string
}

export type wizardSessionStreamLogicType = MakeLogicType<
    wizardSessionStreamLogicValues,
    wizardSessionStreamLogicActions,
    WizardSessionStreamLogicProps,
    wizardSessionStreamLogicMeta
>

/**
 * Subscribes to the SSE stream for a wizard workflow (optionally a specific skill)
 * and exposes the latest WizardSession state plus connection status.
 *
 * Each (workflowId, skillId) pair gets its own logic instance keyed by props, so
 * multiple subscriptions coexist in the app. Omitting skillId pattern-subscribes
 * to every skill under the workflow.
 *
 * When the `onboarding-wizard-sync-mode` flag resolves to `polling` (GROW-118), the
 * EventSource is replaced by a poll loop that re-fetches the latest session via
 * `wizard/sessions/latest/` and feeds it through the same `sessionUpdated` action.
 * The loop (shared with `taskRunStreamLogic`) backs off on consecutive errors, stops on
 * permanent ones (401/403/404), and backs off on consecutive empty responses — so an
 * endpoint with no session (not started yet, or killswitched to 204) winds down instead
 * of being polled at full cadence forever. Consumers stop it via `disconnect`;
 * `installationProgressLogic` does so once a cloud run reaches a terminal state.
 *
 * Usage:
 *
 *   // Listen to any skill under the insights-integration workflow:
 *   const { latestSession } = useValues(
 *       wizardSessionStreamLogic({ workflowId: 'insights-integration' })
 *   )
 *
 *   // Or pin to one skill:
 *   const { latestSession } = useValues(
 *       wizardSessionStreamLogic({ workflowId: 'insights-integration', skillId: 'nextjs' })
 *   )
 */
export const wizardSessionStreamLogic = kea<wizardSessionStreamLogicType>([
    props({} as WizardSessionStreamLogicProps),
    key((props) => `${props.workflowId}::${props.skillId ?? '*'}`),
    path((key) => ['products', 'wizard', 'wizardSessionStreamLogic', key]),
    connect(() => ({
        values: [projectLogic, ['currentProjectId'], featureFlagLogic, ['featureFlags']],
    })),
    actions({
        connect: true,
        disconnect: true,
        sessionUpdated: (session: WizardSessionDTOApi) => ({ session }),
        connectionOpened: true,
        connectionErrored: (error: string) => ({ error }),
    }),
    reducers({
        latestSession: [
            null as WizardSessionDTOApi | null,
            {
                sessionUpdated: (_, { session }) => session,
            },
        ],
        connectionStatus: [
            'idle' as WizardConnectionStatus,
            {
                connect: () => 'connecting',
                connectionOpened: () => 'open',
                connectionErrored: () => 'error',
                disconnect: () => 'closed',
            },
        ],
        lastError: [
            null as string | null,
            {
                connect: () => null,
                sessionUpdated: () => null,
                connectionErrored: (_, { error }) => error,
            },
        ],
    }),
    listeners(({ values, actions, props, cache }) => ({
        connect: () => {
            // A connect (manual or scheduled) supersedes any pending reconnect timer, so the two can
            // never race a second transport open.
            cache.disposables.dispose('session-reconnect')
            // A scheduled retry continues the connect cycle that failed rather than starting a new
            // one, so it must leave the telemetry guards below alone. Resetting them here would emit
            // one transport-error capture per backoff step for as long as the stream stays closed,
            // which is once every 30s per tab and enough to swamp the SSE-vs-polling comparison.
            const resumingCycle = cache.reconnectScheduled === true
            cache.reconnectScheduled = false
            const projectId = values.currentProjectId
            if (projectId === null) {
                actions.connectionErrored('No current project — cannot open wizard session stream.')
                return
            }

            // Per-connect-cycle transport telemetry: one "ready" (first session delivered) and at
            // most one "error" per cycle, so SSE and polling performance stay comparable without
            // an event per delivery or per EventSource retry.
            if (!resumingCycle) {
                cache.connectStartedAt = Date.now()
                cache.transportReadyReported = false
                cache.transportErrorReported = false
            }

            const debugSource = `session ${props.workflowId}::${props.skillId ?? '*'}`

            // Mode is re-sampled on every connect, and the setFeatureFlags listener below reconnects
            // when the resolved mode changes — so a flag flip (or flags arriving after a cold-cache
            // connect) swaps the transport live instead of waiting for a remount.
            const syncMode: WizardSyncMode = resolveWizardSyncMode(
                values.featureFlags[FEATURE_FLAGS.ONBOARDING_WIZARD_SYNC_MODE]
            )
            cache.syncMode = syncMode

            if (syncMode === 'polling') {
                const intervalMs = resolvePollingIntervalMs(
                    getFeatureFlagPayload(FEATURE_FLAGS.ONBOARDING_WIZARD_SYNC_MODE)
                )
                logSyncDebug(debugSource, 'connect', `polling every ~${intervalMs / 1000}s (±20% jitter)`, {
                    mode: 'polling',
                    intervalMs,
                })
                cache.disposables.add(
                    createPollLoop({
                        intervalMs,
                        tick: async () => {
                            // 204 (no session yet, or killswitched) resolves to null — classify as
                            // empty so the loop backs off instead of polling at full cadence forever.
                            const session = await wizardSessionsLatestRetrieve(
                                String(projectId),
                                {
                                    workflow_id: props.workflowId,
                                    skill_id: props.skillId,
                                },
                                { headers: { 'X-Wizard-Poll-Source': 'transport' } }
                            )
                            logSyncDebug(
                                debugSource,
                                'poll',
                                session
                                    ? `poll ok: ${session.run_phase} (${session.skill_id})`
                                    : 'poll ok: no session yet',
                                { mode: 'polling', intervalMs }
                            )
                            // Only on first open / recovery from error — not every tick.
                            if (values.connectionStatus !== 'open') {
                                actions.connectionOpened()
                            }
                            if (!session) {
                                return 'empty'
                            }
                            actions.sessionUpdated(session)
                            return 'ok'
                        },
                        onError: (err) => {
                            logSyncDebug(debugSource, 'error', `poll failed: ${String(err)}`)
                            actions.connectionErrored(`Failed to poll wizard session: ${String(err)}`)
                            return isPermanentPollError(err) ? 'stop' : 'retry'
                        },
                        // Dispose the key so a later tab-visibility resume doesn't restart a loop
                        // that ended itself.
                        onLoopEnd: () => cache.disposables.dispose('session-sync'),
                    }),
                    'session-sync'
                )
                return
            }

            logSyncDebug(debugSource, 'connect', 'opening EventSource', { mode: 'sse' })
            cache.disposables.add((): (() => void) => {
                const url = getWizardSessionsStreamRetrieveUrl(String(projectId), {
                    workflow_id: props.workflowId,
                    skill_id: props.skillId,
                })
                const eventSource = new EventSource(url, { withCredentials: true })

                eventSource.onopen = (): void => {
                    // Mode rides along here too: the connect event can predate the debug panel's
                    // mount (and get dropped), but open/events always land after it.
                    logSyncDebug(debugSource, 'open', 'SSE connection open', { mode: 'sse' })
                    cache.sseReconnectAttempt = 0
                    actions.connectionOpened()
                }
                eventSource.onmessage = (event: MessageEvent<string>): void => {
                    try {
                        const session = JSON.parse(event.data) as WizardSessionDTOApi
                        logSyncDebug(debugSource, 'event', `session → ${session.run_phase} (${session.skill_id})`, {
                            mode: 'sse',
                        })
                        actions.sessionUpdated(session)
                    } catch (err) {
                        logSyncDebug(debugSource, 'error', `failed to parse SSE payload: ${String(err)}`)
                        actions.connectionErrored(`Failed to parse SSE payload: ${String(err)}`)
                    }
                }
                eventSource.onerror = (): void => {
                    // EventSource readyState distinguishes terminal close vs transient retry:
                    //   CLOSED (2)     → browser gave up (any non-200 on a (re)connect lands here)
                    //                    and will NOT auto-reconnect, so schedule our own reconnect
                    //                    with backoff.
                    //   CONNECTING (0) → browser is already retrying; ride it out, surface
                    //                    the state so the UI can show "reconnecting".
                    if (eventSource.readyState === EventSource.CLOSED) {
                        const attempt = ((cache.sseReconnectAttempt as number | undefined) ?? 0) + 1
                        cache.sseReconnectAttempt = attempt
                        const delayMs = sseReconnectDelayMs(attempt)
                        logSyncDebug(
                            debugSource,
                            'error',
                            `SSE closed, reconnecting in ~${Math.round(delayMs / 1000)}s`
                        )
                        actions.connectionErrored(
                            `EventSource connection closed, reconnecting in ~${Math.round(delayMs / 1000)}s`
                        )
                        // The timer rides disposables so an unmount tears it down, and a hidden tab
                        // pauses it instead of reconnecting in the background. connect() re-samples
                        // the sync-mode flag, so a mid-outage flip to polling takes effect here too.
                        cache.disposables.add(() => {
                            const timer = window.setTimeout(() => {
                                cache.reconnectScheduled = true
                                actions.connect()
                            }, delayMs)
                            return () => window.clearTimeout(timer)
                        }, 'session-reconnect')
                    } else {
                        logSyncDebug(debugSource, 'error', 'SSE transport error, reconnecting')
                        actions.connectionErrored('EventSource transport error, reconnecting')
                    }
                }

                return () => eventSource.close()
            }, 'session-sync')
        },
        disconnect: () => {
            logSyncDebug(`session ${props.workflowId}::${props.skillId ?? '*'}`, 'disconnect', 'disconnected')
            cache.disposables.dispose('session-sync')
            cache.disposables.dispose('session-reconnect')
            // The next connect is a new cycle, not a continuation of the one that was torn down.
            cache.reconnectScheduled = false
        },
        sessionUpdated: () => {
            if (cache.transportReadyReported || cache.connectStartedAt === undefined) {
                return
            }
            cache.transportReadyReported = true
            insights.capture('wizard sync transport ready', {
                transport: cache.syncMode ?? 'sse',
                workflow_id: props.workflowId,
                ms_since_connect: Date.now() - cache.connectStartedAt,
            })
        },
        connectionErrored: ({ error }) => {
            if (cache.transportErrorReported || cache.connectStartedAt === undefined) {
                return
            }
            cache.transportErrorReported = true
            insights.capture('wizard sync transport error', {
                transport: cache.syncMode ?? 'sse',
                workflow_id: props.workflowId,
                error,
            })
        },
        // Flags can resolve after a cold-cache connect (insights-js loads them async), and ops can flip
        // the mode mid-incident. Reconnect when the resolved mode differs from the running transport —
        // the keyed disposable makes the swap idempotent.
        [featureFlagLogic.actionTypes.setFeatureFlags]: () => {
            if (cache.syncMode === undefined) {
                return
            }
            if (values.connectionStatus === 'idle' || values.connectionStatus === 'closed') {
                return
            }
            const mode = resolveWizardSyncMode(values.featureFlags[FEATURE_FLAGS.ONBOARDING_WIZARD_SYNC_MODE])
            if (mode !== cache.syncMode) {
                actions.connect()
            }
        },
    })),
])