From b1bfb28006383f3fa698da53d9831a5b73e05a0d Mon Sep 17 00:00:00 2001 From: Jahanzaib Tayyab Date: Sun, 26 Apr 2026 12:39:16 +0500 Subject: [PATCH] feat(core): cap concurrent in-flight requests per provider (#3409) Several users hit "Error: 429 Too many concurrent requests for this model" against rate-limited backends, especially when sub-agents fan out or when /compress runs alongside the main turn. There was no way to gate this on the client; the only available knob was retry-on-429, which both hides the symptom and burns latency. This PR adds a per-provider request-concurrency cap that translates upstream rate limits into client-side back-pressure. Excess callers queue FIFO behind a semaphore; behavior is identical to today when no cap is configured. Wiring: - ContentGeneratorConfig.requestConcurrency -- new optional field (provider-scoped via modelProviders[].generationConfig). - QWEN_REQUEST_CONCURRENCY env var as a fallback when the per-provider field is unset, so the cap can be applied process-wide without editing settings.json. - 0 / undefined / negative => unlimited (default; no behavior change). Implementation: - packages/core/src/utils/concurrencyLimiter.ts -- minimal FIFO counting semaphore. acquire/release pair plus a runExclusive helper. Release is idempotent so accidental double-release does not over-credit. Capacity <= 0 short-circuits to a no-op. - packages/core/src/core/rateLimitedContentGenerator.ts -- wrapper implementing ContentGenerator. Gates generateContent and generateContentStream; intentionally does NOT gate countTokens or embedContent (cheap, often local). For streaming the slot is held until the consumer fully drains the iterator (or it errors, throws during init, or unwinds early via break/return) -- the rate limit the upstream actually counts is concurrent open connections, not request initiation. - packages/core/src/core/contentGenerator.ts -- createContentGenerator applies RateLimitedContentGenerator before LoggingContentGenerator so rate-limit wait time is visible inside the logged request span. Tests - 6 cases in concurrencyLimiter.test.ts (cap respected, FIFO, success + failure release, double-release idempotency, NaN/Infinity coerced to unlimited). - 7 cases in rateLimitedContentGenerator.test.ts (peak in-flight, per-call slot release on rejection, stream slot held until drain, release on init throw, release on consumer-side break, no gating for countTokens/embedContent, useSummarizedThinking pass-through). - 6 cases extending contentGenerator.test.ts pin the resolution precedence: not wrapped by default, wrapped when configured, env fallback, config-over-env, garbage env values treated as unlimited, fractional config values floored. Affected suite: 6,093 core tests pass; full TypeScript build clean. Prepared with assistance from Claude (Anthropic) under human review. --- docs/users/configuration/settings.md | 25 +- packages/cli/src/config/settingsSchema.ts | 11 + .../core/src/core/contentGenerator.test.ts | 130 ++++++++- packages/core/src/core/contentGenerator.ts | 49 +++- .../core/rateLimitedContentGenerator.test.ts | 275 ++++++++++++++++++ .../src/core/rateLimitedContentGenerator.ts | 98 +++++++ .../core/src/utils/concurrencyLimiter.test.ts | 108 +++++++ packages/core/src/utils/concurrencyLimiter.ts | 96 ++++++ .../schemas/settings.schema.json | 4 + 9 files changed, 783 insertions(+), 13 deletions(-) create mode 100644 packages/core/src/core/rateLimitedContentGenerator.test.ts create mode 100644 packages/core/src/core/rateLimitedContentGenerator.ts create mode 100644 packages/core/src/utils/concurrencyLimiter.test.ts create mode 100644 packages/core/src/utils/concurrencyLimiter.ts diff --git a/docs/users/configuration/settings.md b/docs/users/configuration/settings.md index 8ea03777995..00032abac0d 100644 --- a/docs/users/configuration/settings.md +++ b/docs/users/configuration/settings.md @@ -133,17 +133,17 @@ Settings are organized into categories. All settings should be placed within the #### model -| Setting | Type | Description | Default | -| -------------------------------------------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------- | -| `model.name` | string | The Qwen model to use for conversations. | `undefined` | -| `model.maxSessionTurns` | number | Maximum number of user/model/tool turns to keep in a session. -1 means unlimited. | `-1` | -| `model.generationConfig` | object | Advanced overrides passed to the underlying content generator. Supports request controls such as `timeout`, `maxRetries`, `enableCacheControl`, `contextWindowSize` (override model's context window size), `modalities` (override auto-detected input modalities), `customHeaders` (custom HTTP headers for API requests), and `extra_body` (additional body parameters for OpenAI-compatible API requests only), along with fine-tuning knobs under `samplingParams` (for example `temperature`, `top_p`, `max_tokens`). Leave unset to rely on provider defaults. | `undefined` | -| `model.chatCompression.contextPercentageThreshold` | number | Sets the threshold for chat history compression as a percentage of the model's total token limit. This is a value between 0 and 1 that applies to both automatic compression and the manual `/compress` command. For example, a value of `0.6` will trigger compression when the chat history exceeds 60% of the token limit. Use `0` to disable compression entirely. | `0.7` | -| `model.skipNextSpeakerCheck` | boolean | Skip the next speaker check. | `false` | -| `model.skipLoopDetection` | boolean | Disables loop detection checks. Loop detection prevents infinite loops in AI responses but can generate false positives that interrupt legitimate workflows. Enable this option if you experience frequent false positive loop detection interruptions. | `false` | -| `model.skipStartupContext` | boolean | Skips sending the startup workspace context (environment summary and acknowledgement) at the beginning of each session. Enable this if you prefer to provide context manually or want to save tokens on startup. | `false` | -| `model.enableOpenAILogging` | boolean | Enables logging of OpenAI API calls for debugging and analysis. When enabled, API requests and responses are logged to JSON files. | `false` | -| `model.openAILoggingDir` | string | Custom directory path for OpenAI API logs. If not specified, defaults to `logs/openai` in the current working directory. Supports absolute paths, relative paths (resolved from current working directory), and `~` expansion (home directory). | `undefined` | +| Setting | Type | Description | Default | +| -------------------------------------------------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------- | +| `model.name` | string | The Qwen model to use for conversations. | `undefined` | +| `model.maxSessionTurns` | number | Maximum number of user/model/tool turns to keep in a session. -1 means unlimited. | `-1` | +| `model.generationConfig` | object | Advanced overrides passed to the underlying content generator. Supports request controls such as `timeout`, `maxRetries`, `requestConcurrency` (cap concurrent in-flight requests; see below), `enableCacheControl`, `contextWindowSize` (override model's context window size), `modalities` (override auto-detected input modalities), `customHeaders` (custom HTTP headers for API requests), and `extra_body` (additional body parameters for OpenAI-compatible API requests only), along with fine-tuning knobs under `samplingParams` (for example `temperature`, `top_p`, `max_tokens`). Leave unset to rely on provider defaults. | `undefined` | +| `model.chatCompression.contextPercentageThreshold` | number | Sets the threshold for chat history compression as a percentage of the model's total token limit. This is a value between 0 and 1 that applies to both automatic compression and the manual `/compress` command. For example, a value of `0.6` will trigger compression when the chat history exceeds 60% of the token limit. Use `0` to disable compression entirely. | `0.7` | +| `model.skipNextSpeakerCheck` | boolean | Skip the next speaker check. | `false` | +| `model.skipLoopDetection` | boolean | Disables loop detection checks. Loop detection prevents infinite loops in AI responses but can generate false positives that interrupt legitimate workflows. Enable this option if you experience frequent false positive loop detection interruptions. | `false` | +| `model.skipStartupContext` | boolean | Skips sending the startup workspace context (environment summary and acknowledgement) at the beginning of each session. Enable this if you prefer to provide context manually or want to save tokens on startup. | `false` | +| `model.enableOpenAILogging` | boolean | Enables logging of OpenAI API calls for debugging and analysis. When enabled, API requests and responses are logged to JSON files. | `false` | +| `model.openAILoggingDir` | string | Custom directory path for OpenAI API logs. If not specified, defaults to `logs/openai` in the current working directory. Supports absolute paths, relative paths (resolved from current working directory), and `~` expansion (home directory). | `undefined` | **Example model.generationConfig:** @@ -153,6 +153,7 @@ Settings are organized into categories. All settings should be placed within the "generationConfig": { "timeout": 60000, "contextWindowSize": 128000, + "requestConcurrency": 4, "modalities": { "image": true }, @@ -173,6 +174,8 @@ Settings are organized into categories. All settings should be placed within the } ``` +**`requestConcurrency`:** Caps the number of in-flight requests against this provider so the upstream `429 Too many concurrent requests for this model` error never fires — excess callers (typically sub-agents firing in parallel, or `/compress` interleaving with the main turn) wait FIFO instead. Set to `0` (default) for no limit. The environment variable `QWEN_REQUEST_CONCURRENCY` is honored as a fallback when the provider-level value is unset, so the cap can be applied process-wide without editing `settings.json`. + **max_tokens (adaptive output tokens):** When `samplingParams.max_tokens` is not set, Qwen Code uses an adaptive output token strategy to optimize GPU resource usage: diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 9bfa131102f..7b84a9dda58 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -914,6 +914,17 @@ const SETTINGS_SCHEMA = { parentKey: 'generationConfig', showInDialog: false, }, + requestConcurrency: { + type: 'number', + label: 'Request Concurrency', + category: 'Generation Configuration', + requiresRestart: false, + default: undefined as number | undefined, + description: + 'Maximum number of concurrent in-flight requests against this model. Useful when the upstream returns 429 "Too many concurrent requests for this model". 0 / unset means unlimited. Falls back to QWEN_REQUEST_CONCURRENCY env var.', + parentKey: 'generationConfig', + showInDialog: false, + }, enableCacheControl: { type: 'boolean', label: 'Enable Cache Control', diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index bb8e5f7418a..b6b08e5a496 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect, vi, afterEach } from 'vitest'; import { createContentGenerator, createContentGeneratorConfig, @@ -13,6 +13,7 @@ import { import { GoogleGenAI } from '@google/genai'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from './loggingContentGenerator/index.js'; +import { RateLimitedContentGenerator } from './rateLimitedContentGenerator.js'; vi.mock('@google/genai'); @@ -83,6 +84,133 @@ describe('createContentGenerator', () => { }); }); +describe('createContentGenerator request concurrency wrapping (#3409)', () => { + const baseMockConfig = { + getUsageStatisticsEnabled: () => false, + getContentGeneratorConfig: () => ({}), + getCliVersion: () => '1.0.0', + } as unknown as Config; + + const buildGoogleGenAIMock = () => { + const mockGenerator = { models: {} } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + }; + + const originalEnv = process.env['QWEN_REQUEST_CONCURRENCY']; + + afterEach(() => { + if (originalEnv === undefined) { + delete process.env['QWEN_REQUEST_CONCURRENCY']; + } else { + process.env['QWEN_REQUEST_CONCURRENCY'] = originalEnv; + } + }); + + it('does not wrap with rate-limiter when no concurrency is configured', async () => { + buildGoogleGenAIMock(); + delete process.env['QWEN_REQUEST_CONCURRENCY']; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + baseMockConfig, + )) as LoggingContentGenerator; + expect(generator.getWrapped()).not.toBeInstanceOf( + RateLimitedContentGenerator, + ); + }); + + it('wraps with RateLimitedContentGenerator when requestConcurrency is set', async () => { + buildGoogleGenAIMock(); + delete process.env['QWEN_REQUEST_CONCURRENCY']; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + requestConcurrency: 4, + }, + baseMockConfig, + )) as LoggingContentGenerator; + const wrapped = generator.getWrapped(); + expect(wrapped).toBeInstanceOf(RateLimitedContentGenerator); + const limiter = (wrapped as RateLimitedContentGenerator).getLimiter(); + expect(limiter.limit).toBe(4); + }); + + it('falls back to QWEN_REQUEST_CONCURRENCY env var when config is unset', async () => { + buildGoogleGenAIMock(); + process.env['QWEN_REQUEST_CONCURRENCY'] = '6'; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + baseMockConfig, + )) as LoggingContentGenerator; + const wrapped = generator.getWrapped(); + expect(wrapped).toBeInstanceOf(RateLimitedContentGenerator); + expect((wrapped as RateLimitedContentGenerator).getLimiter().limit).toBe(6); + }); + + it('config value wins over the env var', async () => { + buildGoogleGenAIMock(); + process.env['QWEN_REQUEST_CONCURRENCY'] = '6'; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + requestConcurrency: 2, + }, + baseMockConfig, + )) as LoggingContentGenerator; + const wrapped = generator.getWrapped(); + expect(wrapped).toBeInstanceOf(RateLimitedContentGenerator); + expect((wrapped as RateLimitedContentGenerator).getLimiter().limit).toBe(2); + }); + + it.each([ + ['0', 0], + ['-1', 0], + ['nope', 0], + [' ', 0], + ])('treats env var %s as unlimited', async (envVal, _expected) => { + buildGoogleGenAIMock(); + process.env['QWEN_REQUEST_CONCURRENCY'] = envVal; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + baseMockConfig, + )) as LoggingContentGenerator; + expect(generator.getWrapped()).not.toBeInstanceOf( + RateLimitedContentGenerator, + ); + }); + + it('floors fractional config values', async () => { + buildGoogleGenAIMock(); + delete process.env['QWEN_REQUEST_CONCURRENCY']; + const generator = (await createContentGenerator( + { + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + requestConcurrency: 3.7, + }, + baseMockConfig, + )) as LoggingContentGenerator; + const wrapped = generator.getWrapped() as RateLimitedContentGenerator; + expect(wrapped.getLimiter().limit).toBe(3); + }); +}); + describe('createContentGeneratorConfig', () => { const mockConfig = { getProxy: () => undefined, diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 7b161423614..8cf8219ecf0 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -14,6 +14,8 @@ import type { } from '@google/genai'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from './loggingContentGenerator/index.js'; +import { RateLimitedContentGenerator } from './rateLimitedContentGenerator.js'; +import { ConcurrencyLimiter } from '../utils/concurrencyLimiter.js'; import type { ConfigSource, ConfigSourceKind, @@ -83,6 +85,13 @@ export type ContentGeneratorConfig = { timeout?: number; // Timeout configuration in milliseconds maxRetries?: number; // Maximum retries for rate-limit errors retryErrorCodes?: number[]; // Additional error codes that trigger rate-limit retry + /** + * Maximum number of concurrent in-flight requests against the model API. + * Use this when the upstream returns ``429 Too many concurrent requests`` + * for the configured model -- excess callers wait FIFO instead of erroring + * (#3409). ``0`` / ``undefined`` / negative => no limit (default). + */ + requestConcurrency?: number; enableCacheControl?: boolean; // Enable cache control for DashScope providers samplingParams?: { top_p?: number; @@ -356,5 +365,43 @@ export async function createContentGenerator( ); } - return new LoggingContentGenerator(baseGenerator, config, generatorConfig); + // Apply per-provider concurrency limit (#3409) before logging so rate-limit + // back-pressure is visible to the logger as ``in flight`` time, not as a + // separate phase outside of the request span. + const limit = resolveRequestConcurrency(generatorConfig.requestConcurrency); + const wrapped = + limit > 0 + ? new RateLimitedContentGenerator( + baseGenerator, + new ConcurrencyLimiter(limit), + ) + : baseGenerator; + + return new LoggingContentGenerator(wrapped, config, generatorConfig); +} + +/** + * Resolve the effective request-concurrency cap from (in order): + * 1. ``ContentGeneratorConfig.requestConcurrency`` (provider-level setting), + * 2. ``QWEN_REQUEST_CONCURRENCY`` env var. + * Anything <= 0, NaN, or otherwise non-numeric is treated as "unlimited" (0). + */ +function resolveRequestConcurrency(configured: number | undefined): number { + const fromConfig = normalizeConcurrency(configured); + if (fromConfig > 0) { + return fromConfig; + } + const envRaw = process.env['QWEN_REQUEST_CONCURRENCY']; + if (envRaw === undefined) { + return 0; + } + const parsed = Number.parseInt(envRaw.trim(), 10); + return normalizeConcurrency(parsed); +} + +function normalizeConcurrency(value: number | undefined): number { + if (value === undefined || !Number.isFinite(value) || value <= 0) { + return 0; + } + return Math.floor(value); } diff --git a/packages/core/src/core/rateLimitedContentGenerator.test.ts b/packages/core/src/core/rateLimitedContentGenerator.test.ts new file mode 100644 index 00000000000..1cf60d73538 --- /dev/null +++ b/packages/core/src/core/rateLimitedContentGenerator.test.ts @@ -0,0 +1,275 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import { + GenerateContentResponse, + type GenerateContentParameters, + type CountTokensParameters, + type CountTokensResponse, + type EmbedContentParameters, + type EmbedContentResponse, +} from '@google/genai'; +import type { ContentGenerator } from './contentGenerator.js'; +import { RateLimitedContentGenerator } from './rateLimitedContentGenerator.js'; +import { ConcurrencyLimiter } from '../utils/concurrencyLimiter.js'; + +const fakeRequest = (id: string): GenerateContentParameters => ({ + model: 'test-model', + contents: [{ role: 'user', parts: [{ text: id }] }], +}); + +const fakeResponse = (id: string): GenerateContentResponse => { + const response = new GenerateContentResponse(); + response.responseId = id; + return response; +}; + +const requestId = (request: GenerateContentParameters): string => { + const contents = Array.isArray(request.contents) + ? request.contents + : [request.contents]; + const first = contents[0]; + if (typeof first === 'object' && first !== null && 'parts' in first) { + const parts = first.parts; + if (parts && typeof parts[0] === 'object' && 'text' in parts[0]) { + return parts[0].text ?? ''; + } + } + return ''; +}; + +class FakeGenerator implements ContentGenerator { + inFlight = 0; + peakInFlight = 0; + startedIds: string[] = []; + finishedIds: string[] = []; + countTokensCalls = 0; + embedCalls = 0; + generateImpl?: ( + req: GenerateContentParameters, + ) => Promise; + streamImpl?: ( + req: GenerateContentParameters, + ) => AsyncGenerator; + + async generateContent( + request: GenerateContentParameters, + ): Promise { + const id = requestId(request); + this.startedIds.push(id); + this.inFlight++; + this.peakInFlight = Math.max(this.peakInFlight, this.inFlight); + try { + if (this.generateImpl) { + return await this.generateImpl(request); + } + // Default: hold the slot until the test resolves it externally. + await new Promise((resolve) => setTimeout(resolve, 5)); + return fakeResponse(id); + } finally { + this.inFlight--; + this.finishedIds.push(id); + } + } + + async generateContentStream( + request: GenerateContentParameters, + ): Promise> { + const id = requestId(request); + this.startedIds.push(id); + this.inFlight++; + this.peakInFlight = Math.max(this.peakInFlight, this.inFlight); + + const baseImpl = this.streamImpl ?? defaultStreamImpl; + const inner = baseImpl(request); + const onChunkExit = () => { + this.inFlight--; + this.finishedIds.push(id); + }; + async function* outer(): AsyncGenerator { + try { + for await (const chunk of inner) { + yield chunk; + } + } finally { + onChunkExit(); + } + } + return outer(); + } + + async countTokens( + _request: CountTokensParameters, + ): Promise { + this.countTokensCalls++; + return { totalTokens: 1 } as CountTokensResponse; + } + + async embedContent( + _request: EmbedContentParameters, + ): Promise { + this.embedCalls++; + return {} as EmbedContentResponse; + } + + useSummarizedThinking(): boolean { + return false; + } +} + +async function* defaultStreamImpl( + request: GenerateContentParameters, +): AsyncGenerator { + const id = requestId(request); + yield fakeResponse(`${id}-1`); + yield fakeResponse(`${id}-2`); +} + +describe('RateLimitedContentGenerator', () => { + it('caps generateContent to the configured limit', async () => { + const inner = new FakeGenerator(); + const limiter = new ConcurrencyLimiter(2); + const gen = new RateLimitedContentGenerator(inner, limiter); + + const requests = await Promise.all([ + gen.generateContent(fakeRequest('a'), 'p'), + gen.generateContent(fakeRequest('b'), 'p'), + gen.generateContent(fakeRequest('c'), 'p'), + gen.generateContent(fakeRequest('d'), 'p'), + ]); + + expect(requests).toHaveLength(4); + expect(inner.peakInFlight).toBeLessThanOrEqual(2); + expect(inner.startedIds.sort()).toEqual(['a', 'b', 'c', 'd']); + }); + + it('releases the slot after generateContent rejects', async () => { + const inner = new FakeGenerator(); + inner.generateImpl = async () => { + throw new Error('upstream 500'); + }; + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + + await expect(gen.generateContent(fakeRequest('a'), 'p')).rejects.toThrow( + 'upstream 500', + ); + expect(limiter.inFlight).toBe(0); + + // The cap is still honored on the next call. + inner.generateImpl = undefined; + await gen.generateContent(fakeRequest('b'), 'p'); + expect(limiter.inFlight).toBe(0); + }); + + it('releases the slot only after the stream is fully drained', async () => { + const inner = new FakeGenerator(); + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + + const stream1 = await gen.generateContentStream(fakeRequest('a'), 'p'); + expect(limiter.inFlight).toBe(1); + + // A second stream call is parked until the first drains. + let stream2Resolved = false; + const stream2Promise = gen + .generateContentStream(fakeRequest('b'), 'p') + .then((s) => { + stream2Resolved = true; + return s; + }); + + await Promise.resolve(); + expect(stream2Resolved).toBe(false); + + // Drain stream1. + const collected1: string[] = []; + for await (const chunk of stream1) { + collected1.push(chunk.responseId ?? ''); + } + expect(collected1).toEqual(['a-1', 'a-2']); + + const stream2 = await stream2Promise; + const collected2: string[] = []; + for await (const chunk of stream2) { + collected2.push(chunk.responseId ?? ''); + } + expect(collected2).toEqual(['b-1', 'b-2']); + expect(limiter.inFlight).toBe(0); + }); + + it('releases the slot when the stream initialization throws', async () => { + const inner = new FakeGenerator(); + // ``generateContentStream`` may throw synchronously on bad input or + // network setup; the wrapper has to handle both that and async rejection + // identically. + vi.spyOn(inner, 'generateContentStream').mockImplementation(() => { + throw new Error('boom'); + }); + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + + await expect( + gen.generateContentStream(fakeRequest('a'), 'p'), + ).rejects.toThrow('boom'); + expect(limiter.inFlight).toBe(0); + }); + + it('releases the slot when stream iteration aborts mid-flight', async () => { + const inner = new FakeGenerator(); + inner.streamImpl = + async function* (): AsyncGenerator { + yield fakeResponse('x-1'); + yield fakeResponse('x-2'); + yield fakeResponse('x-3'); + }; + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + + const stream = await gen.generateContentStream(fakeRequest('x'), 'p'); + let count = 0; + for await (const _chunk of stream) { + count++; + if (count === 1) { + break; // simulates an early consumer-side abort + } + } + expect(limiter.inFlight).toBe(0); + }); + + it('does not gate countTokens / embedContent', async () => { + const inner = new FakeGenerator(); + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + + // Hold the only slot with a generateContent call. + const hold = gen.generateContent(fakeRequest('a'), 'p'); + + // countTokens / embedContent should run immediately. + const start = Date.now(); + await Promise.all([ + gen.countTokens({} as CountTokensParameters), + gen.embedContent({} as EmbedContentParameters), + ]); + expect(Date.now() - start).toBeLessThan(50); + expect(inner.countTokensCalls).toBe(1); + expect(inner.embedCalls).toBe(1); + + await hold; + }); + + it('forwards useSummarizedThinking without blocking', () => { + const inner = new FakeGenerator(); + const limiter = new ConcurrencyLimiter(1); + const gen = new RateLimitedContentGenerator(inner, limiter); + expect(gen.useSummarizedThinking()).toBe(false); + + const spy = vi.spyOn(inner, 'useSummarizedThinking').mockReturnValue(true); + expect(gen.useSummarizedThinking()).toBe(true); + expect(spy).toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/core/rateLimitedContentGenerator.ts b/packages/core/src/core/rateLimitedContentGenerator.ts new file mode 100644 index 00000000000..30b0e0d0c3f --- /dev/null +++ b/packages/core/src/core/rateLimitedContentGenerator.ts @@ -0,0 +1,98 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + CountTokensParameters, + CountTokensResponse, + EmbedContentParameters, + EmbedContentResponse, + GenerateContentParameters, + GenerateContentResponse, +} from '@google/genai'; +import type { ContentGenerator } from './contentGenerator.js'; +import type { ConcurrencyLimiter } from '../utils/concurrencyLimiter.js'; + +/** + * Wraps a ``ContentGenerator`` and gates concurrent ``generateContent`` / + * ``generateContentStream`` calls behind a semaphore so a configured + * per-provider cap (#3409) prevents the upstream + * ``429 Too many concurrent requests for this model`` error from ever firing. + * + * - ``countTokens`` and ``embedContent`` are intentionally **not** gated: + * they are cheap, often local, and rate-limit responses there are rare. + * - For streaming, the slot is held until the consumer fully drains the + * iterator (or it errors / aborts). This matches what the rate-limit + * actually counts: an open connection, not just request initiation. + * - When the limiter has no cap (``capacity <= 0``), this wrapper is a + * thin pass-through with a single ``await Promise.resolve()`` worth of + * overhead. + */ +export class RateLimitedContentGenerator implements ContentGenerator { + constructor( + private readonly wrapped: ContentGenerator, + private readonly limiter: ConcurrencyLimiter, + ) {} + + /** Exposed for tests / observability. */ + getLimiter(): ConcurrencyLimiter { + return this.limiter; + } + + async generateContent( + request: GenerateContentParameters, + userPromptId: string, + ): Promise { + return this.limiter.runExclusive(() => + this.wrapped.generateContent(request, userPromptId), + ); + } + + async generateContentStream( + request: GenerateContentParameters, + userPromptId: string, + ): Promise> { + const release = await this.limiter.acquire(); + let stream: AsyncGenerator; + try { + stream = await this.wrapped.generateContentStream(request, userPromptId); + } catch (error) { + release(); + throw error; + } + return this.gateStream(stream, release); + } + + countTokens(request: CountTokensParameters): Promise { + return this.wrapped.countTokens(request); + } + + embedContent(request: EmbedContentParameters): Promise { + return this.wrapped.embedContent(request); + } + + useSummarizedThinking(): boolean { + return this.wrapped.useSummarizedThinking(); + } + + /** + * Wrap the underlying iterator so the slot is released exactly once, + * regardless of how iteration ends -- normal completion, thrown error, + * caller-side ``return()`` (e.g. consumer unwound a ``for await`` early), + * or upstream abort. + */ + private async *gateStream( + stream: AsyncGenerator, + release: () => void, + ): AsyncGenerator { + try { + for await (const chunk of stream) { + yield chunk; + } + } finally { + release(); + } + } +} diff --git a/packages/core/src/utils/concurrencyLimiter.test.ts b/packages/core/src/utils/concurrencyLimiter.test.ts new file mode 100644 index 00000000000..3e5a4c4a9d4 --- /dev/null +++ b/packages/core/src/utils/concurrencyLimiter.test.ts @@ -0,0 +1,108 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { ConcurrencyLimiter } from './concurrencyLimiter.js'; + +describe('ConcurrencyLimiter', () => { + it('treats capacity <= 0 as unlimited and resolves immediately', async () => { + const limiter = new ConcurrencyLimiter(0); + const release = await limiter.acquire(); + expect(typeof release).toBe('function'); + expect(limiter.limit).toBe(0); + expect(limiter.inFlight).toBe(0); + release(); + release(); // double-release on the no-op path is safe + }); + + it('caps in-flight count at the configured capacity', async () => { + const limiter = new ConcurrencyLimiter(2); + const r1 = await limiter.acquire(); + const r2 = await limiter.acquire(); + expect(limiter.inFlight).toBe(2); + + let r3Resolved = false; + const r3Promise = limiter.acquire().then((release) => { + r3Resolved = true; + return release; + }); + + // r3 is parked behind the cap. + await Promise.resolve(); + expect(r3Resolved).toBe(false); + expect(limiter.queued).toBe(1); + + r1(); + const r3 = await r3Promise; + expect(r3Resolved).toBe(true); + expect(limiter.inFlight).toBe(2); // r2 + r3 + + r2(); + r3(); + expect(limiter.inFlight).toBe(0); + }); + + it('preserves FIFO ordering of waiters', async () => { + const limiter = new ConcurrencyLimiter(1); + const r1 = await limiter.acquire(); + + const order: number[] = []; + const p2 = limiter.acquire().then((release) => { + order.push(2); + release(); + }); + const p3 = limiter.acquire().then((release) => { + order.push(3); + release(); + }); + const p4 = limiter.acquire().then((release) => { + order.push(4); + release(); + }); + + r1(); + await Promise.all([p2, p3, p4]); + expect(order).toEqual([2, 3, 4]); + }); + + it('runExclusive releases the slot on success and on failure', async () => { + const limiter = new ConcurrencyLimiter(1); + + const ok = await limiter.runExclusive(async () => 'value'); + expect(ok).toBe('value'); + expect(limiter.inFlight).toBe(0); + + await expect( + limiter.runExclusive(async () => { + throw new Error('boom'); + }), + ).rejects.toThrow('boom'); + expect(limiter.inFlight).toBe(0); + }); + + it('release is idempotent so accidental double-releases do not over-credit', async () => { + const limiter = new ConcurrencyLimiter(1); + const release = await limiter.acquire(); + release(); + release(); // second call is a no-op + expect(limiter.inFlight).toBe(0); + + // The next acquire still respects the cap. + const r2 = await limiter.acquire(); + expect(limiter.inFlight).toBe(1); + r2(); + }); + + it('coerces NaN / Infinity capacities to "unlimited"', async () => { + const nanLimiter = new ConcurrencyLimiter(Number.NaN); + expect(nanLimiter.limit).toBe(0); + const infLimiter = new ConcurrencyLimiter(Number.POSITIVE_INFINITY); + expect(infLimiter.limit).toBe(0); + // Both should be effectively no-op. + await nanLimiter.runExclusive(async () => undefined); + await infLimiter.runExclusive(async () => undefined); + }); +}); diff --git a/packages/core/src/utils/concurrencyLimiter.ts b/packages/core/src/utils/concurrencyLimiter.ts new file mode 100644 index 00000000000..96f5a93792d --- /dev/null +++ b/packages/core/src/utils/concurrencyLimiter.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Simple FIFO counting semaphore. + * + * Used to gate concurrent calls to the model API so a configured per-provider + * cap (#3409) translates into back-pressure rather than the upstream + * ``429 Too many concurrent requests`` error -- the model server has no way + * to tell us "wait", so the queue has to live in the client. + * + * - ``capacity <= 0`` is treated as "unlimited" and the limiter becomes a + * pass-through, so callers can wire this in unconditionally. + * - ``acquire`` resolves with a release function. Always invoke release + * exactly once, including on the error path -- ``runExclusive`` does this + * for you. + * - The wait list is FIFO so the order callers acquired the lock matches the + * order they were enqueued; useful for predictable pipeline behavior. + */ +export class ConcurrencyLimiter { + private active = 0; + private readonly waiters: Array<() => void> = []; + private readonly capacity: number; + + constructor(capacity: number) { + this.capacity = Number.isFinite(capacity) && capacity > 0 ? capacity : 0; + } + + /** Number of in-flight slots. */ + get inFlight(): number { + return this.active; + } + + /** Number of callers currently parked waiting for a slot. */ + get queued(): number { + return this.waiters.length; + } + + /** Effective configured limit (``0`` means "unlimited"). */ + get limit(): number { + return this.capacity; + } + + /** + * Acquire a slot. Resolves with a single-use ``release`` function. When + * the limiter is unlimited (``capacity <= 0``) this resolves immediately + * with a no-op release. + */ + async acquire(): Promise<() => void> { + if (this.capacity === 0) { + return () => {}; + } + if (this.active < this.capacity) { + this.active++; + return this.makeRelease(); + } + return new Promise<() => void>((resolve) => { + this.waiters.push(() => { + this.active++; + resolve(this.makeRelease()); + }); + }); + } + + /** + * Run ``fn`` while holding a slot. Always releases the slot, even when + * ``fn`` throws or rejects. Most callers want this rather than the raw + * ``acquire``/release pair. + */ + async runExclusive(fn: () => Promise): Promise { + const release = await this.acquire(); + try { + return await fn(); + } finally { + release(); + } + } + + private makeRelease(): () => void { + let released = false; + return () => { + if (released) { + return; + } + released = true; + this.active--; + const next = this.waiters.shift(); + if (next) { + next(); + } + }; + } +} diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index 13e2306726d..a8ec5862cb6 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -353,6 +353,10 @@ "description": "Maximum number of retries for failed requests.", "type": "number" }, + "requestConcurrency": { + "description": "Maximum number of concurrent in-flight requests against this model. Useful when the upstream returns 429 \"Too many concurrent requests for this model\". 0 / unset means unlimited. Falls back to QWEN_REQUEST_CONCURRENCY env var.", + "type": "number" + }, "enableCacheControl": { "description": "Enable cache control for DashScope providers.", "type": "boolean",