feat: AttachmentContext (#766)

packages/react/src/context/providers/ComposerAttachmentProvider.tsx

@@ -0,0 +1,73 @@
+"use client";
+
+import { type FC, type PropsWithChildren, useEffect, useState } from "react";
+import { StoreApi, create } from "zustand";
+import type { ComposerState } from "../stores";
+import { useThreadContext } from "../react";
+import { AttachmentState } from "../stores/Attachment";
+import {
+  AttachmentContext,
+  AttachmentContextValue,
+} from "../react/AttachmentContext";
+
+type ComposerAttachmentProviderProps = PropsWithChildren<{
+  attachmentIndex: number;
+}>;
+
+const getAttachment = (
+  { attachments }: ComposerState,
+  useAttachment: AttachmentContextValue["useAttachment"] | undefined,
+  partIndex: number,
+) => {
+  let attachment = attachments[partIndex];
+  if (!attachment) return null;
+
+  // if the attachment is the same, don't update
+  const currentState = useAttachment?.getState();
+  if (currentState && currentState.attachment === attachment) return null;
+
+  return Object.freeze({ attachment });
+};
+
+const useComposerAttachmentContext = (partIndex: number) => {
+  const { useComposer } = useThreadContext();
+  const [context] = useState<AttachmentContextValue>(() => {
+    const useAttachment = create<AttachmentState>(
+      () => getAttachment(useComposer.getState(), undefined, partIndex)!,
+    );
+
+    return { useAttachment };
+  });
+
+  useEffect(() => {
+    const syncAttachment = (composer: ComposerState) => {
+      const newState = getAttachment(
+        composer,
+        context.useAttachment,
+        partIndex,
+      );
+      if (!newState) return;
+      (context.useAttachment as unknown as StoreApi<AttachmentState>).setState(
+        newState,
+        true,
+      );
+    };
+
+    syncAttachment(useComposer.getState());
+    return useComposer.subscribe(syncAttachment);
+  }, [context, useComposer, partIndex]);
+
+  return context;
+};
+
+export const ComposerAttachmentProvider: FC<
+  ComposerAttachmentProviderProps
+> = ({ attachmentIndex: partIndex, children }) => {
+  const context = useComposerAttachmentContext(partIndex);
+
+  return (
+    <AttachmentContext.Provider value={context}>
+      {children}
+    </AttachmentContext.Provider>
+  );
+};

packages/react/src/context/react/AttachmentContext.ts

@@ -0,0 +1,26 @@
+"use client";
+
+import { createContext, useContext } from "react";
+import { ReadonlyStore } from "../ReadonlyStore";
+import { AttachmentState } from "../stores/Attachment";
+
+export type AttachmentContextValue = {
+  useAttachment: ReadonlyStore<AttachmentState>;
+};
+
+export const AttachmentContext = createContext<AttachmentContextValue | null>(
+  null,
+);
+
+export function useAttachmentContext(): AttachmentContextValue;
+export function useAttachmentContext(options: {
+  optional: true;
+}): AttachmentContextValue | null;
+export function useAttachmentContext(options?: { optional: true }) {
+  const context = useContext(AttachmentContext);
+  if (!options?.optional && !context)
+    throw new Error(
+      "This component must be used within a ComposerPrimitive.Attachments component.",
+    );
+  return context;
+}

packages/react/src/primitives/composer/ComposerAttachments.tsx

@@ -0,0 +1,85 @@
+"use client";
+
+import { ComponentType, type FC, memo } from "react";
+import { useThreadContext } from "../../context";
+import { useAttachmentContext } from "../../context/react/AttachmentContext";
+import { ComposerAttachmentProvider } from "../../context/providers/ComposerAttachmentProvider";
+import { Attachment } from "../../context/stores/Attachment";
+
+export type ComposerPrimitiveAttachmentsProps = {
+  components:
+    | {
+        Image?: ComponentType | undefined;
+        Document?: ComponentType | undefined;
+        File?: ComponentType | undefined;
+        Fallback?: ComponentType | undefined;
+      }
+    | undefined;
+};
+
+const getComponent = (
+  components: ComposerPrimitiveAttachmentsProps["components"],
+  attachment: Attachment,
+) => {
+  const type = attachment.type;
+  switch (type) {
+    case "image":
+      return components?.Image ?? components?.Fallback;
+    case "document":
+      return components?.Document ?? components?.Fallback;
+    case "file":
+      return components?.File ?? components?.Fallback;
+    default:
+      const _exhaustiveCheck: never = type;
+      throw new Error(`Unknown attachment type: ${_exhaustiveCheck}`);
+  }
+};
+
+const AttachmentComponent: FC<{
+  components: ComposerPrimitiveAttachmentsProps["components"];
+}> = ({ components }) => {
+  const { useAttachment } = useAttachmentContext();
+  const Component = useAttachment((a) =>
+    getComponent(components, a.attachment),
+  );
+
+  if (!Component) return null;
+  return <Component />;
+};
+
+const ComposerAttachmentImpl: FC<
+  ComposerPrimitiveAttachmentsProps & { attachmentIndex: number }
+> = ({ components, attachmentIndex }) => {
+  return (
+    <ComposerAttachmentProvider attachmentIndex={attachmentIndex}>
+      <AttachmentComponent components={components} />
+    </ComposerAttachmentProvider>
+  );
+};
+
+const ComposerAttachment = memo(
+  ComposerAttachmentImpl,
+  (prev, next) =>
+    prev.attachmentIndex === next.attachmentIndex &&
+    prev.components?.Image === next.components?.Image &&
+    prev.components?.Document === next.components?.Document &&
+    prev.components?.File === next.components?.File &&
+    prev.components?.Fallback === next.components?.Fallback,
+);
+
+export const ComposerPrimitiveAttachments: FC<
+  ComposerPrimitiveAttachmentsProps
+> = ({ components }) => {
+  const { useComposer } = useThreadContext();
+  const attachmentsCount = useComposer((s) => s.attachments.length);
+
+  return Array.from({ length: attachmentsCount }, (_, index) => (
+    <ComposerAttachment
+      key={index}
+      attachmentIndex={index}
+      components={components}
+    />
+  ));
+};
+
+ComposerPrimitiveAttachments.displayName = "ComposerPrimitive.Attachments";

packages/react/src/primitives/composer/index.ts

@@ -2,4 +2,5 @@ export { ComposerPrimitiveRoot as Root } from "./ComposerRoot";
 export { ComposerPrimitiveInput as Input } from "./ComposerInput";
 export { ComposerPrimitiveSend as Send } from "./ComposerSend";
 export { ComposerPrimitiveCancel as Cancel } from "./ComposerCancel";
+export { ComposerPrimitiveAttachments as Attachments } from "./ComposerAttachments";
 export { ComposerPrimitiveIf as If } from "./ComposerIf";