1
// Copyright (c) 2025 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
// tailStream is going to be invoked multiple times, but we want to wait
5
// to run the test until all executions are done. Collect promises for
6
// each
7
export let spans = new Map();
8
export let invocationPromises = [];
9
 
10
// tail stream handler function used in several STW instrumentation tests.
11
export const testTailHandler = {
12
  tailStream(event, env, ctx) {
13
    // Capture the top-level span ID from the onset event
14
    const topLevelSpanId = event.event.spanId;
15
 
16
    // For each "onset" event, store a promise which we will resolve when
17
    // we receive the equivalent "outcome" event
18
    let resolveFn;
19
    invocationPromises.push(
20
      new Promise((resolve, reject) => {
21
        resolveFn = resolve;
22
      })
23
    );
24
 
25
    return (event) => {
26
      // For spanOpen events, the new span ID is in event.event.spanId
27
      // For other events, they reference an existing span via event.spanContext.spanId
28
      let spanKey = event.invocationId + event.spanContext.spanId;
29
      switch (event.event.type) {
30
        case 'spanOpen':
31
          // spanOpen creates a new span with ID in event.event.spanId
32
          spanKey = event.invocationId + event.event.spanId;
33
          spans.set(spanKey, {
34
            name: event.event.name,
35
          });
36
          break;
37
        case 'attributes': {
38
          let span = spans.get(spanKey);
39
          if (event.spanContext.spanId === topLevelSpanId) {
40
            // top-level JsRpc method name attribute – transform into a span with the given name
41
            const rpcMethodName = event.event.info.find(
42
              (item) => item['name'] === 'jsrpc.method'
43
            ).value;
44
            span = { name: rpcMethodName };
45
            spans.set(spanKey, span);
46
            break;
47
          }
48
 
49
          // attributes references an existing span via spanContext.spanId
50
          if (!span) {
51
            throw new Error(`Attributes event for unknown span: ${spanKey}`);
52
          }
53
          for (let { name, value } of event.event.info) {
54
            span[name] = value;
55
          }
56
          break;
57
        }
58
        case 'spanClose': {
59
          // spanClose references an existing span via spanContext.spanId
60
          let span = spans.get(spanKey);
61
          if (!span) {
62
            throw new Error(`SpanClose event for unknown span: ${spanKey}`);
63
          }
64
          span['closed'] = true;
65
          break;
66
        }
67
        case 'outcome':
68
          resolveFn();
69
          break;
70
      }
71
    };
72
  },
73
};