@assistant-ui/react-ink 0.0.2 → 0.0.4

package/src/primitives/message/MessageContent.test.tsx

@@ -0,0 +1,83 @@
+import type { ReactElement } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { cleanup, render } from "ink-testing-library";
+import { MessageContent } from "./MessageContent";
+
+const mockUseAui = vi.fn();
+const mockUseAuiState = vi.fn();
+
+type UseAuiStateSelector = Parameters<
+  typeof import("@assistant-ui/store")["useAuiState"]
+>[0];
+
+vi.mock("@assistant-ui/store", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@assistant-ui/store")>();
+  return {
+    ...actual,
+    useAui: () => mockUseAui(),
+    useAuiState: (selector: UseAuiStateSelector) => mockUseAuiState(selector),
+  };
+});
+
+const renderFrame = async (node: ReactElement) => {
+  const instance = render(node);
+  await new Promise((resolve) => setTimeout(resolve, 0));
+  return instance.lastFrame() ?? "";
+};
+
+afterEach(() => {
+  cleanup();
+  vi.clearAllMocks();
+});
+
+describe("MessageContent", () => {
+  it("renders incomplete tool calls as errors instead of running forever", async () => {
+    const state = {
+      tools: { tools: {} },
+      message: {
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: "tool-call-1",
+            toolName: "search",
+            args: { query: "weather" },
+            argsText: '{"query":"weather"}',
+          },
+        ],
+        parts: [
+          {
+            type: "tool-call",
+            toolCallId: "tool-call-1",
+            toolName: "search",
+            args: { query: "weather" },
+            argsText: '{"query":"weather"}',
+            status: {
+              type: "incomplete",
+              reason: "error",
+              error: "request failed",
+            },
+          },
+        ],
+      },
+    };
+
+    mockUseAuiState.mockImplementation((selector: UseAuiStateSelector) =>
+      selector(state as never),
+    );
+    mockUseAui.mockReturnValue({
+      message: () => ({
+        part: () => ({
+          addToolResult: vi.fn(),
+          resumeToolCall: vi.fn(),
+        }),
+      }),
+    });
+
+    const frame = await renderFrame(<MessageContent />);
+
+    expect(frame).toContain("search");
+    expect(frame).toContain("Error:");
+    expect(frame).toContain("request failed");
+    expect(frame).not.toContain("Running...");
+  });
+});

package/src/primitives/message/MessageContent.tsx

@@ -3,16 +3,17 @@ import { Text } from "ink";
 import type {
   ThreadUserMessagePart,
   ThreadAssistantMessagePart,
-  ToolCallMessagePart,
-  DataMessagePart,
+  MessagePartState,
 } from "@assistant-ui/core";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import type {
   ToolCallMessagePartProps,
   DataMessagePartProps,
 } from "../../types";
+import { ToolFallback } from "../toolCall/ToolFallback";
 
 type MessageContentPart = ThreadUserMessagePart | ThreadAssistantMessagePart;
