Skip to content

agent_runner

Reusable end-to-end agent test harnesses.

This module packages the repository's MCP-server-plus-agent integration pattern into a reusable class that downstream extension packages can import directly. The runner owns the full lifecycle of randomized config generation, server startup, agent initialization, prompt execution, and cleanup.

AgentProtocol

Bases: Protocol

Protocol for agent instances managed by AgentTestRunner.

Source code in mada_tools/testing/agent_runner.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
class AgentProtocol(Protocol):
    """Protocol for agent instances managed by `AgentTestRunner`."""

    def __init__(self, *, config_path: str): ...

    async def initialize(self, stack: AsyncExitStack) -> None: ...

    async def process_query(
        self,
        query: str,
        max_tool_calls: int = 10,
        add_tool_context: bool = False,
    ) -> str: ...

AgentTestRunner

Async test harness for starting MCP servers and querying them through an agent.

The runner accepts explicit paths to a server configuration and an agent configuration. During startup it randomizes server ports, rewrites matching MCP URLs in the agent config, launches the configured servers, validates the observed server state, and then initializes the requested agent class.

The default agent class is MultiServerAgent, but tests may substitute any compatible implementation that accepts config_path=..., exposes an async initialize(stack) method, and supports process_query(...).

Source code in mada_tools/testing/agent_runner.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
class AgentTestRunner:
    """Async test harness for starting MCP servers and querying them through an agent.

    The runner accepts explicit paths to a server configuration and an agent
    configuration. During startup it randomizes server ports, rewrites matching
    MCP URLs in the agent config, launches the configured servers, validates the
    observed server state, and then initializes the requested agent class.

    The default agent class is `MultiServerAgent`, but tests may substitute any
    compatible implementation that accepts `config_path=...`, exposes an async
    `initialize(stack)` method, and supports `process_query(...)`.
    """

    def __init__(
        self,
        servers_config_path: Path,
        agent_config_path: Path,
        agent_cls: type[AgentProtocol] = MultiServerAgent,
    ):
        """Initialize the test runner.

        Args:
            servers_config_path: Path to the MCP server configuration JSON.
            agent_config_path: Path to the agent configuration JSON.
            agent_cls: Agent implementation to instantiate after servers are
                running. Defaults to `MultiServerAgent`.
        """
        self.base_servers_config_path = Path(servers_config_path)
        self.base_agent_config_path = Path(agent_config_path)
        self.servers_config_path = self.base_servers_config_path
        self.agent_config_path = self.base_agent_config_path
        self.agent_cls = agent_cls

        self.server_manager = ServerManager(
            state_file=Path.home() / ".mada" / f"server_statuses_{uuid.uuid4().hex}.json"
        )
        self.stack: AsyncExitStack | None = None
        self.agent: AgentProtocol | None = None
        self.servers_config: dict[str, Any] | None = None
        self._generated_config_paths: list[Path] = []

    async def __aenter__(self):
        """Start managed resources when entering an async context."""
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        """Always tear down managed resources when leaving an async context."""
        await self.close()

    async def start(self):
        """Prepare randomized configs, start servers, and initialize the agent.

        Raises:
            FileNotFoundError: If either input configuration path does not
                exist.
            Exception: Propagates any server-startup or agent-initialization
                error after best-effort cleanup.
        """
        if not self.base_servers_config_path.exists():
            raise FileNotFoundError(f"Servers config not found: {self.base_servers_config_path}")

        if not self.base_agent_config_path.exists():
            raise FileNotFoundError(f"Agent config not found: {self.base_agent_config_path}")

        servers_data = json.loads(self.base_servers_config_path.read_text(encoding="utf-8"))
        agent_data = json.loads(self.base_agent_config_path.read_text(encoding="utf-8"))

        port_map = {}

        for name, server in servers_data.get("servers", {}).items():
            if "port" in server:
                port_map[name] = random.randint(1024, 65535)
                server["port"] = port_map[name]

        for name, mcp_server in agent_data.get("mcp_servers", {}).items():
            if name in port_map and "url" in mcp_server:
                mcp_server["url"] = f"http://localhost:{port_map[name]}/mcp"

        self.servers_config_path = self.base_servers_config_path.with_name(
            self.base_servers_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
        )
        self.agent_config_path = self.base_agent_config_path.with_name(
            self.base_agent_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
        )
        self._generated_config_paths = [self.servers_config_path, self.agent_config_path]

        self.servers_config_path.write_text(json.dumps(servers_data, indent=2), encoding="utf-8")
        self.agent_config_path.write_text(json.dumps(agent_data, indent=2), encoding="utf-8")

        self.servers_config = load_server_config(self.servers_config_path)
        self.server_manager.start_servers(self.servers_config_path)

        try:
            active_servers = self.server_manager.state_manager.get_servers(validate=True)
            validate_server_state(self.servers_config["servers"], active_servers)

            self.agent = self.agent_cls(config_path=str(self.agent_config_path))
            self.stack = AsyncExitStack()
            await self.stack.__aenter__()
            await self.agent.initialize(self.stack)

        except Exception:
            await self.close()
            raise

    async def process_query(self, prompt: str) -> str:
        """Process one prompt against the initialized agent.

        Args:
            prompt: Natural-language prompt to send through the managed agent.

        Returns:
            str: Agent response including tool-context annotations.
        """
        assert self.agent is not None, "Call start() first"
        return await self.agent.process_query(prompt, add_tool_context=True)

    async def close(self):
        """Tear down the agent, servers, and temporary config files.

        Cleanup is best-effort so tests do not leak background servers or
        randomized config files even when startup or prompt execution fails.
        """
        try:
            if self.stack is not None:
                await self.stack.aclose()
                self.stack = None
        finally:
            self.server_manager.stop_servers()
            for path in self._generated_config_paths:
                try:
                    path.unlink(missing_ok=True)
                except OSError:
                    pass
            self._generated_config_paths = []
            self.servers_config_path = self.base_servers_config_path
            self.agent_config_path = self.base_agent_config_path

