fix: correctly interrupt from draft trace

Author: niieani

src/v3/Trace.ts

@@ -94,6 +94,14 @@ interface OnEnterComplete<RelationSchemasT>
   transitionToState: 'complete'
 }
 
+export type FinalTransition<RelationSchemasT> = (
+  | (OnEnterInterrupted &
+      Omit<FinalState<RelationSchemasT>, 'transitionFromState'>)
+  | OnEnterComplete<RelationSchemasT>
+) & {
+  lastRelevantSpanAndAnnotation: SpanAndAnnotation<RelationSchemasT> | undefined
+}
+
 interface OnEnterWaitingForInteractive {
   transitionToState: 'waiting-for-interactive'
   transitionFromState: NonTerminalTraceStates
@@ -325,27 +333,25 @@ export class TraceStateMachine<
           }
         }
 
+        // add into span buffer
+        this.#draftBuffer.push(spanAndAnnotation)
+
         // if the entry matches any of the interruptOnSpans criteria,
-        // transition to complete state with the 'matched-on-interrupt' interruptionReason
+        // transition to interrupted state with the correct interruptionReason
         if (this.#context.definition.interruptOnSpans) {
           for (const doesSpanMatch of this.#context.definition
             .interruptOnSpans) {
             if (doesSpanMatch(spanAndAnnotation, this.#context)) {
-              // TODO: do we want to record this span? Nothing else is being recorded at this point and is instead going into the buffer
               return {
-                transitionToState: 'complete',
+                transitionToState: 'interrupted',
                 interruptionReason: doesSpanMatch.requiredSpan
                   ? 'matched-on-required-span-with-error'
                   : 'matched-on-interrupt',
-                lastRequiredSpanAndAnnotation: this.lastRequiredSpan,
-                completeSpanAndAnnotation: this.completeSpan,
               }
             }
           }
         }
 
-        // else, add into span buffer
-        this.#draftBuffer.push(spanAndAnnotation)
         return undefined
       },
 
