538 lines
16 KiB
TypeScript
538 lines
16 KiB
TypeScript
import type { UIMessage } from "ai";
|
|
import type { Dispatch, SetStateAction } from "react";
|
|
import {
|
|
parseAnswerUiBlock,
|
|
type AnswerUiPart,
|
|
applyPermissionResponse,
|
|
agentTaskContextReferenceSchema,
|
|
applyQuestionResponse,
|
|
toTodoUpdate,
|
|
upsertPermission,
|
|
upsertProgress,
|
|
upsertQuestion,
|
|
type AgentChatMessage,
|
|
type AgentPermissionReply,
|
|
type AgentPermissionStatus,
|
|
type AgentQuestionRequest,
|
|
type AgentSessionStreamEvent,
|
|
type AgentStreamRenderState
|
|
} from "@/features/agent";
|
|
|
|
type AgentUiDataParts = {
|
|
session_created: Record<string, unknown>;
|
|
agent_message: Record<string, unknown>;
|
|
task_context: Record<string, unknown>;
|
|
execution_plan: Record<string, unknown>;
|
|
progress: Record<string, unknown>;
|
|
final_answer: Record<string, unknown>;
|
|
credential_refresh_required: Record<string, unknown>;
|
|
credential_refreshed: Record<string, unknown>;
|
|
credential_refresh_failed: Record<string, unknown>;
|
|
todo_update: Record<string, unknown>;
|
|
permission_request: Record<string, unknown>;
|
|
permission_response: Record<string, unknown>;
|
|
question_request: Record<string, unknown>;
|
|
question_response: Record<string, unknown>;
|
|
stream_token: Record<string, unknown>;
|
|
tool_call: Record<string, unknown>;
|
|
frontend_action: Record<string, unknown>;
|
|
frontend_action_result: Record<string, unknown>;
|
|
answer_ui: Record<string, unknown>;
|
|
ui_envelope: Record<string, unknown>;
|
|
session_title: Record<string, unknown>;
|
|
user_input: Record<string, unknown>;
|
|
};
|
|
|
|
export type AgentUiMessage = UIMessage<unknown, AgentUiDataParts>;
|
|
|
|
export type AgentDataPart = Extract<AgentUiMessage["parts"][number], { type: `data-${string}` }>;
|
|
|
|
export type PermissionOverride = {
|
|
status: AgentPermissionStatus;
|
|
error?: string;
|
|
};
|
|
|
|
export type QuestionOverride = {
|
|
status: AgentQuestionRequest["status"];
|
|
answers?: string[][];
|
|
error?: string;
|
|
};
|
|
|
|
type StreamRenderStateSetter = Dispatch<SetStateAction<AgentStreamRenderState>>;
|
|
|
|
export function markStreamRenderPending(
|
|
data: Record<string, unknown>,
|
|
messages: AgentUiMessage[],
|
|
setStreamRenderState: StreamRenderStateSetter
|
|
) {
|
|
const text = getDataString(data, "content");
|
|
if (!text) {
|
|
return;
|
|
}
|
|
|
|
const messageId = getDataString(data, "message_id") ?? getLastAssistantUiMessageId(messages);
|
|
if (!messageId) {
|
|
return;
|
|
}
|
|
|
|
setStreamRenderState((current) => {
|
|
if (current[messageId]?.done === false) {
|
|
return current;
|
|
}
|
|
|
|
return {
|
|
...current,
|
|
[messageId]: { done: false }
|
|
};
|
|
});
|
|
}
|
|
|
|
export function createCompletedStreamRenderState(messages: AgentUiMessage[]): AgentStreamRenderState {
|
|
return messages.reduce<AgentStreamRenderState>((next, message) => {
|
|
if (message.role === "assistant") {
|
|
next[message.id] = { done: true };
|
|
}
|
|
return next;
|
|
}, {});
|
|
}
|
|
|
|
export function markLastAssistantStreamDone(
|
|
messages: AgentUiMessage[],
|
|
setStreamRenderState: StreamRenderStateSetter
|
|
) {
|
|
const messageId = getLastAssistantUiMessageId(messages);
|
|
setStreamRenderState((current) => {
|
|
if (!messageId) {
|
|
return Object.fromEntries(
|
|
Object.entries(current).map(([id, state]) => [id, { ...state, done: true }])
|
|
);
|
|
}
|
|
|
|
return {
|
|
...current,
|
|
[messageId]: { done: true }
|
|
};
|
|
});
|
|
}
|
|
|
|
export function toAgentChatMessages(
|
|
messages: AgentUiMessage[],
|
|
permissionOverrides: Record<string, PermissionOverride>,
|
|
questionOverrides: Record<string, QuestionOverride>
|
|
): AgentChatMessage[] {
|
|
return messages.flatMap((message) => {
|
|
if (message.role !== "user" && message.role !== "assistant") {
|
|
return [];
|
|
}
|
|
|
|
const answerParts = collectAnswerParts(message);
|
|
let next: AgentChatMessage = {
|
|
id: message.id,
|
|
role: message.role,
|
|
content: answerParts ? answerParts.map((part) => part.type === "text" ? part.text : part.block.fallback_text).join("\n\n") : collectMessageText(message),
|
|
...(answerParts ? { answerParts } : {})
|
|
};
|
|
|
|
for (const part of message.parts) {
|
|
if (!isAgentDataPart(part)) {
|
|
continue;
|
|
}
|
|
if (part.type === "data-task_context") {
|
|
const parsed = agentTaskContextReferenceSchema.safeParse(part.data);
|
|
if (parsed.success) next = { ...next, taskContext: parsed.data };
|
|
} else if (part.type === "data-final_answer") {
|
|
const content = getDataString(part.data, "content");
|
|
if (content?.trim()) {
|
|
const refs = part.data?.ui_block_ids;
|
|
const uiBlockIds = Array.isArray(refs) ? [...new Set(refs.filter((id): id is string => typeof id === "string"))].slice(0, 48) : [];
|
|
next = { ...next, finalAnswer: { content, uiBlockIds } };
|
|
}
|
|
} else if (part.type === "data-progress") {
|
|
next = {
|
|
...next,
|
|
progress: upsertProgress(next.progress, part.data)
|
|
};
|
|
} else if (part.type === "data-todo_update") {
|
|
const todoUpdate = toTodoUpdate(part.data);
|
|
if (todoUpdate) {
|
|
next = {
|
|
...next,
|
|
todos: todoUpdate
|
|
};
|
|
}
|
|
} else if (part.type === "data-permission_request") {
|
|
next = {
|
|
...next,
|
|
permissions: upsertPermission(next.permissions, part.data)
|
|
};
|
|
} else if (part.type === "data-permission_response") {
|
|
next = {
|
|
...next,
|
|
permissions: applyPermissionResponse(next.permissions, part.data)
|
|
};
|
|
} else if (part.type === "data-question_request") {
|
|
next = {
|
|
...next,
|
|
questions: upsertQuestion(next.questions, part.data)
|
|
};
|
|
} else if (part.type === "data-question_response") {
|
|
next = {
|
|
...next,
|
|
questions: applyQuestionResponse(next.questions, part.data)
|
|
};
|
|
}
|
|
}
|
|
|
|
if (next.finalAnswer) {
|
|
const selected = new Set(next.finalAnswer.uiBlockIds);
|
|
next.content = [next.finalAnswer.content, ...(next.answerParts ?? []).flatMap((part) => part.type === "answer-ui" && selected.has(part.block.block_id) ? [part.block.fallback_text] : [])].join("\n\n");
|
|
}
|
|
return [applyQuestionOverrides(applyPermissionOverrides(next, permissionOverrides), questionOverrides)];
|
|
});
|
|
}
|
|
|
|
export function toAgentUiMessages(messages: unknown[]): AgentUiMessage[] {
|
|
return messages.flatMap((message) => {
|
|
if (isAgentUiMessage(message)) {
|
|
return [message];
|
|
}
|
|
const legacyMessage = toAgentUiMessageFromLegacy(message);
|
|
return legacyMessage ? [legacyMessage] : [];
|
|
});
|
|
}
|
|
|
|
export function applySessionStreamEvent(messages: AgentUiMessage[], event: AgentSessionStreamEvent): AgentUiMessage[] {
|
|
if (event.type === "state") {
|
|
return toAgentUiMessages(event.messages);
|
|
}
|
|
|
|
if (event.type === "user_input") {
|
|
const id = getDataString(event.data, "id");
|
|
const text = getDataString(event.data, "text");
|
|
if (event.data.status !== "delivered" || !id || !text || messages.some((message) => message.id === id)) return messages;
|
|
const input: AgentUiMessage = { id, role: "user", parts: [{ type: "text", text }] };
|
|
const tail = messages.at(-1);
|
|
return tail?.role === "assistant" ? [...messages.slice(0, -1), input, tail] : [...messages, input];
|
|
}
|
|
|
|
if (event.type === "token") {
|
|
const token = getDataString(event.data, "content");
|
|
return token ? updateLastAssistantUiMessage(messages, (message) => appendTextToUiMessage(message, token)) : messages;
|
|
}
|
|
|
|
if (
|
|
event.type === "final_answer" ||
|
|
event.type === "progress" ||
|
|
event.type === "execution_plan" ||
|
|
event.type === "todo_update" ||
|
|
event.type === "permission_request" ||
|
|
event.type === "permission_response" ||
|
|
event.type === "question_request" ||
|
|
event.type === "question_response" ||
|
|
event.type === "answer_ui" ||
|
|
event.type === "ui_envelope" ||
|
|
event.type === "session_title"
|
|
) {
|
|
return updateLastAssistantUiMessage(messages, (message) =>
|
|
upsertUiDataPart(message, event.type, event.data)
|
|
);
|
|
}
|
|
|
|
if (event.type === "error") {
|
|
const message = getDataString(event.data, "message") ?? "Agent stream failed";
|
|
return updateLastAssistantUiMessage(messages, (item) =>
|
|
appendTextToUiMessage(item, item.parts.some((part) => part.type === "text") ? `\n\n错误：${message}` : `错误：${message}`)
|
|
);
|
|
}
|
|
|
|
return messages;
|
|
}
|
|
|
|
export function toPermissionStatus(reply: AgentPermissionReply): AgentPermissionStatus {
|
|
if (reply === "always") {
|
|
return "approved_always";
|
|
}
|
|
if (reply === "once") {
|
|
return "approved_once";
|
|
}
|
|
return "rejected";
|
|
}
|
|
|
|
export function getDataString(data: unknown, key: string) {
|
|
if (typeof data !== "object" || data === null) {
|
|
return undefined;
|
|
}
|
|
const value = (data as Record<string, unknown>)[key];
|
|
return typeof value === "string" ? value : undefined;
|
|
}
|
|
|
|
export function readBodyString(body: unknown, key: string) {
|
|
if (typeof body !== "object" || body === null) {
|
|
return undefined;
|
|
}
|
|
const value = (body as Record<string, unknown>)[key];
|
|
return typeof value === "string" ? value : undefined;
|
|
}
|
|
|
|
function getLastAssistantUiMessageId(messages: AgentUiMessage[]) {
|
|
for (let index = messages.length - 1; index >= 0; index -= 1) {
|
|
if (messages[index].role === "assistant") {
|
|
return messages[index].id;
|
|
}
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
function isAgentUiMessage(value: unknown): value is AgentUiMessage {
|
|
if (typeof value !== "object" || value === null) {
|
|
return false;
|
|
}
|
|
const message = value as Record<string, unknown>;
|
|
return (
|
|
typeof message.id === "string" &&
|
|
(message.role === "user" || message.role === "assistant") &&
|
|
Array.isArray(message.parts)
|
|
);
|
|
}
|
|
|
|
function toAgentUiMessageFromLegacy(value: unknown): AgentUiMessage | null {
|
|
if (typeof value !== "object" || value === null) {
|
|
return null;
|
|
}
|
|
|
|
const message = value as Record<string, unknown>;
|
|
if (
|
|
typeof message.id !== "string" ||
|
|
(message.role !== "user" && message.role !== "assistant")
|
|
) {
|
|
return null;
|
|
}
|
|
|
|
const parts: AgentUiMessage["parts"] = [];
|
|
if (typeof message.content === "string" && message.content) {
|
|
parts.push({ type: "text", text: message.content });
|
|
}
|
|
appendLegacyDataParts(parts, "progress", message.progress);
|
|
appendLegacyDataParts(parts, "permission_request", message.permissions);
|
|
appendLegacyDataParts(parts, "question_request", message.questions);
|
|
appendLegacyDataParts(parts, "todo_update", message.todos);
|
|
appendLegacyDataParts(parts, "execution_plan", message.executionPlans);
|
|
|
|
return {
|
|
id: message.id,
|
|
role: message.role,
|
|
parts
|
|
} as AgentUiMessage;
|
|
}
|
|
|
|
function appendLegacyDataParts(
|
|
parts: AgentUiMessage["parts"],
|
|
eventType: string,
|
|
value: unknown
|
|
) {
|
|
if (!value) {
|
|
return;
|
|
}
|
|
|
|
const values = Array.isArray(value) ? value : [value];
|
|
values.forEach((item, index) => {
|
|
if (typeof item !== "object" || item === null) {
|
|
return;
|
|
}
|
|
const data = normalizeLegacyDataPart(eventType, item as Record<string, unknown>);
|
|
const id =
|
|
getDataString(data, "id") ??
|
|
getDataString(data, "request_id") ??
|
|
`${eventType}-${index}`;
|
|
parts.push({
|
|
type: `data-${eventType}`,
|
|
id,
|
|
data
|
|
} as AgentDataPart);
|
|
});
|
|
}
|
|
|
|
function normalizeLegacyDataPart(eventType: string, value: Record<string, unknown>) {
|
|
if (eventType === "progress") {
|
|
return {
|
|
...value,
|
|
started_at: value.started_at ?? value.startedAt,
|
|
ended_at: value.ended_at ?? value.endedAt,
|
|
elapsed_ms: value.elapsed_ms ?? value.elapsedMs,
|
|
duration_ms: value.duration_ms ?? value.durationMs
|
|
};
|
|
}
|
|
|
|
if (eventType === "permission_request" || eventType === "question_request") {
|
|
return {
|
|
...value,
|
|
session_id: value.session_id ?? value.sessionId,
|
|
request_id: value.request_id ?? value.requestId,
|
|
created_at: value.created_at ?? value.createdAt
|
|
};
|
|
}
|
|
|
|
if (eventType === "todo_update") {
|
|
return {
|
|
...value,
|
|
session_id: value.session_id ?? value.sessionId,
|
|
message_id: value.message_id ?? value.messageId,
|
|
created_at: value.created_at ?? value.createdAt
|
|
};
|
|
}
|
|
|
|
return value;
|
|
}
|
|
|
|
function updateLastAssistantUiMessage(
|
|
messages: AgentUiMessage[],
|
|
updater: (message: AgentUiMessage) => AgentUiMessage
|
|
) {
|
|
for (let index = messages.length - 1; index >= 0; index -= 1) {
|
|
if (messages[index].role === "assistant") {
|
|
const next = [...messages];
|
|
next[index] = updater(messages[index]);
|
|
return next;
|
|
}
|
|
}
|
|
|
|
return [...messages, updater(createAssistantUiMessage())];
|
|
}
|
|
|
|
function createAssistantUiMessage(): AgentUiMessage {
|
|
return {
|
|
id: `assistant-${Date.now().toString(36)}`,
|
|
role: "assistant",
|
|
parts: []
|
|
} as AgentUiMessage;
|
|
}
|
|
|
|
function appendTextToUiMessage(message: AgentUiMessage, text: string): AgentUiMessage {
|
|
let textIndex = -1;
|
|
for (let index = message.parts.length - 1; index >= 0; index -= 1) {
|
|
if (message.parts[index].type === "text") { textIndex = index; break; }
|
|
}
|
|
if (message.parts.slice(textIndex + 1).some((part) => part.type === "data-answer_ui")) textIndex = -1;
|
|
if (textIndex === -1) {
|
|
return {
|
|
...message,
|
|
parts: [...message.parts, { type: "text", text }]
|
|
};
|
|
}
|
|
|
|
return {
|
|
...message,
|
|
parts: message.parts.map((part, index) =>
|
|
index === textIndex && part.type === "text"
|
|
? { ...part, text: `${part.text}${text}` }
|
|
: part
|
|
)
|
|
};
|
|
}
|
|
|
|
function upsertUiDataPart(
|
|
message: AgentUiMessage,
|
|
eventType: string,
|
|
data: Record<string, unknown>
|
|
): AgentUiMessage {
|
|
const type = `data-${eventType}`;
|
|
const id =
|
|
getDataString(data, "id") ??
|
|
getDataString(data, "request_id") ??
|
|
getDataString(data, "envelope_id") ??
|
|
getDataString(data, "block_id");
|
|
const existingIndex =
|
|
id === undefined
|
|
? -1
|
|
: message.parts.findIndex((part) => part.type === type && "id" in part && part.id === id);
|
|
const part = (id ? { type, id, data } : { type, data }) as AgentDataPart;
|
|
|
|
if (existingIndex === -1) {
|
|
return {
|
|
...message,
|
|
parts: [...message.parts, part]
|
|
};
|
|
}
|
|
|
|
return {
|
|
...message,
|
|
parts: message.parts.map((item, index) => (index === existingIndex ? part : item))
|
|
};
|
|
}
|
|
|
|
function applyPermissionOverrides(
|
|
message: AgentChatMessage,
|
|
permissionOverrides: Record<string, PermissionOverride>
|
|
) {
|
|
if (!message.permissions?.length) {
|
|
return message;
|
|
}
|
|
|
|
return {
|
|
...message,
|
|
permissions: message.permissions.map((permission) => {
|
|
const override = permissionOverrides[permission.requestId];
|
|
return override
|
|
? {
|
|
...permission,
|
|
status: override.status,
|
|
error: override.error,
|
|
repliedAt: override.status === "submitting" || override.status === "error" ? permission.repliedAt : Date.now()
|
|
}
|
|
: permission;
|
|
})
|
|
};
|
|
}
|
|
|
|
function applyQuestionOverrides(
|
|
message: AgentChatMessage,
|
|
questionOverrides: Record<string, QuestionOverride>
|
|
) {
|
|
if (!message.questions?.length) {
|
|
return message;
|
|
}
|
|
|
|
return {
|
|
...message,
|
|
questions: message.questions.map((question) => {
|
|
const override = questionOverrides[question.requestId];
|
|
return override
|
|
? {
|
|
...question,
|
|
status: override.status,
|
|
answers: override.answers ?? question.answers,
|
|
error: override.error,
|
|
repliedAt: override.status === "submitting" || override.status === "error" ? question.repliedAt : Date.now()
|
|
}
|
|
: question;
|
|
})
|
|
};
|
|
}
|
|
|
|
function collectMessageText(message: AgentUiMessage) {
|
|
return message.parts
|
|
.flatMap((part) => (part.type === "text" ? [part.text] : []))
|
|
.join("");
|
|
}
|
|
|
|
function isAgentDataPart(part: AgentUiMessage["parts"][number]): part is AgentDataPart {
|
|
return typeof part.type === "string" && part.type.startsWith("data-") && "data" in part;
|
|
}
|
|
|
|
function collectAnswerParts(message: AgentUiMessage): AnswerUiPart[] | undefined {
|
|
if (!message.parts.some((part) => part.type === "data-answer_ui" || part.type === "data-final_answer")) return undefined;
|
|
const result: AnswerUiPart[] = [];
|
|
for (const part of message.parts) {
|
|
if (part.type === "text" && part.text) result.push({ type: "text", text: part.text });
|
|
if (part.type === "data-answer_ui") {
|
|
const data = typeof part.data === "object" && part.data !== null ? part.data : {};
|
|
const { version, block_id, fallback_text, spec } = data;
|
|
const block = parseAnswerUiBlock({ version, block_id, fallback_text, spec });
|
|
if (block) result.push({ type: "answer-ui", block });
|
|
else result.push({ type: "text", text: typeof fallback_text === "string" ? fallback_text.slice(0, 4000) : "这段结构化内容无法显示。" });
|
|
}
|
|
|
|
}
|
|
return result;
|
|
}
|