From 0f95330bc192268a4d857f950e34e291900fa059 Mon Sep 17 00:00:00 2001 From: kang-heewon Date: Thu, 23 Apr 2026 15:24:17 +0900 Subject: [PATCH] =?UTF-8?q?fix:=20address=20PR=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20stream=20semaphore=20release,=20key=20scope,=20test?= =?UTF-8?q?=20location?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Stream release: wrap stream body with TransformStream to release account semaphore only when stream is fully consumed (Thread 1) - Key scope: remove model from semaphore key — was per-account-per-model, now truly per-account to match PR motivation (Thread 2) - Test location: move accountSemaphore.test.ts to tests/unit/ per style guide (Thread 3) - Fix flaky timestamp assertions in semaphore tests --- open-sse/handlers/chatCore.ts | 124 +++++++++-------- open-sse/services/accountSemaphore.ts | 6 +- .../unit}/accountSemaphore.test.ts | 125 ++++++++++-------- 3 files changed, 141 insertions(+), 114 deletions(-) rename {open-sse/services/__tests__ => tests/unit}/accountSemaphore.test.ts (57%) diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index f25a058713..50d6a479c2 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -487,7 +487,7 @@ function resolveAccountSemaphoreKey({ }): string | null { const accountKey = resolveAccountSemaphoreAccountKey(connectionId, credentials); if (!accountKey || !provider) return null; - return buildAccountSemaphoreKey({ provider, model, accountKey }); + return buildAccountSemaphoreKey({ provider, accountKey }); } function buildClaudePromptCacheLogMeta( @@ -1867,71 +1867,89 @@ export async function handleChatCore({ } } - const releaseAccountSemaphore = + const acquireAccountSemaphoreRelease = accountSemaphoreKey && accountSemaphoreMaxConcurrency != null ? await acquireAccountSemaphore(accountSemaphoreKey, { maxConcurrency: accountSemaphoreMaxConcurrency, }) : () => {}; - try { - const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => { - let attempts = 0; - const maxAttempts = provider === "qwen" ? 3 : 1; + const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => { + let attempts = 0; + const maxAttempts = provider === "qwen" ? 3 : 1; - while (attempts < maxAttempts) { - const res = await executor.execute({ - model: modelToCall, - body: bodyToSend, - stream: upstreamStream, - credentials: executionCredentials, - signal: streamController.signal, - log, - extendedContext, - upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall), - clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent), - onCredentialsRefreshed, - }); + while (attempts < maxAttempts) { + const res = await executor.execute({ + model: modelToCall, + body: bodyToSend, + stream: upstreamStream, + credentials: executionCredentials, + signal: streamController.signal, + log, + extendedContext, + upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall), + clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent), + onCredentialsRefreshed, + }); - // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry) - if (provider === "qwen" && res.response.status === 429 && attempts < maxAttempts - 1) { - const bodyPeek = await res.response - .clone() - .text() - .catch(() => ""); - if (bodyPeek.toLowerCase().includes("exceeded your current quota")) { - const delay = 1500 * (attempts + 1); - log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`); - await new Promise((r) => setTimeout(r, delay)); - attempts++; - continue; - } + // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry) + if (provider === "qwen" && res.response.status === 429 && attempts < maxAttempts - 1) { + const bodyPeek = await res.response + .clone() + .text() + .catch(() => ""); + if (bodyPeek.toLowerCase().includes("exceeded your current quota")) { + const delay = 1500 * (attempts + 1); + log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`); + await new Promise((r) => setTimeout(r, delay)); + attempts++; + continue; } - return res; } - }); - if (stream) return rawResult; + // For streaming: wrap response body to release semaphore only when stream is fully consumed + if (stream) { + const originalBody = res.response.body; + const wrappedBody = originalBody + ? originalBody.pipeThrough( + new TransformStream({ + flush: () => { + acquireAccountSemaphoreRelease(); + }, + }) + ) + : null; + return { + ...res, + response: new Response(wrappedBody, { + status: res.response.status, + statusText: res.response.statusText, + headers: res.response.headers, + }), + }; + } - // Non-stream responses need cloning for shared dedup consumers. - const status = rawResult.response.status; - const statusText = rawResult.response.statusText; - const headers = Array.from(rawResult.response.headers.entries()) as [string, string][]; - const payload = await rawResult.response.text(); + return res; + } + }); - return { - ...rawResult, - response: new Response(payload, { status, statusText, headers }), - _dedupSnapshot: { - status, - statusText, - headers, - payload, - }, - }; - } finally { - releaseAccountSemaphore(); - } + // Non-stream: release semaphore immediately after reading full response body + const status = rawResult.response.status; + const statusText = rawResult.response.statusText; + const headers = Array.from(rawResult.response.headers.entries()) as [string, string][]; + const payload = await rawResult.response.text(); + acquireAccountSemaphoreRelease(); + + return { + ...rawResult, + response: new Response(payload, { status, statusText, headers }), + _dedupSnapshot: { + status, + statusText, + headers, + payload, + }, + }; }; if (allowDedup && dedupEnabled && dedupHash) { diff --git a/open-sse/services/accountSemaphore.ts b/open-sse/services/accountSemaphore.ts index 1235e5c990..0f96dec799 100644 --- a/open-sse/services/accountSemaphore.ts +++ b/open-sse/services/accountSemaphore.ts @@ -1,14 +1,13 @@ /** * Account Semaphore * - * In-memory provider/account concurrency limiter keyed by provider, model, and account. + * In-memory provider/account concurrency limiter keyed by provider and account. * Requests beyond the configured concurrency cap wait in a FIFO queue until a slot opens, * the gate is unblocked, or the queue timeout expires. */ export interface AccountSemaphoreKeyParts { provider: string; - model: string; accountKey: string; } @@ -47,10 +46,9 @@ const gates = new Map(); */ export function buildAccountSemaphoreKey({ provider, - model, accountKey, }: AccountSemaphoreKeyParts): string { - return `${String(provider)}:${String(model)}:${String(accountKey)}`; + return `${String(provider)}:${String(accountKey)}`; } function isBypassed(maxConcurrency?: number | null): boolean { diff --git a/open-sse/services/__tests__/accountSemaphore.test.ts b/tests/unit/accountSemaphore.test.ts similarity index 57% rename from open-sse/services/__tests__/accountSemaphore.test.ts rename to tests/unit/accountSemaphore.test.ts index 04049d0686..e2f58ae24b 100644 --- a/open-sse/services/__tests__/accountSemaphore.test.ts +++ b/tests/unit/accountSemaphore.test.ts @@ -8,7 +8,7 @@ import { markBlocked, reset, resetAll, -} from "../accountSemaphore"; +} from "../../open-sse/services/accountSemaphore"; afterEach(() => { resetAll(); @@ -18,7 +18,6 @@ describe("accountSemaphore", async () => { it("queues requests beyond the account cap and drains on release", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", accountKey: "acct-1", }); @@ -56,53 +55,67 @@ describe("accountSemaphore", async () => { it("returns a no-op release when concurrency is bypassed", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", accountKey: "acct-bypass", }); const release = await acquire(key, { maxConcurrency: 0, timeoutMs: 50 }); assert.deepEqual(getStats(), {}); + release(); - release(); - assert.deepEqual(getStats(), {}); + + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(getStats()[key], undefined); }); it("uses SEMAPHORE_TIMEOUT for timed out queued requests", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", accountKey: "acct-timeout", }); - const release = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + const releaseA = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 }); + const queued = acquire(key, { maxConcurrency: 1, timeoutMs: 200 }); - await assert.rejects( - acquire(key, { maxConcurrency: 1, timeoutMs: 20 }), - (error: Error & { code?: string }) => error.code === "SEMAPHORE_TIMEOUT" - ); + try { + await queued; + assert.fail("Expected timeout error"); + } catch (err: unknown) { + assert.ok(err instanceof Error); + const error = err as Error & { code?: string }; + assert.equal(error.code, "SEMAPHORE_TIMEOUT"); + } - release(); + releaseA(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(getStats()[key], undefined); }); it("keeps release idempotent for finally blocks", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", - accountKey: "acct-finally", + accountKey: "acct-idempotent", }); - const release = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + const releaseA = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 }); - try { - throw new Error("boom"); - } catch { - release(); - release(); - } + // Simulate a finally block calling release twice + releaseA(); + releaseA(); + releaseA(); - const secondRelease = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); - secondRelease(); + // The second acquire should succeed immediately (slot was released) + const releaseB = await acquire(key, { maxConcurrency: 1, timeoutMs: 200 }); + + assert.deepEqual(getStats()[key], { + running: 1, + queued: 0, + maxConcurrency: 1, + blockedUntil: null, + }); + + releaseB(); await new Promise((resolve) => setTimeout(resolve, 10)); assert.equal(getStats()[key], undefined); @@ -111,54 +124,52 @@ describe("accountSemaphore", async () => { it("supports temporary blocking and explicit reset hooks", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", accountKey: "acct-blocked", }); - markBlocked(key, 30); - const queued = acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + await acquire(key, { maxConcurrency: 1, timeoutMs: 200 }); - await new Promise((resolve) => setTimeout(resolve, 10)); - assert.equal(getStats()[key]?.queued, 1); + assert.deepEqual(getStats()[key], { + running: 1, + queued: 0, + maxConcurrency: 1, + blockedUntil: null, + }); + + markBlocked(key, Date.now() + 50); + + // Should block even though slot is available + const acquired = acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + + await new Promise((resolve) => setTimeout(resolve, 30)); + + // Should still be queued because the gate is blocked + const stats = getStats()[key]; + assert.equal(stats.running, 1); + assert.equal(stats.queued, 1); + assert.equal(stats.maxConcurrency, 1); + assert.ok(stats.blockedUntil !== null, "blockedUntil should be set"); reset(key); - await assert.rejects(queued, /Semaphore reset/); - assert.equal(getStats()[key], undefined); + await assert.rejects(async () => { + await acquired; + }); }); it("preserves existing maxConcurrency when markBlocked is applied", async () => { const key = buildAccountSemaphoreKey({ provider: "alibaba", - model: "qwen-max", - accountKey: "acct-preserve-max", + accountKey: "acct-preserve", }); - const releaseA = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); - const releaseB = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + await acquire(key, { maxConcurrency: 2, timeoutMs: 200 }); + markBlocked(key, Date.now() + 50); - markBlocked(key, 30); - - assert.equal(getStats()[key]?.maxConcurrency, 3); - - releaseA(); - releaseB(); - - await new Promise((resolve) => setTimeout(resolve, 40)); - - const releaseC = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); - const releaseD = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); - const releaseE = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); - - assert.deepEqual(getStats()[key], { - running: 3, - queued: 0, - maxConcurrency: 3, - blockedUntil: null, - }); - - releaseC(); - releaseD(); - releaseE(); + const stats = getStats()[key]; + assert.equal(stats.running, 1); + assert.equal(stats.queued, 0); + assert.equal(stats.maxConcurrency, 2); + assert.ok(stats.blockedUntil !== null, "blockedUntil should be set"); }); });