fast-agent-mcp 0.2.40__py3-none-any.whl → 0.2.41__py3-none-any.whl

mcp_agent/llm/model_factory.py

@@ -8,6 +8,7 @@ from mcp_agent.core.exceptions import ModelConfigError
 from mcp_agent.core.request_params import RequestParams
 from mcp_agent.llm.augmented_llm_passthrough import PassthroughLLM
 from mcp_agent.llm.augmented_llm_playback import PlaybackLLM
+from mcp_agent.llm.augmented_llm_silent import SilentLLM
 from mcp_agent.llm.augmented_llm_slow import SlowLLM
 from mcp_agent.llm.provider_types import Provider
 from mcp_agent.llm.providers.augmented_llm_aliyun import AliyunAugmentedLLM
@@ -20,6 +21,7 @@ from mcp_agent.llm.providers.augmented_llm_google_oai import GoogleOaiAugmentedL
 from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_openrouter import OpenRouterAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_tensorzero import TensorZeroAugmentedLLM
+from mcp_agent.llm.providers.augmented_llm_xai import XAIAugmentedLLM
 from mcp_agent.mcp.interfaces import AugmentedLLMProtocol
 
 # from mcp_agent.workflows.llm.augmented_llm_deepseek import DeekSeekAugmentedLLM
