From d2a45f54ebc624b29a44ea1f4f5218e54d674dc8 Mon Sep 17 00:00:00 2001
From: Bob Evans <robert.evans25@gmail.com>
Date: Thu, 10 Oct 2024 17:23:54 -0400
Subject: [PATCH] feat: Removed references to `segment.transaction` and instead
 get transaction either from tracer or just within the function scope.

---
 lib/instrumentation/aws-sdk/v3/bedrock.js    | 12 ++++++------
 lib/instrumentation/core/http-outbound.js    | 18 ++++++++----------
 lib/instrumentation/kafkajs/consumer.js      |  2 +-
 lib/instrumentation/langchain/runnable.js    |  2 +-
 lib/instrumentation/langchain/vectorstore.js |  2 +-
 lib/instrumentation/openai.js                |  2 +-
 lib/instrumentation/undici.js                |  7 +++----
 lib/instrumentation/when/contextualizer.js   |  4 ++--
 lib/llm-events/aws-bedrock/event.js          |  3 ++-
 lib/shim/shim.js                             | 13 +++++++------
 lib/shim/transaction-shim.js                 |  2 +-
 lib/spans/span-event-aggregator.js           | 10 ++++++----
 lib/spans/span-event.js                      |  8 ++++----
 lib/spans/streaming-span-event-aggregator.js | 10 ++++++----
 lib/util/cat.js                              | 12 +++++-------
 15 files changed, 54 insertions(+), 53 deletions(-)

diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js
index 665b2caf3..fa24f0597 100644
--- a/lib/instrumentation/aws-sdk/v3/bedrock.js
+++ b/lib/instrumentation/aws-sdk/v3/bedrock.js
@@ -70,7 +70,7 @@ function recordEvent({ agent, type, msg }) {
  * @param {object} params input params
  * @param {Agent} params.agent NR agent instance
  * @param {TraceSegment} params.segment active segment
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  */
 function addLlmMeta({ agent, segment, transaction }) {
   agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
@@ -91,9 +91,9 @@ function addLlmMeta({ agent, segment, transaction }) {
  * @param {object} params.segment active segment
  * @param {BedrockCommand} params.bedrockCommand parsed input
  * @param {Error|null} params.err error from request if exists
- * @param params.bedrockResponse
- * @param params.shim
- * @param params.transaction
+ * @param {BedrockResponse} params.bedrockResponse bedrock response
+ * @param {Shim} params.shim shim instance
+ * @param {Transaction} params.transaction active transaction
  */
 function recordChatCompletionMessages({
   agent,
@@ -162,8 +162,8 @@ function recordChatCompletionMessages({
  * @param {object} params.segment active segment
  * @param {BedrockCommand} params.bedrockCommand parsed input
  * @param {Error|null} params.err error from request if exists
- * @param params.bedrockResponse
- * @param params.transaction
+ * @param {BedrockResponse} params.bedrockResponse bedrock response
+ * @param {Transaction} params.transaction active transaction
  */
 function recordEmbeddingMessage({
   agent,
diff --git a/lib/instrumentation/core/http-outbound.js b/lib/instrumentation/core/http-outbound.js
index 43c999c43..1e865217a 100644
--- a/lib/instrumentation/core/http-outbound.js
+++ b/lib/instrumentation/core/http-outbound.js
@@ -17,6 +17,7 @@ const http = require('http')
 const synthetics = require('../../synthetics')
 
 const NAMES = require('../../metrics/names')
+const TraceSegment = require('../../transaction/trace/segment')
 const DEFAULT_HOST = 'localhost'
 const DEFAULT_HTTP_PORT = 80
 const DEFAULT_SSL_PORT = 443
@@ -237,7 +238,7 @@ function assignOutgoingHeaders(currentHeaders, outboundHeaders) {
  * @param {string} params.port port of outbound request
  * @param {string} params.hostname host + port of outbound request
  * @param {TraceSegment} params.segment outbound http segment
- * @param {Transaction} params.config agent config
+ * @param {object} params.config agent config
  * @returns {http.IncomingMessage} request actual http outbound request
  */
 function applySegment({ opts, makeRequest, host, port, hostname, segment, config }) {
@@ -314,15 +315,12 @@ function handleError({ transaction, req, error }) {
 /**
  * Ties the response object to the request segment.
  *
- * @param segment.agent
- * @param {object} segment TraceSegment instance
- * @param {string} hostname host of the HTTP request
- * @param host
- * @param {object} res http.ServerResponse
- * @param segment.segment
- * @param segment.transaction
- * @param segment.host
- * @param segment.res
+ * @param {object} params to function
+ * @param {TraceSegment} params.segment TraceSegment instance
+ * @param {string} params.host hostname of the HTTP request
+ * @param {object} params.res http.ServerResponse
+ * @param {Transaction} params.transaction active transaction
+ * @param {Agent} params.agent agent instance
  */
 function handleResponse({ agent, segment, transaction, host, res }) {
   // Add response attributes for spans
diff --git a/lib/instrumentation/kafkajs/consumer.js b/lib/instrumentation/kafkajs/consumer.js
index a1b566b06..327910582 100644
--- a/lib/instrumentation/kafkajs/consumer.js
+++ b/lib/instrumentation/kafkajs/consumer.js
@@ -100,7 +100,7 @@ function handler({ consumer }) {
    *
    * @param {MessageShim} shim instance of shim
    * @param {Array} args arguments passed to the `eachMessage` function of the `consumer.run`
-   * @param tx
+   * @param {Transaction} tx active transaction
    * @returns {MessageSpec} spec for message handling
    */
   return function messageHandler(shim, args, tx) {
diff --git a/lib/instrumentation/langchain/runnable.js b/lib/instrumentation/langchain/runnable.js
index c5b7c6bbc..3cf1be60e 100644
--- a/lib/instrumentation/langchain/runnable.js
+++ b/lib/instrumentation/langchain/runnable.js
@@ -133,7 +133,7 @@ function instrumentStream({ langchain, shim }) {
  * @param {string} params.request the prompt message
  * @param {object} params.metadata metadata for the call
  * @param {Array} params.tags tags for the call
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  */
 function wrapNextHandler({ shim, output, segment, transaction, request, metadata, tags }) {
   shim.wrap(output, 'next', function wrapIterator(shim, orig) {
diff --git a/lib/instrumentation/langchain/vectorstore.js b/lib/instrumentation/langchain/vectorstore.js
index faea02e2e..3d7ae673d 100644
--- a/lib/instrumentation/langchain/vectorstore.js
+++ b/lib/instrumentation/langchain/vectorstore.js
@@ -27,7 +27,7 @@ const LlmErrorMessage = require('../../llm-events/error-message')
  * @param {TraceSegment} params.segment active segment from vector search
  * @param {string} params.pkgVersion langchain version
  * @param {Error} params.err if it exists
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  */
 function recordVectorSearch({
   request,
diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js
index ef732a449..275497bba 100644
--- a/lib/instrumentation/openai.js
+++ b/lib/instrumentation/openai.js
@@ -110,7 +110,7 @@ function addLlmMeta({ agent, transaction }) {
  * @param {object} params.request chat completion params
  * @param {object} params.response chat completion response
  * @param {boolean} [params.err] err if it exists
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  */
 function recordChatCompletionMessages({
   agent,
diff --git a/lib/instrumentation/undici.js b/lib/instrumentation/undici.js
index adfe2182f..1c6b2c7c5 100644
--- a/lib/instrumentation/undici.js
+++ b/lib/instrumentation/undici.js
@@ -48,7 +48,7 @@ module.exports.unsubscribe = function unsubscribe() {
 }
 
 /**
- * Retrieves the current segment in transaction(parent in our context) from executionAsyncResource
+ * Retrieves the active transaction and current segment in transaction(parent in our context) from executionAsyncResource
  * or from `shim.getSegment()` then adds to the executionAsyncResource for future
  * undici requests within same async context.
  *
@@ -64,7 +64,7 @@ module.exports.unsubscribe = function unsubscribe() {
  * it will just call `shim.getSegment()`
  *
  * @param {Shim} shim instance of shim
- * @returns {TraceSegment} parent segment
+ * @returns {Object} { segment, transaction } active transaction and parent segment
  */
 function getParentContext(shim) {
   const { config } = shim.agent
@@ -115,8 +115,7 @@ function addDTHeaders({ transaction, config, request }) {
  * @param {object} params object to fn
  * @param {Shim} params.shim instance of shim
  * @param {object} params.request undici request object
- * @param {object} params.parentSegment current active, about to be parent of external segment
- * @param params.segment
+ * @param {object} params.segment current active, about to be parent of external segment
  */
 function createExternalSegment({ shim, request, segment }) {
   const url = new URL(request.origin + request.path)
diff --git a/lib/instrumentation/when/contextualizer.js b/lib/instrumentation/when/contextualizer.js
index 7ba49fdc8..80b331e50 100644
--- a/lib/instrumentation/when/contextualizer.js
+++ b/lib/instrumentation/when/contextualizer.js
@@ -87,8 +87,8 @@ function bindChild(ctxlzr, next) {
  *
  * @param {Function} prev previous function in chain
  * @param {Function} next next function in chain
- * @param {object} segment proper segment to bind
- * @param transaction
+ * @param {TraceSegment} segment proper segment to bind
+ * @param {Transaction} transaction active transaction
  * @returns {void}
  */
 Contextualizer.link = function link(prev, next, segment, transaction) {
diff --git a/lib/llm-events/aws-bedrock/event.js b/lib/llm-events/aws-bedrock/event.js
index 946ba40af..6e493058a 100644
--- a/lib/llm-events/aws-bedrock/event.js
+++ b/lib/llm-events/aws-bedrock/event.js
@@ -13,7 +13,8 @@ const { DESTINATIONS } = require('../../config/attribute-filter')
  * @property {object} agent A New Relic agent instance.
  * @property {BedrockCommand} bedrockCommand A parsed invoke command.
  * @property {BedrockResponse} bedrockResponse A parsed response from the API.
- * @property {object} segment The current segment for the trace.
+ * @property {TraceSegment} segment The current segment for the trace.
+ * @property {Transaction} transaction Active transaction
  */
 /**
  * @type {LlmEventParams}
diff --git a/lib/shim/shim.js b/lib/shim/shim.js
index bfcdb50b9..fbe960c70 100644
--- a/lib/shim/shim.js
+++ b/lib/shim/shim.js
@@ -758,7 +758,7 @@ function _hasValidCallbackArg(shim, args, specCallback) {
  * @param {boolean} params.shouldCreateSegment Whether the recorder should create a segment
  * @param {Shim} params.shim instance of shim
  * @param {Function} params.fn function being wrapped
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  * @param {string} params.name name of function being wrapped
  * @returns {shim|promise} Returns a shim or promise with recorder segment and
  * bound callbacks, if applicable
@@ -811,7 +811,7 @@ function _doRecord({ segment, transaction, args, segDesc, shouldCreateSegment, s
  * @param {Spec} params.segDesc Segment descriptor spec
  * @param {Shim} params.shim instance of shim
  * @param {Function} params.fn function being wrapped
- * @param params.transaction
+ * @param {Transaction} params.transaction active transaction
  * @param {string} params.name name of function being wrapped
  * @returns {*} return value of wrapped function
  */
@@ -1010,10 +1010,11 @@ function bindSegment(nodule, property, segment, full) {
 }
 
 /**
- *
- * @param root0
- * @param root0.transaction
- * @param root0.segment
+ * Creates the context before binding to context manager
+ * @param {Object} params to function
+ * @param {Transaction} params.transaction active transaction
+ * @param {object} params.segment either a trace segment or object with transaction and segment
+ * @returns {object} ctx new context
  */
 function setContext({ transaction, segment }) {
   let ctx
diff --git a/lib/shim/transaction-shim.js b/lib/shim/transaction-shim.js
index 565d1efbf..ff0be9e24 100644
--- a/lib/shim/transaction-shim.js
+++ b/lib/shim/transaction-shim.js
@@ -195,7 +195,7 @@ function setTransactionName(name) {
  * @param {TraceSegment} [segment]
  *  The trace segment to associate the header data with. If no segment is
  *  provided then the currently active segment is used.
- * @param transaction
+ * @param {Transaction} transaction active transaction
  * @param {string} [transportType]
  *  The transport type that brought the headers. Usually `HTTP` or `HTTPS`.
  */
diff --git a/lib/spans/span-event-aggregator.js b/lib/spans/span-event-aggregator.js
index 32d5f09b3..07df7ee3e 100644
--- a/lib/spans/span-event-aggregator.js
+++ b/lib/spans/span-event-aggregator.js
@@ -61,10 +61,12 @@ class SpanEventAggregator extends EventAggregator {
   /**
    * Attempts to add the given segment to the collection.
    *
-   * @param {TraceSegment}  segment         - The segment to add.
-   * @param {string}        [parentId] - The GUID of the parent span.
-   * @param isRoot
-   * @returns {boolean} True if the segment was added, or false if it was discarded.
+   * @param {object} params to function
+   * @param {TraceSegment} params.segment The segment to add.
+   * @param {string} [params.parentId] - The GUID of the parent span.
+   * @param {boolean} params.isRoot indicates root segment
+   * @param {Transaction} params.transaction active transaction
+   * @returns {boolean} params.root True if the segment was added, or false if it was discarded.
    */
   addSegment({ segment, transaction, parentId, isRoot }) {
     // Check if the priority would be accepted before creating the event object.
diff --git a/lib/spans/span-event.js b/lib/spans/span-event.js
index 027fc6723..3ddb4dff6 100644
--- a/lib/spans/span-event.js
+++ b/lib/spans/span-event.js
@@ -92,10 +92,10 @@ class SpanEvent {
    * The constructed span event will contain extra data depending on the
    * category of the segment.
    *
-   * @param {TraceSegment}  segment         - The segment to turn into a span event.
-   * @param transaction
-   * @param {?string}       [parentId] - The ID of the segment's parent.
-   * @param isRoot
+   * @param {TraceSegment} segment The segment to turn into a span event.
+   * @param {Transaction} transaction active transaction
+   * @param {?string} [parentId] The ID of the segment's parent.
+   * @param {boolean} isRoot indicates root segment
    * @returns {SpanEvent} The constructed event.
    */
   static fromSegment(segment, transaction, parentId = null, isRoot = false) {
diff --git a/lib/spans/streaming-span-event-aggregator.js b/lib/spans/streaming-span-event-aggregator.js
index 5b2c7731c..d282060be 100644
--- a/lib/spans/streaming-span-event-aggregator.js
+++ b/lib/spans/streaming-span-event-aggregator.js
@@ -100,10 +100,12 @@ class StreamingSpanEventAggregator extends Aggregator {
   /**
    * Attempts to add the given segment to the collection.
    *
-   * @param {TraceSegment}  segment         - The segment to add.
-   * @param {string}        [parentId] - The GUID of the parent span.
-   * @param isRoot
-   * @returns {boolean} True if the segment was added, or false if it was discarded.
+   * @param {object} params to function
+   * @param {TraceSegment} params.segment The segment to add.
+   * @param {string} [params.parentId] - The GUID of the parent span.
+   * @param {boolean} params.isRoot indicates root segment
+   * @param {Transaction} params.transaction active transaction
+   * @returns {boolean} params.root True if the segment was added, or false if it was discarded.
    */
   addSegment({ segment, transaction, parentId, isRoot }) {
     if (!this.started) {
diff --git a/lib/util/cat.js b/lib/util/cat.js
index f89c6a5ba..a4e7f9aa7 100644
--- a/lib/util/cat.js
+++ b/lib/util/cat.js
@@ -273,13 +273,11 @@ cat.parseAppData = function parseAppData(config, obfAppData) {
  * Assigns the CAT id, transaction to segment and adds `transaction_guid` when it exists.
  * It also renames the segment name based on the newly decoded app data when host is present
  *
- * @param appData.appData
- * @param {Array} appData decodes CAT app data
- * @param {TraceSegment} segment
- * @param {string} [host] if host is present it will rename segment with app data and host
- * @param appData.segment
- * @param appData.host
- * @param appData.transaction
+ * @param {object} params to function
+ * @param {Array} params.appData decodes CAT app data
+ * @param {TraceSegment} params.segment segment to assign CAT data
+ * @param {Transaction} params.transaction active transaction
+ * @param {string} [params.host] if host is present it will rename segment with app data and host
  */
 cat.assignCatToSegment = function assignCatToSegment({ appData, segment, host, transaction }) {
   if (!Array.isArray(appData) || typeof appData[0] !== 'string') {