From cdf7bc15e91080a400a87811cc269d5832ee683a Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Tue, 8 Sep 2026 16:29:22 -0700
Subject: [PATCH 01/14] refactor(gax): support promises and streams in
 traceAttempt and prevent premature span endings

---
 .../gax/src/observability/TracerHelper.ts     | 124 ++++++--
 core/packages/gax/test/unit/tracerHelper.ts   | 281 ++++++++++++++++++
 2 files changed, 387 insertions(+), 18 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index d3077e8754f..b9309e0bc53 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import {EventEmitter} from 'events';
 import {Span, trace, Tracer} from '@opentelemetry/api';
 
 /**
@@ -67,6 +68,53 @@ export function getGaxTracer(): Tracer {
   return trace.getTracer('google-gax');
 }
 
+/**
+ * Manages span lifecycle for Promise-based operations.
+ *
+ * @param {T} promise - The promise returned from the traced operation.
+ * @param {(err: unknown) => void} recordError - Callback to record errors on the span.
+ * @param {() => void} endSpan - Callback to end the span idempotently.
+ */
+export function handlePromise<T>(
+  promise: T,
+  recordError: (err: unknown) => void,
+  endSpan: () => void,
+): void {
+  Promise.resolve(promise)
+    .then(() => {
+      endSpan();
+      return null;
+    })
+    .catch(err => {
+      recordError(err);
+      endSpan();
+    });
+}
+
+/**
+ * Manages span lifecycle for Stream-based operations.
+ *
+ * @param {EventEmitter} stream - The stream returned from the traced operation.
+ * @param {(err: unknown) => void} recordError - Callback to record errors on the span.
+ * @param {() => void} endSpan - Callback to end the span idempotently.
+ */
+export function handleStream(
+  stream: EventEmitter,
+  recordError: (err: unknown) => void,
+  endSpan: () => void,
+): void {
+  stream.on('error', (err: unknown) => {
+    recordError(err);
+    endSpan();
+  });
+  stream.on('end', () => {
+    endSpan();
+  });
+  stream.on('close', () => {
+    endSpan();
+  });
+}
+
 /**
  * Executes a function within an active OpenTelemetry span, populating standard
  * GCP telemetry attributes and recording errors/exceptions if thrown.
@@ -74,16 +122,33 @@ export function getGaxTracer(): Tracer {
  * @template T
  * @param {DynamicTraceContext} dynamicArgs - Dynamic trace context for the RPC call.
  * @param {StaticTraceContext} staticArgs - Static trace context for the client library.
- * @param {() => Promise<T>} fn - The asynchronous operation to trace.
- * @returns {Promise<T>} The result of the traced operation.
+ * @param {() => T} fn - The operation to trace.
+ * @param {boolean | 'promise' | 'stream'} [isStream=false] - Whether the operation is a stream or a promise.
+ * @returns {T} The result of the traced operation.
  */
-export async function traceAttempt<T = unknown>(
+export function traceAttempt<T extends EventEmitter>(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => T,
+  isStream: true | 'stream',
+): T;
+export function traceAttempt<T>(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => T,
+  isStream?: boolean | 'promise' | 'stream',
+): T;
+export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => Promise<T>,
-): Promise<T> {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  fn: () => any,
+  isStream: boolean | 'promise' | 'stream' = false,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+): any {
+  const isStreamCall = isStream === true || isStream === 'stream';
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
-  return getGaxTracer().startActiveSpan(spanName, {}, async (span: Span) => {
+  return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
     span.setAttributes({
       'gcp.client.service': staticArgs.gcpClientService,
       'gcp.client.version': staticArgs.gcpVersion,
@@ -93,22 +158,45 @@ export async function traceAttempt<T = unknown>(
       'gcp.method.type': dynamicArgs.rpcType,
     });
 
+    let spanEnded = false;
+    const endSpan = () => {
+      if (!spanEnded) {
+        spanEnded = true;
+        span.end();
+      }
+    };
+
+    const recordError = (e: unknown) => {
+      if (e instanceof Error) {
+        span.setAttributes({
+          'error.message': e.message,
+          'error.type': e.constructor?.name ?? e.name,
+        });
+        span.recordException(e);
+        if (e.name) {
+          span.setAttribute('exception.type', e.name);
+        }
+      } else {
+        const message = String(e);
+        span.setAttributes({
+          'error.message': message,
+        });
+        span.recordException(message);
+      }
+    };
+
     try {
-      const result = await fn();
+      const result = fn();
+      if (isStreamCall) {
+        handleStream(result, recordError, endSpan);
+      } else {
+        handlePromise(result, recordError, endSpan);
+      }
       return result;
     } catch (e) {
-      const err = e as Error;
-      span.setAttributes({
-        'error.message': err.message,
-        'error.type': err.constructor?.name ?? err.name,
-      });
-      span.recordException(err);
-      if (err.name) {
-        span.setAttribute('exception.type', err.name);
-      }
+      recordError(e);
+      endSpan();
       throw e;
-    } finally {
-      span.end();
     }
   });
 }
diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index 06515c2f0c9..cdf3976d5f2 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -15,10 +15,13 @@
  */
 
 import * as assert from 'assert';
+import {EventEmitter} from 'events';
 import {describe, it, beforeEach, afterEach} from 'mocha';
 import {
   getGaxTracer,
   traceAttempt,
+  handlePromise,
+  handleStream,
   DynamicTraceContext,
   StaticTraceContext,
 } from '../../src/observability/TracerHelper';
@@ -163,5 +166,283 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans.length, 1);
       assert.strictEqual(spans[0].attributes['gcp.method.type'], 'http');
     });
