@ai-sdk/otel 1.0.0-beta.12 → 1.0.0-beta.123

package/src/{open-telemetry-integration.ts → legacy-open-telemetry.ts}

@@ -1,38 +1,39 @@
-import { LanguageModelV4Prompt } from '@ai-sdk/provider';
+import type { LanguageModelV4Prompt } from '@ai-sdk/provider';
+import type { Context as AISDKContext } from '@ai-sdk/provider-utils';
 import {
-  Attributes,
-  AttributeValue,
   context,
-  Context,
-  Span,
   SpanStatusCode,
   trace,
-  Tracer,
+  type Attributes,
+  type AttributeValue,
+  type Context as OpenTelemetryContext,
+  type Span,
+  type Tracer,
 } from '@opentelemetry/api';
 import type {
-  EmbedFinishEvent,
-  EmbedOnFinishEvent,
-  EmbedOnStartEvent,
+  EmbeddingModelCallEndEvent,
+  EmbedEndEvent,
   EmbedStartEvent,
-  GenerationContext,
-  ObjectOnFinishEvent,
-  ObjectOnStartEvent,
-  ObjectOnStepFinishEvent,
-  ObjectOnStepStartEvent,
-  OnChunkEvent,
-  OnFinishEvent,
-  OnStartEvent,
-  OnStepFinishEvent,
-  OnStepStartEvent,
-  OnToolCallFinishEvent,
-  OnToolCallStartEvent,
-  OutputInterface as Output,
-  RerankFinishEvent,
-  RerankOnFinishEvent,
-  RerankOnStartEvent,
+  EmbeddingModelCallStartEvent,
+  GenerateObjectEndEvent,
+  GenerateObjectStartEvent,
+  GenerateObjectStepEndEvent,
+  GenerateObjectStepStartEvent,
+  GenerateTextAbortEvent,
+  GenerateTextEndEvent,
+  GenerateTextStartEvent,
+  GenerateTextStepEndEvent,
+  GenerateTextStepStartEvent,
+  ToolExecutionEndEvent,
+  ToolExecutionStartEvent,
+  Output,
+  RerankingModelCallEndEvent,
+  RerankEndEvent,
   RerankStartEvent,
-  TelemetryIntegration,
-  TelemetrySettings,
+  RerankingModelCallStartEvent,
+  InferTelemetryEvent,
+  Telemetry,
+  TelemetryOptions,
   ToolSet,
 } from 'ai';
 import { assembleOperationName } from './assemble-operation-name';
