@tanstack/ai 0.6.3 → 0.8.0

package/src/activities/chat/middleware/compose.ts

@@ -0,0 +1,392 @@
+import { aiEventClient } from '@tanstack/ai-event-client'
+import type { StreamChunk } from '../../../types'
+import type {
+  AbortInfo,
+  AfterToolCallInfo,
+  BeforeToolCallDecision,
+  ChatMiddleware,
+  ChatMiddlewareConfig,
+  ChatMiddlewareContext,
+  ErrorInfo,
+  FinishInfo,
+  IterationInfo,
+  ToolCallHookContext,
+  ToolPhaseCompleteInfo,
+  UsageInfo,
+} from './types'
+
+/** Check if a middleware should be skipped for instrumentation events. */
+function shouldSkipInstrumentation(mw: ChatMiddleware): boolean {
+  return mw.name === 'devtools'
+}
+
+/** Build the base context for middleware instrumentation events. */
+function instrumentCtx(ctx: ChatMiddlewareContext) {
+  return {
+    requestId: ctx.requestId,
+    streamId: ctx.streamId,
+    clientId: ctx.conversationId,
+    timestamp: Date.now(),
+  }
+}
+
+/**
+ * Internal middleware runner that manages composed execution of middleware hooks.
+ * Created once per chat() invocation.
+ */
+export class MiddlewareRunner {
+  private readonly middlewares: ReadonlyArray<ChatMiddleware>
+
+  constructor(middlewares: ReadonlyArray<ChatMiddleware>) {
+    this.middlewares = middlewares
+  }
+
+  get hasMiddleware(): boolean {
+    return this.middlewares.length > 0
+  }
+
+  /**
+   * Pipe config through all middleware onConfig hooks in order.
+   * Each middleware receives the merged config from previous middleware.
+   * Partial returns are shallow-merged with the current config.
+   */
+  async runOnConfig(
+    ctx: ChatMiddlewareContext,
+    config: ChatMiddlewareConfig,
+  ): Promise<ChatMiddlewareConfig> {
+    let current = config
+    for (const mw of this.middlewares) {
+      if (mw.onConfig) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        const result = await mw.onConfig(ctx, current)
+        const hasTransform = result !== undefined && result !== null
+        if (hasTransform) {
+          current = { ...current, ...result }
+        }
+        if (!skip) {
+          const base = instrumentCtx(ctx)
+          aiEventClient.emit('middleware:hook:executed', {
+            ...base,
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onConfig',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform,
+          })
+          if (hasTransform) {
+            aiEventClient.emit('middleware:config:transformed', {
+              ...base,
+              middlewareName: mw.name || 'unnamed',
+              iteration: ctx.iteration,
+              changes: result as Record<string, unknown>,
+            })
+          }
+        }
+      }
+    }
+    return current
+  }
+
+  /**
+   * Call onStart on all middleware in order.
+   */
+  async runOnStart(ctx: ChatMiddlewareContext): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onStart) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onStart(ctx)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onStart',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Pipe a single chunk through all middleware onChunk hooks in order.
+   * Returns the resulting chunks (0..N) to yield to the consumer.
+   *
+   * - void: pass through unchanged
+   * - chunk: replace with this chunk
+   * - chunk[]: expand to multiple chunks
+   * - null: drop the chunk entirely
+   */
+  async runOnChunk(
+    ctx: ChatMiddlewareContext,
+    chunk: StreamChunk,
+  ): Promise<Array<StreamChunk>> {
+    let chunks: Array<StreamChunk> = [chunk]
+
+    for (const mw of this.middlewares) {
+      if (!mw.onChunk) continue
+      const skip = shouldSkipInstrumentation(mw)
+
+      const nextChunks: Array<StreamChunk> = []
+      for (const c of chunks) {
+        const result = await mw.onChunk(ctx, c)
+        if (result === null) {
+          // Drop this chunk
+          if (!skip) {
+            aiEventClient.emit('middleware:chunk:transformed', {
+              ...instrumentCtx(ctx),
+              middlewareName: mw.name || 'unnamed',
+              originalChunkType: c.type,
+              resultCount: 0,
+              wasDropped: true,
+            })
+          }
+          continue
+        } else if (result === undefined) {
+          // Pass through — no instrumentation for pass-throughs
+          nextChunks.push(c)
+        } else if (Array.isArray(result)) {
+          // Expand
+          nextChunks.push(...result)
+          if (!skip) {
+            aiEventClient.emit('middleware:chunk:transformed', {
+              ...instrumentCtx(ctx),
+              middlewareName: mw.name || 'unnamed',
+              originalChunkType: c.type,
+              resultCount: result.length,
+              wasDropped: false,
+            })
+          }
+        } else {
+          // Replace
+          nextChunks.push(result)
+          if (!skip) {
+            aiEventClient.emit('middleware:chunk:transformed', {
+              ...instrumentCtx(ctx),
+              middlewareName: mw.name || 'unnamed',
+              originalChunkType: c.type,
+              resultCount: 1,
+              wasDropped: false,
+            })
+          }
+        }
+      }
+      chunks = nextChunks
+    }
+
+    return chunks
+  }
+
+  /**
+   * Run onBeforeToolCall through middleware in order.
+   * Returns the first non-void decision, or undefined to continue normally.
+   */
+  async runOnBeforeToolCall(
+    ctx: ChatMiddlewareContext,
+    hookCtx: ToolCallHookContext,
+  ): Promise<BeforeToolCallDecision> {
+    for (const mw of this.middlewares) {
+      if (mw.onBeforeToolCall) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        const decision = await mw.onBeforeToolCall(ctx, hookCtx)
+        const hasTransform = decision !== undefined && decision !== null
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onBeforeToolCall',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform,
+          })
+        }
+        if (hasTransform) {
+          return decision
+        }
+      }
+    }
+    return undefined
+  }
+
+  /**
+   * Run onAfterToolCall on all middleware in order.
+   */
+  async runOnAfterToolCall(
+    ctx: ChatMiddlewareContext,
+    info: AfterToolCallInfo,
+  ): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onAfterToolCall) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onAfterToolCall(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onAfterToolCall',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onUsage on all middleware in order.
+   */
+  async runOnUsage(
+    ctx: ChatMiddlewareContext,
+    usage: UsageInfo,
+  ): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onUsage) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onUsage(ctx, usage)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onUsage',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onFinish on all middleware in order.
+   */
+  async runOnFinish(
+    ctx: ChatMiddlewareContext,
+    info: FinishInfo,
+  ): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onFinish) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onFinish(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onFinish',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onAbort on all middleware in order.
+   */
+  async runOnAbort(ctx: ChatMiddlewareContext, info: AbortInfo): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onAbort) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onAbort(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onAbort',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onError on all middleware in order.
+   */
+  async runOnError(ctx: ChatMiddlewareContext, info: ErrorInfo): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onError) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onError(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onError',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onIteration on all middleware in order.
+   * Called at the start of each agent loop iteration.
+   */
+  async runOnIteration(
+    ctx: ChatMiddlewareContext,
+    info: IterationInfo,
+  ): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onIteration) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onIteration(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onIteration',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+
+  /**
+   * Run onToolPhaseComplete on all middleware in order.
+   * Called after all tool calls in an iteration have been processed.
+   */
+  async runOnToolPhaseComplete(
+    ctx: ChatMiddlewareContext,
+    info: ToolPhaseCompleteInfo,
+  ): Promise<void> {
+    for (const mw of this.middlewares) {
+      if (mw.onToolPhaseComplete) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.onToolPhaseComplete(ctx, info)
+        if (!skip) {
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onToolPhaseComplete',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+  }
+}

package/src/activities/chat/middleware/index.ts

@@ -0,0 +1,17 @@
+export type {
+  ChatMiddleware,
+  ChatMiddlewareContext,
+  ChatMiddlewarePhase,
+  ChatMiddlewareConfig,
+  ToolCallHookContext,
+  BeforeToolCallDecision,
+  AfterToolCallInfo,
+  IterationInfo,
+  ToolPhaseCompleteInfo,
+  UsageInfo,
+  FinishInfo,
+  AbortInfo,
+  ErrorInfo,
+} from './types'
+
+export { MiddlewareRunner } from './compose'

package/src/activities/chat/middleware/tool-cache-middleware.ts

@@ -0,0 +1,189 @@
+import type { ChatMiddleware } from './types'
+
+/**
+ * A cache entry stored by the tool cache middleware.
+ */
+export interface ToolCacheEntry {
+  result: unknown
+  timestamp: number
+}
+
+/**
+ * Custom storage backend for the tool cache middleware.
+ *
+ * When provided, the middleware delegates all cache operations to this storage
+ * instead of using the built-in in-memory Map. This enables external storage
+ * backends like Redis, localStorage, databases, etc.
+ *
+ * All methods may return a Promise for async storage backends.
+ */
+export interface ToolCacheStorage {
+  getItem: (
+    key: string,
+  ) => ToolCacheEntry | undefined | Promise<ToolCacheEntry | undefined>
+  setItem: (key: string, value: ToolCacheEntry) => void | Promise<void>
+  deleteItem: (key: string) => void | Promise<void>
+}
+
+/**
+ * Options for the tool cache middleware.
+ */
+export interface ToolCacheMiddlewareOptions {
+  /**
+   * Maximum number of entries in the cache.
+   * When exceeded, the oldest entry is evicted (LRU).
+   *
+   * Only applies to the default in-memory storage.
+   * When a custom `storage` is provided, capacity management is the storage's responsibility.
+   *
+   * @default 100
+   */
+  maxSize?: number
+
+  /**
+   * Time-to-live in milliseconds. Entries older than this are not served from cache.
+   * @default Infinity (no expiry)
+   */
+  ttl?: number
+
+  /**
+   * Tool names to cache. If not provided, all tools are cached.
+   */
+  toolNames?: Array<string>
+
+  /**
+   * Custom function to generate a cache key from tool name and args.
+   * Defaults to `JSON.stringify([toolName, args])`.
+   */
+  keyFn?: (toolName: string, args: unknown) => string
+
+  /**
+   * Custom storage backend. When provided, the middleware uses this instead of
+   * the built-in in-memory Map. The storage is responsible for its own capacity
+   * management — the `maxSize` option is ignored.
+   *
+   * @example
+   * ```ts
+   * toolCacheMiddleware({
+   *   storage: {
+   *     getItem: (key) => redisClient.get(key).then(v => v ? JSON.parse(v) : undefined),
+   *     setItem: (key, value) => redisClient.set(key, JSON.stringify(value)),
+   *     deleteItem: (key) => redisClient.del(key),
+   *   },
+   * })
+   * ```
+   */
+  storage?: ToolCacheStorage
+}
+
+function defaultKeyFn(toolName: string, args: unknown): string {
+  return JSON.stringify([toolName, args])
+}
+
+function createDefaultStorage(maxSize: number): ToolCacheStorage {
+  const cache = new Map<string, ToolCacheEntry>()
+
+  return {
+    getItem: (key) => {
+      const entry = cache.get(key)
+      if (entry !== undefined) {
+        // Refresh recency: delete and re-insert so this key becomes newest
+        cache.delete(key)
+        cache.set(key, entry)
+      }
+      return entry
+    },
+    setItem: (key, value) => {
+      // Delete first so re-inserts also refresh recency
+      if (cache.has(key)) {
+        cache.delete(key)
+      } else if (cache.size >= maxSize) {
+        // LRU eviction: Map iteration order is insertion order — first key is least recently used
+        const firstKey = cache.keys().next().value
+        if (firstKey !== undefined) {
+          cache.delete(firstKey)
+        }
+      }
+      cache.set(key, value)
+    },
+    deleteItem: (key) => {
+      cache.delete(key)
+    },
+  }
+}
+
+/**
+ * Creates a middleware that caches tool call results based on tool name + arguments.
+ *
+ * When a tool is called with the same name and arguments as a previous call,
+ * the cached result is returned immediately without executing the tool.
+ *
+ * @example
+ * ```ts
+ * import { chat, toolCacheMiddleware } from '@tanstack/ai'
+ *
+ * const stream = chat({
+ *   adapter,
+ *   messages,
+ *   tools: [weatherTool, stockTool],
+ *   middleware: [
+ *     toolCacheMiddleware({ ttl: 60_000, toolNames: ['getWeather'] }),
+ *   ],
+ * })
+ * ```
+ */
+export function toolCacheMiddleware(
+  options: ToolCacheMiddlewareOptions = {},
+): ChatMiddleware {
+  const {
+    maxSize = 100,
+    ttl = Infinity,
+    toolNames,
+    keyFn = defaultKeyFn,
+    storage = createDefaultStorage(maxSize),
+  } = options
+
+  return {
+    name: 'tool-cache-middleware',
+
+    onBeforeToolCall: async (_ctx, hookCtx) => {
+      if (toolNames && !toolNames.includes(hookCtx.toolName)) {
+        return undefined
+      }
+
+      const key = keyFn(hookCtx.toolName, hookCtx.args)
+      const entry = await storage.getItem(key)
+
+      if (entry) {
+        const age = Date.now() - entry.timestamp
+        if (age < ttl) {
+          return { type: 'skip', result: entry.result }
+        }
+        // Expired — remove
+        await storage.deleteItem(key)
+      }
+
+      return undefined
+    },
+
+    onAfterToolCall: async (_ctx, info) => {
+      if (!info.ok) return
+      if (toolNames && !toolNames.includes(info.toolName)) return
+
+      // Re-derive the key from the raw arguments to match what onBeforeToolCall produces
+      let parsedArgs: unknown
+      try {
+        parsedArgs = JSON.parse(info.toolCall.function.arguments.trim() || '{}')
+      } catch {
+        return
+      }
+
+      const key = keyFn(info.toolName, parsedArgs)
+
+      await storage.setItem(key, {
+        result: info.result,
+        timestamp: Date.now(),
+      })
+    },
+  }
+}