feat(auth): add per-call AWS profile override middleware#205
feat(auth): add per-call AWS profile override middleware#205
Conversation
Adds ProfileOverrideMiddleware that allows routing individual tool calls through dedicated per-profile MCP connections via a `profile` argument. Enabled with `--allow-switch-profile` CLI flag restricted to an explicit allowlist of profile names.
benjstoll
force-pushed
the
feat/profile-override-middleware
branch
2 times, most recently
from
4476898 to
f142bea
Compare
Concurrent tool calls (e.g. from parallel subagents) could race in _get_profile_client, each creating a separate Client for the same profile. The loser's client would leak — connected but never tracked or cleaned up. Wrapping in an asyncio.Lock ensures only one client is created per profile.
benjstoll
force-pushed
the
feat/profile-override-middleware
branch
from
f142bea to
1ba8d99
Compare
| """Intercept ``profile`` and route through a dedicated per-profile client.""" | ||
| arguments = context.message.arguments | ||
| if isinstance(arguments, dict) and 'profile' in arguments: | ||
| profile = arguments['profile'] |
There was a problem hiding this comment.
profile can collide with tools that already have a profile parameter. If a backend tool already has a profile parameter in its schema, this middleware will silently overwrite it in on_list_tools and strip it in on_call_tool. This could break legitimate tool parameters.
| """Forward a tool call through a dedicated per-profile connection.""" | ||
| if profile not in self._allowed_profiles: | ||
| allowed = ', '.join(sorted(self._allowed_profiles)) | ||
| return ToolResult( |
There was a problem hiding this comment.
Why do you return ToolResult(s) on failed validations instead of raising ToolError?
| allowed_profiles = getattr(args, 'allow_switch_profile', None) | ||
| if isinstance(allowed_profiles, list) and allowed_profiles: | ||
| profile_middleware = ProfileOverrideMiddleware( | ||
| allowed_profiles=allowed_profiles, | ||
| service=service, | ||
| region=region, | ||
| metadata=metadata, | ||
| timeout=timeout, | ||
| endpoint=args.endpoint, | ||
| ) | ||
| proxy.add_middleware(profile_middleware) | ||
|
|
There was a problem hiding this comment.
This should be extracted into a separate helper add_profile_override_middleware, to stay consistent with other middleware declarations
| continue | ||
| if 'properties' not in params: | ||
| params['properties'] = {} | ||
| params['properties']['profile'] = { |
There was a problem hiding this comment.
If the upstream tool list is cached or shared, this mutates shared state. This is ok now because AWSProxyToolManager.get_tools() re-fetches, however this creates dependency on the upstream to always produce fresh dicts. I'd suggest deep-copy the parameters dict before mutating. It is cheap insurance and makes the middleware self-contained.
| """Inject ``profile`` into every tool's schema.""" | ||
| tools = await call_next(context) | ||
|
|
||
| for tool in tools: |
There was a problem hiding this comment.
Curious if we could use tool transformations instead of directly manipulating params?
4 participants
Adds ProfileOverrideMiddleware that allows routing individual tool calls through dedicated per-profile MCP connections via a
profileargument. Enabled with--allow-switch-profileCLI flag restricted to an explicit allowlist of profile names.Summary
Changes
validates it against an allowlist, and routes the request through a dedicated per-profile MCP client with its own SigV4-signed transport.
in the finally block).
error handling, and client disconnect logic. 90% branch coverage on the new middleware.
--allow-switch-profile interacts with --profile, including a JSON config example.
User experience
Before: Users who needed to query AWS resources across multiple accounts had to run separate proxy instances per profile, or manually
restart the proxy with a different --profile value.
After: Users pass --allow-switch-profile profile-a profile-b alongside their default --profile. Any tool call can include a profile
argument to route that single request through a dedicated connection signed with the specified profile's credentials. Tool calls without
profile continue to use the default connection. Each profile's connection is created lazily on first use, so there is no startup cost for
unused profiles.
Checklist
If your change doesn't seem to apply, please leave them unchecked.
Is this a breaking change? (Y/N)
Please add details about how this change was tested.
Acknowledgment
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.