Cleanup

2026-09-08 10:47:52 +02:00 · 2025-07-01 01:19:07 +02:00
parent c7bbb26ac0
commit 3d88f75254
2 changed files with 186 additions and 160 deletions
@@ -31,7 +31,7 @@ import type {
 } from "../types.stream.js";

 import {
-  type MutableRefObject,
+  type RefObject,
  useCallback,
  useEffect,
  useMemo,
@@ -316,8 +316,8 @@ function fetchHistory<StateType extends Record<string, unknown>>(
 function useThreadHistory<StateType extends Record<string, unknown>>(
  threadId: string | undefined | null,
  client: Client,
-  clearCallbackRef: MutableRefObject<(() => void) | undefined>,
-  submittingRef: MutableRefObject<boolean>,
+  clearCallbackRef: RefObject<(() => void) | undefined>,
+  submittingRef: RefObject<boolean>,
 ) {
  const [history, setHistory] = useState<ThreadState<StateType>[]>([]);

@@ -515,13 +515,13 @@ export interface UseStreamOptions<
   * Callback that is called when the stream is stopped by the user.
   * Provides a mutate function to update the stream state immediately
   * without requiring a server roundtrip.
-   * 
+   *
   * @example
   * ```typescript
   * onStop: ({ mutate }) => {
   *   mutate((prev) => ({
   *     ...prev,
-   *     ui: prev.ui?.map(component => 
+   *     ui: prev.ui?.map(component =>
   *       component.props.isLoading
   *         ? { ...component, props: { ...component.props, stopped: true, isLoading: false }}
   *         : component
@@ -541,13 +541,6 @@ export interface UseStreamOptions<
   */
  threadId?: string | null;

-  /**
-   * The ID to use when creating a new thread. When provided, this ID will be used
-   * for thread creation when threadId is null. This enables optimistic UI updates
-   * where you know the thread ID before the thread is actually created.
-   */
-  newThreadId?: string;
-
  /**
   * Callback that is called when the thread ID is updated (ie when a new thread is created).
   */
@@ -560,12 +553,12 @@ export interface UseStreamOptions<
   * Initial values to display immediately when loading a thread.
   * Useful for displaying cached thread data while official history loads.
   * These values will be replaced when official thread data is fetched.
-   * 
-   * Note: UI components from initialValues will render immediately if they're 
+   *
+   * Note: UI components from initialValues will render immediately if they're
   * predefined in LoadExternalComponent's components prop, providing instant
   * cached UI display without server fetches.
   */
-  initialValues?: Partial<StateType> | null;
+  initialValues?: StateType | null;
 }

 interface RunMetadataStorage {
@@ -695,6 +688,61 @@ interface SubmitOptions<
   */
  streamSubgraphs?: boolean;
  streamResumable?: boolean;
+  /**
+   * The ID to use when creating a new thread. When provided, this ID will be used
+   * for thread creation when threadId is `null` or `undefined`.
+   * This enables optimistic UI updates where you know the thread ID
+   * before the thread is actually created.
+   */
+  threadId?: string;
+}
+
+function useStreamValuesState<StateType extends Record<string, unknown>>() {
+  type Kind = "stream" | "stop";
+  type Values = StateType | null;
+  type Update = Values | ((prev: Values, kind?: Kind) => Values);
+  type Mutate = Partial<StateType> | ((prev: StateType) => Partial<StateType>);
+
+  const [values, setValues] = useState<[values: StateType, kind: Kind] | null>(
+    null,
+  );
+
+  const setStreamValues = useCallback(
+    (values: Update, kind: Kind = "stream") => {
+      if (typeof values === "function") {
+        setValues((prevTuple) => {
+          const [prevValues, prevKind] = prevTuple ?? [null, "stream"];
+          const next = values(prevValues, prevKind);
+
+          if (next == null) return null;
+          return [next, kind] as [StateType, Kind];
+        });
+
+        return;
+      }
+
+      if (values == null) setValues(null);
+      setValues([values, kind] as [StateType, Kind]);
+    },
+    [],
+  );
+
+  const mutate = useCallback(
+    (kind: Kind, serverValues: StateType) => (update: Mutate) => {
+      setStreamValues((clientValues) => {
+        const prev = { ...serverValues, ...clientValues };
+        const next = typeof update === "function" ? update(prev) : update;
+        return { ...prev, ...next };
+      }, kind);
+    },
+    [setStreamValues],
+  );
+
+  return [values?.[0] ?? null, setStreamValues, mutate] as [
+    Values,
+    (update: Update, kind?: Kind) => void,
+    (kind: Kind, serverValues: StateType) => (update: Mutate) => void,
+  ];
 }

 export function useStream<
@@ -760,7 +808,8 @@ export function useStream<
  const [isLoading, setIsLoading] = useState(false);

  const [streamError, setStreamError] = useState<unknown>(undefined);
-  const [streamValues, setStreamValues] = useState<StateType | null>(null);
+  const [streamValues, setStreamValues, getMutateFn] =
+    useStreamValuesState<StateType>();

  const messageManagerRef = useRef(new MessageTupleManager());
  const submittingRef = useRef(false);
@@ -831,7 +880,9 @@ export function useStream<
  );

  const threadHead: ThreadState<StateType> | undefined = flatHistory.at(-1);
-  const historyValues = threadHead?.values ?? ({} as StateType);
+  const historyValues =
+    threadHead?.values ?? options.initialValues ?? ({} as StateType);
+
  const historyError = (() => {
    const error = threadHead?.tasks?.at(-1)?.error;
    if (error == null) return undefined;
@@ -887,21 +938,6 @@ export function useStream<
    );
  })();

-  // Create a reusable mutate function for both onCustomEvent and onStop callbacks
-  const mutateStreamValues = useCallback(
-    (update: Partial<StateType> | ((prev: StateType) => Partial<StateType>)) => {
-      setStreamValues((prev) => {
-        // should not happen
-        if (prev == null) return prev;
-        return {
-          ...prev,
-          ...(typeof update === "function" ? update(prev) : update),
-        };
-      });
-    },
-    [],
-  );
-
  const stop = () => {
    if (abortRef.current != null) abortRef.current.abort();
    abortRef.current = null;
@@ -912,10 +948,7 @@ export function useStream<
      runMetadataStorage.removeItem(`lg:stream:${threadId}`);
    }

-    // Call onStop callback with mutate function
-    if (options.onStop) {
-      options.onStop({ mutate: mutateStreamValues });
-    }
+    options?.onStop?.({ mutate: getMutateFn("stop", historyValues) });
  };

  async function consumeStream(
@@ -948,7 +981,7 @@ export function useStream<
        if (event === "updates") options.onUpdateEvent?.(data);
        if (event === "custom")
          options.onCustomEvent?.(data, {
-            mutate: mutateStreamValues,
+            mutate: getMutateFn("stream", historyValues),
          });
        if (event === "metadata") options.onMetadataEvent?.(data);
        if (event === "events") options.onLangChainEvent?.(data);
@@ -973,8 +1006,7 @@ export function useStream<
          }

          setStreamValues((streamValues) => {
-            const baseValues = options.initialValues ? { ...historyValues, ...options.initialValues } : historyValues;
-            const values = { ...baseValues, ...streamValues };
+            const values = { ...historyValues, ...streamValues };

            // Assumption: we're concatenating the message
            const messages = getMessages(values).slice();
@@ -991,8 +1023,11 @@ export function useStream<

      // TODO: stream created checkpoints to avoid an unnecessary network request
      const result = await run.onSuccess();
-
-      setStreamValues(null);
+      setStreamValues((values, kind) => {
+        // Do not clear out the user values set on `stop`.
+        if (kind === "stop") return values;
+        return null;
+      });
      if (streamError != null) throw streamError;

      const lastHead = result.at(0);
@@ -1050,27 +1085,23 @@ export function useStream<

      if (newPath != null) setBranch(newPath ?? "");

-      // Assumption: we're setting the initial value
-      // Used for instant feedback
      setStreamValues(() => {
-        const baseValues = options.initialValues ? { ...historyValues, ...options.initialValues } : historyValues;
-
        if (submitOptions?.optimisticValues != null) {
          return {
-            ...baseValues,
+            ...historyValues,
            ...(typeof submitOptions.optimisticValues === "function"
-              ? submitOptions.optimisticValues(baseValues)
+              ? submitOptions.optimisticValues(historyValues)
              : submitOptions.optimisticValues),
          };
        }

-        return baseValues;
+        return historyValues;
      });

      let usableThreadId = threadId;
      if (!usableThreadId) {
        const thread = await client.threads.create({
-          threadId: options.newThreadId,
+          threadId: submitOptions?.threadId,
        });
        onThreadId(thread.thread_id);
        usableThreadId = thread.thread_id;
@@ -1165,7 +1196,7 @@ export function useStream<
  }, [reconnectKey]);

  const error = streamError ?? historyError;
-  const values = streamValues ?? (options.initialValues ? { ...historyValues, ...options.initialValues } : historyValues);
+  const values = streamValues ?? historyValues;

  return {
    get values() {
@@ -6,15 +6,15 @@ import { userEvent } from "@testing-library/user-event";
 import { setupServer } from "msw/node";
 import { http } from "msw";
 import { useStream } from "../react/stream.js";
+import type { Message } from "../types.messages.js";

 import { StateGraph, MessagesAnnotation, START } from "@langchain/langgraph";
 import { MemorySaver } from "@langchain/langgraph-checkpoint";
 import { FakeStreamingChatModel } from "@langchain/core/utils/testing";
-import { AIMessage, BaseMessageLike } from "@langchain/core/messages";
-
-import { Hono } from "hono";
-import { logger } from "hono/logger";
+import { AIMessage } from "@langchain/core/messages";
 import { createEmbedServer } from "@langchain/langgraph-api/experimental/embed";
+import { randomUUID } from "node:crypto";
+import { useState } from "react";

 const threads = (() => {
  const THREADS: Record<
@@ -40,17 +40,14 @@ const checkpointer = new MemorySaver();

 const model = new FakeStreamingChatModel({ responses: [new AIMessage("Hey")] });
 const agent = new StateGraph(MessagesAnnotation)
-  .addNode("agent", async (state: { messages: BaseMessageLike[] }) => {
+  .addNode("agent", async (state: { messages: Message[] }) => {
    const response = await model.invoke(state.messages);
    return { messages: [response] };
  })
  .addEdge(START, "agent")
  .compile();

-const app = new Hono();
-app.use(logger());
-app.route("/", createEmbedServer({ graph: { agent }, checkpointer, threads }));
-
+const app = createEmbedServer({ graph: { agent }, checkpointer, threads });
 const server = setupServer(http.all("*", (ctx) => app.fetch(ctx.request)));

 function TestChatComponent() {
@@ -74,7 +71,12 @@ function TestChatComponent() {
        {isLoading ? "Loading..." : "Not loading"}
      </div>
      {error ? <div data-testid="error">{String(error)}</div> : null}
-      <button data-testid="submit" onClick={() => submit({})}>
+      <button
+        data-testid="submit"
+        onClick={() =>
+          submit({ messages: [{ content: "Hello", type: "human" }] })
+        }
+      >
        Send
      </button>
      <button data-testid="stop" onClick={stop}>
@@ -134,26 +136,35 @@ describe("useStream", () => {
    });
  });

-  it("displays initial values immediately", () => {
-    const initialValues = {
-      messages: [
-        { id: "cached-1", type: "human", content: "Cached user message" },
-        { id: "cached-2", type: "ai", content: "Cached AI response" },
-      ],
-    };
+  it("displays initial values immediately and clears them when submitting", async () => {
+    const user = userEvent.setup();

    function TestCachedComponent() {
-      const { messages, values } = useStream({
-        assistantId: "test-assistant",
+      const { messages, values, submit } = useStream<{
+        messages: Message[];
+      }>({
+        assistantId: "agent",
        apiKey: "test-api-key",
-        initialValues,
+        initialValues: {
+          messages: [
+            { id: "cached-1", type: "human", content: "Cached user message" },
+            { id: "cached-2", type: "ai", content: "Cached AI response" },
+          ],
+        },
      });

      return (
        <div>
          <div data-testid="messages">
            {messages.map((msg, i) => (
-              <div key={msg.id ?? i} data-testid={`message-${i}`}>
+              <div
+                key={msg.id ?? i}
+                data-testid={
+                  msg.id?.includes("cached")
+                    ? `message-cached-${i}`
+                    : `message-${i}`
+                }
+              >
                {typeof msg.content === "string"
                  ? msg.content
                  : JSON.stringify(msg.content)}
@@ -161,6 +172,14 @@ describe("useStream", () => {
            ))}
          </div>
          <div data-testid="values">{JSON.stringify(values)}</div>
+          <button
+            data-testid="submit"
+            onClick={() =>
+              submit({ messages: [{ content: "Hello", type: "human" }] })
+            }
+          >
+            Submit
+          </button>
        </div>
      );
    }
@@ -168,10 +187,10 @@ describe("useStream", () => {
    render(<TestCachedComponent />);

    // Should immediately show cached messages
-    expect(screen.getByTestId("message-0")).toHaveTextContent(
+    expect(screen.getByTestId("message-cached-0")).toHaveTextContent(
      "Cached user message",
    );
-    expect(screen.getByTestId("message-1")).toHaveTextContent(
+    expect(screen.getByTestId("message-cached-1")).toHaveTextContent(
      "Cached AI response",
    );

@@ -179,40 +198,30 @@ describe("useStream", () => {
    expect(screen.getByTestId("values")).toHaveTextContent(
      "Cached user message",
    );
+
+    // Submiting should clear out the cached messages
+    await user.click(screen.getByTestId("submit"));
+
+    // Wait for messages to appear
+    await waitFor(() => {
+      expect(screen.getByTestId("message-0")).toHaveTextContent("Hello");
+      expect(screen.getByTestId("message-1")).toHaveTextContent("Hey");
+    });
  });

-  it("handles null initial values", () => {
-    function TestNullInitialComponent() {
-      const { messages, values } = useStream({
-        assistantId: "test-assistant",
-        apiKey: "test-api-key",
-        initialValues: null,
-      });
+  it("accepts newThreadId option without errors", async () => {
+    const user = userEvent.setup();

-      return (
-        <div>
-          <div data-testid="message-count">{messages.length}</div>
-          <div data-testid="values">{JSON.stringify(values)}</div>
-        </div>
-      );
-    }
+    const spy = vi.fn();
+    const predeterminedThreadId = randomUUID();

-    render(<TestNullInitialComponent />);
-
-    // Should handle null initialValues gracefully
-    expect(screen.getByTestId("message-count")).toHaveTextContent("0");
-    expect(screen.getByTestId("values")).toHaveTextContent("{}");
-  });
-
-  it("accepts newThreadId option without errors", () => {
    // Test that newThreadId option can be passed without causing errors
    function TestNewThreadComponent() {
-      const stream = useStream({
-        assistantId: "test-assistant",
+      const stream = useStream<{ messages: Message[] }>({
+        assistantId: "agent",
        apiKey: "test-api-key",
        threadId: null, // Start with no thread
-        newThreadId: "predetermined-thread-id",
-        onThreadId: () => {}, // Mock callback
+        onThreadId: spy, // Mock callback
      });

      return (
@@ -223,6 +232,14 @@ describe("useStream", () => {
          <div data-testid="thread-id">
            {stream.client ? "Client ready" : "No client"}
          </div>
+          <button
+            data-testid="submit"
+            onClick={() =>
+              stream.submit({}, { threadId: predeterminedThreadId })
+            }
+          >
+            Submit
+          </button>
        </div>
      );
    }
@@ -232,55 +249,20 @@ describe("useStream", () => {
    // Should render without errors
    expect(screen.getByTestId("loading")).toHaveTextContent("Not loading");
    expect(screen.getByTestId("thread-id")).toHaveTextContent("Client ready");
+
+    await user.click(screen.getByTestId("submit"));
+    expect(spy).toHaveBeenCalledWith(predeterminedThreadId);
+    expect(await threads.get(predeterminedThreadId)).toEqual({
+      thread_id: predeterminedThreadId,
+      metadata: {
+        graph_id: "agent",
+        assistant_id: "agent",
+      },
+    });
  });

-  it("shows initial values before any streaming", () => {
-    const initialValues = {
-      messages: [
-        { id: "initial-1", type: "human", content: "Initial message" },
-      ],
-      customField: "initial-value",
-    };
-
-    function TestInitialValuesComponent() {
-      const stream = useStream({
-        assistantId: "test-assistant",
-        apiKey: "test-api-key",
-        initialValues,
-      });
-
-      return (
-        <div>
-          <div data-testid="message-0">
-            {typeof stream.messages[0]?.content === "string"
-              ? stream.messages[0]?.content
-              : JSON.stringify(stream.messages[0]?.content)}
-          </div>
-          <div data-testid="custom-field">{stream.values.customField}</div>
-          <div data-testid="loading">
-            {stream.isLoading ? "Loading" : "Not loading"}
-          </div>
-        </div>
-      );
-    }
-
-    render(<TestInitialValuesComponent />);
-
-    // Should immediately show initial values without any loading
-    expect(screen.getByTestId("message-0")).toHaveTextContent(
-      "Initial message",
-    );
-    expect(screen.getByTestId("custom-field")).toHaveTextContent(
-      "initial-value",
-    );
-    expect(screen.getByTestId("loading")).toHaveTextContent("Not loading");
-  });
-});
-
-describe("useStream onStop callback", () => {
-  const user = userEvent.setup();
-
-  it("calls onStop callback when stop is called", async () => {
+  it("onStop callback is called when stop is called", async () => {
+    const user = userEvent.setup();
    const onStopCallback = vi.fn();

    function TestComponent() {
@@ -317,16 +299,22 @@ describe("useStream onStop callback", () => {
    );
  });

-  it("mutate function updates stream values immediately", async () => {
+  it("onStop mutate function updates stream values immediately", async () => {
+    const user = userEvent.setup();
+
    function TestComponent() {
-      const { submit, stop, values } = useStream({
-        assistantId: "test-assistant",
+      const [stopped, setStopped] = useState(false);
+      const { submit, stop, messages } = useStream<{ messages: Message[] }>({
+        assistantId: "agent",
        apiKey: "test-api-key",
        onStop: ({ mutate }) => {
+          setStopped(true);
          mutate((prev) => ({
            ...prev,
-            stoppedByUser: true,
-            customMessage: "Stream stopped",
+            messages: [
+              ...(prev.messages ?? []),
+              { type: "ai", content: "Stream stopped" },
+            ],
          }));
        },
      });
@@ -334,10 +322,16 @@ describe("useStream onStop callback", () => {
      return (
        <div>
          <div data-testid="stopped-status">
-            {(values as any).stoppedByUser ? "Stopped" : "Not stopped"}
+            {stopped ? "Stopped" : "Not stopped"}
          </div>
-          <div data-testid="custom-message">
-            {(values as any).customMessage || "No message"}
+          <div data-testid="messages">
+            {messages.map((msg, i) => (
+              <div key={msg.id ?? i} data-testid={`message-${i}`}>
+                {typeof msg.content === "string"
+                  ? msg.content
+                  : JSON.stringify(msg.content)}
+              </div>
+            ))}
          </div>
          <button data-testid="submit" onClick={() => submit({})}>
            Send
@@ -355,9 +349,6 @@ describe("useStream onStop callback", () => {
    expect(screen.getByTestId("stopped-status")).toHaveTextContent(
      "Not stopped",
    );
-    expect(screen.getByTestId("custom-message")).toHaveTextContent(
-      "No message",
-    );

    // Start and stop stream
    await user.click(screen.getByTestId("submit"));
@@ -366,13 +357,15 @@ describe("useStream onStop callback", () => {
    // Verify state was updated immediately
    await waitFor(() => {
      expect(screen.getByTestId("stopped-status")).toHaveTextContent("Stopped");
-      expect(screen.getByTestId("custom-message")).toHaveTextContent(
+      expect(screen.getByTestId("message-0")).toHaveTextContent(
        "Stream stopped",
      );
    });
  });

-  it("handles functional updates correctly", async () => {
+  it("onStop handles functional updates correctly", async () => {
+    const user = userEvent.setup();
+
    function TestComponent() {
      const { submit, stop, values } = useStream({
        assistantId: "test-assistant",
@@ -423,7 +416,9 @@ describe("useStream onStop callback", () => {
    });
  });

-  it("is not called when stream completes naturally", async () => {
+  it("onStop is not called when stream completes naturally", async () => {
+    const user = userEvent.setup();
+
    const onStopCallback = vi.fn();

    function TestComponent() {