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/.gitignore b/.gitignore index 53c6d504..c53c86da 100644 --- a/.gitignore +++ b/.gitignore @@ -120,3 +120,10 @@ Dockerfile # Internal issue tracker (not published) .beads/ + +# Kiro specs/steering are managed OUTSIDE the repo (internal assets) — never commit. +.kiro/ + +# Generated CloudFormation package output (account-specific S3 URLs + hashes, +# regenerated every deploy) — never commit. +.packaged-agentcore.yaml diff --git a/docs/api/seller-discovery.md b/docs/api/seller-discovery.md index 9d0dcd7f..e3656852 100644 --- a/docs/api/seller-discovery.md +++ b/docs/api/seller-discovery.md @@ -349,9 +349,65 @@ All errors are logged at `WARNING` or `DEBUG` level. The client never raises `ht !!! warning "Silent Failures" Because the client returns empty results on failure, always check the length of discovery results. An empty list could mean "no sellers match" or "the registry is down." Monitor logs for `WARNING`-level messages to distinguish the two. -## Related +## Connecting to a Seller over OAuth (AgentCore) + +When a discovered seller runs on Amazon Bedrock AgentCore behind a CUSTOM_JWT +authorizer, the registry record's `endpoint_url` is the runtime's OAuth HTTPS +invocations URL (or its ARN) and its `authentication` block advertises the +shared token endpoint + scope. The buyer authenticates with a short-lived +`client_credentials` bearer token — **not** SigV4/boto3, which OAuth runtimes +reject. + +### Minting the token: `OAuthTokenProvider` + +`ad_buyer.auth.oauth_token_provider.OAuthTokenProvider` reads the token endpoint ++ scope from the discovered record, holds the buyer's OWN `client_id`/`secret` +from settings (`BUYER_OAUTH_CLIENT_ID` / `BUYER_OAUTH_CLIENT_SECRET` — never +taken from the registry, never logged), and POSTs `client_credentials` over HTTP +Basic. Tokens are cached by `(token_endpoint, scope)` and refreshed proactively +before `expires_in` and reactively on a 401. + +### Transport selection: `select_seller_client` + +`ad_buyer.registry.transport_selector.select_seller_client(endpoint, authentication=..., region=..., token_provider=...)` +picks the right client from the discovered record: + +| `endpoint_url` | Advertised OAuth? | Client | +|----------------|-------------------|--------| +| `arn:...` | No | `AgentCoreSellerProxy` (SigV4, same-account/dev) | +| `arn:...` | Yes | JWT client — derives the HTTPS invocations URL from the ARN, Bearer over HTTPS | +| `https://...` | Yes | JWT client on that URL | +| `https://...` | No | plain unauthenticated `IABMCPClient` | + +The JWT path returns an `AgentCoreSellerProxy` in **JWT mode**: it preserves the +OpenDirect method surface (`list_products`/`search_products`/`get_product`/ +`check_avails`) that `DealBookingFlow` expects, but frames MCP over a raw HTTPS +POST with a `Authorization: Bearer ` header (with 401 re-mint) instead of +SigV4. + +**Fail-closed:** if the seller advertises OAuth but the buyer has no client +credentials configured, `select_seller_client` raises rather than silently +connecting unauthenticated. + +```python +from ad_buyer.registry import RegistryClient +from ad_buyer.registry.transport_selector import select_seller_client + +registry = RegistryClient(registry_url="https://registry.aamp.iab.com/agent-registry") +record = (await registry.discover_sellers(["ctv"]))[0] -- [Seller Agent Discovery](https://iabtechlab.github.io/seller-agent/api/agent-discovery/) — How seller agents register and expose their agent cards +# The discovered record carries endpoint_url + authentication (token endpoint + scope). +client = select_seller_client( + record.url, + authentication=getattr(record, "authentication", None), +) +# client now speaks the OpenDirect surface; the JWT is minted + refreshed under the hood. +``` + +`crew_tools.py` wires this automatically, routing on the discovered +`SELLER_TOKEN_ENDPOINT`. + +## Related - [Media Kit Discovery](media-kit.md) — Browse seller inventory after discovering sellers - [Authentication](authentication.md) — API key setup for accessing seller endpoints - [A2A Client](a2a-client.md) — Agent-to-agent communication protocol diff --git a/docs/architecture/storage-backends.md b/docs/architecture/storage-backends.md index 9c34cf4a..1316b164 100644 --- a/docs/architecture/storage-backends.md +++ b/docs/architecture/storage-backends.md @@ -48,6 +48,49 @@ checks against the configured database. file (e.g. `DesiredCount: 1` on ECS). Running multiple instances against the same file — including a shared network file system — risks corruption. +## AgentCore deployment: SQLite in-memory by default, durable Postgres opt-in + +When the buyer runs as an AgentCore runtime, `infra/aws/agentcore/deploy.sh` +defaults to `STORAGE_TYPE=sqlite` with an **in-memory** `DATABASE_URL` +(`sqlite:///:memory:`), in PUBLIC network mode. State written by the stores +above therefore does **not** survive a container recycle in that default +deployment. This is the right default for same-account/dev and for the agentic +planning/booking path, where the buyer's own records (deals/orders/negotiations) +are re-derivable from the seller of record and the campaign plan. + +### Durable Postgres (opt-in): `--storage postgres` + +For enterprise/durable deployments the buyer now has a real backend seam and a +`--storage postgres` path (mirrors the seller's group-6 work): + +- **Aurora Serverless v2** (`aurora-postgresql` 16.9, `db.serverless`) deployed + beside Redis in `infra/aws/cloudformation/storage.yaml`, with + `ServerlessV2ScalingConfiguration.MinCapacity: 0` (**scale-to-zero**, bounded + max) so it costs nothing when idle. +- **RDS-managed credentials** — `ManageMasterUserPassword: true` puts the DB + password in an auto-rotating Secrets Manager secret. No plaintext password is + ever in CFN, env, or logs; the runtime reads the secret **by ARN** at startup + (`storage/db_secret.py`) and assembles the connection string in-process. +- **CUSTOMER_VPC mode** — `network-agentcore.yaml` + `main-agentcore.yaml` + provide the VPC interface endpoints (`bedrock-agentcore`/`bedrock-runtime`, + `sts`, `secretsmanager`, ECR, Logs) and the **443 self-ingress** on the + runtime SG required for a private-subnet runtime to reach Aurora and those + endpoints (same reachability lesson as the seller). +- **The backend seam** — `storage/connection_factory.py` selects on + `STORAGE_TYPE`; `storage/pg_connection.py` adapts the sqlite3 API to + psycopg, and `storage/schema_pg.py` derives Postgres DDL from the shared + schema. The 6 connection-owning stores route through the factory; the injected + sub-stores follow for free. `STORAGE_TYPE=hybrid` on the deployed runtime + puts KV state in Aurora with Redis alongside. + +The entrypoint uses `setdefault` semantics so a deploy-supplied `hybrid`/Postgres +value wins while the dev default stays SQLite. As with the seller, product data +is independent of the KV backend. + +If you need durable buyer state and do not want the Aurora path, you can still +run the buyer on ECS with a persistent `DATABASE_URL` (single writer — see the +warning above). + ## Related - [Deal Store](deal-store.md) — full schema and API reference for the primary store 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..b6ecf1f1 100755 --- a/infra/aws/agentcore/deploy.sh +++ b/infra/aws/agentcore/deploy.sh @@ -25,12 +25,43 @@ 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-haiku-4-5-20251001-v1:0}" +# Manager (level-1 orchestrator) runs on a stronger model than the sub-agents. +# Sub-agents (level2/3 research/execution/channel) are the ~165-call fan-out that +# dominates latency, so they default to Haiku 4.5 (fast, temperature-accepting); +# the single manager stays on Sonnet 5 for plan-quality. Override either via +# DEFAULT_LLM_MODEL (sub-agents) / MANAGER_LLM_MODEL (manager) at deploy. +MANAGER_MODEL="${MANAGER_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" +STORAGE_TYPE_ARG="sqlite" DO_TEST=false TEST_ONLY=false DO_CLEANUP=false @@ -41,6 +72,7 @@ while [[ $# -gt 0 ]]; do --region) REGION="$2"; shift 2 ;; --name) AGENT_NAME="$2"; shift 2 ;; --mode) DEPLOY_MODE="$2"; shift 2 ;; + --storage) STORAGE_TYPE_ARG="$2"; shift 2 ;; --profile) AWS_PROFILE="$2"; shift 2 ;; --seller-url) SELLER_AGENT_URL="$2"; shift 2 ;; --test) DO_TEST=true; shift ;; @@ -48,12 +80,18 @@ while [[ $# -gt 0 ]]; do --cleanup) DO_CLEANUP=true; shift ;; --prompt) PROMPT="$2"; shift 2 ;; -h|--help) - echo "Usage: $(basename "$0") [--region REGION] [--name NAME] [--mode MODE] [--profile PROFILE] [--seller-url URL] [--test] [--test-only] [--cleanup] [--prompt JSON]" + echo "Usage: $(basename "$0") [--region REGION] [--name NAME] [--mode MODE] [--storage sqlite|postgres] [--profile PROFILE] [--seller-url URL] [--test] [--test-only] [--cleanup] [--prompt JSON]" exit 0 ;; *) echo "ERROR: Unknown option: $1" >&2; exit 1 ;; esac done +# Validate storage backend +if [[ "${STORAGE_TYPE_ARG}" != "sqlite" && "${STORAGE_TYPE_ARG}" != "postgres" ]]; then + echo "ERROR: --storage must be 'sqlite' (default, in-memory) or 'postgres' (Aurora Serverless v2 + Secrets Manager, CUSTOMER_VPC)" >&2 + exit 1 +fi + # Validate deploy mode (only http supported for now — MCP causes OOM, Issue 19) if [[ "${DEPLOY_MODE}" != "http" ]]; then echo "ERROR: Only --mode http is supported for buyer agent (MCP deferred — Issue 19)" @@ -77,6 +115,156 @@ fi cd "${REPO_ROOT}" +# ============================================================================= +# --storage postgres: deploy Aurora Serverless v2 + Redis + VPC endpoints via +# main-agentcore.yaml, then export the connection facts. Sets (on success): +# VPC_SUBNET_1 / VPC_SUBNET_2 / VPC_SECURITY_GROUP — CUSTOMER_VPC runtime cfg +# DB_SECRET_ARN / AURORA_ENDPOINT / AURORA_PORT / DB_NAME / REDIS_URL — env +# NO plaintext password is ever read or forwarded — the app reads the +# RDS-managed Secrets Manager secret at startup by ARN (Req 12.2/12.7). +# ============================================================================= +VPC_SUBNET_1=""; VPC_SUBNET_2=""; VPC_SECURITY_GROUP="" +DB_SECRET_ARN=""; AURORA_ENDPOINT=""; AURORA_PORT=""; DB_NAME=""; REDIS_URL="" + +deploy_postgres_infrastructure() { + local stack_name="ad-buyer-${ENVIRONMENT:-staging}-agentcore" + local account_id + account_id=$(aws sts get-caller-identity --query Account --output text --region "${REGION}") + local bucket="${TEMPLATE_BUCKET:-ad-buyer-cfn-${account_id}-${REGION}}" + + echo "=============================================" + echo " Deploying buyer durable-storage infrastructure" + echo " Stack : ${stack_name}" + echo "=============================================" + + if ! aws s3 ls "s3://${bucket}" --region "${REGION}" 2>/dev/null; then + echo ">>> Creating template bucket: ${bucket}" + aws s3api create-bucket --bucket "${bucket}" --region "${REGION}" \ + --create-bucket-configuration LocationConstraint="${REGION}" 2>/dev/null \ + || aws s3 mb "s3://${bucket}" --region "${REGION}" + fi + + # Discover the private route table (network.yaml doesn't export it) for the + # S3 gateway endpoint; empty is fine (the endpoint is conditional). + local private_rt="" + private_rt=$(aws ec2 describe-route-tables --region "${REGION}" \ + --filters "Name=tag:Project,Values=ad-buyer-system" \ + --query "RouteTables[?Associations[?Main!=\`true\`]].RouteTableId | [0]" \ + --output text 2>/dev/null || true) + [[ "${private_rt}" == "None" ]] && private_rt="" + + echo ">>> Packaging templates (uploading nested stacks to S3)..." + local packaged="${REPO_ROOT}/.packaged-agentcore.yaml" + aws cloudformation package \ + --template-file "${SCRIPT_DIR}/main-agentcore.yaml" \ + --s3-bucket "${bucket}" \ + --s3-prefix "ad-buyer-system/agentcore" \ + --output-template-file "${packaged}" \ + --region "${REGION}" + + local overrides=("Environment=${ENVIRONMENT:-staging}" "TemplatesBucketName=${bucket}") + [[ -n "${private_rt}" ]] && overrides+=("PrivateRouteTableId=${private_rt}") + + echo ">>> Deploying stack: ${stack_name}" + aws cloudformation deploy \ + --template-file "${packaged}" \ + --stack-name "${stack_name}" \ + --parameter-overrides "${overrides[@]}" \ + --capabilities CAPABILITY_IAM \ + --region "${REGION}" \ + --no-fail-on-empty-changeset + + echo ">>> Reading stack outputs..." + local outputs + outputs=$(aws cloudformation describe-stacks --stack-name "${stack_name}" \ + --region "${REGION}" --query "Stacks[0].Outputs" --output json) + + _out() { echo "${outputs}" | python3 -c "import sys,json; print(next((o['OutputValue'] for o in json.load(sys.stdin) if o['OutputKey']=='$1'),''))"; } + VPC_SECURITY_GROUP=$(_out AgentCoreSecurityGroupId) + VPC_SUBNET_1=$(_out PrivateSubnet1Id) + VPC_SUBNET_2=$(_out PrivateSubnet2Id) + DB_SECRET_ARN=$(_out AuroraSecretArn) + AURORA_ENDPOINT=$(_out AuroraEndpoint) + AURORA_PORT=$(_out AuroraPort) + DB_NAME=$(_out AuroraDatabaseName) + local redis_ep redis_port + redis_ep=$(_out RedisEndpoint); redis_port=$(_out RedisPort) + [[ -n "${redis_ep}" ]] && REDIS_URL="redis://${redis_ep}:${redis_port}/0" + + echo " SG : ${VPC_SECURITY_GROUP}" + echo " Subnets : ${VPC_SUBNET_1}, ${VPC_SUBNET_2}" + echo " DB secret: ${DB_SECRET_ARN}" + echo " Aurora : ${AURORA_ENDPOINT}:${AURORA_PORT}/${DB_NAME}" + echo "✅ Durable-storage infrastructure deployed" +} + +# Grant the runtime execution role read access to the RDS-managed DB secret +# (+ KMS decrypt). Best-effort + idempotent; a separately-named inline policy +# the toolkit's execution-policy rewrite does not clobber. (Req 12.4) +grant_secret_access() { + local role_arn="$1" secret_arn="$2" + [[ -z "${role_arn}" || -z "${secret_arn}" ]] && return 0 + local role_name="${role_arn##*/}" + echo ">>> Granting ${role_name} secretsmanager:GetSecretValue on the DB secret" + aws iam put-role-policy \ + --role-name "${role_name}" \ + --policy-name "BuyerDbSecretAccess" \ + --policy-document "{\"Version\":\"2012-10-17\",\"Statement\":[{\"Effect\":\"Allow\",\"Action\":[\"secretsmanager:GetSecretValue\"],\"Resource\":\"${secret_arn}\"},{\"Effect\":\"Allow\",\"Action\":[\"kms:Decrypt\"],\"Resource\":\"*\",\"Condition\":{\"StringEquals\":{\"kms:ViaService\":\"secretsmanager.${REGION}.amazonaws.com\"}}}]}" \ + --region "${REGION}" 2>&1 || echo " ⚠️ secret grant failed (non-fatal)" +} + +# Resolve the execution role ARN for THIS agent's deployed runtime. Prefers an +# explicit EXECUTION_ROLE_ARN, then the live runtime's actual roleArn from AWS, +# then the agent-specific block in .bedrock_agentcore.yaml. Must resolve the +# role for AGENT_NAME specifically — the yaml holds every agent's config, so a +# bare `grep execution_role | head -1` picks the wrong (first) agent's role. +_resolve_runtime_role_arn() { + local role_arn="${EXECUTION_ROLE_ARN:-}" + if [[ -z "${role_arn}" ]]; then + local rt_id + rt_id=$(aws bedrock-agentcore-control list-agent-runtimes \ + --region "${REGION}" ${AWS_PROFILE:+--profile "${AWS_PROFILE}"} \ + --query "agentRuntimes[?agentRuntimeName=='${AGENT_NAME}'].agentRuntimeId" \ + --output text 2>/dev/null | head -1) + if [[ -n "${rt_id}" ]]; then + role_arn=$(aws bedrock-agentcore-control get-agent-runtime \ + --region "${REGION}" ${AWS_PROFILE:+--profile "${AWS_PROFILE}"} \ + --agent-runtime-id "${rt_id}" --query "roleArn" --output text 2>/dev/null) + fi + fi + if [[ -z "${role_arn}" || "${role_arn}" == "None" ]] && [[ -f .bedrock_agentcore.yaml ]]; then + role_arn=$(awk -v a="${AGENT_NAME}:" ' + $1==a {found=1} + found && /execution_role:/ && /SDKRuntime/ {print $2; exit} + ' .bedrock_agentcore.yaml) + fi + [[ "${role_arn}" == "None" ]] && role_arn="" + printf '%s' "${role_arn}" +} + +# Grant bedrock:CallWithBearerToken to the runtime's execution role (Req 11). +# The runtime mints its own short-lived Bedrock bearer token from this role at +# startup (ad_buyer.llm.bedrock_token) for the Anthropic Messages path. The +# toolkit-created execution role does NOT get this by default (its managed +# policy only has InvokeModel for the SigV4/Converse path), so a runtime on the +# auto-created role 403s on 'bedrock:CallWithBearerToken'. Attached as a +# SEPARATELY-NAMED inline policy so the toolkit's per-deploy execution-policy +# rewrite does not clobber it. Best-effort + idempotent. Mirrors the seller's +# _grant_bedrock_bearer_token_permission so both repos self-sustain identically. +grant_bedrock_bearer_token() { + local role_arn="$1" + [[ -z "${role_arn}" ]] && { echo " ⚠️ Could not resolve runtime role; skipping bedrock:CallWithBearerToken grant." >&2; return 0; } + local role_name="${role_arn##*/}" + echo ">>> Granting ${role_name} bedrock:CallWithBearerToken (self-sustaining Bedrock token)" + aws iam put-role-policy \ + --role-name "${role_name}" \ + --policy-name "BedrockCallWithBearerToken" \ + --policy-document '{"Version":"2012-10-17","Statement":[{"Sid":"BedrockBearerTokenMessagesPath","Effect":"Allow","Action":"bedrock:CallWithBearerToken","Resource":"*"}]}' \ + --region "${REGION}" ${AWS_PROFILE:+--profile "${AWS_PROFILE}"} 2>&1 \ + && echo " ✅ bedrock:CallWithBearerToken granted to ${role_name}" \ + || echo " ⚠️ Could not grant bedrock:CallWithBearerToken to ${role_name} (check IAM perms); runtime token mint may 403." >&2 +} + # ── Cleanup (--cleanup) ──────────────────────────────────────────── if [[ "${DO_CLEANUP}" == "true" ]]; then echo "=============================================" @@ -143,36 +331,126 @@ 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 + + # ── Durable storage (--storage postgres): deploy Aurora+Redis+VPC endpoints, + # then run the runtime in CUSTOMER_VPC mode. Default stays sqlite in-memory. + if [[ "${STORAGE_TYPE_ARG}" == "postgres" ]]; then + deploy_postgres_infrastructure + fi + # Configure echo "" echo ">>> Configuring agent..." + _vpc_args=() + if [[ "${STORAGE_TYPE_ARG}" == "postgres" && -n "${VPC_SECURITY_GROUP}" ]]; then + _vpc_args=(--vpc --subnets "${VPC_SUBNET_1},${VPC_SUBNET_2}" --security-groups "${VPC_SECURITY_GROUP}") + echo " VPC mode: SG=${VPC_SECURITY_GROUP}, Subnets=${VPC_SUBNET_1},${VPC_SUBNET_2}" + fi agentcore configure \ -e src/ad_buyer/interfaces/agentcore/http_main.py \ -n "${AGENT_NAME}" \ -rf infra/aws/agentcore/requirements.txt \ -p HTTP \ -r "${REGION}" \ + ${EXECUTION_ROLE_ARN:+--execution-role "${EXECUTION_ROLE_ARN}"} \ + "${_vpc_args[@]+"${_vpc_args[@]}"}" \ --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}") + # Storage env: default sqlite in-memory (dev); postgres → hybrid with the + # RDS-managed secret ARN (NOT the password) + Aurora/Redis connection facts. + # The app reads the secret from Secrets Manager at startup (Req 12.7). + _storage_env=() + if [[ "${STORAGE_TYPE_ARG}" == "postgres" ]]; then + _storage_env=( + --env "STORAGE_TYPE=hybrid" + --env "DB_SECRET_ARN=${DB_SECRET_ARN}" + --env "AURORA_ENDPOINT=${AURORA_ENDPOINT}" + --env "AURORA_PORT=${AURORA_PORT}" + --env "DB_NAME=${DB_NAME}" + --env "REDIS_URL=${REDIS_URL}" + ) + # Grant the runtime role read access to the DB secret now that configure + # has created/resolved it. Reuse the shared resolver (finds THIS agent's + # role specifically, not the first agent's in the shared yaml). + grant_secret_access "$(_resolve_runtime_role_arn)" "${DB_SECRET_ARN}" + else + _storage_env=( + --env "STORAGE_TYPE=sqlite" + --env "DATABASE_URL=sqlite:///:memory:" + ) + fi agentcore deploy \ --env "DEFAULT_LLM_MODEL=${LLM_MODEL}" \ - --env "MANAGER_LLM_MODEL=${LLM_MODEL}" \ - --env "STORAGE_TYPE=sqlite" \ - --env "DATABASE_URL=sqlite:///:memory:" \ + --env "MANAGER_LLM_MODEL=${MANAGER_MODEL}" \ + --env "PYTHONPATH=/app/src" \ + --env "ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL=${ANTHROPIC_BASE_URL}" \ + "${_key_env[@]+"${_key_env[@]}"}" \ + "${_oauth_env[@]+"${_oauth_env[@]}"}" \ + "${_storage_env[@]+"${_storage_env[@]}"}" \ --env "ANTHROPIC_API_KEY=not-used-with-bedrock" \ --env "SELLER_AGENT_URL=${SELLER_AGENT_URL}" \ --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 "" echo "✅ Deploy complete" + + # Req 11: grant bedrock:CallWithBearerToken to whatever role the runtime + # actually assumed (CFN-provided or toolkit auto-created). Unconditional — + # the runtime ALWAYS mints its own Bedrock token for the Messages path, so + # this must run on every deploy, not only when a seller ARN is wired. The + # role only exists after `agentcore deploy`, hence the grant happens here. + grant_bedrock_bearer_token "$(_resolve_runtime_role_arn)" fi # ── Test (--test or --test-only) ──────────────────────────────────── diff --git a/infra/aws/agentcore/main-agentcore.yaml b/infra/aws/agentcore/main-agentcore.yaml new file mode 100644 index 00000000..1695b209 --- /dev/null +++ b/infra/aws/agentcore/main-agentcore.yaml @@ -0,0 +1,183 @@ +AWSTemplateFormatVersion: "2010-09-09" +Description: > + Ad Buyer System — AgentCore deployment root stack. + Orchestrates the shared Network + Storage stacks and adds AgentCore-specific + security group and VPC endpoints. Does NOT deploy ECS compute. + + Req 12: the buyer's shared network.yaml has no Aurora security group (it + predates durable Postgres), so THIS orchestrator creates the Aurora SG itself + (Option 1) and passes it to the storage + agentcore-network stacks — + network.yaml stays untouched, keeping the ECS path's blast radius at zero. + +# ============================================================================= +# Parameters +# ============================================================================= +Parameters: + Environment: + Type: String + Default: staging + AllowedValues: [staging, production, development] + Description: Deployment environment + TemplatesBucketName: + Type: String + Description: S3 bucket containing nested CloudFormation templates + TemplatesPrefix: + Type: String + Default: ad-buyer-system/cloudformation + Description: S3 key prefix for the shared (network/storage) nested templates + VpcCidr: + Type: String + Default: "10.20.0.0/16" + Description: Must match network.yaml's VpcCidr (its subnet CIDRs are in this range) + PrivateRouteTableId: + Type: String + Default: "" + Description: > + Private route table ID for the S3 gateway endpoint. network.yaml does not + export this, so pass it from the deploy script or leave empty to skip the + S3 gateway endpoint. + MinACU: + Type: Number + Default: 0 + Description: Minimum Aurora Serverless v2 ACU (0 = scale-to-zero) + MaxACU: + Type: Number + Default: 4 + Description: Maximum Aurora Serverless v2 ACU + +# ============================================================================= +# Conditions +# ============================================================================= +Conditions: + HasPrivateRouteTable: !Not [!Equals [!Ref PrivateRouteTableId, ""]] + +# ============================================================================= +# Resources +# ============================================================================= +Resources: + + # --------------------------------------------------------------------------- + # Network stack — VPC, subnets, ALB/ECS/EFS/Redis SGs (shared with ECS path) + # --------------------------------------------------------------------------- + NetworkStack: + Type: AWS::CloudFormation::Stack + Properties: + TemplateURL: ../cloudformation/network.yaml + Parameters: + Environment: !Ref Environment + VpcCidr: !Ref VpcCidr + Tags: + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + + # --------------------------------------------------------------------------- + # Aurora security group (Option 1 — created HERE, not in the shared + # network.yaml, so the ECS path is untouched). Ingress from the AgentCore SG + # is added by the agentcore-network stack. + # --------------------------------------------------------------------------- + DatabaseSecurityGroup: + Type: AWS::EC2::SecurityGroup + Properties: + GroupDescription: !Sub "ad-buyer-${Environment} Aurora Serverless v2" + VpcId: !GetAtt NetworkStack.Outputs.VpcId + Tags: + - Key: Name + Value: !Sub "ad-buyer-${Environment}-aurora-sg" + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + + # --------------------------------------------------------------------------- + # Storage stack — Aurora Serverless v2 (MinACU=0) + Redis + # --------------------------------------------------------------------------- + StorageStack: + Type: AWS::CloudFormation::Stack + DependsOn: NetworkStack + Properties: + TemplateURL: ../cloudformation/storage.yaml + Parameters: + Environment: !Ref Environment + VpcId: !GetAtt NetworkStack.Outputs.VpcId + PrivateSubnet1Id: !GetAtt NetworkStack.Outputs.PrivateSubnet1Id + PrivateSubnet2Id: !GetAtt NetworkStack.Outputs.PrivateSubnet2Id + RedisSecurityGroupId: !GetAtt NetworkStack.Outputs.RedisSecurityGroupId + DatabaseSecurityGroupId: !Ref DatabaseSecurityGroup + MinACU: !Ref MinACU + MaxACU: !Ref MaxACU + Tags: + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + + # --------------------------------------------------------------------------- + # AgentCore network — runtime SG, 443 self-ingress, VPC endpoints, and the + # Aurora/Redis ingress rules from the runtime SG. + # --------------------------------------------------------------------------- + AgentCoreNetworkStack: + Type: AWS::CloudFormation::Stack + DependsOn: + - NetworkStack + - StorageStack + Properties: + TemplateURL: ./network-agentcore.yaml + Parameters: + Environment: !Ref Environment + VpcId: !GetAtt NetworkStack.Outputs.VpcId + PrivateSubnet1Id: !GetAtt NetworkStack.Outputs.PrivateSubnet1Id + PrivateSubnet2Id: !GetAtt NetworkStack.Outputs.PrivateSubnet2Id + PrivateRouteTableId: !If [HasPrivateRouteTable, !Ref PrivateRouteTableId, ""] + DatabaseSecurityGroupId: !Ref DatabaseSecurityGroup + RedisSecurityGroupId: !GetAtt NetworkStack.Outputs.RedisSecurityGroupId + Tags: + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + +# ============================================================================= +# Outputs +# ============================================================================= +Outputs: + AuroraSecretArn: + Description: RDS-managed Secrets Manager secret ARN (Aurora master credential) + Value: !GetAtt StorageStack.Outputs.AuroraSecretArn + + AuroraEndpoint: + Description: Aurora writer endpoint + Value: !GetAtt StorageStack.Outputs.AuroraEndpoint + + AuroraPort: + Description: Aurora port + Value: !GetAtt StorageStack.Outputs.AuroraPort + + AuroraDatabaseName: + Description: Aurora database name + Value: !GetAtt StorageStack.Outputs.AuroraDatabaseName + + RedisEndpoint: + Description: Redis primary endpoint + Value: !GetAtt StorageStack.Outputs.RedisEndpoint + + RedisPort: + Description: Redis primary port + Value: !GetAtt StorageStack.Outputs.RedisPort + + AgentCoreSecurityGroupId: + Description: Security group for AgentCore runtime ENIs + Value: !GetAtt AgentCoreNetworkStack.Outputs.AgentCoreSecurityGroupId + + PrivateSubnet1Id: + Description: Private subnet 1 for AgentCore VPC config + Value: !GetAtt NetworkStack.Outputs.PrivateSubnet1Id + + PrivateSubnet2Id: + Description: Private subnet 2 for AgentCore VPC config + Value: !GetAtt NetworkStack.Outputs.PrivateSubnet2Id + + VpcId: + Description: VPC ID + Value: !GetAtt NetworkStack.Outputs.VpcId diff --git a/infra/aws/agentcore/network-agentcore.yaml b/infra/aws/agentcore/network-agentcore.yaml new file mode 100644 index 00000000..abf41cec --- /dev/null +++ b/infra/aws/agentcore/network-agentcore.yaml @@ -0,0 +1,241 @@ +AWSTemplateFormatVersion: "2010-09-09" +Description: > + AgentCore-specific security group and VPC endpoints (buyer). + Adds ingress rules to existing Aurora/Redis security groups + without modifying network.yaml. Mirror of the seller's + network-agentcore.yaml (Req 12.3 / group 6 parity). + +# ============================================================================= +# Parameters +# ============================================================================= +Parameters: + Environment: + Type: String + Default: staging + AllowedValues: [staging, production] + VpcId: + Type: AWS::EC2::VPC::Id + Description: VPC to deploy into + PrivateSubnet1Id: + Type: AWS::EC2::Subnet::Id + Description: First private subnet for VPC endpoints + PrivateSubnet2Id: + Type: AWS::EC2::Subnet::Id + Description: Second private subnet for VPC endpoints + PrivateRouteTableId: + Type: String + Default: "" + Description: Private route table for S3 gateway endpoint (optional) + DatabaseSecurityGroupId: + Type: AWS::EC2::SecurityGroup::Id + Description: Existing Aurora security group from network stack + RedisSecurityGroupId: + Type: AWS::EC2::SecurityGroup::Id + Description: Existing Redis security group from network stack + +# ============================================================================= +# Conditions +# ============================================================================= +Conditions: + HasPrivateRouteTable: !Not [!Equals [!Ref PrivateRouteTableId, ""]] + +# ============================================================================= +# Resources +# ============================================================================= +Resources: + + # --------------------------------------------------------------------------- + # Security group for AgentCore container ENIs (CUSTOMER_VPC mode) + # --------------------------------------------------------------------------- + AgentCoreSecurityGroup: + Type: AWS::EC2::SecurityGroup + Properties: + GroupDescription: Security group for AgentCore runtime ENIs in CUSTOMER_VPC mode + VpcId: !Ref VpcId + SecurityGroupEgress: + - IpProtocol: tcp + FromPort: 5432 + ToPort: 5432 + DestinationSecurityGroupId: !Ref DatabaseSecurityGroupId + Description: Aurora PostgreSQL + - IpProtocol: tcp + FromPort: 6379 + ToPort: 6379 + DestinationSecurityGroupId: !Ref RedisSecurityGroupId + Description: ElastiCache Redis + - IpProtocol: tcp + FromPort: 443 + ToPort: 443 + CidrIp: "0.0.0.0/0" + Description: HTTPS for VPC endpoints and AWS APIs + Tags: + - Key: Name + Value: !Sub "ad-buyer-${Environment}-agentcore-sg" + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + + # --------------------------------------------------------------------------- + # Ingress rules on EXISTING security groups (no modification to network.yaml) + # --------------------------------------------------------------------------- + # Self-referencing 443 ingress: the interface VPC endpoints below share THIS + # security group, so their ENIs must accept HTTPS from the runtime ENIs (also + # in this SG). Without it the runtime's 443 calls to bedrock/sts/secrets/logs + # endpoints hang — the container never becomes healthy and emits no logs + # (it cannot even reach the CloudWatch Logs endpoint). This is the required + # companion to the 443 egress rule above. + AgentCoreHttpsSelfIngress: + Type: AWS::EC2::SecurityGroupIngress + Properties: + GroupId: !Ref AgentCoreSecurityGroup + IpProtocol: tcp + FromPort: 443 + ToPort: 443 + SourceSecurityGroupId: !Ref AgentCoreSecurityGroup + Description: HTTPS between runtime ENIs and the shared-SG interface endpoints + + AuroraIngressFromAgentCore: + Type: AWS::EC2::SecurityGroupIngress + Properties: + GroupId: !Ref DatabaseSecurityGroupId + IpProtocol: tcp + FromPort: 5432 + ToPort: 5432 + SourceSecurityGroupId: !Ref AgentCoreSecurityGroup + Description: Allow AgentCore ENIs to reach Aurora + + RedisIngressFromAgentCore: + Type: AWS::EC2::SecurityGroupIngress + Properties: + GroupId: !Ref RedisSecurityGroupId + IpProtocol: tcp + FromPort: 6379 + ToPort: 6379 + SourceSecurityGroupId: !Ref AgentCoreSecurityGroup + Description: Allow AgentCore ENIs to reach Redis + + # --------------------------------------------------------------------------- + # VPC Endpoints — required for AgentCore container image pulls and logging + # --------------------------------------------------------------------------- + ECRDkrEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.ecr.dkr" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + ECRApiEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.ecr.api" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + S3GatewayEndpoint: + Type: AWS::EC2::VPCEndpoint + Condition: HasPrivateRouteTable + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.s3" + VpcEndpointType: Gateway + RouteTableIds: + - !Ref PrivateRouteTableId + + CloudWatchLogsEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.logs" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + # --------------------------------------------------------------------------- + # Endpoints a private-subnet CUSTOMER_VPC runtime ALSO needs (Req 12.3). + # Without these a runtime in a private subnet with no NAT egress cannot: + # - bedrock-agentcore : invoke/serve the AgentCore data plane + # - bedrock-runtime : mint the Bedrock bearer token + call the model + # - sts : resolve/verify its execution-role identity + # - secretsmanager : read the RDS-managed Aurora secret at startup + # A missing one fails on a PRECEDING AWS call before the Aurora connection is + # even attempted. Interface endpoints give the private subnet a path to each + # service without a NAT gateway. + # --------------------------------------------------------------------------- + BedrockAgentCoreEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.bedrock-agentcore" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + BedrockRuntimeEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.bedrock-runtime" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + StsEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.sts" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + + SecretsManagerEndpoint: + Type: AWS::EC2::VPCEndpoint + Properties: + VpcId: !Ref VpcId + ServiceName: !Sub "com.amazonaws.${AWS::Region}.secretsmanager" + VpcEndpointType: Interface + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + SecurityGroupIds: + - !Ref AgentCoreSecurityGroup + PrivateDnsEnabled: true + +# ============================================================================= +# Outputs +# ============================================================================= +Outputs: + AgentCoreSecurityGroupId: + Description: AgentCore runtime ENI security group + Value: !Ref AgentCoreSecurityGroup + Export: + Name: !Sub "${AWS::StackName}-AgentCoreSecurityGroupId" diff --git a/infra/aws/agentcore/requirements.txt b/infra/aws/agentcore/requirements.txt index 439b6b8f..5b02f03f 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 @@ -13,5 +28,19 @@ uvicorn>=0.32.0 rich>=13.0.0 sqlalchemy>=2.0.0 aiosqlite>=0.20.0 +# Durable Postgres backend (Scope 2 / Req 12.6-12.7). The PgConnection adapter +# imports psycopg (synchronous v3); db_secret imports boto3 to fetch the +# RDS-managed Secrets Manager credential at startup. Only exercised on the +# --storage postgres path; harmless on the default SQLite path. +psycopg[binary]>=3.1.0 +boto3>=1.34.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/infra/aws/cloudformation/storage.yaml b/infra/aws/cloudformation/storage.yaml index 64e3bec2..c368640d 100644 --- a/infra/aws/cloudformation/storage.yaml +++ b/infra/aws/cloudformation/storage.yaml @@ -37,6 +37,25 @@ Parameters: - cache.t3.medium - cache.r6g.large + # --- Aurora Serverless v2 (durable Postgres, Req 12) --- + DatabaseSecurityGroupId: + Type: AWS::EC2::SecurityGroup::Id + Description: Security group for the Aurora cluster + DBMasterUsername: + Type: String + Default: buyer + Description: Aurora master username (password is RDS-managed in Secrets Manager) + MinACU: + Type: Number + Default: 0 + Description: > + Minimum Aurora Serverless v2 capacity units. 0 = scale-to-zero (auto-pause + after inactivity, cold-resume on connect). Requires Aurora PG >= 16.3. + MaxACU: + Type: Number + Default: 4 + Description: Maximum Aurora Serverless v2 capacity units + # ============================================================================= # Resources # ============================================================================= @@ -75,6 +94,61 @@ Resources: - Key: Environment Value: !Ref Environment + # --------------------------------------------------------------------------- + # Aurora Serverless v2 PostgreSQL (durable buyer store, Req 12) + # Password is RDS-managed in Secrets Manager (ManageMasterUserPassword) and + # auto-rotated — no plaintext password in CFN params, SSM, or runtime env. + # --------------------------------------------------------------------------- + DBSubnetGroup: + Type: AWS::RDS::DBSubnetGroup + Properties: + DBSubnetGroupDescription: !Sub "ad-buyer-${Environment} Aurora subnets" + SubnetIds: + - !Ref PrivateSubnet1Id + - !Ref PrivateSubnet2Id + + AuroraCluster: + Type: AWS::RDS::DBCluster + DeletionPolicy: Snapshot + UpdateReplacePolicy: Snapshot + Properties: + Engine: aurora-postgresql + EngineVersion: "16.9" + DatabaseName: ad_buyer + MasterUsername: !Ref DBMasterUsername + # RDS creates + auto-rotates the master credential in Secrets Manager. + ManageMasterUserPassword: true + DBSubnetGroupName: !Ref DBSubnetGroup + VpcSecurityGroupIds: + - !Ref DatabaseSecurityGroupId + ServerlessV2ScalingConfiguration: + MinCapacity: !Ref MinACU + MaxCapacity: !Ref MaxACU + StorageEncrypted: true + BackupRetentionPeriod: 7 + EnableCloudwatchLogsExports: + - postgresql + Tags: + - Key: Name + Value: !Sub "ad-buyer-${Environment}-aurora" + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + + AuroraInstance1: + Type: AWS::RDS::DBInstance + Properties: + DBClusterIdentifier: !Ref AuroraCluster + DBInstanceClass: db.serverless + Engine: aurora-postgresql + PubliclyAccessible: false + Tags: + - Key: Project + Value: ad-buyer-system + - Key: Environment + Value: !Ref Environment + # ============================================================================= # Outputs # ============================================================================= @@ -90,3 +164,27 @@ Outputs: Value: !GetAtt RedisReplicationGroup.PrimaryEndPoint.Port Export: Name: !Sub "${AWS::StackName}-RedisPort" + + AuroraSecretArn: + Description: ARN of the RDS-managed Secrets Manager secret (master credential) + Value: !GetAtt AuroraCluster.MasterUserSecret.SecretArn + Export: + Name: !Sub "${AWS::StackName}-AuroraSecretArn" + + AuroraEndpoint: + Description: Aurora cluster writer endpoint + Value: !GetAtt AuroraCluster.Endpoint.Address + Export: + Name: !Sub "${AWS::StackName}-AuroraEndpoint" + + AuroraPort: + Description: Aurora cluster port + Value: !GetAtt AuroraCluster.Endpoint.Port + Export: + Name: !Sub "${AWS::StackName}-AuroraPort" + + AuroraDatabaseName: + Description: Aurora database name + Value: ad_buyer + Export: + Name: !Sub "${AWS::StackName}-AuroraDatabaseName" 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/pyproject.toml b/pyproject.toml index 5eb0fd4d..f332bd3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,11 +44,16 @@ redis = [ "redis>=5.0.0", ] postgres = [ - "asyncpg>=0.29.0", + # Synchronous psycopg (v3) — the PgConnection adapter presents the sqlite3 + # API over psycopg for the AgentCore --storage postgres path. boto3 fetches + # the RDS-managed secret at startup (db_secret.py). + "psycopg[binary]>=3.1.0", + "boto3>=1.34.0", ] production = [ "redis>=5.0.0", - "asyncpg>=0.29.0", + "psycopg[binary]>=3.1.0", + "boto3>=1.34.0", ] [project.scripts] 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..d0fac1cb --- /dev/null +++ b/src/ad_buyer/interfaces/agentcore/agentcore_seller_proxy.py @@ -0,0 +1,568 @@ +# 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``, +``get_pricing``) 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 + +# Reuse the REST/ECS client's client-side catalog filter so the proxy applies +# the SAME reconciled vocabulary (adFormat normalization, channel intentionally +# ignored) rather than a hand-rolled copy that silently diverges. Imported as a +# private symbol from the OpenDirect client by design — the two are peer +# client-transport modules and this is the single filter implementation. +from ...clients.opendirect_client import _filter_wire_products as filter_wire_products +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_GET_PRICING = "get_pricing" +_TOOL_CHECK_AVAILS = "check_avails" + + +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}) + raw_items = data.get("products", []) if isinstance(data, dict) else [] + kept = _filter_raw_items(raw_items, filters) + products = _map_products(kept) + 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 [] + kept = _filter_raw_items(raw_items, filters) + products: list[Product] = [] + rejects: list[dict[str, Any]] = [] + for item in kept: + 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), + } + ) + 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}) + raw_items = data.get("products", []) if isinstance(data, dict) else [] + return _map_products(_filter_raw_items(raw_items, filters or {})) + + async def check_avails(self, request: AvailsRequest) -> AvailsResponse: + # Prefer the real seller ``check_avails`` MCP tool (honest availability: + # available_impressions, estimated_cpm, total_cost via the seller's + # catalog avails policy — the same the REST /products/avails route uses). + # A seller runtime that predates that tool returns an error result; fall + # back to ``get_pricing`` (price only; requested volume echoed as + # available) so an older seller still yields a contract-valid response + # rather than raising on the missing tool. + product_id = getattr(request, "product_id", None) or getattr(request, "productId", None) + requested = getattr(request, "requested_impressions", None) or 0 + budget = getattr(request, "budget", None) or 0 + + data = await self._call_tool( + _TOOL_CHECK_AVAILS, + { + "product_id": product_id, + "requested_impressions": int(requested), + "budget": float(budget), + } + if product_id + else {}, + ) + if _is_real_avails(data): + return AvailsResponse.model_validate(data) + + # Fallback: older seller MCP without check_avails -> use get_pricing. + data = await self._call_tool( + _TOOL_GET_PRICING, + {"product_id": product_id, "volume": int(requested)} if product_id else {}, + ) + return AvailsResponse.model_validate( + _avails_payload(data, product_id=product_id, requested_impressions=int(requested)) + ) + + 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 [] + return _map_products(items) + + +def _map_products(items: list[Any]) -> list[Product]: + """Map a list of raw seller-MCP product dicts to OpenDirect ``Product``s. + + Unmappable records are logged and skipped (non-tolerant path); the tolerant + caller (``list_products_tolerant``) maps items itself to collect rejects. + """ + 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 _filter_raw_items( + raw_items: list[Any], filters: dict[str, Any] +) -> list[dict[str, Any]]: + """Filter raw seller-MCP product dicts using the SHARED catalog filter. + + The seller MCP emits ``iab_agentic_primitives`` wire ``Product`` dicts. We + validate each to a wire ``Product`` and delegate to the OpenDirect client's + ``_filter_wire_products`` — the SAME reconciled-vocabulary filter the + REST/ECS ``OpenDirectClient`` uses — then return the surviving raw dicts so + the caller's existing per-item mapping/reject-collection is unchanged. + Filtering on the wire record (not the mapped OpenDirect ``Product``) is + required: the mapped model carries neither ``ad_formats`` nor + ``seller_organization_id``, the fields the filter discriminates on. + + An item that fails wire validation is kept (never silently dropped by the + filter) — mapping downstream decides its fate via the tolerant reject path. + """ + if not filters: + return [i for i in raw_items if isinstance(i, dict)] + + from iab_agentic_primitives.primitives import Product as WireProduct + + indexed: list[tuple[int, WireProduct]] = [] + unvalidatable: list[int] = [] + dicts: list[dict[str, Any]] = [i for i in raw_items if isinstance(i, dict)] + for idx, item in enumerate(dicts): + try: + indexed.append((idx, WireProduct.model_validate(item))) + except Exception: # noqa: BLE001 — keep unvalidatable items for the mapper + unvalidatable.append(idx) + + kept = filter_wire_products([wp for _, wp in indexed], filters) + kept_ids = {id(wp) for wp in kept} + keep_idx = {idx for idx, wp in indexed if id(wp) in kept_ids} | set(unvalidatable) + return [item for idx, item in enumerate(dicts) if idx in keep_idx] + + +def _is_real_avails(data: Any) -> bool: + """True when ``data`` looks like a genuine seller ``check_avails`` result. + + The real tool returns the shared AvailsResponse shape (either snake_case + ``available_impressions`` or the aliased ``availableImpressions``). An error + result (``{"error": ...}``) or a pricing-only dict lacks it, signalling the + caller to fall back to the get_pricing stand-in. + """ + if not isinstance(data, dict) or "error" in data: + return False + return "available_impressions" in data or "availableImpressions" in data + + +def _avails_payload( + data: Any, + *, + product_id: str | None = None, + requested_impressions: int = 0, +) -> dict[str, Any]: + """Map a ``get_pricing`` result into a contract-valid AvailsResponse dict. + + ``AvailsResponse`` requires ``product_id``, ``available_impressions``, + ``estimated_cpm`` and ``total_cost``. The seller's ``get_pricing`` tool + returns ``product_id`` + ``base_cpm``/``final_cpm`` but no availability + volume (the seller MCP has no avails/forecast surface), so: + + - ``estimated_cpm`` = ``final_cpm`` (the tiered/negotiated price the buyer + would pay), falling back to ``base_cpm``. + - ``available_impressions`` = the buyer's requested volume — honest on this + path, where the seller declares no cap. 0 when unknown. + - ``total_cost`` = available_impressions / 1000 * estimated_cpm. + + A pre-existing ``{"available": ...}``-shaped dict (e.g. a future real avails + tool) is passed through untouched. An error/empty result yields a zero-cost + but still contract-valid response so the caller does not raise on a + missing-field ValidationError. + """ + # Pass-through for an already-avails-shaped dict. + if isinstance(data, dict) and "available" in data: + return data + + pid = None + estimated_cpm = 0.0 + if isinstance(data, dict): + pid = data.get("product_id") or data.get("productId") + cpm = data.get("final_cpm") + if cpm is None: + cpm = data.get("base_cpm") + if isinstance(cpm, (int, float)) and not isinstance(cpm, bool): + estimated_cpm = float(cpm) + + pid = pid or product_id or "" + available = max(0, int(requested_impressions or 0)) + total_cost = round(available / 1000 * estimated_cpm, 2) + + return { + "product_id": pid, + "available_impressions": available, + "estimated_cpm": estimated_cpm, + "total_cost": total_cost, + } diff --git a/src/ad_buyer/interfaces/agentcore/crew_tools.py b/src/ad_buyer/interfaces/agentcore/crew_tools.py index 72adf6e0..a9eba34f 100644 --- a/src/ad_buyer/interfaces/agentcore/crew_tools.py +++ b/src/ad_buyer/interfaces/agentcore/crew_tools.py @@ -96,14 +96,22 @@ 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 (sub-agents) and MANAGER_LLM_MODEL (level-1 manager) + # plus ANTHROPIC_COMPATIBLE_LLM_API_BASE_URL so this routes through the + # Anthropic Messages endpoint (no Converse patch). Sub-agents are the + # high-call-count fan-out, so they run on the (typically cheaper/faster) + # DEFAULT model while the single manager can stay on a stronger one via + # MANAGER_LLM_MODEL. The fallbacks are current-generation, Messages- + # supported Claude ids. 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", ) + manager_model = os.environ.get("MANAGER_LLM_MODEL", bedrock_model) from ad_buyer.config.settings import settings as buyer_settings - buyer_settings.manager_llm_model = bedrock_model + buyer_settings.manager_llm_model = manager_model buyer_settings.default_llm_model = bedrock_model # Use pre-extracted brief or parse one from the prompt @@ -124,10 +132,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..8325c806 100644 --- a/src/ad_buyer/interfaces/agentcore/http_main.py +++ b/src/ad_buyer/interfaces/agentcore/http_main.py @@ -55,17 +55,79 @@ 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 +# Durable Postgres opt-in (Scope 2 / Req 12.6-12.7): when the runtime was +# deployed with --storage postgres, deploy.sh forwards NON-secret env +# (DB_SECRET_ARN + AURORA_ENDPOINT + AURORA_PORT + DB_NAME) but never the +# password. At startup we fetch the RDS-managed secret from Secrets Manager, +# assemble the postgresql:// URL, and route the stores to Postgres. If the +# secret is absent/unreadable, resolve_database_url() returns None and we keep +# the SQLite default (fail-safe -- never crash startup on storage selection). +from ad_buyer.storage.db_secret import resolve_database_url # noqa: E402 + +_pg_url = resolve_database_url() +if _pg_url: + os.environ["DATABASE_URL"] = _pg_url + os.environ["STORAGE_TYPE"] = "hybrid" + # Never log _pg_url -- it carries the password. + print("[startup] Durable Postgres storage resolved from DB_SECRET_ARN.", flush=True) +elif os.environ.get("DB_SECRET_ARN"): + # ARN was set but the secret could not be fetched (e.g. the execution role + # lacks secretsmanager:GetSecretValue) -- resolve_database_url already logged + # the specific cause at WARN. Make the fallback unambiguous here. + print( + "[startup] DB_SECRET_ARN is set but the Postgres secret could not be " + "resolved; FELL BACK to default storage (see db_secret WARN above).", + flush=True, ) -except Exception as _exc: - logging.getLogger(__name__).warning("crewai_bedrock_fix patch failed: %s", _exc) +else: + print("[startup] No DB_SECRET_ARN set; using default (SQLite) storage.", flush=True) + +# 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/src/ad_buyer/storage/adserver_store.py b/src/ad_buyer/storage/adserver_store.py index 797a9815..e82e97d4 100644 --- a/src/ad_buyer/storage/adserver_store.py +++ b/src/ad_buyer/storage/adserver_store.py @@ -71,6 +71,7 @@ class AdServerStore: """ def __init__(self, database_url: str) -> None: + self._database_url = database_url self._db_path = self._parse_url(database_url) self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None @@ -91,12 +92,15 @@ def _parse_url(url: str) -> str: def connect(self) -> None: """Open the database connection, set pragmas, and create tables.""" - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute("PRAGMA foreign_keys=ON") - self._conn.execute("PRAGMA busy_timeout=5000") - self._create_tables() + from .connection_factory import ( + apply_sqlite_pragmas, + is_postgres_url, + open_connection, + ) + + self._conn = open_connection(self._database_url, check_same_thread=False) + apply_sqlite_pragmas(self._conn, self._database_url) + self._create_tables(is_postgres=is_postgres_url(self._database_url)) def disconnect(self) -> None: """Close the database connection.""" @@ -104,11 +108,22 @@ def disconnect(self) -> None: self._conn.close() self._conn = None - def _create_tables(self) -> None: - """Create ad server campaign tables and indexes if not present.""" + def _create_tables(self, is_postgres: bool = False) -> None: + """Create ad server campaign tables and indexes if not present. + + On Postgres the inline SQLite DDL is translated via the shared + ``schema_pg._to_postgres`` helper (SQLite path unchanged). + """ + table_ddl = AD_SERVER_CAMPAIGN_TABLE + index_ddls = AD_SERVER_CAMPAIGN_INDEXES + if is_postgres: + from .schema_pg import _to_postgres, _translate_all + + table_ddl = _to_postgres(AD_SERVER_CAMPAIGN_TABLE) + index_ddls = _translate_all(AD_SERVER_CAMPAIGN_INDEXES) cursor = self._conn.cursor() - cursor.execute(AD_SERVER_CAMPAIGN_TABLE) - for idx in AD_SERVER_CAMPAIGN_INDEXES: + cursor.execute(table_ddl) + for idx in index_ddls: cursor.execute(idx) self._conn.commit() diff --git a/src/ad_buyer/storage/audience_audit_log.py b/src/ad_buyer/storage/audience_audit_log.py index 79818bae..32472b0a 100644 --- a/src/ad_buyer/storage/audience_audit_log.py +++ b/src/ad_buyer/storage/audience_audit_log.py @@ -123,40 +123,52 @@ def _parse_url(database_url: str) -> str: return database_url -def _ensure_table(conn: sqlite3.Connection) -> None: +def _ensure_table(conn, is_postgres: bool = False) -> None: """Create `audience_audit_log` if it does not already exist. Used on first connection and as the migration safety-net for the test - that opens a DB created before this change landed. + that opens a DB created before this change landed. On Postgres the DDL is + translated via the shared ``schema_pg._to_postgres`` helper. """ + table_ddl = AUDIENCE_AUDIT_LOG_TABLE + index_ddls = AUDIENCE_AUDIT_LOG_INDEXES + if is_postgres: + from .schema_pg import _to_postgres, _translate_all + + table_ddl = _to_postgres(AUDIENCE_AUDIT_LOG_TABLE) + index_ddls = _translate_all(AUDIENCE_AUDIT_LOG_INDEXES) cursor = conn.cursor() - cursor.execute(AUDIENCE_AUDIT_LOG_TABLE) - for idx in AUDIENCE_AUDIT_LOG_INDEXES: + cursor.execute(table_ddl) + for idx in index_ddls: cursor.execute(idx) conn.commit() -def _get_conn() -> sqlite3.Connection: +def _get_conn(): """Return the shared connection, opening it lazily on first call.""" global _conn if _conn is None: with _conn_lock: if _conn is None: - path = _parse_url(_database_url) - conn = sqlite3.connect(path, check_same_thread=False) - conn.row_factory = sqlite3.Row + from .connection_factory import ( + apply_sqlite_pragmas, + is_postgres_url, + open_connection, + ) + + conn = open_connection(_database_url, check_same_thread=False) # WAL is shared with the rest of the buyer DB so we don't fight - # the deal store on the same file. + # the deal store on the same file. apply_sqlite_pragmas no-ops + # on Postgres and tolerates WAL-rejecting SQLite backends. try: - conn.execute("PRAGMA journal_mode=WAL") + apply_sqlite_pragmas(conn, _database_url) except sqlite3.OperationalError: # `:memory:` and some test backends reject WAL mode -- # fail-open, the table still works in journal mode. pass - conn.execute("PRAGMA busy_timeout=5000") - _ensure_table(conn) + _ensure_table(conn, is_postgres=is_postgres_url(_database_url)) _conn = conn return _conn diff --git a/src/ad_buyer/storage/campaign_store.py b/src/ad_buyer/storage/campaign_store.py index 04d22c4e..a2c3e8dc 100644 --- a/src/ad_buyer/storage/campaign_store.py +++ b/src/ad_buyer/storage/campaign_store.py @@ -35,7 +35,6 @@ from .schema import ( APPROVAL_REQUESTS_INDEXES, APPROVAL_REQUESTS_TABLE, - initialize_schema, ) logger = logging.getLogger(__name__) @@ -58,6 +57,7 @@ class CampaignStore: """ def __init__(self, database_url: str) -> None: + self._database_url = database_url self._db_path = self._parse_url(database_url) self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None @@ -68,13 +68,17 @@ def __init__(self, database_url: str) -> None: def connect(self) -> None: """Open the database connection, set pragmas, and initialize schema.""" - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute("PRAGMA foreign_keys=ON") - self._conn.execute("PRAGMA busy_timeout=5000") - initialize_schema(self._conn) - self._create_campaign_events_table() + from .connection_factory import ( + apply_sqlite_pragmas, + initialize_schema_for, + is_postgres_url, + open_connection, + ) + + self._conn = open_connection(self._database_url, check_same_thread=False) + apply_sqlite_pragmas(self._conn, self._database_url) + initialize_schema_for(self._conn, self._database_url) + self._create_campaign_events_table(is_postgres=is_postgres_url(self._database_url)) def disconnect(self) -> None: """Close the database connection.""" @@ -92,9 +96,13 @@ def _parse_url(url: str) -> str: return url[len("sqlite:///") :] return url - def _create_campaign_events_table(self) -> None: - """Create the campaign_events table if it doesn't exist.""" - self._conn.execute(""" + def _create_campaign_events_table(self, is_postgres: bool = False) -> None: + """Create the campaign_events table if it doesn't exist. + + On Postgres the inline SQLite DDL is translated via the shared + ``schema_pg._to_postgres`` helper (SQLite path unchanged). + """ + table_ddl = """ CREATE TABLE IF NOT EXISTS campaign_events ( event_id TEXT PRIMARY KEY, campaign_id TEXT NOT NULL, @@ -105,14 +113,21 @@ def _create_campaign_events_table(self) -> None: payload TEXT DEFAULT '{}', FOREIGN KEY (campaign_id) REFERENCES campaigns(campaign_id) ) - """) - self._conn.execute( + """ + index_ddls = [ "CREATE INDEX IF NOT EXISTS idx_campaign_events_campaign_id " - "ON campaign_events(campaign_id)" - ) - self._conn.execute( - "CREATE INDEX IF NOT EXISTS idx_campaign_events_timestamp ON campaign_events(timestamp)" - ) + "ON campaign_events(campaign_id)", + "CREATE INDEX IF NOT EXISTS idx_campaign_events_timestamp " + "ON campaign_events(timestamp)", + ] + if is_postgres: + from .schema_pg import _to_postgres, _translate_all + + table_ddl = _to_postgres(table_ddl) + index_ddls = _translate_all(index_ddls) + self._conn.execute(table_ddl) + for idx in index_ddls: + self._conn.execute(idx) self._conn.commit() def _emit_event( diff --git a/src/ad_buyer/storage/connection_factory.py b/src/ad_buyer/storage/connection_factory.py new file mode 100644 index 00000000..4f7dea24 --- /dev/null +++ b/src/ad_buyer/storage/connection_factory.py @@ -0,0 +1,76 @@ +# Donated to IAB Tech Lab + +"""Connection factory — the single seam the buyer stores route through. + +Scope 2 / Req 12.6. By default returns a plain ``sqlite3`` connection (the SQLite +path is byte-for-byte unchanged). When ``database_url`` is a PostgreSQL URL, it +returns the :class:`PgConnection` adapter instead. The stores' connection-open +sites change from ``sqlite3.connect(path, …)`` to ``open_connection(url, …)`` — +one line each — and otherwise keep their existing query code. + +Schema creation is likewise dialect-routed: ``initialize_schema_for`` calls the +SQLite ``schema.initialize_schema`` or the ``schema_pg.initialize_schema_pg``. +""" + +from __future__ import annotations + +import sqlite3 + + +def is_postgres_url(database_url: str | None) -> bool: + return bool(database_url) and database_url.startswith(("postgresql://", "postgres://")) + + +def _sqlite_path(database_url: str) -> str: + """Extract the sqlite file path (matches the stores' historical parsing).""" + if database_url.startswith("sqlite:///"): + return database_url[len("sqlite:///") :] + return database_url + + +def open_connection(database_url: str, *, check_same_thread: bool = False): + """Open a connection for ``database_url``. + + - ``postgresql://…`` / ``postgres://…`` → :class:`PgConnection` (psycopg + wrapped to present the sqlite3 API the stores use). + - anything else → ``sqlite3.connect`` (default, unchanged behavior). + + ``check_same_thread`` is accepted for call-site parity with the previous + ``sqlite3.connect`` calls; it is a no-op for the Postgres adapter. + """ + if is_postgres_url(database_url): + from .pg_connection import PgConnection + + # psycopg accepts the libpq URL directly (strip a SQLAlchemy +driver + # suffix if one slipped through, e.g. postgresql+asyncpg://). + dsn = database_url.replace("+asyncpg", "").replace("+psycopg", "") + return PgConnection(dsn) + + conn = sqlite3.connect(_sqlite_path(database_url), check_same_thread=check_same_thread) + conn.row_factory = sqlite3.Row + return conn + + +def apply_sqlite_pragmas(conn, database_url: str) -> None: + """Apply the WAL/foreign-key/busy-timeout pragmas — SQLite only, no-op on PG. + + Call sites previously ran these unconditionally; routing them here keeps the + Postgres adapter clean (Postgres has no PRAGMA). + """ + if is_postgres_url(database_url): + return + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA busy_timeout=5000") + + +def initialize_schema_for(conn, database_url: str) -> None: + """Create tables using the dialect-appropriate DDL.""" + if is_postgres_url(database_url): + from .schema_pg import initialize_schema_pg + + initialize_schema_pg(conn) + else: + from .schema import initialize_schema + + initialize_schema(conn) diff --git a/src/ad_buyer/storage/db_secret.py b/src/ad_buyer/storage/db_secret.py new file mode 100644 index 00000000..b4bda983 --- /dev/null +++ b/src/ad_buyer/storage/db_secret.py @@ -0,0 +1,60 @@ +# Donated to IAB Tech Lab + +"""Assemble the buyer's Postgres DATABASE_URL from the RDS-managed secret. + +Scope 2 / Req 12.7. On the AgentCore ``--storage postgres`` path the runtime +receives NON-secret env only — ``DB_SECRET_ARN`` + ``AURORA_ENDPOINT`` + +``AURORA_PORT`` + ``DB_NAME`` — never the password. At startup the app fetches +the RDS-managed secret from Secrets Manager by ARN, parses its JSON +(``{username, password, …}``), and assembles a libpq URL. The password is never +logged and never placed in the environment. + +The RDS-managed secret rotates; callers should re-invoke on a connection failure +so a rotated password is picked up (the value is read fresh each call — no +process-level caching of the password). +""" + +from __future__ import annotations + +import json +import logging +import os +from urllib.parse import quote + +logger = logging.getLogger(__name__) + + +def resolve_database_url() -> str | None: + """Return a ``postgresql://`` URL from the RDS-managed secret, or None. + + None means "no Postgres configured" — the caller falls back to its default + (SQLite). Any failure to fetch/parse the secret returns None with a warning + rather than crashing startup. + """ + secret_arn = os.environ.get("DB_SECRET_ARN") + host = os.environ.get("AURORA_ENDPOINT") + port = os.environ.get("AURORA_PORT", "5432") + dbname = os.environ.get("DB_NAME", "ad_buyer") + if not secret_arn or not host: + return None + + try: + import boto3 + + region = os.environ.get("AWS_REGION", "us-west-2") + client = boto3.client("secretsmanager", region_name=region) + resp = client.get_secret_value(SecretId=secret_arn) + secret = json.loads(resp["SecretString"]) + except Exception as exc: # noqa: BLE001 — never crash startup on secret fetch + logger.warning("Could not fetch DB secret (%s); falling back to default storage.", exc) + return None + + username = secret.get("username") + password = secret.get("password") + if not username or not password: + logger.warning("DB secret missing username/password; falling back to default storage.") + return None + + # URL-encode credentials so special characters are safe in the libpq URL. + # NOTE: never log the assembled URL (it contains the password). + return f"postgresql://{quote(username)}:{quote(password)}@{host}:{port}/{dbname}" diff --git a/src/ad_buyer/storage/deal_store.py b/src/ad_buyer/storage/deal_store.py index 628fcd07..8bf15de3 100644 --- a/src/ad_buyer/storage/deal_store.py +++ b/src/ad_buyer/storage/deal_store.py @@ -40,7 +40,6 @@ from .negotiation_store import NegotiationStore from .performance_cache_store import PerformanceCacheStore from .portfolio_metadata_store import PortfolioMetadataStore -from .schema import initialize_schema from .status_transition_store import StatusTransitionStore from .supply_path_template_store import SupplyPathTemplateStore @@ -68,6 +67,7 @@ class DealStore: """ def __init__(self, database_url: str) -> None: + self._database_url = database_url self._db_path = self._parse_url(database_url) self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None @@ -93,12 +93,15 @@ def __init__(self, database_url: str) -> None: def connect(self) -> None: """Open the database connection, set pragmas, and initialize schema.""" - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row # dict-like row access - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute("PRAGMA foreign_keys=ON") - self._conn.execute("PRAGMA busy_timeout=5000") - initialize_schema(self._conn) + from .connection_factory import ( + apply_sqlite_pragmas, + initialize_schema_for, + open_connection, + ) + + self._conn = open_connection(self._database_url, check_same_thread=False) + apply_sqlite_pragmas(self._conn, self._database_url) + initialize_schema_for(self._conn, self._database_url) self._wire_stores() def disconnect(self) -> None: diff --git a/src/ad_buyer/storage/health.py b/src/ad_buyer/storage/health.py index b7d2e745..6c3d724f 100644 --- a/src/ad_buyer/storage/health.py +++ b/src/ad_buyer/storage/health.py @@ -29,20 +29,24 @@ def _sqlite_path(database_url: str) -> str: def probe_database(database_url: str) -> tuple[bool, str | None]: """Attempt a lightweight connection + ``SELECT 1`` against the database. + Routes through the connection factory so a ``postgresql://`` URL is probed + via the psycopg adapter; SQLite behavior is unchanged. + Args: - database_url: SQLite connection string (or plain path). + database_url: SQLite or PostgreSQL connection string (or plain path). Returns: ``(True, None)`` when the database is reachable, otherwise ``(False, error_message)``. """ - db_path = _sqlite_path(database_url) + from .connection_factory import open_connection + try: - conn = sqlite3.connect(db_path) + conn = open_connection(database_url) conn.execute("SELECT 1") conn.close() return True, None - except (sqlite3.Error, OSError) as exc: + except (sqlite3.Error, OSError, Exception) as exc: # noqa: BLE001 return False, str(exc) diff --git a/src/ad_buyer/storage/order_store.py b/src/ad_buyer/storage/order_store.py index 2e5448d8..d1d942c8 100644 --- a/src/ad_buyer/storage/order_store.py +++ b/src/ad_buyer/storage/order_store.py @@ -59,6 +59,10 @@ class OrderStore: """ def __init__(self, database_url: str) -> None: + # Keep the full URL so the connection factory can detect a + # ``postgresql://`` backend; ``_db_path`` preserves the SQLite path for + # any legacy reference. SQLite path is byte-identical by default. + self._database_url = database_url self._db_path = self._parse_url(database_url) self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None @@ -79,12 +83,15 @@ def _parse_url(url: str) -> str: def connect(self) -> None: """Open the database connection, set pragmas, and create tables.""" - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute("PRAGMA foreign_keys=ON") - self._conn.execute("PRAGMA busy_timeout=5000") - self._create_tables() + from .connection_factory import ( + apply_sqlite_pragmas, + is_postgres_url, + open_connection, + ) + + self._conn = open_connection(self._database_url, check_same_thread=False) + apply_sqlite_pragmas(self._conn, self._database_url) + self._create_tables(is_postgres=is_postgres_url(self._database_url)) def disconnect(self) -> None: """Close the database connection.""" @@ -92,11 +99,23 @@ def disconnect(self) -> None: self._conn.close() self._conn = None - def _create_tables(self) -> None: - """Create orders table and indexes if they don't exist.""" + def _create_tables(self, is_postgres: bool = False) -> None: + """Create orders table and indexes if they don't exist. + + On Postgres the inline SQLite DDL is translated via the same + ``schema_pg._to_postgres`` helper used for ``schema.py`` (SQLite path + unchanged). + """ + table_ddl = ORDERS_TABLE + index_ddls = ORDERS_INDEXES + if is_postgres: + from .schema_pg import _to_postgres, _translate_all + + table_ddl = _to_postgres(ORDERS_TABLE) + index_ddls = _translate_all(ORDERS_INDEXES) cursor = self._conn.cursor() - cursor.execute(ORDERS_TABLE) - for idx in ORDERS_INDEXES: + cursor.execute(table_ddl) + for idx in index_ddls: cursor.execute(idx) self._conn.commit() diff --git a/src/ad_buyer/storage/pacing_store.py b/src/ad_buyer/storage/pacing_store.py index 866a97a1..f7a4ec91 100644 --- a/src/ad_buyer/storage/pacing_store.py +++ b/src/ad_buyer/storage/pacing_store.py @@ -73,6 +73,7 @@ class PacingStore: """ def __init__(self, database_url: str) -> None: + self._database_url = database_url self._db_path = self._parse_url(database_url) self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None @@ -93,12 +94,15 @@ def _parse_url(url: str) -> str: def connect(self) -> None: """Open the database connection, set pragmas, and create tables.""" - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute("PRAGMA foreign_keys=ON") - self._conn.execute("PRAGMA busy_timeout=5000") - self._create_tables() + from .connection_factory import ( + apply_sqlite_pragmas, + is_postgres_url, + open_connection, + ) + + self._conn = open_connection(self._database_url, check_same_thread=False) + apply_sqlite_pragmas(self._conn, self._database_url) + self._create_tables(is_postgres=is_postgres_url(self._database_url)) def disconnect(self) -> None: """Close the database connection.""" @@ -106,11 +110,22 @@ def disconnect(self) -> None: self._conn.close() self._conn = None - def _create_tables(self) -> None: - """Create pacing snapshot tables and indexes if not present.""" + def _create_tables(self, is_postgres: bool = False) -> None: + """Create pacing snapshot tables and indexes if not present. + + On Postgres the inline SQLite DDL is translated via the shared + ``schema_pg._to_postgres`` helper (SQLite path unchanged). + """ + table_ddl = PACING_SNAPSHOT_TABLE + index_ddls = PACING_SNAPSHOT_INDEXES + if is_postgres: + from .schema_pg import _to_postgres, _translate_all + + table_ddl = _to_postgres(PACING_SNAPSHOT_TABLE) + index_ddls = _translate_all(PACING_SNAPSHOT_INDEXES) cursor = self._conn.cursor() - cursor.execute(PACING_SNAPSHOT_TABLE) - for idx in PACING_SNAPSHOT_INDEXES: + cursor.execute(table_ddl) + for idx in index_ddls: cursor.execute(idx) self._conn.commit() diff --git a/src/ad_buyer/storage/pg_connection.py b/src/ad_buyer/storage/pg_connection.py new file mode 100644 index 00000000..c79f518f --- /dev/null +++ b/src/ad_buyer/storage/pg_connection.py @@ -0,0 +1,127 @@ +# Donated to IAB Tech Lab + +"""A psycopg connection that presents the ``sqlite3`` API the buyer stores use. + +Scope 2 / Req 12.6 — lets the existing synchronous stores talk to PostgreSQL +with NO change to their query code. The stores call ``conn.execute("… ?", +params)``, iterate dict-like rows, and call ``conn.commit()`` — all of which +this adapter maps onto psycopg. The only per-query dialect delta in this schema +is the ``?`` → ``%s`` paramstyle, translated here at the boundary. + +NOT a general SQLite emulator — it covers exactly the surface the buyer stores +exercise (execute/executemany, fetchone/fetchall, row-by-name access, commit, +rollback, close, context manager, no-op PRAGMA). Deliberately thin so the +maintainers can later replace it with a unified dialect layer. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def _qmark_to_pyformat(sql: str) -> str: + """Translate SQLite ``?`` placeholders to psycopg ``%s``. + + Leaves ``?`` inside single-quoted string literals alone, and escapes any + literal ``%`` so psycopg does not treat it as a format marker. + """ + # Escape existing % so psycopg's own %s parsing is not confused. + sql = sql.replace("%", "%%") + out = [] + in_str = False + for ch in sql: + if ch == "'": + in_str = not in_str + out.append(ch) + elif ch == "?" and not in_str: + out.append("%s") + else: + out.append(ch) + return "".join(out) + + +class _Cursor: + """Wraps a psycopg cursor to accept ``?`` SQL and return dict-like rows.""" + + def __init__(self, cur): + self._cur = cur + + def execute(self, sql: str, params: Any = ()): + self._cur.execute(_qmark_to_pyformat(sql), params or ()) + return self + + def executemany(self, sql: str, seq): + self._cur.executemany(_qmark_to_pyformat(sql), list(seq)) + return self + + def fetchone(self): + return self._cur.fetchone() + + def fetchall(self): + return self._cur.fetchall() + + @property + def rowcount(self): + return self._cur.rowcount + + @property + def lastrowid(self): + # Not used by the buyer stores (they use explicit ids / RETURNING-free + # upserts); return None rather than emulate SQLite rowids. + return None + + def __iter__(self): + return iter(self._cur) + + def close(self): + self._cur.close() + + +class PgConnection: + """A ``sqlite3.Connection``-shaped wrapper over a psycopg connection. + + Rows come back as mappings (``psycopg.rows.dict_row``) so the stores' + ``row["col"]`` access works exactly as with ``sqlite3.Row``. + """ + + def __init__(self, dsn: str): + import psycopg + from psycopg.rows import dict_row + + self._conn = psycopg.connect(dsn, row_factory=dict_row, autocommit=False) + + # -- sqlite3.Connection surface used by the stores -------------------- + def execute(self, sql: str, params: Any = ()): + cur = self._conn.cursor() + cur.execute(_qmark_to_pyformat(sql), params or ()) + return _Cursor(cur) + + def executemany(self, sql: str, seq): + cur = self._conn.cursor() + cur.executemany(_qmark_to_pyformat(sql), list(seq)) + return _Cursor(cur) + + def cursor(self): + return _Cursor(self._conn.cursor()) + + def commit(self): + self._conn.commit() + + def rollback(self): + self._conn.rollback() + + def close(self): + self._conn.close() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + if exc_type is None: + self._conn.commit() + else: + self._conn.rollback() + return False diff --git a/src/ad_buyer/storage/schema_pg.py b/src/ad_buyer/storage/schema_pg.py new file mode 100644 index 00000000..f8ec4042 --- /dev/null +++ b/src/ad_buyer/storage/schema_pg.py @@ -0,0 +1,89 @@ +# Donated to IAB Tech Lab + +"""PostgreSQL DDL for the buyer store, DERIVED from the SQLite DDL in ``schema.py``. + +Scope 2 / Req 12.6 — the AgentCore ``--storage postgres`` path opts into a +durable Postgres backend WITHOUT rewriting the shared SQLite ``schema.py`` (that +dialect unification is deferred to the upstream maintainers). Rather than +hand-maintain a parallel copy of every ``CREATE TABLE`` (which would drift), this +module IMPORTS the SQLite DDL strings from ``schema.py`` and applies a small, +mechanical SQLite→Postgres translation. The SQLite path is untouched. + +Translations (the only dialect deltas in this schema — verified by grep): + - ``INTEGER PRIMARY KEY AUTOINCREMENT`` → ``BIGSERIAL PRIMARY KEY`` + - ``strftime('%Y-%m-%dT%H:%M:%fZ','now')`` default → ``now()`` (timestamptz-ish + text; the app also writes ISO strings explicitly, so the default is a + fallback only) + - ``REAL`` → ``DOUBLE PRECISION`` + - ``INTEGER`` boolean-ish flag columns stay INTEGER (the app stores 0/1) +Runtime queries are already portable: ``?`` paramstyle is handled by the +connection adapter, and the one upsert (``ON CONFLICT(id) DO UPDATE SET +col = excluded.col``) is valid on both SQLite and Postgres. +""" + +from __future__ import annotations + +import re + +from . import schema as _sqlite_schema + + +def _to_postgres(ddl: str) -> str: + """Translate a SQLite CREATE-TABLE/INDEX string to PostgreSQL dialect.""" + out = ddl + # Autoincrement integer PK → BIGSERIAL (Postgres identity) + out = re.sub( + r"INTEGER\s+PRIMARY\s+KEY\s+AUTOINCREMENT", + "BIGSERIAL PRIMARY KEY", + out, + flags=re.IGNORECASE, + ) + # SQLite strftime(...) default → Postgres now() (both yield a timestamp text; + # the app writes explicit ISO strings on insert, so this is only a fallback). + out = re.sub( + r"strftime\('%Y-%m-%dT%H:%M:%fZ',\s*'now'\)", + "now()", + out, + flags=re.IGNORECASE, + ) + # REAL → DOUBLE PRECISION (Postgres has no REAL-with-same-semantics alias we + # want; DOUBLE PRECISION matches the app's float usage). + out = re.sub(r"\bREAL\b", "DOUBLE PRECISION", out) + return out + + +def _translate_all(ddls: list[str]) -> list[str]: + return [_to_postgres(d) for d in ddls] + + +# Ordered DDL: version table, then each table + its indexes, mirroring +# schema.py's create_tables() order. We pull the module's DDL constants by name +# so this list stays in lockstep with the source without duplicating SQL. +def postgres_ddl_statements() -> list[str]: + """Return the full ordered list of Postgres DDL statements.""" + stmts: list[str] = [] + # schema_version first + stmts.append(_to_postgres(_sqlite_schema.SCHEMA_VERSION_TABLE)) + # Every *_TABLE / *_INDEXES constant defined on the schema module, in + # definition order, so new tables added upstream are picked up automatically. + for name in dir(_sqlite_schema): + if name == "SCHEMA_VERSION_TABLE": + continue + val = getattr(_sqlite_schema, name) + if name.endswith("_TABLE") and isinstance(val, str): + stmts.append(_to_postgres(val)) + elif name.endswith("_INDEXES") and isinstance(val, list): + stmts.extend(_translate_all(val)) + return stmts + + +def initialize_schema_pg(conn) -> None: + """Create all tables + indexes on a PostgreSQL connection (adapter-wrapped). + + Idempotent (every statement is ``IF NOT EXISTS``). Uses the same connection + API the SQLite path uses (``execute`` + ``commit``), which the pg adapter + presents. + """ + for stmt in postgres_ddl_statements(): + conn.execute(stmt) + conn.commit() diff --git a/src/ad_buyer/tools/research/avails_check.py b/src/ad_buyer/tools/research/avails_check.py index 415d09ce..97a5a016 100644 --- a/src/ad_buyer/tools/research/avails_check.py +++ b/src/ad_buyer/tools/research/avails_check.py @@ -96,11 +96,18 @@ async def _arun( budget: float | None = None, ) -> str: """Check availability for the specified product.""" + # Parse dates first, in isolation, so a genuine date-format error is + # the ONLY thing reported as one. A Pydantic ValidationError from + # building AvailsResponse downstream is also a ValueError, so a single + # broad ``except ValueError`` here would mislabel a response-validation + # failure (e.g. a missing product_id) as "Error parsing dates". try: - # Parse dates start_dt = datetime.fromisoformat(start_date) end_dt = datetime.fromisoformat(end_date) + except ValueError as e: + return f"Error parsing dates: {e}. Please use YYYY-MM-DD format." + try: # Build request request = AvailsRequest( product_id=product_id, @@ -116,7 +123,9 @@ async def _arun( return self._format_results(product_id, start_date, end_date, avails) except ValueError as e: - return f"Error parsing dates: {e}. Please use YYYY-MM-DD format." + # Not a date error (dates already parsed above): a validation or + # value error from the avails request/response path. + return f"Error checking availability: {e}" except (httpx.HTTPError, OSError) as e: return f"Error checking availability: {e}" diff --git a/tests/integration/agentcore/conftest.py b/tests/integration/agentcore/conftest.py index 4d3299ff..490a86b5 100644 --- a/tests/integration/agentcore/conftest.py +++ b/tests/integration/agentcore/conftest.py @@ -11,3 +11,42 @@ def pytest_addoption(parser): parser.addoption( "--agent-name", action="store", default=None, help="Agent name in .bedrock_agentcore.yaml" ) + + +# --------------------------------------------------------------------------- +# Per-test wall-clock timing +# +# These live tests each drive a full Bedrock crew invoke, so how long each one +# takes is a first-class signal (latency is dominated by LLM round-trips, not +# the framework — see the runtime OTEL logs). Record the call-phase duration of +# every test and print a sorted summary at the end so slow paths are obvious +# without editing individual tests. +# --------------------------------------------------------------------------- + +import pytest # noqa: E402 + +_TEST_DURATIONS: list[tuple[str, float]] = [] + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + """Capture each test's call-phase (body) duration.""" + outcome = yield + report = outcome.get_result() + if report.when == "call": + _TEST_DURATIONS.append((report.nodeid, report.duration)) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + """Print a sorted wall-clock table for the AgentCore live tests.""" + if not _TEST_DURATIONS: + return + rows = sorted(_TEST_DURATIONS, key=lambda r: r[1], reverse=True) + total = sum(d for _, d in rows) + terminalreporter.write_sep("=", "AgentCore live test durations (slowest first)") + for nodeid, dur in rows: + # Trim to the test name (drop the file path prefix) for readability. + name = nodeid.split("::", 1)[-1] if "::" in nodeid else nodeid + terminalreporter.write_line(f" {dur:7.1f}s {name}") + terminalreporter.write_line(f" {'-' * 40}") + terminalreporter.write_line(f" {total:7.1f}s TOTAL ({len(rows)} tests)") diff --git a/tests/integration/agentcore/test_runtime.py b/tests/integration/agentcore/test_runtime.py index 864c2cb2..e0cefc23 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,251 @@ 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, + { + # Search terms are substrings of real seller catalog product + # NAMES ("Sports", "Video") so the discovery path returns + # products deterministically — the seller CSV list_inventory + # filter matches filter_str against the product name only, and + # empirically this wording produced live inv-* results. + "prompt": ( + "Plan a $500K Q4 sports campaign. Search the seller for" + " Sports and Video inventory and list the products found." + ), + "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:]}" + ) + + +class TestCrewExercisesStores: + """Drive each buyer store through a real agent/crew interaction. + + Rather than poke tables directly, each test sends a natural-language prompt + that routes the crew to a specific MCP tool, then asserts (via CloudWatch) + that the tool fired — proving that store's path works end-to-end through + agent interaction on the deployed runtime. + + Store -> representative tool (from mcp_server.py wiring): + DealStore family -> list_deals / search_deals / get_portfolio_summary + CampaignStore -> list_campaigns / get_campaign_status + PacingStore -> check_pacing / get_pacing_report + OrderStore -> list_orders / get_order_status + """ + + @staticmethod + def _log_group(config: "RuntimeConfig") -> str | None: + 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: + 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}") + + # (prompt that should route the crew to the tool, tool name to see in logs) + # + # Prompts are deliberately imperative: they name the store operation and + # explicitly forbid answering from memory, so the crew's LLM calls the tool + # instead of replying conversationally. NL routing is still non-deterministic + # (the test skips-not-fails if no tool fires), but naming the operation + + # "you must call the tool" raises the hit-rate materially over soft prompts. + _STORE_PROBES = [ + ( + "Use your tools to retrieve the deals in my deal library: call the " + "deal-listing tool and the portfolio-summary tool, and report exactly " + "what they return. Do not answer from memory — if a tool exists for " + "this, you must call it.", + ("list_deals", "get_portfolio_summary", "search_deals"), + "DealStore", + ), + ( + "Use your tools to retrieve my current campaigns: call the " + "campaign-listing tool and report each campaign's status exactly as " + "returned. Do not answer from memory — if a tool exists for this, you " + "must call it.", + ("list_campaigns", "get_campaign_status"), + "CampaignStore", + ), + ( + "Use your tools to check pacing on my campaigns: call the pacing tool " + "and return the pacing report exactly as it comes back. Do not answer " + "from memory — if a tool exists for this, you must call it.", + ("check_pacing", "get_pacing_report", "review_budgets"), + "PacingStore", + ), + ( + "Use your tools to retrieve my orders: call the order-listing tool and " + "report each order's current status exactly as returned. Do not answer " + "from memory — if a tool exists for this, you must call it.", + ("list_orders", "get_order_status"), + "OrderStore", + ), + ] + + @pytest.mark.parametrize( + "prompt,tool_names,store", + _STORE_PROBES, + ids=[p[2] for p in _STORE_PROBES], + ) + def test_crew_exercises_store(self, runtime_config, prompt, tool_names, store): + since_ms = int((time.time() - 5) * 1000) + result = invoke_runtime( + runtime_config, + {"prompt": prompt, "routing_mode": "crew"}, + timeout=300, + ) + assert result["success"], f"Invoke failed for {store}: {result['error']}" + + # Assert at least one of the store's tools fired (CloudWatch filter is an + # OR of quoted terms). Skip (not fail) if the crew chose a different + # path -- the store path is best-effort driven by NL, and the assertion + # channel (logs) is not the SUT. + pattern = " ".join(f'"{t}"' for t in tool_names) if len(tool_names) == 1 else ( + "?" + " ?".join(f'"{t}"' for t in tool_names) + ) + text = self._recent_log_text(runtime_config, since_ms, pattern) + if not any(t in text for t in tool_names): + pytest.skip( + f"{store}: crew did not route to {tool_names} for this prompt " + f"(NL routing is non-deterministic; log tail: {text[-300:]})" + ) + # A tool fired -> the store's connect()+query path executed on the + # deployed (Postgres-backed) runtime. diff --git a/tests/integration/test_pg_round_trip.py b/tests/integration/test_pg_round_trip.py new file mode 100644 index 00000000..c529ce07 --- /dev/null +++ b/tests/integration/test_pg_round_trip.py @@ -0,0 +1,213 @@ +# Donated to IAB Tech Lab + +"""Integration round-trip for the buyer Postgres backend (Req 12.6, task 10.10). + +Proves the schema_pg DDL + PgConnection adapter + connection_factory seam work +end-to-end against a REAL PostgreSQL server -- the thing unit tests alone cannot +guarantee (a paramstyle bug, a DDL type mismatch, or the ON CONFLICT upsert can +pass offline and fail on a live server). + +Backend selection (no local PG on the dev laptop): + - Set ``PG_TEST_DATABASE_URL`` (a local Docker Postgres in CI, or the live + Aurora endpoint via a bastion/tunnel) to run these. + - Absent -> the whole module is skipped (offline suites stay green). + +Each test uses a unique table/id so it is safe against a shared database and +cleans up after itself. The SQLite path is covered by the per-store suites; +this file is exclusively the Postgres proof. +""" + +from __future__ import annotations + +import os +import uuid + +import pytest + +PG_URL = os.environ.get("PG_TEST_DATABASE_URL") + +pytestmark = pytest.mark.skipif( + not PG_URL, + reason="PG_TEST_DATABASE_URL not set -- skipping live Postgres round-trip", +) + + +@pytest.fixture +def pg_conn(): + from ad_buyer.storage.connection_factory import open_connection + + conn = open_connection(PG_URL) + yield conn + conn.close() + + +class TestPostgresRoundTrip: + def test_connect_and_select_one(self, pg_conn): + row = pg_conn.execute("SELECT 1 AS one").fetchone() + assert row["one"] == 1 + + def test_qmark_paramstyle_round_trip(self, pg_conn): + t = f"pg_it_{uuid.uuid4().hex[:8]}" + pg_conn.execute(f"CREATE TABLE {t} (id TEXT PRIMARY KEY, data TEXT)") + pg_conn.commit() + try: + # ? placeholders must translate + bind correctly on a real server. + pg_conn.execute(f"INSERT INTO {t} (id, data) VALUES (?, ?)", ("a", "hello")) + pg_conn.commit() + row = pg_conn.execute(f"SELECT data FROM {t} WHERE id = ?", ("a",)).fetchone() + assert row["data"] == "hello" + finally: + pg_conn.execute(f"DROP TABLE {t}") + pg_conn.commit() + + def test_on_conflict_upsert_portable(self, pg_conn): + t = f"pg_it_{uuid.uuid4().hex[:8]}" + pg_conn.execute(f"CREATE TABLE {t} (id TEXT PRIMARY KEY, data TEXT)") + pg_conn.commit() + try: + pg_conn.execute(f"INSERT INTO {t} (id, data) VALUES (?, ?)", ("k", "v1")) + pg_conn.execute( + f"INSERT INTO {t} (id, data) VALUES (?, ?) " + f"ON CONFLICT(id) DO UPDATE SET data = excluded.data", + ("k", "v2"), + ) + pg_conn.commit() + row = pg_conn.execute(f"SELECT data FROM {t} WHERE id = ?", ("k",)).fetchone() + assert row["data"] == "v2" + finally: + pg_conn.execute(f"DROP TABLE {t}") + pg_conn.commit() + + def test_schema_pg_ddl_applies_on_real_server(self, pg_conn): + # The derived Postgres DDL must actually execute against a live server + # (idempotent -- all IF NOT EXISTS). + from ad_buyer.storage.schema_pg import initialize_schema_pg + + initialize_schema_pg(pg_conn) # no raise = valid PG DDL + + +# --------------------------------------------------------------------------- +# ALL stores: prove every store's connect()/DDL works on real Postgres +# --------------------------------------------------------------------------- +class TestAllStoresOnPostgres: + """Connect every buyer store against real Postgres and confirm its tables + are created. This is the comprehensive "all 17 stores" proof: each store's + connect() runs its own DDL through the schema_pg translation on a live + server, so a dialect bug in ANY store's inline DDL surfaces here. + """ + + # Every table the buyer persistence layer defines, and which store creates + # it. The deal facade (DealStore) creates the schema.py family; the 4 + # self-connecting stores + the audit log create their own. + _FACADE_TABLES = [ + "deals", + "negotiation_rounds", + "booking_records", + "jobs", + "events", + "status_transitions", + "portfolio_metadata", + "deal_activations", + "performance_cache", + "creative_assets", + "deal_templates", + "supply_path_templates", + ] + + @staticmethod + def _table_exists(conn, name: str) -> bool: + row = conn.execute( + "SELECT 1 FROM information_schema.tables " + "WHERE table_schema = 'public' AND table_name = ?", + (name,), + ).fetchone() + return row is not None + + def test_deal_facade_and_substores_tables(self): + from ad_buyer.storage.deal_store import DealStore + + store = DealStore(PG_URL) + store.connect() + try: + for tbl in self._FACADE_TABLES: + assert self._table_exists(store._conn, tbl), f"missing facade table: {tbl}" + finally: + store.disconnect() + + def test_campaign_store_tables(self): + from ad_buyer.storage.campaign_store import CampaignStore + + store = CampaignStore(PG_URL) + store.connect() + try: + for tbl in ("campaigns", "campaign_events", "approval_requests"): + assert self._table_exists(store._conn, tbl), f"missing campaign table: {tbl}" + finally: + store.disconnect() + + def test_order_store_tables(self): + from ad_buyer.storage.order_store import OrderStore + + store = OrderStore(PG_URL) + store.connect() + try: + assert self._table_exists(store._conn, "orders") + finally: + store.disconnect() + + def test_pacing_store_tables(self): + from ad_buyer.storage.pacing_store import PacingStore + + store = PacingStore(PG_URL) + store.connect() + try: + assert self._table_exists(store._conn, "pacing_snapshots") + finally: + store.disconnect() + + def test_adserver_store_tables(self): + from ad_buyer.storage.adserver_store import AdServerStore + + store = AdServerStore(PG_URL) + store.connect() + try: + assert self._table_exists(store._conn, "ad_server_campaigns") + finally: + store.disconnect() + + def test_audience_audit_log_table(self): + from ad_buyer.storage import audience_audit_log as aal + + # Module-level connection; configure() points it at PG. + aal.configure(PG_URL) + conn = aal._get_conn() + assert self._table_exists(conn, "audience_audit_log") + + def test_order_store_write_recycle_read(self): + # Full store round-trip: write, disconnect (recycle), reconnect, read. + from ad_buyer.storage.order_store import OrderStore + + oid = f"order:{uuid.uuid4().hex[:8]}" + store = OrderStore(PG_URL) + store.connect() + try: + store._conn.execute( + "INSERT INTO orders (key, data, status) VALUES (?, ?, ?)", + (oid, '{"x": 1}', "pending"), + ) + store._conn.commit() + store.disconnect() + + store2 = OrderStore(PG_URL) + store2.connect() + row = store2._conn.execute( + "SELECT data, status FROM orders WHERE key = ?", (oid,) + ).fetchone() + assert row["status"] == "pending" + store2.disconnect() + finally: + cleanup = OrderStore(PG_URL) + cleanup.connect() + cleanup._conn.execute("DELETE FROM orders WHERE key = ?", (oid,)) + cleanup._conn.commit() + cleanup.disconnect() 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..ebbe59fe --- /dev/null +++ b/tests/unit/agentcore/test_agentcore_seller_proxy.py @@ -0,0 +1,353 @@ +"""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() + 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" + + +class TestCheckAvails: + """Regression: check_avails must call the real get_pricing MCP tool and map + its result into a contract-valid AvailsResponse. The old code called a + phantom ``discover_inventory`` tool the seller MCP does not expose, so + ``AvailsResponse.model_validate`` raised on missing required fields + (product_id / available_impressions / estimated_cpm / total_cost). + """ + + async def test_maps_get_pricing_to_valid_avails(self): + from datetime import datetime + + from ad_buyer.models.opendirect import AvailsRequest + + # get_pricing-shaped result (no available_impressions) -> the proxy's + # check_avails call returns this from the fallback branch (older seller + # without a real check_avails tool), and it maps to a valid AvailsResponse. + fake = _FakeBedrockAgentCore( + { + "product_id": "inv-ctv-apex-sports-nba", + "base_cpm": 35.0, + "final_cpm": 30.0, + "tier_discount": 0.0, + "volume_discount": 0.14, + "rationale": "tiered", + } + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + req = AvailsRequest( + product_id="inv-ctv-apex-sports-nba", + start_date=datetime(2026, 10, 1), + end_date=datetime(2026, 12, 31), + requested_impressions=1_000_000, + ) + avails = await proxy.check_avails(req) + + # Calls the real get_pricing tool (not the phantom discover_inventory). + payload = json.loads(fake.last_kwargs["payload"]) + assert payload["params"]["name"] == "get_pricing" + + # All required AvailsResponse fields are populated -> no ValidationError. + assert avails.product_id == "inv-ctv-apex-sports-nba" + assert avails.available_impressions == 1_000_000 + assert avails.estimated_cpm == 30.0 # final_cpm preferred over base_cpm + assert avails.total_cost == 30_000.0 # 1_000_000 / 1000 * 30.0 + + async def test_uses_real_check_avails_when_available(self): + """When the seller exposes a real check_avails tool (full avails shape), + the proxy validates it directly — no get_pricing fallback.""" + from datetime import datetime + + from ad_buyer.models.opendirect import AvailsRequest + + fake = _FakeBedrockAgentCore( + { + "product_id": "inv-ctv-apex-sports-nba", + "available_impressions": 750_000, + "guaranteed_impressions": 750_000, + "estimated_cpm": 32.0, + "total_cost": 24_000.0, + } + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + req = AvailsRequest( + product_id="inv-ctv-apex-sports-nba", + start_date=datetime(2026, 10, 1), + end_date=datetime(2026, 12, 31), + requested_impressions=1_000_000, + ) + avails = await proxy.check_avails(req) + # Primary tool call is the real check_avails (not get_pricing). + assert json.loads(fake.last_kwargs["payload"])["params"]["name"] == "check_avails" + assert avails.available_impressions == 750_000 + assert avails.estimated_cpm == 32.0 + assert avails.total_cost == 24_000.0 + + async def test_error_result_still_yields_valid_avails(self): + """A seller error/empty pricing result maps to a zero-cost but + still contract-valid AvailsResponse, so the buyer tool does not raise a + misleading validation error (formerly surfaced as 'Error parsing dates').""" + from datetime import datetime + + from ad_buyer.models.opendirect import AvailsRequest + + fake = _FakeBedrockAgentCore({"error": "Product 'nope' not found"}) + proxy = AgentCoreSellerProxy(ARN, client=fake) + req = AvailsRequest( + product_id="nope", + start_date=datetime(2026, 10, 1), + end_date=datetime(2026, 12, 31), + requested_impressions=500_000, + ) + avails = await proxy.check_avails(req) + assert avails.product_id == "nope" # falls back to the requested id + assert avails.estimated_cpm == 0.0 + assert avails.total_cost == 0.0 + + +def _wire_product_fmt(pid: str, ad_formats: list[str], name: str = "P") -> dict: + """A wire product declaring ad_formats (seller channel taxonomy).""" + item = _wire_product(pid, name) + item["ad_formats"] = ad_formats + return item + + +@pytest.mark.asyncio +class TestSearchFilters: + """The proxy delegates catalog filtering to the SHARED product_filter helper + (same one OpenDirectClient uses), filtering wire records BEFORE mapping. + + Regression guard for the original proxy bug: a hand-rolled ``_apply_filters`` + exact-matched ``channel`` against a non-existent ``inventory_type`` attribute, + zeroing out every channel-filtered search ("No products found") even though + the seller returned a full catalog. + """ + + async def test_channel_filter_does_not_zero_the_catalog(self): + # The crew almost always passes a free-text ``channel`` (e.g. "video"). + # Channel is an intentionally IGNORED key — it must NOT exclude anything. + fake = _FakeBedrockAgentCore( + { + "products": [ + _wire_product_fmt("inv-ctv-1", ["ctv"], "CTV One"), + _wire_product_fmt("inv-dig-1", ["display"], "Display One"), + ] + } + ) + proxy = AgentCoreSellerProxy(ARN, region="us-west-2", client=fake) + products = await proxy.search_products({"channel": "video"}) + # Both survive — channel does not filter. (The old bug returned []). + assert {p.name for p in products} == {"CTV One", "Display One"} + + async def test_adformat_vocabulary_reconciles_banner_to_display(self): + # Buyer placement vocab "banner" must reconcile to the seller's "display" + # channel taxonomy via the shared normalizer — the display product is kept. + fake = _FakeBedrockAgentCore( + { + "products": [ + _wire_product_fmt("inv-dig-1", ["display"], "Display One"), + _wire_product_fmt("inv-ctv-1", ["ctv"], "CTV One"), + ] + } + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + products = await proxy.search_products({"adFormat": "banner"}) + assert [p.name for p in products] == ["Display One"] + + async def test_adformat_excludes_declared_mismatch_keeps_undeclared(self): + fake = _FakeBedrockAgentCore( + { + "products": [ + _wire_product_fmt("inv-vid-1", ["video"], "Video One"), + _wire_product_fmt("inv-und-1", [], "Undeclared One"), + _wire_product_fmt("inv-dig-1", ["display"], "Display One"), + ] + } + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + products = await proxy.search_products({"adFormat": "video"}) + # Declared-video kept; undeclared survives (do-not-exclude); display dropped. + assert {p.name for p in products} == {"Video One", "Undeclared One"} + + async def test_no_filters_returns_full_catalog(self): + fake = _FakeBedrockAgentCore( + {"products": [_wire_product_fmt("inv-a", ["ctv"]), _wire_product_fmt("inv-b", ["display"])]} + ) + proxy = AgentCoreSellerProxy(ARN, client=fake) + products = await proxy.search_products({}) + assert len(products) == 2 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= 16.3 (or >=15.7 etc.) + assert (major, minor) >= (16, 3), f"engine {ver} does not support MinCapacity=0" + + def test_storage_exports_secret_and_connection(self): + o = load_cfn_yaml(CFN_DIR / "storage.yaml")["Outputs"] + for k in ("AuroraSecretArn", "AuroraEndpoint", "AuroraPort", "AuroraDatabaseName"): + assert k in o, f"storage.yaml must export {k}" + + # -- network-agentcore.yaml: endpoints + 443 self-ingress -- + def test_agentcore_network_endpoints_and_self_ingress(self): + r = load_cfn_yaml(self.AGENTCORE_DIR / "network-agentcore.yaml")["Resources"] + wanted = { + "BedrockAgentCoreEndpoint": "bedrock-agentcore", + "BedrockRuntimeEndpoint": "bedrock-runtime", + "StsEndpoint": "sts", + "SecretsManagerEndpoint": "secretsmanager", + } + for lid, svc in wanted.items(): + assert lid in r and r[lid]["Type"] == "AWS::EC2::VPCEndpoint" + assert svc in str(r[lid]["Properties"]["ServiceName"]) + # the dark-container fix + assert "AgentCoreHttpsSelfIngress" in r, "443 self-ingress rule required" + + # -- main-agentcore.yaml: orchestrator creates the Aurora SG in-stack -- + def test_main_agentcore_creates_aurora_sg_option1(self): + r = load_cfn_yaml(self.AGENTCORE_DIR / "main-agentcore.yaml")["Resources"] + assert "DatabaseSecurityGroup" in r, ( + "main-agentcore.yaml must create the Aurora SG in-stack (Option 1)" + ) + for stack in ("NetworkStack", "StorageStack", "AgentCoreNetworkStack"): + assert stack in r + + +class TestDeployPostgresPath: + """deploy.sh --storage postgres branch + secret grant wiring (Req 12.4).""" + + DEPLOY = INFRA_ROOT / "aws" / "agentcore" / "deploy.sh" + + @pytest.fixture + def script(self): + return self.DEPLOY.read_text() + + def test_has_storage_arg(self, script): + assert "--storage)" in script and 'STORAGE_TYPE_ARG="sqlite"' in script + + def test_postgres_deploys_infra(self, script): + assert "deploy_postgres_infrastructure" in script + assert "main-agentcore.yaml" in script + + def test_postgres_uses_vpc_mode(self, script): + assert "--vpc" in script and "--security-groups" in script + + def test_postgres_forwards_non_secret_env_only(self, script): + assert "STORAGE_TYPE=hybrid" in script + assert "DB_SECRET_ARN=" in script + # the password must NEVER be interpolated into env + assert "DB_PASSWORD" not in script + + def test_grants_secret_read_to_runtime_role(self, script): + assert "grant_secret_access" in script + assert "secretsmanager:GetSecretValue" in script + + def test_default_stays_sqlite(self, script): + assert "STORAGE_TYPE=sqlite" in script + + class TestTerraformModules: """Validate Terraform module structure.""" diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py index 7384b2fa..27843a6d 100644 --- a/tests/unit/test_llm.py +++ b/tests/unit/test_llm.py @@ -4,12 +4,28 @@ """Unit tests for the custom OpenAI-compatible endpoint alternative (build_llm).""" from ad_buyer.config.settings import Settings -from ad_buyer.llm import _model_accepts_temperature, build_llm +from ad_buyer.llm import ( + _model_accepts_temperature, + _model_uses_max_completion_tokens, + build_llm, +) def _settings(**overrides) -> 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_pg_backend.py b/tests/unit/test_pg_backend.py new file mode 100644 index 00000000..0bcb7866 --- /dev/null +++ b/tests/unit/test_pg_backend.py @@ -0,0 +1,231 @@ +# Donated to IAB Tech Lab + +"""Unit tests for the Scope 2 buyer Postgres backend (Req 12.6/12.7). + +Offline tests (no live database) covering the three dialect seams: + 1. paramstyle translation (pg_connection._qmark_to_pyformat) + 2. DDL translation (schema_pg._to_postgres + inline store DDL) + 3. connection routing (connection_factory.open_connection / *_for) + 4. secret -> URL assembly (db_secret.resolve_database_url) + +The SQLite path is asserted to be byte-identical by default (the factory +returns a plain sqlite3 connection unless the URL is postgresql://). +""" + +from __future__ import annotations + +import sqlite3 + +import pytest + +from ad_buyer.storage import connection_factory as cf +from ad_buyer.storage import schema_pg +from ad_buyer.storage.pg_connection import _qmark_to_pyformat + + +# --------------------------------------------------------------------------- +# 1. Paramstyle translation +# --------------------------------------------------------------------------- +class TestParamstyleTranslation: + def test_simple_placeholders(self): + assert _qmark_to_pyformat("INSERT INTO t VALUES (?, ?)") == ( + "INSERT INTO t VALUES (%s, %s)" + ) + + def test_no_placeholders_unchanged(self): + assert _qmark_to_pyformat("SELECT 1") == "SELECT 1" + + def test_question_mark_inside_string_literal_preserved(self): + # A ? inside a single-quoted literal must NOT become %s. + sql = "SELECT * FROM t WHERE label = 'is it? yes' AND id = ?" + out = _qmark_to_pyformat(sql) + assert "'is it? yes'" in out + assert out.endswith("id = %s") + assert out.count("%s") == 1 + + def test_literal_percent_is_escaped(self): + # Existing % must be escaped so psycopg does not read it as a marker. + out = _qmark_to_pyformat("SELECT * FROM t WHERE name LIKE 'a%'") + assert "%%" in out + + def test_on_conflict_upsert_translates_placeholders_only(self): + # The one upsert in job_store is portable; only ? -> %s changes. + sql = ( + "INSERT INTO jobs (id, data) VALUES (?, ?) " + "ON CONFLICT(id) DO UPDATE SET data = excluded.data" + ) + out = _qmark_to_pyformat(sql) + assert "ON CONFLICT(id) DO UPDATE SET data = excluded.data" in out + assert out.count("%s") == 2 + + +# --------------------------------------------------------------------------- +# 2. DDL translation (schema.py-derived + inline store DDL) +# --------------------------------------------------------------------------- +class TestSchemaPgTranslation: + def test_autoincrement_becomes_bigserial(self): + ddl = "CREATE TABLE t (id INTEGER PRIMARY KEY AUTOINCREMENT, x TEXT)" + out = schema_pg._to_postgres(ddl) + assert "BIGSERIAL PRIMARY KEY" in out + assert "AUTOINCREMENT" not in out + + def test_strftime_default_becomes_now(self): + ddl = "CREATE TABLE t (ts TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')))" + out = schema_pg._to_postgres(ddl) + assert "now()" in out + assert "strftime" not in out + + def test_real_becomes_double_precision(self): + out = schema_pg._to_postgres("CREATE TABLE t (v REAL)") + assert "DOUBLE PRECISION" in out + assert "REAL" not in out + + def test_schema_derived_statements_have_no_sqlite_isms(self): + stmts = schema_pg.postgres_ddl_statements() + assert len(stmts) > 10 + blob = "\n".join(stmts) + assert "AUTOINCREMENT" not in blob + assert "strftime" not in blob + + @pytest.mark.parametrize( + "module_name,table_const,index_const", + [ + ("order_store", "ORDERS_TABLE", "ORDERS_INDEXES"), + ("adserver_store", "AD_SERVER_CAMPAIGN_TABLE", "AD_SERVER_CAMPAIGN_INDEXES"), + ("pacing_store", "PACING_SNAPSHOT_TABLE", "PACING_SNAPSHOT_INDEXES"), + ], + ) + def test_self_connecting_store_inline_ddl_translates( + self, module_name, table_const, index_const + ): + # The 4 self-connecting stores carry their own CREATE TABLE DDL; each + # must translate cleanly via the shared helper (no SQLite-isms left). + import importlib + + mod = importlib.import_module(f"ad_buyer.storage.{module_name}") + table_pg = schema_pg._to_postgres(getattr(mod, table_const)) + assert "AUTOINCREMENT" not in table_pg + assert "strftime" not in table_pg + idx_pg = schema_pg._translate_all(getattr(mod, index_const)) + assert all("strftime" not in i for i in idx_pg) + + +# --------------------------------------------------------------------------- +# 3. Connection routing (the seam) +# --------------------------------------------------------------------------- +class TestConnectionFactoryRouting: + def test_is_postgres_url(self): + assert cf.is_postgres_url("postgresql://h/db") + assert cf.is_postgres_url("postgres://h/db") + assert not cf.is_postgres_url("sqlite:///:memory:") + assert not cf.is_postgres_url(None) + + def test_sqlite_url_returns_plain_sqlite_connection(self): + # SQLite path byte-identical: a real sqlite3.Connection, not the adapter. + conn = cf.open_connection("sqlite:///:memory:") + assert isinstance(conn, sqlite3.Connection) + conn.close() + + def test_apply_pragmas_noop_on_postgres(self): + # Must not touch a PG connection (no PRAGMA in Postgres). + class _Sentinel: + def execute(self, *_a, **_k): + raise AssertionError("PRAGMA should not run on Postgres") + + cf.apply_sqlite_pragmas(_Sentinel(), "postgresql://h/db") # no raise + + def test_apply_pragmas_runs_on_sqlite(self): + conn = sqlite3.connect(":memory:") + cf.apply_sqlite_pragmas(conn, "sqlite:///:memory:") + mode = conn.execute("PRAGMA journal_mode").fetchone()[0] + assert mode is not None + conn.close() + + def test_postgres_url_builds_adapter(self, monkeypatch): + # open_connection on a postgres URL constructs PgConnection; stub psycopg + # so no real server is needed. + import ad_buyer.storage.pg_connection as pgmod + + captured = {} + + class _FakePsycopg: + @staticmethod + def connect(dsn, **kw): + captured["dsn"] = dsn + return object() + + monkeypatch.setattr(pgmod, "psycopg", _FakePsycopg, raising=False) + # PgConnection.__init__ imports psycopg lazily; patch the module import. + import sys + + fake_mod = type(sys)("psycopg") + fake_mod.connect = _FakePsycopg.connect + rows_mod = type(sys)("psycopg.rows") + rows_mod.dict_row = object() + monkeypatch.setitem(sys.modules, "psycopg", fake_mod) + monkeypatch.setitem(sys.modules, "psycopg.rows", rows_mod) + + conn = cf.open_connection("postgresql://u:p@h:5432/ad_buyer") + assert conn.__class__.__name__ == "PgConnection" + assert captured["dsn"].startswith("postgresql://") + + +# --------------------------------------------------------------------------- +# 4. Secret -> URL assembly +# --------------------------------------------------------------------------- +class TestDbSecretResolution: + def test_returns_none_without_env(self, monkeypatch): + for k in ("DB_SECRET_ARN", "AURORA_ENDPOINT"): + monkeypatch.delenv(k, raising=False) + from ad_buyer.storage.db_secret import resolve_database_url + + assert resolve_database_url() is None + + def test_assembles_url_from_secret(self, monkeypatch): + monkeypatch.setenv( + "DB_SECRET_ARN", + "arn:aws:secretsmanager:us-west-2:123456789012:secret:rds!cluster-x", + ) + monkeypatch.setenv("AURORA_ENDPOINT", "aurora.example.us-west-2.rds.amazonaws.com") + monkeypatch.setenv("AURORA_PORT", "5432") + monkeypatch.setenv("DB_NAME", "ad_buyer") + + import ad_buyer.storage.db_secret as ds + + class _FakeClient: + def get_secret_value(self, SecretId): # noqa: N803 + return {"SecretString": '{"username": "admin", "password": "p@ss/w:rd"}'} + + class _FakeBoto3: + @staticmethod + def client(*_a, **_k): + return _FakeClient() + + import sys + + monkeypatch.setitem(sys.modules, "boto3", _FakeBoto3) + url = ds.resolve_database_url() + assert url is not None + assert url.startswith("postgresql://admin:") + assert "@aurora.example.us-west-2.rds.amazonaws.com:5432/ad_buyer" in url + # special chars URL-encoded (no raw / or : in the password segment) + assert "p%40ss" in url + + def test_missing_password_returns_none(self, monkeypatch): + monkeypatch.setenv("DB_SECRET_ARN", "arn:secret") + monkeypatch.setenv("AURORA_ENDPOINT", "h") + import ad_buyer.storage.db_secret as ds + + class _FakeClient: + def get_secret_value(self, SecretId): # noqa: N803 + return {"SecretString": '{"username": "admin"}'} + + class _FakeBoto3: + @staticmethod + def client(*_a, **_k): + return _FakeClient() + + import sys + + monkeypatch.setitem(sys.modules, "boto3", _FakeBoto3) + assert ds.resolve_database_url() is None 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() == {}