feat(core): Allow to pass onSuccess to handleCallbackErrors

I've seen a few places where we are wrapping things that could be sync or async, and we want to do something with the return value. By adding an onSuccess handler to `handelCallbackErrors` we can handle this more generically in the future.

Author: mydea

packages/core/src/utils/anthropic-ai/index.ts

@@ -24,6 +24,7 @@ import {
   GEN_AI_SYSTEM_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { buildMethodPath, getFinalOperationName, getSpanOperation, setTokenUsageAttributes } from '../ai/utils';
+import { handleCallbackErrors } from '../handleCallbackErrors';
 import { ANTHROPIC_AI_INTEGRATION_NAME } from './constants';
 import { instrumentStream } from './streaming';
 import type {
@@ -238,7 +239,7 @@ function instrumentMethod<T extends unknown[], R>(
           op: getSpanOperation(methodPath),
           attributes: requestAttributes as Record<string, SpanAttributeValue>,
         },
-        async (span: Span) => {
+        async span => {
           try {
             if (finalOptions.recordInputs && params) {
               addPrivateRequestAttributes(span, params);
@@ -274,27 +275,27 @@ function instrumentMethod<T extends unknown[], R>(
         op: getSpanOperation(methodPath),
         attributes: requestAttributes as Record<string, SpanAttributeValue>,
       },
-      async (span: Span) => {
-        try {
-          if (finalOptions.recordInputs && args[0] && typeof args[0] === 'object') {
-            addPrivateRequestAttributes(span, args[0] as Record<string, unknown>);
-          }
+      span => {
+        if (finalOptions.recordInputs && params) {
+          addPrivateRequestAttributes(span, params);
+        }
 
-          const result = await originalMethod.apply(context, args);
-          addResponseAttributes(span, result, finalOptions.recordOutputs);
-          return result;
-        } catch (error) {
-          captureException(error, {
-            mechanism: {
-              handled: false,
-              type: 'auto.ai.anthropic',
-              data: {
-                function: methodPath,
+        return handleCallbackErrors(
+          () => originalMethod.apply(context, args),
+          error => {
+            captureException(error, {
+              mechanism: {
+                handled: false,
+                type: 'auto.ai.anthropic',
+                data: {
+                  function: methodPath,
+                },
               },
-            },
-          });
-          throw error;
-        }
+            });
+          },
+          () => {},
+          result => addResponseAttributes(span, result as AnthropicAiResponse, finalOptions.recordOutputs),
+        );
       },
     );
   };

packages/core/src/utils/handleCallbackErrors.ts

@@ -1,5 +1,19 @@
 import { isThenable } from '../utils/is';
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function handleCallbackErrors<Fn extends () => Promise<any>, PromiseValue = Awaited<ReturnType<Fn>>>(
+  fn: Fn,
+  onError: (error: unknown) => void,
+  onFinally?: () => void,
+  onSuccess?: (result: PromiseValue) => void,
+): ReturnType<Fn>;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function handleCallbackErrors<Fn extends () => any>(
+  fn: Fn,
+  onError: (error: unknown) => void,
+  onFinally?: () => void,
+  onSuccess?: (result: ReturnType<Fn>) => void,
+): ReturnType<Fn>;
 /**
  * Wrap a callback function with error handling.
  * If an error is thrown, it will be passed to the `onError` callback and re-thrown.
@@ -14,7 +28,13 @@ import { isThenable } from '../utils/is';
 export function handleCallbackErrors<
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   Fn extends () => any,
->(fn: Fn, onError: (error: unknown) => void, onFinally: () => void = () => {}): ReturnType<Fn> {
+  ValueType = ReturnType<Fn>,
+>(
+  fn: Fn,
+  onError: (error: unknown) => void,
+  onFinally: () => void = () => {},
+  onSuccess: (result: ValueType | Awaited<ValueType>) => void = () => {},
+): ValueType {
   let maybePromiseResult: ReturnType<Fn>;
   try {
     maybePromiseResult = fn();
@@ -24,7 +44,7 @@ export function handleCallbackErrors<
     throw e;
   }
 
-  return maybeHandlePromiseRejection(maybePromiseResult, onError, onFinally);
+  return maybeHandlePromiseRejection(maybePromiseResult, onError, onFinally, onSuccess);
 }
 
 /**
@@ -37,12 +57,14 @@ function maybeHandlePromiseRejection<MaybePromise>(
   value: MaybePromise,
   onError: (error: unknown) => void,
   onFinally: () => void,
+  onSuccess: (result: MaybePromise | Awaited<MaybePromise>) => void,
 ): MaybePromise {
   if (isThenable(value)) {
     // @ts-expect-error - the isThenable check returns the "wrong" type here
     return value.then(
       res => {
         onFinally();
+        onSuccess(res);
         return res;
       },
       e => {
@@ -54,5 +76,6 @@ function maybeHandlePromiseRejection<MaybePromise>(
   }
 
   onFinally();
+  onSuccess(value);
   return value;
 }

packages/core/test/lib/utils/handleCallbackErrors.test.ts

@@ -148,4 +148,94 @@ describe('handleCallbackErrors', () => {
       expect(onFinally).toHaveBeenCalledTimes(1);
     });
   });
+
+  describe('onSuccess', () => {
+    it('triggers after successful sync callback', () => {
+      const onError = vi.fn();
+      const onFinally = vi.fn();
+      const onSuccess = vi.fn();
+
+      const fn = vi.fn(() => 'aa');
+
+      const res = handleCallbackErrors(fn, onError, onFinally, onSuccess);
+
+      expect(res).toBe('aa');
+      expect(fn).toHaveBeenCalledTimes(1);
+      expect(onError).not.toHaveBeenCalled();
+      expect(onFinally).toHaveBeenCalledTimes(1);
+      expect(onSuccess).toHaveBeenCalledTimes(1);
+      expect(onSuccess).toHaveBeenCalledWith('aa');
+    });
+
+    it('does not trigger onSuccess after error in sync callback', () => {
+      const error = new Error('test error');
+
+      const onError = vi.fn();
+      const onFinally = vi.fn();
+      const onSuccess = vi.fn();
+
+      const fn = vi.fn(() => {
+        throw error;
+      });
+
+      expect(() => handleCallbackErrors(fn, onError, onFinally, onSuccess)).toThrow(error);
+
+      expect(fn).toHaveBeenCalledTimes(1);
+      expect(onError).toHaveBeenCalledTimes(1);
+      expect(onError).toHaveBeenCalledWith(error);
+      expect(onFinally).toHaveBeenCalledTimes(1);
+      expect(onSuccess).not.toHaveBeenCalled();
+    });
+
+    it('triggers after successful async callback', async () => {
+      const onError = vi.fn();
+      const onFinally = vi.fn();
+      const onSuccess = vi.fn();
+
+      const fn = vi.fn(async () => 'aa');
+
+      const res = handleCallbackErrors(fn, onError, onFinally, onSuccess);
+
+      expect(res).toBeInstanceOf(Promise);
+
+      expect(fn).toHaveBeenCalledTimes(1);
+      expect(onError).not.toHaveBeenCalled();
+      expect(onFinally).not.toHaveBeenCalled();
+
+      const value = await res;
+      expect(value).toBe('aa');
+
+      expect(onFinally).toHaveBeenCalledTimes(1);
+      expect(onSuccess).toHaveBeenCalled();
+      expect(onSuccess).toHaveBeenCalledWith('aa');
+    });
+
+    it('does not trigger onSuccess after error in async callback', async () => {
+      const onError = vi.fn();
+      const onFinally = vi.fn();
+      const onSuccess = vi.fn();
+
+      const error = new Error('test error');
+
+      const fn = vi.fn(async () => {
+        await new Promise(resolve => setTimeout(resolve, 10));
+        throw error;
+      });
+
+      const res = handleCallbackErrors(fn, onError, onFinally, onSuccess);
+
+      expect(res).toBeInstanceOf(Promise);
+
+      expect(fn).toHaveBeenCalledTimes(1);
+      expect(onError).not.toHaveBeenCalled();
+      expect(onFinally).not.toHaveBeenCalled();
+
+      await expect(res).rejects.toThrow(error);
+
+      expect(onError).toHaveBeenCalledTimes(1);
+      expect(onError).toHaveBeenCalledWith(error);
+      expect(onFinally).toHaveBeenCalledTimes(1);
+      expect(onSuccess).not.toHaveBeenCalled();
+    });
+  });
 });

packages/node/src/integrations/tracing/vercelai/instrumentation.ts

@@ -9,7 +9,6 @@ import {
   getActiveSpan,
   getCurrentScope,
   handleCallbackErrors,
-  isThenable,
   SDK_VERSION,
   withScope,
 } from '@sentry/core';
@@ -212,7 +211,7 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
     this._callbacks.forEach(callback => callback());
     this._callbacks = [];
 
-    function generatePatch(originalMethod: (...args: MethodArgs) => unknown) {
+    const generatePatch = (originalMethod: (...args: MethodArgs) => unknown) => {
       return (...args: MethodArgs) => {
         const existingExperimentalTelemetry = args[0].experimental_telemetry || {};
         const isEnabled = existingExperimentalTelemetry.isEnabled;
@@ -237,20 +236,7 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
         };
 
         return handleCallbackErrors(
-          () => {
-            // @ts-expect-error we know that the method exists
-            const result = originalMethod.apply(this, args);
-
-            if (isThenable(result)) {
-              // check for tool errors when the promise resolves, keep the original promise identity
-              result.then(checkResultForToolErrors, () => {});
-              return result;
-            }
-
-            // check for tool errors when the result is synchronous
-            checkResultForToolErrors(result);
-            return result;
-          },
+          () => originalMethod.apply(this, args),
           error => {
             // This error bubbles up to unhandledrejection handler (if not handled before),
             // where we do not know the active span anymore
@@ -260,9 +246,13 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
               addNonEnumerableProperty(error, '_sentry_active_span', getActiveSpan());
             }
           },
+          () => {},
+          result => {
+            checkResultForToolErrors(result);
+          },
         );
       };
-    }
+    };
 
     // Is this an ESM module?
     // https://tc39.es/ecma262/#sec-module-namespace-objects

packages/remix/src/server/errors.ts

@@ -101,11 +101,10 @@ export function errorHandleDocumentRequestFunction(
   const { request, responseStatusCode, responseHeaders, context, loadContext } = requestContext;
 
   return handleCallbackErrors(
-    () => {
-      return origDocumentRequestFunction.call(this, request, responseStatusCode, responseHeaders, context, loadContext);
-    },
+    () => origDocumentRequestFunction.call(this, request, responseStatusCode, responseHeaders, context, loadContext),
     err => {
-      throw err;
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      captureRemixServerException(err, 'HandleDocumentRequestFunction', request);
     },
   );
 }