fix: harden provider polling recovery

2026-06-09 11:32:53 +08:00
parent 1166811ee4
commit f9da506017
9 changed files with 539 additions and 69 deletions
@@ -0,0 +1,73 @@
+const assert = require("node:assert/strict");
+const { createRequire } = require("node:module");
+
+const nodeRequire = createRequire(__filename);
+
+function loadKeyManagerWithPool(pool) {
+  const dbPath = nodeRequire.resolve("../src/db");
+  const keyManagerPath = nodeRequire.resolve("../src/keyManager");
+  const originalDbModule = nodeRequire.cache[dbPath];
+  const originalKeyManagerModule = nodeRequire.cache[keyManagerPath];
+
+  delete nodeRequire.cache[keyManagerPath];
+  nodeRequire.cache[dbPath] = {
+    id: dbPath,
+    filename: dbPath,
+    loaded: true,
+    exports: {
+      pool,
+      withTransaction: async (fn) => fn(pool),
+    },
+  };
+
+  return {
+    keyManager: nodeRequire("../src/keyManager"),
+    restore() {
+      delete nodeRequire.cache[keyManagerPath];
+      if (originalKeyManagerModule) nodeRequire.cache[keyManagerPath] = originalKeyManagerModule;
+      if (originalDbModule) nodeRequire.cache[dbPath] = originalDbModule;
+      else delete nodeRequire.cache[dbPath];
+    },
+  };
+}
+
+function createReleasePool() {
+  const calls = [];
+  return {
+    calls,
+    async query(sql, params) {
+      calls.push({ sql, params });
+      if (/WITH candidate AS/i.test(sql)) {
+        return {
+          rows: [{
+            id: 10,
+            key_id: 20,
+            lease_user_id: 30,
+            lease_enterprise_id: 40,
+            provider: "dashscope",
+          }],
+        };
+      }
+      if (/UPDATE api_keys SET active_count/i.test(sql)) return { rows: [] };
+      if (/INSERT INTO usage_logs/i.test(sql)) return { rows: [] };
+      throw new Error(`Unexpected SQL: ${sql}`);
+    },
+  };
+}
+
+(async () => {
+  const pool = createReleasePool();
+  const { keyManager, restore } = loadKeyManagerWithPool(pool);
+  try {
+    const result = await keyManager.releaseKey("lease-token-without-user-context");
+
+    assert.equal(result.released, true);
+    const usageLogCall = pool.calls.find((call) => /INSERT INTO usage_logs/i.test(call.sql));
+    assert.deepEqual(usageLogCall.params, [30, 40, 20, 20, "release"]);
+  } finally {
+    restore();
+  }
+})().catch((error) => {
+  console.error(error);
+  process.exitCode = 1;
+});