+
+    it('manages span lifetime for resolved promises', async () => {
+      const result = await traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => Promise.resolve('async-result'),
+      );
+      assert.strictEqual(result, 'async-result');
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(spans[0].events.length, 0);
+    });
+
+    it('does not end span prematurely for pending asynchronous promises', async () => {
+      let resolvePromise: (val: string) => void;
+      const asyncPromise = new Promise<string>(resolve => {
+        resolvePromise = resolve;
+      });
+
+      const resultPromise = traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => asyncPromise,
+      );
+
+      // Verify the span is NOT closed while the promise is pending
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      resolvePromise!('success');
+      const result = await resultPromise;
+      assert.strictEqual(result, 'success');
+
+      // Span should only be ended after resolution
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('supports custom thenables that do not inherit from Promise', async () => {
+      const customThenable = {
+        then(onfulfilled?: (val: unknown) => void) {
+          setTimeout(() => {
+            if (onfulfilled) {
+              onfulfilled('custom-result');
+            }
+          }, 10);
+        },
+      };
+
+      const result = traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => customThenable,
+      );
+      assert.strictEqual(result, customThenable);
+
+      const initialSpans = harness.getSpans('google-gax');
+      assert.strictEqual(initialSpans.length, 0);
+
+      await new Promise(resolve => setTimeout(resolve, 25));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('does not end span prematurely until asynchronous promise rejects', async () => {
+      let rejectPromise: (err: Error) => void;
+      const asyncPromise = new Promise((_resolve, reject) => {
+        rejectPromise = reject;
+      });
+
+      const error = new Error('async promise failure');
+      void traceAttempt(dynamicArgs, staticArgs, () => asyncPromise);
+
+      // Verify the span is NOT closed while the promise is pending
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      rejectPromise!(error);
+      await new Promise(resolve => setTimeout(resolve, 15));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(
+        spans[0].attributes['error.message'],
+        'async promise failure',
+      );
+      assert.strictEqual(spans[0].events.length, 1);
+    });
+
+    it('does not end span prematurely while stream is active and emitting data', () => {
+      const emitter = new EventEmitter();
+      const result = traceAttempt(dynamicArgs, staticArgs, () => emitter, true);
+      assert.strictEqual(result, emitter);
+
+      // Span must not be finished when stream is created
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Emitting data chunks should not close the span
+      emitter.emit('data', 'chunk 1');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      emitter.emit('data', 'chunk 2');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Only when the stream finishes does the span end
+      emitter.emit('end');
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(spans[0].events.length, 0);
+    });
+
+    it('does not end span prematurely until stream emits error event', () => {
+      const emitter = new EventEmitter();
+      traceAttempt(dynamicArgs, staticArgs, () => emitter, true);
+
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      emitter.emit('data', 'chunk');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      const error = new Error('stream failure');
+      emitter.emit('error', error);
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(
+        spans[0].attributes['error.message'],
+        'stream failure',
+      );
+      assert.strictEqual(spans[0].events.length, 1);
+      assert.strictEqual(spans[0].events[0].name, 'exception');
+    });
+
+    it('does not end span prematurely until stream emits close event', () => {
+      const emitter = new EventEmitter();
+      traceAttempt(dynamicArgs, staticArgs, () => emitter, 'stream');
+
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      emitter.emit('data', 'chunk');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      emitter.emit('close');
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+  });
+
+  describe('handlePromise', () => {
+    it('waits for promise resolution before ending span', async () => {
+      let ended = false;
+      let resolvePromise: () => void;
+      const promise = new Promise<void>(resolve => {
+        resolvePromise = resolve;
+      });
+
+      handlePromise(
+        promise,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+      assert.strictEqual(ended, false);
+
+      resolvePromise!();
+      await new Promise(resolve => setTimeout(resolve, 15));
+      assert.strictEqual(ended, true);
+    });
+
+    it('records error and ends span when promise rejects', async () => {
+      let ended = false;
+      let recordedError: unknown;
+      const error = new Error('promise error');
+
+      handlePromise(
+        Promise.reject(error),
+        err => {
+          recordedError = err;
+        },
+        () => {
+          ended = true;
+        },
+      );
+
+      await new Promise(resolve => setTimeout(resolve, 15));
+      assert.strictEqual(ended, true);
+      assert.strictEqual(recordedError, error);
+    });
+
+    it('supports custom thenables', async () => {
+      let ended = false;
+      const thenable = {
+        then(onfulfilled?: (val?: unknown) => unknown) {
+          setTimeout(() => {
+            onfulfilled?.();
+          }, 10);
+        },
+      };
+
+      handlePromise(
+        thenable,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+      assert.strictEqual(ended, false);
+
+      await new Promise(resolve => setTimeout(resolve, 20));
+      assert.strictEqual(ended, true);
+    });
+  });
+
+  describe('handleStream', () => {
+    it('manages stream events and ends span on end', () => {
+      let ended = false;
+      const emitter = new EventEmitter();
+      handleStream(
+        emitter,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+
+      assert.strictEqual(ended, false);
+      emitter.emit('data', 'chunk');
+      assert.strictEqual(ended, false);
+
+      emitter.emit('end');
+      assert.strictEqual(ended, true);
+    });
+
+    it('ends span on stream close', () => {
+      let ended = false;
+      const emitter = new EventEmitter();
+      handleStream(
+        emitter,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+
+      assert.strictEqual(ended, false);
+      emitter.emit('close');
+      assert.strictEqual(ended, true);
+    });
+
+    it('records error and ends span on stream error', () => {
+      let ended = false;
+      let recordedError: unknown;
+      const error = new Error('stream failure');
+      const emitter = new EventEmitter();
+
+      handleStream(
+        emitter,
+        err => {
+          recordedError = err;
+        },
+        () => {
+          ended = true;
+        },
+      );
+
+      assert.strictEqual(ended, false);
+      emitter.emit('error', error);
+      assert.strictEqual(ended, true);
+      assert.strictEqual(recordedError, error);
+    });
   });
 });

From 11597dc0ccbdf78a6e10865cce415f8a45543b95 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Tue, 8 Sep 2026 17:28:13 -0700
Subject: [PATCH 02/14] removes use of any

---
 core/packages/gax/src/observability/TracerHelper.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index b9309e0bc53..e4d7f727b80 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -141,11 +141,9 @@ export function traceAttempt<T>(
 export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  fn: () => any,
+  fn: () => unknown,
   isStream: boolean | 'promise' | 'stream' = false,
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-): any {
+): unknown {
   const isStreamCall = isStream === true || isStream === 'stream';
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
   return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
@@ -188,7 +186,7 @@ export function traceAttempt(
     try {
       const result = fn();
       if (isStreamCall) {
-        handleStream(result, recordError, endSpan);
+        handleStream(result as EventEmitter, recordError, endSpan);
       } else {
         handlePromise(result, recordError, endSpan);
       }

From dc24a7f654ed55c42650714cc5b1cd836355f497 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Tue, 8 Sep 2026 17:38:59 -0700
Subject: [PATCH 03/14] update formatting

---
 core/packages/gax/test/unit/tracerHelper.ts | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index cdf3976d5f2..018c14495a6 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -168,10 +168,8 @@ describe('TracerHelper', () => {
     });
 
     it('manages span lifetime for resolved promises', async () => {
-      const result = await traceAttempt(
-        dynamicArgs,
-        staticArgs,
-        () => Promise.resolve('async-result'),
+      const result = await traceAttempt(dynamicArgs, staticArgs, () =>
+        Promise.resolve('async-result'),
       );
       assert.strictEqual(result, 'async-result');
 

From 638d4002df0c4972df30c6caf0211aa51a4a18f8 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:11:14 -0700
Subject: [PATCH 04/14] refactor(gax): clean up stream listeners and update
 traceAttempt signatures

---
 .../gax/src/observability/TracerHelper.ts     | 53 ++++++++++++-------
 1 file changed, 35 insertions(+), 18 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index e4d7f727b80..b083feff74f 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -16,6 +16,8 @@
 
 import {EventEmitter} from 'events';
 import {Span, trace, Tracer} from '@opentelemetry/api';
+import { CancellableStream } from '../apitypes';
+import { CancellablePromise } from '../call';
 
 /**
  * Static metadata about the Google Cloud client library used to populate
@@ -103,16 +105,32 @@ export function handleStream(
   recordError: (err: unknown) => void,
   endSpan: () => void,
 ): void {
-  stream.on('error', (err: unknown) => {
+  const cleanup = () => {
+    stream.removeListener('error', onError);
+    stream.removeListener('end', onEnd);
+    stream.removeListener('close', onClose);
+    endSpan();
+  };
+
+  const onError = (err: unknown) => {
+    cleanup();
     recordError(err);
     endSpan();
-  });
-  stream.on('end', () => {
+  };
+
+  const onEnd = () => {
+    cleanup();
     endSpan();
-  });
-  stream.on('close', () => {
+  };
+
+  const onClose = () => {
+    cleanup();
     endSpan();
-  });
+  };
+
+  stream.on('error', onError);
+  stream.on('end', onEnd);
+  stream.on('close', onClose);
 }
 
 /**
@@ -123,28 +141,27 @@ export function handleStream(
  * @param {DynamicTraceContext} dynamicArgs - Dynamic trace context for the RPC call.
  * @param {StaticTraceContext} staticArgs - Static trace context for the client library.
  * @param {() => T} fn - The operation to trace.
- * @param {boolean | 'promise' | 'stream'} [isStream=false] - Whether the operation is a stream or a promise.
+ * @param {boolean} [isStreamCall=false] - Whether the operation is a stream or a promise.
  * @returns {T} The result of the traced operation.
  */
 export function traceAttempt<T extends EventEmitter>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => T,
-  isStream: true | 'stream',
+  fn: () => CancellableStream,
+  isStreamCall: true,
 ): T;
 export function traceAttempt<T>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => T,
-  isStream?: boolean | 'promise' | 'stream',
+  fn: () => CancellablePromise<T>,
+  isStreamCall?: false,
 ): T;
 export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => unknown,
-  isStream: boolean | 'promise' | 'stream' = false,
-): unknown {
-  const isStreamCall = isStream === true || isStream === 'stream';
+  fn: () => T,
+  isStreamCall: boolean = false,
+): T {
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
   return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
     span.setAttributes({
@@ -185,9 +202,9 @@ export function traceAttempt(
 
     try {
       const result = fn();
-      if (isStreamCall) {
-        handleStream(result as EventEmitter, recordError, endSpan);
-      } else {
+      if (isStreamCall && result instanceof EventEmitter) {
+        handleStream(result, recordError, endSpan);
+      } else if (!isStreamCall && result instanceof Promise) {
         handlePromise(result, recordError, endSpan);
       }
       return result;

From 12df94f248aeea9959e63149caf58d67fdafbe34 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:20:36 -0700
Subject: [PATCH 05/14] fix(gax): fix span ending in stream cleanup and update
 traceAttempt signatures

---
 core/packages/gax/src/observability/TracerHelper.ts | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index b083feff74f..fd1db7d62e5 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -109,7 +109,6 @@ export function handleStream(
     stream.removeListener('error', onError);
     stream.removeListener('end', onEnd);
     stream.removeListener('close', onClose);
-    endSpan();
   };
 
   const onError = (err: unknown) => {
@@ -150,7 +149,7 @@ export function traceAttempt<T extends EventEmitter>(
   fn: () => CancellableStream,
   isStreamCall: true,
 ): T;
-export function traceAttempt<T>(
+export function traceAttempt<T extends PromiseLike<T>>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
   fn: () => CancellablePromise<T>,
@@ -159,9 +158,9 @@ export function traceAttempt<T>(
 export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => T,
+  fn: () => unknown,
   isStreamCall: boolean = false,
-): T {
+): unknown {
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
   return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
     span.setAttributes({

From 6af7f236ad39737fa6669a8517d566b80d0ea99d Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:21:29 -0700
Subject: [PATCH 06/14] refactor(gax): simplify type parameter on traceAttempt

---
 core/packages/gax/src/observability/TracerHelper.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index fd1db7d62e5..1fd983ae3ad 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -149,7 +149,7 @@ export function traceAttempt<T extends EventEmitter>(
   fn: () => CancellableStream,
   isStreamCall: true,
 ): T;
-export function traceAttempt<T extends PromiseLike<T>>(
+export function traceAttempt<T>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
   fn: () => CancellablePromise<T>,

From aaa6d547365ad8077ad71cacf8347e74d6ca9b8b Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:29:22 -0700
Subject: [PATCH 07/14] refactor(gax): remove apitypes and call imports from
 TracerHelper

---
 core/packages/gax/src/observability/TracerHelper.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index 1fd983ae3ad..999c593e934 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -15,9 +15,7 @@
  */
 
 import {EventEmitter} from 'events';
-import {Span, trace, Tracer} from '@opentelemetry/api';
-import { CancellableStream } from '../apitypes';
-import { CancellablePromise } from '../call';
+import { Span, trace, Tracer } from '@opentelemetry/api';
 
 /**
  * Static metadata about the Google Cloud client library used to populate
@@ -146,13 +144,13 @@ export function handleStream(
 export function traceAttempt<T extends EventEmitter>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => CancellableStream,
+  fn: () => T,
   isStreamCall: true,
 ): T;
 export function traceAttempt<T>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => CancellablePromise<T>,
+  fn: () => T,
   isStreamCall?: false,
 ): T;
 export function traceAttempt(

From b39b884f36f8c62719bc1c59f093df8f1edc2ef6 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:33:20 -0700
Subject: [PATCH 08/14] test(gax): update TracerHelper unit tests and JSDocs

---
 .../gax/src/observability/TracerHelper.ts     | 19 ++--
 core/packages/gax/test/unit/tracerHelper.ts   | 97 +++++++++++++++----
 2 files changed, 90 insertions(+), 26 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index 999c593e934..3a7da500fd8 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -15,7 +15,7 @@
  */
 
 import {EventEmitter} from 'events';
-import { Span, trace, Tracer } from '@opentelemetry/api';
+import {Span, trace, Tracer} from '@opentelemetry/api';
 
 /**
  * Static metadata about the Google Cloud client library used to populate
@@ -71,9 +71,10 @@ export function getGaxTracer(): Tracer {
 /**
  * Manages span lifecycle for Promise-based operations.
  *
+ * @template T
  * @param {T} promise - The promise returned from the traced operation.
- * @param {(err: unknown) => void} recordError - Callback to record errors on the span.
- * @param {() => void} endSpan - Callback to end the span idempotently.
+ * @param {function} recordError - Callback to record errors on the span.
+ * @param {function} endSpan - Callback to end the span idempotently.
  */
 export function handlePromise<T>(
   promise: T,
@@ -92,11 +93,11 @@ export function handlePromise<T>(
 }
 
 /**
- * Manages span lifecycle for Stream-based operations.
+ * Manages span lifecycle for Stream-based operations and cleans up event listeners.
  *
  * @param {EventEmitter} stream - The stream returned from the traced operation.
- * @param {(err: unknown) => void} recordError - Callback to record errors on the span.
- * @param {() => void} endSpan - Callback to end the span idempotently.
+ * @param {function} recordError - Callback to record errors on the span.
+ * @param {function} endSpan - Callback to end the span idempotently.
  */
 export function handleStream(
   stream: EventEmitter,
@@ -137,8 +138,8 @@ export function handleStream(
  * @template T
  * @param {DynamicTraceContext} dynamicArgs - Dynamic trace context for the RPC call.
  * @param {StaticTraceContext} staticArgs - Static trace context for the client library.
- * @param {() => T} fn - The operation to trace.
- * @param {boolean} [isStreamCall=false] - Whether the operation is a stream or a promise.
+ * @param {function} fn - The operation to trace.
+ * @param {boolean} [isStreamCall=false] - Whether the operation is a stream call (true) or promise call (false).
  * @returns {T} The result of the traced operation.
  */
 export function traceAttempt<T extends EventEmitter>(
@@ -157,7 +158,7 @@ export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
   fn: () => unknown,
-  isStreamCall: boolean = false,
+  isStreamCall = false,
 ): unknown {
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
   return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index 018c14495a6..4401b861956 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -204,14 +204,33 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans[0].ended, true);
     });
 
-    it('supports custom thenables that do not inherit from Promise', async () => {
+    it('supports Promise subclasses', async () => {
+      class CustomPromise<T> extends Promise<T> {}
+      const customPromise = new CustomPromise<string>(resolve => {
+        setTimeout(() => {
+          resolve('custom-result');
+        }, 10);
+      });
+
+      const result = traceAttempt(dynamicArgs, staticArgs, () => customPromise);
+      assert.strictEqual(result, customPromise);
+
+      const initialSpans = harness.getSpans('google-gax');
+      assert.strictEqual(initialSpans.length, 0);
+
+      await new Promise(resolve => setTimeout(resolve, 25));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('does not manage promise lifecycle if result does not inherit from Promise', async () => {
       const customThenable = {
         then(onfulfilled?: (val: unknown) => void) {
-          setTimeout(() => {
-            if (onfulfilled) {
-              onfulfilled('custom-result');
-            }
-          }, 10);
+          if (onfulfilled) {
+            onfulfilled('custom-result');
+          }
         },
       };
 
@@ -222,14 +241,10 @@ describe('TracerHelper', () => {
       );
       assert.strictEqual(result, customThenable);
 
-      const initialSpans = harness.getSpans('google-gax');
-      assert.strictEqual(initialSpans.length, 0);
-
-      await new Promise(resolve => setTimeout(resolve, 25));
-
+      await new Promise(resolve => setTimeout(resolve, 20));
+      // Span is not ended because result is not an instanceof Promise
       const spans = harness.getSpans('google-gax');
-      assert.strictEqual(spans.length, 1);
-      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(spans.length, 0);
     });
 
     it('does not end span prematurely until asynchronous promise rejects', async () => {
@@ -305,7 +320,7 @@ describe('TracerHelper', () => {
 
     it('does not end span prematurely until stream emits close event', () => {
       const emitter = new EventEmitter();
-      traceAttempt(dynamicArgs, staticArgs, () => emitter, 'stream');
+      traceAttempt(dynamicArgs, staticArgs, () => emitter, true);
 
       assert.strictEqual(harness.getSpans('google-gax').length, 0);
 
@@ -317,6 +332,33 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans.length, 1);
       assert.strictEqual(spans[0].ended, true);
     });
+
+    it('supports isStreamCall explicitly set to false', async () => {
+      const result = await traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => Promise.resolve('explicit-false'),
+        false,
+      );
+      assert.strictEqual(result, 'explicit-false');
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('does not manage stream lifecycle if isStreamCall is true but result is not an EventEmitter', () => {
+      const nonEmitter = {data: 'not-an-emitter'};
+      const result = traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        () => nonEmitter as any,
+        true,
+      );
+      assert.strictEqual(result, nonEmitter);
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 0);
+    });
   });
 
   describe('handlePromise', () => {
@@ -386,7 +428,7 @@ describe('TracerHelper', () => {
   });
 
   describe('handleStream', () => {
-    it('manages stream events and ends span on end', () => {
+    it('manages stream events, ends span, and cleans up listeners on end', () => {
       let ended = false;
       const emitter = new EventEmitter();
       handleStream(
@@ -398,14 +440,21 @@ describe('TracerHelper', () => {
       );
 
       assert.strictEqual(ended, false);
+      assert.strictEqual(emitter.listenerCount('end'), 1);
+      assert.strictEqual(emitter.listenerCount('close'), 1);
+      assert.strictEqual(emitter.listenerCount('error'), 1);
+
       emitter.emit('data', 'chunk');
       assert.strictEqual(ended, false);
 
       emitter.emit('end');
       assert.strictEqual(ended, true);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
     });
 
-    it('ends span on stream close', () => {
+    it('ends span and cleans up listeners on stream close', () => {
       let ended = false;
       const emitter = new EventEmitter();
       handleStream(
@@ -417,30 +466,44 @@ describe('TracerHelper', () => {
       );
 
       assert.strictEqual(ended, false);
+      assert.strictEqual(emitter.listenerCount('close'), 1);
+
       emitter.emit('close');
       assert.strictEqual(ended, true);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
     });
 
-    it('records error and ends span on stream error', () => {
+    it('records error, ends span, and cleans up listeners on stream error', () => {
       let ended = false;
       let recordedError: unknown;
+      const order: string[] = [];
       const error = new Error('stream failure');
       const emitter = new EventEmitter();
 
       handleStream(
         emitter,
         err => {
+          order.push('recordError');
           recordedError = err;
         },
         () => {
+          order.push('endSpan');
           ended = true;
         },
       );
 
       assert.strictEqual(ended, false);
+      assert.strictEqual(emitter.listenerCount('error'), 1);
+
       emitter.emit('error', error);
       assert.strictEqual(ended, true);
       assert.strictEqual(recordedError, error);
+      assert.deepStrictEqual(order, ['recordError', 'endSpan']);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
     });
   });
 });

From b7a5da0c7031679f6fb70c3b03fab18dec48e9f7 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 11:47:14 -0700
Subject: [PATCH 09/14] fix(gax): end span synchronously when result is not a
 promise or stream

---
 .../gax/src/observability/TracerHelper.ts     |  2 ++
 core/packages/gax/test/unit/tracerHelper.ts   | 28 ++++++-------------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index 3a7da500fd8..06b9aa5a938 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -204,6 +204,8 @@ export function traceAttempt(
         handleStream(result, recordError, endSpan);
       } else if (!isStreamCall && result instanceof Promise) {
         handlePromise(result, recordError, endSpan);
+      } else {
+        endSpan();
       }
       return result;
     } catch (e) {
diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index 4401b861956..7ba985e5710 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -225,26 +225,15 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans[0].ended, true);
     });
 
-    it('does not manage promise lifecycle if result does not inherit from Promise', async () => {
-      const customThenable = {
-        then(onfulfilled?: (val: unknown) => void) {
-          if (onfulfilled) {
-            onfulfilled('custom-result');
-          }
-        },
-      };
+    it('ends span synchronously if result is not a Promise', () => {
+      const syncResult = {data: 'sync-data'};
 
-      const result = traceAttempt(
-        dynamicArgs,
-        staticArgs,
-        () => customThenable,
-      );
-      assert.strictEqual(result, customThenable);
+      const result = traceAttempt(dynamicArgs, staticArgs, () => syncResult);
+      assert.strictEqual(result, syncResult);
 
-      await new Promise(resolve => setTimeout(resolve, 20));
-      // Span is not ended because result is not an instanceof Promise
       const spans = harness.getSpans('google-gax');
-      assert.strictEqual(spans.length, 0);
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
     });
 
     it('does not end span prematurely until asynchronous promise rejects', async () => {
@@ -346,7 +335,7 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans[0].ended, true);
     });
 
-    it('does not manage stream lifecycle if isStreamCall is true but result is not an EventEmitter', () => {
+    it('ends span synchronously if isStreamCall is true but result is not an EventEmitter', () => {
       const nonEmitter = {data: 'not-an-emitter'};
       const result = traceAttempt(
         dynamicArgs,
@@ -357,7 +346,8 @@ describe('TracerHelper', () => {
       );
       assert.strictEqual(result, nonEmitter);
       const spans = harness.getSpans('google-gax');
-      assert.strictEqual(spans.length, 0);
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
     });
   });
 

From 4139ffb8e6405bbe5f7a335bab3f573e06b54552 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 13:44:20 -0700
Subject: [PATCH 10/14] refactor(gax): use GaxCallResult for traceAttempt
 implementation and add tests

---
 .../gax/src/observability/TracerHelper.ts     |  5 ++-
 core/packages/gax/test/unit/tracerHelper.ts   | 42 +++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index 06b9aa5a938..87ab9e7463d 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -16,6 +16,7 @@
 
 import {EventEmitter} from 'events';
 import {Span, trace, Tracer} from '@opentelemetry/api';
+import {GaxCallResult} from '../apitypes';
 
 /**
  * Static metadata about the Google Cloud client library used to populate
@@ -157,9 +158,9 @@ export function traceAttempt<T>(
 export function traceAttempt(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => unknown,
+  fn: () => GaxCallResult,
   isStreamCall = false,
-): unknown {
+): GaxCallResult {
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
   return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
     span.setAttributes({
diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index 7ba985e5710..69abca874d2 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -25,6 +25,11 @@ import {
   DynamicTraceContext,
   StaticTraceContext,
 } from '../../src/observability/TracerHelper';
+import {
+  GaxCallResult,
+  CancellableStream,
+  ResultTuple,
+} from '../../src/apitypes';
 import {OtelHarness} from './otelHarness';
 
 describe('TracerHelper', () => {
@@ -349,6 +354,43 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans.length, 1);
       assert.strictEqual(spans[0].ended, true);
     });
+
+    it('supports GaxCallResult promise operations', async () => {
+      const cancellablePromise = Object.assign(
+        Promise.resolve([{}, undefined, undefined] as ResultTuple),
+        {
+          cancel: () => {},
+        },
+      ) as GaxCallResult;
+
+      const result = traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => cancellablePromise,
+      );
+      assert.strictEqual(result, cancellablePromise);
+      await result;
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('supports GaxCallResult stream operations', () => {
+      const stream = Object.assign(new EventEmitter(), {
+        cancel: () => {},
+      }) as unknown as CancellableStream;
+
+      const result = traceAttempt(dynamicArgs, staticArgs, () => stream, true);
+      assert.strictEqual(result, stream);
+
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+      stream.emit('end');
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
   });
 
   describe('handlePromise', () => {

From 10a3581790e6ed5197da9374775ff4f4b7134b62 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 13:48:24 -0700
Subject: [PATCH 11/14] feat(gax): add GaxCallResult overload for traceAttempt

---
 core/packages/gax/src/observability/TracerHelper.ts | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index 87ab9e7463d..ae40a4f33c2 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -143,6 +143,12 @@ export function handleStream(
  * @param {boolean} [isStreamCall=false] - Whether the operation is a stream call (true) or promise call (false).
  * @returns {T} The result of the traced operation.
  */
+export function traceAttempt(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => GaxCallResult,
+  isStreamCall?: boolean,
+): GaxCallResult;
 export function traceAttempt<T extends EventEmitter>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,

From 474c57cbf41b5b13e3acb33d235edae6a5750c0a Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 18:44:40 -0700
Subject: [PATCH 12/14] fix(gax): support custom thenables and prevent
 duplicate endSpan in handleStream

---
 .../gax/src/observability/TracerHelper.ts     |  89 ++++++--
 core/packages/gax/test/unit/tracerHelper.ts   | 202 ++++++++++++++++++
 2 files changed, 278 insertions(+), 13 deletions(-)

diff --git a/core/packages/gax/src/observability/TracerHelper.ts b/core/packages/gax/src/observability/TracerHelper.ts
index ae40a4f33c2..f52e4862b6e 100644
--- a/core/packages/gax/src/observability/TracerHelper.ts
+++ b/core/packages/gax/src/observability/TracerHelper.ts
@@ -69,6 +69,45 @@ export function getGaxTracer(): Tracer {
   return trace.getTracer('google-gax');
 }
 
+/**
+ * Checks if a value behaves like a Promise or Thenable.
+ *
+ * Note: It is not sufficient to check `result instanceof Promise` because:
+ * 1. Custom classes implementing `CancellablePromise` or Thenables may not
+ *    inherit directly from the native JavaScript `Promise` prototype.
+ * 2. GAX callers or custom callers may return objects such as `OngoingCallPromise`
+ *    that hold the actual promise on a `.promise` property.
+ * 3. Promises originating from different execution realms (such as Node.js vm
+ *    contexts or different package bundles) fail `instanceof Promise` checks.
+ */
+function isPromiseLike<T = unknown>(value: unknown): value is PromiseLike<T> {
+  return (
+    value instanceof Promise ||
+    (value !== null &&
+      (typeof value === 'object' || typeof value === 'function') &&
+      typeof (value as {then?: unknown}).then === 'function')
+  );
+}
+
+/**
+ * Extracts a PromiseLike target from a value, supporting native Promises,
+ * custom Thenables, classes implementing CancellablePromise, and OngoingCallPromise wrappers.
+ */
+function getPromiseTarget<T = unknown>(value: unknown): PromiseLike<T> | null {
+  if (isPromiseLike<T>(value)) {
+    return value;
+  }
+  if (
+    value !== null &&
+    (typeof value === 'object' || typeof value === 'function') &&
+    'promise' in (value as object) &&
+    isPromiseLike<T>((value as {promise?: unknown}).promise)
+  ) {
+    return (value as {promise: PromiseLike<T>}).promise;
+  }
+  return null;
+}
+
 /**
  * Manages span lifecycle for Promise-based operations.
  *
@@ -82,14 +121,25 @@ export function handlePromise<T>(
   recordError: (err: unknown) => void,
   endSpan: () => void,
 ): void {
-  Promise.resolve(promise)
-    .then(() => {
+  let spanEnded = false;
+  const endSpanOnce = () => {
+    if (!spanEnded) {
+      spanEnded = true;
       endSpan();
+    }
+  };
+
+  const target = getPromiseTarget(promise) ?? promise;
+  Promise.resolve(target)
+    .then(() => {
+      endSpanOnce();
       return null;
     })
     .catch(err => {
-      recordError(err);
-      endSpan();
+      if (!spanEnded) {
+        recordError(err);
+        endSpanOnce();
+      }
     });
 }
 
@@ -105,26 +155,35 @@ export function handleStream(
   recordError: (err: unknown) => void,
   endSpan: () => void,
 ): void {
+  let spanEnded = false;
+
   const cleanup = () => {
     stream.removeListener('error', onError);
     stream.removeListener('end', onEnd);
     stream.removeListener('close', onClose);
   };
 
+  const endSpanOnce = () => {
+    if (!spanEnded) {
+      spanEnded = true;
+      cleanup();
+      endSpan();
+    }
+  };
+
   const onError = (err: unknown) => {
-    cleanup();
-    recordError(err);
-    endSpan();
+    if (!spanEnded) {
+      recordError(err);
+      endSpanOnce();
+    }
   };
 
   const onEnd = () => {
-    cleanup();
-    endSpan();
+    endSpanOnce();
   };
 
   const onClose = () => {
-    cleanup();
-    endSpan();
+    endSpanOnce();
   };
 
   stream.on('error', onError);
@@ -207,10 +266,14 @@ export function traceAttempt(
 
     try {
       const result = fn();
+      // Use getPromiseTarget instead of `result instanceof Promise` to ensure custom
+      // thenables, CancellablePromise implementations, and OngoingCallPromise wrappers
+      // are properly tracked rather than leaving spans unclosed or ending them prematurely.
+      const promiseTarget = !isStreamCall ? getPromiseTarget(result) : null;
       if (isStreamCall && result instanceof EventEmitter) {
         handleStream(result, recordError, endSpan);
-      } else if (!isStreamCall && result instanceof Promise) {
-        handlePromise(result, recordError, endSpan);
+      } else if (promiseTarget) {
+        handlePromise(promiseTarget, recordError, endSpan);
       } else {
         endSpan();
       }
diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index 69abca874d2..fab854f2968 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -30,6 +30,7 @@ import {
   CancellableStream,
   ResultTuple,
 } from '../../src/apitypes';
+import {OngoingCallPromise} from '../../src/call';
 import {OtelHarness} from './otelHarness';
 
 describe('TracerHelper', () => {
@@ -230,6 +231,94 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans[0].ended, true);
     });
 
+    it('supports custom thenables implementing CancellablePromise without inheriting from Promise', async () => {
+      class CustomCancellablePromise {
+        private readonly promise: Promise<string>;
+        constructor(executor: (resolve: (val: string) => void) => void) {
+          this.promise = new Promise(executor);
+        }
+        cancel(): void {}
+        then<TResult1 = string, TResult2 = never>(
+          onfulfilled?:
+            ((value: string) => TResult1 | PromiseLike<TResult1>) | null,
+          onrejected?:
+            ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null,
+        ): Promise<TResult1 | TResult2> {
+          return this.promise.then(onfulfilled, onrejected);
+        }
+        catch<TResult = never>(
+          onrejected?:
+            ((reason: unknown) => TResult | PromiseLike<TResult>) | null,
+        ): Promise<string | TResult> {
+          return this.promise.catch(onrejected);
+        }
+      }
+
+      let resolvePromise: (val: string) => void;
+      const customPromise = new CustomCancellablePromise(resolve => {
+        resolvePromise = resolve;
+      });
+
+      // Verify it is NOT an instance of native Promise
+      assert.strictEqual(customPromise instanceof Promise, false);
+
+      const result = traceAttempt(dynamicArgs, staticArgs, () => customPromise);
+      assert.strictEqual(result, customPromise);
+
+      // Verify the span is NOT closed while the custom promise is pending
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      resolvePromise!('custom-success');
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('supports OngoingCallPromise objects whose promise property resolves', async () => {
+      const ongoingCall = new OngoingCallPromise();
+      assert.strictEqual(ongoingCall instanceof Promise, false);
+
+      const result = traceAttempt(dynamicArgs, staticArgs, () => ongoingCall);
+      assert.strictEqual(result, ongoingCall);
+
+      // Verify the span is NOT closed while ongoingCall is in flight
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Complete the call via callback
+      ongoingCall.callback!(null, {data: 'result'});
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+    });
+
+    it('supports OngoingCallPromise objects whose promise property rejects', async () => {
+      const ongoingCall = new OngoingCallPromise();
+      assert.strictEqual(ongoingCall instanceof Promise, false);
+
+      const error = new Error('ongoing call failed');
+      const result = traceAttempt(dynamicArgs, staticArgs, () => ongoingCall);
+      assert.strictEqual(result, ongoingCall);
+
+      // Verify the span is NOT closed while ongoingCall is in flight
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Fail the call via callback
+      ongoingCall.callback!(error);
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(
+        spans[0].attributes['error.message'],
+        'ongoing call failed',
+      );
+    });
+
     it('ends span synchronously if result is not a Promise', () => {
       const syncResult = {data: 'sync-data'};
 
@@ -457,6 +546,48 @@ describe('TracerHelper', () => {
       await new Promise(resolve => setTimeout(resolve, 20));
       assert.strictEqual(ended, true);
     });
+
+    it('supports OngoingCallPromise wrappers in handlePromise', async () => {
+      let ended = false;
+      const ongoingCall = new OngoingCallPromise();
+
+      handlePromise(
+        ongoingCall,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+      assert.strictEqual(ended, false);
+
+      ongoingCall.callback!(null, {data: 'hello'});
+      await new Promise(resolve => setTimeout(resolve, 20));
+      assert.strictEqual(ended, true);
+    });
+
+    it('ensures endSpan is called only once even with thenables that trigger both resolve and reject', async () => {
+      let endSpanCount = 0;
+      const buggyThenable = {
+        then(
+          onfulfilled?: (val?: unknown) => unknown,
+          onrejected?: (err: unknown) => unknown,
+        ) {
+          onfulfilled?.();
+          onrejected?.(new Error('buggy error'));
+        },
+      };
+
+      handlePromise(
+        buggyThenable,
+        () => {},
+        () => {
+          endSpanCount++;
+        },
+      );
+
+      await new Promise(resolve => setTimeout(resolve, 20));
+      assert.strictEqual(endSpanCount, 1);
+    });
   });
 
   describe('handleStream', () => {
@@ -537,5 +668,76 @@ describe('TracerHelper', () => {
       assert.strictEqual(emitter.listenerCount('close'), 0);
       assert.strictEqual(emitter.listenerCount('error'), 0);
     });
+
+    it('ensures endSpan is called only once if stream emits error followed by close', () => {
+      let endSpanCount = 0;
+      let recordedError: unknown;
+      const error = new Error('stream error');
+      const emitter = new EventEmitter();
+
+      handleStream(
+        emitter,
+        err => {
+          recordedError = err;
+        },
+        () => {
+          endSpanCount++;
+        },
+      );
+
+      assert.strictEqual(endSpanCount, 0);
+      emitter.emit('error', error);
+      emitter.emit('close');
+
+      assert.strictEqual(endSpanCount, 1);
+      assert.strictEqual(recordedError, error);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
+    });
+
+    it('ensures endSpan is called only once if stream emits end followed by close', () => {
+      let endSpanCount = 0;
+      const emitter = new EventEmitter();
+
+      handleStream(
+        emitter,
+        () => {},
+        () => {
+          endSpanCount++;
+        },
+      );
+
+      assert.strictEqual(endSpanCount, 0);
+      emitter.emit('end');
+      emitter.emit('close');
+
+      assert.strictEqual(endSpanCount, 1);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
+    });
+
+    it('ensures endSpan is called only once if stream emits close followed by end', () => {
+      let endSpanCount = 0;
+      const emitter = new EventEmitter();
+
+      handleStream(
+        emitter,
+        () => {},
+        () => {
+          endSpanCount++;
+        },
+      );
+
+      assert.strictEqual(endSpanCount, 0);
+      emitter.emit('close');
+      emitter.emit('end');
+
+      assert.strictEqual(endSpanCount, 1);
+      assert.strictEqual(emitter.listenerCount('end'), 0);
+      assert.strictEqual(emitter.listenerCount('close'), 0);
+      assert.strictEqual(emitter.listenerCount('error'), 0);
+    });
   });
 });

From 3be25fef00eabd312bfc35ba59e32d33dc07a217 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Wed, 9 Sep 2026 19:08:26 -0700
Subject: [PATCH 13/14] feat(gax): export traceAttempt from index

---
 core/packages/gax/src/index.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/packages/gax/src/index.ts b/core/packages/gax/src/index.ts
index b0d09dd3b4d..a02da512ae8 100644
--- a/core/packages/gax/src/index.ts
+++ b/core/packages/gax/src/index.ts
@@ -120,7 +120,7 @@ export {
   checkTelemetryEnabled,
 } from './util';
 
-export {StaticTraceContext} from './observability/TracerHelper';
+export {StaticTraceContext, traceAttempt} from './observability/TracerHelper';
 
 export {ServiceError, ChannelCredentials} from '@grpc/grpc-js';
 export {warn} from './warnings';

From 8b2e8a9253e395ad8081e749e18092cee76d2ad3 Mon Sep 17 00:00:00 2001
From: Shivanee Persaud <shivaneep@google.com>
Date: Thu, 10 Sep 2026 11:34:53 -0700
Subject: [PATCH 14/14] test(gax): add unit tests for stream retries and
 listener cleanup in TracerHelper

---
 core/packages/gax/test/unit/tracerHelper.ts | 106 ++++++++++++++++++++
 1 file changed, 106 insertions(+)

diff --git a/core/packages/gax/test/unit/tracerHelper.ts b/core/packages/gax/test/unit/tracerHelper.ts
index fab854f2968..f0bd6e8d058 100644
--- a/core/packages/gax/test/unit/tracerHelper.ts
+++ b/core/packages/gax/test/unit/tracerHelper.ts
@@ -480,6 +480,81 @@ describe('TracerHelper', () => {
       assert.strictEqual(spans.length, 1);
       assert.strictEqual(spans[0].ended, true);
     });
+
+    it('correctly creates separate spans and cleans up listeners across retried stream attempts', () => {
+      const attempt1Stream = new EventEmitter();
+      const attempt2Stream = new EventEmitter();
+      const retryableError = new Error('transient stream failure');
+
+      let attempt = 0;
+      const executeStreamingCall = () => {
+        attempt++;
+        const currentStream = attempt === 1 ? attempt1Stream : attempt2Stream;
+        return traceAttempt(dynamicArgs, staticArgs, () => currentStream, true);
+      };
+
+      // Attempt 1
+      const stream1 = executeStreamingCall();
+      assert.strictEqual(stream1, attempt1Stream);
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Attempt 1 fails with transient error
+      attempt1Stream.emit('error', retryableError);
+
+      const spansAfterAttempt1 = harness.getSpans('google-gax');
+      assert.strictEqual(spansAfterAttempt1.length, 1);
+      assert.strictEqual(spansAfterAttempt1[0].ended, true);
+      assert.strictEqual(
+        spansAfterAttempt1[0].attributes['error.message'],
+        'transient stream failure',
+      );
+      assert.strictEqual(spansAfterAttempt1[0].events.length, 1);
+      assert.strictEqual(attempt1Stream.listenerCount('error'), 0);
+
+      // Attempt 2 (retry)
+      const stream2 = executeStreamingCall();
+      assert.strictEqual(stream2, attempt2Stream);
+      assert.strictEqual(harness.getSpans('google-gax').length, 1);
+
+      // Data chunks received on attempt 2
+      attempt2Stream.emit('data', 'retry chunk 1');
+      assert.strictEqual(harness.getSpans('google-gax').length, 1);
+
+      // Attempt 2 completes successfully
+      attempt2Stream.emit('end');
+
+      const spansAfterAttempt2 = harness.getSpans('google-gax');
+      assert.strictEqual(spansAfterAttempt2.length, 2);
+      assert.strictEqual(spansAfterAttempt2[1].ended, true);
+      assert.strictEqual(spansAfterAttempt2[1].events.length, 0);
+      assert.strictEqual(attempt2Stream.listenerCount('end'), 0);
+    });
+
+    it('keeps span active when a stream handles retries internally before completing', () => {
+      const outerStream = new EventEmitter();
+      const result = traceAttempt(
+        dynamicArgs,
+        staticArgs,
+        () => outerStream,
+        true,
+      );
+      assert.strictEqual(result, outerStream);
+
+      // Initial chunk before internal retry
+      outerStream.emit('data', 'chunk-before-retry');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Internal retry transparently recovers and delivers more data
+      outerStream.emit('data', 'chunk-after-retry');
+      assert.strictEqual(harness.getSpans('google-gax').length, 0);
+
+      // Final completion
+      outerStream.emit('end');
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 1);
+      assert.strictEqual(spans[0].ended, true);
+      assert.strictEqual(spans[0].events.length, 0);
+    });
   });
 
   describe('handlePromise', () => {
@@ -739,5 +814,36 @@ describe('TracerHelper', () => {
       assert.strictEqual(emitter.listenerCount('close'), 0);
       assert.strictEqual(emitter.listenerCount('error'), 0);
     });
+
+    it('does not remove other error or event listeners (such as retry handlers) on cleanup', () => {
+      let ended = false;
+      let otherErrorHandled = false;
+      const error = new Error('retryable error');
+      const emitter = new EventEmitter();
+
+      // Simulate an external retry handler or middleware attached to the stream
+      emitter.on('error', err => {
+        assert.strictEqual(err, error);
+        otherErrorHandled = true;
+      });
+
+      handleStream(
+        emitter,
+        () => {},
+        () => {
+          ended = true;
+        },
+      );
+
+      // Verify two error listeners are present (external retry listener and handleStream listener)
+      assert.strictEqual(emitter.listenerCount('error'), 2);
+
+      emitter.emit('error', error);
+
+      assert.strictEqual(ended, true);
+      assert.strictEqual(otherErrorHandled, true);
+      // handleStream removed its own listener, but the external retry listener is preserved
+      assert.strictEqual(emitter.listenerCount('error'), 1);
+    });
   });
 });