@kortyx/react 0.4.1 → 0.5.0

package/dist/index.js

@@ -41,7 +41,11 @@ function buildAssistantMessage(args) {
     role: "assistant",
     content: plainTextContent,
     debug: args.debug,
-    ...args.trace ? args.trace : {}
+    ...args.trace ? args.trace : {},
+    ...args.checkpoint ? {
+      checkpointId: args.checkpoint.id,
+      checkpointTurnIndex: args.checkpoint.turnIndex
+    } : {}
   };
   return args.pieces.length > 0 ? { ...base, contentPieces: args.pieces } : base;
 }
@@ -51,7 +55,9 @@ var DEFAULT_STORAGE_KEYS = {
   messages: "chat.messages.v1",
   includeHistory: "chat.includeHistory.v1",
   sessionId: "chat.sessionId",
-  workflowId: "chat.workflowId"
+  workflowId: "chat.workflowId",
+  branches: "chat.branches.v1",
+  responseVariants: "chat.responseVariants.v1"
 };
 var getLocalStorage = () => {
   const candidate = globalThis;
@@ -78,13 +84,40 @@ function createBrowserChatStorage(options) {
         const workflowId = storage.getItem(keys.workflowId) ?? "";
         const includeHistoryRaw = storage.getItem(keys.includeHistory);
         const rawMessages = storage.getItem(keys.messages);
+        const rawBranches = storage.getItem(keys.branches);
+        const rawResponseVariants = storage.getItem(keys.responseVariants);
         const parsedMessages = rawMessages ? JSON.parse(rawMessages) : [];
         const messages = Array.isArray(parsedMessages) ? parsedMessages.map((item) => parseMessage(item)).filter((message) => message !== null) : [];
+        const parsedBranches = rawBranches ? JSON.parse(rawBranches) : void 0;
+        const branches = parsedBranches && typeof parsedBranches === "object" ? Object.fromEntries(
+          Object.entries(
+            parsedBranches
+          ).flatMap(([sessionId2, branch]) => {
+            if (typeof branch.sessionId !== "string") return [];
+            const rawBranchMessages = Array.isArray(branch.messages) ? branch.messages : [];
+            const branchMessages = rawBranchMessages.map((item) => parseMessage(item)).filter((message) => message !== null).slice(-maxMessages);
+            return [
+              [
+                sessionId2,
+                {
+                  sessionId: branch.sessionId,
+                  ...typeof branch.parentSessionId === "string" ? { parentSessionId: branch.parentSessionId } : {},
+                  ...typeof branch.forkedFrom === "string" ? { forkedFrom: branch.forkedFrom } : {},
+                  messages: branchMessages,
+                  ...Array.isArray(branch.checkpoints) ? { checkpoints: branch.checkpoints } : {}
+                }
+              ]
+            ];
+          })
+        ) : void 0;
+        const parsedResponseVariants = rawResponseVariants ? JSON.parse(rawResponseVariants) : void 0;
         return {
           sessionId,
           workflowId,
           includeHistory: includeHistoryRaw !== "0",
-          messages
+          messages,
+          ...branches ? { branches } : {},
+          ...Array.isArray(parsedResponseVariants) ? { responseVariants: parsedResponseVariants } : {}
         };
       } catch {
         return {};
@@ -103,6 +136,28 @@ function createBrowserChatStorage(options) {
         storage.setItem(keys.includeHistory, state.includeHistory ? "1" : "0");
         const trimmedMessages = state.messages.slice(-maxMessages).map((message) => serializeMessage(message));
         storage.setItem(keys.messages, JSON.stringify(trimmedMessages));
+        if (state.branches) {
+          const serializedBranches = Object.fromEntries(
+            Object.entries(state.branches).map(([sessionId, branch]) => [
+              sessionId,
+              {
+                ...branch,
+                messages: branch.messages.slice(-maxMessages).map((message) => serializeMessage(message))
+              }
+            ])
+          );
+          storage.setItem(keys.branches, JSON.stringify(serializedBranches));
+        } else {
+          storage.removeItem(keys.branches);
+        }
+        if (state.responseVariants) {
+          storage.setItem(
+            keys.responseVariants,
+            JSON.stringify(state.responseVariants)
+          );
+        } else {
+          storage.removeItem(keys.responseVariants);
+        }
       } catch {
       }
     },
@@ -111,6 +166,8 @@ function createBrowserChatStorage(options) {
         const storage = getLocalStorage();
         if (!storage) return;
         storage.removeItem(keys.messages);
+        storage.removeItem(keys.branches);
+        storage.removeItem(keys.responseVariants);
       } catch {
       }
     }
@@ -118,6 +175,7 @@ function createBrowserChatStorage(options) {
 }
 
 // src/chat-transport.ts
