diff --git a/backend/src/agents/middlewares/clarification_middleware.py b/backend/src/agents/middlewares/clarification_middleware.py index 7a83350..d29987d 100644 --- a/backend/src/agents/middlewares/clarification_middleware.py +++ b/backend/src/agents/middlewares/clarification_middleware.py @@ -5,7 +5,7 @@ from typing import override from langchain.agents import AgentState from langchain.agents.middleware import AgentMiddleware -from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.messages import ToolMessage from langgraph.graph import END from langgraph.prebuilt.tool_node import ToolCallRequest from langgraph.types import Command @@ -118,17 +118,13 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]): name="ask_clarification", ) - ai_response_message = AIMessage(content=formatted_message) - # Return a Command that: - # 1. Adds the formatted tool message (keeping the AI message intact) + # 1. Adds the formatted tool message # 2. Interrupts execution by going to __end__ - # Note: We don't modify the AI message to preserve all fields (reasoning_content, tool_calls, etc.) - # This is especially important for thinking mode where reasoning_content is required - - # Return Command to add the tool message and interrupt + # Note: We don't add an extra AIMessage here - the frontend will detect + # and display ask_clarification tool messages directly return Command( - update={"messages": [tool_message, ai_response_message]}, + update={"messages": [tool_message]}, goto=END, ) diff --git a/frontend/src/components/workspace/messages/message-list.tsx b/frontend/src/components/workspace/messages/message-list.tsx index d74f596..e018335 100644 --- a/frontend/src/components/workspace/messages/message-list.tsx +++ b/frontend/src/components/workspace/messages/message-list.tsx @@ -53,6 +53,17 @@ export function MessageList({ /> ); } + if (group.type === "assistant:clarification") { + const message = group.messages[0]; + if (message && hasContent(message)) { + return ( + + {extractContentFromMessage(message)} + + ); + } + return null; + } if (group.type === "assistant:present-files") { const files: string[] = []; for (const message of group.messages) { diff --git a/frontend/src/core/messages/utils.ts b/frontend/src/core/messages/utils.ts index 75ecbff..3e406ab 100644 --- a/frontend/src/core/messages/utils.ts +++ b/frontend/src/core/messages/utils.ts @@ -14,11 +14,14 @@ interface AssistantMessageGroup extends GenericMessageGroup<"assistant"> {} interface AssistantPresentFilesGroup extends GenericMessageGroup<"assistant:present-files"> {} +interface AssistantClarificationGroup extends GenericMessageGroup<"assistant:clarification"> {} + type MessageGroup = | HumanMessageGroup | AssistantProcessingGroup | AssistantMessageGroup - | AssistantPresentFilesGroup; + | AssistantPresentFilesGroup + | AssistantClarificationGroup; export function groupMessages( messages: Message[], @@ -38,10 +41,28 @@ export function groupMessages( messages: [message], }); } else if (message.type === "tool") { - if ( + // Check if this is a clarification tool message + if (isClarificationToolMessage(message)) { + // Add to processing group if available (to maintain tool call association) + if ( + lastGroup && + lastGroup.type !== "human" && + lastGroup.type !== "assistant" && + lastGroup.type !== "assistant:clarification" + ) { + lastGroup.messages.push(message); + } + // Also create a separate clarification group for prominent display + groups.push({ + id: message.id, + type: "assistant:clarification", + messages: [message], + }); + } else if ( lastGroup && lastGroup.type !== "human" && - lastGroup.type !== "assistant" + lastGroup.type !== "assistant" && + lastGroup.type !== "assistant:clarification" ) { lastGroup.messages.push(message); } else { @@ -190,6 +211,10 @@ export function hasPresentFiles(message: Message) { ); } +export function isClarificationToolMessage(message: Message) { + return message.type === "tool" && message.name === "ask_clarification"; +} + export function extractPresentFilesFromMessage(message: Message) { if (message.type !== "ai" || !hasPresentFiles(message)) { return [];