-
Notifications
You must be signed in to change notification settings - Fork 17
feat(sigv4_helper): inject AWS_REGION in _meta #62
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
arangatang
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the review, glad to see it working! Some comments but nothing fundamental.
mcp_proxy_for_aws/sigv4_helper.py
Outdated
| # Ensure params exists | ||
| # if 'params' not in body: | ||
| # body['params'] = {} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should this be commented out?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe so. There is no MCP request that shouldn't contains a params field.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Then either delete it, or uncomment it. I'd leave it there, fix it and maybe log a warning if this is expected?
mcp_proxy_for_aws/sigv4_helper.py
Outdated
| # Create new content | ||
| new_content = json.dumps(body).encode('utf-8') | ||
|
|
||
| from httpx._content import ByteStream |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Lets not import in the middle of a function.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agreed. Will fix.
mcp_proxy_for_aws/sigv4_helper.py
Outdated
|
|
||
| except (json.JSONDecodeError, KeyError, TypeError) as e: | ||
| # Not a JSON request or invalid format, skip metadata injection | ||
| logger.debug('Skipping metadata injection: %s', e) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Lets make this error level.
mcp_proxy_for_aws/sigv4_helper.py
Outdated
| if '_meta' not in body['params']: | ||
| body['params']['_meta'] = {} | ||
| else: | ||
| print() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why this print?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
oops. remnant from debugging.
| existing_meta = body['params']['_meta'] | ||
|
|
||
| # Merge metadata (existing takes precedence) | ||
| if isinstance(existing_meta, dict): | ||
| body['params']['_meta'] = {**metadata, **existing_meta} | ||
| else: | ||
| body['params']['_meta'] = metadata |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: I think this can be replaced with
body['params']['_meta'] = metadata | body['params']['_meta']
| await _inject_metadata_hook(metadata, request) | ||
|
|
||
| # Clear cached content and read from the new stream | ||
| del request._content |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why the delete here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
After looking at the implementation of I've moved it to request.aread(), I've reached the conclusion that it is indeed unnecessary.sigv4_helper.py, because it is necessary. Since _content functions as a cache, it seems asthough the cached value is read instead of the updated one. Although it doesn't make sense for the caller to need to do this.
|
|
||
| # Merge metadata (existing takes precedence) | ||
| if isinstance(existing_meta, dict): | ||
| body['params']['_meta'] = {**metadata, **existing_meta} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we may need to log if there are failed overwrites here, i.e. if a key was available in both, then this would make it clearer what happened.
|
|
||
|
|
||
| @pytest.mark.asyncio(loop_scope='module') | ||
| async def test_metadata_injection_aws_region( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the integ test!
| assert modified_body['params']['_meta']['field2'] == 'original' | ||
| assert modified_body['params']['_meta']['AWS_REGION'] == 'us-west-2' | ||
|
|
||
| @pytest.mark.asyncio |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Duplicate decorator.
3795fd7 to
b1d59b4
Compare
|
Turn this PR to a draft, as it breaks Sigv4 auth. |
…ta_hook to allow for proper resigning of sigv4 to work
b1d59b4 to
094fb77
Compare
| raise e | ||
|
|
||
|
|
||
| def _resign_request_with_sigv4( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Tbh, if we go with this approach we should not sign it at all before this point and only do signing here. I.e. no need to re-sign.
| # Set the new Content-Length | ||
| request.headers['Content-Length'] = str(len(request.content)) | ||
|
|
||
| logger.info('Headers after cleanup: %s', request.headers) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Don't think we should log this anymore, atleast not with info level.
| service: str, | ||
| profile: Optional[str] = None, | ||
| ) -> None: | ||
| """Re-sign an HTTP request with AWS SigV4 after content modification. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should probably only do signing here and not re-sign. But this changes the core way the proxy works. We need to align on this.
| # Get AWS credentials | ||
| session = create_aws_session(profile) | ||
| credentials = session.get_credentials() | ||
| logger.info('Re-signing request with credentials for access key: %s', credentials.access_key) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should not be logged / be on debug level.
| ) | ||
|
|
||
| # Sign the request | ||
| logger.info('AWS request before signing: %s', aws_request.headers) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same as above, should be debug
| # Sign the request | ||
| logger.info('AWS request before signing: %s', aws_request.headers) | ||
| signer.add_auth(aws_request) | ||
| logger.info('AWS request after signing: %s', aws_request.headers) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same as above, should be debug
|
|
||
| # Update request headers with signed headers | ||
| request.headers.update(dict(aws_request.headers)) | ||
| logger.info('Request headers after re-signing: %s', request.headers) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same as above, should be debug
| ) | ||
| body['params']['_meta'] = {**metadata, **existing_meta} | ||
| else: | ||
| logger.info('Replacing non-dict _meta value with injected metadata') |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What do you mean with non-dict here?
Maybe:
'No conclicting _meta entries exist.`
| # Re-sign the request with the new content | ||
| _resign_request_with_sigv4(request, region, service) | ||
|
|
||
| logger.info('Injected metadata into _meta: %s', body['params']['_meta']) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: MaybeUpdated _meta after injection of additional _meta parameters: <meta>
Summary
Implemented automatic AWS_REGION metadata injection into MCP JSON-RPC requests using
httpxrequest hooks. Added_inject_metadata_hookfunction to sigv4_helper.py that injects metadata into the _meta field of request params, merging with existing metadata when present (existing values take precedence). Includes unit tests and an end-to-end integration test with a new echo_metadata tool in the test MCP server to verify metadata transmission.User experience
AWS_REGION is now automatically injected into all MCP requests transparently - no manual configuration required. The proxy adds region information to the _meta field, making it available to MCP servers without any client-side code changes.
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.