Fix MCP integration and operational issues for successful assessments

mpnikhil · mpnikhil · commit 0f33ad19c51d · 2026-01-12T16:26:28.000-08:00
Resolves critical bugs preventing agent-to-agent shopping workflow:

- Fix MCP 421 errors: Disable DNS rebinding protection in FastMCP for
  trusted Docker network communication (green_agent/src/webshop_mcp/server.py)

- Fix WebShop data access: Add WEBSHOP_DIR environment variable support
  and mount full webshop directory for templates (docker-compose.yml,
  green_agent/src/webshop_patched/engine.py)

- Fix LLM connectivity: Support LLM_MODEL and LLM_API_BASE environment
  variables for Docker deployment (purple_agent/src/shopping_agent.py)

- Fix actions_taken tracking: Add set_action_count() method to sync MCP
  session action counts with SessionState for accurate evaluation
  (green_agent/src/models.py, green_agent/src/agent.py)

All fixes verified with successful 3-task assessment showing proper tool
execution (search, click, checkout) and accurate action tracking.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -21,16 +21,21 @@ services:
       dockerfile: Dockerfile
     image: ghcr.io/mpnikhil/webshop-plus-green:latest
     container_name: webshop-plus-green
+    command: [".venv/bin/python", "src/server.py", "--host", "0.0.0.0", "--port", "8000", "--advertise-host", "green", "--card-url", "http://localhost:8000"]
     ports:
       - "8000:8000"
     environment:
       - LLM_MODEL=${LLM_MODEL:-nebius/Qwen/Qwen3-32B}
       - LLM_API_KEY=${LLM_API_KEY:-}
+      - LLM_API_BASE=${LLM_API_BASE:-}
+      - LLM_TEMPERATURE=${LLM_TEMPERATURE:-0.2}
+      - LLM_MAX_TOKENS=${LLM_MAX_TOKENS:-8192}
       - WEBSHOP_MODE=${WEBSHOP_MODE:-preview}
+      - WEBSHOP_DIR=/app/webshop
       - LOG_LEVEL=${LOG_LEVEL:-INFO}
       - PURPLE_AGENT_URL=http://purple:8001/a2a
     volumes:
-      - ./webshop/data:/app/webshop/data:ro
+      - ./webshop:/app/webshop:ro
     depends_on:
       purple:
         condition: service_healthy
@@ -51,11 +56,15 @@ services:
       dockerfile: Dockerfile
     image: ghcr.io/mpnikhil/webshop-plus-purple:latest
     container_name: webshop-plus-purple
+    command: [".venv/bin/python", "src/server.py", "--host", "0.0.0.0", "--port", "8001", "--card-url", "http://purple:8001"]
     ports:
       - "8001:8001"
     environment:
       - LLM_MODEL=${LLM_MODEL:-nebius/Qwen/Qwen3-32B}
       - LLM_API_KEY=${LLM_API_KEY:-}
