@assistant-ui/react-langchain 0.0.14 → 0.0.17

package/src/streamingTiming.ts

@@ -0,0 +1,85 @@
+"use client";
+
+import type { MessageTiming } from "@assistant-ui/core";
+import {
+  useStreamingTiming,
+  type StreamingTimingAccessors,
+} from "@assistant-ui/core/react";
+import type { LangChainBaseMessage, LangChainContentBlock } from "./types";
+import { getMessageType } from "./convertMessages";
+
+const findAiMessage = (
+  messages: readonly LangChainBaseMessage[],
+  messageId: string,
+): LangChainBaseMessage | undefined =>
+  messages.find((m) => getMessageType(m) === "ai" && m.id === messageId);
+
+const reasoningTextLength = (part: {
+  readonly summary?: ReadonlyArray<{ readonly text?: string }>;
+  readonly reasoning?: string;
+}): number => {
+  if (part.summary && part.summary.length > 0)
+    return part.summary.map((s) => s?.text ?? "").join("\n\n\n").length;
+  return part.reasoning?.length ?? 0;
+};
+
+const getTextLength = (
+  messages: readonly LangChainBaseMessage[],
+  messageId: string,
+): number => {
+  const m = findAiMessage(messages, messageId);
+  if (!m) return 0;
+  const content = m.content;
+  if (typeof content === "string") return content.length;
+  if (!Array.isArray(content)) return 0;
+  let len = 0;
+  for (const part of content as readonly LangChainContentBlock[]) {
+    switch (part.type) {
+      case "text":
+      case "text_delta":
+        if (typeof part.text === "string") len += part.text.length;
+        break;
+      case "thinking":
+        if (typeof part.thinking === "string") len += part.thinking.length;
+        break;
+      case "reasoning":
+        len += reasoningTextLength(part);
+        break;
+    }
+  }
+  return len;
+};
+
+const getToolCallCount = (
+  messages: readonly LangChainBaseMessage[],
+  messageId: string,
+): number => findAiMessage(messages, messageId)?.tool_calls?.length ?? 0;
+
+const getAssistantMessageId = (
+  messages: readonly LangChainBaseMessage[],
+): string | undefined => {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const m = messages[i];
+    if (m && getMessageType(m) === "ai" && m.id) return m.id;
+  }
+  return undefined;
+};
+
+export const langChainStreamingTimingAccessors: StreamingTimingAccessors<LangChainBaseMessage> =
+  {
+    getAssistantMessageId,
+    getTextLength,
+    getToolCallCount,
+  };
+
+/**
+ * Tracks per-message streaming timing for LangChain messages. Delegates to
+ * the shared `useStreamingTiming` primitive in `@assistant-ui/core/react`,
+ * adapted to the `LangChainBaseMessage` shape (`_getType() -> "ai"`, content
+ * blocks including text/thinking/reasoning, `tool_calls`).
+ */
+export const useLangChainStreamingTiming = (
+  messages: readonly LangChainBaseMessage[],
+  isRunning: boolean,
+): Record<string, MessageTiming> =>
+  useStreamingTiming(messages, isRunning, langChainStreamingTimingAccessors);

package/src/types.ts

@@ -1,14 +1,39 @@
-import type { MessageStatus } from "@assistant-ui/core";
+import type {
+  AttachmentAdapter,
+  DictationAdapter,
+  ExternalStoreSharedOptions,
+  FeedbackAdapter,
+  MessageStatus,
+  RealtimeVoiceAdapter,
+  RemoteThreadListAdapter,
+  SpeechSynthesisAdapter,
+} from "@assistant-ui/core";
+import type { AssistantCloud } from "assistant-cloud";
+import type {
+  UseStreamOptions,
+  AssembledToolCall,
+  AnyStream,
+  SubagentDiscoverySnapshot,
+  SubgraphDiscoverySnapshot,
+} from "@langchain/react";
 
 /** Known content block types from @langchain/core messages. */
 export type LangChainContentBlock =
   | { type: "text"; text: string }
   | { type: "text_delta"; text: string }
-  | { type: "image_url"; image_url: string | { url: string } }
+  | { type: "image_url"; image_url: string | { url?: string } }
   | { type: "thinking"; thinking: string }
   | {
       type: "reasoning";
-      summary: Array<{ type: "summary_text"; text: string }>;
+      summary?: Array<{ type: "summary_text"; text?: string }>;
+      reasoning?: string;
+    }
+  | {
+      type: "file";
+      data: string;
+      mime_type: string;
+      source_type?: "base64";
+      metadata?: { filename?: string };
     }
   | { type: "tool_use" | "input_json_delta" };
 
