nuxt-app/server/service/agent/chat-engine.ts

import { type LanguageModel, type ModelMessage, streamText, stepCountIs } from "ai";
import {
  getSessionByIdAndUser,
  getMessagesBySession,
  getMaxSortOrder,
  saveMessage,
  truncateMessagesAfter,
  deleteMessage,
  getMessageById,
  updateMessageContent,
  updateMessageParts,
  touchSession,
  countAssistantMessages,
  updateSession,
} from "#server/service/agent/session";
import { getAgentToolsForChatByAgentId, getAgentToolsBySlugs, executeAgentTool } from "#server/service/agent-tool";
import { getMemoryContext } from "#server/service/agent-document";
import { checkRateLimit, incrementRateLimit } from "#server/service/agent/rate-limit";
import { resolveModelForUser, resolveModelAny, toLanguageModel } from "#server/service/llm/model-resolver";
import { generateSessionTitle } from "./title";
import type { AgentRow } from "./agent";
import type { AgentSessionRow } from "./types";
import type { StoredPart } from "#server/service/agent/stored-part";
import { createStreamBuffer, appendChunk, markBufferDone, removeStreamBuffer } from "#server/service/agent/stream-buffer";
import { registerAbortController, unregisterAbortController } from "#server/service/agent/abort-manager";
import log4js from "logger";

const logger = log4js.getLogger("APP");

export interface ChatEngineParams {
  agent: AgentRow;
  session: AgentSessionRow;
  user: { id: number; role: string | null } | null;
  tempToken: string | null;
  body: {
    content: string;
    editMessageId?: string;
    regenerate?: boolean;
    continueAfterApproval?: boolean;
    approvalToolCallId?: string;
    approved?: boolean;
    approvalReason?: string;
    enableThinking?: boolean;
    enableTools?: boolean;
  };
  ip: string;
  preferredModelId?: number | null;
  requestId?: string;
  registerAbort: (sessionId: string, controller: AbortController) => void;
  unregisterAbort: (sessionId: string) => void;
  setRateLimitHeaders?: (headers: Record<string, string>) => void;
  onClientClose?: (cb: () => void) => void;
}

export interface ChatEngineResult {
  response: Response;
  userMessageId: string | null;
}

function buildModelMessages(
  historyMessages: { id: string; role: string; content: string; parts: string | null }[],
  opts: {
    skipMessageId?: string;
    currentUserContent?: string;
    approvalToolResults?: Record<string, unknown>;
  },
): ModelMessage[] {
  const messages: ModelMessage[] = [];

  for (const m of historyMessages) {
    if (opts.skipMessageId && m.id === opts.skipMessageId) continue;

    if (m.role === "assistant" && m.parts) {
      let parts: StoredPart[];
      try {
        parts = JSON.parse(m.parts);
      } catch {
        parts = [];
      }

      const assistantContent: Array<Record<string, unknown>> = [];

      if (m.content) {
        assistantContent.push({ type: "text", text: m.content });
      }

      for (const p of parts) {
        if (p.type === "tool-call" && p.toolCallId && p.toolName) {
          if (p.state === "approval-requested" && p.approvalId) {
            assistantContent.push({
              type: "tool-approval-request",
              approvalId: p.approvalId,
              toolCallId: p.toolCallId,
            });
          } else if (p.state === "approval-responded" && p.approvalId) {
            const hasResult = opts.approvalToolResults?.[p.toolCallId] !== undefined;
            if (p.approved === true && hasResult) {
              assistantContent.push({
                type: "tool-call",
                toolCallId: p.toolCallId,
                toolName: p.toolName,
                input: p.args ?? {},
              });
            } else {
              assistantContent.push({
                type: "tool-call",
                toolCallId: p.toolCallId,
                toolName: p.toolName,
                input: p.args ?? {},
              });
              assistantContent.push({
                type: "tool-approval-request",
                approvalId: p.approvalId,
                toolCallId: p.toolCallId,
              });
            }
          } else {
            const toolCallPart: Record<string, unknown> = {
              type: "tool-call",
              toolCallId: p.toolCallId,
              toolName: p.toolName,
              input: p.args ?? {},
            };
            if (p.state === "result" && p.result !== undefined) {
              toolCallPart.output = p.result;
            }
            assistantContent.push(toolCallPart);
          }
        } else if (p.type === "reasoning" && p.text) {
          assistantContent.push({ type: "reasoning", text: p.text });
        } else if (p.type === "text" && p.text && !m.content) {
          assistantContent.push({ type: "text", text: p.text });
        }
      }

      if (assistantContent.length > 0) {
        messages.push({ role: "assistant", content: assistantContent } as ModelMessage);
      }

      const toolResults: Array<Record<string, unknown>> = [];
      const approvalResponses: Array<Record<string, unknown>> = [];
      for (const p of parts) {
        if (p.type === "tool-call" && p.toolCallId && p.toolName && p.state === "result" && p.result !== undefined) {
          toolResults.push({
            type: "tool-result",
            toolCallId: p.toolCallId,
            toolName: p.toolName,
            output: { type: "json" as const, value: p.result as unknown },
          });
        }
        if (p.type === "tool-call" && p.state === "approval-responded" && p.approvalId) {
          const tcId = p.toolCallId!;
          const hasResult = opts.approvalToolResults?.[tcId] !== undefined;
          if (p.approved === true && hasResult) {
            toolResults.push({
              type: "tool-result",
              toolCallId: tcId,
              toolName: p.toolName!,
              output: { type: "json" as const, value: opts.approvalToolResults![tcId] as unknown },
            });
          } else {
            approvalResponses.push({
              type: "tool-approval-response",
              approvalId: p.approvalId,
              approved: p.approved ?? false,
              reason: p.approvalReason,
            });
          }
        }
      }

      if (toolResults.length > 0 || approvalResponses.length > 0) {
        messages.push({ role: "tool", content: [...toolResults, ...approvalResponses] } as ModelMessage);
      }
    } else if (m.role === "user") {
      messages.push({ role: "user", content: m.content });
    }
  }

  if (opts.currentUserContent !== undefined) {
    messages.push({ role: "user", content: opts.currentUserContent });
  }

  return messages;
}

