diff --git a/.env.example b/.env.example index 9bfc77d9..473e9925 100644 --- a/.env.example +++ b/.env.example @@ -43,6 +43,28 @@ LLM_MAX_TOKENS=4096 # MANAGER_LLM_MODEL=bedrock/us.anthropic.claude-opus-4-8 # Note: Opus 4.7+ Bedrock IDs are undated (us.anthropic.claude-opus-4-8); older models keep the dated -v1:0 profile form. See https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-4-8.html +# Provider: AWS Bedrock via the Anthropic Messages endpoint (RECOMMENDED for +# Bedrock). Routes Claude through CrewAI's native Anthropic provider against +# Bedrock's /anthropic route, so the Converse toolUse/toolResult sanitizer +# (patches/crewai_bedrock_fix.py) is NOT needed. Only current-generation Claude +# tiers are served on the Messages API (Sonnet 5, Opus 4.7/4.8, Haiku 4.5); +# Sonnet 4.5 is Converse-only. Model ids are the Bedrock model / inference- +# profile ids (no provider prefix). +# ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL=https://bedrock-runtime.us-west-2.amazonaws.com/anthropic +# ANTHROPIC_COMPATIBLE_LLM_API_KEY=your-bedrock-api-key +# DEFAULT_LLM_MODEL=us.anthropic.claude-sonnet-5 +# MANAGER_LLM_MODEL=us.anthropic.claude-sonnet-5 +# MEMORY_LLM_MODEL=us.anthropic.claude-haiku-4-5-20251001-v1:0 +# +# AWS credentials for the deploy CLI (dev account). deploy.sh sources this .env, +# so either an AWS_PROFILE that is logged in, or explicit keys, work: +# AWS_PROFILE=genai +# -- or -- +# AWS_ACCESS_KEY_ID=... +# AWS_SECRET_ACCESS_KEY=... +# AWS_SESSION_TOKEN=... +# AWS_REGION=us-west-2 + # Provider: Custom OpenAI-compatible endpoint (alternative) # For endpoints with no CrewAI native provider prefix of their own — NVIDIA # NIM, Ollama, HuggingFace TGI, vLLM, LM Studio, etc. Set DEFAULT_LLM_MODEL to diff --git a/infra/aws/agentcore/auth-agentcore.yaml b/infra/aws/agentcore/auth-agentcore.yaml new file mode 100644 index 00000000..3601c9c0 --- /dev/null +++ b/infra/aws/agentcore/auth-agentcore.yaml @@ -0,0 +1,198 @@ +AWSTemplateFormatVersion: "2010-09-09" +Description: > + Auth stack for the AAMP buyer AgentCore deployment. Houses the IAM + execution role the buyer runtime assumes, including a SCOPED grant to + invoke seller AgentCore runtimes for buyer->seller (agent-to-agent) calls. + This "auth" stack is the home for auth-related infra (IAM today; Cognito + + gateway authorizer for the cross-org path live alongside it — see NOTE). + +# ============================================================================= +# NOTE — SECURITY MODEL / WHEN TO USE THIS +# ============================================================================= +# The SellerInvoke permission below lets the buyer runtime call the seller's +# AgentCore runtimes directly via bedrock-agentcore:InvokeAgentRuntime, gated +# by AWS IAM. That is appropriate ONLY for: +# +# * Same-account testing — buyer and seller deployed in one AWS account +# (the demo / integration-test setup), OR +# * Customers whose security standards explicitly ALLOW cross-account IAM +# access between the buyer's and seller's accounts (you would then extend +# the resource ARNs / add a cross-account role trust accordingly). +# +# It is NOT the recommended pattern for CROSS-ORGANIZATION buyer<->seller +# interactions (different companies, no shared IAM trust). For that, front the +# seller runtimes with an **AgentCore Gateway + CUSTOM_JWT authorizer backed by +# Cognito** (or the customer's OIDC IdP): the external buyer authenticates with +# a short-lived OAuth2 JWT bearer token at the gateway edge — no IAM in the +# seller's account required. See the `enterprise-auth-gateway` spec. +# +# In short: IAM (this stack) = same-account / trusted cross-account. +# Gateway + CUSTOM_JWT/Cognito = cross-org, zero shared IAM. +# ============================================================================= + +# ============================================================================= +# Parameters +# ============================================================================= +Parameters: + Environment: + Type: String + Default: staging + AllowedValues: [staging, production] + SellerRuntimeNamePrefix: + Type: String + Default: "a4a_aamp_seller_" + Description: > + Name prefix of the seller AgentCore runtimes this buyer may invoke. The + SellerInvoke policy wildcards on runtime/* so the buyer can reach + the seller's http/mcp/a2a runtimes without hardcoding the CLI-assigned + ARN suffixes (which differ per deploy). Keep this as tight as your naming + allows — it is scoped to this account + region + the seller name prefix, + NOT a blanket bedrock-agentcore:* on all runtimes. Match your seller + deploy --name base (e.g. a4a_aamp_seller_ for a4a_aamp_seller_*_mcp). + SellerAccountId: + Type: String + Default: "" + Description: > + Optional. AWS account id that hosts the seller runtimes. Defaults to THIS + account (same-account model). Set this ONLY for a trusted cross-account + setup your security standards permit; the SellerInvoke resource ARNs then + point at that account. (Cross-ORG should use the gateway path instead.) + +# ============================================================================= +# Conditions +# ============================================================================= +Conditions: + # Use this account for the seller ARNs unless a (trusted) cross-account id is given. + UseThisAccountForSeller: !Equals [!Ref SellerAccountId, ""] + +# ============================================================================= +# Resources +# ============================================================================= +Resources: + + # --------------------------------------------------------------------------- + # Buyer AgentCore runtime execution role + # + # Provisioned here (declarative, least-privilege) instead of letting the + # agentcore CLI auto-create a role, so the seller-invoke grant is owned by + # IaC. Pass this role to `agentcore configure --execution-role `. + # --------------------------------------------------------------------------- + BuyerRuntimeExecutionRole: + Type: AWS::IAM::Role + Properties: + RoleName: !Sub "aamp-buyer-${Environment}-runtime-role" + AssumeRolePolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Principal: + Service: bedrock-agentcore.amazonaws.com + Action: sts:AssumeRole + Policies: + # -- Standard AgentCore runtime execution permissions -- + # Mirrors the policy the `agentcore` CLI auto-creates for a runtime + # (ECR image pull, CloudWatch Logs + log-delivery, X-Ray tracing, + # CloudWatch metrics, and AgentCore memory/workload access), so a + # CFN-owned role is a drop-in for the CLI-created one. + - PolicyName: AgentCoreRuntimeExecution + PolicyDocument: + Version: "2012-10-17" + Statement: + - Sid: EcrImagePull + Effect: Allow + Action: + - ecr:BatchGetImage + - ecr:GetDownloadUrlForLayer + - ecr:BatchCheckLayerAvailability + Resource: !Sub "arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/bedrock-agentcore-*" + - Sid: EcrAuth + Effect: Allow + Action: ecr:GetAuthorizationToken + Resource: "*" + - Sid: Logs + Effect: Allow + Action: + - logs:CreateLogGroup + - logs:CreateLogStream + - logs:PutLogEvents + - logs:DescribeLogStreams + - logs:DescribeLogGroups + Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/bedrock-agentcore/runtimes/*" + - Sid: XRayTracing + Effect: Allow + Action: + - xray:PutTraceSegments + - xray:PutTelemetryRecords + - xray:GetSamplingRules + - xray:GetSamplingTargets + Resource: "*" + - Sid: CloudWatchMetrics + Effect: Allow + Action: cloudwatch:PutMetricData + Resource: "*" + - Sid: AgentCoreMemory + Effect: Allow + Action: + - bedrock-agentcore:CreateEvent + - bedrock-agentcore:GetEvent + - bedrock-agentcore:GetMemory + - bedrock-agentcore:GetMemoryRecord + - bedrock-agentcore:ListActors + - bedrock-agentcore:ListEvents + - bedrock-agentcore:ListMemoryRecords + - bedrock-agentcore:ListSessions + Resource: "*" + # -- Bedrock model access (Claude via the Anthropic Messages endpoint) -- + - PolicyName: BedrockInvoke + PolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Action: + - bedrock:InvokeModel + - bedrock:InvokeModelWithResponseStream + # Req 11: the runtime mints a short-lived Bedrock bearer token + # from THIS role at startup for the Anthropic Messages path + # (x-api-key), which is authenticated via CallWithBearerToken + # rather than SigV4 InvokeModel. + - bedrock:CallWithBearerToken + Resource: "*" + # -- SCOPED seller invoke (see NOTE at top of this file) -- + - PolicyName: SellerInvoke + PolicyDocument: + Version: "2012-10-17" + Statement: + - Sid: InvokeSellerRuntimes + Effect: Allow + Action: bedrock-agentcore:InvokeAgentRuntime + # Wildcard on the seller name prefix + its endpoint children, + # scoped to a single account + region. Not bedrock-agentcore:*. + Resource: + - !If + - UseThisAccountForSeller + - !Sub "arn:aws:bedrock-agentcore:${AWS::Region}:${AWS::AccountId}:runtime/${SellerRuntimeNamePrefix}*" + - !Sub "arn:aws:bedrock-agentcore:${AWS::Region}:${SellerAccountId}:runtime/${SellerRuntimeNamePrefix}*" + - !If + - UseThisAccountForSeller + - !Sub "arn:aws:bedrock-agentcore:${AWS::Region}:${AWS::AccountId}:runtime/${SellerRuntimeNamePrefix}*/*" + - !Sub "arn:aws:bedrock-agentcore:${AWS::Region}:${SellerAccountId}:runtime/${SellerRuntimeNamePrefix}*/*" + Tags: + - Key: Name + Value: !Sub "aamp-buyer-${Environment}-runtime-role" + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + +# ============================================================================= +# Outputs +# ============================================================================= +Outputs: + BuyerRuntimeExecutionRoleArn: + Description: > + Execution role ARN for the buyer AgentCore runtime. Pass to + `agentcore configure --execution-role ` so the runtime can invoke + the seller runtimes (same-account / trusted cross-account only). + Value: !GetAtt BuyerRuntimeExecutionRole.Arn + Export: + Name: !Sub "${AWS::StackName}-BuyerRuntimeExecutionRoleArn" diff --git a/infra/aws/agentcore/deploy.sh b/infra/aws/agentcore/deploy.sh index ce69ee4e..f3014208 100755 --- a/infra/aws/agentcore/deploy.sh +++ b/infra/aws/agentcore/deploy.sh @@ -25,10 +25,34 @@ set -euo pipefail +# Load a gitignored .env at repo root if present, so local AWS credentials +# (AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / AWS_SESSION_TOKEN or AWS_PROFILE) +# and the Bedrock endpoint config below can be supplied without exporting them +# by hand. .env is listed in .gitignore and must never be committed. +_SCRIPT_DIR_EARLY="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_REPO_ROOT_EARLY="$(cd "${_SCRIPT_DIR_EARLY}/../../.." && pwd)" +if [[ -f "${_REPO_ROOT_EARLY}/.env" ]]; then + echo ">>> Loading environment from ${_REPO_ROOT_EARLY}/.env" + set -a + # shellcheck disable=SC1091 + source "${_REPO_ROOT_EARLY}/.env" + set +a +fi + REGION="${AWS_REGION:-us-west-2}" AGENT_NAME="${AGENT_NAME:-}" AWS_PROFILE="${AWS_PROFILE:-}" -LLM_MODEL="${DEFAULT_LLM_MODEL:-bedrock/us.amazon.nova-pro-v1:0}" +# Default model: current-generation Claude on Bedrock via the Anthropic +# Messages endpoint (see ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL below). This +# replaces the retired Amazon Nova Pro default. The model id is the Bedrock +# model / inference-profile id passed to the Anthropic provider. +LLM_MODEL="${DEFAULT_LLM_MODEL:-us.anthropic.claude-sonnet-5}" +MEMORY_MODEL="${MEMORY_LLM_MODEL:-us.anthropic.claude-haiku-4-5-20251001-v1:0}" +# Bedrock's Anthropic-compatible (Messages API) base URL + API key. Setting +# these routes Claude through CrewAI's native Anthropic provider against +# Bedrock, so the Converse toolUse/toolResult sanitizer is not applied. +ANTHROPIC_BASE_URL="${ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL:-https://bedrock-runtime.${REGION}.amazonaws.com/anthropic}" +BEDROCK_API_KEY="${ANTHROPIC_COMPATIBLE_LLM_API_KEY:-${AWS_BEARER_TOKEN_BEDROCK:-}}" SELLER_AGENT_URL="${SELLER_AGENT_URL:-}" DEPLOY_MODE="http" DO_TEST=false @@ -143,6 +167,31 @@ if [[ "${TEST_ONLY}" == "false" ]]; then rm -f "${_GENERATED_DOCKERFILE}" fi + # ── Auth stack (opt-in): provision the buyer runtime execution role with a + # scoped grant to invoke the seller runtimes, and pass it to agentcore + # configure. Only when SELLER_AGENT_URL is an AgentCore ARN (buyer->seller + # over InvokeAgentRuntime). For a plain HTTP seller URL or no seller, the CLI + # auto-creates the default role as before. See auth-agentcore.yaml for the + # security model (same-account / trusted cross-account only; cross-org uses + # the gateway + CUSTOM_JWT path). + EXECUTION_ROLE_ARN="" + if [[ "${SELLER_AGENT_URL}" == arn:* ]]; then + AUTH_STACK="aamp-buyer-auth-${ENVIRONMENT:-staging}" + echo "" + echo ">>> Deploying auth stack (buyer runtime role + seller-invoke): ${AUTH_STACK}" + aws cloudformation deploy \ + --stack-name "${AUTH_STACK}" \ + --template-file infra/aws/agentcore/auth-agentcore.yaml \ + --capabilities CAPABILITY_NAMED_IAM \ + --parameter-overrides "Environment=${ENVIRONMENT:-staging}" \ + --region "${REGION}" ${AWS_PROFILE:+--profile "${AWS_PROFILE}"} + EXECUTION_ROLE_ARN=$(aws cloudformation describe-stacks \ + --stack-name "${AUTH_STACK}" --region "${REGION}" ${AWS_PROFILE:+--profile "${AWS_PROFILE}"} \ + --query "Stacks[0].Outputs[?OutputKey=='BuyerRuntimeExecutionRoleArn'].OutputValue" \ + --output text 2>/dev/null || true) + echo " execution role: ${EXECUTION_ROLE_ARN:-}" + fi + # Configure echo "" echo ">>> Configuring agent..." @@ -152,15 +201,39 @@ if [[ "${TEST_ONLY}" == "false" ]]; then -rf infra/aws/agentcore/requirements.txt \ -p HTTP \ -r "${REGION}" \ + ${EXECUTION_ROLE_ARN:+--execution-role "${EXECUTION_ROLE_ARN}"} \ --non-interactive \ --deployment-type container # Deploy echo "" echo ">>> Deploying to AgentCore..." + # Req 11: only bake a Bedrock key when one was explicitly supplied. Otherwise + # omit it so the runtime mints a fresh token from its execution role at startup + # (ad_buyer.llm.bedrock_token) — the CFN runtime role grants CallWithBearerToken. + _key_env=() + if [[ -n "${BEDROCK_API_KEY}" ]]; then + _key_env=(--env "ANTHROPIC_COMPATIBLE_LLM_API_KEY=${BEDROCK_API_KEY}") + fi + # AWS Bedrock AgentCore CUSTOM_JWT buyer->seller path: when the seller + # advertises an OAuth issuer, forward the seller's token endpoint + scope and + # the buyer's OWN client credentials so crew_tools routes through the JWT + # transport (see ad_buyer.registry.transport_selector). Each is forwarded only + # when set — never bake an empty value. The client_secret is passed only if + # present; prefer sourcing it from a secret store in production. + _oauth_env=() + [[ -n "${SELLER_MCP_RUNTIME_ARN:-}" ]] && _oauth_env+=(--env "SELLER_MCP_RUNTIME_ARN=${SELLER_MCP_RUNTIME_ARN}") + [[ -n "${SELLER_TOKEN_ENDPOINT:-}" ]] && _oauth_env+=(--env "SELLER_TOKEN_ENDPOINT=${SELLER_TOKEN_ENDPOINT}") + [[ -n "${SELLER_INVOKE_SCOPE:-}" ]] && _oauth_env+=(--env "SELLER_INVOKE_SCOPE=${SELLER_INVOKE_SCOPE}") + [[ -n "${BUYER_OAUTH_CLIENT_ID:-}" ]] && _oauth_env+=(--env "BUYER_OAUTH_CLIENT_ID=${BUYER_OAUTH_CLIENT_ID}") + [[ -n "${BUYER_OAUTH_CLIENT_SECRET:-}" ]] && _oauth_env+=(--env "BUYER_OAUTH_CLIENT_SECRET=${BUYER_OAUTH_CLIENT_SECRET}") agentcore deploy \ --env "DEFAULT_LLM_MODEL=${LLM_MODEL}" \ --env "MANAGER_LLM_MODEL=${LLM_MODEL}" \ + --env "PYTHONPATH=/app/src" \ + --env "ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL=${ANTHROPIC_BASE_URL}" \ + "${_key_env[@]}" \ + "${_oauth_env[@]}" \ --env "STORAGE_TYPE=sqlite" \ --env "DATABASE_URL=sqlite:///:memory:" \ --env "ANTHROPIC_API_KEY=not-used-with-bedrock" \ @@ -168,7 +241,7 @@ if [[ "${TEST_ONLY}" == "false" ]]; then --env "AWS_REGION=${REGION}" \ --env "AWS_DEFAULT_REGION=${REGION}" \ --env "CREW_MEMORY_ENABLED=true" \ - --env "MEMORY_LLM_MODEL=bedrock/us.amazon.nova-lite-v1:0" \ + --env "MEMORY_LLM_MODEL=${MEMORY_MODEL}" \ --auto-update-on-conflict echo "" diff --git a/infra/aws/agentcore/requirements.txt b/infra/aws/agentcore/requirements.txt index 439b6b8f..3cdf8a98 100644 --- a/infra/aws/agentcore/requirements.txt +++ b/infra/aws/agentcore/requirements.txt @@ -1,7 +1,22 @@ # Requirements for AgentCore CLI deployment path. -# Used by: agentcore configure --requirements-file requirements-agentcore.txt +# Used by: agentcore configure --requirements-file infra/aws/agentcore/requirements.txt +# +# Kept in sync with pyproject.toml [project.dependencies]. The AgentCore +# toolkit-generated Dockerfile installs THIS file (it does not `pip install` +# the ad_buyer package), so every runtime import dependency must be listed +# here — including the git-based shared contract library — or the container +# fails at import with ModuleNotFoundError. bedrock-agentcore -crewai[tools,anthropic,bedrock]==1.10.1 +# Shared IAB agentic contract catalog. The crew path, deals_client, and the +# AgentCoreSellerProxy import iab_agentic_primitives, so it must be in the +# AgentCore image. Git-free GitHub archive tarball (the CodeBuild image has no +# git). Keep in sync with the seller agent's pin. +iab-agentic-primitives[client] @ https://github.com/IABTechLab/iab-agentic-primitives/archive/refs/tags/v0.5.1.tar.gz +# Pin crewai to match pyproject.toml (>=1.14.4). The Bedrock fix patch +# (patches/crewai_bedrock_fix.py) targets 1.14.x and DealBookingFlow relies on +# CrewAI >=1.14 Flow/state behaviors, so an older pin (the previous ==1.10.1) +# would make the AgentCore image diverge from the ECS/source runtime. +crewai[tools,anthropic,bedrock]>=1.14.4,<2.0.0 anthropic>=0.40.0 httpx>=0.27.0 pydantic>=2.0.0 @@ -15,3 +30,11 @@ sqlalchemy>=2.0.0 aiosqlite>=0.20.0 pyyaml>=6.0.0 nest-asyncio>=1.6.0 +# Mints a short-lived Bedrock bearer token from the runtime execution role at +# startup (ad_buyer.llm.bedrock_token) so no expiring token is baked into --env. +aws-bedrock-token-generator>=1.0.0 +# NOTE: the AgentCore build image (astral uv slim) has no `git` binary, so a +# `git+https://` source fails with "Git executable not found". Use the GitHub +# archive tarball instead — a plain HTTPS download that needs no git. Keep this +# pinned to the same tag as pyproject.toml's git ref (v0.5.1). +iab-agentic-primitives[client] @ https://github.com/IABTechLab/iab-agentic-primitives/archive/refs/tags/v0.5.1.tar.gz diff --git a/patches/crewai_agentcore_memory.py b/patches/crewai_agentcore_memory.py index 320c6e03..abac8009 100644 --- a/patches/crewai_agentcore_memory.py +++ b/patches/crewai_agentcore_memory.py @@ -119,7 +119,9 @@ def _patched_memory_init(self, *args, **kwargs): if "llm" not in kwargs or kwargs.get("llm") == "gpt-4o-mini": bedrock_model = os.environ.get( "MEMORY_LLM_MODEL", - os.environ.get("DEFAULT_LLM_MODEL", "bedrock/us.amazon.nova-lite-v1:0"), + os.environ.get( + "DEFAULT_LLM_MODEL", "bedrock/us.anthropic.claude-haiku-4-5-v1:0" + ), ) kwargs["llm"] = bedrock_model try: diff --git a/patches/crewai_bedrock_anthropic_fix.py b/patches/crewai_bedrock_anthropic_fix.py new file mode 100644 index 00000000..203dc1c3 --- /dev/null +++ b/patches/crewai_bedrock_anthropic_fix.py @@ -0,0 +1,60 @@ +"""Patch for CrewAI AnthropicCompletion — Bedrock Anthropic Messages compat. + +When Claude runs on Amazon Bedrock's Anthropic-compatible (Messages) endpoint +(https://bedrock-runtime..amazonaws.com/anthropic), CrewAI's native +Anthropic provider builds tool definitions that include a ``strict`` field. + +CrewAI (>=1.15) unconditionally sets ``"strict": True`` on every tool's +function schema in ``crewai.utilities.agent_utils`` (build_tools), and +``AnthropicCompletion._convert_tools_for_interference`` copies it onto the +Anthropic tool object. Anthropic's own API accepts ``strict``; Amazon Bedrock's +Anthropic-compatible endpoint does NOT, and rejects the request with: + + 400 invalid_request_error: tools.0.custom.strict: Extra inputs are not permitted + +This patch wraps ``_convert_tools_for_interference`` to remove the ``strict`` +key from the produced tool objects, but ONLY when the completion is pointed at +a Bedrock base URL (so behavior against real Anthropic is unchanged). It is the +Messages-endpoint analog of the retired Converse ``crewai_bedrock_fix`` and is +applied only when ``ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL`` is set. + +Usage: + from patches.crewai_bedrock_anthropic_fix import apply_patches + apply_patches() # call once at startup, before building crews +""" + +import logging + +logger = logging.getLogger(__name__) + +_patched = False + + +def _is_bedrock_base_url(url: object) -> bool: + return isinstance(url, str) and ("bedrock-runtime" in url or "bedrock-mantle" in url) + + +def apply_patches() -> None: + """Strip the unsupported ``strict`` tool field on the Bedrock Messages endpoint.""" + global _patched + if _patched: + return + try: + from crewai.llms.providers.anthropic.completion import AnthropicCompletion + except ImportError: + logger.warning("AnthropicCompletion not available — skipping Bedrock Anthropic patch") + return + + original = AnthropicCompletion._convert_tools_for_interference + + def _patched_convert(self, tools): + result = original(self, tools) + if _is_bedrock_base_url(getattr(self, "base_url", None)) and isinstance(result, list): + for tool in result: + if isinstance(tool, dict): + tool.pop("strict", None) + return result + + AnthropicCompletion._convert_tools_for_interference = _patched_convert + _patched = True + logger.info("Patched AnthropicCompletion._convert_tools_for_interference (Bedrock strict-strip)") diff --git a/src/ad_buyer/auth/oauth_token_provider.py b/src/ad_buyer/auth/oauth_token_provider.py new file mode 100644 index 00000000..7808a843 --- /dev/null +++ b/src/ad_buyer/auth/oauth_token_provider.py @@ -0,0 +1,196 @@ +# Donated to IAB Tech Lab + +"""OAuth2 client_credentials token provider for authenticated seller runtimes. + +**Scope: the AWS Bedrock AgentCore deployment path.** This provider exists for +the case where a seller is deployed as a Bedrock AgentCore *runtime* fronted by +a ``CUSTOM_JWT`` authorizer, and advertises an OAuth2 token endpoint + scope in +its AAMP registry record's ``authentication`` object. That is a genuinely +CROSS-ORG, machine-to-machine setting — the buyer and seller are different +organizations with no shared API-key exchange — which is exactly what the OAuth2 +``client_credentials`` grant is the industry-standard answer for. + +It deliberately does NOT replace the existing per-seller API-key path +(:class:`ad_buyer.auth.middleware.AuthMiddleware` + ``ApiKeyStore``). That path +predates AgentCore and targets same-/trusted-account HTTP sellers where a +pre-shared API key is provisioned out of band; there is no issuer to obtain a +token from, so OAuth would have nothing to point at. The ``arn:`` AgentCore +proxy path authenticates with SigV4 (IAM), another standard, so it needs no +bearer token either. OAuth becomes the right tool specifically when a discovered +seller ADVERTISES an issuer — i.e. this AgentCore CUSTOM_JWT path — and the +transport selector routes only those ``https://`` endpoints through here. + +When the buyer discovers such a seller runtime (``auth_required=true`` with an +``authentication`` object: token endpoint + scope), it must present +``Authorization: Bearer `` to invoke it. This provider mints that JWT via +the OAuth2 ``client_credentials`` grant. + +Design (Req 6.6): +- The token ENDPOINT and SCOPE come from the DISCOVERED registry record's + ``authentication`` object — the seller advertises where to authenticate. +- The buyer's OWN ``client_id``/``client_secret`` come from settings/env + (``BUYER_OAUTH_CLIENT_ID`` / ``BUYER_OAUTH_CLIENT_SECRET``) — its identity, + NEVER read from the registry and NEVER logged. +- Tokens are cached keyed by ``(token_endpoint, scope)`` and refreshed + PROACTIVELY before ``expires_in`` (minus a safety skew) and REACTIVELY when a + caller reports a 401 (:meth:`invalidate`). +- The same provider serves a BYO-IdP seller: the discovered token endpoint is + simply the customer's IdP rather than the shared Cognito domain. +""" + +from __future__ import annotations + +import base64 +import logging +import threading +import time +from dataclasses import dataclass + +import httpx + +logger = logging.getLogger(__name__) + +# Refresh this many seconds BEFORE the token's stated expiry, to avoid handing +# out a token that expires mid-flight. +_EXPIRY_SKEW_SECONDS = 60.0 + + +class OAuthTokenError(RuntimeError): + """Raised when a client_credentials token cannot be minted.""" + + +@dataclass(frozen=True) +class _CacheKey: + token_endpoint: str + scope: str + + +@dataclass +class _CachedToken: + access_token: str + expires_at: float # monotonic deadline + + def valid(self, *, now: float) -> bool: + return now < (self.expires_at - _EXPIRY_SKEW_SECONDS) + + +class OAuthTokenProvider: + """Mints + caches OAuth2 client_credentials bearer tokens. + + Args: + client_id: The buyer's own OAuth app-client id. Defaults to + ``settings.buyer_oauth_client_id``. + client_secret: The buyer's own app-client secret. Defaults to + ``settings.buyer_oauth_client_secret``. Never logged. + timeout: HTTP timeout for the token endpoint (seconds). + transport: Optional httpx transport (tests inject a stub endpoint). + """ + + def __init__( + self, + client_id: str | None = None, + client_secret: str | None = None, + *, + timeout: float = 15.0, + transport: httpx.BaseTransport | None = None, + ): + if client_id is None or client_secret is None: + # Lazy import so importing this module never triggers settings load. + from ad_buyer.config.settings import get_settings + + settings = get_settings() + client_id = client_id or settings.buyer_oauth_client_id + client_secret = client_secret or settings.buyer_oauth_client_secret + + self._client_id = client_id + self._client_secret = client_secret + self._timeout = timeout + self._transport = transport + self._cache: dict[_CacheKey, _CachedToken] = {} + self._lock = threading.Lock() + + @property + def configured(self) -> bool: + """True when the buyer has its own client credentials to mint with.""" + return bool(self._client_id and self._client_secret) + + def get_token(self, token_endpoint: str, scope: str = "", *, force: bool = False) -> str: + """Return a valid bearer token for (token_endpoint, scope). + + Reuses a cached token within its lifetime; mints a fresh one otherwise + (or when ``force`` is set). Raises :class:`OAuthTokenError` on failure. + """ + if not self.configured: + raise OAuthTokenError( + "buyer OAuth client credentials are not configured " + "(set BUYER_OAUTH_CLIENT_ID / BUYER_OAUTH_CLIENT_SECRET)" + ) + if not token_endpoint: + raise OAuthTokenError("no token_endpoint provided by the discovered record") + + key = _CacheKey(token_endpoint=token_endpoint, scope=scope or "") + now = time.monotonic() + + with self._lock: + cached = self._cache.get(key) + if cached is not None and not force and cached.valid(now=now): + return cached.access_token + + # Mint outside the lock (network I/O); re-check under the lock on store. + token, expires_in = self._mint(token_endpoint, scope) + with self._lock: + self._cache[key] = _CachedToken( + access_token=token, + expires_at=time.monotonic() + max(float(expires_in), 0.0), + ) + return token + + def invalidate(self, token_endpoint: str, scope: str = "") -> None: + """Drop the cached token so the next :meth:`get_token` re-mints. + + Call this reactively when a downstream call returns 401 with an + otherwise-valid cached token. + """ + key = _CacheKey(token_endpoint=token_endpoint, scope=scope or "") + with self._lock: + self._cache.pop(key, None) + + def _mint(self, token_endpoint: str, scope: str) -> tuple[str, float]: + """POST grant_type=client_credentials (HTTP Basic) and parse the JWT.""" + basic = base64.b64encode( + f"{self._client_id}:{self._client_secret}".encode() + ).decode() + data = {"grant_type": "client_credentials"} + if scope: + data["scope"] = scope + headers = { + "Authorization": f"Basic {basic}", + "Content-Type": "application/x-www-form-urlencoded", + } + try: + with httpx.Client(timeout=self._timeout, transport=self._transport) as client: + resp = client.post(token_endpoint, data=data, headers=headers) + except httpx.HTTPError as exc: + raise OAuthTokenError(f"token endpoint request failed: {exc}") from exc + + if resp.status_code != 200: + # Never echo the response body verbatim (may reflect credentials). + raise OAuthTokenError( + f"token endpoint returned HTTP {resp.status_code} for scope " + f"'{scope or '(none)'}'" + ) + try: + payload = resp.json() + except ValueError as exc: + raise OAuthTokenError("token endpoint returned a non-JSON body") from exc + + token = payload.get("access_token") + if not token: + raise OAuthTokenError("token endpoint response had no access_token") + expires_in = payload.get("expires_in", 3600) + logger.info( + "Minted a client_credentials token (scope=%s, expires_in=%ss).", + scope or "(none)", + expires_in, + ) + return token, float(expires_in) diff --git a/src/ad_buyer/clients/mcp_client.py b/src/ad_buyer/clients/mcp_client.py index 1dd47b65..953e09e8 100644 --- a/src/ad_buyer/clients/mcp_client.py +++ b/src/ad_buyer/clients/mcp_client.py @@ -19,6 +19,21 @@ MCP_SDK_AVAILABLE = False +def _is_unauthorized(exc: Exception) -> bool: + """True when an exception looks like an HTTP 401 from the MCP transport. + + The MCP Streamable-HTTP transport surfaces auth failures as httpx errors or + generic exceptions whose message carries the status; we match either an + httpx 401 or a "401"/"Unauthorized" marker in the string form so a single + reactive re-mint can be attempted before giving up. + """ + resp = getattr(exc, "response", None) + if resp is not None and getattr(resp, "status_code", None) == 401: + return True + text = str(exc).lower() + return "401" in text or "unauthorized" in text + + @dataclass class MCPToolResult: """Result from an MCP tool call.""" @@ -178,25 +193,85 @@ class IABMCPClient: def __init__( self, base_url: str, + *, + auth_token_provider: Any = None, + token_endpoint: str = "", + scope: str = "", ): """Initialize the MCP client. Args: - base_url: Base URL for the MCP server + base_url: Base URL for the MCP server. + auth_token_provider: Optional + :class:`ad_buyer.auth.oauth_token_provider.OAuthTokenProvider`. + When supplied (the AWS Bedrock AgentCore CUSTOM_JWT path, where a + discovered seller advertises an OAuth issuer), the client mints a + client_credentials JWT from ``token_endpoint``/``scope`` and + attaches ``Authorization: Bearer `` to the MCP transport, + re-minting once on a 401. Omitted = no auth header (local dev / + unauthenticated servers). + token_endpoint: OAuth2 token endpoint (from the discovered record's + ``authentication`` object). Required when a provider is given. + scope: OAuth2 scope to request (e.g. ``seller-agent/invoke``). """ self.base_url = base_url.rstrip("/") - self.mcp_url = f"{self.base_url}/mcp/sse" + # Endpoint shapes (verified against the AWS AgentCore MCP docs): + # - An AgentCore CUSTOM_JWT runtime is invoked over HTTPS at + # https://bedrock-agentcore..amazonaws.com/runtimes//invocations?qualifier=DEFAULT + # (SDK/SigV4 is DISALLOWED for OAuth runtimes — a raw HTTPS request + # with a Bearer token is required). That full URL is used AS-IS; the + # Streamable-HTTP transport speaks MCP directly on it. + # - The legacy IAB agentic-direct server exposes MCP at `/mcp/sse`. + # Detect the AgentCore invocations URL and do NOT append `/mcp/sse`. + if "/invocations" in self.base_url or self.base_url.endswith("/invocations"): + self.mcp_url = self.base_url + else: + self.mcp_url = f"{self.base_url}/mcp/sse" self._tools: dict[str, dict] = {} self._session: ClientSession | None = None self._client_ctx = None self._read_stream = None self._write_stream = None self._get_session_id = None + self._auth_provider = auth_token_provider + self._token_endpoint = token_endpoint + self._scope = scope + + def _auth_headers(self, *, force: bool = False) -> dict[str, str]: + """Bearer header for the AgentCore CUSTOM_JWT path, or {} when no auth. + + ``force`` re-mints (used after a 401) rather than reusing the cached token. + """ + if self._auth_provider is None: + return {} + token = self._auth_provider.get_token(self._token_endpoint, self._scope, force=force) + return {"Authorization": f"Bearer {token}"} async def connect(self) -> None: - """Connect to the MCP server and initialize session.""" - # Create streamable HTTP client - self._client_ctx = streamablehttp_client(self.mcp_url) + """Connect to the MCP server and initialize session. + + On the authenticated (AgentCore CUSTOM_JWT) path, a 401 during connect + triggers a single reactive re-mint (the cached token may have expired) + before failing. + """ + try: + await self._connect_once(self._auth_headers()) + except Exception as exc: # noqa: BLE001 — reactively re-mint on a 401 only + if self._auth_provider is not None and _is_unauthorized(exc): + # Drop the cached token and retry once with a fresh JWT. + self._auth_provider.invalidate(self._token_endpoint, self._scope) + await self._connect_once(self._auth_headers(force=True)) + else: + raise + + async def _connect_once(self, headers: dict[str, str]) -> None: + # Create streamable HTTP client (headers carry the bearer JWT when auth + # is configured; empty otherwise). + self._client_ctx = ( + streamablehttp_client(self.mcp_url, headers=headers) + if headers + else streamablehttp_client(self.mcp_url) + ) streams = await self._client_ctx.__aenter__() self._read_stream, self._write_stream, self._get_session_id = streams diff --git a/src/ad_buyer/config/settings.py b/src/ad_buyer/config/settings.py index a1fe8fe2..d267f303 100644 --- a/src/ad_buyer/config/settings.py +++ b/src/ad_buyer/config/settings.py @@ -42,6 +42,19 @@ class Settings(BaseSettings): aamp_registry_url: str = "" aamp_registry_auth_token: str = "" + # Buyer's OWN OAuth2 client_credentials app-client (AWS Bedrock AgentCore + # deployment path only). Used to mint a JWT for a seller that is deployed as + # an AgentCore runtime behind a CUSTOM_JWT authorizer and ADVERTISES an OAuth + # token endpoint in its registry record's `authentication` object (Req 6.6). + # The token ENDPOINT + SCOPE come from that discovered record; the + # client_id/secret are the buyer's own identity and live HERE (never in the + # registry, never logged). Unset = the buyer cannot authenticate to a + # JWT-protected seller. The legacy API-key seller path (AuthMiddleware) and + # the arn: AgentCore SigV4 proxy path are unaffected — OAuth applies ONLY to + # a discovered seller that advertises an issuer. + buyer_oauth_client_id: str = "" + buyer_oauth_client_secret: str = "" + # Seller Agent Endpoints (comma-separated list of MCP/A2A server URLs) # Each endpoint should implement IAB Tech Lab OpenDirect/AdCOM standards seller_endpoints: str = "" @@ -169,11 +182,17 @@ def resolve_seller_base_url(self) -> str: # LLM Settings # Supported providers: anthropic (default), openai, gemini, bedrock # Set DEFAULT_LLM_MODEL to switch provider, e.g.: - # anthropic/claude-sonnet-4-5-20250929 (requires ANTHROPIC_API_KEY) - # openai/gpt-4o (requires OPENAI_API_KEY) - # gemini/gemini-2.5-flash (requires GOOGLE_API_KEY) - # bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 (requires AWS creds) - default_llm_model: str = "anthropic/claude-sonnet-4-5-20250929" + # anthropic/claude-sonnet-5 (requires ANTHROPIC_API_KEY) + # openai/gpt-4o (requires OPENAI_API_KEY) + # gemini/gemini-2.5-flash (requires GOOGLE_API_KEY) + # bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 (Converse; requires AWS creds) + # NOTE: to run Claude on Amazon Bedrock WITHOUT the Converse provider, do + # not use the ``bedrock/`` prefix; instead set + # ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL to Bedrock's /anthropic route (see + # below) with a Messages-supported model id. Only the current-generation + # Claude tiers are served on the Anthropic Messages API (Sonnet 5, + # Opus 4.7/4.8, Haiku 4.5); Sonnet 4.5 is Converse-only. + default_llm_model: str = "anthropic/claude-sonnet-5" manager_llm_model: str = "anthropic/claude-opus-4-8" llm_temperature: float = 0.3 llm_max_tokens: int = 4096 @@ -187,6 +206,17 @@ def resolve_seller_base_url(self) -> str: openai_compatible_llm_api_key: str | None = None openai_compatible_llm_api_base_url: str | None = None + # Alternative: any Anthropic-Messages-wire-compatible endpoint. This is the + # recommended way to run Claude on Amazon Bedrock: set + # ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL=https://bedrock-runtime..amazonaws.com/anthropic + # ANTHROPIC_COMPATIBLE_LLM_API_KEY= + # DEFAULT_LLM_MODEL / MANAGER_LLM_MODEL = the Bedrock model / inference-profile id + # CrewAI's native Anthropic provider then drives messages.create against + # Bedrock, so no Bedrock Converse toolUse/toolResult sanitization is needed. + # (Prefer bedrock-runtime over bedrock-mantle for new apps per AWS.) + anthropic_compatible_llm_api_key: str | None = None + anthropic_compatible_llm_api_base_url: str | None = None + # Database / Storage Configuration database_url: str = "sqlite:///./ad_buyer.db" redis_url: str | None = None diff --git a/src/ad_buyer/interfaces/agentcore/agentcore_seller_proxy.py b/src/ad_buyer/interfaces/agentcore/agentcore_seller_proxy.py new file mode 100644 index 00000000..5162606e --- /dev/null +++ b/src/ad_buyer/interfaces/agentcore/agentcore_seller_proxy.py @@ -0,0 +1,444 @@ +# Author: contributed to IAB Tech Lab (AgentCore adapter) + +"""AgentCore seller proxy — reach a deployed seller runtime by ARN. + +Amazon Bedrock AgentCore exposes a deployed agent as an **ARN**, invoked via +the SigV4-signed ``InvokeAgentRuntime`` data-plane API — NOT as a plain HTTP +URL. The buyer's ``DealBookingFlow`` planning path, however, is written against +the ``OpenDirectClient`` surface (``list_products`` / ``search_products`` / +``get_product`` / ``check_avails``) and expects structured ``Product`` / +``AvailsResponse`` objects. + +``AgentCoreSellerProxy`` bridges the two: it presents the same OpenDirect +method surface the flow needs, but under the hood calls the seller's **MCP** +runtime by ARN via ``bedrock-agentcore:InvokeAgentRuntime``, issuing MCP +``tools/call`` requests (``list_products``, ``get_product_details``, +``discover_inventory``) and mapping the STRUCTURED JSON tool results to the +OpenDirect models at the boundary. + +Why MCP and not A2A for this path: the seller's MCP runtime returns structured +JSON tool results (deterministically parseable into ``Product`` models), +whereas the A2A runtime returns natural-language crew output. A2A is the right +surface for agent-to-agent *negotiation dialogue* (a booking-path concern); the +planning path needs structured product/avails data, so it targets MCP. + +This replaces the previous ``localhost:9999`` stub in ``crew_tools`` for the +deployed (ARN) case. When ``SELLER_AGENT_URL`` is a plain HTTP URL (local dev), +the ordinary ``OpenDirectClient`` is used instead — this proxy is only for the +``arn:`` case. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +from ...clients.contract_mappers import from_wire_product +from ...models.opendirect import AvailsRequest, AvailsResponse, Product + +logger = logging.getLogger(__name__) + +# Seller MCP tool names (see seller-agent CREW_MCP_TOOLS / mcp_server). +_TOOL_LIST_PRODUCTS = "list_products" +_TOOL_GET_PRODUCT = "get_product_details" +_TOOL_DISCOVER = "discover_inventory" + + +class AgentCoreSellerProxy: + """OpenDirect-shaped client backed by InvokeAgentRuntime (seller MCP runtime). + + Only the read surface the planning flow exercises is implemented + (``list_products``, ``list_products_tolerant``, ``get_product``, + ``search_products``, ``check_avails``) plus ``close`` and the async context + manager, matching how ``DealBookingFlow`` / the research crews use the + client. Booking (orders/lines/deals) is not on the planning path and is not + proxied here. + + Args: + runtime_arn: The seller **MCP** AgentCore runtime ARN. + region: AWS region of the runtime. + session_id: Optional runtime session id for multi-call continuity. + client: Optional pre-built boto3 ``bedrock-agentcore`` client (tests + inject a stub; production builds one lazily). + """ + + def __init__( + self, + runtime_arn: str, + *, + region: str | None = None, + session_id: str | None = None, + client: Any = None, + token_provider: Any = None, + token_endpoint: str = "", + scope: str = "", + ) -> None: + self._arn = runtime_arn + self._region = region + self._session_id = session_id + self._client = client + self._mcp_id = 0 + self._mcp_session_id: str | None = None + self._initialized = False + # JWT/HTTPS transport mode (AWS Bedrock AgentCore CUSTOM_JWT path): when + # a token provider is supplied, the seller runtime enforces a JWT + # authorizer, so SigV4/boto3 is DISALLOWED — MCP JSON-RPC is sent over a + # raw HTTPS POST to the verified invocations URL with a Bearer token + # (re-minted once on a 401). Otherwise the legacy SigV4 boto3 path runs. + self._token_provider = token_provider + self._token_endpoint = token_endpoint + self._scope = scope + self._use_jwt = token_provider is not None + + # -- boto3 client (lazy) --------------------------------------------------- + + def _get_client(self) -> Any: + if self._client is None: + import boto3 # local import: only needed on the ARN path + + self._client = boto3.client("bedrock-agentcore", region_name=self._region) + return self._client + + # -- InvokeAgentRuntime + MCP framing -------------------------------------- + + def _next_id(self) -> int: + self._mcp_id += 1 + return self._mcp_id + + # AgentCore MCP runtimes speak Streamable-HTTP: the /mcp path is a + # pass-through of the InvokeAgentRuntime payload, but the MCP transport + # REQUIRES the dual Accept header and a JSON content type, and expects the + # standard initialize -> notifications/initialized handshake before any + # tools/call. boto3 exposes these as first-class header params (accept, + # contentType, mcpProtocolVersion, mcpSessionId). Omitting the Accept + # header yields HTTP 406 (McpRequestUnacceptableException, -32011). + _MCP_ACCEPT = "application/json, text/event-stream" + _MCP_CONTENT_TYPE = "application/json" + _MCP_PROTOCOL_VERSION = "2025-03-26" + + def _invoke_mcp(self, payload: dict[str, Any], *, expect_body: bool = True) -> bytes: + """One InvokeAgentRuntime call carrying an MCP JSON-RPC message. + + SigV4/boto3 by default; a raw HTTPS POST with a Bearer JWT when the + runtime enforces CUSTOM_JWT (``_use_jwt``). Both set the MCP-required + headers, thread the runtime + MCP session ids for microVM stickiness, + and capture the platform ``Mcp-Session-Id`` for reuse. Returns the raw + response body (empty for notifications). + """ + if self._use_jwt: + return self._invoke_mcp_jwt(payload, expect_body=expect_body) + + kwargs: dict[str, Any] = { + "agentRuntimeArn": self._arn, + "payload": json.dumps(payload).encode("utf-8"), + "accept": self._MCP_ACCEPT, + "contentType": self._MCP_CONTENT_TYPE, + "mcpProtocolVersion": self._MCP_PROTOCOL_VERSION, + } + if self._session_id: + kwargs["runtimeSessionId"] = self._session_id + if self._mcp_session_id: + kwargs["mcpSessionId"] = self._mcp_session_id + + resp = self._get_client().invoke_agent_runtime(**kwargs) + + # Capture the platform-issued MCP session id (header) for affinity. + sid = (resp.get("mcpSessionId") if isinstance(resp, dict) else None) or ( + resp.get("ResponseMetadata", {}) + .get("HTTPHeaders", {}) + .get("mcp-session-id") + if isinstance(resp, dict) + else None + ) + if sid: + self._mcp_session_id = sid + + if not expect_body: + return b"" + body = resp["response"] + return body.read() if hasattr(body, "read") else bytes(body) + + def _invoke_mcp_jwt(self, payload: dict[str, Any], *, expect_body: bool = True) -> bytes: + """MCP JSON-RPC over a raw HTTPS POST with a Bearer JWT (CUSTOM_JWT path). + + boto3/SigV4 is DISALLOWED for OAuth runtimes (per the AgentCore docs), so + we POST directly to the verified invocations URL. A 401 triggers a single + reactive re-mint before failing. + """ + import httpx + + from ad_buyer.registry.transport_selector import agentcore_invocations_url + + url = agentcore_invocations_url(self._arn) if self._arn.startswith("arn:") else self._arn + + def _post(force_token: bool) -> httpx.Response: + token = self._token_provider.get_token( + self._token_endpoint, self._scope, force=force_token + ) + headers = { + "Authorization": f"Bearer {token}", + "Accept": self._MCP_ACCEPT, + "Content-Type": self._MCP_CONTENT_TYPE, + } + if self._mcp_session_id: + headers["Mcp-Session-Id"] = self._mcp_session_id + with httpx.Client(timeout=120.0) as client: + return client.post( + url, + content=json.dumps(payload).encode("utf-8"), + headers=headers, + ) + + resp = _post(force_token=False) + if resp.status_code == 401: + # Cached token rejected — drop it and re-mint once. + self._token_provider.invalidate(self._token_endpoint, self._scope) + resp = _post(force_token=True) + + sid = resp.headers.get("mcp-session-id") or resp.headers.get("Mcp-Session-Id") + if sid: + self._mcp_session_id = sid + + if not expect_body: + return b"" + return resp.content + + def _do_initialize(self) -> None: + """Run the MCP initialize + notifications/initialized handshake (sync).""" + init_payload = { + "jsonrpc": "2.0", + "id": self._next_id(), + "method": "initialize", + "params": { + "protocolVersion": self._MCP_PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "ad_buyer.AgentCoreSellerProxy", "version": "1"}, + }, + } + raw = self._invoke_mcp(init_payload) + # Surface a hard protocol/auth error early (406/403 etc. arrive as text). + _raise_for_mcp_error(raw) + # Ack — a notification has no id and expects no body. + self._invoke_mcp( + {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}, + expect_body=False, + ) + self._initialized = True + + async def _call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any: + """Invoke one MCP ``tools/call`` on the seller runtime, return parsed data. + + Ensures the MCP session is initialized, then issues the ``tools/call`` + with the required MCP headers; the streamed response body is parsed back + into the tool's structured result. + """ + import asyncio + + payload = { + "jsonrpc": "2.0", + "id": self._next_id(), + "method": "tools/call", + "params": {"name": name, "arguments": arguments or {}}, + } + + def _invoke() -> bytes: + if not self._initialized: + self._do_initialize() + return self._invoke_mcp(payload) + + raw = await asyncio.get_event_loop().run_in_executor(None, _invoke) + return _parse_mcp_tool_result(raw) + + # -- OpenDirect surface ---------------------------------------------------- + + async def list_products(self, skip: int = 0, top: int = 50, **filters: Any) -> list[Product]: + data = await self._call_tool(_TOOL_LIST_PRODUCTS, {"limit": top}) + products = _extract_products(data) + if filters: + products = _apply_filters(products, filters) + return products[:top] + + async def list_products_tolerant( + self, skip: int = 0, top: int = 50, **filters: Any + ) -> tuple[list[Product], list[dict[str, Any]]]: + """Tolerant variant: per-product mapping, collecting rejects. + + Mirrors ``OpenDirectClient.list_products_tolerant`` so the per-seller + catalog-resolution path (which prefers the tolerant call) works + unchanged over the proxy. + """ + data = await self._call_tool(_TOOL_LIST_PRODUCTS, {"limit": top}) + raw_items = data.get("products", []) if isinstance(data, dict) else [] + products: list[Product] = [] + rejects: list[dict[str, Any]] = [] + for item in raw_items: + try: + products.append(_map_product(item)) + except Exception as exc: # noqa: BLE001 — collect, don't fail the catalog + rejects.append( + { + "product_id": item.get("product_id") if isinstance(item, dict) else None, + "name": item.get("name") if isinstance(item, dict) else None, + "reason": str(exc), + } + ) + if filters: + products = _apply_filters(products, filters) + return products[:top], rejects + + async def get_product(self, product_id: str) -> Product: + data = await self._call_tool(_TOOL_GET_PRODUCT, {"product_id": product_id}) + # get_product_details may return a bare product or {"product": {...}}. + item = data.get("product", data) if isinstance(data, dict) else data + return _map_product(item) + + async def search_products(self, filters: dict[str, Any]) -> list[Product]: + data = await self._call_tool(_TOOL_LIST_PRODUCTS, {"limit": 500}) + return _apply_filters(_extract_products(data), filters or {}) + + async def check_avails(self, request: AvailsRequest) -> AvailsResponse: + product_id = getattr(request, "product_id", None) or getattr(request, "productId", None) + data = await self._call_tool( + _TOOL_DISCOVER, + {"product_id": product_id} if product_id else {}, + ) + return AvailsResponse.model_validate(_avails_payload(data)) + + async def close(self) -> None: # noqa: D401 — parity with OpenDirectClient + """No persistent connection to close (boto3 client is stateless here).""" + return None + + async def __aenter__(self) -> AgentCoreSellerProxy: + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() + + +# -- module-level parsing helpers (unit-testable without AWS) ------------------ + + +def _decode_jsonrpc_body(raw: bytes | str) -> dict | list | None: + """Decode an MCP response body that may be raw JSON or SSE-framed. + + AgentCore MCP runtimes reply with ``Content-Type: text/event-stream`` (we + send ``Accept: application/json, text/event-stream``), so the JSON-RPC + envelope arrives as Server-Sent Events frames:: + + event: message + data: {"jsonrpc":"2.0","id":1,"result":{...}} + + This extracts and parses the ``data:`` payload. A plain-JSON body (no SSE + framing) is parsed directly. Returns the parsed envelope, or ``None`` if the + body is empty / unparseable (callers decide whether that is an error). + """ + if not raw: + return None + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + stripped = text.lstrip() + # Fast path: already JSON. + if stripped[:1] in "{[": + try: + return json.loads(stripped) + except ValueError: + pass + # SSE framing: concatenate all `data:` lines (last complete JSON wins). + data_lines = [ + line[len("data:") :].strip() + for line in text.splitlines() + if line.startswith("data:") + ] + for chunk in reversed(data_lines): + try: + return json.loads(chunk) + except ValueError: + continue + return None + + +def _raise_for_mcp_error(raw: bytes) -> None: + """Raise if a raw MCP response body carries a JSON-RPC ``error`` object. + + AgentCore returns most MCP errors as HTTP 200 with a JSON-RPC error in the + body (e.g. -32011 Accept-header/406, -32603 internal). boto3 raises for the + genuine HTTP-level 4xx/5xx itself; this catches the in-body errors so the + handshake fails loudly instead of proceeding against a dead session. + """ + env = _decode_jsonrpc_body(raw) + if isinstance(env, dict) and env.get("error"): + raise ValueError(f"seller MCP protocol error: {env['error']}") + + +def _parse_mcp_tool_result(raw: bytes) -> Any: + """Parse an MCP ``tools/call`` response body into the tool's structured data. + + The MCP result shape is ``{"result": {"content": [{"type": "text", + "text": ""}], "isError": false}}``. The tool's real payload is the + JSON string inside the first text content block. Also tolerates a bare + JSON-RPC ``{"result": }`` or a plain object. The transport body may be + raw JSON or SSE-framed (see ``_decode_jsonrpc_body``). + """ + envelope = _decode_jsonrpc_body(raw) + if envelope is None: + raw_txt = (raw.decode("utf-8", "replace") if isinstance(raw, bytes) else str(raw))[:200] + raise ValueError(f"seller MCP response was not JSON/SSE: {raw_txt!r}") + + if isinstance(envelope, dict) and envelope.get("error"): + raise ValueError(f"seller MCP error: {envelope['error']}") + + result = envelope.get("result", envelope) if isinstance(envelope, dict) else envelope + # MCP CallToolResult: content[0].text holds the tool's JSON payload. + if isinstance(result, dict) and isinstance(result.get("content"), list): + for block in result["content"]: + if isinstance(block, dict) and block.get("type") == "text" and block.get("text"): + try: + return json.loads(block["text"]) + except ValueError: + return {"raw_text": block["text"]} + return result + + +def _extract_products(data: Any) -> list[Product]: + items = data.get("products", []) if isinstance(data, dict) else [] + out: list[Product] = [] + for item in items: + try: + out.append(_map_product(item)) + except Exception as exc: # noqa: BLE001 + logger.warning("skipping unmappable product from seller MCP: %s", exc) + return out + + +def _map_product(item: dict[str, Any]) -> Product: + """Map a seller MCP product dict to the OpenDirect ``Product`` model. + + The seller MCP tools emit shared ``iab_agentic_primitives`` Product-shaped + records, so the shared-wire mapper (``from_wire_product``) is the source of + truth — the same boundary mapping ``OpenDirectClient`` uses. A record that + does not validate as a shared Product raises; callers on the tolerant path + (``list_products_tolerant`` / ``_extract_products``) collect the reject + rather than failing the whole catalog. + """ + from iab_agentic_primitives.primitives import Product as WireProduct + + return from_wire_product(WireProduct.model_validate(item)) + + +def _apply_filters(products: list[Product], filters: dict[str, Any]) -> list[Product]: + """Client-side filtering (parity with OpenDirectClient's client-side filter).""" + if not filters: + return products + out = products + inv = filters.get("inventory_type") or filters.get("channel") + if inv: + out = [p for p in out if getattr(p, "inventory_type", None) == inv] + return out + + +def _avails_payload(data: Any) -> dict[str, Any]: + """Normalize a discover_inventory result into an AvailsResponse dict.""" + if isinstance(data, dict) and "available" in data: + return data + # Minimal shape when the discover tool returns a summary/list. + return {"available": True, "products": data if isinstance(data, list) else []} diff --git a/src/ad_buyer/interfaces/agentcore/crew_tools.py b/src/ad_buyer/interfaces/agentcore/crew_tools.py index 72adf6e0..62296ebf 100644 --- a/src/ad_buyer/interfaces/agentcore/crew_tools.py +++ b/src/ad_buyer/interfaces/agentcore/crew_tools.py @@ -96,10 +96,13 @@ def run_campaign_plan(prompt: str, brief: dict[str, Any] | None = None) -> dict[ from ad_buyer.flows.deal_booking_flow import DealBookingFlow from ad_buyer.models.flow_state import BookingState - # Override buyer settings to use Bedrock instead of Anthropic. + # Override buyer settings to use the Bedrock-hosted model. Deploy.sh sets + # DEFAULT_LLM_MODEL plus ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL so this + # routes through the Anthropic Messages endpoint (no Converse patch). The + # fallback is a current-generation, Messages-supported Claude id. bedrock_model = os.environ.get( "DEFAULT_LLM_MODEL", - "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-5-v1:0", ) from ad_buyer.config.settings import settings as buyer_settings @@ -124,10 +127,47 @@ def run_campaign_plan(prompt: str, brief: dict[str, Any] | None = None) -> dict[ logger.info("Campaign brief: %s", json.dumps(brief, default=str)) - # Create client — dummy URL since seller is a separate AgentCore runtime + # Seller client selection: + # - Seller advertises OAuth (AWS Bedrock AgentCore CUSTOM_JWT deployment + # path): the seller runtime enforces a JWT authorizer, so SigV4/boto3 is + # rejected. Route through the transport selector, which builds the HTTPS + # invocations URL and attaches a client_credentials Bearer JWT + # (SELLER_TOKEN_ENDPOINT + SELLER_INVOKE_SCOPE come from the discovered + # registry record's `authentication` object; passed here via env for the + # deployed runtime). Works for a seller given by ARN or by https:// URL. + # - ARN without OAuth (same-account/dev deployed seller): the SigV4 + # AgentCoreSellerProxy over InvokeAgentRuntime. + # - Plain HTTP URL (local dev): the ordinary OpenDirectClient. seller_url = os.environ.get("SELLER_AGENT_URL", "http://localhost:8001") - client_url = "http://localhost:9999" if seller_url.startswith("arn:") else seller_url - client = OpenDirectClient(base_url=client_url) + mcp_arn = os.environ.get("SELLER_MCP_RUNTIME_ARN", "") + seller_token_endpoint = os.environ.get("SELLER_TOKEN_ENDPOINT", "") + seller_scope = os.environ.get("SELLER_INVOKE_SCOPE", "") + seller_endpoint = mcp_arn or seller_url + + if seller_token_endpoint: + # CUSTOM_JWT seller → JWT/HTTPS transport (arn: is converted to the + # HTTPS invocations URL inside the selector; https:// used as-is). + from ad_buyer.registry.transport_selector import select_seller_client + + client = select_seller_client( + seller_endpoint, + authentication={ + "token_endpoint": seller_token_endpoint, + "scope": seller_scope, + }, + region=os.environ.get("AWS_REGION"), + ) + elif seller_endpoint.startswith("arn:"): + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import ( + AgentCoreSellerProxy, + ) + + client = AgentCoreSellerProxy( + runtime_arn=seller_endpoint, + region=os.environ.get("AWS_REGION"), + ) + else: + client = OpenDirectClient(base_url=seller_url) flow = DealBookingFlow(client) try: diff --git a/src/ad_buyer/interfaces/agentcore/http_main.py b/src/ad_buyer/interfaces/agentcore/http_main.py index 8b156622..c8d90c10 100644 --- a/src/ad_buyer/interfaces/agentcore/http_main.py +++ b/src/ad_buyer/interfaces/agentcore/http_main.py @@ -55,17 +55,52 @@ os.environ.setdefault("STORAGE_TYPE", "sqlite") os.environ.setdefault("DATABASE_URL", "sqlite:///:memory:") -# Apply CrewAI patches BEFORE any CrewAI imports -try: - from patches.crewai_bedrock_fix import apply_patches as apply_bedrock_patches - - apply_bedrock_patches() -except ImportError as _ie: - logging.getLogger(__name__).warning( - "crewai_bedrock_fix not found: %s (repo_root=%s)", _ie, _repo_root - ) -except Exception as _exc: - logging.getLogger(__name__).warning("crewai_bedrock_fix patch failed: %s", _exc) +# Durable Bedrock auth: if the Anthropic-compatible base URL is a Bedrock +# endpoint, mint a fresh bearer token from the runtime's execution role NOW +# (at startup), authoritative over any baked/stale key. Avoids baking a +# short-lived token into --env at deploy time (which expires and 403s every +# crew call). See ad_buyer.llm.bedrock_token. +from ad_buyer.llm.bedrock_token import ensure_bedrock_token # noqa: E402 + +ensure_bedrock_token() + +# Apply CrewAI patches BEFORE any CrewAI imports. +# +# The Bedrock *Converse* sanitizer (crewai_bedrock_fix) is only needed on the +# legacy Converse path (DEFAULT_LLM_MODEL="bedrock/..."). The recommended setup +# routes Claude through Bedrock's Anthropic Messages endpoint by setting +# ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL, which uses CrewAI's native Anthropic +# provider (messages.create) and does NOT need Converse tool-block +# sanitization. So we apply the Converse patch only when that base URL is unset +# AND the model is a bedrock/ Converse model. +_default_model = os.environ.get("DEFAULT_LLM_MODEL", "") +_on_converse_path = not os.environ.get( + "ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL" +) and _default_model.startswith("bedrock/") +if _on_converse_path: + try: + from patches.crewai_bedrock_fix import apply_patches as apply_bedrock_patches + + apply_bedrock_patches() + except ImportError as _ie: + logging.getLogger(__name__).warning( + "crewai_bedrock_fix not found: %s (repo_root=%s)", _ie, _repo_root + ) + except Exception as _exc: + logging.getLogger(__name__).warning("crewai_bedrock_fix patch failed: %s", _exc) +else: + # Anthropic Messages endpoint path: CrewAI (>=1.15) emits a ``strict`` tool + # field that Bedrock's Anthropic-compatible endpoint rejects. Strip it. + try: + from patches.crewai_bedrock_anthropic_fix import apply_patches as apply_anthropic_patches + + apply_anthropic_patches() + except ImportError as _ie: + logging.getLogger(__name__).warning( + "crewai_bedrock_anthropic_fix not found: %s (repo_root=%s)", _ie, _repo_root + ) + except Exception as _exc: + logging.getLogger(__name__).warning("crewai_bedrock_anthropic_fix patch failed: %s", _exc) try: from patches.crewai_agentcore_memory import apply_patches as apply_memory_patches diff --git a/src/ad_buyer/llm/__init__.py b/src/ad_buyer/llm/__init__.py index 1144940d..d156dbf0 100644 --- a/src/ad_buyer/llm/__init__.py +++ b/src/ad_buyer/llm/__init__.py @@ -27,6 +27,11 @@ # that speaks the OpenAI wire format (NVIDIA NIM, Ollama, HuggingFace TGI, ...). _OPENAI_COMPATIBLE_PROVIDER = "openai" +# CrewAI's native Anthropic SDK client; with a base_url it drives any endpoint +# that speaks the Anthropic Messages wire format (Amazon Bedrock's /anthropic +# route, a self-hosted Anthropic-compatible gateway, ...). +_ANTHROPIC_COMPATIBLE_PROVIDER = "anthropic" + # Anthropic removed sampling parameters (temperature/top_p/top_k) starting # with Opus 4.7; sending temperature to these families returns # "400 invalid_request_error: 'temperature' is deprecated for this model." @@ -50,6 +55,28 @@ def _model_accepts_temperature(model: str) -> bool: return not any(family in model_id for family in _TEMPERATURE_REJECTING_MODEL_FAMILIES) +# OpenAI's newer models (GPT-5.x and the o-series reasoning models) reject the +# legacy ``max_tokens`` parameter — sending it returns +# "400 invalid_request_error: Unsupported parameter: 'max_tokens' is not +# supported with this model." They take ``max_completion_tokens`` instead. +# Matched as substrings so a Bedrock inference-profile id +# ("us.openai.gpt-5.6-sol", ...) or a raw OpenAI id ("gpt-5", "o1", ...) is +# covered. Other models (gpt-4o, gpt-oss, Claude, Nova, ...) keep ``max_tokens``. +_MAX_COMPLETION_TOKENS_MODEL_FAMILIES = ( + "gpt-5", + "gpt-6", + "o1", + "o3", + "o4", +) + + +def _model_uses_max_completion_tokens(model: str) -> bool: + """Return ``True`` if ``model`` needs ``max_completion_tokens`` (not ``max_tokens``).""" + model_id = model.lower() + return any(family in model_id for family in _MAX_COMPLETION_TOKENS_MODEL_FAMILIES) + + def build_llm(model: str, temperature: float, max_tokens: int) -> LLM: """Build an ``LLM`` for ``model``, honoring a custom base URL if configured. @@ -59,10 +86,36 @@ def build_llm(model: str, temperature: float, max_tokens: int) -> LLM: """ settings = get_settings() - kwargs: dict[str, Any] = {"model": model, "max_tokens": max_tokens} + kwargs: dict[str, Any] = {"model": model} + # GPT-5.x / o-series reject legacy ``max_tokens``; they take + # ``max_completion_tokens``. Everything else keeps ``max_tokens``. + if _model_uses_max_completion_tokens(model): + kwargs["max_completion_tokens"] = max_tokens + else: + kwargs["max_tokens"] = max_tokens if _model_accepts_temperature(model): kwargs["temperature"] = temperature + # Amazon Bedrock's Anthropic-compatible endpoint (and any other endpoint + # that speaks the Anthropic Messages wire format). Setting + # ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL routes Claude through CrewAI's + # native Anthropic provider (messages.create) against that base URL, + # instead of the Bedrock Converse provider. This is the path that lets us + # run Claude on Bedrock WITHOUT the Converse toolUse/toolResult sanitizer: + # the Messages API assembles tool turns itself, so the orphaned-block + # ValidationException class does not arise. Use a Messages-supported model + # id (e.g. Claude Sonnet 5 / Opus 4.7+ / Haiku 4.5) and a Bedrock API key. + # For bedrock-runtime, base_url is + # https://bedrock-runtime..amazonaws.com/anthropic + # and the model id is the Bedrock model / inference-profile id. + if settings.anthropic_compatible_llm_api_base_url: + return LLM( + api_key=settings.anthropic_compatible_llm_api_key, + provider=_ANTHROPIC_COMPATIBLE_PROVIDER, + base_url=settings.anthropic_compatible_llm_api_base_url, + **kwargs, + ) + if settings.openai_compatible_llm_api_base_url: return LLM( api_key=settings.openai_compatible_llm_api_key, diff --git a/src/ad_buyer/llm/bedrock_token.py b/src/ad_buyer/llm/bedrock_token.py new file mode 100644 index 00000000..3aeacfc9 --- /dev/null +++ b/src/ad_buyer/llm/bedrock_token.py @@ -0,0 +1,139 @@ +# Donated to IAB Tech Lab + +"""Self-sustaining Bedrock bearer-token minting for the AgentCore runtimes. + +Problem this solves: the Anthropic-compatible path (Claude on Bedrock via the +``/anthropic`` Messages route) authenticates with a Bedrock **API key / bearer +token**, not SigV4 — CrewAI's native Anthropic client sends an ``x-api-key`` +header, so the runtime's IAM execution role cannot be used directly by the SDK. +Baking a token into the runtime ``--env`` at deploy time means it EXPIRES within +hours and every crew call then 403s with ``permission_error: Bearer Token has +expired``. + +Durable fix (Req 11): mint the token from the runtime's own execution role via +``aws-bedrock-token-generator`` (it turns the role's SigV4 credentials into a +short-lived Bedrock bearer token). The role is the single source of truth — no +secret in ``--env``, nothing to hand-rotate. + +Contract: +- On a **Bedrock** base URL, the mint is AUTHORITATIVE: we ALWAYS mint from the + role and overwrite ``ANTHROPIC_COMPATIBLE_LLM_API_KEY``. A stale/baked value + (e.g. an expired token retained across an AgentCore ``--auto-update-on-conflict`` + redeploy) must NEVER shadow a fresh role-minted token. +- On a **non-Bedrock** base URL (real Anthropic / a BYO gateway that needs a + real API key), we RESPECT an operator-supplied ``ANTHROPIC_COMPATIBLE_LLM_API_KEY`` + and never overwrite it. +- ``refresh_bedrock_token()`` re-mints on demand (call it reactively on a + Bedrock 401/403) so a long-lived runtime never serves a stale token. +- Any failure (generator missing, no creds) logs and leaves the key as-is — + ``build_llm``/startup surfaces a clear error rather than crashing the import. +""" + +from __future__ import annotations + +import logging +import os + +logger = logging.getLogger(__name__) + +_API_KEY_ENV = "ANTHROPIC_COMPATIBLE_LLM_API_KEY" +_BASE_URL_ENV = "ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL" + + +def _is_bedrock_base_url(base_url: str | None) -> bool: + """True when the Anthropic-compatible base URL points at Amazon Bedrock. + + deploy.sh sets ``https://bedrock-runtime..amazonaws.com/anthropic``. + We match on the AWS Bedrock host so a real-Anthropic or BYO gateway URL is + NOT treated as Bedrock (those keep operator-supplied keys). + """ + if not base_url: + return False + u = base_url.lower() + return "bedrock-runtime" in u or ("bedrock" in u and "amazonaws.com" in u) + + +def _mint_token(region: str | None) -> str | None: + """Mint a fresh Bedrock bearer token from the execution role, or None.""" + try: + from aws_bedrock_token_generator import provide_token + except ImportError: + logger.warning( + "aws-bedrock-token-generator not installed — cannot mint a Bedrock " + "token from the execution role. Add the dependency to requirements " + "or set %s explicitly.", + _API_KEY_ENV, + ) + return None + + try: + token = provide_token(region=region) if region else provide_token() + except Exception as exc: # noqa: BLE001 — never crash startup on a mint failure + logger.warning("Bedrock token mint from execution role failed: %s", exc) + return None + + if not token: + logger.warning("Bedrock token generator returned an empty token.") + return None + return token + + +def _region(region: str | None) -> str | None: + return region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") + + +def ensure_bedrock_token(region: str | None = None) -> bool: + """Ensure a valid Bedrock bearer token is available for the Anthropic path. + + Called once at entrypoint startup. Returns True if a token is present in the + environment after this call, False otherwise. Never raises. + """ + base_url = os.environ.get(_BASE_URL_ENV) + + # Anthropic-compatible path not configured at all → nothing to do. + if not base_url: + logger.debug("%s unset — Anthropic-compatible path not in use; no mint.", _BASE_URL_ENV) + return bool(os.environ.get(_API_KEY_ENV)) + + # Non-Bedrock base URL (real Anthropic / BYO gateway) → respect an existing + # operator-supplied key; do not mint against a non-Bedrock endpoint. + if not _is_bedrock_base_url(base_url): + if os.environ.get(_API_KEY_ENV): + logger.info("%s set for a non-Bedrock base URL — respecting it.", _API_KEY_ENV) + return True + logger.debug("Non-Bedrock base URL and no key set — leaving key unset.") + return False + + # Bedrock base URL → the mint is AUTHORITATIVE. Always overwrite any baked + # value so a stale/expired token can never shadow the fresh role-minted one. + token = _mint_token(_region(region)) + if token is None: + # Leave whatever is there (may be a baked token); a clear error surfaces + # downstream if it is invalid. We did not make things worse. + return bool(os.environ.get(_API_KEY_ENV)) + + had_prior = bool(os.environ.get(_API_KEY_ENV)) + os.environ[_API_KEY_ENV] = token + logger.info( + "Minted a fresh Bedrock bearer token from the execution role " + "(region=%s, replaced_baked_value=%s).", + _region(region) or "default", + had_prior, + ) + return True + + +def refresh_bedrock_token(region: str | None = None) -> bool: + """Reactively re-mint the Bedrock token (e.g. after a 401/403). + + Only acts on a Bedrock base URL; a no-op (returns False) otherwise. Sets + ``ANTHROPIC_COMPATIBLE_LLM_API_KEY`` to the fresh token on success. + """ + if not _is_bedrock_base_url(os.environ.get(_BASE_URL_ENV)): + return False + token = _mint_token(_region(region)) + if token is None: + return False + os.environ[_API_KEY_ENV] = token + logger.info("Re-minted the Bedrock bearer token from the execution role (refresh).") + return True diff --git a/src/ad_buyer/registry/transport_selector.py b/src/ad_buyer/registry/transport_selector.py new file mode 100644 index 00000000..cca9f755 --- /dev/null +++ b/src/ad_buyer/registry/transport_selector.py @@ -0,0 +1,148 @@ +# Donated to IAB Tech Lab + +"""Transport selection for a discovered seller endpoint. + +Maps a seller's discovered endpoint (and its advertised auth, if any) to the +right client, keeping the three transports the buyer supports behind one +decision point: + +- ``arn:aws:bedrock-agentcore:...`` → :class:`AgentCoreSellerProxy` + (InvokeAgentRuntime, authenticated with **SigV4/IAM** — same-partition path). +- ``https://`` / ``http://`` WITH an advertised OAuth issuer → an authenticated + :class:`IABMCPClient` that mints a client_credentials **JWT** via + :class:`OAuthTokenProvider` and sends ``Authorization: Bearer`` (the AWS + Bedrock AgentCore CUSTOM_JWT deployment path, where a cross-org seller + advertises a token endpoint in its registry ``authentication`` object). +- ``https://`` / ``http://`` WITHOUT an issuer → a plain (unauthenticated) + :class:`IABMCPClient` (local dev / open servers). + +Only the AgentCore CUSTOM_JWT case uses OAuth — see the note in +:mod:`ad_buyer.auth.oauth_token_provider` on why the other paths use SigV4 or a +pre-shared API key instead. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any + +logger = logging.getLogger(__name__) + + +def agentcore_invocations_url(runtime_arn: str, *, qualifier: str = "DEFAULT") -> str: + """Build the HTTPS MCP invocation URL for an AgentCore runtime ARN. + + Verified shape (AWS AgentCore MCP docs): a CUSTOM_JWT (OAuth) runtime is + invoked over HTTPS at + ``https://bedrock-agentcore..amazonaws.com/runtimes//invocations?qualifier=`` + where ``ENCODED_ARN`` replaces ``:``→``%3A`` and ``/``→``%2F``. boto3/SigV4 + is DISALLOWED for OAuth runtimes, so this URL is called with a Bearer token + via the Streamable-HTTP MCP transport. + """ + # arn:aws:bedrock-agentcore:::runtime/ + parts = runtime_arn.split(":") + region = parts[3] if len(parts) > 4 else "" + encoded = runtime_arn.replace(":", "%3A").replace("/", "%2F") + return ( + f"https://bedrock-agentcore.{region}.amazonaws.com" + f"/runtimes/{encoded}/invocations?qualifier={qualifier}" + ) + + +def select_seller_client( + endpoint: str, + *, + authentication: dict[str, Any] | None = None, + region: str | None = None, + token_provider: Any = None, +): + """Return an appropriate seller client for ``endpoint``. + + Routing: + - ``arn:`` WITHOUT advertised OAuth → SigV4 ``AgentCoreSellerProxy`` + (same-account/dev; the seller does not enforce CUSTOM_JWT). + - ``arn:`` WITH advertised OAuth → the runtime enforces CUSTOM_JWT, so + SigV4 is rejected; build the HTTPS invocations URL from the ARN and use + the JWT client (boto3 is disallowed for OAuth runtimes). + - ``https://``/``http://`` WITH advertised OAuth → JWT client on that URL. + - ``https://``/``http://`` WITHOUT OAuth → plain (unauthenticated) client. + + Args: + endpoint: The discovered ``endpoint_url`` (an ``arn:`` or an http(s) URL). + authentication: The discovered record's ``authentication`` object, e.g. + ``{"type": "oauth2", "token_endpoint": "...", "scope": "..."}``. + region: AWS region (ARN path). + token_provider: Optional pre-built + :class:`~ad_buyer.auth.oauth_token_provider.OAuthTokenProvider`. + """ + token_endpoint, scope = _oauth_from_authentication(authentication) + + if endpoint.startswith("arn:") and not token_endpoint: + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import ( + AgentCoreSellerProxy, + ) + + logger.info("Transport: AgentCore SigV4 proxy (arn, no auth) for %s", endpoint) + return AgentCoreSellerProxy( + runtime_arn=endpoint, + region=region or os.environ.get("AWS_REGION"), + ) + + # From here, an OAuth issuer is advertised → JWT/HTTPS path. The seller + # runtime enforces CUSTOM_JWT (SigV4/boto3 disallowed). Use the + # AgentCoreSellerProxy in JWT mode: it keeps the OpenDirect-shaped surface + # DealBookingFlow needs, but frames MCP over a raw HTTPS POST with a Bearer + # token instead of SigV4. An ARN endpoint lets the proxy derive the HTTPS + # invocations URL; a full https:// endpoint is used directly. + if token_endpoint: + from ad_buyer.auth.oauth_token_provider import OAuthTokenProvider + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import ( + AgentCoreSellerProxy, + ) + + provider = token_provider or OAuthTokenProvider() + if not provider.configured: + # The seller REQUIRES auth but the buyer has no client credentials — + # fail loudly rather than silently connecting unauthenticated. + raise RuntimeError( + "seller advertises OAuth auth but the buyer has no client " + "credentials configured (set BUYER_OAUTH_CLIENT_ID / " + "BUYER_OAUTH_CLIENT_SECRET)" + ) + logger.info("Transport: JWT/HTTPS AgentCore proxy (OAuth) for %s", endpoint) + return AgentCoreSellerProxy( + runtime_arn=endpoint, + region=region or os.environ.get("AWS_REGION"), + token_provider=provider, + token_endpoint=token_endpoint, + scope=scope, + ) + + from ad_buyer.clients.mcp_client import IABMCPClient + + logger.info("Transport: plain MCP client (no auth) for %s", endpoint) + return IABMCPClient(base_url=endpoint) + + +def _oauth_from_authentication( + authentication: dict[str, Any] | None, +) -> tuple[str, str]: + """Extract ``(token_endpoint, scope)`` from a discovered auth object. + + Tolerant of the field-name variants a registry record may carry + (``token_endpoint``/``tokenEndpoint``/``token_url``; ``scope``/``scopes``). + Returns ``("", "")`` when no OAuth issuer is advertised. + """ + if not authentication: + return "", "" + token_endpoint = ( + authentication.get("token_endpoint") + or authentication.get("tokenEndpoint") + or authentication.get("token_url") + or "" + ) + scope = authentication.get("scope") or authentication.get("scopes") or "" + if isinstance(scope, (list, tuple)): + scope = " ".join(str(s) for s in scope) + return token_endpoint, scope diff --git a/tests/integration/agentcore/test_runtime.py b/tests/integration/agentcore/test_runtime.py index 864c2cb2..7193ece2 100644 --- a/tests/integration/agentcore/test_runtime.py +++ b/tests/integration/agentcore/test_runtime.py @@ -80,11 +80,18 @@ def runtime_config(request) -> RuntimeConfig: def invoke_runtime( config: RuntimeConfig, payload: dict, - timeout: int = 120, + timeout: int = 240, max_retries: int = 3, retry_wait: int = 30, ) -> dict: - """Invoke the runtime and return parsed response.""" + """Invoke the runtime and return parsed response. + + Default ``timeout`` is generous: these live tests drive the REAL + multi-agent crew on Bedrock Claude (research + memory save/retrieve + + seller tool calls), which routinely runs ~3 min for a full campaign plan, + plus a possible cold-start on the first invoke after a deploy. Crew/ + integration tests raise it further. + """ payload_json = json.dumps(payload) cmd = ["agentcore", "invoke", payload_json] env = os.environ.copy() @@ -189,6 +196,11 @@ def test_campaign_plan_with_budget(self, runtime_config): ), "routing_mode": "crew", }, + # Crew mode fans out a multi-agent crew over multi-turn tool calls on + # Claude Sonnet 5; observed end-to-end latency is ~4 min, well past the + # 120s default. Give real headroom so a correct plan isn't scored as a + # failure purely on latency. + timeout=360, ) assert result["success"], f"Invoke failed: {result['error']}" response = result["response"].lower() @@ -204,6 +216,7 @@ def test_campaign_plan_includes_next_step(self, runtime_config): "prompt": "Plan a $200K Q1 campaign for mobile and display", "routing_mode": "crew", }, + timeout=360, ) assert result["success"], f"Invoke failed: {result['error']}" response = result["response"].lower() @@ -219,6 +232,7 @@ def test_campaign_plan_returns_metadata(self, runtime_config): "prompt": "Plan a $1M Q3 brand awareness campaign", "routing_mode": "crew", }, + timeout=360, ) assert result["success"], f"Invoke failed: {result['error']}" # The raw output should contain campaign plan indicators @@ -226,3 +240,130 @@ def test_campaign_plan_returns_metadata(self, runtime_config): assert any(kw in response.lower() for kw in ["budget", "campaign", "plan", "allocation"]), ( f"No plan data: {response[:300]}" ) + + +# --------------------------------------------------------------------------- +# Buyer -> Seller runtime-to-runtime integration (AgentCoreSellerProxy) +# --------------------------------------------------------------------------- + + +@pytest.mark.agentcore +class TestBuyerSellerIntegration: + """End-to-end buyer->seller over InvokeAgentRuntime (the deployed loop). + + This is the ONE test that requires both runtimes deployed. The broader + buyer<->seller *functional* contract — pricing tiers by buyer identity, the + deal request/booking flow, MCP tools, order lifecycle, negotiation — is + already covered OFFLINE by the existing suites (e.g. seller + ``tests/integration/test_deal_flow_e2e.py`` / ``test_mcp_integration.py``, + buyer ``tests/integration/test_deal_booking_flow.py`` / + ``test_auth_session_negotiation.py``). This test adds only what those cannot: + proof the real cross-*runtime* transport works when deployed. + + The buyer crew's research path calls ``search_advertising_products``, which + (when ``SELLER_AGENT_URL`` is an AgentCore ARN) is backed by + ``AgentCoreSellerProxy`` -> ``bedrock-agentcore:InvokeAgentRuntime`` against + the seller's **MCP** runtime, issuing MCP ``tools/call`` over the + Streamable-HTTP transport (initialize handshake + dual Accept header + + SSE-framed result parsing). + + Where the signal lives: the crew's FINAL invoke response is a budget + *allocation plan* (channels/percentages/rationale) — it does NOT echo the + seller's product ids. The proof of the closed loop is in the runtime's + CloudWatch logs: the buyer logs the seller catalog's REAL ``inv-*`` ids + coming back through the proxy. (Note: ``inv-*`` is the seller *catalog* + namespace; the buyer crew's own *recommendation* ids use ``prod-*`` and are + unrelated — so only an ``inv-*`` id proves the live seller catalog.) + + Requires BOTH runtimes deployed and the buyer runtime's execution role + granted ``InvokeAgentRuntime`` on the seller runtime (see + ``infra/aws/agentcore/auth-agentcore.yaml``). Skips cleanly when the buyer + was not deployed against an ARN seller, or when CloudWatch is unreadable. + """ + + _REAL_INV_RE = re.compile(r"inv-[a-z0-9-]+", re.IGNORECASE) + + @staticmethod + def _log_group(config: "RuntimeConfig") -> str | None: + """DEFAULT runtime log group for the buyer, derived from its ARN.""" + # arn:...:runtime/ -> /aws/bedrock-agentcore/runtimes/-DEFAULT + m = re.search(r"runtime/([^/]+)$", config.arn) + return f"/aws/bedrock-agentcore/runtimes/{m.group(1)}-DEFAULT" if m else None + + def _recent_log_text(self, config: "RuntimeConfig", since_ms: int, pattern: str) -> str: + """Return recent buyer log messages matching a CloudWatch filter pattern. + + Skips (not fails) when logs are unreadable — the runtime behaviour is + still exercised by the invoke; only the assertion channel is missing. + """ + try: + import boto3 + except ImportError: # pragma: no cover + pytest.skip("boto3 unavailable to read CloudWatch logs") + + lg = self._log_group(config) + if not lg: + pytest.skip(f"Could not derive log group from ARN {config.arn}") + + session = boto3.Session(profile_name=config.profile) if config.profile else boto3.Session() + logs = session.client("logs", region_name=config.region) + try: + lines: list[str] = [] + paginator = logs.get_paginator("filter_log_events") + for page in paginator.paginate( + logGroupName=lg, + startTime=since_ms, + filterPattern=pattern, + ): + lines.extend(e.get("message", "") for e in page.get("events", [])) + if len(lines) > 500: + break + return "\n".join(lines) + except logs.exceptions.ResourceNotFoundException: + pytest.skip(f"Log group {lg} not found — runtime not deployed?") + except Exception as e: # noqa: BLE001 — logs are the assertion channel, not the SUT + pytest.skip(f"Could not read CloudWatch logs: {e}") + + def test_crew_reaches_seller_inventory(self, runtime_config): + """Buyer crew discovers real seller inv-* inventory via the proxy. + + Invokes the crew, then asserts the buyer's logs surfaced real ``inv-*`` + ids from the seller catalog through the proxy path. ``inv-*`` is the + seller *catalog* id namespace; the buyer's own crew *recommendations* + use a separate ``prod-*`` id namespace (see + ``tests/integration/test_deal_booking_flow.py``), so the two are NOT + interchangeable and only an ``inv-*`` id proves the seller catalog was + reached live through the proxy. + """ + since_ms = int((time.time() - 5) * 1000) + result = invoke_runtime( + runtime_config, + { + "prompt": ( + "Plan a $500K Q4 automotive campaign across CTV and digital" + " video, and discover seller inventory." + ), + "routing_mode": "crew", + }, + timeout=300, + ) + assert result["success"], f"Invoke failed: {result['error']}" + + # Did the buyer exercise the seller-search tool at all? (skip guard) + tool_text = self._recent_log_text( + runtime_config, since_ms, '"search_advertising_products"' + ) + if "search_advertising_products" not in tool_text: + pytest.skip( + "Buyer did not exercise the seller-search tool (not wired to an " + "ARN seller?); deploy with --seller-url to run." + ) + + # The definitive signal: real seller inv-* ids in the logs. CloudWatch + # filter patterns need the hyphen term quoted. + inv_text = self._recent_log_text(runtime_config, since_ms, '"inv-"') + inv_ids = sorted(set(self._REAL_INV_RE.findall(inv_text))) + assert inv_ids, ( + "Buyer's seller search returned no real inv-* ids — the proxy->MCP " + f"path returned no catalog data. Log tail: {inv_text[-500:]}" + ) diff --git a/tests/integration/test_registry_roundtrip.py b/tests/integration/test_registry_roundtrip.py new file mode 100644 index 00000000..dbc1f9b7 --- /dev/null +++ b/tests/integration/test_registry_roundtrip.py @@ -0,0 +1,132 @@ +# Donated to IAB Tech Lab + +"""Task 7.3 — register → discover → connect ROUND-TRIP (no network, no hosted registry). + +Proves the cross-org discovery wiring end-to-end IN PROCESS: +1. Seed the library's registry double (``create_registry_double``) with a seller + runtime record whose ``endpoint_url`` is its AgentCore invocations URL and + whose ``authentication`` advertises an OAuth2 token endpoint + scope. +2. Discover it through the buyer's :class:`AampRegistryClient` over an httpx + ``ASGITransport`` (the double's FastAPI app, no sockets). +3. Assert the buyer's transport selector (4.3) routes the discovered + OAuth-advertising endpoint to the JWT/HTTPS MCP client, and that an ``arn:`` + endpoint without OAuth routes to the SigV4 proxy (parity case). + +Registry is config-swappable, so prod is only an ``AAMP_REGISTRY_URL`` change — +this test asserts the seam, not a hosted registry. +""" + +import sys +from pathlib import Path + +import httpx +import pytest + +_SRC = Path(__file__).resolve().parents[2] / "src" +if str(_SRC) not in sys.path: + sys.path.insert(0, str(_SRC)) + +from ad_buyer.registry.aamp_client import AampRegistryClient # noqa: E402 +from ad_buyer.registry.transport_selector import ( # noqa: E402 + agentcore_invocations_url, + select_seller_client, +) + +pytestmark = pytest.mark.asyncio + +# A CUSTOM_JWT seller MCP runtime advertised in the registry. Placeholder IDs +# only (never the real account/pool/client). +_SELLER_ARN = ( + "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/aamp_seller_mcp-EXAMPLE01" +) +_TOKEN_ENDPOINT = "https://example-seller.auth.us-west-2.amazoncognito.com/oauth2/token" +_SCOPE = "seller-agent/invoke" + + +def _seed_seller(app, *, endpoint_url: str) -> None: + """Preseed the double's store with one OAuth-advertising seller runtime.""" + app.state.store.agents[1] = { + "id": 1, + "agent_name": "aamp-seller", + "primary_domain": "example.com", + "endpoint_url": endpoint_url, + "protocol_type": "MCP", + "verification_status": "verified", + "domain_verified": True, + "capabilities": ["advertising_inventory"], + "authentication": { + "type": "oauth2", + "token_endpoint": _TOKEN_ENDPOINT, + "scope": _SCOPE, + }, + } + + +async def _discover(app): + """Discover sellers through AampRegistryClient over an in-process ASGI transport.""" + transport = httpx.ASGITransport(app=app) + client = AampRegistryClient( + base_url="http://registry.local", + auth_token="user-token", # a valid token in the double's default table + transport=transport, + ) + return await client.discover_sellers() + + +async def test_roundtrip_discovers_seeded_seller(): + """The seeded seller is discoverable in-process and carries its endpoint.""" + from iab_agentic_primitives.sandbox_registry import create_registry_double + + https_endpoint = agentcore_invocations_url(_SELLER_ARN) + app = create_registry_double() + _seed_seller(app, endpoint_url=https_endpoint) + + cards = await _discover(app) + + assert cards, "discovery returned no sellers from the in-process registry" + assert any(c.url == https_endpoint for c in cards), ( + f"seeded seller endpoint not surfaced; got {[c.url for c in cards]}" + ) + + +async def test_roundtrip_oauth_endpoint_routes_to_jwt_client(): + """An OAuth-advertising https endpoint routes to the JWT/HTTPS client (4.3). + + The selector is fail-closed: when the seller advertises OAuth it REQUIRES + buyer credentials (a token_provider or BUYER_OAUTH_CLIENT_ID/SECRET) and + raises otherwise. We supply a stub provider — the same seam the crew uses — + so the JWT client is built. + """ + https_endpoint = agentcore_invocations_url(_SELLER_ARN) + authentication = { + "type": "oauth2", + "token_endpoint": _TOKEN_ENDPOINT, + "scope": _SCOPE, + } + + class _StubTokenProvider: + configured = True + + def get_token(self): + return "stub-jwt" + + client = select_seller_client( + https_endpoint, + authentication=authentication, + region="us-west-2", + token_provider=_StubTokenProvider(), + ) + # JWT-mode AgentCoreSellerProxy (keeps the OpenDirect surface, MCP over + # Bearer HTTPS). It exposes the JWT-mode marker set by the selector. + assert getattr(client, "_use_jwt", False), ( + f"OAuth endpoint should route to the JWT client, got {type(client).__name__}" + ) + + +async def test_roundtrip_arn_without_oauth_routes_to_sigv4_proxy(): + """Parity case: an arn endpoint WITHOUT advertised OAuth uses the SigV4 proxy.""" + client = select_seller_client(_SELLER_ARN, authentication=None, region="us-west-2") + # SigV4 proxy: NOT in JWT mode. + assert not getattr(client, "_use_jwt", False), ( + "arn without OAuth should use the SigV4 proxy, not the JWT client" + ) diff --git a/tests/unit/agentcore/test_agentcore_seller_proxy.py b/tests/unit/agentcore/test_agentcore_seller_proxy.py new file mode 100644 index 00000000..4833aeef --- /dev/null +++ b/tests/unit/agentcore/test_agentcore_seller_proxy.py @@ -0,0 +1,183 @@ +"""Unit tests for AgentCoreSellerProxy (offline — boto3 client is stubbed). + +Verifies the proxy presents the OpenDirect surface backed by a mocked +InvokeAgentRuntime returning MCP tools/call envelopes, and that the MCP +result parsing / product mapping helpers behave. +""" + +import json + +import pytest + +from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import ( + AgentCoreSellerProxy, + _parse_mcp_tool_result, +) + +ARN = "arn:aws:bedrock-agentcore:us-west-2:111122223333:runtime/seller_mcp-abc123" + + +def _wire_product(pid: str, name: str = "Apex CTV") -> dict: + return { + "product_id": pid, + "seller_organization_id": "org-seller-1", + "name": name, + "base_price": {"amount_micros": 45_000_000, "currency": "USD"}, + "pricing_type": "floor", + "pricing_model": "cpm", + } + + +def _mcp_envelope(payload: dict) -> bytes: + """Wrap a tool payload as an MCP CallToolResult JSON-RPC response body.""" + return json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [{"type": "text", "text": json.dumps(payload)}], + "isError": False, + }, + } + ).encode("utf-8") + + +class _StreamBody: + def __init__(self, data: bytes): + self._data = data + + def read(self) -> bytes: + return self._data + + +class _FakeBedrockAgentCore: + """Stub boto3 bedrock-agentcore client capturing the last invoke.""" + + def __init__(self, response_payload: dict): + self._payload = response_payload + self.last_kwargs: dict | None = None + + def invoke_agent_runtime(self, **kwargs): + self.last_kwargs = kwargs + return {"response": _StreamBody(_mcp_envelope(self._payload))} + + +# --------------------------------------------------------------------------- # +# Parsing helpers +# --------------------------------------------------------------------------- # + + +class TestParseMcpToolResult: + def test_unwraps_content_text_json(self): + raw = _mcp_envelope({"products": [{"product_id": "inv-ctv-1"}]}) + data = _parse_mcp_tool_result(raw) + assert data == {"products": [{"product_id": "inv-ctv-1"}]} + + def test_raises_on_jsonrpc_error(self): + raw = json.dumps({"jsonrpc": "2.0", "id": 1, "error": {"message": "boom"}}).encode() + with pytest.raises(ValueError, match="seller MCP error"): + _parse_mcp_tool_result(raw) + + def test_raises_on_non_json(self): + with pytest.raises(ValueError, match="not JSON"): + _parse_mcp_tool_result(b"nope") + + def test_unwraps_sse_framed_body(self): + """AgentCore MCP replies as text/event-stream; parse the data: frame.""" + inner = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [{"type": "text", "text": json.dumps({"products": [{"product_id": "inv-ctv-1"}]})}], + "isError": False, + }, + } + ) + sse = f"event: message\ndata: {inner}\n\n".encode("utf-8") + data = _parse_mcp_tool_result(sse) + assert data == {"products": [{"product_id": "inv-ctv-1"}]} + + +# --------------------------------------------------------------------------- # +# OpenDirect surface over the mocked runtime +# --------------------------------------------------------------------------- # + + +class _RecordingBedrockAgentCore: + """Stub that records EVERY invoke (to assert the MCP handshake sequence).""" + + def __init__(self, response_payload: dict): + self._payload = response_payload + self.calls: list[dict] = [] + + def invoke_agent_runtime(self, **kwargs): + self.calls.append(kwargs) + return {"response": _StreamBody(_mcp_envelope(self._payload))} + + +class TestMcpTransportContract: + async def test_sets_required_mcp_accept_header(self): + """Regression: omitting the dual Accept header returns HTTP 406 (-32011).""" + fake = _RecordingBedrockAgentCore({"products": []}) + proxy = AgentCoreSellerProxy(ARN, client=fake) + await proxy.list_products() + for call in fake.calls: + assert call["accept"] == "application/json, text/event-stream" + assert call["contentType"] == "application/json" + assert call["mcpProtocolVersion"] == "2025-03-26" + + async def test_initialize_handshake_precedes_tools_call(self): + fake = _RecordingBedrockAgentCore({"products": []}) + proxy = AgentCoreSellerProxy(ARN, client=fake) + await proxy.list_products() + methods = [json.loads(c["payload"])["method"] for c in fake.calls] + assert methods == ["initialize", "notifications/initialized", "tools/call"] + + async def test_initialize_runs_once_across_calls(self): + fake = _RecordingBedrockAgentCore({"products": []}) + proxy = AgentCoreSellerProxy(ARN, client=fake) + await proxy.list_products() + await proxy.list_products() + methods = [json.loads(c["payload"])["method"] for c in fake.calls] + assert methods.count("initialize") == 1 + assert methods.count("tools/call") == 2 + + +class TestListProducts: + async def test_list_products_maps_wire_products(self): + fake = _FakeBedrockAgentCore( + {"products": [_wire_product("inv-ctv-apex-sports-nba"), _wire_product("inv-lin-gnn", "GNN")]} + ) + proxy = AgentCoreSellerProxy(ARN, region="us-west-2", client=fake) + products = await proxy.list_products(top=10) + assert [p.name for p in products] == ["Apex CTV", "GNN"] + # It invoked the seller by ARN via the MCP list_products tool. + payload = json.loads(fake.last_kwargs["payload"]) + assert fake.last_kwargs["agentRuntimeArn"] == ARN + assert payload["method"] == "tools/call" + assert payload["params"]["name"] == "list_products" + + async def test_get_product(self): + fake = _FakeBedrockAgentCore({"product": _wire_product("inv-ctv-apex-series", "Apex Series")}) + proxy = AgentCoreSellerProxy(ARN, client=fake) + product = await proxy.get_product("inv-ctv-apex-series") + assert product.name == "Apex Series" + assert json.loads(fake.last_kwargs["payload"])["params"]["name"] == "get_product_details" + + async def test_list_products_tolerant_collects_rejects(self): + # Second item is missing required fields -> reject, not fatal. + fake = _FakeBedrockAgentCore( + {"products": [_wire_product("inv-ok"), {"name": "broken — no ids"}]} + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + products, rejects = await proxy.list_products_tolerant(top=10) + assert [p.name for p in products] == ["Apex CTV"] + assert len(rejects) == 1 + assert rejects[0]["name"] == "broken — no ids" + + async def test_session_id_threaded(self): + fake = _FakeBedrockAgentCore({"products": []}) + proxy = AgentCoreSellerProxy(ARN, client=fake, session_id="sess-1") + await proxy.list_products() + assert fake.last_kwargs["runtimeSessionId"] == "sess-1" diff --git a/tests/unit/test_agent_hierarchy.py b/tests/unit/test_agent_hierarchy.py index 3981d459..5029baf8 100644 --- a/tests/unit/test_agent_hierarchy.py +++ b/tests/unit/test_agent_hierarchy.py @@ -43,6 +43,22 @@ _create_research_tools, _format_audience_context, ) +from ad_buyer.llm import _model_accepts_temperature + + +def _expected_temp(tuned: float) -> float | None: + """Temperature an agent's LLM should carry for the configured model. + + Agents pass a tuned temperature to ``build_llm``, but newer Claude families + (e.g. claude-sonnet-5, claude-opus-4-7+) REJECT sampling params, so + ``build_llm`` omits temperature for them and the resulting ``LLM`` reports + ``temperature is None``. These tests must therefore assert the tuned value + only when the *configured* model accepts temperature, and ``None`` when it + does not — otherwise they encode a pre-Bedrock-migration assumption and + fail on the default AgentCore model. Ties the assertion to the same + predicate the production factory uses. + """ + return tuned if _model_accepts_temperature(settings.manager_llm_model) else None # Agent constructors now honor settings.crew_memory_enabled @@ -53,6 +69,25 @@ def _force_memory_enabled(monkeypatch): monkeypatch.setattr(settings, "crew_memory_enabled", True) +# The default model is now Claude Sonnet 5, a tier on which Anthropic rejects +# the `temperature` parameter (build_llm therefore omits it, leaving +# llm.temperature is None). This module's assertions are about per-agent +# temperature *wiring* and the L3 Settings: - """Build an isolated Settings instance that ignores any local .env file.""" + """Build an isolated Settings instance that ignores any local .env file. + + Explicitly nulls the custom-endpoint fields unless a test overrides them, + so an ambient repo ``.env`` (or exported ``*_COMPATIBLE_*`` env vars) can't + bleed in and flip which provider branch ``build_llm`` takes. Without this, + a checkout that has a real ``.env`` pointing at the Bedrock Anthropic + endpoint would make the OpenAI-compatible-branch tests see the anthropic + branch instead. + """ overrides.setdefault("anthropic_api_key", "sk-ant-test") + overrides.setdefault("openai_compatible_llm_api_base_url", None) + overrides.setdefault("openai_compatible_llm_api_key", None) + overrides.setdefault("anthropic_compatible_llm_api_base_url", None) + overrides.setdefault("anthropic_compatible_llm_api_key", None) return Settings(_env_file=None, **overrides) @@ -82,6 +98,58 @@ def test_local_ollama_needs_no_key(self, monkeypatch): assert llm.base_url == "http://localhost:11434/v1" +class TestBedrockAnthropicCompatibleEndpoint: + """ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL configured — routes Claude through + CrewAI's native Anthropic provider (Messages API) against a custom base URL + such as Amazon Bedrock's /anthropic endpoint. This is the path that lets + Claude run on Bedrock WITHOUT the Converse toolUse/toolResult sanitizer.""" + + def test_bedrock_messages_routes_via_anthropic_with_base_url(self, monkeypatch): + monkeypatch.setattr( + "ad_buyer.llm.get_settings", + lambda: _settings( + anthropic_compatible_llm_api_key="bedrock-key-test", + anthropic_compatible_llm_api_base_url="https://bedrock-runtime.us-west-2.amazonaws.com/anthropic", + ), + ) + llm = build_llm( + model="us.anthropic.claude-sonnet-5-v1:0", + temperature=0.3, + max_tokens=4096, + ) + assert llm.provider == "anthropic" + assert llm.base_url == "https://bedrock-runtime.us-west-2.amazonaws.com/anthropic" + assert llm.api_key == "bedrock-key-test" + + def test_anthropic_compatible_takes_precedence_over_openai_compatible(self, monkeypatch): + """When both custom endpoints are set, the Anthropic-compatible branch + wins so Claude never gets misrouted to the OpenAI client (Claude is not + served on Bedrock's OpenAI Chat Completions path).""" + monkeypatch.setattr( + "ad_buyer.llm.get_settings", + lambda: _settings( + anthropic_compatible_llm_api_key="bedrock-key-test", + anthropic_compatible_llm_api_base_url="https://bedrock-runtime.us-west-2.amazonaws.com/anthropic", + openai_compatible_llm_api_key="should-not-win", + openai_compatible_llm_api_base_url="https://integrate.api.nvidia.com/v1", + ), + ) + llm = build_llm(model="us.anthropic.claude-sonnet-5-v1:0", temperature=0.3, max_tokens=4096) + assert llm.provider == "anthropic" + assert llm.base_url == "https://bedrock-runtime.us-west-2.amazonaws.com/anthropic" + + def test_sonnet_5_omits_temperature_on_bedrock_messages(self, monkeypatch): + monkeypatch.setattr( + "ad_buyer.llm.get_settings", + lambda: _settings( + anthropic_compatible_llm_api_base_url="https://bedrock-runtime.us-west-2.amazonaws.com/anthropic", + ), + ) + llm = build_llm(model="us.anthropic.claude-sonnet-5-v1:0", temperature=0.3, max_tokens=4096) + assert llm.provider == "anthropic" + assert llm.temperature is None + + class TestModelAcceptsTemperature: """Anthropic rejects the temperature parameter on Opus 4.7+, Sonnet 5+, and Fable/Mythos with a 400 invalid_request_error; older models still @@ -148,3 +216,59 @@ def test_openai_compatible_path_also_omits_temperature(self, monkeypatch): llm = build_llm(model="claude-opus-4-8", temperature=0.3, max_tokens=4096) assert llm.provider == "openai" assert llm.temperature is None + + +class TestModelUsesMaxCompletionTokens: + """OpenAI's GPT-5.x and o-series reasoning models reject the legacy + ``max_tokens`` parameter (400 invalid_request_error: 'max_tokens' is not + supported with this model) and take ``max_completion_tokens`` instead. + Other models keep ``max_tokens``.""" + + def test_gpt_5_uses_max_completion_tokens(self): + assert _model_uses_max_completion_tokens("us.openai.gpt-5.6-sol") is True + + def test_bare_gpt_5_uses_max_completion_tokens(self): + assert _model_uses_max_completion_tokens("openai/gpt-5") is True + + def test_o1_uses_max_completion_tokens(self): + assert _model_uses_max_completion_tokens("o1") is True + + def test_o3_uses_max_completion_tokens(self): + assert _model_uses_max_completion_tokens("openai/o3-mini") is True + + def test_match_is_case_insensitive(self): + assert _model_uses_max_completion_tokens("US.OpenAI.GPT-5.6-SOL") is True + + def test_gpt_4o_uses_legacy_max_tokens(self): + assert _model_uses_max_completion_tokens("openai/gpt-4o") is False + + def test_claude_uses_legacy_max_tokens(self): + assert _model_uses_max_completion_tokens("us.anthropic.claude-sonnet-5") is False + + +class TestMaxTokensRoutedToCorrectParam: + """build_llm routes the token cap to max_completion_tokens for models that + require it, and to max_tokens for everything else.""" + + def test_gpt_5_routes_to_max_completion_tokens(self, monkeypatch): + monkeypatch.setattr( + "ad_buyer.llm.get_settings", + lambda: _settings( + openai_compatible_llm_api_base_url="https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1", + openai_compatible_llm_api_key="bedrock-key", + ), + ) + llm = build_llm(model="us.openai.gpt-5.6-sol", temperature=0.5, max_tokens=4096) + assert llm.max_completion_tokens == 4096 + assert llm.max_tokens is None + + def test_gpt_4o_keeps_max_tokens(self, monkeypatch): + monkeypatch.setattr( + "ad_buyer.llm.get_settings", + lambda: _settings( + openai_compatible_llm_api_base_url="https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1", + openai_compatible_llm_api_key="bedrock-key", + ), + ) + llm = build_llm(model="gpt-4o", temperature=0.5, max_tokens=4096) + assert llm.max_tokens == 4096 diff --git a/tests/unit/test_oauth_token_provider.py b/tests/unit/test_oauth_token_provider.py new file mode 100644 index 00000000..c290ec19 --- /dev/null +++ b/tests/unit/test_oauth_token_provider.py @@ -0,0 +1,138 @@ +# Donated to IAB Tech Lab + +"""Tests for the buyer OAuthTokenProvider (Req 6.6 / task 7.4).""" + +import logging + +import httpx +import pytest + +from ad_buyer.auth.oauth_token_provider import OAuthTokenError, OAuthTokenProvider + +_TOKEN_URL = "https://issuer.example/oauth2/token" +_SCOPE = "seller-agent/invoke" + + +class _Recorder: + """Counts mint calls and lets a test flip the returned token / status.""" + + def __init__(self): + self.calls = 0 + self.token = "jwt-1" + self.status = 200 + self.expires_in = 3600 + self.last_auth_header = None + self.last_body = None + + def handler(self, request: httpx.Request) -> httpx.Response: + self.calls += 1 + self.last_auth_header = request.headers.get("Authorization") + self.last_body = request.content.decode() + if self.status != 200: + return httpx.Response(self.status, json={"error": "invalid_client"}) + return httpx.Response( + 200, + json={ + "access_token": self.token, + "expires_in": self.expires_in, + "token_type": "Bearer", + }, + ) + + +def _provider(recorder: _Recorder, **kw) -> OAuthTokenProvider: + return OAuthTokenProvider( + client_id="buyer-client", + client_secret="buyer-secret", + transport=httpx.MockTransport(recorder.handler), + **kw, + ) + + +def test_mints_a_client_credentials_token(): + rec = _Recorder() + p = _provider(rec) + assert p.get_token(_TOKEN_URL, _SCOPE) == "jwt-1" + assert rec.calls == 1 + # HTTP Basic client_id:secret, grant_type + scope in the form body. + assert rec.last_auth_header.startswith("Basic ") + assert "grant_type=client_credentials" in rec.last_body + assert "scope=seller-agent" in rec.last_body + + +def test_reuses_cached_token_within_lifetime(): + rec = _Recorder() + p = _provider(rec) + p.get_token(_TOKEN_URL, _SCOPE) + p.get_token(_TOKEN_URL, _SCOPE) # cached — no second mint + assert rec.calls == 1 + + +def test_refreshes_before_expiry(): + """A token whose lifetime is within the skew window is re-minted.""" + rec = _Recorder() + rec.expires_in = 10 # < _EXPIRY_SKEW_SECONDS → always considered stale + p = _provider(rec) + p.get_token(_TOKEN_URL, _SCOPE) + rec.token = "jwt-2" + assert p.get_token(_TOKEN_URL, _SCOPE) == "jwt-2" + assert rec.calls == 2 + + +def test_reactive_reminting_after_invalidate(): + rec = _Recorder() + p = _provider(rec) + assert p.get_token(_TOKEN_URL, _SCOPE) == "jwt-1" + p.invalidate(_TOKEN_URL, _SCOPE) # simulate a downstream 401 + rec.token = "jwt-2" + assert p.get_token(_TOKEN_URL, _SCOPE) == "jwt-2" + assert rec.calls == 2 + + +def test_force_bypasses_cache(): + rec = _Recorder() + p = _provider(rec) + p.get_token(_TOKEN_URL, _SCOPE) + p.get_token(_TOKEN_URL, _SCOPE, force=True) + assert rec.calls == 2 + + +def test_distinct_cache_keys_per_scope(): + rec = _Recorder() + p = _provider(rec) + p.get_token(_TOKEN_URL, "scope-a") + p.get_token(_TOKEN_URL, "scope-b") + assert rec.calls == 2 # different scope → separate mint + + +def test_unconfigured_raises(): + p = OAuthTokenProvider(client_id="", client_secret="") + assert p.configured is False + with pytest.raises(OAuthTokenError): + p.get_token(_TOKEN_URL, _SCOPE) + + +def test_missing_token_endpoint_raises(): + rec = _Recorder() + p = _provider(rec) + with pytest.raises(OAuthTokenError): + p.get_token("", _SCOPE) + + +def test_http_error_raises_and_hides_body(): + rec = _Recorder() + rec.status = 401 + p = _provider(rec) + with pytest.raises(OAuthTokenError) as ei: + p.get_token(_TOKEN_URL, _SCOPE) + assert "401" in str(ei.value) + assert "invalid_client" not in str(ei.value) # body not echoed + + +def test_never_logs_the_secret(caplog): + rec = _Recorder() + p = _provider(rec) + with caplog.at_level(logging.DEBUG): + p.get_token(_TOKEN_URL, _SCOPE) + joined = " ".join(r.getMessage() for r in caplog.records) + assert "buyer-secret" not in joined diff --git a/tests/unit/test_settings_lazy_init.py b/tests/unit/test_settings_lazy_init.py index f9bf2fc1..22471771 100644 --- a/tests/unit/test_settings_lazy_init.py +++ b/tests/unit/test_settings_lazy_init.py @@ -16,13 +16,37 @@ import sys from unittest.mock import patch +import pytest + +_SETTINGS_MOD = "ad_buyer.config.settings" + + +@pytest.fixture(autouse=True) +def _restore_settings_module(): + """Restore the original settings module after each test in this file. + + The tests below ``del sys.modules[...]`` and re-import the settings module + to exercise lazy init with an empty cache. Without restoration the rebuilt + module leaks into ``sys.modules``, so later tests in the full suite that do + ``from ad_buyer.config.settings import settings`` bind to a fresh lazy + proxy reading live env — which made + ``test_real_model_path_e2e::test_label_reflects_active_mode`` fail only in + suite order. Snapshot and restore the original module object here so the + reload stays contained to this file. + """ + original = sys.modules.get(_SETTINGS_MOD) + try: + yield + finally: + if original is not None: + sys.modules[_SETTINGS_MOD] = original + def _reload_settings_module(): """Force a fresh import of the settings module so its lru_cache is empty.""" - mod_name = "ad_buyer.config.settings" - if mod_name in sys.modules: - del sys.modules[mod_name] - return importlib.import_module(mod_name) + if _SETTINGS_MOD in sys.modules: + del sys.modules[_SETTINGS_MOD] + return importlib.import_module(_SETTINGS_MOD) def test_importing_settings_does_not_construct_eagerly(): diff --git a/tests/unit/test_transport_selector.py b/tests/unit/test_transport_selector.py new file mode 100644 index 00000000..e4a5d858 --- /dev/null +++ b/tests/unit/test_transport_selector.py @@ -0,0 +1,142 @@ +# Donated to IAB Tech Lab + +"""Tests for the buyer transport selector (task 4.3 / Req 6.3).""" + +import httpx +import pytest + +from ad_buyer.auth.oauth_token_provider import OAuthTokenProvider +from ad_buyer.clients.mcp_client import IABMCPClient +from ad_buyer.registry.transport_selector import ( + _oauth_from_authentication, + agentcore_invocations_url, + select_seller_client, +) + +_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/seller_mcp-abc" +_HTTPS = "https://seller.example/agent" +_AUTH = { + "type": "oauth2", + "token_endpoint": "https://issuer.example/oauth2/token", + "scope": "seller-agent/invoke", +} + + +def _configured_provider(): + return OAuthTokenProvider(client_id="buyer", client_secret="secret") + + +class TestSelectSellerClient: + def test_arn_routes_to_agentcore_proxy(self): + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import AgentCoreSellerProxy + + client = select_seller_client(_ARN, region="us-west-2") + assert isinstance(client, AgentCoreSellerProxy) + + def test_https_with_oauth_routes_to_jwt_proxy(self): + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import AgentCoreSellerProxy + + client = select_seller_client( + _HTTPS, authentication=_AUTH, token_provider=_configured_provider() + ) + assert isinstance(client, AgentCoreSellerProxy) + assert client._use_jwt is True + assert client._token_endpoint == _AUTH["token_endpoint"] + assert client._scope == "seller-agent/invoke" + + def test_https_without_auth_routes_to_plain_client(self): + client = select_seller_client(_HTTPS) + assert isinstance(client, IABMCPClient) + assert client._auth_provider is None + + def test_https_requires_auth_but_unconfigured_raises(self): + unconfigured = OAuthTokenProvider(client_id="", client_secret="") + with pytest.raises(RuntimeError): + select_seller_client(_HTTPS, authentication=_AUTH, token_provider=unconfigured) + + def test_arn_without_auth_routes_to_sigv4_proxy(self): + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import AgentCoreSellerProxy + + client = select_seller_client(_ARN, region="us-west-2") + assert isinstance(client, AgentCoreSellerProxy) + assert client._use_jwt is False + + def test_arn_with_auth_routes_to_jwt_proxy(self): + """A CUSTOM_JWT runtime given by ARN → JWT-mode proxy (SigV4 disallowed).""" + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import AgentCoreSellerProxy + + client = select_seller_client( + _ARN, authentication=_AUTH, token_provider=_configured_provider() + ) + assert isinstance(client, AgentCoreSellerProxy) + assert client._use_jwt is True + + +class TestAgentCoreInvocationsUrl: + def test_encodes_arn_into_verified_shape(self): + url = agentcore_invocations_url(_ARN) + assert url == ( + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/" + "arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime" + "%2Fseller_mcp-abc/invocations?qualifier=DEFAULT" + ) + + def test_full_https_endpoint_used_as_is(self): + """A discovered full invocations URL routes to a JWT-mode proxy that + posts to that URL directly (no ARN encoding, no /mcp/sse).""" + from ad_buyer.interfaces.agentcore.agentcore_seller_proxy import AgentCoreSellerProxy + + full = ( + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/" + "arn%3Aaws%3A.../invocations?qualifier=DEFAULT" + ) + client = select_seller_client( + full, authentication=_AUTH, token_provider=_configured_provider() + ) + assert isinstance(client, AgentCoreSellerProxy) + # The proxy stores the full URL as its "arn"; JWT invoke uses it directly + # (not ARN-encoded) because it does not start with "arn:". + assert client._arn == full + assert not client._arn.startswith("arn:") + + +class TestOAuthExtraction: + def test_none_returns_empty(self): + assert _oauth_from_authentication(None) == ("", "") + + def test_snake_case(self): + assert _oauth_from_authentication( + {"token_endpoint": "u", "scope": "s"} + ) == ("u", "s") + + def test_camel_case_and_token_url(self): + assert _oauth_from_authentication({"tokenEndpoint": "u"})[0] == "u" + assert _oauth_from_authentication({"token_url": "u"})[0] == "u" + + def test_scopes_list_is_space_joined(self): + assert _oauth_from_authentication( + {"token_endpoint": "u", "scopes": ["a", "b"]} + ) == ("u", "a b") + + +class TestJwtMcpClientAuthHeaders: + def test_bearer_header_minted_from_provider(self): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"access_token": "jwt-xyz", "expires_in": 3600}) + + provider = OAuthTokenProvider( + client_id="buyer", + client_secret="secret", + transport=httpx.MockTransport(handler), + ) + client = IABMCPClient( + base_url=_HTTPS, + auth_token_provider=provider, + token_endpoint=_AUTH["token_endpoint"], + scope="seller-agent/invoke", + ) + assert client._auth_headers() == {"Authorization": "Bearer jwt-xyz"} + + def test_no_provider_no_headers(self): + client = IABMCPClient(base_url=_HTTPS) + assert client._auth_headers() == {}