feat(langgraph): add pushMessage method (#1657)

.changeset/calm-doodles-shave.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Add `pushMessage` method for manually publishing to messages stream channel

libs/langgraph/src/graph/index.ts

@@ -16,6 +16,7 @@ export {
 export {
   MessageGraph,
   messagesStateReducer,
+  pushMessage,
   REMOVE_ALL_MESSAGES,
   type Messages,
 } from "./message.js";

libs/langgraph/src/graph/message.test.ts

@@ -223,7 +223,7 @@ describe("pushMessage", () => {
         },
       },
     };
-    pushMessage(message, config, { stateKey: "custom" });
+    pushMessage(message, { ...config, stateKey: "custom" });
   });
 
   it("should push messages in graph", async () => {

libs/langgraph/src/graph/message.ts

@@ -3,9 +3,10 @@ import {
   BaseMessageLike,
   coerceMessageLikeToMessage,
 } from "@langchain/core/messages";
+import type { RunnableConfig } from "@langchain/core/runnables";
+import { AsyncLocalStorageProviderSingleton } from "@langchain/core/singletons";
 import { v4 } from "uuid";
 import { StateGraph } from "./state.js";
-import type { LangGraphRunnableConfig } from "../pregel/runnable_types.js";
 import type { StreamMessagesHandler } from "../pregel/messages.js";
 
 export const REMOVE_ALL_MESSAGES = "__remove_all__";
@@ -100,15 +101,40 @@ export class MessageGraph extends StateGraph<
   }
 }
 
+/**
+ * Manually push a message to a message stream.
+ *
+ * This is useful when you need to push a manually created message before the node
+ * has finished executing.
+ *
+ * When a message is pushed, it will be automatically persisted to the state after the node has finished executing.
+ * To disable persisting, set `options.stateKey` to `null`.
+ *
+ * @param message The message to push. The message must have an ID set, otherwise an error will be thrown.
+ * @param options RunnableConfig / Runtime coming from node context.
+ */
 export function pushMessage(
   message: BaseMessage | BaseMessageLike,
-  config: LangGraphRunnableConfig,
-  options?: { stateKey?: string | null }
-) {
-  let stateKey: string | undefined = options?.stateKey ?? "messages";
-  if (options?.stateKey === null) {
-    stateKey = undefined;
+  options?: RunnableConfig & {
+    /**
+     * The key of the state to push the message to. Set to `null` to avoid persisting.
+     * @default "messages"
+     */
+    stateKey?: string | null;
   }
+) {
+  const rawOptions:
+    | (RunnableConfig & { stateKey?: string | null })
+    | undefined =
+    options ?? AsyncLocalStorageProviderSingleton.getRunnableConfig();
+
+  if (rawOptions == null) {
+    throw new Error("Calling pushMessage outside the context of a graph.");
+  }
+
+  const { stateKey: userStateKey, ...config } = rawOptions;
+  let stateKey: string | undefined = userStateKey ?? "messages";
+  if (userStateKey === null) stateKey = undefined;
 
   // coerce to message
   const validMessage = coerceMessageLikeToMessage(message);

libs/langgraph/src/index.ts

@@ -9,6 +9,7 @@ export * from "./web.js";
 
 export { interrupt } from "./interrupt.js";
 export { writer } from "./writer.js";
+export { pushMessage } from "./graph/message.js";
 export { getStore, getWriter, getConfig } from "./pregel/utils/config.js";
 export { getPreviousState } from "./func/index.js";
 export { getCurrentTaskInput } from "./pregel/utils/config.js";