@langchain/langgraph 0.0.31 → 0.0.33

package/README.md

@@ -4,7 +4,6 @@
 ![Version](https://img.shields.io/npm/v/@langchain/langgraph?logo=npm)  
 [![Downloads](https://img.shields.io/npm/dm/@langchain/langgraph)](https://www.npmjs.com/package/@langchain/langgraph)
 [![Open Issues](https://img.shields.io/github/issues-raw/langchain-ai/langgraphjs)](https://github.com/langchain-ai/langgraphjs/issues)
-[![](https://dcbadge.vercel.app/api/server/6adMQxSpJS?compact=true&style=flat)](https://discord.com/channels/1038097195422978059/1170024642245832774)
 
 ⚡ Building language agents as graphs ⚡
 
@@ -56,55 +55,52 @@ export LANGCHAIN_API_KEY=ls__...
 Now let's define our agent:
 
 ```typescript
-import { HumanMessage, AIMessage } from "@langchain/core/messages";
-import { DynamicStructuredTool } from "@langchain/core/tools";
+import { AIMessage, BaseMessage, HumanMessage } from "@langchain/core/messages";
+import { tool } from "@langchain/core/tools";
 import { z } from "zod";
 import { ChatAnthropic } from "@langchain/anthropic";
-import { END, START, StateGraph, StateGraphArgs } from "@langchain/langgraph";
+import { StateGraph, StateGraphArgs } from "@langchain/langgraph";
 import { MemorySaver } from "@langchain/langgraph";
 import { ToolNode } from "@langchain/langgraph/prebuilt";
 
 // Define the state interface
 interface AgentState {
-  messages: HumanMessage[];
+  messages: BaseMessage[];
 }
 
 // Define the graph state
 const graphState: StateGraphArgs<AgentState>["channels"] = {
   messages: {
-    value: (x: HumanMessage[], y: HumanMessage[]) => x.concat(y),
-    default: () => [],
+    reducer: (x: BaseMessage[], y: BaseMessage[]) => x.concat(y),
   },
 };
 
 // Define the tools for the agent to use
-
-const searchTool = new DynamicStructuredTool({
-  name: "search",
+const weatherTool = tool(async ({ query }) => {
+  // This is a placeholder for the actual implementation
+  if (query.toLowerCase().includes("sf") || query.toLowerCase().includes("san francisco")) {
+    return "It's 60 degrees and foggy."
+  }
+  return "It's 90 degrees and sunny."
+}, {
+  name: "weather",
   description:
-    "Call to surf the web.",
+    "Call to get the current weather for a location.",
   schema: z.object({
     query: z.string().describe("The query to use in your search."),
   }),
-  func: async ({ query }: { query: string }) => {
-    // This is a placeholder for the actual implementation
-    if (query.toLowerCase().includes("sf") || query.toLowerCase().includes("san francisco")) {
-      return "It's 60 degrees and foggy."
-    }
-    return "It's 90 degrees and sunny."
-  },
 });
 
-const tools = [searchTool];
+const tools = [weatherTool];
 const toolNode = new ToolNode<AgentState>(tools);
 
 const model = new ChatAnthropic({
-  model: "claude-3-sonnet-20240229",
+  model: "claude-3-5-sonnet-20240620",
   temperature: 0,
 }).bindTools(tools);
 
 // Define the function that determines whether to continue or not
-function shouldContinue(state: AgentState): "tools" | typeof END {
+function shouldContinue(state: AgentState) {
   const messages = state.messages;
   const lastMessage = messages[messages.length - 1] as AIMessage;
 
@@ -113,7 +109,7 @@ function shouldContinue(state: AgentState): "tools" | typeof END {
     return "tools";
   }
   // Otherwise, we stop (reply to the user)
-  return END;
+  return "__end__";
 }
 
 // Define the function that calls the model
@@ -129,7 +125,7 @@ async function callModel(state: AgentState) {
 const workflow = new StateGraph<AgentState>({ channels: graphState })
   .addNode("agent", callModel)
   .addNode("tools", toolNode)
-  .addEdge(START, "agent")
+  .addEdge("__start__", "agent")
   .addConditionalEdges("agent", shouldContinue)
   .addEdge("tools", "agent");
 
@@ -146,13 +142,21 @@ const finalState = await app.invoke(
   { messages: [new HumanMessage("what is the weather in sf")] },
   { configurable: { thread_id: "42" } }
 );
+
 console.log(finalState.messages[finalState.messages.length - 1].content);
 ```
 
 This will output:
 
 ```
-Based on the search results, I can tell you that the current weather in San Francisco is:\n\nTemperature: 60 degrees Fahrenheit\nConditions: Foggy\n\nSan Francisco is known for its microclimates and frequent fog, especially during the summer months. The temperature of 60°F (about 15.5°C) is quite typical for the city, which tends to have mild temperatures year-round. The fog, often referred to as "Karl the Fog" by locals, is a characteristic feature of San Francisco\'s weather, particularly in the mornings and evenings.\n\nIs there anything else you\'d like to know about the weather in San Francisco or any other location?
+Based on the information I received, the current weather in San Francisco is:
+
+Temperature: 60 degrees Fahrenheit
+Conditions: Foggy
+
+San Francisco is known for its foggy weather, especially during certain times of the year. The moderate temperature of 60°F (about 15.5°C) is quite typical for the city, which generally has mild weather year-round due to its coastal location.
+
+Is there anything else you'd like to know about the weather in San Francisco or any other location?
 ```
 
 Now when we pass the same `"thread_id"`, the conversation context is retained via the saved state (i.e. stored list of messages):
@@ -166,7 +170,16 @@ console.log(nextState.messages[nextState.messages.length - 1].content);
 ```
 
 ```
-Based on the search results, I can tell you that the current weather in New York City is:\n\nTemperature: 90 degrees Fahrenheit (approximately 32.2 degrees Celsius)\nConditions: Sunny\n\nThis weather is quite different from what we just saw in San Francisco. New York is experiencing much warmer temperatures right now. Here are a few points to note:\n\n1. The temperature of 90°F is quite hot, typical of summer weather in New York City.\n2. The sunny conditions suggest clear skies, which is great for outdoor activities but also means it might feel even hotter due to direct sunlight.\n3. This kind of weather in New York often comes with high humidity, which can make it feel even warmer than the actual temperature suggests.\n\nIt's interesting to see the stark contrast between San Francisco's mild, foggy weather and New York's hot, sunny conditions. This difference illustrates how varied weather can be across different parts of the United States, even on the same day.\n\nIs there anything else you'd like to know about the weather in New York or any other location?
+Based on the information I received, the current weather in New York is:
+
+Temperature: 90 degrees Fahrenheit (approximately 32.2 degrees Celsius)
+Conditions: Sunny
+
+New York is experiencing quite warm weather today. A temperature of 90°F is considered hot for most people, and it's significantly warmer than the San Francisco weather we just checked. The sunny conditions suggest it's a clear day without cloud cover, which can make it feel even warmer.
+
+On a day like this in New York, it would be advisable for people to stay hydrated, seek shade when possible, and use sun protection if spending time outdoors.
+
+Is there anything else you'd like to know about the weather in New York or any other location?
 ```
 
 ### Step-by-step Breakdown
@@ -175,7 +188,7 @@ Based on the search results, I can tell you that the current weather in New York
     <summary>Initialize the model and tools.</summary>
 
     - We use `ChatAnthropic` as our LLM. **NOTE:** We need make sure the model knows that it has these tools available to call. We can do this by converting the LangChain tools into the format for Anthropic tool calling using the `.bindTools()` method.
-    - We define the tools we want to use -- a search tool in our case. It is really easy to create your own tools - see documentation [here](https://js.langchain.com/docs/modules/agents/tools/dynamic) on how to do that.
+    - We define the tools we want to use -- a weather tool in our case. See the documentation [here](https://js.langchain.com/docs/modules/agents/tools/dynamic) on how to create your own tools.
    </details>
 
 2. <details>
@@ -197,7 +210,7 @@ Based on the search results, I can tell you that the current weather in New York
 4. <details>
     <summary>Define entry point and graph edges.</summary>
 
-      First, we need to set the entry point for graph execution - `agent` node.
+      First, we need to set the entry point for graph execution - the `agent` node.
 
       Then we define one normal and one conditional edge. A conditional edge means that the destination depends on the contents of the graph's state (`AgentState`). In our case, the destination is not known until the agent (LLM) decides.
 
@@ -225,7 +238,7 @@ Based on the search results, I can tell you that the current weather in New York
         - If `AIMessage` has `tool_calls`, the `"tools"` node executes.
         - The `"agent"` node executes again and returns an `AIMessage`.
 
-    5. Execution progresses to the special `END` value and outputs the final state.
+    5. Execution progresses to the special `__end__` value and outputs the final state.
     As a result, we get a list of all our chat messages as output.
    </details>
 
@@ -235,3 +248,37 @@ Based on the search results, I can tell you that the current weather in New York
 - [How-to Guides](https://langchain-ai.github.io/langgraphjs/how-tos/): Accomplish specific things within LangGraph, from streaming, to adding memory & persistence, to common design patterns (branching, subgraphs, etc.). These are the place to go if you want to copy and run a specific code snippet.
 - [Conceptual Guides](https://langchain-ai.github.io/langgraphjs/concepts/): In-depth explanations of the key concepts and principles behind LangGraph, such as nodes, edges, state and more.
 - [API Reference](https://langchain-ai.github.io/langgraphjs/reference/graphs/): Review important classes and methods, simple examples of how to use the graph and checkpointing APIs, higher-level prebuilt components and more.
+
+## Running Example Juypter Notebooks
+
+Please note that the *.ipynb notebooks in the `examples/` folder require [tslab](https://github.com/yunabe/tslab?tab=readme-ov-file) to be installed. In order to run these notebooks in VSCode, you will also need the [Jupyter](https://marketplace.visualstudio.com/items?itemName=ms-toolsai.jupyter) VSCode Extension installed. After cloning this repository, you can run `yarn build` in the root. You should then be all set!
+
+If you are still having trouble, try adding the following `tsconfig.json` file to the `examples/` directory:
+
+```
+{
+  "compilerOptions": {
+    "esModuleInterop": true,
+    "moduleResolution": "node",
+    "target": "ES2020",
+    "module": "ES2020",
+    "lib": [
+      "ES2020"
+    ],
+    "strict": true,
+    "baseUrl": ".",
+    "paths": {
+      "@langchain/langgraph": [
+        "../langgraph/src"
+      ]
+    }
+  },
+  "include": [
+    "./**/*.ts",
+    "./**/*.tsx"
+  ],
+  "exclude": [
+    "node_modules"
+  ]
+}
+```

package/dist/channels/base.cjs

@@ -5,6 +5,20 @@ const base_js_1 = require("../checkpoint/base.cjs");
 const id_js_1 = require("../checkpoint/id.cjs");
 const errors_js_1 = require("../errors.cjs");
 class BaseChannel {
+    constructor() {
+        Object.defineProperty(this, "ValueType", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "UpdateType", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+    }
 }
 exports.BaseChannel = BaseChannel;
 function emptyChannels(channels, checkpoint) {

package/dist/channels/base.d.ts

@@ -1,6 +1,8 @@
 import { ReadonlyCheckpoint } from "../checkpoint/base.js";
 import { Checkpoint } from "../checkpoint/index.js";
 export declare abstract class BaseChannel<ValueType = unknown, UpdateType = unknown, CheckpointType = unknown> {
+    ValueType: ValueType;
+    UpdateType: UpdateType;
     /**
      * The name of the channel.
      */

package/dist/channels/base.js

@@ -2,6 +2,20 @@ import { deepCopy } from "../checkpoint/base.js";
 import { uuid6 } from "../checkpoint/id.js";
 import { EmptyChannelError } from "../errors.js";
 export class BaseChannel {
+    constructor() {
+        Object.defineProperty(this, "ValueType", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "UpdateType", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+    }
 }
 export function emptyChannels(channels, checkpoint) {
     const newChannels = {};

package/dist/graph/message.d.ts

@@ -2,7 +2,7 @@ import { BaseMessage, BaseMessageLike } from "@langchain/core/messages";
 import { StateGraph } from "./state.js";
 type Messages = Array<BaseMessage | BaseMessageLike> | BaseMessage | BaseMessageLike;
 export declare function messagesStateReducer(left: Messages, right: Messages): BaseMessage[];
-export declare class MessageGraph extends StateGraph<BaseMessage[], Messages> {
+export declare class MessageGraph extends StateGraph<BaseMessage[], BaseMessage[], Messages> {
     constructor();
 }
 export interface MessagesState {

package/dist/graph/state.cjs

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.CompiledStateGraph = exports.StateGraph = void 0;
+exports.CompiledStateGraph = exports.StateGraph = exports.Annotation = void 0;
 const binop_js_1 = require("../channels/binop.cjs");
 const graph_js_1 = require("./graph.cjs");
 const last_value_js_1 = require("../channels/last_value.cjs");
@@ -12,6 +12,16 @@ const utils_js_1 = require("../utils.cjs");
 const constants_js_1 = require("../constants.cjs");
 const errors_js_1 = require("../errors.cjs");
 const ROOT = "__root__";
+function Annotation(annotation) {
+    if (annotation) {
+        return getChannel(annotation);
+    }
+    else {
+        // @ts-expect-error - Annotation without reducer
+        return new last_value_js_1.LastValue();
+    }
+}
+exports.Annotation = Annotation;
 class StateGraph extends graph_js_1.Graph {
     constructor(fields) {
         super();
@@ -28,7 +38,20 @@ class StateGraph extends graph_js_1.Graph {
             writable: true,
             value: new Set()
         });
-        this.channels = _getChannels(fields.channels);
+        if (isStateDefinition(fields)) {
+            this.channels = {};
+            for (const [key, val] of Object.entries(fields)) {
+                if (typeof val === "function") {
+                    this.channels[key] = val();
+                }
+                else {
+                    this.channels[key] = val;
+                }
+            }
+        }
+        else {
+            this.channels = _getChannels(fields.channels);
+        }
         for (const c of Object.values(this.channels)) {
             if (c.lc_graph_name === "BinaryOperatorAggregate") {
                 this.supportMultipleEdges = true;
@@ -147,6 +170,7 @@ function getChannel(reducer) {
         reducer.value) {
         return new binop_js_1.BinaryOperatorAggregate(reducer.value, reducer.default);
     }
+    // @ts-expect-error - Annotation without reducer
     return new last_value_js_1.LastValue();
 }
 class CompiledStateGraph extends graph_js_1.CompiledGraph {
@@ -266,3 +290,13 @@ class CompiledStateGraph extends graph_js_1.CompiledGraph {
     }
 }
 exports.CompiledStateGraph = CompiledStateGraph;
+function isBaseChannel(obj) {
+    return obj != null && typeof obj.lc_graph_name === "string";
+}
+function isStateDefinition(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        !Array.isArray(obj) &&
+        Object.keys(obj).length > 0 &&
+        Object.values(obj).every((v) => typeof v === "function" || isBaseChannel(v)));
+}

package/dist/graph/state.d.ts

@@ -1,9 +1,23 @@
 import { Runnable, RunnableConfig, RunnableLike } from "@langchain/core/runnables";
 import { BaseChannel } from "../channels/base.js";
-import { BinaryOperator } from "../channels/binop.js";
+import { BinaryOperator, BinaryOperatorAggregate } from "../channels/binop.js";
 import { END, CompiledGraph, Graph, START, Branch } from "./graph.js";
+import { LastValue } from "../channels/last_value.js";
 import { BaseCheckpointSaver } from "../checkpoint/base.js";
 import { All } from "../pregel/types.js";
+export declare function Annotation<ValueType>(): LastValue<ValueType>;
+export declare function Annotation<ValueType, UpdateType = ValueType>(annotation: SingleReducer<ValueType, UpdateType>): BinaryOperatorAggregate<ValueType, UpdateType>;
+interface StateDefinition {
+    [key: string]: BaseChannel | (() => BaseChannel);
+}
+type ExtractValueType<C> = C extends BaseChannel ? C["ValueType"] : C extends () => BaseChannel ? ReturnType<C>["ValueType"] : never;
+type ExtractUpdateType<C> = C extends BaseChannel ? C["UpdateType"] : C extends () => BaseChannel ? ReturnType<C>["UpdateType"] : never;
+export type StateInterface<S extends StateDefinition> = {
+    [key in keyof S]: ExtractValueType<S[key]>;
+};
+export type UpdateInterface<S extends StateDefinition> = {
+    [key in keyof S]?: ExtractUpdateType<S[key]>;
+};
 type SingleReducer<ValueType, UpdateType = ValueType> = {
     reducer: BinaryOperator<ValueType, UpdateType>;
     default?: () => ValueType;
@@ -24,24 +38,24 @@ export interface StateGraphArgs<Channels extends object | unknown> {
         __root__: Channels;
     }>;
 }
-export declare class StateGraph<State extends object | unknown, Update extends object | unknown = Partial<Record<keyof State, any>>, N extends string = typeof START> extends Graph<N, State, Update> {
+export declare class StateGraph<SD extends StateDefinition | unknown, S = SD extends StateDefinition ? StateInterface<SD> : SD, U = SD extends StateDefinition ? UpdateInterface<SD> : Partial<S>, N extends string = typeof START> extends Graph<N, S, U> {
     channels: Record<string, BaseChannel>;
     waitingEdges: Set<[N[], N]>;
-    constructor(fields: StateGraphArgs<State>);
+    constructor(fields: SD extends StateDefinition ? SD | StateGraphArgs<S> : StateGraphArgs<S>);
     get allEdges(): Set<[string, string]>;
-    addNode<K extends string>(key: K, action: RunnableLike<State, Update>): StateGraph<State, Update, N | K>;
+    addNode<K extends string>(key: K, action: RunnableLike<S, U>): StateGraph<SD, S, U, N | K>;
     addEdge(startKey: typeof START | N | N[], endKey: N | typeof END): this;
     compile({ checkpointer, interruptBefore, interruptAfter, }?: {
         checkpointer?: BaseCheckpointSaver;
         interruptBefore?: N[] | All;
         interruptAfter?: N[] | All;
-    }): CompiledStateGraph<State, Update, N>;
+    }): CompiledStateGraph<S, U, N>;
 }
-export declare class CompiledStateGraph<State extends object | unknown, Update extends object | unknown = Partial<State>, N extends string = typeof START> extends CompiledGraph<N, State, Update> {
-    builder: StateGraph<State, Update, N>;
+export declare class CompiledStateGraph<S, U, N extends string = typeof START> extends CompiledGraph<N, S, U> {
+    builder: StateGraph<unknown, S, U, N>;
     attachNode(key: typeof START, node?: never): void;
-    attachNode(key: N, node: Runnable<State, Update, RunnableConfig>): void;
+    attachNode(key: N, node: Runnable<S, U, RunnableConfig>): void;
     attachEdge(start: N | N[] | "__start__", end: N | "__end__"): void;
-    attachBranch(start: N | typeof START, name: string, branch: Branch<State, N>): void;
+    attachBranch(start: N | typeof START, name: string, branch: Branch<S, N>): void;
 }
 export {};

package/dist/graph/state.js

@@ -9,6 +9,15 @@ import { RunnableCallable } from "../utils.js";
 import { TAG_HIDDEN } from "../constants.js";
 import { InvalidUpdateError } from "../errors.js";
 const ROOT = "__root__";
+export function Annotation(annotation) {
+    if (annotation) {
+        return getChannel(annotation);
+    }
+    else {
+        // @ts-expect-error - Annotation without reducer
+        return new LastValue();
+    }
+}
 export class StateGraph extends Graph {
     constructor(fields) {
         super();
@@ -25,7 +34,20 @@ export class StateGraph extends Graph {
             writable: true,
             value: new Set()
         });
-        this.channels = _getChannels(fields.channels);
+        if (isStateDefinition(fields)) {
+            this.channels = {};
+            for (const [key, val] of Object.entries(fields)) {
+                if (typeof val === "function") {
+                    this.channels[key] = val();
+                }
+                else {
+                    this.channels[key] = val;
+                }
+            }
+        }
+        else {
+            this.channels = _getChannels(fields.channels);
+        }
         for (const c of Object.values(this.channels)) {
             if (c.lc_graph_name === "BinaryOperatorAggregate") {
                 this.supportMultipleEdges = true;
@@ -143,6 +165,7 @@ function getChannel(reducer) {
         reducer.value) {
         return new BinaryOperatorAggregate(reducer.value, reducer.default);
     }
+    // @ts-expect-error - Annotation without reducer
     return new LastValue();
 }
 export class CompiledStateGraph extends CompiledGraph {
@@ -261,3 +284,13 @@ export class CompiledStateGraph extends CompiledGraph {
         }
     }
 }
+function isBaseChannel(obj) {
+    return obj != null && typeof obj.lc_graph_name === "string";
+}
+function isStateDefinition(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        !Array.isArray(obj) &&
+        Object.keys(obj).length > 0 &&
+        Object.values(obj).every((v) => typeof v === "function" || isBaseChannel(v)));
+}

package/dist/prebuilt/agent_executor.d.ts

@@ -16,5 +16,5 @@ export interface AgentExecutorState {
 export declare function createAgentExecutor({ agentRunnable, tools, }: {
     agentRunnable: Runnable;
     tools: Array<Tool> | ToolExecutor;
-}): import("../graph/state.js").CompiledStateGraph<AgentExecutorState, Partial<Record<keyof AgentExecutorState, any>>, "__start__" | "agent" | "action">;
+}): import("../graph/state.js").CompiledStateGraph<AgentExecutorState, Partial<AgentExecutorState>, "__start__" | "agent" | "action">;
 export {};

package/dist/pregel/index.cjs

@@ -465,7 +465,7 @@ class Pregel extends runnables_1.Runnable {
                 // execute tasks, and wait for one to fail or all to finish.
                 // each task is independent from all other concurrent tasks
                 const tasks = tasksWithConfig.map(([proc, input, updatedConfig]) => () => proc.invoke(input, updatedConfig));
-                await executeTasks(tasks, this.stepTimeout);
+                await executeTasks(tasks, this.stepTimeout, config.signal);
                 // combine pending writes from all tasks
                 const pendingWrites = [];
                 for (const task of nextTasks) {
@@ -536,28 +536,33 @@ class Pregel extends runnables_1.Runnable {
     }
 }
 exports.Pregel = Pregel;
-function timeout(ms) {
-    return new Promise((reject) => {
-        setTimeout(reject, ms);
-    });
-}
-async function executeTasks(tasks, stepTimeout) {
-    // Wrap each task in a Promise that respects the step timeout
-    const wrappedTasks = tasks.map((task) => stepTimeout
-        ? Promise.race([
-            task(),
-            stepTimeout ? timeout(stepTimeout) : Promise.resolve(),
-        ])
-        : task());
-    // Wait for all tasks to settle
-    const results = await Promise.allSettled(wrappedTasks);
-    // Process the results
-    for (const result of results) {
-        if (result.status === "rejected") {
-            // If any task failed, cancel all pending tasks and throw the error
-            throw result.reason;
+async function executeTasks(tasks, stepTimeout, signal) {
+    if (stepTimeout && signal) {
+        if ("any" in AbortSignal) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            signal = AbortSignal.any([
+                signal,
+                AbortSignal.timeout(stepTimeout),
+            ]);
         }
     }
+    else if (stepTimeout) {
+        signal = AbortSignal.timeout(stepTimeout);
+    }
+    // Abort if signal is aborted
+    signal?.throwIfAborted();
+    // Start all tasks
+    const started = tasks.map((task) => task());
+    // Wait for all tasks to settle
+    // If any tasks fail, or signal is aborted, the promise will reject
+    await Promise.all(signal
+        ? [
+            ...started,
+            new Promise((_resolve, reject) => {
+                signal?.addEventListener("abort", () => reject(new Error("Abort")));
+            }),
+        ]
+        : started);
 }
 function _shouldInterrupt(checkpoint, interruptNodes, snapshotChannels, tasks) {
     const anySnapshotChannelUpdated = snapshotChannels.some((chan) => (0, base_js_2.getChannelVersion)(checkpoint, chan) >

package/dist/pregel/index.js

@@ -461,7 +461,7 @@ export class Pregel extends Runnable {
                 // execute tasks, and wait for one to fail or all to finish.
                 // each task is independent from all other concurrent tasks
                 const tasks = tasksWithConfig.map(([proc, input, updatedConfig]) => () => proc.invoke(input, updatedConfig));
-                await executeTasks(tasks, this.stepTimeout);
+                await executeTasks(tasks, this.stepTimeout, config.signal);
                 // combine pending writes from all tasks
                 const pendingWrites = [];
                 for (const task of nextTasks) {
@@ -531,28 +531,33 @@ export class Pregel extends Runnable {
         }
     }
 }
-function timeout(ms) {
-    return new Promise((reject) => {
-        setTimeout(reject, ms);
-    });
-}
-async function executeTasks(tasks, stepTimeout) {
-    // Wrap each task in a Promise that respects the step timeout
-    const wrappedTasks = tasks.map((task) => stepTimeout
-        ? Promise.race([
-            task(),
-            stepTimeout ? timeout(stepTimeout) : Promise.resolve(),
-        ])
-        : task());
-    // Wait for all tasks to settle
-    const results = await Promise.allSettled(wrappedTasks);
-    // Process the results
-    for (const result of results) {
-        if (result.status === "rejected") {
-            // If any task failed, cancel all pending tasks and throw the error
-            throw result.reason;
+async function executeTasks(tasks, stepTimeout, signal) {
+    if (stepTimeout && signal) {
+        if ("any" in AbortSignal) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            signal = AbortSignal.any([
+                signal,
+                AbortSignal.timeout(stepTimeout),
+            ]);
         }
     }
+    else if (stepTimeout) {
+        signal = AbortSignal.timeout(stepTimeout);
+    }
+    // Abort if signal is aborted
+    signal?.throwIfAborted();
+    // Start all tasks
+    const started = tasks.map((task) => task());
+    // Wait for all tasks to settle
+    // If any tasks fail, or signal is aborted, the promise will reject
+    await Promise.all(signal
+        ? [
+            ...started,
+            new Promise((_resolve, reject) => {
+                signal?.addEventListener("abort", () => reject(new Error("Abort")));
+            }),
+        ]
+        : started);
 }
 export function _shouldInterrupt(checkpoint, interruptNodes, snapshotChannels, tasks) {
     const anySnapshotChannelUpdated = snapshotChannels.some((chan) => getChannelVersion(checkpoint, chan) >