@assistant-ui/react-langchain 0.0.15 → 0.0.18

package/src/useLangChainState.test.tsx

@@ -22,7 +22,7 @@ vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
   };
 });
 
-import { useLangChainState } from "./useStreamRuntime";
+import { useLangChainState } from "./hooks";
 
 const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
 

package/src/useLangChainStream.test.tsx

@@ -0,0 +1,42 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { renderHook } from "@testing-library/react";
+import {
+  createRunSelectorAgainst,
+  makeExtras,
+  type Selector,
+} from "./__tests__/langChainTestUtils";
+
+const { mockUseAuiState } = vi.hoisted(() => ({
+  mockUseAuiState: vi.fn(),
+}));
+
+vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    useAuiState: ((selector: Selector) =>
+      mockUseAuiState(selector)) as typeof actual.useAuiState,
+    useAui: (() => ({})) as unknown as typeof actual.useAui,
+  };
+});
+
+import { useLangChainStream } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainStream", () => {
+  it("returns undefined when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainStream());
+    expect(result.current).toBeUndefined();
+  });
+
+  it("returns the stream handle from extras", () => {
+    const stream = { submit: () => {} };
+    runSelectorAgainst(makeExtras({ stream }));
+    const { result } = renderHook(() => useLangChainStream());
+    expect(result.current).toBe(stream);
+  });
+});

package/src/useLangChainSubagents.test.tsx

@@ -0,0 +1,48 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { renderHook } from "@testing-library/react";
+import {
+  createRunSelectorAgainst,
+  makeExtras,
+  type Selector,
+} from "./__tests__/langChainTestUtils";
+
+const { mockUseAuiState } = vi.hoisted(() => ({
+  mockUseAuiState: vi.fn(),
+}));
+
+vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    useAuiState: ((selector: Selector) =>
+      mockUseAuiState(selector)) as typeof actual.useAuiState,
+    useAui: (() => ({})) as unknown as typeof actual.useAui,
+  };
+});
+
+import { useLangChainSubagents } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainSubagents", () => {
+  it("returns the stable empty map when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainSubagents());
+    expect(result.current.size).toBe(0);
+  });
+
+  it("returns the subagents map from extras", () => {
+    const subagents = new Map([["researcher", { namespace: "researcher" }]]);
+    runSelectorAgainst(makeExtras({ subagents }));
+    const { result } = renderHook(() => useLangChainSubagents());
+    expect(result.current).toBe(subagents);
+  });
+
+  it("returns the stable empty map when extras carry no subagents", () => {
+    runSelectorAgainst(makeExtras({ subagents: undefined }));
+    const { result } = renderHook(() => useLangChainSubagents());
+    expect(result.current.size).toBe(0);
+  });
+});

package/src/useLangChainSubgraphs.test.tsx

@@ -0,0 +1,48 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { renderHook } from "@testing-library/react";
+import {
+  createRunSelectorAgainst,
+  makeExtras,
+  type Selector,
+} from "./__tests__/langChainTestUtils";
+
+const { mockUseAuiState } = vi.hoisted(() => ({
+  mockUseAuiState: vi.fn(),
+}));
+
+vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    useAuiState: ((selector: Selector) =>
+      mockUseAuiState(selector)) as typeof actual.useAuiState,
+    useAui: (() => ({})) as unknown as typeof actual.useAui,
+  };
+});
+
+import { useLangChainSubgraphs } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainSubgraphs", () => {
+  it("returns the stable empty map when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainSubgraphs());
+    expect(result.current.size).toBe(0);
+  });
+
+  it("returns the subgraphs map from extras", () => {
+    const subgraphs = new Map([["planner", { namespace: "planner" }]]);
+    runSelectorAgainst(makeExtras({ subgraphs }));
+    const { result } = renderHook(() => useLangChainSubgraphs());
+    expect(result.current).toBe(subgraphs);
+  });
+
+  it("returns the stable empty map when extras carry no subgraphs", () => {
+    runSelectorAgainst(makeExtras({ subgraphs: undefined }));
+    const { result } = renderHook(() => useLangChainSubgraphs());
+    expect(result.current.size).toBe(0);
+  });
+});

package/src/useLangChainToolCalls.test.tsx

@@ -22,7 +22,7 @@ vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
   };
 });
 
-import { useLangChainToolCalls } from "./useStreamRuntime";
+import { useLangChainToolCalls } from "./hooks";
 
 const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
 

package/src/useStreamRuntime.test.tsx