@@ -31,6 +33,7 @@ LLMClass = Union[
     Type[OpenAIAugmentedLLM],
     Type[PassthroughLLM],
     Type[PlaybackLLM],
+    Type[SilentLLM],
     Type[SlowLLM],
     Type[DeepSeekAugmentedLLM],
     Type[OpenRouterAugmentedLLM],
@@ -75,6 +78,7 @@ class ModelFactory:
     """
     DEFAULT_PROVIDERS = {
         "passthrough": Provider.FAST_AGENT,
+        "silent": Provider.FAST_AGENT,
         "playback": Provider.FAST_AGENT,
         "slow": Provider.FAST_AGENT,
         "gpt-4o": Provider.OPENAI,
@@ -106,6 +110,12 @@ class ModelFactory:
         "gemini-2.0-flash": Provider.GOOGLE,
         "gemini-2.5-flash-preview-05-20": Provider.GOOGLE,
         "gemini-2.5-pro-preview-05-06": Provider.GOOGLE,
+        "grok-4": Provider.XAI,
+        "grok-4-0709": Provider.XAI,
+        "grok-3": Provider.XAI,
+        "grok-3-mini": Provider.XAI,
+        "grok-3-fast": Provider.XAI,
+        "grok-3-mini-fast": Provider.XAI,
         "qwen-turbo": Provider.ALIYUN,
         "qwen-plus": Provider.ALIYUN,
         "qwen-max": Provider.ALIYUN,
@@ -140,6 +150,7 @@ class ModelFactory:
         Provider.GENERIC: GenericAugmentedLLM,
         Provider.GOOGLE_OAI: GoogleOaiAugmentedLLM,
         Provider.GOOGLE: GoogleNativeAugmentedLLM,
+        Provider.XAI: XAIAugmentedLLM,
         Provider.OPENROUTER: OpenRouterAugmentedLLM,
         Provider.TENSORZERO: TensorZeroAugmentedLLM,
         Provider.AZURE: AzureOpenAIAugmentedLLM,
@@ -150,6 +161,7 @@ class ModelFactory:
     # This overrides the provider-based class selection
     MODEL_SPECIFIC_CLASSES: Dict[str, LLMClass] = {
         "playback": PlaybackLLM,
+        "silent": SilentLLM,
         "slow": SlowLLM,
     }
 

mcp_agent/llm/provider_key_manager.py

@@ -11,12 +11,8 @@ from pydantic import BaseModel
 from mcp_agent.core.exceptions import ProviderKeyError
 
 PROVIDER_ENVIRONMENT_MAP: Dict[str, str] = {
-    "anthropic": "ANTHROPIC_API_KEY",
-    "openai": "OPENAI_API_KEY",
-    "deepseek": "DEEPSEEK_API_KEY",
-    "google": "GOOGLE_API_KEY",
-    "openrouter": "OPENROUTER_API_KEY",
-    "generic": "GENERIC_API_KEY",
+    # default behaviour in _get_env_key_name is to capitalize the
+    # provider name and suffix "_API_KEY" - so no specific mapping needed unless overriding
     "huggingface": "HF_TOKEN",
 }
 API_KEY_HINT_TEXT = "<your-api-key-here>"
@@ -66,7 +62,14 @@ class ProviderKeyManager:
             ProviderKeyError: If the API key is not found or is invalid
         """
 
+        from mcp_agent.llm.provider_types import Provider
+
         provider_name = provider_name.lower()
+
+        # Fast-agent provider doesn't need external API keys
+        if provider_name == "fast-agent":
+            return ""
+
         api_key = ProviderKeyManager.get_config_file_key(provider_name, config)
         if not api_key:
             api_key = ProviderKeyManager.get_env_var(provider_name)
@@ -75,9 +78,20 @@ class ProviderKeyManager:
             api_key = "ollama"  # Default for generic provider
 
         if not api_key:
+            # Get proper display name for error message
+            try:
+                provider_enum = Provider(provider_name)
+                display_name = provider_enum.display_name
+            except ValueError:
+                # Invalid provider name
+                raise ProviderKeyError(
+                    f"Invalid provider: {provider_name}",
+                    f"'{provider_name}' is not a valid provider name.",
+                )
+
             raise ProviderKeyError(
-                f"{provider_name.title()} API key not configured",
-                f"The {provider_name.title()} API key is required but not set.\n"
+                f"{display_name} API key not configured",
+                f"The {display_name} API key is required but not set.\n"
                 f"Add it to your configuration file under {provider_name}.api_key "
                 f"or set the {ProviderKeyManager.get_env_key_name(provider_name)} environment variable.",
             )

mcp_agent/llm/provider_types.py

@@ -8,15 +8,22 @@ from enum import Enum
 class Provider(Enum):
     """Supported LLM providers"""
 
-    ANTHROPIC = "anthropic"
-    DEEPSEEK = "deepseek"
-    FAST_AGENT = "fast-agent"
-    GENERIC = "generic"
-    GOOGLE_OAI = "googleoai"  # For Google through OpenAI libraries
-    GOOGLE = "google"  # For Google GenAI native library
-    OPENAI = "openai"
-    OPENROUTER = "openrouter"
-    TENSORZERO = "tensorzero"  # For TensorZero Gateway
-    AZURE = "azure"  # Azure OpenAI Service
-    ALIYUN = "aliyun"  # Aliyun Bailian OpenAI Service
-    HUGGINGFACE = "huggingface"  # For HuggingFace MCP connections
+    def __new__(cls, config_name, display_name=None):
+        obj = object.__new__(cls)
+        obj._value_ = config_name
+        obj.display_name = display_name or config_name.title()
+        return obj
+
+    ANTHROPIC = ("anthropic", "Anthropic")
+    DEEPSEEK = ("deepseek", "Deepseek")
+    FAST_AGENT = ("fast-agent", "FastAgent")
+    GENERIC = ("generic", "Generic")
+    GOOGLE_OAI = ("googleoai", "GoogleOAI")  # For Google through OpenAI libraries
+    GOOGLE = ("google", "Google")  # For Google GenAI native library
+    OPENAI = ("openai", "OpenAI")
+    OPENROUTER = ("openrouter", "OpenRouter")
+    TENSORZERO = ("tensorzero", "TensorZero")  # For TensorZero Gateway
+    AZURE = ("azure", "Azure")  # Azure OpenAI Service
+    ALIYUN = ("aliyun", "Aliyun")  # Aliyun Bailian OpenAI Service
+    HUGGINGFACE = ("huggingface", "HuggingFace")  # For HuggingFace MCP connections
+    XAI = ("xai", "XAI")  # For xAI Grok models

mcp_agent/llm/providers/augmented_llm_anthropic.py

@@ -112,7 +112,9 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
                 and event.delta.type == "text_delta"
             ):
                 # Use base class method for token estimation and progress emission
