feat: trace refactor (#1058)

ytallo · web-flow · commit 1be7ecd35adb · 2025-12-11T21:05:36.000-03:00
diff --git a/packages/core/src/call-step-file.ts b/packages/core/src/call-step-file.ts
@@ -70,7 +70,7 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
             processManager.kill()
             const errorMessage = `Step execution timed out after ${timeoutSeconds} seconds`
             logger.error(errorMessage, { step: step.config.name, timeout: timeoutSeconds })
-            await tracer.end({ message: errorMessage })
+            tracer.end({ message: errorMessage })
             trackEvent('step_execution_timeout', {
               stepName: step.config.name,
               traceId,
@@ -94,7 +94,7 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
                   message: err.message,
                 })
 
-                await tracer.end({
+                tracer.end({
                   message: err.message,
                   code: err.code,
                   stack: err.stack?.replace(new RegExp(`${motia.lockedData.baseDir}/`), ''),
@@ -104,35 +104,35 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
 
                 reject(error)
               } else {
-                await tracer.end()
+                tracer.end()
               }
 
               processManager.kill()
             })
             processManager.handler<unknown>('log', async (input: unknown) => logger.log(input))
 
             processManager.handler<StateGetInput, unknown>('state.get', async (input) => {
-              await tracer.stateOperation('get', input)
+              tracer.stateOperation('get', input)
               return motia.state.get(input.traceId, input.key)
             })
 
             processManager.handler<StateSetInput, unknown>('state.set', async (input) => {
-              await tracer.stateOperation('set', { traceId: input.traceId, key: input.key, value: input.value })
+              tracer.stateOperation('set', { traceId: input.traceId, key: input.key, value: input.value })
               return motia.state.set(input.traceId, input.key, input.value)
             })
 
             processManager.handler<StateDeleteInput, unknown>('state.delete', async (input) => {
-              await tracer.stateOperation('delete', input)
+              tracer.stateOperation('delete', input)
               return motia.state.delete(input.traceId, input.key)
             })
 
             processManager.handler<StateClearInput, void>('state.clear', async (input) => {
-              await tracer.stateOperation('clear', input)
+              tracer.stateOperation('clear', input)
               return motia.state.clear(input.traceId)
             })
 
             processManager.handler<StateStreamGetInput>(`state.getGroup`, async (input) => {
-              await tracer.stateOperation('getGroup', input)
+              tracer.stateOperation('getGroup', input)
               return motia.state.getGroup(input.groupId)
             })
 
@@ -149,39 +149,39 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
               const flows = step.config.flows
 
               if (!isAllowedToEmit(step, input.topic)) {
-                await tracer.emitOperation(input.topic, input.data, false)
+                tracer.emitOperation(input.topic, input.data, false)
                 return motia.printer.printInvalidEmit(step, input.topic)
               }
 
-              await tracer.emitOperation(input.topic, input.data, true)
+              tracer.emitOperation(input.topic, input.data, true)
               return motia.eventAdapter.emit({ ...input, traceId, flows, logger, tracer })
             })
 
             Object.entries(streamConfig).forEach(([name, streamFactory]) => {
               const stateStream = streamFactory()
 
               processManager.handler<StateStreamGetInput>(`streams.${name}.get`, async (input) => {
-                await tracer.streamOperation(name, 'get', input)
+                tracer.streamOperation(name, 'get', input)
                 return stateStream.get(input.groupId, input.id)
               })
 
               processManager.handler<StateStreamMutateInput>(`streams.${name}.set`, async (input) => {
-                await tracer.streamOperation(name, 'set', { groupId: input.groupId, id: input.id, data: input.data })
+                tracer.streamOperation(name, 'set', { groupId: input.groupId, id: input.id, data: input.data })
                 return stateStream.set(input.groupId, input.id, input.data)
               })
 
               processManager.handler<StateStreamGetInput>(`streams.${name}.delete`, async (input) => {
-                await tracer.streamOperation(name, 'delete', input)
+                tracer.streamOperation(name, 'delete', input)
                 return stateStream.delete(input.groupId, input.id)
               })
 
               processManager.handler<StateStreamGetInput>(`streams.${name}.getGroup`, async (input) => {
-                await tracer.streamOperation(name, 'getGroup', input)
+                tracer.streamOperation(name, 'getGroup', input)
                 return stateStream.getGroup(input.groupId)
               })
 
               processManager.handler<StateStreamSendInput>(`streams.${name}.send`, async (input) => {
-                await tracer.streamOperation(name, 'send', input)
+                tracer.streamOperation(name, 'send', input)
                 return stateStream.send(input.channel, input.event)
               })
             })
