mcp-use 1.3.13__py3-none-any.whl → 1.4.0__py3-none-any.whl

mcp_use/agents/mcpagent.py

@@ -3,6 +3,12 @@ MCP: Main integration module with customizable system prompt.
 
 This module provides the main MCPAgent class that integrates all components
 to provide a simple interface for using MCP tools with different LLMs.
+
+LangChain 1.0.0 Migration:
+- The agent uses create_agent() from langchain.agents which returns a CompiledStateGraph
+- New methods: astream_simplified() and run_v2() leverage the built-in astream() from
+  CompiledStateGraph which handles the agent loop internally
+- Legacy methods: stream() and run() use manual step-by-step execution for backward compatibility
 """
 
 import logging
@@ -10,17 +16,14 @@ import time
 from collections.abc import AsyncGenerator, AsyncIterator
 from typing import TypeVar
 
-from langchain.agents import AgentExecutor, create_tool_calling_agent
-from langchain.agents.output_parsers.tools import ToolAgentAction
-from langchain.globals import set_debug
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
-from langchain.schema.language_model import BaseLanguageModel
-from langchain_core.agents import AgentAction, AgentFinish
-from langchain_core.exceptions import OutputParserException
+from langchain.agents import create_agent
+from langchain.agents.middleware import ModelCallLimitMiddleware
+from langchain_core.agents import AgentAction
+from langchain_core.globals import set_debug
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
 from langchain_core.runnables.schema import StreamEvent
 from langchain_core.tools import BaseTool
-from langchain_core.utils.input import get_color_mapping
 from pydantic import BaseModel
 
 from mcp_use.agents.adapters.langchain_adapter import LangChainAdapter
@@ -150,7 +153,7 @@ class MCPAgent:
             self.server_manager = ServerManager(self.client, self.adapter)
 
         # State tracking - initialize _tools as empty list
-        self._agent_executor: AgentExecutor | None = None
+        self._agent_executor = None
         self._system_message: SystemMessage | None = None
         self._tools: list[BaseTool] = []
 
@@ -185,7 +188,8 @@ class MCPAgent:
                     logger.info(f"✅ Created {len(self._sessions)} new sessions")
 
                 # Create LangChain tools directly from the client using the adapter
-                self._tools = await self.adapter.create_tools(self.client)
+                await self.adapter.create_all(self.client)
+                self._tools = self.adapter.tools + self.adapter.resources + self.adapter.prompts
                 logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from client")
             else:
                 # Using direct connector - only establish connection
@@ -197,7 +201,10 @@ class MCPAgent:
                         await connector.connect()
 
                 # Create LangChain tools using the adapter with connectors
-                self._tools = await self.adapter._create_tools_from_connectors(connectors_to_use)
+                await self.adapter._create_tools_from_connectors(connectors_to_use)
+                await self.adapter._create_resources_from_connectors(connectors_to_use)
+                await self.adapter._create_prompts_from_connectors(connectors_to_use)
+                self._tools = self.adapter.tools + self.adapter.resources + self.adapter.prompts
                 logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from connectors")
 
             # Get all tools for system message generation
@@ -274,7 +281,7 @@ class MCPAgent:
                 msg for msg in self._conversation_history if not isinstance(msg, SystemMessage)
             ]
 
-    def _create_agent(self) -> AgentExecutor:
+    def _create_agent(self):
         """Create the LangChain agent with the configured system message.
 
         Returns:
@@ -286,42 +293,23 @@ class MCPAgent:
         if self._system_message:
             system_content = self._system_message.content
 
-        if self.memory_enabled:
-            # Query already in chat_history — don't re-inject it
-            prompt = ChatPromptTemplate.from_messages(
-                [
-                    ("system", system_content),
-                    MessagesPlaceholder(variable_name="chat_history"),
-                    ("human", "{input}"),
-                    MessagesPlaceholder(variable_name="agent_scratchpad"),
-                ]
-            )
-        else:
-            # No memory — inject input directly
-            prompt = ChatPromptTemplate.from_messages(
-                [
-                    ("system", system_content),
-                    ("human", "{input}"),
-                    MessagesPlaceholder(variable_name="agent_scratchpad"),
-                ]
-            )
-
         tool_names = [tool.name for tool in self._tools]
         logger.info(f"🧠 Agent ready with tools: {', '.join(tool_names)}")
 
-        # Use the standard create_tool_calling_agent
-        agent = create_tool_calling_agent(llm=self.llm, tools=self._tools, prompt=prompt)
+        # Create middleware to enforce max_steps
+        # ModelCallLimitMiddleware limits the number of model calls, which corresponds to agent steps
+        middleware = [ModelCallLimitMiddleware(run_limit=self.max_steps)]
+
+        # Use the standard create_agent with middleware
+        agent = create_agent(
+            model=self.llm, tools=self._tools, system_prompt=system_content, middleware=middleware, debug=self.verbose
+        )
 
