mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-13 08:15:14 +02:00
feat(cli): add ingest LLM rate-limit governor with paced retries (#261)
* feat(cli): add ingest rate limit governor * feat(cli): wire ingest rate-limit config * feat(cli): report provider rate-limit signals * feat(cli): show ingest rate-limit waits * fix(cli): complete rate-limit event coverage * fix(cli): abort ingest provider calls cleanly * fix(cli): propagate ingest cancellation * fix(cli): reject pre-aborted ingest rate-limit waits * fix(cli): honor Claude rate-limit reset waits * fix(cli): retry thrown Codex rate-limit failures * fix(cli): type Claude rate-limit result details * fix(cli): emit ingest rate-limit countdowns from rejected signals * fix(cli): report ai sdk rate-limit header utilization * fix(cli): gate LLM rate-limit retries on the governor budget The AI SDK and Codex runtimes retried 429 / opaque rate-limit failures up to 6-7 times with no backoff when constructed without a RateLimitGovernor (scan, memory, setup) or with pacing disabled, ignoring Retry-After and worsening the limit. The outer retry loop only cooperates with the governor's pause, so without active pacing there is no backoff to apply. Route the retry bound through a single source: RateLimitGovernor .maxRetryAttempts(), which returns retry.maxAttempts when enabled and 1 (no outer retry) when absent or disabled. All three runtimes (ai-sdk, codex, claude-code) now use it, so ingest.rateLimit.retry.maxAttempts genuinely controls attempts and the hard-coded 6 (plus Codex's off-by-one extra attempt) is gone. Backend-native retry (e.g. the AI SDK's maxRetries) still handles transient 429s. Also correct the ktx.yaml docs for maxWaitMs (caps each wait, not the whole run) and maxAttempts, and sync uv.lock ktx-sl/ktx-daemon to 0.9.0.
This commit is contained in:
Andrey Avtomonov
GitHub
parent
5a8821073b
commit
c3d8cedb0b
35 changed files with 2336 additions and 72 deletions
31
packages/cli/test/context/core/abort.test.ts
Normal file
31
packages/cli/test/context/core/abort.test.ts
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
import { createAbortError, isAbortError, linkAbortSignal, throwIfAborted } from '../../../src/context/core/abort.js';
|
||||
|
||||
describe('abort helpers', () => {
|
||||
it('recognizes DOMException abort errors and common abort-shaped errors', () => {
|
||||
expect(isAbortError(createAbortError())).toBe(true);
|
||||
expect(isAbortError(Object.assign(new Error('cancelled'), { name: 'AbortError' }))).toBe(true);
|
||||
expect(isAbortError(Object.assign(new Error('operation aborted'), { code: 'ABORT_ERR' }))).toBe(true);
|
||||
expect(isAbortError(new Error('ordinary failure'))).toBe(false);
|
||||
});
|
||||
|
||||
it('throws when the provided signal is already aborted', () => {
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
expect(() => throwIfAborted(controller.signal)).toThrow(/Aborted/);
|
||||
});
|
||||
|
||||
it('links a child controller to a parent signal and removes the listener on dispose', () => {
|
||||
const parent = new AbortController();
|
||||
const child = linkAbortSignal(parent.signal);
|
||||
|
||||
expect(child.controller.signal.aborted).toBe(false);
|
||||
parent.abort();
|
||||
expect(child.controller.signal.aborted).toBe(true);
|
||||
|
||||
const removeSpy = vi.spyOn(parent.signal, 'removeEventListener');
|
||||
child.dispose();
|
||||
expect(removeSpy).toHaveBeenCalledWith('abort', expect.any(Function));
|
||||
});
|
||||
});
|
||||
|
|
@ -426,6 +426,177 @@ describe('IngestBundleRunner — Stages 1 → 7', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('uses the rate-limit governor for work-unit start slots', async () => {
|
||||
const deps = makeDeps();
|
||||
const acquireWorkSlot = vi.fn(async () => vi.fn());
|
||||
const runner = buildRunner(deps, {
|
||||
settings: {
|
||||
probeRowCount: 1,
|
||||
memoryIngestionModel: 'test-model',
|
||||
workUnitMaxConcurrency: 2,
|
||||
rateLimitGovernor: { acquireWorkSlot, subscribe: vi.fn(() => vi.fn()) } as never,
|
||||
},
|
||||
});
|
||||
deps.adapter.chunk.mockResolvedValue({
|
||||
workUnits: [
|
||||
{ unitKey: 'u1', rawFiles: ['a.yml'], peerFileIndex: [], dependencyPaths: [] },
|
||||
{ unitKey: 'u2', rawFiles: ['b.yml'], peerFileIndex: [], dependencyPaths: [] },
|
||||
],
|
||||
});
|
||||
(runner as any).stageRawFilesStage1 = vi.fn().mockResolvedValue({
|
||||
currentHashes: new Map([
|
||||
['a.yml', 'h1'],
|
||||
['b.yml', 'h2'],
|
||||
]),
|
||||
rawDirInWorktree: 'raw-sources/c1/fake/s',
|
||||
});
|
||||
(runner as any).resolveStagedDir = vi.fn().mockResolvedValue('/tmp/stage/upload-x');
|
||||
|
||||
await runner.run({
|
||||
jobId: 'j1',
|
||||
connectionId: 'c1',
|
||||
sourceKey: 'fake',
|
||||
trigger: 'upload',
|
||||
bundleRef: { kind: 'upload', uploadId: 'upload-x' },
|
||||
});
|
||||
|
||||
expect(acquireWorkSlot).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('passes the job abort signal into rate-limit work-unit slots', async () => {
|
||||
const deps = makeDeps();
|
||||
const controller = new AbortController();
|
||||
const acquireWorkSlot = vi.fn(async () => vi.fn());
|
||||
const runner = buildRunner(deps, {
|
||||
settings: {
|
||||
probeRowCount: 1,
|
||||
memoryIngestionModel: 'test-model',
|
||||
workUnitMaxConcurrency: 1,
|
||||
rateLimitGovernor: { acquireWorkSlot, subscribe: vi.fn(() => vi.fn()) } as never,
|
||||
},
|
||||
});
|
||||
(runner as any).stageRawFilesStage1 = vi.fn().mockResolvedValue({
|
||||
currentHashes: new Map([['a.yml', 'h1']]),
|
||||
rawDirInWorktree: 'raw-sources/c1/fake/s',
|
||||
});
|
||||
(runner as any).resolveStagedDir = vi.fn().mockResolvedValue('/tmp/stage/upload-x');
|
||||
|
||||
await runner.run(
|
||||
{
|
||||
jobId: 'j1',
|
||||
connectionId: 'c1',
|
||||
sourceKey: 'fake',
|
||||
trigger: 'upload',
|
||||
bundleRef: { kind: 'upload', uploadId: 'upload-x' },
|
||||
},
|
||||
{ jobId: 'j1', abortSignal: controller.signal, startPhase: () => new TestJobContext('j1', null, async () => undefined, async () => undefined) } as any,
|
||||
);
|
||||
|
||||
expect(acquireWorkSlot).toHaveBeenCalledWith(controller.signal);
|
||||
});
|
||||
|
||||
it('does not convert aborted work-unit agent loops into failed work units', async () => {
|
||||
const deps = makeDeps();
|
||||
const controller = new AbortController();
|
||||
deps.agentRunner.runLoop.mockImplementation(async () => {
|
||||
controller.abort();
|
||||
throw new DOMException('Aborted', 'AbortError');
|
||||
});
|
||||
const runner = buildRunner(deps, {
|
||||
settings: {
|
||||
probeRowCount: 1,
|
||||
memoryIngestionModel: 'test-model',
|
||||
workUnitMaxConcurrency: 1,
|
||||
},
|
||||
});
|
||||
(runner as any).stageRawFilesStage1 = vi.fn().mockResolvedValue({
|
||||
currentHashes: new Map([['a.yml', 'h1']]),
|
||||
rawDirInWorktree: 'raw-sources/c1/fake/s',
|
||||
});
|
||||
(runner as any).resolveStagedDir = vi.fn().mockResolvedValue('/tmp/stage/upload-x');
|
||||
|
||||
await expect(
|
||||
runner.run(
|
||||
{
|
||||
jobId: 'j1',
|
||||
connectionId: 'c1',
|
||||
sourceKey: 'fake',
|
||||
trigger: 'upload',
|
||||
bundleRef: { kind: 'upload', uploadId: 'upload-x' },
|
||||
},
|
||||
{ jobId: 'j1', abortSignal: controller.signal, startPhase: () => new TestJobContext('j1', null, async () => undefined, async () => undefined) } as any,
|
||||
),
|
||||
).rejects.toThrow(/Aborted/);
|
||||
|
||||
expect(deps.runsRepo.markFailed).toHaveBeenCalledWith('run-1');
|
||||
expect(deps.reportsRepo.create).not.toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
body: expect.objectContaining({
|
||||
failedWorkUnits: expect.arrayContaining(['u1']),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('emits trace and memory-flow status for rate-limit waits', async () => {
|
||||
const deps = makeDeps();
|
||||
let subscriber: ((state: any) => void) | undefined;
|
||||
const memoryFlow = createMemoryFlowLiveBuffer(bundleReplayInput());
|
||||
const runner = buildRunner(deps, {
|
||||
settings: {
|
||||
probeRowCount: 1,
|
||||
memoryIngestionModel: 'test-model',
|
||||
rateLimitGovernor: {
|
||||
acquireWorkSlot: vi.fn(async () => vi.fn()),
|
||||
subscribe: vi.fn((cb: (state: any) => void) => {
|
||||
subscriber = cb;
|
||||
return vi.fn();
|
||||
}),
|
||||
} as never,
|
||||
},
|
||||
});
|
||||
(runner as any).runInner = async (_job: any, ctx: any) => {
|
||||
subscriber?.({
|
||||
kind: 'wait_tick',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
resumeAtMs: 2_000,
|
||||
remainingMs: 1_000,
|
||||
});
|
||||
ctx.memoryFlow.emit({ type: 'report_created', runId: 'run-1' });
|
||||
return {
|
||||
runId: 'run-1',
|
||||
syncId: 'sync-1',
|
||||
diffSummary: { added: 0, modified: 0, deleted: 0, unchanged: 0 },
|
||||
workUnitCount: 0,
|
||||
failedWorkUnits: [],
|
||||
artifactsWritten: 0,
|
||||
commitSha: null,
|
||||
};
|
||||
};
|
||||
|
||||
await runner.run(
|
||||
{
|
||||
jobId: 'j1',
|
||||
connectionId: 'c1',
|
||||
sourceKey: 'fake',
|
||||
trigger: 'upload',
|
||||
bundleRef: { kind: 'upload', uploadId: 'upload-x' },
|
||||
},
|
||||
{ memoryFlow } as any,
|
||||
);
|
||||
|
||||
expect(memoryFlow.snapshot().events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: 'rate_limit_wait',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
resumeAtMs: 2_000,
|
||||
remainingMs: 1_000,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('fails before squash when reconciliation leaves a touched wiki page with dangling refs', async () => {
|
||||
const deps = makeDeps();
|
||||
let currentToolSession: any = null;
|
||||
|
|
|
|||
|
|
@ -301,6 +301,7 @@ describe('createLocalBundleIngestRuntime', () => {
|
|||
'memoryIngestionModel',
|
||||
'probeRowCount',
|
||||
'profileIngest',
|
||||
'rateLimitGovernor',
|
||||
'workUnitFailureMode',
|
||||
'workUnitMaxConcurrency',
|
||||
'workUnitStepBudget',
|
||||
|
|
|
|||
|
|
@ -146,6 +146,29 @@ describe('memory-flow schemas', () => {
|
|||
expect(parsed.events).toContainEqual({ type: 'stage_skipped', stage: 'actions', reason: 'requires LLM' });
|
||||
});
|
||||
|
||||
it('accepts rate-limit wait replay events', () => {
|
||||
expect(
|
||||
memoryFlowReplayInputSchema.parse({
|
||||
...snapshot(),
|
||||
events: [
|
||||
{
|
||||
type: 'rate_limit_wait',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
resumeAtMs: 2_000,
|
||||
remainingMs: 1_000,
|
||||
},
|
||||
],
|
||||
}).events[0],
|
||||
).toEqual({
|
||||
type: 'rate_limit_wait',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
resumeAtMs: 2_000,
|
||||
remainingMs: 1_000,
|
||||
});
|
||||
});
|
||||
|
||||
it('parses snapshot and closed stream events', () => {
|
||||
expect(memoryFlowStreamEventSchema.parse({ type: 'snapshot', snapshot: snapshot({ status: 'done' }) })).toEqual({
|
||||
type: 'snapshot',
|
||||
|
|
|
|||
|
|
@ -107,6 +107,199 @@ describe('AiSdkKtxLlmRuntime.runAgentLoop', () => {
|
|||
expect(result.error).toBe(err);
|
||||
});
|
||||
|
||||
it('reports AI SDK retry-after rate limits and retries through the governor', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const rateLimitError = Object.assign(new Error('too many requests'), {
|
||||
name: 'TooManyRequestsError',
|
||||
retryAfter: 2,
|
||||
statusCode: 429,
|
||||
});
|
||||
(generateText as any).mockRejectedValueOnce(rateLimitError).mockResolvedValueOnce({
|
||||
text: 'done',
|
||||
toolCalls: [],
|
||||
steps: [],
|
||||
usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
|
||||
});
|
||||
const runtime = new AiSdkKtxLlmRuntime({
|
||||
llmProvider: llmProvider as any,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('natural');
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'anthropic-api',
|
||||
status: 'rejected',
|
||||
retryAfterMs: 2_000,
|
||||
rateLimitType: 'http_429',
|
||||
});
|
||||
expect(waitForReady).toHaveBeenCalledTimes(2);
|
||||
expect(generateText).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('does not retry AI SDK rate limits without a governor', async () => {
|
||||
const rateLimitError = Object.assign(new Error('too many requests'), {
|
||||
name: 'TooManyRequestsError',
|
||||
statusCode: 429,
|
||||
});
|
||||
(generateText as any).mockRejectedValue(rateLimitError);
|
||||
// The beforeEach runtime is constructed without a rateLimitGovernor.
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('error');
|
||||
expect(generateText).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('honors a governor retry budget of one attempt without retrying', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const rateLimitError = Object.assign(new Error('too many requests'), {
|
||||
name: 'TooManyRequestsError',
|
||||
statusCode: 429,
|
||||
});
|
||||
(generateText as any).mockRejectedValue(rateLimitError);
|
||||
const runtime = new AiSdkKtxLlmRuntime({
|
||||
llmProvider: llmProvider as any,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 1 } as never,
|
||||
});
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('error');
|
||||
expect(generateText).toHaveBeenCalledTimes(1);
|
||||
expect(report).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('reports Anthropic API response-header utilization to the governor', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
(generateText as any).mockResolvedValue({
|
||||
text: 'done',
|
||||
toolCalls: [],
|
||||
steps: [],
|
||||
response: {
|
||||
headers: {
|
||||
'anthropic-ratelimit-requests-limit': '100',
|
||||
'anthropic-ratelimit-requests-remaining': '8',
|
||||
'anthropic-ratelimit-input-tokens-limit': '10000',
|
||||
'anthropic-ratelimit-input-tokens-remaining': '9000',
|
||||
},
|
||||
},
|
||||
});
|
||||
const runtime = new AiSdkKtxLlmRuntime({
|
||||
llmProvider: llmProvider as any,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('natural');
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'anthropic-api',
|
||||
status: 'allowed',
|
||||
rateLimitType: 'rpm',
|
||||
utilization: 0.92,
|
||||
});
|
||||
});
|
||||
|
||||
it('reports generic x-ratelimit response-header utilization for Vertex providers', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const vertexProvider = {
|
||||
...llmProvider,
|
||||
getModel: vi.fn().mockReturnValue({ modelId: 'gemini-3-pro', provider: 'google-vertex' }),
|
||||
};
|
||||
(generateText as any).mockResolvedValue({
|
||||
text: 'done',
|
||||
toolCalls: [],
|
||||
steps: [],
|
||||
response: {
|
||||
headers: {
|
||||
'x-ratelimit-limit-requests': '200',
|
||||
'x-ratelimit-remaining-requests': '30',
|
||||
'x-ratelimit-limit-tokens': '100000',
|
||||
'x-ratelimit-remaining-tokens': '4000',
|
||||
},
|
||||
},
|
||||
});
|
||||
const runtime = new AiSdkKtxLlmRuntime({
|
||||
llmProvider: vertexProvider as any,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('natural');
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'vertex',
|
||||
status: 'allowed',
|
||||
rateLimitType: 'tpm',
|
||||
utilization: 0.96,
|
||||
});
|
||||
});
|
||||
|
||||
it('passes abort signals into governor waits and AI SDK generateText calls', async () => {
|
||||
const controller = new AbortController();
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
(generateText as any).mockResolvedValue({ text: 'done', toolCalls: [], steps: [] });
|
||||
const runtime = new AiSdkKtxLlmRuntime({
|
||||
llmProvider: llmProvider as any,
|
||||
rateLimitGovernor: { waitForReady, report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const result = await runtime.runAgentLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
abortSignal: controller.signal,
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('natural');
|
||||
expect(waitForReady).toHaveBeenCalledWith(controller.signal);
|
||||
expect((generateText as any).mock.calls[0][0].abortSignal).toBe(controller.signal);
|
||||
});
|
||||
|
||||
it('returns metrics with stepCount, per-step boundaries, and aggregate token usage', async () => {
|
||||
(generateText as any).mockImplementation(async (opts: any) => {
|
||||
await opts.onStepFinish({});
|
||||
|
|
|
|||
|
|
@ -9,6 +9,14 @@ async function* stream(messages: SDKMessage[]): AsyncGenerator<SDKMessage, void>
|
|||
}
|
||||
}
|
||||
|
||||
function deferred<T>() {
|
||||
let resolve!: (value: T | PromiseLike<T>) => void;
|
||||
const promise = new Promise<T>((innerResolve) => {
|
||||
resolve = innerResolve;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
function initMessage(overrides: Partial<Extract<SDKMessage, { type: 'system'; subtype: 'init' }>> = {}): Extract<
|
||||
SDKMessage,
|
||||
{ type: 'system'; subtype: 'init' }
|
||||
|
|
@ -91,6 +99,247 @@ describe('ClaudeCodeKtxLlmRuntime', () => {
|
|||
});
|
||||
});
|
||||
|
||||
it('waits before Claude Code text generation and reports rate-limit events', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const query = vi.fn((_input: any) =>
|
||||
stream([
|
||||
{
|
||||
type: 'rate_limit_event',
|
||||
rate_limit_info: {
|
||||
status: 'allowed_warning',
|
||||
resetsAt: new Date(2_000).toISOString(),
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 0.91,
|
||||
},
|
||||
} as unknown as SDKMessage,
|
||||
resultMessage({ result: 'ok' }),
|
||||
]),
|
||||
);
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).resolves.toBe('ok');
|
||||
expect(waitForReady).toHaveBeenCalledTimes(1);
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'claude-subscription',
|
||||
status: 'warning',
|
||||
resetAtMs: 2_000,
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 0.91,
|
||||
});
|
||||
});
|
||||
|
||||
it('maps numeric Claude Code reset times from SDK rate-limit events', async () => {
|
||||
const report = vi.fn();
|
||||
const resetAtMs = 1_700_000_000_000;
|
||||
const query = vi.fn((_input: any) =>
|
||||
stream([
|
||||
{
|
||||
type: 'rate_limit_event',
|
||||
rate_limit_info: {
|
||||
status: 'rejected',
|
||||
resetsAt: resetAtMs,
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 1,
|
||||
},
|
||||
} as unknown as SDKMessage,
|
||||
resultMessage({ result: 'ok' }),
|
||||
]),
|
||||
);
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady: vi.fn().mockResolvedValue(undefined), report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).resolves.toBe('ok');
|
||||
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'claude-subscription',
|
||||
status: 'rejected',
|
||||
resetAtMs,
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 1,
|
||||
});
|
||||
});
|
||||
|
||||
it('retries a Claude Code query after an SDK rate-limit result error', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const resetAtMs = 1_700_000_000_000;
|
||||
const query = vi
|
||||
.fn()
|
||||
.mockReturnValueOnce(
|
||||
stream([
|
||||
{
|
||||
type: 'rate_limit_event',
|
||||
rate_limit_info: {
|
||||
status: 'rejected',
|
||||
resetsAt: resetAtMs,
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 1,
|
||||
},
|
||||
} as unknown as SDKMessage,
|
||||
resultMessage({
|
||||
subtype: 'error_during_execution',
|
||||
is_error: true,
|
||||
result: '',
|
||||
errors: ['rate limit retry budget exhausted'],
|
||||
terminal_reason: 'model_error',
|
||||
} as never),
|
||||
]),
|
||||
)
|
||||
.mockReturnValueOnce(stream([resultMessage({ result: 'ok' })]));
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).resolves.toBe('ok');
|
||||
|
||||
expect(query).toHaveBeenCalledTimes(2);
|
||||
expect(waitForReady).toHaveBeenCalledTimes(2);
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'claude-subscription',
|
||||
status: 'rejected',
|
||||
resetAtMs,
|
||||
rateLimitType: 'five_hour',
|
||||
utilization: 1,
|
||||
});
|
||||
});
|
||||
|
||||
it('reports Claude Code api retry messages as warning signals', async () => {
|
||||
const report = vi.fn();
|
||||
const query = vi.fn((_input: any) =>
|
||||
stream([
|
||||
{
|
||||
type: 'system',
|
||||
subtype: 'api_retry',
|
||||
retry_delay_ms: 12_000,
|
||||
} as unknown as SDKMessage,
|
||||
resultMessage({ result: 'ok' }),
|
||||
]),
|
||||
);
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady: vi.fn().mockResolvedValue(undefined), report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await runtime.generateText({ role: 'default', prompt: 'hello' });
|
||||
expect(report).toHaveBeenCalledWith({
|
||||
provider: 'claude-subscription',
|
||||
status: 'warning',
|
||||
retryAfterMs: 12_000,
|
||||
rateLimitType: 'api_retry',
|
||||
});
|
||||
});
|
||||
|
||||
it('passes abort signals into Claude Code governor waits', async () => {
|
||||
const controller = new AbortController();
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const query = vi.fn((_input: any) => stream([resultMessage({ result: 'ok' })]));
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady, report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello', abortSignal: controller.signal })).resolves.toBe('ok');
|
||||
|
||||
expect(waitForReady).toHaveBeenCalledWith(controller.signal);
|
||||
});
|
||||
|
||||
it('interrupts an active Claude Code query when the abort signal fires', async () => {
|
||||
const controller = new AbortController();
|
||||
const streamStarted = deferred<void>();
|
||||
const releaseStream = deferred<void>();
|
||||
const interrupt = vi.fn(() => releaseStream.resolve());
|
||||
const queryResult = {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
streamStarted.resolve();
|
||||
await releaseStream.promise;
|
||||
yield resultMessage({ result: 'ok' });
|
||||
},
|
||||
interrupt,
|
||||
};
|
||||
const query = vi.fn(() => queryResult as never);
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady: vi.fn().mockResolvedValue(undefined), report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const pending = runtime.generateText({ role: 'default', prompt: 'hello', abortSignal: controller.signal });
|
||||
await streamStarted.promise;
|
||||
controller.abort();
|
||||
|
||||
await expect(pending).rejects.toThrow(/Aborted/);
|
||||
expect(interrupt).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('throws abort before starting Claude Code query when the signal is already aborted', async () => {
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
const query = vi.fn((_input: any) => stream([resultMessage({ result: 'ok' })]));
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady: vi.fn().mockResolvedValue(undefined), report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello', abortSignal: controller.signal })).rejects.toThrow(/Aborted/);
|
||||
expect(query).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('treats an interrupted Claude Code stream with no result as abort', async () => {
|
||||
const controller = new AbortController();
|
||||
const streamStarted = deferred<void>();
|
||||
const releaseStream = deferred<void>();
|
||||
const interrupt = vi.fn(() => releaseStream.resolve());
|
||||
const queryResult = {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
streamStarted.resolve();
|
||||
await releaseStream.promise;
|
||||
},
|
||||
interrupt,
|
||||
};
|
||||
const query = vi.fn(() => queryResult as never);
|
||||
const runtime = new ClaudeCodeKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'sonnet' },
|
||||
query,
|
||||
env: {},
|
||||
rateLimitGovernor: { waitForReady: vi.fn().mockResolvedValue(undefined), report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
const pending = runtime.generateText({ role: 'default', prompt: 'hello', abortSignal: controller.signal });
|
||||
await streamStarted.promise;
|
||||
controller.abort();
|
||||
|
||||
await expect(pending).rejects.toThrow(/Aborted/);
|
||||
expect(interrupt).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('validates structured output with the caller schema and whitelists the SDK StructuredOutput tool', async () => {
|
||||
const schema = z.object({ answer: z.string() });
|
||||
const query = vi.fn((_input: any) =>
|
||||
|
|
|
|||
|
|
@ -130,6 +130,150 @@ describe('CodexKtxLlmRuntime', () => {
|
|||
).rejects.toThrow('Codex structured output failed validation');
|
||||
});
|
||||
|
||||
it('reports Codex rate-limit failures and retries with opaque backoff', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const fakeRunner = {
|
||||
runStreamed: vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(events([{ type: 'turn.failed', error: { message: '429 rate limit exceeded' } }]))
|
||||
.mockResolvedValueOnce(
|
||||
events([
|
||||
{ type: 'turn.started' },
|
||||
{ type: 'item.completed', item: { type: 'agent_message', text: 'ok' } },
|
||||
{ type: 'turn.completed' },
|
||||
]),
|
||||
),
|
||||
};
|
||||
const runtime = new CodexKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'codex' },
|
||||
runner: fakeRunner,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).resolves.toBe('ok');
|
||||
expect(report).toHaveBeenCalledWith({ provider: 'codex', status: 'rejected', rateLimitType: 'opaque' });
|
||||
expect(waitForReady).toHaveBeenCalledTimes(2);
|
||||
expect(fakeRunner.runStreamed).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('reports thrown Codex rate-limit failures and retries with opaque backoff', async () => {
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
const report = vi.fn();
|
||||
const fakeRunner = {
|
||||
runStreamed: vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error('ThreadError: 429 rate limit exceeded'))
|
||||
.mockResolvedValueOnce(
|
||||
events([
|
||||
{ type: 'turn.started' },
|
||||
{ type: 'item.completed', item: { type: 'agent_message', text: 'ok' } },
|
||||
{ type: 'turn.completed' },
|
||||
]),
|
||||
),
|
||||
};
|
||||
const runtime = new CodexKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'codex' },
|
||||
runner: fakeRunner,
|
||||
rateLimitGovernor: { waitForReady, report, maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).resolves.toBe('ok');
|
||||
|
||||
expect(report).toHaveBeenCalledWith({ provider: 'codex', status: 'rejected', rateLimitType: 'opaque' });
|
||||
expect(waitForReady).toHaveBeenCalledTimes(2);
|
||||
expect(fakeRunner.runStreamed).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('surfaces Codex rate-limit failures without retrying when no governor is present', async () => {
|
||||
const fakeRunner = runner([{ type: 'turn.failed', error: { message: '429 rate limit exceeded' } }]);
|
||||
const runtime = new CodexKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'codex' },
|
||||
runner: fakeRunner,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello' })).rejects.toThrow(/rate limit/i);
|
||||
expect(fakeRunner.runStreamed).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('passes abort signals into Codex text generation and governor waits', async () => {
|
||||
const controller = new AbortController();
|
||||
const waitForReady = vi.fn().mockResolvedValue(undefined);
|
||||
let observedSignal: AbortSignal | undefined;
|
||||
const fakeRunner = {
|
||||
runStreamed: vi.fn(async (input: { signal?: AbortSignal }) => {
|
||||
observedSignal = input.signal;
|
||||
return events([
|
||||
{ type: 'turn.started' },
|
||||
{ type: 'item.completed', item: { type: 'agent_message', text: 'ok' } },
|
||||
{ type: 'turn.completed' },
|
||||
]);
|
||||
}),
|
||||
};
|
||||
const runtime = new CodexKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'codex' },
|
||||
runner: fakeRunner,
|
||||
rateLimitGovernor: { waitForReady, report: vi.fn(), maxRetryAttempts: () => 6 } as never,
|
||||
});
|
||||
|
||||
await expect(runtime.generateText({ role: 'default', prompt: 'hello', abortSignal: controller.signal })).resolves.toBe('ok');
|
||||
|
||||
expect(waitForReady).toHaveBeenCalledWith(controller.signal);
|
||||
expect(observedSignal).toBe(controller.signal);
|
||||
});
|
||||
|
||||
it('links the parent abort signal into Codex agent-loop streamed runs', async () => {
|
||||
const controller = new AbortController();
|
||||
let releaseStream!: () => void;
|
||||
const streamRelease = new Promise<void>((resolve) => {
|
||||
releaseStream = resolve;
|
||||
});
|
||||
let markRunnerCalled!: () => void;
|
||||
const runnerCalled = new Promise<void>((resolve) => {
|
||||
markRunnerCalled = resolve;
|
||||
});
|
||||
let observedSignal: AbortSignal | undefined;
|
||||
const fakeRunner = {
|
||||
runStreamed: vi.fn(async (input: { signal?: AbortSignal }) => {
|
||||
observedSignal = input.signal;
|
||||
markRunnerCalled();
|
||||
return (async function* () {
|
||||
await streamRelease;
|
||||
yield { type: 'turn.started' };
|
||||
yield { type: 'item.completed', item: { type: 'agent_message', text: 'ok' } };
|
||||
yield { type: 'turn.completed' };
|
||||
})();
|
||||
}),
|
||||
};
|
||||
const runtime = new CodexKtxLlmRuntime({
|
||||
projectDir: '/tmp/project',
|
||||
modelSlots: { default: 'codex' },
|
||||
runner: fakeRunner,
|
||||
});
|
||||
|
||||
const pending = runtime.runAgentLoop({
|
||||
modelRole: 'default',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
abortSignal: controller.signal,
|
||||
});
|
||||
|
||||
await runnerCalled;
|
||||
expect(observedSignal).toBeDefined();
|
||||
expect(observedSignal).not.toBe(controller.signal);
|
||||
controller.abort();
|
||||
expect(observedSignal?.aborted).toBe(true);
|
||||
releaseStream();
|
||||
await expect(pending).resolves.toMatchObject({ stopReason: 'natural' });
|
||||
});
|
||||
|
||||
it('starts and closes a temporary MCP server for tool-backed agent loops', async () => {
|
||||
const close = vi.fn(async () => undefined);
|
||||
const startMcpServer = vi.fn(async () => ({
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
import {
|
||||
createLocalKtxEmbeddingProviderFromConfig,
|
||||
createLocalKtxLlmProviderFromConfig,
|
||||
createLocalKtxLlmRuntimeFromConfig,
|
||||
resolveLocalKtxEmbeddingConfig,
|
||||
resolveLocalKtxLlmConfig,
|
||||
} from '../../../src/context/llm/local-config.js';
|
||||
|
|
@ -129,6 +130,64 @@ describe('local KTX LLM config', () => {
|
|||
vertexFallbackTo5m: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('passes the rate-limit governor into created runtimes', () => {
|
||||
const rateLimitGovernor = {} as never;
|
||||
const createClaudeCodeRuntime = vi.fn(() => ({
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
}));
|
||||
const createCodexRuntime = vi.fn(() => ({
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
}));
|
||||
const createAiSdkRuntime = vi.fn(() => ({
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
}));
|
||||
const createKtxLlmProvider = vi.fn(() => ({
|
||||
getModel: vi.fn(),
|
||||
getModelByName: vi.fn(),
|
||||
cacheMarker: vi.fn(),
|
||||
repairToolCallHandler: vi.fn(),
|
||||
thinkingProviderOptions: vi.fn(),
|
||||
telemetryConfig: vi.fn(),
|
||||
promptCachingConfig: vi.fn(),
|
||||
activeBackend: vi.fn(() => 'anthropic'),
|
||||
}));
|
||||
|
||||
createLocalKtxLlmRuntimeFromConfig(
|
||||
{
|
||||
provider: { backend: 'claude-code' },
|
||||
models: { default: 'sonnet' },
|
||||
promptCaching: undefined,
|
||||
},
|
||||
{ projectDir: '/tmp/project', env: {}, rateLimitGovernor, createClaudeCodeRuntime },
|
||||
);
|
||||
createLocalKtxLlmRuntimeFromConfig(
|
||||
{
|
||||
provider: { backend: 'codex' },
|
||||
models: { default: 'codex' },
|
||||
promptCaching: undefined,
|
||||
},
|
||||
{ projectDir: '/tmp/project', env: {}, rateLimitGovernor, createCodexRuntime },
|
||||
);
|
||||
createLocalKtxLlmRuntimeFromConfig(
|
||||
{
|
||||
provider: { backend: 'anthropic' },
|
||||
models: { default: 'claude-sonnet-4-6' },
|
||||
promptCaching: undefined,
|
||||
},
|
||||
{ env: {}, rateLimitGovernor, createAiSdkRuntime, createKtxLlmProvider: createKtxLlmProvider as never },
|
||||
);
|
||||
|
||||
expect(createClaudeCodeRuntime).toHaveBeenCalledWith(expect.objectContaining({ rateLimitGovernor }));
|
||||
expect(createCodexRuntime).toHaveBeenCalledWith(expect.objectContaining({ rateLimitGovernor }));
|
||||
expect(createAiSdkRuntime).toHaveBeenCalledWith(expect.objectContaining({ rateLimitGovernor }));
|
||||
});
|
||||
});
|
||||
|
||||
describe('local KTX embedding config', () => {
|
||||
|
|
|
|||
278
packages/cli/test/context/llm/rate-limit-governor.test.ts
Normal file
278
packages/cli/test/context/llm/rate-limit-governor.test.ts
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
import {
|
||||
createRateLimitGovernorConfig,
|
||||
RateLimitGovernor,
|
||||
type RateLimitWaitState,
|
||||
} from '../../../src/context/llm/rate-limit-governor.js';
|
||||
|
||||
function testClock(startMs = 1_000) {
|
||||
let nowMs = startMs;
|
||||
return {
|
||||
now: () => nowMs,
|
||||
advance: (ms: number) => {
|
||||
nowMs += ms;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function flushMicrotasks(turns = 10): Promise<void> {
|
||||
for (let i = 0; i < turns; i += 1) {
|
||||
await Promise.resolve();
|
||||
}
|
||||
}
|
||||
|
||||
describe('RateLimitGovernor', () => {
|
||||
it('drops and restores the effective work-unit limit from warning signals', () => {
|
||||
const clock = testClock();
|
||||
const states: RateLimitWaitState[] = [];
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 6, minConcurrencyUnderPressure: 1 }),
|
||||
{ now: clock.now, sleep: async () => undefined, random: () => 0 },
|
||||
);
|
||||
governor.subscribe((state) => states.push(state));
|
||||
|
||||
expect(governor.currentLimit()).toBe(6);
|
||||
governor.report({
|
||||
provider: 'claude-subscription',
|
||||
status: 'warning',
|
||||
utilization: 0.91,
|
||||
rateLimitType: 'five_hour',
|
||||
});
|
||||
expect(governor.currentLimit()).toBe(1);
|
||||
governor.report({
|
||||
provider: 'claude-subscription',
|
||||
status: 'allowed',
|
||||
utilization: 0.2,
|
||||
rateLimitType: 'five_hour',
|
||||
});
|
||||
expect(governor.currentLimit()).toBe(6);
|
||||
expect(states.map((state) => state.kind)).toContain('concurrency_adjusted');
|
||||
});
|
||||
|
||||
it('blocks work slots during a rejected reset window and emits wait states', async () => {
|
||||
const clock = testClock();
|
||||
const states: RateLimitWaitState[] = [];
|
||||
const sleeps: number[] = [];
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 2, waitStateTickMs: 100 }),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (ms) => {
|
||||
sleeps.push(ms);
|
||||
clock.advance(ms);
|
||||
},
|
||||
},
|
||||
);
|
||||
governor.subscribe((state) => states.push(state));
|
||||
|
||||
governor.report({ provider: 'anthropic-api', status: 'rejected', retryAfterMs: 250, rateLimitType: 'rpm' });
|
||||
const release = await governor.acquireWorkSlot();
|
||||
release();
|
||||
|
||||
expect(sleeps).toEqual([100, 100, 50]);
|
||||
expect(states.some((state) => state.kind === 'wait_started' && state.provider === 'anthropic-api')).toBe(true);
|
||||
expect(states.some((state) => state.kind === 'wait_finished' && state.provider === 'anthropic-api')).toBe(true);
|
||||
});
|
||||
|
||||
it('rejects an interrupted wait without consuming a work slot', async () => {
|
||||
const clock = testClock();
|
||||
let abortListener: (() => void) | undefined;
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 1, waitStateTickMs: 100 }),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (_ms, signal) =>
|
||||
new Promise<void>((_resolve, reject) => {
|
||||
abortListener = () => reject(new DOMException('Aborted', 'AbortError'));
|
||||
signal?.addEventListener('abort', abortListener, { once: true });
|
||||
}),
|
||||
},
|
||||
);
|
||||
const controller = new AbortController();
|
||||
|
||||
governor.report({
|
||||
provider: 'claude-subscription',
|
||||
status: 'rejected',
|
||||
resetAtMs: 2_000,
|
||||
rateLimitType: 'five_hour',
|
||||
});
|
||||
const pending = governor.acquireWorkSlot(controller.signal);
|
||||
controller.abort();
|
||||
abortListener?.();
|
||||
|
||||
await expect(pending).rejects.toThrow(/Aborted/);
|
||||
expect(governor.activeSlots()).toBe(0);
|
||||
});
|
||||
|
||||
it('rejects an already-aborted ready wait', async () => {
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 1 }),
|
||||
{ sleep: async () => undefined, random: () => 0 },
|
||||
);
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
await expect(governor.waitForReady(controller.signal)).rejects.toThrow(/Aborted/);
|
||||
});
|
||||
|
||||
it('rejects an already-aborted work slot without consuming capacity', async () => {
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 1 }),
|
||||
{ sleep: async () => undefined, random: () => 0 },
|
||||
);
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
await expect(governor.acquireWorkSlot(controller.signal)).rejects.toThrow(/Aborted/);
|
||||
expect(governor.activeSlots()).toBe(0);
|
||||
});
|
||||
|
||||
it('uses bounded opaque backoff for rejected signals without reset hints', async () => {
|
||||
const clock = testClock();
|
||||
const sleeps: number[] = [];
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({
|
||||
maxConcurrency: 1,
|
||||
retry: { maxAttempts: 3, baseDelayMs: 1_000, maxDelayMs: 60_000, jitter: false },
|
||||
}),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (ms) => {
|
||||
sleeps.push(ms);
|
||||
clock.advance(ms);
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
governor.report({ provider: 'codex', status: 'rejected', rateLimitType: 'opaque' });
|
||||
const release1 = await governor.acquireWorkSlot();
|
||||
release1();
|
||||
governor.report({ provider: 'codex', status: 'rejected', rateLimitType: 'opaque' });
|
||||
const release2 = await governor.acquireWorkSlot();
|
||||
release2();
|
||||
|
||||
expect(sleeps).toEqual([1_000, 2_000]);
|
||||
});
|
||||
|
||||
it('exposes the configured retry budget and disables outer retries when pacing is off', () => {
|
||||
const retry = { maxAttempts: 3, baseDelayMs: 1_000, maxDelayMs: 60_000, jitter: false };
|
||||
const enabled = new RateLimitGovernor(createRateLimitGovernorConfig({ retry }));
|
||||
expect(enabled.maxRetryAttempts()).toBe(3);
|
||||
|
||||
const disabled = new RateLimitGovernor(createRateLimitGovernorConfig({ enabled: false, retry }));
|
||||
expect(disabled.maxRetryAttempts()).toBe(1);
|
||||
});
|
||||
|
||||
it('emits visible wait ticks after a rejected report without a waiting caller', async () => {
|
||||
const clock = testClock();
|
||||
const states: RateLimitWaitState[] = [];
|
||||
const sleeps: number[] = [];
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 4, minConcurrencyUnderPressure: 1, waitStateTickMs: 100 }),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (ms, signal) => {
|
||||
if (signal?.aborted) {
|
||||
throw new DOMException('Aborted', 'AbortError');
|
||||
}
|
||||
sleeps.push(ms);
|
||||
clock.advance(ms);
|
||||
},
|
||||
},
|
||||
);
|
||||
governor.subscribe((state) => states.push(state));
|
||||
|
||||
governor.report({
|
||||
provider: 'claude-subscription',
|
||||
status: 'rejected',
|
||||
resetAtMs: 1_250,
|
||||
rateLimitType: 'five_hour',
|
||||
});
|
||||
await flushMicrotasks();
|
||||
|
||||
expect(sleeps).toEqual([100, 100, 50]);
|
||||
expect(states).toContainEqual(
|
||||
expect.objectContaining({
|
||||
kind: 'wait_started',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
remainingMs: 250,
|
||||
}),
|
||||
);
|
||||
expect(states.filter((state) => state.kind === 'wait_tick')).toHaveLength(3);
|
||||
expect(states).toContainEqual(
|
||||
expect.objectContaining({
|
||||
kind: 'wait_finished',
|
||||
provider: 'claude-subscription',
|
||||
rateLimitType: 'five_hour',
|
||||
remainingMs: 0,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('does not duplicate countdown sleeps when a work slot waits during the same pause', async () => {
|
||||
const clock = testClock();
|
||||
const states: RateLimitWaitState[] = [];
|
||||
const sleeps: number[] = [];
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 2, waitStateTickMs: 100 }),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (ms, signal) => {
|
||||
if (signal?.aborted) {
|
||||
throw new DOMException('Aborted', 'AbortError');
|
||||
}
|
||||
sleeps.push(ms);
|
||||
clock.advance(ms);
|
||||
},
|
||||
},
|
||||
);
|
||||
governor.subscribe((state) => states.push(state));
|
||||
|
||||
governor.report({ provider: 'anthropic-api', status: 'rejected', retryAfterMs: 250, rateLimitType: 'rpm' });
|
||||
const pendingRelease = governor.acquireWorkSlot();
|
||||
await flushMicrotasks();
|
||||
const release = await pendingRelease;
|
||||
release();
|
||||
|
||||
expect(sleeps).toEqual([100, 100, 50]);
|
||||
expect(states.filter((state) => state.kind === 'wait_tick')).toHaveLength(3);
|
||||
expect(governor.activeSlots()).toBe(0);
|
||||
});
|
||||
|
||||
it('stops the visible wait ticker when the last subscriber unsubscribes', async () => {
|
||||
const clock = testClock();
|
||||
let abortCount = 0;
|
||||
const governor = new RateLimitGovernor(
|
||||
createRateLimitGovernorConfig({ maxConcurrency: 1, waitStateTickMs: 100 }),
|
||||
{
|
||||
now: clock.now,
|
||||
random: () => 0,
|
||||
sleep: async (_ms, signal) =>
|
||||
new Promise<void>((_resolve, reject) => {
|
||||
signal?.addEventListener(
|
||||
'abort',
|
||||
() => {
|
||||
abortCount += 1;
|
||||
reject(new DOMException('Aborted', 'AbortError'));
|
||||
},
|
||||
{ once: true },
|
||||
);
|
||||
}),
|
||||
},
|
||||
);
|
||||
const unsubscribe = governor.subscribe(() => undefined);
|
||||
|
||||
governor.report({ provider: 'claude-subscription', status: 'rejected', retryAfterMs: 1_000 });
|
||||
await flushMicrotasks(1);
|
||||
unsubscribe();
|
||||
await flushMicrotasks(1);
|
||||
|
||||
expect(abortCount).toBe(1);
|
||||
});
|
||||
});
|
||||
|
|
@ -50,6 +50,17 @@ connections:
|
|||
maxConcurrency: 1,
|
||||
failureMode: 'continue',
|
||||
},
|
||||
rateLimit: {
|
||||
enabled: true,
|
||||
throttleThreshold: 0.8,
|
||||
minConcurrencyUnderPressure: 1,
|
||||
retry: {
|
||||
maxAttempts: 6,
|
||||
baseDelayMs: 1_000,
|
||||
maxDelayMs: 60_000,
|
||||
jitter: true,
|
||||
},
|
||||
},
|
||||
profile: false,
|
||||
},
|
||||
agent: {
|
||||
|
|
@ -163,6 +174,52 @@ ingest:
|
|||
expect(parseKtxProjectConfig('ingest:\n profile: json\n').ingest.profile).toBe('json');
|
||||
});
|
||||
|
||||
it('defaults ingest rate-limit settings', () => {
|
||||
const config = buildDefaultKtxProjectConfig();
|
||||
expect(config.ingest.rateLimit).toEqual({
|
||||
enabled: true,
|
||||
throttleThreshold: 0.8,
|
||||
minConcurrencyUnderPressure: 1,
|
||||
retry: {
|
||||
maxAttempts: 6,
|
||||
baseDelayMs: 1_000,
|
||||
maxDelayMs: 60_000,
|
||||
jitter: true,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('validates ingest rate-limit retry settings', () => {
|
||||
const config = parseKtxProjectConfig(`
|
||||
llm:
|
||||
provider:
|
||||
backend: none
|
||||
ingest:
|
||||
rateLimit:
|
||||
enabled: true
|
||||
throttleThreshold: 0.7
|
||||
minConcurrencyUnderPressure: 2
|
||||
maxWaitMs: 300000
|
||||
retry:
|
||||
maxAttempts: 4
|
||||
baseDelayMs: 500
|
||||
maxDelayMs: 30000
|
||||
jitter: false
|
||||
`);
|
||||
expect(config.ingest.rateLimit).toEqual({
|
||||
enabled: true,
|
||||
throttleThreshold: 0.7,
|
||||
minConcurrencyUnderPressure: 2,
|
||||
maxWaitMs: 300_000,
|
||||
retry: {
|
||||
maxAttempts: 4,
|
||||
baseDelayMs: 500,
|
||||
maxDelayMs: 30_000,
|
||||
jitter: false,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('parses global Vertex LLM config', () => {
|
||||
const config = parseKtxProjectConfig(`
|
||||
llm:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,17 @@ describe('buildProjectStackSnapshotFields', () => {
|
|||
adapters: [],
|
||||
embeddings: { backend: 'sentence-transformers', dimensions: 384 },
|
||||
workUnits: { stepBudget: 40, maxConcurrency: 1, failureMode: 'continue' },
|
||||
rateLimit: {
|
||||
enabled: true,
|
||||
throttleThreshold: 0.8,
|
||||
minConcurrencyUnderPressure: 1,
|
||||
retry: {
|
||||
maxAttempts: 6,
|
||||
baseDelayMs: 1_000,
|
||||
maxDelayMs: 60_000,
|
||||
jitter: true,
|
||||
},
|
||||
},
|
||||
profile: false,
|
||||
},
|
||||
llm: { provider: { backend: 'none' }, models: {}, promptCaching: {} },
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue