insights/products/streamlit_apps/frontend/streamlitAppEditLogic.ts

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

import { toast } from '@hanzo/elements'

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

import { Breadcrumb } from '~/types'

import {
    getStreamlitAppsUploadVersionCreateUrl,
    streamlitAppsActivateVersionCreate,
    streamlitAppsCreate,
    streamlitAppsDestroy,
    streamlitAppsRetrieve,
    streamlitAppsUpdate,
    streamlitAppsVersionsRetrieve,
} from './generated/api'
import type { AppContractApi, AppVersionContractApi } from './generated/api.schemas'
import { streamlitAppsLogic } from './streamlitAppsLogic'

export interface StreamlitAppEditLogicProps {
    shortId: string
}

// The generated client JSON-encodes its body, so it can't carry the zip. Only the body is
// hand-rolled — the URL still comes from the generated builder.
async function uploadAppVersion(projectId: string, shortId: string, file: File): Promise<AppVersionContractApi> {
    const formData = new FormData()
    formData.append('file', file)
    // nosemgrep: prefer-codegen-api
    return await api.create<AppVersionContractApi>(getStreamlitAppsUploadVersionCreateUrl(projectId, shortId), formData)
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppEditLogicValues {
    currentProjectId: number | null // projectLogic
    breadcrumbs: Breadcrumb[]
    cpuCores: number
    description: string
    memoryGb: number
    name: string
    savedApp: AppContractApi | null
    savedAppLoading: boolean
    streamlitApp: AppContractApi | null
    streamlitAppLoading: boolean
    versions: AppVersionContractApi[]
    versionsLoading: boolean
    zipFile: File | null
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppEditLogicActions {
    deleteApp: () => any
    deleteAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    deleteAppSuccess: (
        savedApp: null,
        payload?: any
    ) => {
        savedApp: null
        payload?: any
    }
    loadStreamlitApp: () => any
    loadStreamlitAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadStreamlitAppSuccess: (
        streamlitApp: AppContractApi | null,
        payload?: any
    ) => {
        streamlitApp: AppContractApi | null
        payload?: any
    }
    loadVersions: () => any
    loadVersionsFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    loadVersionsSuccess: (
        versions: AppVersionContractApi[],
        payload?: any
    ) => {
        versions: AppVersionContractApi[]
        payload?: any
    }
    saveApp: () => any
    saveAppFailure: (
        error: string,
        errorObject?: any
    ) => {
        error: string
        errorObject?: any
    }
    saveAppSuccess: (
        savedApp: AppContractApi,
        payload?: any
    ) => {
        savedApp: AppContractApi
        payload?: any
    }
    setActiveVersionInState: (activeVersion: AppVersionContractApi) => {
        activeVersion: AppVersionContractApi
    }
    setActiveVersionNumber: (versionNumber: number) => {
        versionNumber: number
    }
    setCpuCores: (cpuCores: number) => {
        cpuCores: number
    }
    setDescription: (description: string) => {
        description: string
    }
    setMemoryGb: (memoryGb: number) => {
        memoryGb: number
    }
    setName: (name: string) => {
        name: string
    }
    setZipFile: (file: File | null) => {
        file: File | null
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface streamlitAppEditLogicMeta {
    key: string
    __keaTypeGenInternalSelectorTypes: {
        breadcrumbs: (streamlitApp: AppContractApi | null) => Breadcrumb[]
    }
}

export type streamlitAppEditLogicType = MakeLogicType<
    streamlitAppEditLogicValues,
    streamlitAppEditLogicActions,
    StreamlitAppEditLogicProps,
    streamlitAppEditLogicMeta
>

export const streamlitAppEditLogic = kea<streamlitAppEditLogicType>([
    path(['products', 'streamlit_apps', 'frontend', 'streamlitAppEditLogic']),
    props({} as StreamlitAppEditLogicProps),
    key((props) => props.shortId),

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

    actions({
        setName: (name: string) => ({ name }),
        setDescription: (description: string) => ({ description }),
        setCpuCores: (cpuCores: number) => ({ cpuCores }),
        setMemoryGb: (memoryGb: number) => ({ memoryGb }),
        setZipFile: (file: File | null) => ({ file }),
        setActiveVersionNumber: (versionNumber: number) => ({ versionNumber }),
        // Narrow update; dispatching loadStreamlitAppSuccess would clobber unsaved form edits.
        setActiveVersionInState: (activeVersion: AppVersionContractApi) => ({ activeVersion }),
    }),

    loaders(({ props, values }) => ({
        streamlitApp: [
            null as AppContractApi | null,
            {
                loadStreamlitApp: async () => {
                    if (props.shortId === 'new') {
                        return null
                    }
                    return await streamlitAppsRetrieve(String(values.currentProjectId), props.shortId)
                },
            },
        ],
        versions: [
            [] as AppVersionContractApi[],
            {
                loadVersions: async () => {
                    if (props.shortId === 'new') {
                        return []
                    }
                    const response = await streamlitAppsVersionsRetrieve(String(values.currentProjectId), props.shortId)
                    return response.results
                },
            },
        ],
        savedApp: [
            null as AppContractApi | null,
            {
                saveApp: async () => {
                    const isNew = props.shortId === 'new'
                    const projectId = String(values.currentProjectId)

                    let app: AppContractApi
                    if (isNew) {
                        app = await streamlitAppsCreate(projectId, {
                            name: values.name,
                            description: values.description,
                            cpu_cores: values.cpuCores,
                            memory_gb: values.memoryGb,
                        })
                    } else {
                        app = await streamlitAppsUpdate(projectId, props.shortId, {
                            name: values.name,
                            description: values.description,
                            cpu_cores: values.cpuCores,
                            memory_gb: values.memoryGb,
                        })
                    }

                    if (values.zipFile) {
                        try {
                            await uploadAppVersion(projectId, app.short_id, values.zipFile)
                            app = await streamlitAppsRetrieve(projectId, app.short_id)
                        } catch {
                            toast.error('App created but zip upload failed. You can retry from the edit page.')
                            router.actions.push(urls.streamlitAppEdit(app.short_id))
                            return app
                        }
                    }

                    toast.success(isNew ? 'App created' : 'App saved')
                    if (isNew) {
                        router.actions.push(urls.streamlitApp(app.short_id))
                    }
                    return app
                },
                deleteApp: async () => {
                    await streamlitAppsDestroy(String(values.currentProjectId), props.shortId)
                    toast.success('App deleted')
                    streamlitAppsLogic.findMounted()?.actions.loadStreamlitApps()
                    router.actions.push(urls.streamlitApps())
                    return null
                },
            },
        ],
    })),

    reducers({
        streamlitApp: {
            setActiveVersionInState: (state: AppContractApi | null, { activeVersion }) =>
                state ? { ...state, active_version: activeVersion } : state,
        },
        name: [
            '',
            {
                setName: (_, { name }) => name,
                loadStreamlitAppSuccess: (_, { streamlitApp }) => streamlitApp?.name ?? '',
            },
        ],
        description: [
            '',
            {
                setDescription: (_, { description }) => description,
                loadStreamlitAppSuccess: (_, { streamlitApp }) => streamlitApp?.description ?? '',
            },
        ],
        cpuCores: [
            0.5,
            {
                setCpuCores: (_, { cpuCores }) => cpuCores,
                loadStreamlitAppSuccess: (_, { streamlitApp }) => streamlitApp?.cpu_cores ?? 0.5,
            },
        ],
        memoryGb: [
            1,
            {
                setMemoryGb: (_, { memoryGb }) => memoryGb,
                loadStreamlitAppSuccess: (_, { streamlitApp }) => streamlitApp?.memory_gb ?? 1,
            },
        ],
        zipFile: [
            null as File | null,
            {
                setZipFile: (_, { file }) => file,
            },
        ],
    }),

    selectors({
        breadcrumbs: [
            (s) => [s.streamlitApp],
            (streamlitApp: AppContractApi | null): Breadcrumb[] => [
                {
                    key: 'StreamlitApps',
                    name: 'Apps',
                    path: urls.streamlitApps(),
                },
                {
                    key: ['StreamlitAppEdit', streamlitApp?.short_id || 'new'],
                    name: streamlitApp ? streamlitApp.name : 'New app',
                },
            ],
        ],
    }),

    listeners(({ props, values, actions }) => ({
        saveAppSuccess: ({ savedApp }) => {
            if (savedApp) {
                streamlitAppsLogic.findMounted()?.actions.updateStreamlitApp(savedApp)
            }
        },
        setActiveVersionNumber: async ({ versionNumber }) => {
            if (props.shortId === 'new') {
                return
            }
            try {
                const response = await streamlitAppsActivateVersionCreate(
                    String(values.currentProjectId),
                    props.shortId,
                    { version_number: versionNumber }
                )
                toast.success(`Switched to v${versionNumber}. Restart the app to apply.`)
                const nextActiveVersion = response?.active_version ?? values.streamlitApp?.active_version ?? null
                if (values.streamlitApp && nextActiveVersion) {
                    // Narrow update; see setActiveVersionInState for rationale.
                    actions.setActiveVersionInState(nextActiveVersion)
                    streamlitAppsLogic.findMounted()?.actions.updateStreamlitApp({
                        ...values.streamlitApp,
                        active_version: nextActiveVersion,
                    })
                }
            } catch (error: any) {
                toast.error(
                    `Failed to activate v${versionNumber}: ${error?.detail || error?.message || 'unknown error'}`
                )
            }
        },
    })),

    afterMount(({ props, actions }) => {
        if (props.shortId !== 'new') {
            actions.loadStreamlitApp()
            actions.loadVersions()
        }
    }),
])