+type MessageContentStatePart = MessagePartState;
 
 export type MessageContentProps = {
   renderText?: (props: {
@@ -59,9 +60,12 @@ const ToolUIDisplay = ({
   index,
 }: {
   Fallback:
-    | ((props: { part: ToolCallMessagePart; index: number }) => ReactElement)
+    | ((props: {
+        part: Extract<MessageContentPart, { type: "tool-call" }>;
+        index: number;
+      }) => ReactElement)
     | undefined;
-  part: ToolCallMessagePart;
+  part: Extract<MessageContentStatePart, { type: "tool-call" }>;
   index: number;
 }) => {
   const aui = useAui();
@@ -75,18 +79,17 @@ const ToolUIDisplay = ({
     () => aui.message().part({ index }),
     [aui, index],
   );
+  const toolProps = {
+    ...(part as ToolCallMessagePartProps),
+    addResult: partMethods.addToolResult,
+    resume: partMethods.resumeToolCall,
+  };
 
   if (Render) {
-    return (
-      <Render
-        {...(part as ToolCallMessagePartProps)}
-        addResult={partMethods.addToolResult}
-        resume={partMethods.resumeToolCall}
-      />
-    );
+    return <Render {...toolProps} />;
   }
   if (Fallback) return <Fallback part={part} index={index} />;
-  return null;
+  return <ToolFallback {...toolProps} />;
 };
 
 const DataUIDisplay = ({
@@ -95,9 +98,12 @@ const DataUIDisplay = ({
   index,
 }: {
   Fallback:
-    | ((props: { part: DataMessagePart; index: number }) => ReactElement)
+    | ((props: {
+        part: Extract<MessageContentPart, { type: "data" }>;
+        index: number;
+      }) => ReactElement)
     | undefined;
-  part: DataMessagePart;
+  part: Extract<MessageContentStatePart, { type: "data" }>;
   index: number;
 }) => {
   const Render = useAuiState((s) => {
@@ -119,7 +125,7 @@ export const MessageContent = ({
   renderFile,
   renderData,
 }: MessageContentProps) => {
-  const content = useAuiState((s) => s.message.content);
+  const content = useAuiState((s) => s.message.parts);
 
   return (
     <>

package/src/primitives/message/index.ts

@@ -1,11 +1,17 @@
-export { MessageRoot, type MessageRootProps } from "./MessageRoot";
-export { MessageContent, type MessageContentProps } from "./MessageContent";
 export {
-  MessagePrimitiveParts,
-  MessagePrimitivePartByIndex,
+  MessageRoot as Root,
+  type MessageRootProps as RootProps,
+} from "./MessageRoot";
+export {
+  MessageContent as Content,
+  type MessageContentProps as ContentProps,
+} from "./MessageContent";
+export {
+  MessagePrimitiveParts as Parts,
+  MessagePrimitivePartByIndex as PartByIndex,
 } from "./MessageParts";
-export { MessageIf, type MessageIfProps } from "./MessageIf";
+export { MessageIf as If, type MessageIfProps as IfProps } from "./MessageIf";
 export {
-  MessageAttachments,
-  MessageAttachmentByIndex,
+  MessageAttachments as Attachments,
+  MessageAttachmentByIndex as AttachmentByIndex,
 } from "./MessageAttachments";

package/src/primitives/suggestion/index.ts

@@ -1,12 +1,12 @@
 export {
-  SuggestionTitle,
-  type SuggestionTitleProps,
+  SuggestionTitle as Title,
+  type SuggestionTitleProps as TitleProps,
 } from "./SuggestionTitle";
 export {
-  SuggestionDescription,
-  type SuggestionDescriptionProps,
+  SuggestionDescription as Description,
+  type SuggestionDescriptionProps as DescriptionProps,
 } from "./SuggestionDescription";
 export {
-  SuggestionTrigger,
-  type SuggestionTriggerProps,
+  SuggestionTrigger as Trigger,
+  type SuggestionTriggerProps as TriggerProps,
 } from "./SuggestionTrigger";

package/src/primitives/thread/ThreadEmpty.tsx

@@ -1,5 +1,5 @@
 import type { ReactNode } from "react";
-import { useThreadIsEmpty } from "../../primitive-hooks/useThreadIsEmpty";
+import { useThreadIsEmpty } from "@assistant-ui/core/react";
 
 export type ThreadEmptyProps = {
   children: ReactNode;

package/src/primitives/thread/ThreadMessages.tsx

@@ -1,42 +1,137 @@
-import { type ReactElement } from "react";
+import { type ComponentType, type FC, memo } from "react";
 import { Box } from "ink";
 import type { ThreadMessage } from "@assistant-ui/core";
-import { useAui, useAuiState, AuiProvider, Derived } from "@assistant-ui/store";
+import { useAuiState } from "@assistant-ui/store";
+import { MessageByIndexProvider } from "@assistant-ui/core/react";
+
+type MessageComponents =
+  | {
+      Message: ComponentType;
+      EditComposer?: ComponentType | undefined;
+      UserEditComposer?: ComponentType | undefined;
+      AssistantEditComposer?: ComponentType | undefined;
+      SystemEditComposer?: ComponentType | undefined;
+      UserMessage?: ComponentType | undefined;
+      AssistantMessage?: ComponentType | undefined;
+      SystemMessage?: ComponentType | undefined;
+    }
+  | {
+      Message?: ComponentType | undefined;
+      EditComposer?: ComponentType | undefined;
+      UserEditComposer?: ComponentType | undefined;
+      AssistantEditComposer?: ComponentType | undefined;
+      SystemEditComposer?: ComponentType | undefined;
+      UserMessage: ComponentType;
+      AssistantMessage: ComponentType;
+      SystemMessage?: ComponentType | undefined;
+    };
 
 export type ThreadMessagesProps = {
-  renderMessage: (props: {
-    message: ThreadMessage;
-    index: number;
-  }) => ReactElement;
+  components: MessageComponents;
 };
 
-const MessageScope = ({
-  index,
-  children,
-}: {
-  index: number;
-  children: ReactElement;
+const DEFAULT_SYSTEM_MESSAGE = () => null;
+
+const getComponent = (
+  components: MessageComponents,
+  role: ThreadMessage["role"],
+  isEditing: boolean,
+) => {
+  switch (role) {
+    case "user":
+      if (isEditing) {
+        return (
+          components.UserEditComposer ??
+          components.EditComposer ??
+          components.UserMessage ??
+          (components.Message as ComponentType)
+        );
+      } else {
+        return components.UserMessage ?? (components.Message as ComponentType);
+      }
+    case "assistant":
+      if (isEditing) {
+        return (
+          components.AssistantEditComposer ??
+          components.EditComposer ??
+          components.AssistantMessage ??
+          (components.Message as ComponentType)
+        );
+      } else {
+        return (
+          components.AssistantMessage ?? (components.Message as ComponentType)
+        );
+      }
+    case "system":
+      if (isEditing) {
+        return (
+          components.SystemEditComposer ??
+          components.EditComposer ??
+          components.SystemMessage ??
+          (components.Message as ComponentType) ??
+          DEFAULT_SYSTEM_MESSAGE
+        );
+      } else {
+        return (
+          components.SystemMessage ??
+          (components.Message as ComponentType) ??
+          DEFAULT_SYSTEM_MESSAGE
+        );
+      }
+    default: {
+      const _exhaustiveCheck: never = role;
+      throw new Error(`Unknown message role: ${_exhaustiveCheck}`);
+    }
+  }
+};
+
+const ThreadMessageComponent: FC<{ components: MessageComponents }> = ({
+  components,
 }) => {
-  const aui = useAui({
-    message: Derived({
-      source: "thread",
-      query: { type: "index", index },
-      get: (aui) => aui.threads().thread("main").message({ index }),
-    }),
-  });
-
-  return <AuiProvider value={aui}>{children}</AuiProvider>;
+  const role = useAuiState((s) => s.message.role);
+  const isEditing = useAuiState((s) => s.message.composer.isEditing);
+  const Component = getComponent(components, role, isEditing);
+
+  return <Component />;
 };
 
-export const ThreadMessages = ({ renderMessage }: ThreadMessagesProps) => {
-  const messages = useAuiState((s) => s.thread.messages);
+const isComponentsSame = (prev: MessageComponents, next: MessageComponents) => {
+  return (
+    prev.Message === next.Message &&
+    prev.EditComposer === next.EditComposer &&
+    prev.UserEditComposer === next.UserEditComposer &&
+    prev.AssistantEditComposer === next.AssistantEditComposer &&
+    prev.SystemEditComposer === next.SystemEditComposer &&
+    prev.UserMessage === next.UserMessage &&
+    prev.AssistantMessage === next.AssistantMessage &&
+    prev.SystemMessage === next.SystemMessage
+  );
+};
+
+const ThreadMessageByIndex = memo(
+  ({ index, components }: { index: number; components: MessageComponents }) => {
+    return (
+      <MessageByIndexProvider index={index}>
+        <ThreadMessageComponent components={components} />
+      </MessageByIndexProvider>
+    );
+  },
+  (prev, next) =>
+    prev.index === next.index &&
+    isComponentsSame(prev.components, next.components),
+);
+
+export const ThreadMessages = ({ components }: ThreadMessagesProps) => {
+  const messagesLength = useAuiState((s) => s.thread.messages.length);
 
   return (
     <Box flexDirection="column">
-      {(messages as unknown as ThreadMessage[]).map((message, index) => (
-        <MessageScope key={message.id} index={index}>
-          {renderMessage({ message, index })}
-        </MessageScope>
+      {Array.from({ length: messagesLength }, (_, index) => (
+        <ThreadMessageByIndex
+          key={index}
+          index={index}
+          components={components}
+        />
       ))}
     </Box>
   );

package/src/primitives/thread/index.ts

@@ -1,16 +1,22 @@
-export { ThreadRoot, type ThreadRootProps } from "./ThreadRoot";
-export { ThreadMessages, type ThreadMessagesProps } from "./ThreadMessages";
 export {
-  ThreadPrimitiveMessages,
-  ThreadPrimitiveMessageByIndex,
-} from "@assistant-ui/core/react";
-export { ThreadEmpty, type ThreadEmptyProps } from "./ThreadEmpty";
-export { ThreadIf, type ThreadIfProps } from "./ThreadIf";
+  ThreadRoot as Root,
+  type ThreadRootProps as RootProps,
+} from "./ThreadRoot";
+export {
+  ThreadMessages as Messages,
+  type ThreadMessagesProps as MessagesProps,
+} from "./ThreadMessages";
+export { ThreadPrimitiveMessageByIndex as MessageByIndex } from "@assistant-ui/core/react";
+export {
+  ThreadEmpty as Empty,
+  type ThreadEmptyProps as EmptyProps,
+} from "./ThreadEmpty";
+export { ThreadIf as If, type ThreadIfProps as IfProps } from "./ThreadIf";
 export {
-  ThreadSuggestion,
-  type ThreadSuggestionProps,
+  ThreadSuggestion as Suggestion,
+  type ThreadSuggestionProps as SuggestionProps,
 } from "./ThreadSuggestion";
 export {
-  ThreadPrimitiveSuggestions,
-  ThreadPrimitiveSuggestionByIndex,
+  ThreadPrimitiveSuggestions as Suggestions,
+  ThreadPrimitiveSuggestionByIndex as SuggestionByIndex,
 } from "@assistant-ui/core/react";

package/src/primitives/threadList/index.ts

@@ -1,3 +1,12 @@
-export { ThreadListRoot, type ThreadListRootProps } from "./ThreadListRoot";
-export { ThreadListItems, type ThreadListItemsProps } from "./ThreadListItems";
-export { ThreadListNew, type ThreadListNewProps } from "./ThreadListNew";
+export {
+  ThreadListRoot as Root,
+  type ThreadListRootProps as RootProps,
+} from "./ThreadListRoot";
+export {
+  ThreadListItems as Items,
+  type ThreadListItemsProps as ItemsProps,
+} from "./ThreadListItems";
+export {
+  ThreadListNew as New,
+  type ThreadListNewProps as NewProps,
+} from "./ThreadListNew";

package/src/primitives/threadListItem/index.ts

@@ -1,21 +1,21 @@
 export {
-  ThreadListItemRoot,
-  type ThreadListItemRootProps,
+  ThreadListItemRoot as Root,
+  type ThreadListItemRootProps as RootProps,
 } from "./ThreadListItemRoot";
-export { ThreadListItemPrimitiveTitle as ThreadListItemTitle } from "@assistant-ui/core/react";
+export { ThreadListItemPrimitiveTitle as Title } from "@assistant-ui/core/react";
 export {
-  ThreadListItemTrigger,
-  type ThreadListItemTriggerProps,
+  ThreadListItemTrigger as Trigger,
+  type ThreadListItemTriggerProps as TriggerProps,
 } from "./ThreadListItemTrigger";
 export {
-  ThreadListItemDelete,
-  type ThreadListItemDeleteProps,
+  ThreadListItemDelete as Delete,
+  type ThreadListItemDeleteProps as DeleteProps,
 } from "./ThreadListItemDelete";
 export {
-  ThreadListItemArchive,
-  type ThreadListItemArchiveProps,
+  ThreadListItemArchive as Archive,
+  type ThreadListItemArchiveProps as ArchiveProps,
 } from "./ThreadListItemArchive";
 export {
-  ThreadListItemUnarchive,
-  type ThreadListItemUnarchiveProps,
+  ThreadListItemUnarchive as Unarchive,
+  type ThreadListItemUnarchiveProps as UnarchiveProps,
 } from "./ThreadListItemUnarchive";

package/src/primitives/toolCall/ToolFallback.test.tsx

@@ -0,0 +1,40 @@
+import type { ReactElement } from "react";
+import { afterEach, describe, expect, it } from "vitest";
+import { cleanup, render } from "ink-testing-library";
+import { ToolFallback } from "./ToolFallback";
+
+const renderFrame = async (node: ReactElement) => {
+  const instance = render(node);
+  await new Promise((resolve) => setTimeout(resolve, 0));
+  return instance.lastFrame() ?? "";
+};
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("ToolFallback", () => {
+  it("truncates expanded error output with maxResultLines", async () => {
+    const frame = await renderFrame(
+      <ToolFallback
+        expanded
+        maxResultLines={2}
+        type="tool-call"
+        toolCallId="tool-call-1"
+        toolName="search"
+        args={{}}
+        argsText="{}"
+        result={"line 1\nline 2\nline 3\nline 4"}
+        isError
+        status={{ type: "incomplete", reason: "error", error: "boom" }}
+      />,
+    );
+
+    expect(frame).toContain("Error:");
+    expect(frame).toContain("line 1");
+    expect(frame).toContain("line 2");
+    expect(frame).toContain("... (2 more lines)");
+    expect(frame).not.toContain("line 3");
+    expect(frame).not.toContain("line 4");
+  });
+});