From 43f2b2c288a12aea90d93ae75fbdde8b4adf6b89 Mon Sep 17 00:00:00 2001 From: quiterunner-commits <276358597+quiterunner-commits@users.noreply.github.com> Date: Mon, 31 Aug 2026 01:39:33 +0300 Subject: [PATCH] fix(sse): refuse an AI Horde queue that cannot fit the request budget (#12143) Validado em worktree combinada com typecheck limpo, testes focados verdes, gates de file-size/complexity/cognitive-complexity/cycles OK. Obrigado! --- .../imageGeneration/providers/aihorde.ts | 84 ++++++++++--- tests/unit/aihorde-queue-budget.test.ts | 114 ++++++++++++++++++ 2 files changed, 183 insertions(+), 15 deletions(-) create mode 100644 tests/unit/aihorde-queue-budget.test.ts diff --git a/open-sse/handlers/imageGeneration/providers/aihorde.ts b/open-sse/handlers/imageGeneration/providers/aihorde.ts index 13fa50f5ab..782f641438 100644 --- a/open-sse/handlers/imageGeneration/providers/aihorde.ts +++ b/open-sse/handlers/imageGeneration/providers/aihorde.ts @@ -17,7 +17,11 @@ import { } from "./aihordeMapRequest.ts"; const GENERATE_TIMEOUT_MS = 600_000; -const POLL_INTERVAL_MS = 1_000; +// Опрос начинается частым и разряжается по мере ожидания: короткая очередь +// отдаёт картинку за секунды, а длинная иначе стоила бы Horde сотен запросов +// по общему анонимному ключу (600 опросов на один кадр при полном бюджете). +const POLL_INTERVAL_MIN_MS = 1_000; +const POLL_INTERVAL_MAX_MS = 8_000; // Per-call bound for the Horde API's own submit/check/status/cancel calls // (a fixed, trusted host — no SSRF guard needed, just a hard timeout so a // hung upstream cannot stall a request indefinitely). Individual calls are @@ -119,6 +123,11 @@ async function fetchHordeImageBytes( return value; } +/** Числовое поле ответа Horde: отсутствующее или нечисловое читается как «неизвестно». */ +function numericField(value: unknown): number | null { + return typeof value === "number" && Number.isFinite(value) ? value : null; +} + export async function handleAiHordeImageGeneration({ model, provider, @@ -212,19 +221,23 @@ export async function handleAiHordeImageGeneration({ } let completed = false; + let pollDelayMs = POLL_INTERVAL_MIN_MS; try { while (true) { if (signal?.aborted) throw new Error("Horde image generation cancelled"); if (Date.now() >= deadline) { throw Object.assign(new Error("Horde image generation timed out"), { status: 504 }); } - await sleep(POLL_INTERVAL_MS); - const checkRes = await safeOutboundFetch(`${AI_HORDE_API_BASE}/v2/generate/check/${jobId}`, { - headers: hordeHeaders(apiKey), - signal: signal ?? undefined, - guard: "none", - timeoutMs: boundedTimeoutMs(deadline, HORDE_API_CALL_TIMEOUT_MS), - }); + await sleep(pollDelayMs); + const checkRes = await safeOutboundFetch( + `${AI_HORDE_API_BASE}/v2/generate/check/${jobId}`, + { + headers: hordeHeaders(apiKey), + signal: signal ?? undefined, + guard: "none", + timeoutMs: boundedTimeoutMs(deadline, HORDE_API_CALL_TIMEOUT_MS), + } + ); const check = await safeJson(checkRes); if (!checkRes.ok || !check || typeof check !== "object") { throw Object.assign( @@ -239,14 +252,55 @@ export async function handleAiHordeImageGeneration({ status: 503, }); } - if (!checkObj.done) continue; + if (!checkObj.done) { + // Horde сообщает оценку ожидания в первом же ответе. Если она не + // помещается в остаток бюджета, ждать нечего: запрос всё равно + // упал бы по таймауту, только молча и десятью минутами позже. + // Отказ называет очередь и число воркеров — по ним видно, что + // выручает не терпение, а модель с большим числом воркеров. + const waitSeconds = numericField(checkObj.wait_time); + const remainingMs = deadline - Date.now(); + if (waitSeconds !== null && waitSeconds * 1_000 > remainingMs) { + const queuePosition = numericField(checkObj.queue_position); + const workers = numericField(checkObj.eligible_workers); + const details = [ + `queue wait ~${Math.round(waitSeconds)}s`, + queuePosition !== null ? `position ${queuePosition}` : null, + workers !== null ? `${workers} eligible worker(s)` : null, + `budget ${Math.round(remainingMs / 1_000)}s left`, + ] + .filter(Boolean) + .join(", "); + throw Object.assign( + new Error( + `Horde queue is longer than the request budget (${details}). ` + + `Pick a model with more workers or raise the timeout.` + ), + { status: 504 } + ); + } + // Разрядка опроса: десятая доля оставшегося ожидания, в рамках + // минимума и максимума. Короткая очередь по-прежнему опрашивается + // раз в секунду. + pollDelayMs = + waitSeconds === null + ? POLL_INTERVAL_MIN_MS + : Math.min( + POLL_INTERVAL_MAX_MS, + Math.max(POLL_INTERVAL_MIN_MS, Math.round((waitSeconds * 1_000) / 10)) + ); + continue; + } - const statusRes = await safeOutboundFetch(`${AI_HORDE_API_BASE}/v2/generate/status/${jobId}`, { - headers: hordeHeaders(apiKey), - signal: signal ?? undefined, - guard: "none", - timeoutMs: boundedTimeoutMs(deadline, HORDE_API_CALL_TIMEOUT_MS), - }); + const statusRes = await safeOutboundFetch( + `${AI_HORDE_API_BASE}/v2/generate/status/${jobId}`, + { + headers: hordeHeaders(apiKey), + signal: signal ?? undefined, + guard: "none", + timeoutMs: boundedTimeoutMs(deadline, HORDE_API_CALL_TIMEOUT_MS), + } + ); const status = await safeJson(statusRes); if (!statusRes.ok || !status || typeof status !== "object") { throw Object.assign( diff --git a/tests/unit/aihorde-queue-budget.test.ts b/tests/unit/aihorde-queue-budget.test.ts new file mode 100644 index 0000000000..8f8ac0e782 --- /dev/null +++ b/tests/unit/aihorde-queue-budget.test.ts @@ -0,0 +1,114 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +process.env.DATA_DIR = mkdtempSync(join(tmpdir(), "omniroute-aihorde-queue-")); + +import { handleAiHordeImageGeneration } from "../../open-sse/handlers/imageGeneration/providers/aihorde.ts"; +import { aiHordeImageCatalog } from "../../open-sse/services/aihordeImageCatalog.ts"; + +/** + * Очередь Horde известна с первого ответа — отказывать надо там же. + * + * `/v2/generate/check` возвращает `wait_time` и `queue_position` сразу. Пока + * они не читались, запрос на модель с длинной очередью опрашивал Horde раз в + * секунду весь бюджет и падал по таймауту, ничего не объяснив. Живая проверка + * 2026-08-30: модель `Deliberate` (3 воркера) ответила `wait_time: 1478, + * queue_position: 321` — 25 минут при бюджете в 10. Ждать было бессмысленно + * ещё до первого опроса, а пользователь узнавал об этом через десять минут. + * + * Для сравнения `stable_diffusion` (10 воркеров) в тот же момент отдал картинку + * за 10.4 секунды — то есть отказ должен быть про эту модель и эту очередь, а + * не про Horde вообще, и должен подсказывать, что делать. + */ + +const HORDE_JOB_ID = "queue-budget-job"; + +function stubHordeQueue({ waitTimeSeconds }: { waitTimeSeconds: number }) { + const calls: string[] = []; + globalThis.fetch = (async (input: string | URL, init?: RequestInit) => { + const url = String(input); + const method = (init?.method || "GET").toUpperCase(); + calls.push(`${method} ${url}`); + + if (url.endsWith("/v2/generate/async")) { + return new Response(JSON.stringify({ id: HORDE_JOB_ID, kudos: 6 }), { status: 202 }); + } + if (url.includes("/v2/generate/check/")) { + return new Response( + JSON.stringify({ + done: false, + faulted: false, + is_possible: true, + waiting: 1, + wait_time: waitTimeSeconds, + queue_position: 321, + eligible_workers: 3, + }), + { status: 200 } + ); + } + return new Response("{}", { status: 200 }); + }) as typeof fetch; + return calls; +} + +const originalFetch = globalThis.fetch; + +test.beforeEach(() => { + aiHordeImageCatalog.replace([ + { name: "Deliberate", count: 3, queued: 0, eta: 1478, performance: 1, jobs: 0 }, + ]); +}); + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +test("refuses immediately when the queue cannot fit the remaining budget", async () => { + const calls = stubHordeQueue({ waitTimeSeconds: 1478 }); + + const started = Date.now(); + const result = await handleAiHordeImageGeneration({ + model: "Deliberate", + provider: "aihorde", + body: { model: "aihorde/Deliberate", prompt: "hello world" }, + credentials: { apiKey: "horde-key" }, + timeoutMs: 600_000, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 504); + assert.match(String(result.error), /queue/i); + assert.match(String(result.error), /1478|25/, "в отказе должно быть названо ожидание"); + + assert.ok( + Date.now() - started < 30_000, + "отказ обязан прийти сразу, а не после выработки бюджета" + ); + const checks = calls.filter((c) => c.includes("/v2/generate/check/")); + assert.equal(checks.length, 1, "хватает одного опроса, чтобы узнать очередь"); +}); + +test("keeps waiting when the queue fits the budget", async () => { + // Ожидание заведомо помещается в бюджет: 2 секунды очереди против 6. + const calls = stubHordeQueue({ waitTimeSeconds: 2 }); + + const result = await handleAiHordeImageGeneration({ + model: "Deliberate", + provider: "aihorde", + body: { model: "aihorde/Deliberate", prompt: "hello world" }, + credentials: { apiKey: "horde-key" }, + timeoutMs: 6_000, + }); + + // Заглушка никогда не отвечает done, поэтому запрос доходит до собственного + // таймаута — важно, что он до него дошёл, а не был отбит по очереди. + assert.equal(result.success, false); + assert.ok( + calls.filter((c) => c.includes("/v2/generate/check/")).length > 1, + "короткая очередь не должна приводить к раннему отказу" + ); +});