@@ -18,6 +43,41 @@ export type LangChainToolCall = {
   args: Record<string, unknown>;
 };
 
+/**
+ * A generative UI component the graph accumulates in its state. Read from
+ * `stream.values[uiStateKey]` and rendered via `makeAssistantDataUI`.
+ *
+ * The parent assistant message is carried in `metadata.message_id` by the
+ * Python SDK and in `metadata.id` by the JS SDK; the runtime reads either.
+ */
+export type UIMessage<
+  TName extends string = string,
+  TProps extends Record<string, unknown> = Record<string, unknown>,
+> = {
+  type: "ui";
+  id: string;
+  name: TName;
+  props: TProps;
+  metadata?: {
+    merge?: boolean;
+    run_id?: string;
+    name?: string;
+    tags?: string[];
+    message_id?: string;
+    id?: string;
+    [key: string]: unknown;
+  };
+};
+
+/**
+ * Emitted on the live `custom` channel by the graph's `remove_ui_message`
+ * helper to drop a pushed `UIMessage` by id.
+ */
+export type RemoveUIMessage = {
+  type: "remove-ui";
+  id: string;
+};
+
 /**
  * Minimal duck-typed interface for BaseMessage class instances returned by
  * `useStream`. Used internally by the message converter.
@@ -40,3 +100,87 @@ export type LangChainBaseMessage = {
   /** Present on ToolMessage */
   artifact?: unknown;
 };
+
+export type LangChainRuntimeExtras = {
+  interrupt: { value?: unknown } | undefined;
+  interrupts: readonly { id?: string; value?: unknown }[];
+  toolCalls: readonly AssembledToolCall[];
+  subagents: ReadonlyMap<string, SubagentDiscoverySnapshot>;
+  subgraphs: ReadonlyMap<string, SubgraphDiscoverySnapshot>;
+  stream: AnyStream;
+  error: unknown;
+  submit: (
+    values: Record<string, unknown> | null | undefined,
+    options?: Record<string, unknown>,
+  ) => Promise<void>;
+  respond: (
+    response: unknown,
+    options?: Record<string, unknown>,
+  ) => Promise<void>;
+  respondAll: (
+    responsesById: Record<string, unknown>,
+    options?: Record<string, unknown>,
+  ) => Promise<void>;
+  values: Record<string, unknown>;
+  messagesKey: string;
+};
+
+export type LangChainRuntimeExtraOptions = ExternalStoreSharedOptions & {
+  /**
+   * Called whenever the active thread's canonical (remote) ID changes, so the
+   * value can be treated as a managed/controlled variable (e.g. synced to a URL
+   * query param). Only the settled remote ID is emitted: while a freshly created
+   * thread is still optimistic the value is `undefined`, and the real ID is
+   * emitted once the thread is initialized; the transient local ID is never
+   * surfaced.
+   */
+  onThreadIdChange?: ((threadId: string | undefined) => void) | undefined;
+  cloud?: AssistantCloud | undefined;
+  adapters?:
+    | {
+        attachments?: AttachmentAdapter | undefined;
+        speech?: SpeechSynthesisAdapter | undefined;
+        dictation?: DictationAdapter | undefined;
+        voice?: RealtimeVoiceAdapter | undefined;
+        feedback?: FeedbackAdapter | undefined;
+      }
+    | undefined;
+  /**
+   * When the user sends a new message while previous tool calls are
+   * still pending, automatically submit `tool` messages that cancel
+   * them so the agent's tool-call accounting stays consistent.
+   * Defaults to `true`.
+   */
+  autoCancelPendingToolCalls?: boolean | undefined;
+  /**
+   * Routes the Cancel button's click to `useStream().stop()`. On by
+   * default. Pass `false` to disable the Cancel button.
+   */
+  unstable_allowCancellation?: boolean | undefined;
+  /**
+   * Custom `RemoteThreadListAdapter`. When provided, replaces the
+   * cloud-backed thread list adapter.
+   */
+  unstable_threadListAdapter?: RemoteThreadListAdapter | undefined;
+  /** Custom thread-creation hook, forwarded to the cloud adapter. */
+  create?: (() => Promise<{ externalId: string | undefined }>) | undefined;
+  /** Custom thread-deletion hook, forwarded to the cloud adapter. */
+  delete?: ((threadId: string) => Promise<void>) | undefined;
+  /**
+   * State key the graph accumulates generative `UIMessage`s under. Each UI
+   * is attached to the assistant message identified by its `metadata.message_id`
+   * (Python SDK) or `metadata.id` (JS SDK) and emitted as a `data` part.
+   * Defaults to `"ui"`.
+   */
+  uiStateKey?: string | undefined;
+};
+
+// Distribute the intersection through the union arms of `UseStreamOptions`
+// (`AgentServerOptions | CustomAdapterOptions`). Writing `UseStreamOptions & X`
+// directly collapses arm tracking, so `Omit<…, "cloud">` and the like would
+// produce a flattened structural type that no longer matches either arm.
+export type UseStreamRuntimeOptions = UseStreamOptions extends infer O
+  ? O extends UseStreamOptions
+    ? O & LangChainRuntimeExtraOptions
+    : never
+  : never;

