OrgTJWater/TJWaterFrontend_Refine
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: simplify chat fork flow
Build Push and Deploy / docker-image (push) Successful in 1m29s
Details
Build Push and Deploy / deploy-fallback-log (push) Has been skipped
Details

Browse Source
This commit is contained in:
Huarch
2026-06-08 16:07:39 +08:00
parent 34fd5bfb1a
commit 2691f42581
10 changed files with 274 additions and 656 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/components/chat/hooks/useAgentChatSession.test.tsx
+92 -3
View File
@@ -5,6 +5,7 @@ import { act, renderHook, waitFor } from "@testing-library/react";
import { useAgentChatSession } from "./useAgentChatSession";
import {
abortAgentChat,
forkAgentChat,
replyAgentPermission,
resumeAgentChatStream,
streamAgentChat,
@@ -30,7 +31,6 @@ jest.mock("../chatStorage", () => ({
isTitleManuallyEdited: false,
messages: [],
sessionId: undefined,
branchGroups: [],
})),
deleteChatSession: (...args: unknown[]) => deleteChatSession(...args),
listChatSessions: (...args: unknown[]) => listChatSessions(...args),
@@ -39,7 +39,6 @@ jest.mock("../chatStorage", () => ({
isTitleManuallyEdited: false,
messages: [],
sessionId: "session-loaded",
branchGroups: [],
})),
saveActiveChatState: (...args: unknown[]) => saveActiveChatState(...args),
updateChatSessionTitle: (...args: unknown[]) => updateChatSessionTitle(...args),
@@ -52,10 +51,12 @@ describe("useAgentChatSession", () => {
saveActiveChatState.mockReset();
updateChatSessionTitle.mockReset();
jest.mocked(abortAgentChat).mockReset();
jest.mocked(forkAgentChat).mockReset();
jest.mocked(replyAgentPermission).mockReset();
jest.mocked(resumeAgentChatStream).mockReset();
jest.mocked(streamAgentChat).mockReset();
jest.mocked(abortAgentChat).mockImplementation(async () => undefined);
jest.mocked(forkAgentChat).mockImplementation(async () => "forked-session");
jest.mocked(replyAgentPermission).mockImplementation(async () => undefined);
jest.mocked(resumeAgentChatStream).mockImplementation(async () => undefined);
jest.mocked(streamAgentChat).mockImplementation(async () => undefined);
@@ -596,9 +597,10 @@ describe("useAgentChatSession", () => {
await act(async () => {
await result.current.sendPrompt("重新分析压力异常");
});
const assistantMessageId = result.current.messages[1]?.id ?? "";
await act(async () => {
await result.current.regenerate();
await result.current.regenerate(assistantMessageId);
});
expect(streamAgentChat).toHaveBeenNthCalledWith(
@@ -609,4 +611,91 @@ describe("useAgentChatSession", () => {
}),
);
});
it("replaces the current chain when regenerating a middle assistant message", async () => {
listChatSessions.mockResolvedValue([]);
const { result } = renderHook(() =>
useAgentChatSession({
projectId: "project-1",
onToolCall: jest.fn(),
}),
);
await waitFor(() => expect(result.current.isHydrating).toBe(false));
await act(async () => {
await result.current.sendPrompt("第一轮");
});
await act(async () => {
await result.current.sendPrompt("第二轮");
});
const firstAssistantMessageId = result.current.messages[1]?.id ?? "";
await act(async () => {
await result.current.regenerate(firstAssistantMessageId);
});
expect(result.current.messages).toHaveLength(2);
expect(result.current.messages[0]).toEqual(
expect.objectContaining({
role: "user",
content: "第一轮",
}),
);
expect(result.current.messages[1]).toEqual(
expect.objectContaining({
role: "assistant",
content: "",
}),
);
expect(streamAgentChat).toHaveBeenNthCalledWith(
3,
expect.objectContaining({
message: "第一轮",
regenerateFromMessageIndex: 0,
}),
);
});
it("forks a copied conversation from an assistant message", async () => {
listChatSessions.mockResolvedValue([]);
const { result } = renderHook(() =>
useAgentChatSession({
projectId: "project-1",
onToolCall: jest.fn(),
}),
);
await waitFor(() => expect(result.current.isHydrating).toBe(false));
await act(async () => {
await result.current.sendPrompt("第一轮");
});
const firstAssistantMessageId = result.current.messages[1]?.id ?? "";
await act(async () => {
await result.current.createBranch(firstAssistantMessageId);
});
expect(forkAgentChat).toHaveBeenCalledWith(undefined, 2);
expect(result.current.activeSessionId).toBe("forked-session");
expect(result.current.messages).toHaveLength(2);
expect(result.current.messages[0]).toEqual(
expect.objectContaining({
role: "user",
content: "第一轮",
}),
);
expect(result.current.messages[1]).toEqual(
expect.objectContaining({
role: "assistant",
}),
);
expect(streamAgentChat).toHaveBeenCalledTimes(1);
});
});
src/components/chat/hooks/useAgentChatSession.ts
+58 -213
View File
@@ -18,15 +18,12 @@ import type {
import type {
AgentArtifact,
AgentPermissionRequest,
BranchGroup,
BranchTransition,
ChatProgress,
ChatSessionSummary,
LoadedChatState,
Message,
} from "../GlobalChatbox.types";
import {
cloneBranchGroups,
cloneMessages,
createId,
} from "../GlobalChatbox.utils";
@@ -68,7 +65,6 @@ const createPersistedStateKey = (state: LoadedChatState) =>
isTitleManuallyEdited: state.isTitleManuallyEdited ?? false,
sessionId: state.sessionId ?? null,
messages: state.messages,
branchGroups: state.branchGroups,
});
const upsertProgress = (
@@ -193,13 +189,12 @@ const finalizeAssistantMessageAfterAbort = (message: Message): Message => {
};
};
const createUserMessage = (content: string, branchRootId?: string): Message => {
const createUserMessage = (content: string): Message => {
const id = createId();
return {
id,
role: "user",
content,
branchRootId: branchRootId ?? id,
};
};
@@ -209,9 +204,6 @@ const createAssistantMessage = (): Message => ({
content: "",
});
const messagesEqual = (left: Message[], right: Message[]) =>
JSON.stringify(left) === JSON.stringify(right);
export const useAgentChatSession = ({
projectId,
onToolCall,
@@ -226,15 +218,12 @@ export const useAgentChatSession = ({
const [sessionTitle, setSessionTitle] = useState<string | undefined>(undefined);
const [isSessionTitleManuallyEdited, setIsSessionTitleManuallyEdited] = useState(false);
const [sessionId, setSessionId] = useState<string | undefined>(undefined);
const [branchGroups, setBranchGroups] = useState<BranchGroup[]>([]);
const [chatSessions, setChatSessions] = useState<ChatSessionSummary[]>([]);
const [branchTransition, setBranchTransition] = useState<BranchTransition | null>(null);
const [isStreaming, setIsStreaming] = useState(false);
const [isHydrating, setIsHydrating] = useState(true);
const abortRef = useRef<AbortController | null>(null);
const sessionIdRef = useRef<string | undefined>(undefined);
const messagesRef = useRef<Message[]>([]);
const branchGroupsRef = useRef<BranchGroup[]>([]);
const resumeStreamingSessionRef = useRef<((sessionId: string) => void) | null>(null);
const isSessionTitleManuallyEditedRef = useRef(false);
const cancelPromiseRef = useRef<Promise<void> | null>(null);
@@ -245,7 +234,6 @@ export const useAgentChatSession = ({
title: undefined,
isTitleManuallyEdited: false,
messages: [],
branchGroups: [],
}),
);
@@ -257,9 +245,6 @@ export const useAgentChatSession = ({
messagesRef.current = messages;
}, [messages]);
useEffect(() => {
branchGroupsRef.current = branchGroups;
}, [branchGroups]);
useEffect(() => {
isSessionTitleManuallyEditedRef.current = isSessionTitleManuallyEdited;
@@ -279,17 +264,14 @@ export const useAgentChatSession = ({
isTitleManuallyEdited: false,
messages: [],
sessionId: undefined,
branchGroups: [],
});
hydrationCompletedRef.current = true;
hydrationNonceRef.current += 1;
titleUpdateNonceRef.current += 1;
setBranchTransition(null);
setMessages([]);
setSessionTitle(undefined);
setIsSessionTitleManuallyEdited(false);
setSessionId(undefined);
setBranchGroups([]);
setChatSessions([]);
setIsHydrating(false);
return;
@@ -313,7 +295,6 @@ export const useAgentChatSession = ({
setSessionTitle(loadedState.title);
setIsSessionTitleManuallyEdited(loadedState.isTitleManuallyEdited ?? false);
setSessionId(loadedState.sessionId);
setBranchGroups(loadedState.branchGroups);
setChatSessions(sessions);
if (
loadedState.sessionId &&
@@ -351,7 +332,6 @@ export const useAgentChatSession = ({
isTitleManuallyEdited: isSessionTitleManuallyEdited,
messages,
sessionId,
branchGroups,
};
const currentStateKey = createPersistedStateKey(state);
@@ -381,46 +361,7 @@ export const useAgentChatSession = ({
return () => {
window.clearTimeout(persistTimer);
};
}, [branchGroups, isHydrating, isSessionTitleManuallyEdited, isStreaming, messages, projectId, sessionId, sessionTitle]);
useEffect(() => {
setBranchGroups((prev) => {
let changed = false;
const next = prev.map((group) => {
const rootMessage = messages[group.parentCount];
if (
!rootMessage ||
rootMessage.role !== "user" ||
(rootMessage.branchRootId ?? rootMessage.id) !== group.rootMessageId
) {
return group;
}
const activeBranch = group.branches[group.activeIndex];
if (!activeBranch) {
return group;
}
const nextSuffix = cloneMessages(messages.slice(group.parentCount));
if (
activeBranch.sessionId === sessionId &&
messagesEqual(activeBranch.messages, nextSuffix)
) {
return group;
}
changed = true;
const branches = group.branches.map((branch, index) =>
index === group.activeIndex
? { ...branch, sessionId, messages: nextSuffix }
: branch,
);
return { ...group, branches };
});
return changed ? next : prev;
});
}, [messages, sessionId]);
}, [isHydrating, isSessionTitleManuallyEdited, isStreaming, messages, projectId, sessionId, sessionTitle]);
const appendArtifact = useCallback((messageId: string, artifact: AgentArtifact) => {
setMessages((prev) =>
@@ -479,7 +420,6 @@ export const useAgentChatSession = ({
title: nextTitle,
isTitleManuallyEdited: false,
messages: messagesRef.current,
branchGroups: branchGroupsRef.current,
});
}
if (targetSessionId) {
@@ -643,7 +583,6 @@ export const useAgentChatSession = ({
await cancelPromiseRef.current?.catch(() => undefined);
onBeforeSend?.();
setBranchTransition(null);
const nextUserMessage = userMessage ?? createUserMessage(prompt);
const nextAssistantMessage = assistantMessage ?? createAssistantMessage();
@@ -832,7 +771,6 @@ export const useAgentChatSession = ({
const controller = abortRef.current;
controller?.abort();
setBranchTransition(null);
hydrationNonceRef.current += 1;
titleUpdateNonceRef.current += 1;
sessionIdRef.current = undefined;
@@ -841,13 +779,11 @@ export const useAgentChatSession = ({
isTitleManuallyEdited: false,
messages: [],
sessionId: undefined,
branchGroups: [],
});
setMessages([]);
setSessionTitle("新对话");
setIsSessionTitleManuallyEdited(false);
setSessionId(undefined);
setBranchGroups([]);
setIsStreaming(false);
}, [isHydrating, isStreaming]);
@@ -868,12 +804,10 @@ export const useAgentChatSession = ({
titleUpdateNonceRef.current += 1;
sessionIdRef.current = nextState.sessionId;
lastPersistedStateKeyRef.current = createPersistedStateKey(nextState);
setBranchTransition(null);
setMessages(nextState.messages);
setSessionTitle(nextState.title);
setIsSessionTitleManuallyEdited(nextState.isTitleManuallyEdited ?? false);
setSessionId(nextState.sessionId);
setBranchGroups(nextState.branchGroups);
setChatSessions(sessions);
if (nextState.sessionId && nextState.isStreaming) {
resumeStreamingSession(nextState.sessionId);
@@ -917,14 +851,11 @@ export const useAgentChatSession = ({
isTitleManuallyEdited: false,
messages: [],
sessionId: undefined,
branchGroups: [],
});
setBranchTransition(null);
setMessages([]);
setSessionTitle(undefined);
setIsSessionTitleManuallyEdited(false);
setSessionId(undefined);
setBranchGroups([]);
return;
}
@@ -937,12 +868,10 @@ export const useAgentChatSession = ({
titleUpdateNonceRef.current += 1;
sessionIdRef.current = nextState.sessionId;
lastPersistedStateKeyRef.current = createPersistedStateKey(nextState);
setBranchTransition(null);
setMessages(nextState.messages);
setSessionTitle(nextState.title);
setIsSessionTitleManuallyEdited(nextState.isTitleManuallyEdited ?? false);
setSessionId(nextState.sessionId);
setBranchGroups(nextState.branchGroups);
setChatSessions(sessionsAfterDelete);
} catch (error) {
console.error("[GlobalChatbox] Failed to delete chat session:", error);
@@ -985,183 +914,99 @@ export const useAgentChatSession = ({
title: normalizedTitle,
isTitleManuallyEdited: true,
messages,
branchGroups,
});
}
} catch (error) {
console.error("[GlobalChatbox] Failed to rename chat session:", error);
}
},
[branchGroups, isHydrating, messages],
[isHydrating, messages],
);
const regenerate = useCallback(async () => {
const regenerate = useCallback(async (messageId: string) => {
if (isHydrating || isStreaming || messages.length === 0) return;
let lastUserIndex = messages.length - 1;
while (lastUserIndex >= 0 && messages[lastUserIndex].role !== "user") {
lastUserIndex--;
const targetAssistantIndex = messages.findIndex(
(message) => message.id === messageId && message.role === "assistant",
);
if (targetAssistantIndex < 0) {
return;
}
if (lastUserIndex < 0) return;
let targetUserIndex = targetAssistantIndex - 1;
while (targetUserIndex >= 0 && messages[targetUserIndex].role !== "user") {
targetUserIndex--;
}
const lastUser = messages[lastUserIndex];
const lastUserContent = lastUser.content;
const nextMessages = cloneMessages(messages.slice(0, lastUserIndex));
const nextUserMessage = createUserMessage(
lastUserContent,
lastUser.branchRootId ?? lastUser.id,
);
const nextAssistantMessage = createAssistantMessage();
if (targetUserIndex < 0) return;
setMessages(nextMessages);
await runPrompt({
prompt: lastUserContent,
regenerateFromMessageIndex: lastUserIndex,
preparedMessages: [
...nextMessages,
nextUserMessage,
nextAssistantMessage,
],
userMessage: nextUserMessage,
assistantMessage: nextAssistantMessage,
});
}, [isHydrating, isStreaming, messages, runPrompt]);
const targetUser = messages[targetUserIndex];
const targetUserContent = targetUser.content;
const nextMessages = cloneMessages(messages.slice(0, targetUserIndex));
const nextUserMessage = createUserMessage(targetUserContent);
const nextAssistantMessage = createAssistantMessage();
const editAndResubmit = useCallback(
async (messageId: string, newContent: string) => {
setMessages(nextMessages);
await runPrompt({
prompt: targetUserContent,
regenerateFromMessageIndex: targetUserIndex,
preparedMessages: [
...nextMessages,
nextUserMessage,
nextAssistantMessage,
],
userMessage: nextUserMessage,
assistantMessage: nextAssistantMessage,
});
}, [isHydrating, isStreaming, messages, runPrompt]);
const createBranch = useCallback(
async (messageId: string) => {
if (isHydrating || isStreaming) return;
const trimmedContent = newContent.trim();
if (!trimmedContent) return;
const assistantIndex = messages.findIndex(
(message) => message.id === messageId && message.role === "assistant",
);
if (assistantIndex < 0) return;
const messageIndex = messages.findIndex((m) => m.id === messageId);
if (messageIndex < 0 || messages[messageIndex].role !== "user") return;
const originalMessage = messages[messageIndex];
if (trimmedContent === originalMessage.content.trim()) return;
const rootMessageId = originalMessage.branchRootId ?? originalMessage.id;
const currentSessionId = sessionIdRef.current;
const keepMessageCount = messageIndex;
const prefix = cloneMessages(messages.slice(0, messageIndex));
const originalSuffix = cloneMessages(messages.slice(messageIndex));
const keepMessageCount = assistantIndex + 1;
const copiedMessages = cloneMessages(messages.slice(0, keepMessageCount));
const forkedSessionId = await forkAgentChat(currentSessionId, keepMessageCount);
const nextUserMessage = createUserMessage(trimmedContent, rootMessageId);
const nextAssistantMessage = createAssistantMessage();
const nextSuffix = [nextUserMessage, nextAssistantMessage];
setBranchGroups((prev) => {
const next = cloneBranchGroups(prev);
const groupIndex = next.findIndex(
(group) =>
group.rootMessageId === rootMessageId && group.parentCount === messageIndex,
);
if (groupIndex >= 0) {
const group = next[groupIndex];
group.branches[group.activeIndex] = {
...group.branches[group.activeIndex],
sessionId: currentSessionId,
messages: originalSuffix,
};
group.branches.push({
id: createId(),
label: `分支 ${group.branches.length + 1}`,
sessionId: forkedSessionId,
messages: cloneMessages(nextSuffix),
});
group.activeIndex = group.branches.length - 1;
} else {
next.push({
id: rootMessageId,
rootMessageId,
parentCount: messageIndex,
activeIndex: 1,
branches: [
{
id: createId(),
label: "分支 1",
sessionId: currentSessionId,
messages: originalSuffix,
},
{
id: createId(),
label: "分支 2",
sessionId: forkedSessionId,
messages: cloneMessages(nextSuffix),
},
],
});
}
return next;
});
sessionIdRef.current = forkedSessionId;
setSessionId(forkedSessionId);
await runPrompt({
prompt: trimmedContent,
sessionIdOverride: forkedSessionId,
preparedMessages: [...prefix, ...nextSuffix],
userMessage: nextUserMessage,
assistantMessage: nextAssistantMessage,
});
},
[isHydrating, isStreaming, messages, runPrompt],
);
const cycleBranch = useCallback(
(rootMessageId: string, direction: -1 | 1) => {
if (isHydrating || isStreaming) return;
setBranchGroups((prev) => {
const next = cloneBranchGroups(prev);
const group = next.find((item) => item.rootMessageId === rootMessageId);
if (!group || group.branches.length < 2) {
return prev;
}
const nextIndex =
(group.activeIndex + direction + group.branches.length) % group.branches.length;
const selectedBranch = group.branches[nextIndex];
group.activeIndex = nextIndex;
const nextMessages = [
...cloneMessages(messages.slice(0, group.parentCount)),
...cloneMessages(selectedBranch.messages),
];
setBranchTransition({
rootMessageId,
parentCount: group.parentCount,
activeBranchId: selectedBranch.id,
nonce: Date.now(),
messagesRef.current = copiedMessages;
setMessages(copiedMessages);
setIsSessionTitleManuallyEdited(false);
const forkTitle = sessionTitle ? `${sessionTitle} 副本` : "新对话副本";
setSessionTitle(forkTitle);
try {
await saveActiveChatState({
title: forkTitle,
isTitleManuallyEdited: false,
messages: copiedMessages,
sessionId: forkedSessionId,
});
sessionIdRef.current = selectedBranch.sessionId;
setSessionId(selectedBranch.sessionId);
setMessages(nextMessages);
return next;
});
setChatSessions(await listChatSessions());
} catch (error) {
console.error("[GlobalChatbox] Failed to refresh chat sessions after fork:", error);
}
},
[isHydrating, isStreaming, messages],
[isHydrating, isStreaming, messages, sessionTitle],
);
return {
messages,
chatSessions,
activeSessionId: sessionIdRef.current,
branchGroups,
branchTransition,
isHydrating,
isStreaming,
sessionTitle,
sessionId,
sendPrompt,
regenerate,
editAndResubmit,
cycleBranch,
createBranch,
abort,
replyPermission,
createSession,