Skip to content

feat(auth): add per-call AWS profile override middleware#205

Open
benjstoll wants to merge 16 commits into
aws:mainfrom
benjstoll:feat/profile-override-middleware
Open

feat(auth): add per-call AWS profile override middleware#205
benjstoll wants to merge 16 commits into
aws:mainfrom
benjstoll:feat/profile-override-middleware

Conversation

@benjstoll
Copy link
Copy Markdown

@benjstoll benjstoll commented Mar 24, 2026

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.

Summary

Changes

  • Added ProfileOverrideMiddleware in mcp_proxy_for_aws/middleware/profile_switcher.py that intercepts a profile argument on any tool call,
    validates it against an allowlist, and routes the request through a dedicated per-profile MCP client with its own SigV4-signed transport.
  • Added --allow-switch-profile CLI argument in cli.py that accepts one or more AWS profile names to enable the middleware.
  • Wired the middleware into server.py with proper lifecycle management (lazy client creation, graceful shutdown of per-profile connections
    in the finally block).
  • Added 12 unit tests covering tool schema injection, pass-through behavior, disallowed profiles, argument stripping, connection/tool-call
    error handling, and client disconnect logic. 90% branch coverage on the new middleware.
  • Updated README.md with the new parameter in the configuration table and a "Multi-account access" section explaining how
    --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.

  • I have reviewed the contributing guidelines
  • I have performed a self-review of this change
  • Changes have been tested
  • Changes are documented

Is this a breaking change? (Y/N)

  • Yes
  • No

Please add details about how this change was tested.

  • Did integration tests succeed?
  • If the feature is a new use case, is it necessary to add a new integration test case?

Acknowledgment

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py
Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py Outdated
@benjstoll benjstoll force-pushed the feat/profile-override-middleware branch 3 times, most recently from f142bea to 1ba8d99 Compare April 6, 2026 12:57
Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py Outdated
Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py Outdated
Comment thread mcp_proxy_for_aws/server.py Outdated
Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py Outdated
Comment thread mcp_proxy_for_aws/middleware/profile_switcher.py
@benjstoll benjstoll force-pushed the feat/profile-override-middleware branch from a496fad to 46cfe20 Compare April 13, 2026 15:03
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.
@benjstoll benjstoll force-pushed the feat/profile-override-middleware branch from 46cfe20 to b23a615 Compare April 13, 2026 15:08
@benjstoll benjstoll requested a review from rshevchuk-git April 20, 2026 12:07
@benjstoll benjstoll force-pushed the feat/profile-override-middleware branch from e4b979d to 16e6a12 Compare April 20, 2026 12:11
@anasstahr anasstahr marked this pull request as ready for review May 1, 2026 07:53
@anasstahr anasstahr requested a review from a team as a code owner May 1, 2026 07:53
@anasstahr anasstahr requested a review from arnewouters May 1, 2026 07:53
benjstoll added 2 commits May 1, 2026 10:01
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.
benjstoll and others added 3 commits May 1, 2026 10:01
…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.
@anasstahr
Copy link
Copy Markdown
Contributor

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

@benjstoll
Copy link
Copy Markdown
Author

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!

@rorymcmahon
Copy link
Copy Markdown

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 awslabs.aws-api-mcp-server.

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.

@benjstoll
Copy link
Copy Markdown
Author

Hi @anasstahr, anything else? Looks like this might be good to go

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

5 participants