[refactor] Single source of truth for trace span types

web/oss/src/components/EvalRunDetails/atoms/traces.ts

@@ -3,11 +3,12 @@ import {
     traceEntityAtomFamily,
     transformTracesResponseToTree,
 } from "@agenta/entities/trace"
+import type {TracesResponse} from "@agenta/entities/trace"
 import {uuidToTraceId} from "@agenta/shared/utils"
 import {atomFamily, selectAtom} from "jotai/utils"
 
 import type {TraceData, TraceNode, TraceTree} from "@/oss/lib/evaluations"
-import type {TraceSpanNode, TracesResponse} from "@/oss/services/tracing/types"
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 
 import {resolveInvocationTraceValue} from "../utils/traceValue"
 

web/oss/src/components/SharedDrawers/SessionDrawer/store/sessionDrawerStore.ts

@@ -6,6 +6,7 @@ import {
     transformTracesResponseToTree,
     transformTracingResponse,
 } from "@agenta/entities/trace"
+import type {SpanLink, TracesResponse} from "@agenta/entities/trace"
 import {atom} from "jotai"
 import {atomWithStorage} from "jotai/utils"
 // import {atomWithImmer} from "jotai-immer" // Not using immer for now to keep it simple or use it if complexity grows
@@ -18,7 +19,7 @@ import {AnnotationDto} from "@/oss/lib/hooks/useAnnotations/types"
 import {getNodeById, observabilityTransformer} from "@/oss/lib/traces/observability_helpers"
 import {queryAllAnnotations} from "@/oss/services/annotations/api"
 import {AgentaTreeDTO, TracesWithAnnotations} from "@/oss/services/observability/types"
-import {SpanLink, TraceSpanNode, TracesResponse} from "@/oss/services/tracing/types"
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 import {selectedAppIdAtom} from "@/oss/state/app/selectors/app"
 import {getOrgValues} from "@/oss/state/org"
 import {projectIdAtom} from "@/oss/state/project"

web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceContent/components/OverviewTabItem/index.tsx

@@ -1,6 +1,6 @@
 import {useMemo} from "react"
 
-import {traceSpanMolecule, type TraceSpan as EntityTraceSpan} from "@agenta/entities/trace"
+import {traceSpanMolecule} from "@agenta/entities/trace"
 import {Space} from "antd"
 import {useAtomValue} from "jotai"
 