-                estimated_tokens = self._update_streaming_progress(event.delta.text, model, estimated_tokens)
+                estimated_tokens = self._update_streaming_progress(
+                    event.delta.text, model, estimated_tokens
+                )
 
             # Also check for final message_delta events with actual usage info
             elif (
@@ -285,7 +287,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
                     turn_usage = TurnUsage.from_anthropic(
                         response.usage, model or DEFAULT_ANTHROPIC_MODEL
                     )
-                    self.usage_accumulator.add_turn(turn_usage)
+                    self._finalize_turn_usage(turn_usage)
                 #                    self._show_usage(response.usage, turn_usage)
                 except Exception as e:
                     self.logger.warning(f"Failed to track usage: {e}")
@@ -435,6 +437,9 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         Override this method to use a different LLM.
 
         """
+        # Reset tool call counter for new turn
+        self._reset_turn_tool_calls()
+
         res = await self._anthropic_completion(
             message_param=message_param,
             request_params=request_params,

mcp_agent/llm/providers/augmented_llm_azure.py

@@ -69,7 +69,7 @@ class AzureOpenAIAugmentedLLM(OpenAIAugmentedLLM):
 
             self.get_azure_token = get_azure_token
         else:
-            self.api_key = getattr(azure_cfg, "api_key", None)
+            self.api_key = self._api_key()
             self.resource_name = getattr(azure_cfg, "resource_name", None)
             self.base_url = getattr(azure_cfg, "base_url", None) or (
                 f"https://{self.resource_name}.openai.azure.com/" if self.resource_name else None
@@ -93,6 +93,12 @@ class AzureOpenAIAugmentedLLM(OpenAIAugmentedLLM):
             if not self.resource_name and self.base_url:
                 self.resource_name = _extract_resource_name(self.base_url)
 
+    def _api_key(self):
+        """Override to return 'AzureCredential' when using DefaultAzureCredential"""
+        if self.use_default_cred:
+            return "AzureCredential"
+        return super()._api_key()
+
     def _openai_client(self) -> AsyncOpenAI:
         """
         Returns an AzureOpenAI client, handling both API Key and DefaultAzureCredential.

mcp_agent/llm/providers/augmented_llm_google_native.py

@@ -295,7 +295,7 @@ class GoogleNativeAugmentedLLM(AugmentedLLM[types.Content, types.Content]):
                         turn_usage = TurnUsage.from_google(
                             api_response.usage_metadata, request_params.model
                         )
-                        self.usage_accumulator.add_turn(turn_usage)
+                        self._finalize_turn_usage(turn_usage)
 
                     except Exception as e:
                         self.logger.warning(f"Failed to track usage: {e}")
@@ -439,6 +439,9 @@ class GoogleNativeAugmentedLLM(AugmentedLLM[types.Content, types.Content]):
         """
         Applies the prompt messages and potentially calls the LLM for completion.
         """
+        # Reset tool call counter for new turn
+        self._reset_turn_tool_calls()
+
         request_params = self.get_request_params(
             request_params=request_params
         )  # Get request params

mcp_agent/llm/providers/augmented_llm_openai.py

@@ -108,6 +108,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
     def _openai_client(self) -> AsyncOpenAI:
         try:
             return AsyncOpenAI(api_key=self._api_key(), base_url=self._base_url())
+
         except AuthenticationError as e:
             raise ProviderKeyError(
                 "Invalid OpenAI API key",
@@ -355,7 +356,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
                 try:
                     model_name = self.default_request_params.model or DEFAULT_OPENAI_MODEL
                     turn_usage = TurnUsage.from_openai(response.usage, model_name)
-                    self.usage_accumulator.add_turn(turn_usage)
+                    self._finalize_turn_usage(turn_usage)
                 except Exception as e:
                     self.logger.warning(f"Failed to track usage: {e}")
 
@@ -389,7 +390,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
             messages.append(message)
 
             message_text = message.content
-            if choice.finish_reason in ["tool_calls", "function_call"] and message.tool_calls:
+            if await self._is_tool_stop_reason(choice.finish_reason) and message.tool_calls:
                 if message_text:
                     await self.show_assistant_message(
                         message_text,
@@ -477,12 +478,18 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
 
         return responses
 
+    async def _is_tool_stop_reason(self, finish_reason: str) -> bool:
+        return True
+
     async def _apply_prompt_provider_specific(
         self,
         multipart_messages: List["PromptMessageMultipart"],
         request_params: RequestParams | None = None,
         is_template: bool = False,
     ) -> PromptMessageMultipart:
+        # Reset tool call counter for new turn
+        self._reset_turn_tool_calls()
+
         last_message = multipart_messages[-1]
 
         # Add all previous messages to history (or all messages if last is from assistant)

mcp_agent/llm/providers/augmented_llm_xai.py

@@ -0,0 +1,38 @@
+import os
+
+from mcp_agent.core.request_params import RequestParams
+from mcp_agent.llm.provider_types import Provider
+from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
+
+XAI_BASE_URL = "https://api.x.ai/v1"
+DEFAULT_XAI_MODEL = "grok-3"
+
+
+class XAIAugmentedLLM(OpenAIAugmentedLLM):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(
+            *args, provider=Provider.XAI, **kwargs
+        )  # Properly pass args and kwargs to parent
+
+    def _initialize_default_params(self, kwargs: dict) -> RequestParams:
+        """Initialize xAI parameters"""
+        chosen_model = kwargs.get("model", DEFAULT_XAI_MODEL)
+
+        return RequestParams(
+            model=chosen_model,
+            systemPrompt=self.instruction,
+            parallel_tool_calls=False,
+            max_iterations=10,
+            use_history=True,
+        )
+
+    def _base_url(self) -> str:
+        base_url = os.getenv("XAI_BASE_URL", XAI_BASE_URL)
+        if self.context.config and self.context.config.xai:
+            base_url = self.context.config.xai.base_url
+
+        return base_url
+
+    async def _is_tool_stop_reason(self, finish_reason: str) -> bool:
+        # grok uses Null as the finish reason for tool calls?
+        return await super()._is_tool_stop_reason(finish_reason) or finish_reason is None

mcp_agent/llm/usage_tracking.py

@@ -78,6 +78,9 @@ class TurnUsage(BaseModel):
     tool_use_tokens: int = Field(default=0, description="Tokens used for tool calling prompts")
     reasoning_tokens: int = Field(default=0, description="Tokens used for reasoning/thinking")
 
+    # Tool call count for this turn
+    tool_calls: int = Field(default=0, description="Number of tool calls made in this turn")
+
     # Raw usage data from provider (preserves all original data)
     raw_usage: ProviderUsage
 
@@ -86,7 +89,11 @@ class TurnUsage(BaseModel):
     def current_context_tokens(self) -> int:
         """Current context size after this turn (total input including cache + output)"""
         # For Anthropic: input_tokens + cache_read_tokens represents total input context
-        total_input = self.input_tokens + self.cache_usage.cache_read_tokens + self.cache_usage.cache_write_tokens
+        total_input = (
+            self.input_tokens
+            + self.cache_usage.cache_read_tokens
+            + self.cache_usage.cache_write_tokens
+        )
         return total_input + self.output_tokens
 
     @computed_field
@@ -106,11 +113,20 @@ class TurnUsage(BaseModel):
         """Input tokens to display for 'Last turn' (total submitted tokens)"""
         # For Anthropic: input_tokens excludes cache, so add cache tokens
         if self.provider == Provider.ANTHROPIC:
-            return self.input_tokens + self.cache_usage.cache_read_tokens + self.cache_usage.cache_write_tokens
+            return (
+                self.input_tokens
+                + self.cache_usage.cache_read_tokens
+                + self.cache_usage.cache_write_tokens
+            )
         else:
             # For OpenAI/Google: input_tokens already includes cached tokens
             return self.input_tokens
 
+    def set_tool_calls(self, count: int) -> None:
+        """Set the number of tool calls made in this turn"""
+        # Use object.__setattr__ since this is a Pydantic model
+        object.__setattr__(self, "tool_calls", count)
+
     @classmethod
     def from_anthropic(cls, usage: AnthropicUsage, model: str) -> "TurnUsage":
         # Extract cache tokens with proper null handling
@@ -219,7 +235,9 @@ class UsageAccumulator(BaseModel):
     def cumulative_input_tokens(self) -> int:
         """Total input tokens charged across all turns (including cache tokens)"""
         return sum(
-            turn.input_tokens + turn.cache_usage.cache_read_tokens + turn.cache_usage.cache_write_tokens
+            turn.input_tokens
+            + turn.cache_usage.cache_read_tokens
+            + turn.cache_usage.cache_write_tokens
             for turn in self.turns
         )
 
@@ -247,6 +265,12 @@ class UsageAccumulator(BaseModel):
         """Total tokens written to cache across all turns"""
         return sum(turn.cache_usage.cache_write_tokens for turn in self.turns)
 
+    @computed_field
+    @property
+    def cumulative_tool_calls(self) -> int:
+        """Total tool calls made across all turns"""
+        return sum(turn.tool_calls for turn in self.turns)
+
     @computed_field
     @property
     def cumulative_cache_hit_tokens(self) -> int:
@@ -333,6 +357,7 @@ class UsageAccumulator(BaseModel):
             "cumulative_billing_tokens": self.cumulative_billing_tokens,
             "cumulative_tool_use_tokens": self.cumulative_tool_use_tokens,
             "cumulative_reasoning_tokens": self.cumulative_reasoning_tokens,
+            "cumulative_tool_calls": self.cumulative_tool_calls,
             "current_context_tokens": self.current_context_tokens,
             "context_window_size": self.context_window_size,
             "context_usage_percentage": self.context_usage_percentage,

mcp_agent/mcp/mcp_agent_client_session.py

@@ -78,6 +78,8 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         self.agent_model: str | None = kwargs.pop("agent_model", None)
         # Extract agent_name if provided
         self.agent_name: str | None = kwargs.pop("agent_name", None)
+        # Extract api_key if provided
+        self.api_key: str | None = kwargs.pop("api_key", None)
         # Extract custom elicitation handler if provided
         custom_elicitation_handler = kwargs.pop("elicitation_handler", None)
 

mcp_agent/mcp/mcp_aggregator.py

@@ -221,6 +221,7 @@ class MCPAggregator(ContextDependent):
                     agent_model: str | None = None
                     agent_name: str | None = None
                     elicitation_handler = None
+                    api_key: str | None = None
 
                     # Check if this aggregator is part of an Agent (which has config)
                     # Import here to avoid circular dependency
@@ -230,6 +231,7 @@ class MCPAggregator(ContextDependent):
                         agent_model = self.config.model
                         agent_name = self.config.name
                         elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
 
                     return MCPAgentClientSession(
                         read_stream,
@@ -238,6 +240,7 @@ class MCPAggregator(ContextDependent):
                         server_name=server_name,
                         agent_model=agent_model,
                         agent_name=agent_name,
+                        api_key=api_key,
                         elicitation_handler=elicitation_handler,
                         tool_list_changed_callback=self._handle_tool_list_changed,
                         **kwargs,  # Pass through any additional kwargs like server_config
@@ -292,6 +295,8 @@ class MCPAggregator(ContextDependent):
                     # Get agent's model and name if this aggregator is part of an agent
                     agent_model: str | None = None
                     agent_name: str | None = None
+                    elicitation_handler = None
+                    api_key: str | None = None
 
                     # Check if this aggregator is part of an Agent (which has config)
                     # Import here to avoid circular dependency
@@ -300,6 +305,8 @@ class MCPAggregator(ContextDependent):
                     if isinstance(self, BaseAgent):
                         agent_model = self.config.model
                         agent_name = self.config.name
+                        elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
 
                     return MCPAgentClientSession(
                         read_stream,
@@ -308,6 +315,8 @@ class MCPAggregator(ContextDependent):
                         server_name=server_name,
                         agent_model=agent_model,
                         agent_name=agent_name,
+                        api_key=api_key,
+                        elicitation_handler=elicitation_handler,
                         tool_list_changed_callback=self._handle_tool_list_changed,
                         **kwargs,  # Pass through any additional kwargs like server_config
                     )
@@ -957,58 +966,43 @@ class MCPAggregator(ContextDependent):
 
         async with self._refresh_lock:
             try:
+                # Create a factory function that will include our parameters
+                def create_session(read_stream, write_stream, read_timeout):
+                    # Get agent name if available
+                    agent_model: str | None = None
+                    agent_name: str | None = None
+                    elicitation_handler = None
+                    api_key: str | None = None
+
+                    # Import here to avoid circular dependency
+                    from mcp_agent.agents.base_agent import BaseAgent
+
+                    if isinstance(self, BaseAgent):
+                        agent_model = self.config.model
+                        agent_name = self.config.name
+                        elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
+
+                    return MCPAgentClientSession(
+                        read_stream,
+                        write_stream,
+                        read_timeout,
+                        server_name=server_name,
+                        agent_model=agent_model,
+                        agent_name=agent_name,
+                        api_key=api_key,
+                        elicitation_handler=elicitation_handler,
+                        tool_list_changed_callback=self._handle_tool_list_changed,
+                    )
+
                 # Fetch new tools from the server
                 if self.connection_persistence:
-                    # Create a factory function that will include our parameters
-                    def create_session(read_stream, write_stream, read_timeout):
-                        # Get agent name if available
-                        agent_name: str | None = None
-
-                        # Import here to avoid circular dependency
-                        from mcp_agent.agents.base_agent import BaseAgent
-
-                        if isinstance(self, BaseAgent):
-                            agent_name = self.config.name
-                            elicitation_handler = self.config.elicitation_handler
-
-                        return MCPAgentClientSession(
-                            read_stream,
-                            write_stream,
-                            read_timeout,
-                            server_name=server_name,
-                            agent_name=agent_name,
-                            elicitation_handler=elicitation_handler,
-                            tool_list_changed_callback=self._handle_tool_list_changed,
-                        )
-
                     server_connection = await self._persistent_connection_manager.get_server(
                         server_name, client_session_factory=create_session
                     )
                     tools_result = await server_connection.session.list_tools()
                     new_tools = tools_result.tools or []
                 else:
-                    # Create a factory function for the client session
-                    def create_session(read_stream, write_stream, read_timeout):
-                        # Get agent name if available
-                        agent_name: str | None = None
-
-                        # Import here to avoid circular dependency
-                        from mcp_agent.agents.base_agent import BaseAgent
-
-                        if isinstance(self, BaseAgent):
-                            agent_name = self.config.name
-                            elicitation_handler = self.config.elicitation_handler
-
-                        return MCPAgentClientSession(
-                            read_stream,
-                            write_stream,
-                            read_timeout,
-                            server_name=server_name,
-                            agent_name=agent_name,
-                            elicitation_handler=elicitation_handler,
-                            tool_list_changed_callback=self._handle_tool_list_changed,
-                        )
-
                     async with gen_client(
                         server_name,
                         server_registry=self.context.server_registry,

mcp_agent/mcp/sampling.py

@@ -20,7 +20,7 @@ logger = get_logger(__name__)
 
 
 def create_sampling_llm(
-    params: CreateMessageRequestParams, model_string: str
+    params: CreateMessageRequestParams, model_string: str, api_key: str | None
 ) -> AugmentedLLMProtocol:
     """
     Create an LLM instance for sampling without tools support.
@@ -52,7 +52,7 @@ def create_sampling_llm(
 
     # Create the LLM using the factory
     factory = ModelFactory.create_factory(model_string)
-    llm = factory(agent=agent)
+    llm = factory(agent=agent, api_key=api_key)
 
     # Attach the LLM to the agent
     agent._llm = llm
@@ -77,7 +77,8 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
     Returns:
         A CreateMessageResult containing the LLM's response
     """
-    model = None
+    model: str | None = None
+    api_key: str | None = None
     try:
         # Extract model from server config using type-safe helper
         server_config = get_server_config(mcp_ctx)
@@ -104,13 +105,16 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
                 from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
                 
                 # Try agent's model first (from the session)
-                if (hasattr(mcp_ctx, 'session') and 
-                    isinstance(mcp_ctx.session, MCPAgentClientSession) and
-                    mcp_ctx.session.agent_model):
-                    model = mcp_ctx.session.agent_model
-                    logger.debug(f"Using agent's model for sampling: {model}")
-                else:
-                    # Fall back to system default model
+                if hasattr(mcp_ctx, "session") and isinstance(mcp_ctx.session, MCPAgentClientSession):
+                    if mcp_ctx.session.agent_model:
+                        model = mcp_ctx.session.agent_model
+                        logger.debug(f"Using agent's model for sampling: {model}")
+                    if mcp_ctx.session.api_key:
+                        api_key = mcp_ctx.session.api_key
+                        logger.debug(f"Using agent's API KEY for sampling: {api_key}")
+
+                # Fall back to system default model
+                if model is None:
                     try:
                         if app_context and app_context.config and app_context.config.default_model:
                             model = app_context.config.default_model
@@ -122,7 +126,7 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
             raise ValueError("No model configured for sampling (server config, agent model, or system default)")
 
         # Create an LLM instance
-        llm = create_sampling_llm(params, model)
+        llm = create_sampling_llm(params, model, api_key)
 
         # Extract all messages from the request params
         if not params.messages:

mcp_agent/resources/examples/mcp/elicitations/forms_demo.py

@@ -51,8 +51,6 @@ async def main():
         else:
             console.print("[red]No registration data received[/red]")
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 2: Product Review
         console.print("[bold yellow]Example 2: Product Review Form[/bold yellow]")
         console.print(
@@ -66,8 +64,6 @@ async def main():
             )
             console.print(review_panel)
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 3: Account Settings
         console.print("[bold yellow]Example 3: Account Settings Form[/bold yellow]")
         console.print(
@@ -81,8 +77,6 @@ async def main():
             )
             console.print(settings_panel)
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 4: Service Appointment
         console.print("[bold yellow]Example 4: Service Appointment Booking[/bold yellow]")
         console.print(

mcp_agent/resources/examples/workflows/router.py

@@ -7,6 +7,8 @@ Demonstrates router's ability to either:
 
 import asyncio
 
+from rich.console import Console
+
 from mcp_agent.core.fastagent import FastAgent
 
 # Create the application
@@ -45,7 +47,14 @@ SAMPLE_REQUESTS = [
     agents=["code_expert", "general_assistant", "fetcher"],
 )
 async def main() -> None:
+    console = Console()
+    console.print(
+        "\n[bright_red]Router Workflow Demo[/bright_red]\n\n"
+        "Enter a request to route it to the appropriate agent.\nEnter [bright_red]STOP[/bright_red] to run the demo, [bright_red]EXIT[/bright_red] to leave"
+    )
+
     async with fast.run() as agent:
+        await agent.interactive(agent_name="route")
         for request in SAMPLE_REQUESTS:
             await agent.route(request)