nuxt-app/app/composables/useAgentChat.ts

import { ref } from "vue";
import { parseJsonEventStream, uiMessageChunkSchema } from "ai";
import { useAuthSession } from "./useAuthSession";
import { useAgentRateLimit } from "./useAgentRateLimit";

export type MessagePartType = "text" | "reasoning" | "tool-call" | "tool-result" | "tool-approval";

export interface MessagePart {
  id: string;
  type: MessagePartType;
  text?: string;
  toolName?: string;
  toolCallId?: string;
  args?: unknown;
  result?: unknown;
  state?: "call" | "result" | "approval-requested" | "approval-responded";
  approvalId?: string;
  approved?: boolean;
  approvalReason?: string;
  isAutomaticApproval?: boolean;
  reasoningLoading?: boolean;
  reasoningDuration?: number;
}

export interface AgentMessage {
  id: string;
  role: "user" | "assistant";
  content: string;
  parts?: MessagePart[];
  modelId?: number | null;
  inputTokens?: number | null;
  outputTokens?: number | null;
  createdAt?: string;
  feedback?: "like" | "dislike" | null;
}

export interface UseAgentChatOptions {
  sessionId: () => string | null;
  modelId: () => number | null;
  enableThinking: () => boolean;
  enableTools: () => boolean;
}

