Fix runnable map, start adding tests

langchain-core/src/runnables/base.ts

@@ -2082,16 +2082,7 @@ export class RunnableMap<
           );
         }
       );
-      if (options?.signal) {
-        promises.push(
-          new Promise<never>((_, reject) => {
-            options.signal?.addEventListener("abort", () =>
-              reject(new Error("Aborted"))
-            );
-          })
-        );
-      }
-      await Promise.all(promises);
+      await raceWithSignal(Promise.all(promises), options?.signal);
     } catch (e) {
       await runManager?.handleChainError(e);
       throw e;

langchain-core/src/runnables/tests/runnable_map.test.ts

@@ -67,6 +67,45 @@ test("Test map inference in a sequence", async () => {
   );
 });
 
+test("Test invoke with signal", async () => {
+  const map = RunnableMap.from({
+    question: new RunnablePassthrough(),
+    context: async () => {
+      await new Promise((resolve) => setTimeout(resolve, 500));
+      return "SOME STUFF";
+    },
+  });
+  const controller = new AbortController();
+  await expect(async () => {
+    await Promise.all([
+      map.invoke("testing", {
+        signal: controller.signal,
+      }),
+      new Promise<void>((resolve) => {
+        controller.abort();
+        resolve();
+      }),
+    ]);
+  }).rejects.toThrowError();
+});
+
+test("Test stream with signal", async () => {
+  const map = RunnableMap.from({
+    question: new RunnablePassthrough(),
+    context: async () => {
+      await new Promise((resolve) => setTimeout(resolve, 500));
+      return "SOME STUFF";
+    },
+  });
+  const controller = new AbortController();
+  await expect(async () => {
+    const stream = await map.stream("TESTING", { signal: controller.signal });
+    for await (const _ of stream) {
+      controller.abort();
+    }
+  }).rejects.toThrowError();
+});
+
 test("Should not allow mismatched inputs", async () => {
   const prompt = ChatPromptTemplate.fromTemplate(
     "context: {context}, question: {question}"

langchain-core/src/runnables/tests/runnable_stream_events_v2.test.ts

@@ -2120,3 +2120,24 @@ test("Runnable streamEvents method with text/event-stream encoding", async () =>
 
   expect(decoder.decode(events[3])).toEqual("event: end\n\n");
 });
+
+test("Runnable streamEvents method should respect passed signal", async () => {
+  const r = RunnableLambda.from(reverse);
+
+  const chain = r
+    .withConfig({ runName: "1" })
+    .pipe(r.withConfig({ runName: "2" }))
+    .pipe(r.withConfig({ runName: "3" }));
+
+  const controller = new AbortController();
+  const eventStream = await chain.streamEvents("hello", {
+    version: "v2",
+    signal: controller.signal,
+  });
+  await expect(async () => {
+    for await (const _ of eventStream) {
+      // Abort after the first chunk
+      controller.abort();
+    }
+  }).rejects.toThrowError();
+});

langchain-core/src/utils/signal.ts

@@ -5,12 +5,13 @@ export async function raceWithSignal<T>(
   if (signal === undefined) {
     return promise;
   }
-  if (signal.aborted) {
-    throw new Error("AbortError");
-  }
   return Promise.race([
     promise,
     new Promise<never>((_, reject) => {
+      // Must be inside of the promise to avoid a race condition
+      if (signal.aborted) {
+        return reject(new Error("Aborted"));
+      }
       signal.addEventListener("abort", () => reject(new Error("Aborted")));
     }),
   ]);