Command Dispatcher

async def execute_commands(
    self, 
    commands: List[Command], 
    timeout: float = 6000
) -> Optional[List[Result]]

def generate_error_results(
    self, 
    commands: List[Command], 
    error: Exception
) -> Optional[List[Result]]

Result(
    status=ResultStatus.FAILURE,
    error=f"Error occurred while executing command {command}: {error}",
    result=f"Error occurred while executing command {command}: {error}, "
           f"please retry or execute a different command.",
    call_id=command.call_id
)

from ufo.module.dispatcher import LocalCommandDispatcher
from ufo.client.mcp.mcp_server_manager import MCPServerManager

def _init_context(self) -> None:
    """Initialize context with local dispatcher."""
    super()._init_context()

    # Create MCP server manager
    mcp_server_manager = MCPServerManager()

    # Create local dispatcher
    command_dispatcher = LocalCommandDispatcher(
        session=self,
        mcp_server_manager=mcp_server_manager
    )

    # Attach to context
    self.context.attach_command_dispatcher(command_dispatcher)

from ufo.module.dispatcher import WebSocketCommandDispatcher
from aip.protocol.task_execution import TaskExecutionProtocol

def _init_context(self) -> None:
    """Initialize context with WebSocket dispatcher."""
    super()._init_context()

    # Create WebSocket dispatcher with AIP protocol
    command_dispatcher = WebSocketCommandDispatcher(
        session=self,
        protocol=self.task_protocol  # TaskExecutionProtocol instance
    )

    # Attach to context
    self.context.attach_command_dispatcher(command_dispatcher)

def make_server_response(self, commands: List[Command]) -> ServerMessage:
    """
    Create a server response message for the given commands.
    """
    # Assign unique IDs
    for command in commands:
        command.call_id = str(uuid.uuid4())

    # Extract context
    agent_name = self.session.current_agent_class
    process_name = self.session.context.get(ContextNames.APPLICATION_PROCESS_NAME)
    root_name = self.session.context.get(ContextNames.APPLICATION_ROOT_NAME)
    session_id = self.session.id
    response_id = str(uuid.uuid4())

    # Build AIP message
    return ServerMessage(
        type=ServerMessageType.COMMAND,
        status=TaskStatus.CONTINUE,
        agent_name=agent_name,
        process_name=process_name,
        root_name=root_name,
        actions=commands,
        session_id=session_id,
        task_name=self.session.task,
        timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(),
        response_id=response_id
    )

async def set_result(self, response_id: str, result: ClientMessage) -> None:
    """
    Called by WebSocket handler when client returns a message.
    :param response_id: The ID of the response.
    :param result: The result from the client.
    """
    fut = self.pending.get(response_id)
    if fut and not fut.done():
        fut.set_result(result.action_results)

{
    "status": "failure",  # ResultStatus.FAILURE
    "error": "asyncio.TimeoutError: Command execution timed out",
    "result": "Error occurred while executing command <Command>: TimeoutError, "
              "please retry or execute a different command.",
    "call_id": "cmd_abc123"
}

async def execute_action(self, context: Context) -> None:
    """Execute action with error handling."""
    commands = self.generate_commands()

    results = await context.command_dispatcher.execute_commands(
        commands=commands,
        timeout=30.0
    )

    for command, result in zip(commands, results):
        if result.status == ResultStatus.FAILURE:
            # Log error
            self.logger.error(f"Command {command.tool_name} failed: {result.error}")

            # Decision logic
            if "timeout" in result.error.lower():
                # Retry with longer timeout
                self.retry_count += 1
                if self.retry_count < 3:
                    return await self.execute_action(context)

            elif "connection" in result.error.lower():
                # Switch to alternative approach
                return self.fallback_strategy()

            else:
                # Escalate to error state
                self.transition_to_error_state(result.error)
        else:
            # Process successful result
            self.process_result(result.result)

for command in command_list:
    results = await context.command_dispatcher.execute_commands(
        commands=[command],
        timeout=30.0
    )

    if results[0].status == ResultStatus.SUCCESS:
        # Process result and decide next command
        next_command = self.decide_next_action(results[0])
    else:
        # Handle error and possibly abort
        break

# All commands for a subtask
commands = [
    Command(tool_name="click_element", ...),
    Command(tool_name="type_text", ...),
    Command(tool_name="press_key", ...)
]

results = await context.command_dispatcher.execute_commands(
    commands=commands,
    timeout=60.0
)

