TJWaterAgent/src/routes/chatAuxiliaryRoutes.ts

import { type Router } from "express";
import { z } from "zod";

import { getAgentAuthContext } from "../auth/agentAuth.js";
import { type ChatSessionBridge } from "../chat/sessionBridge.js";
import { logger } from "../logger.js";
import { type ResultReferenceResolver } from "../results/resolver.js";
import { RESULT_REFERENCE_KIND } from "../results/store.js";
import { type SessionMetadataStore } from "../sessions/metadataStore.js";
import { type SessionUiStateStore } from "../sessions/uiStateStore.js";
import { toActorKey, toProjectKey } from "../utils/fileStore.js";
import {
  type ActiveRun,
  type RunStatus,
  cancelBackendTodos,
  completeBackendProgress,
  updateLastAssistantMessage,
} from "./chatUiState.js";

const abortPayloadSchema = z.object({
  session_id: z.string().max(128),
});

type RegisterAuxiliaryRoutesOptions = {
  activeRuns: Map<string, ActiveRun>;
  lastRunStatuses: Map<string, RunStatus>;
  resultReferenceResolver: ResultReferenceResolver;
  sessionBridge: ChatSessionBridge;
  sessionMetadataStore: SessionMetadataStore;
  sessionUiStateStore: SessionUiStateStore;
};

const toSessionUiStateContext = (sessionId: string) => ({
  sessionId,
});

export const registerChatAuxiliaryRoutes = (
  chatRouter: Router,
  {
    activeRuns,
    lastRunStatuses,
    resultReferenceResolver,
    sessionBridge,
    sessionMetadataStore,
    sessionUiStateStore,
  }: RegisterAuxiliaryRoutesOptions,
) => {
  chatRouter.get("/render-ref/:renderRef", async (req, res) => {
    const renderRef = req.params.renderRef?.trim();
    const authContext = getAgentAuthContext(req);
    const userId = authContext.userId;
    const projectId = authContext.projectId;
    const clientSessionId =
      typeof req.query.session_id === "string"
        ? req.query.session_id.trim()
        : undefined;

    if (!renderRef) {
      res.status(400).json({
        message: "render_ref is required",
      });
      return;
    }

    const result = await resultReferenceResolver.getFullAuthorized(
      renderRef,
      {
        actorKey: toActorKey(userId),
        clientSessionId,
        projectId,
      },
      {
        expectedKind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
      },
    );

    if (!result) {
      res.status(404).json({ message: "render_ref not found" });
      return;
    }

    res.json(result);
  });

  chatRouter.post("/abort", async (req, res) => {
    const parsed = abortPayloadSchema.safeParse(req.body);
    if (!parsed.success) {
      res.status(400).json({
        message: "invalid request payload",
        detail: parsed.error.flatten(),
      });
      return;
    }

    try {
      const authContext = getAgentAuthContext(req);
      const projectId = authContext.projectId;
      const userId = authContext.userId;
      const actorKey = toActorKey(userId);
      const projectKey = toProjectKey(projectId);
      const sessionRecord = await sessionMetadataStore.get(
        { actorKey, projectId, projectKey, userId },
        parsed.data.session_id,
      );
      const binding = sessionRecord
        ? await sessionBridge.abort({
            clientSessionId: sessionRecord.sessionId,
            sessionId: sessionRecord.sessionId,
          })
        : null;
      const run = activeRuns.get(parsed.data.session_id);
      if (run && run.status === "running") {
        run.status = "aborted";
        lastRunStatuses.set(parsed.data.session_id, "aborted");
        run.controller.abort();
        run.messages = updateLastAssistantMessage(run.messages, (message) => ({
          ...message,
          content:
            typeof message.content === "string" && message.content.trim()
              ? message.content
              : "⚠️ **请求已中断**",
          isError: true,
          progress: completeBackendProgress(message.progress),
          todos: cancelBackendTodos(message.todos),
        }));
        if (sessionRecord) {
          const currentState = await sessionUiStateStore.read(
            toSessionUiStateContext(sessionRecord.sessionId),
          );
          await sessionUiStateStore.write(toSessionUiStateContext(sessionRecord.sessionId), {
            sessionId: sessionRecord.sessionId,
            isTitleManuallyEdited: currentState?.isTitleManuallyEdited ?? false,
            messages: run.messages,
          });
        }
        for (const subscriber of run.subscribers) {
          subscriber.write("error", {
            session_id: parsed.data.session_id,
            message: "请求已中断",
          });
          subscriber.close();
        }
        run.subscribers.clear();
      }

      if (!binding && !run) {
        res.status(204).end();
        return;
      }

      logger.info(
        {
          clientSessionId: parsed.data.session_id,
          sessionId: binding?.sessionId ?? parsed.data.session_id,
        },
        "aborted chat session by client request",
      );
      res.status(202).json({
        session_id: parsed.data.session_id,
        aborted: true,
      });
    } catch (error) {
      const detail = error instanceof Error ? error.message : String(error);
      logger.error({ err: error }, "chat abort failed");
      res.status(500).json({
        message: "chat abort failed",
        detail,
      });
    }
  });
};