392 lines
10 KiB
TypeScript
392 lines
10 KiB
TypeScript
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);
|
||
}
|