diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 2eb0679d4..6fb701e27 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added a per-provider in-flight request limiter for LLM streams, shared across local OMP processes and configurable by callers with `maxInFlightRequests`. + ## [16.1.11] - 2026-06-21 ### Fixed diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index 3012b51e4..82d22dd29 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -1,3 +1,7 @@ +import * as crypto from "node:crypto"; +import * as fsSync from "node:fs"; +import * as fs from "node:fs/promises"; +import * as path from "node:path"; import type { Effort } from "@oh-my-pi/pi-catalog/effort"; import { isVertexExpressOpenAIUrl, isVertexRawPredictUrl } from "@oh-my-pi/pi-catalog/hosts"; import { @@ -8,7 +12,7 @@ import { resolveWireModelId, } from "@oh-my-pi/pi-catalog/model-thinking"; import { CATALOG_PROVIDERS, type ProviderCatalogEntry } from "@oh-my-pi/pi-catalog/provider-models"; -import { $env, $pickenv, extractHttpStatusFromError } from "@oh-my-pi/pi-utils"; +import { $env, $pickenv, extractHttpStatusFromError, getConfigRootDir, isEnoent, logger } from "@oh-my-pi/pi-utils"; import { getCustomApi } from "./api-registry"; import { AUTH_RETRY_STEPS, isApiKeyResolver, resolveRetryKey } from "./auth-retry"; import { ProviderHttpError } from "./errors"; @@ -73,6 +77,421 @@ function isGoogleVertexAuthenticatedModel(model: Model): boolean { ); } +type ProviderInFlightLease = { + path: string; + heartbeat: NodeJS.Timeout; + flushHeartbeat: () => Promise; +}; + +type ProviderInFlightLeaseInfo = { + pid: number; + timestamp: number; + token: string; +}; +type ProviderInFlightStaleLock = { token: string } | { mtimeMs: number }; + +const PROVIDER_INFLIGHT_LOCK_STALE_MS = 10_000; +const PROVIDER_INFLIGHT_LEASE_STALE_MS = 30_000; +const PROVIDER_INFLIGHT_HEARTBEAT_MS = 5_000; +const PROVIDER_INFLIGHT_SIGNAL_FALLBACK_MS = 250; + +let configuredProviderMaxInFlightRequests: Record = {}; +let providerInFlightRootOverride: string | undefined; + +export function configureProviderMaxInFlightRequests(limits: Record | undefined): void { + configuredProviderMaxInFlightRequests = limits ?? {}; +} + +function resolveProviderInFlightLimit( + provider: string, + options?: Pick, +): number | undefined { + const limits = options?.maxInFlightRequests ?? configuredProviderMaxInFlightRequests; + const value = limits[provider]; + if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) return undefined; + return Math.max(1, Math.floor(value)); +} + +function providerInFlightRoot(): string { + if (providerInFlightRootOverride) return providerInFlightRootOverride; + return path.join(getConfigRootDir(), "run", "provider-inflight"); +} + +function providerInFlightSegment(provider: string): string { + return crypto.createHash("sha256").update(provider).digest("base64url"); +} + +function providerInFlightDir(provider: string): string { + return path.join(providerInFlightRoot(), providerInFlightSegment(provider)); +} + +function providerInFlightSignalPath(provider: string): string { + return path.join(providerInFlightDir(provider), ".wakeup"); +} + +function providerInFlightLockDir(provider: string): string { + return `${providerInFlightDir(provider)}.lock`; +} + +// `process.kill(pid, 0)` may throw for permission/sandbox reasons even when a +// process exists. Treat non-ESRCH failures as alive; timestamp expiry still +// reaps leases whose heartbeat stopped. +function isProcessAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + return (error as NodeJS.ErrnoException).code !== "ESRCH"; + } +} + +async function readProviderInFlightInfo(infoPath: string): Promise { + try { + const content = await fs.readFile(infoPath, "utf-8"); + const parsed = JSON.parse(content) as Partial; + if (typeof parsed.pid !== "number" || typeof parsed.timestamp !== "number" || typeof parsed.token !== "string") { + return null; + } + return { pid: parsed.pid, timestamp: parsed.timestamp, token: parsed.token }; + } catch { + return null; + } +} + +async function writeProviderInFlightInfo(dir: string, token: string): Promise { + const info: ProviderInFlightLeaseInfo = { pid: process.pid, timestamp: Date.now(), token }; + const infoPath = path.join(dir, "info.json"); + const tempPath = path.join(dir, `.info-${process.pid}-${crypto.randomUUID()}.tmp`); + try { + await Bun.write(tempPath, JSON.stringify(info)); + await fs.rename(tempPath, infoPath); + } catch (error) { + await fs.rm(tempPath, { force: true }).catch(() => {}); + throw error; + } +} + +async function isProviderInFlightDirStale(dir: string, staleMs: number): Promise { + const info = await readProviderInFlightInfo(path.join(dir, "info.json")); + if (info) { + if (!isProcessAlive(info.pid)) return true; + return Date.now() - info.timestamp > staleMs; + } + + try { + const stat = await fs.stat(path.join(dir, "info.json")); + return Date.now() - stat.mtimeMs > staleMs; + } catch (error) { + if (!isEnoent(error)) throw error; + } + + try { + const stat = await fs.stat(dir); + return Date.now() - stat.mtimeMs > staleMs; + } catch (error) { + if (isEnoent(error)) return false; + throw error; + } +} + +async function readProviderInFlightStaleLock(lockDir: string): Promise { + const infoPath = path.join(lockDir, "info.json"); + const info = await readProviderInFlightInfo(infoPath); + if (info) return isProcessAlive(info.pid) ? null : { token: info.token }; + + try { + const stat = await fs.stat(lockDir); + return Date.now() - stat.mtimeMs > PROVIDER_INFLIGHT_LOCK_STALE_MS ? { mtimeMs: stat.mtimeMs } : null; + } catch (error) { + if (isEnoent(error)) return null; + throw error; + } +} + +async function releaseProviderInFlightStaleLock(lockDir: string, stale: ProviderInFlightStaleLock): Promise { + if ("token" in stale) { + await releaseProviderInFlightLock(lockDir, stale.token); + return; + } + + const infoPath = path.join(lockDir, "info.json"); + if (await readProviderInFlightInfo(infoPath)) return; + try { + const stat = await fs.stat(lockDir); + if (stat.mtimeMs !== stale.mtimeMs || Date.now() - stat.mtimeMs <= PROVIDER_INFLIGHT_LOCK_STALE_MS) return; + await fs.rm(lockDir, { recursive: true, force: true }); + } catch {} +} + +// Best-effort token-checked release. Untokened calls are used only after stale +// detection concludes the owner is dead or its heartbeat expired. +async function releaseProviderInFlightLock(lockDir: string, token?: string): Promise { + try { + if (token !== undefined) { + const info = await readProviderInFlightInfo(path.join(lockDir, "info.json")); + if (!info || info.token !== token) return; + } + await fs.rm(lockDir, { recursive: true, force: true }); + } catch {} +} + +async function acquireProviderInFlightLock(provider: string, signal?: AbortSignal): Promise<() => Promise> { + const lockDir = providerInFlightLockDir(provider); + await fs.mkdir(path.dirname(lockDir), { recursive: true }); + + while (true) { + if (signal?.aborted) throw signal.reason ?? new Error("Provider request aborted before dispatch"); + try { + await fs.mkdir(lockDir); + const token = crypto.randomUUID(); + try { + await writeProviderInFlightInfo(lockDir, token); + } catch (error) { + await releaseProviderInFlightLock(lockDir); + throw error; + } + return async () => { + await releaseProviderInFlightLock(lockDir, token); + }; + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; + } + + const staleLock = await readProviderInFlightStaleLock(lockDir); + if (staleLock) { + await releaseProviderInFlightStaleLock(lockDir, staleLock); + await signalProviderInFlightWaiters(provider); + continue; + } + + await waitForProviderInFlightSignal(provider, signal); + } +} + +async function cleanupProviderInFlightLeases(providerDir: string): Promise { + let active = 0; + let entries: string[]; + try { + entries = await fs.readdir(providerDir); + } catch (error) { + if (isEnoent(error)) return 0; + throw error; + } + + for (const entry of entries) { + const leaseDir = path.join(providerDir, entry); + let isDirectory = false; + try { + isDirectory = (await fs.stat(leaseDir)).isDirectory(); + } catch (error) { + if (isEnoent(error)) continue; + throw error; + } + if (!isDirectory) continue; + if (await isProviderInFlightDirStale(leaseDir, PROVIDER_INFLIGHT_LEASE_STALE_MS)) { + await fs.rm(leaseDir, { recursive: true, force: true }); + continue; + } + active++; + } + return active; +} + +async function tryAcquireProviderInFlightLease( + provider: string, + limit: number, + signal?: AbortSignal, +): Promise { + const releaseLock = await acquireProviderInFlightLock(provider, signal); + let leaseCreated = false; + try { + const dir = providerInFlightDir(provider); + await fs.mkdir(dir, { recursive: true }); + const active = await cleanupProviderInFlightLeases(dir); + if (active >= limit) return null; + + const leaseDir = path.join(dir, `${process.pid}-${Date.now()}-${crypto.randomUUID()}`); + const token = crypto.randomUUID(); + try { + await fs.mkdir(leaseDir); + await writeProviderInFlightInfo(leaseDir, token); + leaseCreated = true; + } catch (error) { + await removeProviderInFlightLeaseDir(leaseDir).catch(() => {}); + throw error; + } + let heartbeatFlush = Promise.resolve(); + const touchHeartbeat = () => { + heartbeatFlush = heartbeatFlush + .then( + () => writeProviderInFlightInfo(leaseDir, token), + () => writeProviderInFlightInfo(leaseDir, token), + ) + .catch(() => {}); + }; + const heartbeat = setInterval(touchHeartbeat, PROVIDER_INFLIGHT_HEARTBEAT_MS); + heartbeat.unref?.(); + return { path: leaseDir, heartbeat, flushHeartbeat: () => heartbeatFlush }; + } finally { + await releaseLock(); + if (leaseCreated) await signalProviderInFlightWaiters(provider); + } +} + +async function signalProviderInFlightWaiters(provider: string): Promise { + try { + const dir = providerInFlightDir(provider); + await fs.mkdir(dir, { recursive: true }); + await Bun.write(providerInFlightSignalPath(provider), String(Date.now())); + } catch {} +} + +function waitForProviderInFlightSignal(provider: string, signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(signal.reason ?? new Error("Provider request aborted before dispatch")); + const signalPath = providerInFlightSignalPath(provider); + const waitStarted = Date.now(); + const { promise, resolve, reject } = Promise.withResolvers(); + let settled = false; + let watcher: fsSync.FSWatcher | undefined; + const timer = setTimeout(() => finish(resolve), PROVIDER_INFLIGHT_SIGNAL_FALLBACK_MS); + const finish = (settle: () => void) => { + if (settled) return; + settled = true; + clearTimeout(timer); + watcher?.close(); + signal?.removeEventListener("abort", onAbort); + settle(); + }; + const onAbort = () => { + finish(() => reject(signal?.reason ?? new Error("Provider request aborted before dispatch"))); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + try { + watcher = fsSync.watch(providerInFlightDir(provider), (_event, filename) => { + if (filename === ".wakeup" || filename === null) { + finish(resolve); + } + }); + void fs.stat(signalPath).then( + stat => { + if (stat.mtimeMs >= waitStarted) finish(resolve); + }, + error => { + if (!isEnoent(error)) finish(resolve); + }, + ); + } catch { + // Filesystem notifications are best-effort across platforms; the fallback + // timer keeps stale-lock/lease cleanup progressing if an event is dropped. + } + return promise; +} + +async function removeProviderInFlightLeaseDir(leasePath: string): Promise { + for (let attempt = 0; attempt < 3; attempt++) { + try { + await fs.rm(leasePath, { recursive: true, force: true }); + return; + } catch (error) { + if (isEnoent(error)) return; + const code = (error as NodeJS.ErrnoException).code; + if (attempt < 2 && (code === "EBUSY" || code === "ENOTEMPTY" || code === "EPERM")) { + await Bun.sleep(25); + continue; + } + throw error; + } + } +} + +async function releaseProviderInFlightLease(provider: string, lease: ProviderInFlightLease): Promise { + clearInterval(lease.heartbeat); + await lease.flushHeartbeat(); + await removeProviderInFlightLeaseDir(lease.path); + await signalProviderInFlightWaiters(provider); +} + +async function acquireProviderInFlightSlot( + provider: string, + limit: number | undefined, + signal?: AbortSignal, +): Promise<() => Promise> { + if (limit === undefined) return async () => {}; + let loggedWait = false; + while (true) { + if (signal?.aborted) throw signal.reason ?? new Error("Provider request aborted before dispatch"); + const lease = await tryAcquireProviderInFlightLease(provider, limit, signal); + if (lease) return () => releaseProviderInFlightLease(provider, lease); + if (!loggedWait) { + loggedWait = true; + logger.debug("Provider in-flight limit blocked request", { provider, limit }); + } + await waitForProviderInFlightSignal(provider, signal); + } +} + +export const __providerInFlightForTesting = { + setRoot(root: string | undefined): void { + providerInFlightRootOverride = root; + }, + providerDir(provider: string): string { + return providerInFlightDir(provider); + }, + lockDir(provider: string): string { + return providerInFlightLockDir(provider); + }, + async captureStaleLockRelease(provider: string): Promise<(() => Promise) | null> { + const lockDir = providerInFlightLockDir(provider); + const stale = await readProviderInFlightStaleLock(lockDir); + if (!stale) return null; + return () => releaseProviderInFlightStaleLock(lockDir, stale); + }, +}; + +function withProviderInFlightLimit>( + model: Model, + options: TOptions | undefined, + dispatch: () => AssistantMessageEventStream, +): AssistantMessageEventStream { + const limit = resolveProviderInFlightLimit(model.provider, options); + if (limit === undefined) return dispatch(); + + const outer = new AssistantMessageEventStream(); + void (async () => { + let release: (() => Promise) | undefined; + let released = false; + const releaseOnce = async () => { + if (!release || released) return; + released = true; + await release(); + }; + try { + const startedWaitingAt = Date.now(); + release = await acquireProviderInFlightSlot(model.provider, limit, options?.signal); + if (Date.now() - startedWaitingAt >= PROVIDER_INFLIGHT_SIGNAL_FALLBACK_MS) { + logger.debug("Provider in-flight limit wait completed", { provider: model.provider, limit }); + } + if (options?.signal?.aborted) { + throw options.signal.reason ?? new Error("Provider request aborted before dispatch"); + } + const inner = dispatch(); + try { + for await (const event of inner) { + outer.push(event); + if (outer.done) return; + } + if (!outer.done) outer.end(await inner.result()); + } finally { + await releaseOnce(); + } + } catch (error) { + await releaseOnce(); + if (!outer.done) outer.fail(error); + } + })(); + return outer; +} + function createVertexAuthenticatedFetch(options: StreamOptions | undefined): FetchImpl { const baseFetch = options?.fetch ?? fetch; const vertexFetch = async (input: string | URL | Request, init?: RequestInit): Promise => { @@ -227,7 +646,9 @@ export function stream( context: Context, options?: OptionsForApi, ): AssistantMessageEventStream { - return withGeminiThinkingLoopGuard(model, options, opts => streamDispatch(model, context, opts)); + return withGeminiThinkingLoopGuard(model, options, opts => + withProviderInFlightLimit(model, opts, () => streamDispatch(model, context, opts)), + ); } function streamDispatch( @@ -500,14 +921,16 @@ export function streamSimple( // extension-registered APIs can't accidentally override a configured // pi-native transport. if (model.transport === "pi-native") { - return withGeminiThinkingLoopGuard(model, requestOptions, opts => streamPiNative(model, context, opts)); + return withGeminiThinkingLoopGuard(model, requestOptions, opts => + withProviderInFlightLimit(model, opts, () => streamPiNative(model, context, opts)), + ); } // Check custom API registry (extension-provided APIs) const customApiProvider = getCustomApi(model.api); if (customApiProvider) { return withGeminiThinkingLoopGuard(model, requestOptions, opts => - customApiProvider.streamSimple(model, context, opts), + withProviderInFlightLimit(model, opts, () => customApiProvider.streamSimple(model, context, opts)), ); } @@ -531,30 +954,36 @@ export function streamSimple( // GitLab Duo - wraps Anthropic/OpenAI behind GitLab AI Gateway direct access tokens if (isGitLabDuoModel(model)) { - return streamGitLabDuo(model, context, { - ...requestOptions, - apiKey, - }); + return withProviderInFlightLimit(model, requestOptions, () => + streamGitLabDuo(model, context, { + ...requestOptions, + apiKey, + }), + ); } // Kimi Code - route to dedicated handler that wraps OpenAI or Anthropic API if (isKimiModel(model)) { // Pass raw SimpleStreamOptions - streamKimi handles mapping internally - return streamKimi(model as Model<"openai-completions">, context, { - ...requestOptions, - apiKey, - format: requestOptions?.kimiApiFormat ?? "anthropic", - }); + return withProviderInFlightLimit(model, requestOptions, () => + streamKimi(model as Model<"openai-completions">, context, { + ...requestOptions, + apiKey, + format: requestOptions?.kimiApiFormat ?? "anthropic", + }), + ); } // Synthetic - route to dedicated handler that wraps OpenAI or Anthropic API if (isSyntheticModel(model)) { // Pass raw SimpleStreamOptions - streamSynthetic handles mapping internally - return streamSynthetic(model as Model<"openai-completions">, context, { - ...requestOptions, - apiKey, - format: requestOptions?.syntheticApiFormat ?? "openai", // Default to OpenAI format - }); + return withProviderInFlightLimit(model, requestOptions, () => + streamSynthetic(model as Model<"openai-completions">, context, { + ...requestOptions, + apiKey, + format: requestOptions?.syntheticApiFormat ?? "openai", // Default to OpenAI format + }), + ); } const providerOptions = mapOptionsForApi(model, requestOptions, apiKey); return stream(model, context, providerOptions); @@ -774,6 +1203,7 @@ function mapOptionsForApi( streamFirstEventTimeoutMs: options?.streamFirstEventTimeoutMs, streamIdleTimeoutMs: options?.streamIdleTimeoutMs, providerSessionState: options?.providerSessionState, + maxInFlightRequests: options?.maxInFlightRequests, onPayload: options?.onPayload, onResponse: options?.onResponse, onSseEvent: options?.onSseEvent, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 80d1a48a3..2efbe4e3f 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -272,6 +272,14 @@ export interface StreamOptions { * Providers can use this to persist transport/session state between turns. */ providerSessionState?: Map; + /** + * Optional per-provider concurrent request cap for LLM stream calls. Keys are + * provider ids (`model.provider`); positive numeric values cap in-flight + * requests across local OMP processes that share the same config root. Omitted + * providers are unlimited. Non-chat provider APIs that bypass stream helpers + * are not covered. + */ + maxInFlightRequests?: Record; /** * Optional callback for inspecting or replacing provider payloads before sending. * Return undefined to keep the payload unchanged. diff --git a/packages/ai/test/provider-inflight.test.ts b/packages/ai/test/provider-inflight.test.ts new file mode 100644 index 000000000..23eb667b4 --- /dev/null +++ b/packages/ai/test/provider-inflight.test.ts @@ -0,0 +1,270 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { clearCustomApis } from "@oh-my-pi/pi-ai/api-registry"; +import { createMockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock"; +import { + __providerInFlightForTesting, + configureProviderMaxInFlightRequests, + streamSimple, +} from "@oh-my-pi/pi-ai/stream"; +import type { Context } from "@oh-my-pi/pi-ai/types"; + +function context(): Context { + return { + systemPrompt: [], + messages: [{ role: "user", content: "hi", timestamp: 0 }], + }; +} + +let limiterRoot: string | undefined; + +afterEach(async () => { + clearCustomApis(); + configureProviderMaxInFlightRequests(undefined); + __providerInFlightForTesting.setRoot(undefined); + if (limiterRoot !== undefined) { + await fs.rm(limiterRoot, { recursive: true, force: true }); + limiterRoot = undefined; + } +}); + +async function useIsolatedLimiterRoot(): Promise { + limiterRoot = await fs.mkdtemp(path.join(os.tmpdir(), "omp-provider-inflight-test-")); + __providerInFlightForTesting.setRoot(limiterRoot); +} + +function limiterDir(provider: string): string { + return __providerInFlightForTesting.providerDir(provider); +} + +describe("provider in-flight request limits", () => { + beforeEach(async () => { + await useIsolatedLimiterRoot(); + }); + test("serializes concurrent streamSimple calls for the same provider", async () => { + registerMockApi(); + const firstStarted = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + let active = 0; + let maxActive = 0; + let callIndex = 0; + const mock = createMockModel({ + provider: "tests", + handler: async () => { + callIndex++; + active++; + maxActive = Math.max(maxActive, active); + try { + if (callIndex === 1) { + firstStarted.resolve(); + await releaseFirst.promise; + } + return { content: [`reply ${callIndex}`] }; + } finally { + active--; + } + }, + }); + + const first = streamSimple(mock.model, context(), { maxInFlightRequests: { tests: 1 } }); + const firstResult = first.result(); + await firstStarted.promise; + + const second = streamSimple(mock.model, context(), { maxInFlightRequests: { tests: 1 } }); + await Bun.sleep(20); + expect(mock.calls).toHaveLength(1); + + releaseFirst.resolve(); + const [firstMessage, secondMessage] = await Promise.all([firstResult, second.result()]); + + expect(firstMessage.content).toEqual([{ type: "text", text: "reply 1" }]); + expect(secondMessage.content).toEqual([{ type: "text", text: "reply 2" }]); + expect(maxActive).toBe(1); + expect(mock.calls).toHaveLength(2); + }); + + test("removes an aborted queued request without dispatching it", async () => { + registerMockApi(); + const firstStarted = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + let callIndex = 0; + const mock = createMockModel({ + provider: "tests", + handler: async () => { + callIndex++; + if (callIndex === 1) { + firstStarted.resolve(); + await releaseFirst.promise; + } + return { content: [`reply ${callIndex}`] }; + }, + }); + + const first = streamSimple(mock.model, context(), { maxInFlightRequests: { tests: 1 } }); + const firstResult = first.result(); + await firstStarted.promise; + + const controller = new AbortController(); + const second = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + controller.abort(new Error("cancel queued request")); + + await expect(second.result()).rejects.toThrow("cancel queued request"); + expect(mock.calls).toHaveLength(1); + + releaseFirst.resolve(); + await firstResult; + expect(mock.calls).toHaveLength(1); + }); + + test("shares limits with leases created by another process", async () => { + registerMockApi(); + const providerDir = limiterDir("tests"); + const externalLease = path.join(providerDir, "external"); + await fs.mkdir(externalLease, { recursive: true }); + await Bun.write( + path.join(externalLease, "info.json"), + JSON.stringify({ pid: process.pid, timestamp: Date.now(), token: "external" }), + ); + + const controller = new AbortController(); + const mock = createMockModel({ provider: "tests", responses: [{ content: ["reply"] }] }); + const stream = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + + await Bun.sleep(150); + expect(mock.calls).toHaveLength(0); + + await fs.rm(externalLease, { recursive: true, force: true }); + await Bun.write(path.join(providerDir, ".wakeup"), String(Date.now())); + const result = await stream.result(); + expect(result.content).toEqual([{ type: "text", text: "reply" }]); + expect(mock.calls).toHaveLength(1); + }); + + test("does not signal waiters when no slot was freed", async () => { + registerMockApi(); + const providerDir = limiterDir("tests"); + const externalLease = path.join(providerDir, "external"); + await fs.mkdir(externalLease, { recursive: true }); + await Bun.write( + path.join(externalLease, "info.json"), + JSON.stringify({ pid: process.pid, timestamp: Date.now(), token: "external" }), + ); + + const controller = new AbortController(); + const mock = createMockModel({ provider: "tests", responses: [{ content: ["reply"] }] }); + const stream = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + + await Bun.sleep(50); + expect(await Bun.file(path.join(providerDir, ".wakeup")).exists()).toBe(false); + expect(mock.calls).toHaveLength(0); + + controller.abort(new Error("cancel saturated waiter")); + await expect(stream.result()).rejects.toThrow("cancel saturated waiter"); + }); + + test("does not reap a live lock just because its timestamp is old", async () => { + registerMockApi(); + const lockDir = __providerInFlightForTesting.lockDir("tests"); + await fs.mkdir(lockDir, { recursive: true }); + await Bun.write( + path.join(lockDir, "info.json"), + JSON.stringify({ pid: process.pid, timestamp: Date.now() - 60_000, token: "live-lock" }), + ); + + const controller = new AbortController(); + const mock = createMockModel({ provider: "tests", responses: [{ content: ["reply"] }] }); + const stream = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + + await Bun.sleep(150); + expect(mock.calls).toHaveLength(0); + + controller.abort(new Error("cancel lock waiter")); + await expect(stream.result()).rejects.toThrow("cancel lock waiter"); + expect(mock.calls).toHaveLength(0); + }); + + test("treats unreadable fresh lease info as active", async () => { + registerMockApi(); + const providerDir = limiterDir("tests"); + const externalLease = path.join(providerDir, "partial-info"); + await fs.mkdir(externalLease, { recursive: true }); + const old = new Date(Date.now() - 60_000); + await fs.utimes(externalLease, old, old); + await Bun.write(path.join(externalLease, "info.json"), "{"); + + const controller = new AbortController(); + const mock = createMockModel({ provider: "tests", responses: [{ content: ["reply"] }] }); + const stream = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + + await Bun.sleep(150); + expect(mock.calls).toHaveLength(0); + + controller.abort(new Error("cancel partial-info waiter")); + await expect(stream.result()).rejects.toThrow("cancel partial-info waiter"); + expect(mock.calls).toHaveLength(0); + }); + + test("does not delete a fresh lock after observing a stale lock", async () => { + const lockDir = __providerInFlightForTesting.lockDir("tests"); + await fs.mkdir(lockDir, { recursive: true }); + await Bun.write( + path.join(lockDir, "info.json"), + JSON.stringify({ pid: 999999, timestamp: Date.now() - 60_000, token: "stale-lock" }), + ); + const staleRelease = await __providerInFlightForTesting.captureStaleLockRelease("tests"); + expect(staleRelease).not.toBeNull(); + + await fs.rm(lockDir, { recursive: true, force: true }); + await fs.mkdir(lockDir, { recursive: true }); + await Bun.write( + path.join(lockDir, "info.json"), + JSON.stringify({ pid: process.pid, timestamp: Date.now(), token: "fresh-lock" }), + ); + + await staleRelease?.(); + + const remaining = JSON.parse(await Bun.file(path.join(lockDir, "info.json")).text()) as { token: string }; + expect(remaining.token).toBe("fresh-lock"); + }); + + test("does not dispatch when aborted immediately after slot acquisition", async () => { + registerMockApi(); + const controller = new AbortController(); + const mock = createMockModel({ provider: "tests", responses: [{ content: ["reply"] }] }); + const stream = streamSimple(mock.model, context(), { + maxInFlightRequests: { tests: 1 }, + signal: controller.signal, + }); + + controller.abort(new Error("cancel acquired request")); + + await expect(stream.result()).rejects.toThrow("cancel acquired request"); + expect(mock.calls).toHaveLength(0); + }); + + test("uses opaque path segments for provider ids", async () => { + const dir = limiterDir(".."); + const relative = path.relative(limiterRoot!, dir); + + expect(relative).not.toBe(""); + expect(relative.startsWith("..")).toBe(false); + expect(path.isAbsolute(relative)).toBe(false); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 1fc767cea..d2b5a61a8 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `providers.maxInFlightRequests` to cap concurrent LLM requests per provider across local OMP processes from settings. + ## [16.1.11] - 2026-06-21 ### Added diff --git a/packages/coding-agent/src/cli/config-cli.ts b/packages/coding-agent/src/cli/config-cli.ts index c2fd6025e..b9ecf3ca9 100644 --- a/packages/coding-agent/src/cli/config-cli.ts +++ b/packages/coding-agent/src/cli/config-cli.ts @@ -16,6 +16,7 @@ import { Settings, type SettingValue, settings, + validateProviderMaxInFlightRequests, } from "../config/settings"; import { SETTINGS_SCHEMA } from "../config/settings-schema"; import { theme } from "../modes/theme/theme"; @@ -218,6 +219,9 @@ function parseAndSetValue(path: SettingPath, rawValue: string): void { if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) { throw new Error(`Invalid record JSON: ${rawValue}`); } + if (path === "providers.maxInFlightRequests") { + parsed = validateProviderMaxInFlightRequests(parsed); + } parsedValue = parsed; break; } diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 86103e0af..72bb9398d 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -272,6 +272,7 @@ export interface ModelTagsSettings { // under `as const` while still letting SettingValue infer the correct element type. const EMPTY_STRING_ARRAY: string[] = []; const EMPTY_STRING_RECORD: Record = {}; +const EMPTY_NUMBER_RECORD: Record = {}; const DEFAULT_CYCLE_ORDER: string[] = ["smol", "default", "slow"]; const EMPTY_MODEL_TAGS_RECORD: ModelTagsSettings = {}; const HINDSIGHT_RECALL_TYPES_DEFAULT: string[] = ["world", "experience"]; @@ -449,6 +450,18 @@ export const SETTINGS_SCHEMA = { disabledProviders: { type: "array", default: EMPTY_STRING_ARRAY }, + "providers.maxInFlightRequests": { + type: "record", + default: EMPTY_NUMBER_RECORD, + ui: { + tab: "providers", + group: "Services", + label: "Max In-Flight Requests", + description: + 'Maximum concurrent LLM requests per provider id (for example "openai" or "anthropic"), shared across local OMP processes with this config root. Omitted providers are unlimited.', + }, + }, + disabledExtensions: { type: "array", default: EMPTY_STRING_ARRAY }, modelRoles: { type: "record", default: EMPTY_STRING_RECORD }, diff --git a/packages/coding-agent/src/config/settings.ts b/packages/coding-agent/src/config/settings.ts index 58f0398f8..a0972f5e6 100644 --- a/packages/coding-agent/src/config/settings.ts +++ b/packages/coding-agent/src/config/settings.ts @@ -14,6 +14,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; +import { configureProviderMaxInFlightRequests } from "@oh-my-pi/pi-ai/stream"; import { getAgentDbPath, getAgentDir, @@ -105,6 +106,33 @@ function setByPath(obj: RawSettings, segments: string[], value: unknown): void { current[segments[segments.length - 1]] = value; } +export function normalizeProviderMaxInFlightRequests(value: unknown): Record { + if (!value || typeof value !== "object" || Array.isArray(value)) return {}; + const normalized: Record = {}; + for (const [provider, rawLimit] of Object.entries(value)) { + if (typeof rawLimit !== "number" || !Number.isFinite(rawLimit) || rawLimit <= 0) continue; + normalized[provider] = Math.max(1, Math.floor(rawLimit)); + } + return normalized; +} + +export function validateProviderMaxInFlightRequests(value: unknown): Record { + if (!value || typeof value !== "object" || Array.isArray(value)) return {}; + const invalidProviders: string[] = []; + const normalized: Record = {}; + for (const [provider, rawLimit] of Object.entries(value)) { + if (typeof rawLimit !== "number" || !Number.isFinite(rawLimit) || rawLimit <= 0) { + invalidProviders.push(provider); + continue; + } + normalized[provider] = Math.max(1, Math.floor(rawLimit)); + } + if (invalidProviders.length > 0) { + throw new Error(`Provider request limits must be positive numbers: ${invalidProviders.join(", ")}`); + } + return normalized; +} + const PATH_SCOPED_ARRAY_SETTINGS = new Set(["enabledModels", "disabledProviders"]); type PathScopedStringArrayEntry = { path?: unknown; @@ -336,7 +364,7 @@ export class Settings { // Trigger hook if exists const hook = SETTING_HOOKS[path]; if (hook) { - hook(value, prev); + hook(next, prev); } this.#fireEffectiveSettingChanged(path, next, prev); } @@ -1147,6 +1175,9 @@ const SETTING_HOOKS: Partial>> = { appendOnlyModeSignal.fire(value); } }, + "providers.maxInFlightRequests": value => { + configureProviderMaxInFlightRequests(validateProviderMaxInFlightRequests(value)); + }, "hindsight.bankId": () => hindsightScopeSignal.fire(), "hindsight.bankIdPrefix": () => hindsightScopeSignal.fire(), "hindsight.scoping": () => hindsightScopeSignal.fire(), @@ -1211,6 +1242,7 @@ export function resetSettingsForTest(): void { globalInstance = null; globalInstancePromise = null; clearBoundSettingsMethods(); + configureProviderMaxInFlightRequests(undefined); } /** diff --git a/packages/coding-agent/src/modes/components/settings-defs.ts b/packages/coding-agent/src/modes/components/settings-defs.ts index 3c9d54577..582877281 100644 --- a/packages/coding-agent/src/modes/components/settings-defs.ts +++ b/packages/coding-agent/src/modes/components/settings-defs.ts @@ -68,7 +68,16 @@ export interface TextInputSettingDef extends BaseSettingDef { type: "text"; } -export type SettingDef = BooleanSettingDef | EnumSettingDef | SubmenuSettingDef | TextInputSettingDef; +export interface ProviderLimitsSettingDef extends BaseSettingDef { + type: "providerLimits"; +} + +export type SettingDef = + | BooleanSettingDef + | EnumSettingDef + | SubmenuSettingDef + | TextInputSettingDef + | ProviderLimitsSettingDef; // ═══════════════════════════════════════════════════════════════════════════ // Condition Functions @@ -170,6 +179,10 @@ function pathToSettingDef(path: SettingPath): SettingDef | null { return { ...base, type: "text" }; } + if (schemaType === "record") { + return path === "providers.maxInFlightRequests" ? { ...base, type: "providerLimits" } : null; + } + return null; } diff --git a/packages/coding-agent/src/modes/components/settings-selector.ts b/packages/coding-agent/src/modes/components/settings-selector.ts index dbc60a64e..f9ea39a63 100644 --- a/packages/coding-agent/src/modes/components/settings-selector.ts +++ b/packages/coding-agent/src/modes/components/settings-selector.ts @@ -11,6 +11,7 @@ import { Input, matchesKey, parseSgrMouse, + replaceTabs, type SelectItem, SelectList, type SettingItem, @@ -24,7 +25,14 @@ import { visibleWidth, } from "@oh-my-pi/pi-tui"; import type { ShapeTarget } from "@oh-my-pi/snapcompact"; -import { getDefault, type SettingPath, settings } from "../../config/settings"; +import { + getDefault, + getType, + normalizeProviderMaxInFlightRequests, + type SettingPath, + settings, + validateProviderMaxInFlightRequests, +} from "../../config/settings"; import type { SettingTab, StatusLinePreset, @@ -50,6 +58,7 @@ import { getPreset } from "./status-line/presets"; */ class TextInputSubmenu extends Container { #input: Input; + #error: Text; constructor( label: string, @@ -71,11 +80,18 @@ class TextInputSubmenu extends Container { if (currentValue) { this.#input.setValue(currentValue); } + this.#error = new Text("", 0, 0); this.#input.onSubmit = value => { - this.onSubmit(value); // empty string clears the setting + try { + this.onSubmit(value); // empty string clears the setting + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + this.#error.setText(theme.fg("error", truncateToWidth(replaceTabs(message).replace(/[\r\n]+/g, " "), 100))); + } }; this.addChild(this.#input); this.addChild(new Spacer(1)); + this.addChild(this.#error); this.addChild(new Text(theme.fg("dim", " Enter to save · Esc to cancel · Clear field to unset"), 0, 0)); } @@ -217,6 +233,113 @@ class SelectSubmenu extends Container { } } +class ProviderLimitsSubmenu extends Container { + #selectList: SelectList | undefined; + + constructor( + private readonly providers: readonly string[], + private readonly onChange: (value: Record) => void, + private readonly onCancel: () => void, + private readonly requestRender?: () => void, + ) { + super(); + this.#showProviderList(); + } + + #providerIds(): string[] { + const limits = normalizeProviderMaxInFlightRequests(settings.get("providers.maxInFlightRequests")); + return [...new Set([...this.providers, ...Object.keys(limits)])].sort((a, b) => a.localeCompare(b)); + } + + #showProviderList(): void { + this.clear(); + this.addChild(new Text(theme.bold(theme.fg("accent", "Max In-Flight Requests")), 0, 0)); + this.addChild(new Spacer(1)); + this.addChild( + new Text( + theme.fg( + "muted", + "Select a provider, enter a positive number to cap concurrent LLM requests, or clear it for unlimited.", + ), + 0, + 0, + ), + ); + this.addChild(new Spacer(1)); + + const limits = normalizeProviderMaxInFlightRequests(settings.get("providers.maxInFlightRequests")); + const providerItems = this.#providerIds().map((provider): SelectItem => { + const limit = limits[provider]; + return { + value: provider, + label: provider, + description: limit === undefined ? "Unlimited" : `Limit: ${limit}`, + }; + }); + const clearItem: SelectItem[] = + Object.keys(limits).length === 0 + ? [] + : [{ value: "__clear_all", label: "Clear all limits", description: "Make every provider unlimited" }]; + const items = [...providerItems, ...clearItem]; + this.#selectList = new SelectList(items, Math.min(Math.max(items.length, 1), 12), getSelectListTheme()); + this.#selectList.onSelect = item => { + if (item.value === "__clear_all") { + settings.set("providers.maxInFlightRequests", {}); + this.onChange({}); + this.#showProviderList(); + this.requestRender?.(); + return; + } + this.#showProviderEditor(item.value); + }; + this.#selectList.onCancel = this.onCancel; + this.addChild(this.#selectList); + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("dim", " Enter to edit provider · Esc to go back"), 0, 0)); + } + + #showProviderEditor(provider: string): void { + const limits = normalizeProviderMaxInFlightRequests(settings.get("providers.maxInFlightRequests")); + this.clear(); + this.#selectList = undefined; + this.addChild( + new TextInputSubmenu( + `Max In-Flight Requests: ${provider}`, + "Enter a positive number. Decimals round down. Clear the field to make this provider unlimited.", + limits[provider]?.toString() ?? "", + value => { + const next = { ...limits }; + const trimmed = value.trim(); + if (trimmed === "") { + delete next[provider]; + } else { + const limit = Number(trimmed); + if (!Number.isFinite(limit) || limit <= 0) throw new Error("Limit must be a positive number."); + next[provider] = Math.max(1, Math.floor(limit)); + } + const normalized = validateProviderMaxInFlightRequests(next); + settings.set("providers.maxInFlightRequests", normalized); + this.onChange(normalized); + this.#showProviderList(); + this.requestRender?.(); + }, + () => { + this.#showProviderList(); + this.requestRender?.(); + }, + ), + ); + } + + handleInput(data: string): void { + if (this.#selectList) { + this.#selectList.handleInput(data); + return; + } + this.children[0]?.handleInput?.(data); + } +} + let cachedSidebarWidth: number | undefined; /** * Split-sidebar width derived from every group name in the schema (not just @@ -258,6 +381,8 @@ export interface SettingsRuntimeContext { thinkingLevel: ThinkingLevel | undefined; /** Available themes */ availableThemes: string[]; + /** Provider/source ids shown in /model. */ + providers: string[]; /** Working directory for plugins tab */ cwd: string; /** Active model (api + id); resolves what the snapcompact `auto` shape maps to. */ @@ -721,10 +846,20 @@ export class SettingsSelectorComponent implements Component { id: def.path, label: def.label, description: def.description, - currentValue: (currentValue as string) ?? "", + currentValue: this.#formatTextInputValue(def.path, currentValue), submenu: (cv, done) => this.#createTextInput(def, cv, done), changed, }; + + case "providerLimits": + return { + id: def.path, + label: def.label, + description: def.description, + currentValue: this.#formatProviderLimitsValue(currentValue), + submenu: (_cv, done) => this.#createProviderLimitsInput(done), + changed, + }; } } @@ -854,7 +989,7 @@ export class SettingsSelectorComponent implements Component { */ #createTextInput( def: SettingDef & { type: "text" }, - currentValue: string, + _currentValue: string, done: (value?: string) => void, ): Container { this.#textInputActive = true; @@ -865,28 +1000,72 @@ export class SettingsSelectorComponent implements Component { return new TextInputSubmenu( def.label, def.description, - currentValue, + this.#formatTextInputEditValue(def.path, settings.get(def.path)), value => { // Empty string clears the setting; undefined-typed string settings // store "" which the browser.ts expandPath ignores (no-op fallback). this.#setSettingValue(def.path, value); - this.callbacks.onChange(def.path, value); - wrappedDone(value); + this.callbacks.onChange(def.path, settings.get(def.path)); + wrappedDone(this.#formatTextInputValue(def.path, settings.get(def.path))); }, () => wrappedDone(), ); } + #createProviderLimitsInput(done: (value?: string) => void): Container { + return new ProviderLimitsSubmenu( + this.context.providers, + value => { + this.callbacks.onChange("providers.maxInFlightRequests", value); + done(this.#formatProviderLimitsValue(value)); + }, + () => done(), + this.context.requestRender, + ); + } + + #formatProviderLimitsValue(value: unknown): string { + const limits = normalizeProviderMaxInFlightRequests(value); + const entries = Object.entries(limits).sort(([a], [b]) => a.localeCompare(b)); + if (entries.length === 0) return "Unlimited"; + return entries.map(([provider, limit]) => `${provider}: ${limit}`).join(", "); + } + + #formatTextInputValue(path: SettingPath, value: unknown): string { + if (path === "providers.maxInFlightRequests") return this.#formatProviderLimitsValue(value); + return this.#formatTextInputEditValue(path, value); + } + + #formatTextInputEditValue(_path: SettingPath, value: unknown): string { + if (value === undefined || value === null) return ""; + if (typeof value === "object") return JSON.stringify(value); + return String(value); + } + /** * Set a setting value, handling type conversion. */ #setSettingValue(path: SettingPath, value: string): void { - // Handle number conversions const currentValue = settings.get(path); + const schemaType = getType(path); if (path === "compaction.thresholdPercent" && value === "default") { settings.set(path, -1 as never); } else if (path === "compaction.thresholdTokens" && value === "default") { settings.set(path, -1 as never); + } else if (schemaType === "record") { + let parsed: unknown; + try { + parsed = JSON.parse(value || "{}"); + } catch { + throw new Error(`Invalid record JSON for ${path}`); + } + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error(`Invalid record JSON for ${path}`); + } + if (path === "providers.maxInFlightRequests") { + parsed = validateProviderMaxInFlightRequests(parsed); + } + settings.set(path, parsed as never); } else if (typeof currentValue === "number") { settings.set(path, Number(value) as never); } else if (typeof currentValue === "boolean") { diff --git a/packages/coding-agent/src/modes/controllers/selector-controller.ts b/packages/coding-agent/src/modes/controllers/selector-controller.ts index 7ecf68e90..5ea909e17 100644 --- a/packages/coding-agent/src/modes/controllers/selector-controller.ts +++ b/packages/coding-agent/src/modes/controllers/selector-controller.ts @@ -117,6 +117,9 @@ export class SelectorController { availableThinkingLevels: [...this.ctx.session.getAvailableThinkingLevels()], thinkingLevel: this.ctx.session.thinkingLevel, availableThemes, + providers: [...new Set(this.ctx.session.getAvailableModels().map(model => model.provider))].sort( + (a, b) => a.localeCompare(b), + ), cwd: getProjectDir(), model: this.ctx.session.model, imageBudget: this.ctx.ui.imageBudget, diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index a2640858a..05819c48f 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -42,7 +42,7 @@ import { resolveModelRoleValue, } from "./config/model-resolver"; import { loadPromptTemplates as loadPromptTemplatesInternal, type PromptTemplate } from "./config/prompt-templates"; -import { Settings, type SkillsSettings } from "./config/settings"; +import { Settings, type SkillsSettings, validateProviderMaxInFlightRequests } from "./config/settings"; import { CursorExecHandlers } from "./cursor"; import "./discovery"; import { initializeWithSettings } from "./discovery"; @@ -2520,6 +2520,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} ...streamOptions, openrouterVariant: streamOptions?.openrouterVariant ?? openrouterVariant, antigravityEndpointMode: streamOptions?.antigravityEndpointMode ?? antigravityEndpointMode, + maxInFlightRequests: validateProviderMaxInFlightRequests( + streamOptions?.maxInFlightRequests ?? settings.get("providers.maxInFlightRequests"), + ), loopGuard: { enabled: settings.get("model.loopGuard.enabled"), checkAssistantContent: settings.get("model.loopGuard.checkAssistantContent"), diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index a6b91853c..0ef5ff2a5 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -159,7 +159,7 @@ import { import { MODEL_ROLE_IDS, MODEL_ROLES } from "../config/model-roles"; import { expandPromptTemplate, type PromptTemplate } from "../config/prompt-templates"; import type { Settings, SkillsSettings } from "../config/settings"; -import { getDefault, onAppendOnlyModeChanged } from "../config/settings"; +import { getDefault, onAppendOnlyModeChanged, validateProviderMaxInFlightRequests } from "../config/settings"; import { RawSseDebugBuffer } from "../debug/raw-sse-buffer"; import { loadCapability } from "../discovery"; import { expandApplyPatchToEntries, normalizeDiff, normalizeToLF, ParseError, previewPatch, stripBom } from "../edit"; @@ -5298,6 +5298,9 @@ export class AgentSession { ...options, ...(openrouterVariant !== undefined && { openrouterVariant }), ...(antigravityEndpointMode !== undefined && { antigravityEndpointMode }), + maxInFlightRequests: validateProviderMaxInFlightRequests( + options.maxInFlightRequests ?? this.settings.get("providers.maxInFlightRequests"), + ), loopGuard: { enabled: this.settings.get("model.loopGuard.enabled"), checkAssistantContent: this.settings.get("model.loopGuard.checkAssistantContent"), diff --git a/packages/coding-agent/test/config-cli.test.ts b/packages/coding-agent/test/config-cli.test.ts index 76715aa74..f102fac72 100644 --- a/packages/coding-agent/test/config-cli.test.ts +++ b/packages/coding-agent/test/config-cli.test.ts @@ -66,6 +66,46 @@ describe("config CLI schema coverage", () => { expect(parsed.value).toEqual({ default: "claude-opus-4-6" }); }); + it("normalizes valid provider in-flight request limits from JSON objects", async () => { + const logSpy = vi.spyOn(console, "log").mockImplementation(() => {}); + + await runConfigCommand({ + action: "set", + key: "providers.maxInFlightRequests", + value: '{"openai":2.8,"anthropic":1}', + flags: { json: true }, + }); + await runConfigCommand({ action: "get", key: "providers.maxInFlightRequests", flags: { json: true } }); + + const payload = logSpy.mock.calls.at(-1)?.[0]; + expect(typeof payload).toBe("string"); + const parsed = JSON.parse(String(payload)) as { key: string; value: unknown; type: string }; + expect(parsed.key).toBe("providers.maxInFlightRequests"); + expect(parsed.type).toBe("record"); + expect(parsed.value).toEqual({ openai: 2, anthropic: 1 }); + }); + + it("rejects invalid provider in-flight request limit entries", async () => { + vi.spyOn(console, "log").mockImplementation(() => {}); + vi.spyOn(console, "error").mockImplementation(() => {}); + const exitSpy = vi.spyOn(process, "exit").mockImplementation((() => { + throw new Error("process.exit"); + }) as typeof process.exit); + + await expect( + runConfigCommand({ + action: "set", + key: "providers.maxInFlightRequests", + value: '{"openai":"2","anthropic":0}', + flags: { json: true }, + }), + ).rejects.toThrow("process.exit"); + expect(exitSpy).toHaveBeenCalledWith(1); + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining("Provider request limits must be positive numbers: openai, anthropic"), + ); + }); + it("sets and gets array settings as JSON arrays", async () => { const logSpy = vi.spyOn(console, "log").mockImplementation(() => {}); const arrayValue = '["claude-opus-4-6","gpt-5.3-codex"]'; diff --git a/packages/coding-agent/test/modes/components/settings-layout.test.ts b/packages/coding-agent/test/modes/components/settings-layout.test.ts index f9fec2699..a669e6d71 100644 --- a/packages/coding-agent/test/modes/components/settings-layout.test.ts +++ b/packages/coding-agent/test/modes/components/settings-layout.test.ts @@ -76,4 +76,15 @@ describe("settings layout", () => { expect(def.condition?.()).toBe(true); } }); + + it("shows provider request limits as a providers services submenu setting", () => { + const [def] = getSettingsForTab("providers").filter(item => item.path === "providers.maxInFlightRequests"); + + expect(def).toMatchObject({ + path: "providers.maxInFlightRequests", + type: "providerLimits", + tab: "providers", + group: "Services", + }); + }); }); diff --git a/packages/coding-agent/test/modes/components/settings-selector-memory-refresh.test.ts b/packages/coding-agent/test/modes/components/settings-selector-memory-refresh.test.ts index 369a072c9..a288e481f 100644 --- a/packages/coding-agent/test/modes/components/settings-selector-memory-refresh.test.ts +++ b/packages/coding-agent/test/modes/components/settings-selector-memory-refresh.test.ts @@ -45,6 +45,7 @@ function createSelector(onCancel: () => void = () => {}): SettingsSelectorCompon availableThinkingLevels: [], thinkingLevel: undefined, availableThemes: ["dark"], + providers: [], cwd: process.cwd(), }, { diff --git a/packages/coding-agent/test/settings-manager.test.ts b/packages/coding-agent/test/settings-manager.test.ts index dce19a62f..b64ea5c96 100644 --- a/packages/coding-agent/test/settings-manager.test.ts +++ b/packages/coding-agent/test/settings-manager.test.ts @@ -2,6 +2,10 @@ import { afterEach, beforeEach, describe, expect, it } from "bun:test"; import * as fs from "node:fs"; import * as path from "node:path"; import { Effort } from "@oh-my-pi/pi-ai"; +import { clearCustomApis } from "@oh-my-pi/pi-ai/api-registry"; +import { createMockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock"; +import { __providerInFlightForTesting, streamSimple } from "@oh-my-pi/pi-ai/stream"; +import type { Context } from "@oh-my-pi/pi-ai/types"; import { getDefault, getEnumValues, @@ -11,10 +15,18 @@ import { type SettingPath, Settings, } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentStorage } from "@oh-my-pi/pi-coding-agent/session/agent-storage"; import { getProjectAgentDir, TempDir } from "@oh-my-pi/pi-utils"; import { YAML } from "bun"; import { beginSettingsTest, restoreSettingsTestState, type SettingsTestState } from "./helpers/settings-test-state"; +function context(): Context { + return { + systemPrompt: [], + messages: [{ role: "user", content: "hi", timestamp: 0 }], + }; +} + describe("Settings", () => { let settingsState: SettingsTestState | undefined; let tempDir: TempDir; @@ -49,10 +61,14 @@ describe("Settings", () => { return parsed as Record; }; - afterEach(() => { + afterEach(async () => { + clearCustomApis(); + __providerInFlightForTesting.setRoot(undefined); + AgentStorage.resetInstance(); restoreSettingsTestState(settingsState); settingsState = undefined; - tempDir?.removeSync(); + await Bun.sleep(0); + await tempDir?.remove(); }); describe("defaults", () => { it("keeps eight inline images live by default", async () => { @@ -66,6 +82,12 @@ describe("Settings", () => { expect(getDefault("startup.showSplash")).toBe(false); }); + it("defaults provider in-flight request limits to an empty map", async () => { + const settings = Settings.isolated(); + expect(settings.get("providers.maxInFlightRequests")).toEqual({}); + expect(getDefault("providers.maxInFlightRequests")).toEqual({}); + }); + it("exposes all tool calling mode options", () => { const values = getEnumValues("tools.format"); expect(values).toEqual([ @@ -585,5 +607,78 @@ describe("Settings", () => { const settings = await Settings.init({ cwd: projectDir, agentDir }); expect(settings.get("power.sleepPrevention")).toBe("off"); }); + + describe("provider request limits", () => { + it("uses the effective merged value when configuring hooks", async () => { + const settings = Settings.isolated({ "providers.maxInFlightRequests": { openai: 1 } }); + __providerInFlightForTesting.setRoot(tempDir.join("provider-inflight")); + registerMockApi(); + const firstStarted = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + let active = 0; + let maxActive = 0; + let callIndex = 0; + const mock = createMockModel({ + provider: "openai", + handler: async () => { + callIndex++; + active++; + maxActive = Math.max(maxActive, active); + try { + if (callIndex === 1) { + firstStarted.resolve(); + await releaseFirst.promise; + } + return { content: [`reply ${callIndex}`] }; + } finally { + active--; + } + }, + }); + + settings.set("providers.maxInFlightRequests", { openai: 4 }); + + const first = streamSimple(mock.model, context()); + const firstResult = first.result(); + await firstStarted.promise; + const second = streamSimple(mock.model, context()); + await Bun.sleep(20); + + expect(settings.get("providers.maxInFlightRequests")).toEqual({ openai: 1 }); + expect(mock.calls).toHaveLength(1); + + releaseFirst.resolve(); + await Promise.all([firstResult, second.result()]); + expect(maxActive).toBe(1); + }); + + it("rejects invalid provider limits from config.yml", async () => { + await writeSettings({ providers: { maxInFlightRequests: { openai: "2" } } }); + + await expect(Settings.init({ cwd: projectDir, agentDir })).rejects.toThrow( + "Provider request limits must be positive numbers: openai", + ); + }); + + it("rejects invalid provider limits from project settings", async () => { + await Bun.write( + path.join(getProjectAgentDir(projectDir), "settings.json"), + JSON.stringify({ providers: { maxInFlightRequests: { anthropic: 0 } } }), + ); + + await expect(Settings.init({ cwd: projectDir, agentDir, inMemory: true })).rejects.toThrow( + "Provider request limits must be positive numbers: anthropic", + ); + }); + + it("rejects invalid provider limits from config overlays", async () => { + const overlayPath = tempDir.join("overlay.yml"); + await Bun.write(overlayPath, YAML.stringify({ providers: { maxInFlightRequests: { umans: -1 } } })); + + await expect( + Settings.init({ cwd: projectDir, agentDir, inMemory: true, configFiles: [overlayPath] }), + ).rejects.toThrow("Provider request limits must be positive numbers: umans"); + }); + }); }); });