TJWaterFrontend_Refine/src/components/chat/hooks/agentChatSessionState.ts

import type {
  AgentQuestionRequest,
  AgentTodoUpdate,
  PermissionReply,
  StreamEvent,
} from "@/lib/chatStream";
import type {
  AgentPermissionRequest,
  ChatProgress,
  LoadedChatState,
  Message,
} from "../GlobalChatbox.types";
import { createId } from "../GlobalChatbox.utils";

export const createPersistedStateKey = (state: LoadedChatState) =>
  JSON.stringify({
    title: state.title ?? null,
    isTitleManuallyEdited: state.isTitleManuallyEdited ?? false,
    sessionId: state.sessionId ?? null,
    messages: state.messages,
  });

export const upsertProgress = (
  progress: ChatProgress[] | undefined,
  event: StreamEvent & { type: "progress" },
) => {
  const next = [...(progress ?? [])];
  const index = next.findIndex((item) => item.id === event.id);
  const existing = index >= 0 ? next[index] : undefined;
  const now = Date.now();
  const startedAt = event.startedAt ?? existing?.startedAt;
  const isRunning = event.status === "running";
  const endedAt = isRunning ? undefined : event.endedAt ?? existing?.endedAt ?? now;
  const elapsedMs = isRunning
    ? event.elapsedMs ??
      existing?.elapsedMs ??
      (startedAt !== undefined ? Math.max(0, now - startedAt) : undefined)
    : undefined;
  const elapsedSnapshotAt = isRunning
    ? event.elapsedMs !== undefined
      ? now
      : existing?.elapsedSnapshotAt ?? now
    : undefined;
  const durationMs = !isRunning
    ? event.durationMs ??
      existing?.durationMs ??
      (startedAt !== undefined && endedAt !== undefined
        ? Math.max(0, endedAt - startedAt)
        : undefined)
    : undefined;
  const nextItem: ChatProgress = {
    id: event.id,
    phase: event.phase,
    status: event.status,
    title: event.title,
    detail: event.detail,
    startedAt,
    endedAt,
    elapsedMs,
    elapsedSnapshotAt,
    durationMs,
  };
  if (index >= 0) {
    next[index] = nextItem;
  } else {
    next.push(nextItem);
  }
  return next;
};

export const completeRunningProgress = (progress: ChatProgress[] | undefined) =>
  progress?.map((item) => {
    if (item.status !== "running") {
      return item;
    }
    const endedAt = Date.now();
    return {
      ...item,
      status: "completed" as const,
      endedAt,
      elapsedMs: undefined,
      elapsedSnapshotAt: undefined,
      durationMs:
        item.durationMs ??
        (item.startedAt !== undefined
          ? Math.max(0, endedAt - item.startedAt)
          : item.elapsedMs),
    };
  });

export const cancelRunningTodos = (todoUpdate: AgentTodoUpdate | undefined) =>
  todoUpdate
    ? {
        ...todoUpdate,
        todos: todoUpdate.todos.map((todo) =>
          todo.status === "pending" || todo.status === "in_progress"
            ? {
                ...todo,
                status: "cancelled" as const,
                updatedAt: Date.now(),
              }
            : todo,
        ),
      }
    : undefined;

export const upsertPermission = (
  permissions: AgentPermissionRequest[] | undefined,
  event: StreamEvent & { type: "permission_request" },
) => {
  const next = [...(permissions ?? [])];
  const index = next.findIndex((item) => item.requestId === event.requestId);
  const nextItem: AgentPermissionRequest = {
    requestId: event.requestId,
    sessionId: event.sessionId,
    permission: event.permission,
    patterns: event.patterns,
    target: event.target,
    always: event.always,
    tool: event.tool,
    createdAt: event.createdAt,
    status: "pending",
  };
  if (index >= 0) {
    next[index] = {
      ...next[index],
      ...nextItem,
      status: next[index].status === "submitting" ? "submitting" : nextItem.status,
    };
  } else {
    next.push(nextItem);
  }
  return next;
};

export const toPermissionStatus = (reply: PermissionReply): AgentPermissionRequest["status"] => {
  if (reply === "always") return "approved_always";
  if (reply === "once") return "approved_once";
  return "rejected";
};

