@assistant-ui/react/src/primitives/composer/trigger/triggerSelectionResource.test.ts

Open-source TypeScript/React library for building production-grade AI chat experiences

import { describe, expect, it, vi } from "vitest";
import { createTapRoot, flushTapSync, useResource } from "@assistant-ui/tap";
import type {
  Unstable_DirectiveFormatter,
  Unstable_TriggerItem,
} from "@assistant-ui/core";
import type { AssistantClient } from "@assistant-ui/store";
import { TriggerDetectionResource } from "./triggerDetectionResource";
import {
  TriggerSelectionResource,
  type TriggerBehavior,
} from "./triggerSelectionResource";

const item: Unstable_TriggerItem = {
  id: "file",
  type: "file",
  label: "file.txt",
};

const formatter: Unstable_DirectiveFormatter = {
  serialize: () => "@file.txt",
  parse: () => [],
};

const makeAui = (textRef: { value: string }) => {
  const setText = vi.fn((value: string) => {
    textRef.value = value;
  });

  return {
    aui: {
      composer: {
        getState: () => ({ text: textRef.value }),
        setText,
      },
    } as unknown as AssistantClient,
    setText,
  };
};

const render = ({
  behavior,
  text,
  trigger,
}: {
  behavior: TriggerBehavior;
  text: string;
  trigger: { offset: number; query: string; endOffset: number };
}) => {
  const textRef = { value: text };
  const { aui, setText } = makeAui(textRef);
  const setCursorPosition = vi.fn();
  const onSelected = vi.fn();
  const root = createTapRoot(function Root() {
    return useResource(
      TriggerSelectionResource({
        behavior,
        trigger,
        aui,
        triggerChar: "@",
        setCursorPosition,
        onSelected,
      }),
    );
  });

  return { root, setText, setCursorPosition, onSelected };
};

describe("TriggerSelectionResource", () => {
  it("resynchronizes the cursor after inserting a directive", () => {
    const { root, setText, setCursorPosition, onSelected } = render({
      behavior: { kind: "directive", formatter },
      text: "@file tail",
      trigger: { offset: 0, query: "file", endOffset: 5 },
    });

    root.getValue().selectItem(item);

    expect(setText).toHaveBeenCalledWith("@file.txt tail", { origin: "user" });
    expect(setCursorPosition).toHaveBeenCalledWith("@file.txt".length + 1);
    expect(onSelected).toHaveBeenCalledOnce();
  });

  it("reports an inserted directive once a select override finishes the insertion", () => {
    const onInserted = vi.fn();
    const { root, setText, onSelected } = render({
      behavior: { kind: "directive", formatter, onInserted },
      text: "@file tail",
      trigger: { offset: 0, query: "file", endOffset: 5 },
    });
    let finishInsertion: (() => void) | undefined;
    root.getValue().registerSelectItemOverride((_, onOverrideInserted) => {
      finishInsertion = onOverrideInserted;
      return true;
    });

    root.getValue().selectItem(item);

    expect(onInserted).not.toHaveBeenCalled();
    expect(setText).not.toHaveBeenCalled();
    expect(onSelected).toHaveBeenCalledOnce();

    finishInsertion?.();

    expect(onInserted).toHaveBeenCalledExactlyOnceWith(item);
  });

  it("resynchronizes the cursor after removing an action", () => {
    const { root, setText, setCursorPosition } = render({
      behavior: {
        kind: "action",
        formatter,
        onExecute: vi.fn(),
        removeOnExecute: true,
      },
      text: "before @file tail",
      trigger: { offset: 7, query: "file", endOffset: 12 },
    });

    root.getValue().selectItem(item);

    expect(setText).toHaveBeenCalledWith("before tail", { origin: "user" });
    expect(setCursorPosition).toHaveBeenCalledWith(7);
  });

  it("does not rewrite composer text on a second select after insertion", () => {
    const textRef = { value: "@fil" };
    const { aui, setText } = makeAui(textRef);
    const onSelected = vi.fn();
    const root = createTapRoot(function Root() {
      const detection = useResource(
        TriggerDetectionResource({
          text: textRef.value,
          triggerChar: "@",
        }),
      );
      return useResource(
        TriggerSelectionResource({
          behavior: { kind: "directive", formatter },
          trigger: detection.trigger,
          aui,
          triggerChar: "@",
          setCursorPosition: detection.setCursorPosition,
          onSelected,
        }),
      );
    });

    flushTapSync(() => {
      root.getValue().selectItem(item);
    });
    root.getValue().selectItem(item);

    expect(setText).toHaveBeenCalledOnce();
    expect(setText).toHaveBeenCalledWith("@file.txt ", { origin: "user" });
    expect(textRef.value).toBe("@file.txt ");
  });

  it("replaces the matcher endOffset even when query is shorter", () => {
    const { root, setText } = render({
      behavior: { kind: "directive", formatter },
      text: "hello @Example UK tail",
      trigger: { offset: 6, query: "example", endOffset: 17 },
    });

    root.getValue().selectItem(item);

    expect(setText).toHaveBeenCalledWith("hello @file.txt tail", {
      origin: "user",
    });
  });
});