__aenter__() async

Start managed resources when entering an async context.

Source code in mada_tools/testing/agent_runner.py 
80
81
82
83
async def __aenter__(self):
    """Start managed resources when entering an async context."""
    await self.start()
    return self

__aexit__(exc_type, exc, tb) async

Always tear down managed resources when leaving an async context.

Source code in mada_tools/testing/agent_runner.py 
85
86
87
async def __aexit__(self, exc_type, exc, tb):
    """Always tear down managed resources when leaving an async context."""
    await self.close()

__init__(servers_config_path, agent_config_path, agent_cls=MultiServerAgent)

Initialize the test runner.

Parameters:

Name Type Description Default
servers_config_path Path

Path to the MCP server configuration JSON.

 required
agent_config_path Path

Path to the agent configuration JSON.

 required
agent_cls type[AgentProtocol]

Agent implementation to instantiate after servers are running. Defaults to MultiServerAgent.

 MultiServerAgent
Source code in mada_tools/testing/agent_runner.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(
    self,
    servers_config_path: Path,
    agent_config_path: Path,
    agent_cls: type[AgentProtocol] = MultiServerAgent,
):
    """Initialize the test runner.

    Args:
        servers_config_path: Path to the MCP server configuration JSON.
        agent_config_path: Path to the agent configuration JSON.
        agent_cls: Agent implementation to instantiate after servers are
            running. Defaults to `MultiServerAgent`.
    """
    self.base_servers_config_path = Path(servers_config_path)
    self.base_agent_config_path = Path(agent_config_path)
    self.servers_config_path = self.base_servers_config_path
    self.agent_config_path = self.base_agent_config_path
    self.agent_cls = agent_cls

    self.server_manager = ServerManager(
        state_file=Path.home() / ".mada" / f"server_statuses_{uuid.uuid4().hex}.json"
    )
    self.stack: AsyncExitStack | None = None
    self.agent: AgentProtocol | None = None
    self.servers_config: dict[str, Any] | None = None
    self._generated_config_paths: list[Path] = []