# Process all results
for command, result in zip(commands, results):
    if result.status == ResultStatus.FAILURE:
        # One failure might invalidate the whole subtask
        self.handle_subtask_failure(command, result)

# Check state first
check_cmd = Command(tool_name="get_ui_tree", ...)
check_results = await dispatcher.execute_commands([check_cmd])

if check_results[0].status == ResultStatus.SUCCESS:
    ui_tree = check_results[0].result

    # Decide action based on UI state
    if "Login" in ui_tree:
        action_cmd = Command(tool_name="click_element", parameters={"label": "Login"})
    else:
        action_cmd = Command(tool_name="type_text", parameters={"text": "username"})

    # Execute decided action
    await dispatcher.execute_commands([action_cmd])

import asyncio

async def execute_with_retry(
    dispatcher, 
    commands, 
    max_retries=3, 
    base_delay=1.0
):
    """Execute commands with exponential backoff retry."""

    for attempt in range(max_retries):
        results = await dispatcher.execute_commands(commands, timeout=30.0)

        # Check if all succeeded
        all_success = all(r.status == ResultStatus.SUCCESS for r in results)

        if all_success:
            return results

        # Not last attempt - retry with backoff
        if attempt < max_retries - 1:
            delay = base_delay * (2 ** attempt)
            logger.warning(f"Retry attempt {attempt + 1} after {delay}s")
            await asyncio.sleep(delay)

    # All retries exhausted
    return results  # Return last attempt results

# Good: Reuse dispatcher attached to context
results1 = await context.command_dispatcher.execute_commands(commands1)
results2 = await context.command_dispatcher.execute_commands(commands2)

# Bad: Creating new dispatchers
dispatcher1 = LocalCommandDispatcher(session, mcp_manager)
dispatcher2 = LocalCommandDispatcher(session, mcp_manager)

from ufo.module.dispatcher import BasicCommandDispatcher
from aip.messages import Command, Result, ResultStatus
from typing import List, Optional

class CustomCommandDispatcher(BasicCommandDispatcher):
    """
    Custom dispatcher that logs all commands and results.
    """

    def __init__(self, session, log_file: str):
        self.session = session
        self.log_file = log_file

    async def execute_commands(
        self, 
        commands: List[Command], 
        timeout: float = 6000
    ) -> Optional[List[Result]]:
        """Execute with logging."""

        # Log commands
        with open(self.log_file, 'a') as f:
            f.write(f"Executing {len(commands)} commands\n")
            for cmd in commands:
                f.write(f"  {cmd.tool_name}: {cmd.parameters}\n")

        try:
            # Your custom execution logic here
            results = await self.custom_execute(commands, timeout)

            # Log results
            with open(self.log_file, 'a') as f:
                for result in results:
                    f.write(f"  Result: {result.status}\n")

            return results

        except Exception as e:
            # Log error
            with open(self.log_file, 'a') as f:
                f.write(f"  ERROR: {e}\n")

            return self.generate_error_results(commands, e)

    async def custom_execute(
        self, 
        commands: List[Command], 
        timeout: float
    ) -> List[Result]:
        """Implement custom execution logic."""
        # Your implementation here
        pass

from ufo.module.dispatcher import LocalCommandDispatcher, WebSocketCommandDispatcher

def attach_appropriate_dispatcher(session, context):
    """Attach correct dispatcher based on session type."""

    if isinstance(session, ServiceSession):
        # Service session uses WebSocket
        dispatcher = WebSocketCommandDispatcher(
            session=session,
            protocol=session.task_protocol
        )
    else:
        # Interactive session uses local execution
        mcp_manager = MCPServerManager()
        dispatcher = LocalCommandDispatcher(
            session=session,
            mcp_server_manager=mcp_manager
        )

    context.attach_command_dispatcher(dispatcher)

# Check timeout value
results = await dispatcher.execute_commands(commands, timeout=30.0)

# Enable debug logging
logging.getLogger('ufo.module.dispatcher').setLevel(logging.DEBUG)

# For LocalCommandDispatcher
# Check MCP server status
mcp_manager.check_server_health()

# For WebSocketCommandDispatcher
# Check WebSocket connection
if protocol.is_connected():
    print("WebSocket connected")
else:
    print("WebSocket disconnected")

# Check dispatcher type
print(type(context.command_dispatcher))
# Should be LocalCommandDispatcher or WebSocketCommandDispatcher

# Check session type
print(type(session))