From 9287299aaeb7aa91af8e1d26ce615de12748b284 Mon Sep 17 00:00:00 2001
From: sepeterson <10458078+sepeterson@users.noreply.github.com>
Date: Fri, 21 Nov 2025 13:05:35 -0600
Subject: [PATCH] MOB-991 reuse attribute logic, set attribute for timeouts

---
 src/stores/createFirebaseTraceSlice.ts | 36 ++++++++++++++------------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/src/stores/createFirebaseTraceSlice.ts b/src/stores/createFirebaseTraceSlice.ts
index 519cc045b..40366e284 100644
--- a/src/stores/createFirebaseTraceSlice.ts
+++ b/src/stores/createFirebaseTraceSlice.ts
@@ -11,6 +11,7 @@ export enum FIREBASE_TRACES {
 
 export enum FIREBASE_TRACE_ATTRIBUTES {
   ONLINE = "online",
+  DID_TIMEOUT = "did_timeout",
 }
 
 interface TraceData {
@@ -20,10 +21,23 @@ interface TraceData {
 
 const TRACE_TIMEOUT = 10000;
 
+const applyTraceAttributes = (
+  trace: FirebasePerformanceTypes.Trace,
+  attributes: Record<string, string>
+): void => {
+  try {
+    Object.entries( attributes ).forEach( ( [key, value] ) => {
+      trace.putAttribute( key, value );
+    } );
+  } catch ( error ) {
+    logger.error( "Error setting firebase trace attributes", error );
+  }
+};
+
 export interface FirebaseTraceSlice {
   activeFirebaseTraces: Record<string, TraceData>;
-  startFirebaseTrace: ( traceId: string ) => Promise<void>;
-  stopFirebaseTrace: ( traceId: string ) => Promise<void>;
+  startFirebaseTrace: ( traceId: string, attributes: Record<string, string> ) => Promise<void>;
+  stopFirebaseTrace: ( traceId: string, attributes: Record<string, string> ) => Promise<void>;
 }
 
 const createFirebaseTraceSlice: StateCreator<FirebaseTraceSlice>
@@ -35,16 +49,10 @@ const createFirebaseTraceSlice: StateCreator<FirebaseTraceSlice>
     const trace = await perf.startTrace( traceId );
 
     const timeoutId = setTimeout( () => {
-      get().stopFirebaseTrace( traceId );
+      get().stopFirebaseTrace( traceId, { [FIREBASE_TRACE_ATTRIBUTES.DID_TIMEOUT]: "true" } );
     }, TRACE_TIMEOUT );
 
-    try {
-      Object.entries( attributes ).forEach( ( [key, value] ) => {
-        trace.putAttribute( key, value );
-      } );
-    } catch ( error ) {
-      logger.error( "Error setting firebase trace attributes on start", error );
-    }
+    applyTraceAttributes( trace, attributes );
 
     set( state => ( {
       activeFirebaseTraces: {
@@ -61,13 +69,7 @@ const createFirebaseTraceSlice: StateCreator<FirebaseTraceSlice>
       clearTimeout( traceData.timeoutId );
 
       const { trace } = traceData;
-      try {
-        Object.entries( attributes ).forEach( ( [key, value] ) => {
-          trace.putAttribute( key, value );
-        } );
-      } catch ( error ) {
-        logger.error( "Error setting firebase trace attributes on stop", error );
-      }
+      applyTraceAttributes( trace, attributes );
 
       await trace.stop();