@assistant-ui/react 0.0.8 → 0.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.mjs CHANGED
@@ -130,12 +130,14 @@ var ThreadViewport = forwardRef2(({ autoScroll = true, onScroll, children, ...re
130
130
  const ref = useComposedRefs(forwardedRef, divRef);
131
131
  const { useThread } = useAssistantContext();
132
132
  const firstRenderRef = useRef3(true);
133
+ const lastScrollTop = useRef3(0);
133
134
  const scrollToBottom = () => {
134
135
  const div = messagesEndRef.current;
135
136
  if (!div || !autoScroll)
136
137
  return;
137
138
  const behavior = firstRenderRef.current ? "instant" : "auto";
138
139
  firstRenderRef.current = false;
140
+ useThread.setState({ isAtBottom: true });
139
141
  div.scrollIntoView({ behavior });
140
142
  };
141
143
  useOnResizeContent(divRef, () => {
@@ -151,10 +153,12 @@ var ThreadViewport = forwardRef2(({ autoScroll = true, onScroll, children, ...re
151
153
  if (!div)
152
154
  return;
153
155
  const isAtBottom = useThread.getState().isAtBottom;
154
- const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight + 50;
155
- if (newIsAtBottom !== isAtBottom) {
156
+ const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight;
157
+ if (!newIsAtBottom && lastScrollTop.current < div.scrollTop) {
158
+ } else if (newIsAtBottom !== isAtBottom) {
156
159
  useThread.setState({ isAtBottom: newIsAtBottom });
157
160
  }
161
+ lastScrollTop.current = div.scrollTop;
158
162
  };
159
163
  return /* @__PURE__ */ React.createElement(
160
164
  Primitive2.div,
@@ -168,15 +172,19 @@ var ThreadViewport = forwardRef2(({ autoScroll = true, onScroll, children, ...re
168
172
  );
169
173
  });
170
174
 
171
- // src/vercel/useVercelAIBranches.tsx
175
+ // src/adapters/vercel/useVercelAIBranches.tsx
172
176
  import { useCallback, useMemo, useRef as useRef4 } from "react";
173
- var ROOT_ID = "__ROOT_ID__";
177
+
178
+ // src/utils/context/stores/AssistantTypes.ts
179
+ var ROOT_PARENT_ID = "__ROOT_ID__";
180
+
181
+ // src/adapters/vercel/useVercelAIBranches.tsx
174
182
  var UPCOMING_MESSAGE_ID = "__UPCOMING_MESSAGE_ID__";
175
183
  var updateBranchData = (data, messages) => {
176
184
  for (let i = 0; i < messages.length; i++) {
177
185
  const child = messages[i];
178
186
  const childId = child.id;
179
- const parentId = messages[i - 1]?.id ?? ROOT_ID;
187
+ const parentId = messages[i - 1]?.id ?? ROOT_PARENT_ID;
180
188
  data.parentMap.set(childId, parentId);
181
189
  const parentArray = data.branchMap.get(parentId);
182
190
  if (!parentArray) {
@@ -187,32 +195,32 @@ var updateBranchData = (data, messages) => {
187
195
  data.snapshots.set(childId, messages);
188
196
  }
189
197
  };
190
- var getParentId = (data, messages, message) => {
191
- if (message.id === UPCOMING_MESSAGE_ID) {
198
+ var getParentId = (data, messages, messageId) => {
199
+ if (messageId === UPCOMING_MESSAGE_ID) {
192
200
  const parent = messages.at(-1);
193
201
  if (!parent)
194
- return ROOT_ID;
202
+ return ROOT_PARENT_ID;
195
203
  return parent.id;
196
204
  }
197
- const parentId = data.parentMap.get(message.id);
205
+ const parentId = data.parentMap.get(messageId);
198
206
  if (!parentId)
199
207
  throw new Error("Unexpected: Message parent not found");
200
208
  return parentId;
201
209
  };
202
- var getBranchStateImpl = (data, messages, message) => {
203
- const parentId = getParentId(data, messages, message);
210
+ var getBranchStateImpl = (data, messages, messageId) => {
211
+ const parentId = getParentId(data, messages, messageId);
204
212
  const branches = data.branchMap.get(parentId) ?? [];
205
- const branchId = message.id === UPCOMING_MESSAGE_ID ? branches.length : branches.indexOf(message.id);
213
+ const branchId = messageId === UPCOMING_MESSAGE_ID ? branches.length : branches.indexOf(messageId);
206
214
  if (branchId === -1)
207
215
  throw new Error("Unexpected: Message not found in parent children");
208
- const upcomingOffset = message.id === UPCOMING_MESSAGE_ID ? 1 : 0;
216
+ const upcomingOffset = messageId === UPCOMING_MESSAGE_ID ? 1 : 0;
209
217
  return {
210
218
  branchId,
211
219
  branchCount: branches.length + upcomingOffset
212
220
  };
213
221
  };
214
- var switchToBranchImpl = (data, messages, message, branchId) => {
215
- const parentId = getParentId(data, messages, message);
222
+ var switchToBranchImpl = (data, messages, messageId, branchId) => {
223
+ const parentId = getParentId(data, messages, messageId);
216
224
  const branches = data.branchMap.get(parentId);
217
225
  if (!branches)
218
226
  throw new Error("Unexpected: Parent children not found");
@@ -221,20 +229,22 @@ var switchToBranchImpl = (data, messages, message, branchId) => {
221
229
  throw new Error("Unexpected: Requested branch not found");
222
230
  if (branchId < 0 || branchId >= branches.length)
223
231
  throw new Error("Switch to branch called with a branch index out of range");
224
- if (newMessageId === message.id)
232
+ if (newMessageId === messageId)
225
233
  return messages;
226
234
  const snapshot = data.snapshots.get(newMessageId);
227
235
  if (!snapshot)
228
236
  throw new Error("Unexpected: Branch snapshot not found");
229
237
  return snapshot;
230
238
  };
231
- var sliceMessagesUntil = (messages, message) => {
232
- if (message.id === UPCOMING_MESSAGE_ID)
239
+ var sliceMessagesUntil = (messages, messageId) => {
240
+ if (messageId === ROOT_PARENT_ID)
241
+ return [];
242
+ if (messageId === UPCOMING_MESSAGE_ID)
233
243
  return messages;
234
- const messageIdx = messages.findIndex((m) => m.id === message.id);
244
+ const messageIdx = messages.findIndex((m) => m.id === messageId);
235
245
  if (messageIdx === -1)
236
246
  throw new Error("Unexpected: Message not found");
237
- return messages.slice(0, messageIdx);
247
+ return messages.slice(0, messageIdx + 1);
238
248
  };
239
249
  var useVercelAIBranches = (chat, context) => {
240
250
  const data = useRef4({
@@ -244,17 +254,17 @@ var useVercelAIBranches = (chat, context) => {
244
254
  }).current;
245
255
  updateBranchData(data, chat.messages);
246
256
  const getBranchState = useCallback(
247
- (message) => {
248
- return getBranchStateImpl(data, chat.messages, message);
257
+ (messageId) => {
258
+ return getBranchStateImpl(data, chat.messages, messageId);
249
259
  },
250
260
  [data, chat.messages]
251
261
  );
252
262
  const switchToBranch = useCallback(
253
- (message, branchId) => {
263
+ (messageId, branchId) => {
254
264
  const newMessages = switchToBranchImpl(
255
265
  data,
256
266
  chat.messages,
257
- message,
267
+ messageId,
258
268
  branchId
259
269
  );
260
270
  chat.setMessages(newMessages);
@@ -262,27 +272,27 @@ var useVercelAIBranches = (chat, context) => {
262
272
  [data, chat.messages, chat.setMessages]
263
273
  );
264
274
  const reloadMaybe = "reload" in chat ? chat.reload : void 0;
265
- const reloadAt = useCallback(
266
- async (message) => {
275
+ const reload = useCallback(
276
+ async (messageId) => {
267
277
  if (!reloadMaybe)
268
278
  throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
269
- const newMessages = sliceMessagesUntil(chat.messages, message);
279
+ const newMessages = sliceMessagesUntil(chat.messages, messageId);
270
280
  chat.setMessages(newMessages);
271
281
  context.useThread.getState().scrollToBottom();
272
282
  await reloadMaybe();
273
283
  },
274
284
  [context, chat.messages, chat.setMessages, reloadMaybe]
275
285
  );
276
- const editAt = useCallback(
277
- async (message, newMessage) => {
278
- const newMessages = sliceMessagesUntil(chat.messages, message);
286
+ const append = useCallback(
287
+ async (message) => {
288
+ const newMessages = sliceMessagesUntil(chat.messages, message.parentId);
279
289
  chat.setMessages(newMessages);
280
- if (newMessage.content[0]?.type !== "text")
290
+ if (message.content.length !== 1 || message.content[0]?.type !== "text")
281
291
  throw new Error("Only text content is currently supported");
282
292
  context.useThread.getState().scrollToBottom();
283
293
  await chat.append({
284
294
  role: "user",
285
- content: newMessage.content[0].text
295
+ content: message.content[0].text
286
296
  });
287
297
  },
288
298
  [context, chat.messages, chat.setMessages, chat.append]
@@ -291,10 +301,10 @@ var useVercelAIBranches = (chat, context) => {
291
301
  () => ({
292
302
  getBranchState,
293
303
  switchToBranch,
294
- editAt,
295
- reloadAt
304
+ append,
305
+ reload
296
306
  }),
297
- [getBranchState, switchToBranch, editAt, reloadAt]
307
+ [getBranchState, switchToBranch, append, reload]
298
308
  );
299
309
  };
300
310
  var hasUpcomingMessage = (thread) => {
@@ -352,14 +362,13 @@ __export(message_exports, {
352
362
  // src/primitives/message/MessageProvider.tsx
353
363
  import { useMemo as useMemo2, useState } from "react";
354
364
  import { create } from "zustand";
355
- import { useShallow } from "zustand/react/shallow";
356
365
  var getIsLast = (thread, message) => {
357
366
  const hasUpcoming = hasUpcomingMessage(thread);
358
367
  return hasUpcoming ? message.id === UPCOMING_MESSAGE_ID : thread.messages[thread.messages.length - 1]?.id === message.id;
359
368
  };
360
369
  var useMessageContext2 = () => {
361
- const { useBranchObserver } = useAssistantContext();
362
370
  const [context] = useState(() => {
371
+ const { useThread } = useAssistantContext();
363
372
  const useMessage = create(() => ({
364
373
  message: null,
365
374
  isLast: false,
@@ -368,10 +377,6 @@ var useMessageContext2 = () => {
368
377
  setIsCopied: () => {
369
378
  },
370
379
  setIsHovering: () => {
371
- },
372
- branchState: {
373
- branchId: 0,
374
- branchCount: 0
375
380
  }
376
381
  }));
377
382
  const useComposer = create((set, get) => ({
@@ -393,8 +398,8 @@ var useMessageContext2 = () => {
393
398
  const message = useMessage.getState().message;
394
399
  if (message.role !== "user")
395
400
  throw new Error("Editing is only supported for user messages");
396
- useBranchObserver.getState().editAt(message, {
397
- role: "user",
401
+ useThread.getState().append({
402
+ parentId: message.parentId,
398
403
  content: [{ type: "text", text: get().value }]
399
404
  });
400
405
  set({ isEditing: false });
@@ -410,11 +415,8 @@ var MessageProvider = ({
410
415
  message,
411
416
  children
412
417
  }) => {
413
- const { useThread, useBranchObserver } = useAssistantContext();
418
+ const { useThread } = useAssistantContext();
414
419
  const context = useMessageContext2();
415
- const branchState = useBranchObserver(
416
- useShallow((b) => b.getBranchState(message))
417
- );
418
420
  const isLast = useThread((thread) => getIsLast(thread, message));
419
421
  const [isCopied, setIsCopied] = useState(false);
420
422
  const [isHovering, setIsHovering] = useState(false);
@@ -426,12 +428,11 @@ var MessageProvider = ({
426
428
  isCopied,
427
429
  isHovering,
428
430
  setIsCopied,
429
- setIsHovering,
430
- branchState
431
+ setIsHovering
431
432
  },
432
433
  true
433
434
  );
434
- }, [context, message, isLast, isCopied, isHovering, branchState]);
435
+ }, [context, message, isLast, isCopied, isHovering]);
435
436
  return /* @__PURE__ */ React.createElement(MessageContext.Provider, { value: context }, children);
436
437
  };
437
438
 
@@ -466,29 +467,21 @@ var MessageRoot = forwardRef3(
466
467
  // src/primitives/message/MessageIf.tsx
467
468
  var useMessageIf = (props) => {
468
469
  const { useMessage } = useMessageContext();
469
- return useMessage(
470
- ({
471
- message,
472
- isLast,
473
- isCopied,
474
- isHovering,
475
- branchState: { branchCount }
476
- }) => {
477
- if (props.hasBranches === true && branchCount < 2)
478
- return false;
479
- if (props.user && message.role !== "user")
480
- return false;
481
- if (props.assistant && message.role !== "assistant")
482
- return false;
483
- if (props.lastOrHover === true && !isHovering && !isLast)
484
- return false;
485
- if (props.copied === true && !isCopied)
486
- return false;
487
- if (props.copied === false && isCopied)
488
- return false;
489
- return true;
490
- }
491
- );
470
+ return useMessage(({ message, isLast, isCopied, isHovering }) => {
471
+ if (props.hasBranches === true && message.branchCount < 2)
472
+ return false;
473
+ if (props.user && message.role !== "user")
474
+ return false;
475
+ if (props.assistant && message.role !== "assistant")
476
+ return false;
477
+ if (props.lastOrHover === true && !isHovering && !isLast)
478
+ return false;
479
+ if (props.copied === true && !isCopied)
480
+ return false;
481
+ if (props.copied === false && isCopied)
482
+ return false;
483
+ return true;
484
+ });
492
485
  };
493
486
  var MessageIf = ({ children, ...query }) => {
494
487
  const result = useMessageIf(query);
@@ -559,7 +552,12 @@ var ThreadMessages = ({ components }) => {
559
552
  message: {
560
553
  id: UPCOMING_MESSAGE_ID,
561
554
  role: "assistant",
562
- content: [{ type: "text", text: "..." }]
555
+ content: [{ type: "text", text: "..." }],
556
+ parentId: messages.at(-1)?.id ?? ROOT_PARENT_ID,
557
+ // TODO fix these (move upcoming message to AssistantContext)
558
+ branchId: 0,
559
+ branchCount: 1,
560
+ createdAt: /* @__PURE__ */ new Date()
563
561
  }
564
562
  },
565
563
  /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -579,12 +577,11 @@ var ThreadScrollToBottom = forwardRef4(({ onClick, ...rest }, ref) => {
579
577
  const thread = useThread.getState();
580
578
  thread.scrollToBottom();
581
579
  };
582
- if (isAtBottom)
583
- return null;
584
580
  return /* @__PURE__ */ React.createElement(
585
581
  Primitive4.button,
586
582
  {
587
583
  ...rest,
584
+ disabled: isAtBottom,
588
585
  ref,
589
586
  onClick: composeEventHandlers3(onClick, handleScrollToBottom)
590
587
  }
@@ -789,20 +786,17 @@ var useCombinedStore = (stores, selector) => {
789
786
 
790
787
  // src/actions/useGoToNextBranch.tsx
791
788
  var useGoToNextBranch = () => {
792
- const { useThread, useBranchObserver } = useAssistantContext();
789
+ const { useThread } = useAssistantContext();
793
790
  const { useComposer, useMessage } = useMessageContext();
794
791
  const disabled = useCombinedStore(
795
792
  [useThread, useComposer, useMessage],
796
- (t, c, m) => t.isLoading || c.isEditing || m.branchState.branchId + 1 >= m.branchState.branchCount
793
+ (t, c, m) => t.isLoading || c.isEditing || m.message.branchId + 1 >= m.message.branchCount
797
794
  );
798
795
  if (disabled)
799
796
  return null;
800
797
  return () => {
801
- const {
802
- message,
803
- branchState: { branchId }
804
- } = useMessage.getState();
805
- useBranchObserver.getState().switchToBranch(message, branchId + 1);
798
+ const { message } = useMessage.getState();
799
+ useThread.getState().switchToBranch(message.id, message.branchId + 1);
806
800
  };
807
801
  };
808
802
 
@@ -835,20 +829,17 @@ var BranchPickerNext = createActionButton(useGoToNextBranch);
835
829
 
836
830
  // src/actions/useGoToPreviousBranch.tsx
837
831
  var useGoToPreviousBranch = () => {
838
- const { useThread, useBranchObserver } = useAssistantContext();
832
+ const { useThread } = useAssistantContext();
839
833
  const { useComposer, useMessage } = useMessageContext();
840
834
  const disabled = useCombinedStore(
841
835
  [useThread, useComposer, useMessage],
842
- (t, c, m) => t.isLoading || c.isEditing || m.branchState.branchId <= 0
836
+ (t, c, m) => t.isLoading || c.isEditing || m.message.branchId <= 0
843
837
  );
844
838
  if (disabled)
845
839
  return null;
846
840
  return () => {
847
- const {
848
- message,
849
- branchState: { branchId }
850
- } = useMessage.getState();
851
- useBranchObserver.getState().switchToBranch(message, branchId - 1);
841
+ const { message } = useMessage.getState();
842
+ useThread.getState().switchToBranch(message.id, message.branchId - 1);
852
843
  };
853
844
  };
854
845
 
@@ -858,14 +849,14 @@ var BranchPickerPrevious = createActionButton(useGoToPreviousBranch);
858
849
  // src/primitives/branchPicker/BranchPickerCount.tsx
859
850
  var BranchPickerCount = () => {
860
851
  const { useMessage } = useMessageContext();
861
- const branchCount = useMessage((s) => s.branchState.branchCount);
852
+ const branchCount = useMessage((s) => s.message.branchCount);
862
853
  return /* @__PURE__ */ React.createElement(React.Fragment, null, branchCount);
863
854
  };
864
855
 
865
856
  // src/primitives/branchPicker/BranchPickerNumber.tsx
866
857
  var BranchPickerNumber = () => {
867
858
  const { useMessage } = useMessageContext();
868
- const branchId = useMessage((s) => s.branchState.branchId);
859
+ const branchId = useMessage((s) => s.message.branchId);
869
860
  return /* @__PURE__ */ React.createElement(React.Fragment, null, branchId + 1);
870
861
  };
871
862
 
@@ -905,7 +896,7 @@ var ActionBarRoot = forwardRef11(({ hideWhenBusy, autohide, autohideFloat, ...re
905
896
  return "normal" /* Normal */;
906
897
  if (!m.isHovering)
907
898
  return "hidden" /* Hidden */;
908
- if (autohideFloat === "always" || autohideFloat === "single-branch" && m.branchState.branchCount <= 1)
899
+ if (autohideFloat === "always" || autohideFloat === "single-branch" && m.message.branchCount <= 1)
909
900
  return "floating" /* Floating */;
910
901
  return "normal" /* Normal */;
911
902
  }
@@ -943,7 +934,7 @@ var ActionBarCopy = createActionButton(useCopyMessage);
943
934
 
944
935
  // src/actions/useReloadMessage.tsx
945
936
  var useReloadMessage = () => {
946
- const { useThread, useBranchObserver } = useAssistantContext();
937
+ const { useThread } = useAssistantContext();
947
938
  const { useMessage } = useMessageContext();
948
939
  const disabled = useCombinedStore(
949
940
  [useThread, useMessage],
@@ -955,7 +946,7 @@ var useReloadMessage = () => {
955
946
  const message = useMessage.getState().message;
956
947
  if (message.role !== "assistant")
957
948
  throw new Error("Reloading is only supported on assistant messages");
958
- useBranchObserver.getState().reloadAt(message);
949
+ useThread.getState().reload(message.id);
959
950
  };
960
951
  };
961
952
 
@@ -980,10 +971,10 @@ var useBeginMessageEdit = () => {
980
971
  // src/primitives/actionBar/ActionBarEdit.tsx
981
972
  var ActionBarEdit = createActionButton(useBeginMessageEdit);
982
973
 
983
- // src/vercel/VercelAIAssistantProvider.tsx
974
+ // src/adapters/vercel/VercelAIAssistantProvider.tsx
984
975
  import { useCallback as useCallback3, useMemo as useMemo4 } from "react";
985
976
 
986
- // src/vercel/useDummyAIAssistantContext.tsx
977
+ // src/adapters/vercel/useDummyAIAssistantContext.tsx
987
978
  import { useState as useState2 } from "react";
988
979
  import { create as create2 } from "zustand";
989
980
  var useDummyAIAssistantContext = () => {
@@ -998,6 +989,12 @@ var useDummyAIAssistantContext = () => {
998
989
  stop: () => {
999
990
  throw new Error("Not implemented");
1000
991
  },
992
+ switchToBranch: () => {
993
+ throw new Error("Not implemented");
994
+ },
995
+ reload: async () => {
996
+ throw new Error("Not implemented");
997
+ },
1001
998
  isAtBottom: true,
1002
999
  scrollToBottom: () => {
1003
1000
  for (const listener of scrollToBottomListeners) {
@@ -1023,7 +1020,7 @@ var useDummyAIAssistantContext = () => {
1023
1020
  },
1024
1021
  send: () => {
1025
1022
  useThread.getState().append({
1026
- role: "user",
1023
+ parentId: useThread.getState().messages.at(-1)?.id ?? ROOT_PARENT_ID,
1027
1024
  content: [{ type: "text", text: useComposer.getState().value }]
1028
1025
  });
1029
1026
  useComposer.getState().setValue("");
@@ -1032,43 +1029,39 @@ var useDummyAIAssistantContext = () => {
1032
1029
  useThread.getState().stop();
1033
1030
  }
1034
1031
  }));
1035
- const useBranchObserver = create2()(() => ({
1036
- getBranchState: () => ({
1037
- branchId: 0,
1038
- branchCount: 1
1039
- }),
1040
- switchToBranch: () => {
1041
- throw new Error("Not implemented");
1042
- },
1043
- editAt: async () => {
1044
- throw new Error("Not implemented");
1045
- },
1046
- reloadAt: async () => {
1047
- throw new Error("Not implemented");
1048
- }
1049
- }));
1050
- return { useThread, useComposer, useBranchObserver };
1032
+ return { useThread, useComposer };
1051
1033
  });
1052
1034
  return context;
1053
1035
  };
1054
1036
 
1055
- // src/vercel/VercelAIAssistantProvider.tsx
1037
+ // src/adapters/vercel/VercelAIAssistantProvider.tsx
1056
1038
  var ThreadMessageCache = /* @__PURE__ */ new WeakMap();
1057
- var vercelToThreadMessage = (message) => {
1039
+ var vercelToThreadMessage = (message, parentId, branchId, branchCount) => {
1058
1040
  if (message.role !== "user" && message.role !== "assistant")
1059
1041
  throw new Error("Unsupported role");
1060
1042
  return {
1043
+ parentId,
1061
1044
  id: message.id,
1062
1045
  role: message.role,
1063
- content: [{ type: "text", text: message.content }]
1046
+ content: [{ type: "text", text: message.content }],
1047
+ branchId,
1048
+ branchCount,
1049
+ createdAt: message.createdAt ?? /* @__PURE__ */ new Date()
1064
1050
  };
1065
1051
  };
1066
- var vercelToCachedThreadMessages = (messages) => {
1067
- return messages.map((m) => {
1052
+ var vercelToCachedThreadMessages = (messages, getBranchState) => {
1053
+ return messages.map((m, idx) => {
1068
1054
  const cached = ThreadMessageCache.get(m);
1069
- if (cached)
1055
+ const parentId = messages[idx - 1]?.id ?? ROOT_PARENT_ID;
1056
+ const { branchId, branchCount } = getBranchState(m.id);
1057
+ if (cached && cached.parentId === parentId && cached.branchId === branchId && cached.branchCount === branchCount)
1070
1058
  return cached;
1071
- const newMessage = vercelToThreadMessage(m);
1059
+ const newMessage = vercelToThreadMessage(
1060
+ m,
1061
+ parentId,
1062
+ branchId,
1063
+ branchCount
1064
+ );
1072
1065
  ThreadMessageCache.set(m, newMessage);
1073
1066
  return newMessage;
1074
1067
  });
@@ -1079,22 +1072,13 @@ var VercelAIAssistantProvider = ({
1079
1072
  }) => {
1080
1073
  const context = useDummyAIAssistantContext();
1081
1074
  const vercel = "chat" in rest ? rest.chat : rest.assistant;
1075
+ const branches = useVercelAIBranches(vercel, context);
1082
1076
  const messages = useMemo4(() => {
1083
- return vercelToCachedThreadMessages(vercel.messages);
1084
- }, [vercel.messages]);
1085
- const append = useCallback3(
1086
- async (message) => {
1087
- if (message.content[0]?.type !== "text") {
1088
- throw new Error("Only text content is currently supported");
1089
- }
1090
- context.useThread.getState().scrollToBottom();
1091
- await vercel.append({
1092
- role: message.role,
1093
- content: message.content[0].text
1094
- });
1095
- },
1096
- [context, vercel.append]
1097
- );
1077
+ return vercelToCachedThreadMessages(
1078
+ vercel.messages,
1079
+ branches.getBranchState
1080
+ );
1081
+ }, [vercel.messages, branches.getBranchState]);
1098
1082
  const stop = useCallback3(() => {
1099
1083
  const lastMessage = vercel.messages.at(-1);
1100
1084
  vercel.stop();
@@ -1107,10 +1091,12 @@ var VercelAIAssistantProvider = ({
1107
1091
  context.useThread.setState({
1108
1092
  messages,
1109
1093
  isLoading,
1110
- append,
1111
- stop
1094
+ stop,
1095
+ switchToBranch: branches.switchToBranch,
1096
+ append: branches.append,
1097
+ reload: branches.reload
1112
1098
  });
1113
- }, [context, messages, append, stop, isLoading]);
1099
+ }, [context, messages, isLoading, stop, branches]);
1114
1100
  useMemo4(() => {
1115
1101
  context.useComposer.setState({
1116
1102
  canCancel: isLoading,
@@ -1118,34 +1104,35 @@ var VercelAIAssistantProvider = ({
1118
1104
  setValue: vercel.setInput
1119
1105
  });
1120
1106
  }, [context, isLoading, vercel.input, vercel.setInput]);
1121
- const branches = useVercelAIBranches(vercel, context);
1122
- useMemo4(() => {
1123
- context.useBranchObserver.setState(branches, true);
1124
- }, [context, branches]);
1125
1107
  return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
1126
1108
  };
1127
1109
 
1128
- // src/vercel/VercelRSCAssistantProvider.tsx
1110
+ // src/adapters/vercel/VercelRSCAssistantProvider.tsx
1129
1111
  import {
1130
1112
  useCallback as useCallback4,
1131
1113
  useMemo as useMemo5
1132
1114
  } from "react";
1133
1115
  var ThreadMessageCache2 = /* @__PURE__ */ new WeakMap();
1134
- var vercelToThreadMessage2 = (message) => {
1116
+ var vercelToThreadMessage2 = (parentId, message) => {
1135
1117
  if (message.role !== "user" && message.role !== "assistant")
1136
1118
  throw new Error("Unsupported role");
1137
1119
  return {
1120
+ parentId,
1138
1121
  id: message.id,
1139
1122
  role: message.role,
1140
- content: [{ type: "ui", display: message.display }]
1123
+ content: [{ type: "ui", display: message.display }],
1124
+ createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
1125
+ branchId: 0,
1126
+ branchCount: 1
1141
1127
  };
1142
1128
  };
1143
1129
  var vercelToCachedThreadMessages2 = (messages) => {
1144
- return messages.map((m) => {
1130
+ return messages.map((m, idx) => {
1145
1131
  const cached = ThreadMessageCache2.get(m);
1146
- if (cached)
1132
+ const parentId = messages[idx - 1]?.id ?? ROOT_PARENT_ID;
1133
+ if (cached && cached.parentId === parentId)
1147
1134
  return cached;
1148
- const newMessage = vercelToThreadMessage2(m);
1135
+ const newMessage = vercelToThreadMessage2(parentId, m);
1149
1136
  ThreadMessageCache2.set(m, newMessage);
1150
1137
  return newMessage;
1151
1138
  });
@@ -1161,14 +1148,13 @@ var VercelRSCAssistantProvider = ({
1161
1148
  }, [vercelMessages]);
1162
1149
  const append = useCallback4(
1163
1150
  async (message) => {
1151
+ if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? ROOT_PARENT_ID))
1152
+ throw new Error("Unexpected: Message editing is not supported");
1164
1153
  if (message.content[0]?.type !== "text") {
1165
1154
  throw new Error("Only text content is currently supported");
1166
1155
  }
1167
1156
  context.useThread.getState().scrollToBottom();
1168
- await vercelAppend({
1169
- role: message.role,
1170
- content: [{ type: "text", text: message.content[0].text }]
1171
- });
1157
+ await vercelAppend(message);
1172
1158
  },
1173
1159
  [context, vercelAppend]
1174
1160
  );
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@assistant-ui/react",
3
- "version": "0.0.8",
3
+ "version": "0.0.9",
4
4
  "license": "MIT",
5
5
  "exports": {
6
6
  ".": {
@@ -28,7 +28,7 @@
28
28
  "@radix-ui/react-compose-refs": "^1.0.1",
29
29
  "@radix-ui/react-primitive": "^1.0.3",
30
30
  "@radix-ui/react-slot": "^1.0.2",
31
- "ai": "^3.1.12",
31
+ "ai": "^3.1.15",
32
32
  "react-textarea-autosize": "^8.5.3",
33
33
  "zustand": "^4.5.2"
34
34
  },