ai-functions 2.1.3 → 2.4.0

package/src/batch-queue.ts

@@ -26,6 +26,8 @@
 
 import type { FunctionOptions } from './template.js'
 import type { SimpleSchema } from './schema.js'
+import { getLogger } from './logger.js'
+import { policyFor, type BatchTier } from 'language-models'
 
 // ============================================================================
 // Types
@@ -197,6 +199,11 @@ export interface BatchQueueOptions {
 // BatchQueue Implementation
 // ============================================================================
 
+/**
+ * Event handler type for BatchQueue events
+ */
+export type BatchEventHandler<T = unknown> = (data: T) => void
+
 /**
  * BatchQueue collects AI operations for deferred batch execution
  */
@@ -207,6 +214,11 @@ export class BatchQueue {
   private submitted = false
   private job: BatchJob | null = null
 
+  // Error handling properties
+  private _autoSubmitPromise: Promise<void> | null = null
+  private _submissionError: Error | null = null
+  private _eventHandlers: Map<string, Set<BatchEventHandler>> = new Map()
+
   constructor(options: BatchQueueOptions = {}) {
     this.options = {
       provider: 'openai',
@@ -216,6 +228,95 @@ export class BatchQueue {
     }
   }
 
+  /**
+   * Subscribe to batch events
+   * @param event - Event name ('error', 'partial-failure', 'complete')
+   * @param handler - Event handler function
+   */
+  on<T = unknown>(event: string, handler: BatchEventHandler<T>): void {
+    if (!this._eventHandlers.has(event)) {
+      this._eventHandlers.set(event, new Set())
+    }
+    this._eventHandlers.get(event)!.add(handler as BatchEventHandler)
+  }
+
+  /**
+   * Unsubscribe from batch events
+   */
+  off(event: string, handler: BatchEventHandler): void {
+    this._eventHandlers.get(event)?.delete(handler)
+  }
+
+  /**
+   * Emit an event to all subscribed handlers
+   */
+  private emit<T>(event: string, data: T): void {
+    this._eventHandlers.get(event)?.forEach((handler) => handler(data))
+  }
+
+  /**
+   * Get the auto-submit promise (if auto-submit was triggered)
+   */
+  get autoSubmitPromise(): Promise<void> | undefined {
+    return this._autoSubmitPromise ?? undefined
+  }
+
+  /**
+   * Get the last submission error (if any)
+   */
+  get submissionError(): Error | undefined {
+    return this._submissionError ?? undefined
+  }
+
+  /**
+   * Check if there was a submission error
+   */
+  get hasSubmissionError(): boolean {
+    return this._submissionError !== null
+  }
+
+  /**
+   * Await auto-submit completion or failure
+   * Returns a promise that resolves when auto-submit completes or rejects on error
+   */
+  awaitAutoSubmit(): Promise<void> {
+    if (!this._autoSubmitPromise) {
+      return Promise.resolve()
+    }
+    return this._autoSubmitPromise
+  }
+
+  /**
+   * Get items that failed during batch processing
+   */
+  getFailedItems(): BatchItem[] {
+    return this.items.filter((item) => item.status === 'failed')
+  }
+
+  /**
+   * Retry a failed submission
+   * Only available after a failed auto-submit
+   */
+  async retry(): Promise<BatchSubmitResult> {
+    if (!this._submissionError) {
+      throw new Error('No failed submission to retry')
+    }
+
+    // Reset error state and submitted flag
+    this._submissionError = null
+    this.submitted = false
+
+    // Reset item statuses
+    for (const item of this.items) {
+      if (item.status === 'failed') {
+        item.status = 'pending'
+        delete item.error
+      }
+    }
+
+    return this.submit()
+  }
+
   /**
    * Add an item to the batch queue
    * Returns a placeholder that will be resolved after batch completion
@@ -236,9 +337,9 @@ export class BatchQueue {
     const item: BatchItem<T> = {
       id: options?.customId || `item_${++this.idCounter}`,
       prompt,
-      schema: options?.schema,
-      options: options?.options,
-      metadata: options?.metadata,
+      ...(options?.schema !== undefined && { schema: options.schema }),
+      ...(options?.options !== undefined && { options: options.options }),
+      ...(options?.metadata !== undefined && { metadata: options.metadata }),
       status: 'pending',
     }
 
@@ -246,8 +347,50 @@ export class BatchQueue {
 
     // Auto-submit if we hit the limit
     if (this.options.autoSubmit && this.items.length >= (this.options.maxItems || 50000)) {
-      // Fire and forget - user can await the completion promise
-      this.submit().catch(console.error)
+      // Create a trackable promise for auto-submit
+      this._autoSubmitPromise = this.submit()
+        .then(() => {
+          // Success - promise is resolved
+        })
+        .catch((error: Error) => {
+          // Store error and update item statuses
+          this._submissionError = error
+          this.submitted = false // Reset to allow retry
+
+          // Update all items to failed status
+          for (const item of this.items) {
+            if (item.status === 'pending') {
+              item.status = 'failed'
+              item.error = error.message
+            }
+          }
+
+          // Create a synthetic failed job to store error info
+          const errorWithMeta = error as Error & {
+            status?: number
+            headers?: Record<string, string>
+          }
+          this.job = {
+            id: `failed_${Date.now()}`,
+            provider: this.options.provider || 'openai',
+            status: 'failed',
+            totalItems: this.items.length,
+            completedItems: 0,
+            failedItems: this.items.length,
+            createdAt: new Date(),
+            // Add rate limit retry info if available
+            ...(errorWithMeta.headers?.['retry-after'] && {
+              retryAfter: parseInt(errorWithMeta.headers['retry-after'], 10),
+            }),
+          } as BatchJob & { retryAfter?: number }
+
+          // Emit error event
+          this.emit('error', error)
+
+          // Log error and re-throw
+          getLogger().error('Batch auto-submit failed:', error)
+          throw error
+        })
     }
 
     return item
@@ -302,16 +445,35 @@ export class BatchQueue {
     const result = await adapter.submit(this.items, this.options)
     this.job = result.job
 
-    // When completion resolves, update item statuses
+    // When completion resolves, update item statuses and check for partial failures
     result.completion.then((results) => {
+      const failedResults: BatchResult[] = []
+
       for (const r of results) {
         const item = this.items.find((i) => i.id === r.id)
         if (item) {
           item.status = r.status
-          item.result = r.result
-          item.error = r.error
+          if (r.result !== undefined) item.result = r.result
+          if (r.error !== undefined) item.error = r.error
+
+          if (r.status === 'failed') {
+            failedResults.push(r)
+          }
         }
       }
+
+      // Emit partial-failure event if some items failed
+      if (failedResults.length > 0 && failedResults.length < results.length) {
+        this.emit('partial-failure', failedResults)
+      }
+
+      // Emit error event if there were any failures
+      if (failedResults.length > 0) {
+        this.emit('error', new Error(`${failedResults.length} items failed in batch`))
+      }
+
+      // Emit complete event
+      this.emit('complete', results)
     })
 
     return result
@@ -496,6 +658,33 @@ export function hasFlexAdapter(provider: BatchProvider): boolean {
   return flexAdapters[provider] !== null
 }
 
+// ============================================================================
+// Tier eligibility (per-model policy)
+// ============================================================================
+
+/**
+ * List the batch tiers a model is eligible for.
+ *
+ * Reads `ModelPolicy.batchTier` from `language-models` — this is the per-model
+ * policy data, distinct from the runtime adapter registration above.
+ *
+ * @example
+ * ```ts
+ * tiersForModel('sonnet')   // ['immediate', 'batch']
+ * tiersForModel('gpt-4o')   // ['immediate', 'flex', 'batch']
+ * ```
+ */
+export function tiersForModel(alias: string): BatchTier[] {
+  return policyFor(alias).batchTier
+}
+
+/**
+ * Check whether a model is eligible for a given tier.
+ */
+export function modelSupportsTier(alias: string, tier: BatchTier): boolean {
+  return tiersForModel(alias).includes(tier)
+}
+
 // ============================================================================
 // Factory Functions
 // ============================================================================
@@ -590,8 +779,8 @@ export function deferToBatch<T>(
   options?: FunctionOptions & { customId?: string }
 ): BatchItem<T> {
   return batch.add<T>(prompt, {
-    schema,
-    options,
-    customId: options?.customId,
+    ...(schema !== undefined && { schema }),
+    ...(options !== undefined && { options }),
+    ...(options?.customId !== undefined && { customId: options.customId }),
   })
 }

package/src/budget.ts

@@ -124,7 +124,7 @@ const DEFAULT_MODEL_PRICING: Record<string, ModelPricing> = {
   'gpt-4-turbo': { inputPricePerMillion: 10, outputPricePerMillion: 30 },
   'gpt-4': { inputPricePerMillion: 30, outputPricePerMillion: 60 },
   'gpt-3.5-turbo': { inputPricePerMillion: 0.5, outputPricePerMillion: 1.5 },
-  'o1': { inputPricePerMillion: 15, outputPricePerMillion: 60 },
+  o1: { inputPricePerMillion: 15, outputPricePerMillion: 60 },
   'o1-mini': { inputPricePerMillion: 3, outputPricePerMillion: 12 },
   'o1-preview': { inputPricePerMillion: 15, outputPricePerMillion: 60 },
   'o3-mini': { inputPricePerMillion: 1.1, outputPricePerMillion: 4.4 },
@@ -144,7 +144,7 @@ const DEFAULT_MODEL_PRICING: Record<string, ModelPricing> = {
   'gemini-1.5-flash': { inputPricePerMillion: 0.075, outputPricePerMillion: 0.3 },
 
   // Default fallback
-  'default': { inputPricePerMillion: 1, outputPricePerMillion: 3 },
+  default: { inputPricePerMillion: 1, outputPricePerMillion: 3 },
 }
 
 // ============================================================================
@@ -181,7 +181,7 @@ export class TokenCounter {
 
     // Rough estimate: ~4 chars per token for English
     // Unicode characters may use more tokens
-    const unicodeChars = Array.from(text).filter(char => char.charCodeAt(0) > 127).length
+    const unicodeChars = Array.from(text).filter((char) => char.charCodeAt(0) > 127).length
     const asciiChars = charCount - unicodeChars
 
     // ASCII chars: ~4 per token, Unicode: ~2 per token (rough)
@@ -240,7 +240,10 @@ export class BudgetExceededError extends Error {
 export class BudgetTracker {
   private totalInputTokens = 0
   private totalOutputTokens = 0
-  private usageByModel: Record<string, { inputTokens: number; outputTokens: number; cost: number }> = {}
+  private usageByModel: Record<
+    string,
+    { inputTokens: number; outputTokens: number; cost: number }
+  > = {}
   private triggeredThresholds: Set<number> = new Set()
   private requests: StoredRequest[] = []
 
@@ -347,7 +350,9 @@ export class BudgetTracker {
 
       if (projectedCost > this.config.maxCost) {
         throw new BudgetExceededError(
-          `Cost budget exceeded: $${projectedCost.toFixed(4)} would exceed limit of $${this.config.maxCost}`,
+          `Cost budget exceeded: $${projectedCost.toFixed(4)} would exceed limit of $${
+            this.config.maxCost
+          }`,
           'cost',
           this.config.maxCost,
           this.getTotalCost(),
@@ -365,7 +370,9 @@ export class BudgetTracker {
     if (this.config.maxTokens !== undefined) {
       if (this.getTotalTokens() > this.config.maxTokens) {
         throw new BudgetExceededError(
-          `Token budget exceeded: ${this.getTotalTokens()} tokens exceeds limit of ${this.config.maxTokens}`,
+          `Token budget exceeded: ${this.getTotalTokens()} tokens exceeds limit of ${
+            this.config.maxTokens
+          }`,
           'tokens',
           this.config.maxTokens,
           this.getTotalTokens()
@@ -378,7 +385,9 @@ export class BudgetTracker {
       const currentCost = this.getTotalCost()
       if (currentCost > this.config.maxCost) {
         throw new BudgetExceededError(
-          `Cost budget exceeded: $${currentCost.toFixed(4)} exceeds limit of $${this.config.maxCost}`,
+          `Cost budget exceeded: $${currentCost.toFixed(4)} exceeds limit of $${
+            this.config.maxCost
+          }`,
           'cost',
           this.config.maxCost,
           currentCost
@@ -580,11 +589,11 @@ export class RequestContext implements IRequestContext {
 
   constructor(options: RequestContextOptions & { depth?: number } = {}) {
     this.requestId = options.requestId ?? randomUUID()
-    this.userId = options.userId
-    this.tenantId = options.tenantId
-    this.parentRequestId = options.parentRequestId
+    if (options.userId !== undefined) this.userId = options.userId
+    if (options.tenantId !== undefined) this.tenantId = options.tenantId
+    if (options.parentRequestId !== undefined) this.parentRequestId = options.parentRequestId
     this.depth = (options as { depth?: number }).depth ?? 0
-    this.metadata = options.metadata
+    if (options.metadata !== undefined) this.metadata = options.metadata
 
     // Generate trace/span IDs for W3C traceparent
     this.traceId = randomUUID().replace(/-/g, '')
@@ -648,12 +657,13 @@ export class RequestContext implements IRequestContext {
    * Create a RequestContext from trace headers
    */
   static fromHeaders(headers: Record<string, string>): RequestContext {
-    return new RequestContext({
-      requestId: headers['x-request-id'],
-      userId: headers['x-user-id'],
-      tenantId: headers['x-tenant-id'],
-      parentRequestId: headers['x-parent-request-id'],
-    })
+    const opts: RequestContextOptions = {}
+    if (headers['x-request-id'] !== undefined) opts.requestId = headers['x-request-id']
+    if (headers['x-user-id'] !== undefined) opts.userId = headers['x-user-id']
+    if (headers['x-tenant-id'] !== undefined) opts.tenantId = headers['x-tenant-id']
+    if (headers['x-parent-request-id'] !== undefined)
+      opts.parentRequestId = headers['x-parent-request-id']
+    return new RequestContext(opts)
   }
 }
 
@@ -695,7 +705,10 @@ export async function withBudget<T>(
   const { userId, tenantId, ...budgetConfig } = options
 
   const tracker = new BudgetTracker(budgetConfig)
-  const ctx = userId || tenantId ? createRequestContext({ userId, tenantId }) : undefined
+  const ctxOptions: RequestContextOptions = {}
+  if (userId !== undefined) ctxOptions.userId = userId
+  if (tenantId !== undefined) ctxOptions.tenantId = tenantId
+  const ctx = userId || tenantId ? createRequestContext(ctxOptions) : undefined
 
   // Track parent tracker for nested contexts
   const parentTracker = currentBudgetTracker

package/src/cache.ts

@@ -95,6 +95,15 @@ export interface CacheStorage<T> {
 
 /**
  * In-memory cache implementation with TTL and LRU eviction support
+ *
+ * @deprecated Phase C Week 1 — `MemoryCache` has zero production callers in
+ * primitives.org.ai (audited 2026-05-06; see `bd show aip-ibid`). New code
+ * should use `cacheMiddleware` (for `wrapLanguageModel`) or
+ * `embeddingCacheMiddleware` (for `wrapEmbeddingModel`) instead — both
+ * compose with AI SDK 6's `wrapLanguageModel` / `wrapEmbeddingModel` and
+ * carry per-call telemetry (TraceEvent emission) and budget tracking when
+ * paired via `wrapForV3`. The `MemoryCache` class will be removed in the
+ * Phase C semver bump alongside `EmbeddingCache` and `GenerationCache`.
  */
 export class MemoryCache<T> implements CacheStorage<T> {
   private cache: Map<string, CacheEntry<T>> = new Map()
@@ -157,7 +166,7 @@ export class MemoryCache<T> implements CacheStorage<T> {
       createdAt: now,
       lastAccessedAt: now,
       accessCount: 0,
-      expiresAt: ttl ? now + ttl : undefined
+      ...(ttl ? { expiresAt: now + ttl } : {}),
     }
 
     // Check if we need to evict (LRU)
@@ -247,7 +256,7 @@ export class MemoryCache<T> implements CacheStorage<T> {
   dispose(): void {
     if (this.cleanupTimer) {
       clearInterval(this.cleanupTimer)
-      this.cleanupTimer = undefined
+      delete this.cleanupTimer
     }
   }
 
@@ -311,14 +320,12 @@ export class MemoryCache<T> implements CacheStorage<T> {
  * Uses a fast, non-cryptographic hash suitable for cache keys
  */
 export function hashKey(input: unknown): string {
-  const str = typeof input === 'string'
-    ? input
-    : stableStringify(input)
+  const str = typeof input === 'string' ? input : stableStringify(input)
 
   // Simple djb2 hash
   let hash = 5381
   for (let i = 0; i < str.length; i++) {
-    hash = ((hash << 5) + hash) + str.charCodeAt(i)
+    hash = (hash << 5) + hash + str.charCodeAt(i)
     hash = hash & hash // Convert to 32bit integer
   }
 
@@ -338,7 +345,7 @@ function stableStringify(obj: unknown): string {
   }
 
   const keys = Object.keys(obj).sort()
-  const pairs = keys.map(key => {
+  const pairs = keys.map((key) => {
     return JSON.stringify(key) + ':' + stableStringify((obj as Record<string, unknown>)[key])
   })
 
@@ -382,6 +389,14 @@ export interface BatchEmbeddingResult {
 
 /**
  * Specialized cache for embedding vectors
+ *
+ * @deprecated Phase C Week 1 — `EmbeddingCache` has zero production callers
+ * in primitives.org.ai (audited 2026-05-06; see `bd show aip-ibid`). New
+ * code should use `embeddingCacheMiddleware` (for `wrapEmbeddingModel`) —
+ * it composes with AI SDK 6 directly and carries trace + budget telemetry
+ * when paired via `wrapForV3`. Note: `embeddingCacheMiddleware` keys on the
+ * whole batch, not per-text — callers wanting per-text caching should pass
+ * stable per-text batches. Will be removed in the Phase C semver bump.
  */
 export class EmbeddingCache {
   private storage: MemoryCache<number[]>
@@ -418,7 +433,11 @@ export class EmbeddingCache {
   /**
    * Set multiple embeddings at once
    */
-  async setMany(texts: string[], embeddings: number[][], options: EmbeddingCacheOptions): Promise<void> {
+  async setMany(
+    texts: string[],
+    embeddings: number[][],
+    options: EmbeddingCacheOptions
+  ): Promise<void> {
     for (let i = 0; i < texts.length; i++) {
       await this.set(texts[i]!, embeddings[i]!, options)
     }
@@ -459,7 +478,7 @@ export class EmbeddingCache {
       hits: this.stats.hits,
       misses: this.stats.misses,
       hitRate: total > 0 ? this.stats.hits / total : 0,
-      size: this.storage['cache'].size
+      size: this.storage['cache'].size,
     }
   }
 
@@ -502,6 +521,12 @@ export interface GenerationCacheGetOptions {
 
 /**
  * Specialized cache for generation results
+ *
+ * @deprecated Phase C Week 1 — `GenerationCache` has zero production callers
+ * in primitives.org.ai (audited 2026-05-06; see `bd show aip-ibid`). New
+ * code should use `cacheMiddleware` (for `wrapLanguageModel`) — it composes
+ * with AI SDK 6 directly and carries trace + budget telemetry when paired
+ * via `wrapForV3`. Will be removed in the Phase C semver bump.
  */
 export class GenerationCache {
   private storage: MemoryCache<unknown>
@@ -514,7 +539,10 @@ export class GenerationCache {
   /**
    * Get a cached generation result
    */
-  async get<T = unknown>(params: GenerationParams, options?: GenerationCacheGetOptions): Promise<T | undefined> {
+  async get<T = unknown>(
+    params: GenerationParams,
+    options?: GenerationCacheGetOptions
+  ): Promise<T | undefined> {
     if (options?.bypass) {
       return undefined
     }
@@ -548,7 +576,7 @@ export class GenerationCache {
       hits: this.stats.hits,
       misses: this.stats.misses,
       hitRate: total > 0 ? this.stats.hits / total : 0,
-      size: this.storage['cache'].size
+      size: this.storage['cache'].size,
     }
   }
 
@@ -566,19 +594,19 @@ export class GenerationCache {
   private createKey(params: GenerationParams): string {
     const keyParams: Record<string, unknown> = {
       prompt: params.prompt,
-      model: params.model
+      model: params.model,
     }
 
     if (params.system !== undefined) {
-      keyParams.system = params.system
+      keyParams['system'] = params.system
     }
 
     if (params.temperature !== undefined) {
-      keyParams.temperature = params.temperature
+      keyParams['temperature'] = params.temperature
     }
 
     if (params.schemaVersion !== undefined) {
-      keyParams.schemaVersion = params.schemaVersion
+      keyParams['schemaVersion'] = params.schemaVersion
     }
 
     return createCacheKey('generation', keyParams)
@@ -631,7 +659,7 @@ export function withCache<TArgs extends unknown[], TResult>(
     const result = await fn(...args)
 
     // Cache result (don't cache errors - they throw before reaching here)
-    await cache.set(key, result, { ttl })
+    await cache.set(key, result, ttl !== undefined ? { ttl } : undefined)
 
     return result
   }
@@ -644,7 +672,7 @@ export function withCache<TArgs extends unknown[], TResult>(
     const result = await fn(...args)
 
     // Update cache with new result
-    await cache.set(key, result, { ttl })
+    await cache.set(key, result, ttl !== undefined ? { ttl } : undefined)
 
     return result
   }