fix(strands): make integration stateless and fix duplicate tool rendering

Signed-off-by: Tyler Slaton <[email protected]>

Author: tylerslaton

integrations/aws-strands/python/src/ag_ui_strands/agent.py

@@ -8,7 +8,6 @@
 import json
 import uuid
 import asyncio
-import sys
 from typing import AsyncIterator, Any
 from strands import Agent as StrandsAgentCore
 
@@ -93,18 +92,42 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                     has_pending_tool_result = True
                     logger.debug(f"Has pending tool result detected: tool_call_id={getattr(last_msg, 'tool_call_id', 'unknown')}, thread_id={input_data.thread_id}")
 
-            # Get the latest user message
+            # Convert AG-UI messages to Strands format
+            # Strands Message format: {"role": str, "content": List[ContentBlock]}
+            strands_messages = []
+            for msg in input_data.messages:
+                # Skip tool messages as Strands handles them internally
+                if msg.role == "tool":
+                    continue
+
+                # Convert message content to Strands format
+                if isinstance(msg.content, str):
+                    strands_messages.append({
+                        "role": msg.role,
+                        "content": [{"text": msg.content}]
+                    })
+                elif isinstance(msg.content, list):
+                    # Already in content block format
+                    strands_messages.append({
+                        "role": msg.role,
+                        "content": msg.content
+                    })
+
+            # Get the latest user message for state context builder
             user_message = "Hello"
             if input_data.messages:
                 for msg in reversed(input_data.messages):
-                    if (msg.role == "user" or msg.role == "tool") and msg.content:
+                    if msg.role == "user" and msg.content:
                         user_message = msg.content
                         break
 
             # Optionally allow configuration to adjust the outgoing user message
             if self.config.state_context_builder:
                 try:
                     user_message = self.config.state_context_builder(input_data, user_message)
+                    # If state_context_builder modifies the message, update the last user message
+                    if strands_messages and strands_messages[-1]["role"] == "user":
+                        strands_messages[-1]["content"] = [{"text": user_message}]
                 except Exception:
                     # If the builder fails, keep the original message
                     pass
@@ -116,10 +139,14 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
             stop_text_streaming = False
             halt_event_stream = False
 
-            logger.debug(f"Starting agent run: thread_id={input_data.thread_id}, run_id={input_data.run_id}, has_pending_tool_result={has_pending_tool_result}, message_count={len(input_data.messages)}")
+            logger.debug(f"Starting agent run: thread_id={input_data.thread_id}, run_id={input_data.run_id}, has_pending_tool_result={has_pending_tool_result}, message_count={len(input_data.messages)}, strands_message_count={len(strands_messages)}")
 
-            # Stream from Strands agent
-            agent_stream = self.strands_agent.stream_async(user_message)
+            # Reset agent messages to prevent accumulation across threads
+            # The frontend manages conversation history via input_data.messages
+            self.strands_agent.messages = []
+
+            # Stream from Strands agent with full conversation history
+            agent_stream = self.strands_agent.stream_async(strands_messages if strands_messages else user_message)
 
             try:
                 async for event in agent_stream:
@@ -196,35 +223,9 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
 
                             logger.debug(f"Processing tool result: tool_name={tool_name}, result_tool_id={result_tool_id}, has_pending_tool_result={has_pending_tool_result}, thread_id={input_data.thread_id}")
 
-                            if not has_pending_tool_result and not (behavior and behavior.skip_messages_snapshot):
-                                assistant_msg = AssistantMessage(
-                                    id=str(uuid.uuid4()),
-                                    role="assistant",
-                                    tool_calls=[
-                                        ToolCall(
-                                            id=result_tool_id,
-                                            type="function",
-                                            function={
-                                                "name": tool_name or "default_tool",
-                                                "arguments": tool_args or "{}",
-                                            },
-                                        )
-                                    ],
-                                )
-
-                                content_str = json.dumps(result_data) if isinstance(result_data, dict) else str(result_data)
-                                tool_msg = ToolMessage(
-                                    id=str(uuid.uuid4()),
-                                    role="tool",
-                                    content=content_str,
-                                    tool_call_id=result_tool_id,
-                                )
-
-                                all_messages = list(input_data.messages) + [assistant_msg, tool_msg]
-                                yield MessagesSnapshotEvent(
-                                    type=EventType.MESSAGES_SNAPSHOT,
-                                    messages=all_messages
-                                )
+                            # Skip MessagesSnapshotEvent for tool results to avoid duplicates
+                            # The frontend already has the tool call from TOOL_CALL_* events
+                            # and will construct the messages itself
 
                             result_context = ToolResultContext(
                                 input_data=input_data,