improved guard handling

src/mcp_cli/chat/chat_context.py

@@ -236,7 +236,32 @@ def _generate_system_prompt(self) -> None:
         tools_for_prompt = [
             tool.to_llm_format().to_dict() for tool in self.internal_tools
         ]
-        self._system_prompt = generate_system_prompt(tools_for_prompt)
+        server_tool_groups = self._build_server_tool_groups()
+        self._system_prompt = generate_system_prompt(
+            tools=tools_for_prompt,
+            server_tool_groups=server_tool_groups,
+        )
+
+    def _build_server_tool_groups(self) -> list[dict[str, Any]]:
+        """Build server-to-tools grouping for the system prompt."""
+        if not self.server_info:
+            return []
+
+        # Group tools by server namespace
+        server_tools: dict[str, list[str]] = {}
+        for tool_name, namespace in self.tool_to_server_map.items():
+            server_tools.setdefault(namespace, []).append(tool_name)
+
+        groups = []
+        for server in self.server_info:
+            tools = server_tools.get(server.namespace, [])
+            if tools:
+                groups.append({
+                    "name": server.name,
+                    "description": server.display_description,
+                    "tools": sorted(tools),
+                })
+        return groups
 
     async def _initialize_tools(self) -> None:
         """Initialize tool discovery and adaptation."""

src/mcp_cli/chat/conversation.py

@@ -532,9 +532,13 @@ async def _handle_streaming_completion(
 
         try:
             # stream_response returns dict, convert to CompletionResponse
+            messages_for_api = [
+                msg.to_dict() for msg in self.context.conversation_history
+            ]
+            messages_for_api = self._validate_tool_messages(messages_for_api)
             completion_dict = await streaming_handler.stream_response(
                 client=self.context.client,
-                messages=[msg.to_dict() for msg in self.context.conversation_history],
+                messages=messages_for_api,
                 tools=tools,
             )
 
@@ -573,6 +577,7 @@ async def _handle_regular_completion(
             messages_as_dicts = [
                 msg.to_dict() for msg in self.context.conversation_history
             ]
+            messages_as_dicts = self._validate_tool_messages(messages_as_dicts)
             completion_dict = await self.context.client.create_completion(
                 messages=messages_as_dicts,
                 tools=tools,
@@ -588,6 +593,7 @@ async def _handle_regular_completion(
                 messages_as_dicts = [
                     msg.to_dict() for msg in self.context.conversation_history
                 ]
+                messages_as_dicts = self._validate_tool_messages(messages_as_dicts)
                 completion_dict = await self.context.client.create_completion(
                     messages=messages_as_dicts
                 )
@@ -632,6 +638,58 @@ async def _load_tools(self):
             self.context.openai_tools = []
             self.context.tool_name_mapping = {}
 
+    @staticmethod
+    def _validate_tool_messages(messages: list[dict]) -> list[dict]:
+        """Ensure every assistant tool_call_id has a matching tool result.
+
+        Defense-in-depth: repairs orphaned tool_calls before sending to the API.
+        Without this, OpenAI returns a 400 error:
+        "An assistant message with 'tool_calls' must be followed by tool messages
+        responding to each 'tool_call_id'."
+
+        Args:
+            messages: List of message dicts about to be sent to the API.
+
+        Returns:
+            The message list, with placeholder tool results inserted for any
+            orphaned tool_call_ids.
+        """
+        repaired: list[dict] = []
+        i = 0
+        while i < len(messages):
+            msg = messages[i]
+            repaired.append(msg)
+
+            if msg.get("role") == "assistant" and msg.get("tool_calls"):
+                # Collect expected tool_call_ids from this assistant message
+                expected_ids = set()
+                for tc in msg["tool_calls"]:
+                    tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
+                    if tc_id:
+                        expected_ids.add(tc_id)
+
+                # Scan following messages for matching tool results
+                j = i + 1
+                found_ids: set[str] = set()
+                while j < len(messages) and messages[j].get("role") == "tool":
+                    tid = messages[j].get("tool_call_id")
+                    if tid:
+                        found_ids.add(tid)
+                    j += 1
+
+                # Insert placeholders for any missing tool results
+                missing = expected_ids - found_ids
+                for mid in missing:
+                    log.warning(f"Repairing orphaned tool_call_id: {mid}")
+                    repaired.append({
+                        "role": "tool",
+                        "tool_call_id": mid,
+                        "content": "Tool call did not complete.",
+                    })
+
+            i += 1
+        return repaired
+
     def _register_user_literals_from_history(self) -> int:
         """Extract and register numeric literals from recent user messages.
 

src/mcp_cli/chat/system_prompt.py

@@ -2,14 +2,43 @@
 import os
 
 
-def generate_system_prompt(tools=None):
+def _build_server_section(server_tool_groups):
+    """Build the server/tool categorization section for the system prompt."""
+    if not server_tool_groups:
+        return ""
+
+    lines = [
+        "",
+        "**CONNECTED SERVERS & AVAILABLE TOOLS:**",
+        "",
+        "You have access to tools from the following servers. Consider using tools",
+        "from ALL relevant servers when answering a query.",
+        "",
+    ]
+    for group in server_tool_groups:
+        name = group.get("name", "unknown")
+        desc = group.get("description", "")
+        tools = group.get("tools", [])
+        tool_list = ", ".join(tools)
+        lines.append(f"- **{name}** ({desc}): {tool_list}")
+
+    lines.append("")
+    return "\n".join(lines)
+
+
+def generate_system_prompt(tools=None, server_tool_groups=None):
     """Generate a concise system prompt for the assistant.
 
     Note: Tool definitions are passed via the API's tools parameter,
     so we don't duplicate them in the system prompt.
 
     When dynamic tools mode is enabled (MCP_CLI_DYNAMIC_TOOLS=1), generates
     a special prompt explaining the tool discovery workflow.
+
+    Args:
+        tools: List of tool definitions (dicts or ToolInfo objects).
+        server_tool_groups: Optional list of dicts with server/tool grouping,
+            each containing {"name", "description", "tools"}.
     """
     # Check if dynamic tools mode is enabled
     dynamic_mode = os.environ.get("MCP_CLI_DYNAMIC_TOOLS") == "1"
@@ -20,9 +49,13 @@ def generate_system_prompt(tools=None):
     # Count tools for the prompt (tools may be ToolInfo objects or dicts)
     tool_count = len(tools) if tools else 0
 
+    # Build server/tool categorization section
+    server_section = _build_server_section(server_tool_groups)
+
     system_prompt = f"""You are an intelligent assistant with access to {tool_count} tools to help solve user queries effectively.
 
 Use the available tools when appropriate to accomplish tasks. Tools are provided via the API and you can call them as needed.
+{server_section}
 
 **GENERAL GUIDELINES:**