insights/products/ai_observability/frontend/traceMessagesLazyLoaderLogic.ts

import { MakeLogicType, actions, events, kea, listeners, path, reducers, selectors } from 'kea'

import api from 'lib/api'
import { chunk } from 'lib/utils/arrays'

import { InsightsQLQuery, NodeKind } from '~/queries/schema/schema-general'
import { escapeInsightsQLString } from '~/queries/utils'

import { parsePartialJSON } from './utils'

export interface TraceMessages {
    lastInput: unknown
    lastOutput: unknown
    // Latest $ai_generation payloads, used when the preferred $ai_trace
    // state wrapper is missing or resolves to an empty messages array.
    lastInputFallback: unknown
    lastOutputFallback: unknown
}

export interface TraceLazyLoadRequest {
    id: string
    createdAt: string | null
}

const BATCH_MAX_SIZE = 100
const BATCH_DEBOUNCE_MS = 0
// Cap per field in characters. Bounds network payload and Datastore
// aggregation state; truncated JSON is recovered with parsePartialJSON.
const FIELD_TRUNCATE_CHARS = 2000
const TRACE_ID_MAX_LENGTH = 128
// Matches TracesQueryDateRange.CAPTURE_RANGE_MINUTES on the backend runner.
const BATCH_WINDOW_BUFFER_MINUTES = 10

