insights/products/tracing/frontend/tracingOperationSceneLogic.ts

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

import { toast } from '@hanzo/elements'

import api from 'lib/api'
import { dataColorVars } from 'lib/colors'
import { FEATURE_FLAGS } from 'lib/constants'
import { featureFlagLogic, type FeatureFlagsSet } from 'lib/logic/featureFlagLogic'
import { objectsEqual } from 'lib/utils/objects'

import { AggregatedSpanRow, DateRange } from '~/queries/schema/schema-general'

import {
    type DurationHistogramRow,
    type LatencyHeatmapRow,
    pivotDurationHistogram,
    pivotLatencyHeatmap,
    selectionToDurationRange,
    type TracingDurationHistogramData,
    type TracingLatencyHeatmapData,
} from './durationBuckets'
import type { HeatmapBrushSelection } from './heatmapBrush'
import { type DurationRange, operationFilterGroup } from './operationFilters'
import { traceLookupDateRange } from './traceLinks'
import { isUserInitiatedError, NEW_QUERY_STARTED_ERROR_MESSAGE } from './tracingDataLogic'
import { DEFAULT_DATE_RANGE } from './tracingFiltersLogic'
import type { Span } from './types'

export interface TracingOperationSceneLogicProps {
    serviceName: string
    spanName: string
}

// Which chart shows the operation's latency: the static distribution or the over-time heatmap.
export type OperationChartType = 'histogram' | 'heatmap'

