feat(auth): add per-call AWS profile override middleware#205
feat(auth): add per-call AWS profile override middleware#205benjstoll wants to merge 16 commits into
Conversation
f142bea to
1ba8d99
Compare
a496fad to
46cfe20
Compare
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.
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.
…rror on failures Avoids collisions with backend tool arguments by using a namespaced proxy_profile parameter. Errors now raise ToolError instead of returning ToolResult for proper MCP error propagation. Deep-copies tool parameters to prevent mutating shared upstream dicts. Extracts profile override middleware setup into a dedicated helper.
46cfe20 to
b23a615
Compare
e4b979d to
16e6a12
Compare
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.
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.
…rror on failures Avoids collisions with backend tool arguments by using a namespaced proxy_profile parameter. Errors now raise ToolError instead of returning ToolResult for proper MCP error propagation. Deep-copies tool parameters to prevent mutating shared upstream dicts. Extracts profile override middleware setup into a dedicated helper.
|
Hi @benjstoll, could you fix the build errors? We can go ahead with merging this also please run the integration tests and confirm that they're passing |
Looks like it was a bad import, I went ahead and fixed that. Integration tests seems to be passing from my end! |
|
This would unblock us from adopting the managed AWS MCP server. We have 40+ accounts in an Identity Center-managed landing zone and the single-profile-per-session limitation is currently the only reason we're staying on the local The lazy connection creation per-profile and the allowlist approach are both great design choices — no startup cost for unused profiles, and the allowlist gives us a security boundary to prevent agents from accidentally touching production accounts without explicit permission. Looking forward to this landing. |
|
Hi @anasstahr, anything else? Looks like this might be good to go |
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.