fix: expose getCurrentTrace on the Tracer

Author: niieani

src/v3/Trace.ts

@@ -1173,7 +1173,7 @@ export class Trace<
   onEnd(
     traceRecording: TraceRecording<SelectedRelationNameT, RelationSchemasT>,
   ): void {
-    this.traceUtilities.cleanupCurrentTrace(this)
+    this.traceUtilities.onEndTrace(this)
     this.traceUtilities.reportFn(traceRecording, this)
   }
 

src/v3/TraceManager.ts

@@ -9,6 +9,7 @@ import type { Span } from './spanTypes'
 import type { AllPossibleTraces } from './Trace'
 import { Tracer } from './Tracer'
 import type {
+  AllPossibleTraceContexts,
   CompleteTraceDefinition,
   ComputedValueDefinitionInput,
   RelationSchemasBase,
@@ -28,12 +29,11 @@ export class TraceManager<
   private currentTrace: AllPossibleTraces<RelationSchemasT> | undefined =
     undefined
 
-  get currentTracerContext() {
+  get currentTracerContext():
+    | AllPossibleTraceContexts<RelationSchemasT, string>
+    | undefined {
     if (!this.currentTrace) return undefined
-    return {
-      definition: this.currentTrace.definition,
-      input: this.currentTrace.input,
-    }
+    return this.currentTrace
   }
 
   constructor(
@@ -52,7 +52,7 @@ export class TraceManager<
         }
         this.currentTrace = newTrace
       },
-      cleanupCurrentTrace: (traceToCleanUp) => {
+      onEndTrace: (traceToCleanUp) => {
         if (traceToCleanUp === this.currentTrace) {
           this.currentTrace = undefined
         }

src/v3/Tracer.ts

@@ -8,6 +8,7 @@ import {
   type CompleteTraceDefinition,
   type ComputedSpanDefinitionInput,
   type ComputedValueDefinitionInput,
+  type DraftTraceContext,
   type RelationSchemasBase,
   type TraceDefinitionModifications,
   type TraceManagerUtilities,
@@ -219,6 +220,16 @@ export class Tracer<
     trace.transitionDraftToActive(inputAndDefinitionModifications)
   }
 
+  getCurrentTrace = ():
+    | DraftTraceContext<SelectedRelationNameT, RelationSchemasT, VariantsT>
+    | undefined => {
+    const trace = this.traceUtilities.getCurrentTrace()
+    if (!trace || trace.sourceDefinition !== this.definition) {
+      return undefined
+    }
+    return trace
+  }
+
   defineComputedSpan = (
     definition: ComputedSpanDefinitionInput<
       SelectedRelationNameT,

src/v3/types.ts

@@ -171,9 +171,7 @@ export interface TraceManagerUtilities<
     newTrace: AllPossibleTraces<RelationSchemasT>,
     reason: TraceReplaceInterruptionReason,
   ) => void
-  cleanupCurrentTrace: (
-    traceToCleanUp: AllPossibleTraces<RelationSchemasT>,
-  ) => void
+  onEndTrace: (traceToCleanUp: AllPossibleTraces<RelationSchemasT>) => void
   getCurrentTrace: () => AllPossibleTraces<RelationSchemasT> | undefined
 }
 
@@ -574,6 +572,17 @@ export interface DraftTraceContext<
   >
 }
 
+export type AllPossibleTraceContexts<
+  RelationSchemasT,
+  VariantsT extends string,
+> = {
+  [SelectedRelationNameT in keyof RelationSchemasT]: DraftTraceContext<
+    SelectedRelationNameT,
+    RelationSchemasT,
+    VariantsT
+  >
+}[keyof RelationSchemasT]
+
 export interface TraceContext<
   SelectedRelationNameT extends keyof RelationSchemasT,
   RelationSchemasT,