@cartanova/qgrid-cli 2.4.3 → 2.4.5

package/bundle/src/utils/providers/openai/openai-dispatcher.test.ts

@@ -4,7 +4,15 @@ import { QuotaThresholdExceededError } from "../../../application/qgrid/qgrid.ty
 import { type OpenAICredentials } from "../../../application/token/token.types";
 import { type GenerateRequest, type GenerateResult } from "../common/provider-dispatcher";
 import { type CodexAppServerWorker } from "./codex-worker";
-import { ImageGenerationError, OpenAIDispatcher } from "./openai-dispatcher";
+import {
+  estimateOpenAIWorkerRssGiB,
+  ImageGenerationError,
+  makeOpenAIWorkerId,
+  MAX_OPENAI_WORKERS_PER_TOKEN,
+  type OpenAIWorkerPoolConfig,
+  OpenAIDispatcher,
+  resolveOpenAIWorkerPoolConfig,
+} from "./openai-dispatcher";
 import { type OpenAIQuotaUsageResult } from "./openai-quota";
 
 const { readOpenAIQuotaUsageMock, loggerInfoMock, loggerWarnMock } = vi.hoisted(() => ({
@@ -70,7 +78,7 @@ function resultFor(worker: CodexAppServerWorker, threadId = "thread-new"): Gener
     durationMs: 10,
     model: "gpt-5-codex",
     threadCoord: {
-      workerId: worker.tokenId * 10 + worker.workerIndex,
+      workerId: makeOpenAIWorkerId(worker.tokenId, worker.workerIndex),
       threadId,
       epoch: worker.epoch,
     },
@@ -112,6 +120,9 @@ function fakeWorker(
     get isReady() {
       return options.ready ?? true;
     },
+    get hasCapacity() {
+      return !busy;
+    },
     tryAcquireTurn: vi.fn(() => {
       if (busy) return false;
       busy = true;
@@ -156,6 +167,21 @@ function addWorkers(dispatcher: OpenAIDispatcher, workers: CodexAppServerWorker[
   for (const [tokenId, group] of byToken) dispatcher.workerPool.set(tokenId, group);
 }
 
+function autoscaleConfig(
+  overrides: Partial<OpenAIWorkerPoolConfig> = {},
+): OpenAIWorkerPoolConfig {
+  return {
+    autoscale: true,
+    minWorkersPerToken: 1,
+    maxWorkersPerToken: 3,
+    scaleIntervalMs: 5_000,
+    scaleDownIdleMs: 600_000,
+    maxEstimatedRssGiB: 16,
+    minHostAvailableGiB: 20,
+    ...overrides,
+  };
+}
+
 async function waitForQueue(dispatcher: OpenAIDispatcher, length = 1): Promise<void> {
   for (let i = 0; i < 10; i++) {
     await new Promise((resolve) => setTimeout(resolve, 0));
@@ -164,6 +190,159 @@ async function waitForQueue(dispatcher: OpenAIDispatcher, length = 1): Promise<v
   expect(dispatcher.queueLength).toBe(length);
 }
 
+describe("OpenAIDispatcher worker capacity", () => {
+  it("supports 20 workers per token without worker id collisions", () => {
+    expect(MAX_OPENAI_WORKERS_PER_TOKEN).toBe(20);
+
+    const tokenIds = [1, 2, 41, 47, 55, 56, 58];
+    const workerIds = tokenIds.flatMap((tokenId) =>
+      Array.from({ length: MAX_OPENAI_WORKERS_PER_TOKEN }, (_, workerIndex) =>
+        makeOpenAIWorkerId(tokenId, workerIndex),
+      ),
+    );
+
+    expect(new Set(workerIds)).toHaveLength(workerIds.length);
+    expect(makeOpenAIWorkerId(1, 10)).not.toBe(makeOpenAIWorkerId(2, 0));
+  });
+
+  it("defaults to autoscaling from 5 to 15 workers per token", () => {
+    expect(resolveOpenAIWorkerPoolConfig({})).toMatchObject({
+      autoscale: true,
+      minWorkersPerToken: 5,
+      maxWorkersPerToken: 15,
+    });
+  });
+
+  it("keeps an explicit fixed worker mode for emergency rollback", () => {
+    expect(
+      resolveOpenAIWorkerPoolConfig({
+        QGRID_WORKERS_PER_TOKEN: "5",
+        QGRID_OPENAI_AUTOSCALE: "false",
+      }),
+    ).toMatchObject({
+      autoscale: false,
+      minWorkersPerToken: 5,
+      maxWorkersPerToken: 5,
+    });
+  });
+
+  it("estimates worker RSS from the dev0 measurement model", () => {
+    expect(estimateOpenAIWorkerRssGiB(50)).toBeCloseTo(8.56, 2);
+    expect(estimateOpenAIWorkerRssGiB(60)).toBeCloseTo(10.13, 2);
+  });
+});
+
+describe("OpenAIDispatcher autoscaling", () => {
+  it("scales six active tokens from 5 to 15 workers per token", async () => {
+    const dispatcher = new OpenAIDispatcher(
+      autoscaleConfig({
+        minWorkersPerToken: 5,
+        maxWorkersPerToken: 15,
+        maxEstimatedRssGiB: 16,
+      }),
+      () => 64,
+    );
+    vi.spyOn(dispatcher, "spawnWorkers").mockResolvedValue(undefined);
+
+    for (let tokenId = 1; tokenId <= 6; tokenId++) {
+      await dispatcher.onTokenAdded(tokenId, `token-${tokenId}`, credentials(), null, 1);
+      dispatcher.workerPool.set(
+        tokenId,
+        Array.from({ length: 5 }, (_, workerIndex) => fakeWorker({ tokenId, workerIndex })),
+      );
+    }
+
+    vi.spyOn(dispatcher, "spawnSingleWorker").mockImplementation(
+      async (tokenId, tokenName, _credentials, workerIndex) =>
+        fakeWorker({ tokenId, tokenName, workerIndex }),
+    );
+
+    for (let step = 0; step < 12; step++) await dispatcher.scaleUpOneStep();
+
+    expect(dispatcher.workerCount).toBe(90);
+    expect([...dispatcher.workerPool.values()].map((workers) => workers.length)).toEqual([
+      15, 15, 15, 15, 15, 15,
+    ]);
+    expect(estimateOpenAIWorkerRssGiB(dispatcher.workerCount)).toBeLessThan(16);
+  });
+
+  it("requests a scale-up evaluation when a request enters the queue", async () => {
+    const dispatcher = new OpenAIDispatcher(autoscaleConfig(), () => 64);
+    addWorkers(dispatcher, [fakeWorker({ tokenId: 1, busy: true })]);
+    const scaleUpOneStep = vi.spyOn(dispatcher, "scaleUpOneStep").mockResolvedValue(undefined);
+
+    const queued = dispatcher.enqueue(async () => "ok");
+    const rejection = expect(queued).rejects.toThrow("TEST_CLEANUP");
+    await waitForQueue(dispatcher);
+    dispatcher.rejectAllQueued("TEST_CLEANUP");
+
+    await rejection;
+    expect(scaleUpOneStep).toHaveBeenCalledTimes(1);
+  });
+
+  it("adds one worker to each active token in a scale-up step", async () => {
+    const dispatcher = new OpenAIDispatcher(autoscaleConfig(), () => 64);
+    const tokenCredentials = credentials();
+    vi.spyOn(dispatcher, "spawnWorkers").mockResolvedValue(undefined);
+    await dispatcher.onTokenAdded(1, "token", tokenCredentials, null, 1);
+    dispatcher.workerPool.set(1, [fakeWorker({ tokenId: 1, workerIndex: 0 })]);
+    const added = fakeWorker({ tokenId: 1, workerIndex: 1 });
+    const spawnSingleWorker = vi.spyOn(dispatcher, "spawnSingleWorker").mockResolvedValue(added);
+
+    await dispatcher.scaleUpOneStep();
+
+    expect(spawnSingleWorker).toHaveBeenCalledWith(1, "token", tokenCredentials, 1);
+    expect(dispatcher.workerPool.get(1)).toHaveLength(2);
+    expect(dispatcher.workerCount).toBe(2);
+  });
+
+  it("does not scale past the estimated RSS guard", async () => {
+    const dispatcher = new OpenAIDispatcher(
+      autoscaleConfig({ maxEstimatedRssGiB: 0.9 }),
+      () => 64,
+    );
+    vi.spyOn(dispatcher, "spawnWorkers").mockResolvedValue(undefined);
+    await dispatcher.onTokenAdded(1, "token", credentials(), null, 1);
+    dispatcher.workerPool.set(1, [fakeWorker({ tokenId: 1, workerIndex: 0 })]);
+    const spawnSingleWorker = vi.spyOn(dispatcher, "spawnSingleWorker");
+
+    await dispatcher.scaleUpOneStep();
+
+    expect(spawnSingleWorker).not.toHaveBeenCalled();
+    expect(dispatcher.workerCount).toBe(1);
+  });
+
+  it("does not scale when host available memory is below the floor", async () => {
+    const dispatcher = new OpenAIDispatcher(autoscaleConfig(), () => 10);
+    vi.spyOn(dispatcher, "spawnWorkers").mockResolvedValue(undefined);
+    await dispatcher.onTokenAdded(1, "token", credentials(), null, 1);
+    dispatcher.workerPool.set(1, [fakeWorker({ tokenId: 1, workerIndex: 0 })]);
+    const spawnSingleWorker = vi.spyOn(dispatcher, "spawnSingleWorker");
+
+    await dispatcher.scaleUpOneStep();
+
+    expect(spawnSingleWorker).not.toHaveBeenCalled();
+    expect(dispatcher.workerCount).toBe(1);
+  });
+
+  it("removes only idle workers above the configured minimum for inactive tokens", async () => {
+    const dispatcher = new OpenAIDispatcher(autoscaleConfig(), () => 64);
+    vi.spyOn(dispatcher, "spawnWorkers").mockResolvedValue(undefined);
+    await dispatcher.onTokenAdded(1, "token", credentials(), null, 1);
+    const busy = fakeWorker({ tokenId: 1, workerIndex: 0, busy: true });
+    const idle = fakeWorker({ tokenId: 1, workerIndex: 1 });
+    dispatcher.workerPool.set(1, [busy, idle]);
+    dispatcher.onTokenDeactivated(1);
+
+    await dispatcher.scaleDownOneStep();
+    await dispatcher.scaleDownOneStep();
+
+    expect(idle.kill).toHaveBeenCalledTimes(1);
+    expect(busy.kill).not.toHaveBeenCalled();
+    expect(dispatcher.workerPool.get(1)).toEqual([busy]);
+  });
+});
+
 describe("OpenAIDispatcher token updates", () => {
   it("keeps existing workers when only OpenAI auth tokens rotate", async () => {
     const dispatcher = new OpenAIDispatcher();
@@ -195,7 +374,7 @@ describe("OpenAIDispatcher token updates", () => {
     await dispatcher.onTokenUpdated(1, "token", changedIdentity, null);
 
     expect(worker.kill).toHaveBeenCalledTimes(1);
-    expect(spawnWorkers).toHaveBeenCalledWith(1, "token", changedIdentity, null);
+    expect(spawnWorkers).toHaveBeenCalledWith(1, "token", changedIdentity, null, 1);
   });
 
   it("uses tokenId keyed rate-limit cache", async () => {
@@ -227,7 +406,7 @@ describe("OpenAIDispatcher token updates", () => {
     addWorkers(dispatcher, [workerA, workerB]);
 
     await dispatcher.replaceTokens([
-      { id: 1, name: "tok-A", credentials: credentials(), quotaThreshold: 80 },
+      { id: 1, name: "tok-A", credentials: credentials(), quotaThreshold: 80, weight: 1 },
     ]);
     readOpenAIQuotaUsageMock.mockResolvedValueOnce(quotaOk(85));
 
@@ -238,6 +417,108 @@ describe("OpenAIDispatcher token updates", () => {
   });
 });
 
+describe("OpenAIDispatcher weighted routing", () => {
+  it("routes cold OpenAI requests by token weight instead of worker count", async () => {
+    const dispatcher = new OpenAIDispatcher();
+    const a0 = fakeWorker({ tokenId: 1, tokenName: "tok-A", workerIndex: 0 });
+    const a1 = fakeWorker({ tokenId: 1, tokenName: "tok-A", workerIndex: 1 });
+    const b0 = fakeWorker({ tokenId: 2, tokenName: "tok-B", workerIndex: 0 });
+    addWorkers(dispatcher, [a0, a1, b0]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 3);
+    await dispatcher.onTokenUpdated(2, "tok-B", credentials(), null, 1);
+
+    const names: string[] = [];
+    for (let i = 0; i < 4; i++) {
+      const worker = await dispatcher.acquireIdleWorker();
+      names.push(worker!.tokenName);
+      worker!.releaseTurn();
+    }
+
+    expect(names.filter((name) => name === "tok-A")).toHaveLength(3);
+    expect(names.filter((name) => name === "tok-B")).toHaveLength(1);
+  });
+
+  it("does not give a token extra cold turns for having more workers", async () => {
+    const dispatcher = new OpenAIDispatcher();
+    addWorkers(dispatcher, [
+      fakeWorker({ tokenId: 1, tokenName: "tok-A", workerIndex: 0 }),
+      fakeWorker({ tokenId: 1, tokenName: "tok-A", workerIndex: 1 }),
+      fakeWorker({ tokenId: 2, tokenName: "tok-B", workerIndex: 0 }),
+    ]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 1);
+    await dispatcher.onTokenUpdated(2, "tok-B", credentials(), null, 1);
+
+    const names: string[] = [];
+    for (let i = 0; i < 4; i++) {
+      const worker = await dispatcher.acquireIdleWorker();
+      names.push(worker!.tokenName);
+      worker!.releaseTurn();
+    }
+
+    expect(names).toEqual(["tok-A", "tok-B", "tok-A", "tok-B"]);
+  });
+
+  it("skips a high-weight token when all of its workers are busy", async () => {
+    const dispatcher = new OpenAIDispatcher();
+    addWorkers(dispatcher, [
+      fakeWorker({ tokenId: 1, tokenName: "tok-A", busy: true }),
+      fakeWorker({ tokenId: 2, tokenName: "tok-B" }),
+    ]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 100);
+    await dispatcher.onTokenUpdated(2, "tok-B", credentials(), null, 1);
+
+    await expect(dispatcher.acquireIdleWorker()).resolves.toMatchObject({ tokenName: "tok-B" });
+  });
+
+  it("uses weighted selection when draining after a worker release", async () => {
+    const dispatcher = new OpenAIDispatcher();
+    const workerA = fakeWorker({ tokenId: 1, tokenName: "tok-A", busy: true });
+    const workerB = fakeWorker({ tokenId: 2, tokenName: "tok-B", busy: true });
+    addWorkers(dispatcher, [workerA, workerB]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 1);
+    await dispatcher.onTokenUpdated(2, "tok-B", credentials(), null, 3);
+
+    const queued = dispatcher.enqueue(async (worker) => worker.tokenName);
+    await waitForQueue(dispatcher);
+    workerA.releaseTurn();
+    workerB.releaseTurn();
+
+    await dispatcher.drainQueue();
+    await expect(queued).resolves.toBe("tok-B");
+  });
+
+  it("does not advance weighted state for successful thread reuse", async () => {
+    const dispatcher = new OpenAIDispatcher();
+    const workerA = fakeWorker({ tokenId: 1, tokenName: "tok-A" });
+    const workerB = fakeWorker({
+      tokenId: 2,
+      tokenName: "tok-B",
+      threads: ["thread-B"],
+    });
+    addWorkers(dispatcher, [workerA, workerB]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 1);
+    await dispatcher.onTokenUpdated(2, "tok-B", credentials(), null, 1);
+    vi.spyOn(dispatcher, "executeTurn").mockImplementation(
+      async (worker, _req, reuseThreadId) =>
+        resultFor(worker, reuseThreadId ?? "thread-new"),
+    );
+
+    const cold1 = await dispatcher.generate(baseReq());
+    const reused = await dispatcher.generate(
+      baseReq({
+        reuse: { workerId: makeOpenAIWorkerId(2, 0), threadId: "thread-B", epoch: 1 },
+      }),
+    );
+    const cold2 = await dispatcher.generate(baseReq());
+
+    expect([cold1.tokenName, cold2.tokenName]).toEqual(["tok-A", "tok-B"]);
+    expect(reused).toMatchObject({
+      tokenName: "tok-B",
+      threadCoord: { threadId: "thread-B" },
+    });
+  });
+});
+
 describe("OpenAIDispatcher quota threshold gate", () => {
   it("excludes over-threshold tokenIds before idle selection", async () => {
     const dispatcher = new OpenAIDispatcher();
@@ -280,7 +561,7 @@ describe("OpenAIDispatcher quota threshold gate", () => {
     await waitForQueue(dispatcher);
 
     worker.releaseTurn();
-    await dispatcher.drainQueue(worker);
+    await dispatcher.drainQueue();
     await expect(queued).resolves.toBe("tok-A");
   });
 
@@ -310,11 +591,10 @@ describe("OpenAIDispatcher quota threshold gate", () => {
     await waitForQueue(dispatcher);
     readOpenAIQuotaUsageMock.mockReset();
     readOpenAIQuotaUsageMock
-      .mockResolvedValueOnce(quotaOk(90))
       .mockResolvedValueOnce(quotaOk(90))
       .mockResolvedValueOnce(quotaOk(10));
 
-    await dispatcher.drainQueue(workerA);
+    await dispatcher.drainQueue();
 
     await expect(queued).resolves.toBe("tok-B");
   });
@@ -323,11 +603,12 @@ describe("OpenAIDispatcher quota threshold gate", () => {
     const dispatcher = new OpenAIDispatcher();
     const worker = fakeWorker({ tokenId: 1, tokenName: "tok-A" });
     addWorkers(dispatcher, [worker]);
+    await dispatcher.onTokenUpdated(1, "tok-A", credentials(), null, 1);
     const execute = vi.fn(async (w: CodexAppServerWorker) => w.tokenName);
     const queued = dispatcher.enqueue(execute);
     await waitForQueue(dispatcher);
 
-    await Promise.all([dispatcher.drainQueue(worker), dispatcher.drainQueue(worker)]);
+    await Promise.all([dispatcher.drainQueue(), dispatcher.drainQueue()]);
 
     await expect(queued).resolves.toBe("tok-A");
     expect(execute).toHaveBeenCalledTimes(1);
@@ -367,7 +648,9 @@ describe("OpenAIDispatcher quota threshold gate", () => {
       );
 
     const result = await dispatcher.generate(
-      baseReq({ reuse: { workerId: 10, threadId: "thread-1", epoch: 1 } }),
+      baseReq({
+        reuse: { workerId: makeOpenAIWorkerId(1, 0), threadId: "thread-1", epoch: 1 },
+      }),
     );
 
     expect(result.tokenName).toBe("tok-B");
@@ -389,7 +672,9 @@ describe("OpenAIDispatcher quota threshold gate", () => {
       );
 
     const result = await dispatcher.generate(
-      baseReq({ reuse: { workerId: 10, threadId: "thread-1", epoch: 1 } }),
+      baseReq({
+        reuse: { workerId: makeOpenAIWorkerId(1, 0), threadId: "thread-1", epoch: 1 },
+      }),
     );
 
     expect(result.tokenName).toBe("tok-A");
@@ -440,7 +725,11 @@ describe("OpenAIDispatcher quota threshold gate", () => {
     readOpenAIQuotaUsageMock.mockResolvedValueOnce(quotaOk(85));
 
     await expect(
-      dispatcher.generate(baseReq({ reuse: { workerId: 10, threadId: "thread-1", epoch: 1 } })),
+      dispatcher.generate(
+        baseReq({
+          reuse: { workerId: makeOpenAIWorkerId(1, 0), threadId: "thread-1", epoch: 1 },
+        }),
+      ),
     ).rejects.toBeInstanceOf(QuotaThresholdExceededError);
   });
 });
@@ -518,7 +807,10 @@ describe("OpenAIDispatcher image generation routing", () => {
   it("forces cold thread for image requests even when a reuse coordinate is present", async () => {
     const dispatcher = new OpenAIDispatcher();
     const worker = await dispatcher.acquireReuseWorker(
-      baseReq({ imageGeneration: true, reuse: { workerId: 10, threadId: "t1", epoch: 1 } }),
+      baseReq({
+        imageGeneration: true,
+        reuse: { workerId: makeOpenAIWorkerId(1, 0), threadId: "t1", epoch: 1 },
+      }),
     );
     expect(worker).toBeNull();
   });