@assistant-ui/react-ink 0.0.1 → 0.0.4

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");
+  });
+});

package/src/primitives/toolCall/ToolFallback.tsx

@@ -0,0 +1,227 @@
+import { type ReactNode, useMemo } from "react";
+import { Box, Text } from "ink";
+import Spinner from "ink-spinner";
+import type { ToolCallMessagePartStatus } from "@assistant-ui/core";
+import type { ToolCallMessagePartProps } from "../../types";
+
+export type ToolCallStatus =
+  | "running"
+  | "complete"
+  | "error"
+  | "requires-action";
+
+type ToolFallbackBaseProps = Omit<
+  ToolCallMessagePartProps,
+  "addResult" | "resume"
+> &
+  Partial<Pick<ToolCallMessagePartProps, "addResult" | "resume">>;
+
+const STATUS_ICONS: Record<Exclude<ToolCallStatus, "running">, string> = {
+  complete: "+",
+  error: "x",
+  "requires-action": "?",
+};
+
+const STATUS_COLORS: Record<ToolCallStatus, string> = {
+  running: "yellow",
+  complete: "green",
+  error: "red",
+  "requires-action": "cyan",
+};
+
+const prettyPrintArgs = (argsText: string): string => {
+  if (!argsText) return "";
+  try {
+    return JSON.stringify(JSON.parse(argsText), null, 2);
+  } catch {
+    return argsText;
+  }
+};
+
+const formatResult = (result: unknown): string => {
+  if (result === undefined || result === null) return "";
+  if (typeof result === "string") return result;
+  try {
+    return JSON.stringify(result, null, 2);
+  } catch {
+    return String(result);
+  }
+};
+
+const truncate = (text: string, maxLines: number): string => {
+  const lines = text.split("\n");
+  if (lines.length <= maxLines) return text;
+  return (
+    lines.slice(0, maxLines).join("\n") +
+    `\n... (${lines.length - maxLines} more lines)`
+  );
+};
+
+export type ToolFallbackProps = ToolFallbackBaseProps & {
+  /** Force expanded or collapsed. When unset, auto-expands while running, awaiting action, or errored. */
+  expanded?: boolean;
+  /** Maximum lines to show for args when expanded. Defaults to 20. */
+  maxArgLines?: number;
+  /** Maximum lines to show for result when expanded. Defaults to 20. */
+  maxResultLines?: number;
+  /** Maximum lines to show for result preview when collapsed. Defaults to 1. */
+  maxResultPreviewLines?: number;
+  /** Custom header renderer */
+  renderHeader?: (props: {
+    toolName: string;
+    status: ToolCallStatus;
+    expanded: boolean;
+  }) => ReactNode;
+  /** Custom args renderer */
+  renderArgs?: (props: { args: unknown; argsText: string }) => ReactNode;
+  /** Custom result renderer */
+  renderResult?: (props: { result: unknown; isError: boolean }) => ReactNode;
+};
+
+const resolveStatus = (
+  part: {
+    interrupt: unknown;
+    isError: boolean | undefined;
+    result: unknown;
+  },
+  status?: ToolCallMessagePartStatus,
+): ToolCallStatus => {
+  if (status?.type === "requires-action") return "requires-action";
+  if (status?.type === "incomplete") return "error";
+  if (status?.type === "complete") return "complete";
+  if (part.isError) return "error";
+  if (part.result !== undefined) return "complete";
+  if (part.interrupt) return "requires-action";
+  return "running";
+};
+
+const getDisplayResult = (
+  result: unknown,
+  status?: ToolCallMessagePartStatus,
+) => {
+  if (result !== undefined) return result;
+  if (status?.type === "incomplete") return status.error;
+  return undefined;
+};
+
+const getFallbackErrorText = (status?: ToolCallMessagePartStatus) => {
+  if (status?.type === "incomplete" && status.reason === "cancelled") {
+    return "Tool call cancelled";
+  }
+  return "Tool call failed";
+};
+
+const StatusIcon = ({ status }: { status: ToolCallStatus }) => {
+  if (status === "running") {
+    return (
+      <Text color={STATUS_COLORS.running}>
+        <Spinner type="line" />
+      </Text>
+    );
+  }
+  return <Text color={STATUS_COLORS[status]}>{STATUS_ICONS[status]}</Text>;
+};
+
+export const ToolFallback = ({
+  toolName,
+  args,
+  argsText,
+  result: partResult,
+  isError,
+  interrupt,
+  status,
+  expanded: expandedProp,
+  maxArgLines = 20,
+  maxResultLines = 20,
+  maxResultPreviewLines = 1,
+  renderHeader,
+  renderArgs,
+  renderResult,
+}: ToolFallbackProps) => {
+  const displayStatus = resolveStatus(
+    {
+      interrupt,
+      isError,
+      result: partResult,
+    },
+    status,
+  );
+  const expanded =
+    expandedProp ??
+    (displayStatus === "running" ||
+      displayStatus === "requires-action" ||
+      displayStatus === "error");
+
+  const result = getDisplayResult(partResult, status);
+  const resultStr =
+    result !== undefined || displayStatus === "error"
+      ? formatResult(result)
+      : "";
+  const argsDisplay = useMemo(() => prettyPrintArgs(argsText), [argsText]);
+
+  return (
+    <Box flexDirection="column">
+      <Box gap={1}>
+        {renderHeader ? (
+          renderHeader({
+            toolName,
+            status: displayStatus,
+            expanded,
+          })
+        ) : (
+          <>
+            <StatusIcon status={displayStatus} />
+            <Text bold>{toolName}</Text>
+            {!expanded && resultStr ? (
+              <Text dimColor>{truncate(resultStr, maxResultPreviewLines)}</Text>
+            ) : null}
+          </>
+        )}
+      </Box>
+
+      {expanded && (
+        <Box flexDirection="column" marginLeft={2}>
+          {argsText ? (
+            <Box flexDirection="column">
+              <Text dimColor>Args:</Text>
+              {renderArgs ? (
+                renderArgs({ args, argsText })
+              ) : (
+                <Text>{truncate(argsDisplay, maxArgLines)}</Text>
+              )}
+            </Box>
+          ) : null}
+
+          {partResult !== undefined || displayStatus === "error" ? (
+            <Box flexDirection="column">
+              <Text dimColor>
+                {displayStatus === "error" ? "Error:" : "Result:"}
+              </Text>
+              {renderResult ? (
+                renderResult({
+                  result,
+                  isError: displayStatus === "error",
+                })
+              ) : displayStatus === "error" ? (
+                <Text color="red">
+                  {truncate(
+                    resultStr || getFallbackErrorText(status),
+                    maxResultLines,
+                  )}
+                </Text>
+              ) : (
+                <Text>{truncate(resultStr, maxResultLines)}</Text>
+              )}
+            </Box>
+          ) : null}
+
+          {displayStatus === "running" && <Text dimColor>Running...</Text>}
+
+          {displayStatus === "requires-action" && (
+            <Text color="cyan">Waiting for approval...</Text>
+          )}
+        </Box>
+      )}
+    </Box>
+  );
+};

package/src/primitives/toolCall/index.ts

@@ -0,0 +1,5 @@
+export {
+  ToolFallback as Fallback,
+  type ToolFallbackProps as FallbackProps,
+  type ToolCallStatus,
+} from "./ToolFallback";