@@ -13,12 +13,16 @@ export interface ProviderIdentity {
1313}
1414
1515/**
16- * A lane reserves from its own request and token buckets under the same
17- * provider identity, so a person waiting on one embedding never queues behind
18- * a bulk crawl's batches. Cooldown and quota gates stay per identity: a
19- * provider pause or an exhausted balance still stops every lane.
16+ * Every caller reserves from the credential's aggregate buckets, so the
17+ * configured budget is never exceeded. The bulk lane also reserves from a
18+ * bucket capped at {@link BULK_LANE_SHARE} of that budget, which leaves an
19+ * interactive caller headroom instead of a queue behind a crawl's batches.
20+ * Cooldown and quota gates stay per identity: a provider pause or an exhausted
21+ * balance still stops every lane.
2022 */
21- export type ProviderAdmissionLane = 'interactive'
23+ export type ProviderAdmissionLane = 'bulk' | 'interactive'
24+
25+ const BULK_LANE_SHARE = 0.9
2226
2327interface ProviderAdmissionInput extends ProviderIdentity {
2428 inputTokens ?: number
@@ -58,43 +62,48 @@ export async function waitForProviderAdmission(input: ProviderAdmissionInput): P
5862 input . signal ?. throwIfAborted ( )
5963 const deadlineAt = Date . now ( ) + input . maxWaitMs
6064 const key = providerKey ( input )
61- const bucketKey = input . lane ? `${ key } :${ input . lane } ` : key
6265 const requestsPerMinute =
6366 input . operation === 'embedding'
6467 ? envNumber ( env . KB_CONFIG_EMBEDDING_REQUESTS_PER_MINUTE , 600 , { min : 1 } )
6568 : input . operation === 'ocr'
6669 ? envNumber ( env . KB_CONFIG_OCR_REQUESTS_PER_MINUTE , 60 , { min : 1 } )
6770 : envNumber ( env . KB_CONFIG_RERANK_REQUESTS_PER_MINUTE , 60 , { min : 1 } )
71+ const tokenBudget =
72+ input . operation === 'embedding' && input . inputTokens
73+ ? {
74+ cost : input . inputTokens ,
75+ perMinute : envNumber ( env . KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE , 600_000 , { min : 1 } ) ,
76+ }
77+ : undefined
78+ if ( tokenBudget && tokenBudget . cost > tokenBudget . perMinute ) {
79+ throw new Error ( 'Embedding request exceeds the configured per-credential token budget' )
80+ }
81+ const requestBurst = Math . min (
82+ input . operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST ,
83+ requestsPerMinute
84+ )
6885 const reservations : TokenBucketReservation [ ] = [ ]
69- if ( input . operation === 'embedding' && input . inputTokens ) {
70- const tokensPerMinute = envNumber ( env . KB_CONFIG_EMBEDDING_TOKENS_PER_MINUTE , 600_000 , {
71- min : 1 ,
72- } )
73- if ( input . inputTokens > tokensPerMinute ) {
74- throw new Error ( 'Embedding request exceeds the configured per-credential token budget' )
86+ const reserveBuckets = ( bucketKey : string , share : number ) => {
87+ if ( tokenBudget ) {
88+ const maxTokens = Math . floor ( tokenBudget . perMinute * share )
89+ reservations . push ( {
90+ key : `${ bucketKey } :tokens` ,
91+ cost : tokenBudget . cost ,
92+ config : { maxTokens, refillRate : maxTokens / 60 , refillIntervalMs : 1000 } ,
93+ } )
7594 }
7695 reservations . push ( {
77- key : `${ bucketKey } :tokens ` ,
78- cost : input . inputTokens ,
96+ key : `${ bucketKey } :requests ` ,
97+ cost : 1 ,
7998 config : {
80- maxTokens : tokensPerMinute ,
81- refillRate : tokensPerMinute / 60 ,
99+ maxTokens : Math . floor ( requestBurst * share ) ,
100+ refillRate : ( requestsPerMinute * share ) / 60 ,
82101 refillIntervalMs : 1000 ,
83102 } ,
84103 } )
85104 }
86- reservations . push ( {
87- key : `${ bucketKey } :requests` ,
88- cost : 1 ,
89- config : {
90- maxTokens : Math . min (
91- input . operation === 'embedding' ? EMBEDDING_REQUEST_BURST : DEFAULT_REQUEST_BURST ,
92- requestsPerMinute
93- ) ,
94- refillRate : requestsPerMinute / 60 ,
95- refillIntervalMs : 1000 ,
96- } ,
97- } )
105+ reserveBuckets ( key , 1 )
106+ if ( input . lane === 'bulk' ) reserveBuckets ( `${ key } :bulk` , BULK_LANE_SHARE )
98107
99108 /** When the bucket last said capacity returns, so a deadline hit after a sleep reports the wait still left. */
100109 let capacityAvailableAt : number | undefined
0 commit comments