mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-16 08:25: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
|
|
@ -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) =>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue