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

feat(chat): refresh agent credentials during tool calls

Browse Source
This commit is contained in:
jiang
2026-08-06 10:16:50 +08:00
parent b4c96f8524
commit 5037089057
10 changed files with 590 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/components/chat/hooks/useAgentChatSession.actions.test.tsx
+10
View File
@@ -6,6 +6,7 @@ import { useAgentChatSession } from "./useAgentChatSession";
import {
abortAgentChat,
forkAgentChat,
replyAgentCredentialRefresh,
replyAgentPermission,
replyAgentQuestion,
resumeAgentChatStream,
@@ -16,12 +17,19 @@ import type { StreamEvent } from "@/lib/chatStream";
jest.mock("@/lib/chatStream", () => ({
abortAgentChat: jest.fn(async () => undefined),
forkAgentChat: jest.fn(async () => "forked-session"),
replyAgentCredentialRefresh: jest.fn(async () => undefined),
replyAgentPermission: jest.fn(async () => undefined),
replyAgentQuestion: jest.fn(async () => undefined),
resumeAgentChatStream: jest.fn(async () => undefined),
streamAgentChat: jest.fn(async () => undefined),
}));
const mockUpdateSession = jest.fn();
jest.mock("next-auth/react", () => ({
useSession: () => ({ update: mockUpdateSession }),
}));
const listChatSessions = jest.fn();
const deleteChatSession = jest.fn();
const updateChatSessionTitle = jest.fn();
@@ -51,12 +59,14 @@ describe("useAgentChatSession", () => {
updateChatSessionTitle.mockReset();
jest.mocked(abortAgentChat).mockReset();
jest.mocked(forkAgentChat).mockReset();
jest.mocked(replyAgentCredentialRefresh).mockReset();
jest.mocked(replyAgentPermission).mockReset();
jest.mocked(replyAgentQuestion).mockReset();
jest.mocked(resumeAgentChatStream).mockReset();
jest.mocked(streamAgentChat).mockReset();
jest.mocked(abortAgentChat).mockImplementation(async () => undefined);
jest.mocked(forkAgentChat).mockImplementation(async () => "forked-session");
jest.mocked(replyAgentCredentialRefresh).mockImplementation(async () => undefined);
jest.mocked(replyAgentPermission).mockImplementation(async () => undefined);
jest.mocked(replyAgentQuestion).mockImplementation(async () => undefined);
jest.mocked(resumeAgentChatStream).mockImplementation(async () => undefined);
src/components/chat/hooks/useAgentChatSession.lifecycle.test.tsx
+10
View File
@@ -6,6 +6,7 @@ import { useAgentChatSession } from "./useAgentChatSession";
import {
abortAgentChat,
forkAgentChat,
replyAgentCredentialRefresh,
replyAgentPermission,
replyAgentQuestion,
resumeAgentChatStream,
@@ -16,12 +17,19 @@ import type { StreamEvent } from "@/lib/chatStream";
jest.mock("@/lib/chatStream", () => ({
abortAgentChat: jest.fn(async () => undefined),
forkAgentChat: jest.fn(async () => "forked-session"),
replyAgentCredentialRefresh: jest.fn(async () => undefined),
replyAgentPermission: jest.fn(async () => undefined),
replyAgentQuestion: jest.fn(async () => undefined),
resumeAgentChatStream: jest.fn(async () => undefined),
streamAgentChat: jest.fn(async () => undefined),
}));
const mockUpdateSession = jest.fn();
jest.mock("next-auth/react", () => ({
useSession: () => ({ update: mockUpdateSession }),
}));
const listChatSessions = jest.fn();
const deleteChatSession = jest.fn();
const updateChatSessionTitle = jest.fn();
@@ -51,12 +59,14 @@ describe("useAgentChatSession", () => {
updateChatSessionTitle.mockReset();
jest.mocked(abortAgentChat).mockReset();
jest.mocked(forkAgentChat).mockReset();
jest.mocked(replyAgentCredentialRefresh).mockReset();
jest.mocked(replyAgentPermission).mockReset();
jest.mocked(replyAgentQuestion).mockReset();
jest.mocked(resumeAgentChatStream).mockReset();
jest.mocked(streamAgentChat).mockReset();
jest.mocked(abortAgentChat).mockImplementation(async () => undefined);
jest.mocked(forkAgentChat).mockImplementation(async () => "forked-session");
jest.mocked(replyAgentCredentialRefresh).mockImplementation(async () => undefined);
jest.mocked(replyAgentPermission).mockImplementation(async () => undefined);
jest.mocked(replyAgentQuestion).mockImplementation(async () => undefined);
jest.mocked(resumeAgentChatStream).mockImplementation(async () => undefined);
src/components/chat/hooks/useAgentChatSession.ts
+124 -1
View File
@@ -1,9 +1,11 @@
"use client";
import { useCallback, useEffect, useRef, useState } from "react";
import { useSession } from "next-auth/react";
import { abortAgentChat, forkAgentChat, rejectAgentQuestion, replyAgentPermission, replyAgentQuestion, resumeAgentChatStream, streamAgentChat } from "@/lib/chatStream";
import { abortAgentChat, forkAgentChat, rejectAgentQuestion, replyAgentCredentialRefresh, replyAgentPermission, replyAgentQuestion, resumeAgentChatStream, streamAgentChat } from "@/lib/chatStream";
import type { PermissionReply, StreamEvent } from "@/lib/chatStream";
import { useAuthStore } from "@/store/authStore";
import type { AgentArtifact, ChatSessionSummary, Message } from "../GlobalChatbox.types";
import { cloneMessages } from "../globalChatboxUtils";
import { createEmptyChatState, deleteChatSession, listChatSessions, loadChatSessionById, updateChatSessionTitle } from "../chatStorage";
@@ -80,6 +82,7 @@ export const useAgentChatSession = ({
getModel,
getApprovalMode,
}: UseAgentChatSessionOptions) => {
const { update: updateSession } = useSession();
const hydrationNonceRef = useRef(0);
const [messages, setMessages] = useState<Message[]>([]);
@@ -102,6 +105,7 @@ export const useAgentChatSession = ({
content: string;
} | null>(null);
const tokenPlaybackIntervalRef = useRef<number | null>(null);
const credentialRefreshRequestIdsRef = useRef(new Set<string>());
useEffect(() => {
sessionIdRef.current = sessionId;
@@ -289,6 +293,61 @@ export const useAgentChatSession = ({
return assistant?.id ?? fallback;
}, []);
const handleCredentialRefresh = useCallback(
async (event: StreamEvent & { type: "credential_refresh_required" }) => {
if (
!event.sessionId ||
!event.requestId ||
credentialRefreshRequestIdsRef.current.has(event.requestId)
) {
return;
}
credentialRefreshRequestIdsRef.current.add(event.requestId);
try {
const refreshedSession = await updateSession({ forceRefresh: true });
if (
refreshedSession?.error ||
typeof refreshedSession?.accessToken !== "string" ||
!refreshedSession.accessToken
) {
throw new Error("登录凭据续期失败");
}
const authStore = useAuthStore.getState();
authStore.setAccessToken(refreshedSession.accessToken);
authStore.clearSessionExpired();
await replyAgentCredentialRefresh(event.sessionId, event.requestId);
} catch (error) {
useAuthStore.getState().markSessionExpired("refresh_failed");
const assistantMessageId = getLastAssistantMessageId();
if (assistantMessageId) {
const message = error instanceof Error ? error.message : String(error);
setMessages((prev) =>
prev.map((item) =>
item.id === assistantMessageId
? {
...item,
content: item.content || `⚠️ **${message}**`,
isError: true,
progress: upsertProgress(item.progress, {
type: "progress",
sessionId: event.sessionId,
id: `credential-refresh-${event.requestId}`,
phase: "credential_refresh",
status: "error",
title: "登录凭据续期失败",
detail: message,
}),
}
: item,
),
);
}
setIsStreaming(false);
}
},
[getLastAssistantMessageId, updateSession],
);
const applyStreamEvent = useCallback(
(
event: StreamEvent,
@@ -421,6 +480,68 @@ export const useAgentChatSession = ({
assistantMessageId,
),
);
} else if (event.type === "credential_refresh_required") {
setMessages((prev) =>
prev.map((message) =>
message.id === assistantMessageId
? {
...message,
progress: upsertProgress(message.progress, {
type: "progress",
sessionId: event.sessionId,
id: `credential-refresh-${event.requestId}`,
phase: "credential_refresh",
status: "running",
title: "正在续期登录凭据",
detail: `当前工具调用保持等待，最长 ${Math.ceil((event.timeoutMs ?? 30_000) / 1000)} 秒`,
startedAt: Date.now(),
}),
}
: message,
),
);
void handleCredentialRefresh(event);
} else if (event.type === "credential_refreshed") {
setMessages((prev) =>
prev.map((message) =>
message.id === assistantMessageId
? {
...message,
progress: upsertProgress(message.progress, {
type: "progress",
sessionId: event.sessionId,
id: `credential-refresh-${event.requestId}`,
phase: "credential_refresh",
status: "completed",
title: "登录凭据已续期",
}),
}
: message,
),
);
} else if (event.type === "credential_refresh_failed") {
useAuthStore.getState().markSessionExpired("refresh_failed");
setMessages((prev) =>
prev.map((message) =>
message.id === assistantMessageId
? {
...message,
content: message.content || `⚠️ **${event.message}**`,
isError: true,
progress: upsertProgress(message.progress, {
type: "progress",
sessionId: event.sessionId,
id: `credential-refresh-${event.requestId}`,
phase: "credential_refresh",
status: "error",
title: "登录凭据续期失败",
detail: event.message,
}),
}
: message,
),
);
setIsStreaming(false);
} else if (event.type === "done") {
setMessages((prev) =>
prev.map((message) => {
@@ -457,6 +578,7 @@ export const useAgentChatSession = ({
);
setIsStreaming(false);
} else if (event.type === "auth_required") {
useAuthStore.getState().markSessionExpired("unauthorized");
setMessages((prev) =>
prev.map((message) =>
message.id === assistantMessageId
@@ -477,6 +599,7 @@ export const useAgentChatSession = ({
appendArtifact,
flushPendingTokens,
getLastAssistantMessageId,
handleCredentialRefresh,
onToolCall,
queueTokenContent,
],