feat(security-agent-mcp-server): Add custom User-Agent for API usage 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.

Author: jigar-lab

src/security-agent-mcp-server/awslabs/security_agent_mcp_server/aws_client.py

@@ -15,25 +15,70 @@
 """AWS SecurityAgent API client using boto3 SDK."""
 
 import boto3
+import botocore.config
 import json
 import re
 from typing import Any, Optional
 
 
+DEFAULT_MCP_CLIENT_NAME = 'unknown'
+
+
 class SecurityAgentClient:
     """Client for AWS SecurityAgent APIs using boto3."""
 
-    def __init__(self, region: str = 'us-east-1'):
+    def __init__(
+        self,
+        region: str = 'us-east-1',
+        mcp_client_name: str = DEFAULT_MCP_CLIENT_NAME,
+        mcp_client_version: str = '',
+    ):
         """Initialize SecurityAgent client."""
         self.region = region
+        self._mcp_client_name = mcp_client_name
+        self._mcp_client_version = mcp_client_version
+        self._config = self._build_config(mcp_client_name, mcp_client_version)
+
+    def _build_config(
+        self, mcp_client_name: str, mcp_client_version: str
+    ) -> botocore.config.Config:
+        """Build a botocore Config with a custom user_agent_extra string."""
+        try:
+            from importlib.metadata import version as pkg_version
+
+            mcp_server_version = pkg_version('awslabs.security-agent-mcp-server')
+        except Exception:
+            mcp_server_version = 'unknown'
+
+        ua_extra = (
+            f'awslabs-security-agent-mcp-server/{mcp_server_version} md/client#{mcp_client_name}'
+        )
+        if mcp_client_version:
+            ua_extra += f'/{mcp_client_version}'
+
+        return botocore.config.Config(user_agent_extra=ua_extra)
+
+    def set_mcp_client_info(self, mcp_client_name: str, mcp_client_version: str = '') -> None:
+        """Update the MCP client identity and rebuild the botocore config.
+
+        Called after MCP session initialization when clientInfo becomes available.
+        """
+        if (
+            mcp_client_name == self._mcp_client_name
+            and mcp_client_version == self._mcp_client_version
+        ):
+            return  # No change needed
+        self._mcp_client_name = mcp_client_name
+        self._mcp_client_version = mcp_client_version
+        self._config = self._build_config(mcp_client_name, mcp_client_version)
 
     def _get_session(self):
         """Fresh session each call to pick up rotated credentials."""
         return boto3.Session(region_name=self.region)
 
     def _client(self):
-        """Get a fresh securityagent boto3 client."""
-        return self._get_session().client('securityagent')
+        """Get a fresh securityagent boto3 client with custom user-agent."""
+        return self._get_session().client('securityagent', config=self._config)
 
     def call(self, operation: str, params: dict) -> dict:
         """Call any SecurityAgent API operation generically."""
@@ -50,7 +95,7 @@ def call(self, operation: str, params: dict) -> dict:
 
     def get_caller_identity(self) -> dict:
         """Get the current AWS caller identity."""
-        return self._get_session().client('sts').get_caller_identity()
+        return self._get_session().client('sts', config=self._config).get_caller_identity()
 
     def list_agent_spaces(self) -> list[dict]:
         """List all SecurityAgent agent spaces."""
@@ -224,7 +269,7 @@ def create_s3_bucket(self, bucket_name: str) -> str:
 
         Public-access block, SSE-S3, TLS-only policy, and 30-day lifecycle.
         """
-        s3 = self._get_session().client('s3')
+        s3 = self._get_session().client('s3', config=self._config)
         create_args: dict[str, Any] = {'Bucket': bucket_name}
         if self.region != 'us-east-1':
             create_args['CreateBucketConfiguration'] = {'LocationConstraint': self.region}
@@ -287,7 +332,7 @@ def create_s3_bucket(self, bucket_name: str) -> str:
 
     def create_service_role(self, role_name: str, account_id: str, bucket_name: str) -> str:
         """Create IAM service role for Security Agent with S3 + CloudWatch Logs access."""
-        iam = self._get_session().client('iam')
+        iam = self._get_session().client('iam', config=self._config)
 
         trust_policy = json.dumps(
             {
@@ -347,5 +392,5 @@ def create_service_role(self, role_name: str, account_id: str, bucket_name: str)
 
     def upload_to_s3(self, bucket: str, key: str, file_path: str) -> str:
         """Upload a file to S3."""
-        self._get_session().client('s3').upload_file(file_path, bucket, key)
+        self._get_session().client('s3', config=self._config).upload_file(file_path, bucket, key)
         return f's3://{bucket}/{key}'

src/security-agent-mcp-server/awslabs/security_agent_mcp_server/server.py

@@ -17,7 +17,10 @@
 import json
 import os
 import sys
-from awslabs.security_agent_mcp_server.aws_client import SecurityAgentClient
+from awslabs.security_agent_mcp_server.aws_client import (
+    DEFAULT_MCP_CLIENT_NAME,
+    SecurityAgentClient,
+)
 from awslabs.security_agent_mcp_server.consts import (
     DEFAULT_REGION,
     SERVER_INSTRUCTIONS,
@@ -140,6 +143,32 @@ def _client_prefix(ctx: Context) -> str:
         return 'ide'
 
 
+def _ensure_client_ua(ctx: Context) -> None:
+    """Inject MCP clientInfo into the SecurityAgentClient user-agent on first use.
+
+    Called at the start of each tool invocation. Since the module-level _client
+    singleton is created before any MCP session connects, this defers the
+    user-agent configuration until clientInfo is actually available.
+    """
+    try:
+        session = ctx.session
+        if session is None:
+            return
+        client_params = session.client_params
+        if client_params is None:
+            return
+        info = client_params.clientInfo  # type: ignore[union-attr]
+        if info is None:
+            return
+        name = info.name if isinstance(info.name, str) else DEFAULT_MCP_CLIENT_NAME
+        version = (
+            info.version if hasattr(info, 'version') and isinstance(info.version, str) else ''
+        )
+        _client.set_mcp_client_info(name, version)
+    except (AttributeError, TypeError):
+        pass
+
+
 def _ensure_s3_bucket(config: dict, kind: str = 'scans') -> None:
     """Lazily create and register the per-account S3 bucket for the given kind.
 
