[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"
 
@@ -160,22 +161,14 @@ const buildTraceDataFromEntry = (
         },
     }
 
-    // OSS TracesResponse is the same backend payload shape the entities-package
-    // transform expects; align at the boundary, no data is converted.
-    const spanNodes = transformTracesResponseToTree(
-        scopedResponse as unknown as Parameters<typeof transformTracesResponseToTree>[0],
-    )
+    const spanNodes = transformTracesResponseToTree(scopedResponse)
     if (!spanNodes.length) return null
 
     const flat: TraceNode[] = []
     spanNodes.forEach((span) => {
         const inferredTraceId =
             span.trace_id ?? traceId ?? (span.span_id ? `${span.span_id}-trace` : "trace")
-        // `transformTracesResponseToTree` yields the entities-package TraceSpanNode, while
-        // `convertSpanNodeToTraceNode` is written against the structurally-equivalent OSS
-        // TraceSpanNode (same backend span shape). Align the annotation at the boundary; no
-        // data is converted.
-        convertSpanNodeToTraceNode(span as unknown as TraceSpanNode, inferredTraceId, flat)
+        convertSpanNodeToTraceNode(span, inferredTraceId, flat)
     })
 
     const treeEntry: TraceTree = {

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"
@@ -135,16 +136,10 @@ export const sessionTracesAtom = atom<TraceSpanNode[]>((get) => {
     if (!data) return []
 
     const transformed: TraceSpanNode[] = []
-    // entities-package TraceSpanNode is the same backend span shape as the OSS type;
-    // align the annotation at the boundary, no data is converted.
     if (isTracesResponse(data)) {
-        transformed.push(
-            ...(transformTracingResponse(
-                transformTracesResponseToTree(data),
-            ) as unknown as TraceSpanNode[]),
-        )
+        transformed.push(...transformTracingResponse(transformTracesResponseToTree(data)))
     } else if (isSpansResponse(data)) {
-        transformed.push(...(transformTracingResponse(data.spans) as unknown as TraceSpanNode[]))
+        transformed.push(...transformTracingResponse(data.spans))
     }
 
     const filtred = transformed.filter((node) => node.trace_type !== "annotation")

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/TraceContent/components/TraceTypeHeader/index.tsx

@@ -93,13 +93,9 @@ const TraceTypeHeader = ({
             }
         }
 
-        // OSS TraceSpanNode is the same backend span shape as the entities-package type
-        // these helpers expect; align at the boundary, no data is converted.
-        const agData = extractAgData(activeTrace as unknown as Parameters<typeof extractAgData>[0])
+        const agData = extractAgData(activeTrace)
         const hasExtractableData = Boolean(agData?.inputs || agData?.parameters)
-        const hasApp = hasAppReference(
-            activeTrace as unknown as Parameters<typeof hasAppReference>[0],
-        )
+        const hasApp = hasAppReference(activeTrace)
         const isInvocation = INVOCATION_SPAN_TYPES.has(spanType)
 
         // Invocation spans (workflow, task, agent, chain) represent the unit

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

@@ -188,20 +188,12 @@ const TraceHeader = ({
                 })
                 let candidates: TraceSpanNode[] = []
 
-                // entities-package TraceSpanNode is the same backend span shape as the
-                // OSS type; align the annotation at the boundary, no data is converted.
                 if (isTracesResponse(response)) {
-                    candidates = transformTracingResponse(
-                        transformTracesResponseToTree(response),
-                    ) as unknown as TraceSpanNode[]
+                    candidates = transformTracingResponse(transformTracesResponseToTree(response))
                 } else if (isSpansResponse(response)) {
-                    candidates = transformTracingResponse(
-                        response.spans,
-                    ) as unknown as TraceSpanNode[]
+                    candidates = transformTracingResponse(response.spans)
                 } else if (Array.isArray((response as any)?.spans)) {
-                    candidates = transformTracingResponse(
-                        (response as any).spans,
-                    ) as unknown as TraceSpanNode[]
+                    candidates = transformTracingResponse((response as any).spans)
                 }
 
                 if (!candidates.length) return null

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/openInPlayground.ts

@@ -1,4 +1,3 @@
-import type {TraceSpanNode as EntityTraceSpanNode} from "@agenta/entities/trace"
 import type {OpenFromTraceResult} from "@agenta/playground"
 import {playgroundController} from "@agenta/playground"
 import {atom} from "jotai"
@@ -19,11 +18,6 @@ export type {OpenFromTraceResult}
 export const openTraceInPlaygroundAtom = atom(
     null,
     async (_get, set, activeSpan: TraceSpanNode): Promise<OpenFromTraceResult> => {
-        // OSS TraceSpanNode is the same backend span shape as the entities-package type
-        // the controller expects; align at the boundary, no data is converted.
-        return set(
-            playgroundController.actions.openFromTrace,
-            activeSpan as unknown as EntityTraceSpanNode,
-        )
+        return set(playgroundController.actions.openFromTrace, activeSpan)
     },
 )

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">