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 2026-06-05 12:10:27 +02:00 committed by GitHub
parent 5a8821073b
commit c3d8cedb0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 2336 additions and 72 deletions

View file

@ -40,6 +40,7 @@ export interface CuratorPaginationInput {
buildToolSet: (passNumber: number) => KtxRuntimeToolSet;
getReconciliationActions: () => MemoryAction[];
onStepFinish?: (info: { passNumber: number; stepIndex: number; stepBudget: number }) => void;
abortSignal?: AbortSignal;
}
interface CuratorPaginationResult extends ReconciliationOutcome {
@ -243,6 +244,7 @@ export class CuratorPaginationService implements CuratorPaginationPort {
sourceKey: params.input.sourceKey,
jobId: params.input.jobId,
forceRun: params.forceRun,
abortSignal: params.input.abortSignal,
onStepFinish: params.input.onStepFinish
? ({ stepIndex, stepBudget }) =>
params.input.onStepFinish?.({ passNumber: params.passNumber, stepIndex, stepBudget })

View file

@ -21,6 +21,7 @@ export interface RepairFinalGateFailureInput {
repairKind: FinalGateRepairKind;
maxAttempts?: number;
stepBudget?: number;
abortSignal?: AbortSignal;
}
const readRepairFileSchema = z.object({
@ -200,6 +201,7 @@ export async function repairFinalGateFailure(
jobId: input.trace.context.jobId,
repairKind: input.repairKind,
},
abortSignal: input.abortSignal,
}),
);

View file

@ -3,6 +3,7 @@ import { dirname, join } from 'node:path';
import pLimit from 'p-limit';
import { z } from 'zod';
import { type KtxLogger, noopLogger } from '../../context/core/config.js';
import type { RateLimitWaitState } from '../../context/llm/rate-limit-governor.js';
import { createRuntimeToolDescriptorFromAiTool } from '../../context/llm/runtime-tools.js';
import type { KtxRuntimeToolSet } from '../../context/llm/runtime-port.js';
import type { CaptureSession, MemoryAction } from '../../context/memory/types.js';
@ -219,6 +220,10 @@ export class IngestBundleRunner {
}
async run(job: IngestBundleJob, ctx?: IngestJobContext): Promise<IngestBundleResult> {
const unsubscribeRateLimitGovernor = this.subscribeRateLimitGovernor({
trace: this.createTrace(job),
memoryFlow: ctx?.memoryFlow,
});
const key = job.connectionId;
const previous = this.chainByConnection.get(key);
if (previous) {
@ -241,10 +246,72 @@ export class IngestBundleRunner {
ctx?.memoryFlow?.finish('error', [sanitizeMemoryFlowError(error)]);
throw error;
} finally {
unsubscribeRateLimitGovernor();
await this.maybeEmitIngestProfile(job.jobId);
}
}
private formatRateLimitWait(
state: Extract<RateLimitWaitState, { kind: 'wait_tick' | 'wait_started' | 'wait_finished' }>,
): string {
const seconds = Math.ceil(state.remainingMs / 1_000);
const minutes = Math.floor(seconds / 60);
const remainder = seconds % 60;
const duration = minutes > 0 ? `${minutes}m${String(remainder).padStart(2, '0')}s` : `${seconds}s`;
const type = state.rateLimitType ? ` ${state.rateLimitType}` : '';
return `Rate-limited (${state.provider}${type}); resuming in ${duration}; Ctrl+C to stop`;
}
private subscribeRateLimitGovernor(input: {
trace: IngestTraceWriter;
memoryFlow?: MemoryFlowEventSink;
}): () => void {
const governor = this.deps.settings.rateLimitGovernor;
if (!governor) {
return () => undefined;
}
return governor.subscribe((state: RateLimitWaitState) => {
if (state.kind === 'rate_limit_observed') {
void input.trace.event('info', 'rate_limit', 'rate_limit_observed', { ...state });
return;
}
if (state.kind === 'concurrency_adjusted') {
void input.trace.event('info', 'rate_limit', 'concurrency_adjusted', { ...state });
return;
}
void input.trace.event('info', 'rate_limit', state.kind, { ...state });
if (state.kind === 'wait_tick' || state.kind === 'wait_started') {
input.memoryFlow?.emit({
type: 'rate_limit_wait',
provider: state.provider,
...(state.rateLimitType ? { rateLimitType: state.rateLimitType } : {}),
resumeAtMs: state.resumeAtMs,
remainingMs: state.remainingMs,
});
input.memoryFlow?.emit({
type: 'stage_progress',
stage: 'integration',
percent: 50,
message: this.formatRateLimitWait(state),
transient: true,
});
}
});
}
private async withRateLimitWorkSlot<T>(abortSignal: AbortSignal | undefined, fn: () => Promise<T>): Promise<T> {
const governor = this.deps.settings.rateLimitGovernor;
if (!governor) {
return fn();
}
const release = await governor.acquireWorkSlot(abortSignal);
try {
return await fn();
} finally {
release();
}
}
/**
* When profiling is enabled via the `KTX_PROFILE_INGEST` env var or the
* `ingest.profile` config setting read the job's trace + tool transcripts
@ -877,6 +944,7 @@ export class IngestBundleRunner {
includeContextEvidenceTools: boolean;
currentTableExists(tableRef: string): Promise<boolean>;
memoryFlow?: MemoryFlowEventSink;
abortSignal?: AbortSignal;
wuSkillNames: string[];
onStepFinish?: (info: { stepIndex: number; stepBudget: number }) => void;
}): Promise<WorkUnitOutcome> {
@ -1029,6 +1097,7 @@ export class IngestBundleRunner {
jobId: input.job.jobId,
toolFailureCount: (unitKey) => input.transcriptSummaries.get(unitKey)?.fatalErrorCount ?? 0,
onStepFinish: input.onStepFinish,
abortSignal: input.abortSignal,
},
input.wu,
);
@ -1524,7 +1593,8 @@ export class IngestBundleRunner {
try {
await Promise.all(
workUnits.map((wu, index) =>
limitWorkUnit(async () => {
limitWorkUnit(() =>
this.withRateLimitWorkSlot(ctx?.abortSignal, async () => {
const outcome = await runIsolatedWorkUnit({
unitIndex: index,
ingestionBaseSha,
@ -1532,6 +1602,7 @@ export class IngestBundleRunner {
patchDir,
trace: runTrace,
workUnit: wu,
abortSignal: ctx?.abortSignal,
afterSuccess: (child) => copyTransientIngestEvidence(child.workdir, sessionWorktree.workdir),
run: async (child) => {
const scopedWikiService = this.deps.wikiService.forWorktree(child.workdir);
@ -1565,6 +1636,7 @@ export class IngestBundleRunner {
includeContextEvidenceTools: adapter.evidenceIndexing === 'documents' && !!contextReport,
currentTableExists: (tableRef) =>
this.tableRefExistsInSemanticLayer(scopedSemanticLayerService, slConnectionIds, tableRef),
abortSignal: ctx?.abortSignal,
memoryFlow,
wuSkillNames,
onStepFinish: ({ stepIndex, stepBudget }) => {
@ -1594,7 +1666,8 @@ export class IngestBundleRunner {
completedWorkUnits / workUnits.length,
`${completedWorkUnits} of ${workUnits.length} work units complete`,
);
}),
}),
),
),
);
} catch (error) {
@ -1693,6 +1766,7 @@ export class IngestBundleRunner {
reason: context.reason,
maxAttempts: 1,
stepBudget: 12,
abortSignal: ctx?.abortSignal,
});
emitStageProgress(
'integration',
@ -1714,6 +1788,7 @@ export class IngestBundleRunner {
repairKind: 'patch_semantic_gate',
maxAttempts: 1,
stepBudget: 16,
abortSignal: ctx?.abortSignal,
});
emitStageProgress(
'integration',
@ -1993,6 +2068,7 @@ export class IngestBundleRunner {
);
}
: undefined,
abortSignal: ctx?.abortSignal,
});
curatorReport = curatorOutcome.report;
curatorWarnings = curatorOutcome.warnings;
@ -2038,6 +2114,7 @@ export class IngestBundleRunner {
sourceKey: job.sourceKey,
jobId: job.jobId,
force: !!overrideReport,
abortSignal: ctx?.abortSignal,
onStepFinish: stage4
? ({ stepIndex, stepBudget }) => {
emitStageProgress('reconciliation', 85, `Reconciling results: step ${stepIndex}/${stepBudget}`, {
@ -2470,6 +2547,7 @@ export class IngestBundleRunner {
repairKind: 'final_artifact_gate',
maxAttempts: 1,
stepBudget: 16,
abortSignal: ctx?.abortSignal,
});
isolatedDiffSummary.gateRepairAttempts += gateRepair.attempts;

View file

@ -19,6 +19,7 @@ export interface ResolveTextualConflictInput {
reason: string;
maxAttempts?: number;
stepBudget?: number;
abortSignal?: AbortSignal;
}
const readIntegrationFileSchema = z.object({
@ -208,6 +209,7 @@ export async function resolveTextualConflict(
jobId: input.trace.context.jobId,
unitKey: input.unitKey,
},
abortSignal: input.abortSignal,
}),
);

View file

@ -14,6 +14,7 @@ export interface RunIsolatedWorkUnitInput {
patchDir: string;
trace: IngestTraceWriter;
workUnit: WorkUnit;
abortSignal?: AbortSignal;
run(child: IngestSessionWorktree): Promise<WorkUnitOutcome>;
afterSuccess?(child: IngestSessionWorktree): Promise<void>;
}

View file

@ -12,6 +12,7 @@ import type { KtxSemanticLayerComputePort } from '../../context/daemon/semantic-
import { createRuntimeToolDescriptorFromAiTool } from '../../context/llm/runtime-tools.js';
import { createLocalKtxLlmRuntimeFromConfig } from '../../context/llm/local-config.js';
import { KtxIngestEmbeddingPortAdapter } from '../../context/llm/embedding-port.js';
import { createRateLimitGovernorConfig, RateLimitGovernor } from '../../context/llm/rate-limit-governor.js';
import { RuntimeAgentRunner, type AgentRunnerPort, type KtxLlmRuntimePort, type KtxRuntimeToolSet } from '../../context/llm/runtime-port.js';
import type { KtxEmbeddingProvider } from '../../llm/types.js';
import type { KtxLocalProject } from '../../context/project/project.js';
@ -619,7 +620,7 @@ function localIngestLlmProviderGuardMessage(projectDir: string): string {
].join('\n');
}
function resolveAgentRunner(options: CreateLocalBundleIngestRuntimeOptions): {
function resolveAgentRunner(options: CreateLocalBundleIngestRuntimeOptions, rateLimitGovernor: RateLimitGovernor): {
agentRunner: AgentRunnerPort;
llmRuntime?: KtxLlmRuntimePort;
} {
@ -628,6 +629,7 @@ function resolveAgentRunner(options: CreateLocalBundleIngestRuntimeOptions): {
(options.createLlmRuntime ?? createLocalKtxLlmRuntimeFromConfig)(options.project.config.llm, {
projectDir: options.project.projectDir,
env: process.env,
rateLimitGovernor,
}) ??
undefined;
@ -677,7 +679,13 @@ export function createLocalBundleIngestRuntime(
const knowledgeIndex = new LocalKnowledgeIndex(options.project, embedding);
const knowledgeEvents = new NoopKnowledgeEventPort();
const wikiService = new KnowledgeWikiService(rootFileStore, embedding, knowledgeIndex, options.project.git, logger);
const { agentRunner, llmRuntime } = resolveAgentRunner(options);
const rateLimitGovernor = new RateLimitGovernor(
createRateLimitGovernorConfig({
...options.project.config.ingest.rateLimit,
maxConcurrency: options.project.config.ingest.workUnits.maxConcurrency,
}),
);
const { agentRunner, llmRuntime } = resolveAgentRunner(options, rateLimitGovernor);
const promptService = new PromptService({ promptsDir, partials: [], logger });
const storage = new LocalIngestStorage(options.project);
const registry = registerAdapters(options.adapters);
@ -717,6 +725,7 @@ export function createLocalBundleIngestRuntime(
workUnitMaxConcurrency: options.project.config.ingest.workUnits.maxConcurrency,
workUnitStepBudget: options.project.config.ingest.workUnits.stepBudget,
workUnitFailureMode: options.project.config.ingest.workUnits.failureMode,
rateLimitGovernor,
profileIngest: options.project.config.ingest.profile,
ingestTraceLevel: ingestTraceLevelFromEnv(),
},

View file

@ -3,6 +3,7 @@ import { cp, mkdir, rm } from 'node:fs/promises';
import { isAbsolute, resolve } from 'node:path';
import type { KtxSqlQueryExecutorPort } from '../../context/connections/query-executor.js';
import type { KtxLogger } from '../../context/core/config.js';
import { createAbortError, isAbortError } from '../../context/core/abort.js';
import type { KtxSemanticLayerComputePort } from '../../context/daemon/semantic-layer-compute.js';
import type { AgentRunnerPort, KtxLlmRuntimePort } from '../../context/llm/runtime-port.js';
import type { KtxLocalProject } from '../../context/project/project.js';
@ -36,6 +37,7 @@ export interface RunLocalIngestOptions {
queryExecutor?: KtxSqlQueryExecutorPort;
logger?: KtxLogger;
embeddingProvider?: import('../../llm/types.js').KtxEmbeddingProvider | null;
abortSignal?: AbortSignal;
}
export interface LocalIngestResult {
@ -123,10 +125,11 @@ function findAdapter(adapters: SourceAdapter[], source: string): SourceAdapter {
return adapter;
}
function localJobContext(jobId: string, memoryFlow?: MemoryFlowEventSink): IngestJobContext {
function localJobContext(jobId: string, memoryFlow?: MemoryFlowEventSink, abortSignal?: AbortSignal): IngestJobContext {
return {
jobId,
...(memoryFlow ? { memoryFlow } : {}),
...(abortSignal ? { abortSignal } : {}),
startPhase() {
return new LocalIngestPhase();
},
@ -158,6 +161,7 @@ async function runScheduledPullJob(options: {
queryExecutor?: KtxSqlQueryExecutorPort;
logger?: KtxLogger;
embeddingProvider?: import('../../llm/types.js').KtxEmbeddingProvider | null;
abortSignal?: AbortSignal;
}): Promise<LocalIngestResult> {
const runtime = createLocalBundleIngestRuntime(options);
const jobId = options.jobId ?? runtime.nextJobId();
@ -169,7 +173,7 @@ async function runScheduledPullJob(options: {
trigger: options.trigger ?? 'manual_resync',
bundleRef: { kind: 'scheduled_pull', config: options.pullConfig },
},
localJobContext(jobId, options.memoryFlow),
localJobContext(jobId, options.memoryFlow, options.abortSignal),
);
const report = await runtime.store.findByJobId(jobId);
if (!report) {
@ -212,6 +216,7 @@ export async function runLocalIngest(options: RunLocalIngestOptions): Promise<Lo
queryExecutor: options.queryExecutor,
logger: options.logger,
embeddingProvider: options.embeddingProvider,
abortSignal: options.abortSignal,
});
}
@ -223,7 +228,7 @@ export async function runLocalIngest(options: RunLocalIngestOptions): Promise<Lo
trigger: options.trigger ?? (options.sourceDir ? 'upload' : 'manual_resync'),
bundleRef,
},
localJobContext(jobId, options.memoryFlow),
localJobContext(jobId, options.memoryFlow, options.abortSignal),
);
const report = await runtime.store.findByJobId(jobId);
if (!report) {
@ -362,6 +367,9 @@ export async function runLocalMetabaseIngest(
const children: LocalMetabaseFanoutChild[] = [];
for (const childPlan of childPlans) {
if (options.abortSignal?.aborted) {
throw createAbortError();
}
const targetConnectionId = safeSegment('target connection id', childPlan.targetConnectionId);
if (!options.project.config.connections[targetConnectionId]) {
throw new Error(`Target connection "${targetConnectionId}" is not configured in ktx.yaml`);
@ -391,8 +399,12 @@ export async function runLocalMetabaseIngest(
queryExecutor: options.queryExecutor,
logger: options.logger,
embeddingProvider: options.embeddingProvider,
abortSignal: options.abortSignal,
});
} catch (error) {
if (isAbortError(error)) {
throw error;
}
child = await recordLocalMetabaseChildFailure({
project: options.project,
jobId: childJobId,

View file

@ -70,6 +70,13 @@ const memoryFlowEventSchema = z.discriminatedUnion('type', [
message: z.string().min(1),
transient: z.boolean().optional(),
}),
eventSchema({
type: z.literal('rate_limit_wait'),
provider: z.string(),
rateLimitType: z.string().optional(),
resumeAtMs: z.number().int().nonnegative(),
remainingMs: z.number().int().nonnegative(),
}),
eventSchema({
type: z.literal('work_unit_started'),
unitKey: z.string().min(1),

View file

@ -60,6 +60,13 @@ type MemoryFlowEventPayload =
message: string;
transient?: boolean;
}
| {
type: 'rate_limit_wait';
provider: string;
rateLimitType?: string;
resumeAtMs: number;
remainingMs: number;
}
| {
type: 'work_unit_started';
unitKey: string;

View file

@ -5,6 +5,7 @@ import type { KtxFileStorePort } from '../../context/core/file-store.js';
import type { KtxLogger } from '../../context/core/config.js';
import type { SessionOutcome } from '../../context/core/session-worktree.service.js';
import type { AgentRunnerPort, KtxLlmRuntimePort, KtxRuntimeToolSet } from '../../context/llm/runtime-port.js';
import type { RateLimitGovernor } from '../llm/rate-limit-governor.js';
import type { MemoryAction, MemoryKnowledgeSlRefsPort } from '../../context/memory/types.js';
import type { PromptService } from '../../context/prompts/prompt.service.js';
import type { SkillsRegistryService } from '../../context/skills/skills-registry.service.js';
@ -144,6 +145,7 @@ interface IngestSettingsPort {
workUnitMaxConcurrency?: number;
workUnitStepBudget?: number;
workUnitFailureMode?: 'abort' | 'continue';
rateLimitGovernor?: RateLimitGovernor;
/** Print a timing breakdown to stderr at the end of each run (config-driven; see also KTX_PROFILE_INGEST). `'json'` emits the raw structured profile. */
profileIngest?: boolean | 'json';
ingestTraceLevel?: IngestTraceLevel;
@ -323,6 +325,7 @@ export interface CuratorPaginationPort {
buildToolSet: (passNumber: number) => KtxRuntimeToolSet;
getReconciliationActions: () => MemoryAction[];
onStepFinish?: (info: { passNumber: number; stepIndex: number; stepBudget: number }) => void;
abortSignal?: AbortSignal;
}): Promise<ReconciliationOutcome & { report: CuratorPaginationReport; warnings: string[] }>;
}

