@assistant-ui/react 0.15.18 → 0.15.19

package/src/unstable/webmcp/convertTools.ts

@@ -118,6 +118,68 @@ const isStandardSchema = (schema: unknown): schema is StandardSchemaLike =>
   "~standard" in schema &&
   (schema as StandardSchemaLike)["~standard"].version === 1;
 
+const TOOL_ABORTED = Symbol("assistant-ui.webmcp-tool-aborted");
+
+const isThenable = <T>(value: T | PromiseLike<T>): value is PromiseLike<T> =>
+  typeof (value as PromiseLike<T> | null | undefined)?.then === "function";
+
+const raceWithAbort = async <T>(
+  value: PromiseLike<T>,
+  abortSignal: AbortSignal,
+): Promise<T | typeof TOOL_ABORTED> => {
+  let onAbort!: () => void;
+  const abortPromise = new Promise<typeof TOOL_ABORTED>((resolve) => {
+    onAbort = () => resolve(TOOL_ABORTED);
+    if (abortSignal.aborted) {
+      onAbort();
+    } else {
+      abortSignal.addEventListener("abort", onAbort, { once: true });
+    }
+  });
+
+  try {
+    // Unlike assistant-stream's helper, cancellation wins when validation aborts and rejects synchronously.
+    return await Promise.race([abortPromise, value]);
+  } finally {
+    abortSignal.removeEventListener("abort", onAbort);
+  }
+};
+
+// AbortSignal.any sits above the browserslist floor and rejects any input that
+// is not a native AbortSignal, which a navigator.modelContext polyfill's signal
+// is not. The merged signal tracks its inputs only until cleanup runs, where
+// the single-signal path hands the caller the lifecycle signal itself.
+const combineAbortSignals = (
+  callerSignal: AbortSignal,
+  lifecycleSignal: AbortSignal,
+): { signal: AbortSignal; cleanup: () => void } => {
+  const controller = new AbortController();
+  const teardown: (() => void)[] = [];
+  const cleanup = () => {
+    while (teardown.length) teardown.pop()!();
+  };
+  const abort = (reason: unknown) => {
+    cleanup();
+    controller.abort(reason);
+  };
+  const listen = (signal: AbortSignal) => {
+    const onAbort = () => abort(signal.reason);
+    signal.addEventListener("abort", onAbort, { once: true });
+    teardown.push(() => signal.removeEventListener("abort", onAbort));
+  };
+
+  if (callerSignal.aborted) {
+    abort(callerSignal.reason);
+  } else if (lifecycleSignal.aborted) {
+    abort(lifecycleSignal.reason);
+  } else {
+    listen(callerSignal);
+    listen(lifecycleSignal);
+  }
+
+  return { signal: controller.signal, cleanup };
+};
+
 export const toWebMcpTool = (
   name: string,
   getTool: () => Tool<any, any>,
@@ -133,17 +195,33 @@ export const toWebMcpTool = (
     const tool = getTool();
     const args = (rawArgs ?? {}) as Record<string, unknown>;
     const toolCallId = generateId();
+    let cleanup: (() => void) | undefined;
     try {
       const callerSignal = context?.signal;
-      const abortSignal = !callerSignal
-        ? lifecycleSignal
-        : !lifecycleSignal
-          ? callerSignal
-          : AbortSignal.any([callerSignal, lifecycleSignal]);
+      let abortSignal: AbortSignal | undefined;
+      if (callerSignal && lifecycleSignal) {
+        const combined = combineAbortSignals(callerSignal, lifecycleSignal);
+        abortSignal = combined.signal;
+        cleanup = combined.cleanup;
+      } else {
+        abortSignal = callerSignal ?? lifecycleSignal;
+      }
+
+      if (abortSignal?.aborted) {
+        return errorResult("Tool execution was cancelled.");
+      }
+
       let executeFn = tool.execute;
       if (isStandardSchema(tool.parameters)) {
-        let validation = tool.parameters["~standard"].validate(args);
-        validation = await validation;
+        const result = tool.parameters["~standard"].validate(args);
+        const validation = isThenable(result)
+          ? abortSignal
+            ? await raceWithAbort(result, abortSignal)
+            : await result
+          : result;
+        if (validation === TOOL_ABORTED) {
+          return errorResult("Tool execution was cancelled.");
+        }
         if (validation.issues) {
           const issues = validation.issues;
           executeFn =
@@ -175,6 +253,8 @@ export const toWebMcpTool = (
       return await toMcpContent(result, { tool, toolCallId, args });
     } catch (e) {
       return errorResult(e instanceof Error ? e.message : String(e));
+    } finally {
+      cleanup?.();
     }
   },
 });

package/src/unstable/webmcp/useWebMcpProvider.ts

@@ -7,6 +7,7 @@ import type { Tool } from "assistant-stream";
 import { getDefaultWebMcpHost, type WebMcpHost } from "./webmcp-host";
 import { defaultWebMcpFilter, toWebMcpInputSchema } from "./convertTools";
 import { WebMcpRegistrationResource } from "./WebMcpRegistrationResource";
+import { shallowEqual } from "@assistant-ui/store/internal";
 import {
   useModelContextSnapshot,
   type ModelContextSnapshotSource,
@@ -64,10 +65,7 @@ const useStableNames = (names: readonly (string | null)[]) => {
   const [cell] = useState(() => ({ names: EMPTY_NAMES }));
   const next = names.filter((name): name is string => name !== null).sort();
   const previous = cell.names;
-  if (
-    previous.length === next.length &&
-    previous.every((name, index) => name === next[index])
-  ) {
+  if (shallowEqual(previous, next)) {
     return previous;
   }
   cell.names = next;

package/src/utils/Primitive.test.tsx

@@ -1,6 +1,6 @@
 import { describe, expect, it } from "vitest";
 import { renderToStaticMarkup } from "react-dom/server";
-import { Primitive, withRenderProp } from "./Primitive";
+import { Primitive, renderSlot, withRenderProp } from "./Primitive";
 
 const ALL_NODES = [
   "a",
@@ -103,3 +103,48 @@ describe("Primitive", () => {
     });
   });
 });
+
+describe("renderSlot", () => {
+  it("lets supplied children replace the render element's own", () => {
+    expect(
+      renderToStaticMarkup(renderSlot(<em>Original</em>, "Override", {})),
+    ).toBe("<em>Override</em>");
+  });
+
+  it("falls back to the render element's own children", () => {
+    expect(
+      renderToStaticMarkup(renderSlot(<em>Fallback</em>, undefined, {})),
+    ).toBe("<em>Fallback</em>");
+  });
+
+  it("treats null children as supplied rather than absent", () => {
+    expect(renderToStaticMarkup(renderSlot(<em>Fallback</em>, null, {}))).toBe(
+      "<em></em>",
+    );
+  });
+
+  it("merges call-site props into the render element", () => {
+    const html = renderToStaticMarkup(
+      renderSlot(<em className="child" id="kept" />, "text", {
+        className: "parent",
+        id: "dropped",
+      }),
+    );
+
+    expect(html).toContain('id="kept"');
+    expect(html).toContain("parent");
+    expect(html).toContain("child");
+  });
+
+  it("applies the same precedence as withRenderProp", () => {
+    expect(
+      renderToStaticMarkup(
+        renderSlot(<em>Fallback</em>, undefined, { className: "parent" }),
+      ),
+    ).toBe(
+      renderToStaticMarkup(
+        <Primitive.span className="parent" render={<em>Fallback</em>} />,
+      ),
+    );
+  });
+});

package/src/utils/Primitive.tsx

@@ -1,4 +1,5 @@
 import {
+  type ComponentProps,
   type ComponentPropsWithoutRef,
   type ComponentRef,
   type ElementType,
@@ -12,6 +13,7 @@ import {
   isValidElement,
 } from "react";
 import { Primitive as RadixPrimitive } from "radix-ui/internal";
+import { Slot } from "radix-ui";
 
 /**
  * Thin wrapper around Radix `Primitive` that adds `render` prop support.
@@ -64,6 +66,39 @@ type PrimitiveRef<E extends PrimitiveNode> = ComponentRef<
   (typeof RadixPrimitive)[E]
 >;
 
+/**
+ * Composes the children of a `render` element. Outer children win when supplied;
+ * the render element's own children are the fallback.
+ */
+function composeRenderElement(
+  render: ReactElement,
+  children: ReactNode,
+): ReactElement {
+  return cloneElement(
+    render,
+    undefined,
+    children !== undefined
+      ? children
+      : (render.props as { children?: ReactNode }).children,
+  );
+}
+
+/**
+ * Composes a `render` element at a call site that has already computed its props.
+ *
+ * `withRenderProp` wraps a component; this covers the case where the props are
+ * already in hand and the composition happens inline.
+ */
+function renderSlot(
+  render: ReactElement,
+  children: ReactNode,
+  props: ComponentProps<typeof Slot.Root>,
+): ReactElement {
+  return (
+    <Slot.Root {...props}>{composeRenderElement(render, children)}</Slot.Root>
+  );
+}
+
 function withRenderProp<T extends ElementType>(Component: T) {
   const Wrapped = forwardRef<ComponentRef<T>, WithRenderPropRuntimeProps<T>>(
     (
@@ -78,14 +113,9 @@ function withRenderProp<T extends ElementType>(Component: T) {
       const Comp = Component as any;
 
       if (render && isValidElement(render)) {
-        const renderChildren =
-          children !== undefined
-            ? children
-            : ((render.props as Record<string, unknown>).children as ReactNode);
-
         return (
           <Comp {...(rest as any)} asChild ref={ref}>
-            {cloneElement(render, undefined, renderChildren)}
+            {composeRenderElement(render, children)}
           </Comp>
         );
       }
@@ -129,5 +159,5 @@ const Primitive = NODES.reduce(
   },
 );
 
-export { Primitive, withRenderProp };
+export { Primitive, renderSlot, withRenderProp };
 export type { PrimitiveProps, WithRenderPropProps };

package/src/utils/getSelectionMessageId.test.ts

@@ -1,7 +1,7 @@
 /**
  * @vitest-environment jsdom
  */
-import { afterEach, describe, expect, it } from "vitest";
+import { afterEach, assert, describe, expect, it, vi } from "vitest";
 import { getSelectionMessageId } from "./getSelectionMessageId";
 
 const selectText = (start: Text, end = start) => {
@@ -24,6 +24,7 @@ const textNode = (selector: string) => {
 };
 
 afterEach(() => {
+  vi.restoreAllMocks();
   window.getSelection()?.removeAllRanges();
   document.body.replaceChildren();
 });
@@ -116,6 +117,180 @@ describe("getSelectionMessageId", () => {
     ).toBeNull();
   });
 
+  it.each(["", "data-aui-quote-selectable"])(
+    "rejects selections spanning an excluded subtree with root %s",
+    (marker) => {
+      document.body.innerHTML = `
+        <div data-message-id="message-1" ${marker}>
+          <span id="before">before</span>
+          <span data-aui-quote-selectable="false">excluded</span>
+          <span id="after">after</span>
+        </div>
+      `;
+
+      expect(getSelectionMessageId(selectText(textNode("#before")))).toBe(
+        "message-1",
+      );
+      expect(
+        getSelectionMessageId(
+          selectText(textNode("#before"), textNode("#after")),
+        ),
+      ).toBeNull();
+    },
+  );
+
+  it("rejects a paragraph selection that spans an inline exclusion", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1">
+        <p id="text" data-aui-quote-selectable>before <span data-aui-quote-selectable="false">[1]</span> after</p>
+      </div>
+    `;
+
+    const paragraph = document.querySelector("#text");
+    assert(paragraph);
+    const selection = selectText(textNode("#text"));
+    selection.getRangeAt(0).selectNodeContents(paragraph);
+
+    expect(getSelectionMessageId(selection)).toBeNull();
+  });
+
+  it("accepts selections that stop or start at an excluded node boundary", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1" data-aui-quote-selectable>
+        <span id="before">before</span><span id="chip" data-aui-quote-selectable="false">[1]</span><span id="after">after</span>
+      </div>
+    `;
+
+    const chip = document.querySelector("#chip");
+    assert(chip);
+    const before = selectText(textNode("#before"));
+    before.getRangeAt(0).setEndBefore(chip);
+    expect(getSelectionMessageId(before)).toBe("message-1");
+
+    const after = selectText(textNode("#after"));
+    after.getRangeAt(0).setStartAfter(chip);
+    expect(getSelectionMessageId(after)).toBe("message-1");
+  });
+
+  it("preserves a nested quote region inside an excluded subtree", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1">
+        <div data-aui-quote-selectable="false">
+          <p data-aui-quote-selectable>
+            <span id="before">before</span>
+            <span data-aui-quote-selectable="false">[1]</span>
+            <span id="after">after</span>
+          </p>
+        </div>
+      </div>
+    `;
+
+    expect(getSelectionMessageId(selectText(textNode("#before")))).toBe(
+      "message-1",
+    );
+    expect(
+      getSelectionMessageId(
+        selectText(textNode("#before"), textNode("#after")),
+      ),
+    ).toBeNull();
+  });
+
+  it("checks another range outside the active quote region, including its excluded ancestor", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1">
+        <p id="other">before <span id="chip" data-aui-quote-selectable="false">[1]</span> after</p>
+        <p id="active" data-aui-quote-selectable>active text</p>
+      </div>
+    `;
+
+    const other = document.querySelector("#other");
+    assert(other);
+    const range = document.createRange();
+    range.selectNodeContents(other);
+    const selection = selectText(textNode("#active"));
+    const active = selection.getRangeAt(0);
+    vi.spyOn(selection, "rangeCount", "get").mockReturnValue(2);
+    vi.spyOn(selection, "getRangeAt").mockImplementation((index) => {
+      if (index === 0) return range;
+      if (index === 1) return active;
+      throw new DOMException("Range index out of bounds", "IndexSizeError");
+    });
+
+    expect(getSelectionMessageId(selection)).toBeNull();
+
+    range.selectNodeContents(textNode("#chip"));
+    expect(getSelectionMessageId(selection)).toBeNull();
+  });
+
+  it("rejects another range inside the excluded ancestor of the active quote region", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1">
+        <div data-aui-quote-selectable="false">
+          <span id="sibling">excluded prose</span>
+          <p id="active" data-aui-quote-selectable>active text</p>
+        </div>
+      </div>
+    `;
+
+    const range = document.createRange();
+    range.selectNodeContents(textNode("#sibling"));
+    const selection = selectText(textNode("#active"));
+    const active = selection.getRangeAt(0);
+    vi.spyOn(selection, "rangeCount", "get").mockReturnValue(2);
+    vi.spyOn(selection, "getRangeAt").mockImplementation((index) => {
+      if (index === 0) return range;
+      if (index === 1) return active;
+      throw new DOMException("Range index out of bounds", "IndexSizeError");
+    });
+
+    expect(getSelectionMessageId(selection)).toBeNull();
+  });
+
+  it("rejects another range in a different message", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1"><p id="active">active text</p></div>
+      <div data-message-id="message-2">
+        <p id="other" data-aui-quote-selectable="false">excluded prose</p>
+      </div>
+    `;
+
+    const range = document.createRange();
+    range.selectNodeContents(textNode("#other"));
+    const selection = selectText(textNode("#active"));
+    const active = selection.getRangeAt(0);
+    vi.spyOn(selection, "rangeCount", "get").mockReturnValue(2);
+    vi.spyOn(selection, "getRangeAt").mockImplementation((index) => {
+      if (index === 0) return range;
+      if (index === 1) return active;
+      throw new DOMException("Range index out of bounds", "IndexSizeError");
+    });
+
+    expect(getSelectionMessageId(selection)).toBeNull();
+  });
+
+  it("accepts disjoint ranges on either side of an excluded gap", () => {
+    document.body.innerHTML = `
+      <div data-message-id="message-1" data-aui-quote-selectable>
+        <p id="before">before</p>
+        <p data-aui-quote-selectable="false">excluded</p>
+        <p id="after">after</p>
+      </div>
+    `;
+
+    const before = document.createRange();
+    before.selectNodeContents(textNode("#before"));
+    const selection = selectText(textNode("#after"));
+    const after = selection.getRangeAt(0);
+    vi.spyOn(selection, "rangeCount", "get").mockReturnValue(2);
+    vi.spyOn(selection, "getRangeAt").mockImplementation((index) => {
+      if (index === 0) return before;
+      if (index === 1) return after;
+      throw new DOMException("Range index out of bounds", "IndexSizeError");
+    });
+
+    expect(getSelectionMessageId(selection)).toBe("message-1");
+  });
+
   it("rejects selections outside quote-selectable regions when a message opts in", () => {
     document.body.innerHTML = `
       <div data-message-id="message-1">

package/src/utils/getSelectionMessageId.ts

@@ -43,6 +43,17 @@ const findQuoteMarker = (
   return marker;
 };
 
+const intersectsExcluded = (scope: Element, selection: Selection): boolean => {
+  const ranges = Array.from({ length: selection.rangeCount }, (_, i) =>
+    selection.getRangeAt(i),
+  );
+  for (const marker of scope.querySelectorAll(QUOTE_SELECTABLE_SELECTOR)) {
+    if (!isExcluded(marker)) continue;
+    if (ranges.some((range) => range.intersectsNode(marker))) return true;
+  }
+  return false;
+};
+
 export const getSelectionMessageId = (selection: Selection): string | null => {
   const { anchorNode, focusNode } = selection;
   if (!anchorNode || !focusNode) return null;
@@ -63,9 +74,16 @@ export const getSelectionMessageId = (selection: Selection): string | null => {
   if (anchorMarker && isExcluded(anchorMarker)) return null;
   if (focusMarker && isExcluded(focusMarker)) return null;
 
-  if (!hasQuoteSelectableRegion(anchorMessageElement)) return messageId;
+  if (hasQuoteSelectableRegion(anchorMessageElement)) {
+    if (!anchorMarker || anchorMarker !== focusMarker) return null;
+  }
+
+  const scope = anchorMarker ?? anchorMessageElement;
 
-  if (!anchorMarker || anchorMarker !== focusMarker) return null;
+  for (let i = 0; i < selection.rangeCount; i++) {
+    const { commonAncestorContainer } = selection.getRangeAt(i);
+    if (!scope.contains(commonAncestorContainer)) return null;
+  }
 
-  return messageId;
+  return intersectsExcluded(scope, selection) ? null : messageId;
 };