@@ -1241,30 +1247,14 @@ export class Trace<
         transition.transitionToState === 'interrupted' ||
         transition.transitionToState === 'complete'
       ) {
-        const endOfOperationSpan =
-          (transition.transitionToState === 'complete' &&
-            (transition.cpuIdleSpanAndAnnotation ??
-              transition.completeSpanAndAnnotation)) ||
-          lastRelevantSpanAndAnnotation
-
         const traceRecording = createTraceRecording(
           {
             definition: this.definition,
-            // only keep items captured until the endOfOperationSpan
-            recordedItems: endOfOperationSpan
-              ? new Set(
-                  [...this.recordedItems].filter(
-                    (item) =>
-                      item.span.startTime.now + item.span.duration <=
-                      endOfOperationSpan.span.startTime.now +
-                        endOfOperationSpan.span.duration,
-                  ),
-                )
-              : this.recordedItems,
+            recordedItems: this.recordedItems,
             input: this.input,
             recordedItemsByLabel: this.recordedItemsByLabel,
           },
-          transition,
+          { ...transition, lastRelevantSpanAndAnnotation },
         )
         this.onEnd(traceRecording)
 

src/v3/TraceManagerWithDraft.test.ts

@@ -396,7 +396,7 @@ describe('TraceManager', () => {
       // prettier-ignore
       const { spans } = getSpansFromTimeline<TicketIdRelationSchemasFixture>`
         Events: ${Render('start', 0)}-----${Render('interrupt', 0)}-----${Render('end', 0)}
-        Time:   ${0}                      ${100}                      ${200}
+        Time:   ${0}                      ${100}                        ${200}
         `
       processSpans(spans, traceManager)
 
@@ -406,12 +406,12 @@ describe('TraceManager', () => {
         AnyPossibleReportFn<TicketIdRelationSchemasFixture>
       >[0] = reportFn.mock.calls[0][0]
 
-      // there are NO entries in the report because this trace was interrupted before transitioning from draft to active
-      expect(report.entries).toHaveLength(0)
+      // this trace was interrupted before transitioning from draft to active, so the only available entries are up to the interrupt
+      expect(report.entries).toHaveLength(2)
       expect(report.name).toBe('ticket.interrupt-on-basic-operation')
       expect(report.duration).toBeNull()
-      expect(report.status).toBe('interrupted')
       expect(report.interruptionReason).toBe('matched-on-interrupt')
+      expect(report.status).toBe('interrupted')
     })
 
     it('interrupts a draft trace when interrupt() is called with error', () => {

src/v3/recordingComputeUtils.test.ts

@@ -155,6 +155,8 @@ describe('recordingComputeUtils', () => {
         {
           transitionFromState: 'active',
           interruptionReason: 'timeout',
+          transitionToState: 'interrupted',
+          lastRelevantSpanAndAnnotation: undefined,
         },
       )
 
@@ -675,7 +677,11 @@ describe('recordingComputeUtils', () => {
           },
           recordedItemsByLabel: {},
         },
-        { transitionFromState: 'active' },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'complete',
+          lastRelevantSpanAndAnnotation: undefined,
+        },
       )
 
       expect(recording.computedRenderBeaconSpans['test-component']).toEqual({
@@ -708,7 +714,11 @@ describe('recordingComputeUtils', () => {
           },
           recordedItemsByLabel: {},
         },
-        { transitionFromState: 'active' },
+        {
+          transitionFromState: 'active',
+          transitionToState: 'complete',
+          lastRelevantSpanAndAnnotation: undefined,
+        },
       )
 
       // Verify the variant is included in the recording

src/v3/recordingComputeUtils.ts

@@ -3,7 +3,7 @@ import { ensureMatcherFnOrSpecialToken } from './ensureMatcherFn'
 import { type SpanMatcherFn } from './matchSpan'
 import type { SpanAndAnnotation } from './spanAnnotationTypes'
 import type { ActiveTraceInput, DraftTraceInput } from './spanTypes'
-import type { FinalState } from './Trace'
+import type { FinalState, FinalTransition } from './Trace'
 import type { TraceRecording } from './traceRecordingTypes'
 import type { SpecialEndToken, SpecialStartToken, TraceContext } from './types'
 
@@ -223,7 +223,7 @@ function getComputedRenderBeaconSpans<
   RelationSchemasT,
   const VariantsT extends string,
 >(
-  recordedItems: ReadonlySet<SpanAndAnnotation<RelationSchemasT>>,
+  recordedItems: readonly SpanAndAnnotation<RelationSchemasT>[],
   input: ActiveTraceInput<RelationSchemasT[SelectedRelationNameT], VariantsT>,
 ): TraceRecording<
   SelectedRelationNameT,
@@ -380,18 +380,35 @@ export function createTraceRecording<
   context: TraceContext<SelectedRelationNameT, RelationSchemasT, VariantsT>,
   {
     transitionFromState,
+    transitionToState,
     interruptionReason,
     cpuIdleSpanAndAnnotation,
     completeSpanAndAnnotation,
     lastRequiredSpanAndAnnotation,
-  }: FinalState<RelationSchemasT>,
+    lastRelevantSpanAndAnnotation,
+  }: FinalTransition<RelationSchemasT>,
 ): TraceRecording<SelectedRelationNameT, RelationSchemasT> {
   const { definition, recordedItems, input } = context
   const { id, relatedTo, variant } = input
   const { name } = definition
+
+  const endOfOperationSpan =
+    (transitionToState === 'complete' &&
+      (cpuIdleSpanAndAnnotation ?? completeSpanAndAnnotation)) ||
+    lastRelevantSpanAndAnnotation
+
+  // only keep items captured until the endOfOperationSpan or if not available, the lastRelevantSpan
+  const recordedItemsArray = endOfOperationSpan
+    ? [...recordedItems].filter(
+        (item) =>
+          item.span.startTime.now + item.span.duration <=
+          endOfOperationSpan.span.startTime.now +
+            endOfOperationSpan.span.duration,
+      )
+    : [...recordedItems.values()]
+
   // CODE CLEAN UP TODO: let's get this information (wasInterrupted) from up top (in FinalState)
-  const wasInterrupted =
-    interruptionReason && transitionFromState !== 'waiting-for-interactive'
+  const wasInterrupted = transitionToState === 'interrupted'
   const computedSpans = !wasInterrupted
     ? getComputedSpans(context, {
         completeSpanAndAnnotation,
@@ -401,11 +418,9 @@ export function createTraceRecording<
   const computedValues = !wasInterrupted ? getComputedValues(context) : {}
   const computedRenderBeaconSpans =
     !wasInterrupted && isActiveTraceInput(input)
-      ? getComputedRenderBeaconSpans(recordedItems, input)
+      ? getComputedRenderBeaconSpans(recordedItemsArray, input)
       : {}
 
-  const recordedItemsArray = [...recordedItems.values()]
-
   const anyNonSuppressedErrors = recordedItemsArray.some(
     (spanAndAnnotation) =>
       spanAndAnnotation.span.status === 'error' &&

src/v3/requiredSpanWithErrorStatus.ts

@@ -1,5 +1,6 @@
 import type { SpanMatcherFn, SpanMatcherTags } from './matchSpan'
 
+// tag matcher with a special, internal matcher tag, and match on span.status === 'error'
 export function requiredSpanWithErrorStatus<
   const SelectedRelationNameT extends keyof RelationSchemasT,
   const RelationSchemasT,