turn based message model with full tool call history in conversations

Author: emrberk

e2e/utils/aiAssistant.js

@@ -664,15 +664,54 @@ function isTitleRequest(provider, body) {
   )
 }
 
-function requestMatchesQuestion(provider, body, question) {
+function contentToText(value) {
+  if (typeof value === "string") return value
+  if (value == null) return ""
+
+  if (Array.isArray(value)) {
+    return value.map((entry) => contentToText(entry)).join("\n")
+  }
+
+  if (typeof value === "object") {
+    return [
+      value.text,
+      value.content,
+      value.output,
+      value.arguments,
+      value.input,
+    ]
+      .filter((part) => part != null)
+      .map((part) => contentToText(part))
+      .join("\n")
+  }
+
+  return String(value)
+}
+
+function getUserTexts(provider, body) {
   if (provider === "openai") {
-    return body.input?.[0]?.content === question
+    const userInputs = (body.input || []).filter(
+      (item) => item?.role === "user",
+    )
+    return userInputs.map((item) => contentToText(item.content)).filter(Boolean)
   }
-  // Both openai-chat-completions and anthropic use messages array
-  const firstUserMessage = body.messages?.find(
-    (msg) => msg.role === "user" && typeof msg.content === "string",
+
+  const userMessages = (body.messages || []).filter(
+    (msg) => msg?.role === "user",
+  )
+  return userMessages.map((msg) => contentToText(msg.content)).filter(Boolean)
+}
+
+function requestMatchesQuestion(provider, body, question) {
+  const userTexts = getUserTexts(provider, body)
+  if (userTexts.length === 0) return false
+
+  const lastUserText = userTexts[userTexts.length - 1]
+  return (
+    userTexts.some((text) => text === question) ||
+    userTexts.some((text) => text.includes(question)) ||
+    lastUserText.includes(`User request: ${question}`)
   )
-  return firstUserMessage?.content === question
 }
 
 function extractToolOutputContent(provider, body) {
@@ -699,30 +738,68 @@ function extractToolOutputContent(provider, body) {
       msg.content.some((c) => c.type === "tool_result"),
   )
   const latestMessage = toolResultMessages?.[toolResultMessages.length - 1]
-  const latestToolResult = latestMessage?.content?.find(
-    (c) => c.type === "tool_result",
-  )
+  const latestToolResult = [...(latestMessage?.content || [])]
+    .reverse()
+    .find((c) => c.type === "tool_result")
   return latestToolResult?.content || null
 }
 
 function extractAllInputContent(provider, body) {
   if (provider === "openai") {
-    return body.input?.map((item) => item.content || "").join("\n") || ""
+    return (
+      body.input
+        ?.map((item) =>
+          contentToText(item.content ?? item.output ?? item.arguments ?? ""),
+        )
+        .join("\n") || ""
+    )
   }
   // Both openai-chat-completions and anthropic use messages
   return (
-    body.messages
-      ?.map((msg) => {
-        if (typeof msg.content === "string") return msg.content
-        if (Array.isArray(msg.content)) {
-          return msg.content.map((c) => c.text || c.content || "").join("\n")
-        }
-        return ""
-      })
-      .join("\n") || ""
+    body.messages?.map((msg) => contentToText(msg.content)).join("\n") || ""
   )
 }
 
+function normalizeRequestBodyForAssertions(provider, body) {
+  if (!body || typeof body !== "object") return body
+
+  if (provider === "openai") {
+    const input = Array.isArray(body.input) ? body.input : []
+    return {
+      ...body,
+      input: input.filter(
+        (item) =>
+          item?.type !== "function_call" &&
+          item?.type !== "function_call_output",
+      ),
+    }
+  }
+
+  if (provider === "openai-chat-completions") {
+    const messages = Array.isArray(body.messages) ? body.messages : []
+    return {
+      ...body,
+      messages: messages.filter((msg) => msg?.role !== "tool"),
+    }
+  }
+
+  if (provider === "anthropic") {
+    const messages = Array.isArray(body.messages) ? body.messages : []
+    return {
+      ...body,
+      messages: messages.map((msg) => {
+        if (msg?.role !== "user" || !Array.isArray(msg.content)) return msg
+        return {
+          ...msg,
+          content: msg.content.filter((block) => block?.type !== "tool_result"),
+        }
+      }),
+    }
+  }
+
+  return body
+}
+
 // =============================================================================
 // TOOL CALL FLOW
 // =============================================================================
@@ -891,7 +968,10 @@ function createMultiTurnFlow(config) {
 
         const turn = turns[requestCount]
         if (turn) {
-          requestBodies[requestCount] = req.body
+          requestBodies[requestCount] = normalizeRequestBodyForAssertions(
+            provider,
+            req.body,
+          )
 
           if (turn.expectSystemMessage) {
             const allInputContent = extractAllInputContent(provider, req.body)

src/components/AIStatusIndicator/AssistantModes.tsx

@@ -239,7 +239,7 @@ export const buildOperationSections = (
 }
 
 export const formatDurationMs = (ms: number): string | null => {
-  if (ms < 0) return null
+  if (ms <= 0) return null
   if (ms < 1000) return `${ms}ms`
   return `${Math.round(ms / 1000)}s`
 }

src/hooks/useAIQuickActions.ts

@@ -59,6 +59,7 @@ export const useAIQuickActions = () => {
     createConversation,
     openChatWindow,
     addMessage,
+    removeMessages,
     updateMessage,
     updateConversationName,
     persistMessages,
@@ -145,6 +146,7 @@ export const useAIQuickActions = () => {
         }),
         {
           addMessage,
+          removeMessages,
           updateMessage,
           setStatus,
           setIsStreaming,
@@ -186,6 +188,7 @@ export const useAIQuickActions = () => {
       }),
       {
         addMessage,
+        removeMessages,
         updateMessage,
         setStatus,
         setIsStreaming,
@@ -269,6 +272,7 @@ export const useAIQuickActions = () => {
         }),
         {
           addMessage,
+          removeMessages,
           updateMessage,
           setStatus,
           setIsStreaming,
@@ -316,6 +320,7 @@ export const useAIQuickActions = () => {
       }),
       {
         addMessage,
+        removeMessages,
         updateMessage,
         setStatus,
         setIsStreaming,

src/providers/AIConversationProvider/index.tsx

@@ -33,6 +33,7 @@ import { normalizeSql } from "../../utils/aiAssistant"
 import { useDispatch, useSelector } from "react-redux"
 import { trackEvent } from "../../modules/ConsoleEventTracker"
 import { ConsoleEvent } from "../../modules/ConsoleEventTracker/events"
+import { projectConversationTurns } from "../../utils/ai/turnView"
 
 export type AcceptSuggestionParams = {
   conversationId: ConversationId
@@ -95,6 +96,7 @@ type AIConversationContextType = {
     message: Omit<ConversationMessage, "id"> & { id?: string },
     messageIdsToRemove?: string[],
   ) => void
+  removeMessages: (messageIdsToRemove: string[]) => void
   updateMessage: (
     conversationId: ConversationId,
     messageId: string,
@@ -220,13 +222,26 @@ export const AIConversationProvider: React.FC<{
       lastAssistantMessage?: ConversationMessage
     }> => {
       const messages = await aiConversationStore.getMessages(conversationId)
+      const { visibleEntries, previousVisibleUserByAnchorIndex } =
+        projectConversationTurns(messages)
+      const lastEntry = visibleEntries.at(-1)
+      const lastAssistantMessage =
+        lastEntry?.type === "assistantTurn"
+          ? lastEntry.anchorMessage
+          : undefined
+
+      let lastUserMessage: ConversationMessage | undefined
+      if (lastEntry?.type === "assistantTurn") {
+        lastUserMessage = previousVisibleUserByAnchorIndex.get(
+          lastEntry.anchorIndex,
+        )
+      } else if (lastEntry?.type === "user") {
+        lastUserMessage = lastEntry.message
+      }
+
       return {
-        lastUserMessage: messages
-          .filter((m) => m.role === "user" && !m.hideFromUI)
-          .at(-1),
-        lastAssistantMessage: messages
-          .filter((m) => m.role === "assistant" && !m.hideFromUI)
-          .at(-1),
+        lastUserMessage,
+        lastAssistantMessage,
       }
     },
     [],
@@ -399,6 +414,13 @@ export const AIConversationProvider: React.FC<{
     [],
   )
 
+  const removeMessages = useCallback((messageIdsToRemove: string[]) => {
+    if (messageIdsToRemove.length === 0) return
+    setActiveConversationMessages((prev) =>
+      prev.filter((message) => !messageIdsToRemove.includes(message.id)),
+    )
+  }, [])
+
   const updateMessage = useCallback(
     (
       conversationId: ConversationId,
@@ -998,6 +1020,7 @@ export const AIConversationProvider: React.FC<{
         closeHistoryView,
         deleteConversation,
         addMessage,
+        removeMessages,
         updateMessage,
         replaceConversationMessages,
         updateConversationName,

src/providers/AIConversationProvider/messageContent.ts

@@ -1,16 +1,13 @@
 import type { PartitionBy } from "../../utils/questdb"
 import type { QueryKey } from "../../scenes/Editor/Monaco/utils"
+import type { Message } from "../../utils/ai/types"
+import type { TokenUsage } from "../../utils/aiAssistant"
 import type { OperationHistory } from "../AIStatusProvider"
 
 export type { QueryKey }
 
 export type ConversationId = string
 
-export type TokenUsage = {
-  inputTokens: number
-  outputTokens: number
-}
-
 export type SchemaDisplayData = {
   tableName: string
   kind: "table" | "matview" | "view"
@@ -32,22 +29,20 @@ export type UserMessageDisplayType =
   | "schema_explain_request"
   | "health_issue_request"
 
-export type ConversationMessage = {
+export type ConversationMessage = Message & {
   id: string
-  role: "user" | "assistant"
-  content?: string
   timestamp: number
   error?: string
   sql?: string
-  tokenUsage?: TokenUsage // Token usage for current turn in total, including tool calls that we omit from the history after response
+  tokenUsage?: TokenUsage
   previousSQL?: string // SQL before this change (for diff display)
   isRejected?: boolean
+  contentTimestamp?: number // When text content started streaming
   isAccepted?: boolean
-  hideFromUI?: boolean // User messages for accept/reject and compaction result are hidden
-  isCompacted?: boolean // When converted to true, we include it in the history for UI, but do not send to the model anymore
+  hideFromUI?: boolean
+  isCompacted?: boolean
   operationHistory?: OperationHistory
   model?: string
-  // Predefined actions (Fix and Explain)
   displayType?: UserMessageDisplayType
   displayUserMessage?: string
   displaySchemaData?: SchemaDisplayData

src/providers/AIConversationProvider/types.ts

@@ -1,5 +1,3 @@
-import type { ConversationMessage } from "./types"
-
 /**
  * Trims trailing semicolon from SQL for display purposes.
  * Also ensures the result ends with a newline for Monaco diff editor compatibility.
@@ -14,37 +12,3 @@ export const trimSemicolonForDisplay = (
   }
   return trimmed + "\n"
 }
-
-/**
- * Finds the last visible assistant message with an unactioned SQL diff.
- * Returns the message if found, null otherwise.
- *
- */
-export const getLastUnactionedDiff = (
-  messages: ConversationMessage[],
-): ConversationMessage | null => {
-  // Find last visible message
-  const visibleMessages = messages.filter((m) => !m.hideFromUI)
-  if (visibleMessages.length === 0) return null
-
-  const lastVisible = visibleMessages[visibleMessages.length - 1]
-
-  // Check if it's an assistant message with a real SQL change that hasn't been actioned
-  const hasUnactionedDiff =
-    lastVisible.role === "assistant" &&
-    lastVisible.sql !== undefined &&
-    lastVisible.previousSQL !== undefined &&
-    lastVisible.sql.trim() !== lastVisible.previousSQL.trim() &&
-    !lastVisible.isAccepted &&
-    !lastVisible.isRejected
-
-  return hasUnactionedDiff ? lastVisible : null
-}
-
-/**
- * Checks if there's an unactioned diff in the conversation messages.
- * Simple boolean helper wrapping getLastUnactionedDiff.
- */
-export const hasUnactionedDiff = (messages: ConversationMessage[]): boolean => {
-  return getLastUnactionedDiff(messages) !== null
-}