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

import type {
  AgentChatMessage,
  AgentChatProgress,
  AgentInteractionToolRef,
  AgentPermissionRequest,
  AgentPermissionStatus,
  AgentQuestionInfo,
  AgentQuestionRequest,
  AgentTodoItem,
  AgentTodoUpdate
} from "./types";

type RecordValue = Record<string, unknown>;

export function updateMessageById(
  messages: AgentChatMessage[],
  messageId: string,
  updater: (message: AgentChatMessage) => AgentChatMessage
) {
  return messages.map((message) => (message.id === messageId ? updater(message) : message));
}

export function upsertProgress(progress: AgentChatProgress[] | undefined, data: unknown) {
  const payload = asRecord(data);
  const id = readString(payload.id) ?? `progress-${Date.now().toString(36)}`;
  const now = Date.now();
  const index = progress?.findIndex((item) => item.id === id) ?? -1;
  const existing = index >= 0 ? progress?.[index] : undefined;
  const status = readProgressStatus(payload.status);
  const startedAt = readNumber(payload.started_at) ?? readNumber(payload.startedAt) ?? existing?.startedAt;
  const endedAt = readNumber(payload.ended_at) ?? readNumber(payload.endedAt);
  const elapsedMs = readNumber(payload.elapsed_ms) ?? readNumber(payload.elapsedMs);
  const durationMs = readNumber(payload.duration_ms) ?? readNumber(payload.durationMs);
  const nextItem: AgentChatProgress = {
    id,
    phase: readString(payload.phase) ?? "progress",
    status,
    title: readString(payload.title) ?? "正在处理",
    detail: readString(payload.detail),
    startedAt,
    endedAt,
    elapsedMs: status === "running" ? elapsedMs : undefined,
    elapsedSnapshotAt: status === "running" && elapsedMs !== undefined ? now : undefined,
    durationMs: status === "running" ? undefined : durationMs
  };
  const next = [...(progress ?? [])];

  if (index >= 0) {
    next[index] = nextItem;
  } else {
    next.push(nextItem);
  }

  return next;
}

