@assistant-ui/react-langchain 0.0.15 → 0.0.17

package/src/groupUIMessagesByParent.test.ts

@@ -0,0 +1,61 @@
+import { describe, expect, it } from "vitest";
+import { groupUIMessagesByParent } from "./useStreamRuntime";
+import type { UIMessage } from "./types";
+
+const uiMessage = (
+  metadata: UIMessage["metadata"],
+  name = "chart",
+): UIMessage => ({
+  type: "ui",
+  id: "ui-1",
+  name,
+  props: { points: [1, 2, 3] },
+  ...(metadata !== undefined && { metadata }),
+});
+
+describe("groupUIMessagesByParent", () => {
+  it("returns an empty map for non-array state", () => {
+    expect(groupUIMessagesByParent(undefined).size).toBe(0);
+    expect(groupUIMessagesByParent(null).size).toBe(0);
+    expect(groupUIMessagesByParent({ ui: [] }).size).toBe(0);
+  });
+
+  it("groups by metadata.message_id (Python SDK)", () => {
+    const ui = uiMessage({ message_id: "msg-2" });
+    const map = groupUIMessagesByParent([ui]);
+
+    expect(map.get("msg-2")).toEqual([ui]);
+  });
+
+  it("falls back to metadata.id when message_id is absent (JS SDK)", () => {
+    const ui = uiMessage({ id: "msg-2" });
+    const map = groupUIMessagesByParent([ui]);
+
+    expect(map.get("msg-2")).toEqual([ui]);
+  });
+
+  it("prefers message_id over id when both are present", () => {
+    const ui = uiMessage({ message_id: "msg-2", id: "msg-9" });
+    const map = groupUIMessagesByParent([ui]);
+
+    expect(map.get("msg-2")).toEqual([ui]);
+    expect(map.has("msg-9")).toBe(false);
+  });
+
+  it("skips entries with no parent link", () => {
+    const map = groupUIMessagesByParent([
+      uiMessage(undefined),
+      uiMessage({ message_id: "" }),
+    ]);
+
+    expect(map.size).toBe(0);
+  });
+
+  it("collects several UIs under the same parent in order", () => {
+    const a = uiMessage({ message_id: "msg-2" }, "chart");
+    const b = uiMessage({ message_id: "msg-2" }, "table");
+    const map = groupUIMessagesByParent([a, b]);
+
+    expect(map.get("msg-2")).toEqual([a, b]);
+  });
+});

package/src/hooks.ts

