feat: Add botocore instrumentation extension for Bedrock AgentCore services

CHANGELOG.md

@@ -11,3 +11,5 @@ For any change that affects end users of this package, please add an entry under
 If your change does not need a CHANGELOG entry, add the "skip changelog" label to your PR.
 
 ## Unreleased
+- Add botocore instrumentation extension for Bedrock AgentCore services with span attributes
+  ([#490](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/490))

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -20,6 +20,8 @@
 # TODO：Move to Semantic Conventions when these attributes are added.
 AWS_AUTH_ACCESS_KEY: str = "aws.auth.account.access_key"
 AWS_AUTH_REGION: str = "aws.auth.region"
+AWS_AUTH_CREDENTIAL_PROVIDER_ARN: str = "aws.auth.credential_provider.arn"
+AWS_GATEWAY_TARGET_ID = "aws.gateway.target.id"
 AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url"
 AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name"
 AWS_KINESIS_STREAM_ARN: str = "aws.kinesis.stream.arn"
@@ -41,3 +43,10 @@
 AWS_REMOTE_RESOURCE_ACCOUNT_ID: str = "aws.remote.resource.account.id"
 AWS_REMOTE_RESOURCE_REGION: str = "aws.remote.resource.region"
 AWS_SERVICE_TYPE: str = "aws.service.type"
+AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: str = "aws.bedrock.agentcore.runtime.arn"
+AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "aws.bedrock.agentcore.runtime_endpoint.arn"
+AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN: str = "aws.bedrock.agentcore.code_interpreter.arn"
+AWS_BEDROCK_AGENTCORE_BROWSER_ARN: str = "aws.bedrock.agentcore.browser.arn"
+AWS_BEDROCK_AGENTCORE_MEMORY_ARN: str = "aws.bedrock.agentcore.memory.arn"
+AWS_BEDROCK_AGENTCORE_GATEWAY_ARN: str = "aws.bedrock.agentcore.gateway.arn"
+AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "aws.bedrock.agentcore.identity.workload_identity.arn"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_agentcore_patches.py

@@ -0,0 +1,92 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Dict
+
+from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_AUTH_CREDENTIAL_PROVIDER_ARN,
+    AWS_BEDROCK_AGENTCORE_BROWSER_ARN,
+    AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN,
+    AWS_BEDROCK_AGENTCORE_GATEWAY_ARN,
+    AWS_BEDROCK_AGENTCORE_MEMORY_ARN,
+    AWS_BEDROCK_AGENTCORE_RUNTIME_ARN,
+    AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN,
+    AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN,
+    AWS_GATEWAY_TARGET_ID,
+)
+from opentelemetry.instrumentation.botocore.extensions.types import (
+    _AttributeMapT,
+    _AwsSdkExtension,
+    _BotocoreInstrumentorContext,
+    _BotoResultT,
+)
+from opentelemetry.trace.span import Span
+
+GEN_AI_RUNTIME_ID = "gen_ai.runtime.id"
+GEN_AI_BROWSER_ID = "gen_ai.browser.id"
+GEN_AI_CODE_INTERPRETER_ID = "gen_ai.code_interpreter.id"
+GEN_AI_MEMORY_ID = "gen_ai.memory.id"
+GEN_AI_GATEWAY_ID = "gen_ai.gateway.id"
+
+# Mapping of flattened JSON paths to attribute keys
+_ATTRIBUTE_MAPPING = {
+    "agentRuntimeArn": AWS_BEDROCK_AGENTCORE_RUNTIME_ARN,
+    "agentRuntimeEndpointArn": AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN,
+    "agentRuntimeId": GEN_AI_RUNTIME_ID,
+    "browserArn": AWS_BEDROCK_AGENTCORE_BROWSER_ARN,
+    "browserId": GEN_AI_BROWSER_ID,
+    "browserIdentifier": GEN_AI_BROWSER_ID,
+    "codeInterpreterArn": AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN,
+    "codeInterpreterId": GEN_AI_CODE_INTERPRETER_ID,
+    "codeInterpreterIdentifier": GEN_AI_CODE_INTERPRETER_ID,
+    "gatewayArn": AWS_BEDROCK_AGENTCORE_GATEWAY_ARN,
+    "gatewayId": GEN_AI_GATEWAY_ID,
+    "gatewayIdentifier": GEN_AI_GATEWAY_ID,
+    "targetId": AWS_GATEWAY_TARGET_ID,
+    "memory.arn": AWS_BEDROCK_AGENTCORE_MEMORY_ARN,
+    "memory.id": GEN_AI_MEMORY_ID,
+    "memoryId": GEN_AI_MEMORY_ID,
+    "credentialProviderArn": AWS_AUTH_CREDENTIAL_PROVIDER_ARN,
+    "workloadIdentityArn": AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN,
+    "workloadIdentityDetails.workloadIdentityArn": AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN,
+}
+
+
+class _BedrockAgentCoreExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        extracted_attrs = self._extract_attributes(self._call_context.params)
+        attributes.update(extracted_attrs)
+
+    def on_success(
+        self,
+        span: Span,
+        result: _BotoResultT,
+        instrumentor_context: _BotocoreInstrumentorContext,
+    ):
+        if span is None or not span.is_recording():
+            return
+
+        extracted_attrs = self._extract_attributes(result)
+        for attr_name, attr_value in extracted_attrs.items():
+            span.set_attribute(attr_name, attr_value)
+
+    @staticmethod
+    def _extract_attributes(params: Dict[str, Any]):
+        """Extracts all Bedrock AgentCore attributes using mapping-based traversal"""
+        attrs = {}
+        for path, attr_key in _ATTRIBUTE_MAPPING.items():
+            value = _BedrockAgentCoreExtension._get_nested_value(params, path)
+            if value:
+                attrs[attr_key] = value
+        return attrs
+
+    @staticmethod
+    def _get_nested_value(data: Dict[str, Any], path: str):
+        """Get value from nested dictionary using dot notation path"""
+        keys = path.split(".")
+        value = data
+        for key in keys:
+            if isinstance(value, dict) and key in value:
+                value = value[key]
+            else:
+                return None
+        return value

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -23,6 +23,9 @@
     AWS_STEPFUNCTIONS_ACTIVITY_ARN,
     AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
+from amazon.opentelemetry.distro.patches._bedrock_agentcore_patches import (  # noqa # pylint: disable=unused-import
+    _BedrockAgentCoreExtension,
+)
 from amazon.opentelemetry.distro.patches._bedrock_patches import (  # noqa # pylint: disable=unused-import
     _BedrockAgentExtension,
     _BedrockAgentRuntimeExtension,
@@ -247,6 +250,10 @@ def _apply_botocore_bedrock_patch() -> None:  # pylint: disable=too-many-stateme
     _KNOWN_EXTENSIONS["bedrock"] = _lazy_load(".", "_BedrockExtension")
     _KNOWN_EXTENSIONS["bedrock-agent"] = _lazy_load(".", "_BedrockAgentExtension")
     _KNOWN_EXTENSIONS["bedrock-agent-runtime"] = _lazy_load(".", "_BedrockAgentRuntimeExtension")
+    _KNOWN_EXTENSIONS["bedrock-agentcore"] = _lazy_load(".._bedrock_agentcore_patches", "_BedrockAgentCoreExtension")
+    _KNOWN_EXTENSIONS["bedrock-agentcore-control"] = _lazy_load(
+        ".._bedrock_agentcore_patches", "_BedrockAgentCoreExtension"
+    )
 
     # TODO: The following code is to patch bedrock-runtime bugs that are fixed in
     # opentelemetry-instrumentation-botocore==0.56b0 in these PRs:

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py → aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_instrumentation_patch.py

@@ -10,6 +10,24 @@
 
 import opentelemetry.sdk.extension.aws.resource.ec2 as ec2_resource
 import opentelemetry.sdk.extension.aws.resource.eks as eks_resource
+from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_AUTH_CREDENTIAL_PROVIDER_ARN,
+    AWS_BEDROCK_AGENTCORE_BROWSER_ARN,
+    AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN,
+    AWS_BEDROCK_AGENTCORE_GATEWAY_ARN,
+    AWS_BEDROCK_AGENTCORE_MEMORY_ARN,
+    AWS_BEDROCK_AGENTCORE_RUNTIME_ARN,
+    AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN,
+    AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN,
+    AWS_GATEWAY_TARGET_ID,
+)
+from amazon.opentelemetry.distro.patches._bedrock_agentcore_patches import (
+    GEN_AI_BROWSER_ID,
+    GEN_AI_CODE_INTERPRETER_ID,
+    GEN_AI_GATEWAY_ID,
+    GEN_AI_MEMORY_ID,
+    GEN_AI_RUNTIME_ID,
+)
 from amazon.opentelemetry.distro.patches._instrumentation_patch import (
     AWS_GEVENT_PATCH_MODULES,
     apply_instrumentation_patches,
@@ -38,6 +56,24 @@
 _LAMBDA_FUNCTION_NAME: str = "lambdaFunctionName"
 _LAMBDA_SOURCE_MAPPING_ID: str = "lambdaEventSourceMappingID"
 _TABLE_ARN: str = "arn:aws:dynamodb:us-west-2:123456789012:table/testTable"
+_AGENTCORE_RUNTIME_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-123"
+_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime-endpoint/test-endpoint"
+_AGENTCORE_RUNTIME_ID: str = "test-runtime-123"
+_AGENTCORE_BROWSER_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:browser/testBrowser-1234567890"
+_AGENTCORE_BROWSER_ID: str = "testBrowser-1234567890"
+_AGENTCORE_CODE_INTERPRETER_ARN: str = (
+    "arn:aws:bedrock-agentcore:us-east-1:123456789012:code-interpreter/testCodeInt-1234567890"
+)
+_AGENTCORE_CODE_INTERPRETER_ID: str = "testCodeInt-1234567890"
+_AGENTCORE_GATEWAY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/agentGateway-123456789"
+_AGENTCORE_GATEWAY_ID: str = "agentGateway-123456789"
+_AGENTCORE_TARGET_ID: str = "target-123456789"
+_AGENTCORE_MEMORY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:memory/agentMemory-123456789"
+_AGENTCORE_MEMORY_ID: str = "agentMemory-123456789"
+_AGENTCORE_CREDENTIAL_PROVIDER_ARN: str = (
+    "arn:aws:acps:us-east-1:123456789012:token-vault/test-vault/apikeycredentialprovider/test-provider"
+)
+_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:workload-identity/test-wi"
 
 # Patch names
 IMPORTLIB_METADATA_VERSION_PATCH: str = "amazon.opentelemetry.distro._utils.version"
@@ -160,6 +196,14 @@ def _test_unpatched_botocore_instrumentation(self):
             "bedrock-agent-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock Agent Runtime extension"
         )
 
+        # Bedrock AgentCore
+        self.assertFalse("bedrock-agentcore" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore extension")
+
+        # Bedrock AgentCore Control
+        self.assertFalse(
+            "bedrock-agentcore-control" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore Control extension"
+        )
+
         # BedrockRuntime
         self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension")
 
@@ -237,6 +281,41 @@ def _test_patched_botocore_instrumentation(self):
         bedrock_agent_runtime_sucess_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent-runtime")
         self.assertEqual(len(bedrock_agent_runtime_sucess_attributes), 0)
 
+        # Bedrock AgentCore
+        self.assertTrue("bedrock-agentcore" in _KNOWN_EXTENSIONS)
+        bedrock_agentcore_attributes: Dict[str, str] = _do_extract_bedrock_agentcore_attributes()
+        # Runtime attributes
+        self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ARN], _AGENTCORE_RUNTIME_ARN)
+        self.assertEqual(
+            bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN], _AGENTCORE_RUNTIME_ENDPOINT_ARN
+        )
+        self.assertEqual(bedrock_agentcore_attributes[GEN_AI_RUNTIME_ID], _AGENTCORE_RUNTIME_ID)
+        # Browser attributes
+        self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_BROWSER_ARN], _AGENTCORE_BROWSER_ARN)
+        self.assertEqual(bedrock_agentcore_attributes[GEN_AI_BROWSER_ID], _AGENTCORE_BROWSER_ID)
+        # Code interpreter attributes
+        self.assertEqual(
+            bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN], _AGENTCORE_CODE_INTERPRETER_ARN
+        )
+        self.assertEqual(bedrock_agentcore_attributes[GEN_AI_CODE_INTERPRETER_ID], _AGENTCORE_CODE_INTERPRETER_ID)
+        # Gateway attributes
+        self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_GATEWAY_ARN], _AGENTCORE_GATEWAY_ARN)
+        self.assertEqual(bedrock_agentcore_attributes[GEN_AI_GATEWAY_ID], _AGENTCORE_GATEWAY_ID)
+        self.assertEqual(bedrock_agentcore_attributes[AWS_GATEWAY_TARGET_ID], _AGENTCORE_TARGET_ID)
+        # Memory attributes
+        self.assertEqual(bedrock_agentcore_attributes[GEN_AI_MEMORY_ID], _AGENTCORE_MEMORY_ID)
+        self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_MEMORY_ARN], _AGENTCORE_MEMORY_ARN)
+        # Auth and identity attributes
+        self.assertEqual(
+            bedrock_agentcore_attributes[AWS_AUTH_CREDENTIAL_PROVIDER_ARN], _AGENTCORE_CREDENTIAL_PROVIDER_ARN
+        )
+        self.assertEqual(
+            bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN], _AGENTCORE_WORKLOAD_IDENTITY_ARN
+        )
+
+        # Bedrock AgentCore Control
+        self.assertTrue("bedrock-agentcore-control" in _KNOWN_EXTENSIONS)
+
         # BedrockRuntime
         self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
 
