[Chat] Fix tool call positioning in conversation timeline

changelogs/fragments/11115.yml

@@ -0,0 +1,2 @@
+fix:
+- [Chat] Fix tool call positioning in conversation timeline ([#11115](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/11115))

src/plugins/chat/public/services/chat_event_handler.test.ts

@@ -361,6 +361,90 @@ describe('ChatEventHandler', () => {
       const toolMessage = timeline.find((msg) => msg.role === 'tool') as ToolMessage;
       expect(toolMessage.content).toBe('First part\nSecond part');
     });
+
+    it('should attach tool call to last assistant message when it appears after last user message', async () => {
+      const assistantMessageId = 'assistant-msg-123';
+      const toolCallId = 'tool-456';
+
+      // Set up timeline: user message, then assistant message
+      await chatEventHandler.handleEvent({
+        type: EventType.TEXT_MESSAGE_START,
+        messageId: assistantMessageId,
+      } as TextMessageStartEvent);
+
+      await chatEventHandler.handleEvent({
+        type: EventType.TEXT_MESSAGE_END,
+        messageId: assistantMessageId,
+      } as TextMessageEndEvent);
+
+      // Add a user message to timeline manually
+      timeline.unshift({
+        id: 'user-msg-001',
+        role: 'user',
+        content: 'Hello',
+      });
+
+      // Now trigger tool call without parentMessageId
+      // The last assistant message (assistantMessageId) appears AFTER the user message in timeline
+      await chatEventHandler.handleEvent({
+        type: EventType.TOOL_CALL_START,
+        toolCallId,
+        toolCallName: 'test_tool',
+        // No parentMessageId - should use selection strategy
+      } as ToolCallStartEvent);
+
+      // Tool call should be attached to the existing assistant message
+      const assistantMessage = timeline.find(
+        (msg) => msg.id === assistantMessageId
+      ) as AssistantMessage;
+      expect(assistantMessage).toBeDefined();
+      expect(assistantMessage.toolCalls).toHaveLength(1);
+      expect(assistantMessage.toolCalls![0].id).toBe(toolCallId);
+    });
+
+    it('should create fake assistant message when user message appears after last assistant message', async () => {
+      const assistantMessageId = 'assistant-msg-123';
+      const toolCallId = 'tool-456';
+
+      // Set up timeline: assistant message, then user message
+      await chatEventHandler.handleEvent({
+        type: EventType.TEXT_MESSAGE_START,
+        messageId: assistantMessageId,
+      } as TextMessageStartEvent);
+
+      await chatEventHandler.handleEvent({
+        type: EventType.TEXT_MESSAGE_END,
+        messageId: assistantMessageId,
+      } as TextMessageEndEvent);
+
+      // Add a user message AFTER the assistant message
+      timeline.push({
+        id: 'user-msg-002',
+        role: 'user',
+        content: 'Another question',
+      });
+
+      const initialTimelineLength = timeline.length;
+
+      // Trigger tool call without parentMessageId
+      // The last user message appears AFTER the assistant message, so a new fake message should be created
+      await chatEventHandler.handleEvent({
+        type: EventType.TOOL_CALL_START,
+        toolCallId,
+        toolCallName: 'test_tool',
+        // No parentMessageId - should create fake assistant message
+      } as ToolCallStartEvent);
+
+      // Should have created a new fake assistant message
+      expect(timeline.length).toBe(initialTimelineLength + 1);
+
+      // The new message should be an assistant message with the tool call
+      const fakeAssistantMessage = timeline[timeline.length - 1] as AssistantMessage;
+      expect(fakeAssistantMessage.role).toBe('assistant');
+      expect(fakeAssistantMessage.id).toMatch(/^fake-assistant-message-/);
+      expect(fakeAssistantMessage.toolCalls).toHaveLength(1);
+      expect(fakeAssistantMessage.toolCalls![0].id).toBe(toolCallId);
+    });
   });
 
   describe('run state handling', () => {

src/plugins/chat/public/services/chat_event_handler.ts

@@ -210,6 +210,17 @@ export class ChatEventHandler {
 
   /**
    * Handle start of a tool call
+   *
+   * This method determines the correct position in the timeline to place tool calls:
+   * 1. If parentMessageId is provided, attach to that specific message
+   * 2. Otherwise, use a selection strategy to determine placement:
+   *    - If the last assistant text message appears after the last user message,
+   *      attach the tool call to that assistant message
+   *    - If not (e.g., user sent a new message after assistant's response),
+   *      create a new fake assistant message to hold the tool calls
+   *
+   * This ensures tool calls are always associated with the correct assistant response
+   * in the conversation timeline, maintaining proper message ordering.
    */
   private handleToolCallStart(event: ToolCallStartEvent): void {
     const { toolCallId, toolCallName, parentMessageId } = event;
@@ -235,12 +246,46 @@ export class ChatEventHandler {
     // Add to pending map for args accumulation
     this.pendingToolCalls.set(toolCallId, toolCall);
 
-    // Use the last TEXT_MESSAGE_START message ID for association
-    const targetMessageId = parentMessageId || this.lastTextMessageStartId;
+    // Strategy 1: Use explicitly provided parent message ID
+    // This is the most reliable approach when the backend provides it
+    if (parentMessageId) {
+      this.addToolCallToMessage(parentMessageId, toolCall);
+      return;
+    }
+
+    // Strategy 2: Determine placement based on message timeline positions
+    // Check if the last assistant message is still the most recent response
+    const timelineMessages = this.getTimeline();
+    if (this.lastTextMessageStartId) {
+      const lastAssistantTextMessageIndex = timelineMessages.findLastIndex(
+        (message) => message.id === this.lastTextMessageStartId
+      );
+      const lastUserMessageIndex = timelineMessages.findLastIndex(
+        (message) => message.role === 'user'
+      );
 
-    if (targetMessageId) {
-      this.addToolCallToMessage(targetMessageId, toolCall);
+      // If the last assistant message appears after the last user message,
+      // it means this tool call belongs to the current conversation turn
+      if (lastAssistantTextMessageIndex > lastUserMessageIndex) {
+        this.addToolCallToMessage(this.lastTextMessageStartId, toolCall);
+        return;
+      }
     }
+
+    // Strategy 3: Create a new assistant message placeholder
+    // This handles the case where the LLM responds with tool calls but without any text message.
+    // Since there's no TEXT_MESSAGE_START event, we need to create a fake assistant message
+    // to hold the tool calls so they appear in the correct position in the timeline.
+    const fakeAssistantMessageId = `fake-assistant-message-` + new Date().getTime();
+    this.onTimelineUpdate((prev) => {
+      const newMessage: AssistantMessage = {
+        id: fakeAssistantMessageId,
+        role: 'assistant',
+        toolCalls: [toolCall],
+      };
+      return [...prev, newMessage];
+    });
+    this.lastTextMessageStartId = fakeAssistantMessageId;
   }
 
   /**