@@ -181,6 +210,7 @@ async def setup_check(ctx: Context) -> str:
     Verifies agent space and service role are available.
     If not ready, lists existing agent spaces to help with setup.
     """
+    _ensure_client_ua(ctx)
     try:
         config = _state.get_config()
         missing = []
@@ -263,6 +293,7 @@ async def setup(
     - Existing space + new role: setup(agent_space_id='as-xxx')
     - Existing space + existing role: setup(agent_space_id='as-xxx', service_role_arn='arn:...')
     """
+    _ensure_client_ua(ctx)
     try:
         identity = _client.get_caller_identity()
         account_id = identity['Account']
@@ -339,6 +370,7 @@ async def start_security_scan(
     Returns scan_id for polling with get_scan_status. The scan runs server-side.
     Use get_scan_status to check progress and get_scan_findings to retrieve results when complete.
     """
+    _ensure_client_ua(ctx)
     try:
         config = _state.get_config()
         if not config.get('agent_space_id') or not config.get('service_role'):
@@ -393,6 +425,7 @@ async def start_diff_scan(
     the diff patch; the agent focuses on changes while having full source for context.
     No prior scan required.
     """
+    _ensure_client_ua(ctx)
     try:
         config = _state.get_config()
         if not config.get('agent_space_id') or not config.get('service_role'):
@@ -447,6 +480,7 @@ async def start_threat_model_review(
     a threat model job. Returns a scan_id for polling with get_scan_status; retrieve
     identified threats with get_scan_findings. No prior scan required.
     """
+    _ensure_client_ua(ctx)
     try:
         config = _state.get_config()
         if not config.get('agent_space_id') or not config.get('service_role'):
@@ -492,6 +526,7 @@ async def get_scan_status(
     Useful for checking a previous scan from an earlier session, or verifying
     a scan completed after session recovery.
     """
+    _ensure_client_ua(ctx)
     try:
         return json.dumps(await _scanner.get_status(scan_id=scan_id), default=_json_serial)
     except ClientError as e:
@@ -521,6 +556,7 @@ async def get_scan_findings(
 
     Returns findings with title, severity, confidence, file location, and description.
     """
+    _ensure_client_ua(ctx)
     try:
         return json.dumps(
             await _scanner.get_findings(scan_id=scan_id, severity=severity), default=_json_serial
@@ -539,6 +575,7 @@ async def get_scan_findings(
 @mcp.tool()
 async def list_scans(ctx: Context) -> str:
     """List all recent security scans tracked locally with their status."""
+    _ensure_client_ua(ctx)
     try:
         return json.dumps({'scans': _state.list_scans()}, default=_json_serial)
     except ClientError as e:
@@ -558,6 +595,7 @@ async def stop_scan(
     scan_id: str = Field(..., description='The scan ID to stop.'),
 ) -> str:
     """Stop a running security scan."""
+    _ensure_client_ua(ctx)
     try:
         logger.info(f'Stopping scan: {scan_id}')
         return json.dumps(await _scanner.stop_scan(scan_id=scan_id), default=_json_serial)
@@ -588,6 +626,7 @@ async def call_api(
 
     Use get_api_guide to discover available operations and their parameters.
     """
+    _ensure_client_ua(ctx)
     try:
         import re
 
@@ -619,6 +658,7 @@ async def get_api_guide(ctx: Context) -> str:
     Returns operation names dynamically from the service model,
     plus a link to full API documentation with parameter details.
     """
+    _ensure_client_ua(ctx)
     global _cached_operations
     if _cached_operations is None:
         try:

src/security-agent-mcp-server/tests/test_aws_client.py

@@ -457,3 +457,83 @@ def test_batch_get_threats(self, mock_boto3):
         mock_client.batch_get_threats.assert_called_once_with(
             agentSpaceId='as-1', threatIds=['t-1']
         )
+
+
+class TestUserAgentInjection:
+    """Tests for custom User-Agent injection."""
+
+    def test_default_user_agent_has_mcp_server_identifier(self):
+        """SecurityAgentClient includes MCP server identifier in config by default."""
+        client = SecurityAgentClient(region='us-east-1')
+        assert 'awslabs-security-agent-mcp-server' in client._config.user_agent_extra  # type: ignore[attr-defined]
+        assert 'md/client#unknown' in client._config.user_agent_extra  # type: ignore[attr-defined]
+
+    def test_constructor_accepts_mcp_client_info(self):
+        """SecurityAgentClient accepts mcp_client_name and version at construction."""
+        client = SecurityAgentClient(
+            region='us-east-1', mcp_client_name='kiro', mcp_client_version='1.5.0'
+        )
+        assert 'md/client#kiro/1.5.0' in client._config.user_agent_extra  # type: ignore[attr-defined]
+
+    def test_set_mcp_client_info_updates_user_agent(self):
+        """set_mcp_client_info updates the user agent with client name and version."""
+        client = SecurityAgentClient(region='us-east-1')
+        client.set_mcp_client_info('kiro', '1.5.0')
+        assert 'md/client#kiro/1.5.0' in client._config.user_agent_extra  # type: ignore[attr-defined]
+
+    def test_set_mcp_client_info_without_version(self):
+        """set_mcp_client_info works with just client name (no version)."""
+        client = SecurityAgentClient(region='us-east-1')
+        client.set_mcp_client_info('cursor', '')
+        assert 'md/client#cursor' in client._config.user_agent_extra  # type: ignore[attr-defined]
+        assert 'md/client#cursor/' not in client._config.user_agent_extra  # type: ignore[attr-defined]
+
+    def test_set_mcp_client_info_noop_when_same(self):
+        """set_mcp_client_info does not rebuild config if info is unchanged."""
+        client = SecurityAgentClient(region='us-east-1')
+        client.set_mcp_client_info('kiro', '1.5.0')
+        config_after_first = client._config
+        client.set_mcp_client_info('kiro', '1.5.0')
+        assert client._config is config_after_first  # Same object, not rebuilt
+
+    def test_set_mcp_client_info_rebuilds_when_different(self):
+        """set_mcp_client_info rebuilds config when info changes."""
+        client = SecurityAgentClient(region='us-east-1')
+        client.set_mcp_client_info('kiro', '1.0')
+        config_first = client._config
+        client.set_mcp_client_info('claude-code', '2.0')
+        assert client._config is not config_first
+        assert 'md/client#claude-code/2.0' in client._config.user_agent_extra  # type: ignore[attr-defined]
+
+    @patch('awslabs.security_agent_mcp_server.aws_client.boto3')
+    def test_client_passes_config_to_securityagent(self, mock_boto3):
+        """_client() passes the botocore config to boto3 securityagent client."""
+        mock_session = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = SecurityAgentClient(region='us-east-1')
+        client._client()
+        mock_session.client.assert_called_once_with('securityagent', config=client._config)
+
+    @patch('awslabs.security_agent_mcp_server.aws_client.boto3')
+    def test_get_caller_identity_passes_config(self, mock_boto3):
+        """get_caller_identity passes the botocore config to sts client."""
+        mock_session = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = SecurityAgentClient(region='us-east-1')
+        client.get_caller_identity()
+        mock_session.client.assert_called_once_with('sts', config=client._config)
+
+    @patch('awslabs.security_agent_mcp_server.aws_client.boto3')
+    def test_upload_to_s3_passes_config(self, mock_boto3):
+        """upload_to_s3 passes the botocore config to s3 client."""
+        mock_session = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = SecurityAgentClient(region='us-east-1')
+        client.upload_to_s3('bucket', 'key', '/path/to/file')
+        mock_session.client.assert_called_once_with('s3', config=client._config)
+
+    def test_build_config_handles_missing_package_metadata(self):
+        """_build_config falls back to unknown when package metadata unavailable."""
+        with patch('importlib.metadata.version', side_effect=Exception('not found')):
+            client = SecurityAgentClient(region='us-east-1')
+            assert 'awslabs-security-agent-mcp-server/unknown' in client._config.user_agent_extra  # type: ignore[attr-defined]