export const isActionableQuestionRequest = (question: {
  requestId: string;
  tool?: AgentQuestionRequest["tool"];
}) => Boolean(question.requestId && question.requestId !== question.tool?.callID);

export const toQuestionRequest = (
  event: StreamEvent & { type: "question_request" },
  status: AgentQuestionRequest["status"] = "pending",
): AgentQuestionRequest => ({
  requestId: event.requestId,
  sessionId: event.sessionId,
  questions: event.questions,
  tool: event.tool,
  createdAt: event.createdAt,
  status,
});

export const getQuestionContentSignature = (
  questions: AgentQuestionRequest["questions"],
) =>
  JSON.stringify(
    questions.map((question) => ({
      header: question.header,
      question: question.question,
      options: question.options.map((option) => ({
        label: option.label,
        description: option.description,
      })),
      multiple: question.multiple ?? false,
      custom: question.custom !== false,
    })),
  );

export const isSameQuestionRequest = (
  question: AgentQuestionRequest,
  event: StreamEvent & { type: "question_request" },
) => {
  if (question.requestId === event.requestId) return true;
  if (question.tool?.callID && event.tool?.callID) {
    return question.tool.callID === event.tool.callID;
  }
  return (
    question.status === "pending" &&
    question.sessionId === event.sessionId &&
    getQuestionContentSignature(question.questions) ===
      getQuestionContentSignature(event.questions)
  );
};

export const isSameQuestionPair = (
  left: AgentQuestionRequest,
  right: AgentQuestionRequest,
) => {
  if (left.requestId === right.requestId) return true;
  if (left.tool?.callID && right.tool?.callID) {
    return left.tool.callID === right.tool.callID;
  }
  return (
    left.status === "pending" &&
    right.status === "pending" &&
    left.sessionId === right.sessionId &&
    getQuestionContentSignature(left.questions) ===
      getQuestionContentSignature(right.questions)
  );
};

export const dedupeQuestionsAcrossMessages = (messages: Message[]) => {
  const seen: AgentQuestionRequest[] = [];
  let changed = false;
  const nextMessages = messages.map((message) => {
    if (!message.questions?.length) {
      return message;
    }
    const nextQuestions = message.questions.filter((question) => {
      if (seen.some((existing) => isSameQuestionPair(existing, question))) {
        changed = true;
        return false;
      }
      seen.push(question);
      return true;
    });
    if (nextQuestions.length === message.questions.length) {
      return message;
    }
    return {
      ...message,
      questions: nextQuestions.length ? nextQuestions : undefined,
    };
  });
  return changed ? nextMessages : messages;
};

export const upsertQuestionAcrossMessages = (
  messages: Message[],
  event: StreamEvent & { type: "question_request" },
  assistantMessageId: string,
) => {
  let existing: AgentQuestionRequest | undefined;
  for (const message of messages) {
    const match = message.questions?.find((question) =>
      isSameQuestionRequest(question, event),
    );
    if (match) {
      existing = match;
      break;
    }
  }

  const existingStatus: AgentQuestionRequest["status"] | undefined =
    existing?.status === "submitting" ? "submitting" : undefined;
  const nextQuestion =
    existing &&
    isActionableQuestionRequest(existing) &&
    !isActionableQuestionRequest(event)
      ? {
          ...existing,
          sessionId: event.sessionId,
          questions: event.questions,
          tool: event.tool ?? existing.tool,
          createdAt: event.createdAt,
          status: existingStatus ?? existing.status,
        }
      : toQuestionRequest(event, existingStatus ?? "pending");
  const targetMessageId = existing
    ? messages.find((message) =>
        message.questions?.some((question) => isSameQuestionRequest(question, event)),
      )?.id ?? assistantMessageId
    : assistantMessageId;

  return messages.map((message) => {
    const filteredQuestions = message.questions?.filter(
      (question) => !isSameQuestionRequest(question, event),
    );
    if (message.id !== targetMessageId) {
      return filteredQuestions?.length === message.questions?.length
        ? message
        : {
            ...message,
            questions: filteredQuestions?.length ? filteredQuestions : undefined,
      };
    }

    const nextQuestions = [...(filteredQuestions ?? []), nextQuestion];
    return {
      ...message,
      questions: nextQuestions,
    };
  });
};