const SAMPLE_LIMIT = 100

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface tracingOperationSceneLogicValues {
    featureFlags: FeatureFlagsSet; // featureFlagLogic
    chartType: OperationChartType;
    currentSample: Span | null;
    dateRange: DateRange;
    durationSelection: DurationRange | null;
    heatmapEnabled: boolean;
    histogramData: TracingDurationHistogramData;
    latencyHeatmapData: TracingLatencyHeatmapData;
    latencyHeatmapLoaded: boolean;
    operationStats: AggregatedSpanRow | null;
    operationStatsLoading: boolean;
    rawHistogram: DurationHistogramRow[];
    rawHistogramLoading: boolean;
    rawLatencyHeatmap: LatencyHeatmapRow[];
    rawLatencyHeatmapLoading: boolean;
    sampleIndex: number;
    sampleTraceSpans: Span[];
    sampleTraceSpansLoading: boolean;
    samples: Span[];
    samplesHaveMore: boolean;
    samplesLoading: boolean;
    selectedSpanId: string | null;
    serviceName: string;
    spanName: string;
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface tracingOperationSceneLogicActions {
    setFeatureFlags: (flags: string[], variants: Record<string, boolean | string>, payloads?: import("lib/logic/featureFlagLogic").FeatureFlagPayloads | undefined) => {
        flags: string[];
        payloads: import("lib/logic/featureFlagLogic").FeatureFlagPayloads;
        variants: Record<string, boolean | string>;
    }; // featureFlagLogic
    applyHeatmapBrush: (selection: HeatmapBrushSelection) => {
        selection: HeatmapBrushSelection;
    };
    fetchHistogram: (_: void) => void;
    fetchHistogramFailure: (error: string, errorObject?: any) => {
        error: string;
        errorObject?: any;
    };
    fetchHistogramSuccess: (rawHistogram: {
        bucket_ns: number;
        count: number;
        service: string;
    }[], payload?: void) => {
        rawHistogram: {
            bucket_ns: number;
            count: number;
            service: string;
        }[];
        payload?: void;
    };
    fetchLatencyHeatmap: (_: void) => void;
    fetchLatencyHeatmapFailure: (error: string, errorObject?: any) => {
        error: string;
        errorObject?: any;
    };
    fetchLatencyHeatmapSuccess: (rawLatencyHeatmap: {
        bucket_ns: number;
        count: number;
        time: string;
    }[], payload?: void) => {
        rawLatencyHeatmap: {
            bucket_ns: number;
            count: number;
            time: string;
        }[];
        payload?: void;
    };
    fetchSampleTrace: ({ sample }: {
        sample: Span;
    }) => {
        sample: Span;
    };
    fetchSampleTraceFailure: (error: string, errorObject?: any) => {
        error: string;
        errorObject?: any;
    };
    fetchSampleTraceSuccess: (sampleTraceSpans: Span[], payload?: {
        sample: Span;
    }) => {
        sampleTraceSpans: Span[];
        payload?: {
            sample: Span;
        };
    };
    fetchSamples: (_: void) => void;
    fetchSamplesFailure: (error: string, errorObject?: any) => {
        error: string;
        errorObject?: any;
    };
    fetchSamplesSuccess: (samples: Span[], payload?: void) => {
        samples: Span[];
        payload?: void;
    };
    fetchStats: (_: void) => void;
    fetchStatsFailure: (error: string, errorObject?: any) => {
        error: string;
        errorObject?: any;
    };
    fetchStatsSuccess: (operationStats: AggregatedSpanRow | null, payload?: void) => {
        operationStats: AggregatedSpanRow | null;
        payload?: void;
    };
    loadCurrentSampleTrace: () => {
        value: true;
    };
    selectSpan: (spanId: string | null) => {
        spanId: string | null;
    };
    setChartType: (chartType: OperationChartType) => {
        chartType: OperationChartType;
    };
    setDateRange: (dateRange: DateRange) => {
        dateRange: DateRange;
    };
    setDurationSelection: (selection: DurationRange | null) => {
        selection: DurationRange | null;
    };
    setSampleIndex: (index: number) => {
        index: number;
    };
    setSamplesHaveMore: (hasMore: boolean) => {
        hasMore: boolean;
    };
}

// Generated by kea-typegen. Update if you're an agent, ignore if you're human.
export interface tracingOperationSceneLogicMeta {
    key: string;
    __keaTypeGenInternalSelectorTypes: {
        serviceName: (arg: any) => string;
        spanName: (arg: any) => string;
        heatmapEnabled: (featureFlags: FeatureFlagsSet) => boolean;
        histogramData: (rawHistogram: DurationHistogramRow[]) => TracingDurationHistogramData;
        latencyHeatmapData: (rawLatencyHeatmap: LatencyHeatmapRow[]) => TracingLatencyHeatmapData;
        currentSample: (samples: Span[], sampleIndex: number) => Span | null;
    };
}

export type tracingOperationSceneLogicType = MakeLogicType<tracingOperationSceneLogicValues, tracingOperationSceneLogicActions, TracingOperationSceneLogicProps, tracingOperationSceneLogicMeta>;

// One tracked AbortController per loader: re-issuing a fetch aborts the superseded request
// (not just discards its result), and unmount aborts whatever is still in flight.
function trackedSignal(cache: Record<string, any>, key: string): AbortSignal {
    let controller!: AbortController
    cache.disposables.add(
        () => {
            controller = new AbortController()
            return () => controller.abort(NEW_QUERY_STARTED_ERROR_MESSAGE)
        },
        key,
        // In-flight requests must survive tab switches — only supersession or unmount aborts.
        { pauseOnPageHidden: false }
    )
    return controller.signal
}

function isDateRange(value: unknown): value is DateRange {
    if (typeof value !== 'object' || value === null || Array.isArray(value)) {
        return false
    }
    const { date_from, date_to } = value as Record<string, unknown>
    return (
        (date_from === undefined || date_from === null || typeof date_from === 'string') &&
        (date_to === undefined || date_to === null || typeof date_to === 'string')
    )
}

// Returns a valid DateRange from the URL param, or null when it is missing,
// malformed JSON, or a syntactically valid value of the wrong shape.
function parseDateRangeParam(raw: unknown): DateRange | null {
    let value = raw
    if (typeof raw === 'string') {
        try {
            value = JSON.parse(raw)
        } catch {
            return null
        }
    }
    return isDateRange(value) ? value : null
}

export const tracingOperationSceneLogic = kea<tracingOperationSceneLogicType>([
    props({} as TracingOperationSceneLogicProps),
    key(({ serviceName, spanName }) => `${serviceName}//${spanName}`),
    path((key) => ['products', 'tracing', 'frontend', 'tracingOperationSceneLogic', key]),

    connect(() => ({
        values: [featureFlagLogic, ['featureFlags']],
        actions: [featureFlagLogic, ['setFeatureFlags']],
    })),

    actions({
        setDateRange: (dateRange: DateRange) => ({ dateRange }),
        setDurationSelection: (selection: DurationRange | null) => ({ selection }),
        setSampleIndex: (index: number) => ({ index }),
        selectSpan: (spanId: string | null) => ({ spanId }),
        setSamplesHaveMore: (hasMore: boolean) => ({ hasMore }),
        loadCurrentSampleTrace: true,
        setChartType: (chartType: OperationChartType) => ({ chartType }),
        // A completed 2D brush on the heatmap: x → date range, y → duration selection.
        applyHeatmapBrush: (selection: HeatmapBrushSelection) => ({ selection }),
    }),

    loaders(({ props, values, actions, cache }) => ({
        rawHistogram: [
            [] as DurationHistogramRow[],
            {
                fetchHistogram: async (_: void, breakpoint) => {
                    await breakpoint(10) // coalesce same-tick dispatches (URL restore + afterMount)
                    const response = await api.tracing.durationHistogram(
                        {
                            dateRange: values.dateRange,
                            serviceNames: [props.serviceName],
                            filterGroup: operationFilterGroup(props.spanName, null),
                            // Span-level: operations are often child spans, and the samples query
                            // below is span-level too — both must count the same population.
                            rootSpans: false,
                        },
                        trackedSignal(cache, 'fetchHistogram')
                    )
                    breakpoint()
                    return response.results
                },
            },
        ],
        rawLatencyHeatmap: [
            [] as LatencyHeatmapRow[],
            {
                fetchLatencyHeatmap: async (_: void, breakpoint) => {
                    await breakpoint(100)
                    const response = await api.tracing.latencyHeatmap({
                        dateRange: values.dateRange,
                        serviceNames: [props.serviceName],
                        filterGroup: operationFilterGroup(props.spanName, null),
                        // Span-level, like the histogram — operations are often child spans.
                        rootSpans: false,
                    })
                    breakpoint()
                    return response.results
                },
            },
        ],
        samples: [
            [] as Span[],
            {
                fetchSamples: async (_: void, breakpoint) => {
                    await breakpoint(10) // coalesce same-tick dispatches (URL restore + afterMount)
                    const response = await api.tracing.listSpans(
                        {
                            dateRange: values.dateRange,
                            orderBy: 'timestamp',
                            orderDirection: 'DESC',
                            serviceNames: [props.serviceName],
                            filterGroup: operationFilterGroup(props.spanName, values.durationSelection),
                            flatSpans: true,
                            limit: SAMPLE_LIMIT,
                            // Samples only feed the pager header (5 scalar fields); the waterfall
                            // loads the full trace separately, so skip the heavy attribute maps.
                            excludeAttributes: true,
                        },
                        trackedSignal(cache, 'fetchSamples')
                    )
                    breakpoint()
                    actions.setSamplesHaveMore(!!response.hasMore)
                    return response.results as Span[]
                },
            },
        ],
        operationStats: [
            null as AggregatedSpanRow | null,
            {
                fetchStats: async (_: void, breakpoint) => {
                    await breakpoint(10) // coalesce same-tick dispatches (URL restore + afterMount)
                    const response = await api.tracing.aggregate(
                        {
                            dateRange: values.dateRange,
                            serviceNames: [props.serviceName],
                            filterGroup: operationFilterGroup(props.spanName, null),
                        },
                        trackedSignal(cache, 'fetchStats')
                    )
                    breakpoint()
                    return response.results.find((row) => row.name === props.spanName) ?? null
                },
            },
        ],
        sampleTraceSpans: [
            [] as Span[],
            {
                fetchSampleTrace: async ({ sample }: { sample: Span }, breakpoint) => {
                    await breakpoint(100) // debounce rapid pager clicks
                    const response = await api.tracing.getTrace(
                        sample.trace_id,
                        {
                            dateRange: traceLookupDateRange(sample.timestamp),
                        },
                        trackedSignal(cache, 'fetchSampleTrace')
                    )
                    breakpoint()
                    return response.results as Span[]
                },
            },
        ],
    })),

    reducers({
        dateRange: [DEFAULT_DATE_RANGE as DateRange, { setDateRange: (_, { dateRange }) => dateRange }],
        durationSelection: [null as DurationRange | null, { setDurationSelection: (_, { selection }) => selection }],
        chartType: ['histogram' as OperationChartType, { setChartType: (_, { chartType }) => chartType }],
        sampleIndex: [
            0,
            {
                setSampleIndex: (_, { index }) => index,
                setDurationSelection: () => 0,
                setDateRange: () => 0,
            },
        ],
        // Distinguishes "heatmap fetch completed" (possibly empty) from "never fetched". A date
        // change invalidates it so a later switch back to the heatmap refetches for the new range.
        latencyHeatmapLoaded: [
            false,
            {
                fetchLatencyHeatmapSuccess: () => true,
                setDateRange: () => false,
            },
        ],
        selectedSpanId: [null as string | null, { selectSpan: (_, { spanId }) => spanId }],
        samplesHaveMore: [false, { setSamplesHaveMore: (_, { hasMore }) => hasMore }],
        // A stale waterfall must not linger under a new (or empty) sample set.
        sampleTraceSpans: { fetchSamples: () => [] },
    }),

    selectors({
        serviceName: [() => [(_, props) => props.serviceName], (serviceName: string): string => serviceName],
        spanName: [() => [(_, props) => props.spanName], (spanName: string): string => spanName],
        heatmapEnabled: [
            (s) => [s.featureFlags],
            (featureFlags: FeatureFlagsSet): boolean => !!featureFlags[FEATURE_FLAGS.TRACING_LATENCY_HEATMAP],
        ],
        histogramData: [
            (s) => [s.rawHistogram],
            (rows: DurationHistogramRow[]): TracingDurationHistogramData => pivotDurationHistogram(rows, dataColorVars),
        ],
        latencyHeatmapData: [
            (s) => [s.rawLatencyHeatmap],
            (rows: LatencyHeatmapRow[]): TracingLatencyHeatmapData => pivotLatencyHeatmap(rows),
        ],
        currentSample: [
            (s) => [s.samples, s.sampleIndex],
            (samples: Span[], sampleIndex: number): Span | null => samples[sampleIndex] ?? null,
        ],
    }),

    listeners(({ actions, values }) => ({
        setDurationSelection: () => {
            actions.fetchSamples()
        },
        setDateRange: () => {
            actions.fetchHistogram()
            actions.fetchSamples()
            actions.fetchStats()
            if (values.chartType === 'heatmap') {
                actions.fetchLatencyHeatmap()
            }
        },
        setChartType: ({ chartType }) => {
            // Fetch lazily on first switch; date-range changes keep it fresh while active. Gate on
            // "has it ever loaded" rather than row count, so an operation with no data doesn't
            // refetch the same empty result every time the user flips back to the heatmap.
            if (chartType === 'heatmap' && !values.latencyHeatmapLoaded && !values.rawLatencyHeatmapLoading) {
                actions.fetchLatencyHeatmap()
            }
        },
        // Feature flags land asynchronously (insights.js /flags), potentially after the URL
        // restore ran with the flag still off. Re-apply a ?chart=heatmap deep link once the
        // flag makes the heatmap available.
        setFeatureFlags: () => {
            if (
                values.heatmapEnabled &&
                router.values.searchParams.chart === 'heatmap' &&
                values.chartType !== 'heatmap'
            ) {
                actions.setChartType('heatmap')
            }
        },
        applyHeatmapBrush: ({ selection }) => {
            // While a date-range refetch is in flight the rendered heatmap still shows the
            // previous query's buckets. Ignore brushes until it settles so the selection maps
            // against the grid on screen rather than a stale one (mirrors the shared scene).
            if (values.rawLatencyHeatmapLoading) {
                return
            }
            const data = values.latencyHeatmapData
            const dateFrom = data.timeBuckets[selection.x.startIndex]
            if (!dateFrom) {
                return
            }
            // y first: setDateRange's listener refetches the samples, so the selection must be
            // in place when it fires. A full-height brush is a plain time zoom — selection kept.
            const spansAllRows = selection.y.startIndex <= 0 && selection.y.endIndex >= data.bucketsNs.length - 1
            if (!spansAllRows) {
                const range = selectionToDurationRange(data.bucketsNs, selection.y.startIndex, selection.y.endIndex)
                if (range) {
                    actions.setDurationSelection(range)
                }
            }
            // Same bucket-edge convention as the main scene: date_to is the start of the bucket
            // after the last selected column, falling back to the current range end.
            actions.setDateRange({
                date_from: dateFrom,
                date_to: data.timeBuckets[selection.x.endIndex + 1] ?? values.dateRange.date_to ?? null,
            })
        },
        loadCurrentSampleTrace: () => {
            const sample = values.currentSample
            if (!sample) {
                return
            }
            // Adjacent samples often share a trace (a child operation hit N times per request) —
            // re-anchor on the already-loaded trace instead of refetching it, but only when that
            // trace actually contains this span: a trace longer than the ±1h lookup window can
            // share a trace_id yet omit spans that fell outside the window of the earlier fetch.
            if (values.sampleTraceSpans.some((span) => span.span_id === sample.span_id)) {
                actions.selectSpan(sample.span_id)
                return
            }
            actions.fetchSampleTrace({ sample })
        },
        fetchSamplesSuccess: ({ samples }) => {
            if (samples.length === 0) {
                actions.selectSpan(null)
                return
            }
            if (values.sampleIndex >= samples.length) {
                // A deep-linked or stale index past the new result set would render a blank sample.
                actions.setSampleIndex(0)
                return
            }
            actions.loadCurrentSampleTrace()
        },
        setSampleIndex: () => {
            // While samples are (re)loading, the index points into the outgoing result set;
            // fetchSamplesSuccess (or fetchSamplesFailure) loads the trace for the resolved index.
            if (values.samplesLoading) {
                return
            }
            actions.loadCurrentSampleTrace()
        },
        fetchSampleTraceSuccess: () => {
            // Anchor the waterfall on the operation's own span — for child operations the
            // sample sits deep in a larger trace and would otherwise open unfocused at the root.
            actions.selectSpan(values.currentSample?.span_id ?? null)
        },
        fetchHistogramFailure: ({ error }) => {
            if (!isUserInitiatedError(error)) {
                toast.error(`Failed to load latency distribution: ${error}`)
            }
        },
        fetchSamplesFailure: ({ error }) => {
            if (isUserInitiatedError(error)) {
                return
            }
            toast.error(`Failed to load sample traces: ${error}`)
            // The failed refetch already blanked the waterfall (see the sampleTraceSpans reducer);
            // reload the trace for the retained samples so the scene stays usable. A stale index
            // (e.g. restored from the URL mid-load) can point past the retained set, resolving to
            // no current sample — clamp it back to 0 first, mirroring fetchSamplesSuccess.
            if (values.sampleIndex >= values.samples.length) {
                actions.setSampleIndex(0)
                return
            }
            actions.loadCurrentSampleTrace()
        },
        fetchStatsFailure: ({ error }) => {
            if (!isUserInitiatedError(error)) {
                toast.error(`Failed to load operation stats: ${error}`)
            }
        },
        fetchSampleTraceFailure: ({ error }) => {
            if (!isUserInitiatedError(error)) {
                toast.error(`Failed to load trace: ${error}`)
            }
        },
    })),

    actionToUrl(({ values }) => {
        const withSelectionParams = (): [string, Record<string, any>, Record<string, any>, { replace: boolean }] => {
            const { min, max, sample, chart, dateRange, ...rest } = router.values.searchParams
            const params: Record<string, any> = { ...rest }
            if (!objectsEqual(values.dateRange, DEFAULT_DATE_RANGE)) {
                params.dateRange = JSON.stringify(values.dateRange)
            }
            if (values.durationSelection) {
                params.min = values.durationSelection.minNs
                params.max = values.durationSelection.maxNs
            }
            if (values.sampleIndex > 0) {
                params.sample = values.sampleIndex
            }
            if (values.chartType !== 'histogram') {
                params.chart = values.chartType
            }
            return [router.values.location.pathname, params, router.values.hashParams, { replace: true }]
        }
        return {
            setDateRange: withSelectionParams,
            setDurationSelection: withSelectionParams,
            setSampleIndex: withSelectionParams,
            setChartType: withSelectionParams,
        }
    }),

    urlToAction(({ actions, values }) => ({
        '/tracing/operation': (_, searchParams) => {
            // Don't gate on the navigation method: the equality guards below make
            // re-processing our own (REPLACE) actionToUrl writes a no-op, while still
            // restoring state when the scene is entered via a router replacement.
            //
            // Restore the date range first: setDateRange resets sampleIndex.
            if (searchParams.dateRange) {
                const dateRange = parseDateRangeParam(searchParams.dateRange)
                // A malformed or wrong-shaped param yields null — keep the current range.
                if (dateRange && !objectsEqual(dateRange, values.dateRange)) {
                    actions.setDateRange(dateRange)
                }
            } else if (!objectsEqual(values.dateRange, DEFAULT_DATE_RANGE)) {
                // No range in the URL — fall back to the default so a back/forward
                // navigation off a non-default URL doesn't keep a stale range.
                actions.setDateRange(DEFAULT_DATE_RANGE)
            }
            const minNs = parseInt(String(searchParams.min), 10)
            const maxNs = parseInt(String(searchParams.max), 10)
            const selection = !isNaN(minNs) && !isNaN(maxNs) ? { minNs, maxNs } : null
            if (!objectsEqual(selection, values.durationSelection)) {
                actions.setDurationSelection(selection)
            }
            const sample = parseInt(String(searchParams.sample), 10)
            if (!isNaN(sample) && sample >= 0 && sample !== values.sampleIndex) {
                actions.setSampleIndex(sample)
            }
            // With the flag off a ?chart=heatmap deep link degrades to the histogram (and never
            // fetches heatmap data — this restore is the only non-toggle entry point).
            const chartFromUrl: OperationChartType =
                searchParams.chart === 'heatmap' && values.heatmapEnabled ? 'heatmap' : 'histogram'
            if (chartFromUrl !== values.chartType) {
                actions.setChartType(chartFromUrl)
            }
        },
    })),

    afterMount(({ actions, props }) => {
        // A malformed link renders the scene's missing-params fallback; don't query for it.
        if (!props.spanName || !props.serviceName) {
            return
        }
        // URL restore may have dispatched these already in this same tick — the loaders'
        // leading breakpoint coalesces the duplicates into one request each.
        actions.fetchHistogram()
        actions.fetchStats()
        actions.fetchSamples()
    }),
])