+var import_errors = require("@kortyx/core/errors");
 var import_browser = require("@kortyx/stream/browser");
 function createChatTransport(args) {
   return {
@@ -132,7 +190,7 @@ function createChatTransport(args) {
   };
 }
 function createRouteChatTransport(args) {
-  return createChatTransport({
+  const transport = createChatTransport({
     stream: (context) => {
       const body = args.createBody?.(context) ?? {
         sessionId: context.sessionId,
@@ -153,6 +211,47 @@ function createRouteChatTransport(args) {
       });
     }
   });
+  if (!args.checkpointEndpoint) return transport;
+  const fetchJson = async (body) => {
+    const fetchImpl = args.fetchImpl ?? fetch;
+    const response = await fetchImpl(args.checkpointEndpoint, {
+      method: "POST",
+      headers: {
+        "content-type": "application/json",
+        ...args.headers ?? {}
+      },
+      body: JSON.stringify(body)
+    });
+    if (!response.ok) {
+      let failure;
+      let message = `Checkpoint request failed with status ${response.status}`;
+      try {
+        const parsed = await response.json();
+        failure = parsed.failure;
+        if (typeof parsed.error === "string") message = parsed.error;
+      } catch {
+      }
+      throw (0, import_errors.isFailureDescriptor)(failure) ? (0, import_errors.errorFromFailure)(failure) : new import_errors.KortyxError("CHECKPOINT_REQUEST_FAILED", message, {
+        category: "transport",
+        status: response.status,
+        safeMessage: message
+      });
+    }
+    return await response.json();
+  };
+  return {
+    ...transport,
+    listCheckpoints: (sessionId) => fetchJson({ action: "list", sessionId }),
+    rollbackTo: (checkpointId) => fetchJson({
+      action: "rollback",
+      checkpointId
+    }),
+    fork: (checkpointId, options) => fetchJson({
+      action: "fork",
+      checkpointId,
+      ...options?.newSessionId ? { newSessionId: options.newSessionId } : {}
+    })
+  };
 }
 
 // src/create-live-chat-pieces.ts
@@ -262,7 +361,8 @@ function createLiveChatPieces(args) {
       pushPiece({
         id: args.createId(),
         type: "error",
-        content: chunk.message ?? "An error occurred"
+        content: chunk.message ?? "An error occurred",
+        ...chunk.failure ? { failure: chunk.failure } : {}
       });
       return true;
     }
@@ -325,6 +425,7 @@ function toHumanInputPiece(args) {
 }
 
 // src/use-chat.ts
+var import_errors2 = require("@kortyx/core/errors");
 var import_react3 = require("react");
 
 // src/use-chat-stream-debug.ts
@@ -437,6 +538,17 @@ function useStructuredStreams(options) {
   const get = (0, import_react2.useCallback)((streamId) => {
     return itemsRef.current.find((item) => item.streamId === streamId);
   }, []);
+  const deleteStream = (0, import_react2.useCallback)(
+    (streamId) => {
+      const deleted = accumulatorRef.current.delete(streamId);
+      if (!deleted) return false;
+      setNextItems(
+        itemsRef.current.filter((item) => item.streamId !== streamId)
+      );
+      return true;
+    },
+    [setNextItems]
+  );
   const byStreamId = (0, import_react2.useMemo)(
     () => Object.fromEntries(
       items.map((item) => [item.streamId, item.state])
@@ -448,7 +560,8 @@ function useStructuredStreams(options) {
     byStreamId,
     applyStreamChunk,
     clear,
-    get
+    get,
+    delete: deleteStream
   };
 }
 
@@ -481,6 +594,62 @@ function findActiveTextInterrupt(args) {
     (piece) => piece.type === "interrupt" && piece.kind === "text"
   );
 }
+function trimMessagesToCheckpoint(args) {
+  const hasCheckpointMetadata = args.messages.some(
+    (message) => typeof message.checkpointTurnIndex === "number"
+  );
+  if (!hasCheckpointMetadata && args.turnIndex > 0) return args.messages;
+  const keepThroughIndex = args.messages.reduce((lastIndex, message, index) => {
+    return typeof message.checkpointTurnIndex === "number" && message.checkpointTurnIndex <= args.turnIndex ? index : lastIndex;
+  }, -1);
+  return keepThroughIndex >= 0 ? args.messages.slice(0, keepThroughIndex + 1) : [];
+}
+function applyForkedPendingRequests(args) {
+  if (args.pendingRequests.length === 0) return args.messages;
+  const interruptPositions = args.messages.flatMap(
+    (message, messageIndex) => (message.contentPieces ?? []).flatMap(
+      (piece, pieceIndex) => piece.type === "interrupt" ? [{ messageIndex, pieceIndex }] : []
+    )
+  );
+  const positionsToPatch = interruptPositions.slice(
+    -args.pendingRequests.length
+  );
+  if (positionsToPatch.length === 0) return args.messages;
+  return args.messages.map((message, messageIndex) => {
+    if (!message.contentPieces) return message;
+    let changed = false;
+    const contentPieces = message.contentPieces.map((piece, pieceIndex) => {
+      if (piece.type !== "interrupt") return piece;
+      const patchIndex = positionsToPatch.findIndex(
+        (position) => position.messageIndex === messageIndex && position.pieceIndex === pieceIndex
+      );
+      const pending = args.pendingRequests[patchIndex];
+      if (!pending) return piece;
+      changed = true;
+      return {
+        ...piece,
+        resumeToken: pending.token,
+        requestId: pending.requestId,
+        ...pending.schema?.kind ? { kind: pending.schema.kind } : {},
+        ...typeof pending.schema?.multiple === "boolean" ? { multiple: pending.schema.multiple } : {},
+        ...typeof pending.schema?.question === "string" ? { question: pending.schema.question } : {},
+        ...typeof pending.schema?.schemaId === "string" ? { schemaId: pending.schema.schemaId } : {},
+        ...typeof pending.schema?.schemaVersion === "string" ? { schemaVersion: pending.schema.schemaVersion } : {},
+        ...typeof pending.schema?.id === "string" ? { interruptId: pending.schema.id } : {},
+        ...pending.schema?.meta ? { meta: pending.schema.meta } : {},
+        ...pending.options ? { options: pending.options } : {}
+      };
+    });
+    return changed ? { ...message, contentPieces } : message;
+  });
+}
+function latestAssistantMessageId(messages) {
+  for (let index = messages.length - 1; index >= 0; index--) {
+    const message = messages[index];
+    if (message?.role === "assistant") return message.id;
+  }
+  return null;
+}
 function useChat(options) {
   const resolvedStorage = options.storage ?? defaultStorage;
   const transportRef = useLatestRef(options.transport);
@@ -495,13 +664,20 @@ function useChat(options) {
   const [error, setError] = (0, import_react3.useState)(null);
   const [streamContentPieces, setStreamContentPieces] = (0, import_react3.useState)([]);
   const [lastAssistantId, setLastAssistantId] = (0, import_react3.useState)(null);
+  const [checkpoints, setCheckpoints] = (0, import_react3.useState)([]);
   const [includeHistory, setIncludeHistory] = (0, import_react3.useState)(true);
   const [workflowId, setWorkflowId] = (0, import_react3.useState)("");
   const [sessionId, setSessionId] = (0, import_react3.useState)(null);
+  const [branches, setBranches] = (0, import_react3.useState)({});
+  const [responseVariants, setResponseVariants] = (0, import_react3.useState)([]);
   const [didHydrateStorage, setDidHydrateStorage] = (0, import_react3.useState)(false);
   const abortControllerRef = (0, import_react3.useRef)(null);
   const { streamDebug, clearStreamDebug, createRecorder } = useChatStreamDebug();
-  const { applyStreamChunk, clear: clearStructuredStreams } = useStructuredStreams({
+  const {
+    applyStreamChunk,
+    clear: clearStructuredStreams,
+    delete: deleteStructuredStream
+  } = useStructuredStreams({
     createId
   });
   (0, import_react3.useEffect)(() => {
@@ -521,12 +697,28 @@ function useChat(options) {
       if (Array.isArray(storedState.messages) && storedState.messages.length > 0) {
         setMessages(storedState.messages);
       }
+      if (storedState.branches && typeof storedState.branches === "object") {
+        setBranches(
+          storedState.branches
+        );
+      }
+      if (Array.isArray(storedState.responseVariants)) {
+        setResponseVariants(
+          storedState.responseVariants
+        );
+      }
       setDidHydrateStorage(true);
     });
     return () => {
       cancelled = true;
     };
   }, [resolvedStorage]);
+  (0, import_react3.useEffect)(() => {
+    const nextLastAssistantId = latestAssistantMessageId(messages);
+    if (lastAssistantId !== nextLastAssistantId) {
+      setLastAssistantId(nextLastAssistantId);
+    }
+  }, [lastAssistantId, messages]);
   (0, import_react3.useEffect)(() => {
     if (!didHydrateStorage) return;
     const storage = storageRef.current;
@@ -534,21 +726,64 @@ function useChat(options) {
       sessionId,
       workflowId,
       includeHistory,
-      messages
+      messages,
+      branches,
+      responseVariants
     });
   }, [
+    branches,
     didHydrateStorage,
     includeHistory,
     messages,
+    responseVariants,
     sessionId,
     storageRef,
     workflowId
   ]);
+  (0, import_react3.useEffect)(() => {
+    if (!didHydrateStorage || !sessionId) return;
+    setBranches((current) => ({
+      ...current,
+      [sessionId]: {
+        ...current[sessionId] ?? {},
+        sessionId,
+        messages,
+        checkpoints
+      }
+    }));
+  }, [checkpoints, didHydrateStorage, messages, sessionId]);
+  (0, import_react3.useEffect)(() => {
+    if (!didHydrateStorage || !sessionId) return;
+    setResponseVariants(
+      (current) => current.map((group) => {
+        const selectedVariant = group.variants.find(
+          (variant) => variant.id === group.selectedVariantId && variant.sessionId === sessionId
+        );
+        if (!selectedVariant) return group;
+        if (selectedVariant.messages === messages && selectedVariant.checkpoints === checkpoints) {
+          return group;
+        }
+        return {
+          ...group,
+          variants: group.variants.map(
+            (variant) => variant.id === selectedVariant.id ? {
+              ...variant,
+              messages,
+              checkpoints
+            } : variant
+          )
+        };
+      })
+    );
+  }, [checkpoints, didHydrateStorage, messages, sessionId]);
   const clearActiveStream = () => {
     abortControllerRef.current?.abort();
     abortControllerRef.current = null;
     setIsStreaming(false);
     setMessages([]);
+    setCheckpoints([]);
+    setBranches({});
+    setResponseVariants([]);
     setStreamContentPieces([]);
     clearStreamDebug();
     clearStructuredStreams();
@@ -560,6 +795,7 @@ function useChat(options) {
   };
   const resetSession = () => {
     setSessionId(null);
+    setCheckpoints([]);
   };
   const resetChat = () => {
     clearActiveStream();
@@ -577,7 +813,12 @@ function useChat(options) {
   const persistSessionId = (sid) => {
     setSessionId(sid);
   };
-  const resolveSessionId = () => {
+  const resolveSessionId = (preferredSessionId) => {
+    if (preferredSessionId) {
+      if (sessionId !== preferredSessionId)
+        persistSessionId(preferredSessionId);
+      return preferredSessionId;
+    }
     const sid = sessionId ?? createId();
     if (!sessionId) persistSessionId(sid);
     return sid;
@@ -587,13 +828,14 @@ function useChat(options) {
     content: message.content
   });
   const prepareMessagesToSend = async (args) => {
+    const baseMessages = args.baseMessages ?? messages;
     const prepared = options.prepareContextMessages ? await options.prepareContextMessages({
-      messages,
+      messages: baseMessages,
       sessionId: args.sid,
       workflowId,
       reason: args.reason,
       context: requestContext
-    }) : includeHistory ? messages.map(toOutgoingHistoryMessage) : [];
+    }) : includeHistory ? baseMessages.map(toOutgoingHistoryMessage) : [];
     return [...prepared, args.outgoingMessage];
   };
   const streamAssistantResponse = async (args) => {
@@ -602,6 +844,7 @@ function useChat(options) {
     const debug = createRecorder(`${args.debugLabel} sessionId=${args.sid}`);
     let currentTrace;
     let completedTrace;
+    let currentCheckpoint;
     const pieces = createLiveChatPieces({
       createId,
       onChange: setStreamContentPieces,
@@ -637,8 +880,35 @@ function useChat(options) {
           }
           return;
         }
+        if (chunk.type === "checkpoint") {
+          currentCheckpoint = {
+            id: chunk.id,
+            turnIndex: chunk.turnIndex
+          };
+          setCheckpoints((prev) => {
+            const next = {
+              id: chunk.id,
+              sessionId: chunk.sessionId,
+              turnIndex: chunk.turnIndex,
+              createdAt: Date.now(),
+              nodes: [],
+              workflow: workflowId,
+              ...chunk.label ? { label: chunk.label } : {}
+            };
+            return [...prev.filter((item) => item.id !== chunk.id), next].sort(
+              (a, b) => a.turnIndex - b.turnIndex
+            );
+          });
+          return;
+        }
+        if (chunk.type === "structured-data-invalidated") {
+          deleteStructuredStream(chunk.streamId);
+          return;
+        }
         if (chunk.type === "error") {
-          setError(new Error(chunk.message));
+          setError(
+            (0, import_errors2.isFailureDescriptor)(chunk.failure) ? (0, import_errors2.errorFromFailure)(chunk.failure) : new Error(chunk.message)
+          );
         }
         if (chunk.type === "done") {
           completedTrace = currentTrace;
@@ -649,36 +919,54 @@ function useChat(options) {
         });
       }
     });
-    if (args.signal.aborted) return;
+    if (args.signal.aborted) return void 0;
     const assistant = buildAssistantMessage({
       createId,
       pieces: pieces.getPieces(),
       debug: debug.getAll(),
-      trace: completedTrace ?? currentTrace
+      trace: completedTrace ?? currentTrace,
+      checkpoint: currentCheckpoint
     });
-    setMessages((prev) => [...prev, assistant]);
+    setMessages(
+      (prev) => args.messagesBeforeAssistant ? [...args.messagesBeforeAssistant, assistant] : [...prev, assistant]
+    );
     setLastAssistantId(assistant.id);
     clearStructuredStreams();
     setStreamContentPieces([]);
     clearStreamDebug();
+    return assistant;
   };
   const isAbortError = (value, signal) => {
     if (signal.aborted) return true;
     if (!(value instanceof Error)) return false;
     return value.name === "AbortError";
   };
-  const respondToHumanInput = async ({
+  const respondToHumanInputFromBase = async ({
     resumeToken,
     requestId,
     selected,
-    text
+    text,
+    baseMessages,
+    sid
   }) => {
-    if (isStreaming) return;
+    if (isStreaming) return void 0;
     setError(null);
     const label = typeof text === "string" && text.length > 0 ? text : selected.length > 0 ? selected.length === 1 ? selected[0] : selected.join(", ") : "(selection)";
     const userId = createId();
-    const userMsg = { id: userId, role: "user", content: label };
-    setMessages((prev) => [...prev, userMsg]);
+    const userMsg = {
+      id: userId,
+      role: "user",
+      content: label,
+      source: {
+        type: "interrupt-response",
+        resumeToken,
+        requestId,
+        selected,
+        ...text !== void 0 ? { text } : {}
+      }
+    };
+    const nextMessages = [...baseMessages ?? messages, userMsg];
+    setMessages(nextMessages);
     setIsStreaming(true);
     setStreamContentPieces([]);
     clearStreamDebug();
@@ -694,18 +982,25 @@ function useChat(options) {
         content: label,
         metadata: resumePayload
       };
-      const sid = resolveSessionId();
+      const resolvedSessionId = resolveSessionId(sid);
       const messagesToSend = await prepareMessagesToSend({
-        sid,
+        sid: resolvedSessionId,
         outgoingMessage: outgoing,
-        reason: "resume"
+        reason: "resume",
+        baseMessages
       });
-      await streamAssistantResponse({
-        sid,
+      const assistantMessage = await streamAssistantResponse({
+        sid: resolvedSessionId,
         messagesToSend,
+        messagesBeforeAssistant: nextMessages,
         debugLabel: "runChat (resume)",
         signal: abortController.signal
       });
+      return {
+        userMessage: userMsg,
+        ...assistantMessage ? { assistantMessage } : {},
+        messages: assistantMessage ? [...nextMessages, assistantMessage] : nextMessages
+      };
     } catch (err) {
       if (isAbortError(err, abortController.signal)) return;
       const nextError = err instanceof Error ? err : new Error(String(err));
@@ -718,36 +1013,45 @@ function useChat(options) {
       setIsStreaming(false);
     }
   };
+  const respondToHumanInput = async (args) => {
+    await respondToHumanInputFromBase(args);
+  };
   const respondToInterrupt = async (piece, response) => {
     const text = response?.text;
     const selected = response?.selected ?? (typeof text === "string" && text.length > 0 ? [text] : []);
-    return respondToHumanInput({
+    await respondToHumanInput({
       resumeToken: piece.resumeToken,
       requestId: piece.requestId,
       selected,
       ...text !== void 0 ? { text } : {}
     });
   };
-  const send = async (text) => {
+  const sendFromBase = async (text, baseMessages, sid) => {
     const content = text.trim();
-    if (!content || isStreaming) return;
+    if (!content || isStreaming) return void 0;
     setError(null);
     const activeTextInterrupt = findActiveTextInterrupt({
-      messages,
+      messages: baseMessages ?? messages,
       streamContentPieces
     });
     if (activeTextInterrupt) {
-      await respondToHumanInput({
+      return respondToHumanInputFromBase({
         resumeToken: activeTextInterrupt.resumeToken,
         requestId: activeTextInterrupt.requestId,
         selected: [content],
-        text: content
+        text: content,
+        baseMessages,
+        sid
       });
-      return;
     }
     const userId = createId();
-    const userMsg = { id: userId, role: "user", content };
-    setMessages((prev) => [...prev, userMsg]);
+    const userMsg = {
+      id: userId,
+      role: "user",
+      content
+    };
+    const nextMessages = [...baseMessages ?? messages, userMsg];
+    setMessages(nextMessages);
     setIsStreaming(true);
     setStreamContentPieces([]);
     clearStreamDebug();
@@ -755,19 +1059,26 @@ function useChat(options) {
     const abortController = new AbortController();
     abortControllerRef.current = abortController;
     try {
-      const sid = resolveSessionId();
+      const resolvedSessionId = resolveSessionId(sid);
       const messagesToSend = await prepareMessagesToSend({
-        sid,
+        sid: resolvedSessionId,
         outgoingMessage: { role: "user", content },
-        reason: "send"
+        reason: "send",
+        baseMessages
       });
-      await streamAssistantResponse({
-        sid,
+      const assistantMessage = await streamAssistantResponse({
+        sid: resolvedSessionId,
         messagesToSend,
+        messagesBeforeAssistant: nextMessages,
         debugLabel: "runChat (send)",
         openDebugOnInterrupt: true,
         signal: abortController.signal
       });
+      return {
+        userMessage: userMsg,
+        ...assistantMessage ? { assistantMessage } : {},
+        messages: assistantMessage ? [...nextMessages, assistantMessage] : nextMessages
+      };
     } catch (err) {
       if (isAbortError(err, abortController.signal)) return;
       const nextError = err instanceof Error ? err : new Error(String(err));
@@ -780,6 +1091,375 @@ function useChat(options) {
       setIsStreaming(false);
     }
   };
+  const send = async (text) => {
+    await sendFromBase(text);
+  };
+  const requireCheckpointTransport = () => {
+    const transport = transportRef.current;
+    if (!transport.listCheckpoints || !transport.rollbackTo || !transport.fork) {
+      throw new Error(
+        "Checkpoint helpers require a transport with checkpoint methods."
+      );
+    }
+    return transport;
+  };
+  const refreshCheckpoints = async (sid) => {
+    const transport = requireCheckpointTransport();
+    const next = await transport.listCheckpoints(sid);
+    setCheckpoints(next);
+    return next;
+  };
+  const applyStructuredInvalidations = (streamIds) => {
+    for (const streamId of streamIds) {
+      deleteStructuredStream(streamId);
+    }
+  };
+  const rollbackTo = async (id) => {
+    if (isStreaming) {
+      throw new Error("Cannot rollback while a stream is active.");
+    }
+    const transport = requireCheckpointTransport();
+    const result = await transport.rollbackTo(id);
+    applyStructuredInvalidations(result.invalidatedStructuredStreamIds);
+    const refreshedCheckpoints = await refreshCheckpoints(result.sessionId);
+    const target = refreshedCheckpoints.find(
+      (checkpoint) => checkpoint.id === result.head
+    );
+    if (target) {
+      setMessages((current) => {
+        return trimMessagesToCheckpoint({
+          messages: current,
+          turnIndex: target.turnIndex
+        });
+      });
+    }
+    return result;
+  };
+  const fork = async (id, options2) => {
+    if (isStreaming) {
+      throw new Error("Cannot fork while a stream is active.");
+    }
+    const transport = requireCheckpointTransport();
+    const sid = resolveSessionId();
+    let availableCheckpoints = checkpoints.length > 0 ? checkpoints : await refreshCheckpoints(sid);
+    let target = availableCheckpoints.find(
+      (checkpoint) => checkpoint.id === id
+    );
+    if (!target && checkpoints.length > 0) {
+      availableCheckpoints = await refreshCheckpoints(sid);
+      target = availableCheckpoints.find((checkpoint) => checkpoint.id === id);
+    }
+    const result = await transport.fork(id, options2);
+    setSessionId(result.sessionId);
+    const childCheckpoints = await refreshCheckpoints(result.sessionId);
+    const childHead = childCheckpoints.at(-1);
+    const forkTurnIndex = target?.turnIndex ?? childHead?.turnIndex;
+    if (typeof forkTurnIndex === "number") {
+      setMessages(
+        (current) => applyForkedPendingRequests({
+          messages: trimMessagesToCheckpoint({
+            messages: current,
+            turnIndex: forkTurnIndex
+          }),
+          pendingRequests: result.checkpoint?.activePendingRequests ?? []
+        })
+      );
+    }
+    setStreamContentPieces([]);
+    clearStructuredStreams();
+    return result;
+  };
+  const checkpointForMessage = (messageId) => {
+    return messages.find((message) => message.id === messageId)?.checkpointId ?? null;
+  };
+  const variantForMessage = (messageId) => {
+    return responseVariants.find(
+      (group) => group.sourceAssistantMessageId === messageId || group.variants.some(
+        (variant) => variant.assistantMessageId === messageId
+      )
+    ) ?? null;
+  };
+  const findResumeOverrideForFork = (baseMessages, forkedBaseMessages, userMessage) => {
+    if (userMessage.source?.type !== "interrupt-response") return void 0;
+    const userSource = userMessage.source;
+    const interruptPositions = baseMessages.flatMap(
+      (message, messageIndex) => (message.contentPieces ?? []).flatMap(
+        (piece, pieceIndex) => piece.type === "interrupt" ? [{ messageIndex, pieceIndex, piece }] : []
+      )
+    );
+    const matchedPosition = [...interruptPositions].reverse().find(
+      ({ piece }) => piece.resumeToken === userSource.resumeToken || piece.requestId === userSource.requestId
+    );
+    const forkedPiece = typeof matchedPosition?.messageIndex === "number" ? forkedBaseMessages[matchedPosition.messageIndex]?.contentPieces?.[matchedPosition.pieceIndex] : void 0;
+    if (forkedPiece?.type === "interrupt") {
+      return {
+        resumeToken: forkedPiece.resumeToken,
+        requestId: forkedPiece.requestId
+      };
+    }
+    const latestForkedInterrupt = forkedBaseMessages.flatMap((message) => message.contentPieces ?? []).filter((piece) => piece.type === "interrupt").at(-1);
+    return latestForkedInterrupt ? {
+      resumeToken: latestForkedInterrupt.resumeToken,
+      requestId: latestForkedInterrupt.requestId
+    } : void 0;
+  };
+  const resolveRegenerateTarget = async (assistantMessageId) => {
+    const assistantIndex = messages.findIndex(
+      (message) => message.id === assistantMessageId
+    );
+    if (assistantIndex < 0) {
+      throw new Error(`Assistant message "${assistantMessageId}" not found.`);
+    }
+    const assistant = messages[assistantIndex];
+    if (!assistant) {
+      throw new Error(`Assistant message "${assistantMessageId}" not found.`);
+    }
+    if (assistant.role !== "assistant") {
+      throw new Error("Regenerate expects an assistant message id.");
+    }
+    let userIndex = -1;
+    for (let index = assistantIndex - 1; index >= 0; index--) {
+      if (messages[index]?.role === "user") {
+        userIndex = index;
+        break;
+      }
+    }
+    if (userIndex < 0) {
+      throw new Error("No preceding user message found for regenerate.");
+    }
+    const previousUser = messages[userIndex];
+    const sid = resolveSessionId();
+    let availableCheckpoints = checkpoints.length > 0 ? checkpoints : await refreshCheckpoints(sid);
+    const assistantTurn = assistant.checkpointTurnIndex;
+    const targetTurn = typeof assistantTurn === "number" ? assistantTurn - 1 : void 0;
+    let target = targetTurn !== void 0 ? availableCheckpoints.find(
+      (checkpoint) => checkpoint.turnIndex === targetTurn
+    ) : availableCheckpoints.filter(
+      (checkpoint) => checkpoint.turnIndex < Number.MAX_SAFE_INTEGER
+    ).at(-2);
+    if (!target && checkpoints.length > 0) {
+      availableCheckpoints = await refreshCheckpoints(sid);
+      target = targetTurn !== void 0 ? availableCheckpoints.find(
+        (checkpoint) => checkpoint.turnIndex === targetTurn
+      ) : availableCheckpoints.filter(
+        (checkpoint) => checkpoint.turnIndex < Number.MAX_SAFE_INTEGER
+      ).at(-2);
+    }
+    if (!target) {
+      throw new Error("No rollback checkpoint found for regenerate.");
+    }
+    const baseMessages = trimMessagesToCheckpoint({
+      messages,
+      turnIndex: target.turnIndex
+    });
+    return { target, assistant, previousUser, baseMessages };
+  };
+  const replayUserMessage = async (userMessage, baseMessages, editedContent, options2) => {
+    const content = editedContent ?? userMessage.content;
+    if (userMessage.source?.type === "interrupt-response") {
+      return respondToHumanInputFromBase({
+        resumeToken: options2?.resumeOverride?.resumeToken ?? userMessage.source.resumeToken,
+        requestId: options2?.resumeOverride?.requestId ?? userMessage.source.requestId,
+        selected: editedContent !== void 0 ? [editedContent] : userMessage.source.selected,
+        baseMessages,
+        sid: options2?.sid,
+        ...editedContent !== void 0 ? { text: editedContent } : userMessage.source.text !== void 0 ? { text: userMessage.source.text } : {}
+      });
+    }
+    return sendFromBase(content, baseMessages, options2?.sid);
+  };
+  const buildVariant = (args) => ({
+    id: args.id,
+    sessionId: args.sessionId,
+    assistantMessageId: args.assistantMessage.id,
+    createdAt: Date.now(),
+    messages: args.messages,
+    checkpoints: args.checkpoints,
+    ...args.assistantMessage.checkpointId ? { checkpointId: args.assistantMessage.checkpointId } : {},
+    ...typeof args.assistantMessage.checkpointTurnIndex === "number" ? { checkpointTurnIndex: args.assistantMessage.checkpointTurnIndex } : {},
+    ...args.label ? { label: args.label } : {}
+  });
+  const updateVariantSnapshot = (args) => ({
+    ...args.variant,
+    sessionId: args.sessionId,
+    assistantMessageId: args.assistantMessage.id,
+    messages: args.messages,
+    checkpoints: args.checkpoints,
+    ...args.assistantMessage.checkpointId ? { checkpointId: args.assistantMessage.checkpointId } : {},
+    ...typeof args.assistantMessage.checkpointTurnIndex === "number" ? { checkpointTurnIndex: args.assistantMessage.checkpointTurnIndex } : {}
+  });
+  const replaceResponseVariantGroup = (nextGroup) => {
+    setResponseVariants((current) => [
+      ...current.filter((group) => group.id !== nextGroup.id),
+      nextGroup
+    ]);
+  };
+  const regenerateVariant = async (assistantMessageId) => {
+    if (isStreaming) {
+      throw new Error("Cannot regenerate a variant while a stream is active.");
+    }
+    const { target, assistant, previousUser, baseMessages } = await resolveRegenerateTarget(assistantMessageId);
+    const transport = requireCheckpointTransport();
+    const parentSessionId = resolveSessionId();
+    const existingGroup = variantForMessage(assistantMessageId);
+    const sourceAssistantMessageId = existingGroup?.sourceAssistantMessageId ?? assistantMessageId;
+    const existingVariants = existingGroup ? existingGroup.variants.map(
+      (variant) => variant.id === existingGroup.selectedVariantId || variant.assistantMessageId === assistantMessageId ? updateVariantSnapshot({
+        variant,
+        sessionId: parentSessionId,
+        assistantMessage: assistant,
+        messages,
+        checkpoints
+      }) : variant
+    ) : void 0;
+    const originalVariant = existingVariants?.[0] ?? buildVariant({
+      id: `${sourceAssistantMessageId}:original`,
+      sessionId: parentSessionId,
+      assistantMessage: assistant,
+      messages,
+      checkpoints,
+      label: "Original"
+    });
+    const forkResult = await transport.fork(target.id);
+    const forkedSessionId = forkResult.sessionId;
+    setSessionId(forkedSessionId);
+    const childCheckpointsBeforeReplay = await refreshCheckpoints(forkedSessionId);
+    const forkedBaseMessages = applyForkedPendingRequests({
+      messages: baseMessages,
+      pendingRequests: forkResult.checkpoint?.activePendingRequests ?? []
+    });
+    const resumeOverride = findResumeOverrideForFork(
+      baseMessages,
+      forkedBaseMessages,
+      previousUser
+    );
+    setMessages(forkedBaseMessages);
+    setStreamContentPieces([]);
+    clearStructuredStreams();
+    const replay = await replayUserMessage(
+      previousUser,
+      forkedBaseMessages,
+      void 0,
+      {
+        sid: forkedSessionId,
+        resumeOverride
+      }
+    );
+    const childCheckpointsAfterReplay = await refreshCheckpoints(forkedSessionId);
+    const variantMessages = replay?.messages ?? forkedBaseMessages;
+    const variantAssistant = replay?.assistantMessage;
+    if (!variantAssistant) {
+      throw new Error("Regenerate variant did not produce an assistant reply.");
+    }
+    const nextVariant = buildVariant({
+      id: createId(),
+      sessionId: forkedSessionId,
+      assistantMessage: variantAssistant,
+      messages: variantMessages,
+      checkpoints: childCheckpointsAfterReplay.length > 0 ? childCheckpointsAfterReplay : childCheckpointsBeforeReplay,
+      label: `Variant ${(existingGroup?.variants.length ?? 1) + 1}`
+    });
+    const nextVariants = existingVariants ? [...existingVariants, nextVariant] : [originalVariant, nextVariant];
+    const nextGroup = {
+      id: existingGroup?.id ?? createId(),
+      sourceAssistantMessageId,
+      checkpointId: existingGroup?.checkpointId ?? target.id,
+      checkpointTurnIndex: existingGroup?.checkpointTurnIndex ?? target.turnIndex,
+      selectedVariantId: nextVariant.id,
+      variants: nextVariants
+    };
+    replaceResponseVariantGroup(nextGroup);
+    setBranches((current) => ({
+      ...current,
+      [parentSessionId]: {
+        ...current[parentSessionId] ?? {},
+        sessionId: parentSessionId,
+        messages,
+        checkpoints
+      },
+      [forkedSessionId]: {
+        ...current[forkedSessionId] ?? {},
+        sessionId: forkedSessionId,
+        parentSessionId: forkResult.parentSessionId,
+        forkedFrom: forkResult.forkedFrom,
+        messages: variantMessages,
+        checkpoints: nextVariant.checkpoints
+      }
+    }));
+    return nextGroup;
+  };
+  const selectVariant = async (assistantMessageId, variantId) => {
+    const group = variantForMessage(assistantMessageId);
+    if (!group) {
+      throw new Error(
+        `No response variants found for "${assistantMessageId}".`
+      );
+    }
+    const selectedVariant = group.variants.find(
+      (variant) => variant.id === variantId
+    );
+    if (!selectedVariant) {
+      throw new Error(`Response variant "${variantId}" not found.`);
+    }
+    setSessionId(selectedVariant.sessionId);
+    setMessages(selectedVariant.messages);
+    setCheckpoints(selectedVariant.checkpoints);
+    setStreamContentPieces([]);
+    clearStructuredStreams();
+    setBranches((current) => ({
+      ...current,
+      [selectedVariant.sessionId]: {
+        ...current[selectedVariant.sessionId] ?? {},
+        sessionId: selectedVariant.sessionId,
+        messages: selectedVariant.messages,
+        checkpoints: selectedVariant.checkpoints
+      }
+    }));
+    const nextGroup = {
+      ...group,
+      selectedVariantId: selectedVariant.id
+    };
+    replaceResponseVariantGroup(nextGroup);
+  };
+  const regenerate = async (assistantMessageId) => {
+    const { target, previousUser, baseMessages } = await resolveRegenerateTarget(assistantMessageId);
+    await rollbackTo(target.id);
+    setMessages(baseMessages);
+    await replayUserMessage(previousUser, baseMessages);
+  };
+  const regenerateFromCheckpoint = async (checkpointId) => {
+    const sid = resolveSessionId();
+    let availableCheckpoints = checkpoints.length > 0 ? checkpoints : await refreshCheckpoints(sid);
+    let target = availableCheckpoints.find(
+      (checkpoint) => checkpoint.id === checkpointId
+    );
+    if (!target && checkpoints.length > 0) {
+      availableCheckpoints = await refreshCheckpoints(sid);
+      target = availableCheckpoints.find(
+        (checkpoint) => checkpoint.id === checkpointId
+      );
+    }
+    if (!target) {
+      throw new Error(`Checkpoint "${checkpointId}" not found.`);
+    }
+    const baseMessages = trimMessagesToCheckpoint({
+      messages,
+      turnIndex: target.turnIndex
+    });
+    const previousUser = messages.slice(baseMessages.length).find((message) => message.role === "user");
+    if (!previousUser) {
+      throw new Error("No user message found after checkpoint for regenerate.");
+    }
+    await rollbackTo(target.id);
+    setMessages(baseMessages);
+    await replayUserMessage(previousUser, baseMessages);
+  };
+  const retryWithEdit = async (assistantMessageId, newUserContent) => {
+    const { target, previousUser, baseMessages } = await resolveRegenerateTarget(assistantMessageId);
+    await rollbackTo(target.id);
+    setMessages(baseMessages);
+    await replayUserMessage(previousUser, baseMessages, newUserContent);
+  };
   return {
     messages,
     isStreaming,
@@ -788,12 +1468,25 @@ function useChat(options) {
     streamContentPieces,
     streamDebug,
     lastAssistantId,
+    checkpoints,
+    activeSessionId: sessionId,
+    branches,
+    responseVariants,
     workflowId,
     setWorkflowId,
     send,
     respondToHumanInput,
     respondToInterrupt,
     setMessages,
+    checkpointForMessage,
+    rollbackTo,
+    fork,
+    regenerate,
+    regenerateVariant,
+    selectVariant,
+    variantForMessage,
+    regenerateFromCheckpoint,
+    retryWithEdit,
     clearMessages,
     resetSession,
     resetChat,