export const applyQuestionResponse = (
  questions: AgentQuestionRequest[] | undefined,
  event: StreamEvent & { type: "question_response" },
) =>
  (questions ?? []).map((question) =>
    question.requestId === event.requestId
      ? {
          ...question,
          status: event.rejected ? "rejected" as const : "answered" as const,
          answers: event.answers ?? question.answers,
          repliedAt: Date.now(),
          error: undefined,
        }
      : question,
  );

export const createTodoUpdateFromEvent = (
  event: StreamEvent & { type: "todo_update" },
): AgentTodoUpdate => ({
  sessionId: event.sessionId,
  messageId: event.messageId,
  todos: event.todos,
  createdAt: event.createdAt,
});

export const normalizeSessionTodos = (
  messages: Message[],
  nextTodoUpdate?: AgentTodoUpdate,
  targetAssistantMessageId?: string,
) => {
  let latestTodoUpdate = nextTodoUpdate;
  if (!latestTodoUpdate) {
    for (const message of messages) {
      if (message.todos) {
        latestTodoUpdate = message.todos;
      }
    }
  }

  if (!latestTodoUpdate) {
    return messages;
  }

  const targetMessageId =
    targetAssistantMessageId ??
    [...messages].reverse().find((message) => message.role === "assistant")?.id;
  if (!targetMessageId) {
    return messages;
  }

  let changed = false;
  const nextMessages = messages.map((message) => {
    if (message.id === targetMessageId) {
      if (message.todos === latestTodoUpdate) {
        return message;
      }
      changed = true;
      return {
        ...message,
        todos: latestTodoUpdate,
      };
    }
    if (!message.todos) {
      return message;
    }
    changed = true;
    return {
      ...message,
      todos: undefined,
    };
  });

  return changed ? nextMessages : messages;
};

export const abortOpenPermissionsAfterAbort = (
  permissions: AgentPermissionRequest[] | undefined,
) => {
  if (!permissions?.length) return permissions;
  let changed = false;
  const nextPermissions = permissions.map((permission) => {
    if (
      permission.status !== "pending" &&
      permission.status !== "submitting" &&
      permission.status !== "error"
    ) {
      return permission;
    }
    changed = true;
    return {
      ...permission,
      status: "aborted" as const,
      repliedAt: Date.now(),
      error: undefined,
    };
  });
  return changed ? nextPermissions : permissions;
};

export const rejectOpenQuestionsAfterAbort = (
  questions: AgentQuestionRequest[] | undefined,
) => {
  if (!questions?.length) return questions;
  let changed = false;
  const nextQuestions = questions.map((question) => {
    if (
      question.status !== "pending" &&
      question.status !== "submitting" &&
      question.status !== "error"
    ) {
      return question;
    }
    changed = true;
    return {
      ...question,
      status: "rejected" as const,
      repliedAt: Date.now(),
      error: undefined,
    };
  });
  return changed ? nextQuestions : questions;
};

export const finalizeAssistantMessageAfterAbort = (message: Message): Message => {
  const completedProgress = completeRunningProgress(message.progress);
  const cancelledTodos = cancelRunningTodos(message.todos);
  const abortedPermissions = abortOpenPermissionsAfterAbort(message.permissions);
  const rejectedQuestions = rejectOpenQuestionsAfterAbort(message.questions);
  const hasVisibleOutput =
    message.content.trim().length > 0 ||
    Boolean(message.artifacts?.length) ||
    Boolean(abortedPermissions?.length) ||
    Boolean(rejectedQuestions?.length) ||
    Boolean(completedProgress?.length) ||
    Boolean(cancelledTodos);

  if (!hasVisibleOutput) {
    return message;
  }

  return {
    ...message,
    content: message.content || "⚠️ **请求已中断**",
    isError: true,
    progress: completedProgress,
    permissions: abortedPermissions,
    questions: rejectedQuestions,
    todos: cancelledTodos,
  };
};

export const createUserMessage = (content: string): Message => {
  const id = createId();
  return {
    id,
    role: "user",
    content,
  };
};

export const createAssistantMessage = (): Message => ({
  id: createId(),
  role: "assistant",
  content: "",
});