@tanstack/ai 0.32.0 → 0.34.0

package/src/activities/generateVideo/index.ts

@@ -10,8 +10,17 @@
 import { aiEventClient } from '@tanstack/ai-event-client'
 import { toRunErrorPayload } from '../error-payload'
 import { resolveDebugOption } from '../../logger/resolve'
+import {
+  createGenerationContext,
+  runGenerationAbort,
+  runGenerationError,
+  runGenerationFinish,
+  runGenerationStart,
+  runGenerationUsage,
+} from '../middleware'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
+import type { GenerationMiddleware } from '../middleware'
 import type { VideoAdapter } from './adapter'
 import type {
   MediaPrompt,
@@ -160,6 +169,14 @@ export type VideoCreateOptions<
    * control and/or a custom `Logger`.
    */
   debug?: DebugOption
+  /**
+   * Observe-only middleware notified on start, usage, success, and error. Pass
+   * `otelMiddleware()` to emit OpenTelemetry spans, or implement the
+   * `GenerationMiddleware` contract for a custom backend. In streaming mode the
+   * span covers the full create→poll→complete lifecycle; in non-streaming mode
+   * it covers job submission. An abandoned stream fires `onAbort`.
+   */
+  middleware?: Array<GenerationMiddleware>
 } & ({} extends VideoProviderOptions<TAdapter>
     ? {
         /** Provider-specific options for video generation */ modelOptions?: VideoProviderOptions<TAdapter>
@@ -299,14 +316,27 @@ export function generateVideo<
 async function runCreateVideoJob<
   TAdapter extends VideoAdapter<string, any, any, any>,
 >(options: VideoCreateOptions<TAdapter, boolean>): Promise<VideoJobResult> {
-  const { adapter, prompt, size, duration, modelOptions } = options
+  const { adapter, prompt, size, duration, modelOptions, middleware } = options
   const model = adapter.model
+  const requestId = createId('video')
+  const startTime = Date.now()
   const logger: InternalLogger = resolveDebugOption(options.debug)
   const providerName =
     (adapter as { name?: string; provider?: string }).provider ??
     (adapter as { name?: string }).name ??
     'unknown'
 
+  const mwCtx = createGenerationContext({
+    requestId,
+    activity: 'video',
+    provider: adapter.name,
+    model,
+    modelOptions,
+    createId,
+  })
+
+  await runGenerationStart(middleware, mwCtx)
+
   logger.request(`activity=generateVideo provider=${providerName}`, {
     provider: providerName,
     model,
@@ -325,8 +355,17 @@ async function runCreateVideoJob<
       jobId: result.jobId,
       model: result.model,
     })
+    // Non-streaming create only submits the job; usage isn't known until the
+    // job completes via polling, so the span covers submission only.
+    await runGenerationFinish(middleware, mwCtx, {
+      duration: Date.now() - startTime,
+    })
     return result
   } catch (error) {
+    await runGenerationError(middleware, mwCtx, {
+      error,
+      duration: Date.now() - startTime,
+    })
     logger.errors('generateVideo activity failed', {
       error,
       source: 'generateVideo',
@@ -346,9 +385,11 @@ function sleep(ms: number): Promise<void> {
 async function* runStreamingVideoGeneration<
   TAdapter extends VideoAdapter<string, any, any, any>,
 >(options: VideoCreateOptions<TAdapter, true>): AsyncIterable<StreamChunk> {
-  const { adapter, prompt, size, duration, modelOptions } = options
+  const { adapter, prompt, size, duration, modelOptions, middleware } = options
   const model = adapter.model
   const runId = options.runId ?? createId('run')
+  const requestId = createId('video')
+  const obsStartTime = Date.now()
   const pollingInterval = options.pollingInterval ?? 2000
   const maxDuration = options.maxDuration ?? 600_000
   const logger: InternalLogger = resolveDebugOption(options.debug)
@@ -366,6 +407,17 @@ async function* runStreamingVideoGeneration<
     timestamp: Date.now(),
   } as StreamChunk
 
+  const mwCtx = createGenerationContext({
+    requestId,
+    activity: 'video',
+    provider: adapter.name,
+    model,
+    modelOptions,
+    createId,
+  })
+
+  await runGenerationStart(middleware, mwCtx)
+
   logger.request(
     `activity=generateVideo provider=${providerName} stream=true`,
     {
@@ -374,6 +426,9 @@ async function* runStreamingVideoGeneration<
     },
   )
 
+  // Tracks whether a terminal observer event (finish/error) has already fired,
+  // so the `finally` below can fire one on abandonment without double-firing.
+  let settled = false
   try {
     // Create the video generation job
     const jobResult = await adapter.createVideoJob({
@@ -422,6 +477,18 @@ async function* runStreamingVideoGeneration<
           },
         )
 
+        // Fire finish before yielding the terminal chunks: the generation has
+        // succeeded, so a consumer that stops reading after `generation:result`
+        // (without pulling `RUN_FINISHED`) must not trip the abandonment path in
+        // `finally`, which would otherwise report a spurious cancellation.
+        if (urlResult.usage)
+          await runGenerationUsage(middleware, mwCtx, urlResult.usage)
+        await runGenerationFinish(middleware, mwCtx, {
+          duration: Date.now() - obsStartTime,
+          usage: urlResult.usage,
+        })
+        settled = true
+
         yield {
           type: 'CUSTOM',
           name: 'generation:result',
@@ -453,6 +520,14 @@ async function* runStreamingVideoGeneration<
     throw new Error('Video generation timed out')
   } catch (error: unknown) {
     const payload = toRunErrorPayload(error, 'Video generation failed')
+    // Mark settled before firing onError: if a user error-hook throws, the
+    // `finally` below must still not double-fire onAbort over the same op
+    // (which would mask the original error and end the span twice).
+    settled = true
+    await runGenerationError(middleware, mwCtx, {
+      error,
+      duration: Date.now() - obsStartTime,
+    })
     logger.errors('generateVideo activity failed', {
       message: payload.message,
       code: payload.code,
@@ -467,6 +542,17 @@ async function* runStreamingVideoGeneration<
       error: payload,
       timestamp: Date.now(),
     } as StreamChunk
+  } finally {
+    if (!settled) {
+      // The consumer abandoned the stream (broke the `for await` loop or
+      // disconnected) before completion, so the generator is being unwound at
+      // a `yield` without reaching finish/error. Fire `onAbort` — a cancel, not
+      // an error — so otelMiddleware ends its span instead of leaking it.
+      await runGenerationAbort(middleware, mwCtx, {
+        reason: 'Video generation stream abandoned before completion',
+        duration: Date.now() - obsStartTime,
+      })
+    }
   }
 }
 

package/src/activities/middleware/index.ts

@@ -0,0 +1,20 @@
+// Base, activity-agnostic middleware shared by chat and the media activities.
+// The `ChatMiddleware` superset lives at `../chat/middleware`.
+export type {
+  GenerationActivity,
+  GenerationMiddleware,
+  GenerationMiddlewareContext,
+  GenerationUsageInfo,
+  GenerationFinishInfo,
+  GenerationAbortInfo,
+  GenerationErrorInfo,
+  AnyGenerationMiddleware,
+} from './types'
+export {
+  createGenerationContext,
+  runGenerationStart,
+  runGenerationUsage,
+  runGenerationFinish,
+  runGenerationAbort,
+  runGenerationError,
+} from './run'

package/src/activities/middleware/run.ts

@@ -0,0 +1,88 @@
+import type {
+  GenerationAbortInfo,
+  GenerationErrorInfo,
+  GenerationFinishInfo,
+  GenerationMiddleware,
+  GenerationMiddlewareContext,
+  GenerationUsageInfo,
+} from './types'
+
+/**
+ * Build the stable context for a single media-activity call.
+ *
+ * Media activities are always server-side and carry no user runtime context,
+ * so `source` is fixed to `'server'` and `context` to `undefined`.
+ */
+export function createGenerationContext(args: {
+  requestId: string
+  activity: GenerationMiddlewareContext['activity']
+  provider: string
+  model: string
+  modelOptions?: unknown
+  createId: (prefix: string) => string
+}): GenerationMiddlewareContext {
+  return {
+    requestId: args.requestId,
+    activity: args.activity,
+    provider: args.provider,
+    model: args.model,
+    modelOptions: args.modelOptions,
+    source: 'server',
+    createId: args.createId,
+    context: undefined,
+  }
+}
+
+/**
+ * Run a single lifecycle hook across each middleware in registration order,
+ * awaiting each. Exceptions PROPAGATE (matching `chat()` middleware) — a
+ * broken middleware fails the activity rather than being silently swallowed.
+ */
+async function run(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  invoke: (mw: GenerationMiddleware) => void | Promise<void>,
+): Promise<void> {
+  if (!middleware || middleware.length === 0) return
+  for (const mw of middleware) {
+    await invoke(mw)
+  }
+}
+
+export function runGenerationStart(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  ctx: GenerationMiddlewareContext,
+): Promise<void> {
+  return run(middleware, (mw) => mw.onStart?.(ctx))
+}
+
+export function runGenerationUsage(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  ctx: GenerationMiddlewareContext,
+  usage: GenerationUsageInfo,
+): Promise<void> {
+  return run(middleware, (mw) => mw.onUsage?.(ctx, usage))
+}
+
+export function runGenerationFinish(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  ctx: GenerationMiddlewareContext,
+  info: GenerationFinishInfo,
+): Promise<void> {
+  return run(middleware, (mw) => mw.onFinish?.(ctx, info))
+}
+
+export function runGenerationAbort(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  ctx: GenerationMiddlewareContext,
+  info: GenerationAbortInfo,
+): Promise<void> {
+  return run(middleware, (mw) => mw.onAbort?.(ctx, info))
+}
+
+export function runGenerationError(
+  middleware: ReadonlyArray<GenerationMiddleware> | undefined,
+  ctx: GenerationMiddlewareContext,
+  info: GenerationErrorInfo,
+): Promise<void> {
+  return run(middleware, (mw) => mw.onError?.(ctx, info))
+}

package/src/activities/middleware/types.ts

@@ -0,0 +1,173 @@
+import type { TokenUsage } from '../../types'
+
+// ===========================
+// Generation middleware
+// ===========================
+//
+// The base, activity-agnostic middleware contract. Every activity — chat and
+// the media activities — runs middleware that satisfies this shape. `chat()`
+// accepts the richer `ChatMiddleware` superset (it adds config/chunk/tool
+// hooks and capability primitives on top of these lifecycle hooks); media
+// activities accept `GenerationMiddleware` directly.
+//
+// The relationship is intentionally STRUCTURAL, not nominal: `ChatMiddleware`
+// does not `extends GenerationMiddleware`. Chat hooks use function-property
+// syntax, so under `strictFunctionTypes` a narrowed-context subtype would be
+// rejected; declaring it via inheritance would force method syntax and reopen
+// a bivariance hole (a chat hook reading `ctx.messages` slotted where only a
+// base context exists). Instead, the base context/info types are SUPERTYPES
+// (fewer fields) and the chat context/info types are SUBTYPES (more fields),
+// so a single value whose lifecycle hooks are authored against the base — like
+// `otelMiddleware()` — satisfies `GenerationMiddleware & ChatMiddleware` by
+// contravariance, while an arbitrary `ChatMiddleware` is NOT assignable to
+// `GenerationMiddleware`.
+
+/**
+ * The activity an observability event describes.
+ *
+ * Mirrors the public surface a caller reaches for: `'chat'` for `chat()`, and
+ * the media kinds for the `generate*` activities. `'tts'` matches the speech
+ * adapter's kind (the public discriminator avoids inventing a parallel
+ * `'speech'`/`'text'` vocabulary). `otelMiddleware` maps each to its
+ * `gen_ai.operation.name`.
+ */
+export type GenerationActivity =
+  | 'chat'
+  | 'image'
+  | 'video'
+  | 'audio'
+  | 'tts'
+  | 'transcription'
+
+/**
+ * Stable context passed to every {@link GenerationMiddleware} hook. Created
+ * once per activity call and shared across the hooks of that call.
+ *
+ * Carries only fields every activity can honor. `ChatMiddlewareContext`
+ * structurally includes all of these plus chat-only state (messages,
+ * iteration, capabilities, …), which is why a chat middleware that reads those
+ * extra fields is not assignable to `GenerationMiddleware`.
+ */
+export interface GenerationMiddlewareContext<TContext = unknown> {
+  /**
+   * Stable id correlating the `onStart` / `onFinish` / `onError` / `onAbort`
+   * hooks of a single activity call.
+   */
+  requestId: string
+  /** Which activity this call is. Discriminates media from chat. */
+  activity: GenerationActivity
+  /** Provider/adapter name (e.g. `"openai"`). Emitted as `gen_ai.system`. */
+  provider: string
+  /** Model id. Emitted as `gen_ai.request.model`. */
+  model: string
+  /**
+   * Provider-specific options passed to the activity, if any. Typed `unknown`
+   * because each activity's options are strongly typed per model; a supertype
+   * of `ChatMiddlewareContext`'s `modelOptions`.
+   */
+  modelOptions?: unknown
+  /** Where the call originates. Always `'server'` for media activities. */
+  source: 'client' | 'server'
+  /** Generate a unique id with the given prefix. */
+  createId: (prefix: string) => string
+  /** Runtime context provided by the activity options, if any. */
+  context: TContext
+}
+
+// ===========================
+// Hook payloads
+// ===========================
+
+/**
+ * Token usage passed to {@link GenerationMiddleware.onUsage}. Kept as an
+ * interface extending `TokenUsage` to preserve declaration merging for this
+ * publicly exported type.
+ */
+export interface GenerationUsageInfo extends TokenUsage {}
+
+/** Information passed to {@link GenerationMiddleware.onFinish}. */
+export interface GenerationFinishInfo {
+  /** Wall-clock duration of the activity call, in milliseconds. */
+  duration: number
+  /** Unified usage, when the provider reported it. */
+  usage?: TokenUsage | undefined
+}
+
+/** Information passed to {@link GenerationMiddleware.onAbort}. */
+export interface GenerationAbortInfo {
+  /** The reason for the abort, if provided. */
+  reason?: string
+  /** Wall-clock duration until the abort, in milliseconds. */
+  duration: number
+}
+
+/** Information passed to {@link GenerationMiddleware.onError}. */
+export interface GenerationErrorInfo {
+  /** The thrown value (typically an `Error`). */
+  error: unknown
+  /** Wall-clock duration until the failure, in milliseconds. */
+  duration: number
+}
+
+// ===========================
+// Middleware interface
+// ===========================
+
+/**
+ * Activity-agnostic, observe-only middleware.
+ *
+ * A thin lifecycle observer registerable on any activity via its `middleware`
+ * option. Unlike `ChatMiddleware` (which can also rewrite config, chunks, and
+ * tool calls), these hooks only observe — the right fit for the single
+ * request → response shape of media activities. Pass `otelMiddleware()` for
+ * OpenTelemetry, or implement the hooks directly for a custom backend.
+ *
+ * Hooks are awaited in registration order. A hook that throws PROPAGATES and
+ * fails the activity — matching `chat()` middleware semantics. Keep them cheap;
+ * they run inline with the request.
+ *
+ * Exactly one of `onFinish` / `onAbort` / `onError` fires per call.
+ *
+ * @example
+ * ```ts
+ * import { generateImage } from '@tanstack/ai'
+ * import { otelMiddleware } from '@tanstack/ai/middlewares/otel'
+ * import { openaiImage } from '@tanstack/ai-openai'
+ * import { trace } from '@opentelemetry/api'
+ *
+ * await generateImage({
+ *   adapter: openaiImage('gpt-image-1'),
+ *   prompt: 'A serene mountain landscape at sunset',
+ *   middleware: [otelMiddleware({ tracer: trace.getTracer('my-app') })],
+ * })
+ * ```
+ */
+export interface GenerationMiddleware<TContext = unknown> {
+  /** Optional name, surfaced in diagnostics. */
+  name?: string
+  /** Called before the adapter request begins. */
+  onStart?: (ctx: GenerationMiddlewareContext<TContext>) => void | Promise<void>
+  /** Called when the provider reports usage, before `onFinish`. */
+  onUsage?: (
+    ctx: GenerationMiddlewareContext<TContext>,
+    usage: GenerationUsageInfo,
+  ) => void | Promise<void>
+  /** Called after the activity completes successfully. */
+  onFinish?: (
+    ctx: GenerationMiddlewareContext<TContext>,
+    info: GenerationFinishInfo,
+  ) => void | Promise<void>
+  /** Called when the activity is aborted (e.g. an abandoned stream). */
+  onAbort?: (
+    ctx: GenerationMiddlewareContext<TContext>,
+    info: GenerationAbortInfo,
+  ) => void | Promise<void>
+  /** Called when the activity throws before completing. */
+  onError?: (
+    ctx: GenerationMiddlewareContext<TContext>,
+    info: GenerationErrorInfo,
+  ) => void | Promise<void>
+}
+
+/** A `GenerationMiddleware` with a permissive context — for use as a constraint. */
+export type AnyGenerationMiddleware = GenerationMiddleware<any>

package/src/index.ts

@@ -82,6 +82,10 @@ export {
 // Tool call management
 export { ToolCallManager } from './activities/chat/tools/tool-calls'
 
+// Lazy tool discovery (name of the synthetic discovery tool, for custom
+// message-compaction logic that needs to reference it)
+export { DISCOVERY_TOOL_NAME } from './activities/chat/tools/lazy-tool-manager'
+
 // Provider tool type
 export type { ProviderTool } from './tools/provider-tool'
 export { brandProviderTool } from './tools/provider-tool'
@@ -118,6 +122,21 @@ export type {
   ErrorInfo,
 } from './activities/chat/middleware/index'
 
+// Base, activity-agnostic middleware. The observe-only superset that media
+// activities accept via their `middleware` option; `ChatMiddleware` adds the
+// chat-only hooks on top. Pure types only — the `otelMiddleware` value lives at
+// `@tanstack/ai/middlewares/otel` so the root barrel never requires the
+// optional `@opentelemetry/api` peer dependency.
+export type {
+  GenerationMiddleware,
+  GenerationMiddlewareContext,
+  GenerationActivity,
+  GenerationUsageInfo,
+  GenerationFinishInfo,
+  GenerationAbortInfo,
+  GenerationErrorInfo,
+  AnyGenerationMiddleware,
+} from './activities/middleware/index'
 // Capability primitives + middleware builder
 export {
   createCapability,