""" MCPMark Agent Implementation ============================ Unified agent using LiteLLM for all model interactions with minimal MCP support. """ import asyncio import json import os import time from typing import Any, Dict, List, Optional, Callable from pydantic import AnyUrl import httpx import litellm import nest_asyncio from src.logger import get_logger from .base_agent import BaseMCPAgent from .mcp import MCPStdioServer, MCPHttpServer # Apply nested asyncio support nest_asyncio.apply() # Configure LiteLLM litellm.suppress_debug_info = True logger = get_logger(__name__) # To fix the "Object of type AnyUrl is not JSON serializable" error in the find_file_contents function. class CustomJSONEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, AnyUrl): return str(obj) return super().default(obj) class MCPMarkAgent(BaseMCPAgent): """ Unified agent for LLM and MCP server management using LiteLLM. - Anthropic models: Native MCP support via extra_body - Other models: Manual MCP server management with function calling """ MAX_TURNS = 100 SYSTEM_PROMPT = ( "You are a helpful agent that uses tools iteratively to complete the user's task, " 'and when finished, provides the final answer or simply states "Task completed" without further tool calls.' ) COMPACTION_PROMPT = ( "You are performing a CONTEXT CHECKPOINT COMPACTION.\n" "Summarize the conversation so far for another model to continue.\n\n" "Include:\n" "- Current progress and key decisions made\n" "- Important context, constraints, or user preferences\n" "- What remains to be done (clear next steps)\n" "- Any critical data, examples, or references needed to continue\n\n" "Be concise and structured. Do NOT call tools." ) # Per-tool-response summarizer (mirrors mcp-universe's # TOOL_RESPONSE_SUMMARIZER_PROMPT). Applied to browser tool results BEFORE # they enter the history, so the model conditions on — and the trajectory # records — the SAME compressed text. "Extract ALL relevant info" + "keep # structure" is what preserves the exact values a webarena task must read # off the page (counts, votes, prices) through the compression. TOOL_RESPONSE_SUMMARIZER_PROMPT = ( "Extract all information from the tool response that is relevant to the context.\n\n" "Tool Call Context:\n{context}\n\n" "Tool Response:\n{tool_response}\n\n" "Directly output the extracted information. Try to maintain the original " "response structure. Use fewer than 500 words." ) DEFAULT_TIMEOUT = BaseMCPAgent.DEFAULT_TIMEOUT def __init__( self, litellm_input_model_name: str, api_key: str, base_url: str, mcp_service: str, timeout: int = DEFAULT_TIMEOUT, service_config: Optional[Dict[str, Any]] = None, service_config_provider: Optional[Callable[[], Dict[str, Any]]] = None, reasoning_effort: Optional[str] = "default", compaction_token: int = BaseMCPAgent.COMPACTION_DISABLED_TOKEN, extra_body: Optional[Dict[str, Any]] = None, summarize_tool_response: bool = False, ): super().__init__( litellm_input_model_name=litellm_input_model_name, api_key=api_key, base_url=base_url, mcp_service=mcp_service, timeout=timeout, service_config=service_config, service_config_provider=service_config_provider, reasoning_effort=reasoning_effort, compaction_token=compaction_token, extra_body=extra_body, summarize_tool_response=summarize_tool_response, ) logger.debug( "Initialized MCPMarkAgent for '%s' with model '%s' (Claude: %s, Thinking: %s, Reasoning: %s)", mcp_service, litellm_input_model_name, self.is_claude, self.use_claude_thinking, reasoning_effort, ) # ==================== Public Interface Methods ==================== async def execute( self, instruction: str, tool_call_log_file: Optional[str] = None ) -> Dict[str, Any]: """ Execute instruction with the agent. Args: instruction: The instruction/prompt to execute tool_call_log_file: Optional path to log tool calls Returns: Dictionary containing execution results """ start_time = time.time() try: # Reset partial progress for this run self._reset_progress() # Refresh service configuration self._refresh_service_config() # Execute with timeout control async def _execute_with_strategy(): if self.use_claude_thinking: # Claude with thinking -> native Anthropic API with tools return await self._execute_claude_native_with_tools( instruction, tool_call_log_file ) else: # All other cases -> LiteLLM with tools return await self._execute_litellm_with_tools( instruction, tool_call_log_file ) # Apply timeout to the entire execution result = await asyncio.wait_for( _execute_with_strategy(), timeout=self.timeout ) execution_time = time.time() - start_time # Update usage statistics self.usage_tracker.update( success=result["success"], token_usage=result.get("token_usage", {}), turn_count=result.get("turn_count", 0), execution_time=execution_time, ) result["execution_time"] = execution_time return result except Exception as e: execution_time = time.time() - start_time if isinstance(e, asyncio.TimeoutError): error_msg = f"Execution timed out after {self.timeout} seconds" logger.error(error_msg) else: error_msg = f"Agent execution failed: {e}" logger.error(error_msg, exc_info=True) self.usage_tracker.update( success=False, token_usage=self._partial_token_usage or {}, turn_count=self._partial_turn_count or 0, execution_time=execution_time, ) if self._partial_messages: if not self.is_claude: final_msg = self._convert_to_sdk_format(self._partial_messages) else: final_msg = self._partial_messages else: final_msg = [] return { "success": False, "output": final_msg, "token_usage": self._partial_token_usage or {}, "turn_count": self._partial_turn_count or 0, "execution_time": execution_time, "error": error_msg, "litellm_run_model_name": self.litellm_run_model_name, } def execute_sync( self, instruction: str, tool_call_log_file: Optional[str] = None ) -> Dict[str, Any]: """ Synchronous wrapper for execute method. """ return asyncio.run(self.execute(instruction, tool_call_log_file)) def get_usage_stats(self) -> Dict[str, Any]: """Get usage statistics.""" return self.usage_tracker.get_stats() def reset_usage_stats(self): """Reset usage statistics.""" self.usage_tracker.reset() # ==================== Claude Native API Execution Path ==================== async def _execute_claude_native_with_tools( self, instruction: str, tool_call_log_file: Optional[str] = None ) -> Dict[str, Any]: """ Execute Claude with thinking using native Anthropic API. Creates MCP server, gets tools, and executes with thinking. """ logger.debug("Using Claude native API with thinking") thinking_budget = self._get_claude_thinking_budget() # Create and start MCP server mcp_server = await self._create_mcp_server() async with mcp_server: # Get available tools tools = await mcp_server.list_tools() # Convert MCP tools to Anthropic format anthropic_tools = self._convert_to_anthropic_format(tools) # Execute with function calling loop return await self._execute_anthropic_native_tool_loop( instruction, anthropic_tools, mcp_server, thinking_budget, tool_call_log_file, ) async def _call_claude_native_api( self, messages: List[Dict], thinking_budget: int, tools: Optional[List[Dict]] = None, mcp_servers: Optional[List[Dict]] = None, system: Optional[str] = None, ) -> Dict[str, Any]: """ Call Claude's native API directly using httpx. Args: messages: Conversation messages thinking_budget: Token budget for thinking tools: Tool definitions for function calling mcp_servers: MCP server configurations system: System prompt Returns: API response as dictionary """ # Get API base and headers import os api_base = os.getenv("ANTHROPIC_API_BASE", "https://api.anthropic.com") headers = { "x-api-key": self.api_key, "anthropic-version": "2023-06-01", "content-type": "application/json", "anthropic-beta": "context-1m-2025-08-07", # by default } # Build payload max_tokens = max(thinking_budget + 4096, 4096) payload = { "model": self.litellm_input_model_name.replace("anthropic/", ""), "max_tokens": max_tokens, "messages": messages, } # Add thinking configuration if thinking_budget: payload["thinking"] = {"type": "enabled", "budget_tokens": thinking_budget} # Add tools if provided if tools: payload["tools"] = tools payload["tool_choice"] = {"type": "auto"} # Add MCP servers if provided if mcp_servers: headers["anthropic-beta"] = "mcp-client-2025-04-04" payload["mcp_servers"] = mcp_servers # Add system prompt if provided if system: payload["system"] = system # Make the API call async with httpx.AsyncClient() as client: try: response = await client.post( f"{api_base}/v1/messages", headers=headers, json=payload, timeout=self.timeout, ) response.raise_for_status() return response.json(), None except httpx.HTTPStatusError as e: return None, e.response.text except Exception as e: return None, e async def _count_claude_input_tokens( self, messages: List[Dict[str, Any]], tools: Optional[List[Dict]] = None, system: Optional[str] = None, ) -> int: import os api_base = os.getenv("ANTHROPIC_API_BASE", "https://api.anthropic.com") headers = { "x-api-key": self.api_key, "anthropic-version": "2023-06-01", "content-type": "application/json", } payload: Dict[str, Any] = { "model": self.litellm_input_model_name.replace("anthropic/", ""), "messages": messages, } if tools: payload["tools"] = tools if system: payload["system"] = system async with httpx.AsyncClient() as client: response = await client.post( f"{api_base}/v1/messages/count_tokens", headers=headers, json=payload, timeout=self.timeout, ) response.raise_for_status() data = response.json() or {} return int(data.get("input_tokens", 0) or 0) def _extract_litellm_text(self, response: Any) -> str: try: choices = getattr(response, "choices", None) or [] if not choices: return "" msg = getattr(choices[0], "message", None) if msg is not None: return str(getattr(msg, "content", "") or "") return str(getattr(choices[0], "text", "") or "") except Exception: # pragma: no cover - best effort return "" def _extract_anthropic_text(self, response_json: Dict[str, Any]) -> str: pieces: List[str] = [] for block in response_json.get("content", []) or []: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text") if text: pieces.append(str(text)) return "\n".join(pieces).strip() def _merge_usage(self, total_tokens: Dict[str, int], usage: Dict[str, Any]) -> None: try: input_tokens = int(usage.get("input_tokens", 0) or 0) output_tokens = int(usage.get("output_tokens", 0) or 0) total_tokens_count = int( usage.get("total_tokens", 0) or (input_tokens + output_tokens) ) total_tokens["input_tokens"] += input_tokens total_tokens["output_tokens"] += output_tokens total_tokens["total_tokens"] += total_tokens_count except Exception: # pragma: no cover - best effort return async def _maybe_compact_litellm_messages( self, messages: List[Dict[str, Any]], total_tokens: Dict[str, int], tool_call_log_file: Optional[str], current_prompt_tokens: int, ) -> List[Dict[str, Any]]: if not self._compaction_enabled(): return messages if current_prompt_tokens < self.compaction_token: return messages logger.info( f"| [compaction] Triggered at prompt tokens: {current_prompt_tokens:,}" ) if tool_call_log_file: try: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write( f"| [compaction] Triggered at prompt tokens: {current_prompt_tokens:,}\n" ) except Exception: pass compact_messages = [ {"role": "system", "content": self.COMPACTION_PROMPT}, {"role": "user", "content": json.dumps(messages, ensure_ascii=False)}, ] completion_kwargs = { "model": self.litellm_input_model_name, "messages": compact_messages, "api_key": self.api_key, } if self.base_url: completion_kwargs["base_url"] = self.base_url if self.extra_body: completion_kwargs["extra_body"] = self.extra_body response = await litellm.acompletion(**completion_kwargs) usage = getattr(response, "usage", None) if usage: input_tokens = ( getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", None) or 0 ) output_tokens = ( getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", None) or 0 ) total_tokens_count = getattr(usage, "total_tokens", None) if total_tokens_count is None: total_tokens_count = input_tokens + output_tokens total_tokens["input_tokens"] += int(input_tokens or 0) total_tokens["output_tokens"] += int(output_tokens or 0) total_tokens["total_tokens"] += int(total_tokens_count or 0) summary = self._extract_litellm_text(response).strip() or "(no summary)" system_msg = ( messages[0] if messages else {"role": "system", "content": self.SYSTEM_PROMPT} ) first_user = ( messages[1] if len(messages) > 1 else {"role": "user", "content": ""} ) return [ system_msg, first_user, { "role": "user", "content": f"Context summary (auto-compacted due to token limit):\n{summary}", }, ] def _should_summarize_tool(self) -> bool: """Summarize tool responses only for playwright (browser) runs. Browser tool results are giant accessibility trees that are navigation context, not the final answer — safe to compress. Data services (postgres/github/notion/filesystem) often return the exact value the evaluator checks, so they are intentionally NOT summarized. """ return bool(self.summarize_tool_response) and self.mcp_service in ( "playwright", "playwright_webarena", ) async def _summarize_tool_response_litellm( self, tool_response: str, context: str, total_tokens: Dict[str, int], ) -> str: """Compress a tool response via the LLM BEFORE it enters the message history (mirrors mcp-universe's summarize_tool_response). The returned text is appended to ``messages`` — which is exactly what gets saved as the trajectory — so the model conditions on, and the record stores, the SAME content (no post-hoc divergence). On any failure the raw response is returned unchanged so a rollout never breaks. The summarizer's own token usage is folded into total_tokens. """ prompt = self.TOOL_RESPONSE_SUMMARIZER_PROMPT.format( context=context, tool_response=tool_response ) completion_kwargs = { "model": self.litellm_input_model_name, "messages": [{"role": "user", "content": prompt}], "api_key": self.api_key, } if self.base_url: completion_kwargs["base_url"] = self.base_url if self.extra_body: completion_kwargs["extra_body"] = self.extra_body try: response = await asyncio.wait_for( litellm.acompletion(**completion_kwargs), timeout=self.timeout / 2, ) except Exception as e: # noqa: BLE001 - keep raw on any failure logger.warning(f"| [summarize] failed, keeping raw tool response: {e}") return tool_response usage = getattr(response, "usage", None) if usage: input_tokens = getattr(usage, "prompt_tokens", 0) or 0 total_tokens_count = getattr(usage, "total_tokens", 0) or 0 output_tokens = ( total_tokens_count - input_tokens if total_tokens_count > 0 else (getattr(usage, "completion_tokens", 0) or 0) ) total_tokens["input_tokens"] += int(input_tokens) total_tokens["output_tokens"] += int(output_tokens) total_tokens["total_tokens"] += int( total_tokens_count or (input_tokens + output_tokens) ) summary = self._extract_litellm_text(response).strip() return summary or tool_response async def _maybe_compact_anthropic_messages( self, messages: List[Dict[str, Any]], total_tokens: Dict[str, int], thinking_budget: int, tool_call_log_file: Optional[str], current_input_tokens: int, ) -> List[Dict[str, Any]]: if not self._compaction_enabled(): return messages if current_input_tokens < self.compaction_token: return messages logger.info( f"| [compaction] Triggered at input tokens: {current_input_tokens:,}" ) if tool_call_log_file: try: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write( f"| [compaction] Triggered at input tokens: {current_input_tokens:,}\n" ) except Exception: pass compact_messages = [ {"role": "user", "content": self.COMPACTION_PROMPT}, {"role": "user", "content": json.dumps(messages, ensure_ascii=False)}, ] response, error_msg = await self._call_claude_native_api( messages=compact_messages, thinking_budget=thinking_budget, tools=None, system=None, ) if error_msg or not response: logger.warning(f"| [compaction] Failed: {error_msg}") return messages usage = response.get("usage", {}) or {} input_tokens = usage.get("input_tokens", 0) or 0 output_tokens = usage.get("output_tokens", 0) or 0 total_tokens["input_tokens"] += int(input_tokens) total_tokens["output_tokens"] += int(output_tokens) total_tokens["total_tokens"] += int(input_tokens + output_tokens) summary = self._extract_anthropic_text(response) or "(no summary)" first_user = messages[0] if messages else {"role": "user", "content": ""} return [ first_user, { "role": "user", "content": f"Context summary (auto-compacted due to token limit):\n{summary}", }, ] async def _execute_anthropic_native_tool_loop( self, instruction: str, tools: List[Dict], mcp_server: Any, thinking_budget: int, tool_call_log_file: Optional[str] = None, ) -> Dict[str, Any]: """ Execute Claude thinking loop with function calling. Handles thinking blocks, tool calls, and message formatting. """ messages = [{"role": "user", "content": instruction}] total_tokens = { "input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "reasoning_tokens": 0, } turn_count = 0 max_turns = self.MAX_TURNS hit_turn_limit = False ended_normally = False system_text = self.SYSTEM_PROMPT # Record initial state self._update_progress(messages, total_tokens, turn_count) for _ in range(max_turns): turn_count += 1 current_input_tokens = 0 if self._compaction_enabled(): try: current_input_tokens = await self._count_claude_input_tokens( messages=messages, tools=tools, system=system_text, ) except Exception as exc: # noqa: BLE001 logger.debug("Claude token counting failed: %s", exc) messages = await self._maybe_compact_anthropic_messages( messages=messages, total_tokens=total_tokens, thinking_budget=thinking_budget, tool_call_log_file=tool_call_log_file, current_input_tokens=current_input_tokens, ) self._update_progress(messages, total_tokens, turn_count) # Call Claude native API response, error_msg = await self._call_claude_native_api( messages=messages, thinking_budget=thinking_budget, tools=tools, system=system_text, ) if turn_count == 1: self.litellm_run_model_name = response["model"].split("/")[-1] if error_msg: break # Update token usage if "usage" in response: usage = response["usage"] input_tokens = usage.get("input_tokens", 0) output_tokens = usage.get("output_tokens", 0) # Calculate output tokens as total - input for consistency total_tokens_count = output_tokens + input_tokens total_tokens["input_tokens"] += input_tokens total_tokens["output_tokens"] += output_tokens total_tokens["total_tokens"] += total_tokens_count ## TODO: add reasoning tokens for claude # Extract blocks from response blocks = response.get("content", []) tool_uses = [b for b in blocks if b.get("type") == "tool_use"] thinking_blocks = [b for b in blocks if b.get("type") == "thinking"] text_blocks = [b for b in blocks if b.get("type") == "text"] # Log text output for tb in text_blocks: if tb.get("text") and tool_call_log_file: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"{tb['text']}\n") if tb.get("text"): for line in tb["text"].splitlines(): logger.info(f"| {line}") # Build assistant message with all blocks assistant_content = [] # Add thinking blocks for tb in thinking_blocks: assistant_content.append( { "type": "thinking", "thinking": tb.get("thinking", ""), "signature": tb.get("signature", ""), } ) # Add text blocks for tb in text_blocks: if tb.get("text"): assistant_content.append({"type": "text", "text": tb["text"]}) # Add tool_use blocks for tu in tool_uses: assistant_content.append( { "type": "tool_use", "id": tu.get("id"), "name": tu.get("name"), "input": tu.get("input", {}), } ) messages.append({"role": "assistant", "content": assistant_content}) # Update partial progress after assistant response self._update_progress(messages, total_tokens, turn_count) # If no tool calls, we're done if not tool_uses: ended_normally = True break # Execute tools and add results tool_results = [] for tu in tool_uses: name = tu.get("name") inputs = tu.get("input", {}) # Log tool call args_str = json.dumps(inputs, separators=(",", ": ")) display_args = ( args_str[:140] + "..." if len(args_str) > 140 else args_str ) logger.info(f"| \033[1m{name}\033[0m \033[2;37m{display_args}\033[0m") if tool_call_log_file: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"| {name} {args_str}\n") # Execute tool try: result = await asyncio.wait_for( mcp_server.call_tool(name, inputs), timeout=60 ) tool_results.append( { "type": "tool_result", "tool_use_id": tu["id"], "content": [ { "type": "text", "text": json.dumps(result, cls=CustomJSONEncoder), } ], } ) except Exception as e: logger.error(f"Tool call failed: {e}") tool_results.append( { "type": "tool_result", "tool_use_id": tu["id"], "content": [{"type": "text", "text": f"Error: {str(e)}"}], } ) messages.append({"role": "user", "content": tool_results}) # Update partial progress after tool results self._update_progress(messages, total_tokens, turn_count) # Detect if we exited due to hitting the turn limit if not ended_normally: if turn_count >= max_turns: hit_turn_limit = True logger.warning( f"| Max turns ({max_turns}) exceeded; returning failure with partial output." ) if tool_call_log_file: try: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"| Max turns ({max_turns}) exceeded\n") except Exception: pass elif error_msg: logger.warning(f"| {error_msg}\n") if tool_call_log_file: try: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"| {error_msg}\n") except Exception: pass # Display final token usage if total_tokens["total_tokens"] > 0: log_msg = ( f"|\n| Token usage: Total: {total_tokens['total_tokens']:,} | " f"Input: {total_tokens['input_tokens']:,} | " f"Output: {total_tokens['output_tokens']:,}" ) if total_tokens.get("reasoning_tokens", 0) > 0: log_msg += f" | Reasoning: {total_tokens['reasoning_tokens']:,}" logger.info(log_msg) logger.info(f"| Turns: {turn_count}") # Convert messages to SDK format sdk_format_messages = self._convert_to_sdk_format(messages) if hit_turn_limit: return { "success": False, "output": sdk_format_messages, "token_usage": total_tokens, "turn_count": turn_count, "error": f"Max turns ({max_turns}) exceeded", "litellm_run_model_name": self.litellm_run_model_name, } if error_msg: return { "success": False, "output": sdk_format_messages, "token_usage": total_tokens, "turn_count": turn_count, "error": error_msg, "litellm_run_model_name": self.litellm_run_model_name, } return { "success": True, "output": sdk_format_messages, "token_usage": total_tokens, "turn_count": turn_count, "error": None, "litellm_run_model_name": self.litellm_run_model_name, } # ==================== LiteLLM Execution Path ==================== async def _execute_litellm_with_tools( self, instruction: str, tool_call_log_file: Optional[str] = None ) -> Dict[str, Any]: """ Execute with manual MCP server management. Used for all non-Anthropic models and Anthropic models with STDIO services. """ logger.debug("Using manual MCP execution with function calling loop") # Create and start MCP server mcp_server = await self._create_mcp_server() try: async with mcp_server: # Get available tools tools = await mcp_server.list_tools() # Convert MCP tools to OpenAI function format functions = self._convert_to_openai_format(tools) # Execute with function calling loop return await self._execute_litellm_tool_loop( instruction, functions, mcp_server, tool_call_log_file ) except Exception as e: logger.error(f"Manual MCP execution failed: {e}") raise async def _execute_litellm_tool_loop( self, instruction: str, functions: List[Dict], mcp_server: Any, tool_call_log_file: Optional[str] = None, ) -> Dict[str, Any]: """Execute function calling loop with LiteLLM.""" messages = [ {"role": "system", "content": self.SYSTEM_PROMPT}, {"role": "user", "content": instruction}, ] total_tokens = { "input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "reasoning_tokens": 0, } turn_count = 0 max_turns = self.MAX_TURNS # Limit turns to prevent infinite loops consecutive_failures = 0 max_consecutive_failures = 3 hit_turn_limit = False ended_normally = False # Convert functions to tools format for newer models tools = ( [{"type": "function", "function": func} for func in functions] if functions else None ) if tool_call_log_file and tools: max_name_length = ( max(len(tool.get("function", {}).get("name", "")) for tool in tools) if tools else 15 ) with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write("===== Available Tools =====\n") for tool in tools: function_info = tool.get("function", {}) tool_name = function_info.get("name", "N/A") description = function_info.get("description", "N/A") f.write( f"- ToolName: {tool_name:<{max_name_length}} Description: {description}\n" ) f.write("\n===== Execution Logs =====\n") # Record initial state self._update_progress(messages, total_tokens, turn_count) try: while turn_count < max_turns: current_prompt_tokens = 0 if self._compaction_enabled(): current_prompt_tokens = self._count_prompt_tokens_litellm(messages) messages = await self._maybe_compact_litellm_messages( messages=messages, total_tokens=total_tokens, tool_call_log_file=tool_call_log_file, current_prompt_tokens=current_prompt_tokens, ) self._update_progress(messages, total_tokens, turn_count) # Build completion kwargs completion_kwargs = { "model": self.litellm_input_model_name, "messages": messages, "api_key": self.api_key, } # Always use tools format if available - LiteLLM will handle conversion if tools: completion_kwargs["tools"] = tools completion_kwargs["tool_choice"] = "auto" # Add reasoning_effort and base_url if specified if self.reasoning_effort != "default": completion_kwargs["reasoning_effort"] = self.reasoning_effort if self.base_url: completion_kwargs["base_url"] = self.base_url if self.extra_body: completion_kwargs["extra_body"] = self.extra_body try: # Call LiteLLM with timeout for individual call response = await asyncio.wait_for( litellm.acompletion(**completion_kwargs), timeout=self.timeout / 2, # Use half of total timeout ) consecutive_failures = 0 # Reset failure counter on success except asyncio.TimeoutError: logger.warning(f"| ✗ LLM call timed out on turn {turn_count + 1}") consecutive_failures += 1 if consecutive_failures >= max_consecutive_failures: raise Exception( f"Too many consecutive failures ({consecutive_failures})" ) await asyncio.sleep(8**consecutive_failures) # Exponential backoff continue except Exception as e: logger.error(f"| ✗ LLM call failed on turn {turn_count + 1}: {e}") consecutive_failures += 1 if consecutive_failures >= max_consecutive_failures: raise if "ContextWindowExceededError" in str(e): # Best-effort fallback: compact and retry once. messages = await self._maybe_compact_litellm_messages( messages=messages, total_tokens=total_tokens, tool_call_log_file=tool_call_log_file, current_prompt_tokens=self.compaction_token, ) self._update_progress(messages, total_tokens, turn_count) continue elif "RateLimitError" in str(e): await asyncio.sleep(12**consecutive_failures) else: await asyncio.sleep(2**consecutive_failures) continue # Extract actual model name from response (first turn only) if turn_count == 0 and hasattr(response, "model") and response.model: self.litellm_run_model_name = response.model.split("/")[-1] # Update token usage including reasoning tokens if hasattr(response, "usage") and response.usage: input_tokens = response.usage.prompt_tokens or 0 total_tokens_count = response.usage.total_tokens or 0 # Calculate output tokens as total - input for consistency output_tokens = ( total_tokens_count - input_tokens if total_tokens_count > 0 else (response.usage.completion_tokens or 0) ) total_tokens["input_tokens"] += input_tokens total_tokens["output_tokens"] += output_tokens total_tokens["total_tokens"] += total_tokens_count # Extract reasoning tokens if available if hasattr(response.usage, "completion_tokens_details"): details = response.usage.completion_tokens_details if hasattr(details, "reasoning_tokens"): total_tokens["reasoning_tokens"] += ( details.reasoning_tokens or 0 ) # Get response message choices = response.choices if len(choices): message = choices[0].message # deeply dump the message to ensure we capture all fields message_dict = ( message.model_dump() if hasattr(message, "model_dump") else dict(message) ) # Explicitly preserve function_call if present (even if tool_calls exists), # as it may contain provider-specific metadata (e.g. Gemini thought_signature) if hasattr(message, "function_call") and message.function_call: # Ensure it's in the dict if model_dump missed it or it was excluded if ( "function_call" not in message_dict or not message_dict["function_call"] ): fc = message.function_call message_dict["function_call"] = ( fc.model_dump() if hasattr(fc, "model_dump") else fc ) # Log assistant's text content if present if hasattr(message, "content") and message.content: # Display the content with line prefix for line in message.content.splitlines(): logger.info(f"| {line}") # Also log to file if specified if tool_call_log_file: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"{message.content}\n") # Check for tool calls (newer format) if hasattr(message, "tool_calls") and message.tool_calls: messages.append(message_dict) turn_count += 1 # Update progress after assistant with tool calls self._update_progress(messages, total_tokens, turn_count) # Process tool calls for tool_call in message.tool_calls: func_name = tool_call.function.name func_args = json.loads(tool_call.function.arguments) try: result = await asyncio.wait_for( mcp_server.call_tool(func_name, func_args), timeout=60 ) _content = json.dumps(result, cls=CustomJSONEncoder) # Playwright only: compress the browser result before # it enters history, so the model input == the saved # trajectory (fidelity-preserving, keeps tokens down). if self._should_summarize_tool(): _context = json.dumps( {"tool": func_name, "arguments": func_args}, ensure_ascii=False, ) _content = await self._summarize_tool_response_litellm( _content, context=_context, total_tokens=total_tokens, ) messages.append( { "role": "tool", "tool_call_id": tool_call.id, "content": _content, } ) except asyncio.TimeoutError: error_msg = ( f"Tool call '{func_name}' timed out after 60 seconds" ) logger.error(error_msg) messages.append( { "role": "tool", "tool_call_id": tool_call.id, "content": f"Error: {error_msg}", } ) except Exception as e: logger.error(f"Tool call failed: {e}") messages.append( { "role": "tool", "tool_call_id": tool_call.id, "content": f"Error: {str(e)}", } ) # Format arguments for display (truncate if too long) args_str = json.dumps(func_args, separators=(",", ": ")) display_arguments = ( args_str[:140] + "..." if len(args_str) > 140 else args_str ) # Log with ANSI color codes (bold tool name, dim gray arguments) logger.info( f"| \033[1m{func_name}\033[0m \033[2;37m{display_arguments}\033[0m" ) if tool_call_log_file: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"| {func_name} {args_str}\n") # Update progress after tool results appended self._update_progress(messages, total_tokens, turn_count) continue else: # Log end reason if not choices: logger.info( "|\n|\n| Task ended with no messages generated by the model." ) elif choices[0].finish_reason == "stop": logger.info( "|\n|\n| Task ended with the finish reason from messages being 'stop'." ) # No tool/function call, add message and we're done messages.append(message_dict) turn_count += 1 # Update progress before exiting self._update_progress(messages, total_tokens, turn_count) ended_normally = True break except Exception as loop_error: # On any error, return partial conversation, token usage, and turn count logger.error(f"Manual MCP loop failed: {loop_error}", exc_info=True) sdk_format_messages = self._convert_to_sdk_format(messages) return { "success": False, "output": sdk_format_messages, "token_usage": total_tokens, "turn_count": turn_count, "error": str(loop_error), "litellm_run_model_name": self.litellm_run_model_name, } # Detect if we exited due to hitting the turn limit if (not ended_normally) and (turn_count >= max_turns): hit_turn_limit = True logger.warning( f"| Max turns ({max_turns}) exceeded); returning failure with partial output." ) if tool_call_log_file: try: with open(tool_call_log_file, "a", encoding="utf-8") as f: f.write(f"| Max turns ({max_turns}) exceeded\n") except Exception: pass # Display final token usage if total_tokens["total_tokens"] > 0: log_msg = ( f"| Token usage: Total: {total_tokens['total_tokens']:,} | " f"Input: {total_tokens['input_tokens']:,} | " f"Output: {total_tokens['output_tokens']:,}" ) if total_tokens.get("reasoning_tokens", 0) > 0: log_msg += f" | Reasoning: {total_tokens['reasoning_tokens']:,}" logger.info(log_msg) logger.info(f"| Turns: {turn_count}") # Convert messages to SDK format for backward compatibility sdk_format_messages = self._convert_to_sdk_format(messages) return { "success": not hit_turn_limit, "output": sdk_format_messages, "token_usage": total_tokens, "turn_count": turn_count, "error": (f"Max turns ({max_turns}) exceeded" if hit_turn_limit else None), "litellm_run_model_name": self.litellm_run_model_name, } # ==================== MCP Server Management ==================== async def _create_mcp_server(self) -> Any: """Create and return an MCP server instance.""" if self.mcp_service in self.STDIO_SERVICES: return self._create_stdio_server() elif self.mcp_service in self.HTTP_SERVICES: return self._create_http_server() else: raise ValueError(f"Unsupported MCP service: {self.mcp_service}") def _create_stdio_server(self) -> MCPStdioServer: """Create stdio-based MCP server.""" if self.mcp_service == "notion": notion_key = self.service_config.get("notion_key") if not notion_key: raise ValueError("Notion API key required") return MCPStdioServer( command="npx", args=["-y", "@notionhq/notion-mcp-server@1.9.1"], env={ "OPENAPI_MCP_HEADERS": ( '{"Authorization": "Bearer ' + notion_key + '", ' '"Notion-Version": "2022-06-28"}' ) }, ) elif self.mcp_service == "filesystem": test_directory = self.service_config.get("test_directory") if not test_directory: raise ValueError("Test directory required for filesystem service") return MCPStdioServer( command="npx", args=[ "-y", "@modelcontextprotocol/server-filesystem", str(test_directory), ], ) elif self.mcp_service in ["playwright", "playwright_webarena"]: browser = self.service_config.get("browser", "chromium") headless = self.service_config.get("headless", True) viewport_width = self.service_config.get("viewport_width", 1280) viewport_height = self.service_config.get("viewport_height", 720) args = ["-y", "@playwright/mcp@latest"] if headless: args.append("--headless") args.extend( [ "--isolated", "--no-sandbox", "--browser", browser, "--viewport-size", f"{viewport_width},{viewport_height}", ] ) # Real-web playwright only: route the browser through an egress # proxy (e.g. US socks5) when MCPMARK_BROWSER_PROXY is set. WebArena # is localhost docker and must NOT be proxied, so this is gated to # the real-internet "playwright" service. if self.mcp_service == "playwright": _browser_proxy = os.environ.get("MCPMARK_BROWSER_PROXY", "").strip() if _browser_proxy: args.extend(["--proxy-server", _browser_proxy]) return MCPStdioServer(command="npx", args=args) elif self.mcp_service == "postgres": host = self.service_config.get("host", "localhost") port = self.service_config.get("port", 5432) username = self.service_config.get("username") password = self.service_config.get("password") database = self.service_config.get( "current_database" ) or self.service_config.get("database") if not all([username, password, database]): raise ValueError("PostgreSQL requires username, password, and database") database_url = ( f"postgresql://{username}:{password}@{host}:{port}/{database}" ) return MCPStdioServer( command="pipx", args=["run", "postgres-mcp", "--access-mode=unrestricted"], env={"DATABASE_URI": database_url}, ) elif self.mcp_service == "insforge": api_key = self.service_config.get("api_key") backend_url = self.service_config.get("backend_url") if not all([api_key, backend_url]): raise ValueError("Insforge requires api_key and backend_url") return MCPStdioServer( command="npx", args=["-y", "@insforge/mcp@dev"], env={ "INSFORGE_API_KEY": api_key, "INSFORGE_BACKEND_URL": backend_url, }, ) elif self.mcp_service == "github": github_token = self.service_config.get("github_token") if not github_token: raise ValueError("GitHub token required") return MCPStdioServer( command="docker", args=[ "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN", "ghcr.io/github/github-mcp-server:v0.15.0", ], env={"GITHUB_PERSONAL_ACCESS_TOKEN": github_token}, ) else: raise ValueError(f"Unsupported stdio service: {self.mcp_service}") def _create_http_server(self) -> MCPHttpServer: """Create HTTP-based MCP server.""" if self.mcp_service == "supabase": # Use built-in MCP server from Supabase CLI api_url = self.service_config.get("api_url", "http://localhost:54321") api_key = self.service_config.get("api_key", "") if not api_key: raise ValueError( "Supabase requires api_key (use secret key from 'supabase status')" ) # Supabase CLI exposes MCP at /mcp endpoint mcp_url = f"{api_url}/mcp" return MCPHttpServer( url=mcp_url, headers={ "apikey": api_key, "Authorization": f"Bearer {api_key}", }, ) else: raise ValueError(f"Unsupported HTTP service: {self.mcp_service}")