ai-chatbot/lib/kittycad/modeling-session.ts

import "server-only";

import { randomUUID } from "node:crypto";
import { deserialize } from "bson";
import WebSocket, { type RawData } from "ws";
import { getKittyCadProviderConfig } from "./server";

type KittyCadModelingCommand = Record<string, unknown> & {
  type: string;
};

type KittyCadSuccessResponse = {
  request_id?: string | null;
  resp?: {
    data?: Record<string, unknown> | null;
    type?: string | null;
  } | null;
  success: true;
};

type KittyCadFailureResponse = {
  errors?: Record<string, unknown>[];
  request_id?: string | null;
  success: false;
};

export type KittyCadModelingSocketResponse =
  | KittyCadSuccessResponse
  | KittyCadFailureResponse;

type PendingRequest = {
  reject: (error: Error) => void;
  resolve: (value: KittyCadModelingSocketResponse) => void;
  timeout: NodeJS.Timeout;
};

type ModelingSessionState = {
  createdAt: string;
  id: string;
  lastUsedAt: number;
  pending: Map<string, PendingRequest>;
  pingTimer: NodeJS.Timeout;
  socket: WebSocket;
};

const MODELING_COMMAND_TIMEOUT_MS = 20_000;
const MODELING_PING_INTERVAL_MS = 20_000;
const DEFAULT_WS_PATH = "/ws/modeling/commands";

function getModelingSessions() {
  const globalState = globalThis as Record<string, unknown>;
  if (!globalState.__kittyCadModelingSessions) {
    globalState.__kittyCadModelingSessions = new Map<
      string,
      ModelingSessionState
    >();
  }
  return globalState.__kittyCadModelingSessions as Map<
    string,
    ModelingSessionState
  >;
}

function toWebSocketUrl(baseUrl: string) {
  const parsed = new URL(baseUrl);
  parsed.protocol = parsed.protocol === "http:" ? "ws:" : "wss:";
  parsed.pathname = DEFAULT_WS_PATH;
  parsed.search = "";
  parsed.hash = "";
  return parsed.toString();
}

function toBuffer(raw: RawData): Buffer {
  if (typeof raw === "string") {
    return Buffer.from(raw, "utf8");
  }

  if (Buffer.isBuffer(raw)) {
    return raw;
  }

  if (raw instanceof ArrayBuffer) {
    return Buffer.from(raw);
  }

  if (Array.isArray(raw)) {
    return Buffer.concat(raw.map((chunk) => toBuffer(chunk)));
  }

  return Buffer.from(raw);
}

function parseSocketMessage(raw: RawData): KittyCadModelingSocketResponse {
  const buffer = toBuffer(raw);
  const text = buffer.toString("utf8").trim();

  if (text) {
    try {
      return JSON.parse(text) as KittyCadModelingSocketResponse;
    } catch (_error) {
      // Fall through to BSON parsing.
    }
  }

  return deserialize(buffer) as KittyCadModelingSocketResponse;
}

function getSocketErrorMessage(error: unknown) {
  return error instanceof Error ? error.message : "官方建模连接异常断开";
}

function formatModelingErrors(errors: unknown) {
  if (!Array.isArray(errors) || errors.length === 0) {
    return "官方建模命令执行失败";
  }

  return (
    errors
      .map((item) => {
        if (!item || typeof item !== "object") {
          return null;
        }

        const record = item as Record<string, unknown>;
        if (typeof record.message === "string" && record.message.trim()) {
          return record.message.trim();
        }
        if (typeof record.error === "string" && record.error.trim()) {
          return record.error.trim();
        }
        return null;
      })
      .filter((item): item is string => Boolean(item))
      .join("；") || "官方建模命令执行失败"
  );
}

function rejectAllPendingRequests(
  state: ModelingSessionState,
  message: string
) {
  for (const [requestId, pending] of state.pending.entries()) {
    clearTimeout(pending.timeout);
    pending.reject(new Error(message));
    state.pending.delete(requestId);
  }
}

function destroySession(sessionId: string, reason: string) {
  const sessions = getModelingSessions();
  const state = sessions.get(sessionId);
  if (!state) {
    return;
  }

  clearInterval(state.pingTimer);
  rejectAllPendingRequests(state, reason);
  sessions.delete(sessionId);

  if (
    state.socket.readyState === WebSocket.OPEN ||
    state.socket.readyState === WebSocket.CONNECTING
  ) {
    state.socket.close();
  }
}

function getRequiredSession(sessionId: string) {
  const sessions = getModelingSessions();
  const state = sessions.get(sessionId);
  if (!state) {
    throw new Error("官方建模会话不存在或已过期，请重新连接");
  }
  if (state.socket.readyState !== WebSocket.OPEN) {
    destroySession(sessionId, "官方建模会话已断开，请重新连接");
    throw new Error("官方建模会话已断开，请重新连接");
  }
  state.lastUsedAt = Date.now();
  return state;
}