+      - LLM_API_BASE=${LLM_API_BASE:-}
+      - LLM_TEMPERATURE=${LLM_TEMPERATURE:-0.2}
+      - LLM_MAX_TOKENS=${LLM_MAX_TOKENS:-8192}
       - LOG_LEVEL=${LOG_LEVEL:-INFO}
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8001/health"]
diff --git a/green_agent/src/agent.py b/green_agent/src/agent.py
@@ -597,12 +597,19 @@ async def _dispatch_task_to_purple(
                         if "actions" in task_result.result_data:
                             result.actions_taken = len(task_result.result_data["actions"])
                         if "turns_used" in task_result.result_data:
+                            logger.info(
+                                "Setting actions_taken from purple agent result_data",
+                                turns_used=task_result.result_data["turns_used"],
+                                task_id=task.task_id,
+                            )
                             result.actions_taken = task_result.result_data["turns_used"]
 
                     logger.info(
                         "MCP task completed successfully",
                         task_id=task.task_id,
                         result_data=task_result.result_data,
+                        actions_taken=result.actions_taken,
+                        mcp_session_id=mcp_session_id,
                     )
                 else:
                     result.error = task_result.error or "Task failed"
@@ -614,16 +621,41 @@ async def _dispatch_task_to_purple(
 
             # Get final result from MCP session if available
             if mcp_session_id and self._session_manager:
-                if is_session_completed(mcp_session_id):
+                session_completed = is_session_completed(mcp_session_id)
+                logger.info(
+                    "Checking MCP session completion",
+                    mcp_session_id=mcp_session_id,
+                    session_completed=session_completed,
+                    task_id=task.task_id,
+                )
+                if session_completed:
                     mcp_result = get_final_result(mcp_session_id)
+                    logger.info(
+                        "Retrieved MCP final result",
+                        mcp_session_id=mcp_session_id,
+                        has_result=mcp_result is not None,
+                        result_keys=list(mcp_result.keys()) if mcp_result else [],
+                        task_id=task.task_id,
+                    )
                     if mcp_result:
                         # Merge MCP result into task execution result
                         if "turns_used" in mcp_result:
+                            logger.info(
+                                "Setting actions_taken from MCP result",
+                                turns_used=mcp_result["turns_used"],
+                                task_id=task.task_id,
+                            )
                             result.actions_taken = mcp_result["turns_used"]
                         if "success" in mcp_result:
                             result.completed = mcp_result["success"]
                         if "reward" in mcp_result:
                             result.total_reward = mcp_result["reward"]
+                else:
+                    logger.warning(
+                        "MCP session not marked as completed",
+                        mcp_session_id=mcp_session_id,
+                        task_id=task.task_id,
+                    )
 
                 # Clean up MCP session
                 await self._session_manager.cleanup_session(mcp_session_id)
@@ -667,6 +699,10 @@ def _finalize_task(
         memory: Optional[AgentMemory],
     ) -> TaskExecutionResult:
         """Finalize task execution and evaluate."""
+        # Update session with actions_taken from result
+        # (result.actions_taken was set from MCP final result)
+        session.set_action_count(result.actions_taken)
+
         # Complete session
         session_summary = self.state_manager.complete_session(
             session.session_id,
diff --git a/green_agent/src/models.py b/green_agent/src/models.py
@@ -389,6 +389,23 @@ def record_action(self, action: str, observation: str, reward: float = 0.0) -> N
         )
         self.current_observation = observation
 
+    def set_action_count(self, count: int) -> None:
+        """Set the action count (for MCP sessions where actions are tracked externally).
+
+        Args:
+            count: Number of actions to set.
+        """
+        # Clear existing actions and add placeholder records to match count
+        self.actions.clear()
+        for i in range(count):
+            self.actions.append(
+                ActionRecord(
+                    action=f"MCP action {i+1}",
+                    observation="Tracked by MCP session",
+                    reward=0.0,
+                )
+            )
+
     def complete(self) -> None:
         """Mark the session as completed."""
         self.completed = True
diff --git a/green_agent/src/webshop_mcp/server.py b/green_agent/src/webshop_mcp/server.py
@@ -21,6 +21,7 @@
 import structlog
 from bs4 import BeautifulSoup
 from mcp.server.fastmcp import FastMCP
+from mcp.server.transport_security import TransportSecuritySettings
 
 from .session_state import SessionState
 
@@ -190,7 +191,20 @@ def product_item_dict(self) -> dict[str, dict]:
 # Global FastMCP Server
 # =============================================================================
 
-mcp = FastMCP("WebShop MCP Server")
+# Disable DNS rebinding protection for MCP in trusted environment
+# MCP sessions are already isolated via unique session IDs
+# Purple agents are authenticated via A2A protocol before receiving MCP URIs
+transport_security = TransportSecuritySettings(
+    enable_dns_rebinding_protection=False
+)
+
+mcp = FastMCP(
+    "WebShop MCP Server",
+    transport_security=transport_security,
+)
+
+# Add logging for MCP server initialization
+logger.info("FastMCP server created", transport_security_enabled=transport_security.enable_dns_rebinding_protection)
 
 
 # =============================================================================
@@ -713,7 +727,15 @@ def search(query: str) -> dict[str, Any]:
         - turn: Current turn number
         - budget: Remaining budget info
     """
-    state = get_current_state()
+    logger.info("search() tool called", query=query)
+
+    try:
+        state = get_current_state()
+    except Exception as e:
+        logger.exception("Failed to get current state", error=str(e), query=query)
+        raise
+
+    logger.info("Got session state", session_id=state.session_id)
 
     # Check turn limit first
     if state.increment_turn():
@@ -727,8 +749,15 @@ def search(query: str) -> dict[str, Any]:
     })
 
     # Execute search via WebShop
-    webshop = _get_webshop(state.session_id)
-    result = webshop.step(f"search[{query}]")
+    try:
+        logger.info("About to call _get_webshop", session_id=state.session_id)
+        webshop = _get_webshop(state.session_id)
+        logger.info("Got WebShop interface, calling step()")
+        result = webshop.step(f"search[{query}]")
+        logger.info("WebShop step() completed", observation_len=len(result.observation) if hasattr(result, 'observation') else 0)
+    except Exception as e:
+        logger.exception("Error during WebShop search", error=str(e), query=query, session_id=state.session_id)
+        raise
 
     # Update page state
     state.current_page = "search_results"
diff --git a/green_agent/src/webshop_patched/engine.py b/green_agent/src/webshop_patched/engine.py
@@ -19,9 +19,10 @@
 from tqdm import tqdm
 
 
-# Paths - relative to this file
+# Paths - check environment variable first, then calculate relative path
 BASE_DIR = dirname(abspath(__file__))
-WEBSHOP_DIR = join(dirname(dirname(dirname(BASE_DIR))), "webshop")
+# In Docker: /app/webshop, otherwise calculate from BASE_DIR
+WEBSHOP_DIR = os.environ.get("WEBSHOP_DIR", join(dirname(dirname(dirname(BASE_DIR))), "webshop"))
 WEBSHOP_DATA_DIR = join(WEBSHOP_DIR, "data")
 TEMPLATE_DIR = join(WEBSHOP_DIR, "web_agent_site", "templates")
 
diff --git a/purple_agent/src/shopping_agent.py b/purple_agent/src/shopping_agent.py
@@ -49,12 +49,14 @@
 # Default model for the ADK agent using LiteLLM
 # Uses openai/ prefix for LM Studio's OpenAI-compatible API
 # Can be overridden via environment variable
-DEFAULT_MODEL = os.environ.get("ADK_MODEL", "openai/qwen3-coder-30b-a3b-instruct-mlx")
+# Use LLM_MODEL (from Docker) or ADK_MODEL (from local) or default model
+DEFAULT_MODEL = os.environ.get("LLM_MODEL") or os.environ.get("ADK_MODEL", "openai/qwen3-coder-30b-a3b-instruct-mlx")
 
 # LM Studio API base URL (for local development)
 # LM Studio provides an OpenAI-compatible API
-OPENAI_API_BASE = os.environ.get("OPENAI_API_BASE", "http://localhost:1234/v1")
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "lm-studio")
+# Use LLM_API_BASE (from Docker) or OPENAI_API_BASE (from local) or default to localhost
+OPENAI_API_BASE = os.environ.get("LLM_API_BASE") or os.environ.get("OPENAI_API_BASE", "http://localhost:1234/v1")
+OPENAI_API_KEY = os.environ.get("LLM_API_KEY") or os.environ.get("OPENAI_API_KEY", "lm-studio")
 
 # Maximum number of turns the agent can take before timing out
 DEFAULT_MAX_TURNS = 30
