Skip to content

Commit d582e4d

Browse files
feat: make httpx connection pool size configurable via flags and env
The single shared httpx.AsyncClient that serves all proxied MCP tool calls was constructed with a hardcoded httpx.Limits(max_keepalive_connections=1, max_connections=5). MCP clients that fan out parallel aws_* tool calls per turn (Claude Code, opencode — typically 4-8 concurrent) saturate the 5-connection pool, causing intermittent failures. Expose the pool size via new CLI flags --max-connections / --max-keepalive-connections with environment-variable fallbacks MCP_PROXY_MAX_CONNECTIONS / MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS. The two values are threaded from cli -> server -> create_transport_with_sigv4 -> create_sigv4_client, and through the per-profile override middleware. Defaults are unchanged: max_connections=5, max_keepalive_connections=1, so existing deployments behave byte-for-byte as before. A backward- compat test asserts the default limits equal the original hardcoded Limits object. Scope: pool-sizing knob only. Does not touch the auth-path serialization (_sign_request_hook / next(auth_flow)), which is tracked separately in #270. Fixes #295
1 parent ac093f0 commit d582e4d

11 files changed

Lines changed: 239 additions & 2 deletions

File tree

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## Unreleased
99

10+
### Added
11+
12+
- Configurable HTTP connection pool size via `--max-connections` / `--max-keepalive-connections` flags and `MCP_PROXY_MAX_CONNECTIONS` / `MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS` environment variables. Defaults (5/1) preserve existing behavior (#295)
13+
1014
## v1.6.0 (2026-05-29)
1115

1216
### Added

README.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,8 @@ docker build -t mcp-proxy-for-aws .
110110
| `--tool-timeout` | Maximum seconds a tool call may take before being cancelled. When set, returns a graceful error to the agent instead of hanging indefinitely | 300 |No |
111111
| `--skip-auth` | Skip request signing when AWS credentials are unavailable instead of failing | `False` |No |
112112
| `--disable-telemetry` | Disables telemetry data collection | `False` |No |
113+
| `--max-connections` | Maximum number of concurrent HTTP connections in the shared pool used for all proxied tool calls. Increase this when your MCP client fans out parallel `aws_*` tool calls per turn (e.g. Claude Code, opencode). | 5 (falls back to `MCP_PROXY_MAX_CONNECTIONS`) |No |
114+
| `--max-keepalive-connections` | Maximum number of idle keep-alive connections to retain in the pool | 1 (falls back to `MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS`) |No |
113115

114116
### Optional Environment Variables
115117

@@ -129,9 +131,16 @@ export AWS_REGION=<aws_region>
129131

130132
# Multi-profile switching (alternative to --profile flag, useful for plugin integration)
131133
export AWS_MCP_PROXY_PROFILES="prod-readonly dev staging"
134+
135+
# HTTP connection pool sizing (alternative to --max-connections / --max-keepalive-connections)
136+
# Raise the pool size when your MCP client issues many parallel aws_* tool calls per turn.
137+
export MCP_PROXY_MAX_CONNECTIONS=25
138+
export MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS=5
132139
```
133140

134141
> **Note:** `AWS_MCP_PROXY_PROFILES` takes precedence over `--profile` / `AWS_PROFILE` when set.
142+
>
143+
> **Note:** `--max-connections` / `--max-keepalive-connections` take precedence over `MCP_PROXY_MAX_CONNECTIONS` / `MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS` when explicitly passed. Defaults are `5` / `1`, preserving the previous behavior.
135144
136145
### Multi-account access
137146

mcp_proxy_for_aws/cli.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import argparse
1818
import os
1919
from mcp_proxy_for_aws import __version__
20-
from mcp_proxy_for_aws.utils import within_range
20+
from mcp_proxy_for_aws.utils import positive_int, within_range
2121
from typing import Any, Dict, Optional, Sequence
2222

2323

@@ -185,4 +185,22 @@ def parse_args():
185185
help='Skip request signing when AWS credentials are unavailable instead of failing',
186186
)
187187

188+
parser.add_argument(
189+
'--max-connections',
190+
type=positive_int(1),
191+
default=positive_int(1)(os.getenv('MCP_PROXY_MAX_CONNECTIONS', '5')),
192+
help='Maximum number of concurrent HTTP connections in the pool used for all '
193+
'proxied tool calls. Increase this when your MCP client fans out parallel '
194+
'aws_* tool calls per turn (default: 5). '
195+
'Falls back to MCP_PROXY_MAX_CONNECTIONS if not set.',
196+
)
197+
198+
parser.add_argument(
199+
'--max-keepalive-connections',
200+
type=positive_int(0),
201+
default=positive_int(0)(os.getenv('MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS', '1')),
202+
help='Maximum number of idle keep-alive connections to retain in the pool '
203+
'(default: 1). Falls back to MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS if not set.',
204+
)
205+
188206
return parser.parse_args()

mcp_proxy_for_aws/middleware/profile_switcher.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,8 @@ def __init__(
6666
endpoint: str,
6767
disable_telemetry: bool = False,
6868
skip_auth: bool = False,
69+
max_connections: int = 5,
70+
max_keepalive_connections: int = 1,
6971
) -> None:
7072
"""Initialize the middleware with connection and profile configuration."""
7173
super().__init__()
@@ -78,6 +80,8 @@ def __init__(
7880
self._timeout = timeout
7981
self._disable_telemetry = disable_telemetry
8082
self._skip_auth = skip_auth
83+
self._max_connections = max_connections
84+
self._max_keepalive_connections = max_keepalive_connections
8185
self._profile_clients: dict[str, Client] = {}
8286
self._lock = asyncio.Lock()
8387

@@ -172,6 +176,8 @@ async def _get_profile_client(self, profile: str) -> Client:
172176
profile,
173177
self._disable_telemetry,
174178
self._skip_auth,
179+
self._max_connections,
180+
self._max_keepalive_connections,
175181
)
176182
client = Client(transport=transport)
177183
await client.__aenter__()

mcp_proxy_for_aws/server.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,8 @@ async def run_proxy(args) -> None:
108108
default_profile,
109109
args.disable_telemetry,
110110
args.skip_auth,
111+
args.max_connections,
112+
args.max_keepalive_connections,
111113
)
112114
client_factory = AWSMCPProxyClientFactory(transport)
113115

@@ -138,6 +140,8 @@ async def run_proxy(args) -> None:
138140
args.endpoint,
139141
args.disable_telemetry,
140142
args.skip_auth,
143+
args.max_connections,
144+
args.max_keepalive_connections,
141145
)
142146

143147
if args.retries:
@@ -174,6 +178,8 @@ def add_profile_override_middleware(
174178
endpoint: str,
175179
disable_telemetry: bool = False,
176180
skip_auth: bool = False,
181+
max_connections: int = 5,
182+
max_keepalive_connections: int = 1,
177183
) -> ProfileOverrideMiddleware | None:
178184
"""Add profile override middleware to target MCP server.
179185
@@ -188,6 +194,8 @@ def add_profile_override_middleware(
188194
endpoint: The MCP endpoint URL
189195
disable_telemetry: Whether to disable telemetry on profile transports
190196
skip_auth: Whether to skip signing when credentials are unavailable
197+
max_connections: Maximum concurrent connections per profile transport pool
198+
max_keepalive_connections: Maximum idle keep-alive connections per profile pool
191199
192200
Returns:
193201
The ProfileOverrideMiddleware instance if added, None otherwise
@@ -207,6 +215,8 @@ def add_profile_override_middleware(
207215
endpoint=endpoint,
208216
disable_telemetry=disable_telemetry,
209217
skip_auth=skip_auth,
218+
max_connections=max_connections,
219+
max_keepalive_connections=max_keepalive_connections,
210220
)
211221
mcp.add_middleware(middleware)
212222
return middleware

mcp_proxy_for_aws/sigv4_helper.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ def create_sigv4_client(
148148
metadata: Optional[Dict[str, Any]] = None,
149149
disable_telemetry: bool = False,
150150
skip_auth: bool = False,
151+
max_connections: int = 5,
152+
max_keepalive_connections: int = 1,
151153
**kwargs: Any,
152154
) -> httpx.AsyncClient:
153155
"""Create an httpx.AsyncClient with SigV4 authentication.
@@ -161,6 +163,10 @@ def create_sigv4_client(
161163
metadata: Metadata to inject into MCP _meta field
162164
disable_telemetry: Whether to disable telemetry
163165
skip_auth: Whether to skip signing when credentials are unavailable
166+
max_connections: Maximum number of concurrent connections in the pool
167+
(default: 5, preserving historical behavior)
168+
max_keepalive_connections: Maximum number of idle keep-alive connections
169+
to retain (default: 1, preserving historical behavior)
164170
**kwargs: Additional arguments to pass to httpx.AsyncClient
165171
166172
Returns:
@@ -170,7 +176,10 @@ def create_sigv4_client(
170176
client_kwargs = {
171177
'follow_redirects': True,
172178
'timeout': timeout,
173-
'limits': httpx.Limits(max_keepalive_connections=1, max_connections=5),
179+
'limits': httpx.Limits(
180+
max_keepalive_connections=max_keepalive_connections,
181+
max_connections=max_connections,
182+
),
174183
**kwargs,
175184
}
176185

mcp_proxy_for_aws/utils.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ def create_transport_with_sigv4(
7575
profile: Optional[str] = None,
7676
disable_telemetry: bool = False,
7777
skip_auth: bool = False,
78+
max_connections: int = 5,
79+
max_keepalive_connections: int = 1,
7880
) -> StreamableHttpTransport:
7981
"""Create a StreamableHttpTransport with SigV4 authentication.
8082
@@ -87,6 +89,10 @@ def create_transport_with_sigv4(
8789
profile: AWS profile to use (optional)
8890
disable_telemetry: Whether to disable telemetry
8991
skip_auth: Whether to skip signing when credentials are unavailable
92+
max_connections: Maximum number of concurrent connections in the pool
93+
(default: 5, preserving historical behavior)
94+
max_keepalive_connections: Maximum number of idle keep-alive connections
95+
to retain (default: 1, preserving historical behavior)
9096
9197
Returns:
9298
StreamableHttpTransport instance with SigV4 authentication
@@ -107,6 +113,8 @@ def client_factory(
107113
metadata=metadata,
108114
disable_telemetry=disable_telemetry,
109115
skip_auth=skip_auth,
116+
max_connections=max_connections,
117+
max_keepalive_connections=max_keepalive_connections,
110118
auth=auth,
111119
**kw,
112120
)
@@ -247,3 +255,37 @@ def validator(value):
247255
return fvalue
248256

249257
return validator
258+
259+
260+
def positive_int(min_value: int, max_value: Optional[int] = None):
261+
"""Factory function to create integer range validators.
262+
263+
Mirrors ``within_range`` but parses and returns ``int`` instead of ``float``,
264+
which is the correct type for discrete counts such as connection-pool sizes.
265+
266+
Args:
267+
min_value: Minimum allowed value (inclusive)
268+
max_value: Maximum allowed value (inclusive), or None for no upper bound
269+
270+
Returns:
271+
The argparse validator function
272+
273+
Raises:
274+
argparse.ArgumentTypeError: If the value is not an integer within range
275+
"""
276+
277+
def validator(value):
278+
try:
279+
ivalue = int(value)
280+
except (ValueError, TypeError):
281+
raise argparse.ArgumentTypeError(f"'{value}' is not a valid integer")
282+
283+
if ivalue < min_value:
284+
raise argparse.ArgumentTypeError(f"'{value}' must be >= {min_value}")
285+
286+
if max_value is not None and ivalue > max_value:
287+
raise argparse.ArgumentTypeError(f"'{value}' must be <= {max_value}")
288+
289+
return ivalue
290+
291+
return validator

tests/unit/test_cli.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,3 +202,93 @@ def test_parse_args_disable_telemetry_default(self):
202202
args = parse_args()
203203

204204
assert args.disable_telemetry is False
205+
206+
@patch.dict('os.environ', {}, clear=True)
207+
@patch('sys.argv', ['mcp-proxy-for-aws', 'https://test.example.com'])
208+
def test_parse_args_connection_pool_defaults(self):
209+
"""Pool sizing flags default to today's 5/1 behavior when unset."""
210+
args = parse_args()
211+
212+
assert args.max_connections == 5
213+
assert args.max_keepalive_connections == 1
214+
215+
@patch.dict('os.environ', {}, clear=True)
216+
@patch(
217+
'sys.argv',
218+
[
219+
'mcp-proxy-for-aws',
220+
'https://test.example.com',
221+
'--max-connections',
222+
'25',
223+
'--max-keepalive-connections',
224+
'10',
225+
],
226+
)
227+
def test_parse_args_connection_pool_flags(self):
228+
"""Pool sizing flags parse from the command line."""
229+
args = parse_args()
230+
231+
assert args.max_connections == 25
232+
assert args.max_keepalive_connections == 10
233+
234+
@patch.dict(
235+
'os.environ',
236+
{
237+
'MCP_PROXY_MAX_CONNECTIONS': '40',
238+
'MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS': '8',
239+
},
240+
clear=True,
241+
)
242+
@patch('sys.argv', ['mcp-proxy-for-aws', 'https://test.example.com'])
243+
def test_parse_args_connection_pool_from_env(self):
244+
"""Pool sizing falls back to env vars when flags are absent."""
245+
args = parse_args()
246+
247+
assert args.max_connections == 40
248+
assert args.max_keepalive_connections == 8
249+
250+
@patch.dict(
251+
'os.environ',
252+
{
253+
'MCP_PROXY_MAX_CONNECTIONS': '40',
254+
'MCP_PROXY_MAX_KEEPALIVE_CONNECTIONS': '8',
255+
},
256+
clear=True,
257+
)
258+
@patch(
259+
'sys.argv',
260+
[
261+
'mcp-proxy-for-aws',
262+
'https://test.example.com',
263+
'--max-connections',
264+
'25',
265+
'--max-keepalive-connections',
266+
'10',
267+
],
268+
)
269+
def test_parse_args_connection_pool_cli_overrides_env(self):
270+
"""Explicit pool flags take precedence over env vars."""
271+
args = parse_args()
272+
273+
assert args.max_connections == 25
274+
assert args.max_keepalive_connections == 10
275+
276+
@patch.dict('os.environ', {}, clear=True)
277+
@patch(
278+
'sys.argv',
279+
['mcp-proxy-for-aws', 'https://test.example.com', '--max-connections', '0'],
280+
)
281+
def test_parse_args_max_connections_below_minimum(self):
282+
"""--max-connections must be >= 1 (within_range validation)."""
283+
with pytest.raises(SystemExit):
284+
parse_args()
285+
286+
@patch.dict('os.environ', {}, clear=True)
287+
@patch(
288+
'sys.argv',
289+
['mcp-proxy-for-aws', 'https://test.example.com', '--max-keepalive-connections', '-1'],
290+
)
291+
def test_parse_args_max_keepalive_connections_below_minimum(self):
292+
"""--max-keepalive-connections must be >= 0 (within_range validation)."""
293+
with pytest.raises(SystemExit):
294+
parse_args()

tests/unit/test_profile_switcher.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -808,4 +808,6 @@ async def test_all_params_forwarded_to_transport_factory(self, mock_context):
808808
'dev-profile',
809809
True,
810810
True,
811+
5,
812+
1,
811813
)

tests/unit/test_sigv4_helper.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,47 @@ def test_create_sigv4_client_user_agent_includes_client_info_when_telemetry_enab
273273
assert 'my-client/2.0' in user_agent
274274
assert result == mock_client
275275

276+
@patch('mcp_proxy_for_aws.sigv4_helper.get_client_info', return_value=None)
277+
@patch('httpx.AsyncClient')
278+
def test_create_sigv4_client_default_connection_limits(
279+
self, mock_client_class, mock_get_client_info
280+
):
281+
"""Default limits MUST preserve historical 5/1 pool behavior byte-for-byte."""
282+
mock_client = Mock()
283+
mock_client_class.return_value = mock_client
284+
285+
create_sigv4_client(service='test-service', region='test-region')
286+
287+
call_args = mock_client_class.call_args
288+
limits = call_args[1]['limits']
289+
assert isinstance(limits, httpx.Limits)
290+
assert limits.max_connections == 5
291+
assert limits.max_keepalive_connections == 1
292+
# Backward compatibility guard: identical to the original hardcoded Limits.
293+
assert limits == httpx.Limits(max_keepalive_connections=1, max_connections=5)
294+
295+
@patch('mcp_proxy_for_aws.sigv4_helper.get_client_info', return_value=None)
296+
@patch('httpx.AsyncClient')
297+
def test_create_sigv4_client_custom_connection_limits(
298+
self, mock_client_class, mock_get_client_info
299+
):
300+
"""Explicit pool overrides MUST flow into the httpx.AsyncClient limits."""
301+
mock_client = Mock()
302+
mock_client_class.return_value = mock_client
303+
304+
create_sigv4_client(
305+
service='test-service',
306+
region='test-region',
307+
max_connections=25,
308+
max_keepalive_connections=10,
309+
)
310+
311+
call_args = mock_client_class.call_args
312+
limits = call_args[1]['limits']
313+
assert isinstance(limits, httpx.Limits)
314+
assert limits.max_connections == 25
315+
assert limits.max_keepalive_connections == 10
316+
276317

277318
class TestSanitizeHeaders:
278319
"""Test cases for the _sanitize_headers function."""

0 commit comments

Comments
 (0)