fix: harden provider polling recovery

scripts/keyManagerReleaseContract.test.js

@@ -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;
+});

scripts/providerPollLimiterContract.test.js

@@ -0,0 +1,96 @@
+const assert = require("node:assert/strict");
+const { createRequire } = require("node:module");
+
+const nodeRequire = createRequire(__filename);
+
+function loadLimiterWithPool(pool) {
+  const dbPath = nodeRequire.resolve("../src/db");
+  const limiterPath = nodeRequire.resolve("../src/providerPollLimiter");
+  const originalDbModule = nodeRequire.cache[dbPath];
+  const originalLimiterModule = nodeRequire.cache[limiterPath];
+
+  delete nodeRequire.cache[limiterPath];
+  nodeRequire.cache[dbPath] = {
+    id: dbPath,
+    filename: dbPath,
+    loaded: true,
+    exports: { pool },
+  };
+
+  return {
+    limiter: nodeRequire("../src/providerPollLimiter"),
+    restore() {
+      delete nodeRequire.cache[limiterPath];
+      if (originalLimiterModule) nodeRequire.cache[limiterPath] = originalLimiterModule;
+      if (originalDbModule) nodeRequire.cache[dbPath] = originalDbModule;
+      else delete nodeRequire.cache[dbPath];
+    },
+  };
+}
+
+function createPool(options = {}) {
+  const calls = [];
+  return {
+    calls,
+    async query(sql, params = []) {
+      calls.push({ sql, params });
+      if (/CREATE TABLE IF NOT EXISTS generation_provider_poll_slots/i.test(sql)) return { rows: [] };
+      if (/WITH candidate AS/i.test(sql)) {
+        if (options.noAvailableSlot) return { rows: [] };
+        return { rows: [{ scope: params[0], slot_no: 2 }] };
+      }
+      if (/DELETE FROM generation_provider_poll_slots/i.test(sql)) return { rows: [] };
+      throw new Error(`Unexpected SQL: ${sql}`);
+    },
+  };
+}
+
+(async () => {
+  const previousLimit = process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
+  process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = "3";
+
+  const pool = createPool();
+  const { limiter, restore } = loadLimiterWithPool(pool);
+  try {
+    const outcome = await limiter.withProviderPollSlot(101, async () => "polled");
+
+    assert.equal(outcome.acquired, true);
+    assert.equal(outcome.value, "polled");
+
+    const acquireCall = pool.calls.find((call) => /WITH candidate AS/i.test(call.sql));
+    assert.equal(acquireCall.params[1], 3);
+    assert.equal(acquireCall.params[3], 101);
+
+    const releaseCall = pool.calls.find((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql));
+    assert.equal(releaseCall.params[0], acquireCall.params[0]);
+    assert.equal(releaseCall.params[1], 2);
+    assert.equal(releaseCall.params[2], acquireCall.params[2]);
+  } finally {
+    if (previousLimit === undefined) delete process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY;
+    else process.env.TASK_PROVIDER_POLL_MAX_CONCURRENCY = previousLimit;
+    restore();
+  }
+
+  const saturatedPool = createPool({ noAvailableSlot: true });
+  const { limiter: saturatedLimiter, restore: restoreSaturated } = loadLimiterWithPool(saturatedPool);
+  try {
+    let called = false;
+    const outcome = await saturatedLimiter.withProviderPollSlot(202, async () => {
+      called = true;
+      return "should-not-run";
+    });
+
+    assert.equal(outcome.acquired, false);
+    assert.equal(outcome.value, undefined);
+    assert.equal(called, false);
+    assert.equal(
+      saturatedPool.calls.some((call) => /DELETE FROM generation_provider_poll_slots/i.test(call.sql)),
+      false,
+    );
+  } finally {
+    restoreSaturated();
+  }
+})().catch((error) => {
+  console.error(error);
+  process.exitCode = 1;
+});