@@ -0,0 +1,156 @@
+"use client";
+
+import { useAui } from "@assistant-ui/store";
+import type {
+  AssembledToolCall,
+  SubagentDiscoverySnapshot,
+  SubgraphDiscoverySnapshot,
+} from "@langchain/react";
+import { langChainExtras } from "./runtimeExtras";
+import type { LangChainBaseMessage } from "./types";
+
+const EMPTY_TOOL_CALLS: readonly AssembledToolCall[] = [];
+
+const EMPTY_INTERRUPTS: readonly { id?: string; value?: unknown }[] = [];
+
+const EMPTY_SUBAGENTS: ReadonlyMap<string, SubagentDiscoverySnapshot> =
+  new Map();
+
+const EMPTY_SUBGRAPHS: ReadonlyMap<string, SubgraphDiscoverySnapshot> =
+  new Map();
+
+/**
+ * Read the current LangGraph interrupt state from the runtime extras.
+ */
+export const useLangChainInterruptState = () =>
+  langChainExtras.use((e) => e.interrupt, undefined);
+
+/**
+ * Read every interrupt pending at the current checkpoint, each with `id`
+ * and `value`. Defaults to an empty array. Pair with `useLangChainRespondAll`
+ * (keyed by interrupt id) to resolve several at once.
+ */
+export const useLangChainInterrupts = () =>
+  langChainExtras.use(
+    (e) => e.interrupts ?? EMPTY_INTERRUPTS,
+    EMPTY_INTERRUPTS,
+  );
+
+/** Read the last run/hydration error from the runtime extras. */
+export const useLangChainError = () =>
+  langChainExtras.use((e) => e.error, undefined);
+
+/**
+ * Read the root tool calls assembled by `useStream` from the `tools`
+ * channel. Defaults to an empty array, so consumers can `.map` without
+ * a guard. Useful for rendering pending/streamed tool calls and
+ * approval UIs.
+ */
+export const useLangChainToolCalls = () =>
+  langChainExtras.use((e) => e.toolCalls ?? EMPTY_TOOL_CALLS, EMPTY_TOOL_CALLS);
+
+/** Subagents discovered on the current run (multi-agent graphs). */
+export const useLangChainSubagents = () =>
+  langChainExtras.use((e) => e.subagents ?? EMPTY_SUBAGENTS, EMPTY_SUBAGENTS);
+
+/** Subgraphs discovered on the current run. */
+export const useLangChainSubgraphs = () =>
+  langChainExtras.use((e) => e.subgraphs ?? EMPTY_SUBGRAPHS, EMPTY_SUBGRAPHS);
+
+/**
+ * The underlying `useStream` handle, for advanced views — pass it to v1's
+ * `useMessages(stream, target)` / `useToolCalls(stream, target)` with a
+ * subagent/subgraph from `useLangChainSubagents` / `useLangChainSubgraphs`.
+ * `undefined` outside the runtime provider.
+ */
+export const useLangChainStream = () =>
+  langChainExtras.use((e) => e.stream, undefined);
+
+/**
+ * Returns a function to submit raw state updates to the LangGraph agent,
+ * bypassing the normal message flow. Useful for sending interrupt resume
+ * commands.
+ */
+export const useLangChainSubmit = () => {
+  const aui = useAui();
+  return (
+    values: Record<string, unknown> | null | undefined,
+    options?: Record<string, unknown>,
+  ) => langChainExtras.get(aui).submit(values, options);
+};
+
+/**
+ * Resume a LangGraph interrupt with a response payload via
+ * `useStream().respond`. Preferred over `useLangChainSendCommand`; it
+ * carries the response cleanly and handles interrupt namespaces.
+ */
+export const useLangChainRespond = () => {
+  const aui = useAui();
+  return (response: unknown, options?: Record<string, unknown>) =>
+    langChainExtras.get(aui).respond(response, options);
+};
+
+/**
+ * Resume several LangGraph interrupts pending at the same checkpoint in
+ * one run via `useStream().respondAll`. Use when a run pauses on multiple
+ * interrupts at once; sequential `useLangChainRespond` calls can't service
+ * them (the first resume starts a run, stranding the rest).
+ */
+export const useLangChainRespondAll = () => {
+  const aui = useAui();
+  return (
+    responsesById: Record<string, unknown>,
+    options?: Record<string, unknown>,
+  ) => langChainExtras.get(aui).respondAll(responsesById, options);
+};
+
+/**
+ * Submit a list of LangChain-shaped messages on the current thread.
+ * Parity helper for migrating from `useLangGraphSend`. Routes to
+ * `useStream().submit({ [messagesKey]: messages }, options)`.
+ */
+export const useLangChainSend = () => {
+  const aui = useAui();
+  return (
+    messages: readonly LangChainBaseMessage[],
+    options?: Record<string, unknown>,
+  ) => {
+    const extras = langChainExtras.get(aui);
+    return extras.submit({ [extras.messagesKey]: messages }, options);
+  };
+};
+
+/**
+ * Submit a `useStream` command (e.g. interrupt resume). Parity helper
+ * for migrating from `useLangGraphSendCommand`. Note that v1's command
+ * shape (`{ resume?, goto?, update? }`) differs from the legacy
+ * `{ resume: string }` form; to carry a payload, use the input or
+ * `stream.respond` instead.
+ */
+export const useLangChainSendCommand = () => {
+  const submit = useLangChainSubmit();
+  return (command: Record<string, unknown>) => submit(null, { command });
+};
+
+/**
+ * Read a custom LangGraph state key from the current thread. Mirrors
+ * `useStream().values[key]` from `@langchain/react` and updates when the
+ * stream emits new state.
+ *
+ * @example
+ * ```tsx
+ * const todos = useLangChainState<Todo[]>("todos");
+ * const files = useLangChainState<Record<string, string>>("files", {});
+ * ```
+ */
+export function useLangChainState<T>(key: string): T | undefined;
+export function useLangChainState<T>(key: string, defaultValue: T): T;
+export function useLangChainState<T>(
+  key: string,
+  defaultValue?: T,
+): T | undefined {
+  return langChainExtras.use((e) => {
+    const value = e.values[key] as T | undefined;
+    return value !== undefined ? value : defaultValue;
+  }, defaultValue);
+}