package/src/uiMessages.test.ts

@@ -0,0 +1,190 @@
+import { describe, expect, it } from "vitest";
+import {
+  applyUIUpdate,
+  extractUIUpdate,
+  foldUIUpdates,
+  isUIUpdate,
+  mergeUIMessages,
+} from "./uiMessages";
+import type { UIMessage } from "./types";
+
+const ui = (id: string, props: Record<string, unknown> = {}): UIMessage => ({
+  type: "ui",
+  id,
+  name: "chart",
+  props,
+});
+
+describe("isUIUpdate", () => {
+  it("accepts ui and remove-ui objects with a string id", () => {
+    expect(isUIUpdate({ type: "ui", id: "a", name: "x", props: {} })).toBe(
+      true,
+    );
+    expect(isUIUpdate({ type: "remove-ui", id: "a" })).toBe(true);
+  });
+
+  it("accepts an array of updates", () => {
+    expect(
+      isUIUpdate([
+        { type: "ui", id: "a", name: "x", props: {} },
+        { type: "remove-ui", id: "b" },
+      ]),
+    ).toBe(true);
+  });
+
+  it("rejects non-UI values", () => {
+    expect(isUIUpdate(null)).toBe(false);
+    expect(isUIUpdate("ui")).toBe(false);
+    expect(isUIUpdate({ type: "ui", name: "x" })).toBe(false);
+    expect(isUIUpdate({ type: "other", id: "a" })).toBe(false);
+    expect(isUIUpdate([{ type: "ui", id: "a", name: "x", props: {} }, 5])).toBe(
+      false,
+    );
+  });
+
+  it("rejects a ui update missing name or props", () => {
+    expect(isUIUpdate({ type: "ui", id: "a" })).toBe(false);
+    expect(isUIUpdate({ type: "ui", id: "a", name: "x" })).toBe(false);
+    expect(isUIUpdate({ type: "ui", id: "a", props: {} })).toBe(false);
+    expect(isUIUpdate({ type: "ui", id: "a", name: "x", props: null })).toBe(
+      false,
+    );
+  });
+
+  it("rejects an empty array", () => {
+    expect(isUIUpdate([])).toBe(false);
+  });
+});
+
+describe("applyUIUpdate", () => {
+  it("pushes a new UI message", () => {
+    expect(applyUIUpdate([], ui("a"))).toEqual([ui("a")]);
+  });
+
+  it("replaces an existing UI message by id", () => {
+    const result = applyUIUpdate([ui("a", { v: 1 })], ui("a", { v: 2 }));
+    expect(result).toEqual([ui("a", { v: 2 })]);
+  });
+
+  it("merges props when metadata.merge is true", () => {
+    const result = applyUIUpdate([ui("a", { keep: 1, v: 1 })], {
+      ...ui("a", { v: 2 }),
+      metadata: { merge: true },
+    });
+    expect(result).toEqual([
+      { ...ui("a", { keep: 1, v: 2 }), metadata: { merge: true } },
+    ]);
+  });
+
+  it("removes a UI message by id", () => {
+    const result = applyUIUpdate([ui("a"), ui("b")], {
+      type: "remove-ui",
+      id: "a",
+    });
+    expect(result).toEqual([ui("b")]);
+  });
+
+  it("applies an array of updates in order", () => {
+    const result = applyUIUpdate(
+      [ui("a")],
+      [ui("b"), { type: "remove-ui", id: "a" }, ui("c")],
+    );
+    expect(result).toEqual([ui("b"), ui("c")]);
+  });
+
+  it("does not mutate the input list", () => {
+    const input = [ui("a")];
+    applyUIUpdate(input, ui("b"));
+    expect(input).toEqual([ui("a")]);
+  });
+});
+
+describe("extractUIUpdate", () => {
+  it("reads a UI update straight from params.data", () => {
+    const event = {
+      params: { data: { type: "ui", id: "a", name: "x", props: {} } },
+    };
+    expect(extractUIUpdate(event)).toEqual({
+      type: "ui",
+      id: "a",
+      name: "x",
+      props: {},
+    });
+  });
+
+  it("falls back to params.data.payload when data wraps the update", () => {
+    const event = {
+      params: { data: { payload: { type: "remove-ui", id: "a" } } },
+    };
+    expect(extractUIUpdate(event)).toEqual({ type: "remove-ui", id: "a" });
+  });
+
+  it("reads an array of updates", () => {
+    const event = {
+      params: { data: [ui("a"), { type: "remove-ui", id: "b" }] },
+    };
+    expect(extractUIUpdate(event)).toEqual([
+      ui("a"),
+      { type: "remove-ui", id: "b" },
+    ]);
+  });
+
+  it("returns undefined for non-UI custom events and malformed shapes", () => {
+    expect(extractUIUpdate({ params: { data: { foo: 1 } } })).toBeUndefined();
+    expect(extractUIUpdate({ params: {} })).toBeUndefined();
+    expect(extractUIUpdate({})).toBeUndefined();
+    expect(extractUIUpdate(null)).toBeUndefined();
+  });
+});
+
+describe("mergeUIMessages", () => {
+  it("returns the snapshot when there are no live messages", () => {
+    expect(mergeUIMessages([], [ui("a")])).toEqual([ui("a")]);
+  });
+
+  it("returns live messages when the snapshot is not an array", () => {
+    expect(mergeUIMessages([ui("a")], undefined)).toEqual([ui("a")]);
+  });
+
+  it("lets the snapshot win by id", () => {
+    const result = mergeUIMessages([ui("a", { v: 1 })], [ui("a", { v: 2 })]);
+    expect(result).toEqual([ui("a", { v: 2 })]);
+  });
+
+  it("keeps live and snapshot entries with distinct ids", () => {
+    expect(mergeUIMessages([ui("a")], [ui("b")])).toEqual([ui("a"), ui("b")]);
+  });
+});
+
+describe("foldUIUpdates", () => {
+  const evt = (data: unknown) => ({ params: { data } });
+
+  it("folds a sequence of custom events into a UI list", () => {
+    expect(foldUIUpdates([evt(ui("a")), evt(ui("b"))])).toEqual([
+      ui("a"),
+      ui("b"),
+    ]);
+  });
+
+  it("reads updates nested under params.data.payload", () => {
+    expect(foldUIUpdates([evt({ payload: ui("a") })])).toEqual([ui("a")]);
+  });
+
+  it("applies a remove across events", () => {
+    const result = foldUIUpdates([
+      evt(ui("a")),
+      evt(ui("b")),
+      evt({ type: "remove-ui", id: "a" }),
+    ]);
+    expect(result).toEqual([ui("b")]);
+  });
+
+  it("ignores non-UI custom events", () => {
+    const result = foldUIUpdates([evt(ui("a")), evt({ foo: 1 }), evt(ui("b"))]);
+    expect(result).toEqual([ui("a"), ui("b")]);
+  });
+
+  it("returns an empty list for no events", () => {
+    expect(foldUIUpdates([])).toEqual([]);
+  });
+});

