@citeark/agent 0.3.20

package/src/autodl/client.mjs

@@ -0,0 +1,260 @@
+import https from "node:https";
+
+import { CiteArkError } from "../util.mjs";
+
+const RETRYABLE_HTTP_STATUS = new Set([408, 425, 429, 500, 502, 503, 504]);
+const DEFAULT_BASE_URL = "https://api.autodl.com";
+
+export class AutoDlApiError extends CiteArkError {
+  constructor(message, { status = null, code = null, retryable = false, cause } = {}) {
+    super(message, { cause });
+    this.name = "AutoDlApiError";
+    this.status = status;
+    this.code = code;
+    this.retryable = retryable;
+  }
+}
+
+export class AutoDlProClient {
+  constructor({
+    token,
+    baseUrl = DEFAULT_BASE_URL,
+    transport = nodeHttpsJsonRequest,
+    sleep = defaultSleep,
+    timeoutMs = 30_000,
+    maxAttempts = 5,
+  } = {}) {
+    this.token = required(token, "token");
+    this.baseUrl = required(baseUrl, "baseUrl").replace(/\/$/, "");
+    this.transport = transport;
+    this.sleep = sleep;
+    this.timeoutMs = positiveInteger(timeoutMs, "timeoutMs");
+    this.maxAttempts = positiveInteger(maxAttempts, "maxAttempts");
+  }
+
+  async createInstance(input) {
+    const data = await this.#request("POST", "/api/v1/dev/instance/pro/create", {
+      data_center_list: stringArray(input.dataCenterList),
+      req_gpu_amount: positiveInteger(input.gpuCount ?? 1, "gpuCount"),
+      expand_system_disk_by_gb: nonNegativeInteger(
+        input.expandSystemDiskGb ?? 0,
+        "expandSystemDiskGb",
+      ),
+      gpu_spec_uuid: required(input.gpuSpecUuid, "gpuSpecUuid"),
+      image_uuid: required(input.imageUuid, "imageUuid"),
+      cuda_v_from: positiveInteger(input.cudaVersionFrom ?? 118, "cudaVersionFrom"),
+      instance_name: optional(input.instanceName),
+      start_command: optional(input.startCommand),
+    });
+    return required(data, "createInstance.data");
+  }
+
+  async getInstance(instanceUuid) {
+    return await this.#request("GET", "/api/v1/dev/instance/pro/snapshot", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+    });
+  }
+
+  async getStatus(instanceUuid) {
+    return await this.#request("GET", "/api/v1/dev/instance/pro/status", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+    });
+  }
+
+  async listInstances({ pageIndex = 1, pageSize = 100 } = {}) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/list", {
+      page_index: positiveInteger(pageIndex, "pageIndex"),
+      page_size: positiveInteger(pageSize, "pageSize"),
+    });
+  }
+
+  async saveInstanceImage(instanceUuid, imageName) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/image/save", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+      image_name: required(imageName, "imageName"),
+    });
+  }
+
+  async listPrivateImages({ pageIndex = 1, pageSize = 100 } = {}) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/image/private/list", {
+      page_index: positiveInteger(pageIndex, "pageIndex"),
+      page_size: positiveInteger(pageSize, "pageSize"),
+    });
+  }
+
+  async powerOn(instanceUuid, { startCommand } = {}) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/power_on", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+      payload: "gpu",
+      start_command: optional(startCommand),
+    });
+  }
+
+  async powerOff(instanceUuid) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/power_off", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+    });
+  }
+
+  async release(instanceUuid) {
+    return await this.#request("POST", "/api/v1/dev/instance/pro/release", {
+      instance_uuid: required(instanceUuid, "instanceUuid"),
+    });
+  }
+
+  async waitForStatus(instanceUuid, acceptedStatuses, {
+    timeoutMs = 15 * 60_000,
+    pollIntervalMs = 5_000,
+    onStatus = () => {},
+  } = {}) {
+    const accepted = new Set(
+      (Array.isArray(acceptedStatuses) ? acceptedStatuses : [acceptedStatuses])
+        .map((value) => String(value).trim().toLowerCase())
+        .filter(Boolean),
+    );
+    if (!accepted.size) throw new TypeError("AutoDL acceptedStatuses 不能为空");
+    const deadline = Date.now() + positiveInteger(timeoutMs, "timeoutMs");
+    let previous = null;
+    while (Date.now() < deadline) {
+      const status = String(await this.getStatus(instanceUuid)).trim().toLowerCase();
+      if (status !== previous) {
+        previous = status;
+        await onStatus(status);
+      }
+      if (accepted.has(status)) return status;
+      if (new Set(["error", "failed", "released"]).has(status)) {
+        throw new AutoDlApiError(`AutoDL 实例 ${instanceUuid} 进入异常状态：${status}`);
+      }
+      await this.sleep(Math.min(pollIntervalMs, Math.max(0, deadline - Date.now())));
+    }
+    throw new AutoDlApiError(
+      `等待 AutoDL 实例 ${instanceUuid} 进入 ${[...accepted].join("/")} 超时`,
+      { retryable: true },
+    );
+  }
+
+  async #request(method, pathname, body) {
+    let lastError;
+    for (let attempt = 1; attempt <= this.maxAttempts; attempt += 1) {
+      try {
+        const requestUrl = new URL(`${this.baseUrl}${pathname}`);
+        const requestBody = method === "GET" ? undefined : body;
+        if (method === "GET") {
+          for (const [key, value] of Object.entries(body ?? {})) {
+            if (value !== undefined && value !== null) {
+              requestUrl.searchParams.set(key, String(value));
+            }
+          }
+        }
+        const response = await this.transport({
+          url: requestUrl.toString(),
+          method,
+          headers: {
+            authorization: this.token,
+            ...(requestBody === undefined ? {} : { "content-type": "application/json" }),
+          },
+          body: requestBody,
+          timeoutMs: this.timeoutMs,
+        });
+        const payload = isRecord(response.body) ? response.body : {};
+        const apiCode = typeof payload.code === "string" ? payload.code : null;
+        if (response.status >= 200 && response.status < 300 && apiCode === "Success") {
+          return payload.data;
+        }
+        const retryable = RETRYABLE_HTTP_STATUS.has(response.status)
+          || /busy|timeout|temporar|稍后|繁忙|库存/i.test(String(payload.msg ?? ""));
+        throw new AutoDlApiError(
+          `AutoDL API ${pathname} 失败：${String(payload.msg || apiCode || `HTTP ${response.status}`).slice(0, 500)}`,
+          { status: response.status, code: apiCode, retryable },
+        );
+      } catch (error) {
+        lastError = error;
+        const retryable = error instanceof AutoDlApiError ? error.retryable : true;
+        if (!retryable || attempt >= this.maxAttempts) break;
+        await this.sleep(backoffMilliseconds(attempt));
+      }
+    }
+    if (lastError instanceof AutoDlApiError) throw lastError;
+    throw new AutoDlApiError(
+      `AutoDL API 请求失败：${lastError instanceof Error ? lastError.message : String(lastError)}`,
+      { retryable: true, cause: lastError },
+    );
+  }
+}
+
+export async function nodeHttpsJsonRequest({ url, method, headers, body, timeoutMs }) {
+  const serialized = body === undefined ? null : JSON.stringify(body ?? {});
+  return await new Promise((resolve, reject) => {
+    const request = https.request(url, {
+      method,
+      headers: {
+        ...headers,
+        ...(serialized === null ? {} : { "content-length": Buffer.byteLength(serialized) }),
+      },
+      timeout: timeoutMs,
+    }, (response) => {
+      const chunks = [];
+      response.on("data", (chunk) => chunks.push(chunk));
+      response.on("end", () => {
+        const text = Buffer.concat(chunks).toString("utf8");
+        let parsed = {};
+        try {
+          parsed = text ? JSON.parse(text) : {};
+        } catch {
+          parsed = { msg: text.slice(0, 500) };
+        }
+        resolve({ status: response.statusCode ?? 0, body: parsed });
+      });
+    });
+    request.on("timeout", () => request.destroy(new Error("AutoDL API 请求超时")));
+    request.on("error", reject);
+    request.end(serialized ?? undefined);
+  });
+}
+
+function required(value, field) {
+  if (typeof value !== "string" || !value.trim()) {
+    throw new TypeError(`AutoDL 缺少 ${field}`);
+  }
+  return value.trim();
+}
+
+function optional(value) {
+  return typeof value === "string" && value.trim() ? value.trim() : undefined;
+}
+
+function stringArray(value) {
+  if (value === undefined || value === null) return undefined;
+  if (!Array.isArray(value) || value.some((item) => typeof item !== "string" || !item.trim())) {
+    throw new TypeError("AutoDL dataCenterList 必须是非空字符串数组");
+  }
+  return value.map((item) => item.trim());
+}
+
+function positiveInteger(value, field) {
+  const numeric = Number(value);
+  if (!Number.isInteger(numeric) || numeric <= 0) {
+    throw new TypeError(`AutoDL ${field} 必须是正整数`);
+  }
+  return numeric;
+}
+
+function nonNegativeInteger(value, field) {
+  const numeric = Number(value);
+  if (!Number.isInteger(numeric) || numeric < 0) {
+    throw new TypeError(`AutoDL ${field} 必须是非负整数`);
+  }
+  return numeric;
+}
+
+function isRecord(value) {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function backoffMilliseconds(attempt) {
+  return Math.min(15_000, 500 * 2 ** Math.max(0, attempt - 1));
+}
+
+function defaultSleep(milliseconds) {
+  return new Promise((resolve) => setTimeout(resolve, milliseconds));
+}

package/src/autodl/ssh.mjs

@@ -0,0 +1,380 @@
+import ssh2 from "ssh2";
+
+import { CiteArkError } from "../util.mjs";
+
+const { Client } = ssh2;
+
+const DEFAULT_OPERATION_TIMEOUT_MS = 2 * 60_000;
+const DEFAULT_TRANSFER_TIMEOUT_MS = 30 * 60_000;
+const DEFAULT_TRANSFER_CONCURRENCY = 4;
+
+export class AutoDlSshClient {
+  constructor({
+    host,
+    port,
+    username = "root",
+    password,
+    readyTimeoutMs = 60_000,
+    operationTimeoutMs = DEFAULT_OPERATION_TIMEOUT_MS,
+    transferTimeoutMs = DEFAULT_TRANSFER_TIMEOUT_MS,
+    transferConcurrency = DEFAULT_TRANSFER_CONCURRENCY,
+    clientFactory = () => new Client(),
+  } = {}) {
+    this.host = required(host, "host");
+    this.port = positiveInteger(port, "port");
+    this.username = required(username, "username");
+    this.password = required(password, "password");
+    this.readyTimeoutMs = positiveInteger(readyTimeoutMs, "readyTimeoutMs");
+    this.operationTimeoutMs = positiveInteger(operationTimeoutMs, "operationTimeoutMs");
+    this.transferTimeoutMs = positiveInteger(transferTimeoutMs, "transferTimeoutMs");
+    this.transferConcurrency = positiveInteger(
+      transferConcurrency,
+      "transferConcurrency",
+    );
+    this.clientFactory = clientFactory;
+    this.client = null;
+    this.connected = false;
+    this.connecting = null;
+    this.pendingOperations = new Set();
+    this.fingerprint = null;
+  }
+
+  async connect() {
+    if (this.client && this.connected) return this;
+    if (this.connecting) {
+      await this.connecting;
+      return this;
+    }
+    const client = this.clientFactory();
+    this.client = client;
+    const connecting = new Promise((resolve, reject) => {
+      let settled = false;
+      const finish = (callback, value) => {
+        if (settled) return;
+        settled = true;
+        callback(value);
+      };
+      const disconnect = (message, cause) => {
+        const error = new CiteArkError(message, { cause });
+        this.#handleDisconnect(client, error);
+        finish(reject, error);
+      };
+      client.once("ready", () => {
+        if (this.client !== client) {
+          finish(reject, new CiteArkError("AutoDL SSH 连接在就绪前已失效"));
+          return;
+        }
+        this.connected = true;
+        finish(resolve);
+      });
+      client.on("error", (error) => disconnect(
+        `AutoDL SSH 连接错误：${safeMessage(error)}`,
+        error,
+      ));
+      client.once("end", () => disconnect("AutoDL SSH 连接已由远端结束"));
+      client.once("close", () => disconnect("AutoDL SSH 连接已关闭"));
+      try {
+        client.connect({
+          host: this.host,
+          port: this.port,
+          username: this.username,
+          password: this.password,
+          readyTimeout: this.readyTimeoutMs,
+          keepaliveInterval: 15_000,
+          keepaliveCountMax: 4,
+          hostHash: "sha256",
+          hostVerifier: (fingerprint) => {
+            if (!this.fingerprint) {
+              this.fingerprint = fingerprint;
+              return true;
+            }
+            return this.fingerprint === fingerprint;
+          },
+        });
+      } catch (error) {
+        disconnect(
+          `无法连接 AutoDL SSH 节点：${safeMessage(error)}`,
+          error,
+        );
+      }
+    });
+    this.connecting = connecting;
+    try {
+      await connecting;
+      return this;
+    } finally {
+      if (this.connecting === connecting) this.connecting = null;
+    }
+  }
+
+  async exec(command, {
+    timeoutMs = 0,
+    signal,
+    onStdout = () => {},
+    onStderr = () => {},
+    maximumCaptureCharacters = 16_000,
+  } = {}) {
+    const operationTimeoutMs = timeoutMs > 0 ? timeoutMs + 45_000 : null;
+    return await this.#runOperation("AutoDL SSH 远程命令", async (client) => {
+      return await new Promise((resolve, reject) => {
+        if (signal?.aborted) {
+          reject(signal.reason instanceof Error
+            ? signal.reason
+            : new CiteArkError("AutoDL SSH 远程命令已取消"));
+          return;
+        }
+        client.exec(command, (error, stream) => {
+          if (error) {
+            reject(new CiteArkError(`无法启动 AutoDL 远程命令：${safeMessage(error)}`, { cause: error }));
+            return;
+          }
+          let stdout = "";
+          let stderr = "";
+          let timedOut = false;
+          let timer;
+          let abortTimer;
+          const abort = () => {
+            if (timer) clearTimeout(timer);
+            stream.signal("TERM");
+            abortTimer = setTimeout(() => stream.close(), 30_000);
+            abortTimer.unref();
+            reject(signal?.reason instanceof Error
+              ? signal.reason
+              : new CiteArkError("AutoDL SSH 远程命令已取消"));
+          };
+          signal?.addEventListener("abort", abort, { once: true });
+          const append = (current, chunk) => {
+            const combined = `${current}${chunk}`;
+            return combined.slice(-maximumCaptureCharacters);
+          };
+          stream.on("data", (chunk) => {
+            const text = chunk.toString();
+            stdout = append(stdout, text);
+            onStdout(text);
+          });
+          stream.stderr.on("data", (chunk) => {
+            const text = chunk.toString();
+            stderr = append(stderr, text);
+            onStderr(text);
+          });
+          stream.once("error", (streamError) => {
+            if (timer) clearTimeout(timer);
+            if (abortTimer) clearTimeout(abortTimer);
+            signal?.removeEventListener("abort", abort);
+            reject(new CiteArkError(`AutoDL SSH 数据流中断：${safeMessage(streamError)}`, {
+              cause: streamError,
+            }));
+          });
+          stream.once("close", (code, remoteSignal) => {
+            if (timer) clearTimeout(timer);
+            if (abortTimer) clearTimeout(abortTimer);
+            signal?.removeEventListener("abort", abort);
+            resolve({
+              code: Number.isInteger(code) ? code : 1,
+              signal: remoteSignal ?? null,
+              stdout,
+              stderr,
+              timedOut,
+            });
+          });
+          if (timeoutMs > 0) {
+            timer = setTimeout(() => {
+              timedOut = true;
+              stream.signal("TERM");
+              setTimeout(() => stream.close(), 30_000).unref();
+            }, timeoutMs);
+            timer.unref();
+          }
+        });
+      });
+    }, operationTimeoutMs);
+  }
+
+  async putFile(localPath, remotePath, { timeoutMs = this.transferTimeoutMs } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 上传", timeoutMs, async (sftp) => {
+      await callbackPromise((done) => sftp.fastPut(localPath, remotePath, {
+        concurrency: this.transferConcurrency,
+      }, done));
+    });
+  }
+
+  async getFile(remotePath, localPath, {
+    timeoutMs = this.transferTimeoutMs,
+    onProgress,
+  } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 下载", timeoutMs, async (sftp) => {
+      const options = { concurrency: this.transferConcurrency };
+      if (typeof onProgress === "function") {
+        options.step = (transferredBytes, chunkBytes, totalBytes) => {
+          onProgress({ transferredBytes, chunkBytes, totalBytes });
+        };
+      }
+      await callbackPromise((done) => sftp.fastGet(
+        remotePath,
+        localPath,
+        options,
+        done,
+      ));
+    });
+  }
+
+  async readFile(remotePath, {
+    maximumBytes = 4 * 1024 * 1024,
+    timeoutMs = this.operationTimeoutMs,
+  } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 读取", timeoutMs, async (sftp) => {
+      const attributes = await callbackPromise((done) => sftp.stat(remotePath, done));
+      if (Number(attributes?.size) > maximumBytes) {
+        throw new CiteArkError(`AutoDL 远程文件超过读取上限：${remotePath}`);
+      }
+      const bytes = await callbackPromise((done) => sftp.readFile(remotePath, done));
+      if (bytes.length > maximumBytes) {
+        throw new CiteArkError(`AutoDL 远程文件超过读取上限：${remotePath}`);
+      }
+      return bytes;
+    });
+  }
+
+  async writeFile(remotePath, content, {
+    mode = 0o600,
+    timeoutMs = this.operationTimeoutMs,
+  } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 写入", timeoutMs, async (sftp) => {
+      await new Promise((resolve, reject) => {
+        const stream = sftp.createWriteStream(remotePath, { mode });
+        stream.once("error", reject);
+        stream.once("close", resolve);
+        stream.end(content);
+      });
+    });
+  }
+
+  async exists(remotePath, { timeoutMs = this.operationTimeoutMs } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 状态检查", timeoutMs, async (sftp) => {
+      return await new Promise((resolve, reject) => {
+        sftp.stat(remotePath, (error) => {
+          if (!error) return resolve(true);
+          if (error.code === 2 || error.code === "ENOENT") return resolve(false);
+          reject(error);
+        });
+      });
+    });
+  }
+
+  async list(remotePath, { timeoutMs = this.operationTimeoutMs } = {}) {
+    return await this.#runSftpOperation("AutoDL SFTP 目录读取", timeoutMs, async (sftp) => {
+      return await callbackPromise((done) => sftp.readdir(remotePath, done));
+    });
+  }
+
+  close() {
+    const client = this.client;
+    const error = new CiteArkError("AutoDL SSH 客户端已关闭");
+    this.#handleDisconnect(client, error);
+    client?.end();
+  }
+
+  async #runSftpOperation(label, timeoutMs, operation) {
+    return await this.#runOperation(label, async (client) => {
+      const sftp = await callbackPromise((done) => client.sftp(done));
+      try {
+        return await operation(sftp);
+      } finally {
+        sftp.end();
+      }
+    }, timeoutMs);
+  }
+
+  async #runOperation(label, operation, timeoutMs) {
+    await this.connect();
+    const client = this.client;
+    if (!client || !this.connected) {
+      throw new CiteArkError(`${label}无法开始：SSH 连接不可用`);
+    }
+    const boundedTimeoutMs = timeoutMs === null
+      ? null
+      : positiveInteger(timeoutMs ?? this.operationTimeoutMs, "timeoutMs");
+    return await new Promise((resolve, reject) => {
+      let settled = false;
+      let timer;
+      const finish = (callback, value) => {
+        if (settled) return;
+        settled = true;
+        if (timer) clearTimeout(timer);
+        this.pendingOperations.delete(fail);
+        callback(value);
+      };
+      const fail = (error) => finish(reject, error);
+      this.pendingOperations.add(fail);
+      if (boundedTimeoutMs !== null) {
+        timer = setTimeout(() => {
+          const error = new CiteArkError(
+            `${label}超时（${Math.ceil(boundedTimeoutMs / 1000)} 秒）`,
+          );
+          this.#handleDisconnect(client, error);
+          client.destroy();
+        }, boundedTimeoutMs);
+      }
+      Promise.resolve()
+        .then(() => operation(client))
+        .then((value) => finish(resolve, value), fail);
+    });
+  }
+
+  #handleDisconnect(client, error) {
+    if (!client || this.client !== client) return;
+    this.client = null;
+    this.connected = false;
+    const pending = [...this.pendingOperations];
+    this.pendingOperations.clear();
+    for (const reject of pending) reject(error);
+  }
+}
+
+export function parseAutoDlSshConnection(snapshot) {
+  const explicitHost = text(snapshot?.proxy_host);
+  const explicitPort = Number(snapshot?.ssh_port);
+  const command = text(snapshot?.ssh_command);
+  const matched = command?.match(/^ssh\s+-p\s+(\d+)\s+([^@\s]+)@([^\s]+)$/i);
+  const host = explicitHost ?? matched?.[3] ?? null;
+  const port = Number.isInteger(explicitPort) && explicitPort > 0
+    ? explicitPort
+    : Number(matched?.[1]);
+  const username = matched?.[2] ?? "root";
+  const password = text(snapshot?.root_password);
+  if (!host || !Number.isInteger(port) || port <= 0 || !password) {
+    throw new CiteArkError("AutoDL 实例尚未返回完整 SSH 连接信息");
+  }
+  return { host, port, username, password };
+}
+
+async function callbackPromise(register) {
+  return await new Promise((resolve, reject) => {
+    register((error, value) => error ? reject(error) : resolve(value));
+  });
+}
+
+function required(value, field) {
+  if (typeof value !== "string" || !value.trim()) {
+    throw new TypeError(`AutoDL SSH 缺少 ${field}`);
+  }
+  return value.trim();
+}
+
+function positiveInteger(value, field) {
+  const numeric = Number(value);
+  if (!Number.isInteger(numeric) || numeric <= 0) {
+    throw new TypeError(`AutoDL SSH ${field} 必须是正整数`);
+  }
+  return numeric;
+}
+
+function text(value) {
+  return typeof value === "string" && value.trim() ? value.trim() : null;
+}
+
+function safeMessage(error) {
+  return String(error instanceof Error ? error.message : error)
+    .replaceAll(/(password|token|secret)\s*[=:]\s*\S+/gi, "$1=<redacted>")
+    .slice(0, 500);
+}