insights/products/streamlit_apps/frontend/streamlitAppLogic.ts

import {
    MakeLogicType,
    actions,
    afterMount,
    beforeUnmount,
    connect,
    kea,
    key,
    listeners,
    path,
    props,
    reducers,
    selectors,
} from 'kea'
import { loaders } from 'kea-loaders'

import { projectLogic } from 'scenes/projectLogic'
import { urls } from 'scenes/urls'

import { Breadcrumb } from '~/types'

import {
    streamlitAppsConnectInfoRetrieve,
    streamlitAppsRestartCreate,
    streamlitAppsRetrieve,
    streamlitAppsStartCreate,
    streamlitAppsStatusRetrieve,
    streamlitAppsStopCreate,
} from './generated/api'
import type { AppContractApi, AppSandboxContractApi, StreamlitConnectInfoApi } from './generated/api.schemas'
import { StreamlitAppStatus, toStreamlitAppStatus } from './types'

export interface StreamlitAppLogicProps {
    shortId: string
}

const POLL_INTERVAL_MS = 2000
const HEALTH_POLL_INTERVAL_MS = 15000
const TOKEN_REFRESH_RATIO = 0.8
// Floor prevents a malformed expires_in from scheduling an immediate refresh loop.
const MIN_TOKEN_REFRESH_MS = 5_000

