Skip to content

Conversation

@kyoncal
Copy link
Contributor

@kyoncal kyoncal commented Oct 30, 2025

Summary

Implemented automatic AWS_REGION metadata injection into MCP JSON-RPC requests using httpx request hooks. Added _inject_metadata_hook function 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.

  • 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? (Only locally, ACR doesn't work for me on this commit or main head)
  • 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.

@kyoncal kyoncal requested a review from a team as a code owner October 30, 2025 14:42
@kyoncal kyoncal requested review from arangatang and bidesh October 30, 2025 14:42
Copy link
Contributor

@arangatang arangatang left a 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.

Comment on lines 144 to 147
# Ensure params exists
# if 'params' not in body:
# body['params'] = {}
Copy link
Contributor

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?

Copy link
Contributor Author

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.

Copy link
Contributor

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?

# Create new content
new_content = json.dumps(body).encode('utf-8')

from httpx._content import ByteStream
Copy link
Contributor

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. Will fix.


except (json.JSONDecodeError, KeyError, TypeError) as e:
# Not a JSON request or invalid format, skip metadata injection
logger.debug('Skipping metadata injection: %s', e)
Copy link
Contributor

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.

if '_meta' not in body['params']:
body['params']['_meta'] = {}
else:
print()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this print?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oops. remnant from debugging.

Comment on lines 155 to 227
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
Copy link
Contributor

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why the delete here?

Copy link
Contributor Author

@kyoncal kyoncal Oct 30, 2025

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 request.aread(), I've reached the conclusion that it is indeed unnecessary. I've moved it to 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}
Copy link
Contributor

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(
Copy link
Contributor

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicate decorator.

@kyoncal kyoncal marked this pull request as draft November 3, 2025 10:48
@kyoncal
Copy link
Contributor Author

kyoncal commented Nov 3, 2025

Turn this PR to a draft, as it breaks Sigv4 auth.

@kyoncal kyoncal force-pushed the forward-region-via-meta branch from b1d59b4 to 094fb77 Compare November 4, 2025 00:55
@kyoncal kyoncal marked this pull request as ready for review November 4, 2025 00:56
raise e


def _resign_request_with_sigv4(
Copy link
Contributor

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)
Copy link
Contributor

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.
Copy link
Contributor

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)
Copy link
Contributor

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)
Copy link
Contributor

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)
Copy link
Contributor

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)
Copy link
Contributor

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')
Copy link
Contributor

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'])
Copy link
Contributor

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>

@kyoncal kyoncal closed this Nov 4, 2025
@kyoncal kyoncal mentioned this pull request Nov 5, 2025
8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants