refactor: refactored context and injection state

packages/elements/src/chain-of-thought/ChainOfThought.vue

@@ -2,8 +2,7 @@
 import type { HTMLAttributes, Ref } from 'vue'
 import { cn } from '@repo/shadcn-vue/lib/utils'
 import { useVModel } from '@vueuse/core'
-import { provide } from 'vue'
-import { ChainOfThoughtContextKey } from './context'
+import { provideChainOfThought } from './context'
 
 interface ChainOfThoughtProps {
   modelValue?: boolean
@@ -28,7 +27,12 @@ const isOpen = useVModel(props, 'modelValue', emit, {
   passive: true,
 })
 
-provide(ChainOfThoughtContextKey, isOpen as Ref<boolean>)
+provideChainOfThought({
+  isOpen: isOpen as Ref<boolean>,
+  setIsOpen: (open: boolean) => {
+    isOpen.value = open
+  },
+})
 </script>
 
 <template>

packages/elements/src/chain-of-thought/ChainOfThoughtImage.vue

@@ -14,7 +14,7 @@ const props = defineProps<{
     v-bind="$attrs"
   >
     <div
-      class="relative flex max-h-[22rem] items-center justify-center overflow-hidden rounded-lg bg-muted p-3"
+      class="relative flex max-h-88 items-center justify-center overflow-hidden rounded-lg bg-muted p-3"
     >
       <slot />
     </div>

packages/elements/src/chain-of-thought/context.ts

@@ -1,22 +1,9 @@
-import type { InjectionKey, Ref } from 'vue'
-import { inject } from 'vue'
+import type { Ref } from 'vue'
+import { createContext } from 'reka-ui'
 
-export const ChainOfThoughtContextKey: InjectionKey<Ref<boolean>> = Symbol(
-  'ChainOfThoughtContext',
-)
-
-export function useChainOfThought() {
-  const isOpen = inject(ChainOfThoughtContextKey)
-
-  if (!isOpen) {
-    throw new Error(
-      'useChainOfThought must be used within a <ChainOfThought> component',
-    )
-  }
-
-  const setIsOpen = (open: boolean) => {
-    isOpen.value = open
-  }
-
-  return { isOpen, setIsOpen }
+export interface ChainOfThoughtContextValue {
+  isOpen: Ref<boolean>
+  setIsOpen: (open: boolean) => void
 }
+
+export const [useChainOfThought, provideChainOfThought] = createContext<ChainOfThoughtContextValue>('ChainOfThought')

packages/elements/src/code-block/CodeBlock.vue

@@ -3,8 +3,8 @@ import type { BundledLanguage } from 'shiki'
 import type { HTMLAttributes } from 'vue'
 import { cn } from '@repo/shadcn-vue/lib/utils'
 import { reactiveOmit, useDebounceFn } from '@vueuse/core'
-import { computed, onBeforeUnmount, provide, ref, watch } from 'vue'
-import { CodeBlockKey } from './context'
+import { computed, onBeforeUnmount, ref, watch } from 'vue'
+import { provideCodeBlock } from './context'
 import { highlightCode } from './utils'
 
 const props = withDefaults(
@@ -24,10 +24,6 @@ const delegatedProps = reactiveOmit(props, 'code', 'language', 'showLineNumbers'
 const html = ref('')
 const darkHtml = ref('')
 
-provide(CodeBlockKey, {
-  code: computed(() => props.code),
-})
-
 let requestId = 0
 let isUnmounted = false
 
@@ -48,6 +44,10 @@ const updateHighlight = useDebounceFn(async (code: string, language: BundledLang
   }
 }, 100)
 
+provideCodeBlock({
+  code: computed(() => props.code),
+})
+
 watch(
   () => [props.code, props.language, props.showLineNumbers] as const,
   ([code, language, showLineNumbers]) => {

packages/elements/src/code-block/CodeBlockCopyButton.vue

@@ -4,8 +4,8 @@ import { Button } from '@repo/shadcn-vue/components/ui/button'
 import { cn } from '@repo/shadcn-vue/lib/utils'
 import { reactiveOmit } from '@vueuse/core'
 import { CheckIcon, CopyIcon } from 'lucide-vue-next'
-import { computed, inject, onBeforeUnmount, ref } from 'vue'
-import { CodeBlockKey } from './context'
+import { computed, onBeforeUnmount, ref } from 'vue'
+import { useCodeBlock } from './context'
 
 const props = withDefaults(
   defineProps<{
@@ -24,12 +24,7 @@ const emit = defineEmits<{
 
 const delegatedProps = reactiveOmit(props, 'timeout', 'class')
 
-const context = inject(CodeBlockKey)
-
-if (!context)
-  throw new Error('CodeBlockCopyButton must be used within a <CodeBlock />')
-
-const { code } = context
+const { code } = useCodeBlock()
 
 const isCopied = ref(false)
 let resetTimer: ReturnType<typeof setTimeout> | undefined

packages/elements/src/code-block/context.ts

@@ -1,7 +1,8 @@
-import type { ComputedRef, InjectionKey } from 'vue'
+import type { ComputedRef } from 'vue'
+import { createContext } from 'reka-ui'
 
-export interface CodeBlockContext {
+export interface CodeBlockContextValue {
   code: ComputedRef<string>
 }
 
-export const CodeBlockKey: InjectionKey<CodeBlockContext> = Symbol('CodeBlock')
+export const [useCodeBlock, provideCodeBlock] = createContext<CodeBlockContextValue>('CodeBlock')

packages/elements/src/confirmation/Confirmation.vue

@@ -4,16 +4,16 @@ import type { HTMLAttributes } from 'vue'
 import type { ToolUIPartApproval } from './context'
 import { Alert } from '@repo/shadcn-vue/components/ui/alert'
 import { cn } from '@repo/shadcn-vue/lib/utils'
-import { provide, toRef } from 'vue'
-import { ConfirmationKey } from './context'
+import { toRef } from 'vue'
+import { provideConfirmation } from './context'
 
 const props = defineProps<{
   approval?: ToolUIPartApproval
   state: ToolUIPart['state']
   class?: HTMLAttributes['class']
 }>()
 
-provide(ConfirmationKey, {
+provideConfirmation({
   approval: toRef(props, 'approval'),
   state: toRef(props, 'state'),
 })

packages/elements/src/confirmation/ConfirmationAccepted.vue

@@ -1,7 +1,7 @@
 <script setup lang="ts">
-import { useConfirmationContext } from './context'
+import { useConfirmation } from './context'
 
-const { approval, state } = useConfirmationContext()
+const { approval, state } = useConfirmation()
 </script>
 
 <template>

packages/elements/src/confirmation/ConfirmationActions.vue

@@ -1,13 +1,13 @@
 <script setup lang="ts">
 import type { HTMLAttributes } from 'vue'
 import { cn } from '@repo/shadcn-vue/lib/utils'
-import { useConfirmationContext } from './context'
+import { useConfirmation } from './context'
 
 const props = defineProps<{
   class?: HTMLAttributes['class']
 }>()
 
-const { state } = useConfirmationContext()
+const { state } = useConfirmation()
 </script>
 
 <template>

packages/elements/src/confirmation/ConfirmationRejected.vue

@@ -1,7 +1,7 @@
 <script setup lang="ts">
-import { useConfirmationContext } from './context'
+import { useConfirmation } from './context'
 
-const { approval, state } = useConfirmationContext()
+const { approval, state } = useConfirmation()
 </script>
 
 <template>

packages/elements/src/confirmation/ConfirmationRequest.vue

@@ -1,7 +1,7 @@
 <script setup lang="ts">
-import { useConfirmationContext } from './context'
+import { useConfirmation } from './context'
 
-const { state } = useConfirmationContext()
+const { state } = useConfirmation()
 </script>
 
 <template>

packages/elements/src/confirmation/context.ts

@@ -1,6 +1,6 @@
 import type { ToolUIPart } from 'ai'
-import type { InjectionKey, Ref } from 'vue'
-import { inject } from 'vue'
+import type { Ref } from 'vue'
+import { createContext } from 'reka-ui'
 
 export type ToolUIPartApproval
   = | {
@@ -35,12 +35,4 @@ export interface ConfirmationContextValue {
   state: Ref<ToolUIPart['state']>
 }
 
-export const ConfirmationKey: InjectionKey<ConfirmationContextValue>
-  = Symbol('ConfirmationContext')
-
-export function useConfirmationContext() {
-  const context = inject<ConfirmationContextValue | null>(ConfirmationKey, null)
-  if (!context)
-    throw new Error('Confirmation components must be used within <Confirmation>')
-  return context
-}
+export const [useConfirmation, provideConfirmation] = createContext<ConfirmationContextValue>('Confirmation')

packages/elements/src/context/Context.vue

@@ -2,8 +2,8 @@
 import type { LanguageModelUsage } from 'ai'
 import type { ModelId } from './context'
 import { HoverCard } from '@repo/shadcn-vue/components/ui/hover-card'
-import { computed, provide } from 'vue'
-import { ContextKey } from './context'
+import { computed } from 'vue'
+import { provideContext } from './context'
 
 interface Props {
   usedTokens: number
@@ -14,7 +14,7 @@ interface Props {
 
 const props = defineProps<Props>()
 
-provide(ContextKey, {
+provideContext({
   usedTokens: computed(() => props.usedTokens),
   maxTokens: computed(() => props.maxTokens),
   usage: computed(() => props.usage),

packages/elements/src/context/context.ts

@@ -1,6 +1,6 @@
 import type { LanguageModelUsage } from 'ai'
-import type { ComputedRef, InjectionKey } from 'vue'
-import { inject } from 'vue'
+import type { ComputedRef } from 'vue'
+import { createContext } from 'reka-ui'
 
 export type ModelId = string
 
@@ -11,13 +11,4 @@ export interface ContextContextValue {
   modelId: ComputedRef<ModelId | undefined>
 }
 
-export const ContextKey: InjectionKey<ContextContextValue>
-  = Symbol('ContextContext')
-
-export function useContextValue(): ContextContextValue {
-  const context = inject<ContextContextValue>(ContextKey)
-  if (!context) {
-    throw new Error('Context components must be used within Context')
-  }
-  return context
-}
+export const [useContextValue, provideContext] = createContext<ContextContextValue>('Context')

packages/elements/src/message/MessageBranch.vue

@@ -1,9 +1,8 @@
 <script setup lang="ts">
 import type { HTMLAttributes, VNode } from 'vue'
-import type { MessageBranchContextType } from './context'
 import { cn } from '@repo/shadcn-vue/lib/utils'
-import { provide, readonly, ref } from 'vue'
-import { MessageBranchKey } from './context'
+import { readonly, ref } from 'vue'
+import { provideMessageBranch } from './context'
 
 interface Props {
   defaultBranch?: number
@@ -44,16 +43,14 @@ function setBranches(count: number) {
   totalBranches.value = count
 }
 
-const contextValue: MessageBranchContextType = {
+provideMessageBranch({
   currentBranch: readonly(currentBranch),
   totalBranches: readonly(totalBranches),
   goToPrevious,
   goToNext,
   branches,
   setBranches,
-}
-
-provide(MessageBranchKey, contextValue)
+})
 </script>
 
 <template>

packages/elements/src/message/MessageBranchContent.vue

@@ -2,7 +2,7 @@
 import type { HTMLAttributes } from 'vue'
 import { cn } from '@repo/shadcn-vue/lib/utils'
 import { computed, Fragment, isVNode, onMounted, useSlots, watch } from 'vue'
-import { useMessageBranchContext } from './context'
+import { useMessageBranch } from './context'
 
 interface Props {
   class?: HTMLAttributes['class']
@@ -11,7 +11,7 @@ interface Props {
 const props = defineProps<Props>()
 const slots = useSlots()
 
-const { currentBranch, setBranches } = useMessageBranchContext()
+const { currentBranch, setBranches } = useMessageBranch()
 
 const branchVNodes = computed(() => {
   const nodes = slots.default?.() ?? []

packages/elements/src/message/MessageBranchNext.vue

@@ -1,9 +1,9 @@
 <script setup lang="ts">
 import { Button } from '@repo/shadcn-vue/components/ui/button'
 import { ChevronRightIcon } from 'lucide-vue-next'
-import { useMessageBranchContext } from './context'
+import { useMessageBranch } from './context'
 
-const { goToNext, totalBranches } = useMessageBranchContext()
+const { goToNext, totalBranches } = useMessageBranch()
 </script>
 
 <template>

packages/elements/src/message/MessageBranchPage.vue

@@ -2,15 +2,15 @@
 import type { HTMLAttributes } from 'vue'
 import { ButtonGroupText } from '@repo/shadcn-vue/components/ui/button-group'
 import { cn } from '@repo/shadcn-vue/lib/utils'
-import { useMessageBranchContext } from './context'
+import { useMessageBranch } from './context'
 
 interface Props {
   class?: HTMLAttributes['class']
 }
 
 const props = defineProps<Props>()
 
-const { currentBranch, totalBranches } = useMessageBranchContext()
+const { currentBranch, totalBranches } = useMessageBranch()
 </script>
 
 <template>

packages/elements/src/message/MessageBranchPrevious.vue

@@ -1,9 +1,9 @@
 <script setup lang="ts">
 import { Button } from '@repo/shadcn-vue/components/ui/button'
 import { ChevronLeftIcon } from 'lucide-vue-next'
-import { useMessageBranchContext } from './context'
+import { useMessageBranch } from './context'
 
-const { goToPrevious, totalBranches } = useMessageBranchContext()
+const { goToPrevious, totalBranches } = useMessageBranch()
 </script>
 
 <template>

packages/elements/src/message/MessageBranchSelector.vue

@@ -1,14 +1,14 @@
 <script setup lang="ts">
 import type { UIMessage } from 'ai'
 import { ButtonGroup } from '@repo/shadcn-vue/components/ui/button-group'
-import { useMessageBranchContext } from './context'
+import { useMessageBranch } from './context'
 
 interface Props {
   from: UIMessage['role']
 }
 defineProps<Props>()
 
-const { totalBranches } = useMessageBranchContext()
+const { totalBranches } = useMessageBranch()
 </script>
 
 <template>

packages/elements/src/message/context.ts

@@ -1,5 +1,5 @@
-import type { InjectionKey, Ref, VNode } from 'vue'
-import { inject } from 'vue'
+import type { Ref, VNode } from 'vue'
+import { createContext } from 'reka-ui'
 
 export interface MessageBranchContextType<T = VNode[]> {
   currentBranch: Readonly<Ref<number>>
@@ -10,13 +10,4 @@ export interface MessageBranchContextType<T = VNode[]> {
   setBranches: (count: number) => void
 }
 
-export const MessageBranchKey: InjectionKey<MessageBranchContextType>
-  = Symbol('MessageBranch')
-
-export function useMessageBranchContext(): MessageBranchContextType {
-  const ctx = inject(MessageBranchKey)
-  if (!ctx) {
-    throw new Error('Message Branch components must be used within Message Branch')
-  }
-  return ctx
-}
+export const [useMessageBranch, provideMessageBranch] = createContext<MessageBranchContextType>('MessageBranch')