feat: useXContext optional mode (#438)

Yonom · Jul 9, 2024 · bc6b5d4 · bc6b5d4
1 parent c7ba6a2
commit bc6b5d4
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 15 deletions.
diff --git a/packages/react/src/context/react/AssistantContext.ts b/packages/react/src/context/react/AssistantContext.ts
@@ -16,11 +16,15 @@ export const AssistantContext = createContext<AssistantContextValue | null>(
   null,
 );
 
-export const useAssistantContext = (): AssistantContextValue => {
+export function useAssistantContext(): AssistantContextValue;
+export function useAssistantContext(options: {
+  optional: true;
+}): AssistantContextValue | null;
+export function useAssistantContext(options?: { optional: true }) {
   const context = useContext(AssistantContext);
-  if (!context)
+  if (!options?.optional && !context)
     throw new Error(
       "This component must be used within an AssistantRuntimeProvider.",
     );
   return context;
-};
+}
diff --git a/packages/react/src/context/react/ComposerContext.ts b/packages/react/src/context/react/ComposerContext.ts
@@ -1,5 +1,5 @@
-import { useContext, useMemo } from "react";
-import { MessageContext } from "./MessageContext";
+import { useMemo } from "react";
+import { useMessageContext } from "./MessageContext";
 import { useThreadContext } from "./ThreadContext";
 import type { ComposerState } from "../stores/Composer";
 import type { EditComposerState } from "../stores/EditComposer";
@@ -12,7 +12,7 @@ export type ComposerContextValue = {
 
 export const useComposerContext = (): ComposerContextValue => {
   const { useComposer } = useThreadContext();
-  const { useEditComposer } = useContext(MessageContext) ?? {};
+  const { useEditComposer } = useMessageContext({ optional: true }) ?? {};
   return useMemo(
     () => ({
       useComposer: (useEditComposer ?? useComposer) as ReadonlyStore<

diff --git a/packages/react/src/context/react/ContentPartContext.ts b/packages/react/src/context/react/ContentPartContext.ts
@@ -12,11 +12,15 @@ export const ContentPartContext = createContext<ContentPartContextValue | null>(
   null,
 );
 
-export const useContentPartContext = (): ContentPartContextValue => {
+export function useContentPartContext(): ContentPartContextValue;
+export function useContentPartContext(options: {
+  optional: true;
+}): ContentPartContextValue | null;
+export function useContentPartContext(options?: { optional: true }) {
   const context = useContext(ContentPartContext);
-  if (!context)
+  if (!options?.optional && !context)
     throw new Error(
       "This component can only be used inside a component passed to <MessagePrimitive.Content components={...} >.",
     );
   return context;
-};
+}
diff --git a/packages/react/src/context/react/MessageContext.ts b/packages/react/src/context/react/MessageContext.ts
@@ -14,11 +14,15 @@ export type MessageContextValue = {
 
 export const MessageContext = createContext<MessageContextValue | null>(null);
 
-export const useMessageContext = () => {
+export function useMessageContext(): MessageContextValue;
+export function useMessageContext(options: {
+  optional: true;
+}): MessageContextValue | null;
+export function useMessageContext(options?: { optional: true }) {
   const context = useContext(MessageContext);
-  if (!context)
+  if (!options?.optional && !context)
     throw new Error(
       "This component can only be used inside a component passed to <ThreadPrimitive.Messages components={...} />.",
     );
   return context;
-};
+}
diff --git a/packages/react/src/context/react/ThreadContext.ts b/packages/react/src/context/react/ThreadContext.ts
@@ -18,11 +18,15 @@ export type ThreadContextValue = {
 
 export const ThreadContext = createContext<ThreadContextValue | null>(null);
 
-export const useThreadContext = (): ThreadContextValue => {
+export function useThreadContext(): ThreadContextValue;
+export function useThreadContext(options: {
+  optional: true;
+}): ThreadContextValue | null;
+export function useThreadContext(options?: { optional: true }) {
   const context = useContext(ThreadContext);
-  if (!context)
+  if (!options?.optional && !context)
     throw new Error(
       "This component must be used within an AssistantRuntimeProvider.",
     );
   return context;
-};
+}