# Tool Call Monitor

This document contains code snippets for monitoring and formatting tool calls in an MCP client. These can be incorporated into other projects that need to track and display tool invocations.

## Core Components

### 1. Console Setup

```python
from rich.console import Console
import sys

# Global Rich Console instance with color support
console = Console(color_system="truecolor")

def log_system(msg: str, title: str = None) -> None:
    """Log a system message with optional title."""
    if title:
        console.print(f"[bold blue]{title}[/bold blue]: {msg}")
    else:
        console.print(f"[bold blue]system[/bold blue]: {msg}")
    
    # Ensure output is flushed immediately
    sys.stdout.flush()
    console.file.flush() if hasattr(console, "file") else None
```

### 2. Tool Output Formatter

```python
def format_tool_output(result):
    """Format tool outputs into readable text."""
    # Handle error objects
    if hasattr(result, "content"):
        return result.content
    
    # Handle dictionary responses
    if isinstance(result, dict):
        if result.get("isError") is True:
            return f"ERROR: {result.get('content', 'Unknown error')}"
        return result.get("content", str(result))
    
    # Handle string responses
    if isinstance(result, str):
        return result.replace("\\n", "\n")
    
    # Default: convert to string
    return str(result).replace("\\n", "\n")
```

### 3. Tool Wrapper with Logging

```python
import json

def wrap_tool(tool):
    """Wrap a tool for logging with tidier output."""
    # Clean tool name if needed
    tool_name = tool.name
    if tool_name.startswith("[Tool]"):
        tool_name = tool_name.replace("[Tool]", "").strip()
    
    updates = {}
    
    # Define a wrapper function for both sync and async invocation
    def log_and_call(func):
        def wrapper(call_args, config=None):
            args_only = call_args.get("args", {})
            log_system(
                f"▶ {tool_name} called with args: {json.dumps(args_only, indent=2)}"
            )
            sys.stdout.flush()
            
            try:
                result = func(call_args, config)
                formatted = format_tool_output(result)
                log_system(f"◀ {tool_name} output: {formatted}")
                sys.stdout.flush()
                
                return result
            except Exception as e:
                error_msg = f"Tool execution failed: {str(e)}"
                log_system(f"✖ Error: {error_msg}")
                sys.stdout.flush()
                # Return error in a format your framework expects
                return {"isError": True, "content": error_msg}
        
        return wrapper
    
    # Apply wrappers to sync and async methods if they exist
    if hasattr(tool, "invoke"):
        updates["invoke"] = log_and_call(tool.invoke)
    
    if hasattr(tool, "ainvoke"):
        orig_ainvoke = tool.ainvoke
        
        async def ainvoke_wrapper(call_args, config=None):
            args_only = call_args.get("args", {})
            log_system(
                f"▶ {tool_name} called with args: {json.dumps(args_only, indent=2)}"
            )
            sys.stdout.flush()
            
            try:
                result = await orig_ainvoke(call_args, config)
                formatted = format_tool_output(result)
                log_system(f"◀ {tool_name} output: {formatted}")
                sys.stdout.flush()
                
                return result
            except Exception as e:
                error_msg = f"Tool execution failed: {str(e)}"
                log_system(f"✖ Error: {error_msg}")
                sys.stdout.flush()
                # Return error in a format your framework expects
                return {"isError": True, "content": error_msg}
        
        updates["ainvoke"] = ainvoke_wrapper
    
    if updates:
        return tool.copy(update=updates)
    else:
        log_system(
            f"Warning: {tool_name} has no invoke or ainvoke method; cannot wrap for logging."
        )
        return tool
```

### 4. Usage Example

```python
# Assuming you have a list of tools from your MCP client
raw_tools = await load_mcp_tools(session)  # Your tool loading function

# Wrap each tool with logging functionality
wrapped_tools = [wrap_tool(tool) for tool in raw_tools]

# Now when tools are invoked, you'll see formatted input/output
# Example output:
# ▶ add_item called with args: {
#   "index": 0,
#   "content": "x = Int('x')"
# }
# ◀ add_item output: Item added at index 0
```

### 5. Tool Statistics Tracking (Optional)

```python
from collections import defaultdict

class ToolStats:
    """Track tool usage statistics."""
    _instance = None
    
    def __init__(self):
        self.tool_calls = defaultdict(int)
        self.total_calls = 0
        self.enabled = True
    
    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance
    
    def record_tool_call(self, tool_name: str):
        """Record a tool call."""
        if self.enabled:
            self.tool_calls[tool_name] += 1
            self.total_calls += 1
    
    def get_stats(self):
        """Get tool usage statistics."""
        return dict(self.tool_calls)
```

### 6. Integration with Tool Wrapper

To add statistics tracking to the tool wrapper, modify the wrapper functions:

```python
# Inside the wrapper function, after extracting args:
# Record tool usage statistics
tool_stats = ToolStats.get_instance()
tool_stats.record_tool_call(tool_name)
```

### 7. Display Tool Statistics

```python
from rich.table import Table

def display_tool_stats():
    """Display tool usage statistics."""
    tool_stats = ToolStats.get_instance()
    
    if tool_stats.enabled and tool_stats.total_calls > 0:
        table = Table(title="Tool Usage Statistics")
        table.add_column("Tool Name", style="cyan")
        table.add_column("Calls", style="yellow")
        
        # Sort tools by number of calls (descending)
        sorted_tools = sorted(
            tool_stats.tool_calls.items(), 
            key=lambda x: x[1], 
            reverse=True
        )
        
        for tool_name, count in sorted_tools:
            table.add_row(tool_name, str(count))
        
        table.add_row("TOTAL", str(tool_stats.total_calls), style="bold")
        
        console.print("\n")
        console.print(table)
```

## Key Features

1. **Formatted Input Display**: Tool arguments are shown as indented JSON
2. **Formatted Output Display**: Tool results are processed for readability
3. **Error Handling**: Gracefully handles and displays tool execution errors
4. **Async Support**: Works with both synchronous and asynchronous tools
5. **Visual Indicators**: Uses symbols (▶, ◀, ✖) for clear visual separation
6. **Immediate Flushing**: Ensures output appears immediately in the console
7. **Optional Statistics**: Track which tools are used and how often

## Customization Options

- Change the visual indicators (▶, ◀, ✖) to your preference
- Modify the color scheme in Rich Console styling
- Add timestamp information to each log
- Filter or truncate very long outputs
- Add different formatting for specific tool types