// "2026-04-11T19:20:55.828Z" → "2026-04-11 19:20:55" (Datastore toDateTime format, UTC).
function formatInsightsqlDateTime(d: Date): string {
    return d.toISOString().replace('T', ' ').slice(0, 19)
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface traceMessagesLazyLoaderLogicValues {
    getTraceMessages: (traceId: string) => TraceMessages | null | undefined
    loadingTraceIds: Set<string>
    messagesByTraceId: Record<string, TraceMessages | null>
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface traceMessagesLazyLoaderLogicActions {
    ensureTraceMessagesLoaded: (requests: TraceLazyLoadRequest[]) => {
        requests: TraceLazyLoadRequest[]
    }
    loadTraceMessagesBatchFailure: (requestedTraceIds: string[]) => {
        requestedTraceIds: string[]
    }
    loadTraceMessagesBatchSuccess: (
        results: Record<string, TraceMessages>,
        requestedTraceIds: string[]
    ) => {
        requestedTraceIds: string[]
        results: Record<string, TraceMessages>
    }
    markTraceIdsLoading: (traceIds: string[]) => {
        traceIds: string[]
    }
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface traceMessagesLazyLoaderLogicMeta {
    __keaTypeGenInternalSelectorTypes: {
        getTraceMessages: (
            messagesByTraceId: Record<string, TraceMessages | null>
        ) => (traceId: string) => TraceMessages | null | undefined
    }
}

export type traceMessagesLazyLoaderLogicType = MakeLogicType<
    traceMessagesLazyLoaderLogicValues,
    traceMessagesLazyLoaderLogicActions,
    Record<string, any>,
    traceMessagesLazyLoaderLogicMeta
>

export const traceMessagesLazyLoaderLogic = kea<traceMessagesLazyLoaderLogicType>([
    path(['products', 'ai_observability', 'frontend', 'traceMessagesLazyLoaderLogic']),

    actions({
        ensureTraceMessagesLoaded: (requests: TraceLazyLoadRequest[]) => ({ requests }),
        markTraceIdsLoading: (traceIds: string[]) => ({ traceIds }),
        loadTraceMessagesBatchSuccess: (results: Record<string, TraceMessages>, requestedTraceIds: string[]) => ({
            results,
            requestedTraceIds,
        }),
        loadTraceMessagesBatchFailure: (requestedTraceIds: string[]) => ({ requestedTraceIds }),
    }),

    reducers({
        messagesByTraceId: [
            {} as Record<string, TraceMessages | null>,
            {
                loadTraceMessagesBatchSuccess: (state, { results, requestedTraceIds }) => {
                    const next = { ...state }
                    for (const traceId of requestedTraceIds) {
                        next[traceId] = results[traceId] ?? {
                            lastInput: null,
                            lastOutput: null,
                            lastInputFallback: null,
                            lastOutputFallback: null,
                        }
                    }
                    return next
                },
                loadTraceMessagesBatchFailure: (state, { requestedTraceIds }) => {
                    const next = { ...state }
                    for (const traceId of requestedTraceIds) {
                        next[traceId] = null
                    }
                    return next
                },
            },
        ],

        loadingTraceIds: [
            new Set<string>(),
            {
                markTraceIdsLoading: (state, { traceIds }) => {
                    const next = new Set(state)
                    for (const traceId of traceIds) {
                        if (traceId) {
                            next.add(traceId)
                        }
                    }
                    return next
                },
                loadTraceMessagesBatchSuccess: (state, { requestedTraceIds }) => {
                    const next = new Set(state)
                    for (const traceId of requestedTraceIds) {
                        next.delete(traceId)
                    }
                    return next
                },
                loadTraceMessagesBatchFailure: (state, { requestedTraceIds }) => {
                    const next = new Set(state)
                    for (const traceId of requestedTraceIds) {
                        next.delete(traceId)
                    }
                    return next
                },
            },
        ],
    }),

    selectors({
        getTraceMessages: [
            (s) => [s.messagesByTraceId],
            (
                messagesByTraceId: Record<string, TraceMessages | null>
            ): ((traceId: string) => TraceMessages | null | undefined) => {
                return (traceId: string) => messagesByTraceId[traceId]
            },
        ],
    }),

    listeners(({ actions, values, cache }) => ({
        ensureTraceMessagesLoaded: ({ requests }) => {
            const uncached = requests.filter(
                (r) => r.id && values.messagesByTraceId[r.id] === undefined && !values.loadingTraceIds.has(r.id)
            )

            if (uncached.length === 0) {
                return
            }

            actions.markTraceIdsLoading(uncached.map((r) => r.id))

            const pendingTraces = cache.pendingTraces as Map<string, TraceLazyLoadRequest>
            for (const r of uncached) {
                pendingTraces.set(r.id, r)
            }

            if (cache.batchTimer) {
                return
            }

            cache.batchTimer = setTimeout(async () => {
                const requestedTraces = Array.from(pendingTraces.values())
                pendingTraces.clear()
                cache.batchTimer = null

                if (requestedTraces.length === 0) {
                    return
                }

                const chunks = chunk(requestedTraces, BATCH_MAX_SIZE)

                await Promise.allSettled(
                    chunks.map(async (batch) => {
                        try {
                            // Drop IDs that are too long to inline safely or whose
                            // createdAt we can't parse into a scan window anchor.
                            const safe = batch.flatMap((r) => {
                                const createdAtMs = r.createdAt ? Date.parse(r.createdAt) : NaN
                                return r.id && r.id.length <= TRACE_ID_MAX_LENGTH && Number.isFinite(createdAtMs)
                                    ? [{ id: r.id, createdAtMs }]
                                    : []
                            })
                            // Mark the dropped entries as failed upfront so they
                            // aren't silently counted as "successfully loaded with
                            // no data" alongside the real response.
                            const safeIdSet = new Set(safe.map((s) => s.id))
                            const invalidIds = batch.map((r) => r.id).filter((id) => !safeIdSet.has(id))
                            if (invalidIds.length > 0) {
                                actions.loadTraceMessagesBatchFailure(invalidIds)
                            }
                            if (safe.length === 0) {
                                return
                            }

                            // Union window across the batch (min/max createdAt ± buffer).
                            // Stays narrow when rows are clustered in time, so Datastore
                            // can prune partitions regardless of the UI date filter.
                            const bufferMs = BATCH_WINDOW_BUFFER_MINUTES * 60 * 1000
                            const createdAtList = safe.map((s) => s.createdAtMs)
                            const fromStr = formatInsightsqlDateTime(new Date(Math.min(...createdAtList) - bufferMs))
                            const toStr = formatInsightsqlDateTime(new Date(Math.max(...createdAtList) + bufferMs))

                            // Inlining IDs + timestamps rather than using a values dict: we
                            // can't combine `{values}` with `{filters}`-style placeholders
                            // because parse_select eagerly resolves all placeholders against
                            // the values dict before find_placeholders runs.
                            // escapeInsightsQLString returns the value already wrapped in single quotes.
                            const idList = safe.map((s) => escapeInsightsQLString(s.id)).join(',')
                            const query: InsightsQLQuery = {
                                kind: NodeKind.InsightsQLQuery,
                                query: `
                                    SELECT
                                        trace_id,
                                        anyIf(
                                            substringUTF8(input_state, 1, ${FIELD_TRUNCATE_CHARS}),
                                            event = '$ai_trace'
                                                AND length(input_state) > 0
                                        ) AS last_input,
                                        anyIf(
                                            substringUTF8(output_state, 1, ${FIELD_TRUNCATE_CHARS}),
                                            event = '$ai_trace'
                                                AND length(output_state) > 0
                                        ) AS last_output,
                                        argMaxIf(
                                            substringUTF8(input, 1, ${FIELD_TRUNCATE_CHARS}),
                                            timestamp,
                                            event = '$ai_generation'
                                                AND length(input) > 0
                                        ) AS last_input_fallback,
                                        argMaxIf(
                                            substringUTF8(output_choices, 1, ${FIELD_TRUNCATE_CHARS}),
                                            timestamp,
                                            event = '$ai_generation'
                                                AND length(output_choices) > 0
                                        ) AS last_output_fallback
                                    FROM insights.ai_events AS ai_events
                                    WHERE event IN ('$ai_trace', '$ai_generation')
                                      AND timestamp >= toDateTime('${fromStr}', 'UTC')
                                      AND timestamp <= toDateTime('${toStr}', 'UTC')
                                      AND trace_id IN (${idList})
                                    GROUP BY trace_id
                                `,
                            }

                            const response = await api.query(query)
                            const results: Record<string, TraceMessages> = {}

                            for (const row of response.results || []) {
                                const [traceId, lastInput, lastOutput, lastInputFallback, lastOutputFallback] = row as [
                                    string,
                                    unknown,
                                    unknown,
                                    unknown,
                                    unknown,
                                ]
                                if (traceId) {
                                    results[traceId] = {
                                        lastInput: parseTruncatedJson(lastInput),
                                        lastOutput: parseTruncatedJson(lastOutput),
                                        lastInputFallback: parseTruncatedJson(lastInputFallback),
                                        lastOutputFallback: parseTruncatedJson(lastOutputFallback),
                                    }
                                }
                            }

                            actions.loadTraceMessagesBatchSuccess(
                                results,
                                safe.map((s) => s.id)
                            )
                        } catch (error) {
                            console.warn('Error loading trace messages batch', error)
                            actions.loadTraceMessagesBatchFailure(batch.map((r) => r.id))
                        }
                    })
                )
            }, BATCH_DEBOUNCE_MS)
        },
    })),

    events(({ cache }) => ({
        afterMount: () => {
            cache.pendingTraces = new Map<string, TraceLazyLoadRequest>()
            cache.batchTimer = null
        },
        beforeUnmount: () => {
            if (cache.batchTimer) {
                clearTimeout(cache.batchTimer)
                cache.batchTimer = null
            }
            cache.pendingTraces?.clear?.()
        },
    })),
])

// Parse Datastore-truncated JSON: strict parse, then partial-JSON recovery,
// then raw string so the normalizer can still show something.
function parseTruncatedJson(value: unknown): unknown {
    if (typeof value !== 'string') {
        return value
    }
    if (value.length === 0) {
        return null
    }
    try {
        return JSON.parse(value)
    } catch {
        try {
            return parsePartialJSON(value)
        } catch {
            return value
        }
    }
}