@tanstack/ai 0.14.0 → 0.16.0

package/src/middlewares/otel.ts

@@ -0,0 +1,861 @@
+import {
+  SpanKind,
+  SpanStatusCode,
+  context as otelContext,
+  trace as otelTrace,
+} from '@opentelemetry/api'
+import type {
+  AttributeValue,
+  Exception,
+  Meter,
+  Span,
+  SpanOptions,
+  Tracer,
+} from '@opentelemetry/api'
+import type {
+  ChatMiddleware,
+  ChatMiddlewareContext,
+} from '../activities/chat/middleware/types'
+
+/**
+ * Scope (role) of an OTel span emitted by this middleware.
+ *
+ * - `chat` — the root span for a single `chat()` call
+ * - `iteration` — one per agent-loop iteration (one model call)
+ * - `tool` — one per tool execution inside an iteration
+ */
+export type OtelSpanScope = 'chat' | 'iteration' | 'tool'
+
+/**
+ * Alias retained for backwards compatibility. Prefer {@link OtelSpanScope}.
+ *
+ * @deprecated Use `OtelSpanScope` instead — the name shadows OTel's built-in
+ * `SpanKind` which is also imported by integrations of this middleware.
+ */
+export type OtelSpanKind = OtelSpanScope
+
+/**
+ * Span metadata passed to `spanNameFormatter`, `attributeEnricher`,
+ * `onBeforeSpanStart`, and `onSpanEnd`. Discriminated by `kind` so that
+ * tool-only fields narrow automatically inside callback bodies.
+ */
+export type OtelSpanInfo<TScope extends OtelSpanScope = OtelSpanScope> =
+  TScope extends 'chat'
+    ? { kind: 'chat'; ctx: ChatMiddlewareContext }
+    : TScope extends 'iteration'
+      ? { kind: 'iteration'; ctx: ChatMiddlewareContext; iteration: number }
+      : TScope extends 'tool'
+        ? {
+            kind: 'tool'
+            ctx: ChatMiddlewareContext
+            iteration: number
+            toolName: string
+            toolCallId: string
+          }
+        : never
+
+export interface OtelMiddlewareOptions {
+  /** OTel `Tracer` used to start root, iteration, and tool spans. */
+  tracer: Tracer
+  /**
+   * Optional OTel `Meter`. When provided, the middleware records
+   * `gen_ai.client.operation.duration` and `gen_ai.client.token.usage`
+   * histograms. Omit to disable metrics without disabling tracing.
+   */
+  meter?: Meter
+  /**
+   * When `true`, prompt and completion content is attached to iteration spans
+   * as `gen_ai.*.message` / `gen_ai.choice` events. Defaults to `false` so
+   * that PII never lands on a span by accident.
+   */
+  captureContent?: boolean
+  /**
+   * Invoked on every captured content string before it lands on a span.
+   * Return a redacted version. If this function throws, the middleware emits
+   * the literal sentinel `"[redaction_failed]"` instead of the original text
+   * — it never falls back to raw content.
+   */
+  redact?: (text: string) => string
+  /**
+   * Maximum characters kept in the per-iteration assistant text buffer used
+   * to emit `gen_ai.choice` events. Extra characters are truncated with a
+   * trailing `"…"` marker. Defaults to 100 000. Set to `0` to disable the
+   * cap. Exporters typically truncate long attribute values anyway.
+   */
+  maxContentLength?: number
+  /** Override the default span name for each `kind`. */
+  spanNameFormatter?: (info: OtelSpanInfo) => string
+  /** Add extra attributes to each span. */
+  attributeEnricher?: (info: OtelSpanInfo) => Record<string, AttributeValue>
+  /** Mutate `SpanOptions` immediately before `tracer.startSpan(...)`. */
+  onBeforeSpanStart?: (info: OtelSpanInfo, options: SpanOptions) => SpanOptions
+  /** Fires just before every `span.end()`. */
+  onSpanEnd?: (info: OtelSpanInfo, span: Span) => void
+}
+
+interface RequestState {
+  rootSpan: Span
+  currentIterationSpan: Span | null
+  toolSpans: Map<string, { span: Span; toolName: string }>
+  iterationCount: number
+  assistantTextBuffer: string
+  assistantTextBufferTruncated: boolean
+  startTime: number
+}
+
+const stateByCtx = new WeakMap<ChatMiddlewareContext, RequestState>()
+
+const DEFAULT_MAX_CONTENT_LENGTH = 100_000
+const REDACTION_FAILED_SENTINEL = '[redaction_failed]'
+
+function serializeContent(content: unknown): string {
+  if (typeof content === 'string') return content
+  if (!Array.isArray(content)) return ''
+  const parts: Array<string> = []
+  for (const part of content) {
+    if (!part || typeof part !== 'object') continue
+    const type = (part as { type?: string }).type
+    switch (type) {
+      case 'text':
+        parts.push(
+          (
+            (part as { text?: string }).text ??
+            (part as { content?: string }).content ??
+            ''
+          ).toString(),
+        )
+        break
+      case 'image':
+        parts.push('[image]')
+        break
+      case 'audio':
+        parts.push('[audio]')
+        break
+      case 'video':
+        parts.push('[video]')
+        break
+      case 'document':
+        parts.push('[document]')
+        break
+      default:
+        parts.push(`[${type ?? 'unknown'}]`)
+    }
+  }
+  return parts.join(' ')
+}
+
+function messageEventName(role: string): string {
+  switch (role) {
+    case 'user':
+      return 'gen_ai.user.message'
+    case 'assistant':
+      return 'gen_ai.assistant.message'
+    case 'tool':
+      return 'gen_ai.tool.message'
+    case 'system':
+      return 'gen_ai.system.message'
+    default:
+      return `gen_ai.${role}.message`
+  }
+}
+
+function errorMessage(err: unknown): string | undefined {
+  if (err instanceof Error) return err.message
+  if (typeof err === 'string') return err
+  if (err && typeof err === 'object' && 'message' in err) {
+    const m = (err as { message?: unknown }).message
+    if (typeof m === 'string') return m
+  }
+  return undefined
+}
+
+function errorTypeName(err: unknown): string {
+  if (err instanceof Error) return err.name || 'Error'
+  if (err && typeof err === 'object' && 'name' in err) {
+    const n = (err as { name?: unknown }).name
+    if (typeof n === 'string') return n
+  }
+  return 'Error'
+}
+
+function safeCall<T>(label: string, fn: () => T): T | undefined {
+  try {
+    return fn()
+  } catch (err) {
+    // Keep middleware non-fatal, but surface callback failures so that broken
+    // extension points (attributeEnricher, spanNameFormatter, onSpanEnd, ...)
+    // are observable. Matches the guarantee documented in docs/advanced/otel.md.
+    console.warn(`[otelMiddleware] ${label} failed`, err)
+    return undefined
+  }
+}
+
+export function otelMiddleware(options: OtelMiddlewareOptions): ChatMiddleware {
+  const {
+    tracer,
+    meter,
+    captureContent = false,
+    redact = (s) => s,
+    maxContentLength = DEFAULT_MAX_CONTENT_LENGTH,
+    spanNameFormatter,
+    attributeEnricher,
+    onBeforeSpanStart,
+    onSpanEnd,
+  } = options
+
+  const durationHistogram = meter?.createHistogram(
+    'gen_ai.client.operation.duration',
+    {
+      description: 'GenAI client operation duration',
+      unit: 's',
+    },
+  )
+  const tokenHistogram = meter?.createHistogram('gen_ai.client.token.usage', {
+    description: 'GenAI client token usage',
+    unit: '{token}',
+  })
+
+  // Redact user content, failing closed to a sentinel string instead of ever
+  // letting raw text through. Callers that pass `captureContent: true` with a
+  // third-party PII redactor depend on this invariant.
+  const redactContent = (text: string): string => {
+    try {
+      return redact(text)
+    } catch (err) {
+      console.warn('[otelMiddleware] otel.redact failed', err)
+      return REDACTION_FAILED_SENTINEL
+    }
+  }
+
+  const appendAssistantText = (state: RequestState, delta: string): void => {
+    if (maxContentLength > 0) {
+      if (state.assistantTextBufferTruncated) return
+      const remaining = maxContentLength - state.assistantTextBuffer.length
+      if (remaining <= 0) {
+        state.assistantTextBufferTruncated = true
+        state.assistantTextBuffer += '…'
+        return
+      }
+      if (delta.length > remaining) {
+        state.assistantTextBuffer += delta.slice(0, remaining) + '…'
+        state.assistantTextBufferTruncated = true
+        return
+      }
+    }
+    state.assistantTextBuffer += delta
+  }
+
+  const closeIterationSpan = (
+    state: RequestState,
+    ctx: ChatMiddlewareContext,
+  ): void => {
+    if (!state.currentIterationSpan) return
+    const span = state.currentIterationSpan
+    const iteration = state.iterationCount - 1
+    safeCall('otel.onSpanEnd', () =>
+      onSpanEnd?.(
+        { kind: 'iteration', ctx, iteration } as OtelSpanInfo<'iteration'>,
+        span,
+      ),
+    )
+    span.end()
+    state.currentIterationSpan = null
+  }
+
+  return {
+    name: 'otel',
+
+    onStart(ctx) {
+      safeCall('otel.onStart', () => {
+        const info: OtelSpanInfo<'chat'> = { kind: 'chat', ctx }
+        const name =
+          safeCall('otel.spanNameFormatter', () => spanNameFormatter?.(info)) ??
+          `chat ${ctx.model}`
+        const baseOptions: SpanOptions = {
+          kind: SpanKind.INTERNAL,
+          attributes: {
+            'gen_ai.system': ctx.provider,
+            'gen_ai.request.model': ctx.model,
+            // NOTE: `gen_ai.operation.name` is deliberately NOT set on the
+            // root span. The root represents a `chat()` invocation that may
+            // span multiple model calls; only iteration spans correspond to
+            // a single chat operation. Backends that map `operation.name=chat`
+            // to a "generation" event (e.g. PostHog LLM Analytics) would
+            // otherwise emit a duplicate generation for the wrapper span.
+          },
+        }
+        const spanOptions =
+          safeCall('otel.onBeforeSpanStart', () =>
+            onBeforeSpanStart?.(info, baseOptions),
+          ) ?? baseOptions
+        const rootSpan = tracer.startSpan(name, spanOptions)
+
+        const enriched = safeCall('otel.attributeEnricher', () =>
+          attributeEnricher?.(info),
+        )
+        if (enriched) rootSpan.setAttributes(enriched)
+
+        stateByCtx.set(ctx, {
+          rootSpan,
+          currentIterationSpan: null,
+          toolSpans: new Map(),
+          iterationCount: 0,
+          assistantTextBuffer: '',
+          assistantTextBufferTruncated: false,
+          startTime: Date.now(),
+        })
+      })
+    },
+
+    onConfig(ctx, config) {
+      if (ctx.phase !== 'beforeModel') return
+      safeCall('otel.onConfig', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        // The previous iteration's span stays open through tool execution and
+        // onUsage so that tool spans nest under it and token attributes land
+        // on it. Close it here, just before opening the next iteration.
+        closeIterationSpan(state, ctx)
+
+        const info: OtelSpanInfo<'iteration'> = {
+          kind: 'iteration',
+          ctx,
+          iteration: ctx.iteration,
+        }
+        const name =
+          safeCall('otel.spanNameFormatter', () => spanNameFormatter?.(info)) ??
+          `chat ${ctx.model} #${ctx.iteration}`
+
+        const baseAttrs: Record<string, AttributeValue> = {
+          'gen_ai.system': ctx.provider,
+          'gen_ai.operation.name': 'chat',
+          'gen_ai.request.model': ctx.model,
+          'tanstack.ai.iteration': ctx.iteration,
+        }
+        if (config.temperature !== undefined)
+          baseAttrs['gen_ai.request.temperature'] = config.temperature
+        if (config.topP !== undefined)
+          baseAttrs['gen_ai.request.top_p'] = config.topP
+        if (config.maxTokens !== undefined)
+          baseAttrs['gen_ai.request.max_tokens'] = config.maxTokens
+
+        const baseOptions: SpanOptions = {
+          kind: SpanKind.CLIENT,
+          attributes: baseAttrs,
+        }
+        const spanOptions =
+          safeCall('otel.onBeforeSpanStart', () =>
+            onBeforeSpanStart?.(info, baseOptions),
+          ) ?? baseOptions
+
+        const parentCtx = otelTrace.setSpan(
+          otelContext.active(),
+          state.rootSpan,
+        )
+        let iterSpan!: Span
+        otelContext.with(parentCtx, () => {
+          // Pass the parent context explicitly as the 3rd arg — this is a
+          // real-OTel-compatible way to ensure the span is parented to
+          // `rootSpan` even when the host app has not registered a context
+          // manager (e.g. in tests or minimal setups).
+          iterSpan = tracer.startSpan(name, spanOptions, parentCtx)
+        })
+
+        const enriched = safeCall('otel.attributeEnricher', () =>
+          attributeEnricher?.(info),
+        )
+        if (enriched) iterSpan.setAttributes(enriched)
+
+        state.currentIterationSpan = iterSpan
+        state.assistantTextBuffer = ''
+        state.assistantTextBufferTruncated = false
+
+        if (captureContent) {
+          // Span events follow the original GenAI semconv (one event per
+          // message). Backends that read events get content this way.
+          for (const sys of config.systemPrompts) {
+            iterSpan.addEvent('gen_ai.system.message', {
+              content: redactContent(sys),
+            })
+          }
+          for (const m of config.messages) {
+            const body = serializeContent(m.content)
+            if (body.length === 0) continue
+            iterSpan.addEvent(messageEventName(m.role), {
+              content: redactContent(body),
+            })
+          }
+
+          // Also emit the current GenAI-semconv attribute form
+          // (`gen_ai.input.messages`) — backends like PostHog read prompt
+          // content from this attribute, not from span events.
+          const inputMessages: Array<{ role: string; content: string }> = []
+          for (const sys of config.systemPrompts) {
+            inputMessages.push({
+              role: 'system',
+              content: redactContent(sys),
+            })
+          }
+          for (const m of config.messages) {
+            const body = serializeContent(m.content)
+            if (body.length === 0) continue
+            inputMessages.push({
+              role: m.role,
+              content: redactContent(body),
+            })
+          }
+          if (inputMessages.length > 0) {
+            const inputJson = JSON.stringify(inputMessages)
+            // Current OTel GenAI semconv — Sentry / PostHog / Datadog read
+            // prompt content from this attribute.
+            iterSpan.setAttribute('gen_ai.input.messages', inputJson)
+            // Langfuse-native attribute. Highest priority in Langfuse's OTLP
+            // ingestion (checked before events and gen_ai.input.messages) so
+            // the Input panel populates reliably. Harmless to other backends —
+            // the attribute is namespaced and unrecognised keys are ignored.
+            iterSpan.setAttribute('langfuse.observation.input', inputJson)
+
+            // Mirror the first iteration's input onto the root span and at
+            // trace level so Langfuse fills Input on the trace card and the
+            // chat-level observation. Later iterations append tool-call /
+            // assistant messages that are useful per-iteration but noise at
+            // the chat / trace level.
+            if (state.iterationCount === 0) {
+              state.rootSpan.setAttribute(
+                'langfuse.observation.input',
+                inputJson,
+              )
+              state.rootSpan.setAttribute('langfuse.trace.input', inputJson)
+            }
+          }
+        }
+
+        state.iterationCount += 1
+      })
+      return undefined
+    },
+
+    onChunk(ctx, chunk) {
+      safeCall('otel.onChunk', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        if (captureContent && chunk.type === 'TEXT_MESSAGE_CONTENT') {
+          appendAssistantText(state, chunk.delta)
+        }
+
+        if (chunk.type !== 'RUN_FINISHED') return
+        const span = state.currentIterationSpan
+        if (!span) return
+
+        if (chunk.finishReason) {
+          span.setAttribute('gen_ai.response.finish_reasons', [
+            chunk.finishReason,
+          ])
+        }
+        if (chunk.model) span.setAttribute('gen_ai.response.model', chunk.model)
+
+        // Set usage attributes on the iteration span directly from the chunk
+        // so they're available before `onUsage` fires. Histogram recording is
+        // deliberately NOT done here — the chat runner always invokes
+        // `runOnUsage` when `chunk.usage` is present, and `onUsage` is the
+        // canonical place for the metric. Recording in both would double-count.
+        if (chunk.usage) {
+          span.setAttributes({
+            'gen_ai.usage.input_tokens': chunk.usage.promptTokens,
+            'gen_ai.usage.output_tokens': chunk.usage.completionTokens,
+          })
+        }
+
+        if (captureContent && state.assistantTextBuffer.length > 0) {
+          const completion = redactContent(state.assistantTextBuffer)
+          const outputJson = JSON.stringify([
+            { role: 'assistant', content: completion },
+          ])
+          // Event form (older semconv) — kept for backends that consume it.
+          span.addEvent('gen_ai.choice', { content: completion })
+          // Attribute form (current semconv) — required by backends like
+          // PostHog that read completion content from `gen_ai.output.messages`.
+          span.setAttribute('gen_ai.output.messages', outputJson)
+          // Langfuse-native attribute (highest priority in Langfuse mapping).
+          span.setAttribute('langfuse.observation.output', outputJson)
+          // Mirror to the root span and trace card. Each iteration overwrites,
+          // so the final iteration's completion lands on the root — which is
+          // the final answer the user saw, not an intermediate tool-call turn.
+          state.rootSpan.setAttribute('langfuse.observation.output', outputJson)
+          state.rootSpan.setAttribute('langfuse.trace.output', outputJson)
+          state.assistantTextBuffer = ''
+          state.assistantTextBufferTruncated = false
+        }
+
+        // Intentionally leave the iteration span open: tool spans started
+        // after `RUN_FINISHED` (tool_calls finishReason) must nest under it,
+        // and `onUsage` may still fire. The span is closed in `onConfig` when
+        // the next iteration starts, or in `onFinish` / `onError` / `onAbort`.
+      })
+      return undefined
+    },
+
+    onUsage(ctx, usage) {
+      safeCall('otel.onUsage', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        // Always record the token histogram — metrics don't depend on having
+        // an iteration span, and skipping here would drop metric data if an
+        // adapter emits `onUsage` outside the iteration window.
+        if (tokenHistogram) {
+          const metricAttrs = {
+            'gen_ai.system': ctx.provider,
+            'gen_ai.operation.name': 'chat',
+            'gen_ai.request.model': ctx.model,
+          }
+          tokenHistogram.record(usage.promptTokens, {
+            ...metricAttrs,
+            'gen_ai.token.type': 'input',
+          })
+          tokenHistogram.record(usage.completionTokens, {
+            ...metricAttrs,
+            'gen_ai.token.type': 'output',
+          })
+        }
+
+        const span = state.currentIterationSpan ?? state.rootSpan
+        span.setAttributes({
+          'gen_ai.usage.input_tokens': usage.promptTokens,
+          'gen_ai.usage.output_tokens': usage.completionTokens,
+        })
+      })
+    },
+
+    onBeforeToolCall(ctx, hookCtx) {
+      safeCall('otel.onBeforeToolCall', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+        const parent = state.currentIterationSpan ?? state.rootSpan
+
+        const info: OtelSpanInfo<'tool'> = {
+          kind: 'tool',
+          ctx,
+          toolName: hookCtx.toolName,
+          toolCallId: hookCtx.toolCallId,
+          iteration: state.iterationCount - 1,
+        }
+        const name =
+          safeCall('otel.spanNameFormatter', () => spanNameFormatter?.(info)) ??
+          `execute_tool ${hookCtx.toolName}`
+
+        const baseAttrs: Record<string, AttributeValue> = {
+          'gen_ai.tool.name': hookCtx.toolName,
+          'gen_ai.tool.call.id': hookCtx.toolCallId,
+          'gen_ai.tool.type': 'function',
+        }
+        const baseOptions: SpanOptions = {
+          kind: SpanKind.INTERNAL,
+          attributes: baseAttrs,
+        }
+        const spanOptions =
+          safeCall('otel.onBeforeSpanStart', () =>
+            onBeforeSpanStart?.(info, baseOptions),
+          ) ?? baseOptions
+
+        const parentCtx = otelTrace.setSpan(otelContext.active(), parent)
+        let toolSpan!: Span
+        otelContext.with(parentCtx, () => {
+          toolSpan = tracer.startSpan(name, spanOptions, parentCtx)
+        })
+
+        const enriched = safeCall('otel.attributeEnricher', () =>
+          attributeEnricher?.(info),
+        )
+        if (enriched) toolSpan.setAttributes(enriched)
+
+        // Stamp the tool args onto the tool span so backends that render an
+        // input panel per span (e.g. PostHog) have something to show.
+        if (captureContent) {
+          const argsBody =
+            typeof hookCtx.args === 'string'
+              ? hookCtx.args
+              : (safeCall('otel.serializeToolArgs', () =>
+                  JSON.stringify(hookCtx.args ?? null),
+                ) ?? '[unserializable_tool_args]')
+          const redactedArgs = redactContent(argsBody)
+          const toolInputJson = JSON.stringify([
+            { role: 'tool', content: redactedArgs },
+          ])
+          toolSpan.setAttribute('gen_ai.input.messages', toolInputJson)
+          // Langfuse-native (highest priority in Langfuse mapping).
+          toolSpan.setAttribute('langfuse.observation.input', toolInputJson)
+        }
+
+        state.toolSpans.set(hookCtx.toolCallId, {
+          span: toolSpan,
+          toolName: hookCtx.toolName,
+        })
+      })
+      return undefined
+    },
+
+    onAfterToolCall(ctx, info) {
+      safeCall('otel.onAfterToolCall', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+        const entry = state.toolSpans.get(info.toolCallId)
+        if (!entry) return
+        const { span: toolSpan } = entry
+
+        const outcome = info.ok ? 'success' : 'error'
+        toolSpan.setAttribute('tanstack.ai.tool.outcome', outcome)
+
+        if (!info.ok && info.error !== undefined) {
+          toolSpan.recordException(info.error as Exception)
+          toolSpan.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: errorMessage(info.error),
+          })
+        }
+
+        if (captureContent) {
+          // Serialization can throw on circular refs or `BigInt` values. If it
+          // does, fall back to a sentinel so the rest of this handler (span
+          // end, onSpanEnd, toolSpans cleanup) still runs — otherwise the tool
+          // span would dangle until the onFinish/onError sweep.
+          const body =
+            typeof info.result === 'string'
+              ? info.result
+              : (safeCall('otel.serializeToolResult', () =>
+                  JSON.stringify(info.result ?? null),
+                ) ?? '[unserializable_tool_result]')
+          const redactedBody = redactContent(body)
+          if (state.currentIterationSpan) {
+            state.currentIterationSpan.addEvent('gen_ai.tool.message', {
+              content: redactedBody,
+              tool_call_id: info.toolCallId,
+            })
+          }
+          // Output panel of the tool span itself — `gen_ai.output.messages` is
+          // what current GenAI semconv consumers (e.g. PostHog) read.
+          const toolOutputJson = JSON.stringify([
+            { role: 'tool', content: redactedBody },
+          ])
+          toolSpan.setAttribute('gen_ai.output.messages', toolOutputJson)
+          // Langfuse-native (highest priority in Langfuse mapping).
+          toolSpan.setAttribute('langfuse.observation.output', toolOutputJson)
+        }
+
+        safeCall('otel.onSpanEnd', () =>
+          onSpanEnd?.(
+            {
+              kind: 'tool',
+              ctx,
+              toolName: info.toolName,
+              toolCallId: info.toolCallId,
+              iteration: state.iterationCount - 1,
+            } as OtelSpanInfo<'tool'>,
+            toolSpan,
+          ),
+        )
+        toolSpan.end()
+        state.toolSpans.delete(info.toolCallId)
+      })
+    },
+
+    onError(ctx, info) {
+      safeCall('otel.onError', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        const errType = errorTypeName(info.error)
+        const message = errorMessage(info.error)
+        const exception = info.error as Exception
+
+        if (state.currentIterationSpan) {
+          state.currentIterationSpan.recordException(exception)
+          state.currentIterationSpan.setStatus({
+            code: SpanStatusCode.ERROR,
+            message,
+          })
+          safeCall('otel.onSpanEnd', () =>
+            onSpanEnd?.(
+              {
+                kind: 'iteration',
+                ctx,
+                iteration: state.iterationCount - 1,
+              } as OtelSpanInfo<'iteration'>,
+              state.currentIterationSpan!,
+            ),
+          )
+          state.currentIterationSpan.end()
+          state.currentIterationSpan = null
+        }
+
+        for (const [id, entry] of state.toolSpans) {
+          const { span, toolName } = entry
+          span.recordException(exception)
+          span.setStatus({ code: SpanStatusCode.ERROR, message })
+          safeCall('otel.onSpanEnd', () =>
+            onSpanEnd?.(
+              {
+                kind: 'tool',
+                ctx,
+                toolCallId: id,
+                toolName,
+                iteration: state.iterationCount - 1,
+              } as OtelSpanInfo<'tool'>,
+              span,
+            ),
+          )
+          span.end()
+          state.toolSpans.delete(id)
+        }
+
+        state.rootSpan.recordException(exception)
+        state.rootSpan.setStatus({ code: SpanStatusCode.ERROR, message })
+
+        if (durationHistogram) {
+          durationHistogram.record(info.duration / 1000, {
+            'gen_ai.system': ctx.provider,
+            'gen_ai.operation.name': 'chat',
+            'gen_ai.request.model': ctx.model,
+            'error.type': errType,
+          })
+        }
+
+        safeCall('otel.onSpanEnd', () =>
+          onSpanEnd?.({ kind: 'chat', ctx }, state.rootSpan),
+        )
+        state.rootSpan.end()
+        stateByCtx.delete(ctx)
+      })
+    },
+
+    onAbort(ctx, info) {
+      safeCall('otel.onAbort', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        const closeCancelled = (span: Span): void => {
+          // `gen_ai.completion.reason` is not part of the GenAI semconv; use a
+          // TanStack-namespaced attribute so downstream exporters don't treat
+          // it as standard. The span status still carries the error code.
+          span.setAttribute('tanstack.ai.completion.reason', 'cancelled')
+          span.setStatus({ code: SpanStatusCode.ERROR, message: 'cancelled' })
+        }
+
+        if (state.currentIterationSpan) {
+          closeCancelled(state.currentIterationSpan)
+          safeCall('otel.onSpanEnd', () =>
+            onSpanEnd?.(
+              {
+                kind: 'iteration',
+                ctx,
+                iteration: state.iterationCount - 1,
+              } as OtelSpanInfo<'iteration'>,
+              state.currentIterationSpan!,
+            ),
+          )
+          state.currentIterationSpan.end()
+          state.currentIterationSpan = null
+        }
+        for (const [id, entry] of state.toolSpans) {
+          const { span, toolName } = entry
+          closeCancelled(span)
+          safeCall('otel.onSpanEnd', () =>
+            onSpanEnd?.(
+              {
+                kind: 'tool',
+                ctx,
+                toolCallId: id,
+                toolName,
+                iteration: state.iterationCount - 1,
+              } as OtelSpanInfo<'tool'>,
+              span,
+            ),
+          )
+          span.end()
+          state.toolSpans.delete(id)
+        }
+        closeCancelled(state.rootSpan)
+
+        if (durationHistogram) {
+          durationHistogram.record(info.duration / 1000, {
+            'gen_ai.system': ctx.provider,
+            'gen_ai.operation.name': 'chat',
+            'gen_ai.request.model': ctx.model,
+            'error.type': 'cancelled',
+          })
+        }
+
+        safeCall('otel.onSpanEnd', () =>
+          onSpanEnd?.({ kind: 'chat', ctx }, state.rootSpan),
+        )
+        state.rootSpan.end()
+        stateByCtx.delete(ctx)
+      })
+    },
+
+    onFinish(ctx, info) {
+      safeCall('otel.onFinish', () => {
+        const state = stateByCtx.get(ctx)
+        if (!state) return
+
+        // Close any tool spans that never received `onAfterToolCall` (adapter
+        // quirk). Done before the iteration span so the hierarchy is closed
+        // in depth-first order.
+        for (const [id, entry] of state.toolSpans) {
+          const { span, toolName } = entry
+          span.setAttribute('tanstack.ai.tool.outcome', 'unknown')
+          safeCall('otel.onSpanEnd', () =>
+            onSpanEnd?.(
+              {
+                kind: 'tool',
+                ctx,
+                toolCallId: id,
+                toolName,
+                iteration: state.iterationCount - 1,
+              } as OtelSpanInfo<'tool'>,
+              span,
+            ),
+          )
+          span.end()
+          state.toolSpans.delete(id)
+        }
+
+        // The final iteration's span is still open because we keep it open
+        // through tool execution and `onUsage`. Close it now.
+        closeIterationSpan(state, ctx)
+
+        if (durationHistogram) {
+          durationHistogram.record(info.duration / 1000, {
+            'gen_ai.system': ctx.provider,
+            'gen_ai.operation.name': 'chat',
+            'gen_ai.request.model': ctx.model,
+          })
+        }
+
+        if (info.usage) {
+          state.rootSpan.setAttributes({
+            'gen_ai.usage.input_tokens': info.usage.promptTokens,
+            'gen_ai.usage.output_tokens': info.usage.completionTokens,
+          })
+        }
+        if (info.finishReason) {
+          state.rootSpan.setAttribute('gen_ai.response.finish_reasons', [
+            info.finishReason,
+          ])
+        }
+        state.rootSpan.setAttribute(
+          'tanstack.ai.iterations',
+          state.iterationCount,
+        )
+
+        safeCall('otel.onSpanEnd', () =>
+          onSpanEnd?.({ kind: 'chat', ctx }, state.rootSpan),
+        )
+        state.rootSpan.end()
+        stateByCtx.delete(ctx)
+      })
+    },
+  }
+}