View file

@ -1,4 +1,5 @@
import type { KtxModelRole } from '../../../llm/types.js';
import { isAbortError } from '../../core/abort.js';
import type { AgentRunnerPort, KtxRuntimeToolSet, RunLoopMetrics } from '../../../context/llm/runtime-port.js';
import type { CaptureSession, MemoryAction } from '../../../context/memory/types.js';
import { listTouchedSlSources, type TouchedSlSource } from '../../../context/tools/touched-sl-sources.js';
@ -28,6 +29,7 @@ export interface WorkUnitExecutionDeps {
connectionId: string;
jobId: string;
onStepFinish?: (info: { stepIndex: number; stepBudget: number }) => void;
abortSignal?: AbortSignal;
toolFailureCount?: (unitKey: string) => number;
}
@ -106,8 +108,12 @@ export async function executeWorkUnit(deps: WorkUnitExecutionDeps, wu: WorkUnit)
jobId: deps.jobId,
},
onStepFinish: deps.onStepFinish,
abortSignal: deps.abortSignal,
});
} catch (error) {
if (isAbortError(error)) {
throw error;
}
return failWithResetFromCurrentHead(error instanceof Error ? error.message : String(error));
}

View file

@ -16,6 +16,7 @@ export interface ReconciliationContext {
jobId: string;
force?: boolean;
onStepFinish?: (info: { stepIndex: number; stepBudget: number }) => void;
abortSignal?: AbortSignal;
forceRun?: boolean;
}
@ -40,6 +41,7 @@ export async function runReconciliationStage4(ctx: ReconciliationContext): Promi
stepBudget: ctx.stepBudget,
telemetryTags: { operationName: 'ingest-bundle-reconcile', source: ctx.sourceKey, jobId: ctx.jobId },
onStepFinish: ctx.onStepFinish,
abortSignal: ctx.abortSignal,
});
return { skipped: false, stopReason: run.stopReason, error: run.error, ...(run.metrics ? { metrics: run.metrics } : {}) };
}

View file

@ -220,5 +220,6 @@ export interface IngestJobPhase {
export interface IngestJobContext {
jobId: string;
memoryFlow?: MemoryFlowEventSink;
abortSignal?: AbortSignal;
startPhase(weight: number): IngestJobPhase;
}