@@ -56,13 +57,13 @@ function recordSpanError(span: Span, error: unknown): void {
 }
 
 function shouldRecord(
-  telemetry: TelemetrySettings | undefined,
-): telemetry is TelemetrySettings {
-  return telemetry?.isEnabled === true;
+  telemetry: TelemetryOptions | undefined,
+): telemetry is TelemetryOptions {
+  return telemetry?.isEnabled !== false;
 }
 
 function selectAttributes(
-  telemetry: TelemetrySettings | undefined,
+  telemetry: TelemetryOptions | undefined,
   attributes: Record<
     string,
     | AttributeValue
@@ -110,9 +111,9 @@ function selectAttributes(
 
 interface OtelStepStartEvent<
   TOOLS extends ToolSet = ToolSet,
-  CONTEXT extends GenerationContext<TOOLS> = GenerationContext<TOOLS>,
-  OUTPUT extends Output = Output,
-> extends OnStepStartEvent<TOOLS, CONTEXT, OUTPUT> {
+  RUNTIME_CONTEXT extends AISDKContext = AISDKContext,
+  OUTPUT extends Output.Output = Output.Output,
+> extends GenerateTextStepStartEvent<TOOLS, RUNTIME_CONTEXT, OUTPUT> {
   readonly promptMessages?: LanguageModelV4Prompt;
   readonly stepTools?: ReadonlyArray<Record<string, unknown>>;
   readonly stepToolChoice?: unknown;
@@ -120,19 +121,19 @@ interface OtelStepStartEvent<
 
 interface CallState {
   operationId: string;
-  telemetry: TelemetrySettings | undefined;
+  telemetry: TelemetryOptions | undefined;
   rootSpan: Span | undefined;
-  rootContext: Context | undefined;
+  rootContext: OpenTelemetryContext | undefined;
   stepSpan: Span | undefined;
-  stepContext: Context | undefined;
-  embedSpans: Map<string, { span: Span; context: Context }>;
-  rerankSpan: { span: Span; context: Context } | undefined;
-  toolSpans: Map<string, { span: Span; context: Context }>;
+  stepContext: OpenTelemetryContext | undefined;
+  embedSpans: Map<string, { span: Span; context: OpenTelemetryContext }>;
+  rerankSpan: { span: Span; context: OpenTelemetryContext } | undefined;
+  toolSpans: Map<string, { span: Span; context: OpenTelemetryContext }>;
   baseTelemetryAttributes: Attributes;
   settings: Record<string, unknown>;
 }
 
-export class OpenTelemetryIntegration implements TelemetryIntegration {
+export class LegacyOpenTelemetry implements Telemetry {
   private readonly callStates = new Map<string, CallState>();
 
   /**
@@ -174,25 +175,45 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     return context.with(toolSpanEntry.context, execute);
   }
 
+  /**
+   * Runs the provider `doGenerate`/`doStream` call with the active legacy
+   * model-call context.
+   */
+  executeLanguageModelCall<T>({
+    callId,
+    execute,
+  }: {
+    callId: string;
+    execute: () => PromiseLike<T>;
+  }): PromiseLike<T> {
+    const stepContext = this.getCallState(callId)?.stepContext;
+
+    if (stepContext == null) {
+      return execute();
+    }
+
+    return context.with(stepContext, execute);
+  }
+
   onStart(
     event:
-      | OnStartEvent<ToolSet, Output>
-      | ObjectOnStartEvent
-      | EmbedOnStartEvent
-      | RerankOnStartEvent,
+      | InferTelemetryEvent<GenerateTextStartEvent>
+      | InferTelemetryEvent<GenerateObjectStartEvent>
+      | InferTelemetryEvent<EmbedStartEvent>
+      | InferTelemetryEvent<RerankStartEvent>,
   ): void {
-    if (event.isEnabled !== true) return;
-
     if (
       event.operationId === 'ai.embed' ||
       event.operationId === 'ai.embedMany'
     ) {
-      this.onEmbedOperationStart(event as EmbedOnStartEvent);
+      this.onEmbedOperationStart(event as InferTelemetryEvent<EmbedStartEvent>);
       return;
     }
 
     if (event.operationId === 'ai.rerank') {
-      this.onRerankOperationStart(event as RerankOnStartEvent);
+      this.onRerankOperationStart(
+        event as InferTelemetryEvent<RerankStartEvent>,
+      );
       return;
     }
 
@@ -200,20 +221,22 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
       event.operationId === 'ai.generateObject' ||
       event.operationId === 'ai.streamObject'
     ) {
-      this.onObjectOperationStart(event as ObjectOnStartEvent);
+      this.onObjectOperationStart(
+        event as InferTelemetryEvent<GenerateObjectStartEvent>,
+      );
       return;
     }
 
-    this.onGenerateStart(event as OnStartEvent<ToolSet, Output>);
+    this.onGenerateStart(event as InferTelemetryEvent<GenerateTextStartEvent>);
   }
 
-  private onGenerateStart(event: OnStartEvent<ToolSet, Output>): void {
-    const telemetry: TelemetrySettings = {
-      isEnabled: event.isEnabled,
+  private onGenerateStart(
+    event: InferTelemetryEvent<GenerateTextStartEvent>,
+  ): void {
+    const telemetry: TelemetryOptions = {
       recordInputs: event.recordInputs,
       recordOutputs: event.recordOutputs,
       functionId: event.functionId,
-      metadata: event.metadata,
     };
 
     const settings: Record<string, unknown> = {
@@ -230,9 +253,9 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
 
     const baseTelemetryAttributes = getBaseTelemetryAttributes({
       model: { provider: event.provider, modelId: event.modelId },
-      telemetry,
       headers: event.headers,
       settings,
+      context: event.runtimeContext as Record<string, unknown> | undefined,
     });
 
     const attributes = selectAttributes(telemetry, {
@@ -246,8 +269,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
       'ai.prompt': {
         input: () =>
           JSON.stringify({
-            system: event.system,
-            prompt: event.prompt,
+            system: event.instructions,
             messages: event.messages,
           }),
       },
@@ -271,13 +293,13 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     });
   }
 
-  private onObjectOperationStart(event: ObjectOnStartEvent): void {
-    const telemetry: TelemetrySettings = {
-      isEnabled: event.isEnabled,
+  private onObjectOperationStart(
+    event: InferTelemetryEvent<GenerateObjectStartEvent>,
+  ): void {
+    const telemetry: TelemetryOptions = {
       recordInputs: event.recordInputs,
       recordOutputs: event.recordOutputs,
       functionId: event.functionId,
-      metadata: event.metadata,
     };
 
     const settings: Record<string, unknown> = {
@@ -293,9 +315,9 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
 
     const baseTelemetryAttributes = getBaseTelemetryAttributes({
       model: { provider: event.provider, modelId: event.modelId },
-      telemetry,
       headers: event.headers,
       settings,
+      context: undefined,
     });
 
     const attributes = selectAttributes(telemetry, {
@@ -339,7 +361,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
   }
 
   /** @deprecated */
-  onObjectStepStart(event: ObjectOnStepStartEvent): void {
+  onObjectStepStart(event: GenerateObjectStepStartEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan || !state.rootContext) return;
 
@@ -390,7 +412,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
   }
 
   /** @deprecated */
-  onObjectStepFinish(event: ObjectOnStepFinishEvent): void {
+  onObjectStepEnd(event: GenerateObjectStepEndEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.stepSpan) return;
 
@@ -418,8 +440,10 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
         'ai.usage.inputTokens': event.usage.inputTokens,
         'ai.usage.outputTokens': event.usage.outputTokens,
         'ai.usage.totalTokens': event.usage.totalTokens,
-        'ai.usage.reasoningTokens': event.usage.reasoningTokens,
-        'ai.usage.cachedInputTokens': event.usage.cachedInputTokens,
+        'ai.usage.reasoningTokens':
+          event.usage.outputTokenDetails?.reasoningTokens,
+        'ai.usage.cachedInputTokens':
+          event.usage.inputTokenDetails?.cacheReadTokens,
 
         'gen_ai.response.finish_reasons': [event.finishReason],
         'gen_ai.response.id': event.response.id,
@@ -443,13 +467,13 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     state.stepContext = undefined;
   }
 
-  private onEmbedOperationStart(event: EmbedOnStartEvent): void {
-    const telemetry: TelemetrySettings = {
-      isEnabled: event.isEnabled,
+  private onEmbedOperationStart(
+    event: InferTelemetryEvent<EmbedStartEvent>,
+  ): void {
+    const telemetry: TelemetryOptions = {
       recordInputs: event.recordInputs,
       recordOutputs: event.recordOutputs,
       functionId: event.functionId,
-      metadata: event.metadata,
     };
 
     const settings: Record<string, unknown> = {
@@ -458,9 +482,9 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
 
     const baseTelemetryAttributes = getBaseTelemetryAttributes({
       model: { provider: event.provider, modelId: event.modelId },
-      telemetry,
       headers: event.headers,
       settings,
+      context: undefined,
     });
 
     const value = event.value;
@@ -503,7 +527,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     });
   }
 
-  onStepStart(event: OtelStepStartEvent<ToolSet, Output>): void {
+  onStepStart(event: OtelStepStartEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan || !state.rootContext) return;
 
@@ -568,7 +592,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     state.stepContext = trace.setSpan(state.rootContext, state.stepSpan);
   }
 
-  onToolCallStart(event: OnToolCallStartEvent<ToolSet>): void {
+  onToolExecutionStart(event: ToolExecutionStartEvent<ToolSet>): void {
     const state = this.getCallState(event.callId);
     if (!state?.stepContext) return;
 
@@ -600,7 +624,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     });
   }
 
-  onToolCallFinish(event: OnToolCallFinishEvent<ToolSet>): void {
+  onToolExecutionEnd(event: ToolExecutionEndEvent<ToolSet>): void {
     const state = this.getCallState(event.callId);
     if (!state) return;
 
@@ -610,31 +634,33 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     const { span } = toolSpanEntry;
     const { telemetry } = state;
 
-    if (event.success) {
+    const { toolOutput } = event;
+    if (toolOutput.type === 'tool-result') {
       try {
         span.setAttributes(
           selectAttributes(telemetry, {
             'ai.toolCall.result': {
-              output: () => JSON.stringify(event.output),
+              output: () => JSON.stringify(toolOutput.output),
             },
           }),
         );
-      } catch (_ignored) {
+      } catch {
         // JSON.stringify might fail for non-serializable results
       }
     } else {
-      recordSpanError(span, event.error);
+      recordSpanError(span, toolOutput.error);
     }
 
     span.end();
     state.toolSpans.delete(event.toolCall.toolCallId);
   }
 
-  onStepFinish(event: OnStepFinishEvent<ToolSet>): void {
+  onStepEnd(event: GenerateTextStepEndEvent<ToolSet>): void {
     const state = this.getCallState(event.callId);
     if (!state?.stepSpan) return;
 
     const { telemetry } = state;
+    const isStreamText = state.operationId === 'ai.streamText';
 
     state.stepSpan.setAttributes(
       selectAttributes(telemetry, {
@@ -681,12 +707,23 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
         'ai.response.providerMetadata': event.providerMetadata
           ? JSON.stringify(event.providerMetadata)
           : undefined,
+        'ai.response.msToFirstChunk': isStreamText
+          ? event.performance.timeToFirstOutputMs
+          : undefined,
+        'ai.response.msToFinish': isStreamText
+          ? event.performance.responseTimeMs
+          : undefined,
+        'ai.response.avgOutputTokensPerSecond': isStreamText
+          ? event.performance.effectiveOutputTokensPerSecond
+          : undefined,
 
         'ai.usage.inputTokens': event.usage.inputTokens,
         'ai.usage.outputTokens': event.usage.outputTokens,
         'ai.usage.totalTokens': event.usage.totalTokens,
-        'ai.usage.reasoningTokens': event.usage.reasoningTokens,
-        'ai.usage.cachedInputTokens': event.usage.cachedInputTokens,
+        'ai.usage.reasoningTokens':
+          event.usage.outputTokenDetails?.reasoningTokens,
+        'ai.usage.cachedInputTokens':
+          event.usage.inputTokenDetails?.cacheReadTokens,
         'ai.usage.inputTokenDetails.noCacheTokens':
           event.usage.inputTokenDetails?.noCacheTokens,
         'ai.usage.inputTokenDetails.cacheReadTokens':
@@ -706,17 +743,36 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
       }),
     );
 
+    if (isStreamText && event.performance.timeToFirstOutputMs != null) {
+      state.stepSpan.addEvent('ai.stream.firstChunk', {
+        'ai.response.msToFirstChunk': event.performance.timeToFirstOutputMs,
+      });
+    }
+
+    if (isStreamText) {
+      state.stepSpan.addEvent('ai.stream.finish', {
+        'ai.response.msToFinish': event.performance.responseTimeMs,
+        'ai.response.avgOutputTokensPerSecond':
+          event.performance.effectiveOutputTokensPerSecond,
+      });
+    }
+
     state.stepSpan.end();
     state.stepSpan = undefined;
     state.stepContext = undefined;
   }
 
-  onFinish(
+  /** @deprecated Use `onStepEnd` instead. */
+  onStepFinish(event: GenerateTextStepEndEvent<ToolSet>): void {
+    this.onStepEnd(event);
+  }
+
+  onEnd(
     event:
-      | OnFinishEvent<ToolSet>
-      | ObjectOnFinishEvent<unknown>
-      | EmbedOnFinishEvent
-      | RerankOnFinishEvent,
+      | GenerateTextEndEvent<ToolSet>
+      | GenerateObjectEndEvent<unknown>
+      | EmbedEndEvent
+      | RerankEndEvent,
   ): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan) return;
@@ -725,12 +781,12 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
       state.operationId === 'ai.embed' ||
       state.operationId === 'ai.embedMany'
     ) {
-      this.onEmbedOperationFinish(event as EmbedOnFinishEvent);
+      this.onEmbedOperationEnd(event as EmbedEndEvent);
       return;
     }
 
     if (state.operationId === 'ai.rerank') {
-      this.onRerankOperationFinish(event as RerankOnFinishEvent);
+      this.onRerankOperationEnd(event as RerankEndEvent);
       return;
     }
 
@@ -738,14 +794,14 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
       state.operationId === 'ai.generateObject' ||
       state.operationId === 'ai.streamObject'
     ) {
-      this.onObjectOperationFinish(event as ObjectOnFinishEvent<unknown>);
+      this.onObjectOperationEnd(event as GenerateObjectEndEvent<unknown>);
       return;
     }
 
-    this.onGenerateFinish(event as OnFinishEvent<ToolSet>);
+    this.onGenerateEnd(event as GenerateTextEndEvent<ToolSet>);
   }
 
-  private onGenerateFinish(event: OnFinishEvent<ToolSet>): void {
+  private onGenerateEnd(event: GenerateTextEndEvent<ToolSet>): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan) return;
 
@@ -759,8 +815,8 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
         },
         'ai.response.reasoning': {
           output: () =>
-            event.reasoning.length > 0
-              ? event.reasoning
+            event.finalStep.reasoning.length > 0
+              ? event.finalStep.reasoning
                   .filter(part => 'text' in part)
                   .map(part => part.text)
                   .join('\n')
@@ -790,25 +846,27 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
                 )
               : undefined,
         },
-        'ai.response.providerMetadata': event.providerMetadata
-          ? JSON.stringify(event.providerMetadata)
+        'ai.response.providerMetadata': event.finalStep.providerMetadata
+          ? JSON.stringify(event.finalStep.providerMetadata)
           : undefined,
 
-        'ai.usage.inputTokens': event.totalUsage.inputTokens,
-        'ai.usage.outputTokens': event.totalUsage.outputTokens,
-        'ai.usage.totalTokens': event.totalUsage.totalTokens,
-        'ai.usage.reasoningTokens': event.totalUsage.reasoningTokens,
-        'ai.usage.cachedInputTokens': event.totalUsage.cachedInputTokens,
+        'ai.usage.inputTokens': event.usage.inputTokens,
+        'ai.usage.outputTokens': event.usage.outputTokens,
+        'ai.usage.totalTokens': event.usage.totalTokens,
+        'ai.usage.reasoningTokens':
+          event.usage.outputTokenDetails?.reasoningTokens,
+        'ai.usage.cachedInputTokens':
+          event.usage.inputTokenDetails?.cacheReadTokens,
         'ai.usage.inputTokenDetails.noCacheTokens':
-          event.totalUsage.inputTokenDetails?.noCacheTokens,
+          event.usage.inputTokenDetails?.noCacheTokens,
         'ai.usage.inputTokenDetails.cacheReadTokens':
-          event.totalUsage.inputTokenDetails?.cacheReadTokens,
+          event.usage.inputTokenDetails?.cacheReadTokens,
         'ai.usage.inputTokenDetails.cacheWriteTokens':
-          event.totalUsage.inputTokenDetails?.cacheWriteTokens,
+          event.usage.inputTokenDetails?.cacheWriteTokens,
         'ai.usage.outputTokenDetails.textTokens':
-          event.totalUsage.outputTokenDetails?.textTokens,
+          event.usage.outputTokenDetails?.textTokens,
         'ai.usage.outputTokenDetails.reasoningTokens':
-          event.totalUsage.outputTokenDetails?.reasoningTokens,
+          event.usage.outputTokenDetails?.reasoningTokens,
       }),
     );
 
@@ -816,7 +874,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     this.cleanupCallState(event.callId);
   }
 
-  private onObjectOperationFinish(event: ObjectOnFinishEvent<unknown>): void {
+  private onObjectOperationEnd(event: GenerateObjectEndEvent<unknown>): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan) return;
 
@@ -836,8 +894,10 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
         'ai.usage.inputTokens': event.usage.inputTokens,
         'ai.usage.outputTokens': event.usage.outputTokens,
         'ai.usage.totalTokens': event.usage.totalTokens,
-        'ai.usage.reasoningTokens': event.usage.reasoningTokens,
-        'ai.usage.cachedInputTokens': event.usage.cachedInputTokens,
+        'ai.usage.reasoningTokens':
+          event.usage.outputTokenDetails?.reasoningTokens,
+        'ai.usage.cachedInputTokens':
+          event.usage.inputTokenDetails?.cacheReadTokens,
       }),
     );
 
@@ -845,7 +905,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     this.cleanupCallState(event.callId);
   }
 
-  private onEmbedOperationFinish(event: EmbedOnFinishEvent): void {
+  private onEmbedOperationEnd(event: EmbedEndEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan) return;
 
@@ -874,7 +934,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     this.cleanupCallState(event.callId);
   }
 
-  onEmbedStart(event: EmbedStartEvent): void {
+  onEmbedStart(event: EmbeddingModelCallStartEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan || !state.rootContext) return;
 
@@ -904,7 +964,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     });
   }
 
-  onEmbedFinish(event: EmbedFinishEvent): void {
+  onEmbedEnd(event: EmbeddingModelCallEndEvent): void {
     const state = this.getCallState(event.callId);
     if (!state) return;
 
@@ -928,13 +988,13 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     state.embedSpans.delete(event.embedCallId);
   }
 
-  private onRerankOperationStart(event: RerankOnStartEvent): void {
-    const telemetry: TelemetrySettings = {
-      isEnabled: event.isEnabled,
+  private onRerankOperationStart(
+    event: InferTelemetryEvent<RerankStartEvent>,
+  ): void {
+    const telemetry: TelemetryOptions = {
       recordInputs: event.recordInputs,
       recordOutputs: event.recordOutputs,
       functionId: event.functionId,
-      metadata: event.metadata,
     };
 
     const settings: Record<string, unknown> = {
@@ -943,9 +1003,9 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
 
     const baseTelemetryAttributes = getBaseTelemetryAttributes({
       model: { provider: event.provider, modelId: event.modelId },
-      telemetry,
       headers: event.headers,
       settings,
+      context: undefined,
     });
 
     const attributes = selectAttributes(telemetry, {
@@ -977,7 +1037,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     });
   }
 
-  private onRerankOperationFinish(event: RerankOnFinishEvent): void {
+  private onRerankOperationEnd(event: RerankEndEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan) return;
 
@@ -985,7 +1045,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     this.cleanupCallState(event.callId);
   }
 
-  onRerankStart(event: RerankStartEvent): void {
+  onRerankStart(event: RerankingModelCallStartEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rootSpan || !state.rootContext) return;
 
@@ -1012,7 +1072,7 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     state.rerankSpan = { span: rerankSpan, context: rerankContext };
   }
 
-  onRerankFinish(event: RerankFinishEvent): void {
+  onRerankEnd(event: RerankingModelCallEndEvent): void {
     const state = this.getCallState(event.callId);
     if (!state?.rerankSpan) return;
 
@@ -1032,37 +1092,33 @@ export class OpenTelemetryIntegration implements TelemetryIntegration {
     state.rerankSpan = undefined;
   }
 
-  onChunk(event: OnChunkEvent<ToolSet>): void {
-    const chunk = event.chunk as {
-      type: string;
-      callId?: unknown;
-      attributes?: unknown;
-    };
+  onAbort(event: GenerateTextAbortEvent<ToolSet>): void {
+    const state = this.getCallState(event.callId);
+    if (!state?.rootSpan) return;
 
-    if (typeof chunk.callId !== 'string') {
-      return;
+    for (const { span: toolSpan } of state.toolSpans.values()) {
+      toolSpan.end();
     }
+    state.toolSpans.clear();
 
-    if (
-      chunk.type !== 'ai.stream.firstChunk' &&
-      chunk.type !== 'ai.stream.finish'
-    ) {
-      return;
+    if (state.stepSpan) {
+      state.stepSpan.end();
+      state.stepSpan = undefined;
+      state.stepContext = undefined;
     }
 
-    const state = this.getCallState(chunk.callId);
-    if (!state?.stepSpan) return;
-
-    const attributes = Object.fromEntries(
-      Object.entries(
-        (chunk.attributes as Record<string, unknown>) ?? {},
-      ).filter(([, value]) => value != null),
-    ) as Attributes;
+    for (const { span: embedSpan } of state.embedSpans.values()) {
+      embedSpan.end();
+    }
+    state.embedSpans.clear();
 
-    state.stepSpan.addEvent(chunk.type, attributes);
-    if (Object.keys(attributes).length > 0) {
-      state.stepSpan.setAttributes(attributes);
+    if (state.rerankSpan) {
+      state.rerankSpan.span.end();
+      state.rerankSpan = undefined;
     }
+
+    state.rootSpan.end();
+    this.cleanupCallState(event.callId);
   }
 
   onError(error: unknown): void {