@@ -0,0 +1,232 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
+import { useAui } from "@assistant-ui/store";
+import type { LangChainBaseMessage } from "./types";
+import type { ReactNode } from "react";
+
+const { mockUseChannel, mockUseStream, streamController } = vi.hoisted(() => ({
+  mockUseChannel: vi.fn(() => []),
+  mockUseStream: vi.fn(),
+  streamController: Symbol("STREAM_CONTROLLER"),
+}));
+
+vi.mock("@langchain/react", () => ({
+  STREAM_CONTROLLER: streamController,
+  useChannel: mockUseChannel,
+  useStream: mockUseStream,
+}));
+
+import { useStreamRuntime } from "./useStreamRuntime";
+
+type MockStream = {
+  messages: LangChainBaseMessage[];
+  isLoading: boolean;
+  isThreadLoading: boolean;
+  values: Record<string, unknown>;
+  interrupts: unknown[];
+  toolCalls: unknown[];
+  subagents: unknown[];
+  subgraphs: unknown[];
+  error: unknown;
+  submit: ReturnType<typeof vi.fn>;
+  respond: ReturnType<typeof vi.fn>;
+  respondAll: ReturnType<typeof vi.fn>;
+  interrupt: ReturnType<typeof vi.fn>;
+  stop: ReturnType<typeof vi.fn>;
+  client: Record<string, unknown>;
+  [streamController]: {
+    messageMetadataStore: {
+      getSnapshot: ReturnType<typeof vi.fn>;
+    };
+  };
+};
+
+const message = (
+  id: string,
+  type: "human" | "ai",
+  content: string,
+): LangChainBaseMessage & { id: string } => ({
+  id,
+  _getType: () => type,
+  content,
+});
+
+const createMockStream = (
+  messages: LangChainBaseMessage[] = [],
+): MockStream => ({
+  messages,
+  isLoading: false,
+  isThreadLoading: false,
+  values: {},
+  interrupts: [],
+  toolCalls: [],
+  subagents: [],
+  subgraphs: [],
+  error: undefined,
+  submit: vi.fn(async () => {}),
+  respond: vi.fn(),
+  respondAll: vi.fn(),
+  interrupt: vi.fn(),
+  stop: vi.fn(),
+  client: {},
+  [streamController]: {
+    messageMetadataStore: {
+      getSnapshot: vi.fn(),
+    },
+  },
+});
+
+const renderRuntime = (stream: MockStream) => {
+  mockUseStream.mockReturnValue(stream);
+  return renderHook(() => useStreamRuntime({ apiUrl: "/api" } as never));
+};
+
+const renderAui = (stream: MockStream) => {
+  const runtimeHook = renderRuntime(stream);
+  const Wrapper = ({ children }: { children: ReactNode }) => (
+    <AssistantRuntimeProvider runtime={runtimeHook.result.current}>
+      {children}
+    </AssistantRuntimeProvider>
+  );
+  Wrapper.displayName = "TestWrapper";
+  const auiHook = renderHook(() => useAui(), { wrapper: Wrapper });
+  return {
+    auiResult: auiHook.result,
+    rerender: () => {
+      runtimeHook.rerender();
+      auiHook.rerender();
+    },
+  };
+};
+
+const getText = (aui: ReturnType<typeof useAui>) =>
+  aui
+    .thread()
+    .getState()
+    .messages.map((m) =>
+      m.content
+        .filter((part) => part.type === "text")
+        .map((part) => part.text)
+        .join(""),
+    );
+
+describe("useStreamRuntime staged messages", () => {
+  it("stages a new user message without submitting when startRun is false", async () => {
+    const stream = createMockStream([message("u1", "human", "earlier")]);
+    const { auiResult } = renderAui(stream);
+
+    await act(async () => {
+      auiResult.current.thread().append({
+        role: "user",
+        content: [{ type: "text", text: "draft" }],
+        startRun: false,
+      });
+    });
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual(["earlier", "draft"]);
+    });
+    expect(stream.submit).not.toHaveBeenCalled();
+  });
+
+  it("keeps a staged edit truncated when stream messages update before promotion", async () => {
+    const stream = createMockStream([
+      message("u1", "human", "first"),
+      message("a1", "ai", "first answer"),
+      message("u2", "human", "second"),
+    ]);
+    const { auiResult, rerender } = renderAui(stream);
+
+    await act(async () => {
+      auiResult.current.thread().append({
+        role: "user",
+        parentId: "u1",
+        content: [{ type: "text", text: "edited" }],
+        startRun: false,
+      });
+    });
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual(["first", "edited"]);
+    });
+
+    stream.messages = [
+      message("u1", "human", "first"),
+      message("a1", "ai", "first answer from refresh"),
+      message("u2", "human", "second from refresh"),
+    ];
+    rerender();
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual(["first", "edited"]);
+    });
+    expect(stream.submit).not.toHaveBeenCalled();
+  });
+
+  it("keeps later staged messages visible after promoting one staged parent", async () => {
+    const stream = createMockStream([message("u1", "human", "earlier")]);
+    const { auiResult, rerender } = renderAui(stream);
+
+    await act(async () => {
+      auiResult.current.thread().append({
+        role: "user",
+        content: [{ type: "text", text: "first staged" }],
+        startRun: false,
+      });
+      auiResult.current.thread().append({
+        role: "user",
+        content: [{ type: "text", text: "second staged" }],
+        startRun: false,
+      });
+    });
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual([
+        "earlier",
+        "first staged",
+        "second staged",
+      ]);
+    });
+
+    const firstStagedId = auiResult.current.thread().getState().messages[1]!.id;
+    await act(async () => {
+      await auiResult.current.thread().startRun({
+        parentId: firstStagedId,
+        sourceId: null,
+        runConfig: {},
+      });
+    });
+
+    expect(stream.submit).toHaveBeenCalledWith(
+      {
+        messages: [
+          expect.objectContaining({
+            id: firstStagedId,
+            type: "human",
+            content: "first staged",
+          }),
+        ],
+      },
+      undefined,
+    );
+
+    stream.messages = [
+      message("u1", "human", "earlier"),
+      message(firstStagedId, "human", "first staged"),
+      message("a1", "ai", "answer"),
+    ];
+    rerender();
+
+    await waitFor(() => {
+      expect(getText(auiResult.current)).toEqual([
+        "earlier",
+        "first staged",
+        "answer",
+        "second staged",
+      ]);
+    });
+  });
+});