diff --git a/src/prerna/engine/impl/model/AbstractModelEngine.java b/src/prerna/engine/impl/model/AbstractModelEngine.java index 25f5d465383..a1427dc0ad8 100644 --- a/src/prerna/engine/impl/model/AbstractModelEngine.java +++ b/src/prerna/engine/impl/model/AbstractModelEngine.java @@ -151,7 +151,7 @@ public AskModelEngineResponse askRoom(String question, Room room, AbstractMessag } else { room.setMessages(messageList); } - String messageJson = RoomMessageStore.providerMessageHistory(room, messageList); + String messageJson = RoomMessageStore.serializeMessageHistory(room, messageList); question = messageJson; parameters.put("message_json", messageJson); diff --git a/src/prerna/engine/impl/model/Room.java b/src/prerna/engine/impl/model/Room.java index a287a88d8e4..2412a0ab393 100644 --- a/src/prerna/engine/impl/model/Room.java +++ b/src/prerna/engine/impl/model/Room.java @@ -289,7 +289,7 @@ public synchronized ResponseMessage ask(InputMessage msg, IModelEngine modelEngi // drop orphan tool_use (cancel mid-tool, crash) before building the outbound // branch so providers do not reject the next payload - RoomMessageStore.normalizeForProviderPayload(this); + RoomMessageStore.sanitizeRoomMessages(this); // Set model type and add message to history msg.setModel(modelEngine); @@ -298,13 +298,20 @@ public synchronized ResponseMessage ask(InputMessage msg, IModelEngine modelEngi // first check that messages is not empty. otherwise its the first message of // the thread and parent is null if (!messages.isEmpty()) { - // if a parent message id is passed in, validate it exists and use it. - if (parentMessageId != null && !parentMessageId.isEmpty()) { - msg.setParentMessageId(parentMessageId); + String requestedParentMessageId = parentMessageId == null ? null : parentMessageId.trim(); + boolean requestedParentExists = requestedParentMessageId != null + && !requestedParentMessageId.isEmpty() + && messages.stream().anyMatch(message -> message != null + && requestedParentMessageId.equals(message.getMessageId())); + if (requestedParentExists) { + msg.setParentMessageId(requestedParentMessageId); } else { - // if no parent message id is passed in, use the last message as the parent. AbstractMessage lastMsg = messages.get(messages.size() - 1); msg.setParentMessageId(lastMsg.getMessageId()); + if (requestedParentMessageId != null && !requestedParentMessageId.isEmpty()) { + classLogger.warn("Room {} requested parent message {} was not found after sanitization; using {}", + getId(), requestedParentMessageId, lastMsg.getMessageId()); + } } } else { msg.setParentMessageId(null); // first message diff --git a/src/prerna/engine/impl/model/RoomMessageStore.java b/src/prerna/engine/impl/model/RoomMessageStore.java index 2ed76fefa54..f6bf06993de 100644 --- a/src/prerna/engine/impl/model/RoomMessageStore.java +++ b/src/prerna/engine/impl/model/RoomMessageStore.java @@ -138,22 +138,23 @@ public static boolean refreshFromHotProjection(Room room) { public static String messageHistoryWithNewMessage(Room room, AbstractMessage newMessage) { List branch = MessageUtils.getMessageBranchWithNewMessage(room.getMessages(), newMessage); - validateProviderPayload(room, branch); - return MessageUtils.toJsonArrayWithImageData(branch); + return serializeMessageBranch(room, branch); } public static String currentMessageHistory(Room room) { List branch = MessageUtils.getMessageBranchWithNewMessage(room.getMessages(), null); - validateProviderPayload(room, branch); - return MessageUtils.toJsonArrayWithImageData(branch); + return serializeMessageBranch(room, branch); } - public static String providerMessageHistory(Room room, List messages) { - validateProviderPayload(room, messages); - return MessageUtils.toJsonArrayWithImageData(messages); + public static String serializeMessageHistory(Room room, List messages) { + return serializeMessageBranch(room, messages); } - public static void normalizeForProviderPayload(Room room) { + /** + * Sanitizes the room's in-memory message tree. The healed tree is persisted by + * the next normal successful room write; this method does not write storage. + */ + public static void sanitizeRoomMessages(Room room) { List messages = room.getMessages(); List sanitized = MessageUtils.sanitizeOrphanToolCalls(messages, room); if (sanitized != messages) { @@ -162,6 +163,13 @@ public static void normalizeForProviderPayload(Room room) { validateForPersistence(room, room.getMessages()); } + /** Sanitizes and validates an isolated branch before serializing model history. */ + private static String serializeMessageBranch(Room room, List branch) { + List sanitizedBranch = MessageUtils.sanitizedCopy(branch, room); + validateMessageBranch(room, sanitizedBranch); + return MessageUtils.toJsonArrayWithImageData(sanitizedBranch); + } + public static boolean persist(Room room, String userId) { try (RoomMutationLock ignored = acquireMutationLock(room)) { String messageHistory = room.getMessagesAsString(); @@ -258,7 +266,7 @@ private static void validateForPersistence(Room room, List mess } } - private static void validateProviderPayload(Room room, List messages) { + private static void validateMessageBranch(Room room, List messages) { validateForPersistence(room, messages); Set toolCallIds = new HashSet<>(); @@ -282,13 +290,13 @@ private static void validateProviderPayload(Room room, List mes if (!toolCallIds.containsAll(toolResultIds)) { Set unmatched = new HashSet<>(toolResultIds); unmatched.removeAll(toolCallIds); - throw new IllegalStateException("Room message payload contains tool results without tool calls: " + throw new IllegalStateException("Room message branch contains tool results without tool calls: " + unmatched); } if (!toolResultIds.containsAll(toolCallIds)) { Set unmatched = new HashSet<>(toolCallIds); unmatched.removeAll(toolResultIds); - throw new IllegalStateException("Room message payload contains unresolved tool calls: " + unmatched); + throw new IllegalStateException("Room message branch contains unresolved tool calls: " + unmatched); } } diff --git a/src/prerna/engine/impl/model/message/MessageUtils.java b/src/prerna/engine/impl/model/message/MessageUtils.java index 9fc1e3f46bf..5ce7deb837b 100644 --- a/src/prerna/engine/impl/model/message/MessageUtils.java +++ b/src/prerna/engine/impl/model/message/MessageUtils.java @@ -32,9 +32,12 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.TreeSet; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -395,64 +398,151 @@ public static List fromJsonArrayPreservingToolState(String json } /** - * Truncate at the first ResponseMessage with a TOOL_CALL id that never gets a matching TOOL_RESULT. - * Providers reject unpaired tool_use, so this rewinds in-memory to the last clean turn boundary. - * Pure read-side; persisted JSON is untouched until the next normal turn rewrites it. + * Strip only the specific TOOL_CALL / TOOL_RESULT parts that violate the model-history pairing rule. + * Model APIs reject unpaired tool_use and unpaired tool_result, but removing the whole message throws away + * perfectly good sibling parts -- e.g. a parallel batch where one tool_use is answered and another is not + * (HITL ask-one/answer-other, or a crash mid-batch) would otherwise lose the answered call and orphan its + * result. Instead drop just the offending parts; a message that still has content survives in place, and only + * a message left without any valid parts is removed (with any surviving child re-linked to + * the removed message's parent so the branch chain stays intact). + * Not pure read-side: surviving messages are mutated in place (parts and/or parentMessageId), and the healed + * state is persisted on the next normal turn via RoomMessageStore.sanitizeRoomMessages. */ public static List sanitizeOrphanToolCalls(List messages, Room room) { - if (messages == null || messages.size() < 2) { - return messages != null ? messages : new ArrayList<>(); + if (messages == null) { + return new ArrayList<>(); } - // pass 1: collect (messageIndex, [callIds]) for every TOOL_CALL message, and the full set of TOOL_RESULT ids - List toolCallSlots = new ArrayList<>(); - List> toolCallIdsPerSlot = new ArrayList<>(); - java.util.Set resultIdsSeen = new java.util.HashSet<>(); - for (int i = 0; i < messages.size(); i++) { - AbstractMessage m = messages.get(i); + if (messages.isEmpty()) { + return messages; + } + // pass 1: collect the full set of tool_use ids and tool_result ids seen across the branch + Set toolUseIds = new HashSet<>(); + Set toolResultIds = new HashSet<>(); + for (AbstractMessage m : messages) { if (m == null) continue; for (MessagePart part : m.getParts()) { if (part instanceof ToolCallMessagePart tcp) { Map tc = tcp.getToolCall(); String id = tc != null ? asStringOrNull(tc.get("id")) : null; - if (id != null && !id.isBlank()) { - if (toolCallSlots.isEmpty() || toolCallSlots.get(toolCallSlots.size() - 1)[0] != i) { - toolCallSlots.add(new int[] { i }); - toolCallIdsPerSlot.add(new ArrayList<>()); - } - toolCallIdsPerSlot.get(toolCallIdsPerSlot.size() - 1).add(id); - } + if (id != null && !id.isBlank()) toolUseIds.add(id); } else if (part instanceof ToolResultMessagePart trp) { ToolResultPart tr = trp.getToolResult(); String id = tr != null ? tr.getToolCallId() : null; - if (id != null && !id.isBlank()) resultIdsSeen.add(id); + if (id != null && !id.isBlank()) toolResultIds.add(id); } } } - if (toolCallSlots.isEmpty()) return messages; - - // pass 2: first slot with any unanswered id is the truncation point - int truncateAt = -1; - List orphanIds = null; - for (int s = 0; s < toolCallSlots.size(); s++) { - List unmatched = null; - for (String id : toolCallIdsPerSlot.get(s)) { - if (!resultIdsSeen.contains(id)) { - if (unmatched == null) unmatched = new ArrayList<>(); - unmatched.add(id); + // orphans in both directions -- model history requires bidirectional pairing + Set orphanUseIds = new HashSet<>(toolUseIds); + orphanUseIds.removeAll(toolResultIds); + Set orphanResultIds = new HashSet<>(toolResultIds); + orphanResultIds.removeAll(toolUseIds); + if (orphanUseIds.isEmpty() && orphanResultIds.isEmpty()) return messages; + + // pass 2: strip orphan parts; a message emptied of meaningful content is marked for removal + Set removedIndices = new TreeSet<>(); + Map removedIdToParentId = new HashMap<>(); + for (int i = 0; i < messages.size(); i++) { + AbstractMessage m = messages.get(i); + if (m == null) continue; + List parts = m.getParts(); + boolean hasOrphan = false; + for (MessagePart part : parts) { + if (isOrphanToolPart(part, orphanUseIds, orphanResultIds)) { + hasOrphan = true; + break; } } - if (unmatched != null && !unmatched.isEmpty()) { - truncateAt = toolCallSlots.get(s)[0]; - orphanIds = unmatched; - break; + if (!hasOrphan) continue; + + List kept = new ArrayList<>(parts.size()); + for (MessagePart part : parts) { + if (part != null && !isOrphanToolPart(part, orphanUseIds, orphanResultIds)) kept.add(part); + } + if (hasRemainingPart(kept)) { + m.setParts(kept); + } else { + removedIndices.add(i); + if (m.getMessageId() != null) removedIdToParentId.put(m.getMessageId(), m.getParentMessageId()); + } + } + + // Resolve transitively so consecutive removals collapse into the nearest surviving ancestor + for (String id : new ArrayList<>(removedIdToParentId.keySet())) { + String parent = removedIdToParentId.get(id); + Set guard = new HashSet<>(); + while (parent != null && removedIdToParentId.containsKey(parent) && guard.add(parent)) { + parent = removedIdToParentId.get(parent); + } + removedIdToParentId.put(id, parent); + } + + // Build sanitized list; re-link parentMessageId for any surviving child of a removed message + List sanitized = new ArrayList<>(messages.size() - removedIndices.size()); + for (int i = 0; i < messages.size(); i++) { + if (removedIndices.contains(i)) continue; + AbstractMessage m = messages.get(i); + if (m != null && m.getParentMessageId() != null + && removedIdToParentId.containsKey(m.getParentMessageId())) { + m.setParentMessageId(removedIdToParentId.get(m.getParentMessageId())); } + sanitized.add(m); } - if (truncateAt < 0) return messages; String roomId = room != null ? room.getId() : ""; - classLogger.warn("sanitizeOrphanToolCalls: room {} truncating {} message(s) at index {} -- unpaired tool_use ids: {}", - roomId, (messages.size() - truncateAt), truncateAt, orphanIds); - return new ArrayList<>(messages.subList(0, truncateAt)); + classLogger.warn("sanitizeOrphanToolCalls: room {} stripped orphan tool_use ids {} and tool_result ids {}; " + + "removed {} emptied message(s) at indices {}", + roomId, orphanUseIds, orphanResultIds, removedIndices.size(), removedIndices); + return sanitized; + } + + /** + * Creates an isolated, sanitized copy of a message branch. The source messages + * are not mutated, so branch-specific repair cannot damage shared room + * ancestors that remain valid on a sibling branch. + * + * @param messages message branch to copy + * @param room room context used to rehydrate copied messages + * @return sanitized copies of the supplied messages + */ + public static List sanitizedCopy(List messages, Room room) { + if (messages == null || messages.isEmpty()) { + return new ArrayList<>(); + } + List copies = new ArrayList<>(messages.size()); + for (AbstractMessage message : messages) { + if (message == null) { + copies.add(null); + continue; + } + copies.add(fromJson(GSON_FOR_DB.toJson(message), room)); + } + return sanitizeOrphanToolCalls(copies, room); + } + + /** An orphan part is a TOOL_CALL whose id has no matching result, or a TOOL_RESULT whose id has no matching call. */ + private static boolean isOrphanToolPart(MessagePart part, Set orphanUseIds, Set orphanResultIds) { + if (part instanceof ToolCallMessagePart tcp) { + Map tc = tcp.getToolCall(); + String id = tc != null ? asStringOrNull(tc.get("id")) : null; + return id != null && orphanUseIds.contains(id); + } + if (part instanceof ToolResultMessagePart trp) { + ToolResultPart tr = trp.getToolResult(); + String id = tr != null ? tr.getToolCallId() : null; + return id != null && orphanResultIds.contains(id); + } + return false; + } + + /** A message survives whenever at least one valid, non-orphan part remains. */ + private static boolean hasRemainingPart(List parts) { + for (MessagePart part : parts) { + if (part != null) { + return true; + } + } + return false; } // --- Core two serialization methods ---