From e281036495580bfd988437cb123ff46b3199c50d Mon Sep 17 00:00:00 2001
From: Ashlee Radka <ashlee@vellum.ai>
Date: Wed, 25 Feb 2026 20:16:37 -0500
Subject: [PATCH] fix: omit textSegments when truncated, handle legacy tool
 results in rehydrate

Co-Authored-By: Claude <noreply@anthropic.com>
---
 assistant/src/daemon/handlers/sessions.ts  | 35 +++++++++++++++++++---
 assistant/src/memory/conversation-store.ts | 21 ++++++++++++-
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/assistant/src/daemon/handlers/sessions.ts b/assistant/src/daemon/handlers/sessions.ts
index cd3fb599fd5..58bb59f52f6 100644
--- a/assistant/src/daemon/handlers/sessions.ts
+++ b/assistant/src/daemon/handlers/sessions.ts
@@ -764,8 +764,8 @@ export function handleHistoryRequest(
       timestamp: m.timestamp,
       ...(truncatedToolCalls.length > 0 ? { toolCalls: truncatedToolCalls, toolCallsBeforeText: m.toolCallsBeforeText } : {}),
       ...(attachments ? { attachments } : {}),
-      ...(m.textSegments.length > 0 ? { textSegments: m.textSegments } : {}),
-      ...(m.contentOrder.length > 0 ? { contentOrder: m.contentOrder } : {}),
+      ...(!wasTruncated && m.textSegments.length > 0 ? { textSegments: m.textSegments } : {}),
+      ...(!wasTruncated && m.contentOrder.length > 0 ? { contentOrder: m.contentOrder } : {}),
       ...(filteredSurfaces.length > 0 ? { surfaces: filteredSurfaces } : {}),
       ...(m.subagentNotification ? { subagentNotification: m.subagentNotification } : {}),
       ...(wasTruncated ? { wasTruncated: true } : {}),
@@ -927,8 +927,35 @@ export function handleMessageContentRequest(
     const content = JSON.parse(dbMessage.content);
     const rendered = renderHistoryContent(content);
     text = rendered.text || undefined;
-    if (rendered.toolCalls.length > 0) {
-      toolCalls = rendered.toolCalls.map((tc) => ({
+    let mergedToolCalls = rendered.toolCalls;
+
+    // Handle legacy conversations where tool_result blocks are stored in the
+    // following user message rather than inline with the assistant message.
+    // This mirrors the mergeToolResults logic used by handleHistoryRequest.
+    if (dbMessage.role === 'assistant' && mergedToolCalls.some((tc) => tc.result === undefined)) {
+      const nextMsg = conversationStore.getNextMessage(msg.sessionId, dbMessage.createdAt);
+      if (nextMsg && nextMsg.role === 'user') {
+        try {
+          const nextContent = JSON.parse(nextMsg.content);
+          const nextRendered = renderHistoryContent(nextContent);
+          if (nextRendered.text.trim() === '' && nextRendered.toolCalls.length > 0) {
+            for (const resultEntry of nextRendered.toolCalls) {
+              const unresolved = mergedToolCalls.find((tc) => tc.result === undefined);
+              if (unresolved) {
+                unresolved.result = resultEntry.result;
+                unresolved.isError = resultEntry.isError;
+                if (resultEntry.imageData) unresolved.imageData = resultEntry.imageData;
+              }
+            }
+          }
+        } catch {
+          // Next message isn't valid JSON — skip merging
+        }
+      }
+    }
+
+    if (mergedToolCalls.length > 0) {
+      toolCalls = mergedToolCalls.map((tc) => ({
         name: tc.name,
         input: tc.input,
         ...(tc.result !== undefined ? { result: tc.result } : {}),
diff --git a/assistant/src/memory/conversation-store.ts b/assistant/src/memory/conversation-store.ts
index 2f3bc1308c7..cbb17b3b2a4 100644
--- a/assistant/src/memory/conversation-store.ts
+++ b/assistant/src/memory/conversation-store.ts
@@ -1,4 +1,4 @@
-import { and, asc, count, desc, eq, inArray, isNull, lt, lte, ne, sql } from 'drizzle-orm';
+import { and, asc, count, desc, eq, gt, inArray, isNull, lt, lte, ne, sql } from 'drizzle-orm';
 import { v4 as uuid } from 'uuid';
 import { z } from 'zod';
 
@@ -338,6 +338,25 @@ export function getMessageById(messageId: string, conversationId?: string): Mess
   return row ? parseMessage(row) : null;
 }
 
+/**
+ * Get the next message in a conversation after a given message (by timestamp).
+ * Used for legacy tool_result merging in the rehydrate endpoint.
+ */
+export function getNextMessage(conversationId: string, afterTimestamp: number): MessageRow | null {
+  const db = getDb();
+  const row = db
+    .select()
+    .from(messages)
+    .where(and(
+      eq(messages.conversationId, conversationId),
+      gt(messages.createdAt, afterTimestamp),
+    ))
+    .orderBy(asc(messages.createdAt))
+    .limit(1)
+    .get();
+  return row ? parseMessage(row) : null;
+}
+
 export interface PaginatedMessagesResult {
   messages: MessageRow[];
   /** Whether older messages exist beyond the returned page. */