feat(core): Allow to pass onSuccess to handleCallbackErrors

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,30 @@
 import { isThenable } from '../utils/is';
 
+/* eslint-disable */
+// Vendor "Awaited" in to be TS 3.8 compatible
+type AwaitedPromise<T> = T extends null | undefined
+  ? T // special case for `null | undefined` when not in `--strictNullChecks` mode
+  : T extends object & { then(onfulfilled: infer F, ...args: infer _): any } // `await` only unwraps object types with a callable `then`. Non-object types are not unwrapped
+    ? F extends (value: infer V, ...args: infer _) => any // if the argument to `then` is callable, extracts the first argument
+      ? V // normally this would recursively unwrap, but this is not possible in TS3.8
+      : never // the argument to `then` was not callable
+    : T; // non-object or non-thenable
+/* eslint-enable */
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function handleCallbackErrors<Fn extends () => Promise<any>, PromiseValue = AwaitedPromise<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 +39,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 | AwaitedPromise<ValueType>) => void = () => {},
+): ValueType {
   let maybePromiseResult: ReturnType<Fn>;
   try {
     maybePromiseResult = fn();
@@ -24,7 +55,7 @@ export function handleCallbackErrors<
     throw e;
   }
 
-  return maybeHandlePromiseRejection(maybePromiseResult, onError, onFinally);
+  return maybeHandlePromiseRejection(maybePromiseResult, onError, onFinally, onSuccess);
 }
 
 /**
@@ -37,12 +68,14 @@ function maybeHandlePromiseRejection<MaybePromise>(
   value: MaybePromise,
   onError: (error: unknown) => void,
   onFinally: () => void,
+  onSuccess: (result: MaybePromise | AwaitedPromise<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 +87,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,
   getClient,
   handleCallbackErrors,
-  isThenable,
   SDK_VERSION,
   withScope,
 } from '@sentry/core';
@@ -238,19 +237,7 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
           };
 
           return handleCallbackErrors(
-            () => {
-              const result = Reflect.apply(target, thisArg, 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;
-            },
+            () => Reflect.apply(target, thisArg, args),
             error => {
               // This error bubbles up to unhandledrejection handler (if not handled before),
               // where we do not know the active span anymore
@@ -260,6 +247,10 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
                 addNonEnumerableProperty(error, '_sentry_active_span', getActiveSpan());
               }
             },
+            () => {},
+            result => {
+              checkResultForToolErrors(result);
+            },
           );
         },
       });