export function completeRunningProgress(progress: AgentChatProgress[] | undefined) {
  return 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 function toTodoUpdate(data: unknown): AgentTodoUpdate | null {
  const payload = asRecord(data);
  const sessionId = readString(payload.session_id) ?? readString(payload.sessionId);
  const rawTodos = Array.isArray(payload.todos) ? payload.todos : [];

  if (!sessionId || rawTodos.length === 0) {
    return null;
  }

  return {
    sessionId,
    messageId: readString(payload.message_id) ?? readString(payload.messageId),
    todos: rawTodos.map(toTodoItem).filter((item): item is AgentTodoItem => Boolean(item)),
    createdAt: readNumber(payload.created_at) ?? readNumber(payload.createdAt) ?? Date.now()
  };
}

export function cancelRunningTodos(todoUpdate: AgentTodoUpdate | undefined) {
  return 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 function upsertPermission(
  permissions: AgentPermissionRequest[] | undefined,
  data: unknown
) {
  const payload = asRecord(data);
  const requestId = readString(payload.request_id) ?? readString(payload.requestId);
  const sessionId = readString(payload.session_id) ?? readString(payload.sessionId);

  if (!requestId || !sessionId) {
    return permissions;
  }

  const nextItem: AgentPermissionRequest = {
    requestId,
    sessionId,
    permission: readString(payload.permission) ?? "unknown",
    patterns: readStringArray(payload.patterns),
    target: readString(payload.target),
    always: readStringArray(payload.always),
    tool: toToolRef(payload.tool),
    createdAt: readNumber(payload.created_at) ?? readNumber(payload.createdAt) ?? Date.now(),
    status: "pending"
  };
  const next = [...(permissions ?? [])];
  const index = next.findIndex((permission) => permission.requestId === requestId);

  if (index >= 0) {
    next[index] = {
      ...next[index],
      ...nextItem,
      status: next[index].status === "submitting" ? "submitting" : nextItem.status
    };
  } else {
    next.push(nextItem);
  }

  return next;
}

export function applyPermissionResponse(
  permissions: AgentPermissionRequest[] | undefined,
  data: unknown
) {
  const payload = asRecord(data);
  const requestId = readString(payload.request_id) ?? readString(payload.requestId);
  if (!requestId) {
    return permissions;
  }

  return permissions?.map((permission) =>
    permission.requestId === requestId
      ? {
          ...permission,
          status: toPermissionStatus(readString(payload.reply)),
          repliedAt: Date.now(),
          error: undefined
        }
      : permission
  );
}

export function upsertQuestion(questions: AgentQuestionRequest[] | undefined, data: unknown) {
  const payload = asRecord(data);
  const requestId = readString(payload.request_id) ?? readString(payload.requestId);
  const sessionId = readString(payload.session_id) ?? readString(payload.sessionId);
  const rawQuestions = Array.isArray(payload.questions) ? payload.questions : [];

  if (!requestId || !sessionId || rawQuestions.length === 0) {
    return questions;
  }

  const nextItem: AgentQuestionRequest = {
    requestId,
    sessionId,
    questions: rawQuestions.map(toQuestionInfo).filter((item): item is AgentQuestionInfo => Boolean(item)),
    tool: toToolRef(payload.tool),
    createdAt: readNumber(payload.created_at) ?? readNumber(payload.createdAt) ?? Date.now(),
    status: "pending"
  };
  const next = [...(questions ?? [])];
  const index = next.findIndex((question) => isSameQuestion(question, nextItem));

  if (index >= 0) {
    next[index] = {
      ...next[index],
      ...nextItem,
      status: next[index].status === "submitting" ? "submitting" : nextItem.status
    };
  } else {
    next.push(nextItem);
  }

  return next;
}

export function applyQuestionResponse(questions: AgentQuestionRequest[] | undefined, data: unknown) {
  const payload = asRecord(data);
  const requestId = readString(payload.request_id) ?? readString(payload.requestId);
  if (!requestId) {
    return questions;
  }

  const answers = Array.isArray(payload.answers)
    ? payload.answers.map((answer) => (Array.isArray(answer) ? answer.map(String) : []))
    : undefined;
  const rejected = payload.rejected === true;

  return questions?.map((question) =>
    question.requestId === requestId
      ? {
          ...question,
          status: rejected ? ("rejected" as const) : ("answered" as const),
          answers: answers ?? question.answers,
          repliedAt: Date.now(),
          error: undefined
        }
      : question
  );
}

export function finalizeAssistantMessageAfterAbort(message: AgentChatMessage): AgentChatMessage {
  const progress = completeRunningProgress(message.progress);
  const todos = cancelRunningTodos(message.todos);

  return {
    ...message,
    content: message.content || "请求已中止。",
    progress,
    todos,
    permissions: message.permissions?.map((permission) =>
      permission.status === "pending" || permission.status === "submitting" || permission.status === "error"
        ? {
            ...permission,
            status: "aborted" as const,
            repliedAt: Date.now(),
            error: undefined
          }
        : permission
    ),
    questions: message.questions?.map((question) =>
      question.status === "pending" || question.status === "submitting" || question.status === "error"
        ? {
            ...question,
            status: "rejected" as const,
            repliedAt: Date.now(),
            error: undefined
          }
        : question
    )
  };
}

function toPermissionStatus(reply: string | undefined): AgentPermissionStatus {
  if (reply === "always") {
    return "approved_always";
  }
  if (reply === "once") {
    return "approved_once";
  }
  return "rejected";
}

function isSameQuestion(left: AgentQuestionRequest, right: AgentQuestionRequest) {
  if (left.requestId === right.requestId) {
    return true;
  }
  return Boolean(left.tool?.callID && right.tool?.callID && left.tool.callID === right.tool.callID);
}

function toQuestionInfo(value: unknown): AgentQuestionInfo | null {
  const payload = asRecord(value);
  const question = readString(payload.question);
  if (!question) {
    return null;
  }

  return {
    header: readString(payload.header) ?? "",
    question,
    options: Array.isArray(payload.options)
      ? payload.options
          .map((option) => {
            const item = asRecord(option);
            const label = readString(item.label);
            return label
              ? {
                  label,
                  description: readString(item.description) ?? ""
                }
              : null;
          })
          .filter((item): item is { label: string; description: string } => Boolean(item))
      : [],
    multiple: typeof payload.multiple === "boolean" ? payload.multiple : undefined,
    custom: typeof payload.custom === "boolean" ? payload.custom : undefined
  };
}

function toTodoItem(value: unknown): AgentTodoItem | null {
  const payload = asRecord(value);
  const id = readString(payload.id);
  const content = readString(payload.content);
  const status = readTodoStatus(payload.status);

  if (!id || !content || !status) {
    return null;
  }

  return {
    id,
    content,
    status,
    priority: readTodoPriority(payload.priority),
    createdAt: readNumber(payload.created_at) ?? readNumber(payload.createdAt),
    updatedAt: readNumber(payload.updated_at) ?? readNumber(payload.updatedAt)
  };
}

function toToolRef(value: unknown): AgentInteractionToolRef | undefined {
  const payload = asRecord(value);
  const messageID = readString(payload.messageID);
  const callID = readString(payload.callID);
  return messageID && callID ? { messageID, callID } : undefined;
}

function readProgressStatus(value: unknown): AgentChatProgress["status"] {
  return value === "completed" || value === "error" ? value : "running";
}

function readTodoStatus(value: unknown): AgentTodoItem["status"] | undefined {
  if (value === "pending" || value === "in_progress" || value === "completed" || value === "cancelled") {
    return value;
  }
  return undefined;
}

function readTodoPriority(value: unknown): AgentTodoItem["priority"] | undefined {
  if (value === "low" || value === "medium" || value === "high") {
    return value;
  }
  return undefined;
}

function asRecord(value: unknown): RecordValue {
  return typeof value === "object" && value !== null && !Array.isArray(value)
    ? (value as RecordValue)
    : {};
}

function readString(value: unknown) {
  return typeof value === "string" && value.trim() ? value.trim() : undefined;
}

function readNumber(value: unknown) {
  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
}

function readStringArray(value: unknown) {
  return Array.isArray(value) ? value.map((item) => String(item).trim()).filter(Boolean) : [];
}