@lota-sdk/ui 0.3.0 → 0.3.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/ui",
-  "version": "0.3.0",
+  "version": "0.3.2",
   "type": "module",
   "main": "./src/index.ts",
   "types": "./src/index.ts",
@@ -18,14 +18,14 @@
     "lint": "bunx oxlint --fix -c ../oxlint.config.ts src",
     "format": "bunx oxfmt src",
     "typecheck": "bunx tsgo --noEmit",
-    "test:unit": "bun test ../tests/unit/ui"
+    "test:unit": "bun run scripts/test-unit.ts"
   },
   "publishConfig": {
     "access": "public",
     "registry": "https://registry.npmjs.org/"
   },
   "dependencies": {
-    "@lota-sdk/shared": "0.3.0",
+    "@lota-sdk/shared": "0.3.2",
     "ai": "^6.0.145"
   },
   "peerDependencies": {

package/src/chat/attachments.ts

@@ -14,7 +14,15 @@ type UploadedAttachmentPayload = {
   url: string
 }
 
+function hasLocalFile(value: FileUIPart): value is FileUIPart & { file: File } {
+  return 'file' in value && value.file instanceof File
+}
+
 async function toUploadFile(filePart: FileUIPart): Promise<File> {
+  if (hasLocalFile(filePart)) {
+    return filePart.file
+  }
+
   if (!filePart.url) {
     throw new Error('Attachment payload is missing file data.')
   }

package/src/chat/index.ts

@@ -1,4 +1,3 @@
 export * from './attachments'
 export * from './message-parts'
 export * from './messages'
-export * from './transport'

package/src/runtime/active-thread-helpers.ts

@@ -1,8 +1,8 @@
-import type { FileUIPart, UIMessage } from 'ai'
+import type { UIMessage } from 'ai'
 
-export type ComposerMessage<TFile extends FileUIPart = FileUIPart> = { text: string; files?: TFile[] }
+export type ComposerMessage<TFile = File> = { text: string; files?: TFile[] }
 
-export function validateComposerMessage<TFile extends FileUIPart = FileUIPart>(
+export function validateComposerMessage<TFile = File>(
   message: ComposerMessage<TFile>,
 ): { textContent: string; files: TFile[] } | null {
   const textContent = message.text.trim()

package/src/runtime/index.ts

@@ -1,7 +1,6 @@
 export * from './active-thread-helpers'
 export * from './session-common'
 export * from './use-active-thread-session'
-export * from './use-lota-chat'
 export * from './use-session-messages'
 export * from './use-thread-chat'
 export * from './use-thread-management'

package/src/runtime/use-active-thread-session.ts

@@ -1,5 +1,5 @@
 import { getMessageCreatedAt } from '@lota-sdk/shared'
-import type { FileUIPart, UIMessage } from 'ai'
+import type { UIMessage } from 'ai'
 import { useCallback, useEffect, useMemo, useRef } from 'react'
 
 import { mergeChatMessages } from '../chat/messages'
@@ -21,7 +21,7 @@ interface ActiveThreadLike {
 export interface UseActiveThreadSessionOptions<
   TMessage extends UIMessage & { metadata?: { createdAt?: unknown } },
   TThread extends ActiveThreadLike,
-  TFile extends FileUIPart = FileUIPart,
+  TFile = File,
   TTrackerState = never,
 > {
   threadId: string | null
@@ -49,7 +49,7 @@ export interface UseActiveThreadSessionOptions<
   isThreadBusy?: boolean
 }
 
-export interface UseActiveThreadSessionReturn<TMessage extends UIMessage, TFile extends FileUIPart = FileUIPart> {
+export interface UseActiveThreadSessionReturn<TMessage extends UIMessage, TFile = File> {
   messages: TMessage[]
   status: ChatStatus
   error: Error | undefined
@@ -67,7 +67,7 @@ export interface UseActiveThreadSessionReturn<TMessage extends UIMessage, TFile
 export function useActiveThreadSession<
   TMessage extends UIMessage & { metadata?: { createdAt?: unknown } },
   TThread extends ActiveThreadLike,
-  TFile extends FileUIPart = FileUIPart,
+  TFile = File,
   TTrackerState = never,
 >({
   threadId,

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> {
@@ -58,6 +59,9 @@ export function useThreadChat<UI_MESSAGE extends UIMessage>({
   const onDataRef = useRef(onData)
   const onMessagesSettledRef = useRef(onMessagesSettled)
   const setThreadRunningStateRef = useRef(setThreadRunningState)
+  const attemptedResumeForChatRef = useRef<string | null>(null)
+  const skipNextResumeForChatRef = useRef<string | null>(null)
+  const locallyStartedChatRef = useRef<string | null>(null)
 
   useEffect(() => {
     onDataRef.current = onData
@@ -78,13 +82,13 @@ export function useThreadChat<UI_MESSAGE extends UIMessage>({
     status,
     error,
     clearError,
-    sendMessage,
-    regenerate,
+    sendMessage: sendChatMessage,
+    regenerate: regenerateChat,
     stop: stopChatStream,
     setMessages,
     addToolApprovalResponse,
     resumeStream,
-  } = useChat<UI_MESSAGE>({
+  } = useSnapshottingChat<UI_MESSAGE>({
     id: chatId,
     transport,
     messages: initialMessages,
@@ -98,10 +102,17 @@ export function useThreadChat<UI_MESSAGE extends UIMessage>({
     experimental_throttle: experimentalThrottle,
     sendAutomaticallyWhen,
     onFinish: () => {
+      if (locallyStartedChatRef.current === chatId) {
+        skipNextResumeForChatRef.current = chatId
+        locallyStartedChatRef.current = null
+      }
       setThreadRunningStateRef.current?.(false)
       onMessagesSettledRef.current?.()
     },
     onError: (error) => {
+      if (locallyStartedChatRef.current === chatId) {
+        locallyStartedChatRef.current = null
+      }
       if (!(error instanceof Error && error.name === 'AbortError')) {
         notifyError?.('Chat stream failed.', toError(error))
       }
@@ -113,13 +124,54 @@ export function useThreadChat<UI_MESSAGE extends UIMessage>({
   // Manual resume with ref guard — prevents React StrictMode double-fire.
   // We do NOT pass `resume` to useChat because the AI SDK's built-in effect
   // has no StrictMode guard, causing two concurrent GET streams.
-  const resumedForChatRef = useRef<string | null>(null)
   useEffect(() => {
-    if (!resume || status !== 'ready' || resumedForChatRef.current === chatId) return
-    resumedForChatRef.current = chatId
+    if (!resume || status !== 'ready') return
+    if (skipNextResumeForChatRef.current === chatId) {
+      skipNextResumeForChatRef.current = null
+      return
+    }
+    if (attemptedResumeForChatRef.current === chatId) return
+    attemptedResumeForChatRef.current = chatId
     void resumeStream()
   }, [resume, status, chatId, resumeStream])
 
+  useEffect(() => {
+    if (status === 'ready' && resume) return
+    if (attemptedResumeForChatRef.current === chatId) {
+      attemptedResumeForChatRef.current = null
+    }
+  }, [chatId, resume, status])
+
+  const sendMessage = useCallback<UseChatHelpers<UI_MESSAGE>['sendMessage']>(
+    async (message, options) => {
+      locallyStartedChatRef.current = chatId
+      try {
+        return await sendChatMessage(message, options)
+      } catch (error) {
+        if (locallyStartedChatRef.current === chatId) {
+          locallyStartedChatRef.current = null
+        }
+        throw error
+      }
+    },
+    [chatId, sendChatMessage],
+  )
+
+  const regenerate = useCallback<UseChatHelpers<UI_MESSAGE>['regenerate']>(
+    async (options) => {
+      locallyStartedChatRef.current = chatId
+      try {
+        return await regenerateChat(options)
+      } catch (error) {
+        if (locallyStartedChatRef.current === chatId) {
+          locallyStartedChatRef.current = null
+        }
+        throw error
+      }
+    },
+    [chatId, regenerateChat],
+  )
+
   const stop = useCallback(async () => {
     const stopRequest = threadId && stopRemoteRun ? stopRemoteRun(threadId) : Promise.resolve(null)
     const [stopResponse] = await Promise.allSettled([stopRequest, stopChatStream()])

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

@@ -5,7 +5,7 @@ const PLACEHOLDER_THREAD_PREFIX = '__placeholder_'
 export interface ThreadManagementItem<TAgentId extends string = string> {
   threadId: string
   title: string
-  status: 'regular' | 'archived'
+  status: 'active' | 'archived'
   type: string
   isMutable: boolean
   agentId?: TAgentId | null
@@ -62,7 +62,7 @@ export function findReusableRegularGroupThreadId<TThread extends ThreadManagemen
   defaultThreadTitle: string,
 ): string | null {
   const draft = threads.find(
-    (thread) => thread.type === 'group' && thread.status === 'regular' && thread.title.trim() === defaultThreadTitle,
+    (thread) => thread.type === 'group' && thread.status === 'active' && thread.title.trim() === defaultThreadTitle,
   )
 
   return draft?.threadId ?? null
@@ -245,6 +245,7 @@ export function useThreadManagement<TThread extends ThreadManagementItem<TAgentI
       return
     }
 
+    if (!enabled) return
     if (isCreatingThreadRef.current) return
     if (isCreatingThread) return
     isCreatingThreadRef.current = true
@@ -257,6 +258,7 @@ export function useThreadManagement<TThread extends ThreadManagementItem<TAgentI
         isCreatingThreadRef.current = false
       })
   }, [
+    enabled,
     currentThreadId,
     handleCreateNewThread,
     isCreatingThread,

package/src/chat/transport.ts

@@ -1,8 +0,0 @@
-import { DefaultChatTransport } from 'ai'
-import type { HttpChatTransportInitOptions, UIMessage } from 'ai'
-
-export type ChatTransportOptions<UI_MESSAGE extends UIMessage> = HttpChatTransportInitOptions<UI_MESSAGE>
-
-export function createChatTransport<UI_MESSAGE extends UIMessage>(options: ChatTransportOptions<UI_MESSAGE>) {
-  return new DefaultChatTransport<UI_MESSAGE>(options)
-}

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

@@ -1,13 +0,0 @@
-import type { UIMessage } from 'ai'
-
-import type { UseThreadChatOptions, UseThreadChatReturn } from './use-thread-chat'
-import { useThreadChat } from './use-thread-chat'
-
-export type UseLotaChatOptions<UI_MESSAGE extends UIMessage> = UseThreadChatOptions<UI_MESSAGE>
-export type UseLotaChatReturn<UI_MESSAGE extends UIMessage> = UseThreadChatReturn<UI_MESSAGE>
-
-export function useLotaChat<UI_MESSAGE extends UIMessage>(
-  options: UseLotaChatOptions<UI_MESSAGE>,
-): UseLotaChatReturn<UI_MESSAGE> {
-  return useThreadChat(options)
-}