fix(js): keep traceable wrappedFunc returnValue props (#787)

iterables can have other properties

js/src/singletons/types.ts

@@ -32,11 +32,10 @@ type UnionToIntersection<U> = (U extends any ? (x: U) => void : never) extends (
   ? I
   : never;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-
 export type TraceableFunction<Func extends (...args: any[]) => any> =
   // function overloads are represented as intersections rather than unions
   // matches the behavior introduced in https://github.com/microsoft/TypeScript/pull/54448
-  Func extends {
+  (Func extends {
     (...args: infer A1): infer R1;
     (...args: infer A2): infer R2;
     (...args: infer A3): infer R3;
@@ -70,6 +69,9 @@ export type TraceableFunction<Func extends (...args: any[]) => any> =
         (...args: infer A1): infer R1;
       }
     ? UnionToIntersection<WrapArgReturnPair<[A1, R1]>>
-    : never;
+    : never) & {
+    // Other properties of Func
+    [K in keyof Func]: Func[K];
+  };
 
 export type RunTreeLike = RunTree;

js/src/tests/traceable.test.ts

@@ -513,6 +513,44 @@ describe("async generators", () => {
       },
     });
   });
+
+  test("iterable with props", async () => {
+    const { client, callSpy } = mockClient();
+
+    const iterableTraceable = traceable(
+      function iterableWithProps() {
+        return {
+          *[Symbol.asyncIterator]() {
+            yield 0;
+          },
+          prop: "value",
+        };
+      },
+      {
+        client,
+        tracingEnabled: true,
+      }
+    );
+
+    const numbers: number[] = [];
+    const iterableWithProps = await iterableTraceable();
+    for await (const num of iterableWithProps) {
+      numbers.push(num);
+    }
+
+    expect(numbers).toEqual([0]);
+
+    expect(iterableWithProps.prop).toBe("value");
+    expect(getAssumedTreeFromCalls(callSpy.mock.calls)).toMatchObject({
+      nodes: ["iterableWithProps:0"],
+      edges: [],
+      data: {
+        "iterableWithProps:0": {
+          outputs: { outputs: [0] },
+        },
+      },
+    });
+  });
 });
 
 describe("deferred input", () => {

js/src/traceable.ts

@@ -434,14 +434,13 @@ export function traceable<Func extends (...args: any[]) => any>(
         return chunks;
       }
 
-      async function* wrapAsyncGeneratorForTracing(
-        iterable: AsyncIterable<unknown>,
+      async function* wrapAsyncIteratorForTracing(
+        iterator: AsyncIterator<unknown, unknown, undefined>,
         snapshot: ReturnType<typeof AsyncLocalStorage.snapshot> | undefined
       ) {
         let finished = false;
         const chunks: unknown[] = [];
         try {
-          const iterator = iterable[Symbol.asyncIterator]();
           while (true) {
             const { value, done } = await (snapshot
               ? snapshot(() => iterator.next())
@@ -464,6 +463,15 @@ export function traceable<Func extends (...args: any[]) => any>(
           await handleEnd();
         }
       }
+      function wrapAsyncGeneratorForTracing(
+        iterable: AsyncIterable<unknown>,
+        snapshot: ReturnType<typeof AsyncLocalStorage.snapshot> | undefined
+      ) {
+        const iterator = iterable[Symbol.asyncIterator]();
+        const wrappedIterator = wrapAsyncIteratorForTracing(iterator, snapshot);
+        iterable[Symbol.asyncIterator] = () => wrappedIterator;
+        return iterable;
+      }
 
       async function handleEnd() {
         const onEnd = config?.on_end;