close() async

Tear down the agent, servers, and temporary config files.

Cleanup is best-effort so tests do not leak background servers or randomized config files even when startup or prompt execution fails.

Source code in mada_tools/testing/agent_runner.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
async def close(self):
    """Tear down the agent, servers, and temporary config files.

    Cleanup is best-effort so tests do not leak background servers or
    randomized config files even when startup or prompt execution fails.
    """
    try:
        if self.stack is not None:
            await self.stack.aclose()
            self.stack = None
    finally:
        self.server_manager.stop_servers()
        for path in self._generated_config_paths:
            try:
                path.unlink(missing_ok=True)
            except OSError:
                pass
        self._generated_config_paths = []
        self.servers_config_path = self.base_servers_config_path
        self.agent_config_path = self.base_agent_config_path

process_query(prompt) async

Process one prompt against the initialized agent.

Parameters:

Name Type Description Default
prompt str

Natural-language prompt to send through the managed agent.

 required

Returns:

Name Type Description
str str

Agent response including tool-context annotations.
Source code in mada_tools/testing/agent_runner.py 
145
146
147
148
149
150
151
152
153
154
155
async def process_query(self, prompt: str) -> str:
    """Process one prompt against the initialized agent.

    Args:
        prompt: Natural-language prompt to send through the managed agent.

    Returns:
        str: Agent response including tool-context annotations.
    """
    assert self.agent is not None, "Call start() first"
    return await self.agent.process_query(prompt, add_tool_context=True)

start() async

Prepare randomized configs, start servers, and initialize the agent.

Raises:

Type Description
FileNotFoundError

If either input configuration path does not exist.
Exception

Propagates any server-startup or agent-initialization error after best-effort cleanup.
Source code in mada_tools/testing/agent_runner.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
async def start(self):
    """Prepare randomized configs, start servers, and initialize the agent.

    Raises:
        FileNotFoundError: If either input configuration path does not
            exist.
        Exception: Propagates any server-startup or agent-initialization
            error after best-effort cleanup.
    """
    if not self.base_servers_config_path.exists():
        raise FileNotFoundError(f"Servers config not found: {self.base_servers_config_path}")

    if not self.base_agent_config_path.exists():
        raise FileNotFoundError(f"Agent config not found: {self.base_agent_config_path}")

    servers_data = json.loads(self.base_servers_config_path.read_text(encoding="utf-8"))
    agent_data = json.loads(self.base_agent_config_path.read_text(encoding="utf-8"))

    port_map = {}

    for name, server in servers_data.get("servers", {}).items():
        if "port" in server:
            port_map[name] = random.randint(1024, 65535)
            server["port"] = port_map[name]

    for name, mcp_server in agent_data.get("mcp_servers", {}).items():
        if name in port_map and "url" in mcp_server:
            mcp_server["url"] = f"http://localhost:{port_map[name]}/mcp"

    self.servers_config_path = self.base_servers_config_path.with_name(
        self.base_servers_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
    )
    self.agent_config_path = self.base_agent_config_path.with_name(
        self.base_agent_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
    )
    self._generated_config_paths = [self.servers_config_path, self.agent_config_path]

    self.servers_config_path.write_text(json.dumps(servers_data, indent=2), encoding="utf-8")
    self.agent_config_path.write_text(json.dumps(agent_data, indent=2), encoding="utf-8")

    self.servers_config = load_server_config(self.servers_config_path)
    self.server_manager.start_servers(self.servers_config_path)

    try:
        active_servers = self.server_manager.state_manager.get_servers(validate=True)
        validate_server_state(self.servers_config["servers"], active_servers)

        self.agent = self.agent_cls(config_path=str(self.agent_config_path))
        self.stack = AsyncExitStack()
        await self.stack.__aenter__()
        await self.agent.initialize(self.stack)

    except Exception:
        await self.close()
        raise