export function useAgentChat(options: UseAgentChatOptions) {
  const { sessionId, modelId, enableThinking, enableTools } = options;
  const auth = useAuthSession();
  const rateLimit = useAgentRateLimit();

  const messages = ref<AgentMessage[]>([]);
  const isLoading = ref(false);
  const errorMessage = ref("");
  const isStopped = ref(false);

  let abortController: AbortController | null = null;

  function generateId(): string {
    return Date.now().toString(36) + Math.random().toString(36).slice(2);
  }

  function getOrCreateLastPart(msg: AgentMessage, type: MessagePartType): MessagePart | null {
    if (!msg.parts) msg.parts = [];
    const last = msg.parts[msg.parts.length - 1];
    if (last && last.type === type) return last;
    return null;
  }

  function appendPart(msg: AgentMessage, part: MessagePart) {
    if (!msg.parts) msg.parts = [];
    msg.parts.push(part);
  }

  function updateLastReasoningDuration(msg: AgentMessage) {
    if (!msg.parts) return;
    for (let i = msg.parts.length - 1; i >= 0; i--) {
      const p = msg.parts[i];
      if (!p) continue;
      if (p.type === "reasoning" && p.reasoningLoading) {
        p.reasoningLoading = false;
        return;
      }
    }
  }

  function cleanupInProgressToolCalls(msg: AgentMessage) {
    if (!msg.parts) return;
    msg.parts = msg.parts.filter(
      (p) => p.type !== "tool-call" || p.state === "result" || p.state === "approval-responded",
    );
  }

  async function loadMessages(sid: string) {
    messages.value = [];
    errorMessage.value = "";
    try {
      const res = await $fetch<{ code: number; data: { messages: AgentMessage[] } }>(
        `/api/agent/sessions/${sid}/messages`,
        { method: "GET" },
      );
      messages.value = (res.data?.messages ?? []).map((m) => ({
        ...m,
        parts: m.parts ? (typeof m.parts === "string" ? JSON.parse(m.parts) : m.parts) : undefined,
      }));
    } catch {
      messages.value = [];
    }
  }

  function buildRequestBody(content: string, opts?: { editMessageId?: string; regenerate?: boolean }) {
    return {
      sessionId: sessionId(),
      content,
      modelId: modelId(),
      enableThinking: enableThinking(),
      enableTools: enableTools(),
      ...opts,
    };
  }

  async function processStream(res: Response, assistantIdx: number) {
    if (!res.body) throw new Error("响应体为空");

    let reasoningStartTime: number | null = null;

    const chunkStream = parseJsonEventStream({
      stream: res.body,
      schema: uiMessageChunkSchema,
    });

    const reader = chunkStream.getReader();
    for (;;) {
      const { done, value: parsed } = await reader.read();
      if (done) break;
      if (!parsed.success) continue;
      const chunk = parsed.value;

      switch (chunk.type) {
        case "reasoning-start": {
          const msg = messages.value[assistantIdx];
          if (!msg) break;
          if (reasoningStartTime === null) reasoningStartTime = Date.now();
          appendPart(msg, { id: generateId(), type: "reasoning", text: "", reasoningLoading: true });
          break;
        }
        case "reasoning-delta": {
          const msg = messages.value[assistantIdx];
          if (!msg) break;
          const part = getOrCreateLastPart(msg, "reasoning");
          if (part) {
            part.text = (part.text ?? "") + chunk.delta;
          }
          break;
        }
        case "reasoning-end": {
          const msg = messages.value[assistantIdx];
          if (msg) updateLastReasoningDuration(msg);
          break;
        }
        case "text-delta": {
          const msg = messages.value[assistantIdx];
          if (!msg) break;
          updateLastReasoningDuration(msg);
          let part = getOrCreateLastPart(msg, "text");
          if (!part) {
            part = { id: generateId(), type: "text", text: "" };
            appendPart(msg, part);
          }
          part.text = (part.text ?? "") + chunk.delta;
          msg.content += chunk.delta;
          break;
        }
        case "error": {
          errorMessage.value = chunk.errorText || "流式响应出错";
          break;
        }
        case "tool-input-available": {
          const msg = messages.value[assistantIdx];
          if (!msg) break;
          updateLastReasoningDuration(msg);
          appendPart(msg, {
            id: generateId(),
            type: "tool-call",
            toolName: chunk.toolName,
            toolCallId: chunk.toolCallId,
            args: chunk.input,
            state: "call",
          });
          break;
        }
        case "tool-output-available": {
          const msg = messages.value[assistantIdx];
          if (!msg || !msg.parts) break;
          const callPart = msg.parts.find((p) => p.type === "tool-call" && p.toolCallId === chunk.toolCallId);
          if (callPart) {
            callPart.result = chunk.output;
            callPart.state = "result";
          }
          break;
        }
        case "tool-output-denied": {
          const msg = messages.value[assistantIdx];
          if (!msg || !msg.parts) break;
          const callPart = msg.parts.find((p) => p.type === "tool-call" && p.toolCallId === chunk.toolCallId);
          if (callPart) {
            callPart.state = "result";
            callPart.result = "工具执行被拒绝";
          }
          break;
        }
        case "tool-approval-request": {
          const msg = messages.value[assistantIdx];
          if (!msg) break;
          updateLastReasoningDuration(msg);
          const isAutomatic = !!(chunk as any).isAutomatic;
          const existingPart = msg.parts?.find((p) => p.type === "tool-call" && p.toolCallId === chunk.toolCallId);
          if (existingPart) {
            existingPart.state = "approval-requested";
            existingPart.approvalId = chunk.approvalId;
            existingPart.isAutomaticApproval = isAutomatic;
          } else {
            appendPart(msg, {
              id: generateId(),
              type: "tool-call",
              toolName: (chunk as any).toolName,
              toolCallId: chunk.toolCallId,
              args: (chunk as any).input,
              state: "approval-requested",
              approvalId: chunk.approvalId,
              isAutomaticApproval: isAutomatic,
            });
          }
          break;
        }
        case "tool-approval-response": {
          const msg = messages.value[assistantIdx];
          if (!msg || !msg.parts) break;
          const part = msg.parts.find((p) => p.type === "tool-call" && p.approvalId === chunk.approvalId);
          if (part) {
            part.state = "approval-responded";
            part.approved = chunk.approved;
            if (chunk.reason) part.approvalReason = chunk.reason;
          }
          break;
        }
        default:
          break;
      }
    }

    const msg = messages.value[assistantIdx];
    if (msg) updateLastReasoningDuration(msg);
  }

  function validateAssistantContent(assistantIdx: number) {
    const finalMsg = messages.value[assistantIdx];
    if (finalMsg && !errorMessage.value) {
      const hasText = finalMsg.parts?.some((p) => p.type === "text" && p.text);
      const hasToolCall = finalMsg.parts?.some((p) => p.type === "tool-call");
      const hasPendingApproval = finalMsg.parts?.some(
        (p) => p.state === "approval-requested" && !p.isAutomaticApproval,
      );
      if (hasPendingApproval) return;
      if (!hasText && !hasToolCall) {
        errorMessage.value = "模型未返回任何内容（可能已达到工具调用次数上限或模型无响应）";
        messages.value.splice(assistantIdx, 1);
      } else if (!hasText && hasToolCall) {
        finalMsg.parts?.push({
          id: generateId(),
          type: "text",
          text: "（已达到工具调用次数上限，模型未能生成最终回答。以上是工具调用的尝试记录。）",
        });
      }
    }
  }

  async function send(
    content: string,
    opts?: { editMessageId?: string; regenerate?: boolean },
  ) {
    const trimmed = content.trim();
    const sid = sessionId();
    if (!trimmed || sid === null || isLoading.value) return;

    errorMessage.value = "";
    isStopped.value = false;

    let userMsgIdx = -1;
    if (!opts?.editMessageId && !opts?.regenerate) {
      messages.value.push({
        id: generateId(),
        role: "user",
        content: trimmed,
      });
      userMsgIdx = messages.value.length - 1;
    }

    const assistantMsg: AgentMessage = {
      id: generateId(),
      role: "assistant",
      content: "",
      parts: [],
    };
    messages.value.push(assistantMsg);
    const assistantIdx = messages.value.length - 1;

    isLoading.value = true;
    abortController = new AbortController();

    try {
      const res = await fetch("/api/agent/chat", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(buildRequestBody(trimmed, opts)),
        signal: abortController.signal,
      });

      if (!auth.loggedIn.value) {
        rateLimit.updateFromHeaders(res.headers);
      }

      const dbUserMessageId = res.headers.get("X-User-Message-Id");
      if (dbUserMessageId && userMsgIdx >= 0) {
        const userMsg = messages.value[userMsgIdx];
        if (userMsg) {
          userMsg.id = dbUserMessageId;
        }
      }

      if (!res.ok) {
        const errText = await res.text();
        throw new Error(errText || `请求失败 (${res.status})`);
      }

      await processStream(res, assistantIdx);
      validateAssistantContent(assistantIdx);
    } catch (err: any) {
      if (err.name === "AbortError") {
        isStopped.value = true;
        const msg = messages.value[assistantIdx];
        if (msg) {
          cleanupInProgressToolCalls(msg);
        }
      } else {
        errorMessage.value = err.message || "请求失败";
        const msg = messages.value[assistantIdx];
        if (msg && !msg.content && (!msg.parts || msg.parts.length === 0)) {
          messages.value.splice(assistantIdx, 1);
        }
      }
    } finally {
      isLoading.value = false;
      abortController = null;
    }
  }

  async function respondToApproval(toolCallId: string, approved: boolean, reason?: string) {
    if (isLoading.value) return;

    const assistantMsg = messages.value.find((m) =>
      m.parts?.some((p) => p.toolCallId === toolCallId && p.state === "approval-requested"),
    );
    if (!assistantMsg) return;

    const approvalPart = assistantMsg.parts?.find(
      (p) => p.toolCallId === toolCallId && p.state === "approval-requested",
    );
    if (!approvalPart) return;

    try {
      await $fetch("/api/agent/chat/tool-approve", {
        method: "POST",
        body: {
          sessionId: sessionId(),
          toolCallId,
          approved,
          reason,
        },
      });
    } catch {
    }

    approvalPart.state = "approval-responded";
    approvalPart.approved = approved;
    approvalPart.approvalReason = reason;

    isLoading.value = true;
    abortController = new AbortController();

    try {
      const res = await fetch("/api/agent/chat", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(buildRequestBody("", { regenerate: false })),
        signal: abortController.signal,
      });

      if (!res.ok) {
        const errText = await res.text();
        throw new Error(errText || `请求失败 (${res.status})`);
      }

      await processStream(res, messages.value.indexOf(assistantMsg));
      validateAssistantContent(messages.value.indexOf(assistantMsg));
    } catch (err: any) {
      if (err.name !== "AbortError") {
        errorMessage.value = err.message || "请求失败";
      }
    } finally {
      isLoading.value = false;
      abortController = null;
    }
  }

  async function sendFeedback(messageId: string, feedback: "like" | "dislike") {
    if (!auth.loggedIn.value) return;
    try {
      await $fetch("/api/agent/feedback", {
        method: "POST",
        body: { messageId, feedback },
      });
      const msg = messages.value.find((m) => m.id === messageId);
      if (msg) {
        msg.feedback = msg.feedback === feedback ? null : feedback;
      }
    } catch {
    }
  }

  function stopGeneration() {
    if (abortController) {
      abortController.abort();
      abortController = null;
    }
  }

  function clear() {
    messages.value = [];
    errorMessage.value = "";
    isStopped.value = false;
  }

  return {
    messages,
    isLoading,
    errorMessage,
    isStopped,
    rateLimit,
    send,
    stopGeneration,
    clear,
    loadMessages,
    respondToApproval,
    sendFeedback,
  };
}