@assistant-ui/react-langchain 0.0.14 → 0.0.17

package/src/useLangChainRespondAll.test.tsx

@@ -0,0 +1,52 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { renderHook } from "@testing-library/react";
+import { makeExtras } from "./__tests__/langChainTestUtils";
+
+const { mockUseAui } = vi.hoisted(() => ({
+  mockUseAui: vi.fn(),
+}));
+
+vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    useAui: (() => mockUseAui()) as unknown as typeof actual.useAui,
+  };
+});
+
+import { useLangChainRespondAll } from "./hooks";
+
+describe("useLangChainRespondAll", () => {
+  it("forwards the responses and options to stream.respondAll", async () => {
+    const respondAll = vi.fn().mockResolvedValue(undefined);
+    mockUseAui.mockReturnValue({
+      thread: () => ({
+        getState: () => ({ extras: makeExtras({ respondAll }) }),
+      }),
+    });
+
+    const { result } = renderHook(() => useLangChainRespondAll());
+    await result.current(
+      { a: { approved: true }, b: { approved: false } },
+      { metadata: { src: "ui" } },
+    );
+
+    expect(respondAll).toHaveBeenCalledWith(
+      { a: { approved: true }, b: { approved: false } },
+      { metadata: { src: "ui" } },
+    );
+  });
+
+  it("throws when extras are absent or unbranded", () => {
+    mockUseAui.mockReturnValue({
+      thread: () => ({ getState: () => ({ extras: { respondAll: vi.fn() } }) }),
+    });
+
+    const { result } = renderHook(() => useLangChainRespondAll());
+    expect(() => result.current({ a: { approved: true } })).toThrow(
+      "useStreamRuntime",
+    );
+  });
+});

package/src/useLangChainState.test.tsx

@@ -2,8 +2,12 @@
 
 import { describe, expect, it, vi } from "vitest";
 import { renderHook } from "@testing-library/react";
+import {
+  createRunSelectorAgainst,
+  makeExtras,
+  type Selector,
+} from "./__tests__/langChainTestUtils";
 
-type Selector = (state: unknown) => unknown;
 const { mockUseAuiState } = vi.hoisted(() => ({
   mockUseAuiState: vi.fn(),
 }));
@@ -18,27 +22,9 @@ vi.mock(import("@assistant-ui/store"), async (importOriginal) => {
   };
 });
 
-import { useLangChainState } from "./useStreamRuntime";
+import { useLangChainState } from "./hooks";
 
-// The real guard symbol in useStreamRuntime is module-private. To stand in for
-// runtime-produced extras, we use a Proxy whose `has` trap claims any symbol
-// with the description "langchain-runtime-extras" is present.
-const makeExtras = (values: Record<string, unknown>) =>
-  new Proxy(
-    { values },
-    {
-      has: (target, key) =>
-        (typeof key === "symbol" &&
-          key.description === "langchain-runtime-extras") ||
-        Reflect.has(target, key),
-    },
-  );
-
-const runSelectorAgainst = (extras: unknown) => {
-  mockUseAuiState.mockImplementationOnce((selector: Selector) =>
-    selector({ thread: { extras } }),
-  );
-};
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
 
 describe("useLangChainState", () => {
   it("returns undefined when extras are absent", () => {
@@ -54,7 +40,9 @@ describe("useLangChainState", () => {
   });
 
   it("reads the value for the given key from extras.values", () => {
-    runSelectorAgainst(makeExtras({ todos: [{ id: "a" }], count: 7 }));
+    runSelectorAgainst(
+      makeExtras({ values: { todos: [{ id: "a" }], count: 7 } }),
+    );
     const { result } = renderHook(() =>
       useLangChainState<Array<{ id: string }>>("todos"),
     );
@@ -62,26 +50,26 @@ describe("useLangChainState", () => {
   });
 
   it("falls back to defaultValue when the key is missing from values", () => {
-    runSelectorAgainst(makeExtras({ unrelated: "x" }));
+    runSelectorAgainst(makeExtras({ values: { unrelated: "x" } }));
     const { result } = renderHook(() => useLangChainState<number>("count", 99));
     expect(result.current).toBe(99);
   });
 
   it("returns undefined when the key is missing and no default is given", () => {
-    runSelectorAgainst(makeExtras({ unrelated: "x" }));
+    runSelectorAgainst(makeExtras({ values: { unrelated: "x" } }));
     const { result } = renderHook(() => useLangChainState<number>("count"));
     expect(result.current).toBeUndefined();
   });
 
   it("returns the stored value (not default) when both are present", () => {
-    runSelectorAgainst(makeExtras({ count: 0 }));
+    runSelectorAgainst(makeExtras({ values: { count: 0 } }));
     const { result } = renderHook(() => useLangChainState<number>("count", 99));
     // 0 is a valid value; must be preserved over the default.
     expect(result.current).toBe(0);
   });
 
   it("preserves explicit null over defaultValue", () => {
-    runSelectorAgainst(makeExtras({ flag: null }));
+    runSelectorAgainst(makeExtras({ values: { flag: null } }));
     const { result } = renderHook(() =>
       useLangChainState<string | null>("flag", "default"),
     );

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

@@ -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 { useLangChainToolCalls } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainToolCalls", () => {
+  it("returns an empty array when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainToolCalls());
+    expect(result.current).toEqual([]);
+  });
+
+  it("returns the assembled tool calls from extras", () => {
+    const toolCalls = [{ id: "a", name: "search", args: { q: "x" } }];
+    runSelectorAgainst(makeExtras({ toolCalls }));
+    const { result } = renderHook(() => useLangChainToolCalls());
+    expect(result.current).toBe(toolCalls);
+  });
+
+  it("returns an empty array when extras carry no tool calls", () => {
+    runSelectorAgainst(makeExtras({ toolCalls: undefined }));
+    const { result } = renderHook(() => useLangChainToolCalls());
+    expect(result.current).toEqual([]);
+  });
+});

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",
+      ]);
+    });
+  });
+});