feat(security-agent-mcp-server): Add custom User-Agent - #4319
Conversation
5f53f7f to
7f6e306
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4319 +/- ##
=======================================
Coverage 93.12% 93.13%
=======================================
Files 1036 1036
Lines 87307 87358 +51
Branches 14085 14090 +5
=======================================
+ Hits 81308 81359 +51
Misses 3638 3638
Partials 2361 2361 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
b67f874 to
fe66696
Compare
479bcda to
904af8e
Compare
scottschreckengaust
left a comment
There was a problem hiding this comment.
Thanks for adding usage telemetry — the lazy-injection approach is reasonable and every failure path degrades gracefully (no crashes). I reviewed at high effort and verified the technical claims below empirically. Two items keep the PR from delivering its own stated contract and are worth fixing before merge; the rest are cleanup/altitude.
Should fix (contract gaps):
get_api_guidesends an untagged User-Agent. It builds its ownsession.client('securityagent')withoutconfig=, so it's the one.client()site the PR didn't thread the custom UA through — contradicting "passes it to every.client()call." The_ensure_client_ua(ctx)at the top of that tool is also inert, since it mutates the_clientsingleton this tool bypasses.- Version reports
unknownfrom any uninstalled/editable checkout.importlib.metadata.version('awslabs.security-agent-mcp-server')raisesPackageNotFoundErrorwhen run from source (verified in this repo), silently degrading tounknown.__version__ = '0.1.2'is already exported by__init__.py, and ~30 other servers in this repo build the UA from that constant. Reusing it also removes the per-rebuild metadata lookup.
Consider (cleanup / consistency):
3. clientInfo.name is injected into the UA unsanitized, unlike _client_prefix which normalizes it — a name like Claude Code produces a broken UA token (botocore does not sanitize user_agent_extra; confirmed).
4. _ensure_client_ua(ctx) is duplicated as the first line of all 11 tools, and duplicates the clientInfo-extraction logic already in _client_prefix (with a divergent fallback: unknown vs ide).
5. The shared mutable _client._config is safe under today's stdio transport but would cross-contaminate the UA across concurrent sessions under HTTP/SSE.
Inline comments below. Nothing here is a crash; posting as a Comment, not a blocking Request Changes.
1b88b92 to
1f99fec
Compare
…tracking Add a custom User-Agent string to all boto3 API calls made by the security-agent-mcp-server. The User-Agent includes: - MCP server name and version (awslabs-security-agent-mcp-server/X.Y.Z) - Calling MCP client name and version (e.g. kiro/1.5.0, claude-code/1.0) This enables service-side tracking of API usage originating from the local MCP server, broken down by which IDE/tool invoked it. The client identity is extracted from the MCP protocol clientInfo field sent during the initialize handshake. Falls back to unknown when clientInfo is not available. Changes: - aws_client.py: Added botocore.config.Config with user_agent_extra to all boto3 client calls (securityagent, sts, s3, iam). Added set_mcp_client_info() for lazy injection after MCP init. - server.py: Added _ensure_client_ua(ctx) that extracts clientInfo from the MCP session context and injects it into the AWS client on the first tool invocation.
… refactor to decorator, fix missed client call - Use __version__ constant instead of importlib.metadata.version() for reliability - Sanitize mcp_client_name (lowercase, replace spaces) before injecting into UA - Add concurrency note about per-session isolation for future HTTP/SSE transport - Refactor _ensure_client_ua from per-tool boilerplate to @ensure_client_ua decorator - Fix get_api_guide to pass config=_client._config for tagged UA - Remove obsolete importlib fallback test
1f99fec to
f815bc6
Compare
Summary
Adds a custom User-Agent string to all boto3 API calls from the security-agent-mcp-server, enabling service-side tracking of MCP usage by IDE/tool.
User-Agent format
Follows the repo-wide
md/awslabs#mcp#<server-name>#<version>convention, with an additionalmd/client#<name>/<version>suffix for IDE tracking.How it works
aws_client.py:SecurityAgentClientbuilds abotocore.config.Config(user_agent_extra=...)and passes it to every.client()call (securityagent, sts, s3, iam). Uses the package__version__constant directly.server.py:@ensure_client_uadecorator extractsclientInfofrom the MCP session context (sent by the IDE duringinitialize) and callsset_mcp_client_info()on the shared AWS client singleton.Design decisions
__version__constant: Uses the package-level__version__instead ofimportlib.metadata.version()to avoid silent degradation to "unknown" in editable/source installs.clientInfo.nameis normalized with.lower().replace(" ", "-")before injection into the UA string, preventing malformed tokens (e.g. "Claude Code" → "claude-code").@ensure_client_uaapplied to all tool functions as a single choke point — no per-tool boilerplate, and new tools cannot forget it.SecurityAgentClientsingleton is created at module level before any MCP session connects, soclientInfois injected on first tool invocation.set_mcp_client_info()short-circuits if the info has not changed.unknownwhenclientInfois unavailable._configwould need per-session isolation.Tracking values by IDE
Tested
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.