export async function createKittyCadModelingSession() {
  const config = await getKittyCadProviderConfig();
  const sessionId = randomUUID();
  const createdAt = new Date().toISOString();
  const socketUrl = toWebSocketUrl(config.baseUrl);

  const sessionState = await new Promise<ModelingSessionState>(
    (resolve, reject) => {
      const socket = new WebSocket(socketUrl, {
        headers: {
          Authorization: `Bearer ${String(config.apiKey || "").replace(/\s+/g, "")}`,
        },
      });

      const cleanup = () => {
        socket.removeListener("error", handleError);
        socket.removeListener("unexpected-response", handleUnexpectedResponse);
        socket.removeListener("open", handleOpen);
      };

      const handleError = (error: Error) => {
        cleanup();
        reject(error);
      };

      const handleUnexpectedResponse = (
        _request: unknown,
        response: {
          on: (event: string, listener: (chunk?: Buffer) => void) => void;
          statusCode?: number;
          statusMessage?: string;
        }
      ) => {
        const chunks: Buffer[] = [];
        response.on("data", (chunk) => {
          if (chunk) {
            chunks.push(Buffer.from(chunk));
          }
        });
        response.on("end", () => {
          cleanup();
          const responseText = Buffer.concat(chunks).toString("utf8").trim();
          reject(
            new Error(
              responseText ||
                `官方建模连接失败（HTTP ${response.statusCode || 502}${response.statusMessage ? ` ${response.statusMessage}` : ""}）`
            )
          );
        });
      };

      const handleOpen = () => {
        cleanup();

        const pingTimer = setInterval(() => {
          if (socket.readyState === WebSocket.OPEN) {
            socket.send(JSON.stringify({ type: "ping" }));
          }
        }, MODELING_PING_INTERVAL_MS);

        resolve({
          createdAt,
          id: sessionId,
          lastUsedAt: Date.now(),
          pending: new Map(),
          pingTimer,
          socket,
        });
      };

      socket.once("error", handleError);
      socket.once("unexpected-response", handleUnexpectedResponse);
      socket.once("open", handleOpen);
    }
  );

  sessionState.socket.on("message", (raw: RawData) => {
    try {
      const response = parseSocketMessage(raw);
      const requestId = String(response.request_id || "").trim();
      if (!requestId) {
        return;
      }

      const pending = sessionState.pending.get(requestId);
      if (!pending) {
        return;
      }

      clearTimeout(pending.timeout);
      sessionState.pending.delete(requestId);
      pending.resolve(response);
    } catch (error) {
      console.error("[KittyCADModeling] failed to parse socket message", error);
    }
  });

  sessionState.socket.on("close", () => {
    destroySession(sessionId, "官方建模会话已关闭，请重新连接");
  });

  sessionState.socket.on("error", (error: Error) => {
    destroySession(sessionId, getSocketErrorMessage(error));
  });

  getModelingSessions().set(sessionId, sessionState);

  return {
    createdAt,
    sessionId,
  };
}

export async function sendKittyCadModelingCommand(params: {
  cmd: KittyCadModelingCommand;
  cmdId?: string;
  sessionId: string;
}) {
  const state = getRequiredSession(params.sessionId);
  const requestId = String(params.cmdId || "").trim() || randomUUID();

  if (state.pending.has(requestId)) {
    throw new Error("官方建模命令 ID 冲突，请重试");
  }

  const payload = {
    cmd: params.cmd,
    cmd_id: requestId,
    type: "modeling_cmd_req",
  };

  return await new Promise<KittyCadModelingSocketResponse>(
    (resolve, reject) => {
      const timeout = setTimeout(() => {
        state.pending.delete(requestId);
        reject(new Error("官方建模命令超时，请稍后重试"));
      }, MODELING_COMMAND_TIMEOUT_MS);

      state.pending.set(requestId, {
        reject,
        resolve,
        timeout,
      });

      state.socket.send(JSON.stringify(payload), (error?: Error) => {
        if (!error) {
          return;
        }

        clearTimeout(timeout);
        state.pending.delete(requestId);
        reject(new Error(getSocketErrorMessage(error)));
      });
    }
  );
}

export async function takeKittyCadModelingSnapshot(params: {
  format?: "jpeg" | "png";
  sessionId: string;
}) {
  const format = params.format || "png";
  const response = await sendKittyCadModelingCommand({
    cmd: {
      format,
      type: "take_snapshot",
    },
    sessionId: params.sessionId,
  });

  if (!response.success) {
    throw new Error(formatModelingErrors(response.errors));
  }

  const contents = String(response.resp?.data?.contents || "").trim();
  if (!contents) {
    return {
      requestId: response.request_id || null,
      snapshotUrl: null,
    };
  }

  const mimeType = format === "jpeg" ? "image/jpeg" : "image/png";
  return {
    requestId: response.request_id || null,
    snapshotUrl: `data:${mimeType};base64,${contents}`,
  };
}

export function closeKittyCadModelingSession(sessionId: string) {
  destroySession(sessionId, "官方建模会话已手动关闭");
  return {
    sessionId,
  };
}

export function isKittyCadModelingCommandSuccessful(
  response: KittyCadModelingSocketResponse
): response is KittyCadSuccessResponse {
  return response.success;
}

export function getKittyCadModelingCommandError(response: unknown) {
  if (!response || typeof response !== "object") {
    return "官方建模命令执行失败";
  }

  const record = response as KittyCadFailureResponse;
  return formatModelingErrors(record.errors);
}