feat(vcr): add support for proxying aws bedrock-runtime requests (#241)

* merge

* more formatting

* modified solution

* comment

* fmt

* just use an open source library instead

* readme

* release note

* update comment

* try fixing mypy issues

Author: sabrenner

README.md

@@ -146,6 +146,13 @@ The cassettes are matched based on the path, method, and body of the request. To
 
 Optionally specifying whatever mounted path is used for the cassettes directory. The test agent comes with a default set of cassettes for OpenAI, Azure OpenAI, and DeepSeek.
 
+#### AWS Services
+AWS service proxying, specifically recording cassettes for the first time, requires a `AWS_SECRET_ACCESS_KEY` environment variable to be set for the container running the test agent. This is used to recalculate the AWS signature for the request, as the one generated client-side likely used `{test-agent-host}:{test-agent-port}/vcr/{aws-service}` as the host, and the signature will mismatch that on the actual AWS service.
+
+Additionally, the `AWS_REGION` environment variable can be set, defaulting to `us-east-1`.
+
+To add a new AWS service to proxy, add an entry in the `PROVIDER_BASE_URLS` for its provider url, and an entry in the `AWS_SERVICES` dictionary for the service name, since they are not always a one-to-one mapping with the implied provider url (e.g, `https://bedrock-runtime.{AWS_REGION}.amazonaws.com` is the provider url, but the service name is `bedrock`, as `bedrock` also has multiple sub services, like `converse`).
+
 #### Usage in clients
 
 To use this feature in your client, you can use the `/vcr/{provider}` endpoint to proxy requests to the provider API.

ddapm_test_agent/vcr_proxy.py

@@ -1,30 +1,64 @@
 import hashlib
 import json
+import logging
 import os
 import re
+from typing import Any
+from typing import Dict
 from typing import Optional
 from urllib.parse import urljoin
 
 from aiohttp.web import Request
 from aiohttp.web import Response
 import requests
+from requests_aws4auth import AWS4Auth
 import vcr
 
 
+logger = logging.getLogger(__name__)
+
+
+#  Used for AWS signature recalculation for aws services initial proxying
+AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")
+AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
+
+
 def url_path_join(base_url: str, path: str) -> str:
     """Join a base URL with a path, handling slashes automatically."""
     return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
 
 
+AWS_SERVICES = {
+    "bedrock-runtime": "bedrock",
+}
+
+
 PROVIDER_BASE_URLS = {
     "openai": "https://api.openai.com/v1",
     "azure-openai": "https://dd.openai.azure.com/",
     "deepseek": "https://api.deepseek.com/",
     "anthropic": "https://api.anthropic.com/",
     "datadog": "https://api.datadoghq.com/",
     "genai": "https://generativelanguage.googleapis.com/",
+    "bedrock-runtime": f"https://bedrock-runtime.{AWS_REGION}.amazonaws.com",
 }
 
+CASSETTE_FILTER_HEADERS = [
+    "authorization",
+    "OpenAI-Organization",
+    "api-key",
+    "x-api-key",
+    "dd-api-key",
+    "dd-application-key",
+    "x-goog-api-key",
+    "x-amz-security-token",
+    "x-amz-content-sha256",
+    "x-amz-date",
+    "x-amz-user-agent",
+    "amz-sdk-invocation-id",
+    "amz-sdk-request",
+]
+
 NORMALIZERS = [
     (
         r"--form-data-boundary-[^\r\n]+",
@@ -65,22 +99,29 @@ def normalize_multipart_body(body: bytes) -> str:
             return f"[binary_data_{hex_digest}]"
 
 
+def parse_authorization_header(auth_header: str) -> Dict[str, str]:
+    """Parse AWS Authorization header to extract components"""
+    if not auth_header.startswith("AWS4-HMAC-SHA256 "):
+        return {}
+
+    auth_parts = auth_header[len("AWS4-HMAC-SHA256 ") :].split(",")
+    parsed = {}
+
+    for part in auth_parts:
+        key, value = part.split("=", 1)
+        parsed[key.strip()] = value.strip()
+
+    return parsed
+
+
 def get_vcr(subdirectory: str, vcr_cassettes_directory: str) -> vcr.VCR:
     cassette_dir = os.path.join(vcr_cassettes_directory, subdirectory)
 
     return vcr.VCR(
         cassette_library_dir=cassette_dir,
         record_mode="once",
         match_on=["path", "method"],
-        filter_headers=[
-            "authorization",
-            "OpenAI-Organization",
-            "api-key",
-            "x-api-key",
-            "dd-api-key",
-            "dd-application-key",
-            "x-goog-api-key",
-        ],
+        filter_headers=CASSETTE_FILTER_HEADERS,
     )
 
 
@@ -125,31 +166,47 @@ async def proxy_request(request: Request, vcr_cassettes_directory: str) -> Respo
     body_bytes = await request.read()
 
     vcr_cassette_prefix = request.pop("vcr_cassette_prefix", None)
-
     cassette_name = generate_cassette_name(path, request.method, body_bytes, vcr_cassette_prefix)
+
+    request_kwargs: Dict[str, Any] = {
+        "method": request.method,
+        "url": target_url,
+        "headers": headers,
+        "data": body_bytes,
+        "cookies": dict(request.cookies),
+        "allow_redirects": False,
+        "stream": True,
+    }
+
+    if provider in AWS_SERVICES and not os.path.exists(os.path.join(vcr_cassettes_directory, provider, cassette_name)):
+        if not AWS_SECRET_ACCESS_KEY:
+            return Response(
+                body="AWS_SECRET_ACCESS_KEY environment variable not set for aws signature recalculation",
+                status=400,
+            )
+
+        auth_header = request.headers.get("Authorization", "")
+        auth_parts = parse_authorization_header(auth_header)
+        aws_access_key = auth_parts.get("Credential", "").split("/")[0]
+
+        auth = AWS4Auth(aws_access_key, AWS_SECRET_ACCESS_KEY, AWS_REGION, AWS_SERVICES[provider])
+        request_kwargs["auth"] = auth
+
     with get_vcr(provider, vcr_cassettes_directory).use_cassette(f"{cassette_name}.yaml"):
-        oai_response = requests.request(
-            method=request.method,
-            url=target_url,
-            headers=headers,
-            data=body_bytes,
-            cookies=dict(request.cookies),
-            allow_redirects=False,
-            stream=True,
-        )
+        provider_response = requests.request(**request_kwargs)
 
     # Extract content type without charset
-    content_type = oai_response.headers.get("content-type", "")
+    content_type = provider_response.headers.get("content-type", "")
     if ";" in content_type:
         content_type = content_type.split(";")[0].strip()
 
     response = Response(
-        body=oai_response.content,
-        status=oai_response.status_code,
+        body=provider_response.content,
+        status=provider_response.status_code,
         content_type=content_type,
     )
 
-    for key, value in oai_response.headers.items():
+    for key, value in provider_response.headers.items():
         if key.lower() not in (
             "content-length",
             "transfer-encoding",

flake.nix

@@ -58,6 +58,7 @@
               requests
               yarl
               vcrpy
+              requests-aws4auth
               protobuf
               opentelemetry-proto
               grpcio

releasenotes/notes/vcr-aws-bedrock-proxy-3bf018b1712d5105.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    vcr: adds support for proxying aws bedrock runtime. to record cassettes for the first time, the `AWS_SECRET_ACCESS_KEY` environment variable must be set for the container running the test agent, for request signature recalculation. Additionally, `AWS_REGION` can be set, defaulting to `us-east-1`.

setup.py

@@ -37,6 +37,7 @@
         "typing_extensions",
         "yarl",
         "vcrpy",
+        "requests-aws4auth",
         # ddtrace libraries officially support opentelemetry-proto 1.33.1
         # which implements the v1.7.0 spec
         "opentelemetry-proto>1.33.0,<1.37.0",