@assistant-ui/react 0.0.8 → 0.0.9

Sign up to get free protection for your applications and to get access to all the features.
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
  },