package/src/index.ts

@@ -1,19 +1,34 @@
+export { useStreamRuntime } from "./useStreamRuntime";
+
 export {
-  useStreamRuntime,
   useLangChainError,
+  useLangChainInterrupts,
   useLangChainInterruptState,
+  useLangChainRespond,
+  useLangChainRespondAll,
   useLangChainSend,
   useLangChainSendCommand,
   useLangChainState,
+  useLangChainStream,
+  useLangChainSubagents,
+  useLangChainSubgraphs,
   useLangChainSubmit,
   useLangChainToolCalls,
-} from "./useStreamRuntime";
-export type { UseStreamRuntimeOptions } from "./useStreamRuntime";
+} from "./hooks";
 
 export { convertLangChainBaseMessage } from "./convertMessages";
+export { useLangChainStreamingTiming } from "./streamingTiming";
+
+export type {
+  SubagentDiscoverySnapshot,
+  SubgraphDiscoverySnapshot,
+} from "@langchain/react";
 
 export type {
   LangChainBaseMessage,
   LangChainContentBlock,
   LangChainToolCall,
+  RemoveUIMessage,
+  UIMessage,
+  UseStreamRuntimeOptions,
 } from "./types";

package/src/resolveForkCheckpoint.test.ts

@@ -0,0 +1,202 @@
+import { describe, expect, it, vi } from "vitest";
+import type { LangChainBaseMessage } from "./types";
+import type { ForkCheckpointClient } from "./findForkCheckpointInHistory";
+import { resolveForkCheckpoint } from "./resolveForkCheckpoint";
+
+const msg = (id: string): LangChainBaseMessage => ({
+  _getType: () => "human",
+  content: "",
+  id,
+});
+
+type HistoryState = {
+  values: Record<string, unknown>;
+  checkpoint: { checkpoint_id?: string };
+};
+
+const makeClient = (history: HistoryState[]): ForkCheckpointClient => ({
+  threads: {
+    getHistory: vi.fn(async () => history as never),
+  },
+});
+
+const metadata = (entries: Record<string, string | undefined>) =>
+  new Map(
+    Object.entries(entries).map(([id, parentCheckpointId]) => [
+      id,
+      { parentCheckpointId },
+    ]),
+  );
+
+describe("resolveForkCheckpoint", () => {
+  it("uses the head metadata fast-path without hitting getHistory", async () => {
+    const client = makeClient([]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      "h1",
+      "a1",
+      metadata({ a1: "cp-head" }),
+      "messages",
+    );
+
+    expect(result).toBe("cp-head");
+    expect(client.threads.getHistory).not.toHaveBeenCalled();
+  });
+
+  it("falls back to history search when sourceId is not the head message", async () => {
+    const client = makeClient([
+      {
+        values: { messages: [msg("h1")] },
+        checkpoint: { checkpoint_id: "cp-h1" },
+      },
+    ]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1"), msg("h2")],
+      "h1",
+      "h1",
+      metadata({ h2: "cp-head" }),
+      "messages",
+    );
+
+    expect(result).toBe("cp-h1");
+    expect(client.threads.getHistory).toHaveBeenCalled();
+  });
+
+  it("falls back to history search when the head metadata has no checkpoint", async () => {
+    const client = makeClient([
+      {
+        values: { messages: [msg("h1")] },
+        checkpoint: { checkpoint_id: "cp-h1" },
+      },
+    ]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      "h1",
+      "a1",
+      metadata({}),
+      "messages",
+    );
+
+    expect(result).toBe("cp-h1");
+    expect(client.threads.getHistory).toHaveBeenCalled();
+  });
+
+  it("matches the parent prefix through history search", async () => {
+    const client = makeClient([
+      {
+        values: { messages: [msg("h1"), msg("a1")] },
+        checkpoint: { checkpoint_id: "cp-full" },
+      },
+      {
+        values: { messages: [msg("h1")] },
+        checkpoint: { checkpoint_id: "cp-parent" },
+      },
+    ]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      "h1",
+      undefined,
+      undefined,
+      "messages",
+    );
+
+    expect(result).toBe("cp-parent");
+  });
+
+  it("returns null when a non-null parentId is not in the messages", async () => {
+    const client = makeClient([]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      "missing",
+      undefined,
+      undefined,
+      "messages",
+    );
+
+    expect(result).toBeNull();
+    expect(client.threads.getHistory).not.toHaveBeenCalled();
+  });
+
+  it("forks the first message from the initial empty-message checkpoint", async () => {
+    const client = makeClient([
+      {
+        values: { messages: [msg("h1"), msg("a1")] },
+        checkpoint: { checkpoint_id: "cp-full" },
+      },
+      {
+        values: { messages: [] },
+        checkpoint: { checkpoint_id: "cp-initial" },
+      },
+    ]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      null,
+      "h1",
+      metadata({}),
+      "messages",
+    );
+
+    expect(result).toBe("cp-initial");
+  });
+
+  it("returns null when no checkpoint resolves", async () => {
+    const client = makeClient([
+      {
+        values: { messages: [msg("other")] },
+        checkpoint: { checkpoint_id: "cp-other" },
+      },
+    ]);
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1")],
+      "h1",
+      undefined,
+      undefined,
+      "messages",
+    );
+
+    expect(result).toBeNull();
+  });
+
+  it("swallows history search errors and returns null", async () => {
+    const client: ForkCheckpointClient = {
+      threads: {
+        getHistory: vi.fn(async () => {
+          throw new Error("network");
+        }),
+      },
+    };
+
+    const result = await resolveForkCheckpoint(
+      client,
+      "thread-1",
+      [msg("h1"), msg("a1")],
+      "h1",
+      undefined,
+      undefined,
+      "messages",
+    );
+
+    expect(result).toBeNull();
+  });
+});

