@lota-sdk/ui 0.3.1 → 0.3.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/ui",
-  "version": "0.3.1",
+  "version": "0.3.2",
   "type": "module",
   "main": "./src/index.ts",
   "types": "./src/index.ts",
@@ -25,7 +25,7 @@
     "registry": "https://registry.npmjs.org/"
   },
   "dependencies": {
-    "@lota-sdk/shared": "0.3.1",
+    "@lota-sdk/shared": "0.3.2",
     "ai": "^6.0.145"
   },
   "peerDependencies": {

package/src/runtime/use-snapshotting-chat.ts

@@ -0,0 +1,249 @@
+import type { UseChatHelpers } from '@ai-sdk/react'
+import { AbstractChat } from 'ai'
+import type { ChatInit, ChatState, ChatStatus, UIMessage } from 'ai'
+import { useCallback, useEffect, useRef, useSyncExternalStore } from 'react'
+
+function throttleCallback(callback: () => void, waitMs?: number): () => void {
+  if (waitMs == null) {
+    return callback
+  }
+
+  let timeoutId: ReturnType<typeof setTimeout> | null = null
+  let hasPendingCall = false
+
+  return () => {
+    if (timeoutId !== null) {
+      hasPendingCall = true
+      return
+    }
+
+    callback()
+    timeoutId = setTimeout(() => {
+      timeoutId = null
+
+      if (!hasPendingCall) {
+        return
+      }
+
+      hasPendingCall = false
+      callback()
+    }, waitMs)
+  }
+}
+
+class SnapshottingChatState<UI_MESSAGE extends UIMessage> implements ChatState<UI_MESSAGE> {
+  #messages: UI_MESSAGE[]
+  #status: ChatStatus = 'ready'
+  #error: Error | undefined = undefined
+
+  #messagesCallbacks = new Set<() => void>()
+  #statusCallbacks = new Set<() => void>()
+  #errorCallbacks = new Set<() => void>()
+
+  constructor(initialMessages: UI_MESSAGE[] = []) {
+    this.#messages = initialMessages
+  }
+
+  get status(): ChatStatus {
+    return this.#status
+  }
+
+  set status(newStatus: ChatStatus) {
+    this.#status = newStatus
+    this.#callStatusCallbacks()
+  }
+
+  get error(): Error | undefined {
+    return this.#error
+  }
+
+  set error(newError: Error | undefined) {
+    this.#error = newError
+    this.#callErrorCallbacks()
+  }
+
+  get messages(): UI_MESSAGE[] {
+    return this.#messages
+  }
+
+  set messages(newMessages: UI_MESSAGE[]) {
+    this.#messages = [...newMessages]
+    this.#callMessagesCallbacks()
+  }
+
+  pushMessage = (message: UI_MESSAGE) => {
+    const lastMessage = this.#messages.at(-1)
+
+    if (lastMessage?.role === 'assistant' && message.role === 'assistant' && lastMessage.id !== message.id) {
+      // The AI SDK reuses one mutable assistant message object across consecutive `start` events
+      // in a single stream. When a second assistant starts, clear inherited parts so the new
+      // message begins empty instead of briefly rendering the previous assistant's content.
+      message.parts = [] as UI_MESSAGE['parts']
+    }
+
+    // Snapshot on insert so later stream mutations cannot rewrite already-rendered messages.
+    this.#messages = this.#messages.concat(this.snapshot(message))
+    this.#callMessagesCallbacks()
+  }
+
+  popMessage = () => {
+    this.#messages = this.#messages.slice(0, -1)
+    this.#callMessagesCallbacks()
+  }
+
+  replaceMessage = (index: number, message: UI_MESSAGE) => {
+    this.#messages = [...this.#messages.slice(0, index), this.snapshot(message), ...this.#messages.slice(index + 1)]
+    this.#callMessagesCallbacks()
+  }
+
+  snapshot = <T>(value: T): T => structuredClone(value)
+
+  '~registerMessagesCallback' = (onChange: () => void, throttleWaitMs?: number): (() => void) => {
+    const callback = throttleCallback(onChange, throttleWaitMs)
+    this.#messagesCallbacks.add(callback)
+    return () => {
+      this.#messagesCallbacks.delete(callback)
+    }
+  }
+
+  '~registerStatusCallback' = (onChange: () => void): (() => void) => {
+    this.#statusCallbacks.add(onChange)
+    return () => {
+      this.#statusCallbacks.delete(onChange)
+    }
+  }
+
+  '~registerErrorCallback' = (onChange: () => void): (() => void) => {
+    this.#errorCallbacks.add(onChange)
+    return () => {
+      this.#errorCallbacks.delete(onChange)
+    }
+  }
+
+  #callMessagesCallbacks = () => {
+    this.#messagesCallbacks.forEach((callback) => callback())
+  }
+
+  #callStatusCallbacks = () => {
+    this.#statusCallbacks.forEach((callback) => callback())
+  }
+
+  #callErrorCallbacks = () => {
+    this.#errorCallbacks.forEach((callback) => callback())
+  }
+}
+
+class SnapshottingChat<UI_MESSAGE extends UIMessage> extends AbstractChat<UI_MESSAGE> {
+  #state: SnapshottingChatState<UI_MESSAGE>
+
+  constructor({ messages, ...init }: ChatInit<UI_MESSAGE>) {
+    const state = new SnapshottingChatState(messages)
+    super({ ...init, state })
+    this.#state = state
+  }
+
+  '~registerMessagesCallback' = (onChange: () => void, throttleWaitMs?: number): (() => void) =>
+    this.#state['~registerMessagesCallback'](onChange, throttleWaitMs)
+
+  '~registerStatusCallback' = (onChange: () => void): (() => void) => this.#state['~registerStatusCallback'](onChange)
+
+  '~registerErrorCallback' = (onChange: () => void): (() => void) => this.#state['~registerErrorCallback'](onChange)
+}
+
+type UseSnapshottingChatOptions<UI_MESSAGE extends UIMessage> = ChatInit<UI_MESSAGE> & {
+  experimental_throttle?: number
+  resume?: boolean
+}
+
+export function useSnapshottingChat<UI_MESSAGE extends UIMessage = UIMessage>({
+  experimental_throttle: throttleWaitMs,
+  resume = false,
+  ...options
+}: UseSnapshottingChatOptions<UI_MESSAGE> = {}): UseChatHelpers<UI_MESSAGE> {
+  const callbacksRef = useRef({
+    onToolCall: options.onToolCall,
+    onData: options.onData,
+    onFinish: options.onFinish,
+    onError: options.onError,
+    sendAutomaticallyWhen: options.sendAutomaticallyWhen,
+  })
+
+  callbacksRef.current = {
+    onToolCall: options.onToolCall,
+    onData: options.onData,
+    onFinish: options.onFinish,
+    onError: options.onError,
+    sendAutomaticallyWhen: options.sendAutomaticallyWhen,
+  }
+
+  const optionsWithCallbacks: ChatInit<UI_MESSAGE> = {
+    ...options,
+    onToolCall: (arg) => callbacksRef.current.onToolCall?.(arg),
+    onData: (arg) => callbacksRef.current.onData?.(arg),
+    onFinish: (arg) => callbacksRef.current.onFinish?.(arg),
+    onError: (arg) => callbacksRef.current.onError?.(arg),
+    sendAutomaticallyWhen: (arg) => callbacksRef.current.sendAutomaticallyWhen?.(arg) ?? false,
+  }
+
+  const chatRef = useRef<SnapshottingChat<UI_MESSAGE>>(new SnapshottingChat(optionsWithCallbacks))
+
+  if (chatRef.current.id !== (options.id ?? chatRef.current.id)) {
+    chatRef.current = new SnapshottingChat(optionsWithCallbacks)
+  }
+
+  const subscribeToMessages = useCallback(
+    (update: () => void) => chatRef.current['~registerMessagesCallback'](update, throttleWaitMs),
+    [throttleWaitMs, chatRef.current.id],
+  )
+
+  const messages = useSyncExternalStore(
+    subscribeToMessages,
+    () => chatRef.current.messages,
+    () => chatRef.current.messages,
+  )
+
+  const status = useSyncExternalStore(
+    chatRef.current['~registerStatusCallback'],
+    () => chatRef.current.status,
+    () => chatRef.current.status,
+  )
+
+  const error = useSyncExternalStore(
+    chatRef.current['~registerErrorCallback'],
+    () => chatRef.current.error,
+    () => chatRef.current.error,
+  )
+
+  const setMessages = useCallback(
+    (messagesParam: UI_MESSAGE[] | ((messages: UI_MESSAGE[]) => UI_MESSAGE[])) => {
+      if (typeof messagesParam === 'function') {
+        messagesParam = messagesParam(chatRef.current.messages)
+      }
+
+      chatRef.current.messages = messagesParam
+    },
+    [],
+  )
+
+  useEffect(() => {
+    if (resume) {
+      chatRef.current.resumeStream()
+    }
+  }, [resume])
+
+  return {
+    id: chatRef.current.id,
+    messages,
+    setMessages,
+    sendMessage: chatRef.current.sendMessage,
+    regenerate: chatRef.current.regenerate,
+    clearError: chatRef.current.clearError,
+    stop: chatRef.current.stop,
+    error,
+    resumeStream: chatRef.current.resumeStream,
+    status,
+    addToolResult: chatRef.current.addToolOutput,
+    addToolOutput: chatRef.current.addToolOutput,
+    addToolApprovalResponse: chatRef.current.addToolApprovalResponse,
+  }
+}

package/src/runtime/use-thread-chat.ts

@@ -1,9 +1,10 @@
-import { useChat } from '@ai-sdk/react'
 import type { UseChatHelpers } from '@ai-sdk/react'
 import { lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai'
 import type { ChatInit, UIMessage } from 'ai'
 import { useCallback, useEffect, useRef } from 'react'
 
+import { useSnapshottingChat } from './use-snapshotting-chat'
+
 type ChatTransport<UI_MESSAGE extends UIMessage> = NonNullable<ChatInit<UI_MESSAGE>['transport']>
 
 export interface UseThreadChatOptions<UI_MESSAGE extends UIMessage> {
@@ -87,7 +88,7 @@ export function useThreadChat<UI_MESSAGE extends UIMessage>({
     setMessages,
     addToolApprovalResponse,
     resumeStream,
-  } = useChat<UI_MESSAGE>({
+  } = useSnapshottingChat<UI_MESSAGE>({
     id: chatId,
     transport,
     messages: initialMessages,