refactor: Updated shim.recordConsume to use shim.record and added ability to invoke an after hook with callback args (#2207)

Author: bizob2828

lib/instrumentation/amqplib/amqplib.js

@@ -266,18 +266,20 @@ function wrapModel(shim, Model, promiseMode) {
       destinationName: shim.FIRST,
       callback: setCallback(shim, promiseMode),
       promise: promiseMode,
-      messageHandler: function handleConsumedMessage(shim, fn, name, message) {
+      after: function handleConsumedMessage({ shim, result, args, segment }) {
+        if (!shim.agent.config.message_tracer.segment_parameters.enabled) {
+          shim.logger.trace('Not capturing segment parameters')
+          return
+        }
+
         // the message is the param when using the promised based model
-        message = promiseMode ? message : message[1]
+        const message = promiseMode ? result : args?.[1]
         if (!message) {
           shim.logger.trace('No results from consume.')
           return null
         }
         const parameters = getParametersFromMessage(message)
-
-        const headers = message?.properties?.headers
-
-        return { parameters, headers }
+        shim.copySegmentParameters(segment, parameters)
       }
     })
   )
@@ -312,12 +314,10 @@ function wrapModel(shim, Model, promiseMode) {
  *  Extracts the appropriate messageHandler parameters for the consume method.
  *
  *  @param {Shim} shim instance of shim
- *  @param {object} _consumer not used
- *  @param {string} _name not used
  *  @param {Array} args arguments passed to the consume method
  *  @returns {object} message params
  */
-function describeMessage(shim, _consumer, _name, args) {
+function describeMessage(shim, args) {
   const [message] = args
 
   if (!message?.properties) {

lib/instrumentation/aws-sdk/v3/bedrock.js

@@ -197,8 +197,7 @@ function getBedrockSpec({ commandName }, shim, _original, _name, args) {
   return new RecorderSpec({
     promise: true,
     name: `Llm/${modelType}/Bedrock/${commandName}`,
-    // eslint-disable-next-line max-params
-    after: (shim, _fn, _fnName, err, response, segment) => {
+    after: ({ shim, error: err, result: response, segment }) => {
       const passThroughParams = {
         shim,
         err,

lib/instrumentation/core/inspector.js

@@ -16,7 +16,7 @@ function initialize(agent, inspector, name, shim) {
   shim.wrap(sessionProto, 'post', function wrapPost(shim, fn) {
     return function wrappedPost() {
       const args = shim.argsToArray.apply(shim, arguments)
-      shim.bindCallbackSegment(args, shim.LAST)
+      shim.bindCallbackSegment(null, args, shim.LAST)
       return fn.apply(this, args)
     }
   })

lib/instrumentation/langchain/runnable.js

@@ -61,8 +61,7 @@ function instrumentInvokeChain({ langchain, shim }) {
       return new RecorderSpec({
         name: `${LANGCHAIN.CHAIN}/${fnName}`,
         promise: true,
-        // eslint-disable-next-line max-params
-        after(_shim, _fn, _name, err, output, segment) {
+        after({ error: err, result: output, segment }) {
           recordChatCompletionEvents({
             segment,
             messages: [output],
@@ -97,8 +96,7 @@ function instrumentStream({ langchain, shim }) {
       return new RecorderSpec({
         name: `${LANGCHAIN.CHAIN}/${fnName}`,
         promise: true,
-        // eslint-disable-next-line max-params
-        after(_shim, _fn, _name, err, output, segment) {
+        after({ error: err, result: output, segment }) {
           // Input error occurred which means a stream was not created.
           // Skip instrumenting streaming and create Llm Events from
           // the data we have

lib/instrumentation/langchain/tools.js

@@ -31,8 +31,7 @@ module.exports = function initialize(shim, tools) {
     return new RecorderSpec({
       name: `${LANGCHAIN.TOOL}/${name}`,
       promise: true,
-      // eslint-disable-next-line max-params
-      after(_shim, _fn, _name, err, output, segment) {
+      after({ error: err, result: output, segment }) {
         const metadata = mergeMetadata(instanceMeta, paramsMeta)
         const tags = mergeTags(instanceTags, paramsTags)
         segment.end()

lib/instrumentation/langchain/vectorstore.js

@@ -89,8 +89,7 @@ module.exports = function initialize(shim, vectorstores) {
       return new RecorderSpec({
         name: `${LANGCHAIN.VECTORSTORE}/${fnName}`,
         promise: true,
-        // eslint-disable-next-line max-params
-        after(_shim, _fn, _name, err, output, segment) {
+        after({ error: err, result: output, segment }) {
           if (!output) {
             // If we get an error, it is possible that `output = null`.
             // In that case, we define it to be an empty array.

lib/instrumentation/memcached.js

@@ -76,7 +76,7 @@ module.exports = function initialize(agent, memcached, moduleName, shim) {
     return new OperationSpec({
       name: metacall.type || 'Unknown',
       callback: function wrapCallback(shim, fn, fnName, opSegment) {
-        shim.bindCallbackSegment(metacall, 'callback', opSegment)
+        shim.bindCallbackSegment(null, metacall, 'callback', opSegment)
       },
       parameters
     })

lib/instrumentation/openai.js

@@ -262,8 +262,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
       return new RecorderSpec({
         name: OPENAI.COMPLETION,
         promise: true,
-        // eslint-disable-next-line max-params
-        after(_shim, _fn, _name, err, response, segment) {
+        after({ error: err, result: response, segment }) {
           if (request.stream) {
             instrumentStream({ agent, shim, request, response, segment })
           } else {
@@ -294,8 +293,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
       return new RecorderSpec({
         name: OPENAI.EMBEDDING,
         promise: true,
-        // eslint-disable-next-line max-params
-        after(_shim, _fn, _name, err, response, segment) {
+        after({ error: err, result: response, segment }) {
           addLlmMeta({ agent, segment })
 
           if (!response) {

lib/instrumentation/redis.js

@@ -49,7 +49,7 @@ function registerInternalSendCommand(shim, proto) {
         parameters,
         callback: function bindCallback(shim, _f, _n, segment) {
           if (shim.isFunction(commandObject.callback)) {
-            shim.bindCallbackSegment(commandObject, 'callback', segment)
+            shim.bindCallbackSegment(null, commandObject, 'callback', segment)
           } else {
             const self = this
             commandObject.callback = shim.bindSegment(
@@ -87,9 +87,9 @@ function registerSendCommand(shim, proto) {
       callback: function bindCallback(shim, _f, _n, segment) {
         const last = args[args.length - 1]
         if (shim.isFunction(last)) {
-          shim.bindCallbackSegment(args, shim.LAST, segment)
+          shim.bindCallbackSegment(null, args, shim.LAST, segment)
         } else if (shim.isArray(last) && shim.isFunction(last[last.length - 1])) {
-          shim.bindCallbackSegment(last, shim.LAST, segment)
+          shim.bindCallbackSegment(null, last, shim.LAST, segment)
         }
       }
     })

lib/instrumentation/superagent.js

@@ -50,7 +50,7 @@ function wrapCallback(shim, callback) {
   return function wrappedCallback() {
     const segment = shim.getSegment(this)
     if (segment && segment.transaction.isActive()) {
-      shim.bindCallbackSegment(this, '_callback', segment)
+      shim.bindCallbackSegment(null, this, '_callback', segment)
     }
     return callback.apply(this, arguments)
   }