@@ -33,9 +33,7 @@ const OverviewTabItem = ({
     const entityWithDrillIn = traceSpanMolecule as typeof traceSpanMolecule & {
         drillIn: NonNullable<typeof traceSpanMolecule.drillIn>
     }
-    // OSS TraceSpanNode is the same backend span shape as the entities-package type
-    // the drill-in API expects; align at the boundary, no data is converted.
-    const drillInSpan = activeTrace as unknown as EntityTraceSpan
+    const drillInSpan = activeTrace
     const metaConfig = useAtomValue(spanMetaConfigurationAtomFamily(activeTrace))
     const inputsFromSelectors = useAtomValue(spanDataInputsAtomFamily(activeTrace))
     const outputsFromSelectors = useAtomValue(spanDataOutputsAtomFamily(activeTrace))

web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceTree/assets/spanVisibility.ts

@@ -1,4 +1,6 @@
-import {SpanCategory, StatusCode, TraceSpanNode} from "@/oss/services/tracing/types"
+import type {SpanCategory} from "@agenta/entities/trace"
+
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 
 /**
  * Span tree visibility modes.
@@ -29,14 +31,14 @@ export const SPAN_VISIBILITY_OPTIONS: SpanVisibilityOption[] = [
  * category here to change what counts as a key span across the whole app.
  */
 export const KEY_SPAN_TYPES: ReadonlySet<SpanCategory> = new Set([
-    SpanCategory.AGENT,
-    SpanCategory.LLM,
-    SpanCategory.CHAT,
-    SpanCategory.COMPLETION,
-    SpanCategory.TOOL,
-    SpanCategory.EMBEDDING,
-    SpanCategory.QUERY,
-    SpanCategory.RERANK,
+    "agent",
+    "llm",
+    "chat",
+    "completion",
+    "tool",
+    "embedding",
+    "query",
+    "rerank",
 ])
 
 /**
@@ -78,7 +80,7 @@ export const keySpanRules: KeySpanRule[] = [
     {
         id: "errored-span",
         description: "Spans that ended in an error",
-        test: (node) => node.status_code === StatusCode.STATUS_CODE_ERROR,
+        test: (node) => node.status_code === "STATUS_CODE_ERROR",
     },
 ]
 

web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceTree/index.tsx

@@ -16,7 +16,7 @@ import {useLocalStorage} from "usehooks-ts"
 import CustomTreeComponent from "@/oss/components/CustomUIs/CustomTreeComponent"
 import {filterTree} from "@/oss/components/pages/observability/assets/utils"
 import AvatarTreeContent from "@/oss/components/pages/observability/components/AvatarTreeContent"
-import {StatusCode, TraceSpanNode} from "@/oss/services/tracing/types"
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 import {
     formattedSpanCostAtomFamily,
     formattedSpanLatencyAtomFamily,
@@ -50,7 +50,7 @@ export const TreeContent = ({value, settings}: {value: TraceSpanNode; settings:
                 <Tooltip title={span_name} mouseEnterDelay={0.25}>
                     <Typography.Text
                         className={
-                            status_code === StatusCode.STATUS_CODE_ERROR
+                            status_code === "STATUS_CODE_ERROR"
                                 ? `${treeTitleClass} text-[var(--ag-c-D61010)] font-[500]`
                                 : treeTitleClass
                         }

web/oss/src/components/SharedDrawers/TraceDrawer/store/traceDrawerStore.ts

@@ -3,6 +3,7 @@ import {
     transformTracesResponseToTree,
     transformTracingResponse,
 } from "@agenta/entities/trace"
+import type {SpanLink, TracesResponse} from "@agenta/entities/trace"
 import {atom} from "jotai"
 import {atomWithStorage} from "jotai/utils"
 import {atomWithImmer} from "jotai-immer"
@@ -14,7 +15,7 @@ import {AnnotationDto} from "@/oss/lib/hooks/useAnnotations/types"
 import {getNodeById, observabilityTransformer} from "@/oss/lib/traces/observability_helpers"
 import {queryAllAnnotations} from "@/oss/services/annotations/api"
 import {AgentaTreeDTO, TracesWithAnnotations} from "@/oss/services/observability/types"
-import {SpanLink, TracesResponse, type TraceSpanNode} from "@/oss/services/tracing/types"
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 import {getOrgValues} from "@/oss/state/org"
 import {projectIdAtom} from "@/oss/state/project"
 

web/oss/src/components/pages/observability/assets/constants.ts

@@ -1,3 +1,6 @@
+import type {ComponentType} from "react"
+
+import type {SpanCategory} from "@agenta/entities/trace"
 import {
     ArrowBendRightDownIcon,
     ArrowBendRightUpIcon,
@@ -17,10 +20,10 @@ import {
     TreeStructureIcon,
     TreeViewIcon,
     WarningOctagonIcon,
+    type IconProps,
 } from "@phosphor-icons/react"
 
 import {FilterMenuNode} from "@/oss/components/Filters/types"
-import {SpanCategory} from "@/oss/services/tracing/types"
 
 import {
     COLLECTION_MEMBERSHIP_OPS,
@@ -699,63 +702,67 @@ export const FILTER_COLUMNS: FilterMenuNode[] = [
     // },
 ]
 
-export const spanTypeStyles = {
-    [SpanCategory.AGENT]: {
+/** Keyed by the canonical SpanCategory union, so a new backend category fails the build. */
+export const spanTypeStyles: Record<
+    SpanCategory,
+    {bgColor: string; color: string; icon: ComponentType<IconProps>}
+> = {
+    ["agent"]: {
         bgColor: "var(--ant-blue-1)",
         color: "var(--ant-blue-5)",
         icon: Gear,
     },
-    [SpanCategory.WORKFLOW]: {
+    ["workflow"]: {
         color: "var(--ant-color-text-secondary)",
         bgColor: "var(--ant-color-fill-secondary)",
         icon: TreeStructureIcon,
     },
-    [SpanCategory.CHAIN]: {
+    ["chain"]: {
         bgColor: "var(--ant-blue-1)",
         color: "var(--ant-blue-5)",
         icon: Gear,
     },
-    [SpanCategory.TASK]: {
+    ["task"]: {
         bgColor: "var(--ag-zinc-2)",
         color: "var(--ant-color-text-secondary)",
         icon: TreeStructureIcon,
     },
-    [SpanCategory.TOOL]: {
+    ["tool"]: {
         bgColor: "var(--ant-purple-1)",
         color: "var(--ant-purple-5)",
         icon: Download,
     },
-    [SpanCategory.EMBEDDING]: {
+    ["embedding"]: {
         bgColor: "var(--ant-gold-1)",
         color: "var(--ant-gold-7)",
         icon: LineSegments,
     },
-    [SpanCategory.COMPLETION]: {
+    ["completion"]: {
         bgColor: "var(--ant-cyan-1)",
         color: "var(--ant-cyan-6)",
         icon: Sparkle,
     },
-    [SpanCategory.QUERY]: {
+    ["query"]: {
         bgColor: "var(--ant-gold-1)",
         color: "var(--ant-gold-7)",
         icon: LineSegments,
     },
-    [SpanCategory.CHAT]: {
+    ["chat"]: {
         bgColor: "var(--ant-cyan-1)",
         color: "var(--ant-cyan-6)",
         icon: Sparkle,
     },
-    [SpanCategory.RERANK]: {
+    ["rerank"]: {
         bgColor: "var(--ant-gold-1)",
         color: "var(--ant-gold-7)",
         icon: LineSegments,
     },
-    [SpanCategory.LLM]: {
+    ["llm"]: {
         bgColor: "var(--ant-cyan-1)",
         color: "var(--ant-cyan-6)",
         icon: Sparkle,
     },
-    [SpanCategory.UNDEFINED]: {
+    ["unknown"]: {
         bgColor: "var(--ag-zinc-1)",
         color: "var(--ant-color-text-secondary)",
         icon: TreeStructureIcon,

web/oss/src/components/pages/observability/assets/exportUtils.ts

@@ -1,7 +1,8 @@
+import type {TraceSpan} from "@agenta/entities/trace"
 import {formatCurrency, formatLatency, formatTokenUsage} from "@agenta/shared/utils"
 
 import {formatDay} from "@/oss/lib/helpers/dateTimeHelper"
-import {TraceSpan, TraceSpanNode} from "@/oss/services/tracing/types"
+import type {TraceSpanNode} from "@/oss/services/tracing/types"
 import {
     getAgDataInputs,
     getAgDataOutputs,

web/oss/src/components/pages/observability/components/AvatarTreeContent.tsx

@@ -1,12 +1,11 @@
+import type {SpanCategory} from "@agenta/entities/trace"
 import {Avatar} from "antd"
 
-import {SpanCategory} from "@/oss/services/tracing/types"
-
 import {spanTypeStyles} from "../assets/constants"
 import {AvatarTreeContentProps} from "../assets/types"
 
 export const statusMapper = (span: SpanCategory | null | undefined) => {
-    const {bgColor, color, icon: Icon} = spanTypeStyles[span ?? "undefined"]
+    const {bgColor, color, icon: Icon} = spanTypeStyles[span ?? "unknown"] ?? spanTypeStyles.unknown
     return {
         bgColor,
         color,

web/oss/src/components/pages/observability/components/NodeNameCell.tsx

@@ -1,18 +1,17 @@
 import {memo} from "react"
 
+import type {SpanCategory} from "@agenta/entities/trace"
 import {Typography} from "antd"
 
-import {SpanCategory} from "@/oss/services/tracing/types"
-
 import {spanTypeStyles} from "../assets/constants"
 
 interface Props {
-    name?: string
-    type?: SpanCategory
+    name?: string | null
+    type?: SpanCategory | null
 }
 
 const NodeNameCell = memo(({name, type}: Props) => {
-    const {icon: Icon} = spanTypeStyles[type ?? "undefined"]
+    const {icon: Icon} = spanTypeStyles[type ?? "unknown"] ?? spanTypeStyles.unknown
 
     return (
         <div className="flex items-center gap-1 min-w-0">

web/oss/src/components/pages/observability/components/StatusRenderer.tsx

@@ -1,11 +1,10 @@
+import type {StatusCode} from "@agenta/entities/trace"
 import {CheckCircleOutlined, CloseCircleOutlined, InfoCircleOutlined} from "@ant-design/icons"
 import {Space, Tag, TagProps, Tooltip} from "antd"
 
-import {StatusCode} from "@/oss/services/tracing/types"
-
 export const statusMapper = (status: StatusCode) => {
     switch (status) {
-        case StatusCode.STATUS_CODE_ERROR:
+        case "STATUS_CODE_ERROR":
             return {
                 label: "failure",
                 color: "error",
@@ -26,13 +25,13 @@ const StatusRenderer = ({
     showMore = false,
     tagProps,
 }: {
-    status?: StatusCode
-    message?: string
+    status?: StatusCode | null
+    message?: string | null
     showMore?: boolean
     tagProps?: TagProps
 }) => {
-    const {label, color, icon} = statusMapper(status || StatusCode.STATUS_CODE_UNSET)
-    const errorMsg = status === StatusCode.STATUS_CODE_ERROR ? message : null
+    const {label, color, icon} = statusMapper(status || "STATUS_CODE_UNSET")
+    const errorMsg = status === "STATUS_CODE_ERROR" ? message : null
 
     const {bordered, variant, ...restTagProps} = tagProps || {}
     const resolvedVariant = variant ?? (bordered === false ? "filled" : undefined)

web/oss/src/components/pages/observability/components/TimestampCell.tsx

@@ -5,7 +5,7 @@ import {useAtomValue} from "jotai"
 import {formattedTimestampAtomFamily} from "@/oss/state/newObservability"
 
 interface Props {
-    timestamp?: string | number
+    timestamp?: string | number | null
 }
 
 const TimestampCell = memo(({timestamp}: Props) => {