From 3d88f752547007fd14503ed7e53a24317675a208 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 1 Jul 2025 01:19:07 +0200 Subject: [PATCH] Cleanup --- libs/sdk-js/src/react/stream.tsx | 131 ++++++++++------ libs/sdk-js/src/tests/stream.test.tsx | 215 +++++++++++++------------- 2 files changed, 186 insertions(+), 160 deletions(-) diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx index 21da1c6ec..3d51497b4 100644 --- a/libs/sdk-js/src/react/stream.tsx +++ b/libs/sdk-js/src/react/stream.tsx @@ -31,7 +31,7 @@ import type { } from "../types.stream.js"; import { - type MutableRefObject, + type RefObject, useCallback, useEffect, useMemo, @@ -316,8 +316,8 @@ function fetchHistory>( function useThreadHistory>( threadId: string | undefined | null, client: Client, - clearCallbackRef: MutableRefObject<(() => void) | undefined>, - submittingRef: MutableRefObject, + clearCallbackRef: RefObject<(() => void) | undefined>, + submittingRef: RefObject, ) { const [history, setHistory] = useState[]>([]); @@ -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 | 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>() { + type Kind = "stream" | "stop"; + type Values = StateType | null; + type Update = Values | ((prev: Values, kind?: Kind) => Values); + type Mutate = Partial | ((prev: StateType) => Partial); + + 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(undefined); - const [streamValues, setStreamValues] = useState(null); + const [streamValues, setStreamValues, getMutateFn] = + useStreamValuesState(); const messageManagerRef = useRef(new MessageTupleManager()); const submittingRef = useRef(false); @@ -831,7 +880,9 @@ export function useStream< ); const threadHead: ThreadState | 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 | ((prev: StateType) => Partial)) => { - 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() { diff --git a/libs/sdk-js/src/tests/stream.test.tsx b/libs/sdk-js/src/tests/stream.test.tsx index a5588d805..e7597d69e 100644 --- a/libs/sdk-js/src/tests/stream.test.tsx +++ b/libs/sdk-js/src/tests/stream.test.tsx @@ -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"} {error ?
{String(error)}
: null} - ); } @@ -168,10 +187,10 @@ describe("useStream", () => { render(); // 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 ( -
-
{messages.length}
-
{JSON.stringify(values)}
-
- ); - } + const spy = vi.fn(); + const predeterminedThreadId = randomUUID(); - render(); - - // 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", () => {
{stream.client ? "Client ready" : "No client"}
+ ); } @@ -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 ( -
-
- {typeof stream.messages[0]?.content === "string" - ? stream.messages[0]?.content - : JSON.stringify(stream.messages[0]?.content)} -
-
{stream.values.customField}
-
- {stream.isLoading ? "Loading" : "Not loading"} -
-
- ); - } - - render(); - - // 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 (
- {(values as any).stoppedByUser ? "Stopped" : "Not stopped"} + {stopped ? "Stopped" : "Not stopped"}
-
- {(values as any).customMessage || "No message"} +
+ {messages.map((msg, i) => ( +
+ {typeof msg.content === "string" + ? msg.content + : JSON.stringify(msg.content)} +
+ ))}