export async function executeChat(params: ChatEngineParams): Promise<ChatEngineResult> {
  const { agent, session, user, body, ip, requestId = "-" } = params;
  const sessionId = session.id;
  const {
    content,
    editMessageId,
    regenerate,
    continueAfterApproval,
    approvalToolCallId,
    approved,
    approvalReason,
    enableThinking: bodyEnableThinking,
    enableTools: bodyEnableTools,
  } = body;

  const isApprovalContinue = continueAfterApproval === true && approvalToolCallId;

  if (!user) {
    const rateLimit = checkRateLimit(agent.slug, sessionId, ip);
    if (rateLimit.blocked) {
      params.setRateLimitHeaders?.({
        "X-RateLimit-Session-Remaining": String(rateLimit.sessionRemaining),
        "X-RateLimit-Ip-Remaining": String(rateLimit.ipRemaining),
      });
      throw createError({
        statusCode: 429,
        statusMessage:
          rateLimit.sessionRemaining === 0 ? "当前会话已达回复上限，请登录后继续" : "今日回复次数已达上限，请登录或明日再试",
      });
    }
    params.setRateLimitHeaders?.({
      "X-RateLimit-Session-Remaining": String(rateLimit.sessionRemaining),
      "X-RateLimit-Ip-Remaining": String(rateLimit.ipRemaining),
    });
  }

  let resolvedModel;
  if (user) {
    const targetModelId = session.modelId ?? params.preferredModelId ?? agent.defaultModelId;
    if (!targetModelId) {
      throw createError({ statusCode: 400, statusMessage: "未配置模型，请先在设置中选择模型" });
    }
    resolvedModel = await resolveModelForUser(targetModelId, user.id);
  } else {
    const targetModelId = session.modelId ?? agent.defaultModelId;
    if (!targetModelId) {
      throw createError({ statusCode: 400, statusMessage: "系统未配置默认模型" });
    }
    resolvedModel = await resolveModelAny(targetModelId);
  }

  if (!resolvedModel) {
    throw createError({ statusCode: 404, statusMessage: "模型不存在或无权访问" });
  }

  if (!resolvedModel.provider.apiKey) {
    throw createError({ statusCode: 400, statusMessage: "供应商未配置 API Key" });
  }

  const languageModel: LanguageModel = toLanguageModel(resolvedModel);

  const baseSystemPrompt = session.systemPrompt?.trim() || agent.systemPrompt || undefined;

  let systemPrompt = baseSystemPrompt;
  if (user) {
    const memoryContext = await getMemoryContext(user.id);
    if (memoryContext) {
      const memoryInstruction =
        "\n\n# 持久记忆（已注入上下文）\n以下 <memory> 标签内的信息是你的持久记忆，每次对话开始时自动加载，始终是最新的。\n**禁止调用 document 工具的 read 或 list 操作来读取这些记忆** — 它们已经在你的上下文中了。\n直接使用 <memory> 中的内容回答用户问题即可。\n只有需要创建新记忆或修改已有记忆时，才使用 document 工具。";
      systemPrompt = systemPrompt
        ? `${systemPrompt}${memoryInstruction}\n${memoryContext}`
        : `${memoryInstruction}\n${memoryContext}`;
    }
  }

  const sessionEnableTools = session.enableTools !== null ? session.enableTools === 1 : agent.enableTools === 1;
  const toolsEnabled = bodyEnableTools !== undefined ? bodyEnableTools : sessionEnableTools;

  const { tools } = await getAgentToolsForChatByAgentId({
    agentId: agent.id,
    userId: user?.id ?? null,
    userRole: (user?.role as any) ?? null,
    enableTools: toolsEnabled,
  });

  let historyMessages = await getMessagesBySession(sessionId, { limit: 100, latest: true });

  let userMessage;
  let shouldGenerateTitle = false;

  if (isApprovalContinue) {
    userMessage = historyMessages.filter((m) => m.role === "user").pop() ?? historyMessages[0];
    const assistantCount = await countAssistantMessages(sessionId);
    if (assistantCount <= 1) {
      shouldGenerateTitle = true;
    }
  } else if (editMessageId) {
    const editMsg = await getMessageById(editMessageId);
    if (!editMsg || editMsg.sessionId !== sessionId) {
      throw createError({ statusCode: 400, statusMessage: "编辑的消息不存在" });
    }
    await truncateMessagesAfter(sessionId, editMsg.sortOrder);
    await updateMessageContent(editMessageId, content.trim());
    historyMessages = await getMessagesBySession(sessionId, { limit: 100, latest: true });
    userMessage = editMsg;
  } else if (regenerate) {
    const lastAssistant = historyMessages.filter((m) => m.role === "assistant").pop();
    if (lastAssistant) {
      await deleteMessage(lastAssistant.id);
      historyMessages = await getMessagesBySession(sessionId, { limit: 100, latest: true });
    }
    const lastUser = historyMessages.filter((m) => m.role === "user").pop();
    if (lastUser) {
      await updateMessageContent(lastUser.id, content.trim());
      userMessage = lastUser;
    } else {
      const sortOrder = (await getMaxSortOrder(sessionId)) + 1;
      userMessage = await saveMessage({
        sessionId,
        role: "user",
        content: content.trim(),
        sortOrder,
      });
      shouldGenerateTitle = true;
    }
  } else {
    const sortOrder = (await getMaxSortOrder(sessionId)) + 1;
    userMessage = await saveMessage({
      sessionId,
      role: "user",
      content: content.trim(),
      sortOrder,
    });
    const assistantCount = await countAssistantMessages(sessionId);
    if (assistantCount === 0) {
      shouldGenerateTitle = true;
    }
  }

  const approvalToolResults: Record<string, unknown> = {};
  if (isApprovalContinue) {
    const pendingSlugs = new Set<string>();
    for (const m of historyMessages) {
      if (!m.parts) continue;
      let parts: StoredPart[];
      try {
        parts = JSON.parse(m.parts);
      } catch {
        continue;
      }
      for (const p of parts) {
        if (p.type !== "tool-call" || !p.toolCallId || !p.toolName) continue;
        if (p.state !== "approval-responded" || p.approved !== true) continue;
        if (approvalToolResults[p.toolCallId] !== undefined) continue;
        pendingSlugs.add(p.toolName);
      }
    }
    const toolMap = await getAgentToolsBySlugs([...pendingSlugs]);

    for (const m of historyMessages) {
      if (!m.parts) continue;
      let parts: StoredPart[];
      try {
        parts = JSON.parse(m.parts);
      } catch {
        continue;
      }
      for (const p of parts) {
        if (p.type !== "tool-call" || !p.toolCallId || !p.toolName) continue;
        if (p.state !== "approval-responded" || p.approved !== true) continue;
        if (approvalToolResults[p.toolCallId] !== undefined) continue;
        const agentTool = toolMap.get(p.toolName);
        if (!agentTool) {
          approvalToolResults[p.toolCallId] = { error: `Tool ${p.toolName} not found` };
          continue;
        }
        try {
          const result = await executeAgentTool(agentTool.id, p.args, user?.id ?? null, p.approved);
          approvalToolResults[p.toolCallId] = result;
          logger.info("[%s] [APPROVAL-EXEC] tool=%s result=%s", requestId, p.toolName, JSON.stringify(result).slice(0, 200));
        } catch (err) {
          approvalToolResults[p.toolCallId] = { error: err instanceof Error ? err.message : String(err) };
          logger.error("[%s] [APPROVAL-EXEC] tool=%s error=%s", requestId, p.toolName, approvalToolResults[p.toolCallId]);
        }
      }
    }
  }

  const modelMessages = buildModelMessages(historyMessages, {
    skipMessageId: isApprovalContinue ? undefined : userMessage?.id,
    currentUserContent: isApprovalContinue ? undefined : content.trim(),
    approvalToolResults,
  });

  const sessionEnableThinking =
    session.enableThinking !== null ? session.enableThinking === 1 : agent.enableThinking === 1;
  const thinkingEnabled = bodyEnableThinking !== undefined ? bodyEnableThinking : sessionEnableThinking;

  const maxSteps = agent.maxStepCount ?? 8;

  logger.info(
    "[%s] [CHAT-ENGINE] agent=%s userId=%s sessionId=%s thinking=%s tools=%d maxSteps=%d edit=%s regen=%s approval=%s",
    requestId,
    agent.slug,
    user?.id ?? "temp",
    sessionId,
    thinkingEnabled ? "on" : "off",
    Object.keys(tools).length,
    maxSteps,
    editMessageId ? "yes" : "no",
    regenerate ? "yes" : "no",
    isApprovalContinue ? "yes" : "no",
  );

  const streamBuffer = createStreamBuffer(sessionId, {
    modelId: resolvedModel.dbId,
    userMessageId: userMessage?.id ?? null,
  });

  const serverAbortController = new AbortController();
  params.registerAbort(sessionId, serverAbortController);

  let streamedText = "";
  let streamedReasoning = "";
  let streamedToolCalls: Array<Record<string, unknown>> = [];
  let chunkPartId = 0;

  params.onClientClose?.(() => {
    if (!serverAbortController.signal.aborted) {
      logger.info("[%s] [CHAT-ENGINE] client disconnected, aborting streamText", requestId);
      serverAbortController.abort();
    }
  });

  let approvalPrefixChunk: Uint8Array | null = null;
  if (isApprovalContinue && Object.keys(approvalToolResults).length > 0) {
    const prefixChunks = Object.entries(approvalToolResults).map(
      ([toolCallId, output]) =>
        `data: ${JSON.stringify({
          type: "tool-output-available",
          toolCallId,
          output,
        })}\n\n`,
    );
    approvalPrefixChunk = new TextEncoder().encode(prefixChunks.join(""));
    appendChunk(sessionId, approvalPrefixChunk);
    logger.info("[%s] [APPROVAL-PREFIX] written %d bytes to buffer before stream start", requestId, approvalPrefixChunk.length);
  }

  const result = streamText({
    model: languageModel,
    system: systemPrompt,
    messages: modelMessages,
    abortSignal: serverAbortController.signal,
    ...(Object.keys(tools).length > 0
      ? { tools, stopWhen: stepCountIs(maxSteps) }
      : {}),
    ...(thinkingEnabled
      ? {
          providerOptions: {
            openaiCompatible: { reasoningEffort: "high" },
          },
        }
      : {}),
    onError: (errorData) => {
      const err = errorData?.error;
      const errMsg = err instanceof Error ? err.message : String(err ?? "未知错误");
      const errCause = err instanceof Error ? JSON.stringify(err.cause ?? null, null, 2) : "";
      logger.error("[%s] [CHAT-ENGINE] streamText error: %s\ncause: %s", requestId, errMsg, errCause);
    },
    onChunk: ({ chunk }) => {
      if (chunk.type === "text-delta") {
        streamedText += chunk.text;
      } else if (chunk.type === "reasoning-delta") {
        streamedReasoning += chunk.text;
      } else if (chunk.type === "tool-call") {
        streamedToolCalls.push({
          id: `p_${chunkPartId++}`,
          type: "tool-call",
          toolCallId: chunk.toolCallId,
          toolName: chunk.toolName,
          args: chunk.input,
          state: "call",
        });
      } else if (chunk.type === "tool-result") {
        const tcIdx = streamedToolCalls.findIndex((tc) => tc.toolCallId === chunk.toolCallId);
        const tcEntry = tcIdx >= 0 ? streamedToolCalls[tcIdx] : undefined;
        if (tcEntry) {
          tcEntry.result = chunk.output;
          tcEntry.state = "result";
        }
      }
    },
    onFinish: async ({ finishReason, usage, steps, text: assistantText }) => {
      try {
        if (serverAbortController.signal.aborted) {
          logger.info("[%s] [CHAT-ENGINE] onFinish skipped due to abort", requestId);
          params.unregisterAbort(sessionId);
          return;
        }
        logger.info(
          "[%s] [CHAT-ENGINE] finished: reason=%s steps=%d inputTokens=%d outputTokens=%d",
          requestId,
          finishReason,
          steps.length,
          usage?.inputTokens ?? 0,
          usage?.outputTokens ?? 0,
        );

        const parts: Array<Record<string, unknown>> = [];
        for (const step of steps) {
          const toolResultMap = new Map(step.toolResults.map((tr) => [tr.toolCallId, tr.output]));
          const contentArr = step.content as Array<Record<string, unknown>>;
          const approvalToolCallIds = new Set<string>();
          for (const cp of contentArr) {
            if (cp.type === "tool-approval-request" && (cp as any).approvalId && (cp as any).toolCall) {
              approvalToolCallIds.add((cp as any).toolCall.toolCallId);
            }
          }
          for (const cp of contentArr) {
            if (cp.type === "reasoning" && (cp as any).text) {
              parts.push({ id: `p_${parts.length}`, type: "reasoning", text: (cp as any).text });
            } else if (cp.type === "text" && (cp as any).text) {
              parts.push({ id: `p_${parts.length}`, type: "text", text: (cp as any).text });
            } else if (cp.type === "tool-approval-request" && (cp as any).approvalId && (cp as any).toolCall) {
              const tc = (cp as any).toolCall;
              const toolResult = toolResultMap.get(tc.toolCallId);
              if (toolResult === undefined) {
                parts.push({
                  id: `p_${parts.length}`,
                  type: "tool-call",
                  toolCallId: tc.toolCallId,
                  toolName: tc.toolName,
                  args: tc.input,
                  state: "approval-requested",
                  approvalId: (cp as any).approvalId,
                });
              } else {
                parts.push({
                  id: `p_${parts.length}`,
                  type: "tool-call",
                  toolCallId: tc.toolCallId,
                  toolName: tc.toolName,
                  args: tc.input,
                  result: toolResult,
                  state: "result",
                  approvalId: (cp as any).approvalId,
                });
              }
            } else if (cp.type === "tool-call") {
              const tc = cp as any;
              if (approvalToolCallIds.has(tc.toolCallId)) continue;
              const toolResult = toolResultMap.get(tc.toolCallId);
              parts.push({
                id: `p_${parts.length}`,
                type: "tool-call",
                toolCallId: tc.toolCallId,
                toolName: tc.toolName,
                args: tc.input,
                result: toolResult,
                state: "result",
              });
            }
          }
        }

        if (isApprovalContinue) {
          const lastAssistant = historyMessages.filter((m) => m.role === "assistant").pop();
          if (lastAssistant) {
            let existingParts: Array<Record<string, unknown>> = [];
            try {
              existingParts = lastAssistant.parts ? JSON.parse(lastAssistant.parts) : [];
            } catch {
              existingParts = [];
            }
            logger.info("[%s] [APPROVAL-CONTINUE] merged %d existing parts with %d new parts", requestId, existingParts.length, parts.length);
            for (const p of existingParts) {
              if (p.type === "tool-call" && p.state === "approval-responded") {
                const tcId = p.toolCallId as string | undefined;
                if (p.approved === true && tcId && approvalToolResults[tcId] !== undefined) {
                  p.result = approvalToolResults[tcId];
                  p.state = "result";
                } else if (p.approved === false) {
                  p.result = "工具执行被拒绝";
                  p.state = "result";
                }
              }
            }
            const newParts = parts.filter(
              (np) => np.type !== "tool-call" || !existingParts.some((ep) => ep.toolCallId === np.toolCallId),
            );
            const idOffset = existingParts.length;
            for (let i = 0; i < newParts.length; i++) {
              newParts[i]!.id = `p_${idOffset + i}`;
            }
            const mergedParts = [...existingParts, ...newParts];
            const mergedContent = (lastAssistant.content || "") + assistantText;
            await updateMessageContent(lastAssistant.id, mergedContent);
            await updateMessageParts(lastAssistant.id, JSON.stringify(mergedParts));
          }
        } else {
          const assistantSortOrder = (await getMaxSortOrder(sessionId)) + 1;
          await saveMessage({
            sessionId,
            role: "assistant",
            content: assistantText,
            parts: parts.length > 0 ? JSON.stringify(parts) : null,
            modelId: resolvedModel.dbId,
            inputTokens: usage?.inputTokens ?? null,
            outputTokens: usage?.outputTokens ?? null,
            sortOrder: assistantSortOrder,
          });
        }

        await touchSession(sessionId);

        if (!user) {
          incrementRateLimit(agent.slug, sessionId, ip);
          const rateLimit = checkRateLimit(agent.slug, sessionId, ip);
          params.setRateLimitHeaders?.({
            "X-RateLimit-Session-Remaining": String(rateLimit.sessionRemaining),
            "X-RateLimit-Ip-Remaining": String(rateLimit.ipRemaining),
          });
        }

        if (shouldGenerateTitle && assistantText.trim().length > 0) {
          generateSessionTitle({
            sessionId,
            agent,
            firstUserMessage: content.trim(),
            firstAssistantMessage: assistantText,
            userId: user?.id ?? null,
          }).catch((e) => {
            logger.error("[CHAT-ENGINE] title gen failed: %s", e?.message ?? e);
          });
        }
        params.unregisterAbort(sessionId);
      } catch (e) {
        logger.error("[%s] [CHAT-ENGINE] onFinish error: %s", requestId, e instanceof Error ? e.message : String(e));
      }
    },
    onAbort: async () => {
      try {
        const partialText = streamedText;
        logger.info("[%s] [CHAT-ENGINE] aborted: streamedTextLen=%d streamedToolCalls=%d", requestId, partialText.length, streamedToolCalls.length);

        const parts: Array<Record<string, unknown>> = [];
        if (streamedReasoning.trim()) {
          parts.push({ id: `p_${parts.length}`, type: "reasoning", text: streamedReasoning });
        }
        if (partialText.trim()) {
          parts.push({ id: `p_${parts.length}`, type: "text", text: partialText });
        }
        for (const tc of streamedToolCalls) {
          parts.push({
            id: `p_${parts.length}`,
            type: "tool-call",
            toolCallId: tc.toolCallId,
            toolName: tc.toolName,
            args: tc.args,
            result: tc.result,
            state: tc.result !== undefined ? "result" : "call",
          });
        }

        if (isApprovalContinue) {
          const lastAssistant = historyMessages.filter((m) => m.role === "assistant").pop();
          if (lastAssistant) {
            let existingParts: Array<Record<string, unknown>> = [];
            try {
              existingParts = lastAssistant.parts ? JSON.parse(lastAssistant.parts) : [];
            } catch {
              existingParts = [];
            }
            for (const p of existingParts) {
              if (p.type === "tool-call" && p.state === "approval-responded") {
                const tcId = p.toolCallId as string | undefined;
                if (p.approved === true && tcId && approvalToolResults[tcId] !== undefined) {
                  p.result = approvalToolResults[tcId];
                  p.state = "result";
                } else if (p.approved === false) {
                  p.result = "工具执行被拒绝";
                  p.state = "result";
                }
              }
            }
            const newParts = parts.filter(
              (np) => np.type !== "tool-call" || !existingParts.some((ep) => ep.toolCallId === np.toolCallId),
            );
            const idOffset = existingParts.length;
            for (let i = 0; i < newParts.length; i++) {
              newParts[i]!.id = `p_${idOffset + i}`;
            }
            const mergedParts = [...existingParts, ...newParts];
            const mergedContent = (lastAssistant.content || "") + partialText;
            await updateMessageContent(lastAssistant.id, mergedContent);
            await updateMessageParts(lastAssistant.id, JSON.stringify(mergedParts));
          }
        } else {
          const hasContent = partialText.trim().length > 0 || parts.length > 0;
          if (hasContent) {
            const assistantSortOrder = (await getMaxSortOrder(sessionId)) + 1;
            await saveMessage({
              sessionId,
              role: "assistant",
              content: partialText,
              parts: parts.length > 0 ? JSON.stringify(parts) : null,
              modelId: resolvedModel.dbId,
              sortOrder: assistantSortOrder,
            });
          }
        }

        await touchSession(sessionId);

        if (!user) {
          incrementRateLimit(agent.slug, sessionId, ip);
        }

        markBufferDone(sessionId);
        params.unregisterAbort(sessionId);
        logger.info("[%s] [CHAT-ENGINE] abort save complete: contentLen=%d partsCount=%d", requestId, partialText.length, parts.length);
      } catch (abortErr) {
        logger.error("[%s] [CHAT-ENGINE] onAbort error: %s", requestId, abortErr instanceof Error ? abortErr.message : String(abortErr));
        markBufferDone(sessionId);
        params.unregisterAbort(sessionId);
      }
    },
  });

  const responseHeaders: Record<string, string> = {
    "X-User-Message-Id": userMessage?.id ?? "",
  };

  const response = result.toUIMessageStreamResponse({
    sendReasoning: true,
    headers: responseHeaders,
  });

  function wrapStreamWithBuffer(originalBody: ReadableStream<Uint8Array>): ReadableStream<Uint8Array> {
    return new ReadableStream({
      start(controller) {
        const reader = originalBody.getReader();
        let clientDisconnected = false;
        function pump() {
          reader.read().then(({ done, value }) => {
            if (done) {
              markBufferDone(sessionId);
              if (!clientDisconnected) {
                try { controller.close(); } catch { /* already closed */ }
              }
              return;
            }
            if (value) {
              appendChunk(sessionId, value);
              if (!clientDisconnected) {
                try {
                  controller.enqueue(value);
                } catch {
                  clientDisconnected = true;
                }
              }
            }
            pump();
          }).catch((err) => {
            if (serverAbortController.signal.aborted) {
              logger.info("[%s] [CHAT-ENGINE] stream reader aborted by server abort", requestId);
            } else {
              logger.error("[%s] [CHAT-ENGINE] buffer stream read error: %s", requestId, err instanceof Error ? err.message : String(err));
            }
            markBufferDone(sessionId);
            if (!clientDisconnected) {
              try { controller.close(); } catch { /* already closed */ }
            }
          });
        }
        pump();
      },
    });
  }

  if (continueAfterApproval && approvalPrefixChunk && Object.keys(approvalToolResults).length > 0) {
    const originalBody = response.body;
    if (originalBody) {
      const prefix = approvalPrefixChunk;

      const transformed = new ReadableStream({
        async start(controller) {
          controller.enqueue(prefix);
          const reader = originalBody.getReader();
          let clientDisconnected = false;
          try {
            while (true) {
              const { done, value } = await reader.read();
              if (done) break;
              if (value) {
                appendChunk(sessionId, value);
                if (!clientDisconnected) {
                  try {
                    controller.enqueue(value);
                  } catch {
                    clientDisconnected = true;
                  }
                }
              }
            }
          } catch (err) {
            logger.error("[%s] [CHAT-ENGINE] stream read error: %s", requestId, err instanceof Error ? err.message : String(err));
            const errorChunk = new TextEncoder().encode(
              `data: ${JSON.stringify({ type: "error", errorText: "流式响应中断" })}\n\n`,
            );
            appendChunk(sessionId, errorChunk);
            if (!clientDisconnected) {
              try { controller.enqueue(errorChunk); } catch { /* client gone */ }
            }
          } finally {
            markBufferDone(sessionId);
            if (!clientDisconnected) {
              try { controller.close(); } catch { /* already closed */ }
            }
          }
        },
      });

      return {
        response: new Response(transformed, {
          status: response.status,
          statusText: response.statusText,
          headers: response.headers,
        }),
        userMessageId: userMessage?.id ?? null,
      };
    }
  }

  if (response.body) {
    const bufferedBody = wrapStreamWithBuffer(response.body);
    return {
      response: new Response(bufferedBody, {
        status: response.status,
        statusText: response.statusText,
        headers: response.headers,
      }),
      userMessageId: userMessage?.id ?? null,
    };
  }

  return {
    response,
    userMessageId: userMessage?.id ?? null,
  };
}