package/src/resolveForkCheckpoint.ts

@@ -0,0 +1,49 @@
+import type { LangChainBaseMessage } from "./types";
+import {
+  findForkCheckpointInHistory,
+  type ForkCheckpointClient,
+} from "./findForkCheckpointInHistory";
+
+type MessageMetadataSnapshot =
+  | ReadonlyMap<string, { readonly parentCheckpointId: string | undefined }>
+  | undefined;
+
+/**
+ * Hydration seeds every message with the head's parent checkpoint, so only the
+ * head's recorded fork checkpoint is reliable; older turns are matched against
+ * server history by message id. A `null` `parentId` edits the first human
+ * message and forks from the thread's initial (empty-message) checkpoint.
+ */
+export const resolveForkCheckpoint = async (
+  client: ForkCheckpointClient,
+  threadId: string,
+  messages: readonly LangChainBaseMessage[],
+  parentId: string | null,
+  sourceId: string | null | undefined,
+  metadata: MessageMetadataSnapshot,
+  messagesKey: string,
+): Promise<string | null> => {
+  const lastMessage = messages[messages.length - 1];
+  let checkpointId: string | null =
+    sourceId != null && lastMessage?.id === sourceId
+      ? (metadata?.get(sourceId)?.parentCheckpointId ?? null)
+      : null;
+
+  if (!checkpointId) {
+    const parentIndex =
+      parentId == null ? -1 : messages.findIndex((m) => m.id === parentId);
+    if (parentId != null && parentIndex === -1) return null;
+    try {
+      checkpointId = await findForkCheckpointInHistory(
+        client,
+        threadId,
+        messages.slice(0, parentIndex + 1),
+        messagesKey,
+      );
+    } catch {
+      return null;
+    }
+  }
+
+  return checkpointId;
+};