package/src/uiMessages.ts

@@ -0,0 +1,91 @@
+import type { RemoveUIMessage, UIMessage } from "./types";
+
+export type UIUpdate = UIMessage | RemoveUIMessage;
+
+export const isUIUpdate = (
+  value: unknown,
+): value is UIUpdate | readonly UIUpdate[] => {
+  if (Array.isArray(value)) return value.length > 0 && value.every(isUIUpdate);
+  if (value == null || typeof value !== "object") return false;
+  const v = value as {
+    type?: unknown;
+    id?: unknown;
+    name?: unknown;
+    props?: unknown;
+  };
+  if (typeof v.id !== "string") return false;
+  if (v.type === "remove-ui") return true;
+  return (
+    v.type === "ui" &&
+    typeof v.name === "string" &&
+    typeof v.props === "object" &&
+    v.props !== null
+  );
+};
+
+export const applyUIUpdate = (
+  list: readonly UIMessage[],
+  update: UIUpdate | readonly UIUpdate[],
+): UIMessage[] => {
+  const events = Array.isArray(update) ? update : [update as UIUpdate];
+  let next = list.slice();
+  for (const event of events) {
+    if (event.type === "remove-ui") {
+      next = next.filter((ui) => ui.id !== event.id);
+      continue;
+    }
+    const index = next.findIndex((ui) => ui.id === event.id);
+    if (index === -1) {
+      next.push(event);
+      continue;
+    }
+    next[index] =
+      event.metadata?.merge === true
+        ? { ...event, props: { ...next[index]!.props, ...event.props } }
+        : event;
+  }
+  return next;
+};
+
+/**
+ * Pulls a UI update out of a raw `custom`-channel event. The graph writes the
+ * `UIMessage` straight to the channel, so it lands at `params.data`; some
+ * transports wrap it one level deeper at `params.data.payload`.
+ */
+export const extractUIUpdate = (
+  event: unknown,
+): UIUpdate | readonly UIUpdate[] | undefined => {
+  const data = (event as { params?: { data?: unknown } } | null)?.params?.data;
+  if (isUIUpdate(data)) return data;
+  const payload = (data as { payload?: unknown } | undefined)?.payload;
+  if (isUIUpdate(payload)) return payload;
+  return undefined;
+};
+
+export const foldUIUpdates = (events: readonly unknown[]): UIMessage[] => {
+  let acc: UIMessage[] = [];
+  for (const event of events) {
+    const update = extractUIUpdate(event);
+    if (update) acc = applyUIUpdate(acc, update);
+  }
+  return acc;
+};
+
+/**
+ * Merges live-streamed UI with the state snapshot. The snapshot is
+ * authoritative by id: once a UI lands in graph state it supersedes its live
+ * copy. A consequence is that a live `remove-ui` is overridden while the
+ * snapshot still contains that id; the removal only takes visible effect once
+ * the snapshot catches up.
+ */
+export const mergeUIMessages = (
+  live: readonly UIMessage[],
+  snapshot: unknown,
+): UIMessage[] => {
+  const byId = new Map<string, UIMessage>();
+  for (const ui of live) byId.set(ui.id, ui);
+  if (Array.isArray(snapshot)) {
+    for (const ui of snapshot as UIMessage[]) byId.set(ui.id, ui);
+  }
+  return [...byId.values()];
+};