-        # Use the standard AgentExecutor with callbacks
-        executor = AgentExecutor(
-            agent=agent,
-            tools=self._tools,
-            max_iterations=self.max_steps,
-            verbose=self.verbose,
-            callbacks=self.callbacks,
+        logger.debug(
+            f"Created agent with max_steps={self.max_steps} (via ModelCallLimitMiddleware) "
+            f"and {len(self.callbacks)} callbacks"
         )
-        logger.debug(f"Created agent executor with max_iterations={self.max_steps} and {len(self.callbacks)} callbacks")
-        return executor
+        return agent
 
     def get_conversation_history(self) -> list[BaseMessage]:
         """Get the current conversation history.
@@ -393,16 +381,14 @@ class MCPAgent:
 
     async def _consume_and_return(
         self,
-        generator: AsyncGenerator[tuple[AgentAction, str] | str | T, None],
+        generator: AsyncGenerator[str | T, None],
     ) -> tuple[str | T, int]:
-        """Consume the generator and return the final result.
+        """Consume the stream generator and return the final result.
 
-        This method manually iterates through the generator to consume the steps.
-        In Python, async generators cannot return values directly, so we expect
-        the final result to be yielded as a special marker.
+        This is used by the run() method with the astream implementation.
 
         Args:
-            generator: The async generator that yields steps and a final result.
+            generator: The async generator from astream.
 
         Returns:
             A tuple of (final_result, steps_taken). final_result can be a string
@@ -411,416 +397,12 @@ class MCPAgent:
         final_result = ""
         steps_taken = 0
         async for item in generator:
-            # If it's a string, it's the final result (regular output)
-            if isinstance(item, str):
-                final_result = item
-                break
-            # If it's not a tuple, it might be structured output (Pydantic model)
-            elif not isinstance(item, tuple):
-                final_result = item
-                break
-            # Otherwise it's a step tuple, just consume it
-            else:
-                steps_taken += 1
+            # The last item yielded is always the final result
+            final_result = item
+        # Count steps as the number of tools used during execution
+        steps_taken = len(self.tools_used_names)
         return final_result, steps_taken
 
-    @telemetry("agent_stream")
-    async def stream(
-        self,
-        query: str,
-        max_steps: int | None = None,
-        manage_connector: bool = True,
-        external_history: list[BaseMessage] | None = None,
-        track_execution: bool = True,
-        output_schema: type[T] | None = None,
-    ) -> AsyncGenerator[tuple[AgentAction, str] | str | T, None]:
-        """Run the agent and yield intermediate steps as an async generator.
-
-        Args:
-            query: The query to run.
-            max_steps: Optional maximum number of steps to take.
-            manage_connector: Whether to handle the connector lifecycle internally.
-            external_history: Optional external history to use instead of the
-                internal conversation history.
-            track_execution: Whether to track execution for telemetry.
-            output_schema: Optional Pydantic BaseModel class for structured output.
-                If provided, the agent will attempt structured output at finish points
-                and continue execution if required information is missing.
-
-        Yields:
-            Intermediate steps as (AgentAction, str) tuples, followed by the final result.
-            If output_schema is provided, yields structured output as instance of the schema.
-        """
-        # Delegate to remote agent if in remote mode
-        if self._is_remote and self._remote_agent:
-            async for item in self._remote_agent.stream(
-                query, max_steps, manage_connector, external_history, track_execution, output_schema
-            ):
-                yield item
-            return
-
-        result = ""
-        initialized_here = False
-        start_time = time.time()
-        steps_taken = 0
-        success = False
-
-        # Schema-aware setup for structured output
-        structured_llm = None
-        schema_description = ""
-        if output_schema:
-            query = self._enhance_query_with_schema(query, output_schema)
-            structured_llm = self.llm.with_structured_output(output_schema)
-            # Get schema description for feedback
-            schema_fields = []
-            try:
-                for field_name, field_info in output_schema.model_fields.items():
-                    description = getattr(field_info, "description", "") or field_name
-                    required = not hasattr(field_info, "default") or field_info.default is None
-                    schema_fields.append(f"- {field_name}: {description} {'(required)' if required else '(optional)'}")
-
-                schema_description = "\n".join(schema_fields)
-            except Exception as e:
-                logger.warning(f"Could not extract schema details: {e}")
-                schema_description = f"Schema: {output_schema.__name__}"
-
-        try:
-            # Initialize if needed
-            if manage_connector and not self._initialized:
-                await self.initialize()
-                initialized_here = True
-            elif not self._initialized and self.auto_initialize:
-                await self.initialize()
-                initialized_here = True
-
-            # Check if initialization succeeded
-            if not self._agent_executor:
-                raise RuntimeError("MCP agent failed to initialize")
-
-            steps = max_steps or self.max_steps
-            if self._agent_executor:
-                self._agent_executor.max_iterations = steps
-
-            display_query = query[:50].replace("\n", " ") + "..." if len(query) > 50 else query.replace("\n", " ")
-            logger.info(f"💬 Received query: '{display_query}'")
-
-            # Use the provided history or the internal history
-            history_to_use = external_history if external_history is not None else self._conversation_history
-
-            # Convert messages to format expected by LangChain agent input
-            # Exclude the main system message as it's part of the agent's prompt
-            langchain_history = []
-            for msg in history_to_use:
-                if isinstance(msg, HumanMessage):
-                    langchain_history.append(msg)
-                elif isinstance(msg, AIMessage):
-                    langchain_history.append(msg)
-
-            intermediate_steps: list[tuple[AgentAction, str]] = []
-            inputs = {"input": query, "chat_history": langchain_history}
-
-            # Construct a mapping of tool name to tool for easy lookup
-            name_to_tool_map = {tool.name: tool for tool in self._tools}
-            color_mapping = get_color_mapping([tool.name for tool in self._tools], excluded_colors=["green", "red"])
-
-            logger.info(f"🏁 Starting agent execution with max_steps={steps}")
-
-            # Track whether agent finished successfully vs reached max iterations
-            agent_finished_successfully = False
-            result = None
-
-            # Create a run manager with our callbacks if we have any - ONCE for the entire execution
-            run_manager = None
-            if self.callbacks:
-                # Create an async callback manager with our callbacks
-                from langchain_core.callbacks.manager import AsyncCallbackManager
-
-                callback_manager = AsyncCallbackManager.configure(
-                    inheritable_callbacks=self.callbacks,
-                    local_callbacks=self.callbacks,
-                )
-                # Create a run manager for this chain execution
-                run_manager = await callback_manager.on_chain_start(
-                    {"name": "MCPAgent (mcp-use)"},
-                    inputs,
-                )
-
-            for step_num in range(steps):
-                steps_taken = step_num + 1
-                # --- Check for tool updates if using server manager ---
-                if self.use_server_manager and self.server_manager:
-                    current_tools = self.server_manager.tools
-                    current_tool_names = {tool.name for tool in current_tools}
-                    existing_tool_names = {tool.name for tool in self._tools}
-
-                    if current_tool_names != existing_tool_names:
-                        logger.info(
-                            f"🔄 Tools changed before step {step_num + 1}, updating agent."
-                            f"New tools: {', '.join(current_tool_names)}"
-                        )
-                        self._tools = current_tools
-                        # Regenerate system message with ALL current tools
-                        await self._create_system_message_from_tools(self._tools)
-                        # Recreate the agent executor with the new tools and system message
-                        self._agent_executor = self._create_agent()
-                        self._agent_executor.max_iterations = steps
-                        # Update maps for this iteration
-                        name_to_tool_map = {tool.name: tool for tool in self._tools}
-                        color_mapping = get_color_mapping(
-                            [tool.name for tool in self._tools], excluded_colors=["green", "red"]
-                        )
-
-                logger.info(f"👣 Step {step_num + 1}/{steps}")
-
-                # --- Plan and execute the next step ---
-                try:
-                    retry_count = 0
-                    next_step_output = None
-
-                    while retry_count <= self.max_retries_per_step:
-                        try:
-                            # Use the internal _atake_next_step which handles planning and execution
-                            # This requires providing the necessary context like maps and intermediate steps
-                            next_step_output = await self._agent_executor._atake_next_step(
-                                name_to_tool_map=name_to_tool_map,
-                                color_mapping=color_mapping,
-                                inputs=inputs,
-                                intermediate_steps=intermediate_steps,
-                                run_manager=run_manager,
-                            )
-
-                            # If we get here, the step succeeded, break out of retry loop
-                            break
-
-                        except Exception as e:
-                            if not self.retry_on_error or retry_count >= self.max_retries_per_step:
-                                logger.error(f"❌ Validation error during step {step_num + 1}: {e}")
-                                result = f"Agent stopped due to a validation error: {str(e)}"
-                                success = False
-                                yield result
-                                return
-
-                            retry_count += 1
-                            logger.warning(
-                                f"⚠️ Validation error, retrying ({retry_count}/{self.max_retries_per_step}): {e}"
-                            )
-
-                            # Create concise feedback for the LLM about the validation error
-                            error_message = f"Error: {str(e)}"
-                            inputs["input"] = error_message
-
-                            # Continue to next iteration of retry loop
-                            continue
-
-                    # Process the output
-                    if isinstance(next_step_output, AgentFinish):
-                        logger.info(f"✅ Agent finished at step {step_num + 1}")
-                        agent_finished_successfully = True
-                        output_value = next_step_output.return_values.get("output", "No output generated")
-                        result = self._normalize_output(output_value)
-                        # End the chain if we have a run manager
-                        if run_manager:
-                            await run_manager.on_chain_end({"output": result})
-
-                        # If structured output is requested, attempt to create it
-                        if output_schema and structured_llm:
-                            try:
-                                logger.info("🔧 Attempting structured output...")
-                                structured_result = await self._attempt_structured_output(
-                                    result, structured_llm, output_schema, schema_description
-                                )
-
-                                # Add the final response to conversation history if memory is enabled
-                                if self.memory_enabled:
-                                    self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
-
-                                logger.info("✅ Structured output successful")
-                                success = True
-                                yield structured_result
-                                return
-
-                            except Exception as e:
-                                logger.warning(f"⚠️ Structured output failed: {e}")
-                                # Continue execution to gather missing information
-                                missing_info_prompt = f"""
-                                The current result cannot be formatted into the required structure.
-                                Error: {str(e)}
-
-                                Current information: {result}
-
-                                Please continue working to gather the missing information needed for:
-                                {schema_description}
-
-                                Focus on finding the specific missing details.
-                                """
-
-                                # Add this as feedback and continue the loop
-                                inputs["input"] = missing_info_prompt
-                                if self.memory_enabled:
-                                    self.add_to_history(HumanMessage(content=missing_info_prompt))
-
-                                logger.info("🔄 Continuing execution to gather missing information...")
-                                continue
-                        else:
-                            # Regular execution without structured output
-                            break
-
-                    # If it's actions/steps, add to intermediate steps and yield them
-                    intermediate_steps.extend(next_step_output)
-
-                    # Yield each step and track tool usage
-                    for agent_step in next_step_output:
-                        yield agent_step
-                        action, observation = agent_step
-                        reasoning = getattr(action, "log", "")
-                        if reasoning:
-                            reasoning_str = reasoning.replace("\n", " ")
-                            if len(reasoning_str) > 300:
-                                reasoning_str = reasoning_str[:297] + "..."
-                            logger.info(f"💭 Reasoning: {reasoning_str}")
-                        tool_name = action.tool
-                        self.tools_used_names.append(tool_name)
-                        tool_input_str = str(action.tool_input)
-                        # Truncate long inputs for readability
-                        if len(tool_input_str) > 100:
-                            tool_input_str = tool_input_str[:97] + "..."
-                        logger.info(f"🔧 Tool call: {tool_name} with input: {tool_input_str}")
-                        # Truncate long outputs for readability
-                        observation_str = str(observation)
-                        if len(observation_str) > 100:
-                            observation_str = observation_str[:97] + "..."
-                        observation_str = observation_str.replace("\n", " ")
-                        logger.info(f"📄 Tool result: {observation_str}")
-
-                    # Check for return_direct on the last action taken
-                    if len(next_step_output) > 0:
-                        last_step: tuple[AgentAction, str] = next_step_output[-1]
-                        tool_return = self._agent_executor._get_tool_return(last_step)
-                        if tool_return is not None:
-                            logger.info(f"🏆 Tool returned directly at step {step_num + 1}")
-                            agent_finished_successfully = True
-                            result = tool_return.return_values.get("output", "No output generated")
-                            result = self._normalize_output(result)
-                            break
-
-                except OutputParserException as e:
-                    logger.error(f"❌ Output parsing error during step {step_num + 1}: {e}")
-                    result = f"Agent stopped due to a parsing error: {str(e)}"
-                    if run_manager:
-                        await run_manager.on_chain_error(e)
-                    break
-                except Exception as e:
-                    logger.error(f"❌ Error during agent execution step {step_num + 1}: {e}")
-                    import traceback
-
-                    traceback.print_exc()
-                    # End the chain with error if we have a run manager
-                    if run_manager:
-                        await run_manager.on_chain_error(e)
-                    result = f"Agent stopped due to an error: {str(e)}"
-                    break
-
-            # --- Loop finished ---
-            if not result:
-                if agent_finished_successfully:
-                    # Agent finished successfully but returned empty output
-                    result = "Agent completed the task successfully."
-                    logger.info("✅ Agent finished successfully with empty output")
-                else:
-                    # Agent actually reached max iterations
-                    logger.warning(f"⚠️ Agent stopped after reaching max iterations ({steps})")
-                    result = f"Agent stopped after reaching the maximum number of steps ({steps})."
-                    if run_manager:
-                        await run_manager.on_chain_end({"output": result})
-
-            # If structured output was requested but not achieved, attempt one final time
-            if output_schema and structured_llm and not success:
-                try:
-                    logger.info("🔧 Final attempt at structured output...")
-                    structured_result = await self._attempt_structured_output(
-                        result, structured_llm, output_schema, schema_description
-                    )
-
-                    # Add the final response to conversation history if memory is enabled
-                    if self.memory_enabled:
-                        self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
-
-                    logger.info("✅ Final structured output successful")
-                    success = True
-                    yield structured_result
-                    return
-
-                except Exception as e:
-                    logger.error(f"❌ Final structured output attempt failed: {e}")
-                    raise RuntimeError(f"Failed to generate structured output after {steps} steps: {str(e)}") from e
-
-            if self.memory_enabled:
-                self.add_to_history(HumanMessage(content=query))
-
-            if self.memory_enabled and not output_schema:
-                self.add_to_history(AIMessage(content=self._normalize_output(result)))
-
-            logger.info(f"🎉 Agent execution complete in {time.time() - start_time} seconds")
-            if not success:
-                success = True
-
-            # Yield the final result (only for non-structured output)
-            if not output_schema:
-                yield result
-
-        except Exception as e:
-            logger.error(f"❌ Error running query: {e}")
-            if initialized_here and manage_connector:
-                logger.info("🧹 Cleaning up resources after initialization error in stream")
-                await self.close()
-            raise
-
-        finally:
-            # Track comprehensive execution data
-            execution_time_ms = int((time.time() - start_time) * 1000)
-
-            server_count = 0
-            if self.client:
-                server_count = len(self.client.get_all_active_sessions())
-            elif self.connectors:
-                server_count = len(self.connectors)
-
-            conversation_history_length = len(self._conversation_history) if self.memory_enabled else 0
-
-            # Safely access _tools in case initialization failed
-            tools_available = getattr(self, "_tools", [])
-
-            if track_execution:
-                self.telemetry.track_agent_execution(
-                    execution_method="stream",
-                    query=query,
-                    success=success,
-                    model_provider=self._model_provider,
-                    model_name=self._model_name,
-                    server_count=server_count,
-                    server_identifiers=[connector.public_identifier for connector in self.connectors],
-                    total_tools_available=len(tools_available),
-                    tools_available_names=[tool.name for tool in tools_available],
-                    max_steps_configured=self.max_steps,
-                    memory_enabled=self.memory_enabled,
-                    use_server_manager=self.use_server_manager,
-                    max_steps_used=max_steps,
-                    manage_connector=manage_connector,
-                    external_history_used=external_history is not None,
-                    steps_taken=steps_taken,
-                    tools_used_count=len(self.tools_used_names),
-                    tools_used_names=self.tools_used_names,
-                    response=result,
-                    execution_time_ms=execution_time_ms,
-                    error_type=None if success else "execution_error",
-                    conversation_history_length=conversation_history_length,
-                )
-
-            # Clean up if necessary (e.g., if not using client-managed sessions)
-            if manage_connector and not self.client and initialized_here:
-                logger.info("🧹 Closing agent after stream completion")
-                await self.close()
-
     @telemetry("agent_run")
     async def run(
         self,
@@ -830,23 +412,15 @@ class MCPAgent:
         external_history: list[BaseMessage] | None = None,
         output_schema: type[T] | None = None,
     ) -> str | T:
-        """Run a query using the MCP tools and return the final result.
-
-        This method uses the streaming implementation internally and returns
-        the final result after consuming all intermediate steps. If output_schema
-        is provided, the agent will be schema-aware and return structured output.
+        """Run a query using LangChain 1.0.0's agent and return the final result.
 
         Args:
             query: The query to run.
             max_steps: Optional maximum number of steps to take.
             manage_connector: Whether to handle the connector lifecycle internally.
-                If True, this method will connect, initialize, and disconnect from
-                the connector automatically. If False, the caller is responsible
-                for managing the connector lifecycle.
             external_history: Optional external history to use instead of the
                 internal conversation history.
             output_schema: Optional Pydantic BaseModel class for structured output.
-                If provided, the agent will attempt to return an instance of this model.
 
         Returns:
             The result of running the query as a string, or if output_schema is provided,
@@ -882,8 +456,8 @@ class MCPAgent:
             query, max_steps, manage_connector, external_history, track_execution=False, output_schema=output_schema
         )
         error = None
-        steps_taken = 0
         result = None
+        steps_taken = 0
         try:
             result, steps_taken = await self._consume_and_return(generator)
 
@@ -983,6 +557,329 @@ class MCPAgent:
 
         return enhanced_query
 
+    @telemetry("agent_stream")
+    async def stream(
+        self,
+        query: str,
+        max_steps: int | None = None,
+        manage_connector: bool = True,
+        external_history: list[BaseMessage] | None = None,
+        track_execution: bool = True,
+        output_schema: type[T] | None = None,
+    ) -> AsyncGenerator[tuple[AgentAction, str] | str | T, None]:
+        """Async generator using LangChain 1.0.0's create_agent and astream.
+
+        This method leverages the LangChain 1.0.0 API where create_agent returns
+        a CompiledStateGraph that handles the agent loop internally via astream.
+
+        **Tool Updates with Server Manager:**
+        When using server_manager mode, this method handles dynamic tool updates:
+        - **Before execution:** Updates are applied immediately to the new stream
+        - **During execution:** When tools change, we wait for a "safe restart point"
+          (after tool results complete), then interrupt the stream, recreate the agent
+          with new tools, and resume execution with accumulated messages.
+        - **Safe restart points:** Only restart after tool results to ensure message
+          pairs (tool_use + tool_result) are complete, satisfying LLM API requirements.
+        - **Max restarts:** Limited to 3 restarts to prevent infinite loops
+
+        This interrupt-and-restart approach ensures that tools added mid-execution
+        (e.g., via connect_to_mcp_server) are immediately available to the agent,
+        maintaining the same behavior as the legacy implementation while respecting
+        API constraints.
+
+        Args:
+            query: The query to run.
+            manage_connector: Whether to handle the connector lifecycle internally.
+            external_history: Optional external history to use instead of the
+                internal conversation history.
+            output_schema: Optional Pydantic BaseModel class for structured output.
+
+        Yields:
+            Intermediate steps and final result from the agent execution.
+        """
+        # Delegate to remote agent if in remote mode
+        if self._is_remote and self._remote_agent:
+            async for item in self._remote_agent.stream(query, max_steps, external_history, output_schema):
+                yield item
+            return
+
+        initialized_here = False
+        start_time = time.time()
+        success = False
+        final_output = None
+        steps_taken = 0
+
+        try:
+            # 1. Initialize if needed
+            if manage_connector and not self._initialized:
+                await self.initialize()
+                initialized_here = True
+            elif not self._initialized and self.auto_initialize:
+                await self.initialize()
+                initialized_here = True
+
+            if not self._agent_executor:
+                raise RuntimeError("MCP agent failed to initialize")
+
+            # Check for tool updates before starting execution (if using server manager)
+            if self.use_server_manager and self.server_manager:
+                current_tools = self.server_manager.tools
+                current_tool_names = {tool.name for tool in current_tools}
+                existing_tool_names = {tool.name for tool in self._tools}
+
+                if current_tool_names != existing_tool_names:
+                    logger.info(
+                        f"🔄 Tools changed before execution, updating agent. New tools: {', '.join(current_tool_names)}"
+                    )
+                    self._tools = current_tools
+                    # Regenerate system message with ALL current tools
+                    await self._create_system_message_from_tools(self._tools)
+                    # Recreate the agent executor with the new tools and system message
+                    self._agent_executor = self._create_agent()
+
+            # 2. Build inputs for the agent
+            history_to_use = external_history if external_history is not None else self._conversation_history
+
+            # Convert messages to format expected by LangChain agent
+            langchain_history = []
+            for msg in history_to_use:
+                if isinstance(msg, HumanMessage | AIMessage):
+                    langchain_history.append(msg)
+
+            inputs = {"messages": [*langchain_history, HumanMessage(content=query)]}
+
+            display_query = query[:50].replace("\n", " ") + "..." if len(query) > 50 else query.replace("\n", " ")
+            logger.info(f"💬 Received query: '{display_query}'")
+            logger.info("🏁 Starting agent execution")
+
+            # 3. Stream using the built-in astream from CompiledStateGraph
+            # The agent graph handles the loop internally
+            # With dynamic tool reload: if tools change mid-execution, we interrupt and restart
+            max_restarts = 3  # Prevent infinite restart loops
+            restart_count = 0
+            accumulated_messages = list(langchain_history) + [HumanMessage(content=query)]
+            pending_tool_calls = {}  # Map tool_call_id -> AgentAction
+
+            while restart_count <= max_restarts:
+                # Update inputs with accumulated messages
+                inputs = {"messages": accumulated_messages}
+                should_restart = False
+
+                async for chunk in self._agent_executor.astream(
+                    inputs,
+                    stream_mode="updates",  # Get updates as they happen
+                    config={"callbacks": self.callbacks},
+                ):
+                    # chunk is a dict with node names as keys
+                    # The agent node will have 'messages' with the AI response
+                    # The tools node will have 'messages' with tool calls and results
+
+                    for node_name, node_output in chunk.items():
+                        logger.debug(f"📦 Node '{node_name}' output: {node_output}")
+
+                        # Extract messages from the node output and accumulate them
+                        if node_output is not None and "messages" in node_output:
+                            messages = node_output["messages"]
+                            if not isinstance(messages, list):
+                                messages = [messages]
+
+                            # Add new messages to accumulated messages for potential restart
+                            for msg in messages:
+                                if msg not in accumulated_messages:
+                                    accumulated_messages.append(msg)
+                            for message in messages:
+                                # Track tool calls
+                                if hasattr(message, "tool_calls") and message.tool_calls:
+                                    # Extract text content from message for the log
+                                    log_text = ""
+                                    if hasattr(message, "content"):
+                                        if isinstance(message.content, str):
+                                            log_text = message.content
+                                        elif isinstance(message.content, list):
+                                            # Extract text blocks from content array
+                                            text_parts = [
+                                                block.get("text", "") if isinstance(block, dict) else str(block)
+                                                for block in message.content
+                                                if isinstance(block, dict) and block.get("type") == "text"
+                                            ]
+                                            log_text = "\n".join(text_parts)
+
+                                    for tool_call in message.tool_calls:
+                                        tool_name = tool_call.get("name", "unknown")
+                                        tool_input = tool_call.get("args", {})
+                                        tool_call_id = tool_call.get("id")
+
+                                        action = AgentAction(tool=tool_name, tool_input=tool_input, log=log_text)
+                                        if tool_call_id:
+                                            pending_tool_calls[tool_call_id] = action
+
+                                        self.tools_used_names.append(tool_name)
+                                        steps_taken += 1
+
+                                        tool_input_str = str(tool_input)
+                                        if len(tool_input_str) > 100:
+                                            tool_input_str = tool_input_str[:97] + "..."
+                                        logger.info(f"🔧 Tool call: {tool_name} with input: {tool_input_str}")
+
+                                # Track tool results and yield AgentStep
+                                if hasattr(message, "type") and message.type == "tool":
+                                    observation = message.content
+                                    tool_call_id = getattr(message, "tool_call_id", None)
+
+                                    if tool_call_id and tool_call_id in pending_tool_calls:
+                                        action = pending_tool_calls.pop(tool_call_id)
+                                        yield (action, str(observation))
+
+                                    observation_str = str(observation)
+                                    if len(observation_str) > 100:
+                                        observation_str = observation_str[:97] + "..."
+                                    observation_str = observation_str.replace("\n", " ")
+                                    logger.info(f"📄 Tool result: {observation_str}")
+
+                                    # --- Check for tool updates after tool results (safe restart point) ---
+                                    if self.use_server_manager and self.server_manager:
+                                        current_tools = self.server_manager.tools
+                                        current_tool_names = {tool.name for tool in current_tools}
+                                        existing_tool_names = {tool.name for tool in self._tools}
+
+                                        if current_tool_names != existing_tool_names:
+                                            logger.info(
+                                                f"🔄 Tools changed during execution. "
+                                                f"New tools: {', '.join(current_tool_names)}"
+                                            )
+                                            self._tools = current_tools
+                                            # Regenerate system message with ALL current tools
+                                            await self._create_system_message_from_tools(self._tools)
+                                            # Recreate the agent executor with the new tools and system message
+                                            self._agent_executor = self._create_agent()
+
+                                            # Set restart flag - safe to restart now after tool results
+                                            should_restart = True
+                                            restart_count += 1
+                                            logger.info(
+                                                f"🔃 Restarting execution with updated tools "
+                                                f"(restart {restart_count}/{max_restarts})"
+                                            )
+                                            break  # Break out of the message loop
+
+                                # Track final AI message (without tool calls = final response)
+                                if isinstance(message, AIMessage) and not getattr(message, "tool_calls", None):
+                                    final_output = self._normalize_output(message.content)
+                                    logger.info("✅ Agent finished with output")
+
+                        # Break out of node loop if restarting
+                        if should_restart:
+                            break
+
+                    # Break out of chunk loop if restarting
+                    if should_restart:
+                        break
+
+                # Check if we should restart or if execution completed
+                if not should_restart:
+                    # Execution completed successfully without tool changes
+                    break
+
+                # If we've hit max restarts, log warning and continue
+                if restart_count > max_restarts:
+                    logger.warning(f"⚠️ Max restarts ({max_restarts}) reached. Continuing with current tools.")
+                    break
+
+            # 4. Update conversation history
+            if self.memory_enabled:
+                self.add_to_history(HumanMessage(content=query))
+                if final_output:
+                    self.add_to_history(AIMessage(content=final_output))
+
+            # 5. Handle structured output if requested
+            if output_schema and final_output:
+                try:
+                    logger.info("🔧 Attempting structured output...")
+                    structured_llm = self.llm.with_structured_output(output_schema)
+
+                    # Get schema description
+                    schema_fields = []
+                    for field_name, field_info in output_schema.model_fields.items():
+                        description = getattr(field_info, "description", "") or field_name
+                        required = not hasattr(field_info, "default") or field_info.default is None
+                        schema_fields.append(
+                            f"- {field_name}: {description} " + ("(required)" if required else "(optional)")
+                        )
+                    schema_description = "\n".join(schema_fields)
+
+                    structured_result = await self._attempt_structured_output(
+                        final_output, structured_llm, output_schema, schema_description
+                    )
+
+                    if self.memory_enabled:
+                        self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
+
+                    logger.info("✅ Structured output successful")
+                    success = True
+                    yield structured_result
+                    return
+                except Exception as e:
+                    logger.error(f"❌ Structured output failed: {e}")
+                    raise RuntimeError(f"Failed to generate structured output: {str(e)}") from e
+
+            # 6. Yield final result
+            logger.info(f"🎉 Agent execution complete in {time.time() - start_time:.2f} seconds")
+            success = True
+            yield final_output or "No output generated"
+
+        except Exception as e:
+            logger.error(f"❌ Error running query: {e}")
+            if initialized_here and manage_connector:
+                logger.info("🧹 Cleaning up resources after error")
+                await self.close()
+            raise
+
+        finally:
+            # Track comprehensive execution data
+            execution_time_ms = int((time.time() - start_time) * 1000)
+
+            server_count = 0
+            if self.client:
+                server_count = len(self.client.get_all_active_sessions())
+            elif self.connectors:
+                server_count = len(self.connectors)
+
+            conversation_history_length = len(self._conversation_history) if self.memory_enabled else 0
+
+            # Safely access _tools in case initialization failed
+            tools_available = getattr(self, "_tools", [])
+
+            if track_execution:
+                self.telemetry.track_agent_execution(
+                    execution_method="stream",
+                    query=query,
+                    success=success,
+                    model_provider=self._model_provider,
+                    model_name=self._model_name,
+                    server_count=server_count,
+                    server_identifiers=[connector.public_identifier for connector in self.connectors],
+                    total_tools_available=len(tools_available),
+                    tools_available_names=[tool.name for tool in tools_available],
+                    max_steps_configured=self.max_steps,
+                    memory_enabled=self.memory_enabled,
+                    use_server_manager=self.use_server_manager,
+                    max_steps_used=max_steps,
+                    manage_connector=manage_connector,
+                    external_history_used=external_history is not None,
+                    steps_taken=steps_taken,
+                    tools_used_count=len(self.tools_used_names),
+                    tools_used_names=self.tools_used_names,
+                    response=final_output,
+                    execution_time_ms=execution_time_ms,
+                    error_type=None if success else "execution_error",
+                    conversation_history_length=conversation_history_length,
+                )
+
+            # Clean up if necessary
+            if manage_connector and not self.client and initialized_here:
+                logger.info("🧹 Closing agent after stream completion")
+                await self.close()
+
     async def _generate_response_chunks_async(
         self,
         query: str,
@@ -1012,19 +909,21 @@ class MCPAgent:
             raise RuntimeError("MCP agent failed to initialise – call initialise() first?")
 
         # 2. Build inputs --------------------------------------------------------
-        effective_max_steps = max_steps or self.max_steps
-        self._agent_executor.max_iterations = effective_max_steps
+        self.max_steps = max_steps or self.max_steps
 
+        # 3. Build inputs --------------------------------------------------------
         history_to_use = external_history if external_history is not None else self._conversation_history
         inputs = {"input": query, "chat_history": history_to_use}
 
         # 3. Stream & diff -------------------------------------------------------
-        async for event in self._agent_executor.astream_events(inputs):
+        async for event in self._agent_executor.astream_events(inputs, config={"callbacks": self.callbacks}):
             if event.get("event") == "on_chain_end":
                 output = event["data"]["output"]
                 if isinstance(output, list):
                     for message in output:
-                        if not isinstance(message, ToolAgentAction):
+                        # Filter out ToolMessage (equivalent to old ToolAgentAction)
+                        # to avoid adding intermediate tool execution details to history
+                        if isinstance(message, BaseMessage) and not isinstance(message, ToolMessage):
                             self.add_to_history(message)
             yield event