@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.js CHANGED
@@ -168,12 +168,14 @@ var ThreadViewport = (0, import_react5.forwardRef)(({ autoScroll = true, onScrol
168
168
  const ref = (0, import_react_compose_refs.useComposedRefs)(forwardedRef, divRef);
169
169
  const { useThread } = useAssistantContext();
170
170
  const firstRenderRef = (0, import_react5.useRef)(true);
171
+ const lastScrollTop = (0, import_react5.useRef)(0);
171
172
  const scrollToBottom = () => {
172
173
  const div = messagesEndRef.current;
173
174
  if (!div || !autoScroll)
174
175
  return;
175
176
  const behavior = firstRenderRef.current ? "instant" : "auto";
176
177
  firstRenderRef.current = false;
178
+ useThread.setState({ isAtBottom: true });
177
179
  div.scrollIntoView({ behavior });
178
180
  };
179
181
  useOnResizeContent(divRef, () => {
@@ -189,10 +191,12 @@ var ThreadViewport = (0, import_react5.forwardRef)(({ autoScroll = true, onScrol
189
191
  if (!div)
190
192
  return;
191
193
  const isAtBottom = useThread.getState().isAtBottom;
192
- const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight + 50;
193
- if (newIsAtBottom !== isAtBottom) {
194
+ const newIsAtBottom = div.scrollHeight - div.scrollTop <= div.clientHeight;
195
+ if (!newIsAtBottom && lastScrollTop.current < div.scrollTop) {
196
+ } else if (newIsAtBottom !== isAtBottom) {
194
197
  useThread.setState({ isAtBottom: newIsAtBottom });
195
198
  }
199
+ lastScrollTop.current = div.scrollTop;
196
200
  };
197
201
  return /* @__PURE__ */ React.createElement(
198
202
  import_react_primitive2.Primitive.div,
@@ -206,15 +210,19 @@ var ThreadViewport = (0, import_react5.forwardRef)(({ autoScroll = true, onScrol
206
210
  );
207
211
  });
208
212
 
209
- // src/vercel/useVercelAIBranches.tsx
213
+ // src/adapters/vercel/useVercelAIBranches.tsx
210
214
  var import_react6 = require("react");
211
- var ROOT_ID = "__ROOT_ID__";
215
+
216
+ // src/utils/context/stores/AssistantTypes.ts
217
+ var ROOT_PARENT_ID = "__ROOT_ID__";
218
+
219
+ // src/adapters/vercel/useVercelAIBranches.tsx
212
220
  var UPCOMING_MESSAGE_ID = "__UPCOMING_MESSAGE_ID__";
213
221
  var updateBranchData = (data, messages) => {
214
222
  for (let i = 0; i < messages.length; i++) {
215
223
  const child = messages[i];
216
224
  const childId = child.id;
217
- const parentId = messages[i - 1]?.id ?? ROOT_ID;
225
+ const parentId = messages[i - 1]?.id ?? ROOT_PARENT_ID;
218
226
  data.parentMap.set(childId, parentId);
219
227
  const parentArray = data.branchMap.get(parentId);
220
228
  if (!parentArray) {
@@ -225,32 +233,32 @@ var updateBranchData = (data, messages) => {
225
233
  data.snapshots.set(childId, messages);
226
234
  }
227
235
  };
228
- var getParentId = (data, messages, message) => {
229
- if (message.id === UPCOMING_MESSAGE_ID) {
236
+ var getParentId = (data, messages, messageId) => {
237
+ if (messageId === UPCOMING_MESSAGE_ID) {
230
238
  const parent = messages.at(-1);
231
239
  if (!parent)
232
- return ROOT_ID;
240
+ return ROOT_PARENT_ID;
233
241
  return parent.id;
234
242
  }
235
- const parentId = data.parentMap.get(message.id);
243
+ const parentId = data.parentMap.get(messageId);
236
244
  if (!parentId)
237
245
  throw new Error("Unexpected: Message parent not found");
238
246
  return parentId;
239
247
  };
240
- var getBranchStateImpl = (data, messages, message) => {
241
- const parentId = getParentId(data, messages, message);
248
+ var getBranchStateImpl = (data, messages, messageId) => {
249
+ const parentId = getParentId(data, messages, messageId);
242
250
  const branches = data.branchMap.get(parentId) ?? [];
243
- const branchId = message.id === UPCOMING_MESSAGE_ID ? branches.length : branches.indexOf(message.id);
251
+ const branchId = messageId === UPCOMING_MESSAGE_ID ? branches.length : branches.indexOf(messageId);
244
252
  if (branchId === -1)
245
253
  throw new Error("Unexpected: Message not found in parent children");
246
- const upcomingOffset = message.id === UPCOMING_MESSAGE_ID ? 1 : 0;
254
+ const upcomingOffset = messageId === UPCOMING_MESSAGE_ID ? 1 : 0;
247
255
  return {
248
256
  branchId,
249
257
  branchCount: branches.length + upcomingOffset
250
258
  };
251
259
  };
252
- var switchToBranchImpl = (data, messages, message, branchId) => {
253
- const parentId = getParentId(data, messages, message);
260
+ var switchToBranchImpl = (data, messages, messageId, branchId) => {
261
+ const parentId = getParentId(data, messages, messageId);
254
262
  const branches = data.branchMap.get(parentId);
255
263
  if (!branches)
256
264
  throw new Error("Unexpected: Parent children not found");
@@ -259,20 +267,22 @@ var switchToBranchImpl = (data, messages, message, branchId) => {
259
267
  throw new Error("Unexpected: Requested branch not found");
260
268
  if (branchId < 0 || branchId >= branches.length)
261
269
  throw new Error("Switch to branch called with a branch index out of range");
262
- if (newMessageId === message.id)
270
+ if (newMessageId === messageId)
263
271
  return messages;
264
272
  const snapshot = data.snapshots.get(newMessageId);
265
273
  if (!snapshot)
266
274
  throw new Error("Unexpected: Branch snapshot not found");
267
275
  return snapshot;
268
276
  };
269
- var sliceMessagesUntil = (messages, message) => {
270
- if (message.id === UPCOMING_MESSAGE_ID)
277
+ var sliceMessagesUntil = (messages, messageId) => {
278
+ if (messageId === ROOT_PARENT_ID)
279
+ return [];
280
+ if (messageId === UPCOMING_MESSAGE_ID)
271
281
  return messages;
272
- const messageIdx = messages.findIndex((m) => m.id === message.id);
282
+ const messageIdx = messages.findIndex((m) => m.id === messageId);
273
283
  if (messageIdx === -1)
274
284
  throw new Error("Unexpected: Message not found");
275
- return messages.slice(0, messageIdx);
285
+ return messages.slice(0, messageIdx + 1);
276
286
  };
277
287
  var useVercelAIBranches = (chat, context) => {
278
288
  const data = (0, import_react6.useRef)({
@@ -282,17 +292,17 @@ var useVercelAIBranches = (chat, context) => {
282
292
  }).current;
283
293
  updateBranchData(data, chat.messages);
284
294
  const getBranchState = (0, import_react6.useCallback)(
285
- (message) => {
286
- return getBranchStateImpl(data, chat.messages, message);
295
+ (messageId) => {
296
+ return getBranchStateImpl(data, chat.messages, messageId);
287
297
  },
288
298
  [data, chat.messages]
289
299
  );
290
300
  const switchToBranch = (0, import_react6.useCallback)(
291
- (message, branchId) => {
301
+ (messageId, branchId) => {
292
302
  const newMessages = switchToBranchImpl(
293
303
  data,
294
304
  chat.messages,
295
- message,
305
+ messageId,
296
306
  branchId
297
307
  );
298
308
  chat.setMessages(newMessages);
@@ -300,27 +310,27 @@ var useVercelAIBranches = (chat, context) => {
300
310
  [data, chat.messages, chat.setMessages]
301
311
  );
302
312
  const reloadMaybe = "reload" in chat ? chat.reload : void 0;
303
- const reloadAt = (0, import_react6.useCallback)(
304
- async (message) => {
313
+ const reload = (0, import_react6.useCallback)(
314
+ async (messageId) => {
305
315
  if (!reloadMaybe)
306
316
  throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
307
- const newMessages = sliceMessagesUntil(chat.messages, message);
317
+ const newMessages = sliceMessagesUntil(chat.messages, messageId);
308
318
  chat.setMessages(newMessages);
309
319
  context.useThread.getState().scrollToBottom();
310
320
  await reloadMaybe();
311
321
  },
312
322
  [context, chat.messages, chat.setMessages, reloadMaybe]
313
323
  );
314
- const editAt = (0, import_react6.useCallback)(
315
- async (message, newMessage) => {
316
- const newMessages = sliceMessagesUntil(chat.messages, message);
324
+ const append = (0, import_react6.useCallback)(
325
+ async (message) => {
326
+ const newMessages = sliceMessagesUntil(chat.messages, message.parentId);
317
327
  chat.setMessages(newMessages);
318
- if (newMessage.content[0]?.type !== "text")
328
+ if (message.content.length !== 1 || message.content[0]?.type !== "text")
319
329
  throw new Error("Only text content is currently supported");
320
330
  context.useThread.getState().scrollToBottom();
321
331
  await chat.append({
322
332
  role: "user",
323
- content: newMessage.content[0].text
333
+ content: message.content[0].text
324
334
  });
325
335
  },
326
336
  [context, chat.messages, chat.setMessages, chat.append]
@@ -329,10 +339,10 @@ var useVercelAIBranches = (chat, context) => {
329
339
  () => ({
330
340
  getBranchState,
331
341
  switchToBranch,
332
- editAt,
333
- reloadAt
342
+ append,
343
+ reload
334
344
  }),
335
- [getBranchState, switchToBranch, editAt, reloadAt]
345
+ [getBranchState, switchToBranch, append, reload]
336
346
  );
337
347
  };
338
348
  var hasUpcomingMessage = (thread) => {
@@ -390,14 +400,13 @@ __export(message_exports, {
390
400
  // src/primitives/message/MessageProvider.tsx
391
401
  var import_react9 = require("react");
392
402
  var import_zustand = require("zustand");
393
- var import_shallow = require("zustand/react/shallow");
394
403
  var getIsLast = (thread, message) => {
395
404
  const hasUpcoming = hasUpcomingMessage(thread);
396
405
  return hasUpcoming ? message.id === UPCOMING_MESSAGE_ID : thread.messages[thread.messages.length - 1]?.id === message.id;
397
406
  };
398
407
  var useMessageContext2 = () => {
399
- const { useBranchObserver } = useAssistantContext();
400
408
  const [context] = (0, import_react9.useState)(() => {
409
+ const { useThread } = useAssistantContext();
401
410
  const useMessage = (0, import_zustand.create)(() => ({
402
411
  message: null,
403
412
  isLast: false,
@@ -406,10 +415,6 @@ var useMessageContext2 = () => {
406
415
  setIsCopied: () => {
407
416
  },
408
417
  setIsHovering: () => {
409
- },
410
- branchState: {
411
- branchId: 0,
412
- branchCount: 0
413
418
  }
414
419
  }));
415
420
  const useComposer = (0, import_zustand.create)((set, get) => ({
@@ -431,8 +436,8 @@ var useMessageContext2 = () => {
431
436
  const message = useMessage.getState().message;
432
437
  if (message.role !== "user")
433
438
  throw new Error("Editing is only supported for user messages");
434
- useBranchObserver.getState().editAt(message, {
435
- role: "user",
439
+ useThread.getState().append({
440
+ parentId: message.parentId,
436
441
  content: [{ type: "text", text: get().value }]
437
442
  });
438
443
  set({ isEditing: false });
@@ -448,11 +453,8 @@ var MessageProvider = ({
448
453
  message,
449
454
  children
450
455
  }) => {
451
- const { useThread, useBranchObserver } = useAssistantContext();
456
+ const { useThread } = useAssistantContext();
452
457
  const context = useMessageContext2();
453
- const branchState = useBranchObserver(
454
- (0, import_shallow.useShallow)((b) => b.getBranchState(message))
455
- );
456
458
  const isLast = useThread((thread) => getIsLast(thread, message));
457
459
  const [isCopied, setIsCopied] = (0, import_react9.useState)(false);
458
460
  const [isHovering, setIsHovering] = (0, import_react9.useState)(false);
@@ -464,12 +466,11 @@ var MessageProvider = ({
464
466
  isCopied,
465
467
  isHovering,
466
468
  setIsCopied,
467
- setIsHovering,
468
- branchState
469
+ setIsHovering
469
470
  },
470
471
  true
471
472
  );
472
- }, [context, message, isLast, isCopied, isHovering, branchState]);
473
+ }, [context, message, isLast, isCopied, isHovering]);
473
474
  return /* @__PURE__ */ React.createElement(MessageContext.Provider, { value: context }, children);
474
475
  };
475
476
 
@@ -502,29 +503,21 @@ var MessageRoot = (0, import_react10.forwardRef)(
502
503
  // src/primitives/message/MessageIf.tsx
503
504
  var useMessageIf = (props) => {
504
505
  const { useMessage } = useMessageContext();
505
- return useMessage(
506
- ({
507
- message,
508
- isLast,
509
- isCopied,
510
- isHovering,
511
- branchState: { branchCount }
512
- }) => {
513
- if (props.hasBranches === true && branchCount < 2)
514
- return false;
515
- if (props.user && message.role !== "user")
516
- return false;
517
- if (props.assistant && message.role !== "assistant")
518
- return false;
519
- if (props.lastOrHover === true && !isHovering && !isLast)
520
- return false;
521
- if (props.copied === true && !isCopied)
522
- return false;
523
- if (props.copied === false && isCopied)
524
- return false;
525
- return true;
526
- }
527
- );
506
+ return useMessage(({ message, isLast, isCopied, isHovering }) => {
507
+ if (props.hasBranches === true && message.branchCount < 2)
508
+ return false;
509
+ if (props.user && message.role !== "user")
510
+ return false;
511
+ if (props.assistant && message.role !== "assistant")
512
+ return false;
513
+ if (props.lastOrHover === true && !isHovering && !isLast)
514
+ return false;
515
+ if (props.copied === true && !isCopied)
516
+ return false;
517
+ if (props.copied === false && isCopied)
518
+ return false;
519
+ return true;
520
+ });
528
521
  };
529
522
  var MessageIf = ({ children, ...query }) => {
530
523
  const result = useMessageIf(query);
@@ -595,7 +588,12 @@ var ThreadMessages = ({ components }) => {
595
588
  message: {
596
589
  id: UPCOMING_MESSAGE_ID,
597
590
  role: "assistant",
598
- content: [{ type: "text", text: "..." }]
591
+ content: [{ type: "text", text: "..." }],
592
+ parentId: messages.at(-1)?.id ?? ROOT_PARENT_ID,
593
+ // TODO fix these (move upcoming message to AssistantContext)
594
+ branchId: 0,
595
+ branchCount: 1,
596
+ createdAt: /* @__PURE__ */ new Date()
599
597
  }
600
598
  },
601
599
  /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -613,12 +611,11 @@ var ThreadScrollToBottom = (0, import_react11.forwardRef)(({ onClick, ...rest },
613
611
  const thread = useThread.getState();
614
612
  thread.scrollToBottom();
615
613
  };
616
- if (isAtBottom)
617
- return null;
618
614
  return /* @__PURE__ */ React.createElement(
619
615
  import_react_primitive4.Primitive.button,
620
616
  {
621
617
  ...rest,
618
+ disabled: isAtBottom,
622
619
  ref,
623
620
  onClick: (0, import_primitive3.composeEventHandlers)(onClick, handleScrollToBottom)
624
621
  }
@@ -812,20 +809,17 @@ var useCombinedStore = (stores, selector) => {
812
809
 
813
810
  // src/actions/useGoToNextBranch.tsx
814
811
  var useGoToNextBranch = () => {
815
- const { useThread, useBranchObserver } = useAssistantContext();
812
+ const { useThread } = useAssistantContext();
816
813
  const { useComposer, useMessage } = useMessageContext();
817
814
  const disabled = useCombinedStore(
818
815
  [useThread, useComposer, useMessage],
819
- (t, c, m) => t.isLoading || c.isEditing || m.branchState.branchId + 1 >= m.branchState.branchCount
816
+ (t, c, m) => t.isLoading || c.isEditing || m.message.branchId + 1 >= m.message.branchCount
820
817
  );
821
818
  if (disabled)
822
819
  return null;
823
820
  return () => {
824
- const {
825
- message,
826
- branchState: { branchId }
827
- } = useMessage.getState();
828
- useBranchObserver.getState().switchToBranch(message, branchId + 1);
821
+ const { message } = useMessage.getState();
822
+ useThread.getState().switchToBranch(message.id, message.branchId + 1);
829
823
  };
830
824
  };
831
825
 
@@ -856,20 +850,17 @@ var BranchPickerNext = createActionButton(useGoToNextBranch);
856
850
 
857
851
  // src/actions/useGoToPreviousBranch.tsx
858
852
  var useGoToPreviousBranch = () => {
859
- const { useThread, useBranchObserver } = useAssistantContext();
853
+ const { useThread } = useAssistantContext();
860
854
  const { useComposer, useMessage } = useMessageContext();
861
855
  const disabled = useCombinedStore(
862
856
  [useThread, useComposer, useMessage],
863
- (t, c, m) => t.isLoading || c.isEditing || m.branchState.branchId <= 0
857
+ (t, c, m) => t.isLoading || c.isEditing || m.message.branchId <= 0
864
858
  );
865
859
  if (disabled)
866
860
  return null;
867
861
  return () => {
868
- const {
869
- message,
870
- branchState: { branchId }
871
- } = useMessage.getState();
872
- useBranchObserver.getState().switchToBranch(message, branchId - 1);
862
+ const { message } = useMessage.getState();
863
+ useThread.getState().switchToBranch(message.id, message.branchId - 1);
873
864
  };
874
865
  };
875
866
 
@@ -879,14 +870,14 @@ var BranchPickerPrevious = createActionButton(useGoToPreviousBranch);
879
870
  // src/primitives/branchPicker/BranchPickerCount.tsx
880
871
  var BranchPickerCount = () => {
881
872
  const { useMessage } = useMessageContext();
882
- const branchCount = useMessage((s) => s.branchState.branchCount);
873
+ const branchCount = useMessage((s) => s.message.branchCount);
883
874
  return /* @__PURE__ */ React.createElement(React.Fragment, null, branchCount);
884
875
  };
885
876
 
886
877
  // src/primitives/branchPicker/BranchPickerNumber.tsx
887
878
  var BranchPickerNumber = () => {
888
879
  const { useMessage } = useMessageContext();
889
- const branchId = useMessage((s) => s.branchState.branchId);
880
+ const branchId = useMessage((s) => s.message.branchId);
890
881
  return /* @__PURE__ */ React.createElement(React.Fragment, null, branchId + 1);
891
882
  };
892
883
 
@@ -922,7 +913,7 @@ var ActionBarRoot = (0, import_react20.forwardRef)(({ hideWhenBusy, autohide, au
922
913
  return "normal" /* Normal */;
923
914
  if (!m.isHovering)
924
915
  return "hidden" /* Hidden */;
925
- if (autohideFloat === "always" || autohideFloat === "single-branch" && m.branchState.branchCount <= 1)
916
+ if (autohideFloat === "always" || autohideFloat === "single-branch" && m.message.branchCount <= 1)
926
917
  return "floating" /* Floating */;
927
918
  return "normal" /* Normal */;
928
919
  }
@@ -960,7 +951,7 @@ var ActionBarCopy = createActionButton(useCopyMessage);
960
951
 
961
952
  // src/actions/useReloadMessage.tsx
962
953
  var useReloadMessage = () => {
963
- const { useThread, useBranchObserver } = useAssistantContext();
954
+ const { useThread } = useAssistantContext();
964
955
  const { useMessage } = useMessageContext();
965
956
  const disabled = useCombinedStore(
966
957
  [useThread, useMessage],
@@ -972,7 +963,7 @@ var useReloadMessage = () => {
972
963
  const message = useMessage.getState().message;
973
964
  if (message.role !== "assistant")
974
965
  throw new Error("Reloading is only supported on assistant messages");
975
- useBranchObserver.getState().reloadAt(message);
966
+ useThread.getState().reload(message.id);
976
967
  };
977
968
  };
978
969
 
@@ -997,10 +988,10 @@ var useBeginMessageEdit = () => {
997
988
  // src/primitives/actionBar/ActionBarEdit.tsx
998
989
  var ActionBarEdit = createActionButton(useBeginMessageEdit);
999
990
 
1000
- // src/vercel/VercelAIAssistantProvider.tsx
991
+ // src/adapters/vercel/VercelAIAssistantProvider.tsx
1001
992
  var import_react22 = require("react");
1002
993
 
1003
- // src/vercel/useDummyAIAssistantContext.tsx
994
+ // src/adapters/vercel/useDummyAIAssistantContext.tsx
1004
995
  var import_react21 = require("react");
1005
996
  var import_zustand2 = require("zustand");
1006
997
  var useDummyAIAssistantContext = () => {
@@ -1015,6 +1006,12 @@ var useDummyAIAssistantContext = () => {
1015
1006
  stop: () => {
1016
1007
  throw new Error("Not implemented");
1017
1008
  },
1009
+ switchToBranch: () => {
1010
+ throw new Error("Not implemented");
1011
+ },
1012
+ reload: async () => {
1013
+ throw new Error("Not implemented");
1014
+ },
1018
1015
  isAtBottom: true,
1019
1016
  scrollToBottom: () => {
1020
1017
  for (const listener of scrollToBottomListeners) {
@@ -1040,7 +1037,7 @@ var useDummyAIAssistantContext = () => {
1040
1037
  },
1041
1038
  send: () => {
1042
1039
  useThread.getState().append({
1043
- role: "user",
1040
+ parentId: useThread.getState().messages.at(-1)?.id ?? ROOT_PARENT_ID,
1044
1041
  content: [{ type: "text", text: useComposer.getState().value }]
1045
1042
  });
1046
1043
  useComposer.getState().setValue("");
@@ -1049,43 +1046,39 @@ var useDummyAIAssistantContext = () => {
1049
1046
  useThread.getState().stop();
1050
1047
  }
1051
1048
  }));
1052
- const useBranchObserver = (0, import_zustand2.create)()(() => ({
1053
- getBranchState: () => ({
1054
- branchId: 0,
1055
- branchCount: 1
1056
- }),
1057
- switchToBranch: () => {
1058
- throw new Error("Not implemented");
1059
- },
1060
- editAt: async () => {
1061
- throw new Error("Not implemented");
1062
- },
1063
- reloadAt: async () => {
1064
- throw new Error("Not implemented");
1065
- }
1066
- }));
1067
- return { useThread, useComposer, useBranchObserver };
1049
+ return { useThread, useComposer };
1068
1050
  });
1069
1051
  return context;
1070
1052
  };
1071
1053
 
1072
- // src/vercel/VercelAIAssistantProvider.tsx
1054
+ // src/adapters/vercel/VercelAIAssistantProvider.tsx
1073
1055
  var ThreadMessageCache = /* @__PURE__ */ new WeakMap();
1074
- var vercelToThreadMessage = (message) => {
1056
+ var vercelToThreadMessage = (message, parentId, branchId, branchCount) => {
1075
1057
  if (message.role !== "user" && message.role !== "assistant")
1076
1058
  throw new Error("Unsupported role");
1077
1059
  return {
1060
+ parentId,
1078
1061
  id: message.id,
1079
1062
  role: message.role,
1080
- content: [{ type: "text", text: message.content }]
1063
+ content: [{ type: "text", text: message.content }],
1064
+ branchId,
1065
+ branchCount,
1066
+ createdAt: message.createdAt ?? /* @__PURE__ */ new Date()
1081
1067
  };
1082
1068
  };
1083
- var vercelToCachedThreadMessages = (messages) => {
1084
- return messages.map((m) => {
1069
+ var vercelToCachedThreadMessages = (messages, getBranchState) => {
1070
+ return messages.map((m, idx) => {
1085
1071
  const cached = ThreadMessageCache.get(m);
1086
- if (cached)
1072
+ const parentId = messages[idx - 1]?.id ?? ROOT_PARENT_ID;
1073
+ const { branchId, branchCount } = getBranchState(m.id);
1074
+ if (cached && cached.parentId === parentId && cached.branchId === branchId && cached.branchCount === branchCount)
1087
1075
  return cached;
1088
- const newMessage = vercelToThreadMessage(m);
1076
+ const newMessage = vercelToThreadMessage(
1077
+ m,
1078
+ parentId,
1079
+ branchId,
1080
+ branchCount
1081
+ );
1089
1082
  ThreadMessageCache.set(m, newMessage);
1090
1083
  return newMessage;
1091
1084
  });
@@ -1096,22 +1089,13 @@ var VercelAIAssistantProvider = ({
1096
1089
  }) => {
1097
1090
  const context = useDummyAIAssistantContext();
1098
1091
  const vercel = "chat" in rest ? rest.chat : rest.assistant;
1092
+ const branches = useVercelAIBranches(vercel, context);
1099
1093
  const messages = (0, import_react22.useMemo)(() => {
1100
- return vercelToCachedThreadMessages(vercel.messages);
1101
- }, [vercel.messages]);
1102
- const append = (0, import_react22.useCallback)(
1103
- async (message) => {
1104
- if (message.content[0]?.type !== "text") {
1105
- throw new Error("Only text content is currently supported");
1106
- }
1107
- context.useThread.getState().scrollToBottom();
1108
- await vercel.append({
1109
- role: message.role,
1110
- content: message.content[0].text
1111
- });
1112
- },
1113
- [context, vercel.append]
1114
- );
1094
+ return vercelToCachedThreadMessages(
1095
+ vercel.messages,
1096
+ branches.getBranchState
1097
+ );
1098
+ }, [vercel.messages, branches.getBranchState]);
1115
1099
  const stop = (0, import_react22.useCallback)(() => {
1116
1100
  const lastMessage = vercel.messages.at(-1);
1117
1101
  vercel.stop();
@@ -1124,10 +1108,12 @@ var VercelAIAssistantProvider = ({
1124
1108
  context.useThread.setState({
1125
1109
  messages,
1126
1110
  isLoading,
1127
- append,
1128
- stop
1111
+ stop,
1112
+ switchToBranch: branches.switchToBranch,
1113
+ append: branches.append,
1114
+ reload: branches.reload
1129
1115
  });
1130
- }, [context, messages, append, stop, isLoading]);
1116
+ }, [context, messages, isLoading, stop, branches]);
1131
1117
  (0, import_react22.useMemo)(() => {
1132
1118
  context.useComposer.setState({
1133
1119
  canCancel: isLoading,
@@ -1135,31 +1121,32 @@ var VercelAIAssistantProvider = ({
1135
1121
  setValue: vercel.setInput
1136
1122
  });
1137
1123
  }, [context, isLoading, vercel.input, vercel.setInput]);
1138
- const branches = useVercelAIBranches(vercel, context);
1139
- (0, import_react22.useMemo)(() => {
1140
- context.useBranchObserver.setState(branches, true);
1141
- }, [context, branches]);
1142
1124
  return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
1143
1125
  };
1144
1126
 
1145
- // src/vercel/VercelRSCAssistantProvider.tsx
1127
+ // src/adapters/vercel/VercelRSCAssistantProvider.tsx
1146
1128
  var import_react23 = require("react");
1147
1129
  var ThreadMessageCache2 = /* @__PURE__ */ new WeakMap();
1148
- var vercelToThreadMessage2 = (message) => {
1130
+ var vercelToThreadMessage2 = (parentId, message) => {
1149
1131
  if (message.role !== "user" && message.role !== "assistant")
1150
1132
  throw new Error("Unsupported role");
1151
1133
  return {
1134
+ parentId,
1152
1135
  id: message.id,
1153
1136
  role: message.role,
1154
- content: [{ type: "ui", display: message.display }]
1137
+ content: [{ type: "ui", display: message.display }],
1138
+ createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
1139
+ branchId: 0,
1140
+ branchCount: 1
1155
1141
  };
1156
1142
  };
1157
1143
  var vercelToCachedThreadMessages2 = (messages) => {
1158
- return messages.map((m) => {
1144
+ return messages.map((m, idx) => {
1159
1145
  const cached = ThreadMessageCache2.get(m);
1160
- if (cached)
1146
+ const parentId = messages[idx - 1]?.id ?? ROOT_PARENT_ID;
1147
+ if (cached && cached.parentId === parentId)
1161
1148
  return cached;
1162
- const newMessage = vercelToThreadMessage2(m);
1149
+ const newMessage = vercelToThreadMessage2(parentId, m);
1163
1150
  ThreadMessageCache2.set(m, newMessage);
1164
1151
  return newMessage;
1165
1152
  });
@@ -1175,14 +1162,13 @@ var VercelRSCAssistantProvider = ({
1175
1162
  }, [vercelMessages]);
1176
1163
  const append = (0, import_react23.useCallback)(
1177
1164
  async (message) => {
1165
+ if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? ROOT_PARENT_ID))
1166
+ throw new Error("Unexpected: Message editing is not supported");
1178
1167
  if (message.content[0]?.type !== "text") {
1179
1168
  throw new Error("Only text content is currently supported");
1180
1169
  }
1181
1170
  context.useThread.getState().scrollToBottom();
1182
- await vercelAppend({
1183
- role: message.role,
1184
- content: [{ type: "text", text: message.content[0].text }]
1185
- });
1171
+ await vercelAppend(message);
1186
1172
  },
1187
1173
  [context, vercelAppend]
1188
1174
  );