package/src/useLangChainError.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 { useLangChainError } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainError", () => {
+  it("returns undefined when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainError());
+    expect(result.current).toBeUndefined();
+  });
+
+  it("returns the error value when extras carry error", () => {
+    const error = new Error("boom");
+    runSelectorAgainst(makeExtras({ error }));
+    const { result } = renderHook(() => useLangChainError());
+    expect(result.current).toBe(error);
+  });
+
+  it("returns undefined when extras carry no error", () => {
+    runSelectorAgainst(makeExtras({ error: undefined }));
+    const { result } = renderHook(() => useLangChainError());
+    expect(result.current).toBeUndefined();
+  });
+});

package/src/useLangChainInterrupts.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 { useLangChainInterrupts } from "./hooks";
+
+const runSelectorAgainst = createRunSelectorAgainst(mockUseAuiState);
+
+describe("useLangChainInterrupts", () => {
+  it("returns an empty array when extras are absent", () => {
+    runSelectorAgainst(undefined);
+    const { result } = renderHook(() => useLangChainInterrupts());
+    expect(result.current).toEqual([]);
+  });
+
+  it("returns the pending interrupts from extras", () => {
+    const interrupts = [{ id: "i1", value: { q: "ok?" } }];
+    runSelectorAgainst(makeExtras({ interrupts }));
+    const { result } = renderHook(() => useLangChainInterrupts());
+    expect(result.current).toBe(interrupts);
+  });
+
+  it("returns an empty array when extras carry no interrupts", () => {
+    runSelectorAgainst(makeExtras({ interrupts: undefined }));
+    const { result } = renderHook(() => useLangChainInterrupts());
+    expect(result.current).toEqual([]);
+  });
+});

package/src/useLangChainRespond.test.tsx

@@ -0,0 +1,49 @@
+// @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 { useLangChainRespond } from "./hooks";
+
+describe("useLangChainRespond", () => {
+  it("forwards the response and options to stream.respond", async () => {
+    const respond = vi.fn().mockResolvedValue(undefined);
+    mockUseAui.mockReturnValue({
+      thread: () => ({
+        getState: () => ({ extras: makeExtras({ respond }) }),
+      }),
+    });
+
+    const { result } = renderHook(() => useLangChainRespond());
+    await result.current({ approved: true }, { interruptId: "x" });
+
+    expect(respond).toHaveBeenCalledWith(
+      { approved: true },
+      { interruptId: "x" },
+    );
+  });
+
+  it("throws when extras are absent or unbranded", () => {
+    mockUseAui.mockReturnValue({
+      thread: () => ({ getState: () => ({ extras: { respond: vi.fn() } }) }),
+    });
+
+    const { result } = renderHook(() => useLangChainRespond());
+    expect(() => result.current({ approved: true })).toThrow(
+      "useStreamRuntime",
+    );
+  });
+});