Clean up OTel instrumentation for AskAi (#2144)

* Clean up OTel instrumentation

* Add comment

* Fix naming

* Run prettier

* Fix duplicate spans and other adjustments

* Fix setting conversationId

* Change naming from threadId to conversationId

* formatting

* Fix message rendering

* Simplify streaming message

* Run prettier

* Move StreamTransformerBase to core package

* Combine aiProviderStore and chat.store

* Fix tests

Author: reakaleek

Directory.Packages.props

@@ -25,7 +25,7 @@
     <PackageVersion Include="AWSSDK.SQS" Version="4.0.2" />
     <PackageVersion Include="AWSSDK.S3" Version="4.0.7.14" />
     <PackageVersion Include="Elastic.OpenTelemetry" Version="1.1.0" />
-    <PackageVersion Include="OpenTelemetry.Instrumentation.AspNetCore" Version="1.12.0" />
+    <PackageVersion Include="OpenTelemetry.Instrumentation.AspNetCore" Version="1.13.0" />
     <PackageVersion Include="KubernetesClient" Version="17.0.14" />
     <PackageVersion Include="Elastic.Aspire.Hosting.Elasticsearch" Version="9.3.0" />
     <PackageVersion Include="Elastic.Clients.Elasticsearch" Version="9.1.4" />
@@ -76,10 +76,10 @@
   <ItemGroup>
     <PackageVersion Include="Microsoft.Extensions.Http.Resilience" Version="9.7.0" />
     <PackageVersion Include="Microsoft.Extensions.ServiceDiscovery" Version="9.4.0" />
-    <PackageVersion Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" Version="1.12.0" />
-    <PackageVersion Include="OpenTelemetry.Extensions.Hosting" Version="1.12.0" />
-    <PackageVersion Include="OpenTelemetry.Instrumentation.Http" Version="1.12.0" />
-    <PackageVersion Include="OpenTelemetry.Instrumentation.Runtime" Version="1.12.0" />
+    <PackageVersion Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" Version="1.13.1" />
+    <PackageVersion Include="OpenTelemetry.Extensions.Hosting" Version="1.13.1" />
+    <PackageVersion Include="OpenTelemetry.Instrumentation.Http" Version="1.13.0" />
+    <PackageVersion Include="OpenTelemetry.Instrumentation.Runtime" Version="1.13.0" />
   </ItemGroup>
   <!-- Test packages -->
   <ItemGroup>
@@ -99,4 +99,4 @@
     </PackageVersion>
     <PackageVersion Include="xunit.v3" Version="2.0.2" />
   </ItemGroup>
-</Project>
+</Project>

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/AiProviderSelector.tsx

@@ -1,5 +1,5 @@
 /** @jsxImportSource @emotion/react */
-import { useAiProviderStore } from './aiProviderStore'
+import { useChatActions, useAiProvider, type AiProvider } from './chat.store'
 import { EuiRadioGroup } from '@elastic/eui'
 import type { EuiRadioGroupOption } from '@elastic/eui'
 import { css } from '@emotion/react'
@@ -22,16 +22,15 @@ const options: EuiRadioGroupOption[] = [
 ]
 
 export const AiProviderSelector = () => {
-    const { provider, setProvider } = useAiProviderStore()
+    const provider = useAiProvider()
+    const { setAiProvider } = useChatActions()
 
     return (
         <div css={containerStyles}>
             <EuiRadioGroup
                 options={options}
                 idSelected={provider}
-                onChange={(id) =>
-                    setProvider(id as 'AgentBuilder' | 'LlmGateway')
-                }
+                onChange={(id) => setAiProvider(id as AiProvider)}
                 name="aiProvider"
                 legend={{
                     children: 'AI Provider',

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/AskAiEvent.ts

@@ -4,8 +4,8 @@ import * as z from 'zod'
 // Event type constants for type-safe referencing
 export const EventTypes = {
     CONVERSATION_START: 'conversation_start',
-    CHUNK: 'chunk',
-    CHUNK_COMPLETE: 'chunk_complete',
+    MESSAGE_CHUNK: 'message_chunk',
+    MESSAGE_COMPLETE: 'message_complete',
     SEARCH_TOOL_CALL: 'search_tool_call',
     TOOL_CALL: 'tool_call',
     TOOL_RESULT: 'tool_result',
@@ -23,14 +23,14 @@ export const ConversationStartEventSchema = z.object({
 })
 
 export const ChunkEventSchema = z.object({
-    type: z.literal(EventTypes.CHUNK),
+    type: z.literal(EventTypes.MESSAGE_CHUNK),
     id: z.string(),
     timestamp: z.number(),
     content: z.string(),
 })
 
 export const ChunkCompleteEventSchema = z.object({
-    type: z.literal(EventTypes.CHUNK_COMPLETE),
+    type: z.literal(EventTypes.MESSAGE_COMPLETE),
     id: z.string(),
     timestamp: z.number(),
     fullContent: z.string(),

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/Chat.test.tsx

@@ -9,9 +9,11 @@ jest.mock('./chat.store', () => ({
         getState: jest.fn(),
     },
     useChatMessages: jest.fn(() => []),
+    useAiProvider: jest.fn(() => 'LlmGateway'),
     useChatActions: jest.fn(() => ({
         submitQuestion: jest.fn(),
         clearChat: jest.fn(),
+        setAiProvider: jest.fn(),
     })),
 }))
 
@@ -94,14 +96,14 @@ describe('Chat Component', () => {
                 id: '1',
                 type: 'user' as const,
                 content: 'What is Elasticsearch?',
-                threadId: 'thread-1',
+                conversationId: 'thread-1',
                 timestamp: Date.now(),
             },
             {
                 id: '2',
                 type: 'ai' as const,
                 content: 'Elasticsearch is a search engine...',
-                threadId: 'thread-1',
+                conversationId: 'thread-1',
                 timestamp: Date.now(),
                 status: 'complete' as const,
             },
@@ -245,14 +247,14 @@ describe('Chat Component', () => {
                     id: '1',
                     type: 'user' as const,
                     content: 'Question',
-                    threadId: 'thread-1',
+                    conversationId: 'thread-1',
                     timestamp: Date.now(),
                 },
                 {
                     id: '2',
                     type: 'ai' as const,
                     content: 'Answer...',
-                    threadId: 'thread-1',
+                    conversationId: 'thread-1',
                     timestamp: Date.now(),
                     status: 'streaming' as const,
                 },

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/ChatMessage.test.tsx

@@ -13,7 +13,7 @@ describe('ChatMessage Component', () => {
             id: '1',
             type: 'user',
             content: 'What is Elasticsearch?',
-            threadId: 'thread-1',
+            conversationId: 'thread-1',
             timestamp: Date.now(),
         }
 
@@ -44,7 +44,7 @@ describe('ChatMessage Component', () => {
             id: '2',
             type: 'ai',
             content: 'Elasticsearch is a distributed search engine...',
-            threadId: 'thread-1',
+            conversationId: 'thread-1',
             timestamp: Date.now(),
             status: 'complete',
         }
@@ -102,7 +102,7 @@ describe('ChatMessage Component', () => {
             id: '3',
             type: 'ai',
             content: 'Elasticsearch is...',
-            threadId: 'thread-1',
+            conversationId: 'thread-1',
             timestamp: Date.now(),
             status: 'streaming',
         }
@@ -142,7 +142,7 @@ describe('ChatMessage Component', () => {
             id: '4',
             type: 'ai',
             content: 'Previous content...',
-            threadId: 'thread-1',
+            conversationId: 'thread-1',
             timestamp: Date.now(),
             status: 'error',
         }
@@ -176,7 +176,7 @@ describe('ChatMessage Component', () => {
             id: '5',
             type: 'ai',
             content: '# Heading\n\n**Bold text** and *italic*',
-            threadId: 'thread-1',
+            conversationId: 'thread-1',
             timestamp: Date.now(),
             status: 'complete',
         }

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/ChatMessage.tsx

@@ -63,13 +63,6 @@ interface ChatMessageProps {
     onRetry?: () => void
 }
 
-const getAccumulatedContent = (messages: AskAiEvent[]) => {
-    return messages
-        .filter((m) => m.type === 'chunk')
-        .map((m) => m.content)
-        .join('')
-}
-
 const splitContentAndReferences = (
     content: string
 ): { mainContent: string; referencesJson: string | null } => {
@@ -144,7 +137,7 @@ const computeAiStatus = (
                 m.type === EventTypes.SEARCH_TOOL_CALL ||
                 m.type === EventTypes.TOOL_CALL ||
                 m.type === EventTypes.TOOL_RESULT ||
-                m.type === EventTypes.CHUNK
+                m.type === EventTypes.MESSAGE_CHUNK
         )
         .sort((a, b) => a.timestamp - b.timestamp)
 
@@ -166,9 +159,9 @@ const computeAiStatus = (
         case EventTypes.TOOL_RESULT:
             return STATUS_MESSAGES.ANALYZING
 
-        case EventTypes.CHUNK: {
+        case EventTypes.MESSAGE_CHUNK: {
             const allContent = events
-                .filter((m) => m.type === EventTypes.CHUNK)
+                .filter((m) => m.type === EventTypes.MESSAGE_CHUNK)
                 .map((m) => m.content)
                 .join('')
 
@@ -279,9 +272,9 @@ export const ChatMessage = ({
         )
     }
 
-    const content =
-        streamingContent ||
-        (events.length > 0 ? getAccumulatedContent(events) : message.content)
+    // Use streamingContent during streaming, otherwise use message.content from store
+    // message.content is updated atomically with status when CONVERSATION_END arrives
+    const content = streamingContent || message.content
 
     const hasError = message.status === 'error' || !!error
 

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/StreamingAiMessage.tsx

@@ -3,10 +3,9 @@ import { ChatMessage } from './ChatMessage'
 import {
     ChatMessage as ChatMessageType,
     useChatActions,
-    useThreadId,
+    useConversationId,
 } from './chat.store'
 import { useAskAi } from './useAskAi'
-import * as React from 'react'
 import { useEffect, useRef } from 'react'
 
 interface StreamingAiMessageProps {
@@ -22,20 +21,20 @@ export const StreamingAiMessage = ({
         updateAiMessage,
         hasMessageBeenSent,
         markMessageAsSent,
-        setThreadId,
+        setConversationId,
     } = useChatActions()
-    const threadId = useThreadId()
+    const conversationId = useConversationId()
     const contentRef = useRef('')
 
     const { events, sendQuestion } = useAskAi({
-        threadId: threadId ?? undefined,
+        conversationId: conversationId ?? undefined,
         onEvent: (event) => {
             if (event.type === EventTypes.CONVERSATION_START) {
                 // Capture conversationId from backend on first request
-                if (event.conversationId && !threadId) {
-                    setThreadId(event.conversationId)
+                if (event.conversationId && !conversationId) {
+                    setConversationId(event.conversationId)
                 }
-            } else if (event.type === EventTypes.CHUNK) {
+            } else if (event.type === EventTypes.MESSAGE_CHUNK) {
                 contentRef.current += event.content
             } else if (event.type === EventTypes.ERROR) {
                 // Handle error events from the stream
@@ -45,10 +44,21 @@ export const StreamingAiMessage = ({
                     'error'
                 )
             } else if (event.type === EventTypes.CONVERSATION_END) {
-                updateAiMessage(message.id, contentRef.current, 'complete')
+                updateAiMessage(
+                    message.id,
+                    message.content || contentRef.current,
+                    'complete'
+                )
             }
         },
-        onError: () => {
+        onError: (error) => {
+            console.error('[AI Provider] Error in StreamingAiMessage:', {
+                messageId: message.id,
+                errorMessage: error.message,
+                errorStack: error.stack,
+                errorName: error.name,
+                fullError: error,
+            })
             updateAiMessage(
                 message.id,
                 message.content || 'Error occurred',
@@ -78,15 +88,16 @@ export const StreamingAiMessage = ({
         markMessageAsSent,
     ])
 
+    // Always use contentRef.current if it has content (regardless of status)
+    // This way we don't need to save to message.content and can just use streamingContent
+    const streamingContentToPass =
+        isLast && contentRef.current ? contentRef.current : undefined
+
     return (
         <ChatMessage
             message={message}
             events={isLast ? events : []}
-            streamingContent={
-                isLast && message.status === 'streaming'
-                    ? contentRef.current
-                    : undefined
-            }
+            streamingContent={streamingContentToPass}
         />
     )
 }

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/aiProviderResolver.ts

@@ -94,7 +94,7 @@ describe('chat.store', () => {
 
         // Verify fresh state
         expect(chatStore.getState().chatMessages).toHaveLength(0)
-        expect(chatStore.getState().threadId).toBeNull()
+        expect(chatStore.getState().conversationId).toBeNull()
 
         // Start new conversation
         act(() => {