next-tjwater-drainage-frontend/features/workbench/hooks/agent-session-message-state.ts

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;
}