Hi,

I'm switching to very similar solutions like jotai-ai (ai sdk to unified api call + jotai for client-side state management) from ai/rsc recently.

Thanks for providing excellent reference implementations. But when try to use it in my case, I feel it's not that easy to adapt widely. I have implemented extended types based CoreMessage (which also could be seen as DB message or server state), but both original useChat and current chatAtoms are all difficult to customize client side message state, I almost need to write almost everything. Actually, the latest reference example ai-chatbot has been using CoreMessage for server state, Message (from @ai/ui-utils) for client state (which should be better to name as UIMessage). I hope the new design would achieve that easier.

This PR implements some basic MVP to demonstrate my ideas:

1. Split main user interface to two functions makeChatAtoms and useChat. Let useChat expose useful and similar handlers/states like upstream useChat (ai/react), and left makeChatAtoms contains low level atoms for advanced users. Indeed, makeChatAtoms only controls how client interactive with remote stream data (append and receive new messages), and left all others into useChat in order that advanced users could implement theirs inputAtoms, sumbitAtoms. Built-in useChat could be just an official reference implementation has the same features with upstream projects.

2. Accept a user defined messagesAtom from outside. Like

const messagesAtom = atom<Message[]>([]);

const {
  // state data containers,
  isLoadingAtom,
  errorAtom,
  dataAtom,

  // actions
  stopAtom,
  appendAtom,
  reloadAtom,

  // handlers
  onErrorAtom,
  onResponseAtom,
  onToolCallAtom,
  onFinishAtom,
} = makeChatAtoms({ messagesAtom });

In this way, users could decide how to prepare their UI messages, whatever derived from other states or initialed from fetch or even SSR.

// Initial messages in SSR approach
// so `makeChatAtoms` doesn't need to consider 
// how to deal with initial messages or initial input.
import { useHydrateAtoms } from 'jotai/utils';

const messagesAtom = atom<Message[]>([]);

const Component = ({ initChatMesssages }:{ initChatMesssages: Message[] }) => {
  useHydrateAtoms([[messagesAtom, initChatMesssages]])

  const messages = useAtomValue(messagesAtom)

  return <Chat>{messages}<Chat/>
}

3. It would be easier to introduce more features and keep flexibility by using this architecture. We're able to add more states like firstTokenReceived or action like deleteAtom in the future without breaking up signatures of useChat.

Also, there are several issues I haven't confirmed and figured out clearly:

1. In current implements, all deprecated API and parameters are literally deprecated. And the ai project has already released alpha candidates of the next v4 version.

2. Current callback handler like onErrorAtom, ..., onResponseAtom is initialed by makeChatAtoms and cloud be updated by useChat. I'm not sure whether it's a correct design or proper usage of atoms. For now, it's just implemented to fit upstream UI test cases easily.

3. I hope we could add a generic type interface for Body into makeChatAtoms. This data type can be shared between server endpoint and user client. It's more safe and ergonomic to make sure correct data being sent and parsed.

That's all. Thanks for your efforts!

Reject or feedback is all appreciate :). Have a nice day 🎉.

Sincerely,
Lanqing