@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.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
  );