@@ -854,6 +933,31 @@ def _do_extract_lambda_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
+def _do_extract_bedrock_agentcore_attributes() -> Dict[str, str]:
+    service_name: str = "bedrock-agentcore"
+    params: Dict[str, Any] = {
+        "agentRuntimeArn": _AGENTCORE_RUNTIME_ARN,
+        "agentRuntimeEndpointArn": _AGENTCORE_RUNTIME_ENDPOINT_ARN,
+        "agentRuntimeId": _AGENTCORE_RUNTIME_ID,
+        "browserArn": _AGENTCORE_BROWSER_ARN,
+        "browserId": _AGENTCORE_BROWSER_ID,
+        "browserIdentifier": _AGENTCORE_BROWSER_ID,
+        "codeInterpreterArn": _AGENTCORE_CODE_INTERPRETER_ARN,
+        "codeInterpreterId": _AGENTCORE_CODE_INTERPRETER_ID,
+        "codeInterpreterIdentifier": _AGENTCORE_CODE_INTERPRETER_ID,
+        "gatewayArn": _AGENTCORE_GATEWAY_ARN,
+        "gatewayId": _AGENTCORE_GATEWAY_ID,
+        "gatewayIdentifier": _AGENTCORE_GATEWAY_ID,
+        "targetId": _AGENTCORE_TARGET_ID,
+        "memoryId": _AGENTCORE_MEMORY_ID,
+        "credentialProviderArn": _AGENTCORE_CREDENTIAL_PROVIDER_ARN,
+        "workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN,
+        "memory": {"arn": _AGENTCORE_MEMORY_ARN, "id": _AGENTCORE_MEMORY_ID},
+        "workloadIdentityDetails": {"workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN},
+    }
+    return _do_extract_attributes(service_name, params)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params