@@ -203,19 +203,19 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
 
               if (code !== 0 && code !== null) {
                 const error = { message: `Process exited with code ${code}`, code }
-                await tracer.end(error)
+                tracer.end(error)
                 trackEvent('step_execution_error', { stepName: step.config.name, traceId, code })
                 reject(`Process exited with code ${code}`)
               } else {
-                await tracer.end()
+                tracer.end()
                 resolve(result)
               }
             })
 
             processManager.onProcessError(async (error) => {
               if (timeoutId) clearTimeout(timeoutId)
               processManager.close()
-              await tracer.end({
+              tracer.end({
                 message: error.message,
                 code: error.code,
                 stack: error.stack,
@@ -236,7 +236,7 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
           })
           .catch(async (error) => {
             if (timeoutId) clearTimeout(timeoutId)
-            await tracer.end({
+            tracer.end({
               message: error.message,
               code: error.code,
               stack: error.stack,
@@ -253,7 +253,7 @@ export const callStepFile = <TData>(options: CallStepFileOptions, motia: Motia):
       })
     } catch (error: unknown) {
       const err = error as Error & { code?: string }
-      await tracer.end({
+      tracer.end({
         message: err.message,
         code: err.code,
         stack: err.stack,
diff --git a/packages/core/src/observability/create-trace.ts b/packages/core/src/observability/create-trace.ts
@@ -7,7 +7,6 @@ export const createTrace = (traceGroup: TraceGroup, step: Step) => {
   const trace: Trace = {
     id,
     name: step.config.name,
-    correlationId: traceGroup.correlationId,
     parentTraceId: traceGroup.id,
     status: 'running',
     startTime: Date.now(),
@@ -16,8 +15,5 @@ export const createTrace = (traceGroup: TraceGroup, step: Step) => {
     events: [],
   }
 
-  traceGroup.metadata.totalSteps++
-  traceGroup.metadata.activeSteps++
-
   return trace
 }
diff --git a/packages/core/src/observability/stream-tracer.ts b/packages/core/src/observability/stream-tracer.ts
@@ -33,39 +33,6 @@ export class StreamTracer implements Tracer {
     this.trace.error = err
 
     await this.manager.updateTrace()
-    await this.recomputeTraceGroupStatus()
-    await this.manager.updateTraceGroup()
-  }
-
-  private async recomputeTraceGroupStatus() {
-    const allTracesFromStorage = await this.manager.getAllTracesForGroup()
-
-    const currentTraceIndex = allTracesFromStorage.findIndex((t) => t.id === this.trace.id)
-    const allTraces = [...allTracesFromStorage]
-
-    if (currentTraceIndex >= 0) {
-      allTraces[currentTraceIndex] = this.trace
-    } else {
-      allTraces.push(this.trace)
-    }
-
-    const completedCount = allTraces.filter((t) => t.status === 'completed').length
-    const failedCount = allTraces.filter((t) => t.status === 'failed').length
-    const runningCount = allTraces.filter((t) => t.status === 'running').length
-
-    this.traceGroup.metadata.completedSteps = completedCount
-    this.traceGroup.metadata.activeSteps = runningCount
-
-    if (failedCount > 0) {
-      this.traceGroup.status = 'failed'
-    } else if (runningCount === 0 && completedCount > 0) {
-      this.traceGroup.status = 'completed'
-      if (!this.traceGroup.endTime) {
-        this.traceGroup.endTime = Date.now()
-      }
-    } else {
-      this.traceGroup.status = 'running'
-    }
   }
 
   async stateOperation(operation: StateOperation, input: unknown) {
@@ -106,9 +73,7 @@ export class StreamTracer implements Tracer {
         lastEvent.data.data = input.data
         lastEvent.maxTimestamp = Date.now()
 
-        this.traceGroup.lastActivity = lastEvent.maxTimestamp
         await this.manager.updateTrace()
-        await this.manager.updateTraceGroup()
 
         return
       }
@@ -133,9 +98,7 @@ export class StreamTracer implements Tracer {
 
   private async addEvent(event: TraceEvent) {
     this.trace.events.push(event)
-    this.traceGroup.lastActivity = event.timestamp
 
     await this.manager.updateTrace()
-    await this.manager.updateTraceGroup()
   }
 }
diff --git a/packages/core/src/observability/trace-manager.ts b/packages/core/src/observability/trace-manager.ts
@@ -9,7 +9,6 @@ export class TraceManager {
     private readonly trace: Trace,
   ) {
     this.updateTrace().catch(() => {})
-    this.updateTraceGroup().catch(() => {})
   }
 
   async updateTrace() {
diff --git a/packages/core/src/observability/tracer.ts b/packages/core/src/observability/tracer.ts
@@ -40,20 +40,14 @@ export class BaseTracerFactory implements TracerFactory {
     const traceGroup: TraceGroup = {
       id: traceId,
       name: step.config.name,
-      lastActivity: Date.now(),
-      metadata: {
-        completedSteps: 0,
-        activeSteps: 0,
-        totalSteps: 0,
-      },
-      correlationId: undefined,
-      status: 'running',
       startTime: Date.now(),
     }
 
     const trace = createTrace(traceGroup, step)
     const manager = new TraceManager(this.traceStream, this.traceGroupStream, traceGroup, trace)
 
+    await manager.updateTraceGroup()
+
     return new StreamTracer(manager, traceGroup, trace, logger)
   }
 
diff --git a/packages/core/src/observability/types.ts b/packages/core/src/observability/types.ts
@@ -2,17 +2,8 @@ import type { StepConfig } from '../types'
 
 export interface TraceGroup {
   id: string
-  correlationId: string | undefined
   name: string
-  status: 'running' | 'completed' | 'failed'
   startTime: number
-  endTime?: number
-  lastActivity: number
-  metadata: {
-    completedSteps: number
-    activeSteps: number
-    totalSteps: number
-  }
 }
 
 export type TraceError = {
diff --git a/packages/e2e/tests/traces/traces.spec.ts b/packages/e2e/tests/traces/traces.spec.ts
@@ -10,7 +10,10 @@ test.describe('Traces tests', () => {
       await workbench.open()
       await workbench.navigateToTraces()
       await tracesPage.verifyTracesInterface()
-      await expect(tracesPage.traceDetailsContainer).toHaveText('Select a trace or trace group to view the timeline')
+      await expect(tracesPage.traceDetailsContainer).toBeVisible()
+      await expect(tracesPage.traceDetailsContainer).toContainText(
+        /Select a trace or trace group to view the timeline|0ms/,
+      )
     })
 
     await test.step('Execute basic tutorial flow', async () => {
diff --git a/plugins/plugin-observability/src/components/trace-timeline.tsx b/plugins/plugin-observability/src/components/trace-timeline.tsx
@@ -2,6 +2,7 @@ import { Button } from '@motiadev/ui'
 import { Minus, Plus } from 'lucide-react'
 import type React from 'react'
 import { memo, useMemo, useState } from 'react'
+import { deriveTraceGroup } from '../hooks/use-derive-trace-group'
 import { useGetEndTime } from '../hooks/use-get-endtime'
 import { useTracesStream } from '../hooks/use-traces-stream'
 import { formatDuration } from '../lib/utils'
@@ -22,9 +23,25 @@ interface TraceTimelineComponentProps {
 const TraceTimelineComponent: React.FC<TraceTimelineComponentProps> = memo(({ groupId }) => {
   useTracesStream()
 
-  const traceGroups = useObservabilityStore((state) => state.traceGroups)
+  const traceGroupMetas = useObservabilityStore((state) => state.traceGroupMetas)
   const traces = useObservabilityStore((state) => state.traces)
-  const group = useMemo(() => traceGroups.find((traceGroup) => traceGroup.id === groupId), [traceGroups, groupId])
+  const group = useMemo(() => {
+    const meta = traceGroupMetas.find((m) => m.id === groupId)
+    if (!meta) return null
+    if (traces.length === 0) {
+      return {
+        ...meta,
+        status: 'running' as const,
+        lastActivity: meta.startTime,
+        metadata: {
+          completedSteps: 0,
+          activeSteps: 0,
+          totalSteps: 0,
+        },
+      }
+    }
+    return deriveTraceGroup(meta, traces)
+  }, [traceGroupMetas, traces, groupId])
 
   const endTime = useGetEndTime(group)
   const [zoom, setZoom] = useState(100)
diff --git a/plugins/plugin-observability/src/hooks/use-derive-trace-group.ts b/plugins/plugin-observability/src/hooks/use-derive-trace-group.ts
@@ -0,0 +1,30 @@
+import type { Trace, TraceGroup, TraceGroupMeta } from '../types/observability'
+
+export const deriveTraceGroup = (meta: TraceGroupMeta, traces: Trace[]): TraceGroup => {
+  const completedSteps = traces.filter((t) => t.status === 'completed').length
+  const failedCount = traces.filter((t) => t.status === 'failed').length
+  const runningCount = traces.filter((t) => t.status === 'running').length
+
+  const status = failedCount > 0 ? 'failed' : runningCount > 0 ? 'running' : 'completed'
+
+  const endTimes = traces.filter((t) => t.endTime).map((t) => t.endTime!)
+  const endTime = status === 'completed' && endTimes.length > 0 ? Math.max(...endTimes) : undefined
+
+  const lastActivity = Math.max(
+    meta.startTime,
+    ...traces.map((t) => t.endTime || t.startTime),
+    ...traces.flatMap((t) => t.events.map((e) => e.timestamp)),
+  )
+
+  return {
+    ...meta,
+    status,
+    endTime,
+    lastActivity,
+    metadata: {
+      completedSteps,
+      activeSteps: runningCount,
+      totalSteps: traces.length,
+    },
+  }
+}
diff --git a/plugins/plugin-observability/src/hooks/use-filtered-trace-groups.ts b/plugins/plugin-observability/src/hooks/use-filtered-trace-groups.ts
@@ -1,17 +1,34 @@
 import { useMemo } from 'react'
 import { useObservabilityStore } from '../stores/use-observability-store'
+import { deriveTraceGroup } from './use-derive-trace-group'
 
 export const useFilteredTraceGroups = () => {
-  const traceGroups = useObservabilityStore((state) => state.traceGroups)
+  const traceGroupMetas = useObservabilityStore((state) => state.traceGroupMetas)
+  const tracesByGroupId = useObservabilityStore((state) => state.tracesByGroupId)
   const search = useObservabilityStore((state) => state.search)
 
-  return useMemo(
-    () =>
-      traceGroups.filter(
-        (group) =>
-          group.name.toLowerCase().includes(search.toLowerCase()) ||
-          group.id.toLowerCase().includes(search.toLowerCase()),
-      ),
-    [traceGroups, search],
-  )
+  return useMemo(() => {
+    const traceGroups = traceGroupMetas.map((meta) => {
+      const traces = tracesByGroupId[meta.id] || []
+      if (traces.length === 0) {
+        return {
+          ...meta,
+          status: 'running' as const,
+          lastActivity: meta.startTime,
+          metadata: {
+            completedSteps: 0,
+            activeSteps: 0,
+            totalSteps: 0,
+          },
+        }
+      }
+      return deriveTraceGroup(meta, traces)
+    })
+
+    return traceGroups.filter(
+      (group) =>
+        group.name.toLowerCase().includes(search.toLowerCase()) ||
+        group.id.toLowerCase().includes(search.toLowerCase()),
+    )
+  }, [traceGroupMetas, tracesByGroupId, search])
 }
diff --git a/plugins/plugin-observability/src/hooks/use-trace-groups-stream.ts b/plugins/plugin-observability/src/hooks/use-trace-groups-stream.ts
diff --git a/plugins/plugin-observability/src/hooks/use-traces-stream.ts b/plugins/plugin-observability/src/hooks/use-traces-stream.ts
diff --git a/plugins/plugin-observability/src/stores/use-observability-store.ts b/plugins/plugin-observability/src/stores/use-observability-store.ts
diff --git a/plugins/plugin-observability/src/types/observability.ts b/plugins/plugin-observability/src/types/observability.ts

Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,6 @@ export class TraceManager {`
`9`	`9`	`private readonly trace: Trace,`
`10`	`10`	`) {`
`11`	`11`	`this.updateTrace().catch(() => {})`
`12`		`- this.updateTraceGroup().catch(() => {})`
`13`	`12`	`}`
`14`	`13`
`15`	`14`	`async updateTrace() {`