package/src/runtimeExtras.ts

@@ -0,0 +1,5 @@
+import { createRuntimeExtras } from "@assistant-ui/core/internal";
+import type { LangChainRuntimeExtras } from "./types";
+
+export const langChainExtras =
+  createRuntimeExtras<LangChainRuntimeExtras>("useStreamRuntime");

package/src/streamingTiming.test.ts

@@ -0,0 +1,118 @@
+import { describe, expect, it } from "vitest";
+import type { LangChainBaseMessage } from "./types";
+import { langChainStreamingTimingAccessors } from "./streamingTiming";
+
+const ai = (
+  fields: Partial<LangChainBaseMessage> & { id: string; content: unknown },
+): LangChainBaseMessage => ({
+  _getType: () => "ai",
+  ...fields,
+});
+
+const human = (id: string, content: unknown): LangChainBaseMessage => ({
+  _getType: () => "human",
+  id,
+  content,
+});
+
+const { getAssistantMessageId, getTextLength, getToolCallCount } =
+  langChainStreamingTimingAccessors;
+
+describe("langChainStreamingTimingAccessors", () => {
+  it("resolves the last assistant message id via _getType()", () => {
+    expect(
+      getAssistantMessageId([human("u1", "hi"), ai({ id: "a1", content: "" })]),
+    ).toBe("a1");
+    expect(
+      getAssistantMessageId([
+        ai({ id: "a1", content: "" }),
+        human("u2", "hi"),
+        ai({ id: "a2", content: "" }),
+      ]),
+    ).toBe("a2");
+  });
+
+  it("returns undefined when there is no assistant message", () => {
+    expect(getAssistantMessageId([human("u1", "hi")])).toBeUndefined();
+  });
+
+  it("measures string content length", () => {
+    expect(getTextLength([ai({ id: "a1", content: "hello" })], "a1")).toBe(5);
+  });
+
+  it("sums text and thinking lengths across content blocks", () => {
+    expect(
+      getTextLength(
+        [
+          ai({
+            id: "a1",
+            content: [
+              { type: "thinking", thinking: "hmm" },
+              { type: "text", text: "answer" },
+            ],
+          }),
+        ],
+        "a1",
+      ),
+    ).toBe("hmm".length + "answer".length);
+  });
+
+  it("measures reasoning blocks (summary and reasoning fields)", () => {
+    // reasoning via summary_text entries, joined like contentToParts does.
+    expect(
+      getTextLength(
+        [
+          ai({
+            id: "a1",
+            content: [
+              {
+                type: "reasoning",
+                summary: [
+                  { type: "summary_text", text: "step one" },
+                  { type: "summary_text", text: "step two" },
+                ],
+              },
+            ],
+          }),
+        ],
+        "a1",
+      ),
+    ).toBe("step one\n\n\nstep two".length);
+
+    // reasoning via the bare `reasoning` field when no summary is present.
+    expect(
+      getTextLength(
+        [
+          ai({
+            id: "a2",
+            content: [{ type: "reasoning", reasoning: "deduced" }],
+          }),
+        ],
+        "a2",
+      ),
+    ).toBe("deduced".length);
+  });
+
+  it("counts tool_calls on the assistant message", () => {
+    expect(
+      getToolCallCount(
+        [
+          ai({
+            id: "a1",
+            content: "",
+            tool_calls: [
+              { id: "t1", name: "search", args: {} },
+              { id: "t2", name: "fetch", args: {} },
+            ],
+          }),
+        ],
+        "a1",
+      ),
+    ).toBe(2);
+  });
+
+  it("returns 0 when the message is missing or not assistant", () => {
+    expect(getTextLength([human("u1", "hi")], "a1")).toBe(0);
+    expect(getToolCallCount([human("u1", "hi")], "a1")).toBe(0);
+  });
+});