feat: initialize drainage network frontend
This commit is contained in:
@@ -0,0 +1,484 @@
|
||||
import type { UIMessage } from "ai";
|
||||
import type { Dispatch, MutableRefObject, SetStateAction } from "react";
|
||||
import type {
|
||||
AgentChatMessage,
|
||||
AgentPermissionReply,
|
||||
AgentPermissionStatus,
|
||||
AgentQuestionRequest,
|
||||
AgentStreamRenderState
|
||||
} from "@/features/agent";
|
||||
import type { AgentSessionStreamEvent } from "@/features/agent/api/client";
|
||||
import {
|
||||
applyPermissionResponse,
|
||||
applyQuestionResponse,
|
||||
toTodoUpdate,
|
||||
upsertPermission,
|
||||
upsertProgress,
|
||||
upsertQuestion
|
||||
} from "@/features/agent/session-state";
|
||||
|
||||
type AgentUiDataParts = {
|
||||
progress: 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>;
|
||||
ui_envelope: Record<string, unknown>;
|
||||
session_title: 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>>;
|
||||
type StreamRenderChunkIdRef = MutableRefObject<number>;
|
||||
|
||||
export function appendStreamRenderToken(
|
||||
data: Record<string, unknown>,
|
||||
messages: AgentUiMessage[],
|
||||
setStreamRenderState: StreamRenderStateSetter,
|
||||
chunkIdRef: StreamRenderChunkIdRef
|
||||
) {
|
||||
const text = getDataString(data, "content");
|
||||
if (!text) {
|
||||
return;
|
||||
}
|
||||
|
||||
const messageId = getDataString(data, "message_id") ?? getLastAssistantUiMessageId(messages);
|
||||
if (!messageId) {
|
||||
return;
|
||||
}
|
||||
|
||||
const chunkId = chunkIdRef.current;
|
||||
chunkIdRef.current += 1;
|
||||
|
||||
setStreamRenderState((current) => {
|
||||
const previous = current[messageId] ?? { chunks: [], done: false };
|
||||
return {
|
||||
...current,
|
||||
[messageId]: {
|
||||
chunks: [...previous.chunks, { id: chunkId, text }],
|
||||
done: false
|
||||
}
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
export function createCompletedStreamRenderState(messages: AgentUiMessage[]): AgentStreamRenderState {
|
||||
return messages.reduce<AgentStreamRenderState>((next, message) => {
|
||||
if (message.role === "assistant") {
|
||||
next[message.id] = {
|
||||
chunks: [],
|
||||
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]: {
|
||||
chunks: current[messageId]?.chunks ?? [],
|
||||
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 [];
|
||||
}
|
||||
|
||||
let next: AgentChatMessage = {
|
||||
id: message.id,
|
||||
role: message.role,
|
||||
content: collectMessageText(message)
|
||||
};
|
||||
|
||||
for (const part of message.parts) {
|
||||
if (!isAgentDataPart(part)) {
|
||||
continue;
|
||||
}
|
||||
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)
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
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 === "token") {
|
||||
const token = getDataString(event.data, "content");
|
||||
return token ? updateLastAssistantUiMessage(messages, (message) => appendTextToUiMessage(message, token)) : messages;
|
||||
}
|
||||
|
||||
if (
|
||||
event.type === "progress" ||
|
||||
event.type === "todo_update" ||
|
||||
event.type === "permission_request" ||
|
||||
event.type === "permission_response" ||
|
||||
event.type === "question_request" ||
|
||||
event.type === "question_response" ||
|
||||
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);
|
||||
|
||||
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 {
|
||||
const textIndex = message.parts.findIndex((part) => part.type === "text");
|
||||
if (textIndex === -1) {
|
||||
return {
|
||||
...message,
|
||||
parts: [{ type: "text", text }, ...message.parts]
|
||||
};
|
||||
}
|
||||
|
||||
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");
|
||||
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;
|
||||
}
|
||||
Reference in New Issue
Block a user