type TimerKey = 'pollTimer' | 'healthPollTimer' | 'tokenRefreshTimer'

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppLogicValues {
    currentProjectId: number | null // projectLogic
    appStatus: StreamlitAppStatus
    breadcrumbs: Breadcrumb[]
    connectError: string | null
    connectInfo: StreamlitConnectInfoApi | null
    connectInfoLoading: boolean
    iframeSrc: string | null
    isHealthPolling: boolean
    isPolling: boolean
    sandboxStatus: AppSandboxContractApi | null
    sandboxStatusLoading: boolean
    streamlitApp: AppContractApi | null
    streamlitAppLoading: boolean
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppLogicActions {
    clearTokenRefresh: () => {
        value: true
    }
    loadConnectInfo: () => any
    loadConnectInfoFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadConnectInfoSuccess: (
        connectInfo: StreamlitConnectInfoApi,
        payload?: any
    ) => {
        connectInfo: StreamlitConnectInfoApi
        payload?: any
    }
    loadSandboxStatus: () => any
    loadSandboxStatusFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadSandboxStatusSuccess: (
        sandboxStatus: {
            last_activity_at: string | null
            last_error: string
            restart_count: number
            started_at: string | null
            status: string
            version_number: number | null
        },
        payload?: any
    ) => {
        sandboxStatus: {
            last_activity_at: string | null
            last_error: string
            restart_count: number
            started_at: string | null
            status: string
            version_number: number | null
        }
        payload?: any
    }
    loadStreamlitApp: () => any
    loadStreamlitAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadStreamlitAppSuccess: (
        streamlitApp: AppContractApi,
        payload?: any
    ) => {
        streamlitApp: AppContractApi
        payload?: any
    }
    restartApp: () => any
    restartAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    restartAppSuccess: (
        sandboxStatus: AppSandboxContractApi | null,
        payload?: any
    ) => {
        sandboxStatus: AppSandboxContractApi | null
        payload?: any
    }
    scheduleTokenRefresh: (expiresIn: number) => {
        expiresIn: number
    }
    startApp: () => any
    startAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    startAppSuccess: (
        sandboxStatus: AppSandboxContractApi | null,
        payload?: any
    ) => {
        sandboxStatus: AppSandboxContractApi | null
        payload?: any
    }
    startHealthPolling: () => {
        value: true
    }
    startPolling: () => {
        value: true
    }
    stopApp: () => any
    stopAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    stopAppSuccess: (
        sandboxStatus: AppSandboxContractApi | null,
        payload?: any
    ) => {
        sandboxStatus: AppSandboxContractApi | null
        payload?: any
    }
    stopHealthPolling: () => {
        value: true
    }
    stopPolling: () => {
        value: true
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppLogicMeta {
    key: string
    __keaTypeGenInternalSelectorTypes: {
        breadcrumbs: (streamlitApp: AppContractApi | null) => Breadcrumb[]
        appStatus: (sandboxStatus: AppSandboxContractApi | null) => StreamlitAppStatus
        iframeSrc: (
            appStatus: 'error' | 'running' | 'starting' | 'stopped' | 'stopping',
            connectInfo: StreamlitConnectInfoApi | null
        ) => string | null
    }
}

export type streamlitAppLogicType = MakeLogicType<
    streamlitAppLogicValues,
    streamlitAppLogicActions,
    StreamlitAppLogicProps,
    streamlitAppLogicMeta
>

export const streamlitAppLogic = kea<streamlitAppLogicType>([
    path(['products', 'streamlit_apps', 'frontend', 'streamlitAppLogic']),
    props({} as StreamlitAppLogicProps),
    key((props) => props.shortId),

    connect(() => ({
        values: [projectLogic, ['currentProjectId']],
    })),

    actions({
        startPolling: true,
        stopPolling: true,
        startHealthPolling: true,
        stopHealthPolling: true,
        scheduleTokenRefresh: (expiresIn: number) => ({ expiresIn }),
        clearTokenRefresh: true,
    }),

    loaders(({ props, values }) => ({
        streamlitApp: [
            null as AppContractApi | null,
            {
                loadStreamlitApp: async () => {
                    return await streamlitAppsRetrieve(String(values.currentProjectId), props.shortId)
                },
            },
        ],
        sandboxStatus: [
            null as AppSandboxContractApi | null,
            {
                loadSandboxStatus: async () => {
                    const status = await streamlitAppsStatusRetrieve(String(values.currentProjectId), props.shortId)
                    // The status action omits version_number when stopped; the embedded sandbox
                    // always carries it, so normalize to one shape.
                    return { ...status, version_number: status.version_number ?? null }
                },
                startApp: async () => {
                    const app = await streamlitAppsStartCreate(String(values.currentProjectId), props.shortId)
                    return app.sandbox ?? null
                },
                stopApp: async () => {
                    const app = await streamlitAppsStopCreate(String(values.currentProjectId), props.shortId)
                    return app.sandbox ?? null
                },
                restartApp: async () => {
                    const app = await streamlitAppsRestartCreate(String(values.currentProjectId), props.shortId)
                    return app.sandbox ?? null
                },
            },
        ],
        connectInfo: [
            null as StreamlitConnectInfoApi | null,
            {
                loadConnectInfo: async () => {
                    return await streamlitAppsConnectInfoRetrieve(String(values.currentProjectId), props.shortId)
                },
            },
        ],
    })),

    reducers({
        isPolling: [
            false,
            {
                startPolling: () => true,
                stopPolling: () => false,
            },
        ],
        isHealthPolling: [
            false,
            {
                startHealthPolling: () => true,
                stopHealthPolling: () => false,
            },
        ],
        connectError: [
            null as string | null,
            {
                loadConnectInfoFailure: (_, { error }) => String(error),
                loadConnectInfoSuccess: () => null,
                startApp: () => null,
                restartApp: () => null,
            },
        ],
        connectInfo: {
            // Restart mints a new sandbox — clear so we refetch fresh tokens.
            restartApp: () => null,
            restartAppSuccess: () => null,
        },
    }),

    selectors({
        breadcrumbs: [
            (s) => [s.streamlitApp],
            (streamlitApp: AppContractApi | null): Breadcrumb[] => [
                {
                    key: 'StreamlitApps',
                    name: 'Apps',
                    path: urls.streamlitApps(),
                },
                {
                    key: ['StreamlitApp', streamlitApp?.short_id || 'new'],
                    name: streamlitApp?.name || 'Loading...',
                },
            ],
        ],
        appStatus: [
            (s) => [s.sandboxStatus],
            (sandboxStatus: AppSandboxContractApi | null): StreamlitAppStatus =>
                toStreamlitAppStatus(sandboxStatus?.status),
        ],
        iframeSrc: [
            (s) => [s.appStatus, s.connectInfo],
            (appStatus: StreamlitAppStatus, connectInfo: StreamlitConnectInfoApi | null): string | null => {
                if (appStatus !== 'running' || !connectInfo) {
                    return null
                }
                return connectInfo.iframe_url
            },
        ],
    }),

    listeners(({ actions, values, cache }) => {
        // Idempotent — kea listeners can fire repeatedly and timers would stack.
        // clearInterval also accepts timeout IDs, so this covers both.
        const clearCached = (key: TimerKey): void => {
            if (cache[key]) {
                clearInterval(cache[key])
                cache[key] = null
            }
        }
        return {
            loadStreamlitAppSuccess: ({ streamlitApp }) => {
                if (!streamlitApp) {
                    return
                }
                // Seed sandboxStatus from app data so appStatus is correct immediately.
                if (streamlitApp.sandbox) {
                    actions.loadSandboxStatusSuccess(streamlitApp.sandbox)
                }
                const status = streamlitApp.sandbox?.status ?? 'stopped'
                if (status === 'starting') {
                    actions.startPolling()
                } else if (status === 'running') {
                    actions.loadConnectInfo()
                }
            },
            startAppSuccess: () => {
                actions.startPolling()
            },
            startPolling: () => {
                clearCached('pollTimer')
                cache.pollTimer = setInterval(() => {
                    actions.loadSandboxStatus()
                }, POLL_INTERVAL_MS)
            },
            stopPolling: () => clearCached('pollTimer'),
            startHealthPolling: () => {
                clearCached('healthPollTimer')
                cache.healthPollTimer = setInterval(() => {
                    actions.loadSandboxStatus()
                }, HEALTH_POLL_INTERVAL_MS)
            },
            stopHealthPolling: () => clearCached('healthPollTimer'),
            loadSandboxStatusSuccess: ({ sandboxStatus }) => {
                if (!sandboxStatus) {
                    return
                }
                if (sandboxStatus.status === 'running') {
                    actions.stopPolling()
                    if (!values.isHealthPolling) {
                        actions.startHealthPolling()
                    }
                    if (!values.connectInfo) {
                        actions.loadConnectInfo()
                    }
                } else if (sandboxStatus.status === 'error' || sandboxStatus.status === 'stopped') {
                    actions.stopPolling()
                    actions.stopHealthPolling()
                    actions.clearTokenRefresh()
                }
            },
            loadConnectInfoSuccess: ({ connectInfo }) => {
                if (connectInfo) {
                    actions.scheduleTokenRefresh(connectInfo.expires_in)
                }
            },
            scheduleTokenRefresh: ({ expiresIn }) => {
                actions.clearTokenRefresh()
                const refreshMs = Math.max(expiresIn * TOKEN_REFRESH_RATIO * 1000, MIN_TOKEN_REFRESH_MS)
                cache.tokenRefreshTimer = setTimeout(() => {
                    actions.loadConnectInfo()
                }, refreshMs)
            },
            loadConnectInfoFailure: () => {
                // Don't reschedule — a 5s retry loop would hammer a broken sandbox.
                actions.clearTokenRefresh()
            },
            clearTokenRefresh: () => clearCached('tokenRefreshTimer'),
            stopAppSuccess: () => {
                actions.stopPolling()
                actions.stopHealthPolling()
                actions.clearTokenRefresh()
            },
            restartAppSuccess: () => {
                actions.startPolling()
                actions.clearTokenRefresh()
            },
        }
    }),

    afterMount(({ actions }) => {
        actions.loadStreamlitApp()
    }),

    beforeUnmount(({ actions }) => {
        actions.stopPolling()
        actions.stopHealthPolling()
        actions.clearTokenRefresh()
    }),
])